From e9891f90c65cbe35b89d073ae2c32deb6ac6c756 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 16 May 2026 20:23:01 +0530 Subject: [PATCH 0001/1087] Fix projectService.ts --- wren-ui/src/apollo/server/services/projectService.ts | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index e6e70e014e..0ac684677e 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -1,3 +1,4 @@ +#ai import crypto from 'crypto'; import * as fs from 'fs'; import path from 'path'; @@ -44,11 +45,11 @@ export interface ProjectData { connectionInfo: WREN_AI_CONNECTION_INFO; } -export interface ProjectRecommendationQuestionsResult { +export type ProjectRecommendationQuestionsResult = { status: RecommendQuestionResultStatus; questions: RecommendationQuestion[]; - error: WrenAIError; -} + error: WrenAIError | null; +}; export interface IProjectService { createProject: (projectData: ProjectData) => Promise; updateProject: ( From 0dd63f2172a1b5fed16b96cd266f2dc588f3f11b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 17 May 2026 19:49:45 +0530 Subject: [PATCH 0002/1087] Updated indexing and qdrant changes --- .../pipelines/indexing/historical_question.py | 3 + .../src/providers/document_store/qdrant.py | 111 ++++++++++++------ 2 files changed, 78 insertions(+), 36 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/historical_question.py b/wren-ai-service/src/pipelines/indexing/historical_question.py index ff30d91f0f..95515f7ac4 100644 --- a/wren-ai-service/src/pipelines/indexing/historical_question.py +++ b/wren-ai-service/src/pipelines/indexing/historical_question.py @@ -111,6 +111,9 @@ def chunk( @observe(capture_input=False, capture_output=False) async def embedding(chunk: Dict[str, Any], embedder: Any) -> Dict[str, Any]: + if not chunk["documents"]: + return chunk + return await embedder.run(documents=chunk["documents"]) diff --git a/wren-ai-service/src/providers/document_store/qdrant.py b/wren-ai-service/src/providers/document_store/qdrant.py index b90961c456..5b7d8ff282 100644 --- a/wren-ai-service/src/providers/document_store/qdrant.py +++ b/wren-ai-service/src/providers/document_store/qdrant.py @@ -21,6 +21,7 @@ from haystack_integrations.document_stores.qdrant.filters import ( convert_filters_to_qdrant, ) +from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.http import models as rest from tqdm import tqdm @@ -30,6 +31,21 @@ logger = logging.getLogger("wren-ai-service") +def _env_flag(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _is_missing_collection_error(err: Exception) -> bool: + if not isinstance(err, UnexpectedResponse): + return False + + status_code = getattr(err, "status_code", None) + return status_code == 404 and "doesn't exist" in str(err) + + def convert_haystack_documents_to_qdrant_points( documents: List[Document], *, @@ -172,27 +188,36 @@ async def _query_by_embedding( ) -> List[Document]: qdrant_filters = convert_filters_to_qdrant(filters) - points = await self.async_client.search( - collection_name=self.index, - query_vector=rest.NamedVector( - name=DENSE_VECTORS_NAME if self.use_sparse_embeddings else "", - vector=query_embedding, - ), - search_params=( - rest.SearchParams( - quantization=rest.QuantizationSearchParams( - rescore=True, - oversampling=3.0, - ), + try: + points = await self.async_client.search( + collection_name=self.index, + query_vector=rest.NamedVector( + name=DENSE_VECTORS_NAME if self.use_sparse_embeddings else "", + vector=query_embedding, + ), + search_params=( + rest.SearchParams( + quantization=rest.QuantizationSearchParams( + rescore=True, + oversampling=3.0, + ), + ) + if len(query_embedding) + >= 1024 # reference: https://qdrant.tech/articles/binary-quantization/#when-should-you-not-use-bq + else None + ), + query_filter=qdrant_filters, + limit=top_k, + with_vectors=return_embedding, + ) + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, returning no documents", + self.index, ) - if len(query_embedding) - >= 1024 # reference: https://qdrant.tech/articles/binary-quantization/#when-should-you-not-use-bq - else None - ), - query_filter=qdrant_filters, - limit=top_k, - with_vectors=return_embedding, - ) + return [] + raise results = [ convert_qdrant_point_to_haystack_document( point, use_sparse_embeddings=self.use_sparse_embeddings @@ -218,12 +243,21 @@ async def _query_by_filters( points_list = [] offset = None while True: - points = await self.async_client.scroll( - collection_name=self.index, - offset=offset, - scroll_filter=qdrant_filters, - limit=top_k, - ) + try: + points = await self.async_client.scroll( + collection_name=self.index, + offset=offset, + scroll_filter=qdrant_filters, + limit=top_k, + ) + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, returning no documents", + self.index, + ) + return [] + raise points_list.extend(points[0]) if points[1] is None: break @@ -262,11 +296,20 @@ async def count_documents(self, filters: Optional[Dict[str, Any]] = None) -> int else: qdrant_filters = convert_filters_to_qdrant(filters) - return ( - await self.async_client.count( - collection_name=self.index, count_filter=qdrant_filters - ) - ).count + try: + return ( + await self.async_client.count( + collection_name=self.index, count_filter=qdrant_filters + ) + ).count + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, returning count 0", + self.index, + ) + return 0 + raise async def write_documents( self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.FAIL @@ -378,11 +421,7 @@ def __init__( if os.getenv("EMBEDDING_MODEL_DIMENSION") else 0 ), - recreate_index: bool = ( - bool(os.getenv("SHOULD_FORCE_DEPLOY")) - if os.getenv("SHOULD_FORCE_DEPLOY") - else False - ), + recreate_index: bool = _env_flag("SHOULD_FORCE_DEPLOY"), **_, ): self._location = location From c74c686676b69e4c913c2dc12b908c8dbe356832 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 17 May 2026 20:36:14 +0530 Subject: [PATCH 0003/1087] Updated project services and qdrant changes --- wren-ai-service/src/providers/document_store/qdrant.py | 8 ++++++++ wren-ui/src/apollo/server/services/projectService.ts | 1 - 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/providers/document_store/qdrant.py b/wren-ai-service/src/providers/document_store/qdrant.py index 5b7d8ff282..5c7894493d 100644 --- a/wren-ai-service/src/providers/document_store/qdrant.py +++ b/wren-ai-service/src/providers/document_store/qdrant.py @@ -289,6 +289,14 @@ async def delete_documents(self, filters: Optional[Dict[str, Any]] = None): logger.warning( "Called QdrantDocumentStore.delete_documents() on a non-existing ID", ) + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, skipping delete", + self.index, + ) + return + raise async def count_documents(self, filters: Optional[Dict[str, Any]] = None) -> int: if not filters: diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 0ac684677e..88f0915082 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -1,4 +1,3 @@ -#ai import crypto from 'crypto'; import * as fs from 'fs'; import path from 'path'; From f2c9a2c21ba8a374fbe6bee86882bfe5427a4022 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 17 May 2026 23:54:31 +0530 Subject: [PATCH 0004/1087] Updated litellm --- wren-ai-service/src/providers/embedder/litellm.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index 4d051e3284..c027b88252 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -61,7 +61,7 @@ async def run(self, text: str): response = await aembedding( model=self._model, - input=[text_to_embed], + input=text_to_embed, api_key=self._api_key, api_base=self._api_base_url, timeout=self._timeout, @@ -97,23 +97,20 @@ def __init__( async def _embed_batch( self, texts_to_embed: List[str], batch_size: int ) -> Tuple[List[List[float]], Dict[str, Any]]: - async def embed_single_batch(batch: List[str]) -> Any: + # Some OpenAI-compatible local embedding servers accept scalar string input + # but fail on array input. Embed documents individually to avoid that path. + async def embed_single_text(text: str) -> Any: return await aembedding( model=self._model, - input=batch, + input=text, api_key=self._api_key, api_base=self._api_base_url, timeout=self._timeout, **self._kwargs, ) - batches = [ - texts_to_embed[i : i + batch_size] - for i in range(0, len(texts_to_embed), batch_size) - ] - responses = await asyncio.gather( - *[embed_single_batch(batch) for batch in batches] + *[embed_single_text(text) for text in texts_to_embed] ) all_embeddings = [] From 80042d1bbee2ce0ee1d4724deaa123c75374a213 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 00:20:28 +0530 Subject: [PATCH 0005/1087] Updated litellm code --- .../src/providers/embedder/litellm.py | 42 +++++++++++++++---- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index c027b88252..8b301e5ea0 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -6,7 +6,6 @@ import backoff import openai from haystack import Document, component -from litellm import aembedding from src.core.provider import EmbedderProvider from src.providers.loader import provider @@ -15,6 +14,35 @@ logger = logging.getLogger("wren-ai-service") +def _normalize_model_name(model: str, api_base_url: Optional[str]) -> str: + # OpenAI-compatible local servers often expect the raw model name and will + # reject litellm-style "openai/" prefixes. + if api_base_url and model.startswith("openai/"): + return model.split("/", 1)[1] + return model + + +async def _create_embedding( + *, + model: str, + input_text: str, + api_key: Optional[str], + api_base_url: Optional[str], + timeout: Optional[float], + **kwargs, +): + client = openai.AsyncOpenAI( + api_key=api_key, + base_url=api_base_url, + timeout=timeout, + ) + return await client.embeddings.create( + model=_normalize_model_name(model, api_base_url), + input=input_text, + **kwargs, + ) + + def _prepare_texts_to_embed(documents: List[Document]) -> List[str]: """ Prepare the texts to embed by concatenating the Document text with the metadata fields to embed. @@ -59,11 +87,11 @@ async def run(self, text: str): # replace newlines, which can negatively affect performance. text_to_embed = text.replace("\n", " ") - response = await aembedding( + response = await _create_embedding( model=self._model, - input=text_to_embed, + input_text=text_to_embed, api_key=self._api_key, - api_base=self._api_base_url, + api_base_url=self._api_base_url, timeout=self._timeout, **self._kwargs, ) @@ -100,11 +128,11 @@ async def _embed_batch( # Some OpenAI-compatible local embedding servers accept scalar string input # but fail on array input. Embed documents individually to avoid that path. async def embed_single_text(text: str) -> Any: - return await aembedding( + return await _create_embedding( model=self._model, - input=text, + input_text=text, api_key=self._api_key, - api_base=self._api_base_url, + api_base_url=self._api_base_url, timeout=self._timeout, **self._kwargs, ) From 910eaba7171c3dc036c228a4d48e690255299d1b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 00:35:07 +0530 Subject: [PATCH 0006/1087] Updated litellm code db --- wren-ai-service/src/providers/embedder/litellm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index 8b301e5ea0..6bafdb73f8 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -101,7 +101,7 @@ async def run(self, text: str): "usage": dict(response.usage) if hasattr(response, "usage") else {}, } - return {"embedding": response.data[0]["embedding"], "meta": meta} + return {"embedding": response.data[0].embedding, "meta": meta} @component From 267c33ce3e72c90cb3eacc2251453a250c61475c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 00:51:13 +0530 Subject: [PATCH 0007/1087] litellm --- wren-ai-service/src/providers/embedder/litellm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index 6bafdb73f8..1657fd6230 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -145,7 +145,10 @@ async def embed_single_text(text: str) -> Any: meta: Dict[str, Any] = {} for response in responses: - embeddings = [el["embedding"] for el in response.data] + embeddings = [ + el.embedding if hasattr(el, "embedding") else el["embedding"] + for el in response.data + ] all_embeddings.extend(embeddings) if "model" not in meta: From 9d81e274a17aa1c6b3ff7d527bc0dcf91256fd9d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 13:42:24 +0530 Subject: [PATCH 0008/1087] litellm table --- .../pipelines/indexing/table_description.py | 28 +++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 6da100868f..8e1b875b49 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -18,9 +18,28 @@ logger = logging.getLogger("wren-ai-service") +MAX_TABLE_DESCRIPTION_COLUMNS = 200 +MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH = 4000 + @component class TableDescriptionChunker: + def _truncate_description(self, description: str) -> str: + if len(description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: + return description + + return description[:MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH].rstrip() + "..." + + def _format_columns(self, columns: List[str]) -> str: + if len(columns) <= MAX_TABLE_DESCRIPTION_COLUMNS: + return ", ".join(columns) + + remaining_columns = len(columns) - MAX_TABLE_DESCRIPTION_COLUMNS + truncated_columns = columns[:MAX_TABLE_DESCRIPTION_COLUMNS] + [ + f"... (+{remaining_columns} more columns)" + ] + return ", ".join(truncated_columns) + @component.output_types(documents=List[Document]) def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None): def _additional_meta() -> Dict[str, Any]: @@ -67,8 +86,10 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: return [ { "name": resource["name"], - "description": resource["properties"].get("description", ""), - "columns": ", ".join(resource["columns"]), + "description": self._truncate_description( + resource["properties"].get("description", "") + ), + "columns": self._format_columns(resource["columns"]), } for resource in resources if resource["name"] is not None @@ -94,6 +115,9 @@ def chunk( @observe(capture_input=False, capture_output=False) async def embedding(chunk: Dict[str, Any], embedder: Any) -> Dict[str, Any]: + if not chunk["documents"]: + return chunk + return await embedder.run(documents=chunk["documents"]) From 94176d9deb01415b39ce6c3bf2f6d6f9c671b40f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 13:42:49 +0530 Subject: [PATCH 0009/1087] litellm table --- .../src/providers/embedder/litellm.py | 54 +++++++++++-------- 1 file changed, 33 insertions(+), 21 deletions(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index 1657fd6230..a171daf81b 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -58,6 +58,14 @@ def _prepare_texts_to_embed(documents: List[Document]) -> List[str]: return texts_to_embed +def _iter_batches(items: List[str], batch_size: int) -> List[List[str]]: + effective_batch_size = max(batch_size, 1) + return [ + items[index : index + effective_batch_size] + for index in range(0, len(items), effective_batch_size) + ] + + @component class AsyncTextEmbedder: def __init__( @@ -137,30 +145,31 @@ async def embed_single_text(text: str) -> Any: **self._kwargs, ) - responses = await asyncio.gather( - *[embed_single_text(text) for text in texts_to_embed] - ) - all_embeddings = [] meta: Dict[str, Any] = {} - for response in responses: - embeddings = [ - el.embedding if hasattr(el, "embedding") else el["embedding"] - for el in response.data - ] - all_embeddings.extend(embeddings) - - if "model" not in meta: - meta["model"] = response.model - if "usage" not in meta: - meta["usage"] = ( - dict(response.usage) if hasattr(response, "usage") else {} - ) - else: - if hasattr(response, "usage"): - meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens - meta["usage"]["total_tokens"] += response.usage.total_tokens + for batch in _iter_batches(texts_to_embed, batch_size): + responses = await asyncio.gather( + *[embed_single_text(text) for text in batch] + ) + + for response in responses: + embeddings = [ + el.embedding if hasattr(el, "embedding") else el["embedding"] + for el in response.data + ] + all_embeddings.extend(embeddings) + + if "model" not in meta: + meta["model"] = response.model + if "usage" not in meta: + meta["usage"] = ( + dict(response.usage) if hasattr(response, "usage") else {} + ) + else: + if hasattr(response, "usage"): + meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens + meta["usage"]["total_tokens"] += response.usage.total_tokens return all_embeddings, meta @@ -177,6 +186,9 @@ async def run(self, documents: List[Document]): "In case you want to embed a string, please use the AsyncTextEmbedder." ) + if not documents: + return {"documents": documents, "meta": {}} + texts_to_embed = _prepare_texts_to_embed(documents=documents) embeddings, meta = await self._embed_batch( From 92f0dc081d8fe6fb5e8e424f2d1bfe5763787983 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 16:27:08 +0530 Subject: [PATCH 0010/1087] litellm table-desc --- .../src/pipelines/indexing/db_schema.py | 136 +++++++++++++++--- .../apollo/server/resolvers/modelResolver.ts | 20 ++- .../server/resolvers/projectResolver.ts | 5 +- 3 files changed, 134 insertions(+), 27 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 394d087b46..197d8def8c 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -25,9 +25,90 @@ logger = logging.getLogger("wren-ai-service") +MAX_DB_SCHEMA_COMMENT_LENGTH = 4000 +MAX_DB_SCHEMA_STATEMENT_LENGTH = 8000 +MAX_DB_SCHEMA_DOCUMENT_LENGTH = 12000 + @component class DDLChunker: + def _truncate_text(self, text: str, max_length: int) -> str: + if len(text) <= max_length: + return text + + return text[:max_length].rstrip() + "..." + + def _serialize_table_columns_payload(self, columns: List[dict]) -> str: + return str({"type": "TABLE_COLUMNS", "columns": columns}) + + def _fit_table_column_command(self, command: dict) -> dict: + if ( + len(self._serialize_table_columns_payload([command])) + <= MAX_DB_SCHEMA_DOCUMENT_LENGTH + ): + return command + + fitted_command = {**command} + comment = fitted_command.get("comment", "") + if comment: + base_length = len( + self._serialize_table_columns_payload( + [{**fitted_command, "comment": ""}] + ) + ) + available_comment_length = max( + MAX_DB_SCHEMA_COMMENT_LENGTH // 8, + MAX_DB_SCHEMA_DOCUMENT_LENGTH - base_length - 3, + ) + fitted_command["comment"] = self._truncate_text( + comment, + available_comment_length, + ) + + if ( + len(self._serialize_table_columns_payload([fitted_command])) + <= MAX_DB_SCHEMA_DOCUMENT_LENGTH + ): + return fitted_command + + fitted_command["comment"] = "" + return fitted_command + + def _build_table_column_payloads( + self, + model_name: str, + commands: List[dict], + column_batch_size: int, + ) -> List[Dict[str, str]]: + batches: List[List[dict]] = [] + current_batch: List[dict] = [] + effective_batch_size = max(column_batch_size, 1) + + for command in [self._fit_table_column_command(command) for command in commands]: + candidate_batch = current_batch + [command] + candidate_payload = self._serialize_table_columns_payload(candidate_batch) + + if current_batch and ( + len(current_batch) >= effective_batch_size + or len(candidate_payload) > MAX_DB_SCHEMA_DOCUMENT_LENGTH + ): + batches.append(current_batch) + current_batch = [command] + continue + + current_batch = candidate_batch + + if current_batch: + batches.append(current_batch) + + return [ + { + "name": model_name, + "payload": self._serialize_table_columns_payload(batch), + } + for batch in batches + ] + @component.output_types(documents=List[Document]) async def run( self, @@ -134,9 +215,15 @@ def _model_command(model: Dict[str, Any]) -> dict: model_properties = { "alias": clean_display_name(properties.get("displayName", "")), - "description": properties.get("description", ""), + "description": self._truncate_text( + properties.get("description", ""), + MAX_DB_SCHEMA_COMMENT_LENGTH, + ), } - comment = f"\n/* {str(model_properties)} */\n" + comment = self._truncate_text( + f"\n/* {str(model_properties)} */\n", + MAX_DB_SCHEMA_COMMENT_LENGTH, + ) table_name = model["name"] payload = { @@ -155,10 +242,14 @@ def _column_command(column: Dict[str, Any], model: Dict[str, Any]) -> dict: for helper in helper.COLUMN_COMMENT_HELPERS.values() if helper.condition(column) ] + comment = self._truncate_text( + "".join(comments), + MAX_DB_SCHEMA_COMMENT_LENGTH, + ) return { "type": "COLUMN", - "comment": "".join(comments), + "comment": comment, "name": column["name"], "data_type": column["type"], "is_primary_key": column["name"] == model["primaryKey"], @@ -193,7 +284,10 @@ def _relationship_command( return { "type": "FOREIGN_KEY", - "comment": f'-- {{"condition": {condition}, "joinType": {join_type}}}\n ', + "comment": self._truncate_text( + f'-- {{"condition": {condition}, "joinType": {join_type}}}\n ', + MAX_DB_SCHEMA_COMMENT_LENGTH, + ), "constraint": fk_constraint, "tables": models, } @@ -210,18 +304,11 @@ def _column_batch( filtered = [command for command in commands if command is not None] - return [ - { - "name": model["name"], - "payload": str( - { - "type": "TABLE_COLUMNS", - "columns": filtered[i : i + column_batch_size], - } - ), - } - for i in range(0, len(filtered), column_batch_size) - ] + return self._build_table_column_payloads( + model["name"], + filtered, + column_batch_size, + ) # A map to store model primary keys for foreign key relationships primary_keys_map = {model["name"]: model["primaryKey"] for model in models} @@ -237,11 +324,15 @@ def _convert_views(self, views: List[Dict[str, Any]]) -> List[Dict[str, str]]: def _payload(view: Dict[str, Any]) -> dict: return { "type": "VIEW", - "comment": f"/* {view['properties']} */\n" - if "properties" in view - else "", + "comment": self._truncate_text( + f"/* {view['properties']} */\n" if "properties" in view else "", + MAX_DB_SCHEMA_COMMENT_LENGTH, + ), "name": view["name"], - "statement": view["statement"], + "statement": self._truncate_text( + view["statement"], + MAX_DB_SCHEMA_STATEMENT_LENGTH, + ), } return [ @@ -252,7 +343,7 @@ def _convert_metrics(self, metrics: List[Dict[str, Any]]) -> List[Dict[str, str] def _create_column(name: str, data_type: str, comment: str) -> dict: return { "type": "COLUMN", - "comment": comment, + "comment": self._truncate_text(comment, MAX_DB_SCHEMA_COMMENT_LENGTH), "name": name, "data_type": data_type, } @@ -315,6 +406,9 @@ async def chunk( @observe(capture_input=False, capture_output=False) async def embedding(chunk: Dict[str, Any], embedder: Any) -> Dict[str, Any]: + if not chunk["documents"]: + return chunk + return await embedder.run(documents=chunk["documents"]) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 6facf755ff..59ad3cd7bb 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -235,8 +235,9 @@ export class ModelResolver { args.force, ); - // only generating for user's data source - if (project.sampleDataset === null) { + // Recommendation generation depends on a successful deployment because + // question validation calls previewSql against the deployed manifest. + if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { await ctx.projectService.generateProjectRecommendationQuestions(); } return deployRes; @@ -812,7 +813,7 @@ export class ModelResolver { // create view const project = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.deployService.getLastDeployment(project.id); + const manifest = await this.getLastDeployedManifest(ctx, project.id); // get sql statement of a response const response = await ctx.askingService.getResponse(responseId); @@ -941,7 +942,7 @@ export class ModelResolver { const project = projectId ? await ctx.projectService.getProjectById(parseInt(projectId)) : await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.deployService.getLastDeployment(project.id); + const manifest = await this.getLastDeployedManifest(ctx, project.id); return await ctx.queryService.preview(sql, { project, limit: limit, @@ -1089,6 +1090,17 @@ export class ModelResolver { }; } + private async getLastDeployedManifest(ctx: IContext, projectId: number) { + const deployment = await ctx.deployService.getLastDeployment(projectId); + if (!deployment?.manifest) { + throw new Error( + 'Project has not been deployed successfully yet. Deploy the model before previewing or validating SQL.', + ); + } + + return deployment.manifest; + } + private validateTableExist( tableName: string, dataSourceTables: CompactTable[], diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 0d6d6ba0e0..4b5c440aa8 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -637,8 +637,9 @@ export class ProjectResolver { const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy(manifest, project.id); - // only generating for user's data source - if (project.sampleDataset === null) { + // Recommendation generation depends on a successful deployment because + // question validation calls previewSql against the deployed manifest. + if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { await ctx.projectService.generateProjectRecommendationQuestions(); } return deployRes; From 8ce8d8812164af6630f2b79381f63286e9a1cadb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 16:46:24 +0530 Subject: [PATCH 0011/1087] litellm test --- .../src/providers/embedder/litellm.py | 27 +++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index a171daf81b..df06cd69e2 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -13,6 +13,8 @@ logger = logging.getLogger("wren-ai-service") +DEFAULT_MAX_EMBED_INPUT_CHARS = 1800 + def _normalize_model_name(model: str, api_base_url: Optional[str]) -> str: # OpenAI-compatible local servers often expect the raw model name and will @@ -43,7 +45,16 @@ async def _create_embedding( ) -def _prepare_texts_to_embed(documents: List[Document]) -> List[str]: +def _truncate_text_for_embedding(text: str, max_input_chars: int) -> str: + if len(text) <= max_input_chars: + return text + + return text[:max_input_chars].rstrip() + "..." + + +def _prepare_texts_to_embed( + documents: List[Document], max_input_chars: int +) -> List[str]: """ Prepare the texts to embed by concatenating the Document text with the metadata fields to embed. """ @@ -54,6 +65,7 @@ def _prepare_texts_to_embed(documents: List[Document]) -> List[str]: # copied from OpenAI embedding_utils (https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py) # replace newlines, which can negatively affect performance. text_to_embed = text_to_embed.replace("\n", " ") + text_to_embed = _truncate_text_for_embedding(text_to_embed, max_input_chars) texts_to_embed.append(text_to_embed) return texts_to_embed @@ -74,12 +86,14 @@ def __init__( api_key: Optional[str] = None, api_base_url: Optional[str] = None, timeout: Optional[float] = None, + max_input_chars: int = DEFAULT_MAX_EMBED_INPUT_CHARS, **kwargs, ): self._api_key = api_key self._model = model self._api_base_url = api_base_url self._timeout = timeout + self._max_input_chars = max(max_input_chars, 1) self._kwargs = kwargs @component.output_types(embedding=List[float], meta=Dict[str, Any]) @@ -94,6 +108,10 @@ async def run(self, text: str): # copied from OpenAI embedding_utils (https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py) # replace newlines, which can negatively affect performance. text_to_embed = text.replace("\n", " ") + text_to_embed = _truncate_text_for_embedding( + text_to_embed, + self._max_input_chars, + ) response = await _create_embedding( model=self._model, @@ -121,6 +139,7 @@ def __init__( api_key: Optional[str] = None, api_base_url: Optional[str] = None, timeout: Optional[float] = None, + max_input_chars: int = DEFAULT_MAX_EMBED_INPUT_CHARS, **kwargs, ): self._api_key = api_key @@ -128,6 +147,7 @@ def __init__( self._batch_size = batch_size self._api_base_url = api_base_url self._timeout = timeout + self._max_input_chars = max(max_input_chars, 1) self._kwargs = kwargs async def _embed_batch( @@ -189,7 +209,10 @@ async def run(self, documents: List[Document]): if not documents: return {"documents": documents, "meta": {}} - texts_to_embed = _prepare_texts_to_embed(documents=documents) + texts_to_embed = _prepare_texts_to_embed( + documents=documents, + max_input_chars=self._max_input_chars, + ) embeddings, meta = await self._embed_batch( texts_to_embed=texts_to_embed, From b06f4adaef04e26b5a10ef3dc8dfe5d736f3908c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 16:59:57 +0530 Subject: [PATCH 0012/1087] litellm test wren --- .../src/providers/embedder/litellm.py | 51 +++++++++++++++---- .../apollo/server/adaptors/wrenAIAdaptor.ts | 7 ++- 2 files changed, 47 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index df06cd69e2..7587556236 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -13,7 +13,8 @@ logger = logging.getLogger("wren-ai-service") -DEFAULT_MAX_EMBED_INPUT_CHARS = 1800 +DEFAULT_MAX_EMBED_INPUT_CHARS = 900 +MIN_EMBED_INPUT_CHARS = 128 def _normalize_model_name(model: str, api_base_url: Optional[str]) -> str: @@ -52,6 +53,19 @@ def _truncate_text_for_embedding(text: str, max_input_chars: int) -> str: return text[:max_input_chars].rstrip() + "..." +def _is_input_too_large_error(error: Exception) -> bool: + error_message = str(error).lower() + return any( + phrase in error_message + for phrase in [ + "too large to process", + "context size has been exceeded", + "physical batch size", + "input (", + ] + ) + + def _prepare_texts_to_embed( documents: List[Document], max_input_chars: int ) -> List[str]: @@ -156,14 +170,33 @@ async def _embed_batch( # Some OpenAI-compatible local embedding servers accept scalar string input # but fail on array input. Embed documents individually to avoid that path. async def embed_single_text(text: str) -> Any: - return await _create_embedding( - model=self._model, - input_text=text, - api_key=self._api_key, - api_base_url=self._api_base_url, - timeout=self._timeout, - **self._kwargs, - ) + candidate_text = text + while True: + try: + return await _create_embedding( + model=self._model, + input_text=candidate_text, + api_key=self._api_key, + api_base_url=self._api_base_url, + timeout=self._timeout, + **self._kwargs, + ) + except openai.APIError as error: + if ( + not _is_input_too_large_error(error) + or len(candidate_text) <= MIN_EMBED_INPUT_CHARS + ): + raise + + next_max_chars = max(len(candidate_text) // 2, MIN_EMBED_INPUT_CHARS) + logger.warning( + "Embedding input exceeded provider limits; retrying with %s characters", + next_max_chars, + ) + candidate_text = _truncate_text_for_embedding( + candidate_text, + next_max_chars, + ) all_embeddings = [] meta: Dict[str, Any] = {} diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 2d0675633e..1877f1da42 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -801,8 +801,11 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { `${this.wrenAIBaseEndpoint}/v1/semantics-preparations/${deployId}/status`, ); if (res.data.error) { - // passing AI response error string to catch block - throw new Error(res.data.error); + const error = + typeof res.data.error === 'string' + ? res.data.error + : res.data.error.message || JSON.stringify(res.data.error); + throw new Error(error); } return res.data?.status.toUpperCase() as WrenAISystemStatus; } catch (err: any) { From 138ef67ed21bed9c2e00756b60db83d0630ca9cd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 17:29:36 +0530 Subject: [PATCH 0013/1087] Question generation --- .../v1/services/question_recommendation.py | 196 ++++++++++++++++-- 1 file changed, 182 insertions(+), 14 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 6033237a45..8aa5ced38e 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -13,6 +13,25 @@ logger = logging.getLogger("wren-ai-service") +DEFAULT_RECOMMENDATION_CONTEXT_ITEMS = 8 +DEFAULT_RECOMMENDATION_CONTEXT_CHARS = 12000 +DEFAULT_VALIDATION_CONTEXT_ITEMS = 4 +DEFAULT_VALIDATION_CONTEXT_CHARS = 6000 +STRICT_VALIDATION_CONTEXT_ITEMS = 2 +STRICT_VALIDATION_CONTEXT_CHARS = 2500 +DEFAULT_VALIDATION_SQL_SAMPLE_ITEMS = 2 +DEFAULT_VALIDATION_SQL_SAMPLE_CHARS = 2000 +STRICT_VALIDATION_SQL_SAMPLE_ITEMS = 0 +STRICT_VALIDATION_SQL_SAMPLE_CHARS = 0 +DEFAULT_VALIDATION_INSTRUCTION_ITEMS = 6 +DEFAULT_VALIDATION_INSTRUCTION_CHARS = 1500 +STRICT_VALIDATION_INSTRUCTION_ITEMS = 2 +STRICT_VALIDATION_INSTRUCTION_CHARS = 500 +DEFAULT_VALIDATION_SQL_FUNCTION_ITEMS = 12 +DEFAULT_VALIDATION_SQL_FUNCTION_CHARS = 2500 +STRICT_VALIDATION_SQL_FUNCTION_ITEMS = 0 +STRICT_VALIDATION_SQL_FUNCTION_CHARS = 0 + class QuestionRecommendation: class Error(BaseModel): @@ -42,6 +61,87 @@ def __init__( self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval + def _truncate_text(self, value: str, max_chars: int) -> str: + if max_chars <= 0: + return "" + if len(value) <= max_chars: + return value + return value[:max_chars].rstrip() + "..." + + def _limit_text_items( + self, items: list[str], max_items: int, max_chars: int + ) -> list[str]: + if max_items <= 0 or max_chars <= 0: + return [] + + limited_items: list[str] = [] + remaining_chars = max_chars + + for item in items: + if len(limited_items) >= max_items or remaining_chars <= 0: + break + + truncated = self._truncate_text(item, remaining_chars) + if not truncated: + break + + limited_items.append(truncated) + remaining_chars -= len(truncated) + + return limited_items + + def _limit_dict_items( + self, + items: list[dict], + key: str, + max_items: int, + max_chars: int, + ) -> list[dict]: + if max_items <= 0 or max_chars <= 0: + return [] + + limited_items: list[dict] = [] + remaining_chars = max_chars + + for item in items: + if len(limited_items) >= max_items or remaining_chars <= 0: + break + + value = str(item.get(key, "")) + truncated_value = self._truncate_text(value, remaining_chars) + if not truncated_value: + break + + next_item = {**item, key: truncated_value} + limited_items.append(next_item) + remaining_chars -= len(truncated_value) + + return limited_items + + def _limit_sql_functions( + self, functions: list, max_items: int, max_chars: int + ) -> list: + serialized_functions = [str(function) for function in functions] + limited_values = self._limit_text_items( + serialized_functions, + max_items=max_items, + max_chars=max_chars, + ) + limited_count = len(limited_values) + return functions[:limited_count] + + def _is_context_size_error(self, error: Exception) -> bool: + error_message = str(error).lower() + return any( + phrase in error_message + for phrase in [ + "context size has been exceeded", + "too large to process", + "maximum context length", + "prompt is too long", + ] + ) + def _handle_exception( self, event_id: str, @@ -121,19 +221,83 @@ async def _instructions_retrieval() -> list[dict]: else: sql_knowledge = None - generated_sql = await self._pipelines["sql_generation"].run( - query=candidate["question"], - contexts=table_ddls, - project_id=project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - allow_data_preview=allow_data_preview, - sql_knowledge=sql_knowledge, - ) + validation_attempts = [ + { + "contexts": self._limit_text_items( + table_ddls, + max_items=DEFAULT_VALIDATION_CONTEXT_ITEMS, + max_chars=DEFAULT_VALIDATION_CONTEXT_CHARS, + ), + "sql_samples": self._limit_dict_items( + sql_samples, + key="sql", + max_items=DEFAULT_VALIDATION_SQL_SAMPLE_ITEMS, + max_chars=DEFAULT_VALIDATION_SQL_SAMPLE_CHARS, + ), + "instructions": self._limit_dict_items( + instructions, + key="instruction", + max_items=DEFAULT_VALIDATION_INSTRUCTION_ITEMS, + max_chars=DEFAULT_VALIDATION_INSTRUCTION_CHARS, + ), + "sql_functions": self._limit_sql_functions( + sql_functions, + max_items=DEFAULT_VALIDATION_SQL_FUNCTION_ITEMS, + max_chars=DEFAULT_VALIDATION_SQL_FUNCTION_CHARS, + ), + }, + { + "contexts": self._limit_text_items( + table_ddls, + max_items=STRICT_VALIDATION_CONTEXT_ITEMS, + max_chars=STRICT_VALIDATION_CONTEXT_CHARS, + ), + "sql_samples": self._limit_dict_items( + sql_samples, + key="sql", + max_items=STRICT_VALIDATION_SQL_SAMPLE_ITEMS, + max_chars=STRICT_VALIDATION_SQL_SAMPLE_CHARS, + ), + "instructions": self._limit_dict_items( + instructions, + key="instruction", + max_items=STRICT_VALIDATION_INSTRUCTION_ITEMS, + max_chars=STRICT_VALIDATION_INSTRUCTION_CHARS, + ), + "sql_functions": self._limit_sql_functions( + sql_functions, + max_items=STRICT_VALIDATION_SQL_FUNCTION_ITEMS, + max_chars=STRICT_VALIDATION_SQL_FUNCTION_CHARS, + ), + }, + ] + + generated_sql = None + for attempt_index, attempt in enumerate(validation_attempts): + try: + generated_sql = await self._pipelines["sql_generation"].run( + query=candidate["question"], + contexts=attempt["contexts"], + project_id=project_id, + sql_samples=attempt["sql_samples"], + instructions=attempt["instructions"], + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=attempt["sql_functions"], + allow_data_preview=allow_data_preview, + sql_knowledge=sql_knowledge, + ) + break + except Exception as error: + is_last_attempt = attempt_index == len(validation_attempts) - 1 + if is_last_attempt or not self._is_context_size_error(error): + raise + + logger.warning( + "Request %s: SQL validation prompt exceeded context window; retrying with reduced context", + request_id, + ) post_process = generated_sql["post_process"] @@ -207,7 +371,11 @@ async def recommend(self, input: Request, **kwargs) -> Event: ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) - table_ddls = [document.get("table_ddl") for document in documents] + table_ddls = self._limit_text_items( + [document.get("table_ddl") for document in documents], + max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, + max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, + ) request = { "contexts": table_ddls, From c5b2233ef8ce0a326a65afcf24e3816165493808 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 18:54:20 +0530 Subject: [PATCH 0014/1087] updated Sql --- wren-ai-service/src/pipelines/common.py | 21 +++-- .../generation/followup_sql_generation.py | 20 +++-- .../pipelines/generation/sql_correction.py | 26 ++++-- .../pipelines/generation/sql_regeneration.py | 23 ++++- .../src/pipelines/generation/utils/sql.py | 86 +++++++++++++------ .../retrieval/db_schema_retrieval.py | 5 +- .../src/apollo/server/adaptors/ibisAdaptor.ts | 43 ++++++++-- 7 files changed, 171 insertions(+), 53 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index f6114d63b1..940fa66eb7 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -4,11 +4,21 @@ from haystack import Document, component +def normalize_data_type(data_type: Any) -> str: + if data_type is None: + return "" + return str(data_type).strip() + + def get_engine_supported_data_type(data_type: str) -> str: """ This function makes sure downstream ai pipeline get column data types in a format that is supported by the data engine. """ - match data_type.upper(): + normalized_data_type = normalize_data_type(data_type) + if not normalized_data_type: + return "UNKNOWN" + + match normalized_data_type.upper(): case "BPCHAR" | "NAME" | "UUID" | "INET": return "VARCHAR" case "OID": @@ -24,7 +34,7 @@ def get_engine_supported_data_type(data_type: str) -> str: case "INT64": return "BIGINT" case _: - return data_type.upper() + return normalized_data_type.upper() def build_table_ddl( @@ -36,16 +46,17 @@ def build_table_ddl( for column in content["columns"]: if column["type"] == "COLUMN": + column_data_type = normalize_data_type(column.get("data_type")) if ( (not columns or (columns and column["name"] in columns)) - and column["data_type"].lower() + and column_data_type.lower() != "unknown" # quick fix: filtering out UNKNOWN column type ): if "This column is a Calculated Field" in column["comment"]: has_calculated_field = True - if column["data_type"].lower() == "json": + if column_data_type.lower() == "json": has_json_field = True - column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column_data_type)}" if column["is_primary_key"]: column_ddl += " PRIMARY KEY" columns_ddl.append(column_ddl) 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 35cfb8fccf..8c02759758 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -34,6 +34,9 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. +### TARGET DATA SOURCE ### +{{ data_source }} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -92,6 +95,7 @@ def prompt( documents: list[str], sql_generation_reasoning: str, prompt_builder: PromptBuilder, + data_source: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -102,6 +106,7 @@ def prompt( ) -> dict: _prompt = prompt_builder.run( query=query, + data_source=data_source, documents=documents, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( @@ -113,7 +118,9 @@ def prompt( else "" ), metric_instructions=( - get_metric_instructions(sql_knowledge) if has_metric else "" + get_metric_instructions(sql_knowledge, data_source=data_source) + if has_metric + else "" ), json_field_instructions=( get_json_field_instructions(sql_knowledge) if has_json_field else "" @@ -131,10 +138,14 @@ async def generate_sql_in_followup( generator: Any, histories: list[AskHistory], generator_name: str, + data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: history_messages = construct_ask_history_messages(histories) - current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) + current_system_prompt = get_sql_generation_system_prompt( + sql_knowledge, + data_source=data_source, + ) return await generator( prompt=prompt.get("prompt"), history_messages=history_messages, @@ -211,10 +222,7 @@ async def run( ): logger.info("Follow-Up SQL Generation pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = await retrieve_metadata(project_id or "", self._retriever) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..791b797f0d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -25,8 +25,14 @@ logger = logging.getLogger("wren-ai-service") -def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) -> str: - text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) +def get_sql_correction_system_prompt( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: + text_to_sql_rules = get_text_to_sql_rules( + sql_knowledge, + data_source=data_source, + ) return f""" ### TASK ### @@ -52,6 +58,9 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) sql_correction_user_prompt_template = """ +### TARGET DATA SOURCE ### +{{ data_source }} + {% if documents %} ### DATABASE SCHEMA ### {% for document in documents %} @@ -87,10 +96,12 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, + data_source: str, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( + data_source=data_source, documents=documents, invalid_generation_result=invalid_generation_result, instructions=construct_instructions( @@ -107,9 +118,13 @@ async def generate_sql_correction( prompt: dict, generator: Any, generator_name: str, + data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_correction_system_prompt(sql_knowledge) + current_system_prompt = get_sql_correction_system_prompt( + sql_knowledge, + data_source=data_source, + ) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -178,10 +193,7 @@ async def run( ): logger.info("SQLCorrection pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = await retrieve_metadata(project_id or "", self._retriever) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 4b7284aa26..26bd957a7e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -29,8 +29,12 @@ def get_sql_regeneration_system_prompt( sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, ) -> str: - text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) + text_to_sql_rules = get_text_to_sql_rules( + sql_knowledge, + data_source=data_source, + ) return f""" ### TASK ### @@ -51,6 +55,9 @@ def get_sql_regeneration_system_prompt( sql_regeneration_user_prompt_template = """ +### TARGET DATA SOURCE ### +{{ data_source }} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -107,6 +114,7 @@ def prompt( sql_generation_reasoning: str, sql: str, prompt_builder: PromptBuilder, + data_source: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -117,6 +125,7 @@ def prompt( ) -> dict: _prompt = prompt_builder.run( sql=sql, + data_source=data_source, documents=documents, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( @@ -128,7 +137,9 @@ def prompt( else "" ), metric_instructions=( - get_metric_instructions(sql_knowledge) if has_metric else "" + get_metric_instructions(sql_knowledge, data_source=data_source) + if has_metric + else "" ), json_field_instructions=( get_json_field_instructions(sql_knowledge) if has_json_field else "" @@ -145,9 +156,13 @@ async def regenerate_sql( prompt: dict, generator: Any, generator_name: str, + data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_regeneration_system_prompt(sql_knowledge) + current_system_prompt = get_sql_regeneration_system_prompt( + sql_knowledge, + data_source=data_source, + ) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -197,6 +212,7 @@ async def run( contexts: list[str], sql_generation_reasoning: str, sql: str, + data_source: str = "local_file", sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, @@ -222,6 +238,7 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, + "data_source": data_source, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..300dde6502 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -17,6 +17,10 @@ logger = logging.getLogger("wren-ai-service") +def normalize_data_source(data_source: str | None) -> str: + return (data_source or "").strip().upper() + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -189,13 +193,12 @@ async def _classify_generation_result( - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds -- ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query - - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) +- When you need to cast a date/time related field, CAST it to a temporal type that is supported by the target data source and consistent with the SQL FUNCTIONS section. + - example 1: CAST(properties_closedate AS TIMESTAMP) + - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP) - If the user asks for a specific date, please give the date range in SQL query - example: "What is the total revenue for the month of 2024-11-01?" - - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" + - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP) AND CAST(r.PurchaseTimestamp AS TIMESTAMP) < CAST('2024-11-02 00:00:00' AS TIMESTAMP)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. @@ -222,6 +225,20 @@ async def _classify_generation_result( - For the ranking problem, you must add the ranking column to the final SELECT clause. """ +_MSSQL_TEXT_TO_SQL_RULES = """ +### MSSQL-SPECIFIC RULES ### +- The target database is MSSQL. +- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. +- Prefer GETDATE() for the current timestamp and CAST(GETDATE() AS DATE) when you need the current date only. +- For relative time windows, use DATEADD together with GETDATE(). +- For previous calendar month boundaries, prefer: + - start_of_previous_month: DATEADD(month, DATEDIFF(month, 0, GETDATE()) - 1, 0) + - start_of_current_month: DATEADD(month, DATEDIFF(month, 0, GETDATE()), 0) +- For month bucketing, prefer DATEADD(month, DATEDIFF(month, 0, ), 0). If DATETRUNC is available in your server version, you may use DATETRUNC(month, ), but prefer DATEADD/DATEDIFF when uncertain. +- When a temporal cast is required, prefer DATETIME2. Use DATETIMEOFFSET only when timezone-aware semantics are explicitly required by the question. +- Keep relative date logic simple and native to MSSQL. Never emit INTERVAL-like expressions for MSSQL. +""" + _DEFAULT_CALCULATED_FIELD_INSTRUCTIONS = """ #### Instructions for Calculated Field #### @@ -348,17 +365,9 @@ async def _classify_generation_result( 1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. 2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. 3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. -So utilize those metric components in the SQL generation process to give an answer like this: - -SQL Query: -SELECT - CustomerId, - PriceSum AS TotalRevenue -FROM - Revenue -WHERE - PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND - PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) +So utilize those metric components in the SQL generation process to give an answer using the date functions that are valid for the target data source and listed in the SQL FUNCTIONS section. + +For example, the SQL should filter PurchaseTimestamp to the previous calendar month using the dialect-appropriate month-boundary functions from the SQL FUNCTIONS section. """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -461,13 +470,24 @@ def _extract_from_sql_knowledge( return value if value and value.strip() else default_value -def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: +def _append_data_source_rules(base_rules: str, data_source: str | None = None) -> str: + normalized_data_source = normalize_data_source(data_source) + if normalized_data_source == "MSSQL": + return f"{base_rules}\n\n{_MSSQL_TEXT_TO_SQL_RULES}" + return base_rules + + +def get_text_to_sql_rules( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: if sql_knowledge is not None: - return _extract_from_sql_knowledge( + base_rules = _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) + return _append_data_source_rules(base_rules, data_source) - return _DEFAULT_TEXT_TO_SQL_RULES + return _append_data_source_rules(_DEFAULT_TEXT_TO_SQL_RULES, data_source) def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: @@ -481,13 +501,25 @@ def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) return _DEFAULT_CALCULATED_FIELD_INSTRUCTIONS -def get_metric_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: +def get_metric_instructions( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: + instructions = _DEFAULT_METRIC_INSTRUCTIONS if sql_knowledge is not None: - return _extract_from_sql_knowledge( + instructions = _extract_from_sql_knowledge( sql_knowledge, "metric_instructions", _DEFAULT_METRIC_INSTRUCTIONS ) - return _DEFAULT_METRIC_INSTRUCTIONS + if normalize_data_source(data_source) == "MSSQL": + instructions += """ + +#### MSSQL Metric Notes #### +- When filtering metrics by month or other relative date windows in MSSQL, use DATEADD/DATEDIFF or other MSSQL-native date functions from the SQL FUNCTIONS section. +- Do not use DATE_TRUNC, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries. +""" + + return instructions def get_json_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: @@ -499,8 +531,14 @@ def get_json_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> st return _DEFAULT_JSON_FIELD_INSTRUCTIONS -def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) -> str: - text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) +def get_sql_generation_system_prompt( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: + text_to_sql_rules = get_text_to_sql_rules( + sql_knowledge, + data_source=data_source, + ) return f""" You are a helpful assistant that converts natural language queries into ANSI SQL queries. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..7649b702f2 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -18,6 +18,7 @@ build_table_ddl, clean_up_new_lines, get_engine_supported_data_type, + normalize_data_type, ) from src.utils import trace_cost from src.web.v1.services.ask import AskHistory @@ -101,9 +102,9 @@ def _build_metric_ddl(content: dict) -> str: columns_ddl = [ - f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + f"{column['comment']}{column['name']} {get_engine_supported_data_type(normalize_data_type(column.get('data_type')))}" for column in content["columns"] - if column["data_type"].lower() + if normalize_data_type(column.get("data_type")).lower() != "unknown" # quick fix: filtering out UNKNOWN column type ] diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 0d2d4a9c5e..3b16be0c34 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -280,7 +280,11 @@ export class IbisAdaptor implements IIbisAdaptor { ); return res.data; } catch (e) { - logger.debug(`Dry plan error: ${e.response?.data || e.message}`); + logger.debug( + `Dry plan error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error during dry plan execution'); } } @@ -322,7 +326,9 @@ export class IbisAdaptor implements IIbisAdaptor { override: res.headers['x-cache-override'] === 'true', }; } catch (e) { - logger.debug(`Query error: ${e.response?.data || e.message}`); + logger.debug( + `Query error: ${this.stringifyDebugValue(e.response?.data || e.message)}`, + ); this.throwError(e, 'Error querying ibis server'); } } @@ -351,7 +357,11 @@ export class IbisAdaptor implements IIbisAdaptor { processTime: response.headers['x-process-time'], }; } catch (err) { - logger.debug(`Dry run error: ${err.response?.data || err.message}`); + logger.debug( + `Dry run error: ${this.stringifyDebugValue( + err.response?.data || err.message, + )}`, + ); this.throwError(err, 'Error during dry run execution'); } } @@ -395,7 +405,11 @@ export class IbisAdaptor implements IIbisAdaptor { ); return await getTablesByConnectionInfo(ibisConnectionInfo); } catch (e) { - logger.debug(`Get tables error: ${e.response?.data || e.message}`); + logger.debug( + `Get tables error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error getting table from ibis server'); } } @@ -417,7 +431,11 @@ export class IbisAdaptor implements IIbisAdaptor { ); return res.data; } catch (e) { - logger.debug(`Get constraints error: ${e.response?.data || e.message}`); + logger.debug( + `Get constraints error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error getting constraint from ibis server'); } } @@ -511,11 +529,24 @@ export class IbisAdaptor implements IIbisAdaptor { ); return res.data; } catch (e) { - logger.debug(`Get version error: ${e.response?.data || e.message}`); + logger.debug( + `Get version error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error getting version from ibis server'); } } + private stringifyDebugValue(value: any): string { + if (typeof value === 'string') return value; + try { + return JSON.stringify(value); + } catch { + return String(value); + } + } + private updateConnectionInfo(connectionInfo: any) { if ( config.otherServiceUsingDocker && From e2dfbaeb67a9ce1b5ff55d254928598ce606df9d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 19:04:41 +0530 Subject: [PATCH 0015/1087] updated Sql generation --- .../pipelines/generation/sql_generation.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..5f9a820506 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,6 +28,9 @@ sql_generation_user_prompt_template = """ +### TARGET DATA SOURCE ### +{{ data_source }} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -87,6 +90,7 @@ def prompt( query: str, documents: list[str], prompt_builder: PromptBuilder, + data_source: str, sql_generation_reasoning: str | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -98,6 +102,7 @@ def prompt( ) -> dict: _prompt = prompt_builder.run( query=query, + data_source=data_source, documents=documents, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( @@ -109,7 +114,9 @@ def prompt( else "" ), metric_instructions=( - get_metric_instructions(sql_knowledge) if has_metric else "" + get_metric_instructions(sql_knowledge, data_source=data_source) + if has_metric + else "" ), json_field_instructions=( get_json_field_instructions(sql_knowledge) if has_json_field else "" @@ -126,9 +133,13 @@ async def generate_sql( prompt: dict, generator: Any, generator_name: str, + data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) + current_system_prompt = get_sql_generation_system_prompt( + sql_knowledge, + data_source=data_source, + ) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -205,10 +216,7 @@ async def run( ): logger.info("SQL Generation pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = await retrieve_metadata(project_id or "", self._retriever) return await self._pipe.execute( ["post_process"], From 606ce2d6e65c1753e880e97196177afcc1662692 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 19:45:43 +0530 Subject: [PATCH 0016/1087] updated Sql add --- .../src/pipelines/generation/utils/sql.py | 6 ++++ .../components/pages/home/prompt/index.tsx | 6 ++-- wren-ui/src/hooks/useAdjustAnswer.tsx | 10 ++++++ wren-ui/src/hooks/useAskPrompt.tsx | 31 ++++++++++++++++--- wren-ui/src/pages/home/[id].tsx | 17 ++++++++-- 5 files changed, 60 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 300dde6502..8c567b926c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -229,12 +229,17 @@ async def _classify_generation_result( ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. +- DO NOT use YEAR(...), MONTH(...), DAY(...), or DATEPART(...) in generated SQL unless the exact function is explicitly listed as supported in the SQL FUNCTIONS section. Prefer range predicates or DATEADD/DATEDIFF bucket expressions instead. - Prefer GETDATE() for the current timestamp and CAST(GETDATE() AS DATE) when you need the current date only. - For relative time windows, use DATEADD together with GETDATE(). - For previous calendar month boundaries, prefer: - start_of_previous_month: DATEADD(month, DATEDIFF(month, 0, GETDATE()) - 1, 0) - start_of_current_month: DATEADD(month, DATEDIFF(month, 0, GETDATE()), 0) - For month bucketing, prefer DATEADD(month, DATEDIFF(month, 0, ), 0). If DATETRUNC is available in your server version, you may use DATETRUNC(month, ), but prefer DATEADD/DATEDIFF when uncertain. +- For year bucketing, prefer DATEADD(year, DATEDIFF(year, 0, ), 0) instead of YEAR(...). +- For filtering a specific year such as 2025, prefer a closed-open range: + - >= CAST('2025-01-01 00:00:00' AS DATETIME2) + - AND < CAST('2026-01-01 00:00:00' AS DATETIME2) - When a temporal cast is required, prefer DATETIME2. Use DATETIMEOFFSET only when timezone-aware semantics are explicitly required by the question. - Keep relative date logic simple and native to MSSQL. Never emit INTERVAL-like expressions for MSSQL. """ @@ -517,6 +522,7 @@ def get_metric_instructions( #### MSSQL Metric Notes #### - When filtering metrics by month or other relative date windows in MSSQL, use DATEADD/DATEDIFF or other MSSQL-native date functions from the SQL FUNCTIONS section. - Do not use DATE_TRUNC, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries. +- Avoid YEAR(...), MONTH(...), DAY(...), and DATEPART(...) in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. """ return instructions diff --git a/wren-ui/src/components/pages/home/prompt/index.tsx b/wren-ui/src/components/pages/home/prompt/index.tsx index c65a0ed2f6..4ecc56c650 100644 --- a/wren-ui/src/components/pages/home/prompt/index.tsx +++ b/wren-ui/src/components/pages/home/prompt/index.tsx @@ -118,8 +118,10 @@ export default forwardRef(function Prompt(props, ref) { // create thread response for text to sql const intentSQLAnswer = async () => { - onCreateResponse && - (await onCreateResponse({ question, taskId: askingTask?.queryId })); + const taskId = askingTask?.queryId; + if (!taskId || !question) return; + + onCreateResponse && (await onCreateResponse({ question, taskId })); setShowResult(false); }; diff --git a/wren-ui/src/hooks/useAdjustAnswer.tsx b/wren-ui/src/hooks/useAdjustAnswer.tsx index 832af9d7fa..4a4bdbdcd4 100644 --- a/wren-ui/src/hooks/useAdjustAnswer.tsx +++ b/wren-ui/src/hooks/useAdjustAnswer.tsx @@ -96,6 +96,8 @@ export default function useAdjustAnswer(threadId?: number) { responseId: number, input: { tables: string[]; sqlGenerationReasoning: string }, ) => { + if (!responseId) return; + const response = await adjustThreadResponse({ variables: { responseId, @@ -108,6 +110,8 @@ export default function useAdjustAnswer(threadId?: number) { // start polling new thread response const nextThreadResponse = response.data?.adjustThreadResponse; + if (!nextThreadResponse?.id) return; + await fetchThreadResponse({ variables: { responseId: nextThreadResponse.id }, }); @@ -121,12 +125,16 @@ export default function useAdjustAnswer(threadId?: number) { }; const onAdjustSQL = async (responseId: number, sql: string) => { + if (!responseId) return; + const response = await adjustThreadResponse({ variables: { responseId, data: { sql } }, }); // update thread cache const nextThreadResponse = response.data?.adjustThreadResponse; + if (!nextThreadResponse) return; + handleUpdateThreadCache( threadId, nextThreadResponse, @@ -150,6 +158,8 @@ export default function useAdjustAnswer(threadId?: number) { const onReRun = async (threadResponse: ThreadResponse) => { const responseId = threadResponse.id; + if (!responseId) return; + await rerunAdjustmentTask({ variables: { responseId } }); await fetchThreadResponse({ variables: { responseId } }); }; diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index ac45713287..ce0451eb94 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -208,6 +208,8 @@ export default function useAskPrompt(threadId?: number) { ); const startRecommendedQuestions = useCallback(async () => { + if (!originalQuestion?.trim()) return; + const previousQuestions = [ // slice the last 5 questions in threadQuestions ...uniq(threadQuestions).slice(-5), @@ -216,14 +218,21 @@ export default function useAskPrompt(threadId?: number) { const response = await createInstantRecommendedQuestions({ variables: { data: { previousQuestions } }, }); + const taskId = response.data?.createInstantRecommendedQuestions?.id; + if (!taskId) return; + fetchInstantRecommendedQuestions({ - variables: { taskId: response.data.createInstantRecommendedQuestions.id }, + variables: { taskId }, }); - }, [originalQuestion]); + }, [originalQuestion, threadQuestions]); const checkFetchAskingStreamTask = useCallback( (task: AskingTask) => { - if (!askingStreamTask && task.status === AskingTaskStatus.PLANNING) { + if ( + !askingStreamTask && + task?.queryId && + task.status === AskingTaskStatus.PLANNING + ) { fetchAskingStreamTask(task.queryId); } }, @@ -274,15 +283,22 @@ export default function useAskPrompt(threadId?: number) { }; const onReRun = async (threadResponse: ThreadResponse) => { + if (!threadResponse?.id) return; + askingStreamTaskResult.reset(); setOriginalQuestion(threadResponse.question); try { const response = await rerunAskingTask({ variables: { responseId: threadResponse.id }, }); + const taskId = response.data?.rerunAskingTask?.id; + if (!taskId) return; + const { data } = await fetchAskingTask({ - variables: { taskId: response.data.rerunAskingTask.id }, + variables: { taskId }, }); + if (!data?.askingTask) return; + // update the asking task in cache manually handleUpdateRerunAskingTaskCache( threadId, @@ -302,8 +318,11 @@ export default function useAskPrompt(threadId?: number) { const response = await createAskingTask({ variables: { data: { question: value, threadId } }, }); + const taskId = response.data?.createAskingTask?.id; + if (!taskId) return; + await fetchAskingTask({ - variables: { taskId: response.data.createAskingTask.id }, + variables: { taskId }, }); } catch (error) { console.error(error); @@ -311,6 +330,8 @@ export default function useAskPrompt(threadId?: number) { }; const onFetching = async (queryId: string) => { + if (!queryId) return; + await fetchAskingTask({ variables: { taskId: queryId }, }); diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index e27fe1f8b2..77c1d2650a 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -190,11 +190,15 @@ export default function HomeThread() { }; const onGenerateThreadResponseAnswer = async (responseId: number) => { + if (!responseId) return; + await generateThreadResponseAnswer({ variables: { responseId } }); fetchThreadResponse({ variables: { responseId } }); }; const onGenerateThreadResponseChart = async (responseId: number) => { + if (!responseId) return; + await generateThreadResponseChart({ variables: { responseId } }); fetchThreadResponse({ variables: { responseId } }); }; @@ -203,6 +207,8 @@ export default function HomeThread() { responseId: number, data: AdjustThreadResponseChartInput, ) => { + if (!responseId) return; + await adjustThreadResponseChart({ variables: { responseId, data }, }); @@ -210,6 +216,8 @@ export default function HomeThread() { }; const onGenerateThreadRecommendedQuestions = async () => { + if (!threadId) return; + await generateThreadRecommendationQuestions({ variables: { threadId } }); fetchThreadRecommendationQuestions({ variables: { threadId } }); }; @@ -221,8 +229,9 @@ export default function HomeThread() { (response) => response?.askingTask && !getIsFinished(response?.askingTask?.status), ); - if (unfinishedAskingResponse) { - askPrompt.onFetching(unfinishedAskingResponse?.askingTask?.queryId); + const unfinishedTaskId = unfinishedAskingResponse?.askingTask?.queryId; + if (unfinishedAskingResponse && unfinishedTaskId) { + askPrompt.onFetching(unfinishedTaskId); return; } @@ -297,7 +306,9 @@ export default function HomeThread() { try { askPrompt.onStopPolling(); - const threadId = thread.id; + const threadId = thread?.id; + if (!threadId) return; + await createThreadResponse({ variables: { threadId, data: payload }, }); From 2659e673e8b3098a873aa8d06f12165547ec4b7b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 20:28:13 +0530 Subject: [PATCH 0017/1087] updated questions --- .../server/backgrounds/recommend-question.ts | 14 ++++ .../apollo/server/services/askingService.ts | 9 +++ .../server/services/askingTaskTracker.ts | 65 ++++++++++++++++++- wren-ui/src/common.ts | 10 ++- 4 files changed, 94 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts index da33377deb..a649bdaf4b 100644 --- a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts +++ b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts @@ -29,6 +29,7 @@ export class ProjectRecommendQuestionBackgroundTracker { private runningJobs = new Set(); private telemetry: ITelemetry; private logger: Logger; + private initialized = false; constructor({ telemetry, @@ -150,6 +151,10 @@ export class ProjectRecommendQuestionBackgroundTracker { } public async initialize() { + if (this.initialized) { + return; + } + const projects = await this.projectRepository.findAll(); for (const project of projects) { if ( @@ -159,6 +164,8 @@ export class ProjectRecommendQuestionBackgroundTracker { this.addTask(project); } } + + this.initialized = true; } public taskKey(project: Project) { @@ -179,6 +186,7 @@ export class ThreadRecommendQuestionBackgroundTracker { private runningJobs = new Set(); private telemetry: ITelemetry; private logger: Logger; + private initialized = false; constructor({ telemetry, @@ -299,6 +307,10 @@ export class ThreadRecommendQuestionBackgroundTracker { } public async initialize() { + if (this.initialized) { + return; + } + const threads = await this.threadRepository.findAll(); for (const thread of threads) { if ( @@ -309,6 +321,8 @@ export class ThreadRecommendQuestionBackgroundTracker { this.addTask(thread); } } + + this.initialized = true; } public taskKey(thread: Thread) { diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 82c9375856..b047312bb3 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -416,6 +416,7 @@ export class AskingService implements IAskingService { private askingTaskTracker: IAskingTaskTracker; private askingTaskRepository: IAskingTaskRepository; private adjustmentBackgroundTracker: AdjustmentBackgroundTaskTracker; + private initialized = false; constructor({ telemetry, @@ -563,6 +564,12 @@ export class AskingService implements IAskingService { } public async initialize() { + if (this.initialized) { + return; + } + + await this.askingTaskTracker.initialize(); + // list thread responses from database // filter status not finalized and put them into background tracker const threadResponses = await this.threadResponseRepository.findAll(); @@ -579,6 +586,8 @@ export class AskingService implements IAskingService { for (const threadResponse of unfininshedBreakdownThreadResponses) { this.breakdownBackgroundTracker.addTask(threadResponse); } + + this.initialized = true; } /** diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index a2f18a61f0..204a5cf305 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -45,6 +45,7 @@ export interface IAskingTaskTracker { getAskingResult(queryId: string): Promise; getAskingResultById(id: number): Promise; cancelAskingTask(queryId: string): Promise; + initialize(): Promise; bindThreadResponse( id: number, queryId: string, @@ -64,6 +65,7 @@ export class AskingTaskTracker implements IAskingTaskTracker { private runningJobs = new Set(); private threadResponseRepository: IThreadResponseRepository; private viewRepository: IViewRepository; + private initialized = false; constructor({ wrenAIAdaptor, @@ -112,6 +114,12 @@ export class AskingTaskTracker implements IAskingTaskTracker { queryId, lastPolled: Date.now(), question: input.query, + result: { + type: null, + status: AskResultStatus.UNDERSTANDING, + response: null, + error: null, + }, isFinalized: false, rerunFromCancelled: input.rerunFromCancelled, } as TrackedTask; @@ -140,7 +148,16 @@ export class AskingTaskTracker implements IAskingTaskTracker { // update the query id in database await this.askingTaskRepository.updateOne(input.previousTaskId, { queryId, + detail: task.result, + }); + } else { + const createdTask = await this.askingTaskRepository.createOne({ + queryId, + question: input.query, + detail: task.result, }); + task.taskId = createdTask.id; + this.trackedTasksById.set(createdTask.id, task); } logger.info(`Created asking task with queryId: ${queryId}`); @@ -167,7 +184,11 @@ export class AskingTaskTracker implements IAskingTaskTracker { } // If not in memory or no result yet, check the database - return this.getAskingResultFromDB({ queryId }); + const result = await this.getAskingResultFromDB({ queryId }); + if (result && !this.isTaskFinalized(result.status)) { + this.restoreTrackedTask(result); + } + return result; } public async getAskingResultById( @@ -178,13 +199,38 @@ export class AskingTaskTracker implements IAskingTaskTracker { return this.getAskingResult(task.queryId); } - return this.getAskingResultFromDB({ taskId: id }); + const result = await this.getAskingResultFromDB({ taskId: id }); + if (result && !this.isTaskFinalized(result.status)) { + this.restoreTrackedTask(result); + } + return result; } public async cancelAskingTask(queryId: string): Promise { await this.wrenAIAdaptor.cancelAsk(queryId); } + public async initialize(): Promise { + if (this.initialized) return; + + const taskRecords = await this.askingTaskRepository.findAll(); + taskRecords.forEach((taskRecord) => { + const detail = taskRecord.detail as AskResult | undefined; + if (!taskRecord.queryId || !detail || this.isTaskFinalized(detail.status)) { + return; + } + + this.restoreTrackedTask({ + ...detail, + queryId: taskRecord.queryId, + question: taskRecord.question, + taskId: taskRecord.id, + }); + }); + + this.initialized = true; + } + public stopPolling(): void { if (this.pollingIntervalId) { clearInterval(this.pollingIntervalId); @@ -468,4 +514,19 @@ export class AskingTaskTracker implements IAskingTaskTracker { return false; } + + private restoreTrackedTask(result: TrackedAskingResult) { + const restoredTask: TrackedTask = { + queryId: result.queryId, + taskId: result.taskId, + lastPolled: Date.now(), + question: result.question, + result, + isFinalized: false, + }; + this.trackedTasks.set(result.queryId, restoredTask); + if (result.taskId) { + this.trackedTasksById.set(result.taskId, restoredTask); + } + } } diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 6ad7384bf8..404543ab93 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -221,5 +221,11 @@ export const initComponents = () => { }; }; -// singleton components -export const components = initComponents(); +declare global { + // eslint-disable-next-line no-var + var __wrenComponents: ReturnType | undefined; +} + +// Keep a single server-side component graph across Next.js dev reloads. +export const components = + globalThis.__wrenComponents || (globalThis.__wrenComponents = initComponents()); From 5317a98e2b4331477af54d5b1c863d2e98f9e2fa Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 21:14:06 +0530 Subject: [PATCH 0018/1087] updated questions recommadation --- wren-ui/src/hooks/useAdjustAnswer.tsx | 52 ++++++++++-- wren-ui/src/hooks/useAskPrompt.tsx | 113 ++++++++++++++++++++------ wren-ui/src/pages/home/[id].tsx | 95 ++++++++++++++++++---- 3 files changed, 208 insertions(+), 52 deletions(-) diff --git a/wren-ui/src/hooks/useAdjustAnswer.tsx b/wren-ui/src/hooks/useAdjustAnswer.tsx index 4a4bdbdcd4..a2d2680be8 100644 --- a/wren-ui/src/hooks/useAdjustAnswer.tsx +++ b/wren-ui/src/hooks/useAdjustAnswer.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo } from 'react'; +import { useCallback, useEffect, useMemo, useRef } from 'react'; import { cloneDeep } from 'lodash'; import { ApolloClient, NormalizedCacheObject } from '@apollo/client'; import { THREAD } from '@/apollo/client/graphql/home'; @@ -71,9 +71,10 @@ export default function useAdjustAnswer(threadId?: number) { onError: (error) => console.error(error), }); const [fetchThreadResponse, threadResponseResult] = - useThreadResponseLazyQuery({ - pollInterval: 1000, - }); + useThreadResponseLazyQuery(); + const threadResponsePollingRef = useRef | null>( + null, + ); const loading = adjustThreadResponseResult.loading; @@ -87,9 +88,38 @@ export default function useAdjustAnswer(threadId?: number) { }; }, [adjustmentTask]); + const stopThreadResponsePolling = useCallback(() => { + if (threadResponsePollingRef.current) { + clearInterval(threadResponsePollingRef.current); + threadResponsePollingRef.current = null; + } + }, []); + + const startThreadResponsePolling = useCallback( + async (responseId?: number) => { + if (!responseId) return; + + stopThreadResponsePolling(); + + const run = async () => { + try { + await fetchThreadResponse({ + variables: { responseId }, + }); + } catch (error) { + console.error(error); + } + }; + + await run(); + threadResponsePollingRef.current = setInterval(run, 1000); + }, + [fetchThreadResponse, stopThreadResponsePolling], + ); + useEffect(() => { const isFinished = getIsFinished(adjustmentTask?.status); - if (isFinished) threadResponseResult.stopPolling(); + if (isFinished) stopThreadResponsePolling(); }, [adjustmentTask?.status]); const onAdjustReasoningSteps = async ( @@ -112,9 +142,7 @@ export default function useAdjustAnswer(threadId?: number) { const nextThreadResponse = response.data?.adjustThreadResponse; if (!nextThreadResponse?.id) return; - await fetchThreadResponse({ - variables: { responseId: nextThreadResponse.id }, - }); + await startThreadResponsePolling(nextThreadResponse.id); // update new thread response to cache handleUpdateThreadCache( @@ -161,9 +189,15 @@ export default function useAdjustAnswer(threadId?: number) { if (!responseId) return; await rerunAdjustmentTask({ variables: { responseId } }); - await fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; + useEffect(() => { + return () => { + stopThreadResponsePolling(); + }; + }, [stopThreadResponsePolling]); + return { data, loading, diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index ce0451eb94..0c247dbfa2 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -1,4 +1,4 @@ -import { useCallback, useEffect, useMemo, useState } from 'react'; +import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { cloneDeep, uniq } from 'lodash'; import { AdjustmentTask, @@ -169,18 +169,20 @@ export default function useAskPrompt(threadId?: number) { const [rerunAskingTask] = useRerunAskingTaskMutation({ onError: (error) => console.error(error), }); - const [fetchAskingTask, askingTaskResult] = useAskingTaskLazyQuery({ - pollInterval: 1000, - }); + const [fetchAskingTask, askingTaskResult] = useAskingTaskLazyQuery(); const [fetchAskingStreamTask, askingStreamTaskResult] = useAskingStreamTask(); const [createInstantRecommendedQuestions] = useCreateInstantRecommendedQuestionsMutation({ onError: (error) => console.error(error), }); const [fetchInstantRecommendedQuestions, instantRecommendedQuestionsResult] = - useInstantRecommendedQuestionsLazyQuery({ - pollInterval: 1000, - }); + useInstantRecommendedQuestionsLazyQuery(); + const askingTaskPollingRef = useRef | null>( + null, + ); + const recommendedPollingRef = useRef | null>( + null, + ); const askingTask = useMemo( () => askingTaskResult.data?.askingTask || null, @@ -197,6 +199,64 @@ export default function useAskPrompt(threadId?: number) { const loading = askingStreamTaskResult.loading; + const stopAskingTaskPolling = useCallback(() => { + if (askingTaskPollingRef.current) { + clearInterval(askingTaskPollingRef.current); + askingTaskPollingRef.current = null; + } + }, []); + + const stopRecommendedPolling = useCallback(() => { + if (recommendedPollingRef.current) { + clearInterval(recommendedPollingRef.current); + recommendedPollingRef.current = null; + } + }, []); + + const startAskingTaskPolling = useCallback( + async (taskId?: string) => { + if (!taskId) return; + + stopAskingTaskPolling(); + + const run = async () => { + try { + await fetchAskingTask({ + variables: { taskId }, + }); + } catch (error) { + console.error(error); + } + }; + + await run(); + askingTaskPollingRef.current = setInterval(run, 1000); + }, + [fetchAskingTask, stopAskingTaskPolling], + ); + + const startRecommendedPolling = useCallback( + async (taskId?: string) => { + if (!taskId) return; + + stopRecommendedPolling(); + + const run = async () => { + try { + await fetchInstantRecommendedQuestions({ + variables: { taskId }, + }); + } catch (error) { + console.error(error); + } + }; + + await run(); + recommendedPollingRef.current = setInterval(run, 1000); + }, + [fetchInstantRecommendedQuestions, stopRecommendedPolling], + ); + const data = useMemo( () => ({ originalQuestion, @@ -221,10 +281,8 @@ export default function useAskPrompt(threadId?: number) { const taskId = response.data?.createInstantRecommendedQuestions?.id; if (!taskId) return; - fetchInstantRecommendedQuestions({ - variables: { taskId }, - }); - }, [originalQuestion, threadQuestions]); + await startRecommendedPolling(taskId); + }, [originalQuestion, threadQuestions, startRecommendedPolling]); const checkFetchAskingStreamTask = useCallback( (task: AskingTask) => { @@ -241,7 +299,7 @@ export default function useAskPrompt(threadId?: number) { useEffect(() => { const isFinished = getIsFinished(askingTask?.status); - if (isFinished) askingTaskResult.stopPolling(); + if (isFinished) stopAskingTaskPolling(); // handle update cache for preparing component if (isNeedPreparing(askingTask)) { @@ -260,8 +318,9 @@ export default function useAskPrompt(threadId?: number) { }, [askingTask?.type]); useEffect(() => { - if (isRecommendedFinished(recommendedQuestions?.status)) - instantRecommendedQuestionsResult.stopPolling(); + if (isRecommendedFinished(recommendedQuestions?.status)) { + stopRecommendedPolling(); + } }, [recommendedQuestions]); useEffect(() => { @@ -277,6 +336,7 @@ export default function useAskPrompt(threadId?: number) { await cancelAskingTask({ variables: { taskId } }).catch((error) => console.error(error), ); + stopAskingTaskPolling(); // waiting for polling fetching stop await nextTick(1000); } @@ -294,11 +354,11 @@ export default function useAskPrompt(threadId?: number) { const taskId = response.data?.rerunAskingTask?.id; if (!taskId) return; - const { data } = await fetchAskingTask({ - variables: { taskId }, - }); + const { data } = await fetchAskingTask({ variables: { taskId } }); if (!data?.askingTask) return; + await startAskingTaskPolling(taskId); + // update the asking task in cache manually handleUpdateRerunAskingTaskCache( threadId, @@ -321,9 +381,7 @@ export default function useAskPrompt(threadId?: number) { const taskId = response.data?.createAskingTask?.id; if (!taskId) return; - await fetchAskingTask({ - variables: { taskId }, - }); + await startAskingTaskPolling(taskId); } catch (error) { console.error(error); } @@ -332,16 +390,21 @@ export default function useAskPrompt(threadId?: number) { const onFetching = async (queryId: string) => { if (!queryId) return; - await fetchAskingTask({ - variables: { taskId: queryId }, - }); + await startAskingTaskPolling(queryId); }; - const onStopPolling = () => askingTaskResult.stopPolling(); + const onStopPolling = () => stopAskingTaskPolling(); const onStopStreaming = () => askingStreamTaskResult.reset(); - const onStopRecommend = () => instantRecommendedQuestionsResult.stopPolling(); + const onStopRecommend = () => stopRecommendedPolling(); + + useEffect(() => { + return () => { + stopAskingTaskPolling(); + stopRecommendedPolling(); + }; + }, [stopAskingTaskPolling, stopRecommendedPolling]); const onStoreThreadQuestions = (questions: string[]) => setThreadQuestions(questions); diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index 77c1d2650a..2e4800f3d1 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -2,10 +2,10 @@ import { useRouter } from 'next/router'; import { useParams } from 'next/navigation'; import { ComponentRef, + useRef, useCallback, useEffect, useMemo, - useRef, useState, } from 'react'; import { isEmpty } from 'lodash'; @@ -123,7 +123,6 @@ export default function HomeThread() { }); const [fetchThreadResponse, threadResponseResult] = useThreadResponseLazyQuery({ - pollInterval: 1000, onCompleted(next) { const nextResponse = next.threadResponse; updateThreadQuery((prev) => ({ @@ -146,9 +145,13 @@ export default function HomeThread() { const [ fetchThreadRecommendationQuestions, threadRecommendationQuestionsResult, - ] = useGetThreadRecommendationQuestionsLazyQuery({ - pollInterval: 1000, - }); + ] = useGetThreadRecommendationQuestionsLazyQuery(); + const threadResponsePollingRef = useRef | null>( + null, + ); + const threadRecommendationPollingRef = useRef< + ReturnType | null + >(null); const [generateThreadResponseAnswer] = useGenerateThreadResponseAnswerMutation({ @@ -183,6 +186,64 @@ export default function HomeThread() { [pollingResponse], ); + const stopThreadResponsePolling = useCallback(() => { + if (threadResponsePollingRef.current) { + clearInterval(threadResponsePollingRef.current); + threadResponsePollingRef.current = null; + } + }, []); + + const startThreadResponsePolling = useCallback( + async (responseId?: number) => { + if (!responseId) return; + + stopThreadResponsePolling(); + + const run = async () => { + try { + await fetchThreadResponse({ + variables: { responseId }, + }); + } catch (error) { + console.error(error); + } + }; + + await run(); + threadResponsePollingRef.current = setInterval(run, 1000); + }, + [fetchThreadResponse, stopThreadResponsePolling], + ); + + const stopThreadRecommendationPolling = useCallback(() => { + if (threadRecommendationPollingRef.current) { + clearInterval(threadRecommendationPollingRef.current); + threadRecommendationPollingRef.current = null; + } + }, []); + + const startThreadRecommendationPolling = useCallback( + async (nextThreadId?: number) => { + if (!nextThreadId) return; + + stopThreadRecommendationPolling(); + + const run = async () => { + try { + await fetchThreadRecommendationQuestions({ + variables: { threadId: nextThreadId }, + }); + } catch (error) { + console.error(error); + } + }; + + await run(); + threadRecommendationPollingRef.current = setInterval(run, 1000); + }, + [fetchThreadRecommendationQuestions, stopThreadRecommendationPolling], + ); + const onFixSQLStatement = async (responseId: number, sql: string) => { await updateThreadResponse({ variables: { where: { id: responseId }, data: { sql } }, @@ -193,14 +254,14 @@ export default function HomeThread() { if (!responseId) return; await generateThreadResponseAnswer({ variables: { responseId } }); - fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; const onGenerateThreadResponseChart = async (responseId: number) => { if (!responseId) return; await generateThreadResponseChart({ variables: { responseId } }); - fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; const onAdjustThreadResponseChart = async ( @@ -212,14 +273,14 @@ export default function HomeThread() { await adjustThreadResponseChart({ variables: { responseId, data }, }); - fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; const onGenerateThreadRecommendedQuestions = async () => { if (!threadId) return; await generateThreadRecommendationQuestions({ variables: { threadId } }); - fetchThreadRecommendationQuestions({ variables: { threadId } }); + await startThreadRecommendationPolling(threadId); }; const handleUnfinishedTasks = useCallback( @@ -244,12 +305,10 @@ export default function HomeThread() { canFetchThreadResponse(unfinishedThreadResponse?.askingTask) && unfinishedThreadResponse ) { - fetchThreadResponse({ - variables: { responseId: unfinishedThreadResponse.id }, - }); + startThreadResponsePolling(unfinishedThreadResponse.id); } }, - [askPrompt, fetchThreadResponse], + [askPrompt, startThreadResponsePolling], ); // store thread questions for instant recommended questions @@ -264,13 +323,13 @@ export default function HomeThread() { // stop all requests when change thread useEffect(() => { if (threadId !== null) { - fetchThreadRecommendationQuestions({ variables: { threadId } }); + startThreadRecommendationPolling(threadId); setShowRecommendedQuestions(true); } return () => { askPrompt.onStopPolling(); - threadResponseResult.stopPolling(); - threadRecommendationQuestionsResult.stopPolling(); + stopThreadResponsePolling(); + stopThreadRecommendationPolling(); $prompt.current?.close(); }; }, [threadId]); @@ -284,7 +343,7 @@ export default function HomeThread() { useEffect(() => { if (isPollingResponseFinished) { - threadResponseResult.stopPolling(); + stopThreadResponsePolling(); setShowRecommendedQuestions(true); } }, [isPollingResponseFinished]); @@ -298,7 +357,7 @@ export default function HomeThread() { useEffect(() => { if (isRecommendedFinished(recommendedQuestions?.status)) { - threadRecommendationQuestionsResult.stopPolling(); + stopThreadRecommendationPolling(); } }, [recommendedQuestions]); From 0ee26baef0eeb3c32c956184be08e81667d02ea5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 18 May 2026 21:58:04 +0530 Subject: [PATCH 0019/1087] updated llm --- wren-ai-service/src/providers/llm/__init__.py | 68 ++++++++++++++++--- 1 file changed, 57 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/providers/llm/__init__.py b/wren-ai-service/src/providers/llm/__init__.py index c3f2a47b3f..0469adf7e8 100644 --- a/wren-ai-service/src/providers/llm/__init__.py +++ b/wren-ai-service/src/providers/llm/__init__.py @@ -221,7 +221,50 @@ def build_chunk(chunk: Any) -> StreamingChunk: return chunk_message -def convert_message_to_openai_format(message: ChatMessage) -> Dict[str, str]: +def _get_message_role_value(message: Any) -> str: + role = getattr(message, "role", ChatRole.USER) + return role.value if hasattr(role, "value") else str(role) + + +def _get_message_text_content(message: Any) -> Optional[str]: + try: + content = getattr(message, "content", None) + if isinstance(content, str) and content: + return content + except AttributeError: + # Haystack 2.x removed `.content` in favor of `.text`. + pass + + text = getattr(message, "text", None) + if isinstance(text, str) and text: + return text + + raw_content = getattr(message, "_content", None) + if isinstance(raw_content, str) and raw_content: + return raw_content + + if isinstance(raw_content, list): + text_parts = [] + for part in raw_content: + if isinstance(part, str) and part: + text_parts.append(part) + continue + + part_text = getattr(part, "text", None) + if isinstance(part_text, str) and part_text: + text_parts.append(part_text) + if text_parts: + return "\n".join(text_parts) + + return None + + +def _get_message_image_url(message: Any) -> Optional[str]: + image_url = getattr(message, "image_url", None) + return image_url if isinstance(image_url, str) and image_url else None + + +def convert_message_to_openai_format(message: Any) -> Dict[str, Any]: """ Convert a message to the format expected by OpenAI's Chat API. @@ -232,21 +275,24 @@ def convert_message_to_openai_format(message: ChatMessage) -> Dict[str, str]: - `content` - `name` (optional) """ - openai_msg = {"role": message.role.value} + openai_msg = {"role": _get_message_role_value(message)} + message_text = _get_message_text_content(message) + image_url = _get_message_image_url(message) - if message.content and hasattr(message, "image_url") and message.image_url: + if message_text and image_url: openai_msg["content"] = [ - {"type": "text", "text": message.content}, - {"type": "image_url", "image_url": {"url": message.image_url}}, + {"type": "text", "text": message_text}, + {"type": "image_url", "image_url": {"url": image_url}}, ] - elif message.content: - openai_msg["content"] = message.content - elif hasattr(message, "image_url") and message.image_url: + elif message_text: + openai_msg["content"] = message_text + elif image_url: openai_msg["content"] = [ - {"type": "image_url", "image_url": {"url": message.image_url}} + {"type": "image_url", "image_url": {"url": image_url}} ] - if hasattr(message, "name") and message.name: - openai_msg["name"] = message.name + name = getattr(message, "name", None) + if name: + openai_msg["name"] = name return openai_msg From ed64d8ca849324b941cbd7d879b0bd1c71a7706d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 19 May 2026 19:10:23 +0530 Subject: [PATCH 0020/1087] updated questions --- .../server/services/askingTaskTracker.ts | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index 204a5cf305..a7a1598b9f 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -1,4 +1,5 @@ import { getLogger } from '@server/utils'; +import { isEqual } from 'lodash'; import { AskResult, AskResultType, @@ -507,12 +508,10 @@ export class AskingTaskTracker implements IAskingTaskTracker { previousResult: AskResult, newResult: AskResult, ): boolean { - // check status change - if (previousResult?.status !== newResult.status) { - return true; - } - - return false; + return !isEqual( + this.getComparableAskResult(previousResult), + this.getComparableAskResult(newResult), + ); } private restoreTrackedTask(result: TrackedAskingResult) { @@ -521,7 +520,7 @@ export class AskingTaskTracker implements IAskingTaskTracker { taskId: result.taskId, lastPolled: Date.now(), question: result.question, - result, + result: this.getComparableAskResult(result), isFinalized: false, }; this.trackedTasks.set(result.queryId, restoredTask); @@ -529,4 +528,19 @@ export class AskingTaskTracker implements IAskingTaskTracker { this.trackedTasksById.set(result.taskId, restoredTask); } } + + private getComparableAskResult(result: AskResult | TrackedAskingResult) { + return { + status: result?.status ?? null, + type: result?.type ?? null, + response: result?.response ?? null, + error: result?.error ?? null, + rephrasedQuestion: result?.rephrasedQuestion ?? null, + intentReasoning: result?.intentReasoning ?? null, + sqlGenerationReasoning: result?.sqlGenerationReasoning ?? null, + retrievedTables: result?.retrievedTables ?? null, + invalidSql: result?.invalidSql ?? null, + traceId: result?.traceId ?? null, + } as AskResult; + } } From 83e6479176534310abb926aff4ccca9037eb7c17 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 19 May 2026 21:44:55 +0530 Subject: [PATCH 0021/1087] updated questions recommad --- wren-ai-service/src/pipelines/common.py | 4 ++ .../pipelines/generation/sql_correction.py | 5 +- .../retrieval/db_schema_retrieval.py | 50 ++++++++++++++++--- .../historical_question_retrieval.py | 10 +++- .../src/pipelines/retrieval/instructions.py | 24 +++++++++ .../retrieval/sql_pairs_retrieval.py | 7 +++ wren-ai-service/src/web/v1/services/ask.py | 18 +++++-- .../src/web/v1/services/ask_feedback.py | 15 ++++-- .../apollo/server/adaptors/wrenAIAdaptor.ts | 5 +- .../apollo/server/services/askingService.ts | 6 +++ .../apollo/server/services/deployService.ts | 1 + .../apollo/server/services/projectService.ts | 1 + 12 files changed, 129 insertions(+), 17 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 940fa66eb7..825e264e98 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -88,6 +88,10 @@ async def retrieve_metadata(project_id: str, retriever) -> dict[str, Any]: result = await retriever.run(query_embedding=[], filters=filters) documents = result["documents"] + if not documents and project_id: + result = await retriever.run(query_embedding=[], filters=None) + documents = result["documents"] + # only one document for a project, thus we can return the first one if documents: doc = documents[0] diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 791b797f0d..d3d3986205 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -83,7 +83,10 @@ def get_sql_correction_system_prompt( {% endif %} ### QUESTION ### -SQL: {{ invalid_generation_result.sql }} +{% if invalid_generation_result.original_sql %} +Original SQL: {{ invalid_generation_result.original_sql }} +{% endif %} +Invalid SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} Let's think step by step. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 7649b702f2..de5e3051d6 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -141,7 +141,7 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> d async def table_retrieval( embedding: dict, project_id: str, tables: list[str], table_retriever: Any ) -> dict: - filters = { + base_filters = { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, @@ -149,23 +149,48 @@ async def table_retrieval( } if project_id: - filters["conditions"].append( + base_filters["conditions"].append( {"field": "project_id", "operator": "==", "value": project_id} ) if embedding: + result = await table_retriever.run( + query_embedding=embedding.get("embedding"), + filters=base_filters, + ) + if result.get("documents") or not project_id: + return result + fallback_filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + ], + } return await table_retriever.run( query_embedding=embedding.get("embedding"), - filters=filters, + filters=fallback_filters, ) else: - filters["conditions"].append( + base_filters["conditions"].append( {"field": "name", "operator": "in", "value": tables} ) + result = await table_retriever.run( + query_embedding=[], + filters=base_filters, + ) + if result.get("documents") or not project_id: + return result + fallback_filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "name", "operator": "in", "value": tables}, + ], + } return await table_retriever.run( query_embedding=[], - filters=filters, + filters=fallback_filters, ) @@ -199,7 +224,20 @@ async def dbschema_retrieval( ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] + if results.get("documents") or not project_id: + return results["documents"] + + fallback_filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } + fallback_results = await dbschema_retriever.run( + query_embedding=[], filters=fallback_filters + ) + return fallback_results["documents"] return [] diff --git a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py index 0dcbc839ab..77911ef659 100644 --- a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py @@ -51,7 +51,10 @@ async def count_documents( else None ) - return await view_questions_store.count_documents(filters=filters) + count = await view_questions_store.count_documents(filters=filters) + if count == 0 and project_id: + count = await view_questions_store.count_documents(filters=None) + return count @observe(capture_input=False, capture_output=False) @@ -84,6 +87,11 @@ async def retrieval( query_embedding=embedding.get("embedding"), filters=filters, ) + if not view_question_res.get("documents") and project_id: + view_question_res = await view_questions_retriever.run( + query_embedding=embedding.get("embedding"), + filters=None, + ) return dict(documents=view_question_res.get("documents")) return {} diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 86c17e93de..22e688159e 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -70,6 +70,8 @@ async def count_documents( else None ) document_count = await store.count_documents(filters=filters) + if document_count == 0 and project_id: + document_count = await store.count_documents(filters=None) return document_count @@ -102,6 +104,17 @@ async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: query_embedding=embedding.get("embedding"), filters=filters, ) + if not res.get("documents") and project_id: + fallback_filters = { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": False}, + ], + } + res = await retriever.run( + query_embedding=embedding.get("embedding"), + filters=fallback_filters, + ) return dict(documents=res.get("documents")) @@ -156,6 +169,17 @@ async def default_instructions( query_embedding=None, filters=filters, ) + if not _res.get("documents") and project_id: + fallback_filters = { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": True}, + ], + } + _res = await retriever.run( + query_embedding=None, + filters=fallback_filters, + ) res = scope_filter.run( documents=_res.get("documents"), diff --git a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py index 3fe44f32eb..c8d2997c17 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py @@ -49,6 +49,8 @@ async def count_documents( else None ) document_count = await store.count_documents(filters=filters) + if document_count == 0 and project_id: + document_count = await store.count_documents(filters=None) return document_count @@ -78,6 +80,11 @@ async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: query_embedding=embedding.get("embedding"), filters=filters, ) + if not res.get("documents") and project_id: + res = await retriever.run( + query_embedding=embedding.get("embedding"), + filters=None, + ) return dict(documents=res.get("documents")) return {} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index aa26fa3f81..45c31cf8be 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -518,6 +518,7 @@ async def ask( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + sql_diagnosis_reasoning = None current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -545,16 +546,21 @@ async def ask( "post_process" ].get("reasoning") + correction_error_message = error_message + if sql_diagnosis_reasoning: + correction_error_message = ( + f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" + ) + sql_correction_results = await self._pipelines[ "sql_correction" ].run( contexts=table_ddls, instructions=instructions, invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "original_sql": original_sql, + "sql": invalid_sql, + "error": correction_error_message, }, project_id=ask_request.project_id, use_dry_plan=use_dry_plan, @@ -579,6 +585,10 @@ async def ask( failed_dry_run_result = sql_correction_results["post_process"][ "invalid_generation_result" ] + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) if api_results: if not self._is_stopped(query_id, self._ask_results): diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 9c5b06a772..25044de18c 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -206,6 +206,7 @@ async def ask_feedback( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + sql_diagnosis_reasoning = None self._ask_feedback_results[ query_id @@ -222,21 +223,27 @@ async def ask_feedback( original_sql=original_sql, invalid_sql=invalid_sql, error_message=error_message, + language=ask_feedback_request.configurations.language, ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") + correction_error_message = error_message + if sql_diagnosis_reasoning: + correction_error_message = ( + f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" + ) + sql_correction_results = await self._pipelines[ "sql_correction" ].run( contexts=table_ddls, instructions=instructions, invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "original_sql": original_sql, + "sql": invalid_sql, + "error": correction_error_message, }, project_id=ask_feedback_request.project_id, sql_functions=sql_functions, diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 1877f1da42..6d5b46d88e 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -237,6 +237,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, { query: input.query, id: input.deployId, + project_id: input.projectId, histories: this.transformHistoryInput(input.histories), configurations: input.configurations, }); @@ -333,13 +334,14 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } public async deploy(deployData: DeployData): Promise { - const { manifest, hash } = deployData; + const { manifest, hash, projectId } = deployData; try { const res = await axios.post( `${this.wrenAIBaseEndpoint}/v1/semantics-preparations`, { mdl: JSON.stringify(manifest), id: hash, + project_id: projectId.toString(), }, ); const deployId = res.data.id; @@ -373,6 +375,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { const body = { mdl: JSON.stringify(input.manifest), previous_questions: input.previousQuestions, + project_id: input.projectId, max_questions: input.maxQuestions, max_categories: input.maxCategories, configuration: input.configuration, diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index b047312bb3..12add0a435 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -58,6 +58,7 @@ export interface Task { export interface AskingPayload { threadId?: number; language: string; + projectId?: number; } export interface AskingTaskInput { @@ -545,6 +546,7 @@ export class AskingService implements IAskingService { const questions = slicedThreadResponses.map(({ question }) => question); const recommendQuestionData: RecommendationQuestionsInput = { manifest, + projectId: project.id.toString(), previousQuestions: questions, ...this.getThreadRecommendationQuestionsConfig(project), }; @@ -601,6 +603,8 @@ export class AskingService implements IAskingService { threadResponseId?: number, ): Promise { const { threadId, language } = payload; + const projectId = + payload.projectId ?? (await this.projectService.getCurrentProject()).id; const deployId = await this.getDeployId(); // if it's a follow-up question, then the input will have a threadId @@ -613,6 +617,7 @@ export class AskingService implements IAskingService { query: input.question, histories, deployId, + projectId: projectId.toString(), configurations: { language }, rerunFromCancelled, previousTaskId, @@ -1013,6 +1018,7 @@ export class AskingService implements IAskingService { const response = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, + projectId: project.id.toString(), previousQuestions: input.previousQuestions, ...this.getThreadRecommendationQuestionsConfig(project), }); diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index a6885146cb..3c751d3156 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -102,6 +102,7 @@ export class DeployService implements IDeployService { await this.wrenAIAdaptor.deploy({ manifest, hash, + projectId, }); // update deploy status diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 88f0915082..3f1b74318b 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -140,6 +140,7 @@ export class ProjectService implements IProjectService { const recommendQuestionResult = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, + projectId: project.id.toString(), ...this.getProjectRecommendationQuestionsConfig(project), }); From 92fc662222b86e5978d4f7fd8ba76c75b7b3762d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 19 May 2026 21:53:55 +0530 Subject: [PATCH 0022/1087] updat questions recommad --- wren-ui/jest.config.js | 1 + wren-ui/src/apollo/server/models/adaptor.ts | 2 ++ 2 files changed, 3 insertions(+) diff --git a/wren-ui/jest.config.js b/wren-ui/jest.config.js index 9cac8ffa54..2bdc6395ce 100644 --- a/wren-ui/jest.config.js +++ b/wren-ui/jest.config.js @@ -3,6 +3,7 @@ module.exports = { preset: 'ts-jest', testEnvironment: 'node', moduleNameMapper: { + '^@/(.*)$': '/src/$1', '^@server/(.*)$': '/src/apollo/server/$1', }, modulePathIgnorePatterns: ['/e2e/'], diff --git a/wren-ui/src/apollo/server/models/adaptor.ts b/wren-ui/src/apollo/server/models/adaptor.ts index 8135c07160..6a4fad2a73 100644 --- a/wren-ui/src/apollo/server/models/adaptor.ts +++ b/wren-ui/src/apollo/server/models/adaptor.ts @@ -51,6 +51,7 @@ export enum WrenAILanguage { export interface DeployData { manifest: Manifest; hash: string; + projectId: number; } // ask @@ -73,6 +74,7 @@ export interface ProjectConfigurations { export interface AskInput { query: string; deployId: string; + projectId?: string; histories?: ThreadResponse[]; configurations?: ProjectConfigurations; } From b550791983b6c197dbfac1690968ce308e7df96e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 00:04:09 +0530 Subject: [PATCH 0023/1087] updated one --- .../src/providers/embedder/litellm.py | 153 ++++++++++++++++-- 1 file changed, 140 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index 7587556236..affc633b9c 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -1,8 +1,11 @@ import asyncio +import json import logging import os +from types import SimpleNamespace from typing import Any, Dict, List, Optional, Tuple +import aiohttp import backoff import openai from haystack import Document, component @@ -17,6 +20,10 @@ MIN_EMBED_INPUT_CHARS = 128 +class EmbeddingRequestError(Exception): + pass + + def _normalize_model_name(model: str, api_base_url: Optional[str]) -> str: # OpenAI-compatible local servers often expect the raw model name and will # reject litellm-style "openai/" prefixes. @@ -25,6 +32,100 @@ def _normalize_model_name(model: str, api_base_url: Optional[str]) -> str: return model +def _should_use_minimal_http_client(api_base_url: Optional[str]) -> bool: + if not api_base_url: + return False + + return "api.openai.com" not in api_base_url.lower() + + +def _build_embedding_meta(response: Any) -> Dict[str, Any]: + usage = getattr(response, "usage", {}) or {} + usage_dict = dict(usage) if isinstance(usage, dict) or hasattr(usage, "__iter__") else {} + + return { + "model": getattr(response, "model", ""), + "usage": usage_dict, + } + + +def _get_usage_value(usage: Any, key: str) -> int: + if isinstance(usage, dict): + return usage.get(key, 0) or 0 + + return getattr(usage, key, 0) or 0 + + +def _coerce_embedding_response(payload: Dict[str, Any]) -> Any: + data = payload.get("data") + if not data and payload.get("embedding") is not None: + data = [{"embedding": payload["embedding"]}] + + if not isinstance(data, list) or not data: + raise EmbeddingRequestError( + "Embedding provider returned an invalid response payload." + ) + + normalized_data = [] + for item in data: + embedding = item.get("embedding") if isinstance(item, dict) else None + if embedding is None: + raise EmbeddingRequestError( + "Embedding provider response did not include an embedding." + ) + normalized_data.append(SimpleNamespace(embedding=embedding)) + + return SimpleNamespace( + model=payload.get("model", ""), + data=normalized_data, + usage=payload.get("usage", {}) or {}, + ) + + +async def _create_embedding_via_http( + *, + model: str, + input_text: str, + api_key: Optional[str], + api_base_url: str, + timeout: Optional[float], + **kwargs, +): + endpoint = f"{remove_trailing_slash(api_base_url)}/embeddings" + headers = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + payload = { + "model": _normalize_model_name(model, api_base_url), + "input": input_text, + } + payload.update({key: value for key, value in kwargs.items() if value is not None}) + + client_timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None + async with aiohttp.ClientSession() as session: + async with session.post( + endpoint, + json=payload, + headers=headers, + timeout=client_timeout, + ) as response: + body = await response.text() + if response.status >= 400: + raise EmbeddingRequestError( + f"Embedding request failed with status {response.status}: {body}" + ) + + try: + payload = json.loads(body) + except json.JSONDecodeError as error: + raise EmbeddingRequestError( + "Embedding provider returned a non-JSON response." + ) from error + + return _coerce_embedding_response(payload) + + async def _create_embedding( *, model: str, @@ -34,6 +135,16 @@ async def _create_embedding( timeout: Optional[float], **kwargs, ): + if _should_use_minimal_http_client(api_base_url): + return await _create_embedding_via_http( + model=model, + input_text=input_text, + api_key=api_key, + api_base_url=api_base_url, + timeout=timeout, + **kwargs, + ) + client = openai.AsyncOpenAI( api_key=api_key, base_url=api_base_url, @@ -111,7 +222,12 @@ def __init__( self._kwargs = kwargs @component.output_types(embedding=List[float], meta=Dict[str, Any]) - @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) + @backoff.on_exception( + backoff.expo, + (aiohttp.ClientError, asyncio.TimeoutError, EmbeddingRequestError, openai.APIError), + max_time=60.0, + max_tries=3, + ) async def run(self, text: str): if not isinstance(text, str): raise TypeError( @@ -136,10 +252,7 @@ async def run(self, text: str): **self._kwargs, ) - meta = { - "model": response.model, - "usage": dict(response.usage) if hasattr(response, "usage") else {}, - } + meta = _build_embedding_meta(response) return {"embedding": response.data[0].embedding, "meta": meta} @@ -181,7 +294,12 @@ async def embed_single_text(text: str) -> Any: timeout=self._timeout, **self._kwargs, ) - except openai.APIError as error: + except ( + aiohttp.ClientError, + asyncio.TimeoutError, + EmbeddingRequestError, + openai.APIError, + ) as error: if ( not _is_input_too_large_error(error) or len(candidate_text) <= MIN_EMBED_INPUT_CHARS @@ -214,20 +332,29 @@ async def embed_single_text(text: str) -> Any: all_embeddings.extend(embeddings) if "model" not in meta: - meta["model"] = response.model + meta["model"] = getattr(response, "model", "") if "usage" not in meta: - meta["usage"] = ( - dict(response.usage) if hasattr(response, "usage") else {} - ) + meta["usage"] = _build_embedding_meta(response)["usage"] else: if hasattr(response, "usage"): - meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens - meta["usage"]["total_tokens"] += response.usage.total_tokens + meta["usage"]["prompt_tokens"] += _get_usage_value( + response.usage, + "prompt_tokens", + ) + meta["usage"]["total_tokens"] += _get_usage_value( + response.usage, + "total_tokens", + ) return all_embeddings, meta @component.output_types(documents=List[Document], meta=Dict[str, Any]) - @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) + @backoff.on_exception( + backoff.expo, + (aiohttp.ClientError, asyncio.TimeoutError, EmbeddingRequestError, openai.APIError), + max_time=60.0, + max_tries=3, + ) async def run(self, documents: List[Document]): if ( not isinstance(documents, list) From 3a0fb75fc643870acede658c90d03f3175e9ebd4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 00:32:07 +0530 Subject: [PATCH 0024/1087] updated one home --- .../apollo/client/graphql/home.generated.ts | 16 +++++++------- wren-ui/src/apollo/client/graphql/home.ts | 6 ++--- .../apollo/server/resolvers/askingResolver.ts | 22 +++++++++++++++---- wren-ui/src/apollo/server/schema.ts | 6 ++--- 4 files changed, 32 insertions(+), 18 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/home.generated.ts b/wren-ui/src/apollo/client/graphql/home.generated.ts index 5a477d401a..0f1e6921b3 100644 --- a/wren-ui/src/apollo/client/graphql/home.generated.ts +++ b/wren-ui/src/apollo/client/graphql/home.generated.ts @@ -23,7 +23,7 @@ export type SuggestedQuestionsQueryVariables = Types.Exact<{ [key: string]: neve export type SuggestedQuestionsQuery = { __typename?: 'Query', suggestedQuestions: { __typename?: 'SuggestedQuestionResponse', questions: Array<{ __typename?: 'SuggestedQuestion', label: string, question: string } | null> } }; export type AskingTaskQueryVariables = Types.Exact<{ - taskId: Types.Scalars['String']; + taskId?: Types.InputMaybe; }>; @@ -144,7 +144,7 @@ export type CreateInstantRecommendedQuestionsMutationVariables = Types.Exact<{ export type CreateInstantRecommendedQuestionsMutation = { __typename?: 'Mutation', createInstantRecommendedQuestions: { __typename?: 'Task', id: string } }; export type InstantRecommendedQuestionsQueryVariables = Types.Exact<{ - taskId: Types.Scalars['String']; + taskId?: Types.InputMaybe; }>; @@ -197,7 +197,7 @@ export type AdjustThreadResponseChartMutationVariables = Types.Exact<{ export type AdjustThreadResponseChartMutation = { __typename?: 'Mutation', adjustThreadResponseChart: { __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null } }; export type AdjustmentTaskQueryVariables = Types.Exact<{ - taskId: Types.Scalars['String']; + taskId?: Types.InputMaybe; }>; @@ -391,7 +391,7 @@ export type SuggestedQuestionsQueryHookResult = ReturnType; export type SuggestedQuestionsQueryResult = Apollo.QueryResult; export const AskingTaskDocument = gql` - query AskingTask($taskId: String!) { + query AskingTask($taskId: String) { askingTask(taskId: $taskId) { ...CommonAskingTask } @@ -960,7 +960,7 @@ export type CreateInstantRecommendedQuestionsMutationHookResult = ReturnType; export type CreateInstantRecommendedQuestionsMutationOptions = Apollo.BaseMutationOptions; export const InstantRecommendedQuestionsDocument = gql` - query InstantRecommendedQuestions($taskId: String!) { + query InstantRecommendedQuestions($taskId: String) { instantRecommendedQuestions(taskId: $taskId) { ...CommonRecommendedQuestionsTask } @@ -1225,9 +1225,9 @@ export type AdjustThreadResponseChartMutationHookResult = ReturnType; export type AdjustThreadResponseChartMutationOptions = Apollo.BaseMutationOptions; export const AdjustmentTaskDocument = gql` - query AdjustmentTask($taskId: String!) { + query AdjustmentTask($taskId: String) { adjustmentTask(taskId: $taskId) { - queryId + queryId status error { code @@ -1330,4 +1330,4 @@ export function useRerunAdjustmentTaskMutation(baseOptions?: Apollo.MutationHook } export type RerunAdjustmentTaskMutationHookResult = ReturnType; export type RerunAdjustmentTaskMutationResult = Apollo.MutationResult; -export type RerunAdjustmentTaskMutationOptions = Apollo.BaseMutationOptions; \ No newline at end of file +export type RerunAdjustmentTaskMutationOptions = Apollo.BaseMutationOptions; diff --git a/wren-ui/src/apollo/client/graphql/home.ts b/wren-ui/src/apollo/client/graphql/home.ts index 1fe6086912..477e7b2717 100644 --- a/wren-ui/src/apollo/client/graphql/home.ts +++ b/wren-ui/src/apollo/client/graphql/home.ts @@ -164,7 +164,7 @@ export const SUGGESTED_QUESTIONS = gql` `; export const ASKING_TASK = gql` - query AskingTask($taskId: String!) { + query AskingTask($taskId: String) { askingTask(taskId: $taskId) { ...CommonAskingTask } @@ -317,7 +317,7 @@ export const CREATE_INSTANT_RECOMMENDED_QUESTIONS = gql` `; export const INSTANT_RECOMMENDED_QUESTIONS = gql` - query InstantRecommendedQuestions($taskId: String!) { + query InstantRecommendedQuestions($taskId: String) { instantRecommendedQuestions(taskId: $taskId) { ...CommonRecommendedQuestionsTask } @@ -389,7 +389,7 @@ export const ADJUST_THREAD_RESPONSE_CHART = gql` `; export const ADJUSTMENT_TASK = gql` - query AdjustmentTask($taskId: String!) { + query AdjustmentTask($taskId: String) { adjustmentTask(taskId: $taskId) { queryId status diff --git a/wren-ui/src/apollo/server/resolvers/askingResolver.ts b/wren-ui/src/apollo/server/resolvers/askingResolver.ts index e674c0e278..b4b54ffdfa 100644 --- a/wren-ui/src/apollo/server/resolvers/askingResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/askingResolver.ts @@ -19,6 +19,7 @@ import { safeFormatSQL } from '@server/utils/sqlFormat'; import { AskingDetailTaskInput, constructCteSql, + RecommendQuestionResultStatus, ThreadRecommendQuestionResult, } from '../services/askingService'; import { @@ -78,7 +79,7 @@ export interface RecommendedQuestionsTask { category: string; sql: string; }[]; - status: RecommendationQuestionStatus; + status: RecommendationQuestionStatus | RecommendQuestionResultStatus; error: WrenAIError | null; } @@ -205,10 +206,13 @@ export class AskingResolver { public async getAskingTask( _root: any, - args: { taskId: string }, + args: { taskId?: string | null }, ctx: IContext, ): Promise { const { taskId } = args; + if (!taskId) { + return null; + } const askingService = ctx.askingService; const askResult = await askingService.getAskingTask(taskId); @@ -543,10 +547,13 @@ export class AskingResolver { public async getAdjustmentTask( _root: any, - args: { taskId: string }, + args: { taskId?: string | null }, ctx: IContext, ): Promise { const { taskId } = args; + if (!taskId) { + return null; + } const askingService = ctx.askingService; const adjustmentTask = await askingService.getAdjustmentTask(taskId); return { @@ -665,10 +672,17 @@ export class AskingResolver { public async getInstantRecommendedQuestions( _root: any, - args: { taskId: string }, + args: { taskId?: string | null }, ctx: IContext, ): Promise { const { taskId } = args; + if (!taskId) { + return { + questions: [], + status: RecommendQuestionResultStatus.NOT_STARTED, + error: null, + }; + } const askingService = ctx.askingService; const result = await askingService.getInstantRecommendedQuestions(taskId); return { diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 64698ad56d..8ad044fb2c 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -1129,7 +1129,7 @@ export const typeDefs = gql` view(where: ViewWhereUniqueInput!): ViewInfo! # Ask - askingTask(taskId: String!): AskingTask + askingTask(taskId: String): AskingTask suggestedQuestions: SuggestedQuestionResponse! threads: [Thread!]! thread(threadId: Int!): DetailedThread! @@ -1137,7 +1137,7 @@ export const typeDefs = gql` nativeSql(responseId: Int!): String! # Adjustment - adjustmentTask(taskId: String!): AdjustmentTask + adjustmentTask(taskId: String): AdjustmentTask # Settings settings: Settings! @@ -1151,7 +1151,7 @@ export const typeDefs = gql` # Recommendation questions getThreadRecommendationQuestions(threadId: Int!): RecommendedQuestionsTask! getProjectRecommendationQuestions: RecommendedQuestionsTask! - instantRecommendedQuestions(taskId: String!): RecommendedQuestionsTask! + instantRecommendedQuestions(taskId: String): RecommendedQuestionsTask! # Dashboard dashboardItems: [DashboardItem!]! From 6b07f406c9ca43088846fafb7f50a0c2538eb1f8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 01:00:33 +0530 Subject: [PATCH 0025/1087] updated schema --- wren-ui/src/apollo/client/graphql/home.generated.ts | 6 +++--- wren-ui/src/apollo/client/graphql/home.ts | 2 +- wren-ui/src/apollo/server/resolvers/askingResolver.ts | 5 ++++- wren-ui/src/apollo/server/schema.ts | 2 +- 4 files changed, 9 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/home.generated.ts b/wren-ui/src/apollo/client/graphql/home.generated.ts index 0f1e6921b3..a8ecfbd961 100644 --- a/wren-ui/src/apollo/client/graphql/home.generated.ts +++ b/wren-ui/src/apollo/client/graphql/home.generated.ts @@ -42,11 +42,11 @@ export type ThreadQueryVariables = Types.Exact<{ export type ThreadQuery = { __typename?: 'Query', thread: { __typename?: 'DetailedThread', id: number, responses: Array<{ __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null }> } }; export type ThreadResponseQueryVariables = Types.Exact<{ - responseId: Types.Scalars['Int']; + responseId?: Types.InputMaybe; }>; -export type ThreadResponseQuery = { __typename?: 'Query', threadResponse: { __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null } }; +export type ThreadResponseQuery = { __typename?: 'Query', threadResponse?: { __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null } | null }; export type CreateAskingTaskMutationVariables = Types.Exact<{ data: Types.AskingTaskInput; @@ -499,7 +499,7 @@ export type ThreadQueryHookResult = ReturnType; export type ThreadLazyQueryHookResult = ReturnType; export type ThreadQueryResult = Apollo.QueryResult; export const ThreadResponseDocument = gql` - query ThreadResponse($responseId: Int!) { + query ThreadResponse($responseId: Int) { threadResponse(responseId: $responseId) { ...CommonResponse } diff --git a/wren-ui/src/apollo/client/graphql/home.ts b/wren-ui/src/apollo/client/graphql/home.ts index 477e7b2717..762e41d00a 100644 --- a/wren-ui/src/apollo/client/graphql/home.ts +++ b/wren-ui/src/apollo/client/graphql/home.ts @@ -194,7 +194,7 @@ export const THREAD = gql` `; export const THREAD_RESPONSE = gql` - query ThreadResponse($responseId: Int!) { + query ThreadResponse($responseId: Int) { threadResponse(responseId: $responseId) { ...CommonResponse } diff --git a/wren-ui/src/apollo/server/resolvers/askingResolver.ts b/wren-ui/src/apollo/server/resolvers/askingResolver.ts index b4b54ffdfa..98c15c342d 100644 --- a/wren-ui/src/apollo/server/resolvers/askingResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/askingResolver.ts @@ -624,10 +624,13 @@ export class AskingResolver { public async getResponse( _root: any, - args: { responseId: number }, + args: { responseId?: number | null }, ctx: IContext, ): Promise { const { responseId } = args; + if (!responseId) { + return null; + } const askingService = ctx.askingService; const response = await askingService.getResponse(responseId); diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 8ad044fb2c..cc52002589 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -1133,7 +1133,7 @@ export const typeDefs = gql` suggestedQuestions: SuggestedQuestionResponse! threads: [Thread!]! thread(threadId: Int!): DetailedThread! - threadResponse(responseId: Int!): ThreadResponse! + threadResponse(responseId: Int): ThreadResponse nativeSql(responseId: Int!): String! # Adjustment From 65c5cb407bf7d8be8788dfa640d13fd08120f2e2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 01:59:54 +0530 Subject: [PATCH 0026/1087] updated schema db --- .../generation/question_recommendation.py | 10 ++++ .../v1/services/question_recommendation.py | 51 ++++++++++++++++--- wren-ui/src/apollo/server/config.ts | 2 +- .../apollo/server/services/askingService.ts | 1 + 4 files changed, 56 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index a6e7c17b02..9b2556f69a 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -16,6 +16,13 @@ logger = logging.getLogger("wren-ai-service") +DEFAULT_QUESTION_CATEGORIES = [ + "Descriptive Questions", + "Segmentation Questions", + "Comparative Questions", + "Data Quality/Accuracy Questions", +] + system_prompt = """ You are an expert in data analysis and SQL query generation. Given a data model specification, optionally a user's question, and a list of categories, your task is to generate insightful, specific questions that can be answered using the provided data model. Each question should be accompanied by a brief explanation of its relevance or importance. @@ -66,6 +73,9 @@ 5. **General Guidelines for All Questions:** - Ensure questions can be answered using the data model. + - Use only the tables, fields, and relationships that are explicitly present in the provided database schema. + - Do not invent tables, columns, business entities, or time dimensions that are not present in the schema. + - Keep the question grounded in the deployed database domain shown by the schema context. - Mix simple and complex questions. - Avoid open-ended questions - each should have a definite answer. - Incorporate time-based analysis where relevant. diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 8aa5ced38e..4432c65ce6 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -31,6 +31,12 @@ DEFAULT_VALIDATION_SQL_FUNCTION_CHARS = 2500 STRICT_VALIDATION_SQL_FUNCTION_ITEMS = 0 STRICT_VALIDATION_SQL_FUNCTION_CHARS = 0 +DEFAULT_QUESTION_CATEGORIES = [ + "Descriptive Questions", + "Segmentation Questions", + "Comparative Questions", + "Data Quality/Accuracy Questions", +] class QuestionRecommendation: @@ -142,6 +148,31 @@ def _is_context_size_error(self, error: Exception) -> bool: ] ) + def _get_target_categories( + self, + requested_categories: list[str], + max_categories: int, + ) -> list[str]: + categories = requested_categories or DEFAULT_QUESTION_CATEGORIES + return categories[:max_categories] + + def _get_underfilled_categories( + self, + response_questions: dict[str, list[dict]], + requested_categories: list[str], + max_categories: int, + max_questions: int, + ) -> list[str]: + target_categories = self._get_target_categories( + requested_categories=requested_categories, + max_categories=max_categories, + ) + return [ + category + for category in target_categories + if len(response_questions.get(category, [])) < max_questions + ] + def _handle_exception( self, event_id: str, @@ -333,6 +364,7 @@ class Request(BaseRequest): event_id: str mdl: str previous_questions: list[str] = [] + categories: list[str] = [] max_questions: int = 5 max_categories: int = 3 regenerate: bool = False @@ -380,6 +412,10 @@ async def recommend(self, input: Request, **kwargs) -> Event: request = { "contexts": table_ddls, "previous_questions": input.previous_questions, + "categories": self._get_target_categories( + requested_categories=input.categories, + max_categories=input.max_categories, + ), "language": input.configurations.language, "max_questions": input.max_questions, "max_categories": input.max_categories, @@ -394,17 +430,18 @@ async def recommend(self, input: Request, **kwargs) -> Event: resource.trace_id = trace_id response = resource.response - categories_count = { - category: input.max_questions - len(questions) - for category, questions in response["questions"].items() - if len(questions) < input.max_questions - } - categories = list(categories_count.keys()) - need_regenerate = len(categories) > 0 and input.regenerate + categories = self._get_underfilled_categories( + response_questions=response["questions"], + requested_categories=request["categories"], + max_categories=input.max_categories, + max_questions=input.max_questions, + ) + need_regenerate = bool(categories) and input.regenerate resource.status = "generating" if need_regenerate else "finished" if resource.status == "finished": + resource.request_from = input.request_from return resource.with_metadata() await self._recommend( diff --git a/wren-ui/src/apollo/server/config.ts b/wren-ui/src/apollo/server/config.ts index ad43e2da84..bafebe866b 100644 --- a/wren-ui/src/apollo/server/config.ts +++ b/wren-ui/src/apollo/server/config.ts @@ -147,7 +147,7 @@ const config = { threadRecommendationQuestionsMaxQuestions: process.env .THREAD_RECOMMENDATION_QUESTIONS_MAX_QUESTIONS ? parseInt(process.env.THREAD_RECOMMENDATION_QUESTIONS_MAX_QUESTIONS) - : 1, + : 3, }; export function getConfig(): IConfig { diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 12add0a435..09d6b7a38e 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1198,6 +1198,7 @@ export class AskingService implements IAskingService { return { maxCategories: config.threadRecommendationQuestionMaxCategories, maxQuestions: config.threadRecommendationQuestionsMaxQuestions, + regenerate: true, configuration: { language: WrenAILanguage[project.language] || WrenAILanguage.EN, }, From a68bd3190e216c57909c327f20b962cef448a128 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 12:17:46 +0530 Subject: [PATCH 0027/1087] updated --- wren-ui/src/common.ts | 14 ++++++++++++-- .../hooks/useRecommendedQuestionsInstruction.tsx | 7 +++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 404543ab93..5d45b01d12 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -43,6 +43,10 @@ import { DashboardCacheBackgroundTracker, } from './apollo/server/backgrounds'; import { SqlPairService } from './apollo/server/services/sqlPairService'; +import { + disposeComponentGraph, + isReusableComponentGraph, +} from './componentGraph'; export const serverConfig = getConfig(); @@ -227,5 +231,11 @@ declare global { } // Keep a single server-side component graph across Next.js dev reloads. -export const components = - globalThis.__wrenComponents || (globalThis.__wrenComponents = initComponents()); +const existingComponents = globalThis.__wrenComponents; + +if (!isReusableComponentGraph(existingComponents)) { + disposeComponentGraph(existingComponents); + globalThis.__wrenComponents = initComponents(); +} + +export const components = globalThis.__wrenComponents; diff --git a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx index 9a04a605a3..e3c2b85c5c 100644 --- a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx +++ b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx @@ -63,6 +63,9 @@ export default function useRecommendedQuestionsInstruction() { const fetchRecommendationQuestionsData = async () => { const result = await fetchRecommendationQuestions(); const data = result.data?.getProjectRecommendationQuestions; + if (!data) { + return; + } // for existing projects that do not have to generate recommended questions yet if (isRecommendedFinished(data.status)) { @@ -79,6 +82,10 @@ export default function useRecommendedQuestionsInstruction() { }, []); useEffect(() => { + if (!recommendedQuestionsTask) { + return; + } + if (isRecommendedFinished(recommendedQuestionsTask?.status)) { recommendationQuestionsResult.stopPolling(); From d07f30e635eacf1a09b8c4c04bfcf3544b2782b6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 12:30:26 +0530 Subject: [PATCH 0028/1087] updated database --- .../pipelines/indexing/test_db_schema.py | 71 ++++++++++++++++++- .../indexing/test_table_description.py | 24 +++++++ wren-ui/src/common.ts | 44 ++++++++++-- 3 files changed, 134 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py index 20bd8ac682..5bc8c1303f 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py @@ -5,7 +5,11 @@ from haystack import Document from pytest_mock import MockFixture -from src.pipelines.indexing.db_schema import DBSchema, DDLChunker +from src.pipelines.indexing.db_schema import ( + MAX_DB_SCHEMA_DOCUMENT_LENGTH, + DBSchema, + DDLChunker, +) @pytest.mark.asyncio @@ -444,6 +448,71 @@ async def test_column_batch_size(): ) +@pytest.mark.asyncio +async def test_long_model_description_is_truncated(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "user", + "properties": { + "displayName": "user", + "description": "x" * 5000, + }, + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=1) + + assert len(actual["documents"]) == 1 + document: Document = actual["documents"][0] + assert len(document.content) < 5000 + assert "..." in document.content + + +@pytest.mark.asyncio +async def test_table_columns_are_bounded_to_document_length(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "user", + "columns": [ + { + "name": f"column_{index}", + "type": "VARCHAR", + "properties": { + "displayName": f"column_{index}", + "description": "x" * 6000, + }, + } + for index in range(2) + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=50) + table_column_documents = [ + document + for document in actual["documents"] + if "TABLE_COLUMNS" in document.content + ] + + assert len(table_column_documents) >= 1 + assert all( + len(document.content) <= MAX_DB_SCHEMA_DOCUMENT_LENGTH + for document in table_column_documents + ) + + @pytest.mark.asyncio async def test_view(): chunker = DDLChunker() diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index 7def966fdb..1585db75d5 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -126,6 +126,30 @@ def test_table_description_missing_description(): assert document.content == str({"name": "user", "description": "", "columns": ""}) +def test_table_description_truncates_long_column_lists(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + { + "name": "user", + "columns": [ + {"name": f"column_{index}"} + for index in range(205) + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = chunker.run(mdl) + + assert len(actual["documents"]) == 1 + document: Document = actual["documents"][0] + assert "... (+5 more columns)" in document.content + + @pytest.mark.asyncio async def test_pipeline_run(mocker: MockFixture): test_mdl = { diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 5d45b01d12..35f6262acb 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -43,13 +43,49 @@ import { DashboardCacheBackgroundTracker, } from './apollo/server/backgrounds'; import { SqlPairService } from './apollo/server/services/sqlPairService'; -import { - disposeComponentGraph, - isReusableComponentGraph, -} from './componentGraph'; export const serverConfig = getConfig(); +type Initializable = { + initialize?: unknown; +}; + +type ReusableComponentGraph = { + askingTaskTracker?: Initializable; + askingService?: Initializable; + projectRecommendQuestionBackgroundTracker?: Initializable; + threadRecommendQuestionBackgroundTracker?: Initializable; + knex?: { + destroy?: () => unknown; + }; +}; + +const hasInitialize = ( + value: Initializable | null | undefined, +): value is { initialize: () => Promise | void } => { + return typeof value?.initialize === 'function'; +}; + +const isReusableComponentGraph = ( + graph?: ReusableComponentGraph, +): boolean => { + return Boolean( + graph && + hasInitialize(graph.askingTaskTracker) && + hasInitialize(graph.askingService) && + hasInitialize(graph.projectRecommendQuestionBackgroundTracker) && + hasInitialize(graph.threadRecommendQuestionBackgroundTracker), + ); +}; + +const disposeComponentGraph = (graph?: ReusableComponentGraph): void => { + if (typeof graph?.knex?.destroy !== 'function') { + return; + } + + void Promise.resolve(graph.knex.destroy()).catch(() => undefined); +}; + export const initComponents = () => { const telemetry = new PostHogTelemetry(); const knex = bootstrapKnex({ From a92cb1dc9bb4c43a16ce42e34430e7888bdbbf0a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 14:10:33 +0530 Subject: [PATCH 0029/1087] updated database sql --- .../generation/followup_sql_generation.py | 4 +- .../pipelines/generation/sql_correction.py | 4 +- .../pipelines/generation/sql_generation.py | 4 +- .../pipelines/generation/sql_regeneration.py | 4 +- .../src/pipelines/generation/utils/sql.py | 11 ++++++ wren-ai-service/src/utils.py | 38 ++++++++++++++++++- 6 files changed, 56 insertions(+), 9 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 8c02759758..6e295933cf 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -12,13 +12,13 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, + get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -189,7 +189,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_generation_system_prompt(None), - generation_kwargs=SQL_GENERATION_MODEL_KWARGS, + generation_kwargs=get_sql_generation_model_kwargs(llm_provider), ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index d3d3986205..f57a1e7b45 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -13,9 +13,9 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -169,7 +169,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_correction_system_prompt(None), - generation_kwargs=SQL_GENERATION_MODEL_KWARGS, + generation_kwargs=get_sql_generation_model_kwargs(llm_provider), ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 5f9a820506..58c89442c7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -12,12 +12,12 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, + get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -183,7 +183,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_generation_system_prompt(None), - generation_kwargs=SQL_GENERATION_MODEL_KWARGS, + generation_kwargs=get_sql_generation_model_kwargs(llm_provider), ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 26bd957a7e..562e35df5b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -12,12 +12,12 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, + get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -193,7 +193,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_regeneration_system_prompt(None), - generation_kwargs=SQL_GENERATION_MODEL_KWARGS, + generation_kwargs=get_sql_generation_model_kwargs(llm_provider), ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 8c567b926c..2c8caded0f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -11,6 +11,7 @@ Engine, clean_generation_result, ) +from src.core.provider import LLMProvider from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.web.v1.services.ask import AskHistory @@ -585,6 +586,16 @@ class SqlGenerationResult(BaseModel): } +def get_sql_generation_model_kwargs(llm_provider: LLMProvider) -> dict: + model_kwargs = llm_provider.get_model_kwargs() or {} + response_format = model_kwargs.get("response_format", {}) + + if isinstance(response_format, dict) and response_format.get("type") == "text": + return {} + + return SQL_GENERATION_MODEL_KWARGS + + def construct_instructions( instructions: list[dict] | None = None, ): diff --git a/wren-ai-service/src/utils.py b/wren-ai-service/src/utils.py index d368080c3c..e6820b5efc 100644 --- a/wren-ai-service/src/utils.py +++ b/wren-ai-service/src/utils.py @@ -217,4 +217,40 @@ def extract_braces_content(resp: str) -> str: Returns the JSON string including braces, or the original string if no match is found. """ match = re.search(r"```json\s*(\{.*?\})\s*```", resp, re.DOTALL) - return match.group(1) if match else resp + if match: + return match.group(1) + + start = resp.find("{") + if start == -1: + return resp + + depth = 0 + in_string = False + escaped = False + + for index in range(start, len(resp)): + char = resp[index] + + if escaped: + escaped = False + continue + + if char == "\\": + escaped = True + continue + + if char == '"': + in_string = not in_string + continue + + if in_string: + continue + + if char == "{": + depth += 1 + elif char == "}": + depth -= 1 + if depth == 0: + return resp[start : index + 1] + + return resp From 85e12e2dc104474985b90e416dabe4208e53df45 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 14:40:36 +0530 Subject: [PATCH 0030/1087] update database sql --- .../src/pipelines/generation/utils/sql.py | 91 +++++- wren-ai-service/src/web/v1/services/ask.py | 285 +++++++++++------- .../apollo/server/services/askingService.ts | 8 +- wren-ui/src/common.ts | 5 + 4 files changed, 283 insertions(+), 106 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2c8caded0f..71f8d80ba8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, List import aiohttp @@ -22,6 +23,77 @@ def normalize_data_source(data_source: str | None) -> str: return (data_source or "").strip().upper() +def _rewrite_temporal_bucket_functions(sql: str) -> str: + expression_pattern = r'((?:"[^"]+"(?:\."[^"]+")?)|(?:[A-Za-z_][A-Za-z0-9_\.]*))' + replacements = [ + ( + re.compile( + rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f'DATEADD(year, DATEDIFF(year, 0, {m.group(1)}), 0)', + ), + ( + re.compile( + rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f'DATEADD(month, DATEDIFF(month, 0, {m.group(1)}), 0)', + ), + ( + re.compile( + rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f'DATEADD(day, DATEDIFF(day, 0, {m.group(1)}), 0)', + ), + ( + re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f'DATEADD(year, DATEDIFF(year, 0, {m.group(1)}), 0)', + ), + ( + re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f'DATEADD(month, DATEDIFF(month, 0, {m.group(1)}), 0)', + ), + ( + re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f'DATEADD(day, DATEDIFF(day, 0, {m.group(1)}), 0)', + ), + ( + re.compile( + rf"DATETRUNC\(\s*month\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f'DATEADD(month, DATEDIFF(month, 0, {m.group(1)}), 0)', + ), + ( + re.compile( + rf"DATETRUNC\(\s*year\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f'DATEADD(year, DATEDIFF(year, 0, {m.group(1)}), 0)', + ), + ] + + rewritten = sql + for pattern, replacement in replacements: + rewritten = pattern.sub(replacement, rewritten) + + return rewritten + + +def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: + normalized = sql + + if normalize_data_source(data_source) == "MSSQL": + normalized = re.sub( + r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE + ) + normalized = _rewrite_temporal_bucket_functions(normalized) + + return re.sub(r"\s+", " ", normalized).strip() + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -49,6 +121,10 @@ async def run( "sql" ] + cleaned_generation_result = normalize_generation_result_sql( + cleaned_generation_result, data_source=data_source + ) + ( valid_generation_result, invalid_generation_result, @@ -84,6 +160,9 @@ async def _classify_generation_result( ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} + generation_result = normalize_generation_result_sql( + generation_result, data_source=data_source + ) use_dry_run = not allow_data_preview async with aiohttp.ClientSession() as session: @@ -125,8 +204,12 @@ async def _classify_generation_result( } else: error_message = addition.get("error_message", "") + normalized_error_sql = normalize_generation_result_sql( + addition.get("error_sql", generation_result), + data_source=data_source, + ) invalid_generation_result = { - "sql": addition.get("error_sql", generation_result), + "sql": normalized_error_sql, "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") @@ -155,8 +238,12 @@ async def _classify_generation_result( if error_message == "" else "PREVIEW_FAILED" ) + normalized_error_sql = normalize_generation_result_sql( + addition.get("error_sql", generation_result), + data_source=data_source, + ) invalid_generation_result = { - "sql": addition.get("error_sql", generation_result), + "sql": normalized_error_sql, "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 45c31cf8be..f34b31a72d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,5 +1,6 @@ import asyncio import logging +import re from typing import Dict, List, Literal, Optional from cachetools import TTLCache @@ -105,6 +106,7 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, + pipeline_timeout_seconds: int = 90, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -119,6 +121,7 @@ def __init__( self._allow_sql_diagnosis = allow_sql_diagnosis self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval self._enable_column_pruning = enable_column_pruning + self._pipeline_timeout_seconds = pipeline_timeout_seconds self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries @@ -130,6 +133,34 @@ def _is_stopped(self, query_id: str, container: dict): return False + def _is_greeting_query(self, query: str) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + greeting_patterns = { + "hi", + "hello", + "hey", + "hii", + "hola", + "good morning", + "good afternoon", + "good evening", + "how are you", + "thanks", + "thank you", + } + return normalized in greeting_patterns + + async def _run_with_timeout(self, label: str, coroutine): + try: + return await asyncio.wait_for( + coroutine, + timeout=self._pipeline_timeout_seconds, + ) + except TimeoutError as exc: + raise TimeoutError( + f"{label} timed out after {self._pipeline_timeout_seconds} seconds" + ) from exc + @observe(name="Ask Question") @trace_metadata async def ask( @@ -189,9 +220,36 @@ async def ask( is_followup=True if histories else False, ) - historical_question = await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, + if self._is_greeting_query(user_query): + asyncio.create_task( + self._pipelines["user_guide_assistance"].run( + query=( + f'The user said "{user_query}". ' + "Reply with a short greeting and ask how you can help " + "with their PCB database or Wren AI." + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) + ) + + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + trace_id=trace_id, + is_followup=True if histories else False, + general_type="USER_GUIDE", + ) + results["metadata"]["type"] = "GENERAL" + return results + + historical_question = await self._run_with_timeout( + "Historical question retrieval", + self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ), ) # we only return top 1 result @@ -213,15 +271,18 @@ async def ask( sql_generation_reasoning = "" else: # Run both pipeline operations concurrently - sql_samples_task, instructions_task = await asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", + sql_samples_task, instructions_task = await self._run_with_timeout( + "SQL pair and instruction retrieval", + asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), ), ) @@ -235,13 +296,16 @@ async def ask( if self._allow_intent_classification: intent_classification_result = ( - await self._pipelines["intent_classification"].run( - query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, + await self._run_with_timeout( + "Intent classification", + self._pipelines["intent_classification"].run( + query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, + ), ) ).get("post_process", {}) intent = intent_classification_result.get("intent") @@ -343,11 +407,14 @@ async def ask( is_followup=True if histories else False, ) - retrieval_result = await self._pipelines["db_schema_retrieval"].run( - query=user_query, - histories=histories, - project_id=ask_request.project_id, - enable_column_pruning=enable_column_pruning, + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=enable_column_pruning, + ), ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -392,25 +459,31 @@ async def ask( if histories: sql_generation_reasoning = ( - await self._pipelines["followup_sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, + await self._run_with_timeout( + "Follow-up SQL generation reasoning", + self._pipelines["followup_sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), ) ).get("post_process", {}) else: sql_generation_reasoning = ( - await self._pipelines["sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, + await self._run_with_timeout( + "SQL generation reasoning", + self._pipelines["sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), ) ).get("post_process", {}) @@ -438,19 +511,21 @@ async def ask( ) if allow_sql_functions_retrieval: - sql_functions = await self._pipelines[ - "sql_functions_retrieval" - ].run( - project_id=ask_request.project_id, + sql_functions = await self._run_with_timeout( + "SQL functions retrieval", + self._pipelines["sql_functions_retrieval"].run( + project_id=ask_request.project_id, + ), ) else: sql_functions = [] if allow_sql_knowledge_retrieval: - sql_knowledge = await self._pipelines[ - "sql_knowledge_retrieval" - ].run( - project_id=ask_request.project_id, + sql_knowledge = await self._run_with_timeout( + "SQL knowledge retrieval", + self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_request.project_id, + ), ) has_calculated_field = _retrieval_result.get( @@ -460,41 +535,43 @@ async def ask( has_json_field = _retrieval_result.get("has_json_field", False) if histories: - text_to_sql_generation_results = await self._pipelines[ - "followup_sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, + text_to_sql_generation_results = await self._run_with_timeout( + "Follow-up SQL generation", + self._pipelines["followup_sql_generation"].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), ) else: - text_to_sql_generation_results = await self._pipelines[ - "sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, + text_to_sql_generation_results = await self._run_with_timeout( + "SQL generation", + self._pipelines["sql_generation"].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -533,14 +610,15 @@ async def ask( ) if allow_sql_diagnosis: - sql_diagnosis_results = await self._pipelines[ - "sql_diagnosis" - ].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_request.configurations.language, + sql_diagnosis_results = await self._run_with_timeout( + "SQL diagnosis", + self._pipelines["sql_diagnosis"].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_request.configurations.language, + ), ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" @@ -552,21 +630,22 @@ async def ask( f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" ) - sql_correction_results = await self._pipelines[ - "sql_correction" - ].run( - contexts=table_ddls, - instructions=instructions, - invalid_generation_result={ - "original_sql": original_sql, - "sql": invalid_sql, - "error": correction_error_message, - }, - project_id=ask_request.project_id, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, + sql_correction_results = await self._run_with_timeout( + "SQL correction", + self._pipelines["sql_correction"].run( + contexts=table_ddls, + instructions=instructions, + invalid_generation_result={ + "original_sql": original_sql, + "sql": invalid_sql, + "error": correction_error_message, + }, + project_id=ask_request.project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + ), ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 09d6b7a38e..f1de143943 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -570,7 +570,13 @@ export class AskingService implements IAskingService { return; } - await this.askingTaskTracker.initialize(); + if (typeof this.askingTaskTracker?.initialize === 'function') { + await this.askingTaskTracker.initialize(); + } else { + logger.warn( + 'Asking task tracker does not expose initialize(); skipping tracker restoration', + ); + } // list thread responses from database // filter status not finalized and put them into background tracker diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 35f6262acb..d7074eb8ac 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -69,10 +69,15 @@ const hasInitialize = ( const isReusableComponentGraph = ( graph?: ReusableComponentGraph, ): boolean => { + const nestedAskingTaskTracker = ( + graph?.askingService as { askingTaskTracker?: Initializable } | undefined + )?.askingTaskTracker; + return Boolean( graph && hasInitialize(graph.askingTaskTracker) && hasInitialize(graph.askingService) && + hasInitialize(nestedAskingTaskTracker) && hasInitialize(graph.projectRecommendQuestionBackgroundTracker) && hasInitialize(graph.threadRecommendQuestionBackgroundTracker), ); From 7139d55c01593d2a66f8c259d48d4bd67ff04e44 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 15:03:11 +0530 Subject: [PATCH 0031/1087] update sql --- .../src/pipelines/generation/utils/sql.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 71f8d80ba8..65e320cb1c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -31,47 +31,47 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f'DATEADD(year, DATEDIFF(year, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f'DATEADD(month, DATEDIFF(month, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f'DATEADD(day, DATEDIFF(day, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f'DATEADD(year, DATEDIFF(year, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f'DATEADD(month, DATEDIFF(month, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f'DATEADD(day, DATEDIFF(day, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile( - rf"DATETRUNC\(\s*month\s*,\s*{expression_pattern}\s*\)", + rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f'DATEADD(month, DATEDIFF(month, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( - rf"DATETRUNC\(\s*year\s*,\s*{expression_pattern}\s*\)", + rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f'DATEADD(year, DATEDIFF(year, 0, {m.group(1)}), 0)', + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ] From 6fa00a7ed4efd6854fe79451f3dd096bc5d7d3fa Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 16:07:38 +0530 Subject: [PATCH 0032/1087] update sql db --- wren-ai-service/src/config.py | 2 +- .../src/pipelines/generation/utils/sql.py | 151 +++++++++++++++--- 2 files changed, 133 insertions(+), 20 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index c5acf4ae47..20637032a4 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -41,7 +41,7 @@ class Settings(BaseSettings): allow_sql_generation_reasoning: bool = Field(default=True) allow_sql_functions_retrieval: bool = Field(default=True) allow_sql_diagnosis: bool = Field(default=True) - allow_sql_knowledge_retrieval: bool = Field(default=False) + allow_sql_knowledge_retrieval: bool = Field(default=True) max_histories: int = Field(default=5) max_sql_correction_retries: int = Field(default=3) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 65e320cb1c..27ce218899 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,6 @@ import logging import re +from datetime import datetime, timedelta from typing import Any, Dict, List import aiohttp @@ -23,6 +24,109 @@ def normalize_data_source(data_source: str | None) -> str: return (data_source or "").strip().upper() +def _format_timestamp_literal(value: datetime) -> str: + return value.strftime("'%Y-%m-%d %H:%M:%S'") + + +def _add_months(value: datetime, months: int) -> datetime: + month_index = value.month - 1 + months + year = value.year + month_index // 12 + month = month_index % 12 + 1 + day = min( + value.day, + [ + 31, + 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28, + 31, + 30, + 31, + 30, + 31, + 31, + 30, + 31, + 30, + 31, + ][month - 1], + ) + return value.replace(year=year, month=month, day=day) + + +def _start_of_month(value: datetime) -> datetime: + return value.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + + +def _replace_relative_getdate_calls(sql: str, now: datetime) -> str: + def replace_month_offset(match: re.Match[str]) -> str: + months = int(match.group(1)) + return _format_timestamp_literal(_add_months(now, months)) + + def replace_year_offset(match: re.Match[str]) -> str: + years = int(match.group(1)) + return _format_timestamp_literal(_add_months(now, years * 12)) + + def replace_day_offset(match: re.Match[str]) -> str: + days = int(match.group(1)) + return _format_timestamp_literal(now + timedelta(days=days)) + + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_month_offset, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*year\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_year_offset, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*day\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_day_offset, + sql, + flags=re.IGNORECASE, + ) + + current_month_start = _format_timestamp_literal(_start_of_month(now)) + previous_month_start = _format_timestamp_literal( + _start_of_month(_add_months(now, -1)) + ) + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*,\s*0\s*\)", + current_month_start, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*-\s*1\s*,\s*0\s*\)", + previous_month_start, + sql, + flags=re.IGNORECASE, + ) + return sql + + +def _rewrite_mssql_bucket_functions(sql: str) -> str: + expression_pattern = r'((?:"[^"]+"(?:\."[^"]+")?)|(?:[A-Za-z_][A-Za-z0-9_\.]*))' + + sql = re.sub( + rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", + lambda m: ( + f"(DATEPART('YEAR', {m.group(1)}) * 100 + DATEPART('MONTH', {m.group(1)}))" + ), + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", + lambda m: f"DATEPART('YEAR', {m.group(1)})", + sql, + flags=re.IGNORECASE, + ) + return sql + + def _rewrite_temporal_bucket_functions(sql: str) -> str: expression_pattern = r'((?:"[^"]+"(?:\."[^"]+")?)|(?:[A-Za-z_][A-Za-z0-9_\.]*))' replacements = [ @@ -86,9 +190,18 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = sql if normalize_data_source(data_source) == "MSSQL": + now = datetime.now() normalized = re.sub( r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE ) + normalized = re.sub( + r"CAST\(\s*('(?:[^']|'')*')\s+AS\s+DATETIME(?:2|OFFSET)\s*\)", + r"\1", + normalized, + flags=re.IGNORECASE, + ) + normalized = _replace_relative_getdate_calls(normalized, now) + normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) return re.sub(r"\s+", " ", normalized).strip() @@ -316,20 +429,20 @@ async def _classify_generation_result( _MSSQL_TEXT_TO_SQL_RULES = """ ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. -- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. -- DO NOT use YEAR(...), MONTH(...), DAY(...), or DATEPART(...) in generated SQL unless the exact function is explicitly listed as supported in the SQL FUNCTIONS section. Prefer range predicates or DATEADD/DATEDIFF bucket expressions instead. -- Prefer GETDATE() for the current timestamp and CAST(GETDATE() AS DATE) when you need the current date only. -- For relative time windows, use DATEADD together with GETDATE(). -- For previous calendar month boundaries, prefer: - - start_of_previous_month: DATEADD(month, DATEDIFF(month, 0, GETDATE()) - 1, 0) - - start_of_current_month: DATEADD(month, DATEDIFF(month, 0, GETDATE()), 0) -- For month bucketing, prefer DATEADD(month, DATEDIFF(month, 0, ), 0). If DATETRUNC is available in your server version, you may use DATETRUNC(month, ), but prefer DATEADD/DATEDIFF when uncertain. -- For year bucketing, prefer DATEADD(year, DATEDIFF(year, 0, ), 0) instead of YEAR(...). +- The planner in this environment accepts DATEPART with quoted date-part literals, for example DATEPART('YEAR', "created_at"). +- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. +- DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. +- Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. +- For month bucketing, prefer separate year/month fields: + - DATEPART('YEAR', ) AS "year" + - DATEPART('MONTH', ) AS "month" + Then GROUP BY and ORDER BY the same year/month expressions. +- For year bucketing, prefer DATEPART('YEAR', ). - For filtering a specific year such as 2025, prefer a closed-open range: - - >= CAST('2025-01-01 00:00:00' AS DATETIME2) - - AND < CAST('2026-01-01 00:00:00' AS DATETIME2) -- When a temporal cast is required, prefer DATETIME2. Use DATETIMEOFFSET only when timezone-aware semantics are explicitly required by the question. -- Keep relative date logic simple and native to MSSQL. Never emit INTERVAL-like expressions for MSSQL. + - >= '2025-01-01 00:00:00' + - AND < '2026-01-01 00:00:00' +- When a temporal cast is required, keep literal timestamps as plain ISO strings if the column is already datetime-like. +- Keep MSSQL date logic simple and planner-safe. Never emit DATEADD/DATEDIFF fallback expressions unless the SQL FUNCTIONS section explicitly requires them. """ @@ -532,9 +645,9 @@ async def _classify_generation_result( ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; -otherwise, you will put the relative timeframe in the SQL query. +2. Explicitly state the following information in the reasoning plan: +if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; +if the user uses a relative timeframe and Current Time is provided in the input, you will resolve it into an absolute time frame in the SQL query using exact dates rather than relative date arithmetic. 3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. @@ -608,9 +721,9 @@ def get_metric_instructions( instructions += """ #### MSSQL Metric Notes #### -- When filtering metrics by month or other relative date windows in MSSQL, use DATEADD/DATEDIFF or other MSSQL-native date functions from the SQL FUNCTIONS section. -- Do not use DATE_TRUNC, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries. -- Avoid YEAR(...), MONTH(...), DAY(...), and DATEPART(...) in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. +- Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. +- Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. +- For month trend metrics, prefer DATEPART('YEAR', ) and DATEPART('MONTH', ) as separate grouped columns. """ return instructions From f050ebe09226b1bb760a0292dc0c0ad5885f9355 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 16:56:08 +0530 Subject: [PATCH 0033/1087] update sql database --- wren-ai-service/src/providers/llm/litellm.py | 11 ++- wren-ai-service/src/web/v1/services/ask.py | 74 ++++++++++++++----- .../server/services/askingTaskTracker.ts | 2 + 3 files changed, 66 insertions(+), 21 deletions(-) diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index e94918b8c5..7c1becdc58 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -103,6 +103,11 @@ async def _run( **combined_generation_kwargs, **(generation_kwargs or {}), } + should_stream = ( + streaming_callback is not None + and query_id is not None + and generation_kwargs.pop("stream", True) + ) allowed_openai_params = generation_kwargs.get( "allowed_openai_params", [] @@ -112,7 +117,7 @@ async def _run( completion = await self._router.acompletion( model=self._model, messages=openai_formatted_messages, - stream=streaming_callback is not None, + stream=should_stream, allowed_openai_params=allowed_openai_params, mock_testing_fallbacks=self._enable_fallback_testing, **generation_kwargs, @@ -125,13 +130,13 @@ async def _run( api_version=self._api_version, timeout=self._timeout, messages=openai_formatted_messages, - stream=streaming_callback is not None, + stream=should_stream, allowed_openai_params=allowed_openai_params, **generation_kwargs, ) completions: List[ChatMessage] = [] - if streaming_callback is not None: + if should_stream: num_responses = generation_kwargs.pop("n", 1) if num_responses > 1: raise ValueError( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f34b31a72d..8d73d57b74 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -115,6 +115,9 @@ def __init__( self._ask_results: Dict[str, AskResultResponse] = TTLCache( maxsize=maxsize, ttl=ttl ) + self._general_streaming_results: Dict[str, str] = TTLCache( + maxsize=maxsize, ttl=ttl + ) self._allow_sql_generation_reasoning = allow_sql_generation_reasoning self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_intent_classification = allow_intent_classification @@ -161,6 +164,28 @@ async def _run_with_timeout(self, label: str, coroutine): f"{label} timed out after {self._pipeline_timeout_seconds} seconds" ) from exc + def _build_greeting_response(self, query: str) -> str: + return ( + f"Hi. I can help with questions about your PCB database and Wren AI.\n\n" + f"Try a data question like:\n" + f"- Show repair trends for the last 12 months\n" + f"- Compare average debug hours by product family\n" + f"- Which failure codes occur most often?\n\n" + f"If you want, ask a database question directly instead of `{query}`." + ) + + def _extract_pipeline_reply(self, result: dict, key: str) -> str: + payload = result.get(key) + if isinstance(payload, tuple): + payload = payload[0] + + if isinstance(payload, dict): + replies = payload.get("replies") or [] + if replies and isinstance(replies[0], str): + return replies[0] + + return "" + @observe(name="Ask Question") @trace_metadata async def ask( @@ -221,17 +246,8 @@ async def ask( ) if self._is_greeting_query(user_query): - asyncio.create_task( - self._pipelines["user_guide_assistance"].run( - query=( - f'The user said "{user_query}". ' - "Reply with a short greeting and ask how you can help " - "with their PCB database or Wren AI." - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, - ) + self._general_streaming_results[query_id] = ( + self._build_greeting_response(user_query) ) self._ask_results[query_id] = AskResultResponse( @@ -318,7 +334,8 @@ async def ask( user_query = rephrased_question if intent == "MISLEADING_QUERY": - asyncio.create_task( + general_result = await self._run_with_timeout( + "Misleading assistance", self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, @@ -326,14 +343,18 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "misleading_assistance" ) ) self._ask_results[query_id] = AskResultResponse( status="finished", - type="GENERAL", + type="MISLEADING_QUERY", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, @@ -343,7 +364,8 @@ async def ask( results["metadata"]["type"] = "MISLEADING_QUERY" return results elif intent == "GENERAL": - asyncio.create_task( + general_result = await self._run_with_timeout( + "Data assistance", self._pipelines["data_assistance"].run( query=user_query, histories=histories, @@ -351,8 +373,12 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "data_assistance" ) ) @@ -368,12 +394,17 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results elif intent == "USER_GUIDE": - asyncio.create_task( + general_result = await self._run_with_timeout( + "User guide assistance", self._pipelines["user_guide_assistance"].run( query=user_query, language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "user_guide_assistance" ) ) @@ -757,6 +788,13 @@ async def get_ask_streaming_result( self, query_id: str, ): + if general_response := self._general_streaming_results.get(query_id): + event = SSEEvent( + data=SSEEvent.SSEEventMessage(message=general_response), + ) + yield event.serialize() + return + if self._ask_results.get(query_id): _pipeline_name = "" if self._ask_results.get(query_id).type == "GENERAL": diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index a7a1598b9f..aaf015d53e 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -350,6 +350,8 @@ export class AskingTaskTracker implements IAskingTaskTracker { }, }, ); + } else { + await this.updateTaskInDatabase({ queryId }, task); } this.runningJobs.delete(queryId); return; From 1b725c7f14c7540b4de649222d659b906b92f863 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 17:42:46 +0530 Subject: [PATCH 0034/1087] updated --- .../src/pipelines/generation/utils/sql.py | 72 +++++++++++++++---- 1 file changed, 57 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 27ce218899..5aa2d1486a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -113,14 +113,14 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", lambda m: ( - f"(DATEPART('YEAR', {m.group(1)}) * 100 + DATEPART('MONTH', {m.group(1)}))" + f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" ), sql, flags=re.IGNORECASE, ) sql = re.sub( rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", sql, flags=re.IGNORECASE, ) @@ -135,47 +135,89 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ] @@ -429,15 +471,15 @@ async def _classify_generation_result( _MSSQL_TEXT_TO_SQL_RULES = """ ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. -- The planner in this environment accepts DATEPART with quoted date-part literals, for example DATEPART('YEAR', "created_at"). +- Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - - DATEPART('YEAR', ) AS "year" - - DATEPART('MONTH', ) AS "month" + - DATEPART(YEAR, ) AS "year" + - DATEPART(MONTH, ) AS "month" Then GROUP BY and ORDER BY the same year/month expressions. -- For year bucketing, prefer DATEPART('YEAR', ). +- For year bucketing, prefer DATEPART(YEAR, ). - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' - AND < '2026-01-01 00:00:00' @@ -723,7 +765,7 @@ def get_metric_instructions( #### MSSQL Metric Notes #### - Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. - Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. -- For month trend metrics, prefer DATEPART('YEAR', ) and DATEPART('MONTH', ) as separate grouped columns. +- For month trend metrics, prefer DATEPART(YEAR, ) and DATEPART(MONTH, ) as separate grouped columns. """ return instructions From 02866e87f54fd71fa068d9193fad65a590593d3b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 18:52:13 +0530 Subject: [PATCH 0035/1087] updated llm --- wren-ai-service/src/providers/llm/litellm.py | 27 +++++++++++++++++--- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index 7c1becdc58..584dc7e837 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -70,6 +70,23 @@ def get_generator( **(self._model_kwargs or {}), } + def _normalize_generation_kwargs( + kwargs: Optional[Dict[str, Any]], + ) -> Dict[str, Any]: + normalized = dict(kwargs or {}) + response_format = normalized.get("response_format") + + # Plain text is the default chat-completions behavior. + # Some OpenAI-compatible endpoints reject an explicit + # {"type": "text"} payload or serialize it incorrectly. + if ( + isinstance(response_format, dict) + and response_format.get("type") == "text" + ): + normalized.pop("response_format", None) + + return normalized + @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) async def _run( prompt: str, @@ -99,10 +116,12 @@ async def _run( convert_message_to_openai_format(message) for message in messages ] - generation_kwargs = { - **combined_generation_kwargs, - **(generation_kwargs or {}), - } + generation_kwargs = _normalize_generation_kwargs( + { + **combined_generation_kwargs, + **(generation_kwargs or {}), + } + ) should_stream = ( streaming_callback is not None and query_id is not None From 6aaf37021c5e0f3946f65ea4458282afc5ca9347 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 19:07:55 +0530 Subject: [PATCH 0036/1087] updated llm ask --- wren-ai-service/src/web/v1/services/__init__.py | 10 +--------- wren-ai-service/src/web/v1/services/ask.py | 4 ++++ 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/__init__.py b/wren-ai-service/src/web/v1/services/__init__.py index 5c76e4f3fb..250cffb2dc 100644 --- a/wren-ai-service/src/web/v1/services/__init__.py +++ b/wren-ai-service/src/web/v1/services/__init__.py @@ -56,7 +56,7 @@ def serialize(self): # for POST, PATCH, UPDATE, DELETE requests class BaseRequest(BaseModel): - _query_id: str | None = None + query_id: Optional[str] = Field(default=None, exclude=True) project_id: Optional[str] = None thread_id: Optional[str] = None configurations: Configuration = Field( @@ -65,14 +65,6 @@ class BaseRequest(BaseModel): ) request_from: Literal["ui", "api"] = "ui" - @property - def query_id(self) -> str: - return self._query_id - - @query_id.setter - def query_id(self, query_id: str): - self._query_id = query_id - # Put the services imports here to avoid circular imports and make them accessible directly to the rest of packages from .ask import AskService # noqa: E402 diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8d73d57b74..39129057be 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -205,6 +205,10 @@ async def ask( } query_id = ask_request.query_id + if not query_id: + raise ValueError("query_id is required for ask service execution") + + logger.info(f"Ask pipeline started for query_id: {query_id}") histories = ask_request.histories[: self._max_histories][ ::-1 ] # reverse the order of histories From 54e34f8b82a487dafdbcd54d2125d7e3ad32ace9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 19:58:26 +0530 Subject: [PATCH 0037/1087] ask promt --- wren-ui/src/hooks/useAskPrompt.tsx | 10 ++++++++-- wren-ui/src/pages/home/[id].tsx | 7 ++++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index 0c247dbfa2..77cb216753 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -169,14 +169,20 @@ export default function useAskPrompt(threadId?: number) { const [rerunAskingTask] = useRerunAskingTaskMutation({ onError: (error) => console.error(error), }); - const [fetchAskingTask, askingTaskResult] = useAskingTaskLazyQuery(); + const [fetchAskingTask, askingTaskResult] = useAskingTaskLazyQuery({ + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', + }); const [fetchAskingStreamTask, askingStreamTaskResult] = useAskingStreamTask(); const [createInstantRecommendedQuestions] = useCreateInstantRecommendedQuestionsMutation({ onError: (error) => console.error(error), }); const [fetchInstantRecommendedQuestions, instantRecommendedQuestionsResult] = - useInstantRecommendedQuestionsLazyQuery(); + useInstantRecommendedQuestionsLazyQuery({ + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', + }); const askingTaskPollingRef = useRef | null>( null, ); diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index 2e4800f3d1..f0e6773209 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -123,6 +123,8 @@ export default function HomeThread() { }); const [fetchThreadResponse, threadResponseResult] = useThreadResponseLazyQuery({ + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', onCompleted(next) { const nextResponse = next.threadResponse; updateThreadQuery((prev) => ({ @@ -145,7 +147,10 @@ export default function HomeThread() { const [ fetchThreadRecommendationQuestions, threadRecommendationQuestionsResult, - ] = useGetThreadRecommendationQuestionsLazyQuery(); + ] = useGetThreadRecommendationQuestionsLazyQuery({ + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', + }); const threadResponsePollingRef = useRef | null>( null, ); From 6aba6bbc3818e566a004272326283804ecf5b2eb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 20:40:20 +0530 Subject: [PATCH 0038/1087] ask promt ask --- wren-ai-service/src/web/v1/routers/ask.py | 47 ++++++++++++++++------ wren-ai-service/src/web/v1/services/ask.py | 2 +- 2 files changed, 36 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/web/v1/routers/ask.py b/wren-ai-service/src/web/v1/routers/ask.py index 84bc5dbb37..8e63f711d1 100644 --- a/wren-ai-service/src/web/v1/routers/ask.py +++ b/wren-ai-service/src/web/v1/routers/ask.py @@ -1,7 +1,8 @@ +import asyncio import uuid from dataclasses import asdict -from fastapi import APIRouter, BackgroundTasks, Depends +from fastapi import APIRouter, Depends from fastapi.responses import StreamingResponse from src.globals import ( @@ -25,21 +26,47 @@ @router.post("/asks") async def ask( ask_request: AskRequest, - background_tasks: BackgroundTasks, service_container: ServiceContainer = Depends(get_service_container), service_metadata: ServiceMetadata = Depends(get_service_metadata), ) -> AskResponse: query_id = str(uuid.uuid4()) ask_request.query_id = query_id - service_container.ask_service._ask_results[query_id] = AskResultResponse( + ask_service = service_container.ask_service + ask_service._ask_results[query_id] = AskResultResponse( status="understanding", ) - background_tasks.add_task( - service_container.ask_service.ask, - ask_request, - service_metadata=asdict(service_metadata), + if ask_service._is_greeting_query(ask_request.query): + ask_service._general_streaming_results[query_id] = ( + ask_service._build_greeting_response(ask_request.query) + ) + ask_service._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + ) + return AskResponse(query_id=query_id) + + task = asyncio.create_task( + ask_service.ask( + ask_request, + service_metadata=asdict(service_metadata), + ) ) + + def _handle_task_done(completed_task: asyncio.Task): + try: + completed_task.result() + except Exception: + # ask() already captures and records task failures, but we still + # log unexpected task-level exceptions instead of dropping them. + import logging + + logging.getLogger("wren-ai-service").exception( + "Unhandled exception in ask background task for query_id %s", + query_id, + ) + + task.add_done_callback(_handle_task_done) return AskResponse(query_id=query_id) @@ -47,14 +74,10 @@ async def ask( async def stop_ask( query_id: str, stop_ask_request: StopAskRequest, - background_tasks: BackgroundTasks, service_container: ServiceContainer = Depends(get_service_container), ) -> StopAskResponse: stop_ask_request.query_id = query_id - background_tasks.add_task( - service_container.ask_service.stop_ask, - stop_ask_request, - ) + service_container.ask_service.stop_ask(stop_ask_request) return StopAskResponse(query_id=query_id) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 39129057be..be1f614c7c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -76,7 +76,7 @@ class _AskResultResponse(BaseModel): rephrased_question: Optional[str] = None intent_reasoning: Optional[str] = None sql_generation_reasoning: Optional[str] = None - type: Optional[Literal["GENERAL", "TEXT_TO_SQL"]] = None + type: Optional[Literal["GENERAL", "TEXT_TO_SQL", "MISLEADING_QUERY"]] = None retrieved_tables: Optional[List[str]] = None response: Optional[List[AskResult]] = None invalid_sql: Optional[str] = None From 8c411f196332f88cb0ecf6c1fb2e82713d378931 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 20 May 2026 21:46:13 +0530 Subject: [PATCH 0039/1087] update promt ask --- .../src/pipelines/generation/utils/sql.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5aa2d1486a..29df20b5dc 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -113,14 +113,14 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", lambda m: ( - f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" + f"(DATEPART('YEAR', {m.group(1)}) * 100 + DATEPART('MONTH', {m.group(1)}))" ), sql, flags=re.IGNORECASE, ) sql = re.sub( rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", sql, flags=re.IGNORECASE, ) @@ -135,89 +135,89 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ] @@ -471,15 +471,15 @@ async def _classify_generation_result( _MSSQL_TEXT_TO_SQL_RULES = """ ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. -- Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). +- Prefer native T-SQL date bucket syntax such as DATEPART('YEAR', "created_at") and DATEPART('MONTH', "created_at"). - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - - DATEPART(YEAR, ) AS "year" - - DATEPART(MONTH, ) AS "month" + - DATEPART('YEAR', ) AS "year" + - DATEPART('MONTH', ) AS "month" Then GROUP BY and ORDER BY the same year/month expressions. -- For year bucketing, prefer DATEPART(YEAR, ). +- For year bucketing, prefer DATEPART('YEAR', ). - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' - AND < '2026-01-01 00:00:00' @@ -765,7 +765,7 @@ def get_metric_instructions( #### MSSQL Metric Notes #### - Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. - Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. -- For month trend metrics, prefer DATEPART(YEAR, ) and DATEPART(MONTH, ) as separate grouped columns. +- For month trend metrics, prefer DATEPART('YEAR', ) and DATEPART('MONTH', ) as separate grouped columns. """ return instructions From 53141d1cfd7404a19207bdbfb7771da1fcfc8988 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 21 May 2026 00:11:14 +0530 Subject: [PATCH 0040/1087] update chart --- .../pipelines/generation/chart_generation.py | 2 + .../src/pipelines/generation/utils/chart.py | 244 +++++++++++++++++- wren-ui/src/components/chart/handler.ts | 61 +++-- 3 files changed, 277 insertions(+), 30 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/chart_generation.py b/wren-ai-service/src/pipelines/generation/chart_generation.py index 6daca5ec17..1c22161977 100644 --- a/wren-ai-service/src/pipelines/generation/chart_generation.py +++ b/wren-ai-service/src/pipelines/generation/chart_generation.py @@ -97,12 +97,14 @@ def post_process( vega_schema: Dict[str, Any], remove_data_from_chart_schema: bool, preprocess_data: dict, + query: str, post_processor: ChartGenerationPostProcessor, ) -> dict: return post_processor.run( generate_chart.get("replies"), vega_schema, preprocess_data["sample_data"], + query, remove_data_from_chart_schema, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 5d06b949a9..aee5d66d13 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -1,4 +1,6 @@ import logging +import re +from copy import deepcopy from typing import Any, Dict, Literal, Optional import orjson @@ -11,6 +13,214 @@ logger = logging.getLogger("wren-ai-service") +def _humanize_title(name: str) -> str: + return name.replace("_", " ").strip().title() + + +def _detect_requested_chart_type(query: str | None) -> str: + normalized = (query or "").lower() + checks = [ + ("grouped_bar", ["grouped bar"]), + ("stacked_bar", ["stacked bar"]), + ("multi_line", ["multi line", "multi-line"]), + ("line", ["line chart", "line graph", "line plot"]), + ("bar", ["bar chart", "bar graph", "column chart"]), + ("pie", ["pie chart", "donut chart", "doughnut chart"]), + ("area", ["area chart", "area graph"]), + ] + for chart_type, patterns in checks: + if any(pattern in normalized for pattern in patterns): + return chart_type + if re.search(r"\b(chart|graph|plot|visuali[sz](?:e|ation)?)\b", normalized): + return "bar" + return "" + + +def _match_column_name(field: str, columns: list[str]) -> str: + if field in columns: + return field + + lowered = {column.lower(): column for column in columns} + normalized = lowered.get(field.lower()) + if normalized: + return normalized + + compact = re.sub(r"[\s_]+", "", field.lower()) + for column in columns: + if re.sub(r"[\s_]+", "", column.lower()) == compact: + return column + + return field + + +def _normalize_chart_schema_fields(chart_schema: dict, columns: list[str]) -> dict: + normalized = deepcopy(chart_schema) + encoding = normalized.get("encoding", {}) + + for key in ("x", "y", "color", "xOffset", "theta"): + axis = encoding.get(key) + if isinstance(axis, dict) and axis.get("field"): + axis["field"] = _match_column_name(axis["field"], columns) + + for transform in normalized.get("transform", []) or []: + if isinstance(transform, dict) and isinstance(transform.get("fold"), list): + transform["fold"] = [ + _match_column_name(field, columns) for field in transform["fold"] + ] + + return normalized + + +def _infer_column_types(sample_data: list[dict]) -> dict[str, list[str]]: + if not sample_data: + return {"quantitative": [], "temporal": [], "nominal": []} + + df = pd.DataFrame(sample_data) + quantitative: list[str] = [] + temporal: list[str] = [] + nominal: list[str] = [] + + for column in df.columns: + values = df[column].dropna() + if values.empty: + nominal.append(column) + continue + + column_name = str(column).lower() + numeric_values = pd.to_numeric( + values.astype(str).str.replace(",", "", regex=False), + errors="coerce", + ) + temporal_values = pd.to_datetime(values, errors="coerce") + is_temporal_name = bool( + re.search(r"(date|time|month|year|day|created|updated)", column_name) + ) + + if numeric_values.notna().all() and not is_temporal_name: + quantitative.append(column) + elif temporal_values.notna().all() or is_temporal_name: + temporal.append(column) + else: + nominal.append(column) + + return { + "quantitative": quantitative, + "temporal": temporal, + "nominal": nominal, + } + + +def _build_fallback_chart_schema( + query: str | None, + chart_type: str, + sample_data: list[dict], +) -> dict: + if not sample_data: + return {} + + columns = list(sample_data[0].keys()) + inferred = _infer_column_types(sample_data) + quantitative = inferred["quantitative"] + temporal = inferred["temporal"] + nominal = inferred["nominal"] + + title = _humanize_title(query or "Chart") + + def axis(field: str, field_type: str) -> dict: + base = {"field": field, "type": field_type, "title": _humanize_title(field)} + if field_type == "temporal": + base["timeUnit"] = "yearmonth" + return base + + if chart_type == "pie": + color_field = nominal[0] if nominal else columns[0] + theta_field = quantitative[0] if quantitative else ( + columns[1] if len(columns) > 1 else columns[0] + ) + return { + "title": title, + "mark": {"type": "arc"}, + "encoding": { + "theta": axis(theta_field, "quantitative"), + "color": axis(color_field, "nominal"), + }, + } + + if chart_type in {"line", "area", "multi_line"}: + y_field = quantitative[0] if quantitative else columns[-1] + if {"year", "month"}.issubset({c.lower() for c in columns}): + month_field = next(c for c in columns if c.lower() == "month") + encoding = { + "x": axis(month_field, "ordinal"), + "y": axis(y_field, "quantitative"), + } + years = [c for c in columns if c.lower() == "year"] + if years: + encoding["color"] = axis(years[0], "nominal") + return { + "title": title, + "mark": {"type": "area" if chart_type == "area" else "line"}, + "encoding": encoding, + } + + x_field = temporal[0] if temporal else (nominal[0] if nominal else columns[0]) + x_type = "temporal" if x_field in temporal else "ordinal" + return { + "title": title, + "mark": {"type": "area" if chart_type == "area" else "line"}, + "encoding": { + "x": axis(x_field, x_type), + "y": axis(y_field, "quantitative"), + }, + } + + x_field = nominal[0] if nominal else (temporal[0] if temporal else columns[0]) + y_field = quantitative[0] if quantitative else (columns[1] if len(columns) > 1 else columns[0]) + x_type = "nominal" if x_field in nominal else ("temporal" if x_field in temporal else "ordinal") + encoding = { + "x": axis(x_field, x_type), + "y": axis(y_field, "quantitative"), + } + if chart_type == "grouped_bar" and len(nominal) > 1: + encoding["xOffset"] = axis(nominal[1], "nominal") + encoding["color"] = axis(nominal[1], "nominal") + elif nominal: + encoding["color"] = axis(nominal[0], "nominal") + + mark = {"type": "bar"} + if chart_type == "stacked_bar": + encoding["y"]["stack"] = "zero" + + return { + "title": title, + "mark": mark, + "encoding": encoding, + } + + +def _is_schema_compatible_with_sample_data( + chart_schema: dict, + sample_data: list[dict], +) -> bool: + if not chart_schema or not sample_data: + return False + + columns = set(sample_data[0].keys()) + encoding = chart_schema.get("encoding", {}) + for key in ("x", "y", "color", "xOffset", "theta"): + axis = encoding.get(key) + if isinstance(axis, dict) and axis.get("field") and axis["field"] not in columns: + return False + + for transform in chart_schema.get("transform", []) or []: + if isinstance(transform, dict): + for field in transform.get("fold", []) or []: + if field not in columns: + return False + + return True + + chart_generation_instructions = """ ### INSTRUCTIONS ### @@ -288,17 +498,28 @@ def run( replies: str, vega_schema: Dict[str, Any], sample_data: list[dict], + query: Optional[str] = None, remove_data_from_chart_schema: Optional[bool] = True, ): try: generation_result = orjson.loads(replies[0]) reasoning = generation_result.get("reasoning", "") - chart_type = generation_result.get("chart_type", "") + requested_chart_type = _detect_requested_chart_type(query) + chart_type = requested_chart_type or generation_result.get("chart_type", "") if chart_schema := generation_result.get("chart_schema", {}): # sometimes the chart_schema is still in string format if isinstance(chart_schema, str): chart_schema = orjson.loads(chart_schema) + chart_schema = _normalize_chart_schema_fields( + chart_schema, list(sample_data[0].keys()) if sample_data else [] + ) + + if not _is_schema_compatible_with_sample_data(chart_schema, sample_data): + chart_schema = _build_fallback_chart_schema( + query, chart_type or "bar", sample_data + ) + chart_schema[ "$schema" ] = "https://vega.github.io/schema/vega-lite/v5.json" @@ -319,29 +540,40 @@ def run( return { "results": { - "chart_schema": {}, + "chart_schema": _build_fallback_chart_schema( + query, chart_type or "bar", sample_data + ), "reasoning": reasoning, "chart_type": chart_type, } } except ValidationError as e: logger.exception(f"Vega-lite schema is not valid: {e}") + fallback_schema = _build_fallback_chart_schema( + query, + _detect_requested_chart_type(query) or "", + sample_data, + ) return { "results": { - "chart_schema": {}, + "chart_schema": fallback_schema, "reasoning": "", - "chart_type": "", + "chart_type": _detect_requested_chart_type(query) or "", } } except Exception as e: logger.exception(f"JSON deserialization failed: {e}") + fallback_chart_type = _detect_requested_chart_type(query) or "" + fallback_schema = _build_fallback_chart_schema( + query, fallback_chart_type, sample_data + ) return { "results": { - "chart_schema": {}, + "chart_schema": fallback_schema, "reasoning": "", - "chart_type": "", + "chart_type": fallback_chart_type, } } diff --git a/wren-ui/src/components/chart/handler.ts b/wren-ui/src/components/chart/handler.ts index bb150b9878..acc27e6886 100644 --- a/wren-ui/src/components/chart/handler.ts +++ b/wren-ui/src/components/chart/handler.ts @@ -415,30 +415,31 @@ export default class ChartSpecHandler { return encoding[axis]?.title || undefined; } - private transformDataValues( - data: DataSpec, - encoding: { - x?: { type?: string; field?: string }; - y?: { type?: string; field?: string }; - }, - ) { - // If axis x is temporal - if (encoding?.x?.type === 'temporal') { - const transformedValues = data.values.map((val) => ({ - ...val, - [encoding.x.field]: this.transformTemporalValue(val[encoding.x.field]), - })); - return { ...data, values: transformedValues }; - } - // If axis y is temporal - if (encoding?.y?.type === 'temporal') { - const transformedValues = data.values.map((val) => ({ - ...val, - [encoding.y.field]: this.transformTemporalValue(val[encoding.y.field]), - })); - return { ...data, values: transformedValues }; - } - return data; + private transformDataValues(data: DataSpec, encoding?: EncodingSpec) { + const encodingKeys = ['x', 'y', 'theta', 'color', 'xOffset']; + const transformedValues = data.values.map((val) => { + const next = { ...val }; + + encodingKeys.forEach((key) => { + const axis = encoding?.[key] as + | { type?: string; field?: string } + | undefined; + if (!axis || typeof axis.field !== 'string') return; + + if (axis.type === 'temporal') { + next[axis.field] = this.transformTemporalValue(val[axis.field]); + return; + } + + if (axis.type === 'quantitative') { + next[axis.field] = this.transformQuantitativeValue(val[axis.field]); + } + }); + + return next; + }); + + return { ...data, values: transformedValues }; } private transformTemporalValue(value: string | any) { @@ -453,6 +454,18 @@ export default class ChartSpecHandler { } return strValue; } + + private transformQuantitativeValue(value: any) { + if (value === null || value === undefined || value === '') { + return value; + } + if (typeof value === 'number') { + return value; + } + + const numericValue = Number(String(value).replace(/,/g, '')); + return Number.isFinite(numericValue) ? numericValue : value; + } } export const convertToChartType = ( From 00e127fdb073cc2e187d2f43dbc26ecbe79045fa Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 21 May 2026 21:59:04 +0530 Subject: [PATCH 0041/1087] update response --- .../src/pipelines/generation/utils/sql.py | 46 +++++++- wren-ai-service/src/web/v1/services/ask.py | 104 +++++++++++------- wren-ui/src/hooks/useAskPrompt.tsx | 37 ++++++- wren-ui/src/pages/home/[id].tsx | 43 +++++++- 4 files changed, 175 insertions(+), 55 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 29df20b5dc..306140a1fc 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -21,7 +21,12 @@ def normalize_data_source(data_source: str | None) -> str: - return (data_source or "").strip().upper() + normalized = (data_source or "").strip().upper().replace("-", "_").replace( + " ", "_" + ) + if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: + return "MSSQL" + return normalized def _format_timestamp_literal(value: datetime) -> str: @@ -108,7 +113,7 @@ def replace_day_offset(match: re.Match[str]) -> str: def _rewrite_mssql_bucket_functions(sql: str) -> str: - expression_pattern = r'((?:"[^"]+"(?:\."[^"]+")?)|(?:[A-Za-z_][A-Za-z0-9_\.]*))' + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", @@ -128,7 +133,7 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: def _rewrite_temporal_bucket_functions(sql: str) -> str: - expression_pattern = r'((?:"[^"]+"(?:\."[^"]+")?)|(?:[A-Za-z_][A-Za-z0-9_\.]*))' + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" replacements = [ ( re.compile( @@ -163,6 +168,20 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), lambda m: f"DATEPART('DAY', {m.group(1)})", ), + ( + re.compile( + rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART('MONTH', {m.group(1)})", + ), + ( + re.compile( + rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART('YEAR', {m.group(1)})", + ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", @@ -198,6 +217,27 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: ), lambda m: f"DATEPART('DAY', {m.group(1)})", ), + ( + re.compile( + rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART('YEAR', {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART('MONTH', {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART('DAY', {m.group(1)})", + ), ( re.compile( rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index be1f614c7c..cb96d5dab6 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -14,6 +14,10 @@ logger = logging.getLogger("wren-ai-service") +async def _return_value(value): + return value + + class AskHistory(BaseModel): sql: str question: str @@ -493,34 +497,52 @@ async def ask( ) if histories: - sql_generation_reasoning = ( - await self._run_with_timeout( - "Follow-up SQL generation reasoning", - self._pipelines["followup_sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), + try: + sql_generation_reasoning = ( + await self._run_with_timeout( + "Follow-up SQL generation reasoning", + self._pipelines[ + "followup_sql_generation_reasoning" + ].run( + query=user_query, + contexts=table_ddls, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), + ) + ).get("post_process", {}) + except Exception as reasoning_error: + logger.warning( + "Follow-up SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", + query_id, + reasoning_error, ) - ).get("post_process", {}) + sql_generation_reasoning = "" else: - sql_generation_reasoning = ( - await self._run_with_timeout( - "SQL generation reasoning", - self._pipelines["sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), + try: + sql_generation_reasoning = ( + await self._run_with_timeout( + "SQL generation reasoning", + self._pipelines["sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), + ) + ).get("post_process", {}) + except Exception as reasoning_error: + logger.warning( + "SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", + query_id, + reasoning_error, ) - ).get("post_process", {}) + sql_generation_reasoning = "" self._ask_results[query_id] = AskResultResponse( status="planning", @@ -545,23 +567,25 @@ async def ask( is_followup=True if histories else False, ) - if allow_sql_functions_retrieval: - sql_functions = await self._run_with_timeout( - "SQL functions retrieval", - self._pipelines["sql_functions_retrieval"].run( - project_id=ask_request.project_id, + sql_functions, sql_knowledge = await self._run_with_timeout( + "SQL helper retrieval", + asyncio.gather( + ( + self._pipelines["sql_functions_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_functions_retrieval + else _return_value([]) ), - ) - else: - sql_functions = [] - - if allow_sql_knowledge_retrieval: - sql_knowledge = await self._run_with_timeout( - "SQL knowledge retrieval", - self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_request.project_id, + ( + self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_knowledge_retrieval + else _return_value(None) ), - ) + ), + ) has_calculated_field = _retrieval_result.get( "has_calculated_field", False diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index 77cb216753..f80fed59cb 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -30,6 +30,9 @@ export interface AskPromptData { recommendedQuestions?: RecommendedQuestionsTask; } +const ASKING_TASK_POLL_INTERVAL_MS = 2000; +const RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS = 2000; + export const getIsFinished = (status: AskingTaskStatus) => [ AskingTaskStatus.FINISHED, @@ -183,12 +186,14 @@ export default function useAskPrompt(threadId?: number) { fetchPolicy: 'network-only', nextFetchPolicy: 'network-only', }); - const askingTaskPollingRef = useRef | null>( + const askingTaskPollingRef = useRef | null>( null, ); - const recommendedPollingRef = useRef | null>( + const askingTaskPollingSessionRef = useRef(0); + const recommendedPollingRef = useRef | null>( null, ); + const recommendedPollingSessionRef = useRef(0); const askingTask = useMemo( () => askingTaskResult.data?.askingTask || null, @@ -206,15 +211,17 @@ export default function useAskPrompt(threadId?: number) { const loading = askingStreamTaskResult.loading; const stopAskingTaskPolling = useCallback(() => { + askingTaskPollingSessionRef.current += 1; if (askingTaskPollingRef.current) { - clearInterval(askingTaskPollingRef.current); + clearTimeout(askingTaskPollingRef.current); askingTaskPollingRef.current = null; } }, []); const stopRecommendedPolling = useCallback(() => { + recommendedPollingSessionRef.current += 1; if (recommendedPollingRef.current) { - clearInterval(recommendedPollingRef.current); + clearTimeout(recommendedPollingRef.current); recommendedPollingRef.current = null; } }, []); @@ -224,19 +231,28 @@ export default function useAskPrompt(threadId?: number) { if (!taskId) return; stopAskingTaskPolling(); + const pollingSessionId = askingTaskPollingSessionRef.current; const run = async () => { + if (askingTaskPollingSessionRef.current !== pollingSessionId) return; + try { await fetchAskingTask({ variables: { taskId }, }); } catch (error) { console.error(error); + } finally { + if (askingTaskPollingSessionRef.current === pollingSessionId) { + askingTaskPollingRef.current = setTimeout( + run, + ASKING_TASK_POLL_INTERVAL_MS, + ); + } } }; await run(); - askingTaskPollingRef.current = setInterval(run, 1000); }, [fetchAskingTask, stopAskingTaskPolling], ); @@ -246,19 +262,28 @@ export default function useAskPrompt(threadId?: number) { if (!taskId) return; stopRecommendedPolling(); + const pollingSessionId = recommendedPollingSessionRef.current; const run = async () => { + if (recommendedPollingSessionRef.current !== pollingSessionId) return; + try { await fetchInstantRecommendedQuestions({ variables: { taskId }, }); } catch (error) { console.error(error); + } finally { + if (recommendedPollingSessionRef.current === pollingSessionId) { + recommendedPollingRef.current = setTimeout( + run, + RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS, + ); + } } }; await run(); - recommendedPollingRef.current = setInterval(run, 1000); }, [fetchInstantRecommendedQuestions, stopRecommendedPolling], ); diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index f0e6773209..c09ca7c842 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -70,6 +70,9 @@ const getThreadResponseIsFinished = (threadResponse: ThreadResponse) => { return isAnswerFinished !== false && isChartFinished !== false; }; +const THREAD_RESPONSE_POLL_INTERVAL_MS = 2000; +const THREAD_RECOMMENDATION_POLL_INTERVAL_MS = 2000; + export default function HomeThread() { const $prompt = useRef>(null); const router = useRouter(); @@ -151,12 +154,14 @@ export default function HomeThread() { fetchPolicy: 'network-only', nextFetchPolicy: 'network-only', }); - const threadResponsePollingRef = useRef | null>( + const threadResponsePollingRef = useRef | null>( null, ); + const threadResponsePollingSessionRef = useRef(0); const threadRecommendationPollingRef = useRef< - ReturnType | null + ReturnType | null >(null); + const threadRecommendationPollingSessionRef = useRef(0); const [generateThreadResponseAnswer] = useGenerateThreadResponseAnswerMutation({ @@ -192,8 +197,9 @@ export default function HomeThread() { ); const stopThreadResponsePolling = useCallback(() => { + threadResponsePollingSessionRef.current += 1; if (threadResponsePollingRef.current) { - clearInterval(threadResponsePollingRef.current); + clearTimeout(threadResponsePollingRef.current); threadResponsePollingRef.current = null; } }, []); @@ -203,26 +209,36 @@ export default function HomeThread() { if (!responseId) return; stopThreadResponsePolling(); + const pollingSessionId = threadResponsePollingSessionRef.current; const run = async () => { + if (threadResponsePollingSessionRef.current !== pollingSessionId) return; + try { await fetchThreadResponse({ variables: { responseId }, }); } catch (error) { console.error(error); + } finally { + if (threadResponsePollingSessionRef.current === pollingSessionId) { + threadResponsePollingRef.current = setTimeout( + run, + THREAD_RESPONSE_POLL_INTERVAL_MS, + ); + } } }; await run(); - threadResponsePollingRef.current = setInterval(run, 1000); }, [fetchThreadResponse, stopThreadResponsePolling], ); const stopThreadRecommendationPolling = useCallback(() => { + threadRecommendationPollingSessionRef.current += 1; if (threadRecommendationPollingRef.current) { - clearInterval(threadRecommendationPollingRef.current); + clearTimeout(threadRecommendationPollingRef.current); threadRecommendationPollingRef.current = null; } }, []); @@ -232,19 +248,34 @@ export default function HomeThread() { if (!nextThreadId) return; stopThreadRecommendationPolling(); + const pollingSessionId = threadRecommendationPollingSessionRef.current; const run = async () => { + if ( + threadRecommendationPollingSessionRef.current !== pollingSessionId + ) { + return; + } + try { await fetchThreadRecommendationQuestions({ variables: { threadId: nextThreadId }, }); } catch (error) { console.error(error); + } finally { + if ( + threadRecommendationPollingSessionRef.current === pollingSessionId + ) { + threadRecommendationPollingRef.current = setTimeout( + run, + THREAD_RECOMMENDATION_POLL_INTERVAL_MS, + ); + } } }; await run(); - threadRecommendationPollingRef.current = setInterval(run, 1000); }, [fetchThreadRecommendationQuestions, stopThreadRecommendationPolling], ); From a3851e7b0b80ff165bcb89be69940032cbc97180 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 21 May 2026 22:54:23 +0530 Subject: [PATCH 0042/1087] update questions --- wren-ui/src/hooks/useAdjustAnswer.tsx | 70 +++++++++++++-- wren-ui/src/hooks/useAskPrompt.tsx | 106 ++++++++++++++++++++-- wren-ui/src/pages/home/[id].tsx | 121 ++++++++++++++++++++++++-- 3 files changed, 280 insertions(+), 17 deletions(-) diff --git a/wren-ui/src/hooks/useAdjustAnswer.tsx b/wren-ui/src/hooks/useAdjustAnswer.tsx index a2d2680be8..bec1fc39a4 100644 --- a/wren-ui/src/hooks/useAdjustAnswer.tsx +++ b/wren-ui/src/hooks/useAdjustAnswer.tsx @@ -15,6 +15,9 @@ import { ThreadResponse, } from '@/apollo/client/graphql/__types__'; +const ADJUSTMENT_POLL_INTERVAL_MS = 2000; +const ADJUSTMENT_POLL_MAX_INTERVAL_MS = 10000; + export const getIsFinished = (status: AskingTaskStatus) => [ AskingTaskStatus.FINISHED, @@ -72,9 +75,14 @@ export default function useAdjustAnswer(threadId?: number) { }); const [fetchThreadResponse, threadResponseResult] = useThreadResponseLazyQuery(); - const threadResponsePollingRef = useRef | null>( + const threadResponsePollingRef = useRef | null>( null, ); + const threadResponsePollingSessionRef = useRef(0); + const threadResponsePollingTargetRef = useRef(null); + const threadResponsePollingRequestRef = useRef | null>(null); + const threadResponsePollingDelayRef = useRef(ADJUSTMENT_POLL_INTERVAL_MS); + const lastAdjustmentTaskFingerprintRef = useRef(null); const loading = adjustThreadResponseResult.loading; @@ -89,30 +97,56 @@ export default function useAdjustAnswer(threadId?: number) { }, [adjustmentTask]); const stopThreadResponsePolling = useCallback(() => { + threadResponsePollingSessionRef.current += 1; + threadResponsePollingTargetRef.current = null; if (threadResponsePollingRef.current) { - clearInterval(threadResponsePollingRef.current); + clearTimeout(threadResponsePollingRef.current); threadResponsePollingRef.current = null; } + threadResponsePollingDelayRef.current = ADJUSTMENT_POLL_INTERVAL_MS; }, []); const startThreadResponsePolling = useCallback( async (responseId?: number) => { if (!responseId) return; + if ( + threadResponsePollingTargetRef.current === responseId && + (threadResponsePollingRequestRef.current || threadResponsePollingRef.current) + ) { + return; + } stopThreadResponsePolling(); + threadResponsePollingTargetRef.current = responseId; + const pollingSessionId = threadResponsePollingSessionRef.current; const run = async () => { + if (threadResponsePollingSessionRef.current !== pollingSessionId) return; + if (threadResponsePollingRequestRef.current) { + await threadResponsePollingRequestRef.current; + if (threadResponsePollingSessionRef.current !== pollingSessionId) return; + } + try { - await fetchThreadResponse({ + const request = fetchThreadResponse({ variables: { responseId }, - }); + }).then(() => undefined); + threadResponsePollingRequestRef.current = request; + await request; } catch (error) { console.error(error); + } finally { + threadResponsePollingRequestRef.current = null; + if (threadResponsePollingSessionRef.current === pollingSessionId) { + threadResponsePollingRef.current = setTimeout( + run, + threadResponsePollingDelayRef.current, + ); + } } }; await run(); - threadResponsePollingRef.current = setInterval(run, 1000); }, [fetchThreadResponse, stopThreadResponsePolling], ); @@ -122,6 +156,32 @@ export default function useAdjustAnswer(threadId?: number) { if (isFinished) stopThreadResponsePolling(); }, [adjustmentTask?.status]); + useEffect(() => { + const fingerprint = JSON.stringify({ + queryId: adjustmentTask?.queryId || null, + status: adjustmentTask?.status || null, + sql: adjustmentTask?.sql || null, + errorCode: adjustmentTask?.error?.code || null, + invalidSql: adjustmentTask?.invalidSql || null, + }); + + if (lastAdjustmentTaskFingerprintRef.current === fingerprint) { + threadResponsePollingDelayRef.current = Math.min( + threadResponsePollingDelayRef.current * 2, + ADJUSTMENT_POLL_MAX_INTERVAL_MS, + ); + } else { + threadResponsePollingDelayRef.current = ADJUSTMENT_POLL_INTERVAL_MS; + lastAdjustmentTaskFingerprintRef.current = fingerprint; + } + }, [ + adjustmentTask?.queryId, + adjustmentTask?.status, + adjustmentTask?.sql, + adjustmentTask?.error?.code, + adjustmentTask?.invalidSql, + ]); + const onAdjustReasoningSteps = async ( responseId: number, input: { tables: string[]; sqlGenerationReasoning: string }, diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index f80fed59cb..b028604794 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -32,6 +32,8 @@ export interface AskPromptData { const ASKING_TASK_POLL_INTERVAL_MS = 2000; const RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS = 2000; +const ASKING_TASK_POLL_MAX_INTERVAL_MS = 10000; +const RECOMMENDED_QUESTIONS_POLL_MAX_INTERVAL_MS = 10000; export const getIsFinished = (status: AskingTaskStatus) => [ @@ -190,10 +192,20 @@ export default function useAskPrompt(threadId?: number) { null, ); const askingTaskPollingSessionRef = useRef(0); + const askingTaskPollingTargetRef = useRef(null); + const askingTaskPollingRequestRef = useRef | null>(null); + const askingTaskPollingDelayRef = useRef(ASKING_TASK_POLL_INTERVAL_MS); + const lastAskingTaskFingerprintRef = useRef(null); const recommendedPollingRef = useRef | null>( null, ); const recommendedPollingSessionRef = useRef(0); + const recommendedPollingTargetRef = useRef(null); + const recommendedPollingRequestRef = useRef | null>(null); + const recommendedPollingDelayRef = useRef( + RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS, + ); + const lastRecommendedFingerprintRef = useRef(null); const askingTask = useMemo( () => askingTaskResult.data?.askingTask || null, @@ -212,41 +224,59 @@ export default function useAskPrompt(threadId?: number) { const stopAskingTaskPolling = useCallback(() => { askingTaskPollingSessionRef.current += 1; + askingTaskPollingTargetRef.current = null; if (askingTaskPollingRef.current) { clearTimeout(askingTaskPollingRef.current); askingTaskPollingRef.current = null; } + askingTaskPollingDelayRef.current = ASKING_TASK_POLL_INTERVAL_MS; }, []); const stopRecommendedPolling = useCallback(() => { recommendedPollingSessionRef.current += 1; + recommendedPollingTargetRef.current = null; if (recommendedPollingRef.current) { clearTimeout(recommendedPollingRef.current); recommendedPollingRef.current = null; } + recommendedPollingDelayRef.current = RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS; }, []); const startAskingTaskPolling = useCallback( async (taskId?: string) => { if (!taskId) return; + if ( + askingTaskPollingTargetRef.current === taskId && + (askingTaskPollingRequestRef.current || askingTaskPollingRef.current) + ) { + return; + } stopAskingTaskPolling(); + askingTaskPollingTargetRef.current = taskId; const pollingSessionId = askingTaskPollingSessionRef.current; const run = async () => { if (askingTaskPollingSessionRef.current !== pollingSessionId) return; + if (askingTaskPollingRequestRef.current) { + await askingTaskPollingRequestRef.current; + if (askingTaskPollingSessionRef.current !== pollingSessionId) return; + } try { - await fetchAskingTask({ + const request = fetchAskingTask({ variables: { taskId }, - }); + }).then(() => undefined); + askingTaskPollingRequestRef.current = request; + await request; } catch (error) { console.error(error); } finally { + askingTaskPollingRequestRef.current = null; if (askingTaskPollingSessionRef.current === pollingSessionId) { askingTaskPollingRef.current = setTimeout( run, - ASKING_TASK_POLL_INTERVAL_MS, + askingTaskPollingDelayRef.current, ); } } @@ -260,24 +290,38 @@ export default function useAskPrompt(threadId?: number) { const startRecommendedPolling = useCallback( async (taskId?: string) => { if (!taskId) return; + if ( + recommendedPollingTargetRef.current === taskId && + (recommendedPollingRequestRef.current || recommendedPollingRef.current) + ) { + return; + } stopRecommendedPolling(); + recommendedPollingTargetRef.current = taskId; const pollingSessionId = recommendedPollingSessionRef.current; const run = async () => { if (recommendedPollingSessionRef.current !== pollingSessionId) return; + if (recommendedPollingRequestRef.current) { + await recommendedPollingRequestRef.current; + if (recommendedPollingSessionRef.current !== pollingSessionId) return; + } try { - await fetchInstantRecommendedQuestions({ + const request = fetchInstantRecommendedQuestions({ variables: { taskId }, - }); + }).then(() => undefined); + recommendedPollingRequestRef.current = request; + await request; } catch (error) { console.error(error); } finally { + recommendedPollingRequestRef.current = null; if (recommendedPollingSessionRef.current === pollingSessionId) { recommendedPollingRef.current = setTimeout( run, - RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS, + recommendedPollingDelayRef.current, ); } } @@ -341,6 +385,34 @@ export default function useAskPrompt(threadId?: number) { } }, [askingTask?.status, threadId, checkFetchAskingStreamTask]); + useEffect(() => { + const fingerprint = JSON.stringify({ + queryId: askingTask?.queryId || null, + status: askingTask?.status || null, + type: askingTask?.type || null, + candidateCount: askingTask?.candidates?.length || 0, + errorCode: askingTask?.error?.code || null, + traceId: askingTask?.traceId || null, + }); + + if (lastAskingTaskFingerprintRef.current === fingerprint) { + askingTaskPollingDelayRef.current = Math.min( + askingTaskPollingDelayRef.current * 2, + ASKING_TASK_POLL_MAX_INTERVAL_MS, + ); + } else { + askingTaskPollingDelayRef.current = ASKING_TASK_POLL_INTERVAL_MS; + lastAskingTaskFingerprintRef.current = fingerprint; + } + }, [ + askingTask?.queryId, + askingTask?.status, + askingTask?.type, + askingTask?.candidates?.length, + askingTask?.error?.code, + askingTask?.traceId, + ]); + useEffect(() => { // handle instant recommended questions if (isNeedRecommendedQuestions(askingTask)) { @@ -348,6 +420,28 @@ export default function useAskPrompt(threadId?: number) { } }, [askingTask?.type]); + useEffect(() => { + const fingerprint = JSON.stringify({ + status: recommendedQuestions?.status || null, + count: recommendedQuestions?.questions?.length || 0, + errorCode: recommendedQuestions?.error?.code || null, + }); + + if (lastRecommendedFingerprintRef.current === fingerprint) { + recommendedPollingDelayRef.current = Math.min( + recommendedPollingDelayRef.current * 2, + RECOMMENDED_QUESTIONS_POLL_MAX_INTERVAL_MS, + ); + } else { + recommendedPollingDelayRef.current = RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS; + lastRecommendedFingerprintRef.current = fingerprint; + } + }, [ + recommendedQuestions?.status, + recommendedQuestions?.questions?.length, + recommendedQuestions?.error?.code, + ]); + useEffect(() => { if (isRecommendedFinished(recommendedQuestions?.status)) { stopRecommendedPolling(); diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index c09ca7c842..e5b0b05095 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -72,6 +72,8 @@ const getThreadResponseIsFinished = (threadResponse: ThreadResponse) => { const THREAD_RESPONSE_POLL_INTERVAL_MS = 2000; const THREAD_RECOMMENDATION_POLL_INTERVAL_MS = 2000; +const THREAD_RESPONSE_POLL_MAX_INTERVAL_MS = 10000; +const THREAD_RECOMMENDATION_POLL_MAX_INTERVAL_MS = 10000; export default function HomeThread() { const $prompt = useRef>(null); @@ -158,10 +160,22 @@ export default function HomeThread() { null, ); const threadResponsePollingSessionRef = useRef(0); + const threadResponsePollingTargetRef = useRef(null); + const threadResponsePollingRequestRef = useRef | null>(null); + const threadResponsePollingDelayRef = useRef(THREAD_RESPONSE_POLL_INTERVAL_MS); + const lastThreadResponseFingerprintRef = useRef(null); const threadRecommendationPollingRef = useRef< ReturnType | null >(null); const threadRecommendationPollingSessionRef = useRef(0); + const threadRecommendationPollingTargetRef = useRef(null); + const threadRecommendationPollingRequestRef = useRef | null>( + null, + ); + const threadRecommendationPollingDelayRef = useRef( + THREAD_RECOMMENDATION_POLL_INTERVAL_MS, + ); + const lastThreadRecommendationFingerprintRef = useRef(null); const [generateThreadResponseAnswer] = useGenerateThreadResponseAnswerMutation({ @@ -198,33 +212,51 @@ export default function HomeThread() { const stopThreadResponsePolling = useCallback(() => { threadResponsePollingSessionRef.current += 1; + threadResponsePollingTargetRef.current = null; if (threadResponsePollingRef.current) { clearTimeout(threadResponsePollingRef.current); threadResponsePollingRef.current = null; } + threadResponsePollingDelayRef.current = THREAD_RESPONSE_POLL_INTERVAL_MS; }, []); const startThreadResponsePolling = useCallback( async (responseId?: number) => { if (!responseId) return; + if ( + threadResponsePollingTargetRef.current === responseId && + (threadResponsePollingRequestRef.current || threadResponsePollingRef.current) + ) { + return; + } stopThreadResponsePolling(); + threadResponsePollingTargetRef.current = responseId; const pollingSessionId = threadResponsePollingSessionRef.current; const run = async () => { if (threadResponsePollingSessionRef.current !== pollingSessionId) return; + if (threadResponsePollingRequestRef.current) { + await threadResponsePollingRequestRef.current; + if (threadResponsePollingSessionRef.current !== pollingSessionId) { + return; + } + } try { - await fetchThreadResponse({ + const request = fetchThreadResponse({ variables: { responseId }, - }); + }).then(() => undefined); + threadResponsePollingRequestRef.current = request; + await request; } catch (error) { console.error(error); } finally { + threadResponsePollingRequestRef.current = null; if (threadResponsePollingSessionRef.current === pollingSessionId) { threadResponsePollingRef.current = setTimeout( run, - THREAD_RESPONSE_POLL_INTERVAL_MS, + threadResponsePollingDelayRef.current, ); } } @@ -237,17 +269,28 @@ export default function HomeThread() { const stopThreadRecommendationPolling = useCallback(() => { threadRecommendationPollingSessionRef.current += 1; + threadRecommendationPollingTargetRef.current = null; if (threadRecommendationPollingRef.current) { clearTimeout(threadRecommendationPollingRef.current); threadRecommendationPollingRef.current = null; } + threadRecommendationPollingDelayRef.current = + THREAD_RECOMMENDATION_POLL_INTERVAL_MS; }, []); const startThreadRecommendationPolling = useCallback( async (nextThreadId?: number) => { if (!nextThreadId) return; + if ( + threadRecommendationPollingTargetRef.current === nextThreadId && + (threadRecommendationPollingRequestRef.current || + threadRecommendationPollingRef.current) + ) { + return; + } stopThreadRecommendationPolling(); + threadRecommendationPollingTargetRef.current = nextThreadId; const pollingSessionId = threadRecommendationPollingSessionRef.current; const run = async () => { @@ -256,20 +299,31 @@ export default function HomeThread() { ) { return; } + if (threadRecommendationPollingRequestRef.current) { + await threadRecommendationPollingRequestRef.current; + if ( + threadRecommendationPollingSessionRef.current !== pollingSessionId + ) { + return; + } + } try { - await fetchThreadRecommendationQuestions({ + const request = fetchThreadRecommendationQuestions({ variables: { threadId: nextThreadId }, - }); + }).then(() => undefined); + threadRecommendationPollingRequestRef.current = request; + await request; } catch (error) { console.error(error); } finally { + threadRecommendationPollingRequestRef.current = null; if ( threadRecommendationPollingSessionRef.current === pollingSessionId ) { threadRecommendationPollingRef.current = setTimeout( run, - THREAD_RECOMMENDATION_POLL_INTERVAL_MS, + threadRecommendationPollingDelayRef.current, ); } } @@ -384,6 +438,38 @@ export default function HomeThread() { } }, [isPollingResponseFinished]); + useEffect(() => { + const fingerprint = JSON.stringify({ + id: pollingResponse?.id || null, + askingStatus: pollingResponse?.askingTask?.status || null, + askingType: pollingResponse?.askingTask?.type || null, + answerStatus: pollingResponse?.answerDetail?.status || null, + chartStatus: pollingResponse?.chartDetail?.status || null, + breakdownStatus: pollingResponse?.breakdownDetail?.status || null, + adjustmentStatus: pollingResponse?.adjustmentTask?.status || null, + sql: pollingResponse?.sql || null, + }); + + if (lastThreadResponseFingerprintRef.current === fingerprint) { + threadResponsePollingDelayRef.current = Math.min( + threadResponsePollingDelayRef.current * 2, + THREAD_RESPONSE_POLL_MAX_INTERVAL_MS, + ); + } else { + threadResponsePollingDelayRef.current = THREAD_RESPONSE_POLL_INTERVAL_MS; + lastThreadResponseFingerprintRef.current = fingerprint; + } + }, [ + pollingResponse?.id, + pollingResponse?.askingTask?.status, + pollingResponse?.askingTask?.type, + pollingResponse?.answerDetail?.status, + pollingResponse?.chartDetail?.status, + pollingResponse?.breakdownDetail?.status, + pollingResponse?.adjustmentTask?.status, + pollingResponse?.sql, + ]); + const recommendedQuestions = useMemo( () => threadRecommendationQuestionsResult.data @@ -397,6 +483,29 @@ export default function HomeThread() { } }, [recommendedQuestions]); + useEffect(() => { + const fingerprint = JSON.stringify({ + status: recommendedQuestions?.status || null, + count: recommendedQuestions?.questions?.length || 0, + errorCode: recommendedQuestions?.error?.code || null, + }); + + if (lastThreadRecommendationFingerprintRef.current === fingerprint) { + threadRecommendationPollingDelayRef.current = Math.min( + threadRecommendationPollingDelayRef.current * 2, + THREAD_RECOMMENDATION_POLL_MAX_INTERVAL_MS, + ); + } else { + threadRecommendationPollingDelayRef.current = + THREAD_RECOMMENDATION_POLL_INTERVAL_MS; + lastThreadRecommendationFingerprintRef.current = fingerprint; + } + }, [ + recommendedQuestions?.status, + recommendedQuestions?.questions?.length, + recommendedQuestions?.error?.code, + ]); + const onCreateResponse = async (payload: CreateThreadResponseInput) => { try { askPrompt.onStopPolling(); From aa0c2ea917f7d296726146530ad31c96cf45bfd4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 21 May 2026 23:50:50 +0530 Subject: [PATCH 0043/1087] update questions request --- .../src/apollo/server/backgrounds/chart.ts | 6 +- .../textBasedAnswerBackgroundTracker.ts | 153 ++++++++++-------- .../apollo/server/services/askingService.ts | 5 +- .../apollo/server/services/queryService.ts | 53 +++++- wren-ui/src/pages/home/[id].tsx | 15 +- 5 files changed, 152 insertions(+), 80 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index 8034d67c8c..01f3039cfa 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -42,7 +42,7 @@ export class ChartBackgroundTracker { this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadResponseRepository = threadResponseRepository; - this.intervalTime = 1000; + this.intervalTime = 2000; this.start(); } @@ -92,6 +92,7 @@ export class ChartBackgroundTracker { await this.threadResponseRepository.updateOne(threadResponse.id, { chartDetail: updatedChartDetail, }); + threadResponse.chartDetail = updatedChartDetail; // remove the task from tracker if it is finalized if (isFinalized(result.status)) { @@ -164,7 +165,7 @@ export class ChartAdjustmentBackgroundTracker { this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadResponseRepository = threadResponseRepository; - this.intervalTime = 1000; + this.intervalTime = 2000; this.start(); } @@ -215,6 +216,7 @@ export class ChartAdjustmentBackgroundTracker { await this.threadResponseRepository.updateOne(threadResponse.id, { chartDetail: updatedChartDetail, }); + threadResponse.chartDetail = updatedChartDetail; // remove the task from tracker if it is finalized if (isFinalized(result.status)) { diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index 91f14bcab1..79b769e18c 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -17,6 +17,8 @@ import { getLogger } from '@server/utils'; const logger = getLogger('TextBasedAnswerBackgroundTracker'); logger.level = 'debug'; +const ANSWER_PREVIEW_LIMIT = 200; + export class TextBasedAnswerBackgroundTracker { // tasks is a kv pair of task id and thread response private tasks: Record = {}; @@ -46,7 +48,7 @@ export class TextBasedAnswerBackgroundTracker { this.projectService = projectService; this.deployService = deployService; this.queryService = queryService; - this.intervalTime = 1000; + this.intervalTime = 2000; this.start(); } @@ -62,85 +64,100 @@ export class TextBasedAnswerBackgroundTracker { } this.runningJobs.add(threadResponse.id); - // update the status to fetching data - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: { - ...threadResponse.answerDetail, + const answerDetail = threadResponse.answerDetail; + + if ( + !answerDetail.queryId && + answerDetail.status !== ThreadResponseAnswerStatus.FETCHING_DATA + ) { + const fetchingDetail = { + ...answerDetail, status: ThreadResponseAnswerStatus.FETCHING_DATA, - }, - }); - - // get sql data - const project = await this.projectService.getCurrentProject(); - const deployment = await this.deployService.getLastDeployment( - project.id, - ); - const mdl = deployment.manifest; - let data: PreviewDataResponse; - try { - data = (await this.queryService.preview(threadResponse.sql, { - project, - manifest: mdl, - modelingOnly: false, - limit: 500, - })) as PreviewDataResponse; - } catch (error) { - logger.error(`Error when query sql data: ${error}`); + }; await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: { + answerDetail: fetchingDetail, + }); + threadResponse.answerDetail = fetchingDetail; + + const project = await this.projectService.getCurrentProject(); + const deployment = await this.deployService.getLastDeployment( + project.id, + ); + const mdl = deployment.manifest; + let data: PreviewDataResponse; + try { + data = (await this.queryService.preview(threadResponse.sql, { + project, + manifest: mdl, + modelingOnly: false, + limit: ANSWER_PREVIEW_LIMIT, + })) as PreviewDataResponse; + } catch (error) { + logger.error(`Error when query sql data: ${error}`); + const failedDetail = { ...threadResponse.answerDetail, status: ThreadResponseAnswerStatus.FAILED, error: error?.extensions || error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: failedDetail, + }); + threadResponse.answerDetail = failedDetail; + throw error; + } + + const response = await this.wrenAIAdaptor.createTextBasedAnswer({ + query: threadResponse.question, + sql: threadResponse.sql, + sqlData: data, + threadId: threadResponse.threadId.toString(), + configurations: { + language: WrenAILanguage[project.language] || WrenAILanguage.EN, }, }); - throw error; - } - // request AI service - const response = await this.wrenAIAdaptor.createTextBasedAnswer({ - query: threadResponse.question, - sql: threadResponse.sql, - sqlData: data, - threadId: threadResponse.threadId.toString(), - configurations: { - language: WrenAILanguage[project.language] || WrenAILanguage.EN, - }, - }); - - // update the status to preprocessing - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: { + const preprocessingDetail = { ...threadResponse.answerDetail, + queryId: response.queryId, status: ThreadResponseAnswerStatus.PREPROCESSING, - }, - }); - - // polling query id to check the status - let result: TextBasedAnswerResult; - do { - result = await this.wrenAIAdaptor.getTextBasedAnswerResult( - response.queryId, - ); + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: preprocessingDetail, + }); + threadResponse.answerDetail = preprocessingDetail; + this.runningJobs.delete(threadResponse.id); + return; + } + + if ( + answerDetail.queryId && + answerDetail.status === ThreadResponseAnswerStatus.PREPROCESSING + ) { + const result: TextBasedAnswerResult = + await this.wrenAIAdaptor.getTextBasedAnswerResult( + answerDetail.queryId, + ); + if (result.status === TextBasedAnswerStatus.PREPROCESSING) { - await new Promise((resolve) => setTimeout(resolve, 500)); + this.runningJobs.delete(threadResponse.id); + return; } - } while (result.status === TextBasedAnswerStatus.PREPROCESSING); - - // update the status to final - const updatedAnswerDetail = { - queryId: response.queryId, - status: - result.status === TextBasedAnswerStatus.SUCCEEDED - ? ThreadResponseAnswerStatus.STREAMING - : ThreadResponseAnswerStatus.FAILED, - numRowsUsedInLLM: result.numRowsUsedInLLM, - error: result.error, - }; - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: updatedAnswerDetail, - }); - - delete this.tasks[threadResponse.id]; + + const updatedAnswerDetail = { + queryId: answerDetail.queryId, + status: + result.status === TextBasedAnswerStatus.SUCCEEDED + ? ThreadResponseAnswerStatus.STREAMING + : ThreadResponseAnswerStatus.FAILED, + numRowsUsedInLLM: result.numRowsUsedInLLM, + error: result.error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: updatedAnswerDetail, + }); + threadResponse.answerDetail = updatedAnswerDetail; + delete this.tasks[threadResponse.id]; + } // Mark the job as finished this.runningJobs.delete(threadResponse.id); diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index f1de143943..f3196d1f9e 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -297,12 +297,12 @@ class BreakdownBackgroundTracker { }: { telemetry: PostHogTelemetry; wrenAIAdaptor: IWrenAIAdaptor; - threadResponseRepository: IThreadResponseRepository; + threadResponseRepository: IThreadResponseRepository; }) { this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadResponseRepository = threadResponseRepository; - this.intervalTime = 1000; + this.intervalTime = 2000; this.start(); } @@ -349,6 +349,7 @@ class BreakdownBackgroundTracker { await this.threadResponseRepository.updateOne(threadResponse.id, { breakdownDetail: updatedBreakdownDetail, }); + threadResponse.breakdownDetail = updatedBreakdownDetail; // remove the task from tracker if it is finalized if (isFinalized(result.status)) { diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 08c78d6132..24d5d6dd50 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -76,6 +76,34 @@ export interface IQueryService { ): Promise; } +const normalizePreviewSqlForIbis = ( + sql: string, + dataSource: DataSourceName, + limit?: number, +): { sql: string; limit?: number } => { + if (dataSource !== DataSourceName.MSSQL) { + return { sql, limit }; + } + + const topMatch = sql.match(/^\s*SELECT\s+(DISTINCT\s+)?TOP\s*\(?\s*(\d+)\s*\)?\s+/i); + if (!topMatch) { + return { sql, limit }; + } + + const distinctClause = topMatch[1] || ''; + const topLimit = Number(topMatch[2]); + const normalizedSql = sql.replace( + /^\s*SELECT\s+(DISTINCT\s+)?TOP\s*\(?\s*\d+\s*\)?\s+/i, + `SELECT ${distinctClause}`, + ); + + return { + sql: normalizedSql, + limit: + limit && limit > 0 ? Math.min(limit, topLimit) : topLimit, + }; +}; + export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; @@ -194,19 +222,23 @@ export class QueryService implements IQueryService { connectionInfo: any, mdl: Manifest, ): Promise { + const normalizedQuery = normalizePreviewSqlForIbis(sql, dataSource).sql; const event = TelemetryEvent.IBIS_DRY_RUN; try { - const res = await this.ibisAdaptor.dryRun(sql, { + const res = await this.ibisAdaptor.dryRun(normalizedQuery, { dataSource, connectionInfo, mdl, }); - this.sendIbisEvent(event, res, { dataSource, sql }); + this.sendIbisEvent(event, res, { dataSource, sql: normalizedQuery }); return { correlationId: res.correlationId, }; } catch (err: any) { - this.sendIbisFailedEvent(event, err, { dataSource, sql }); + this.sendIbisFailedEvent(event, err, { + dataSource, + sql: normalizedQuery, + }); throw err; } } @@ -220,17 +252,21 @@ export class QueryService implements IQueryService { refresh?: boolean, cacheEnabled?: boolean, ): Promise { + const normalizedPreview = normalizePreviewSqlForIbis(sql, dataSource, limit); const event = TelemetryEvent.IBIS_QUERY; try { - const res = await this.ibisAdaptor.query(sql, { + const res = await this.ibisAdaptor.query(normalizedPreview.sql, { dataSource, connectionInfo, mdl, - limit, + limit: normalizedPreview.limit, refresh, cacheEnabled, }); - this.sendIbisEvent(event, res, { dataSource, sql }); + this.sendIbisEvent(event, res, { + dataSource, + sql: normalizedPreview.sql, + }); const data = this.transformDataType(res); return { correlationId: res.correlationId, @@ -241,7 +277,10 @@ export class QueryService implements IQueryService { ...data, }; } catch (err: any) { - this.sendIbisFailedEvent(event, err, { dataSource, sql }); + this.sendIbisFailedEvent(event, err, { + dataSource, + sql: normalizedPreview.sql, + }); throw err; } } diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index e5b0b05095..c3b3cb48ee 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -413,8 +413,21 @@ export default function HomeThread() { // stop all requests when change thread useEffect(() => { if (threadId !== null) { - startThreadRecommendationPolling(threadId); setShowRecommendedQuestions(true); + void (async () => { + try { + const result = await fetchThreadRecommendationQuestions({ + variables: { threadId }, + }); + const status = + result.data?.getThreadRecommendationQuestions?.status || null; + if (status && !isRecommendedFinished(status)) { + await startThreadRecommendationPolling(threadId); + } + } catch (error) { + console.error(error); + } + })(); } return () => { askPrompt.onStopPolling(); From 5093e6514d35f649176a636ada1c7bdebf656759 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 22 May 2026 00:27:50 +0530 Subject: [PATCH 0044/1087] update request --- .../apollo/server/services/askingService.ts | 39 +++++++++++++++++++ .../pages/home/promptThread/AnswerResult.tsx | 37 +++++++++++++++++- 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index f3196d1f9e..f0c3789b63 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -102,6 +102,21 @@ export enum ThreadResponseAnswerStatus { INTERRUPTED = 'INTERRUPTED', } +const isAnswerGenerationInProgress = ( + status?: ThreadResponseAnswerStatus | string | null, +) => + ([ + ThreadResponseAnswerStatus.NOT_STARTED, + ThreadResponseAnswerStatus.FETCHING_DATA, + ThreadResponseAnswerStatus.PREPROCESSING, + ThreadResponseAnswerStatus.STREAMING, + ] as string[]).includes(status || ""); + +const isChartGenerationInProgress = (status?: ChartStatus | string | null) => + ([ChartStatus.FETCHING, ChartStatus.GENERATING] as string[]).includes( + status || "", + ); + // adjustment input export interface AdjustmentReasoningInput { tables: string[]; @@ -847,6 +862,13 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${threadResponseId} not found`); } + if (isAnswerGenerationInProgress(threadResponse.answerDetail?.status)) { + logger.debug( + `Thread response ${threadResponseId} answer generation already in progress, skipping duplicate request`, + ); + return threadResponse; + } + // update with initial status const updatedThreadResponse = await this.threadResponseRepository.updateOne( threadResponse.id, @@ -875,6 +897,13 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${threadResponseId} not found`); } + if (isChartGenerationInProgress(threadResponse.chartDetail?.status)) { + logger.debug( + `Thread response ${threadResponseId} chart generation already in progress, skipping duplicate request`, + ); + return threadResponse; + } + // 1. create a task on AI service to generate the chart const response = await this.wrenAIAdaptor.generateChart({ query: threadResponse.question, @@ -912,6 +941,16 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${threadResponseId} not found`); } + if ( + isChartGenerationInProgress(threadResponse.chartDetail?.status) && + threadResponse.chartDetail?.adjustment + ) { + logger.debug( + `Thread response ${threadResponseId} chart adjustment already in progress, skipping duplicate request`, + ); + return threadResponse; + } + // 1. create a task on AI service to adjust the chart const response = await this.wrenAIAdaptor.adjustChart({ query: threadResponse.question, diff --git a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx index 6a0e628206..bcfe7535cc 100644 --- a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx +++ b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo } from 'react'; +import { useEffect, useMemo, useRef } from 'react'; import { isEmpty, debounce } from 'lodash'; import clsx from 'clsx'; import { Button, Typography, Tabs, Tag, Tooltip } from 'antd'; @@ -25,6 +25,7 @@ import ChartAnswer from '@/components/pages/home/promptThread/ChartAnswer'; import Preparation from '@/components/pages/home/preparation'; import { AskingTaskStatus, + ChartTaskStatus, ThreadResponse, ThreadResponseAnswerDetail, ThreadResponseAnswerStatus, @@ -186,6 +187,19 @@ const isNeedGenerateAnswer = (answerDetail: ThreadResponseAnswerDetail) => { return answerDetail?.queryId === null && !isFinished && !isProcessing; }; +const isAnswerGenerationInProgress = ( + status?: ThreadResponseAnswerStatus | null, +) => + [ + ThreadResponseAnswerStatus.NOT_STARTED, + ThreadResponseAnswerStatus.FETCHING_DATA, + ThreadResponseAnswerStatus.PREPROCESSING, + ThreadResponseAnswerStatus.STREAMING, + ].includes(status); + +const isChartGenerationActive = (status?: ChartTaskStatus | null) => + [ChartTaskStatus.FETCHING, ChartTaskStatus.GENERATING].includes(status); + export default function AnswerResult(props: Props) { const { threadResponse, isLastThreadResponse, isOpeningQuestion } = props; @@ -213,6 +227,7 @@ export default function AnswerResult(props: Props) { view, adjustment, } = threadResponse; + const autoGenerateAnswerRef = useRef(null); const resultStyle = isLastThreadResponse ? { minHeight: 'calc(100vh - (194px))' } @@ -235,10 +250,17 @@ export default function AnswerResult(props: Props) { // initialize generate answer useEffect(() => { if (isBreakdownOnly) return; + if ( + autoGenerateAnswerRef.current === id && + isAnswerGenerationInProgress(answerDetail?.status) + ) { + return; + } if ( canGenerateAnswer(askingTask, adjustmentTask) && isNeedGenerateAnswer(answerDetail) ) { + autoGenerateAnswerRef.current = id; const debouncedGenerateAnswer = debounce( () => { onGenerateTextBasedAnswer(id); @@ -254,14 +276,25 @@ export default function AnswerResult(props: Props) { }; } }, [ + id, isBreakdownOnly, askingTask?.status, adjustmentTask?.status, answerDetail?.status, ]); + useEffect(() => { + if (getAnswerIsFinished(answerDetail?.status)) { + autoGenerateAnswerRef.current = null; + } + }, [answerDetail?.status]); + const onTabClick = (activeKey: string) => { - if (activeKey === ANSWER_TAB_KEYS.CHART && !threadResponse.chartDetail) { + if ( + activeKey === ANSWER_TAB_KEYS.CHART && + !threadResponse.chartDetail && + !isChartGenerationActive(threadResponse.chartDetail?.status) + ) { onGenerateChartAnswer(id); } }; From 465c1ed9a43f74bebdbcd9b87b0d2edc24e74641 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 22 May 2026 15:17:19 +0530 Subject: [PATCH 0045/1087] update SQL --- .../pipelines/generation/sql_regeneration.py | 2 + .../src/pipelines/generation/utils/sql.py | 52 +++++++++---------- 2 files changed, 28 insertions(+), 26 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 562e35df5b..7cf5787bc5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -172,11 +172,13 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, + data_source: str, project_id: str | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, + data_source=data_source, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 306140a1fc..543b916657 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -118,14 +118,14 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", lambda m: ( - f"(DATEPART('YEAR', {m.group(1)}) * 100 + DATEPART('MONTH', {m.group(1)}))" + f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" ), sql, flags=re.IGNORECASE, ) sql = re.sub( rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", sql, flags=re.IGNORECASE, ) @@ -140,124 +140,124 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"DATEPART(YEAR, {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"DATEPART(MONTH, {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DATEPART(DAY, {m.group(1)})", ), ] @@ -511,15 +511,15 @@ async def _classify_generation_result( _MSSQL_TEXT_TO_SQL_RULES = """ ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. -- Prefer native T-SQL date bucket syntax such as DATEPART('YEAR', "created_at") and DATEPART('MONTH', "created_at"). +- Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - - DATEPART('YEAR', ) AS "year" - - DATEPART('MONTH', ) AS "month" + - DATEPART(YEAR, ) AS "year" + - DATEPART(MONTH, ) AS "month" Then GROUP BY and ORDER BY the same year/month expressions. -- For year bucketing, prefer DATEPART('YEAR', ). +- For year bucketing, prefer DATEPART(YEAR, ). - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' - AND < '2026-01-01 00:00:00' @@ -805,7 +805,7 @@ def get_metric_instructions( #### MSSQL Metric Notes #### - Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. - Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. -- For month trend metrics, prefer DATEPART('YEAR', ) and DATEPART('MONTH', ) as separate grouped columns. +- For month trend metrics, prefer DATEPART(YEAR, ) and DATEPART(MONTH, ) as separate grouped columns. """ return instructions From 4db3f9d0b69cc8fc7717a8cbf6653fbced05e770 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 22 May 2026 17:41:08 +0530 Subject: [PATCH 0046/1087] update chart --- .../src/apollo/server/backgrounds/chart.ts | 170 +++++++++++++----- .../server/services/askingTaskTracker.ts | 49 ++++- wren-ui/src/components/chart/handler.ts | 37 ++++ 3 files changed, 210 insertions(+), 46 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index 01f3039cfa..39b9dc23de 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -22,8 +22,13 @@ const isFinalized = (status: ChartStatus) => { ); }; +const MIN_POLL_DELAY = 2000; +const MAX_POLL_DELAY = 10000; + export class ChartBackgroundTracker { private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private threadResponseRepository: IThreadResponseRepository; @@ -56,6 +61,10 @@ export class ChartBackgroundTracker { return; } + if (Date.now() < (this.nextPollAt[threadResponse.id] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(threadResponse.id); @@ -67,8 +76,17 @@ export class ChartBackgroundTracker { chartDetail.queryId, ); + const statusChanged = chartDetail.status !== result.status; + this.scheduleNextPoll(threadResponse.id, result.status, statusChanged); + + if (isFinalized(result.status) && !statusChanged) { + this.finalizeTask(threadResponse, result); + this.runningJobs.delete(threadResponse.id); + return; + } + // check if status change - if (chartDetail.status === result.status) { + if (!statusChanged) { // mark the job as finished logger.debug( `Job ${threadResponse.id} chart status not changed, finished`, @@ -96,27 +114,7 @@ export class ChartBackgroundTracker { // remove the task from tracker if it is finalized if (isFinalized(result.status)) { - const eventProperties = { - question: threadResponse.question, - error: result.error, - }; - if (result.status === ChartStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_CHART, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_CHART, - eventProperties, - WrenService.AI, - false, - ); - } - logger.debug( - `Job ${threadResponse.id} chart is finalized, removing`, - ); - delete this.tasks[threadResponse.id]; + this.finalizeTask(threadResponse, result); } // mark the job as finished @@ -138,15 +136,60 @@ export class ChartBackgroundTracker { public addTask(threadResponse: ThreadResponse) { this.tasks[threadResponse.id] = threadResponse; + this.nextPollAt[threadResponse.id] = Date.now(); + this.pollDelay[threadResponse.id] = MIN_POLL_DELAY; } public getTasks() { return this.tasks; } + + private scheduleNextPoll( + taskId: number, + status: ChartStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskId] = Number.MAX_SAFE_INTEGER; + return; + } + + const baseDelay = status === ChartStatus.FETCHING ? MIN_POLL_DELAY : 3000; + this.pollDelay[taskId] = resultChanged + ? baseDelay + : Math.min( + Math.max((this.pollDelay[taskId] || baseDelay) * 1.5, baseDelay), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskId] = Date.now() + this.pollDelay[taskId]; + } + + private finalizeTask(threadResponse: ThreadResponse, result) { + const eventProperties = { + question: threadResponse.question, + error: result.error, + }; + if (result.status === ChartStatus.FINISHED) { + this.telemetry.sendEvent(TelemetryEvent.HOME_ANSWER_CHART, eventProperties); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_CHART, + eventProperties, + WrenService.AI, + false, + ); + } + logger.debug(`Job ${threadResponse.id} chart is finalized, removing`); + delete this.tasks[threadResponse.id]; + delete this.nextPollAt[threadResponse.id]; + delete this.pollDelay[threadResponse.id]; + } } export class ChartAdjustmentBackgroundTracker { private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private threadResponseRepository: IThreadResponseRepository; @@ -179,6 +222,10 @@ export class ChartAdjustmentBackgroundTracker { return; } + if (Date.now() < (this.nextPollAt[threadResponse.id] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(threadResponse.id); @@ -190,8 +237,17 @@ export class ChartAdjustmentBackgroundTracker { chartDetail.queryId, ); + const statusChanged = chartDetail.status !== result.status; + this.scheduleNextPoll(threadResponse.id, result.status, statusChanged); + + if (isFinalized(result.status) && !statusChanged) { + this.finalizeTask(threadResponse, result); + this.runningJobs.delete(threadResponse.id); + return; + } + // check if status change - if (chartDetail.status === result.status) { + if (!statusChanged) { // mark the job as finished logger.debug( `Job ${threadResponse.id} chart status not changed, finished`, @@ -220,27 +276,7 @@ export class ChartAdjustmentBackgroundTracker { // remove the task from tracker if it is finalized if (isFinalized(result.status)) { - const eventProperties = { - question: threadResponse.question, - error: result.error, - }; - if (result.status === ChartStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_ADJUST_CHART, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_ADJUST_CHART, - eventProperties, - WrenService.AI, - false, - ); - } - logger.debug( - `Job ${threadResponse.id} chart is finalized, removing`, - ); - delete this.tasks[threadResponse.id]; + this.finalizeTask(threadResponse, result); } // mark the job as finished @@ -262,9 +298,55 @@ export class ChartAdjustmentBackgroundTracker { public addTask(threadResponse: ThreadResponse) { this.tasks[threadResponse.id] = threadResponse; + this.nextPollAt[threadResponse.id] = Date.now(); + this.pollDelay[threadResponse.id] = MIN_POLL_DELAY; } public getTasks() { return this.tasks; } + + private scheduleNextPoll( + taskId: number, + status: ChartStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskId] = Number.MAX_SAFE_INTEGER; + return; + } + + const baseDelay = status === ChartStatus.FETCHING ? MIN_POLL_DELAY : 3000; + this.pollDelay[taskId] = resultChanged + ? baseDelay + : Math.min( + Math.max((this.pollDelay[taskId] || baseDelay) * 1.5, baseDelay), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskId] = Date.now() + this.pollDelay[taskId]; + } + + private finalizeTask(threadResponse: ThreadResponse, result) { + const eventProperties = { + question: threadResponse.question, + error: result.error, + }; + if (result.status === ChartStatus.FINISHED) { + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_ADJUST_CHART, + eventProperties, + ); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_ADJUST_CHART, + eventProperties, + WrenService.AI, + false, + ); + } + logger.debug(`Job ${threadResponse.id} chart is finalized, removing`); + delete this.tasks[threadResponse.id]; + delete this.nextPollAt[threadResponse.id]; + delete this.pollDelay[threadResponse.id]; + } } diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index aaf015d53e..41438e7945 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -22,6 +22,8 @@ interface TrackedTask { queryId: string; taskId?: number; lastPolled: number; + nextPollAt: number; + pollDelay: number; question?: string; result?: AskResult; isFinalized: boolean; @@ -56,6 +58,8 @@ export interface IAskingTaskTracker { } export class AskingTaskTracker implements IAskingTaskTracker { + private readonly minPollDelay = 2000; + private readonly maxPollDelay = 8000; private wrenAIAdaptor: IWrenAIAdaptor; private askingTaskRepository: IAskingTaskRepository; private trackedTasks: Map = new Map(); @@ -114,6 +118,8 @@ export class AskingTaskTracker implements IAskingTaskTracker { const task = { queryId, lastPolled: Date.now(), + nextPollAt: Date.now(), + pollDelay: this.minPollDelay, question: input.query, result: { type: null, @@ -296,16 +302,22 @@ export class AskingTaskTracker implements IAskingTaskTracker { return; } + if (now < task.nextPollAt) { + return; + } + // Mark the job as running this.runningJobs.add(queryId); // Poll for updates - logger.info(`Polling for updates for task ${queryId}`); + logger.debug(`Polling for updates for task ${queryId}`); const result = await this.wrenAIAdaptor.getAskResult(queryId); task.lastPolled = now; + const resultChanged = this.isResultChanged(task.result, result); + this.scheduleNextPoll(task, result.status, resultChanged); // if result is not changed, we don't need to update the database - if (!this.isResultChanged(task.result, result)) { + if (!resultChanged) { this.runningJobs.delete(queryId); return; } @@ -521,6 +533,8 @@ export class AskingTaskTracker implements IAskingTaskTracker { queryId: result.queryId, taskId: result.taskId, lastPolled: Date.now(), + nextPollAt: Date.now(), + pollDelay: this.minPollDelay, question: result.question, result: this.getComparableAskResult(result), isFinalized: false, @@ -545,4 +559,35 @@ export class AskingTaskTracker implements IAskingTaskTracker { traceId: result?.traceId ?? null, } as AskResult; } + + private scheduleNextPoll( + task: TrackedTask, + status: AskResultStatus, + resultChanged: boolean, + ) { + if (this.isTaskFinalized(status)) { + task.nextPollAt = Number.MAX_SAFE_INTEGER; + return; + } + + const baseDelay = this.getBasePollDelay(status); + task.pollDelay = resultChanged + ? baseDelay + : Math.min(Math.max(task.pollDelay * 1.5, baseDelay), this.maxPollDelay); + task.nextPollAt = Date.now() + task.pollDelay; + } + + private getBasePollDelay(status: AskResultStatus) { + if ( + [ + AskResultStatus.UNDERSTANDING, + AskResultStatus.SEARCHING, + AskResultStatus.PLANNING, + ].includes(status) + ) { + return this.minPollDelay; + } + + return 3000; + } } diff --git a/wren-ui/src/components/chart/handler.ts b/wren-ui/src/components/chart/handler.ts index acc27e6886..48300930dc 100644 --- a/wren-ui/src/components/chart/handler.ts +++ b/wren-ui/src/components/chart/handler.ts @@ -185,6 +185,7 @@ export default class ChartSpecHandler { // avoid mutating the original spec const clonedSpec = cloneDeep(spec); + this.normalizeSpecFields(clonedSpec); this.parseSpec(clonedSpec); } @@ -258,6 +259,42 @@ export default class ChartSpecHandler { } } + private normalizeSpecFields(spec: TopLevelSpec) { + const values = ((this.data as any)?.values || []) as Record[]; + if (!values.length) return; + + const columns = Object.keys(values[0]); + const normalizeField = (field?: string) => { + if (!field) return field; + if (columns.includes(field)) return field; + + const lowered = columns.find( + (column) => column.toLowerCase() === field.toLowerCase(), + ); + if (lowered) return lowered; + + const compactField = field.replace(/[\s_]+/g, '').toLowerCase(); + return ( + columns.find( + (column) => + column.replace(/[\s_]+/g, '').toLowerCase() === compactField, + ) || field + ); + }; + + const encoding = (spec as any).encoding as EncodingSpec; + ['x', 'y', 'theta', 'color', 'xOffset'].forEach((key) => { + const axis = encoding?.[key] as { field?: string } | undefined; + if (axis?.field) axis.field = normalizeField(axis.field); + }); + + (spec.transform || []).forEach((transform: any) => { + if (Array.isArray(transform?.fold)) { + transform.fold = transform.fold.map(normalizeField); + } + }); + } + private addMark(mark: MarkSpec) { let additionalProps = {}; From fad57b487fd7dfdb9b1d0d028039a4febf9ca251 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 22 May 2026 19:19:56 +0530 Subject: [PATCH 0047/1087] update chart polling --- .../server/backgrounds/recommend-question.ts | 203 ++++++++++++------ .../apollo/server/services/askingService.ts | 41 ++++ .../apollo/server/services/projectService.ts | 17 ++ wren-ui/src/hooks/useAskPrompt.tsx | 44 +++- 4 files changed, 236 insertions(+), 69 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts index a649bdaf4b..e68395e076 100644 --- a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts +++ b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts @@ -20,9 +20,14 @@ const isFinalized = (status: RecommendationQuestionStatus) => { ].includes(status); }; +const MIN_POLL_DELAY = 2000; +const MAX_POLL_DELAY = 10000; + export class ProjectRecommendQuestionBackgroundTracker { // tasks is a kv pair of task id and thread response private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private projectRepository: IProjectRepository; @@ -58,6 +63,10 @@ export class ProjectRecommendQuestionBackgroundTracker { return; } + if (Date.now() < (this.nextPollAt[this.taskKey(project)] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(this.taskKey(project)); @@ -68,11 +77,19 @@ export class ProjectRecommendQuestionBackgroundTracker { project.queryId, ); + const changed = + project.questionsStatus !== result.status || + result.response?.questions.length !== (project.questions || []).length; + this.scheduleNextPoll(this.taskKey(project), result.status, changed); + + if (isFinalized(result.status) && !changed) { + this.finalizeTask(project, result); + this.runningJobs.delete(this.taskKey(project)); + return; + } + // check if status change - if ( - project.questionsStatus === result.status && - result.response?.questions.length === (project.questions || []).length - ) { + if (!changed) { // mark the job as finished this.logger.debug( `${loggerPrefix}job ${this.taskKey(project)} status not changed, returning question count: ${result.response?.questions.length || 0}`, @@ -82,10 +99,7 @@ export class ProjectRecommendQuestionBackgroundTracker { } // update database - if ( - result.status !== project.questionsStatus || - result.response?.questions.length !== (project.questions || []).length - ) { + if (changed) { this.logger.debug( `${loggerPrefix}job ${this.taskKey(project)} have changes, returning question count: ${result.response?.questions.length || 0}, updating`, ); @@ -100,30 +114,7 @@ export class ProjectRecommendQuestionBackgroundTracker { // remove the task from tracker if it is finalized if (isFinalized(result.status)) { - const eventProperties = { - projectId: project.id, - projectType: project.type, - status: result.status, - questions: project.questions, - error: result.error, - }; - if (result.status === RecommendationQuestionStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, - eventProperties, - WrenService.AI, - false, - ); - } - this.logger.debug( - `${loggerPrefix}job ${this.taskKey(project)} is finalized, removing`, - ); - delete this.tasks[this.taskKey(project)]; + this.finalizeTask(project, result); } // mark the job as finished @@ -144,6 +135,8 @@ export class ProjectRecommendQuestionBackgroundTracker { public addTask(project: Project) { this.tasks[this.taskKey(project)] = project; + this.nextPollAt[this.taskKey(project)] = Date.now(); + this.pollDelay[this.taskKey(project)] = MIN_POLL_DELAY; } public getTasks() { @@ -175,11 +168,60 @@ export class ProjectRecommendQuestionBackgroundTracker { public isExist(project: Project) { return this.tasks[this.taskKey(project)]; } + + private scheduleNextPoll( + taskKey: number, + status: RecommendationQuestionStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskKey] = Number.MAX_SAFE_INTEGER; + return; + } + + this.pollDelay[taskKey] = resultChanged + ? MIN_POLL_DELAY + : Math.min( + Math.max((this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, MIN_POLL_DELAY), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskKey] = Date.now() + this.pollDelay[taskKey]; + } + + private finalizeTask(project: Project, result) { + const eventProperties = { + projectId: project.id, + projectType: project.type, + status: result.status, + questions: project.questions, + error: result.error, + }; + if (result.status === RecommendationQuestionStatus.FINISHED) { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, + eventProperties, + ); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, + eventProperties, + WrenService.AI, + false, + ); + } + const taskKey = this.taskKey(project); + this.logger.debug(`${loggerPrefix}job ${taskKey} is finalized, removing`); + delete this.tasks[taskKey]; + delete this.nextPollAt[taskKey]; + delete this.pollDelay[taskKey]; + } } export class ThreadRecommendQuestionBackgroundTracker { // tasks is a kv pair of task id and thread response private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private threadRepository: IThreadRepository; @@ -215,6 +257,10 @@ export class ThreadRecommendQuestionBackgroundTracker { return; } + if (Date.now() < (this.nextPollAt[this.taskKey(thread)] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(this.taskKey(thread)); @@ -225,11 +271,19 @@ export class ThreadRecommendQuestionBackgroundTracker { thread.queryId, ); + const changed = + thread.questionsStatus !== result.status || + result.response?.questions.length !== (thread.questions || []).length; + this.scheduleNextPoll(this.taskKey(thread), result.status, changed); + + if (isFinalized(result.status) && !changed) { + this.finalizeTask(thread, result); + this.runningJobs.delete(this.taskKey(thread)); + return; + } + // check if status change - if ( - thread.questionsStatus === result.status && - result.response?.questions.length === (thread.questions || []).length - ) { + if (!changed) { // mark the job as finished this.logger.debug( `${loggerPrefix}job ${this.taskKey(thread)} status not changed, returning question count: ${result.response?.questions.length || 0}`, @@ -239,10 +293,7 @@ export class ThreadRecommendQuestionBackgroundTracker { } // update database - if ( - result.status !== thread.questionsStatus || - result.response?.questions.length !== (thread.questions || []).length - ) { + if (changed) { this.logger.debug( `${loggerPrefix}job ${this.taskKey(thread)} have changes, returning question count: ${result.response?.questions.length || 0}, updating`, ); @@ -257,29 +308,7 @@ export class ThreadRecommendQuestionBackgroundTracker { // remove the task from tracker if it is finalized if (isFinalized(result.status)) { - const eventProperties = { - thread_id: thread.id, - status: result.status, - questions: thread.questions, - error: result.error, - }; - if (result.status === RecommendationQuestionStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, - eventProperties, - WrenService.AI, - false, - ); - } - this.logger.debug( - `${loggerPrefix}job ${this.taskKey(thread)} is finalized, removing`, - ); - delete this.tasks[this.taskKey(thread)]; + this.finalizeTask(thread, result); } // mark the job as finished @@ -300,6 +329,8 @@ export class ThreadRecommendQuestionBackgroundTracker { public addTask(thread: Thread) { this.tasks[this.taskKey(thread)] = thread; + this.nextPollAt[this.taskKey(thread)] = Date.now(); + this.pollDelay[this.taskKey(thread)] = MIN_POLL_DELAY; } public getTasks() { @@ -332,4 +363,50 @@ export class ThreadRecommendQuestionBackgroundTracker { public isExist(thread: Thread) { return this.tasks[this.taskKey(thread)]; } + + private scheduleNextPoll( + taskKey: number, + status: RecommendationQuestionStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskKey] = Number.MAX_SAFE_INTEGER; + return; + } + + this.pollDelay[taskKey] = resultChanged + ? MIN_POLL_DELAY + : Math.min( + Math.max((this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, MIN_POLL_DELAY), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskKey] = Date.now() + this.pollDelay[taskKey]; + } + + private finalizeTask(thread: Thread, result) { + const eventProperties = { + thread_id: thread.id, + status: result.status, + questions: thread.questions, + error: result.error, + }; + if (result.status === RecommendationQuestionStatus.FINISHED) { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, + eventProperties, + ); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, + eventProperties, + WrenService.AI, + false, + ); + } + const taskKey = this.taskKey(thread); + this.logger.debug(`${loggerPrefix}job ${taskKey} is finalized, removing`); + delete this.tasks[taskKey]; + delete this.nextPollAt[taskKey]; + delete this.pollDelay[taskKey]; + } } diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index f0c3789b63..b9b0545c0f 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -433,6 +433,8 @@ export class AskingService implements IAskingService { private askingTaskTracker: IAskingTaskTracker; private askingTaskRepository: IAskingTaskRepository; private adjustmentBackgroundTracker: AdjustmentBackgroundTaskTracker; + private instantRecommendationJobs = new Map>(); + private threadRecommendationJobs = new Map>(); private initialized = false; constructor({ @@ -536,6 +538,26 @@ export class AskingService implements IAskingService { public async generateThreadRecommendationQuestions( threadId: number, + ): Promise { + const existingJob = this.threadRecommendationJobs.get(threadId); + if (existingJob) { + logger.debug( + `thread "${threadId}" recommended questions are already being requested, reusing in-flight job`, + ); + return existingJob; + } + + const job = this.doGenerateThreadRecommendationQuestions(threadId); + this.threadRecommendationJobs.set(threadId, job); + try { + return await job; + } finally { + this.threadRecommendationJobs.delete(threadId); + } + } + + private async doGenerateThreadRecommendationQuestions( + threadId: number, ): Promise { const thread = await this.threadRepository.findOneBy({ id: threadId }); if (!thread) { @@ -1058,6 +1080,25 @@ export class AskingService implements IAskingService { public async createInstantRecommendedQuestions( input: InstantRecommendedQuestionsInput, + ): Promise { + const key = JSON.stringify(input.previousQuestions || []); + const existingJob = this.instantRecommendationJobs.get(key); + if (existingJob) { + logger.debug('instant recommended questions are already being requested'); + return existingJob; + } + + const job = this.doCreateInstantRecommendedQuestions(input); + this.instantRecommendationJobs.set(key, job); + try { + return await job; + } finally { + this.instantRecommendationJobs.delete(key); + } + } + + private async doCreateInstantRecommendedQuestions( + input: InstantRecommendedQuestionsInput, ): Promise { const project = await this.projectService.getCurrentProject(); const { manifest } = await this.deployService.getLastDeployment(project.id); diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 3f1b74318b..fc9cd98bb8 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -88,6 +88,7 @@ export class ProjectService implements IProjectService { private mdlService: IMDLService; private wrenAIAdaptor: IWrenAIAdaptor; private projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; + private projectRecommendationJob: Promise | null = null; constructor({ projectRepository, metadataService, @@ -132,6 +133,22 @@ export class ProjectService implements IProjectService { } public async generateProjectRecommendationQuestions(): Promise { + if (this.projectRecommendationJob) { + logger.debug( + 'project recommended questions are already being requested, reusing in-flight job', + ); + return this.projectRecommendationJob; + } + + this.projectRecommendationJob = this.doGenerateProjectRecommendationQuestions(); + try { + return await this.projectRecommendationJob; + } finally { + this.projectRecommendationJob = null; + } + } + + private async doGenerateProjectRecommendationQuestions(): Promise { const project = await this.getCurrentProject(); if (!project) { throw new Error(`Project not found`); diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index b028604794..a524ee7b3c 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -206,6 +206,8 @@ export default function useAskPrompt(threadId?: number) { RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS, ); const lastRecommendedFingerprintRef = useRef(null); + const recommendedCreationKeyRef = useRef(null); + const recommendedCreationRequestRef = useRef | null>(null); const askingTask = useMemo( () => askingTaskResult.data?.askingTask || null, @@ -240,6 +242,8 @@ export default function useAskPrompt(threadId?: number) { recommendedPollingRef.current = null; } recommendedPollingDelayRef.current = RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS; + recommendedCreationKeyRef.current = null; + recommendedCreationRequestRef.current = null; }, []); const startAskingTaskPolling = useCallback( @@ -350,14 +354,42 @@ export default function useAskPrompt(threadId?: number) { ...uniq(threadQuestions).slice(-5), originalQuestion, ]; - const response = await createInstantRecommendedQuestions({ - variables: { data: { previousQuestions } }, + const creationKey = JSON.stringify({ + askingQueryId: askingTask?.queryId || null, + previousQuestions, }); - const taskId = response.data?.createInstantRecommendedQuestions?.id; - if (!taskId) return; - await startRecommendedPolling(taskId); - }, [originalQuestion, threadQuestions, startRecommendedPolling]); + if (recommendedCreationKeyRef.current === creationKey) { + if (recommendedCreationRequestRef.current) { + await recommendedCreationRequestRef.current; + } + return; + } + + recommendedCreationKeyRef.current = creationKey; + const request = (async () => { + const response = await createInstantRecommendedQuestions({ + variables: { data: { previousQuestions } }, + }); + const taskId = response.data?.createInstantRecommendedQuestions?.id; + if (!taskId) return; + + await startRecommendedPolling(taskId); + })(); + + recommendedCreationRequestRef.current = request; + try { + await request; + } finally { + recommendedCreationRequestRef.current = null; + } + }, [ + originalQuestion, + threadQuestions, + askingTask?.queryId, + createInstantRecommendedQuestions, + startRecommendedPolling, + ]); const checkFetchAskingStreamTask = useCallback( (task: AskingTask) => { From 279464cb4e301261f42592c78424b48963b6bf5f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 22 May 2026 22:00:17 +0530 Subject: [PATCH 0048/1087] update sql --- .../src/pipelines/generation/utils/sql.py | 33 +++++++++++++++---- 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 543b916657..9d5f869426 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -268,6 +268,26 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: return rewritten +def _rewrite_mssql_timestamp_casts(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + timestamp_function_pattern = re.compile( + rf"\bTO_TIMESTAMP(?:_(?:MILLIS|SECONDS|MICROS|NANOS))?\(\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ) + timestamp_cast_pattern = re.compile( + r"CAST\(\s*((?:[^()]|\([^()]*\))+?)\s+AS\s+TIMESTAMP\s*\)", + re.IGNORECASE, + ) + + rewritten = timestamp_function_pattern.sub( + lambda m: f"CAST({m.group(1)} AS DATETIME)", sql + ) + rewritten = timestamp_cast_pattern.sub( + lambda m: f"CAST({m.group(1)} AS DATETIME)", rewritten + ) + return rewritten + + def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = sql @@ -283,6 +303,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> flags=re.IGNORECASE, ) normalized = _replace_relative_getdate_calls(normalized, now) + normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -472,10 +493,10 @@ async def _classify_generation_result( - Use "lower(.) = lower()" when: - The user requests an exact, specific value. - There is no ambiguity or pattern in the value. -- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query - - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds - - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds - - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds +- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to a temporal type first before using it in the query. + - For engines that list these functions in SQL FUNCTIONS, use TO_TIMESTAMP_MILLIS("") if the timestamp_column is in milliseconds. + - For engines that list these functions in SQL FUNCTIONS, use TO_TIMESTAMP_SECONDS("") if the timestamp_column is in seconds. + - For engines that list these functions in SQL FUNCTIONS, use TO_TIMESTAMP_MICROS("") if the timestamp_column is in microseconds. - When you need to cast a date/time related field, CAST it to a temporal type that is supported by the target data source and consistent with the SQL FUNCTIONS section. - example 1: CAST(properties_closedate AS TIMESTAMP) - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP) @@ -512,7 +533,7 @@ async def _classify_generation_result( ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. - Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). -- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, or :: casts. +- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: @@ -523,7 +544,7 @@ async def _classify_generation_result( - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' - AND < '2026-01-01 00:00:00' -- When a temporal cast is required, keep literal timestamps as plain ISO strings if the column is already datetime-like. +- When a temporal cast is required, use CAST( AS DATETIME), or keep literal timestamps as plain ISO strings if the column is already datetime-like. - Keep MSSQL date logic simple and planner-safe. Never emit DATEADD/DATEDIFF fallback expressions unless the SQL FUNCTIONS section explicitly requires them. """ From fd285e52ccc0f4e7c6594f0b202571e464b5228c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 25 May 2026 19:21:31 +0530 Subject: [PATCH 0049/1087] update sql cache --- .../src/apollo/server/backgrounds/chart.ts | 45 ++++++++-- .../dashboardCacheBackgroundTracker.ts | 14 +++- .../server/backgrounds/recommend-question.ts | 41 +++++++-- .../apollo/server/services/askingService.ts | 44 +++++++--- .../apollo/server/services/projectService.ts | 10 ++- wren-ui/src/common.ts | 84 ++++++++++++++----- 6 files changed, 190 insertions(+), 48 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index 39b9dc23de..35c9c22b79 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -34,6 +34,7 @@ export class ChartBackgroundTracker { private threadResponseRepository: IThreadResponseRepository; private runningJobs = new Set(); private telemetry: PostHogTelemetry; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -52,8 +53,11 @@ export class ChartBackgroundTracker { } private start() { + if (this.intervalId) { + return; + } logger.info('Chart background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map( (threadResponse) => async () => { // check if same job is running @@ -77,7 +81,11 @@ export class ChartBackgroundTracker { ); const statusChanged = chartDetail.status !== result.status; - this.scheduleNextPoll(threadResponse.id, result.status, statusChanged); + this.scheduleNextPoll( + threadResponse.id, + result.status, + statusChanged, + ); if (isFinalized(result.status) && !statusChanged) { this.finalizeTask(threadResponse, result); @@ -134,6 +142,14 @@ export class ChartBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(threadResponse: ThreadResponse) { this.tasks[threadResponse.id] = threadResponse; this.nextPollAt[threadResponse.id] = Date.now(); @@ -170,7 +186,10 @@ export class ChartBackgroundTracker { error: result.error, }; if (result.status === ChartStatus.FINISHED) { - this.telemetry.sendEvent(TelemetryEvent.HOME_ANSWER_CHART, eventProperties); + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_CHART, + eventProperties, + ); } else { this.telemetry.sendEvent( TelemetryEvent.HOME_ANSWER_CHART, @@ -195,6 +214,7 @@ export class ChartAdjustmentBackgroundTracker { private threadResponseRepository: IThreadResponseRepository; private runningJobs = new Set(); private telemetry: PostHogTelemetry; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -213,8 +233,11 @@ export class ChartAdjustmentBackgroundTracker { } private start() { + if (this.intervalId) { + return; + } logger.info('Chart adjustment background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map( (threadResponse) => async () => { // check if same job is running @@ -238,7 +261,11 @@ export class ChartAdjustmentBackgroundTracker { ); const statusChanged = chartDetail.status !== result.status; - this.scheduleNextPoll(threadResponse.id, result.status, statusChanged); + this.scheduleNextPoll( + threadResponse.id, + result.status, + statusChanged, + ); if (isFinalized(result.status) && !statusChanged) { this.finalizeTask(threadResponse, result); @@ -296,6 +323,14 @@ export class ChartAdjustmentBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(threadResponse: ThreadResponse) { this.tasks[threadResponse.id] = threadResponse; this.nextPollAt[threadResponse.id] = Date.now(); diff --git a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts index baa496b21f..bad4c69b13 100644 --- a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts @@ -25,6 +25,7 @@ export class DashboardCacheBackgroundTracker { private deployService: IDeployService; private queryService: IQueryService; private runningJobs = new Set(); + private intervalId?: NodeJS.Timeout; constructor({ dashboardRepository, @@ -52,12 +53,23 @@ export class DashboardCacheBackgroundTracker { } private start(): void { + if (this.intervalId) { + return; + } logger.info('Dashboard cache background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { this.checkAndRefreshCaches(); }, this.intervalTime); } + public stop(): void { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + private async checkAndRefreshCaches(): Promise { try { // Get all dashboards with cache enabled diff --git a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts index e68395e076..4831425afa 100644 --- a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts +++ b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts @@ -35,6 +35,7 @@ export class ProjectRecommendQuestionBackgroundTracker { private telemetry: ITelemetry; private logger: Logger; private initialized = false; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -55,8 +56,11 @@ export class ProjectRecommendQuestionBackgroundTracker { } public start() { + if (this.intervalId) { + return; + } this.logger.info('Recommend question background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map((project) => async () => { // check if same job is running if (this.runningJobs.has(this.taskKey(project))) { @@ -79,7 +83,8 @@ export class ProjectRecommendQuestionBackgroundTracker { const changed = project.questionsStatus !== result.status || - result.response?.questions.length !== (project.questions || []).length; + result.response?.questions.length !== + (project.questions || []).length; this.scheduleNextPoll(this.taskKey(project), result.status, changed); if (isFinalized(result.status) && !changed) { @@ -133,6 +138,14 @@ export class ProjectRecommendQuestionBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(project: Project) { this.tasks[this.taskKey(project)] = project; this.nextPollAt[this.taskKey(project)] = Date.now(); @@ -182,7 +195,10 @@ export class ProjectRecommendQuestionBackgroundTracker { this.pollDelay[taskKey] = resultChanged ? MIN_POLL_DELAY : Math.min( - Math.max((this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, MIN_POLL_DELAY), + Math.max( + (this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, + MIN_POLL_DELAY, + ), MAX_POLL_DELAY, ); this.nextPollAt[taskKey] = Date.now() + this.pollDelay[taskKey]; @@ -229,6 +245,7 @@ export class ThreadRecommendQuestionBackgroundTracker { private telemetry: ITelemetry; private logger: Logger; private initialized = false; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -249,8 +266,11 @@ export class ThreadRecommendQuestionBackgroundTracker { } public start() { + if (this.intervalId) { + return; + } this.logger.info('Recommend question background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map((thread) => async () => { // check if same job is running if (this.runningJobs.has(this.taskKey(thread))) { @@ -327,6 +347,14 @@ export class ThreadRecommendQuestionBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(thread: Thread) { this.tasks[this.taskKey(thread)] = thread; this.nextPollAt[this.taskKey(thread)] = Date.now(); @@ -377,7 +405,10 @@ export class ThreadRecommendQuestionBackgroundTracker { this.pollDelay[taskKey] = resultChanged ? MIN_POLL_DELAY : Math.min( - Math.max((this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, MIN_POLL_DELAY), + Math.max( + (this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, + MIN_POLL_DELAY, + ), MAX_POLL_DELAY, ); this.nextPollAt[taskKey] = Date.now() + this.pollDelay[taskKey]; diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index b9b0545c0f..c0e06972c3 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -105,16 +105,18 @@ export enum ThreadResponseAnswerStatus { const isAnswerGenerationInProgress = ( status?: ThreadResponseAnswerStatus | string | null, ) => - ([ - ThreadResponseAnswerStatus.NOT_STARTED, - ThreadResponseAnswerStatus.FETCHING_DATA, - ThreadResponseAnswerStatus.PREPROCESSING, - ThreadResponseAnswerStatus.STREAMING, - ] as string[]).includes(status || ""); + ( + [ + ThreadResponseAnswerStatus.NOT_STARTED, + ThreadResponseAnswerStatus.FETCHING_DATA, + ThreadResponseAnswerStatus.PREPROCESSING, + ThreadResponseAnswerStatus.STREAMING, + ] as string[] + ).includes(status || ''); const isChartGenerationInProgress = (status?: ChartStatus | string | null) => ([ChartStatus.FETCHING, ChartStatus.GENERATING] as string[]).includes( - status || "", + status || '', ); // adjustment input @@ -312,7 +314,7 @@ class BreakdownBackgroundTracker { }: { telemetry: PostHogTelemetry; wrenAIAdaptor: IWrenAIAdaptor; - threadResponseRepository: IThreadResponseRepository; + threadResponseRepository: IThreadResponseRepository; }) { this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; @@ -449,6 +451,9 @@ export class AskingService implements IAskingService { queryService, mdlService, askingTaskTracker, + chartBackgroundTracker, + chartAdjustmentBackgroundTracker, + threadRecommendQuestionBackgroundTracker, }: { telemetry: PostHogTelemetry; wrenAIAdaptor: IWrenAIAdaptor; @@ -461,6 +466,9 @@ export class AskingService implements IAskingService { queryService: IQueryService; mdlService: IMDLService; askingTaskTracker: IAskingTaskTracker; + chartBackgroundTracker?: ChartBackgroundTracker; + chartAdjustmentBackgroundTracker?: ChartAdjustmentBackgroundTracker; + threadRecommendQuestionBackgroundTracker?: ThreadRecommendQuestionBackgroundTracker; }) { this.wrenAIAdaptor = wrenAIAdaptor; this.deployService = deployService; @@ -483,18 +491,22 @@ export class AskingService implements IAskingService { deployService, queryService, }); - this.chartBackgroundTracker = new ChartBackgroundTracker({ - telemetry, - wrenAIAdaptor, - threadResponseRepository, - }); + this.chartBackgroundTracker = + chartBackgroundTracker ?? + new ChartBackgroundTracker({ + telemetry, + wrenAIAdaptor, + threadResponseRepository, + }); this.chartAdjustmentBackgroundTracker = + chartAdjustmentBackgroundTracker ?? new ChartAdjustmentBackgroundTracker({ telemetry, wrenAIAdaptor, threadResponseRepository, }); this.threadRecommendQuestionBackgroundTracker = + threadRecommendQuestionBackgroundTracker ?? new ThreadRecommendQuestionBackgroundTracker({ telemetry, wrenAIAdaptor, @@ -512,6 +524,12 @@ export class AskingService implements IAskingService { this.askingTaskTracker = askingTaskTracker; } + public dispose(): void { + this.chartBackgroundTracker.stop(); + this.chartAdjustmentBackgroundTracker.stop(); + this.threadRecommendQuestionBackgroundTracker.stop(); + } + public async getThreadRecommendationQuestions( threadId: number, ): Promise { diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index fc9cd98bb8..c620fcb62c 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -95,24 +95,31 @@ export class ProjectService implements IProjectService { mdlService, wrenAIAdaptor, telemetry, + projectRecommendQuestionBackgroundTracker, }: { projectRepository: IProjectRepository; metadataService: IDataSourceMetadataService; mdlService: IMDLService; wrenAIAdaptor: IWrenAIAdaptor; telemetry: ITelemetry; + projectRecommendQuestionBackgroundTracker?: ProjectRecommendQuestionBackgroundTracker; }) { this.projectRepository = projectRepository; this.metadataService = metadataService; this.mdlService = mdlService; this.wrenAIAdaptor = wrenAIAdaptor; this.projectRecommendQuestionBackgroundTracker = + projectRecommendQuestionBackgroundTracker ?? new ProjectRecommendQuestionBackgroundTracker({ projectRepository, telemetry, wrenAIAdaptor, }); } + + public dispose(): void { + this.projectRecommendQuestionBackgroundTracker.stop(); + } public async updateProject( projectId: number, projectData: Partial, @@ -140,7 +147,8 @@ export class ProjectService implements IProjectService { return this.projectRecommendationJob; } - this.projectRecommendationJob = this.doGenerateProjectRecommendationQuestions(); + this.projectRecommendationJob = + this.doGenerateProjectRecommendationQuestions(); try { return await this.projectRecommendationJob; } finally { diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index d7074eb8ac..e924374435 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -41,6 +41,8 @@ import { ProjectRecommendQuestionBackgroundTracker, ThreadRecommendQuestionBackgroundTracker, DashboardCacheBackgroundTracker, + ChartBackgroundTracker, + ChartAdjustmentBackgroundTracker, } from './apollo/server/backgrounds'; import { SqlPairService } from './apollo/server/services/sqlPairService'; @@ -50,11 +52,18 @@ type Initializable = { initialize?: unknown; }; +type Disposable = { + dispose?: unknown; + stop?: unknown; +}; + type ReusableComponentGraph = { askingTaskTracker?: Initializable; - askingService?: Initializable; - projectRecommendQuestionBackgroundTracker?: Initializable; - threadRecommendQuestionBackgroundTracker?: Initializable; + askingService?: Initializable & Disposable; + projectService?: Disposable; + projectRecommendQuestionBackgroundTracker?: Initializable & Disposable; + threadRecommendQuestionBackgroundTracker?: Initializable & Disposable; + dashboardCacheBackgroundTracker?: Disposable; knex?: { destroy?: () => unknown; }; @@ -66,9 +75,7 @@ const hasInitialize = ( return typeof value?.initialize === 'function'; }; -const isReusableComponentGraph = ( - graph?: ReusableComponentGraph, -): boolean => { +const isReusableComponentGraph = (graph?: ReusableComponentGraph): boolean => { const nestedAskingTaskTracker = ( graph?.askingService as { askingTaskTracker?: Initializable } | undefined )?.askingTaskTracker; @@ -84,11 +91,25 @@ const isReusableComponentGraph = ( }; const disposeComponentGraph = (graph?: ReusableComponentGraph): void => { - if (typeof graph?.knex?.destroy !== 'function') { - return; - } + const disposables = [ + graph?.askingService, + graph?.projectService, + graph?.projectRecommendQuestionBackgroundTracker, + graph?.threadRecommendQuestionBackgroundTracker, + graph?.dashboardCacheBackgroundTracker, + ]; - void Promise.resolve(graph.knex.destroy()).catch(() => undefined); + disposables.forEach((disposable) => { + if (typeof disposable?.dispose === 'function') { + disposable.dispose(); + } else if (typeof disposable?.stop === 'function') { + disposable.stop(); + } + }); + + if (typeof graph?.knex?.destroy === 'function') { + void Promise.resolve(graph.knex.destroy()).catch(() => undefined); + } }; export const initComponents = () => { @@ -132,6 +153,32 @@ export const initComponents = () => { ibisServerEndpoint: serverConfig.ibisServerEndpoint, }); + // background trackers + const projectRecommendQuestionBackgroundTracker = + new ProjectRecommendQuestionBackgroundTracker({ + telemetry, + wrenAIAdaptor, + projectRepository, + }); + const threadRecommendQuestionBackgroundTracker = + new ThreadRecommendQuestionBackgroundTracker({ + telemetry, + wrenAIAdaptor, + threadRepository, + }); + const chartBackgroundTracker = new ChartBackgroundTracker({ + telemetry, + wrenAIAdaptor, + threadResponseRepository, + }); + const chartAdjustmentBackgroundTracker = new ChartAdjustmentBackgroundTracker( + { + telemetry, + wrenAIAdaptor, + threadResponseRepository, + }, + ); + // services const metadataService = new DataSourceMetadataService({ ibisAdaptor, @@ -161,6 +208,7 @@ export const initComponents = () => { mdlService, wrenAIAdaptor, telemetry, + projectRecommendQuestionBackgroundTracker, }); const askingTaskTracker = new AskingTaskTracker({ wrenAIAdaptor, @@ -180,6 +228,9 @@ export const initComponents = () => { mdlService, askingTaskTracker, askingTaskRepository, + chartBackgroundTracker, + chartAdjustmentBackgroundTracker, + threadRecommendQuestionBackgroundTracker, }); const dashboardService = new DashboardService({ projectService, @@ -196,19 +247,6 @@ export const initComponents = () => { wrenAIAdaptor, }); - // background trackers - const projectRecommendQuestionBackgroundTracker = - new ProjectRecommendQuestionBackgroundTracker({ - telemetry, - wrenAIAdaptor, - projectRepository, - }); - const threadRecommendQuestionBackgroundTracker = - new ThreadRecommendQuestionBackgroundTracker({ - telemetry, - wrenAIAdaptor, - threadRepository, - }); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ dashboardRepository, dashboardItemRepository, From e5258df9c87ae55058bf37e80ac20944db9970f2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 25 May 2026 19:55:24 +0530 Subject: [PATCH 0050/1087] updated question --- .../apollo/server/services/askingService.ts | 1 + .../server/services/askingTaskTracker.ts | 18 ++++++++++++++---- wren-ui/src/common.ts | 3 ++- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index c0e06972c3..bf5cd38eeb 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -528,6 +528,7 @@ export class AskingService implements IAskingService { this.chartBackgroundTracker.stop(); this.chartAdjustmentBackgroundTracker.stop(); this.threadRecommendQuestionBackgroundTracker.stop(); + this.askingTaskTracker.stopPolling(); } public async getThreadRecommendationQuestions( diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index 41438e7945..abb0c53585 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -49,6 +49,7 @@ export interface IAskingTaskTracker { getAskingResultById(id: number): Promise; cancelAskingTask(queryId: string): Promise; initialize(): Promise; + stopPolling(): void; bindThreadResponse( id: number, queryId: string, @@ -58,15 +59,15 @@ export interface IAskingTaskTracker { } export class AskingTaskTracker implements IAskingTaskTracker { - private readonly minPollDelay = 2000; - private readonly maxPollDelay = 8000; + private readonly minPollDelay = 5000; + private readonly maxPollDelay = 30000; private wrenAIAdaptor: IWrenAIAdaptor; private askingTaskRepository: IAskingTaskRepository; private trackedTasks: Map = new Map(); private trackedTasksById: Map = new Map(); private pollingInterval: number; private memoryRetentionTime: number; - private pollingIntervalId: NodeJS.Timeout; + private pollingIntervalId?: NodeJS.Timeout; private runningJobs = new Set(); private threadResponseRepository: IThreadResponseRepository; private viewRepository: IViewRepository; @@ -223,7 +224,11 @@ export class AskingTaskTracker implements IAskingTaskTracker { const taskRecords = await this.askingTaskRepository.findAll(); taskRecords.forEach((taskRecord) => { const detail = taskRecord.detail as AskResult | undefined; - if (!taskRecord.queryId || !detail || this.isTaskFinalized(detail.status)) { + if ( + !taskRecord.queryId || + !detail || + this.isTaskFinalized(detail.status) + ) { return; } @@ -241,6 +246,7 @@ export class AskingTaskTracker implements IAskingTaskTracker { public stopPolling(): void { if (this.pollingIntervalId) { clearInterval(this.pollingIntervalId); + this.pollingIntervalId = undefined; } } @@ -269,6 +275,10 @@ export class AskingTaskTracker implements IAskingTaskTracker { } private startPolling(): void { + if (this.pollingIntervalId) { + return; + } + this.pollingIntervalId = setInterval(() => { this.pollTasks(); }, this.pollingInterval); diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index e924374435..cdaee2242e 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -58,7 +58,7 @@ type Disposable = { }; type ReusableComponentGraph = { - askingTaskTracker?: Initializable; + askingTaskTracker?: Initializable & Disposable; askingService?: Initializable & Disposable; projectService?: Disposable; projectRecommendQuestionBackgroundTracker?: Initializable & Disposable; @@ -93,6 +93,7 @@ const isReusableComponentGraph = (graph?: ReusableComponentGraph): boolean => { const disposeComponentGraph = (graph?: ReusableComponentGraph): void => { const disposables = [ graph?.askingService, + graph?.askingTaskTracker, graph?.projectService, graph?.projectRecommendQuestionBackgroundTracker, graph?.threadRecommendQuestionBackgroundTracker, From 0181c11c1ae83ed7cbeec3a14321bf8a832857d9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 25 May 2026 20:37:58 +0530 Subject: [PATCH 0051/1087] updated question db --- .../src/pipelines/generation/utils/sql.py | 36 +++++++++++++++++++ .../server/services/askingTaskTracker.ts | 4 +-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 9d5f869426..ed5e11c07f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -288,6 +288,40 @@ def _rewrite_mssql_timestamp_casts(sql: str) -> str: return rewritten +def _rewrite_mssql_datepart_alias_references(sql: str) -> str: + datepart_alias_pattern = re.compile( + r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+\"([^\"]+)\"", + re.IGNORECASE, + ) + aliases: dict[str, str] = {} + + for match in datepart_alias_pattern.finditer(sql): + expression = match.group(1) + alias = match.group(4) + aliases[alias.lower()] = expression + + if not aliases: + return sql + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match) -> str: + body = match.group("body") + for alias, expression in aliases.items(): + body = re.sub( + rf'"{re.escape(alias)}"', + expression, + body, + flags=re.IGNORECASE, + ) + return f"{match.group(1)}{body}" + + return clause_pattern.sub(replace_clause, sql) + + def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = sql @@ -306,6 +340,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) + normalized = _rewrite_mssql_datepart_alias_references(normalized) return re.sub(r"\s+", " ", normalized).strip() @@ -540,6 +575,7 @@ async def _classify_generation_result( - DATEPART(YEAR, ) AS "year" - DATEPART(MONTH, ) AS "month" Then GROUP BY and ORDER BY the same year/month expressions. +- Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the DATEPART(...) expression instead. - For year bucketing, prefer DATEPART(YEAR, ). - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index abb0c53585..04c2cf8707 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -59,8 +59,8 @@ export interface IAskingTaskTracker { } export class AskingTaskTracker implements IAskingTaskTracker { - private readonly minPollDelay = 5000; - private readonly maxPollDelay = 30000; + private readonly minPollDelay = 10000; + private readonly maxPollDelay = 60000; private wrenAIAdaptor: IWrenAIAdaptor; private askingTaskRepository: IAskingTaskRepository; private trackedTasks: Map = new Map(); From 29e609f3c723f5f1483a3744c7b4076efcaa53d4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 26 May 2026 19:28:41 +0530 Subject: [PATCH 0052/1087] updated question db2 --- .../src/pipelines/generation/utils/sql.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ed5e11c07f..dea8b3a4c9 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -288,6 +288,16 @@ def _rewrite_mssql_timestamp_casts(sql: str) -> str: return rewritten +def _rewrite_mssql_to_unixtime(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + to_unixtime_pattern = re.compile( + rf"\bTO_UNIXTIME\(\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ) + + return to_unixtime_pattern.sub(lambda m: m.group(1), sql) + + def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+\"([^\"]+)\"", @@ -337,6 +347,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> flags=re.IGNORECASE, ) normalized = _replace_relative_getdate_calls(normalized, now) + normalized = _rewrite_mssql_to_unixtime(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -568,8 +579,9 @@ async def _classify_generation_result( ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. - Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). -- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. +- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. +- Do not calculate duration with DATEDIFF/date_diff/datediff. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section does not list a duration function, return the timestamps separately instead of inventing a duration function. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - DATEPART(YEAR, ) AS "year" From 3e136d022d622b72824352e507f3c06a3929011f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 26 May 2026 19:47:41 +0530 Subject: [PATCH 0053/1087] updated question db --- .../src/pipelines/generation/utils/sql.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index dea8b3a4c9..4b5e1d6e8f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -298,6 +298,27 @@ def _rewrite_mssql_to_unixtime(sql: str) -> str: return to_unixtime_pattern.sub(lambda m: m.group(1), sql) +def _rewrite_mssql_timestamp_subtraction(sql: str) -> str: + expression_pattern = r"((?:[^(),+\-]|\([^()]*\))+?)" + timestamp_subtraction_pattern = re.compile( + rf"{expression_pattern}\s*-\s*{expression_pattern}\s+AS\s+(\"[^\"]+\")", + re.IGNORECASE, + ) + + def replace_subtraction(match: re.Match[str]) -> str: + left = match.group(1).strip() + right = match.group(2).strip() + alias = match.group(3) + alias_text = alias.strip('"').lower() + + if not any(token in alias_text for token in ("duration", "turnaround")): + return match.group(0) + + return f"DATEDIFF('second', {right}, {left}) AS {alias}" + + return timestamp_subtraction_pattern.sub(replace_subtraction, sql) + + def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+\"([^\"]+)\"", @@ -348,6 +369,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> ) normalized = _replace_relative_getdate_calls(normalized, now) normalized = _rewrite_mssql_to_unixtime(normalized) + normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -581,7 +603,7 @@ async def _classify_generation_result( - Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. -- Do not calculate duration with DATEDIFF/date_diff/datediff. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section does not list a duration function, return the timestamps separately instead of inventing a duration function. +- Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - DATEPART(YEAR, ) AS "year" From 6989e2fb4664b5aacedf390029eca7f6633073a4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 26 May 2026 23:42:46 +0530 Subject: [PATCH 0054/1087] updated question data --- .../src/pipelines/generation/utils/sql.py | 23 ++++++++++++++++--- 1 file changed, 20 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 4b5e1d6e8f..158ff4c18e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -321,14 +321,14 @@ def replace_subtraction(match: re.Match[str]) -> str: def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( - r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+\"([^\"]+)\"", + r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", re.IGNORECASE, ) aliases: dict[str, str] = {} for match in datepart_alias_pattern.finditer(sql): expression = match.group(1) - alias = match.group(4) + alias = match.group(4) or match.group(5) or match.group(6) aliases[alias.lower()] = expression if not aliases: @@ -341,13 +341,30 @@ def _rewrite_mssql_datepart_alias_references(sql: str) -> str: def replace_clause(match: re.Match) -> str: body = match.group("body") + placeholders: dict[str, str] = {} for alias, expression in aliases.items(): + placeholder = f"__WREN_MSSQL_DATEPART_ALIAS_{len(placeholders)}__" + placeholders[placeholder] = expression body = re.sub( rf'"{re.escape(alias)}"', - expression, + placeholder, body, flags=re.IGNORECASE, ) + body = re.sub( + rf"\[{re.escape(alias)}\]", + placeholder, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"\b{re.escape(alias)}\b", + placeholder, + body, + flags=re.IGNORECASE, + ) + for placeholder, expression in placeholders.items(): + body = body.replace(placeholder, expression) return f"{match.group(1)}{body}" return clause_pattern.sub(replace_clause, sql) From 826aca22e1d77471ceeabdde3de91d8c1b1eaa1f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 26 May 2026 23:54:47 +0530 Subject: [PATCH 0055/1087] updated question databases --- .../apollo/server/services/queryService.ts | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 24d5d6dd50..e5c165e7cd 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -57,6 +57,46 @@ export interface ValidateResponse { message?: string; } +const rewriteMssqlDatepartAliasReferences = (sql: string): string => { + const aliases: Record = {}; + const aliasPattern = + /\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/gi; + + for (const match of sql.matchAll(aliasPattern)) { + const expression = match[1]; + const alias = match[4] || match[5] || match[6]; + aliases[alias.toLowerCase()] = expression; + } + + if (!Object.keys(aliases).length) { + return sql; + } + + const clausePattern = + /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; + + return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { + let body = groups?.body || ''; + const placeholders: Record = {}; + + Object.entries(aliases).forEach(([alias, expression]) => { + const placeholder = `__WREN_MSSQL_DATEPART_ALIAS_${Object.keys(placeholders).length}__`; + placeholders[placeholder] = expression; + const escapedAlias = alias.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + + body = body.replace(new RegExp(`"${escapedAlias}"`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\[${escapedAlias}\\]`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\b${escapedAlias}\\b`, 'gi'), placeholder); + }); + + Object.entries(placeholders).forEach(([placeholder, expression]) => { + body = body.replaceAll(placeholder, expression); + }); + + return `${clause}${body}`; + }); +}; + export interface IQueryService { preview( sql: string, @@ -85,6 +125,8 @@ const normalizePreviewSqlForIbis = ( return { sql, limit }; } + sql = rewriteMssqlDatepartAliasReferences(sql); + const topMatch = sql.match(/^\s*SELECT\s+(DISTINCT\s+)?TOP\s*\(?\s*(\d+)\s*\)?\s+/i); if (!topMatch) { return { sql, limit }; From d0f222ca38da4be87e0f211c632eb4d1aad7ff32 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 27 May 2026 01:25:46 +0530 Subject: [PATCH 0056/1087] updated question query --- .../src/apollo/server/adaptors/ibisAdaptor.ts | 54 ++++++++++++++++++- .../apollo/server/services/queryService.ts | 7 ++- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 3b16be0c34..b41db25972 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -164,6 +164,56 @@ const dataSourceUrlMap: Record = { [SupportedDataSource.DATABRICKS]: 'databricks', }; +const rewriteMssqlDatepartAliasReferences = ( + sql: string, + dataSource: DataSourceName, +): string => { + if (dataSource !== DataSourceName.MSSQL) { + return sql; + } + + sql = sql.replace(/\\"/g, '"'); + + const aliases: Record = {}; + const aliasPattern = + /\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:"([^"]+)"|\\+"([^"]+)\\+"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/gi; + + for (const match of sql.matchAll(aliasPattern)) { + const expression = match[1]; + const alias = match[4] || match[5] || match[6] || match[7]; + aliases[alias.toLowerCase()] = expression; + } + + if (!Object.keys(aliases).length) { + return sql; + } + + const clausePattern = + /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; + + return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { + let body = groups?.body || ''; + const placeholders: Record = {}; + + Object.entries(aliases).forEach(([alias, expression]) => { + const placeholder = `__WREN_MSSQL_DATEPART_ALIAS_${Object.keys(placeholders).length}__`; + placeholders[placeholder] = expression; + const escapedAlias = alias.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + + body = body.replace(new RegExp(`"${escapedAlias}"`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\\\+"${escapedAlias}\\\\+"`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\[${escapedAlias}\\]`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\b${escapedAlias}\\b`, 'gi'), placeholder); + }); + + Object.entries(placeholders).forEach(([placeholder, expression]) => { + body = body.replaceAll(placeholder, expression); + }); + + return `${clause}${body}`; + }); +}; + export interface TableResponse { tables: CompactTable[]; } @@ -270,7 +320,7 @@ export class IbisAdaptor implements IIbisAdaptor { public async getNativeSql(options: IbisDryPlanOptions): Promise { const { dataSource, mdl, sql } = options; const body = { - sql, + sql: rewriteMssqlDatepartAliasReferences(sql, dataSource), manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), }; try { @@ -294,6 +344,7 @@ export class IbisAdaptor implements IIbisAdaptor { options: IbisQueryOptions, ): Promise { const { dataSource, mdl } = options; + query = rewriteMssqlDatepartAliasReferences(query, dataSource); const connectionInfo = this.updateConnectionInfo(options.connectionInfo); const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo); const queryString = this.buildQueryString(options); @@ -338,6 +389,7 @@ export class IbisAdaptor implements IIbisAdaptor { options: IbisQueryOptions, ): Promise { const { dataSource, mdl } = options; + query = rewriteMssqlDatepartAliasReferences(query, dataSource); const connectionInfo = this.updateConnectionInfo(options.connectionInfo); const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo); const body = { diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index e5c165e7cd..b3460f988e 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -58,13 +58,15 @@ export interface ValidateResponse { } const rewriteMssqlDatepartAliasReferences = (sql: string): string => { + sql = sql.replace(/\\"/g, '"'); + const aliases: Record = {}; const aliasPattern = - /\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/gi; + /\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:"([^"]+)"|\\+"([^"]+)\\+"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/gi; for (const match of sql.matchAll(aliasPattern)) { const expression = match[1]; - const alias = match[4] || match[5] || match[6]; + const alias = match[4] || match[5] || match[6] || match[7]; aliases[alias.toLowerCase()] = expression; } @@ -85,6 +87,7 @@ const rewriteMssqlDatepartAliasReferences = (sql: string): string => { const escapedAlias = alias.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); body = body.replace(new RegExp(`"${escapedAlias}"`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\\\+"${escapedAlias}\\\\+"`, 'gi'), placeholder); body = body.replace(new RegExp(`\\[${escapedAlias}\\]`, 'gi'), placeholder); body = body.replace(new RegExp(`\\b${escapedAlias}\\b`, 'gi'), placeholder); }); From 13e308580328565863f1372ebb085faea3614ded Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 27 May 2026 01:43:26 +0530 Subject: [PATCH 0057/1087] updated question querys --- .../src/apollo/server/adaptors/ibisAdaptor.ts | 36 ++++++++++++++----- .../apollo/server/services/queryService.ts | 36 ++++++++++++++----- 2 files changed, 56 insertions(+), 16 deletions(-) diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index b41db25972..862f9f3732 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -175,14 +175,34 @@ const rewriteMssqlDatepartAliasReferences = ( sql = sql.replace(/\\"/g, '"'); const aliases: Record = {}; - const aliasPattern = - /\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:"([^"]+)"|\\+"([^"]+)\\+"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/gi; - - for (const match of sql.matchAll(aliasPattern)) { - const expression = match[1]; - const alias = match[4] || match[5] || match[6] || match[7]; - aliases[alias.toLowerCase()] = expression; - } + const aliasTargetPattern = + String.raw`(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))`; + const aliasPatterns = [ + new RegExp( + String.raw`\b(DATEPART\(\s*(?:YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b((?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b(DATE_PART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b(EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + ]; + + aliasPatterns.forEach((aliasPattern) => { + for (const match of sql.matchAll(aliasPattern)) { + const expression = match[1]; + const alias = match[3] || match[4] || match[5]; + aliases[alias.toLowerCase()] = expression; + } + }); if (!Object.keys(aliases).length) { return sql; diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index b3460f988e..e39e81362d 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -61,14 +61,34 @@ const rewriteMssqlDatepartAliasReferences = (sql: string): string => { sql = sql.replace(/\\"/g, '"'); const aliases: Record = {}; - const aliasPattern = - /\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:"([^"]+)"|\\+"([^"]+)\\+"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/gi; - - for (const match of sql.matchAll(aliasPattern)) { - const expression = match[1]; - const alias = match[4] || match[5] || match[6] || match[7]; - aliases[alias.toLowerCase()] = expression; - } + const aliasTargetPattern = + String.raw`(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))`; + const aliasPatterns = [ + new RegExp( + String.raw`\b(DATEPART\(\s*(?:YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b((?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b(DATE_PART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b(EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + ]; + + aliasPatterns.forEach((aliasPattern) => { + for (const match of sql.matchAll(aliasPattern)) { + const expression = match[1]; + const alias = match[3] || match[4] || match[5]; + aliases[alias.toLowerCase()] = expression; + } + }); if (!Object.keys(aliases).length) { return sql; From 32fd438d66ddf1d8d9af66b1d0a43b73358aae58 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 27 May 2026 20:29:40 +0530 Subject: [PATCH 0058/1087] updated question sql --- .../src/pipelines/generation/utils/sql.py | 107 ++++++++++++++++-- 1 file changed, 100 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 158ff4c18e..f85f5a76e4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -20,6 +20,93 @@ logger = logging.getLogger("wren-ai-service") +def _parse_sql_json_payload(payload: str) -> str | None: + try: + parsed_payload = orjson.loads(payload) + except orjson.JSONDecodeError: + return None + + if isinstance(parsed_payload, dict) and isinstance(parsed_payload.get("sql"), str): + return parsed_payload["sql"] + + return None + + +def _extract_json_object_with_sql(result: str) -> str | None: + sql_key_match = re.search(r'"sql"\s*:', result, flags=re.IGNORECASE) + if not sql_key_match: + return None + + start = result.rfind("{", 0, sql_key_match.start()) + if start == -1: + return None + + depth = 0 + in_string = False + escape_next = False + for index, char in enumerate(result[start:], start=start): + if escape_next: + escape_next = False + continue + if char == "\\" and in_string: + escape_next = True + continue + if char == '"': + in_string = not in_string + continue + if in_string: + continue + if char == "{": + depth += 1 + elif char == "}": + depth -= 1 + if depth == 0: + return result[start : index + 1] + + return None + + +def _extract_select_statement(result: str) -> str | None: + sql_match = re.search(r"\b(?:WITH|SELECT)\b", result, flags=re.IGNORECASE) + if not sql_match: + return None + + statement = result[sql_match.start() :].strip() + semicolon_index = statement.find(";") + if semicolon_index >= 0: + statement = statement[:semicolon_index] + + return statement + + +def extract_sql_generation_result(result: str) -> str: + fenced_blocks = re.findall( + r"```(?:json|sql)?\s*(.*?)```", result, flags=re.IGNORECASE | re.DOTALL + ) + for block in fenced_blocks: + if sql := _parse_sql_json_payload(block.strip()): + return clean_generation_result(sql) + if sql := _extract_select_statement(block): + return clean_generation_result(sql) + + cleaned_result = clean_generation_result(result) + if sql := _parse_sql_json_payload(cleaned_result): + return clean_generation_result(sql) + + if json_payload := _extract_json_object_with_sql(result): + if sql := _parse_sql_json_payload(json_payload): + return clean_generation_result(sql) + + if sql := _extract_select_statement(result): + return clean_generation_result(sql) + + return cleaned_result + + +def is_select_statement(sql: str) -> bool: + return bool(re.match(r"^\s*(?:WITH|SELECT)\b", sql, flags=re.IGNORECASE)) + + def normalize_data_source(data_source: str | None) -> str: normalized = (data_source or "").strip().upper().replace("-", "_").replace( " ", "_" @@ -414,18 +501,24 @@ async def run( allow_data_preview: bool = False, ) -> dict: try: - cleaned_generation_result = clean_generation_result(replies[0]) - - # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' - if cleaned_generation_result.startswith("{"): - cleaned_generation_result = orjson.loads(cleaned_generation_result)[ - "sql" - ] + cleaned_generation_result = extract_sql_generation_result(replies[0]) cleaned_generation_result = normalize_generation_result_sql( cleaned_generation_result, data_source=data_source ) + if not is_select_statement(cleaned_generation_result): + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "DRY_RUN", + "error": "Generated response did not contain a SQL SELECT statement.", + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, From a7e0831ecfc347126252fe383121a775fd01301a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 27 May 2026 20:43:43 +0530 Subject: [PATCH 0059/1087] updated question sqll --- .../src/pipelines/generation/sql_diagnosis.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py index 3f22b9d512..129536cabe 100644 --- a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py +++ b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py @@ -92,7 +92,17 @@ async def generate_sql_diagnosis( async def post_process( generate_sql_diagnosis: dict, ) -> str: - return orjson.loads(generate_sql_diagnosis.get("replies")[0]) + reply = (generate_sql_diagnosis.get("replies") or [""])[0] + try: + parsed_reply = orjson.loads(reply) + except orjson.JSONDecodeError: + return {"reasoning": reply.strip()} + + if isinstance(parsed_reply, dict): + reasoning = parsed_reply.get("reasoning", "") + return {"reasoning": reasoning if isinstance(reasoning, str) else ""} + + return {"reasoning": str(parsed_reply)} ## End of Pipeline From 872e6da54fd9929052537670421a9ee293651159 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 27 May 2026 21:02:46 +0530 Subject: [PATCH 0060/1087] updated question sq --- wren-ai-service/src/providers/llm/litellm.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index 584dc7e837..1b32bd9f05 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -19,6 +19,12 @@ from src.utils import extract_braces_content, remove_trailing_slash +def normalize_litellm_model_name(model: str, api_base: Optional[str] = None) -> str: + if api_base and "/" not in model: + return f"openai/{model}" + return model + + @provider("litellm_llm") class LitellmLLMProvider(LLMProvider): def __init__( @@ -36,7 +42,7 @@ def __init__( fallback_testing: bool = False, **_, ): - self._model = model + self._model = normalize_litellm_model_name(model, api_base) # TODO: remove _api_key, _api_base, _api_version in the future, as it is not used in litellm self._api_key = os.getenv(api_key_name) if api_key_name else None self._api_base = remove_trailing_slash(api_base) if api_base else None From 725565d97ae0a26da8c86191dfa8cb89447d39b1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 27 May 2026 21:47:38 +0530 Subject: [PATCH 0061/1087] updated chart --- .../generation/followup_sql_generation.py | 9 +++++++++ .../src/pipelines/generation/sql_correction.py | 13 +++++++++++++ .../src/pipelines/generation/sql_generation.py | 9 +++++++++ .../src/pipelines/generation/utils/sql.py | 17 +++++++++++++++++ 4 files changed, 48 insertions(+) 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 6e295933cf..55e5ddbd83 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -15,6 +15,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -42,6 +43,13 @@ {{ document }} {% endfor %} +### VALID TABLE NAMES ### +Only use these exact table names from the schema. Do not invent, rename, singularize, +pluralize, or add catalog/schema prefixes unless the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -108,6 +116,7 @@ def prompt( query=query, data_source=data_source, documents=documents, + valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f57a1e7b45..32babb4850 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_names, get_sql_generation_model_kwargs, get_text_to_sql_rules, ) @@ -68,6 +69,17 @@ def get_sql_correction_system_prompt( {% endfor %} {% endif %} +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only use these exact table names from the schema. If the invalid SQL references a +table not listed here, replace it with the closest listed table only when the schema +clearly supports the user's request. Do not invent, rename, singularize, pluralize, +or add catalog/schema prefixes unless the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + {% if sql_functions %} ### SQL FUNCTIONS ### {% for function in sql_functions %} @@ -106,6 +118,7 @@ def prompt( _prompt = prompt_builder.run( data_source=data_source, documents=documents, + valid_table_names=construct_valid_table_names(documents), invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 58c89442c7..9c247af46e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -36,6 +37,13 @@ {{ document }} {% endfor %} +### VALID TABLE NAMES ### +Only use these exact table names from the schema. Do not invent, rename, singularize, +pluralize, or add catalog/schema prefixes unless the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -104,6 +112,7 @@ def prompt( query=query, data_source=data_source, documents=documents, + valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index f85f5a76e4..75b862a5c2 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1091,6 +1091,23 @@ def construct_instructions( return _instructions +def construct_valid_table_names(documents: list[Any] | None = None) -> list[str]: + table_names = [] + for document in documents or []: + content = getattr(document, "content", document) + if not isinstance(content, str): + continue + + for match in re.finditer( + r"\bCREATE\s+TABLE\s+([`\"\[]?)([A-Za-z_][A-Za-z0-9_.$]*)\1", + content, + flags=re.IGNORECASE, + ): + table_names.append(match.group(2)) + + return sorted(set(table_names)) + + def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: From 5794c50a76b5aa8aa199468b623a6de2266bc742 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 00:43:32 +0530 Subject: [PATCH 0062/1087] updated chart changes --- .../pipelines/generation/sql_generation.py | 6 ++ .../retrieval/db_schema_retrieval.py | 60 ++++++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 41 +++++++++++++ 3 files changed, 106 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 9c247af46e..efc6d93fa9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -83,6 +83,12 @@ ### QUESTION ### User's Question: {{ query }} +### PCB ANALYTICS TERM MAPPING ### +If the user asks about PCB repair trends, repair volume, repair counts, debug hours, +turnaround time, resolved entries, or failure category, map those business terms to +the closest explicit table and column names in DATABASE SCHEMA and VALID TABLE NAMES. +Do not answer with general guidance when a SQL aggregation, comparison, trend, or chart is requested. + {% if sql_generation_reasoning %} ### REASONING PLAN ### {{ sql_generation_reasoning }} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index de5e3051d6..cedf8e0163 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -122,6 +122,35 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline +def expand_business_terms_for_retrieval(query: str) -> str: + normalized = (query or "").lower() + pcb_terms = { + "pcb", + "repair", + "debug", + "turnaround", + "failure", + "resolved", + "trend", + "volume", + "count", + "average", + "chart", + } + if not any(term in normalized for term in pcb_terms): + return query + + return "\n".join( + [ + query, + "PCB repair debug analytics aliases:", + "repair trends repair volume repair counts debug entries debug fixes", + "average debug hours turnaround time resolved entries failure category failure code", + "monthly trend quarter grouped by month bar chart line chart", + ] + ) + + @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: @@ -131,6 +160,7 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> d previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query + query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -239,7 +269,35 @@ async def dbschema_retrieval( ) return fallback_results["documents"] - return [] + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + logger.info( + "No table-description matches found; falling back to deployed schema for project_id %s", + project_id, + ) + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + if results.get("documents") or not project_id: + return results.get("documents", []) + + fallback_results = await dbschema_retriever.run( + query_embedding=[], + filters={ + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + }, + ) + return fallback_results.get("documents", []) @observe() diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cb96d5dab6..8b8f92a112 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -157,6 +157,33 @@ def _is_greeting_query(self, query: str) -> bool: } return normalized in greeting_patterns + def _is_data_analysis_query(self, query: str) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + analysis_terms = { + "average", + "avg", + "bar chart", + "chart", + "compare", + "count", + "debug", + "failure", + "group", + "grouped", + "monthly", + "pcb", + "quarter", + "repair", + "resolved", + "trend", + "turnaround", + "volume", + } + return any(term in normalized for term in analysis_terms) + async def _run_with_timeout(self, label: str, coroutine): try: return await asyncio.wait_for( @@ -338,6 +365,17 @@ async def ask( ) intent_reasoning = intent_classification_result.get("reasoning") + if intent in {"GENERAL", "MISLEADING_QUERY"} and ( + self._is_data_analysis_query(user_query) + or self._is_data_analysis_query(rephrased_question or "") + ): + logger.info( + "Overriding intent %s to TEXT_TO_SQL for analytics query: %s", + intent, + user_query, + ) + intent = "TEXT_TO_SQL" + if rephrased_question: user_query = rephrased_question @@ -461,6 +499,9 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] + logger.info( + "Retrieved tables for query_id %s: %s", query_id, table_names + ) if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") From 9131578e0e80588a03a228f98522fe36ad295fe8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 00:52:48 +0530 Subject: [PATCH 0063/1087] updated chart change --- .../retrieval/db_schema_retrieval.py | 70 ++++++++++++++++++- 1 file changed, 68 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index cedf8e0163..7e191773d2 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -151,6 +151,60 @@ def expand_business_terms_for_retrieval(query: str) -> str: ) +def candidate_pcb_table_names(query: str) -> list[str]: + normalized = (query or "").lower() + if not any( + term in normalized + for term in [ + "pcb", + "repair", + "debug", + "turnaround", + "failure", + "resolved", + "trend", + "volume", + "count", + "average", + "chart", + ] + ): + return [] + + names = [ + "dbo.DebugEntries", + "DebugEntries", + "dbo.DebugFixLogs", + "DebugFixLogs", + "dbo.failure_patterns", + "failure_patterns", + "dbo.pcb_tags", + "pcb_tags", + "dbo.batch_records", + "batch_records", + "dbo.risk_scores", + "risk_scores", + ] + + if "failure" in normalized: + names = [ + "dbo.failure_patterns", + "failure_patterns", + "dbo.pcb_tags", + "pcb_tags", + ] + names + + if "turnaround" in normalized or "debug" in normalized or "resolved" in normalized: + names = [ + "dbo.DebugEntries", + "DebugEntries", + "dbo.DebugFixLogs", + "DebugFixLogs", + ] + names + + return list(dict.fromkeys(names)) + + @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: @@ -226,7 +280,7 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, project_id: str, dbschema_retriever: Any + query: str, table_retrieval: dict, project_id: str, dbschema_retriever: Any ) -> list[Document]: tables = table_retrieval.get("documents", []) table_names = [] @@ -269,20 +323,32 @@ async def dbschema_retrieval( ) return fallback_results["documents"] + candidate_names = candidate_pcb_table_names(query) filters = { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, ], } + if candidate_names: + filters["conditions"].append( + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": table_name} + for table_name in candidate_names + ], + } + ) if project_id: filters["conditions"].append( {"field": "project_id", "operator": "==", "value": project_id} ) logger.info( - "No table-description matches found; falling back to deployed schema for project_id %s", + "No table-description matches found; falling back to deployed schema for project_id %s with candidate tables: %s", project_id, + candidate_names, ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) if results.get("documents") or not project_id: From 53d2c3c51a33cea908c4933bdba160ab2f704502 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 01:04:49 +0530 Subject: [PATCH 0064/1087] updated chart changed --- .../retrieval/db_schema_retrieval.py | 72 ++----------------- 1 file changed, 5 insertions(+), 67 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 7e191773d2..2c6a6f322d 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -134,8 +134,12 @@ def expand_business_terms_for_retrieval(query: str) -> str: "trend", "volume", "count", + "counts", "average", + "avg", "chart", + "month", + "monthly", } if not any(term in normalized for term in pcb_terms): return query @@ -151,60 +155,6 @@ def expand_business_terms_for_retrieval(query: str) -> str: ) -def candidate_pcb_table_names(query: str) -> list[str]: - normalized = (query or "").lower() - if not any( - term in normalized - for term in [ - "pcb", - "repair", - "debug", - "turnaround", - "failure", - "resolved", - "trend", - "volume", - "count", - "average", - "chart", - ] - ): - return [] - - names = [ - "dbo.DebugEntries", - "DebugEntries", - "dbo.DebugFixLogs", - "DebugFixLogs", - "dbo.failure_patterns", - "failure_patterns", - "dbo.pcb_tags", - "pcb_tags", - "dbo.batch_records", - "batch_records", - "dbo.risk_scores", - "risk_scores", - ] - - if "failure" in normalized: - names = [ - "dbo.failure_patterns", - "failure_patterns", - "dbo.pcb_tags", - "pcb_tags", - ] + names - - if "turnaround" in normalized or "debug" in normalized or "resolved" in normalized: - names = [ - "dbo.DebugEntries", - "DebugEntries", - "dbo.DebugFixLogs", - "DebugFixLogs", - ] + names - - return list(dict.fromkeys(names)) - - @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: @@ -323,32 +273,20 @@ async def dbschema_retrieval( ) return fallback_results["documents"] - candidate_names = candidate_pcb_table_names(query) filters = { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, ], } - if candidate_names: - filters["conditions"].append( - { - "operator": "OR", - "conditions": [ - {"field": "name", "operator": "==", "value": table_name} - for table_name in candidate_names - ], - } - ) if project_id: filters["conditions"].append( {"field": "project_id", "operator": "==", "value": project_id} ) logger.info( - "No table-description matches found; falling back to deployed schema for project_id %s with candidate tables: %s", + "No table-description matches found; falling back to all deployed schema for project_id %s", project_id, - candidate_names, ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) if results.get("documents") or not project_id: From f81d301bf354d4c9d29a0a41feaaba792a43871e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 17:51:57 +0530 Subject: [PATCH 0065/1087] updated month --- .../src/pipelines/generation/utils/sql.py | 99 ++++++++++++++++++- 1 file changed, 98 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 75b862a5c2..ad44031008 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -406,6 +406,93 @@ def replace_subtraction(match: re.Match[str]) -> str: return timestamp_subtraction_pattern.sub(replace_subtraction, sql) +def _infer_mssql_timestamp_expression(sql: str) -> str | None: + timestamp_column_pattern = re.compile( + r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|opened_at|closed_at|completed_at|resolved_at)")', + re.IGNORECASE, + ) + if match := timestamp_column_pattern.search(sql): + return match.group(0) + + table_pattern = re.compile( + r'\bFROM\s+("[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)', + re.IGNORECASE, + ) + if match := table_pattern.search(sql): + table_name = match.group(1) + if any( + token in table_name.strip('"[]').lower() + for token in ("repair", "ticket", "debug", "event", "log") + ): + return f'{table_name}."created_at"' + + return None + + +def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: + timestamp_expression = _infer_mssql_timestamp_expression(sql) + if not timestamp_expression: + return sql + + invented_date_identifier_pattern = re.compile( + r'(? str: + timestamp_expression = _infer_mssql_timestamp_expression(sql) + if not timestamp_expression: + return sql + + bucket_expressions = { + "year": f"DATEPART(YEAR, {timestamp_expression})", + "month": f"DATEPART(MONTH, {timestamp_expression})", + "day": f"DATEPART(DAY, {timestamp_expression})", + } + rewritten = sql + + for bucket, expression in bucket_expressions.items(): + select_identifier_pattern = re.compile( + rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', + re.IGNORECASE, + ) + + def replace_select_identifier(match: re.Match[str]) -> str: + prefix = match.group("prefix") + suffix = match.group("suffix") + return f'{prefix}{expression} AS "{bucket}"{suffix}' + + rewritten = select_identifier_pattern.sub( + replace_select_identifier, rewritten + ) + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match[str]) -> str: + body = match.group("body") + for bucket, expression in bucket_expressions.items(): + body = re.sub( + rf'"{bucket}"', + expression, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"\[{bucket}\]", + expression, + body, + flags=re.IGNORECASE, + ) + return f"{match.group(1)}{body}" + + return clause_pattern.sub(replace_clause, rewritten) + + def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", @@ -445,7 +532,7 @@ def replace_clause(match: re.Match) -> str: flags=re.IGNORECASE, ) body = re.sub( - rf"\b{re.escape(alias)}\b", + rf"(? normalized = _rewrite_mssql_to_unixtime(normalized) normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) + normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) @@ -712,6 +801,10 @@ async def _classify_generation_result( - The target database is MSSQL. - Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. +- DO NOT use JSON extraction functions or operators such as JSON_VALUE, JSON_QUERY, JSON_EXTRACT, JSON_EXTRACT_SCALAR, JSON_EXTRACT_ARRAY, json_value, json_extract, ->, or ->>. The MSSQL Wren/Ibis runtime does not support them. +- If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. +- Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. +- For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. @@ -1042,6 +1135,10 @@ def get_sql_generation_system_prompt( 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. 4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +6. YOU MUST ONLY use table names and column names that are explicitly present in the DATABASE SCHEMA or VALID TABLE NAMES sections. +7. NEVER invent generic table names such as repair_logs, repair_log, sales_data, orders, users, tickets, events, or transactions unless that exact table name is present in the DATABASE SCHEMA or VALID TABLE NAMES sections. +8. If the user asks about a business concept such as repairs, PCB, cost, turnaround time, or volume, map it to the closest explicit table and column names from the provided schema. Do not create a new table name from the business concept. +9. Do not prefix table names with catalog or schema names unless the DATABASE SCHEMA or VALID TABLE NAMES section shows the table name with that exact prefix. {text_to_sql_rules} From 98cf974280868a00f726457005c08eea88261382 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 18:13:00 +0530 Subject: [PATCH 0066/1087] updated month log --- wren-ai-service/src/pipelines/generation/utils/sql.py | 8 +++++--- wren-ai-service/src/pipelines/retrieval/sql_executor.py | 8 ++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ad44031008..ac654229c5 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -15,7 +15,6 @@ ) from src.core.provider import LLMProvider from src.pipelines.retrieval.sql_knowledge import SqlKnowledge -from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -408,7 +407,7 @@ def replace_subtraction(match: re.Match[str]) -> str: def _infer_mssql_timestamp_expression(sql: str) -> str | None: timestamp_column_pattern = re.compile( - r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|opened_at|closed_at|completed_at|resolved_at)")', + r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at)")', re.IGNORECASE, ) if match := timestamp_column_pattern.search(sql): @@ -420,6 +419,9 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: ) if match := table_pattern.search(sql): table_name = match.group(1) + normalized_table_name = table_name.strip('"[]').lower() + if "report" in normalized_table_name: + return f'{table_name}."generated_at"' if any( token in table_name.strip('"[]').lower() for token in ("repair", "ticket", "debug", "event", "log") @@ -1206,7 +1208,7 @@ def construct_valid_table_names(documents: list[Any] | None = None) -> list[str] def construct_ask_history_messages( - histories: list[AskHistory] | list[dict], + histories: list[Any] | list[dict], ) -> list[ChatMessage]: messages = [] for history in histories: diff --git a/wren-ai-service/src/pipelines/retrieval/sql_executor.py b/wren-ai-service/src/pipelines/retrieval/sql_executor.py index b41151469f..3530154d7f 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_executor.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_executor.py @@ -10,14 +10,16 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import normalize_generation_result_sql logger = logging.getLogger("wren-ai-service") @component class DataFetcher: - def __init__(self, engine: Engine): + def __init__(self, engine: Engine, data_source: str | None = None): self._engine = engine + self._data_source = data_source @component.output_types( results=Optional[Dict[str, Any]], @@ -28,6 +30,7 @@ async def run( project_id: str | None = None, limit: int = 500, ): + sql = normalize_generation_result_sql(sql, data_source=self._data_source) async with aiohttp.ClientSession() as session: _, data, addition = await self._engine.execute_sql( sql, @@ -64,10 +67,11 @@ class SQLExecutor(BasicPipeline): def __init__( self, engine: Engine, + data_source: str | None = None, **kwargs, ): self._components = { - "data_fetcher": DataFetcher(engine=engine), + "data_fetcher": DataFetcher(engine=engine, data_source=data_source), } super().__init__( From 66738254cbc663260efac7070ac9d34758848de5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 18:49:22 +0530 Subject: [PATCH 0067/1087] updated sql pair --- .../src/pipelines/generation/utils/sql.py | 472 +----------------- .../src/pipelines/retrieval/sql_executor.py | 2 +- .../src/pipelines/sql_normalizer.py | 469 +++++++++++++++++ 3 files changed, 474 insertions(+), 469 deletions(-) create mode 100644 wren-ai-service/src/pipelines/sql_normalizer.py diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ac654229c5..3851867f11 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,6 +1,5 @@ import logging import re -from datetime import datetime, timedelta from typing import Any, Dict, List import aiohttp @@ -14,6 +13,10 @@ clean_generation_result, ) from src.core.provider import LLMProvider +from src.pipelines.sql_normalizer import ( + normalize_data_source, + normalize_generation_result_sql, +) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge logger = logging.getLogger("wren-ai-service") @@ -106,473 +109,6 @@ def is_select_statement(sql: str) -> bool: return bool(re.match(r"^\s*(?:WITH|SELECT)\b", sql, flags=re.IGNORECASE)) -def normalize_data_source(data_source: str | None) -> str: - normalized = (data_source or "").strip().upper().replace("-", "_").replace( - " ", "_" - ) - if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: - return "MSSQL" - return normalized - - -def _format_timestamp_literal(value: datetime) -> str: - return value.strftime("'%Y-%m-%d %H:%M:%S'") - - -def _add_months(value: datetime, months: int) -> datetime: - month_index = value.month - 1 + months - year = value.year + month_index // 12 - month = month_index % 12 + 1 - day = min( - value.day, - [ - 31, - 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28, - 31, - 30, - 31, - 30, - 31, - 31, - 30, - 31, - 30, - 31, - ][month - 1], - ) - return value.replace(year=year, month=month, day=day) - - -def _start_of_month(value: datetime) -> datetime: - return value.replace(day=1, hour=0, minute=0, second=0, microsecond=0) - - -def _replace_relative_getdate_calls(sql: str, now: datetime) -> str: - def replace_month_offset(match: re.Match[str]) -> str: - months = int(match.group(1)) - return _format_timestamp_literal(_add_months(now, months)) - - def replace_year_offset(match: re.Match[str]) -> str: - years = int(match.group(1)) - return _format_timestamp_literal(_add_months(now, years * 12)) - - def replace_day_offset(match: re.Match[str]) -> str: - days = int(match.group(1)) - return _format_timestamp_literal(now + timedelta(days=days)) - - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_month_offset, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*year\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_year_offset, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*day\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_day_offset, - sql, - flags=re.IGNORECASE, - ) - - current_month_start = _format_timestamp_literal(_start_of_month(now)) - previous_month_start = _format_timestamp_literal( - _start_of_month(_add_months(now, -1)) - ) - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*,\s*0\s*\)", - current_month_start, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*-\s*1\s*,\s*0\s*\)", - previous_month_start, - sql, - flags=re.IGNORECASE, - ) - return sql - - -def _rewrite_mssql_bucket_functions(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - - sql = re.sub( - rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: ( - f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" - ), - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"DATEPART(YEAR, {m.group(1)})", - sql, - flags=re.IGNORECASE, - ) - return sql - - -def _rewrite_temporal_bucket_functions(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - replacements = [ - ( - re.compile( - rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(DAY, {m.group(1)})", - ), - ( - re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(YEAR, {m.group(1)})", - ), - ( - re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(MONTH, {m.group(1)})", - ), - ( - re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(DAY, {m.group(1)})", - ), - ( - re.compile( - rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", - ), - ( - re.compile( - rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", - ), - ( - re.compile( - rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", - ), - ( - re.compile( - rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(DAY, {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(DAY, {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"DATEPART(DAY, {m.group(1)})", - ), - ] - - rewritten = sql - for pattern, replacement in replacements: - rewritten = pattern.sub(replacement, rewritten) - - return rewritten - - -def _rewrite_mssql_timestamp_casts(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - timestamp_function_pattern = re.compile( - rf"\bTO_TIMESTAMP(?:_(?:MILLIS|SECONDS|MICROS|NANOS))?\(\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ) - timestamp_cast_pattern = re.compile( - r"CAST\(\s*((?:[^()]|\([^()]*\))+?)\s+AS\s+TIMESTAMP\s*\)", - re.IGNORECASE, - ) - - rewritten = timestamp_function_pattern.sub( - lambda m: f"CAST({m.group(1)} AS DATETIME)", sql - ) - rewritten = timestamp_cast_pattern.sub( - lambda m: f"CAST({m.group(1)} AS DATETIME)", rewritten - ) - return rewritten - - -def _rewrite_mssql_to_unixtime(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - to_unixtime_pattern = re.compile( - rf"\bTO_UNIXTIME\(\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ) - - return to_unixtime_pattern.sub(lambda m: m.group(1), sql) - - -def _rewrite_mssql_timestamp_subtraction(sql: str) -> str: - expression_pattern = r"((?:[^(),+\-]|\([^()]*\))+?)" - timestamp_subtraction_pattern = re.compile( - rf"{expression_pattern}\s*-\s*{expression_pattern}\s+AS\s+(\"[^\"]+\")", - re.IGNORECASE, - ) - - def replace_subtraction(match: re.Match[str]) -> str: - left = match.group(1).strip() - right = match.group(2).strip() - alias = match.group(3) - alias_text = alias.strip('"').lower() - - if not any(token in alias_text for token in ("duration", "turnaround")): - return match.group(0) - - return f"DATEDIFF('second', {right}, {left}) AS {alias}" - - return timestamp_subtraction_pattern.sub(replace_subtraction, sql) - - -def _infer_mssql_timestamp_expression(sql: str) -> str | None: - timestamp_column_pattern = re.compile( - r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at)")', - re.IGNORECASE, - ) - if match := timestamp_column_pattern.search(sql): - return match.group(0) - - table_pattern = re.compile( - r'\bFROM\s+("[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)', - re.IGNORECASE, - ) - if match := table_pattern.search(sql): - table_name = match.group(1) - normalized_table_name = table_name.strip('"[]').lower() - if "report" in normalized_table_name: - return f'{table_name}."generated_at"' - if any( - token in table_name.strip('"[]').lower() - for token in ("repair", "ticket", "debug", "event", "log") - ): - return f'{table_name}."created_at"' - - return None - - -def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: - timestamp_expression = _infer_mssql_timestamp_expression(sql) - if not timestamp_expression: - return sql - - invented_date_identifier_pattern = re.compile( - r'(? str: - timestamp_expression = _infer_mssql_timestamp_expression(sql) - if not timestamp_expression: - return sql - - bucket_expressions = { - "year": f"DATEPART(YEAR, {timestamp_expression})", - "month": f"DATEPART(MONTH, {timestamp_expression})", - "day": f"DATEPART(DAY, {timestamp_expression})", - } - rewritten = sql - - for bucket, expression in bucket_expressions.items(): - select_identifier_pattern = re.compile( - rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', - re.IGNORECASE, - ) - - def replace_select_identifier(match: re.Match[str]) -> str: - prefix = match.group("prefix") - suffix = match.group("suffix") - return f'{prefix}{expression} AS "{bucket}"{suffix}' - - rewritten = select_identifier_pattern.sub( - replace_select_identifier, rewritten - ) - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match[str]) -> str: - body = match.group("body") - for bucket, expression in bucket_expressions.items(): - body = re.sub( - rf'"{bucket}"', - expression, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"\[{bucket}\]", - expression, - body, - flags=re.IGNORECASE, - ) - return f"{match.group(1)}{body}" - - return clause_pattern.sub(replace_clause, rewritten) - - -def _rewrite_mssql_datepart_alias_references(sql: str) -> str: - datepart_alias_pattern = re.compile( - r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", - re.IGNORECASE, - ) - aliases: dict[str, str] = {} - - for match in datepart_alias_pattern.finditer(sql): - expression = match.group(1) - alias = match.group(4) or match.group(5) or match.group(6) - aliases[alias.lower()] = expression - - if not aliases: - return sql - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match) -> str: - body = match.group("body") - placeholders: dict[str, str] = {} - for alias, expression in aliases.items(): - placeholder = f"__WREN_MSSQL_DATEPART_ALIAS_{len(placeholders)}__" - placeholders[placeholder] = expression - body = re.sub( - rf'"{re.escape(alias)}"', - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"\[{re.escape(alias)}\]", - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"(? str: - normalized = sql - - if normalize_data_source(data_source) == "MSSQL": - now = datetime.now() - normalized = re.sub( - r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE - ) - normalized = re.sub( - r"CAST\(\s*('(?:[^']|'')*')\s+AS\s+DATETIME(?:2|OFFSET)\s*\)", - r"\1", - normalized, - flags=re.IGNORECASE, - ) - normalized = _replace_relative_getdate_calls(normalized, now) - normalized = _rewrite_mssql_to_unixtime(normalized) - normalized = _rewrite_mssql_timestamp_subtraction(normalized) - normalized = _rewrite_mssql_timestamp_casts(normalized) - normalized = _rewrite_mssql_invented_date_identifiers(normalized) - normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) - normalized = _rewrite_mssql_bucket_functions(normalized) - normalized = _rewrite_temporal_bucket_functions(normalized) - normalized = _rewrite_mssql_datepart_alias_references(normalized) - - return re.sub(r"\s+", " ", normalized).strip() - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): diff --git a/wren-ai-service/src/pipelines/retrieval/sql_executor.py b/wren-ai-service/src/pipelines/retrieval/sql_executor.py index 3530154d7f..b6f50b7826 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_executor.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_executor.py @@ -10,7 +10,7 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import normalize_generation_result_sql +from src.pipelines.sql_normalizer import normalize_generation_result_sql logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/sql_normalizer.py b/wren-ai-service/src/pipelines/sql_normalizer.py new file mode 100644 index 0000000000..b8071e268c --- /dev/null +++ b/wren-ai-service/src/pipelines/sql_normalizer.py @@ -0,0 +1,469 @@ +import re +from datetime import datetime, timedelta + + +def normalize_data_source(data_source: str | None) -> str: + normalized = (data_source or "").strip().upper().replace("-", "_").replace( + " ", "_" + ) + if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: + return "MSSQL" + return normalized + + +def _format_timestamp_literal(value: datetime) -> str: + return value.strftime("'%Y-%m-%d %H:%M:%S'") + + +def _add_months(value: datetime, months: int) -> datetime: + month_index = value.month - 1 + months + year = value.year + month_index // 12 + month = month_index % 12 + 1 + day = min( + value.day, + [ + 31, + 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28, + 31, + 30, + 31, + 30, + 31, + 31, + 30, + 31, + 30, + 31, + ][month - 1], + ) + return value.replace(year=year, month=month, day=day) + + +def _start_of_month(value: datetime) -> datetime: + return value.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + + +def _replace_relative_getdate_calls(sql: str, now: datetime) -> str: + def replace_month_offset(match: re.Match[str]) -> str: + months = int(match.group(1)) + return _format_timestamp_literal(_add_months(now, months)) + + def replace_year_offset(match: re.Match[str]) -> str: + years = int(match.group(1)) + return _format_timestamp_literal(_add_months(now, years * 12)) + + def replace_day_offset(match: re.Match[str]) -> str: + days = int(match.group(1)) + return _format_timestamp_literal(now + timedelta(days=days)) + + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_month_offset, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*year\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_year_offset, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*day\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_day_offset, + sql, + flags=re.IGNORECASE, + ) + + current_month_start = _format_timestamp_literal(_start_of_month(now)) + previous_month_start = _format_timestamp_literal( + _start_of_month(_add_months(now, -1)) + ) + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*,\s*0\s*\)", + current_month_start, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*-\s*1\s*,\s*0\s*\)", + previous_month_start, + sql, + flags=re.IGNORECASE, + ) + return sql + + +def _rewrite_mssql_bucket_functions(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + + sql = re.sub( + rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", + lambda m: ( + f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" + ), + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", + lambda m: f"DATEPART(YEAR, {m.group(1)})", + sql, + flags=re.IGNORECASE, + ) + return sql + + +def _rewrite_temporal_bucket_functions(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + replacements = [ + ( + re.compile( + rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile( + rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ] + + rewritten = sql + for pattern, replacement in replacements: + rewritten = pattern.sub(replacement, rewritten) + + return rewritten + + +def _rewrite_mssql_timestamp_casts(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + timestamp_function_pattern = re.compile( + rf"\bTO_TIMESTAMP(?:_(?:MILLIS|SECONDS|MICROS|NANOS))?\(\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ) + timestamp_cast_pattern = re.compile( + r"CAST\(\s*((?:[^()]|\([^()]*\))+?)\s+AS\s+TIMESTAMP\s*\)", + re.IGNORECASE, + ) + + rewritten = timestamp_function_pattern.sub( + lambda m: f"CAST({m.group(1)} AS DATETIME)", sql + ) + rewritten = timestamp_cast_pattern.sub( + lambda m: f"CAST({m.group(1)} AS DATETIME)", rewritten + ) + return rewritten + + +def _rewrite_mssql_to_unixtime(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + to_unixtime_pattern = re.compile( + rf"\bTO_UNIXTIME\(\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ) + + return to_unixtime_pattern.sub(lambda m: m.group(1), sql) + + +def _rewrite_mssql_timestamp_subtraction(sql: str) -> str: + expression_pattern = r"((?:[^(),+\-]|\([^()]*\))+?)" + timestamp_subtraction_pattern = re.compile( + rf"{expression_pattern}\s*-\s*{expression_pattern}\s+AS\s+(\"[^\"]+\")", + re.IGNORECASE, + ) + + def replace_subtraction(match: re.Match[str]) -> str: + left = match.group(1).strip() + right = match.group(2).strip() + alias = match.group(3) + alias_text = alias.strip('"').lower() + + if not any(token in alias_text for token in ("duration", "turnaround")): + return match.group(0) + + return f"DATEDIFF('second', {right}, {left}) AS {alias}" + + return timestamp_subtraction_pattern.sub(replace_subtraction, sql) + + +def _infer_mssql_timestamp_expression(sql: str) -> str | None: + timestamp_column_pattern = re.compile( + r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at)")', + re.IGNORECASE, + ) + if match := timestamp_column_pattern.search(sql): + return match.group(0) + + table_pattern = re.compile( + r'\bFROM\s+("[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)', + re.IGNORECASE, + ) + if match := table_pattern.search(sql): + table_name = match.group(1) + normalized_table_name = table_name.strip('"[]').lower() + if "report" in normalized_table_name: + return f'{table_name}."generated_at"' + if any( + token in normalized_table_name + for token in ("repair", "ticket", "debug", "event", "log") + ): + return f'{table_name}."created_at"' + + return None + + +def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: + timestamp_expression = _infer_mssql_timestamp_expression(sql) + if not timestamp_expression: + return sql + + invented_date_identifier_pattern = re.compile( + r'(? str: + timestamp_expression = _infer_mssql_timestamp_expression(sql) + if not timestamp_expression: + return sql + + bucket_expressions = { + "year": f"DATEPART(YEAR, {timestamp_expression})", + "month": f"DATEPART(MONTH, {timestamp_expression})", + "day": f"DATEPART(DAY, {timestamp_expression})", + } + rewritten = sql + + for bucket, expression in bucket_expressions.items(): + select_identifier_pattern = re.compile( + rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', + re.IGNORECASE, + ) + + def replace_select_identifier(match: re.Match[str]) -> str: + prefix = match.group("prefix") + suffix = match.group("suffix") + return f'{prefix}{expression} AS "{bucket}"{suffix}' + + rewritten = select_identifier_pattern.sub( + replace_select_identifier, rewritten + ) + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match[str]) -> str: + body = match.group("body") + for bucket, expression in bucket_expressions.items(): + body = re.sub( + rf'"{bucket}"', + expression, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"\[{bucket}\]", + expression, + body, + flags=re.IGNORECASE, + ) + return f"{match.group(1)}{body}" + + return clause_pattern.sub(replace_clause, rewritten) + + +def _rewrite_mssql_datepart_alias_references(sql: str) -> str: + datepart_alias_pattern = re.compile( + r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", + re.IGNORECASE, + ) + aliases: dict[str, str] = {} + + for match in datepart_alias_pattern.finditer(sql): + expression = match.group(1) + alias = match.group(4) or match.group(5) or match.group(6) + aliases[alias.lower()] = expression + + if not aliases: + return sql + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match) -> str: + body = match.group("body") + placeholders: dict[str, str] = {} + for alias, expression in aliases.items(): + placeholder = f"__WREN_MSSQL_DATEPART_ALIAS_{len(placeholders)}__" + placeholders[placeholder] = expression + body = re.sub( + rf'"{re.escape(alias)}"', + placeholder, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"\[{re.escape(alias)}\]", + placeholder, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"(? str: + normalized = sql + + if normalize_data_source(data_source) == "MSSQL": + now = datetime.now() + normalized = re.sub( + r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE + ) + normalized = re.sub( + r"CAST\(\s*('(?:[^']|'')*')\s+AS\s+DATETIME(?:2|OFFSET)\s*\)", + r"\1", + normalized, + flags=re.IGNORECASE, + ) + normalized = _replace_relative_getdate_calls(normalized, now) + normalized = _rewrite_mssql_to_unixtime(normalized) + normalized = _rewrite_mssql_timestamp_subtraction(normalized) + normalized = _rewrite_mssql_timestamp_casts(normalized) + normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) + normalized = _rewrite_mssql_bucket_functions(normalized) + normalized = _rewrite_temporal_bucket_functions(normalized) + normalized = _rewrite_mssql_datepart_alias_references(normalized) + + return re.sub(r"\s+", " ", normalized).strip() From ac9be6ac44f08161ff4f93e4f748ba4624a33c47 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 19:17:30 +0530 Subject: [PATCH 0068/1087] updated sql pairs --- .../src/pipelines/generation/utils/sql.py | 472 +++++++++++++++++- .../src/pipelines/retrieval/sql_executor.py | 3 +- 2 files changed, 470 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 3851867f11..daa6b802ec 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,6 @@ import logging import re +from datetime import datetime, timedelta from typing import Any, Dict, List import aiohttp @@ -13,10 +14,6 @@ clean_generation_result, ) from src.core.provider import LLMProvider -from src.pipelines.sql_normalizer import ( - normalize_data_source, - normalize_generation_result_sql, -) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge logger = logging.getLogger("wren-ai-service") @@ -109,6 +106,473 @@ def is_select_statement(sql: str) -> bool: return bool(re.match(r"^\s*(?:WITH|SELECT)\b", sql, flags=re.IGNORECASE)) +def normalize_data_source(data_source: str | None) -> str: + normalized = (data_source or "").strip().upper().replace("-", "_").replace( + " ", "_" + ) + if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: + return "MSSQL" + return normalized + + +def _format_timestamp_literal(value: datetime) -> str: + return value.strftime("'%Y-%m-%d %H:%M:%S'") + + +def _add_months(value: datetime, months: int) -> datetime: + month_index = value.month - 1 + months + year = value.year + month_index // 12 + month = month_index % 12 + 1 + day = min( + value.day, + [ + 31, + 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28, + 31, + 30, + 31, + 30, + 31, + 31, + 30, + 31, + 30, + 31, + ][month - 1], + ) + return value.replace(year=year, month=month, day=day) + + +def _start_of_month(value: datetime) -> datetime: + return value.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + + +def _replace_relative_getdate_calls(sql: str, now: datetime) -> str: + def replace_month_offset(match: re.Match[str]) -> str: + months = int(match.group(1)) + return _format_timestamp_literal(_add_months(now, months)) + + def replace_year_offset(match: re.Match[str]) -> str: + years = int(match.group(1)) + return _format_timestamp_literal(_add_months(now, years * 12)) + + def replace_day_offset(match: re.Match[str]) -> str: + days = int(match.group(1)) + return _format_timestamp_literal(now + timedelta(days=days)) + + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_month_offset, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*year\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_year_offset, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*day\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", + replace_day_offset, + sql, + flags=re.IGNORECASE, + ) + + current_month_start = _format_timestamp_literal(_start_of_month(now)) + previous_month_start = _format_timestamp_literal( + _start_of_month(_add_months(now, -1)) + ) + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*,\s*0\s*\)", + current_month_start, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*-\s*1\s*,\s*0\s*\)", + previous_month_start, + sql, + flags=re.IGNORECASE, + ) + return sql + + +def _rewrite_mssql_bucket_functions(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + + sql = re.sub( + rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", + lambda m: ( + f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" + ), + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", + lambda m: f"DATEPART(YEAR, {m.group(1)})", + sql, + flags=re.IGNORECASE, + ) + return sql + + +def _rewrite_temporal_bucket_functions(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + replacements = [ + ( + re.compile( + rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile( + rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(YEAR, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(MONTH, {m.group(1)})", + ), + ( + re.compile( + rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ), + lambda m: f"DATEPART(DAY, {m.group(1)})", + ), + ] + + rewritten = sql + for pattern, replacement in replacements: + rewritten = pattern.sub(replacement, rewritten) + + return rewritten + + +def _rewrite_mssql_timestamp_casts(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + timestamp_function_pattern = re.compile( + rf"\bTO_TIMESTAMP(?:_(?:MILLIS|SECONDS|MICROS|NANOS))?\(\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ) + timestamp_cast_pattern = re.compile( + r"CAST\(\s*((?:[^()]|\([^()]*\))+?)\s+AS\s+TIMESTAMP\s*\)", + re.IGNORECASE, + ) + + rewritten = timestamp_function_pattern.sub( + lambda m: f"CAST({m.group(1)} AS DATETIME)", sql + ) + rewritten = timestamp_cast_pattern.sub( + lambda m: f"CAST({m.group(1)} AS DATETIME)", rewritten + ) + return rewritten + + +def _rewrite_mssql_to_unixtime(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + to_unixtime_pattern = re.compile( + rf"\bTO_UNIXTIME\(\s*{expression_pattern}\s*\)", + re.IGNORECASE, + ) + + return to_unixtime_pattern.sub(lambda m: m.group(1), sql) + + +def _rewrite_mssql_timestamp_subtraction(sql: str) -> str: + expression_pattern = r"((?:[^(),+\-]|\([^()]*\))+?)" + timestamp_subtraction_pattern = re.compile( + rf"{expression_pattern}\s*-\s*{expression_pattern}\s+AS\s+(\"[^\"]+\")", + re.IGNORECASE, + ) + + def replace_subtraction(match: re.Match[str]) -> str: + left = match.group(1).strip() + right = match.group(2).strip() + alias = match.group(3) + alias_text = alias.strip('"').lower() + + if not any(token in alias_text for token in ("duration", "turnaround")): + return match.group(0) + + return f"DATEDIFF('second', {right}, {left}) AS {alias}" + + return timestamp_subtraction_pattern.sub(replace_subtraction, sql) + + +def _infer_mssql_timestamp_expression(sql: str) -> str | None: + timestamp_column_pattern = re.compile( + r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at)")', + re.IGNORECASE, + ) + if match := timestamp_column_pattern.search(sql): + return match.group(0) + + table_pattern = re.compile( + r'\bFROM\s+("[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)', + re.IGNORECASE, + ) + if match := table_pattern.search(sql): + table_name = match.group(1) + normalized_table_name = table_name.strip('"[]').lower() + if "report" in normalized_table_name: + return f'{table_name}."generated_at"' + if any( + token in normalized_table_name + for token in ("repair", "ticket", "debug", "event", "log") + ): + return f'{table_name}."created_at"' + + return None + + +def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: + timestamp_expression = _infer_mssql_timestamp_expression(sql) + if not timestamp_expression: + return sql + + invented_date_identifier_pattern = re.compile( + r'(? str: + timestamp_expression = _infer_mssql_timestamp_expression(sql) + if not timestamp_expression: + return sql + + bucket_expressions = { + "year": f"DATEPART(YEAR, {timestamp_expression})", + "month": f"DATEPART(MONTH, {timestamp_expression})", + "day": f"DATEPART(DAY, {timestamp_expression})", + } + rewritten = sql + + for bucket, expression in bucket_expressions.items(): + select_identifier_pattern = re.compile( + rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', + re.IGNORECASE, + ) + + def replace_select_identifier(match: re.Match[str]) -> str: + prefix = match.group("prefix") + suffix = match.group("suffix") + return f'{prefix}{expression} AS "{bucket}"{suffix}' + + rewritten = select_identifier_pattern.sub( + replace_select_identifier, rewritten + ) + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match[str]) -> str: + body = match.group("body") + for bucket, expression in bucket_expressions.items(): + body = re.sub( + rf'"{bucket}"', + expression, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"\[{bucket}\]", + expression, + body, + flags=re.IGNORECASE, + ) + return f"{match.group(1)}{body}" + + return clause_pattern.sub(replace_clause, rewritten) + + +def _rewrite_mssql_datepart_alias_references(sql: str) -> str: + datepart_alias_pattern = re.compile( + r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", + re.IGNORECASE, + ) + aliases: dict[str, str] = {} + + for match in datepart_alias_pattern.finditer(sql): + expression = match.group(1) + alias = match.group(4) or match.group(5) or match.group(6) + aliases[alias.lower()] = expression + + if not aliases: + return sql + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match) -> str: + body = match.group("body") + placeholders: dict[str, str] = {} + for alias, expression in aliases.items(): + placeholder = f"__WREN_MSSQL_DATEPART_ALIAS_{len(placeholders)}__" + placeholders[placeholder] = expression + body = re.sub( + rf'"{re.escape(alias)}"', + placeholder, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"\[{re.escape(alias)}\]", + placeholder, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"(? str: + normalized = sql + + if normalize_data_source(data_source) == "MSSQL": + now = datetime.now() + normalized = re.sub( + r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE + ) + normalized = re.sub( + r"CAST\(\s*('(?:[^']|'')*')\s+AS\s+DATETIME(?:2|OFFSET)\s*\)", + r"\1", + normalized, + flags=re.IGNORECASE, + ) + normalized = _replace_relative_getdate_calls(normalized, now) + normalized = _rewrite_mssql_to_unixtime(normalized) + normalized = _rewrite_mssql_timestamp_subtraction(normalized) + normalized = _rewrite_mssql_timestamp_casts(normalized) + normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) + normalized = _rewrite_mssql_bucket_functions(normalized) + normalized = _rewrite_temporal_bucket_functions(normalized) + normalized = _rewrite_mssql_datepart_alias_references(normalized) + + return re.sub(r"\s+", " ", normalized).strip() + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): diff --git a/wren-ai-service/src/pipelines/retrieval/sql_executor.py b/wren-ai-service/src/pipelines/retrieval/sql_executor.py index b6f50b7826..129c4992c3 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_executor.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_executor.py @@ -10,7 +10,6 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline -from src.pipelines.sql_normalizer import normalize_generation_result_sql logger = logging.getLogger("wren-ai-service") @@ -30,6 +29,8 @@ async def run( project_id: str | None = None, limit: int = 500, ): + from src.pipelines.generation.utils.sql import normalize_generation_result_sql + sql = normalize_generation_result_sql(sql, data_source=self._data_source) async with aiohttp.ClientSession() as session: _, data, addition = await self._engine.execute_sql( From 649a5c56f14e0b27243b235511cc31d17a26b7a6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 19:52:19 +0530 Subject: [PATCH 0069/1087] updated sql pairs log --- wren-ai-service/src/pipelines/generation/utils/sql.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index daa6b802ec..e863c69358 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -794,8 +794,10 @@ async def _classify_generation_result( - DON'T USE "TO_CHAR" function in the generated SQL query. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -- For the ranking problem, you must add the ranking column to the final SELECT clause. +- For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. +- For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. +- Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. +- If a ranking helper is required internally, compute it in a subquery/CTE and filter on it, but omit it from the final SELECT unless explicitly requested. """ _MSSQL_TEXT_TO_SQL_RULES = """ @@ -816,6 +818,7 @@ async def _classify_generation_result( Then GROUP BY and ORDER BY the same year/month expressions. - Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the DATEPART(...) expression instead. - For year bucketing, prefer DATEPART(YEAR, ). +- For top/bottom N questions in MSSQL, prefer SELECT TOP (N) with ORDER BY over DENSE_RANK/ROW_NUMBER when the user did not explicitly request ranks. - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' - AND < '2026-01-01 00:00:00' @@ -1026,8 +1029,8 @@ async def _classify_generation_result( 2. Explicitly state the following information in the reasoning plan: if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; if the user uses a relative timeframe and Current Time is provided in the input, you will resolve it into an absolute time frame in the SQL query using exact dates rather than relative date arithmetic. -3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. +3. For top/bottom N questions, plan to order by the relevant metric and limit the result to N rows. Do not add a rank column unless the user explicitly asks to see ranks. +4. For questions like "top 10 common failures", the final table should contain the grouped business field and its count/metric, not helper ranking columns. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. From 559e9f49e128f8b60096d33a0edf9190ef27c933 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 29 May 2026 23:19:51 +0530 Subject: [PATCH 0070/1087] migration --- wren-ui/knexfile.js | 54 +++- wren-ui/package.json | 2 + wren-ui/src/apollo/server/config.ts | 31 +++ wren-ui/src/apollo/server/utils/knex.ts | 83 ++++++- wren-ui/src/common.ts | 8 + wren-ui/tools/knex.js | 316 +++++++++++++++++++++++- 6 files changed, 474 insertions(+), 20 deletions(-) diff --git a/wren-ui/knexfile.js b/wren-ui/knexfile.js index 72c26263a4..c63ed48de4 100644 --- a/wren-ui/knexfile.js +++ b/wren-ui/knexfile.js @@ -1,14 +1,66 @@ // Update with your config settings. +const normalizeDbType = (dbType) => + (dbType || 'sqlite').trim().toLowerCase().replace(/[-_ ]/g, ''); + +const parseBooleanUrlParam = (searchParams, key, fallback) => { + const value = searchParams.get(key); + if (value === null) return fallback; + return value.toLowerCase() === 'true'; +}; + +const getMssqlConnection = () => { + if (process.env.MSSQL_URL) { + const url = new URL(process.env.MSSQL_URL); + return { + server: url.hostname, + port: url.port ? parseInt(url.port) : 1433, + database: decodeURIComponent(url.pathname.replace(/^\//, '')), + user: decodeURIComponent(url.username), + password: decodeURIComponent(url.password), + options: { + encrypt: parseBooleanUrlParam(url.searchParams, 'encrypt', false), + trustServerCertificate: parseBooleanUrlParam( + url.searchParams, + 'trustServerCertificate', + true, + ), + }, + }; + } + + return { + server: process.env.MSSQL_HOST || 'localhost', + port: process.env.MSSQL_PORT ? parseInt(process.env.MSSQL_PORT) : 1433, + database: process.env.MSSQL_DATABASE || 'wren_ui', + user: process.env.MSSQL_USER, + password: process.env.MSSQL_PASSWORD, + options: { + encrypt: process.env.MSSQL_ENCRYPT === 'true', + trustServerCertificate: + process.env.MSSQL_TRUST_SERVER_CERTIFICATE !== 'false', + }, + }; +}; + +const dbType = normalizeDbType(process.env.DB_TYPE); + /** * @type { Object. } */ -if (process.env.DB_TYPE === 'pg') { +if (dbType === 'pg' || dbType === 'postgres' || dbType === 'postgresql') { console.log('Using Postgres'); module.exports = { client: 'pg', connection: process.env.PG_URL, }; +} else if (dbType === 'mssql' || dbType === 'sqlserver') { + console.log('Using MSSQL'); + module.exports = { + client: 'mssql', + connection: getMssqlConnection(), + pool: { min: 2, max: 10 }, + }; } else { console.log('Using SQLite'); module.exports = { diff --git a/wren-ui/package.json b/wren-ui/package.json index 3f76b0e43b..9e5af55473 100644 --- a/wren-ui/package.json +++ b/wren-ui/package.json @@ -11,6 +11,7 @@ "test:e2e": "npx playwright install chromium && npx playwright test", "check-types": "tsc --noEmit", "migrate": "yarn knex migrate:latest", + "migrate:sqlite-to-mssql": "node -e \"process.env.MIGRATE_SQLITE_TO_MSSQL='true'; require('./tools/knex.js')\"", "rollback": "yarn knex migrate:rollback", "generate-gql": "yarn graphql-codegen --config codegen.yaml" }, @@ -34,6 +35,7 @@ "pg-cursor": "^2.7.4", "posthog-node": "^4.3.2", "sql-formatter": "^15.3.0", + "tedious": "^18.6.1", "uuid": "^11.1.0" }, "devDependencies": { diff --git a/wren-ui/src/apollo/server/config.ts b/wren-ui/src/apollo/server/config.ts index bafebe866b..836551f2c7 100644 --- a/wren-ui/src/apollo/server/config.ts +++ b/wren-ui/src/apollo/server/config.ts @@ -9,6 +9,15 @@ export interface IConfig { // pg pgUrl?: string; debug?: boolean; + // mssql + mssqlUrl?: string; + mssqlHost?: string; + mssqlPort?: number; + mssqlDatabase?: string; + mssqlUser?: string; + mssqlPassword?: string; + mssqlEncrypt?: boolean; + mssqlTrustServerCertificate?: boolean; // sqlite sqliteFile?: string; @@ -59,6 +68,13 @@ const defaultConfig = { pgUrl: 'postgres://postgres:postgres@localhost:5432/admin_ui', debug: false, + // mssql + mssqlHost: 'localhost', + mssqlPort: 1433, + mssqlDatabase: 'wren_ui', + mssqlEncrypt: false, + mssqlTrustServerCertificate: true, + // sqlite sqliteFile: './db.sqlite3', @@ -88,6 +104,21 @@ const config = { // pg pgUrl: process.env.PG_URL, debug: process.env.DEBUG === 'true', + // mssql + mssqlUrl: process.env.MSSQL_URL, + mssqlHost: process.env.MSSQL_HOST, + mssqlPort: process.env.MSSQL_PORT + ? parseInt(process.env.MSSQL_PORT) + : undefined, + mssqlDatabase: process.env.MSSQL_DATABASE, + mssqlUser: process.env.MSSQL_USER, + mssqlPassword: process.env.MSSQL_PASSWORD, + mssqlEncrypt: process.env.MSSQL_ENCRYPT + ? process.env.MSSQL_ENCRYPT === 'true' + : undefined, + mssqlTrustServerCertificate: process.env.MSSQL_TRUST_SERVER_CERTIFICATE + ? process.env.MSSQL_TRUST_SERVER_CERTIFICATE === 'true' + : undefined, // sqlite sqliteFile: process.env.SQLITE_FILE, diff --git a/wren-ui/src/apollo/server/utils/knex.ts b/wren-ui/src/apollo/server/utils/knex.ts index b7c74bba53..c5c5a91263 100644 --- a/wren-ui/src/apollo/server/utils/knex.ts +++ b/wren-ui/src/apollo/server/utils/knex.ts @@ -2,11 +2,67 @@ interface KnexOptions { dbType: string; pgUrl?: string; debug?: boolean; + mssqlUrl?: string; + mssqlHost?: string; + mssqlPort?: number; + mssqlDatabase?: string; + mssqlUser?: string; + mssqlPassword?: string; + mssqlEncrypt?: boolean; + mssqlTrustServerCertificate?: boolean; sqliteFile?: string; } +const normalizeDbType = (dbType?: string) => + (dbType || 'sqlite').trim().toLowerCase().replace(/[-_ ]/g, ''); + +const parseBooleanUrlParam = ( + searchParams: URLSearchParams, + key: string, + fallback: boolean, +) => { + const value = searchParams.get(key); + if (value === null) return fallback; + return value.toLowerCase() === 'true'; +}; + +const getMssqlConnection = (options: KnexOptions) => { + if (options.mssqlUrl) { + const url = new URL(options.mssqlUrl); + return { + server: url.hostname, + port: url.port ? parseInt(url.port) : 1433, + database: decodeURIComponent(url.pathname.replace(/^\//, '')), + user: decodeURIComponent(url.username), + password: decodeURIComponent(url.password), + options: { + encrypt: parseBooleanUrlParam(url.searchParams, 'encrypt', false), + trustServerCertificate: parseBooleanUrlParam( + url.searchParams, + 'trustServerCertificate', + true, + ), + }, + }; + } + + return { + server: options.mssqlHost, + port: options.mssqlPort || 1433, + database: options.mssqlDatabase, + user: options.mssqlUser, + password: options.mssqlPassword, + options: { + encrypt: options.mssqlEncrypt ?? false, + trustServerCertificate: options.mssqlTrustServerCertificate ?? true, + }, + }; +}; + export const bootstrapKnex = (options: KnexOptions) => { - if (options.dbType === 'pg') { + const dbType = normalizeDbType(options.dbType); + + if (dbType === 'pg' || dbType === 'postgres' || dbType === 'postgresql') { const { pgUrl, debug } = options; console.log('using pg'); /* eslint-disable @typescript-eslint/no-var-requires */ @@ -16,15 +72,26 @@ export const bootstrapKnex = (options: KnexOptions) => { debug, pool: { min: 2, max: 10 }, }); - } else { - console.log('using sqlite'); + } + + if (dbType === 'mssql' || dbType === 'sqlserver') { + console.log('using mssql'); /* eslint-disable @typescript-eslint/no-var-requires */ return require('knex')({ - client: 'better-sqlite3', - connection: { - filename: options.sqliteFile, - }, - useNullAsDefault: true, + client: 'mssql', + connection: getMssqlConnection(options), + debug: options.debug, + pool: { min: 2, max: 10 }, }); } + + console.log('using sqlite'); + /* eslint-disable @typescript-eslint/no-var-requires */ + return require('knex')({ + client: 'better-sqlite3', + connection: { + filename: options.sqliteFile, + }, + useNullAsDefault: true, + }); }; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index cdaee2242e..33be8d9465 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -119,6 +119,14 @@ export const initComponents = () => { dbType: serverConfig.dbType, pgUrl: serverConfig.pgUrl, debug: serverConfig.debug, + mssqlUrl: serverConfig.mssqlUrl, + mssqlHost: serverConfig.mssqlHost, + mssqlPort: serverConfig.mssqlPort, + mssqlDatabase: serverConfig.mssqlDatabase, + mssqlUser: serverConfig.mssqlUser, + mssqlPassword: serverConfig.mssqlPassword, + mssqlEncrypt: serverConfig.mssqlEncrypt, + mssqlTrustServerCertificate: serverConfig.mssqlTrustServerCertificate, sqliteFile: serverConfig.sqliteFile, }); diff --git a/wren-ui/tools/knex.js b/wren-ui/tools/knex.js index 0324ea66f4..27beb1b0ec 100644 --- a/wren-ui/tools/knex.js +++ b/wren-ui/tools/knex.js @@ -1,10 +1,81 @@ +const fs = require('fs'); +const path = require('path'); + const DB_TYPE = process.env.DB_TYPE; // export DB_TYPE=pg const PG_URL = process.env.PG_URL; const DEBUG = process.env.DEBUG === 'true'; // export DEBUG=true const SQLITE_FILE = process.env.SQLITE_FILE; // export SQLITE_FILE=./db.sqlite3 -const getKnex = () => { - if (DB_TYPE === 'pg') { +const APP_TABLE_ORDER = [ + 'project', + 'model', + 'model_column', + 'model_nested_column', + 'relation', + 'metric', + 'metric_measure', + 'view', + 'deploy_log', + 'thread', + 'thread_response', + 'schema_change', + 'learning', + 'dashboard', + 'dashboard_item', + 'sql_pair', + 'instruction', + 'dashboard_item_refresh_job', + 'asking_task', + 'api_history', +]; + +const normalizeDbType = (dbType) => + (dbType || 'sqlite').trim().toLowerCase().replace(/[-_ ]/g, ''); + +const parseBooleanUrlParam = (searchParams, key, fallback) => { + const value = searchParams.get(key); + if (value === null) return fallback; + return value.toLowerCase() === 'true'; +}; + +const getMssqlConnection = () => { + if (process.env.MSSQL_URL) { + const url = new URL(process.env.MSSQL_URL); + return { + server: url.hostname, + port: url.port ? parseInt(url.port) : 1433, + database: decodeURIComponent(url.pathname.replace(/^\//, '')), + user: decodeURIComponent(url.username), + password: decodeURIComponent(url.password), + options: { + encrypt: parseBooleanUrlParam(url.searchParams, 'encrypt', false), + trustServerCertificate: parseBooleanUrlParam( + url.searchParams, + 'trustServerCertificate', + true, + ), + }, + }; + } + + return { + server: process.env.MSSQL_HOST || 'localhost', + port: process.env.MSSQL_PORT ? parseInt(process.env.MSSQL_PORT) : 1433, + database: process.env.MSSQL_DATABASE || 'wren_ui', + user: process.env.MSSQL_USER, + password: process.env.MSSQL_PASSWORD, + options: { + encrypt: process.env.MSSQL_ENCRYPT === 'true', + trustServerCertificate: + process.env.MSSQL_TRUST_SERVER_CERTIFICATE !== 'false', + }, + }; +}; + +const getKnex = (options = {}) => { + const dbType = normalizeDbType(options.dbType || DB_TYPE); + + if (dbType === 'pg' || dbType === 'postgres' || dbType === 'postgresql') { console.log('using pg'); /* eslint-disable @typescript-eslint/no-var-requires */ return require('knex')({ @@ -13,20 +84,238 @@ const getKnex = () => { debug: DEBUG, pool: { min: 2, max: 10 }, }); - } else { - console.log('using sqlite'); + } + + if (dbType === 'mssql' || dbType === 'sqlserver') { + console.log('using mssql'); /* eslint-disable @typescript-eslint/no-var-requires */ return require('knex')({ - client: 'better-sqlite3', - connection: { - filename: SQLITE_FILE, - }, - useNullAsDefault: true, + client: 'mssql', + connection: getMssqlConnection(), + debug: DEBUG, + pool: { min: 2, max: 10 }, }); } + + console.log('using sqlite'); + /* eslint-disable @typescript-eslint/no-var-requires */ + return require('knex')({ + client: 'better-sqlite3', + connection: { + filename: options.sqliteFile || SQLITE_FILE, + }, + useNullAsDefault: true, + }); +}; + +const getSqliteFile = () => { + const appRoot = path.resolve(__dirname, '..'); + const candidates = [ + SQLITE_FILE, + path.join(appRoot, 'data', 'db.sqlite3'), + path.join(appRoot, 'db.sqlite3'), + ].filter(Boolean); + + const sqliteFile = candidates.find((candidate) => fs.existsSync(candidate)); + if (!sqliteFile) { + throw new Error( + `SQLite source database not found. Set SQLITE_FILE explicitly. Checked: ${candidates.join( + ', ', + )}`, + ); + } + return sqliteFile; +}; + +const getSourceTables = async (sourceDb) => { + const rows = await sourceDb.raw(` + SELECT name + FROM sqlite_master + WHERE type = 'table' + AND name NOT LIKE 'sqlite_%' + AND name NOT IN ('knex_migrations', 'knex_migrations_lock') + `); + const tableNames = rows.map((row) => row.name); + const orderedTables = APP_TABLE_ORDER.filter((table) => + tableNames.includes(table), + ); + const remainingTables = tableNames + .filter((table) => !orderedTables.includes(table)) + .sort(); + return [...orderedTables, ...remainingTables]; +}; + +const getTargetColumns = async (targetDb, tableName) => { + const rows = await targetDb('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ TABLE_SCHEMA: 'dbo', TABLE_NAME: tableName }) + .orderBy('ORDINAL_POSITION', 'asc'); + return rows.map((row) => row.COLUMN_NAME); +}; + +const parseCount = (row) => Number(row.count || row.Count || row[''] || 0); + +const getTableCount = async (db, tableName) => { + const [row] = await db(tableName).count({ count: '*' }); + return parseCount(row); +}; + +const ensureTargetIsEmpty = async (targetDb, tableNames) => { + const nonEmptyTables = []; + for (const tableName of tableNames) { + if (!(await targetDb.schema.hasTable(tableName))) { + continue; + } + const count = await getTableCount(targetDb, tableName); + if (count > 0) { + nonEmptyTables.push(`${tableName} (${count})`); + } + } + + if (nonEmptyTables.length > 0 && process.env.MIGRATE_OVERWRITE !== 'true') { + throw new Error( + `Refusing to copy into a non-empty MSSQL database. Non-empty tables: ${nonEmptyTables.join( + ', ', + )}. Set MIGRATE_OVERWRITE=true to delete target Wren UI rows first.`, + ); + } +}; + +const setForeignKeysEnabled = async (targetDb, enabled) => { + const rows = await targetDb + .select( + targetDb.raw( + "QUOTENAME(SCHEMA_NAME(schema_id)) + '.' + QUOTENAME(name) AS full_name", + ), + ) + .from('sys.tables') + .where({ is_ms_shipped: 0 }); + + for (const row of rows) { + const fullName = row.full_name || row.fullName; + if (enabled) { + await targetDb.raw( + `ALTER TABLE ${fullName} WITH CHECK CHECK CONSTRAINT ALL`, + ); + } else { + await targetDb.raw(`ALTER TABLE ${fullName} NOCHECK CONSTRAINT ALL`); + } + } +}; + +const clearTargetTables = async (targetDb, tableNames) => { + for (const tableName of [...tableNames].reverse()) { + if (await targetDb.schema.hasTable(tableName)) { + await targetDb(tableName).delete(); + } + } +}; + +const copyTable = async (sourceDb, targetDb, tableName) => { + if (!(await targetDb.schema.hasTable(tableName))) { + console.log(`Skipping ${tableName}: target table does not exist`); + return { tableName, sourceCount: 0, targetCount: 0 }; + } + + const sourceRows = await sourceDb(tableName).select('*'); + if (sourceRows.length === 0) { + console.log(`Copied ${tableName}: 0 rows`); + return { + tableName, + sourceCount: 0, + targetCount: await getTableCount(targetDb, tableName), + }; + } + + const targetColumns = await getTargetColumns(targetDb, tableName); + const commonColumns = targetColumns.filter((column) => + Object.prototype.hasOwnProperty.call(sourceRows[0], column), + ); + const rows = sourceRows.map((row) => + Object.fromEntries(commonColumns.map((column) => [column, row[column]])), + ); + const chunkSize = Math.max( + 1, + Math.floor(1800 / Math.max(commonColumns.length, 1)), + ); + const hasId = commonColumns.includes('id'); + + if (hasId) { + await targetDb.raw(`SET IDENTITY_INSERT [${tableName}] ON`); + } + + try { + for (let index = 0; index < rows.length; index += chunkSize) { + await targetDb(tableName).insert(rows.slice(index, index + chunkSize)); + } + } finally { + if (hasId) { + await targetDb.raw(`SET IDENTITY_INSERT [${tableName}] OFF`); + } + } + + const targetCount = await getTableCount(targetDb, tableName); + console.log(`Copied ${tableName}: ${sourceRows.length} rows`); + return { tableName, sourceCount: sourceRows.length, targetCount }; +}; + +const migrateSqliteToMssql = async () => { + const sqliteFile = getSqliteFile(); + const sourceDb = getKnex({ dbType: 'sqlite', sqliteFile }); + const targetDb = getKnex({ dbType: 'mssql' }); + + try { + const migrationsDir = path.resolve(__dirname, '..', 'migrations'); + console.log(`Migrating Wren UI application tables from ${sqliteFile}`); + console.log('Running MSSQL schema migrations'); + await targetDb.migrate.latest({ directory: migrationsDir }); + + const tableNames = await getSourceTables(sourceDb); + await ensureTargetIsEmpty(targetDb, tableNames); + + await targetDb.transaction(async (trx) => { + await setForeignKeysEnabled(trx, false); + try { + if (process.env.MIGRATE_OVERWRITE === 'true') { + await clearTargetTables(trx, tableNames); + } + + const verification = []; + for (const tableName of tableNames) { + verification.push(await copyTable(sourceDb, trx, tableName)); + } + + const mismatches = verification.filter( + ({ sourceCount, targetCount }) => sourceCount !== targetCount, + ); + if (mismatches.length > 0) { + throw new Error( + `Row-count verification failed: ${mismatches + .map( + ({ tableName, sourceCount, targetCount }) => + `${tableName} sqlite=${sourceCount} mssql=${targetCount}`, + ) + .join(', ')}`, + ); + } + } finally { + await setForeignKeysEnabled(trx, true); + } + }); + + console.log('SQLite to MSSQL migration completed successfully.'); + } finally { + await sourceDb.destroy(); + await targetDb.destroy(); + } }; const main = async () => { + if (process.env.MIGRATE_SQLITE_TO_MSSQL === 'true') { + await migrateSqliteToMssql(); + return; + } + const knex = getKnex(); const query = knex.queryBuilder(); @@ -36,7 +325,12 @@ const main = async () => { .whereIn('id', [7, 8]); console.log(projects); - process.exit(0); + await knex.destroy(); }; -main(); +main() + .then(() => process.exit(0)) + .catch((error) => { + console.error(error); + process.exit(1); + }); From 66aba3cb28759c305cf7033ad908a327fdaf669c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 1 Jun 2026 14:41:07 +0530 Subject: [PATCH 0071/1087] main --- wren-ai-service/src/__main__.py | 52 +++++++++++++++++++++++++++++++-- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py index de141c3fde..6ff2d11165 100644 --- a/wren-ai-service/src/__main__.py +++ b/wren-ai-service/src/__main__.py @@ -1,10 +1,14 @@ from contextlib import asynccontextmanager +from importlib.util import find_spec +from pathlib import Path import uvicorn from fastapi import FastAPI from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware +from fastapi.openapi.docs import get_swagger_ui_html from fastapi.responses import ORJSONResponse, RedirectResponse +from fastapi.staticfiles import StaticFiles from langfuse.decorators import langfuse_context from src.config import settings @@ -24,6 +28,22 @@ ) +def get_local_swagger_static_dir() -> Path | None: + litellm_spec = find_spec("litellm") + if not litellm_spec or not litellm_spec.submodule_search_locations: + return None + + swagger_dir = ( + Path(next(iter(litellm_spec.submodule_search_locations))) / "proxy" / "swagger" + ) + required_assets = ("swagger-ui-bundle.js", "swagger-ui.css", "favicon.ico") + + if all((swagger_dir / asset).is_file() for asset in required_assets): + return swagger_dir + + return None + + # https://fastapi.tiangolo.com/advanced/events/#lifespan @asynccontextmanager async def lifespan(app: FastAPI): @@ -42,10 +62,19 @@ async def lifespan(app: FastAPI): app = FastAPI( title="wren-ai-service API Docs", lifespan=lifespan, + docs_url=None, redoc_url=None, default_response_class=ORJSONResponse, ) +swagger_static_dir = get_local_swagger_static_dir() +if swagger_static_dir: + app.mount( + "/_docs/static", + StaticFiles(directory=swagger_static_dir), + name="swagger-static", + ) + app.add_middleware( CORSMiddleware, allow_origins=["*"], @@ -81,6 +110,25 @@ def root(): return RedirectResponse(url="/docs") +@app.get("/docs", include_in_schema=False) +def swagger_ui_html(): + kwargs = { + "openapi_url": app.openapi_url, + "title": f"{app.title} - Swagger UI", + } + + if swagger_static_dir: + kwargs.update( + { + "swagger_js_url": "/_docs/static/swagger-ui-bundle.js", + "swagger_css_url": "/_docs/static/swagger-ui.css", + "swagger_favicon_url": "/_docs/static/favicon.ico", + } + ) + + return get_swagger_ui_html(**kwargs) + + @app.get("/health") def health(): return {"status": "ok"} @@ -95,6 +143,6 @@ def health(): reload_includes=["src/**/*.py", ".env.dev", "config.yaml"], reload_excludes=["tests/**/*.py", "eval/**/*.py"], workers=1, - loop="uvloop", - http="httptools", + loop="auto", + http="auto", ) From ee693959dc95c53a0dd67d98d878a8db6acec63a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 1 Jun 2026 15:37:43 +0530 Subject: [PATCH 0072/1087] ask --- wren-ai-service/src/web/v1/services/ask.py | 99 +++++++++++++++++++++- 1 file changed, 95 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 8b8f92a112..e28801f425 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -184,6 +184,49 @@ def _is_data_analysis_query(self, query: str) -> bool: } return any(term in normalized for term in analysis_terms) + def _get_unqueryable_metric_message( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + normalized_schema = re.sub(r"\s+", " ", " ".join(table_ddls).lower()) + + if not normalized_query: + return None + + first_pass_yield_terms = ( + "first pass yield", + "first-pass yield", + "first_pass_yield", + "fpy", + ) + if not any(term in normalized_query for term in first_pass_yield_terms): + return None + + required_field_patterns = ( + r"\bfirst[_ ]?pass[_ ]?yield\b", + r"\bfpy\b", + r"\battempt\b", + r"\battempt[_ ]?number\b", + r"\bfirst[_ ]?attempt\b", + r"\bpass[_ ]?fail\b", + r"\byield\b", + ) + has_required_field = any( + re.search(pattern, normalized_schema) + for pattern in required_field_patterns + ) + + if has_required_field: + return None + + return ( + "The schema does not expose first-pass yield, attempt number, " + "first-attempt result, or pass/fail fields as queryable columns. " + "I cannot calculate First Pass Yield from only generic JSON/text " + "fields such as data. Add those fields as first-class columns or " + "calculated fields, then ask again." + ) + async def _run_with_timeout(self, label: str, coroutine): try: return await asyncio.wait_for( @@ -321,6 +364,7 @@ async def ask( ] sql_generation_reasoning = "" else: + original_user_query = user_query # Run both pipeline operations concurrently sql_samples_task, instructions_task = await self._run_with_timeout( "SQL pair and instruction retrieval", @@ -364,9 +408,12 @@ async def ask( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") + is_original_analytics_query = self._is_data_analysis_query( + original_user_query + ) if intent in {"GENERAL", "MISLEADING_QUERY"} and ( - self._is_data_analysis_query(user_query) + is_original_analytics_query or self._is_data_analysis_query(rephrased_question or "") ): logger.info( @@ -376,7 +423,16 @@ async def ask( ) intent = "TEXT_TO_SQL" - if rephrased_question: + if is_original_analytics_query: + if rephrased_question and rephrased_question != user_query: + logger.info( + "Ignoring rephrased analytics query from intent classification. original=%s rephrased=%s", + original_user_query, + rephrased_question, + ) + user_query = original_user_query + rephrased_question = original_user_query + elif rephrased_question: user_query = rephrased_question if intent == "MISLEADING_QUERY": @@ -498,11 +554,39 @@ async def ask( ) documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] - table_ddls = [document.get("table_ddl") for document in documents] + table_ddls = [ + document.get("table_ddl", "") or "" for document in documents + ] logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) + if unqueryable_metric_message := self._get_unqueryable_metric_message( + user_query, table_ddls + ): + logger.info( + "ask pipeline - NO_RELEVANT_SQL due to unqueryable metric: %s", + user_query, + ) + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_SQL", + message=unqueryable_metric_message, + ), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = unqueryable_metric_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): @@ -689,7 +773,14 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + if failed_dry_run_result["type"] in { + "TIME_OUT", + "UNSUPPORTED_SQL", + }: + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) break original_sql = failed_dry_run_result["original_sql"] From f5f74bb4d42ad5ce6582807645d40b14c87b2776 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 1 Jun 2026 15:53:18 +0530 Subject: [PATCH 0073/1087] asking --- .../src/pipelines/generation/utils/sql.py | 107 +++++++++++++++++- 1 file changed, 105 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index e863c69358..303aea43d2 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -436,11 +436,52 @@ def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: if not timestamp_expression: return sql + qualified_invented_date_identifier_pattern = re.compile( + r'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' + r'(?:"(?:RepairDate|repair_date|repairDate|date|month_date|event_date)"|\[(?:RepairDate|repair_date|repairDate|date|month_date|event_date)\]|(?:RepairDate|repair_date|repairDate|month_date|event_date))', + re.IGNORECASE, + ) invented_date_identifier_pattern = re.compile( r'(? str: + if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): + return sql + + invented_failure_pattern_id_pattern = re.compile( + r'(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)', + re.IGNORECASE, + ) + + return invented_failure_pattern_id_pattern.sub( + '"dbo_repair_logs"."failure_code"', sql + ) + + +def contains_unsupported_mssql_json_access(sql: str) -> bool: + if re.search(r"(?:->>|->)", sql): + return True + + unsupported_json_functions = ( + "JSON_VALUE", + "JSON_QUERY", + "JSON_EXTRACT", + "JSON_EXTRACT_SCALAR", + "JSON_EXTRACT_ARRAY", + "LAX_BOOL", + "LAX_FLOAT64", + "LAX_INT64", + "LAX_STRING", + ) + function_pattern = r"\b(?:{})\s*\(".format("|".join(unsupported_json_functions)) + return bool(re.search(function_pattern, sql, flags=re.IGNORECASE)) def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: @@ -456,16 +497,28 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: rewritten = sql for bucket, expression in bucket_expressions.items(): + qualified_bucket_pattern = re.compile( + rf'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' + rf'(?:"{bucket}"|\[{bucket}\])', + re.IGNORECASE, + ) select_identifier_pattern = re.compile( rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', re.IGNORECASE, ) + select_qualified_identifier_pattern = re.compile( + rf'(?P\bSELECT\s+|,\s*){qualified_bucket_pattern.pattern}(?P\s*(?:,|\bFROM\b))', + re.IGNORECASE, + ) def replace_select_identifier(match: re.Match[str]) -> str: prefix = match.group("prefix") suffix = match.group("suffix") return f'{prefix}{expression} AS "{bucket}"{suffix}' + rewritten = select_qualified_identifier_pattern.sub( + replace_select_identifier, rewritten + ) rewritten = select_identifier_pattern.sub( replace_select_identifier, rewritten ) @@ -478,6 +531,12 @@ def replace_select_identifier(match: re.Match[str]) -> str: def replace_clause(match: re.Match[str]) -> str: body = match.group("body") for bucket, expression in bucket_expressions.items(): + qualified_bucket_pattern = re.compile( + rf'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' + rf'(?:"{bucket}"|\[{bucket}\])', + re.IGNORECASE, + ) + body = qualified_bucket_pattern.sub(expression, body) body = re.sub( rf'"{bucket}"', expression, @@ -565,6 +624,9 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_invented_repair_relationship_identifiers( + normalized + ) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -610,6 +672,27 @@ async def run( }, } + if normalize_data_source( + data_source + ) == "MSSQL" and contains_unsupported_mssql_json_access( + cleaned_generation_result + ): + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "UNSUPPORTED_SQL", + "error": ( + "Generated SQL uses JSON extraction, but the MSSQL " + "Wren/Ibis runtime does not support JSON operators " + "or JSON extraction functions. Use only first-class " + "columns exposed in the schema." + ), + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, @@ -794,6 +877,7 @@ async def _classify_generation_result( - DON'T USE "TO_CHAR" function in the generated SQL query. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. +- Never invent foreign key columns or relationship fields such as "FailurePatternID", "FailurePatternId", "TicketID", or "ID" unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. - For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. - For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. - Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. @@ -807,8 +891,10 @@ async def _classify_generation_result( - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. - DO NOT use JSON extraction functions or operators such as JSON_VALUE, JSON_QUERY, JSON_EXTRACT, JSON_EXTRACT_SCALAR, JSON_EXTRACT_ARRAY, json_value, json_extract, ->, or ->>. The MSSQL Wren/Ibis runtime does not support them. - If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. +- If a requested metric such as debug hours, risk score, repair cost, or turnaround time is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. - Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. - For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. +- For repair counts grouped by failure category, use explicit exposed fields such as "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. @@ -1110,7 +1196,24 @@ def get_metric_instructions( return instructions -def get_json_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: +_MSSQL_JSON_FIELD_INSTRUCTIONS = """ +#### MSSQL JSON Field Instructions #### +- The target runtime cannot execute JSON extraction from generic JSON/text columns. +- Do not use JSON operators or functions such as ->, ->>, JSON_VALUE, JSON_QUERY, + JSON_EXTRACT, JSON_EXTRACT_SCALAR, LAX_STRING, LAX_INT64, LAX_FLOAT64, or LAX_BOOL. +- If the requested value is only inside a generic JSON/text column such as "data", + do not infer or extract it. Use only first-class columns and calculated fields + that are explicitly exposed in the DATABASE SCHEMA. +""" + + +def get_json_field_instructions( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: + if normalize_data_source(data_source) == "MSSQL": + return _MSSQL_JSON_FIELD_INSTRUCTIONS + if sql_knowledge is not None: return _extract_from_sql_knowledge( sql_knowledge, "json_field_instructions", _DEFAULT_JSON_FIELD_INSTRUCTIONS From 8fcac75219326aa82d091eb8766b2a8323d4e5b0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 1 Jun 2026 16:19:57 +0530 Subject: [PATCH 0074/1087] asking que --- .../src/apollo/server/adaptors/ibisAdaptor.ts | 77 +------ .../apollo/server/services/queryService.ts | 66 +----- .../apollo/server/utils/mssqlSqlNormalizer.ts | 189 ++++++++++++++++++ 3 files changed, 195 insertions(+), 137 deletions(-) create mode 100644 wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 862f9f3732..a6adb47b1d 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -6,6 +6,7 @@ import { Manifest } from '@server/mdl/type'; import * as Errors from '@server/utils/error'; import { getConfig } from '@server/config'; import { toDockerHost } from '@server/utils'; +import { normalizeMssqlSqlForIbis } from '@server/utils/mssqlSqlNormalizer'; import { CompactColumn, CompactTable, @@ -164,76 +165,6 @@ const dataSourceUrlMap: Record = { [SupportedDataSource.DATABRICKS]: 'databricks', }; -const rewriteMssqlDatepartAliasReferences = ( - sql: string, - dataSource: DataSourceName, -): string => { - if (dataSource !== DataSourceName.MSSQL) { - return sql; - } - - sql = sql.replace(/\\"/g, '"'); - - const aliases: Record = {}; - const aliasTargetPattern = - String.raw`(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))`; - const aliasPatterns = [ - new RegExp( - String.raw`\b(DATEPART\(\s*(?:YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b((?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b(DATE_PART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b(EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - ]; - - aliasPatterns.forEach((aliasPattern) => { - for (const match of sql.matchAll(aliasPattern)) { - const expression = match[1]; - const alias = match[3] || match[4] || match[5]; - aliases[alias.toLowerCase()] = expression; - } - }); - - if (!Object.keys(aliases).length) { - return sql; - } - - const clausePattern = - /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; - - return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { - let body = groups?.body || ''; - const placeholders: Record = {}; - - Object.entries(aliases).forEach(([alias, expression]) => { - const placeholder = `__WREN_MSSQL_DATEPART_ALIAS_${Object.keys(placeholders).length}__`; - placeholders[placeholder] = expression; - const escapedAlias = alias.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); - - body = body.replace(new RegExp(`"${escapedAlias}"`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\\\+"${escapedAlias}\\\\+"`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\[${escapedAlias}\\]`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\b${escapedAlias}\\b`, 'gi'), placeholder); - }); - - Object.entries(placeholders).forEach(([placeholder, expression]) => { - body = body.replaceAll(placeholder, expression); - }); - - return `${clause}${body}`; - }); -}; - export interface TableResponse { tables: CompactTable[]; } @@ -340,7 +271,7 @@ export class IbisAdaptor implements IIbisAdaptor { public async getNativeSql(options: IbisDryPlanOptions): Promise { const { dataSource, mdl, sql } = options; const body = { - sql: rewriteMssqlDatepartAliasReferences(sql, dataSource), + sql: normalizeMssqlSqlForIbis(sql, dataSource), manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), }; try { @@ -364,7 +295,7 @@ export class IbisAdaptor implements IIbisAdaptor { options: IbisQueryOptions, ): Promise { const { dataSource, mdl } = options; - query = rewriteMssqlDatepartAliasReferences(query, dataSource); + query = normalizeMssqlSqlForIbis(query, dataSource); const connectionInfo = this.updateConnectionInfo(options.connectionInfo); const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo); const queryString = this.buildQueryString(options); @@ -409,7 +340,7 @@ export class IbisAdaptor implements IIbisAdaptor { options: IbisQueryOptions, ): Promise { const { dataSource, mdl } = options; - query = rewriteMssqlDatepartAliasReferences(query, dataSource); + query = normalizeMssqlSqlForIbis(query, dataSource); const connectionInfo = this.updateConnectionInfo(options.connectionInfo); const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo); const body = { diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index e39e81362d..269da4cbba 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -9,6 +9,7 @@ import { IbisResponse, } from '../adaptors/ibisAdaptor'; import { getLogger } from '@server/utils'; +import { normalizeMssqlSqlForIbis } from '@server/utils/mssqlSqlNormalizer'; import { Project } from '../repositories'; import { PostHogTelemetry, TelemetryEvent } from '../telemetry/telemetry'; @@ -57,69 +58,6 @@ export interface ValidateResponse { message?: string; } -const rewriteMssqlDatepartAliasReferences = (sql: string): string => { - sql = sql.replace(/\\"/g, '"'); - - const aliases: Record = {}; - const aliasTargetPattern = - String.raw`(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))`; - const aliasPatterns = [ - new RegExp( - String.raw`\b(DATEPART\(\s*(?:YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b((?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b(DATE_PART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b(EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - ]; - - aliasPatterns.forEach((aliasPattern) => { - for (const match of sql.matchAll(aliasPattern)) { - const expression = match[1]; - const alias = match[3] || match[4] || match[5]; - aliases[alias.toLowerCase()] = expression; - } - }); - - if (!Object.keys(aliases).length) { - return sql; - } - - const clausePattern = - /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; - - return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { - let body = groups?.body || ''; - const placeholders: Record = {}; - - Object.entries(aliases).forEach(([alias, expression]) => { - const placeholder = `__WREN_MSSQL_DATEPART_ALIAS_${Object.keys(placeholders).length}__`; - placeholders[placeholder] = expression; - const escapedAlias = alias.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); - - body = body.replace(new RegExp(`"${escapedAlias}"`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\\\+"${escapedAlias}\\\\+"`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\[${escapedAlias}\\]`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\b${escapedAlias}\\b`, 'gi'), placeholder); - }); - - Object.entries(placeholders).forEach(([placeholder, expression]) => { - body = body.replaceAll(placeholder, expression); - }); - - return `${clause}${body}`; - }); -}; - export interface IQueryService { preview( sql: string, @@ -148,7 +86,7 @@ const normalizePreviewSqlForIbis = ( return { sql, limit }; } - sql = rewriteMssqlDatepartAliasReferences(sql); + sql = normalizeMssqlSqlForIbis(sql, dataSource); const topMatch = sql.match(/^\s*SELECT\s+(DISTINCT\s+)?TOP\s*\(?\s*(\d+)\s*\)?\s+/i); if (!topMatch) { diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts new file mode 100644 index 0000000000..5d2dcf0fa6 --- /dev/null +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -0,0 +1,189 @@ +import { DataSourceName } from '@server/types'; + +const escapeRegex = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + +const inferMssqlTimestampExpression = (sql: string): string => { + const qualifiedTimestamp = sql.match( + /"([^"]+)"\."(created_at|updated_at|generated_at|created_date|date)"/i, + ); + if (qualifiedTimestamp) { + return qualifiedTimestamp[0]; + } + + const fromTable = sql.match(/\bFROM\s+"([^"]+)"/i); + if (fromTable) { + return `"${fromTable[1]}"."created_at"`; + } + + const bracketedFromTable = sql.match(/\bFROM\s+\[([^\]]+)\]/i); + if (bracketedFromTable) { + return `"${bracketedFromTable[1]}"."created_at"`; + } + + return '"created_at"'; +}; + +const replaceInventedDateFields = (sql: string): string => { + const timestampExpression = inferMssqlTimestampExpression(sql); + const inventedDateFields = [ + 'RepairDate', + 'repairDate', + 'repair_date', + 'Repair_Date', + 'EventDate', + 'event_date', + 'Date', + 'date', + ]; + + inventedDateFields.forEach((field) => { + const escaped = escapeRegex(field); + sql = sql.replace( + new RegExp(String.raw`(?:"[^"]+"\.)"${escaped}"`, 'gi'), + timestampExpression, + ); + sql = sql.replace(new RegExp(String.raw`"${escaped}"`, 'gi'), timestampExpression); + sql = sql.replace(new RegExp(String.raw`\[${escaped}\]`, 'gi'), timestampExpression); + }); + + return sql; +}; + +const replaceInventedTimeBuckets = (sql: string): string => { + const timestampExpression = inferMssqlTimestampExpression(sql); + const bucketExpressions: Record = { + YEAR: `DATEPART(YEAR, ${timestampExpression})`, + MONTH: `DATEPART(MONTH, ${timestampExpression})`, + DAY: `DATEPART(DAY, ${timestampExpression})`, + }; + + sql = sql.replace(/\bSELECT\b(?.*?)(?=\bFROM\b)/is, (match, _body, _offset, _source, groups) => { + let body = groups?.body || ''; + Object.entries(bucketExpressions).forEach(([bucket, expression]) => { + const alias = bucket.toLowerCase(); + body = body.replace( + new RegExp( + String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"${bucket}"|(?:\[[^\]]+\]\.)\[${bucket}\]|"${bucket}"|\[${bucket}\])(?=\s*(?:,|$))`, + 'gi', + ), + `$1 ${expression} AS "${alias}"`, + ); + }); + return `SELECT${body}`; + }); + + Object.entries(bucketExpressions).forEach(([bucket, expression]) => { + sql = sql.replace( + new RegExp(String.raw`(?:"[^"]+"\.)"${bucket}"`, 'gi'), + expression, + ); + sql = sql.replace( + new RegExp(String.raw`(?:\[[^\]]+\]\.)\[${bucket}\]`, 'gi'), + expression, + ); + }); + + const clausePattern = + /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; + sql = sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { + let body = groups?.body || ''; + Object.entries(bucketExpressions).forEach(([bucket, expression]) => { + body = body.replace(new RegExp(String.raw`"${bucket}"`, 'gi'), expression); + body = body.replace(new RegExp(String.raw`\[${bucket}\]`, 'gi'), expression); + }); + return `${clause}${body}`; + }); + + return sql; +}; + +export const normalizeMssqlGeneratedSqlFields = ( + sql: string, + dataSource: DataSourceName, +): string => { + if (dataSource !== DataSourceName.MSSQL) { + return sql; + } + + sql = sql.replace(/\\"/g, '"'); + sql = replaceInventedDateFields(sql); + sql = replaceInventedTimeBuckets(sql); + return sql; +}; + +export const rewriteMssqlDatepartAliasReferences = ( + sql: string, + dataSource: DataSourceName, +): string => { + if (dataSource !== DataSourceName.MSSQL) { + return sql; + } + + sql = sql.replace(/\\"/g, '"'); + + const aliases: Record = {}; + const aliasTargetPattern = + String.raw`(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))`; + const aliasPatterns = [ + new RegExp( + String.raw`\b(DATEPART\(\s*(?:YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b((?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b(DATE_PART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + new RegExp( + String.raw`\b(EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + 'gi', + ), + ]; + + aliasPatterns.forEach((aliasPattern) => { + for (const match of sql.matchAll(aliasPattern)) { + const expression = match[1]; + const alias = match[3] || match[4] || match[5]; + aliases[alias.toLowerCase()] = expression; + } + }); + + if (!Object.keys(aliases).length) { + return sql; + } + + const clausePattern = + /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; + + return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { + let body = groups?.body || ''; + const placeholders: Record = {}; + + Object.entries(aliases).forEach(([alias, expression]) => { + const placeholder = `__WREN_MSSQL_DATEPART_ALIAS_${Object.keys(placeholders).length}__`; + placeholders[placeholder] = expression; + const escapedAlias = escapeRegex(alias); + + body = body.replace(new RegExp(`"${escapedAlias}"`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\\\+"${escapedAlias}\\\\+"`, 'gi'), placeholder); + body = body.replace(new RegExp(`\\[${escapedAlias}\\]`, 'gi'), placeholder); + }); + + Object.entries(placeholders).forEach(([placeholder, expression]) => { + body = body.replaceAll(placeholder, expression); + }); + + return `${clause}${body}`; + }); +}; + +export const normalizeMssqlSqlForIbis = ( + sql: string, + dataSource: DataSourceName, +): string => { + sql = normalizeMssqlGeneratedSqlFields(sql, dataSource); + return rewriteMssqlDatepartAliasReferences(sql, dataSource); +}; From cc4b5b52c8c824b6de5438cf6c272768045359f5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 1 Jun 2026 17:28:52 +0530 Subject: [PATCH 0075/1087] asking questions --- .../src/pipelines/generation/utils/sql.py | 122 ++++- .../pipelines/generation/test_sql_utils.py | 493 ++++++++++++++++++ .../apollo/server/utils/mssqlSqlNormalizer.ts | 78 ++- 3 files changed, 681 insertions(+), 12 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 303aea43d2..c75539dab6 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -407,7 +407,7 @@ def replace_subtraction(match: re.Match[str]) -> str: def _infer_mssql_timestamp_expression(sql: str) -> str | None: timestamp_column_pattern = re.compile( - r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at)")', + r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at|DateIn|DateOut|FailedAt)")', re.IGNORECASE, ) if match := timestamp_column_pattern.search(sql): @@ -419,12 +419,16 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: ) if match := table_pattern.search(sql): table_name = match.group(1) - normalized_table_name = table_name.strip('"[]').lower() + raw_table_name = table_name.strip('"[]') + normalized_table_name = raw_table_name.lower() + quoted_table_name = f'"{raw_table_name}"' + if normalized_table_name == "dbo_debugentries": + return f'{quoted_table_name}."DateIn"' if "report" in normalized_table_name: return f'{table_name}."generated_at"' if any( token in normalized_table_name - for token in ("repair", "ticket", "debug", "event", "log") + for token in ("repair", "ticket", "event", "log") ): return f'{table_name}."created_at"' @@ -452,18 +456,105 @@ def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: def _rewrite_mssql_invented_repair_relationship_identifiers(sql: str) -> str: - if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): - return sql + rewritten = sql - invented_failure_pattern_id_pattern = re.compile( - r'(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)', - re.IGNORECASE, - ) + if re.search(r"\bdbo_repair_logs\b", rewritten, flags=re.IGNORECASE): + invented_failure_pattern_id_pattern = re.compile( + r'(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)', + re.IGNORECASE, + ) + rewritten = invented_failure_pattern_id_pattern.sub( + '"dbo_repair_logs"."failure_code"', rewritten + ) + + if re.search(r"\bdbo_DebugEntries\b", rewritten, flags=re.IGNORECASE) and re.search( + r"\bdbo_failure_patterns\b", rewritten, flags=re.IGNORECASE + ): + invented_debug_failure_pattern_pattern = re.compile( + r'(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)', + re.IGNORECASE, + ) + rewritten = invented_debug_failure_pattern_pattern.sub( + '"dbo_DebugEntries"."FailureSys"', rewritten + ) + + debug_id_identifier = ( + r'(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)' + r'\s*\.\s*(?:"DebugEntryId"|\[DebugEntryId\]|DebugEntryId)' + ) + failure_pattern_id_identifier = ( + r'(?:"dbo_failure_patterns"|\[dbo_failure_patterns\]|dbo_failure_patterns)' + r'\s*\.\s*(?:"id"|\[id\]|id)' + ) + debug_id_to_failure_pattern_pattern = re.compile( + rf"{debug_id_identifier}\s*=\s*{failure_pattern_id_identifier}", + re.IGNORECASE, + ) + failure_pattern_to_debug_id_pattern = re.compile( + rf"{failure_pattern_id_identifier}\s*=\s*{debug_id_identifier}", + re.IGNORECASE, + ) + rewritten = debug_id_to_failure_pattern_pattern.sub( + '"dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id"', + rewritten, + ) + rewritten = failure_pattern_to_debug_id_pattern.sub( + '"dbo_failure_patterns"."id" = "dbo_DebugEntries"."FailureSys"', + rewritten, + ) - return invented_failure_pattern_id_pattern.sub( - '"dbo_repair_logs"."failure_code"', sql + return rewritten + + +def _rewrite_mssql_invented_pcb_throughput_identifiers(sql: str) -> str: + rewritten = sql + manufacturing_unit_pattern = ( + r'(?:"ManufacturingUnit"|"Manufacturing_Unit"|"manufacturing_unit"|' + r'\[ManufacturingUnit\]|\[Manufacturing_Unit\]|\[manufacturing_unit\]|' + r'ManufacturingUnit|Manufacturing_Unit|manufacturing_unit)' ) + if re.search(r"\bdbo_DebugEntries\b", rewritten, flags=re.IGNORECASE): + debug_table_pattern = ( + r'(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)\s*\.\s*' + ) + rewritten = re.sub( + rf"{debug_table_pattern}{manufacturing_unit_pattern}", + '"dbo_DebugEntries"."BusinessUnit"', + rewritten, + flags=re.IGNORECASE, + ) + + if re.search(r"\bdbo_repair_logs\b", rewritten, flags=re.IGNORECASE) and re.search( + manufacturing_unit_pattern, rewritten, flags=re.IGNORECASE + ): + rewritten = re.sub( + r'(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)', + '"dbo_DebugEntries"', + rewritten, + flags=re.IGNORECASE, + ) + rewritten = re.sub( + rf'"dbo_DebugEntries"\s*\.\s*{manufacturing_unit_pattern}', + '"dbo_DebugEntries"."BusinessUnit"', + rewritten, + flags=re.IGNORECASE, + ) + rewritten = re.sub( + r'"dbo_DebugEntries"\s*\.\s*(?:"id"|\[id\]|id)', + '"dbo_DebugEntries"."DebugEntryId"', + rewritten, + flags=re.IGNORECASE, + ) + rewritten = re.sub( + r'"dbo_DebugEntries"\s*\.\s*(?:"created_at"|"updated_at"|\[created_at\]|\[updated_at\]|created_at|updated_at)', + '"dbo_DebugEntries"."DateIn"', + rewritten, + flags=re.IGNORECASE, + ) + + return rewritten + def contains_unsupported_mssql_json_access(sql: str) -> bool: if re.search(r"(?:->>|->)", sql): @@ -627,6 +718,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_invented_repair_relationship_identifiers( normalized ) + normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -895,6 +987,14 @@ async def _classify_generation_result( - Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. - For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. - For repair counts grouped by failure category, use explicit exposed fields such as "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. +- For PCB/debug-entry failure charts, do not join "dbo_DebugEntries"."DebugEntryId" to "dbo_failure_patterns"."id"; those fields have incompatible types. If both "dbo_DebugEntries"."FailureSys" and "dbo_failure_patterns"."id" exist, join "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". +- For PCB synced database questions: + - Use "dbo_DebugEntries" for debug/PCB event records when the schema contains it. + - Use columns such as "Material", "WorkOrder", "SerialNumber", "FailedAt", "DateIn", "DateOut", "Hours", "Priority", "Actions", "Notes", and "FailureSys" only when they appear in the schema. + - Use "dbo_failure_patterns" for failure names, categories, severity, trend, occurrence counts, daily pattern summaries, and cost impact when those columns appear in the schema. + - For throughput trends across manufacturing/business units, use "dbo_DebugEntries"."BusinessUnit" as the unit dimension and a real debug-entry timestamp such as "dbo_DebugEntries"."DateIn" or "dbo_DebugEntries"."FailedAt" for the trend bucket. Do not use "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH", or invented manufacturing/date fields. + - For top/common PCB failure questions, prefer grouping by "dbo_failure_patterns"."name" or "dbo_failure_patterns"."category" and counting "dbo_DebugEntries"."DebugEntryId" after joining "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". + - If a useful aggregate already exists in "dbo_failure_patterns" such as "occurrences", it can be used directly for top failure pattern questions without joining event rows. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py new file mode 100644 index 0000000000..0489449625 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -0,0 +1,493 @@ +from src.pipelines.generation.utils.sql import ( + contains_unsupported_mssql_json_access, + construct_valid_table_names, + extract_sql_generation_result, + get_json_field_instructions, + get_metric_instructions, + normalize_data_source, + normalize_generation_result_sql, + get_sql_generation_system_prompt, + get_text_to_sql_rules, +) + + +def test_construct_valid_table_names_from_schema_documents(): + documents = [ + 'CREATE TABLE repair_logs ("id" INTEGER);', + '/* comment */ CREATE TABLE "employees" ("emp_no" INTEGER);', + ] + + assert construct_valid_table_names(documents) == ["employees", "repair_logs"] + + +def test_extract_sql_generation_result_from_json_payload(): + result = '{"sql": "SELECT COUNT(*) AS repair_count FROM repairs;"}' + + assert ( + extract_sql_generation_result(result) + == "SELECT COUNT(*) AS repair_count FROM repairs" + ) + + +def test_extract_sql_generation_result_from_prose_wrapped_sql(): + result = ( + "The SQL query is: SELECT DATEPART(YEAR, created_at) AS year, " + "COUNT(*) AS repair_count FROM repairs GROUP BY DATEPART(YEAR, created_at);" + ) + + assert extract_sql_generation_result(result) == ( + "SELECT DATEPART(YEAR, created_at) AS year, COUNT(*) AS repair_count " + "FROM repairs GROUP BY DATEPART(YEAR, created_at)" + ) + + +def test_extract_sql_generation_result_from_fenced_sql(): + result = """ + Here is the query: + ```sql + SELECT id FROM repairs; + ``` + """ + + assert extract_sql_generation_result(result) == "SELECT id FROM repairs" + + +def test_extract_sql_generation_result_from_prose_wrapped_json(): + result = 'Here is the result:\n{"sql": "SELECT id FROM repairs;"}' + + assert extract_sql_generation_result(result) == "SELECT id FROM repairs" + + +def test_get_text_to_sql_rules_adds_mssql_specific_constraints(): + rules = get_text_to_sql_rules(data_source="MSSQL") + + assert "The target database is MSSQL." in rules + assert "DATEPART(YEAR, )" in rules + assert "DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET" in rules + assert "TO_UNIXTIME" in rules + assert "Do not subtract timestamp/date columns directly" in rules + assert "TO_TIMESTAMP_MILLIS" in rules + assert "DO NOT use PostgreSQL-style or Trino-style date syntax" in rules + assert "DO NOT use JSON extraction functions or operators" in rules + assert "JSON_VALUE" in rules + assert "JSON_EXTRACT" in rules + assert "->>" in rules + assert "do not assume keys inside it are queryable" in rules + assert "Never invent JSON-derived columns" in rules + assert "Resolve relative time phrases" in rules + assert "Do not include helper ranking columns" in rules + assert "prefer SELECT TOP (N)" in rules + assert "FailurePatternID" in rules + assert "failure_code" in rules + assert "CURRENT_DATE - INTERVAL '1 month'" not in rules + + +def test_get_json_field_instructions_for_mssql_disables_json_extraction(): + instructions = get_json_field_instructions(data_source="MSSQL") + + assert "cannot execute JSON extraction" in instructions + assert "JSON_VALUE" in instructions + assert "->>" in instructions + assert "Use only first-class columns" in instructions + assert "LAX_STRING(JSON_QUERY" not in instructions + + +def test_contains_unsupported_mssql_json_access_detects_json_syntax(): + assert contains_unsupported_mssql_json_access( + 'SELECT "data" ->> \'AttemptNumber\' FROM "dbo_repair_logs"' + ) + assert contains_unsupported_mssql_json_access( + 'SELECT JSON_VALUE("data", \'$.AttemptNumber\') FROM "dbo_repair_logs"' + ) + assert not contains_unsupported_mssql_json_access( + 'SELECT "created_at", "status" FROM "dbo_repair_logs"' + ) + + +def test_get_metric_instructions_for_mssql_avoids_date_trunc_example(): + instructions = get_metric_instructions(data_source="MSSQL") + + assert "DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month')" not in instructions + assert "Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE" in instructions + assert "DATEPART(YEAR, )" in instructions + + +def test_get_sql_generation_system_prompt_uses_data_source_specific_rules(): + prompt = get_sql_generation_system_prompt(data_source="MSSQL") + + assert "The target database is MSSQL." in prompt + assert "DATEPART(YEAR, )" in prompt + + +def test_normalize_generation_result_sql_rewrites_common_mssql_time_patterns(): + sql = """ + SELECT + DATEPART(YEAR, "created_at") AS "year", + DATEPART(MONTH, "created_at") AS "month", + COUNT("id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY DATEPART(YEAR, "created_at"), DATEPART(MONTH, "created_at") + ORDER BY "year" ASC NULLS LAST, "month" ASC NULLS LAST + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert 'DATEPART(\'YEAR\', "created_at")' not in normalized + assert 'DATEPART(\'MONTH\', "created_at")' not in normalized + assert "NULLS LAST" not in normalized + assert 'DATEPART(YEAR, "created_at")' in normalized + assert 'DATEPART(MONTH, "created_at")' in normalized + + +def test_normalize_generation_result_sql_rewrites_common_mssql_dateadd_patterns(): + sql = """ + SELECT + DATEADD(month, DATEDIFF(month, 0, "created_at"), 0) AS "month_start", + COUNT("id") AS "repair_count" + FROM "dbo_repair_logs" + WHERE "created_at" >= DATEADD(month, -12, GETDATE()) + AND "created_at" < DATEADD(month, DATEDIFF(month, 0, GETDATE()), 0) + GROUP BY DATEADD(month, DATEDIFF(month, 0, "created_at"), 0) + ORDER BY "month_start" ASC NULLS LAST + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "DATEADD(" not in normalized + assert "DATEDIFF(" not in normalized + assert "NULLS LAST" not in normalized + assert 'DATEPART(YEAR, "created_at")' in normalized + assert 'DATEPART(MONTH, "created_at")' in normalized + + +def test_normalize_data_source_maps_sql_server_aliases_to_mssql(): + assert normalize_data_source("sqlserver") == "MSSQL" + assert normalize_data_source("SQL Server") == "MSSQL" + + +def test_normalize_generation_result_sql_rewrites_nested_temporal_patterns_for_mssql(): + sql = """ + SELECT + DATE_PART('YEAR', CAST("created_at" AS DATETIME)) AS "year", + DATE_TRUNC('MONTH', CAST("created_at" AS DATETIME)) AS "month_bucket", + EXTRACT(DAY FROM CAST("created_at" AS DATETIME)) AS "day_of_month" + FROM "dbo_repair_logs" + ORDER BY "month_bucket" ASC NULLS LAST + """ + + normalized = normalize_generation_result_sql(sql, data_source="sqlserver") + + assert "DATE_PART(" not in normalized + assert "DATE_TRUNC(" not in normalized + assert "EXTRACT(" not in normalized + assert "NULLS LAST" not in normalized + assert 'DATEPART(YEAR, CAST("created_at" AS DATETIME))' in normalized + assert 'DATEPART(MONTH, CAST("created_at" AS DATETIME))' in normalized + assert 'DATEPART(DAY, CAST("created_at" AS DATETIME))' in normalized + + +def test_normalize_generation_result_sql_rewrites_to_timestamp_for_mssql(): + sql = """ + SELECT + DATEPART(YEAR, TO_TIMESTAMP("created_at")) AS "year", + COUNT("id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY DATEPART(YEAR, TO_TIMESTAMP("created_at")) + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "TO_TIMESTAMP(" not in normalized + assert 'DATEPART(YEAR, CAST("created_at" AS DATETIME))' in normalized + + +def test_normalize_generation_result_sql_rewrites_to_timestamp_variants_for_mssql(): + sql = """ + SELECT + TO_TIMESTAMP_MILLIS("created_at_ms") AS "created_at", + TO_TIMESTAMP_SECONDS("closed_at_sec") AS "closed_at" + FROM "dbo_repair_logs" + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "TO_TIMESTAMP_MILLIS(" not in normalized + assert "TO_TIMESTAMP_SECONDS(" not in normalized + assert 'CAST("created_at_ms" AS DATETIME)' in normalized + assert 'CAST("closed_at_sec" AS DATETIME)' in normalized + + +def test_normalize_generation_result_sql_rewrites_mssql_datepart_alias_references(): + sql = """ + SELECT + DATEPART(YEAR, "created_at") AS "YEAR", + DATEPART(MONTH, "created_at") AS "MONTH", + COUNT("id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY "YEAR", "MONTH" + ORDER BY "YEAR" ASC, "MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert 'GROUP BY "YEAR"' not in normalized + assert 'ORDER BY "YEAR"' not in normalized + assert 'DATEPART(YEAR, "created_at") AS "YEAR"' in normalized + assert ( + 'GROUP BY DATEPART(YEAR, "created_at"), DATEPART(MONTH, "created_at")' + in normalized + ) + assert ( + 'ORDER BY DATEPART(YEAR, "created_at") ASC, DATEPART(MONTH, "created_at") ASC' + in normalized + ) + + +def test_normalize_generation_result_sql_rewrites_unquoted_mssql_datepart_alias_references(): + sql = """ + SELECT + DATEPART(YEAR, "created_at") AS YEAR, + DATEPART(MONTH, "created_at") AS MONTH, + COUNT("id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY YEAR, MONTH + ORDER BY YEAR ASC, MONTH ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "GROUP BY YEAR" not in normalized + assert "ORDER BY YEAR" not in normalized + assert 'DATEPART(YEAR, "created_at") AS YEAR' in normalized + assert ( + 'GROUP BY DATEPART(YEAR, "created_at"), DATEPART(MONTH, "created_at")' + in normalized + ) + assert ( + 'ORDER BY DATEPART(YEAR, "created_at") ASC, DATEPART(MONTH, "created_at") ASC' + in normalized + ) + + +def test_normalize_generation_result_sql_rewrites_invented_repair_date_for_mssql(): + sql = """ + SELECT + DATEPART(YEAR, "RepairDate") AS "YEAR", + DATEPART(MONTH, "RepairDate") AS "MONTH", + COUNT("dbo_repair_logs"."id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY DATEPART(YEAR, "RepairDate"), DATEPART(MONTH, "RepairDate") + ORDER BY "YEAR" ASC, "MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert '"RepairDate"' not in normalized + assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "YEAR"' in normalized + assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "MONTH"' in normalized + assert ( + 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at")' + in normalized + ) + assert ( + 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + in normalized + ) + + +def test_normalize_generation_result_sql_rewrites_qualified_invented_repair_date_for_mssql(): + sql = """ + SELECT + DATEPART(YEAR, "dbo_repair_logs"."RepairDate") AS "YEAR", + DATEPART(MONTH, "dbo_repair_logs"."RepairDate") AS "MONTH", + COUNT("dbo_repair_logs"."id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY DATEPART(YEAR, "dbo_repair_logs"."RepairDate"), + DATEPART(MONTH, "dbo_repair_logs"."RepairDate") + ORDER BY "YEAR" ASC, "MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "RepairDate" not in normalized + assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "YEAR"' in normalized + assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "MONTH"' in normalized + + +def test_normalize_generation_result_sql_rewrites_invented_repair_failure_pattern_id_for_mssql(): + sql = """ + SELECT + "dbo_failure_patterns"."category" AS "failure_category", + COUNT("dbo_repair_logs"."id") AS "repair_count" + FROM "dbo_repair_logs" + JOIN "dbo_failure_patterns" + ON "dbo_repair_logs"."FailurePatternID" = "dbo_failure_patterns"."id" + GROUP BY "dbo_failure_patterns"."category" + ORDER BY "repair_count" DESC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "FailurePatternID" not in normalized + assert ( + '"dbo_repair_logs"."failure_code" = "dbo_failure_patterns"."id"' + in normalized + ) + + +def test_normalize_generation_result_sql_rewrites_debug_entry_failure_pattern_join_for_mssql(): + sql = """ + SELECT TOP 10 + "dbo_failure_patterns"."category" AS "FailureCategory", + COUNT_BIG(1) AS "FailureCount" + FROM "dbo_DebugEntries" + INNER JOIN "dbo_failure_patterns" + ON "dbo_DebugEntries"."DebugEntryId" = "dbo_failure_patterns"."id" + GROUP BY "dbo_failure_patterns"."category" + ORDER BY "FailureCount" DESC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert '"dbo_DebugEntries"."DebugEntryId" = "dbo_failure_patterns"."id"' not in normalized + assert ( + '"dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id"' + in normalized + ) + + +def test_normalize_generation_result_sql_rewrites_pcb_throughput_repair_log_fields_for_mssql(): + sql = """ + SELECT + "dbo_repair_logs"."ManufacturingUnit" AS "manufacturing_unit", + "dbo_repair_logs"."MONTH", + COUNT("dbo_repair_logs"."id") AS "throughput" + FROM "dbo_repair_logs" + GROUP BY "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH" + ORDER BY "dbo_repair_logs"."MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "dbo_repair_logs" not in normalized + assert "ManufacturingUnit" not in normalized + assert '"dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit"' in normalized + assert '"dbo_DebugEntries"."DebugEntryId"' in normalized + assert 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month"' in normalized + assert 'GROUP BY "dbo_DebugEntries"."BusinessUnit"' in normalized + assert 'ORDER BY DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' in normalized + + +def test_normalize_generation_result_sql_rewrites_bare_month_field_for_mssql(): + sql = """ + SELECT + "MONTH", + COUNT("dbo_repair_logs"."id") AS "repair_count" + FROM "dbo_repair_logs" + WHERE "dbo_repair_logs"."created_at" >= '2025-05-01 00:00:00' + GROUP BY "MONTH" + ORDER BY "MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert 'SELECT "MONTH"' not in normalized + assert 'GROUP BY "MONTH"' not in normalized + assert 'ORDER BY "MONTH"' not in normalized + assert ( + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' + in normalized + ) + assert 'GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized + assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized + + +def test_normalize_generation_result_sql_rewrites_qualified_month_field_for_mssql(): + sql = """ + SELECT + "dbo_repair_logs"."MONTH", + COUNT("dbo_repair_logs"."id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY "dbo_repair_logs"."MONTH" + ORDER BY "dbo_repair_logs"."MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert '"dbo_repair_logs"."MONTH"' not in normalized + assert ( + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' + in normalized + ) + assert 'GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized + assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized + + +def test_normalize_generation_result_sql_rewrites_bare_year_for_report_charts(): + sql = """ + SELECT + "YEAR", + COUNT("dbo_reports"."id") AS "report_count" + FROM "dbo_reports" + GROUP BY "YEAR" + ORDER BY "YEAR" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert 'SELECT "YEAR"' not in normalized + assert 'GROUP BY "YEAR"' not in normalized + assert 'ORDER BY "YEAR"' not in normalized + assert 'DATEPART(YEAR, "dbo_reports"."generated_at") AS "year"' in normalized + assert 'GROUP BY DATEPART(YEAR, "dbo_reports"."generated_at")' in normalized + assert 'ORDER BY DATEPART(YEAR, "dbo_reports"."generated_at") ASC' in normalized + + +def test_normalize_generation_result_sql_rewrites_timestamp_casts_for_mssql(): + sql = """ + SELECT COUNT("id") + FROM "dbo_repair_logs" + WHERE CAST("created_at" AS TIMESTAMP) >= CAST('2026-01-01 00:00:00' AS TIMESTAMP) + """ + + normalized = normalize_generation_result_sql(sql, data_source="sqlserver") + + assert " AS TIMESTAMP" not in normalized + assert 'CAST("created_at" AS DATETIME)' in normalized + assert "CAST('2026-01-01 00:00:00' AS DATETIME)" in normalized + + +def test_normalize_generation_result_sql_strips_to_unixtime_for_mssql(): + sql = """ + SELECT + TO_UNIXTIME(CAST("created_at" AS TIMESTAMP)) AS "created_at_unix", + AVG("repair_cost") AS "avg_repair_cost" + FROM "dbo_repair_logs" + GROUP BY TO_UNIXTIME(CAST("created_at" AS TIMESTAMP)) + ORDER BY "created_at_unix" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "TO_UNIXTIME(" not in normalized + assert 'CAST("created_at" AS DATETIME) AS "created_at_unix"' in normalized + assert 'GROUP BY CAST("created_at" AS DATETIME)' in normalized + + +def test_normalize_generation_result_sql_rewrites_timestamp_subtraction_for_mssql(): + sql = """ + SELECT + "updated_at" - "created_at" AS "turnaround_seconds", + "repair_cost" + FROM "dbo_repair_logs" + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert '"updated_at" - "created_at"' not in normalized + assert ( + 'DATEDIFF(\'second\', "created_at", "updated_at") AS "turnaround_seconds"' + in normalized + ) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 5d2dcf0fa6..77d586b05e 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -4,7 +4,7 @@ const escapeRegex = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$ const inferMssqlTimestampExpression = (sql: string): string => { const qualifiedTimestamp = sql.match( - /"([^"]+)"\."(created_at|updated_at|generated_at|created_date|date)"/i, + /"([^"]+)"\."(created_at|updated_at|generated_at|created_date|date|DateIn|DateOut|FailedAt)"/i, ); if (qualifiedTimestamp) { return qualifiedTimestamp[0]; @@ -12,11 +12,17 @@ const inferMssqlTimestampExpression = (sql: string): string => { const fromTable = sql.match(/\bFROM\s+"([^"]+)"/i); if (fromTable) { + if (fromTable[1].toLowerCase() === 'dbo_debugentries') { + return `"${fromTable[1]}"."DateIn"`; + } return `"${fromTable[1]}"."created_at"`; } const bracketedFromTable = sql.match(/\bFROM\s+\[([^\]]+)\]/i); if (bracketedFromTable) { + if (bracketedFromTable[1].toLowerCase() === 'dbo_debugentries') { + return `"${bracketedFromTable[1]}"."DateIn"`; + } return `"${bracketedFromTable[1]}"."created_at"`; } @@ -97,6 +103,74 @@ const replaceInventedTimeBuckets = (sql: string): string => { return sql; }; +const replaceBadFailurePatternJoins = (sql: string): string => { + if ( + !/\bdbo_DebugEntries\b/i.test(sql) || + !/\bdbo_failure_patterns\b/i.test(sql) + ) { + return sql; + } + + const debugTable = String.raw`(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)`; + const failurePatternTable = String.raw`(?:"dbo_failure_patterns"|\[dbo_failure_patterns\]|dbo_failure_patterns)`; + const debugEntryId = String.raw`${debugTable}\s*\.\s*(?:"DebugEntryId"|\[DebugEntryId\]|DebugEntryId)`; + const debugFailureSys = '"dbo_DebugEntries"."FailureSys"'; + const failurePatternId = String.raw`${failurePatternTable}\s*\.\s*(?:"id"|\[id\]|id)`; + const normalizedFailurePatternId = '"dbo_failure_patterns"."id"'; + + sql = sql.replace( + new RegExp(String.raw`${debugEntryId}\s*=\s*${failurePatternId}`, 'gi'), + `${debugFailureSys} = ${normalizedFailurePatternId}`, + ); + sql = sql.replace( + new RegExp(String.raw`${failurePatternId}\s*=\s*${debugEntryId}`, 'gi'), + `${normalizedFailurePatternId} = ${debugFailureSys}`, + ); + sql = sql.replace( + new RegExp( + String.raw`${debugTable}\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)`, + 'gi', + ), + debugFailureSys, + ); + + return sql; +}; + +const replacePcbThroughputFields = (sql: string): string => { + const manufacturingUnitField = + String.raw`(?:"ManufacturingUnit"|"Manufacturing_Unit"|"manufacturing_unit"|\[ManufacturingUnit\]|\[Manufacturing_Unit\]|\[manufacturing_unit\]|ManufacturingUnit|Manufacturing_Unit|manufacturing_unit)`; + + if (/\bdbo_DebugEntries\b/i.test(sql)) { + const debugTable = String.raw`(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)`; + sql = sql.replace( + new RegExp(String.raw`${debugTable}\s*\.\s*${manufacturingUnitField}`, 'gi'), + '"dbo_DebugEntries"."BusinessUnit"', + ); + } + + if (/\bdbo_repair_logs\b/i.test(sql) && new RegExp(manufacturingUnitField, 'i').test(sql)) { + sql = sql.replace( + /(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)/gi, + '"dbo_DebugEntries"', + ); + sql = sql.replace( + new RegExp(String.raw`"dbo_DebugEntries"\s*\.\s*${manufacturingUnitField}`, 'gi'), + '"dbo_DebugEntries"."BusinessUnit"', + ); + sql = sql.replace( + /"dbo_DebugEntries"\s*\.\s*(?:"id"|\[id\]|id)/gi, + '"dbo_DebugEntries"."DebugEntryId"', + ); + sql = sql.replace( + /"dbo_DebugEntries"\s*\.\s*(?:"created_at"|"updated_at"|\[created_at\]|\[updated_at\]|created_at|updated_at)/gi, + '"dbo_DebugEntries"."DateIn"', + ); + } + + return sql; +}; + export const normalizeMssqlGeneratedSqlFields = ( sql: string, dataSource: DataSourceName, @@ -107,7 +181,9 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = sql.replace(/\\"/g, '"'); sql = replaceInventedDateFields(sql); + sql = replacePcbThroughputFields(sql); sql = replaceInventedTimeBuckets(sql); + sql = replaceBadFailurePatternJoins(sql); return sql; }; From b9b67f21bbbdf8f6f24ccdcddcdb086d01a1aafb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 1 Jun 2026 17:51:47 +0530 Subject: [PATCH 0076/1087] month --- .../src/pipelines/generation/utils/sql.py | 4 ++-- .../pipelines/generation/test_sql_utils.py | 21 +++++++++++++++++++ .../apollo/server/utils/mssqlSqlNormalizer.ts | 10 ++++++--- 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c75539dab6..8df9fd7fec 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -590,7 +590,7 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: for bucket, expression in bucket_expressions.items(): qualified_bucket_pattern = re.compile( rf'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' - rf'(?:"{bucket}"|\[{bucket}\])', + rf'(?:"{bucket}"|\[{bucket}\]|{bucket})', re.IGNORECASE, ) select_identifier_pattern = re.compile( @@ -624,7 +624,7 @@ def replace_clause(match: re.Match[str]) -> str: for bucket, expression in bucket_expressions.items(): qualified_bucket_pattern = re.compile( rf'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' - rf'(?:"{bucket}"|\[{bucket}\])', + rf'(?:"{bucket}"|\[{bucket}\]|{bucket})', re.IGNORECASE, ) body = qualified_bucket_pattern.sub(expression, body) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 0489449625..fea082a408 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -425,6 +425,27 @@ def test_normalize_generation_result_sql_rewrites_qualified_month_field_for_mssq assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized +def test_normalize_generation_result_sql_rewrites_unquoted_qualified_month_field_for_mssql(): + sql = """ + SELECT + dbo_repair_logs.MONTH, + COUNT(dbo_repair_logs.id) AS "repair_count" + FROM dbo_repair_logs + GROUP BY dbo_repair_logs.MONTH + ORDER BY dbo_repair_logs.MONTH ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "dbo_repair_logs.MONTH" not in normalized + assert ( + 'DATEPART(MONTH, dbo_repair_logs."created_at") AS "month"' + in normalized + ) + assert 'GROUP BY DATEPART(MONTH, dbo_repair_logs."created_at")' in normalized + assert 'ORDER BY DATEPART(MONTH, dbo_repair_logs."created_at") ASC' in normalized + + def test_normalize_generation_result_sql_rewrites_bare_year_for_report_charts(): sql = """ SELECT diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 77d586b05e..62829fcbf7 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -69,7 +69,7 @@ const replaceInventedTimeBuckets = (sql: string): string => { const alias = bucket.toLowerCase(); body = body.replace( new RegExp( - String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"${bucket}"|(?:\[[^\]]+\]\.)\[${bucket}\]|"${bucket}"|\[${bucket}\])(?=\s*(?:,|$))`, + String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"?${bucket}"?|(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})|\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b|"${bucket}"|\[${bucket}\])(?=\s*(?:,|$))`, 'gi', ), `$1 ${expression} AS "${alias}"`, @@ -80,11 +80,15 @@ const replaceInventedTimeBuckets = (sql: string): string => { Object.entries(bucketExpressions).forEach(([bucket, expression]) => { sql = sql.replace( - new RegExp(String.raw`(?:"[^"]+"\.)"${bucket}"`, 'gi'), + new RegExp(String.raw`(?:"[^"]+"\.)"?${bucket}"?`, 'gi'), expression, ); sql = sql.replace( - new RegExp(String.raw`(?:\[[^\]]+\]\.)\[${bucket}\]`, 'gi'), + new RegExp(String.raw`(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})`, 'gi'), + expression, + ); + sql = sql.replace( + new RegExp(String.raw`\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b`, 'gi'), expression, ); }); From 8b972104462196fa27f04af7bea9868f7d22f8db Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 1 Jun 2026 18:28:48 +0530 Subject: [PATCH 0077/1087] month/day --- .../src/pipelines/generation/utils/sql.py | 18 +++++++++++++++ .../pipelines/generation/test_sql_utils.py | 23 +++++++++++++++++++ .../apollo/server/utils/mssqlSqlNormalizer.ts | 19 +++++++++++++++ 3 files changed, 60 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 8df9fd7fec..151c22bc40 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -556,6 +556,23 @@ def _rewrite_mssql_invented_pcb_throughput_identifiers(sql: str) -> str: return rewritten +def _rewrite_mssql_repair_log_throughput_shape(sql: str) -> str: + if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\bavg_turnaround_time\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\brepair_count\b|\bthroughput\b", sql, flags=re.IGNORECASE): + return sql + + return ( + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' + 'COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" ' + 'FROM "dbo_DebugEntries" ' + 'GROUP BY "dbo_DebugEntries"."BusinessUnit" ' + 'ORDER BY "throughput" DESC' + ) + + def contains_unsupported_mssql_json_access(sql: str) -> bool: if re.search(r"(?:->>|->)", sql): return True @@ -718,6 +735,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_invented_repair_relationship_identifiers( normalized ) + normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index fea082a408..99198911d3 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -380,6 +380,29 @@ def test_normalize_generation_result_sql_rewrites_pcb_throughput_repair_log_fiel assert 'ORDER BY DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' in normalized +def test_normalize_generation_result_sql_rewrites_repair_log_turnaround_throughput_shape_for_mssql(): + sql = """ + SELECT + board_model AS unit_name, + COUNT(*) AS repair_count, + AVG((DATEPART(DAY, updated_at) - DATEPART(DAY, created_at))) AS avg_turnaround_time + FROM dbo_repair_logs + GROUP BY board_model + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "DATEPART(DAY" not in normalized + assert "avg_turnaround_time" not in normalized + assert "dbo_repair_logs" not in normalized + assert ( + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' + 'COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput"' + in normalized + ) + assert 'FROM "dbo_DebugEntries"' in normalized + + def test_normalize_generation_result_sql_rewrites_bare_month_field_for_mssql(): sql = """ SELECT diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 62829fcbf7..bcdd543be3 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -175,6 +175,24 @@ const replacePcbThroughputFields = (sql: string): string => { return sql; }; +const replaceRepairLogThroughputShape = (sql: string): string => { + if ( + !/\bdbo_repair_logs\b/i.test(sql) || + !/\bavg_turnaround_time\b/i.test(sql) || + !/\b(?:repair_count|throughput)\b/i.test(sql) + ) { + return sql; + } + + return [ + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name",', + 'COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput"', + 'FROM "dbo_DebugEntries"', + 'GROUP BY "dbo_DebugEntries"."BusinessUnit"', + 'ORDER BY "throughput" DESC', + ].join(' '); +}; + export const normalizeMssqlGeneratedSqlFields = ( sql: string, dataSource: DataSourceName, @@ -185,6 +203,7 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = sql.replace(/\\"/g, '"'); sql = replaceInventedDateFields(sql); + sql = replaceRepairLogThroughputShape(sql); sql = replacePcbThroughputFields(sql); sql = replaceInventedTimeBuckets(sql); sql = replaceBadFailurePatternJoins(sql); From 705c5f0693446c85df20635621912155a60ae442 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 2 Jun 2026 13:34:47 +0530 Subject: [PATCH 0078/1087] date --- .../pipelines/generation/sql_correction.py | 7 + .../src/pipelines/generation/utils/sql.py | 140 ++++++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 1 + .../pipelines/generation/test_sql_utils.py | 38 +++++ 4 files changed, 186 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 32babb4850..223e50c100 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -95,6 +95,9 @@ def get_sql_correction_system_prompt( {% endif %} ### QUESTION ### +{% if query %} +User's Question: {{ query }} +{% endif %} {% if invalid_generation_result.original_sql %} Original SQL: {{ invalid_generation_result.original_sql }} {% endif %} @@ -112,10 +115,12 @@ def prompt( invalid_generation_result: Dict, prompt_builder: PromptBuilder, data_source: str, + query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( + query=query, data_source=data_source, documents=documents, valid_table_names=construct_valid_table_names(documents), @@ -206,6 +211,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + query: str | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -216,6 +222,7 @@ async def run( inputs={ "invalid_generation_result": invalid_generation_result, "documents": contexts, + "query": query, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 151c22bc40..0d81ce86da 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -354,6 +354,50 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: return rewritten +def _qualify_mssql_temporal_expression(expression: str, sql: str) -> str: + expression = expression.strip() + if "." in expression or expression.startswith(('"', "[")): + return expression + + if re.search(r"\bdbo_DebugEntries\b", sql, flags=re.IGNORECASE) and re.fullmatch( + r"(?:DateIn|DateOut|FailedAt)", expression, flags=re.IGNORECASE + ): + canonical_columns = { + "datein": "DateIn", + "dateout": "DateOut", + "failedat": "FailedAt", + } + return f'"dbo_DebugEntries"."{canonical_columns[expression.lower()]}"' + + return f'"{expression}"' + + +def _rewrite_mssql_to_date_buckets(sql: str) -> str: + expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" + + def make_day_bucket(expression: str) -> str: + timestamp_expression = _qualify_mssql_temporal_expression(expression, sql) + return ( + f"(DATEPART(YEAR, {timestamp_expression}) * 10000 + " + f"DATEPART(MONTH, {timestamp_expression}) * 100 + " + f"DATEPART(DAY, {timestamp_expression}))" + ) + + rewritten = re.sub( + rf"\bTO_DATE\(\s*{expression_pattern}\s*,\s*'YYYY-MM-DD'\s*\)", + lambda match: make_day_bucket(match.group(1)), + sql, + flags=re.IGNORECASE, + ) + rewritten = re.sub( + rf"\bDATE\(\s*{expression_pattern}\s*\)", + lambda match: make_day_bucket(match.group(1)), + rewritten, + flags=re.IGNORECASE, + ) + return rewritten + + def _rewrite_mssql_timestamp_casts(sql: str) -> str: expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" timestamp_function_pattern = re.compile( @@ -713,6 +757,100 @@ def replace_clause(match: re.Match) -> str: return clause_pattern.sub(replace_clause, sql) +def _split_top_level_select_items(select_body: str) -> list[str]: + items: list[str] = [] + current: list[str] = [] + depth = 0 + in_single_quote = False + in_double_quote = False + + for char in select_body: + if char == "'" and not in_double_quote: + in_single_quote = not in_single_quote + elif char == '"' and not in_single_quote: + in_double_quote = not in_double_quote + elif not in_single_quote and not in_double_quote: + if char == "(": + depth += 1 + elif char == ")" and depth > 0: + depth -= 1 + elif char == "," and depth == 0: + items.append("".join(current).strip()) + current = [] + continue + current.append(char) + + if current: + items.append("".join(current).strip()) + + return items + + +def _rewrite_mssql_temporal_bucket_alias_references(sql: str) -> str: + select_match = re.search( + r"\bSELECT\b(?P.*?)(?=\bFROM\b)", + sql, + flags=re.IGNORECASE | re.DOTALL, + ) + if not select_match: + return sql + + aliases: dict[str, str] = {} + for item in _split_top_level_select_items(select_match.group("body")): + if not re.search(r"\bDATEPART\s*\(", item, flags=re.IGNORECASE): + continue + + alias_match = re.search( + r"\s+(?:AS\s+)?(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))\s*$", + item, + flags=re.IGNORECASE, + ) + if not alias_match: + continue + + alias = alias_match.group(1) or alias_match.group(2) or alias_match.group(3) + expression = item[: alias_match.start()].strip() + aliases[alias.lower()] = expression + + if not aliases: + return sql + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match[str]) -> str: + body = match.group("body") + placeholders: dict[str, str] = {} + for alias, expression in aliases.items(): + placeholder = f"__WREN_MSSQL_TEMPORAL_BUCKET_ALIAS_{len(placeholders)}__" + placeholders[placeholder] = expression + body = re.sub( + rf'"{re.escape(alias)}"', + placeholder, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"\[{re.escape(alias)}\]", + placeholder, + body, + flags=re.IGNORECASE, + ) + body = re.sub( + rf"(? str: normalized = sql @@ -730,6 +868,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _replace_relative_getdate_calls(normalized, now) normalized = _rewrite_mssql_to_unixtime(normalized) normalized = _rewrite_mssql_timestamp_subtraction(normalized) + normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) normalized = _rewrite_mssql_invented_repair_relationship_identifiers( @@ -741,6 +880,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) + normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) return re.sub(r"\s+", " ", normalized).strip() diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e28801f425..682d53c600 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -836,6 +836,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + query=user_query, ), ) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 99198911d3..6d1e317484 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -503,6 +503,44 @@ def test_normalize_generation_result_sql_rewrites_timestamp_casts_for_mssql(): assert "CAST('2026-01-01 00:00:00' AS DATETIME)" in normalized +def test_normalize_generation_result_sql_rewrites_to_date_bucket_for_mssql(): + sql = """ + SELECT + TO_DATE(DateIn, 'YYYY-MM-DD') EntryDate, + COUNT(*) Throughput + FROM dbo_DebugEntries + GROUP BY EntryDate + ORDER BY EntryDate ASC NULLS LAST + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "TO_DATE(" not in normalized + assert "NULLS LAST" not in normalized + assert "GROUP BY EntryDate" not in normalized + assert "ORDER BY EntryDate" not in normalized + assert 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn")' in normalized + assert 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn")' in normalized + assert 'DATEPART(DAY, "dbo_DebugEntries"."DateIn")' in normalized + + +def test_normalize_generation_result_sql_rewrites_date_function_for_mssql(): + sql = """ + SELECT + DATE(DateIn) AS EntryDate, + COUNT(*) AS Throughput + FROM dbo_DebugEntries + GROUP BY EntryDate + ORDER BY EntryDate ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "DATE(DateIn)" not in normalized + assert "GROUP BY EntryDate" not in normalized + assert 'DATEPART(DAY, "dbo_DebugEntries"."DateIn")' in normalized + + def test_normalize_generation_result_sql_strips_to_unixtime_for_mssql(): sql = """ SELECT From ed7b87803994b3b1e14d534e0c8292ccc7f4cfaa Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 2 Jun 2026 13:53:01 +0530 Subject: [PATCH 0079/1087] dates --- .../src/pipelines/generation/utils/sql.py | 98 +++++++++++++++ .../pipelines/generation/test_sql_utils.py | 53 ++++++++ .../apollo/server/utils/mssqlSqlNormalizer.ts | 113 ++++++++++++++++-- 3 files changed, 253 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 0d81ce86da..b2d1c62428 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -198,6 +198,50 @@ def replace_day_offset(match: re.Match[str]) -> str: return sql +def _replace_relative_current_date_calls(sql: str, now: datetime) -> str: + def replace_date_sub_interval(match: re.Match[str]) -> str: + amount = int(match.group("amount")) + unit = match.group("unit").lower() + if unit.startswith("month"): + return _format_timestamp_literal(_add_months(now, -amount)) + if unit.startswith("year"): + return _format_timestamp_literal(_add_months(now, -amount * 12)) + if unit.startswith("day"): + return _format_timestamp_literal(now - timedelta(days=amount)) + return match.group(0) + + def replace_date_sub_unit_amount(match: re.Match[str]) -> str: + unit = match.group("unit").lower() + amount = int(match.group("amount")) + if unit.startswith("month"): + return _format_timestamp_literal(_add_months(now, -amount)) + if unit.startswith("year"): + return _format_timestamp_literal(_add_months(now, -amount * 12)) + if unit.startswith("day"): + return _format_timestamp_literal(now - timedelta(days=amount)) + return match.group(0) + + sql = re.sub( + r"\bDATE_SUB\(\s*CURRENT_DATE(?:\(\))?\s*,\s*INTERVAL\s+(?P\d+)\s+(?PYEAR|MONTH|DAY)S?\s*\)", + replace_date_sub_interval, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"\bDATE_SUB\(\s*'?(?PYEAR|MONTH|DAY)'?\s*,\s*(?P\d+)\s*,\s*CURRENT_DATE(?:\(\))?\s*\)", + replace_date_sub_unit_amount, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"\bCURRENT_DATE(?:\(\))?\b", + _format_timestamp_literal(now), + sql, + flags=re.IGNORECASE, + ) + return sql + + def _rewrite_mssql_bucket_functions(sql: str) -> str: expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" @@ -617,6 +661,57 @@ def _rewrite_mssql_repair_log_throughput_shape(sql: str) -> str: ) +def _rewrite_mssql_invented_failure_category(sql: str) -> str: + if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\bfailure_category\b", sql, flags=re.IGNORECASE): + return sql + + failure_code_expression = '"dbo_repair_logs"."failure_code"' + rewritten = re.sub( + r'(?P\bSELECT\s+|,\s*)(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)(?P\s*(?:,|\bFROM\b))', + rf'\g{failure_code_expression} AS "failure_category"\g', + sql, + flags=re.IGNORECASE, + ) + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match[str]) -> str: + body = re.sub( + r'(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)', + failure_code_expression, + match.group("body"), + flags=re.IGNORECASE, + ) + return f"{match.group(1)}{body}" + + return clause_pattern.sub(replace_clause, rewritten) + + +def _rewrite_mssql_invented_report_fields(sql: str) -> str: + if not re.search(r"\bdbo_reports\b", sql, flags=re.IGNORECASE): + return sql + + report_table = r'(?:"dbo_reports"|\[dbo_reports\]|dbo_reports)' + rewritten = re.sub( + rf"(?:(?:{report_table})\s*\.\s*)?(?:\"filters\"|\[filters\]|\bfilters\b)", + '"dbo_reports"."data"', + sql, + flags=re.IGNORECASE, + ) + rewritten = re.sub( + rf"(?:(?:{report_table})\s*\.\s*)?(?:\"report_size\"|\"file_size\"|\[report_size\]|\[file_size\]|\breport_size\b|\bfile_size\b)", + '"dbo_reports"."size_bytes"', + rewritten, + flags=re.IGNORECASE, + ) + return rewritten + + def contains_unsupported_mssql_json_access(sql: str) -> bool: if re.search(r"(?:->>|->)", sql): return True @@ -866,6 +961,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> flags=re.IGNORECASE, ) normalized = _replace_relative_getdate_calls(normalized, now) + normalized = _replace_relative_current_date_calls(normalized, now) normalized = _rewrite_mssql_to_unixtime(normalized) normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) @@ -876,6 +972,8 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> ) normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) + normalized = _rewrite_mssql_invented_failure_category(normalized) + normalized = _rewrite_mssql_invented_report_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6d1e317484..e51f1da679 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -541,6 +541,59 @@ def test_normalize_generation_result_sql_rewrites_date_function_for_mssql(): assert 'DATEPART(DAY, "dbo_DebugEntries"."DateIn")' in normalized +def test_normalize_generation_result_sql_rewrites_date_sub_for_mssql(): + sql = """ + SELECT + DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "MONTH", + COUNT(*) AS "repair_volume" + FROM "dbo_repair_logs" + WHERE "dbo_repair_logs"."created_at" >= DATE_SUB(CURRENT_DATE, INTERVAL 12 MONTH) + GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at") + ORDER BY "MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "DATE_SUB(" not in normalized + assert "CURRENT_DATE" not in normalized + assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized + + +def test_normalize_generation_result_sql_rewrites_repair_log_failure_category_for_mssql(): + sql = """ + SELECT + failure_category, + COUNT(*) AS repair_count + FROM dbo_repair_logs + GROUP BY failure_category + ORDER BY repair_count DESC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "failure_category," not in normalized + assert "GROUP BY failure_category" not in normalized + assert '"dbo_repair_logs"."failure_code" AS "failure_category"' in normalized + assert 'GROUP BY "dbo_repair_logs"."failure_code"' in normalized + + +def test_normalize_generation_result_sql_rewrites_report_hallucinated_fields_for_mssql(): + sql = """ + SELECT + COUNT(*) total_reports, + SUM(CASE WHEN (filters LIKE '%raw%data%file%') THEN 1 ELSE 0 END) raw_data_files_included, + AVG((CASE WHEN (filters LIKE '%raw%data%file%') THEN file_size ELSE null END)) avg_file_size_with_raw_data + FROM dbo_reports + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "filters" not in normalized + assert "THEN file_size" not in normalized + assert '"dbo_reports"."data" LIKE' in normalized + assert '"dbo_reports"."size_bytes"' in normalized + + def test_normalize_generation_result_sql_strips_to_unixtime_for_mssql(): sql = """ SELECT diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index bcdd543be3..3ac3a3b4f4 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -2,6 +2,46 @@ import { DataSourceName } from '@server/types'; const escapeRegex = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); +const formatTimestampLiteral = (date: Date) => { + const pad = (value: number) => String(value).padStart(2, '0'); + return `'${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())} ${pad(date.getHours())}:${pad(date.getMinutes())}:${pad(date.getSeconds())}'`; +}; + +const addMonths = (date: Date, months: number) => { + const next = new Date(date); + next.setMonth(next.getMonth() + months); + return next; +}; + +const replaceRelativeCurrentDateCalls = (sql: string): string => { + const now = new Date(); + const relativeLiteral = (unit: string, amount: number) => { + const normalizedUnit = unit.toLowerCase(); + if (normalizedUnit.startsWith('month')) { + return formatTimestampLiteral(addMonths(now, -amount)); + } + if (normalizedUnit.startsWith('year')) { + return formatTimestampLiteral(addMonths(now, -amount * 12)); + } + if (normalizedUnit.startsWith('day')) { + const next = new Date(now); + next.setDate(next.getDate() - amount); + return formatTimestampLiteral(next); + } + return null; + }; + + sql = sql.replace( + /\bDATE_SUB\(\s*CURRENT_DATE(?:\(\))?\s*,\s*INTERVAL\s+(\d+)\s+(YEAR|MONTH|DAY)S?\s*\)/gi, + (match, amount, unit) => relativeLiteral(unit, Number(amount)) || match, + ); + sql = sql.replace( + /\bDATE_SUB\(\s*'?(YEAR|MONTH|DAY)'?\s*,\s*(\d+)\s*,\s*CURRENT_DATE(?:\(\))?\s*\)/gi, + (match, unit, amount) => relativeLiteral(unit, Number(amount)) || match, + ); + return sql.replace(/\bCURRENT_DATE(?:\(\))?\b/gi, formatTimestampLiteral(now)); +}; + const inferMssqlTimestampExpression = (sql: string): string => { const qualifiedTimestamp = sql.match( /"([^"]+)"\."(created_at|updated_at|generated_at|created_date|date|DateIn|DateOut|FailedAt)"/i, @@ -10,20 +50,19 @@ const inferMssqlTimestampExpression = (sql: string): string => { return qualifiedTimestamp[0]; } - const fromTable = sql.match(/\bFROM\s+"([^"]+)"/i); + const fromTable = sql.match(/\bFROM\s+(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/i); if (fromTable) { - if (fromTable[1].toLowerCase() === 'dbo_debugentries') { - return `"${fromTable[1]}"."DateIn"`; + const tableName = fromTable[1] || fromTable[2] || fromTable[3]; + if (tableName.toLowerCase() === 'dbo_debugentries') { + return `"${tableName}"."DateIn"`; } - return `"${fromTable[1]}"."created_at"`; - } - - const bracketedFromTable = sql.match(/\bFROM\s+\[([^\]]+)\]/i); - if (bracketedFromTable) { - if (bracketedFromTable[1].toLowerCase() === 'dbo_debugentries') { - return `"${bracketedFromTable[1]}"."DateIn"`; + if (tableName.toLowerCase() === 'dbo_reports') { + return `"${tableName}"."generated_at"`; + } + if (tableName.toLowerCase() === 'dbo_repair_logs') { + return `"${tableName}"."created_at"`; } - return `"${bracketedFromTable[1]}"."created_at"`; + return `"${tableName}"."created_at"`; } return '"created_at"'; @@ -193,6 +232,55 @@ const replaceRepairLogThroughputShape = (sql: string): string => { ].join(' '); }; +const replaceInventedFailureCategory = (sql: string): string => { + if (!/\bdbo_repair_logs\b/i.test(sql) || !/\bfailure_category\b/i.test(sql)) { + return sql; + } + + const failureCodeExpression = '"dbo_repair_logs"."failure_code"'; + sql = sql.replace( + /\bSELECT\b(?.*?)(?=\bFROM\b)/is, + (match, _body, _offset, _source, groups) => { + let body = groups?.body || ''; + body = body.replace( + /(^|,)\s*(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)(?=\s*(?:,|$))/gi, + `$1 ${failureCodeExpression} AS "failure_category"`, + ); + return `SELECT${body}`; + }, + ); + + const clausePattern = + /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; + return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { + const body = (groups?.body || '').replace( + /(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)/gi, + failureCodeExpression, + ); + return `${clause}${body}`; + }); +}; + +const replaceInventedReportFields = (sql: string): string => { + if (!/\bdbo_reports\b/i.test(sql)) { + return sql; + } + + const reportTable = String.raw`(?:"dbo_reports"|\[dbo_reports\]|dbo_reports)`; + sql = sql.replace( + new RegExp(String.raw`(?:(?:${reportTable})\s*\.\s*)?(?:"filters"|\[filters\]|\bfilters\b)`, 'gi'), + '"dbo_reports"."data"', + ); + sql = sql.replace( + new RegExp( + String.raw`(?:(?:${reportTable})\s*\.\s*)?(?:"report_size"|"file_size"|\[report_size\]|\[file_size\]|\breport_size\b|\bfile_size\b)`, + 'gi', + ), + '"dbo_reports"."size_bytes"', + ); + return sql; +}; + export const normalizeMssqlGeneratedSqlFields = ( sql: string, dataSource: DataSourceName, @@ -202,9 +290,12 @@ export const normalizeMssqlGeneratedSqlFields = ( } sql = sql.replace(/\\"/g, '"'); + sql = replaceRelativeCurrentDateCalls(sql); sql = replaceInventedDateFields(sql); sql = replaceRepairLogThroughputShape(sql); sql = replacePcbThroughputFields(sql); + sql = replaceInventedFailureCategory(sql); + sql = replaceInventedReportFields(sql); sql = replaceInventedTimeBuckets(sql); sql = replaceBadFailurePatternJoins(sql); return sql; From 75144d13f7d6611458fdf945bac7a0bb6bd5294b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 2 Jun 2026 15:24:25 +0530 Subject: [PATCH 0080/1087] day --- .../src/pipelines/generation/utils/sql.py | 17 ++++++++++++-- .../pipelines/generation/test_sql_utils.py | 23 +++++++++++++++++++ .../apollo/server/utils/mssqlSqlNormalizer.ts | 6 ++++- 3 files changed, 43 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 b2d1c62428..977aa70fa7 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -513,12 +513,12 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: if normalized_table_name == "dbo_debugentries": return f'{quoted_table_name}."DateIn"' if "report" in normalized_table_name: - return f'{table_name}."generated_at"' + return f'{quoted_table_name}."generated_at"' if any( token in normalized_table_name for token in ("repair", "ticket", "event", "log") ): - return f'{table_name}."created_at"' + return f'{quoted_table_name}."created_at"' return None @@ -753,6 +753,10 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', re.IGNORECASE, ) + select_bare_identifier_pattern = re.compile( + rf"(?P\bSELECT\s+|,\s*){bucket}(?P\s*(?:,|\bFROM\b))", + re.IGNORECASE, + ) select_qualified_identifier_pattern = re.compile( rf'(?P\bSELECT\s+|,\s*){qualified_bucket_pattern.pattern}(?P\s*(?:,|\bFROM\b))', re.IGNORECASE, @@ -769,6 +773,9 @@ def replace_select_identifier(match: re.Match[str]) -> str: rewritten = select_identifier_pattern.sub( replace_select_identifier, rewritten ) + rewritten = select_bare_identifier_pattern.sub( + replace_select_identifier, rewritten + ) clause_pattern = re.compile( r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", @@ -796,6 +803,12 @@ def replace_clause(match: re.Match[str]) -> str: body, flags=re.IGNORECASE, ) + body = re.sub( + rf"(? { const alias = bucket.toLowerCase(); body = body.replace( new RegExp( - String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"?${bucket}"?|(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})|\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b|"${bucket}"|\[${bucket}\])(?=\s*(?:,|$))`, + String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"?${bucket}"?|(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})|\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b|"${bucket}"|\[${bucket}\]|\b${bucket}\b)(?=\s*(?:,|$))`, 'gi', ), `$1 ${expression} AS "${alias}"`, @@ -139,6 +139,10 @@ const replaceInventedTimeBuckets = (sql: string): string => { Object.entries(bucketExpressions).forEach(([bucket, expression]) => { body = body.replace(new RegExp(String.raw`"${bucket}"`, 'gi'), expression); body = body.replace(new RegExp(String.raw`\[${bucket}\]`, 'gi'), expression); + body = body.replace( + new RegExp(String.raw`(? Date: Tue, 2 Jun 2026 16:52:32 +0530 Subject: [PATCH 0081/1087] days --- .../src/pipelines/generation/utils/sql.py | 67 ++++++++++++++++++- .../pipelines/generation/test_sql_utils.py | 67 +++++++++++++++++++ 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 977aa70fa7..cd9776848d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -115,6 +115,15 @@ def normalize_data_source(data_source: str | None) -> str: return normalized +def _is_local_sql_data_source(data_source: str | None) -> bool: + return normalize_data_source(data_source) in { + "DUCKDB", + "LOCAL_FILE", + "SQLITE", + "SQLITE3", + } + + def _format_timestamp_literal(value: datetime) -> str: return value.strftime("'%Y-%m-%d %H:%M:%S'") @@ -661,6 +670,29 @@ def _rewrite_mssql_repair_log_throughput_shape(sql: str) -> str: ) +def _rewrite_mssql_repair_log_turnaround_trend_shape(sql: str) -> str: + if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\bavg_turnaround_time\b|\bturnaround\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\bMONTH\b|DATEPART\(\s*MONTH", sql, flags=re.IGNORECASE): + return sql + + return ( + 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' + 'AVG(DATEDIFF(\'second\', "dbo_repair_logs"."created_at", ' + '"dbo_repair_logs"."updated_at")) AS "avg_turnaround_seconds" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' + 'AND "dbo_repair_logs"."updated_at" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + ) + + def _rewrite_mssql_invented_failure_category(sql: str) -> str: if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): return sql @@ -959,10 +991,38 @@ def replace_clause(match: re.Match[str]) -> str: return clause_pattern.sub(replace_clause, sql) +def _references_known_hallucination_prone_schema(sql: str) -> bool: + return bool( + re.search( + r"\b(?:dbo_repair_logs|dbo_DebugEntries|dbo_reports)\b", + sql, + flags=re.IGNORECASE, + ) + ) + + +def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: + normalized = _replace_relative_current_date_calls(sql, now) + normalized = _rewrite_mssql_to_date_buckets(normalized) + normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_invented_repair_relationship_identifiers(normalized) + normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) + normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) + normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) + normalized = _rewrite_mssql_invented_failure_category(normalized) + normalized = _rewrite_mssql_invented_report_fields(normalized) + normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) + normalized = _rewrite_temporal_bucket_functions(normalized) + normalized = _rewrite_mssql_datepart_alias_references(normalized) + normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) + return normalized + + def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = sql + normalized_data_source = normalize_data_source(data_source) - if normalize_data_source(data_source) == "MSSQL": + if normalized_data_source == "MSSQL": now = datetime.now() normalized = re.sub( r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE @@ -983,6 +1043,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_invented_repair_relationship_identifiers( normalized ) + normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_invented_failure_category(normalized) @@ -992,6 +1053,10 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) + elif _is_local_sql_data_source( + normalized_data_source + ) and _references_known_hallucination_prone_schema(normalized): + normalized = _rewrite_known_schema_hallucinations(normalized, datetime.now()) return re.sub(r"\s+", " ", normalized).strip() diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 90228e7dbc..e1801778af 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -403,6 +403,29 @@ def test_normalize_generation_result_sql_rewrites_repair_log_turnaround_throughp assert 'FROM "dbo_DebugEntries"' in normalized +def test_normalize_generation_result_sql_rewrites_repair_log_turnaround_month_trend_shape_for_mssql(): + sql = """ + SELECT + MONTH, + AVG(avg_turnaround_time) AS avg_turnaround_time + FROM dbo_repair_logs + GROUP BY MONTH + ORDER BY MONTH ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "avg_turnaround_time" not in normalized + assert "GROUP BY MONTH" not in normalized + assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year"' in normalized + assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' in normalized + assert ( + 'AVG(DATEDIFF(\'second\', "dbo_repair_logs"."created_at", ' + '"dbo_repair_logs"."updated_at")) AS "avg_turnaround_seconds"' + in normalized + ) + + def test_normalize_generation_result_sql_rewrites_bare_month_field_for_mssql(): sql = """ SELECT @@ -492,6 +515,50 @@ def test_normalize_generation_result_sql_rewrites_bare_unquoted_month_field_for_ assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized +def test_normalize_generation_result_sql_rewrites_bare_month_field_for_local_file(): + sql = """ + SELECT + MONTH, + COUNT(*) AS repair_volume + FROM dbo_repair_logs + GROUP BY MONTH + ORDER BY MONTH ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="local_file") + + assert "SELECT MONTH" not in normalized + assert "GROUP BY MONTH" not in normalized + assert "ORDER BY MONTH" not in normalized + assert ( + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' + in normalized + ) + assert 'GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized + assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized + + +def test_normalize_generation_result_sql_rewrites_bare_month_field_for_sqlite(): + sql = """ + SELECT + "MONTH", + COUNT("dbo_repair_logs"."id") AS "repair_count" + FROM "dbo_repair_logs" + GROUP BY "MONTH" + ORDER BY "MONTH" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="sqlite") + + assert 'SELECT "MONTH"' not in normalized + assert 'GROUP BY "MONTH"' not in normalized + assert 'ORDER BY "MONTH"' not in normalized + assert ( + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' + in normalized + ) + + def test_normalize_generation_result_sql_rewrites_bare_year_for_report_charts(): sql = """ SELECT From ffbf21f0a6649fcc1560749be5cd729b5de48ea9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 2 Jun 2026 17:12:05 +0530 Subject: [PATCH 0082/1087] line --- .../src/pipelines/generation/utils/sql.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index cd9776848d..c78a4f1d7d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -115,15 +115,6 @@ def normalize_data_source(data_source: str | None) -> str: return normalized -def _is_local_sql_data_source(data_source: str | None) -> bool: - return normalize_data_source(data_source) in { - "DUCKDB", - "LOCAL_FILE", - "SQLITE", - "SQLITE3", - } - - def _format_timestamp_literal(value: datetime) -> str: return value.strftime("'%Y-%m-%d %H:%M:%S'") @@ -1053,9 +1044,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) - elif _is_local_sql_data_source( - normalized_data_source - ) and _references_known_hallucination_prone_schema(normalized): + elif _references_known_hallucination_prone_schema(normalized): normalized = _rewrite_known_schema_hallucinations(normalized, datetime.now()) return re.sub(r"\s+", " ", normalized).strip() @@ -1304,6 +1293,9 @@ async def _classify_generation_result( - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. - Never invent foreign key columns or relationship fields such as "FailurePatternID", "FailurePatternId", "TicketID", or "
ID" unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. +- Never invent time bucket columns such as "MONTH", "YEAR", "DAY", "month", "year", or "date" unless that exact column appears in the DATABASE SCHEMA. For monthly, yearly, or daily trends, apply a supported date/time bucket function from SQL FUNCTIONS to a real timestamp column from the selected table. +- For synced repair-log schemas, if "dbo_repair_logs" contains "created_at" and the user asks for monthly repair volume or repair trends, count repair rows and bucket "dbo_repair_logs"."created_at". Do not select, group by, or order by "dbo_repair_logs"."MONTH" or bare "MONTH" unless the schema explicitly contains that column. +- For repair counts grouped by failure category in synced repair-log schemas, use "dbo_repair_logs"."failure_code" when that column appears in the schema. Do not invent "failure_category" unless it appears in the DATABASE SCHEMA. - For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. - For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. - Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. From 159a7d36badb8db570af7df0de6c11ae39d5c138 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 3 Jun 2026 00:55:51 +0530 Subject: [PATCH 0083/1087] rbac implementation --- wren-ui/docs/rbac-architecture.md | 101 ++++++ .../20250602000000_create_rbac_tables.js | 76 +++++ wren-ui/src/apollo/client/graphql/rbac.ts | 169 ++++++++++ wren-ui/src/apollo/server/models/index.ts | 1 + wren-ui/src/apollo/server/models/rbac.ts | 38 +++ .../src/apollo/server/repositories/index.ts | 1 + .../server/repositories/rbacRepository.ts | 160 ++++++++++ wren-ui/src/apollo/server/resolvers.ts | 19 ++ .../apollo/server/resolvers/rbacResolver.ts | 149 +++++++++ wren-ui/src/apollo/server/schema.ts | 98 ++++++ wren-ui/src/apollo/server/services/index.ts | 1 + .../src/apollo/server/services/rbacService.ts | 297 ++++++++++++++++++ wren-ui/src/apollo/server/types/context.ts | 8 + wren-ui/src/common.ts | 16 + wren-ui/src/components/HeaderBar.tsx | 8 + .../components/pages/administration/types.tsx | 51 +++ .../src/components/sidebar/Administration.tsx | 73 +++++ wren-ui/src/components/sidebar/index.tsx | 5 + .../src/pages/administration/assignments.tsx | 281 +++++++++++++++++ wren-ui/src/pages/administration/index.tsx | 14 + wren-ui/src/pages/administration/roles.tsx | 245 +++++++++++++++ wren-ui/src/pages/administration/users.tsx | 296 +++++++++++++++++ wren-ui/src/pages/api/graphql.ts | 8 + wren-ui/src/utils/enum/menu.ts | 3 + wren-ui/src/utils/enum/path.ts | 4 + 25 files changed, 2122 insertions(+) create mode 100644 wren-ui/docs/rbac-architecture.md create mode 100644 wren-ui/migrations/20250602000000_create_rbac_tables.js create mode 100644 wren-ui/src/apollo/client/graphql/rbac.ts create mode 100644 wren-ui/src/apollo/server/models/rbac.ts create mode 100644 wren-ui/src/apollo/server/repositories/rbacRepository.ts create mode 100644 wren-ui/src/apollo/server/resolvers/rbacResolver.ts create mode 100644 wren-ui/src/apollo/server/services/rbacService.ts create mode 100644 wren-ui/src/components/pages/administration/types.tsx create mode 100644 wren-ui/src/components/sidebar/Administration.tsx create mode 100644 wren-ui/src/pages/administration/assignments.tsx create mode 100644 wren-ui/src/pages/administration/index.tsx create mode 100644 wren-ui/src/pages/administration/roles.tsx create mode 100644 wren-ui/src/pages/administration/users.tsx diff --git a/wren-ui/docs/rbac-architecture.md b/wren-ui/docs/rbac-architecture.md new file mode 100644 index 0000000000..cec918bbfe --- /dev/null +++ b/wren-ui/docs/rbac-architecture.md @@ -0,0 +1,101 @@ +# RBAC Foundation + +This document describes the application-level RBAC foundation in Wren UI. It does not enforce permissions yet; it establishes durable role, user, and user-role assignment primitives for future governance work. + +## Scope + +Implemented: + +- Roles: `Admin`, `Manager`, `Analyst`, `Viewer`, plus custom roles. +- Users with local identity metadata. +- User-role mappings. +- GraphQL APIs for role, user, and assignment management. +- Administration UI for user management, role management, and role assignment. + +Deferred: + +- Authorization middleware. +- Governance policies. +- Data scoping. +- SQL validation. +- Schema/table-level permissions. +- Teams, LDAP, and Azure AD synchronization. + +## Database Model + +Tables: + +- `roles` + - `id` + - `name` + - `description` + - timestamps +- `users` + - `id` + - `name` + - `email` + - `external_id` + - `identity_provider` + - `is_active` + - timestamps +- `user_roles` + - `id` + - `user_id` + - `role_id` + - timestamps + +`external_id` and `identity_provider` are intentionally present now so Teams, LDAP, and Azure AD integrations can later attach external identities without replacing the RBAC tables. + +## Backend Layers + +- Migration: `migrations/20250602000000_create_rbac_tables.js` +- Models: `src/apollo/server/models/rbac.ts` +- Repositories: `src/apollo/server/repositories/rbacRepository.ts` +- Service: `src/apollo/server/services/rbacService.ts` +- Resolver: `src/apollo/server/resolvers/rbacResolver.ts` +- GraphQL schema: `src/apollo/server/schema.ts` + +The service owns validation and duplicate checks. The repository owns persistence and joined user-role mapping queries. + +## GraphQL API + +Queries: + +- `roles` +- `users` +- `userRoleMappings` + +Mutations: + +- `createRole` +- `updateRole` +- `createUser` +- `updateUser` +- `assignRoleToUser` +- `updateUserRoles` +- `removeRoleFromUser` + +## UI + +Navigation: + +- Header tab: `Admin` +- Sidebar section: `Administration` + +Screens: + +- `/administration/users` +- `/administration/roles` +- `/administration/assignments` + +The UI uses the existing Next.js, Apollo Client, Ant Design, and `SiderLayout`/`PageLayout` patterns. + +## Future Permission Model + +Future schema-level or table-level permissions should be added as separate tables referencing `roles.id`, for example: + +- `role_schema_permissions` +- `role_table_permissions` +- `role_policy_bindings` + +This keeps identity and assignment management stable while allowing governance policies to evolve independently. diff --git a/wren-ui/migrations/20250602000000_create_rbac_tables.js b/wren-ui/migrations/20250602000000_create_rbac_tables.js new file mode 100644 index 0000000000..92bd99a2bc --- /dev/null +++ b/wren-ui/migrations/20250602000000_create_rbac_tables.js @@ -0,0 +1,76 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + await knex.schema.createTable('roles', (table) => { + table.increments('id').primary(); + table.string('name', 80).notNullable().unique(); + table + .text('description') + .nullable() + .comment('Human-readable role purpose and future governance notes'); + table.timestamps(true, true); + }); + + await knex.schema.createTable('users', (table) => { + table.increments('id').primary(); + table.string('name', 160).notNullable(); + table.string('email', 320).notNullable().unique(); + table.string('external_id', 255).nullable().unique(); + table.string('identity_provider', 80).nullable(); + table.boolean('is_active').notNullable().defaultTo(true); + table.timestamps(true, true); + }); + + await knex.schema.createTable('user_roles', (table) => { + table.increments('id').primary(); + table.integer('user_id').notNullable(); + table.integer('role_id').notNullable(); + table.timestamps(true, true); + + table.foreign('user_id').references('users.id').onDelete('CASCADE'); + table.foreign('role_id').references('roles.id').onDelete('CASCADE'); + table.unique(['user_id', 'role_id']); + }); + + const now = new Date().toISOString(); + await knex('roles').insert([ + { + name: 'Admin', + description: 'Full administration access foundation role.', + created_at: now, + updated_at: now, + }, + { + name: 'Manager', + description: + 'Manages users, assignments, and future governance workflows.', + created_at: now, + updated_at: now, + }, + { + name: 'Analyst', + description: 'Creates and analyzes project content.', + created_at: now, + updated_at: now, + }, + { + name: 'Viewer', + description: + 'Read-only foundation role for future permission enforcement.', + created_at: now, + updated_at: now, + }, + ]); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex.schema.dropTableIfExists('user_roles'); + await knex.schema.dropTableIfExists('users'); + await knex.schema.dropTableIfExists('roles'); +}; diff --git a/wren-ui/src/apollo/client/graphql/rbac.ts b/wren-ui/src/apollo/client/graphql/rbac.ts new file mode 100644 index 0000000000..6da6060285 --- /dev/null +++ b/wren-ui/src/apollo/client/graphql/rbac.ts @@ -0,0 +1,169 @@ +import { gql } from '@apollo/client'; + +export const ROLE_FIELDS = gql` + fragment RoleFields on Role { + id + name + description + createdAt + updatedAt + } +`; + +export const USER_FIELDS = gql` + fragment UserFields on User { + id + name + email + externalId + identityProvider + isActive + createdAt + updatedAt + } +`; + +export const LIST_RBAC_USERS = gql` + query RbacUsers { + users { + ...UserFields + roles { + ...RoleFields + } + } + roles { + ...RoleFields + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const LIST_RBAC_ROLES = gql` + query RbacRoles { + roles { + ...RoleFields + users { + ...UserFields + } + } + users { + ...UserFields + } + } + + ${ROLE_FIELDS} + ${USER_FIELDS} +`; + +export const LIST_USER_ROLE_MAPPINGS = gql` + query UserRoleMappings { + userRoleMappings { + id + userId + roleId + createdAt + updatedAt + user { + ...UserFields + roles { + ...RoleFields + } + } + role { + ...RoleFields + } + } + users { + ...UserFields + roles { + ...RoleFields + } + } + roles { + ...RoleFields + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const CREATE_ROLE = gql` + mutation CreateRole($data: CreateRoleInput!) { + createRole(data: $data) { + ...RoleFields + } + } + + ${ROLE_FIELDS} +`; + +export const UPDATE_ROLE = gql` + mutation UpdateRole($where: RoleWhereInput!, $data: UpdateRoleInput!) { + updateRole(where: $where, data: $data) { + ...RoleFields + } + } + + ${ROLE_FIELDS} +`; + +export const CREATE_USER = gql` + mutation CreateUser($data: CreateUserInput!) { + createUser(data: $data) { + ...UserFields + roles { + ...RoleFields + } + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const UPDATE_USER = gql` + mutation UpdateUser($where: UserWhereInput!, $data: UpdateUserInput!) { + updateUser(where: $where, data: $data) { + ...UserFields + roles { + ...RoleFields + } + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const ASSIGN_ROLE_TO_USER = gql` + mutation AssignRoleToUser($data: UserRoleInput!) { + assignRoleToUser(data: $data) { + id + userId + roleId + } + } +`; + +export const UPDATE_USER_ROLES = gql` + mutation UpdateUserRoles($data: UpdateUserRolesInput!) { + updateUserRoles(data: $data) { + ...UserFields + roles { + ...RoleFields + } + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const REMOVE_ROLE_FROM_USER = gql` + mutation RemoveRoleFromUser($data: UserRoleInput!) { + removeRoleFromUser(data: $data) + } +`; diff --git a/wren-ui/src/apollo/server/models/index.ts b/wren-ui/src/apollo/server/models/index.ts index 349625fd76..2a3049bb90 100644 --- a/wren-ui/src/apollo/server/models/index.ts +++ b/wren-ui/src/apollo/server/models/index.ts @@ -2,3 +2,4 @@ export * from './model'; export * from './instruction'; export * from './adaptor'; export * from './dashboard'; +export * from './rbac'; diff --git a/wren-ui/src/apollo/server/models/rbac.ts b/wren-ui/src/apollo/server/models/rbac.ts new file mode 100644 index 0000000000..52b1a98ce9 --- /dev/null +++ b/wren-ui/src/apollo/server/models/rbac.ts @@ -0,0 +1,38 @@ +export interface CreateRoleInput { + name: string; + description?: string | null; +} + +export interface UpdateRoleInput { + id: number; + name?: string | null; + description?: string | null; +} + +export interface CreateUserInput { + name: string; + email: string; + externalId?: string | null; + identityProvider?: string | null; + isActive?: boolean; + roleIds?: number[]; +} + +export interface UpdateUserInput { + id: number; + name?: string | null; + email?: string | null; + externalId?: string | null; + identityProvider?: string | null; + isActive?: boolean | null; +} + +export interface UserRoleInput { + userId: number; + roleId: number; +} + +export interface UpdateUserRolesInput { + userId: number; + roleIds: number[]; +} diff --git a/wren-ui/src/apollo/server/repositories/index.ts b/wren-ui/src/apollo/server/repositories/index.ts index 0dd1cc5905..28e92cce08 100644 --- a/wren-ui/src/apollo/server/repositories/index.ts +++ b/wren-ui/src/apollo/server/repositories/index.ts @@ -19,3 +19,4 @@ export * from './askingTaskRepository'; export * from './instructionRepository'; export * from './apiHistoryRepository'; export * from './dashboardItemRefreshJobRepository'; +export * from './rbacRepository'; diff --git a/wren-ui/src/apollo/server/repositories/rbacRepository.ts b/wren-ui/src/apollo/server/repositories/rbacRepository.ts new file mode 100644 index 0000000000..d742af6579 --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/rbacRepository.ts @@ -0,0 +1,160 @@ +import { Knex } from 'knex'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; + +export interface Role { + id: number; + name: string; + description?: string | null; + createdAt: string; + updatedAt: string; +} + +export interface RbacUser { + id: number; + name: string; + email: string; + externalId?: string | null; + identityProvider?: string | null; + isActive: boolean; + createdAt: string; + updatedAt: string; +} + +export interface UserRole { + id: number; + userId: number; + roleId: number; + createdAt: string; + updatedAt: string; +} + +export interface UserRoleMapping extends UserRole { + user: RbacUser; + role: Role; +} + +export interface IRoleRepository extends IBasicRepository {} + +export interface IUserRepository extends IBasicRepository {} + +export interface IUserRoleRepository extends IBasicRepository { + findMappings(queryOptions?: IQueryOptions): Promise; + findMappingsByUserId( + userId: number, + queryOptions?: IQueryOptions, + ): Promise; + findMappingsByRoleId( + roleId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export class RoleRepository + extends BaseRepository + implements IRoleRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'roles' }); + } +} + +export class UserRepository + extends BaseRepository + implements IUserRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'users' }); + } +} + +export class UserRoleRepository + extends BaseRepository + implements IUserRoleRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'user_roles' }); + } + + public async findMappings(queryOptions?: IQueryOptions) { + return this.queryMappings({}, queryOptions); + } + + public async findMappingsByUserId( + userId: number, + queryOptions?: IQueryOptions, + ) { + return this.queryMappings({ userId }, queryOptions); + } + + public async findMappingsByRoleId( + roleId: number, + queryOptions?: IQueryOptions, + ) { + return this.queryMappings({ roleId }, queryOptions); + } + + private async queryMappings( + filter: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const query = executer('user_roles') + .select( + 'user_roles.*', + 'users.id as user__id', + 'users.name as user__name', + 'users.email as user__email', + 'users.external_id as user__external_id', + 'users.identity_provider as user__identity_provider', + 'users.is_active as user__is_active', + 'users.created_at as user__created_at', + 'users.updated_at as user__updated_at', + 'roles.id as role__id', + 'roles.name as role__name', + 'roles.description as role__description', + 'roles.created_at as role__created_at', + 'roles.updated_at as role__updated_at', + ) + .join('users', 'user_roles.user_id', 'users.id') + .join('roles', 'user_roles.role_id', 'roles.id') + .orderBy('users.email') + .orderBy('roles.name'); + + if (filter.userId) { + query.where('user_roles.user_id', filter.userId); + } + if (filter.roleId) { + query.where('user_roles.role_id', filter.roleId); + } + + const rows = await query; + return rows.map((row) => ({ + id: row.id, + userId: row.user_id, + roleId: row.role_id, + createdAt: row.created_at, + updatedAt: row.updated_at, + user: { + id: row.user__id, + name: row.user__name, + email: row.user__email, + externalId: row.user__external_id, + identityProvider: row.user__identity_provider, + isActive: row.user__is_active, + createdAt: row.user__created_at, + updatedAt: row.user__updated_at, + }, + role: { + id: row.role__id, + name: row.role__name, + description: row.role__description, + createdAt: row.role__created_at, + updatedAt: row.role__updated_at, + }, + })); + } +} diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 4c4c9af187..e5e6d5a309 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -8,6 +8,7 @@ import { DashboardResolver } from './resolvers/dashboardResolver'; import { SqlPairResolver } from './resolvers/sqlPairResolver'; import { InstructionResolver } from './resolvers/instructionResolver'; import { ApiHistoryResolver } from './resolvers/apiHistoryResolver'; +import { RbacResolver } from './resolvers/rbacResolver'; import { convertColumnType } from '@server/utils'; import { DialectSQLScalar } from './scalars'; @@ -20,6 +21,7 @@ const dashboardResolver = new DashboardResolver(); const sqlPairResolver = new SqlPairResolver(); const instructionResolver = new InstructionResolver(); const apiHistoryResolver = new ApiHistoryResolver(); +const rbacResolver = new RbacResolver(); const resolvers = { JSON: GraphQLJSON, DialectSQL: DialectSQLScalar, @@ -75,6 +77,11 @@ const resolvers = { // API History apiHistory: apiHistoryResolver.getApiHistory, + + // Administration / RBAC + roles: rbacResolver.listRoles, + users: rbacResolver.listUsers, + userRoleMappings: rbacResolver.listUserRoleMappings, }, Mutation: { deploy: modelResolver.deploy, @@ -177,6 +184,15 @@ const resolvers = { createInstruction: instructionResolver.createInstruction, updateInstruction: instructionResolver.updateInstruction, deleteInstruction: instructionResolver.deleteInstruction, + + // Administration / RBAC + createRole: rbacResolver.createRole, + updateRole: rbacResolver.updateRole, + createUser: rbacResolver.createUser, + updateUser: rbacResolver.updateUser, + assignRoleToUser: rbacResolver.assignRoleToUser, + updateUserRoles: rbacResolver.updateUserRoles, + removeRoleFromUser: rbacResolver.removeRoleFromUser, }, ThreadResponse: askingResolver.getThreadResponseNestedResolver(), DetailStep: askingResolver.getDetailStepNestedResolver(), @@ -196,6 +212,9 @@ const resolvers = { // Add ApiHistoryResponse nested resolvers ApiHistoryResponse: apiHistoryResolver.getApiHistoryNestedResolver(), + + Role: rbacResolver.getRoleNestedResolver(), + User: rbacResolver.getUserNestedResolver(), }; export default resolvers; diff --git a/wren-ui/src/apollo/server/resolvers/rbacResolver.ts b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts new file mode 100644 index 0000000000..26f4f761ab --- /dev/null +++ b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts @@ -0,0 +1,149 @@ +import { IContext } from '@server/types'; +import { + CreateRoleInput, + CreateUserInput, + UpdateRoleInput, + UpdateUserInput, + UpdateUserRolesInput, + UserRoleInput, +} from '@server/models'; +import { + RbacUser, + Role, + UserRole, + UserRoleMapping, +} from '@server/repositories'; +import { RbacUserWithRoles, RoleWithUsers } from '@server/services'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('RbacResolver'); +logger.level = 'debug'; + +export class RbacResolver { + constructor() { + this.listRoles = this.listRoles.bind(this); + this.listUsers = this.listUsers.bind(this); + this.listUserRoleMappings = this.listUserRoleMappings.bind(this); + this.createRole = this.createRole.bind(this); + this.updateRole = this.updateRole.bind(this); + this.createUser = this.createUser.bind(this); + this.updateUser = this.updateUser.bind(this); + this.assignRoleToUser = this.assignRoleToUser.bind(this); + this.updateUserRoles = this.updateUserRoles.bind(this); + this.removeRoleFromUser = this.removeRoleFromUser.bind(this); + } + + public getRoleNestedResolver() { + return { + users: async (role: RoleWithUsers, _args: any, ctx: IContext) => { + if (role.users) return role.users; + const mappings = await ctx.rbacService.getUserRoleMappings(); + return mappings + .filter((mapping) => mapping.roleId === role.id) + .map((mapping) => mapping.user); + }, + }; + } + + public getUserNestedResolver() { + return { + roles: async (user: RbacUserWithRoles, _args: any, ctx: IContext) => { + if (user.roles) return user.roles; + const mappings = await ctx.rbacService.getUserRoleMappings(); + return mappings + .filter((mapping) => mapping.userId === user.id) + .map((mapping) => mapping.role); + }, + }; + } + + public async listRoles( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + try { + return await ctx.rbacService.listRoles(); + } catch (error) { + logger.error(`Error listing roles: ${error}`); + throw error; + } + } + + public async listUsers( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + try { + return await ctx.rbacService.listUsers(); + } catch (error) { + logger.error(`Error listing users: ${error}`); + throw error; + } + } + + public async listUserRoleMappings( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + return ctx.rbacService.getUserRoleMappings(); + } + + public async createRole( + _root: any, + args: { data: CreateRoleInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.createRole(args.data); + } + + public async updateRole( + _root: any, + args: { where: { id: number }; data: Omit }, + ctx: IContext, + ): Promise { + return ctx.rbacService.updateRole({ id: args.where.id, ...args.data }); + } + + public async createUser( + _root: any, + args: { data: CreateUserInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.createUser(args.data); + } + + public async updateUser( + _root: any, + args: { where: { id: number }; data: Omit }, + ctx: IContext, + ): Promise { + return ctx.rbacService.updateUser({ id: args.where.id, ...args.data }); + } + + public async assignRoleToUser( + _root: any, + args: { data: UserRoleInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.assignRoleToUser(args.data); + } + + public async updateUserRoles( + _root: any, + args: { data: UpdateUserRolesInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.updateUserRoles(args.data); + } + + public async removeRoleFromUser( + _root: any, + args: { data: UserRoleInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.removeRoleFromUser(args.data); + } +} diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index cc52002589..c5a4cbb63e 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -1110,6 +1110,90 @@ export const typeDefs = gql` id: Int! } + type Role { + id: Int! + name: String! + description: String + users: [User!]! + createdAt: String! + updatedAt: String! + } + + type User { + id: Int! + name: String! + email: String! + externalId: String + identityProvider: String + isActive: Boolean! + roles: [Role!]! + createdAt: String! + updatedAt: String! + } + + type UserRole { + id: Int! + userId: Int! + roleId: Int! + createdAt: String! + updatedAt: String! + } + + type UserRoleMapping { + id: Int! + userId: Int! + roleId: Int! + user: User! + role: Role! + createdAt: String! + updatedAt: String! + } + + input RoleWhereInput { + id: Int! + } + + input UserWhereInput { + id: Int! + } + + input CreateRoleInput { + name: String! + description: String + } + + input UpdateRoleInput { + name: String + description: String + } + + input CreateUserInput { + name: String! + email: String! + externalId: String + identityProvider: String + isActive: Boolean + roleIds: [Int!] + } + + input UpdateUserInput { + name: String + email: String + externalId: String + identityProvider: String + isActive: Boolean + } + + input UserRoleInput { + userId: Int! + roleId: Int! + } + + input UpdateUserRolesInput { + userId: Int! + roleIds: [Int!]! + } + # Query and Mutation type Query { # On Boarding Steps @@ -1167,6 +1251,11 @@ export const typeDefs = gql` filter: ApiHistoryFilterInput pagination: ApiHistoryPaginationInput! ): ApiHistoryPaginatedResponse! + + # Administration / RBAC + roles: [Role!]! + users: [User!]! + userRoleMappings: [UserRoleMapping!]! } type Mutation { @@ -1312,5 +1401,14 @@ export const typeDefs = gql` data: UpdateInstructionInput! ): Instruction! deleteInstruction(where: InstructionWhereInput!): Boolean! + + # Administration / RBAC + createRole(data: CreateRoleInput!): Role! + updateRole(where: RoleWhereInput!, data: UpdateRoleInput!): Role! + createUser(data: CreateUserInput!): User! + updateUser(where: UserWhereInput!, data: UpdateUserInput!): User! + assignRoleToUser(data: UserRoleInput!): UserRole! + updateUserRoles(data: UpdateUserRolesInput!): User! + removeRoleFromUser(data: UserRoleInput!): Boolean! } `; diff --git a/wren-ui/src/apollo/server/services/index.ts b/wren-ui/src/apollo/server/services/index.ts index ce1bb78175..1a7c8dadb0 100644 --- a/wren-ui/src/apollo/server/services/index.ts +++ b/wren-ui/src/apollo/server/services/index.ts @@ -8,3 +8,4 @@ export * from './metadataService'; export * from './dashboardService'; export * from './askingTaskTracker'; export * from './instructionService'; +export * from './rbacService'; diff --git a/wren-ui/src/apollo/server/services/rbacService.ts b/wren-ui/src/apollo/server/services/rbacService.ts new file mode 100644 index 0000000000..76064df4b5 --- /dev/null +++ b/wren-ui/src/apollo/server/services/rbacService.ts @@ -0,0 +1,297 @@ +import { Knex } from 'knex'; +import { + CreateRoleInput, + CreateUserInput, + UpdateRoleInput, + UpdateUserInput, + UpdateUserRolesInput, + UserRoleInput, +} from '@server/models'; +import { + IRoleRepository, + IUserRepository, + IUserRoleRepository, + RbacUser, + Role, + UserRole, + UserRoleMapping, +} from '@server/repositories'; + +export interface RbacUserWithRoles extends RbacUser { + roles: Role[]; +} + +export interface RoleWithUsers extends Role { + users: RbacUser[]; +} + +export interface IRbacService { + listRoles(): Promise; + createRole(input: CreateRoleInput): Promise; + updateRole(input: UpdateRoleInput): Promise; + listUsers(): Promise; + createUser(input: CreateUserInput): Promise; + updateUser(input: UpdateUserInput): Promise; + assignRoleToUser(input: UserRoleInput): Promise; + updateUserRoles(input: UpdateUserRolesInput): Promise; + removeRoleFromUser(input: UserRoleInput): Promise; + getUserRoleMappings(): Promise; +} + +const EMAIL_PATTERN = /^[^\s@]+@[^\s@]+\.[^\s@]+$/; + +export class RbacService implements IRbacService { + private readonly roleRepository: IRoleRepository; + private readonly userRepository: IUserRepository; + private readonly userRoleRepository: IUserRoleRepository; + + constructor({ + roleRepository, + userRepository, + userRoleRepository, + }: { + roleRepository: IRoleRepository; + userRepository: IUserRepository; + userRoleRepository: IUserRoleRepository; + }) { + this.roleRepository = roleRepository; + this.userRepository = userRepository; + this.userRoleRepository = userRoleRepository; + } + + public async listRoles(): Promise { + const roles = await this.roleRepository.findAll({ order: 'name' }); + const mappings = await this.userRoleRepository.findMappings(); + return roles.map((role) => ({ + ...role, + users: mappings + .filter((mapping) => mapping.roleId === role.id) + .map((mapping) => mapping.user), + })); + } + + public async createRole(input: CreateRoleInput): Promise { + const name = this.validateRoleName(input.name); + await this.assertUniqueRoleName(name); + const now = new Date().toISOString(); + return this.roleRepository.createOne({ + name, + description: this.normalizeNullable(input.description), + createdAt: now, + updatedAt: now, + }); + } + + public async updateRole(input: UpdateRoleInput): Promise { + const role = await this.getRoleOrThrow(input.id); + const data: Partial = { updatedAt: new Date().toISOString() }; + + if (input.name !== undefined && input.name !== null) { + const name = this.validateRoleName(input.name); + await this.assertUniqueRoleName(name, role.id); + data.name = name; + } + if (input.description !== undefined) { + data.description = this.normalizeNullable(input.description); + } + + return this.roleRepository.updateOne(role.id, data); + } + + public async listUsers(): Promise { + const users = await this.userRepository.findAll({ order: 'email' }); + const mappings = await this.userRoleRepository.findMappings(); + return users.map((user) => ({ + ...user, + roles: mappings + .filter((mapping) => mapping.userId === user.id) + .map((mapping) => mapping.role), + })); + } + + public async createUser(input: CreateUserInput): Promise { + const name = this.validateRequiredText(input.name, 'User name'); + const email = this.validateEmail(input.email); + await this.assertUniqueUserEmail(email); + const now = new Date().toISOString(); + + const tx = await this.userRepository.transaction(); + try { + const user = await this.userRepository.createOne( + { + name, + email, + externalId: this.normalizeNullable(input.externalId), + identityProvider: this.normalizeNullable(input.identityProvider), + isActive: input.isActive ?? true, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + + if (input.roleIds?.length) { + await this.createUserRoleAssignments(user.id, input.roleIds, tx); + } + + await tx.commit(); + return user; + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async updateUser(input: UpdateUserInput): Promise { + const user = await this.getUserOrThrow(input.id); + const data: Partial = { updatedAt: new Date().toISOString() }; + + if (input.name !== undefined && input.name !== null) { + data.name = this.validateRequiredText(input.name, 'User name'); + } + if (input.email !== undefined && input.email !== null) { + const email = this.validateEmail(input.email); + await this.assertUniqueUserEmail(email, user.id); + data.email = email; + } + if (input.externalId !== undefined) { + data.externalId = this.normalizeNullable(input.externalId); + } + if (input.identityProvider !== undefined) { + data.identityProvider = this.normalizeNullable(input.identityProvider); + } + if (input.isActive !== undefined && input.isActive !== null) { + data.isActive = input.isActive; + } + + return this.userRepository.updateOne(user.id, data); + } + + public async assignRoleToUser(input: UserRoleInput): Promise { + await this.getUserOrThrow(input.userId); + await this.getRoleOrThrow(input.roleId); + const existing = await this.userRoleRepository.findOneBy(input); + if (existing) return existing; + + const now = new Date().toISOString(); + return this.userRoleRepository.createOne({ + ...input, + createdAt: now, + updatedAt: now, + }); + } + + public async updateUserRoles( + input: UpdateUserRolesInput, + ): Promise { + const user = await this.getUserOrThrow(input.userId); + const roleIds = this.uniqueIds(input.roleIds); + const tx = await this.userRoleRepository.transaction(); + + try { + await this.userRoleRepository.deleteAllBy({ userId: user.id }, { tx }); + if (roleIds.length) { + await this.createUserRoleAssignments(user.id, roleIds, tx); + } + await tx.commit(); + } catch (error) { + await tx.rollback(); + throw error; + } + + const mappings = await this.userRoleRepository.findMappingsByUserId( + user.id, + ); + return { ...user, roles: mappings.map((mapping) => mapping.role) }; + } + + public async removeRoleFromUser(input: UserRoleInput): Promise { + const existing = await this.userRoleRepository.findOneBy(input); + if (!existing) return true; + await this.userRoleRepository.deleteOne(existing.id); + return true; + } + + public async getUserRoleMappings(): Promise { + return this.userRoleRepository.findMappings(); + } + + private async createUserRoleAssignments( + userId: number, + roleIds: number[], + tx: Knex.Transaction, + ): Promise { + const uniqueRoleIds = this.uniqueIds(roleIds); + for (const roleId of uniqueRoleIds) { + await this.getRoleOrThrow(roleId); + } + const now = new Date().toISOString(); + await this.userRoleRepository.createMany( + uniqueRoleIds.map((roleId) => ({ + userId, + roleId, + createdAt: now, + updatedAt: now, + })), + { tx }, + ); + } + + private async getRoleOrThrow(id: number): Promise { + const role = await this.roleRepository.findOneBy({ id }); + if (!role) throw new Error(`Role ${id} was not found.`); + return role; + } + + private async getUserOrThrow(id: number): Promise { + const user = await this.userRepository.findOneBy({ id }); + if (!user) throw new Error(`User ${id} was not found.`); + return user; + } + + private async assertUniqueRoleName(name: string, exceptId?: number) { + const roles = await this.roleRepository.findAll(); + const duplicate = roles.find( + (role) => + role.name.toLowerCase() === name.toLowerCase() && role.id !== exceptId, + ); + if (duplicate) throw new Error(`Role "${name}" already exists.`); + } + + private async assertUniqueUserEmail(email: string, exceptId?: number) { + const users = await this.userRepository.findAll(); + const duplicate = users.find( + (user) => + user.email.toLowerCase() === email.toLowerCase() && + user.id !== exceptId, + ); + if (duplicate) throw new Error(`User "${email}" already exists.`); + } + + private validateRoleName(name: string): string { + return this.validateRequiredText(name, 'Role name'); + } + + private validateEmail(email: string): string { + const normalized = this.validateRequiredText(email, 'Email').toLowerCase(); + if (!EMAIL_PATTERN.test(normalized)) { + throw new Error('A valid email address is required.'); + } + return normalized; + } + + private validateRequiredText(value: string, label: string): string { + const normalized = `${value || ''}`.trim(); + if (!normalized) throw new Error(`${label} is required.`); + return normalized; + } + + private normalizeNullable(value?: string | null): string | null { + const normalized = `${value || ''}`.trim(); + return normalized || null; + } + + private uniqueIds(ids: number[]): number[] { + return Array.from(new Set((ids || []).filter(Boolean))); + } +} diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index a037ec75ee..f92562a6e0 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -20,6 +20,9 @@ import { IInstructionRepository, IApiHistoryRepository, IDashboardItemRefreshJobRepository, + IRoleRepository, + IUserRepository, + IUserRoleRepository, } from '@server/repositories'; import { IQueryService, @@ -30,6 +33,7 @@ import { IProjectService, IDashboardService, IInstructionService, + IRbacService, } from '@server/services'; import { ITelemetry } from '@server/telemetry/telemetry'; import { @@ -59,6 +63,7 @@ export interface IContext { dashboardService: IDashboardService; sqlPairService: ISqlPairService; instructionService: IInstructionService; + rbacService: IRbacService; // repository projectRepository: IProjectRepository; @@ -76,6 +81,9 @@ export interface IContext { instructionRepository: IInstructionRepository; apiHistoryRepository: IApiHistoryRepository; dashboardItemRefreshJobRepository: IDashboardItemRefreshJobRepository; + roleRepository: IRoleRepository; + userRepository: IUserRepository; + userRoleRepository: IUserRoleRepository; // background trackers projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 33be8d9465..18ef4da4d2 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -19,6 +19,9 @@ import { InstructionRepository, ApiHistoryRepository, DashboardItemRefreshJobRepository, + RoleRepository, + UserRepository, + UserRoleRepository, } from '@server/repositories'; import { WrenEngineAdaptor, @@ -35,6 +38,7 @@ import { DashboardService, AskingTaskTracker, InstructionService, + RbacService, } from '@server/services'; import { PostHogTelemetry } from './apollo/server/telemetry/telemetry'; import { @@ -150,6 +154,9 @@ export const initComponents = () => { const apiHistoryRepository = new ApiHistoryRepository(knex); const dashboardItemRefreshJobRepository = new DashboardItemRefreshJobRepository(knex); + const roleRepository = new RoleRepository(knex); + const userRepository = new UserRepository(knex); + const userRoleRepository = new UserRoleRepository(knex); // adaptors const wrenEngineAdaptor = new WrenEngineAdaptor({ @@ -255,6 +262,11 @@ export const initComponents = () => { instructionRepository, wrenAIAdaptor, }); + const rbacService = new RbacService({ + roleRepository, + userRepository, + userRoleRepository, + }); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ dashboardRepository, @@ -288,6 +300,9 @@ export const initComponents = () => { apiHistoryRepository, instructionRepository, dashboardItemRefreshJobRepository, + roleRepository, + userRepository, + userRoleRepository, // adaptors wrenEngineAdaptor, @@ -304,6 +319,7 @@ export const initComponents = () => { dashboardService, sqlPairService, instructionService, + rbacService, askingTaskTracker, // background trackers diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index a3a0292c49..528dfdab5b 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -79,6 +79,14 @@ export default function HeaderBar() { > API + router.push(Path.AdministrationUsers)} + > + Admin + )} diff --git a/wren-ui/src/components/pages/administration/types.tsx b/wren-ui/src/components/pages/administration/types.tsx new file mode 100644 index 0000000000..896bf8cf40 --- /dev/null +++ b/wren-ui/src/components/pages/administration/types.tsx @@ -0,0 +1,51 @@ +import { Tag } from 'antd'; + +export interface Role { + id: number; + name: string; + description?: string | null; + createdAt: string; + updatedAt: string; + users?: User[]; +} + +export interface User { + id: number; + name: string; + email: string; + externalId?: string | null; + identityProvider?: string | null; + isActive: boolean; + roles?: Role[]; + createdAt: string; + updatedAt: string; +} + +export interface UserRoleMapping { + id: number; + userId: number; + roleId: number; + user: User; + role: Role; + createdAt: string; + updatedAt: string; +} + +export const RoleTags = ({ roles = [] }: { roles?: Role[] }) => { + if (!roles.length) return No roles; + return ( + <> + {roles.map((role) => ( + + {role.name} + + ))} + + ); +}; + +export const StatusTag = ({ active }: { active: boolean }) => ( + + {active ? 'Active' : 'Inactive'} + +); diff --git a/wren-ui/src/components/sidebar/Administration.tsx b/wren-ui/src/components/sidebar/Administration.tsx new file mode 100644 index 0000000000..ecec7dd7ea --- /dev/null +++ b/wren-ui/src/components/sidebar/Administration.tsx @@ -0,0 +1,73 @@ +import Link from 'next/link'; +import { useRouter } from 'next/router'; +import styled from 'styled-components'; +import TeamOutlined from '@ant-design/icons/TeamOutlined'; +import IdcardOutlined from '@ant-design/icons/IdcardOutlined'; +import SafetyCertificateOutlined from '@ant-design/icons/SafetyCertificateOutlined'; +import { Path, MENU_KEY } from '@/utils/enum'; +import SidebarMenu from '@/components/sidebar/SidebarMenu'; + +const Layout = styled.div` + padding: 16px 0; + position: absolute; + z-index: 1; + left: 0; + top: 0; + width: 100%; + background-color: var(--gray-2); + overflow: hidden; +`; + +const MENU_KEY_MAP = { + [Path.AdministrationUsers]: MENU_KEY.ADMIN_USERS, + [Path.AdministrationRoles]: MENU_KEY.ADMIN_ROLES, + [Path.AdministrationAssignments]: MENU_KEY.ADMIN_ASSIGNMENTS, +}; + +const linkStyle = { color: 'inherit', transition: 'none' }; + +export default function Administration() { + const router = useRouter(); + + const menuItems = [ + { + label: ( + + User Management + + ), + icon: , + key: MENU_KEY.ADMIN_USERS, + className: 'pl-4', + }, + { + label: ( + + Role Management + + ), + icon: , + key: MENU_KEY.ADMIN_ROLES, + className: 'pl-4', + }, + { + label: ( + + User Role Assignment + + ), + icon: , + key: MENU_KEY.ADMIN_ASSIGNMENTS, + className: 'pl-4', + }, + ]; + + return ( + + + + ); +} diff --git a/wren-ui/src/components/sidebar/index.tsx b/wren-ui/src/components/sidebar/index.tsx index 442c7255ed..0ed182995b 100644 --- a/wren-ui/src/components/sidebar/index.tsx +++ b/wren-ui/src/components/sidebar/index.tsx @@ -9,6 +9,7 @@ import Home, { Props as HomeSidebarProps } from './Home'; import Modeling, { Props as ModelingSidebarProps } from './Modeling'; import Knowledge from './Knowledge'; import APIManagement from './APIManagement'; +import Administration from './Administration'; import LearningSection from '@/components/learning'; const Layout = styled.div` @@ -68,6 +69,10 @@ const DynamicSidebar = ( return ; } + if (pathname.startsWith(Path.Administration)) { + return ; + } + return null; }; diff --git a/wren-ui/src/pages/administration/assignments.tsx b/wren-ui/src/pages/administration/assignments.tsx new file mode 100644 index 0000000000..252691c282 --- /dev/null +++ b/wren-ui/src/pages/administration/assignments.tsx @@ -0,0 +1,281 @@ +import { useEffect, useState } from 'react'; +import { useMutation, useQuery } from '@apollo/client'; +import { + Button, + Form, + Modal, + Popconfirm, + Select, + Table, + TableColumnsType, + Typography, + message, +} from 'antd'; +import SafetyCertificateOutlined from '@ant-design/icons/SafetyCertificateOutlined'; +import EditOutlined from '@ant-design/icons/EditOutlined'; +import DeleteOutlined from '@ant-design/icons/DeleteOutlined'; +import SiderLayout from '@/components/layouts/SiderLayout'; +import PageLayout from '@/components/layouts/PageLayout'; +import { + ASSIGN_ROLE_TO_USER, + LIST_USER_ROLE_MAPPINGS, + REMOVE_ROLE_FROM_USER, + UPDATE_USER_ROLES, +} from '@/apollo/client/graphql/rbac'; +import { + Role, + RoleTags, + User, + UserRoleMapping, +} from '@/components/pages/administration/types'; +import { getAbsoluteTime } from '@/utils/time'; + +const { Text } = Typography; + +type AssignmentModalState = { + visible: boolean; + user?: User; +}; + +const AssignmentModal = ({ + users, + roles, + state, + loading, + onClose, + onSubmit, +}: { + users: User[]; + roles: Role[]; + state: AssignmentModalState; + loading: boolean; + onClose: () => void; + onSubmit: (values: any, user?: User) => Promise; +}) => { + const [form] = Form.useForm(); + const isUpdate = !!state.user; + + useEffect(() => { + if (!state.visible) return; + form.setFieldsValue({ + userId: state.user?.id, + roleId: undefined, + roleIds: state.user?.roles?.map((role) => role.id) || [], + }); + }, [form, state.visible, state.user]); + + const submit = async () => { + const values = await form.validateFields(); + await onSubmit(values, state.user); + form.resetFields(); + onClose(); + }; + + return ( + form.resetFields()} + > +
+ + ({ + label: role.name, + value: role.id, + }))} + /> + + ) : ( + +
+ + + + ); +} diff --git a/wren-ui/src/pages/administration/index.tsx b/wren-ui/src/pages/administration/index.tsx new file mode 100644 index 0000000000..653c921b87 --- /dev/null +++ b/wren-ui/src/pages/administration/index.tsx @@ -0,0 +1,14 @@ +import { useEffect } from 'react'; +import { useRouter } from 'next/router'; +import PageLoading from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; + +export default function AdministrationIndex() { + const router = useRouter(); + + useEffect(() => { + router.replace(Path.AdministrationUsers); + }, [router]); + + return ; +} diff --git a/wren-ui/src/pages/administration/roles.tsx b/wren-ui/src/pages/administration/roles.tsx new file mode 100644 index 0000000000..a17f1457e1 --- /dev/null +++ b/wren-ui/src/pages/administration/roles.tsx @@ -0,0 +1,245 @@ +import { useEffect, useState } from 'react'; +import { useMutation, useQuery } from '@apollo/client'; +import { + Button, + Form, + Input, + Modal, + Table, + TableColumnsType, + Tag, + Typography, + message, +} from 'antd'; +import IdcardOutlined from '@ant-design/icons/IdcardOutlined'; +import EditOutlined from '@ant-design/icons/EditOutlined'; +import SiderLayout from '@/components/layouts/SiderLayout'; +import PageLayout from '@/components/layouts/PageLayout'; +import { + CREATE_ROLE, + LIST_RBAC_ROLES, + UPDATE_ROLE, +} from '@/apollo/client/graphql/rbac'; +import { Role, User } from '@/components/pages/administration/types'; +import { getAbsoluteTime } from '@/utils/time'; + +const { Paragraph, Text } = Typography; + +type RoleModalState = { + visible: boolean; + role?: Role; +}; + +const RoleModal = ({ + state, + loading, + onClose, + onSubmit, +}: { + state: RoleModalState; + loading: boolean; + onClose: () => void; + onSubmit: (values: any, role?: Role) => Promise; +}) => { + const [form] = Form.useForm(); + const isEdit = !!state.role; + + useEffect(() => { + if (!state.visible) return; + form.setFieldsValue({ + name: state.role?.name, + description: state.role?.description, + }); + }, [form, state.visible, state.role]); + + const submit = async () => { + const values = await form.validateFields(); + await onSubmit(values, state.role); + form.resetFields(); + onClose(); + }; + + return ( + form.resetFields()} + > + + + + + + + + + + ); +}; + +const AssignedUsers = ({ users = [] }: { users?: User[] }) => { + if (!users.length) return No users assigned; + return ( + <> + {users.map((user) => ( + + {user.name} - {user.email} + + ))} + + ); +}; + +export default function RoleManagement() { + const [modalState, setModalState] = useState({ + visible: false, + }); + const { data, loading, refetch } = useQuery(LIST_RBAC_ROLES, { + fetchPolicy: 'cache-and-network', + }); + const roles: Role[] = data?.roles || []; + + const mutationOptions = { + onError: (error) => message.error(error.message), + }; + const [createRole, createRoleState] = useMutation( + CREATE_ROLE, + mutationOptions, + ); + const [updateRole, updateRoleState] = useMutation( + UPDATE_ROLE, + mutationOptions, + ); + + const closeModal = () => setModalState({ visible: false }); + + const submitRole = async (values: any, role?: Role) => { + if (role) { + await updateRole({ + variables: { + where: { id: role.id }, + data: { + name: values.name, + description: values.description || null, + }, + }, + }); + message.success('Successfully updated role.'); + } else { + await createRole({ + variables: { + data: { + name: values.name, + description: values.description || null, + }, + }, + }); + message.success('Successfully created role.'); + } + await refetch(); + }; + + const columns: TableColumnsType = [ + { + title: 'Role', + dataIndex: 'name', + width: 220, + render: (name) => {name}, + }, + { + title: 'Description', + dataIndex: 'description', + render: (description) => ( + + {description || No description} + + ), + }, + { + title: 'Assigned users', + dataIndex: 'users', + render: (users) => , + }, + { + title: 'Created', + dataIndex: 'createdAt', + width: 180, + render: (value) => ( + {getAbsoluteTime(value)} + ), + }, + { + title: 'Actions', + width: 110, + align: 'center', + render: (_, record) => ( + + ), + }, + ]; + + return ( + + + + Role Management + + } + description="Maintain reusable role definitions. Permission enforcement can be added later without changing these assignments." + titleExtra={ + + } + > +
+ + + + ); +} diff --git a/wren-ui/src/pages/administration/users.tsx b/wren-ui/src/pages/administration/users.tsx new file mode 100644 index 0000000000..ad4d605bde --- /dev/null +++ b/wren-ui/src/pages/administration/users.tsx @@ -0,0 +1,296 @@ +import { useEffect, useState } from 'react'; +import { useMutation, useQuery } from '@apollo/client'; +import { + Button, + Form, + Input, + Modal, + Select, + Switch, + Table, + TableColumnsType, + Typography, + message, +} from 'antd'; +import TeamOutlined from '@ant-design/icons/TeamOutlined'; +import EditOutlined from '@ant-design/icons/EditOutlined'; +import SiderLayout from '@/components/layouts/SiderLayout'; +import PageLayout from '@/components/layouts/PageLayout'; +import { + CREATE_USER, + LIST_RBAC_USERS, + UPDATE_USER, + UPDATE_USER_ROLES, +} from '@/apollo/client/graphql/rbac'; +import { + Role, + RoleTags, + StatusTag, + User, +} from '@/components/pages/administration/types'; +import { getAbsoluteTime } from '@/utils/time'; + +const { Text } = Typography; + +type UserModalState = { + visible: boolean; + user?: User; +}; + +const UserModal = ({ + roles, + state, + loading, + onClose, + onSubmit, +}: { + roles: Role[]; + state: UserModalState; + loading: boolean; + onClose: () => void; + onSubmit: (values: any, user?: User) => Promise; +}) => { + const [form] = Form.useForm(); + const isEdit = !!state.user; + + useEffect(() => { + if (!state.visible) return; + form.setFieldsValue({ + name: state.user?.name, + email: state.user?.email, + externalId: state.user?.externalId, + identityProvider: state.user?.identityProvider, + isActive: state.user?.isActive ?? true, + roleIds: state.user?.roles?.map((role) => role.id) || [], + }); + }, [form, state.visible, state.user]); + + const submit = async () => { + const values = await form.validateFields(); + await onSubmit(values, state.user); + form.resetFields(); + onClose(); + }; + + return ( + form.resetFields()} + > +
+ + + + + + + + + + + + + + + + + + ); +}; + +export default function UserManagement() { + const [modalState, setModalState] = useState({ + visible: false, + }); + const { data, loading, refetch } = useQuery(LIST_RBAC_USERS, { + fetchPolicy: 'cache-and-network', + }); + const users: User[] = data?.users || []; + const roles: Role[] = data?.roles || []; + + const mutationOptions = { + onError: (error) => message.error(error.message), + }; + const [createUser, createUserState] = useMutation( + CREATE_USER, + mutationOptions, + ); + const [updateUser, updateUserState] = useMutation( + UPDATE_USER, + mutationOptions, + ); + const [updateUserRoles, updateRolesState] = useMutation( + UPDATE_USER_ROLES, + mutationOptions, + ); + + const closeModal = () => setModalState({ visible: false }); + + const submitUser = async (values: any, user?: User) => { + if (user) { + await updateUser({ + variables: { + where: { id: user.id }, + data: { + name: values.name, + email: values.email, + externalId: values.externalId || null, + identityProvider: values.identityProvider || null, + isActive: values.isActive, + }, + }, + }); + await updateUserRoles({ + variables: { + data: { userId: user.id, roleIds: values.roleIds || [] }, + }, + }); + message.success('Successfully updated user.'); + } else { + await createUser({ + variables: { + data: { + name: values.name, + email: values.email, + externalId: values.externalId || null, + identityProvider: values.identityProvider || null, + isActive: values.isActive, + roleIds: values.roleIds || [], + }, + }, + }); + message.success('Successfully created user.'); + } + await refetch(); + }; + + const columns: TableColumnsType = [ + { + title: 'User', + dataIndex: 'name', + render: (_, record) => ( +
+
{record.name}
+ {record.email} +
+ ), + }, + { + title: 'Status', + dataIndex: 'isActive', + width: 120, + render: (active) => , + }, + { + title: 'Assigned roles', + dataIndex: 'roles', + render: (assignedRoles) => , + }, + { + title: 'Identity provider', + dataIndex: 'identityProvider', + width: 180, + render: (value) => value || Local, + }, + { + title: 'Created', + dataIndex: 'createdAt', + width: 180, + render: (value) => ( + {getAbsoluteTime(value)} + ), + }, + { + title: 'Actions', + width: 110, + align: 'center', + render: (_, record) => ( + + ), + }, + ]; + + return ( + + + + User Management + + } + description="Create users, maintain local identity metadata, and review assigned foundation roles." + titleExtra={ + + } + > +
+ + + + ); +} diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index 60ed8334f4..b31de21a75 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -47,6 +47,9 @@ const bootstrapServer = async () => { instructionRepository, apiHistoryRepository, dashboardItemRefreshJobRepository, + roleRepository, + userRepository, + userRoleRepository, // adaptors wrenEngineAdaptor, ibisAdaptor, @@ -62,6 +65,7 @@ const bootstrapServer = async () => { sqlPairService, instructionService, + rbacService, // background trackers projectRecommendQuestionBackgroundTracker, threadRecommendQuestionBackgroundTracker, @@ -142,6 +146,7 @@ const bootstrapServer = async () => { dashboardService, sqlPairService, instructionService, + rbacService, // repository projectRepository, modelRepository, @@ -158,6 +163,9 @@ const bootstrapServer = async () => { instructionRepository, apiHistoryRepository, dashboardItemRefreshJobRepository, + roleRepository, + userRepository, + userRoleRepository, // background trackers projectRecommendQuestionBackgroundTracker, threadRecommendQuestionBackgroundTracker, diff --git a/wren-ui/src/utils/enum/menu.ts b/wren-ui/src/utils/enum/menu.ts index 2d5e55c354..4ab0941256 100644 --- a/wren-ui/src/utils/enum/menu.ts +++ b/wren-ui/src/utils/enum/menu.ts @@ -3,4 +3,7 @@ export enum MENU_KEY { INSTRUCTIONS = 'instructions', API_HISTORY = 'api-history', API_REFERENCE = 'api-reference', + ADMIN_USERS = 'admin-users', + ADMIN_ROLES = 'admin-roles', + ADMIN_ASSIGNMENTS = 'admin-assignments', } diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index 10f5fa8091..54c0eb91f0 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -12,4 +12,8 @@ export enum Path { KnowledgeInstructions = '/knowledge/instructions', APIManagement = '/api-management', APIManagementHistory = '/api-management/history', + Administration = '/administration', + AdministrationUsers = '/administration/users', + AdministrationRoles = '/administration/roles', + AdministrationAssignments = '/administration/assignments', } From 33ad5c36ddfecda2a200825efdeab5fc8b00214e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 3 Jun 2026 02:58:09 +0530 Subject: [PATCH 0084/1087] index --- wren-ui/src/apollo/client/index.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/client/index.ts b/wren-ui/src/apollo/client/index.ts index df6d75ffde..c36af89098 100644 --- a/wren-ui/src/apollo/client/index.ts +++ b/wren-ui/src/apollo/client/index.ts @@ -1,7 +1,8 @@ import { ApolloClient, HttpLink, InMemoryCache, from } from '@apollo/client'; -import { onError } from '@apollo/client/link/error'; import errorHandler from '@/utils/errorHandler'; +const { onError } = require('@apollo/client/link/error/error.cjs'); + const apolloErrorLink = onError((error) => errorHandler(error)); const httpLink = new HttpLink({ From b74fcff747f0a99bb01b544ae3c260c765803968 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 3 Jun 2026 13:33:22 +0530 Subject: [PATCH 0085/1087] apollo --- wren-ui/docs/rbac-architecture.md | 70 ++- ...0_create_organization_membership_tables.js | 174 ++++++ wren-ui/src/apollo/client/graphql/rbac.ts | 142 ++++- wren-ui/src/apollo/server/models/rbac.ts | 38 ++ .../server/repositories/rbacRepository.ts | 352 ++++++++++++ wren-ui/src/apollo/server/resolvers.ts | 8 + .../apollo/server/resolvers/rbacResolver.ts | 100 +++- wren-ui/src/apollo/server/schema.ts | 79 +++ .../src/apollo/server/services/rbacService.ts | 515 +++++++++++++++++- wren-ui/src/apollo/server/types/context.ts | 10 + wren-ui/src/apollo/server/utils/auth.ts | 54 ++ wren-ui/src/apollo/server/utils/index.ts | 1 + wren-ui/src/common.ts | 16 + wren-ui/src/components/HeaderBar.tsx | 42 +- .../components/pages/administration/types.tsx | 39 ++ .../src/components/sidebar/Administration.tsx | 2 +- wren-ui/src/hooks/useAuth.tsx | 134 +++++ wren-ui/src/pages/_app.tsx | 15 +- wren-ui/src/pages/accept-invitation.tsx | 92 ++++ .../src/pages/administration/assignments.tsx | 210 ++----- .../src/pages/api/auth/accept-invitation.ts | 17 + wren-ui/src/pages/api/auth/bootstrap.ts | 17 + wren-ui/src/pages/api/auth/login.ts | 17 + wren-ui/src/pages/api/auth/logout.ts | 18 + wren-ui/src/pages/api/auth/me.ts | 21 + wren-ui/src/pages/api/auth/status.ts | 11 + wren-ui/src/pages/api/graphql.ts | 99 ++-- wren-ui/src/pages/login.tsx | 112 ++++ wren-ui/src/utils/enum/path.ts | 2 + 29 files changed, 2156 insertions(+), 251 deletions(-) create mode 100644 wren-ui/migrations/20250603000000_create_organization_membership_tables.js create mode 100644 wren-ui/src/apollo/server/utils/auth.ts create mode 100644 wren-ui/src/hooks/useAuth.tsx create mode 100644 wren-ui/src/pages/accept-invitation.tsx create mode 100644 wren-ui/src/pages/api/auth/accept-invitation.ts create mode 100644 wren-ui/src/pages/api/auth/bootstrap.ts create mode 100644 wren-ui/src/pages/api/auth/login.ts create mode 100644 wren-ui/src/pages/api/auth/logout.ts create mode 100644 wren-ui/src/pages/api/auth/me.ts create mode 100644 wren-ui/src/pages/api/auth/status.ts create mode 100644 wren-ui/src/pages/login.tsx diff --git a/wren-ui/docs/rbac-architecture.md b/wren-ui/docs/rbac-architecture.md index cec918bbfe..caa594f115 100644 --- a/wren-ui/docs/rbac-architecture.md +++ b/wren-ui/docs/rbac-architecture.md @@ -1,20 +1,22 @@ # RBAC Foundation -This document describes the application-level RBAC foundation in Wren UI. It does not enforce permissions yet; it establishes durable role, user, and user-role assignment primitives for future governance work. +This document describes the application-level organization and member RBAC foundation in Wren UI. It establishes durable organization, user, member, invitation, session, and role primitives for future governance work. ## Scope Implemented: - Roles: `Admin`, `Manager`, `Analyst`, `Viewer`, plus custom roles. -- Users with local identity metadata. -- User-role mappings. -- GraphQL APIs for role, user, and assignment management. -- Administration UI for user management, role management, and role assignment. +- Organizations with active members. +- Users with local identity metadata and password hashes for local login. +- Organization-member role assignments. +- Member invitations and local auth sessions. +- GraphQL APIs for role, member, invitation, and assignment management. +- Administration UI for member management, role management, and role assignment. Deferred: -- Authorization middleware. +- Authorization middleware for the existing AI/query APIs. - Governance policies. - Data scoping. - SQL validation. @@ -43,6 +45,38 @@ Tables: - `user_id` - `role_id` - timestamps +- `organizations` + - `id` + - `name` + - `slug` + - external identity fields + - `is_active` + - timestamps +- `organization_members` + - `id` + - `organization_id` + - `user_id` + - `role_id` + - `status` + - `joined_at` + - timestamps +- `member_invitations` + - `id` + - `organization_id` + - `role_id` + - `email` + - `token` + - `status` + - `expires_at` + - timestamps +- `auth_sessions` + - `id` + - `user_id` + - `organization_member_id` + - `token` + - `expires_at` + - `revoked_at` + - timestamps `external_id` and `identity_provider` are intentionally present now so Teams, LDAP, and Azure AD integrations can later attach external identities without replacing the RBAC tables. @@ -64,6 +98,11 @@ Queries: - `roles` - `users` - `userRoleMappings` +- `organizations` +- `organizationMembers` +- `memberInvitations` +- `currentSession` +- `bootstrapStatus` Mutations: @@ -74,6 +113,16 @@ Mutations: - `assignRoleToUser` - `updateUserRoles` - `removeRoleFromUser` +- `inviteMember` +- `updateMember` +- `updateMemberRole` + +Admin-only behavior: + +- Inviting members. +- Creating and editing roles. +- Editing members and updating member roles. +- Legacy user-role mutation paths are also guarded for Admin members. ## UI @@ -84,12 +133,19 @@ Navigation: Screens: -- `/administration/users` +- `/administration/users` (Member Management) - `/administration/roles` - `/administration/assignments` The UI uses the existing Next.js, Apollo Client, Ant Design, and `SiderLayout`/`PageLayout` patterns. +Authentication routes: + +- `/login` +- `/accept-invitation?token=...` + +The first Admin can bootstrap the first organization when no active Admin member exists. + ## Future Permission Model Future schema-level or table-level permissions should be added as separate tables referencing `roles.id`, for example: diff --git a/wren-ui/migrations/20250603000000_create_organization_membership_tables.js b/wren-ui/migrations/20250603000000_create_organization_membership_tables.js new file mode 100644 index 0000000000..ae917bd847 --- /dev/null +++ b/wren-ui/migrations/20250603000000_create_organization_membership_tables.js @@ -0,0 +1,174 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + const hasUsers = await knex.schema.hasTable('users'); + const hasRoles = await knex.schema.hasTable('roles'); + + if (hasUsers) { + const hasPasswordHash = await knex.schema.hasColumn( + 'users', + 'password_hash', + ); + const hasLastLoginAt = await knex.schema.hasColumn( + 'users', + 'last_login_at', + ); + await knex.schema.alterTable('users', (table) => { + if (!hasPasswordHash) { + table.string('password_hash', 255).nullable(); + } + if (!hasLastLoginAt) { + table.timestamp('last_login_at').nullable(); + } + }); + } + + const hasOrganizations = await knex.schema.hasTable('organizations'); + if (!hasOrganizations) { + await knex.schema.createTable('organizations', (table) => { + table.increments('id').primary(); + table.string('name', 160).notNullable().unique(); + table.string('slug', 180).notNullable().unique(); + table.string('external_id', 255).nullable().unique(); + table.string('identity_provider', 80).nullable(); + table.boolean('is_active').notNullable().defaultTo(true); + table.timestamps(true, true); + }); + } + + const hasOrganizationMembers = await knex.schema.hasTable( + 'organization_members', + ); + if (!hasOrganizationMembers) { + await knex.schema.createTable('organization_members', (table) => { + table.increments('id').primary(); + table.integer('organization_id').notNullable(); + table.integer('user_id').notNullable(); + table.integer('role_id').notNullable(); + table.string('status', 40).notNullable().defaultTo('active'); + table.integer('invited_by_member_id').nullable(); + table.timestamp('joined_at').nullable(); + table.timestamps(true, true); + + table + .foreign('organization_id') + .references('organizations.id') + .onDelete('CASCADE'); + table.foreign('user_id').references('users.id').onDelete('CASCADE'); + table.foreign('role_id').references('roles.id').onDelete('RESTRICT'); + table.unique(['organization_id', 'user_id']); + }); + } + + const hasMemberInvitations = await knex.schema.hasTable('member_invitations'); + if (!hasMemberInvitations) { + await knex.schema.createTable('member_invitations', (table) => { + table.increments('id').primary(); + table.integer('organization_id').notNullable(); + table.integer('role_id').notNullable(); + table.string('email', 320).notNullable(); + table.string('name', 160).nullable(); + table.string('token', 128).notNullable().unique(); + table.string('status', 40).notNullable().defaultTo('pending'); + table.integer('invited_by_member_id').nullable(); + table.timestamp('expires_at').notNullable(); + table.timestamp('accepted_at').nullable(); + table.timestamps(true, true); + + table + .foreign('organization_id') + .references('organizations.id') + .onDelete('CASCADE'); + table.foreign('role_id').references('roles.id').onDelete('RESTRICT'); + table.unique(['organization_id', 'email', 'status']); + }); + } + + const hasAuthSessions = await knex.schema.hasTable('auth_sessions'); + if (!hasAuthSessions) { + await knex.schema.createTable('auth_sessions', (table) => { + table.increments('id').primary(); + table.integer('user_id').notNullable(); + table.integer('organization_member_id').notNullable(); + table.string('token', 128).notNullable().unique(); + table.timestamp('expires_at').notNullable(); + table.timestamp('revoked_at').nullable(); + table.timestamps(true, true); + + table.foreign('user_id').references('users.id').onDelete('CASCADE'); + table + .foreign('organization_member_id') + .references('organization_members.id') + .onDelete('CASCADE'); + }); + } + + if (hasOrganizations || !hasUsers || !hasRoles) return; + + const now = new Date().toISOString(); + const [organization] = await knex('organizations') + .insert({ + name: 'Default organization', + slug: 'default', + is_active: true, + created_at: now, + updated_at: now, + }) + .returning('*'); + + const adminRole = await knex('roles').where({ name: 'Admin' }).first(); + const roles = await knex('roles'); + const users = await knex('users'); + const userRoles = await knex('user_roles'); + + const roleById = new Map(roles.map((role) => [role.id, role])); + const firstRoleByUserId = new Map(); + userRoles.forEach((mapping) => { + if (!firstRoleByUserId.has(mapping.user_id)) { + firstRoleByUserId.set(mapping.user_id, mapping.role_id); + } + }); + + for (const user of users) { + const roleId = firstRoleByUserId.get(user.id) || adminRole?.id; + if (!roleById.has(roleId)) continue; + await knex('organization_members').insert({ + organization_id: organization.id, + user_id: user.id, + role_id: roleId, + status: 'active', + joined_at: now, + created_at: now, + updated_at: now, + }); + } +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex.schema.dropTableIfExists('auth_sessions'); + await knex.schema.dropTableIfExists('member_invitations'); + await knex.schema.dropTableIfExists('organization_members'); + await knex.schema.dropTableIfExists('organizations'); + + const hasUsers = await knex.schema.hasTable('users'); + if (hasUsers) { + const hasPasswordHash = await knex.schema.hasColumn( + 'users', + 'password_hash', + ); + const hasLastLoginAt = await knex.schema.hasColumn( + 'users', + 'last_login_at', + ); + await knex.schema.alterTable('users', (table) => { + if (hasPasswordHash) table.dropColumn('password_hash'); + if (hasLastLoginAt) table.dropColumn('last_login_at'); + }); + } +}; diff --git a/wren-ui/src/apollo/client/graphql/rbac.ts b/wren-ui/src/apollo/client/graphql/rbac.ts index 6da6060285..9816fa8abd 100644 --- a/wren-ui/src/apollo/client/graphql/rbac.ts +++ b/wren-ui/src/apollo/client/graphql/rbac.ts @@ -23,13 +23,68 @@ export const USER_FIELDS = gql` } `; +export const ORGANIZATION_FIELDS = gql` + fragment OrganizationFields on Organization { + id + name + slug + isActive + createdAt + updatedAt + } +`; + +export const MEMBER_FIELDS = gql` + fragment MemberFields on OrganizationMember { + id + organizationId + userId + roleId + status + joinedAt + createdAt + updatedAt + user { + ...UserFields + } + role { + ...RoleFields + } + organization { + ...OrganizationFields + } + } +`; + +export const INVITATION_FIELDS = gql` + fragment InvitationFields on MemberInvitation { + id + organizationId + roleId + email + name + token + status + expiresAt + acceptedAt + createdAt + updatedAt + role { + ...RoleFields + } + organization { + ...OrganizationFields + } + } +`; + export const LIST_RBAC_USERS = gql` query RbacUsers { - users { - ...UserFields - roles { - ...RoleFields - } + organizationMembers { + ...MemberFields + } + memberInvitations { + ...InvitationFields } roles { ...RoleFields @@ -38,48 +93,34 @@ export const LIST_RBAC_USERS = gql` ${USER_FIELDS} ${ROLE_FIELDS} + ${ORGANIZATION_FIELDS} + ${MEMBER_FIELDS} + ${INVITATION_FIELDS} `; export const LIST_RBAC_ROLES = gql` query RbacRoles { roles { ...RoleFields - users { - ...UserFields - } } - users { - ...UserFields + organizationMembers { + ...MemberFields } } ${ROLE_FIELDS} ${USER_FIELDS} + ${ORGANIZATION_FIELDS} + ${MEMBER_FIELDS} `; export const LIST_USER_ROLE_MAPPINGS = gql` query UserRoleMappings { - userRoleMappings { - id - userId - roleId - createdAt - updatedAt - user { - ...UserFields - roles { - ...RoleFields - } - } - role { - ...RoleFields - } + organizationMembers { + ...MemberFields } - users { - ...UserFields - roles { - ...RoleFields - } + memberInvitations { + ...InvitationFields } roles { ...RoleFields @@ -88,6 +129,9 @@ export const LIST_USER_ROLE_MAPPINGS = gql` ${USER_FIELDS} ${ROLE_FIELDS} + ${ORGANIZATION_FIELDS} + ${MEMBER_FIELDS} + ${INVITATION_FIELDS} `; export const CREATE_ROLE = gql` @@ -167,3 +211,41 @@ export const REMOVE_ROLE_FROM_USER = gql` removeRoleFromUser(data: $data) } `; + +export const INVITE_MEMBER = gql` + mutation InviteMember($data: InviteMemberInput!) { + inviteMember(data: $data) { + ...InvitationFields + } + } + + ${ROLE_FIELDS} + ${ORGANIZATION_FIELDS} + ${INVITATION_FIELDS} +`; + +export const UPDATE_MEMBER = gql` + mutation UpdateMember($data: UpdateMemberInput!) { + updateMember(data: $data) { + ...MemberFields + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} + ${ORGANIZATION_FIELDS} + ${MEMBER_FIELDS} +`; + +export const UPDATE_MEMBER_ROLE = gql` + mutation UpdateMemberRole($data: UpdateMemberRoleInput!) { + updateMemberRole(data: $data) { + ...MemberFields + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} + ${ORGANIZATION_FIELDS} + ${MEMBER_FIELDS} +`; diff --git a/wren-ui/src/apollo/server/models/rbac.ts b/wren-ui/src/apollo/server/models/rbac.ts index 52b1a98ce9..e35defbc2f 100644 --- a/wren-ui/src/apollo/server/models/rbac.ts +++ b/wren-ui/src/apollo/server/models/rbac.ts @@ -12,6 +12,7 @@ export interface UpdateRoleInput { export interface CreateUserInput { name: string; email: string; + password?: string | null; externalId?: string | null; identityProvider?: string | null; isActive?: boolean; @@ -36,3 +37,40 @@ export interface UpdateUserRolesInput { userId: number; roleIds: number[]; } + +export interface LoginInput { + email: string; + password: string; +} + +export interface BootstrapAdminInput { + organizationName: string; + name: string; + email: string; + password: string; +} + +export interface InviteMemberInput { + organizationId?: number | null; + email: string; + name?: string | null; + roleId: number; +} + +export interface AcceptInvitationInput { + token: string; + name?: string | null; + password: string; +} + +export interface UpdateMemberInput { + id: number; + name?: string | null; + roleId?: number | null; + status?: string | null; +} + +export interface UpdateMemberRoleInput { + memberId: number; + roleId: number; +} diff --git a/wren-ui/src/apollo/server/repositories/rbacRepository.ts b/wren-ui/src/apollo/server/repositories/rbacRepository.ts index d742af6579..c3c3bd6e3f 100644 --- a/wren-ui/src/apollo/server/repositories/rbacRepository.ts +++ b/wren-ui/src/apollo/server/repositories/rbacRepository.ts @@ -17,9 +17,11 @@ export interface RbacUser { id: number; name: string; email: string; + passwordHash?: string | null; externalId?: string | null; identityProvider?: string | null; isActive: boolean; + lastLoginAt?: string | null; createdAt: string; updatedAt: string; } @@ -37,6 +39,67 @@ export interface UserRoleMapping extends UserRole { role: Role; } +export interface Organization { + id: number; + name: string; + slug: string; + externalId?: string | null; + identityProvider?: string | null; + isActive: boolean; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationMember { + id: number; + organizationId: number; + userId: number; + roleId: number; + status: string; + invitedByMemberId?: number | null; + joinedAt?: string | null; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationMemberMapping extends OrganizationMember { + organization: Organization; + user: RbacUser; + role: Role; +} + +export interface MemberInvitation { + id: number; + organizationId: number; + roleId: number; + email: string; + name?: string | null; + token: string; + status: string; + invitedByMemberId?: number | null; + expiresAt: string; + acceptedAt?: string | null; + createdAt: string; + updatedAt: string; +} + +export interface MemberInvitationMapping extends MemberInvitation { + organization: Organization; + role: Role; + invitedBy?: OrganizationMemberMapping | null; +} + +export interface AuthSession { + id: number; + userId: number; + organizationMemberId: number; + token: string; + expiresAt: string; + revokedAt?: string | null; + createdAt: string; + updatedAt: string; +} + export interface IRoleRepository extends IBasicRepository {} export interface IUserRepository extends IBasicRepository {} @@ -53,6 +116,46 @@ export interface IUserRoleRepository extends IBasicRepository { ): Promise; } +export interface IOrganizationRepository + extends IBasicRepository {} + +export interface IOrganizationMemberRepository + extends IBasicRepository { + findMappings( + queryOptions?: IQueryOptions, + ): Promise; + findMappingsByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ): Promise; + findMappingById( + id: number, + queryOptions?: IQueryOptions, + ): Promise; + findActiveMappingByUserId( + userId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export interface IMemberInvitationRepository + extends IBasicRepository { + findMappings( + queryOptions?: IQueryOptions, + ): Promise; + findMappingByToken( + token: string, + queryOptions?: IQueryOptions, + ): Promise; +} + +export interface IAuthSessionRepository extends IBasicRepository { + findActiveByToken( + token: string, + queryOptions?: IQueryOptions, + ): Promise; +} + export class RoleRepository extends BaseRepository implements IRoleRepository @@ -158,3 +261,252 @@ export class UserRoleRepository })); } } + +export class OrganizationRepository + extends BaseRepository + implements IOrganizationRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organizations' }); + } +} + +export class OrganizationMemberRepository + extends BaseRepository + implements IOrganizationMemberRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization_members' }); + } + + public async findMappings(queryOptions?: IQueryOptions) { + return this.queryMappings({}, queryOptions); + } + + public async findMappingsByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ) { + return this.queryMappings({ organizationId }, queryOptions); + } + + public async findMappingById(id: number, queryOptions?: IQueryOptions) { + const [mapping] = await this.queryMappings({ id }, queryOptions); + return mapping || null; + } + + public async findActiveMappingByUserId( + userId: number, + queryOptions?: IQueryOptions, + ) { + const [mapping] = await this.queryMappings( + { userId, status: 'active' }, + queryOptions, + ); + return mapping || null; + } + + private async queryMappings( + filter: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const query = executer('organization_members') + .select( + 'organization_members.*', + 'organizations.id as organization__id', + 'organizations.name as organization__name', + 'organizations.slug as organization__slug', + 'organizations.external_id as organization__external_id', + 'organizations.identity_provider as organization__identity_provider', + 'organizations.is_active as organization__is_active', + 'organizations.created_at as organization__created_at', + 'organizations.updated_at as organization__updated_at', + 'users.id as user__id', + 'users.name as user__name', + 'users.email as user__email', + 'users.password_hash as user__password_hash', + 'users.external_id as user__external_id', + 'users.identity_provider as user__identity_provider', + 'users.is_active as user__is_active', + 'users.last_login_at as user__last_login_at', + 'users.created_at as user__created_at', + 'users.updated_at as user__updated_at', + 'roles.id as role__id', + 'roles.name as role__name', + 'roles.description as role__description', + 'roles.created_at as role__created_at', + 'roles.updated_at as role__updated_at', + ) + .join( + 'organizations', + 'organization_members.organization_id', + 'organizations.id', + ) + .join('users', 'organization_members.user_id', 'users.id') + .join('roles', 'organization_members.role_id', 'roles.id') + .orderBy('users.email'); + + if (filter.id) query.where('organization_members.id', filter.id); + if (filter.organizationId) { + query.where( + 'organization_members.organization_id', + filter.organizationId, + ); + } + if (filter.userId) + query.where('organization_members.user_id', filter.userId); + if (filter.status) + query.where('organization_members.status', filter.status); + + const rows = await query; + return rows.map((row) => ({ + id: row.id, + organizationId: row.organization_id, + userId: row.user_id, + roleId: row.role_id, + status: row.status, + invitedByMemberId: row.invited_by_member_id, + joinedAt: row.joined_at, + createdAt: row.created_at, + updatedAt: row.updated_at, + organization: { + id: row.organization__id, + name: row.organization__name, + slug: row.organization__slug, + externalId: row.organization__external_id, + identityProvider: row.organization__identity_provider, + isActive: row.organization__is_active, + createdAt: row.organization__created_at, + updatedAt: row.organization__updated_at, + }, + user: { + id: row.user__id, + name: row.user__name, + email: row.user__email, + passwordHash: row.user__password_hash, + externalId: row.user__external_id, + identityProvider: row.user__identity_provider, + isActive: row.user__is_active, + lastLoginAt: row.user__last_login_at, + createdAt: row.user__created_at, + updatedAt: row.user__updated_at, + }, + role: { + id: row.role__id, + name: row.role__name, + description: row.role__description, + createdAt: row.role__created_at, + updatedAt: row.role__updated_at, + }, + })); + } +} + +export class MemberInvitationRepository + extends BaseRepository + implements IMemberInvitationRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'member_invitations' }); + } + + public async findMappings(queryOptions?: IQueryOptions) { + return this.queryMappings({}, queryOptions); + } + + public async findMappingByToken(token: string, queryOptions?: IQueryOptions) { + const [mapping] = await this.queryMappings({ token }, queryOptions); + return mapping || null; + } + + private async queryMappings( + filter: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const query = executer('member_invitations') + .select( + 'member_invitations.*', + 'organizations.id as organization__id', + 'organizations.name as organization__name', + 'organizations.slug as organization__slug', + 'organizations.external_id as organization__external_id', + 'organizations.identity_provider as organization__identity_provider', + 'organizations.is_active as organization__is_active', + 'organizations.created_at as organization__created_at', + 'organizations.updated_at as organization__updated_at', + 'roles.id as role__id', + 'roles.name as role__name', + 'roles.description as role__description', + 'roles.created_at as role__created_at', + 'roles.updated_at as role__updated_at', + ) + .join( + 'organizations', + 'member_invitations.organization_id', + 'organizations.id', + ) + .join('roles', 'member_invitations.role_id', 'roles.id') + .orderBy('member_invitations.created_at', 'desc'); + + if (filter.token) query.where('member_invitations.token', filter.token); + if (filter.organizationId) { + query.where('member_invitations.organization_id', filter.organizationId); + } + if (filter.status) query.where('member_invitations.status', filter.status); + + const rows = await query; + return rows.map((row) => ({ + id: row.id, + organizationId: row.organization_id, + roleId: row.role_id, + email: row.email, + name: row.name, + token: row.token, + status: row.status, + invitedByMemberId: row.invited_by_member_id, + expiresAt: row.expires_at, + acceptedAt: row.accepted_at, + createdAt: row.created_at, + updatedAt: row.updated_at, + organization: { + id: row.organization__id, + name: row.organization__name, + slug: row.organization__slug, + externalId: row.organization__external_id, + identityProvider: row.organization__identity_provider, + isActive: row.organization__is_active, + createdAt: row.organization__created_at, + updatedAt: row.organization__updated_at, + }, + role: { + id: row.role__id, + name: row.role__name, + description: row.role__description, + createdAt: row.role__created_at, + updatedAt: row.role__updated_at, + }, + invitedBy: null, + })); + } +} + +export class AuthSessionRepository + extends BaseRepository + implements IAuthSessionRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'auth_sessions' }); + } + + public async findActiveByToken(token: string, queryOptions?: IQueryOptions) { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const result = await executer('auth_sessions') + .where({ token }) + .whereNull('revoked_at') + .where('expires_at', '>', new Date().toISOString()) + .limit(1); + return result?.[0] ? this.transformFromDBData(result[0]) : null; + } +} diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index e5e6d5a309..7ab2a4ab6d 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -79,6 +79,11 @@ const resolvers = { apiHistory: apiHistoryResolver.getApiHistory, // Administration / RBAC + bootstrapStatus: rbacResolver.bootstrapStatus, + currentSession: rbacResolver.currentSession, + organizations: rbacResolver.listOrganizations, + organizationMembers: rbacResolver.listMembers, + memberInvitations: rbacResolver.listInvitations, roles: rbacResolver.listRoles, users: rbacResolver.listUsers, userRoleMappings: rbacResolver.listUserRoleMappings, @@ -193,6 +198,9 @@ const resolvers = { assignRoleToUser: rbacResolver.assignRoleToUser, updateUserRoles: rbacResolver.updateUserRoles, removeRoleFromUser: rbacResolver.removeRoleFromUser, + inviteMember: rbacResolver.inviteMember, + updateMember: rbacResolver.updateMember, + updateMemberRole: rbacResolver.updateMemberRole, }, ThreadResponse: askingResolver.getThreadResponseNestedResolver(), DetailStep: askingResolver.getDetailStepNestedResolver(), diff --git a/wren-ui/src/apollo/server/resolvers/rbacResolver.ts b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts index 26f4f761ab..5bb9966bdc 100644 --- a/wren-ui/src/apollo/server/resolvers/rbacResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts @@ -2,6 +2,9 @@ import { IContext } from '@server/types'; import { CreateRoleInput, CreateUserInput, + InviteMemberInput, + UpdateMemberInput, + UpdateMemberRoleInput, UpdateRoleInput, UpdateUserInput, UpdateUserRolesInput, @@ -12,8 +15,15 @@ import { Role, UserRole, UserRoleMapping, + Organization, + OrganizationMemberMapping, + MemberInvitationMapping, } from '@server/repositories'; -import { RbacUserWithRoles, RoleWithUsers } from '@server/services'; +import { + AuthSessionResult, + RbacUserWithRoles, + RoleWithUsers, +} from '@server/services'; import { getLogger } from '@server/utils'; const logger = getLogger('RbacResolver'); @@ -24,6 +34,11 @@ export class RbacResolver { this.listRoles = this.listRoles.bind(this); this.listUsers = this.listUsers.bind(this); this.listUserRoleMappings = this.listUserRoleMappings.bind(this); + this.bootstrapStatus = this.bootstrapStatus.bind(this); + this.currentSession = this.currentSession.bind(this); + this.listOrganizations = this.listOrganizations.bind(this); + this.listMembers = this.listMembers.bind(this); + this.listInvitations = this.listInvitations.bind(this); this.createRole = this.createRole.bind(this); this.updateRole = this.updateRole.bind(this); this.createUser = this.createUser.bind(this); @@ -31,6 +46,9 @@ export class RbacResolver { this.assignRoleToUser = this.assignRoleToUser.bind(this); this.updateUserRoles = this.updateUserRoles.bind(this); this.removeRoleFromUser = this.removeRoleFromUser.bind(this); + this.inviteMember = this.inviteMember.bind(this); + this.updateMember = this.updateMember.bind(this); + this.updateMemberRole = this.updateMemberRole.bind(this); } public getRoleNestedResolver() { @@ -91,12 +109,50 @@ export class RbacResolver { return ctx.rbacService.getUserRoleMappings(); } + public async bootstrapStatus(_root: any, _args: any, ctx: IContext) { + return ctx.rbacService.getBootstrapStatus(); + } + + public async currentSession( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + return ctx.currentUser + ? ({ ...ctx.currentUser } as AuthSessionResult) + : null; + } + + public async listOrganizations( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + return ctx.rbacService.listOrganizations(); + } + + public async listMembers( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + return ctx.rbacService.listMembers(ctx.currentUser); + } + + public async listInvitations( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + return ctx.rbacService.listInvitations(ctx.currentUser); + } + public async createRole( _root: any, args: { data: CreateRoleInput }, ctx: IContext, ): Promise { - return ctx.rbacService.createRole(args.data); + return ctx.rbacService.createRole(args.data, ctx.currentUser); } public async updateRole( @@ -104,7 +160,10 @@ export class RbacResolver { args: { where: { id: number }; data: Omit }, ctx: IContext, ): Promise { - return ctx.rbacService.updateRole({ id: args.where.id, ...args.data }); + return ctx.rbacService.updateRole( + { id: args.where.id, ...args.data }, + ctx.currentUser, + ); } public async createUser( @@ -112,6 +171,7 @@ export class RbacResolver { args: { data: CreateUserInput }, ctx: IContext, ): Promise { + this.assertAdmin(ctx); return ctx.rbacService.createUser(args.data); } @@ -120,6 +180,7 @@ export class RbacResolver { args: { where: { id: number }; data: Omit }, ctx: IContext, ): Promise { + this.assertAdmin(ctx); return ctx.rbacService.updateUser({ id: args.where.id, ...args.data }); } @@ -128,6 +189,7 @@ export class RbacResolver { args: { data: UserRoleInput }, ctx: IContext, ): Promise { + this.assertAdmin(ctx); return ctx.rbacService.assignRoleToUser(args.data); } @@ -136,6 +198,7 @@ export class RbacResolver { args: { data: UpdateUserRolesInput }, ctx: IContext, ): Promise { + this.assertAdmin(ctx); return ctx.rbacService.updateUserRoles(args.data); } @@ -144,6 +207,37 @@ export class RbacResolver { args: { data: UserRoleInput }, ctx: IContext, ): Promise { + this.assertAdmin(ctx); return ctx.rbacService.removeRoleFromUser(args.data); } + + public async inviteMember( + _root: any, + args: { data: InviteMemberInput }, + ctx: IContext, + ) { + return ctx.rbacService.inviteMember(args.data, ctx.currentUser); + } + + public async updateMember( + _root: any, + args: { data: UpdateMemberInput }, + ctx: IContext, + ) { + return ctx.rbacService.updateMember(args.data, ctx.currentUser); + } + + public async updateMemberRole( + _root: any, + args: { data: UpdateMemberRoleInput }, + ctx: IContext, + ) { + return ctx.rbacService.updateMemberRole(args.data, ctx.currentUser); + } + + private assertAdmin(ctx: IContext) { + if (ctx.currentUser?.role.name !== 'Admin') { + throw new Error('Admin role is required for this action.'); + } + } } diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index c5a4cbb63e..f63517edb7 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -1126,6 +1126,7 @@ export const typeDefs = gql` externalId: String identityProvider: String isActive: Boolean! + lastLoginAt: String roles: [Role!]! createdAt: String! updatedAt: String! @@ -1149,6 +1150,56 @@ export const typeDefs = gql` updatedAt: String! } + type Organization { + id: Int! + name: String! + slug: String! + externalId: String + identityProvider: String + isActive: Boolean! + createdAt: String! + updatedAt: String! + } + + type OrganizationMember { + id: Int! + organizationId: Int! + userId: Int! + roleId: Int! + status: String! + invitedByMemberId: Int + joinedAt: String + organization: Organization! + user: User! + role: Role! + createdAt: String! + updatedAt: String! + } + + type MemberInvitation { + id: Int! + organizationId: Int! + roleId: Int! + email: String! + name: String + token: String! + status: String! + invitedByMemberId: Int + expiresAt: String! + acceptedAt: String + organization: Organization! + role: Role! + createdAt: String! + updatedAt: String! + } + + type CurrentSession { + user: User! + member: OrganizationMember! + organization: Organization! + role: Role! + } + input RoleWhereInput { id: Int! } @@ -1170,6 +1221,7 @@ export const typeDefs = gql` input CreateUserInput { name: String! email: String! + password: String externalId: String identityProvider: String isActive: Boolean @@ -1194,6 +1246,25 @@ export const typeDefs = gql` roleIds: [Int!]! } + input InviteMemberInput { + organizationId: Int + email: String! + name: String + roleId: Int! + } + + input UpdateMemberInput { + id: Int! + name: String + roleId: Int + status: String + } + + input UpdateMemberRoleInput { + memberId: Int! + roleId: Int! + } + # Query and Mutation type Query { # On Boarding Steps @@ -1253,6 +1324,11 @@ export const typeDefs = gql` ): ApiHistoryPaginatedResponse! # Administration / RBAC + bootstrapStatus: JSON! + currentSession: CurrentSession + organizations: [Organization!]! + organizationMembers: [OrganizationMember!]! + memberInvitations: [MemberInvitation!]! roles: [Role!]! users: [User!]! userRoleMappings: [UserRoleMapping!]! @@ -1410,5 +1486,8 @@ export const typeDefs = gql` assignRoleToUser(data: UserRoleInput!): UserRole! updateUserRoles(data: UpdateUserRolesInput!): User! removeRoleFromUser(data: UserRoleInput!): Boolean! + inviteMember(data: InviteMemberInput!): MemberInvitation! + updateMember(data: UpdateMemberInput!): OrganizationMember! + updateMemberRole(data: UpdateMemberRoleInput!): OrganizationMember! } `; diff --git a/wren-ui/src/apollo/server/services/rbacService.ts b/wren-ui/src/apollo/server/services/rbacService.ts index 76064df4b5..a716c2dfc6 100644 --- a/wren-ui/src/apollo/server/services/rbacService.ts +++ b/wren-ui/src/apollo/server/services/rbacService.ts @@ -1,7 +1,15 @@ import { Knex } from 'knex'; +import bcrypt from 'bcryptjs'; +import { randomBytes } from 'crypto'; import { + AcceptInvitationInput, + BootstrapAdminInput, CreateRoleInput, CreateUserInput, + InviteMemberInput, + LoginInput, + UpdateMemberInput, + UpdateMemberRoleInput, UpdateRoleInput, UpdateUserInput, UpdateUserRolesInput, @@ -11,6 +19,15 @@ import { IRoleRepository, IUserRepository, IUserRoleRepository, + IOrganizationRepository, + IOrganizationMemberRepository, + IMemberInvitationRepository, + IAuthSessionRepository, + AuthSession, + MemberInvitation, + MemberInvitationMapping, + Organization, + OrganizationMemberMapping, RbacUser, Role, UserRole, @@ -26,9 +43,14 @@ export interface RoleWithUsers extends Role { } export interface IRbacService { + getBootstrapStatus(): Promise<{ required: boolean }>; + bootstrapAdmin(input: BootstrapAdminInput): Promise; + login(input: LoginInput): Promise; + logout(token: string): Promise; + getSession(token?: string | null): Promise; listRoles(): Promise; - createRole(input: CreateRoleInput): Promise; - updateRole(input: UpdateRoleInput): Promise; + createRole(input: CreateRoleInput, actor?: AuthActor | null): Promise; + updateRole(input: UpdateRoleInput, actor?: AuthActor | null): Promise; listUsers(): Promise; createUser(input: CreateUserInput): Promise; updateUser(input: UpdateUserInput): Promise; @@ -36,27 +58,227 @@ export interface IRbacService { updateUserRoles(input: UpdateUserRolesInput): Promise; removeRoleFromUser(input: UserRoleInput): Promise; getUserRoleMappings(): Promise; + listOrganizations(): Promise; + listMembers(actor?: AuthActor | null): Promise; + listInvitations(actor?: AuthActor | null): Promise; + inviteMember( + input: InviteMemberInput, + actor?: AuthActor | null, + ): Promise; + acceptInvitation(input: AcceptInvitationInput): Promise; + updateMember( + input: UpdateMemberInput, + actor?: AuthActor | null, + ): Promise; + updateMemberRole( + input: UpdateMemberRoleInput, + actor?: AuthActor | null, + ): Promise; } const EMAIL_PATTERN = /^[^\s@]+@[^\s@]+\.[^\s@]+$/; +const ADMIN_ROLE_NAME = 'Admin'; +const SESSION_TTL_MS = 1000 * 60 * 60 * 24 * 7; +const INVITATION_TTL_MS = 1000 * 60 * 60 * 24 * 7; + +export interface AuthActor { + user: RbacUser; + member: OrganizationMemberMapping; + organization: Organization; + role: Role; +} + +export interface AuthSessionResult extends AuthActor { + session: AuthSession; +} export class RbacService implements IRbacService { private readonly roleRepository: IRoleRepository; private readonly userRepository: IUserRepository; private readonly userRoleRepository: IUserRoleRepository; + private readonly organizationRepository: IOrganizationRepository; + private readonly organizationMemberRepository: IOrganizationMemberRepository; + private readonly memberInvitationRepository: IMemberInvitationRepository; + private readonly authSessionRepository: IAuthSessionRepository; constructor({ roleRepository, userRepository, userRoleRepository, + organizationRepository, + organizationMemberRepository, + memberInvitationRepository, + authSessionRepository, }: { roleRepository: IRoleRepository; userRepository: IUserRepository; userRoleRepository: IUserRoleRepository; + organizationRepository: IOrganizationRepository; + organizationMemberRepository: IOrganizationMemberRepository; + memberInvitationRepository: IMemberInvitationRepository; + authSessionRepository: IAuthSessionRepository; }) { this.roleRepository = roleRepository; this.userRepository = userRepository; this.userRoleRepository = userRoleRepository; + this.organizationRepository = organizationRepository; + this.organizationMemberRepository = organizationMemberRepository; + this.memberInvitationRepository = memberInvitationRepository; + this.authSessionRepository = authSessionRepository; + } + + public async getBootstrapStatus(): Promise<{ required: boolean }> { + const adminRole = await this.roleRepository.findOneBy({ + name: ADMIN_ROLE_NAME, + }); + if (!adminRole) return { required: true }; + + const activeMembers = + await this.organizationMemberRepository.findMappings(); + return { + required: !activeMembers.some( + (member) => + member.status === 'active' && member.role.name === ADMIN_ROLE_NAME, + ), + }; + } + + public async bootstrapAdmin( + input: BootstrapAdminInput, + ): Promise { + const status = await this.getBootstrapStatus(); + if (!status.required) { + throw new Error('An Admin member already exists.'); + } + + const name = this.validateRequiredText(input.name, 'Name'); + const email = this.validateEmail(input.email); + const passwordHash = await this.hashPassword(input.password); + const role = await this.getRoleByNameOrThrow(ADMIN_ROLE_NAME); + const organizationName = this.validateRequiredText( + input.organizationName, + 'Organization name', + ); + const now = new Date().toISOString(); + const tx = await this.userRepository.transaction(); + + try { + const organization = await this.organizationRepository.createOne( + { + name: organizationName, + slug: await this.uniqueOrganizationSlug(organizationName), + isActive: true, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + const user = await this.userRepository.createOne( + { + name, + email, + passwordHash, + identityProvider: 'local', + isActive: true, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + const member = await this.organizationMemberRepository.createOne( + { + organizationId: organization.id, + userId: user.id, + roleId: role.id, + status: 'active', + joinedAt: now, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + const session = await this.createSession(user.id, member.id, tx); + await tx.commit(); + return { + user, + member: { ...member, user, role, organization }, + role, + organization, + session, + }; + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async login(input: LoginInput): Promise { + const email = this.validateEmail(input.email); + const user = await this.userRepository.findOneBy({ email }); + if (!user?.passwordHash) throw new Error('Invalid email or password.'); + const passwordMatches = await bcrypt.compare( + input.password || '', + user.passwordHash, + ); + if (!passwordMatches) throw new Error('Invalid email or password.'); + if (!user.isActive) throw new Error('This user is inactive.'); + + const member = + await this.organizationMemberRepository.findActiveMappingByUserId( + user.id, + ); + if (!member) throw new Error('This user is not an active member.'); + + const now = new Date().toISOString(); + const tx = await this.userRepository.transaction(); + try { + await this.userRepository.updateOne( + user.id, + { lastLoginAt: now, updatedAt: now }, + { tx }, + ); + const session = await this.createSession(user.id, member.id, tx); + await tx.commit(); + return { + user: { ...user, lastLoginAt: now }, + member, + role: member.role, + organization: member.organization, + session, + }; + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async logout(token: string): Promise { + const session = await this.authSessionRepository.findOneBy({ token }); + if (!session) return true; + await this.authSessionRepository.updateOne(session.id, { + revokedAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }); + return true; + } + + public async getSession(token?: string | null) { + if (!token) return null; + const session = await this.authSessionRepository.findActiveByToken(token); + if (!session) return null; + const member = await this.organizationMemberRepository.findMappingById( + session.organizationMemberId, + ); + if (!member || member.status !== 'active' || !member.user.isActive) { + return null; + } + return { + user: member.user, + member, + role: member.role, + organization: member.organization, + session, + }; } public async listRoles(): Promise { @@ -70,7 +292,11 @@ export class RbacService implements IRbacService { })); } - public async createRole(input: CreateRoleInput): Promise { + public async createRole( + input: CreateRoleInput, + actor?: AuthActor | null, + ): Promise { + this.assertAdmin(actor); const name = this.validateRoleName(input.name); await this.assertUniqueRoleName(name); const now = new Date().toISOString(); @@ -82,7 +308,11 @@ export class RbacService implements IRbacService { }); } - public async updateRole(input: UpdateRoleInput): Promise { + public async updateRole( + input: UpdateRoleInput, + actor?: AuthActor | null, + ): Promise { + this.assertAdmin(actor); const role = await this.getRoleOrThrow(input.id); const data: Partial = { updatedAt: new Date().toISOString() }; @@ -113,6 +343,9 @@ export class RbacService implements IRbacService { const name = this.validateRequiredText(input.name, 'User name'); const email = this.validateEmail(input.email); await this.assertUniqueUserEmail(email); + const passwordHash = input.password + ? await this.hashPassword(input.password) + : null; const now = new Date().toISOString(); const tx = await this.userRepository.transaction(); @@ -121,6 +354,7 @@ export class RbacService implements IRbacService { { name, email, + passwordHash, externalId: this.normalizeNullable(input.externalId), identityProvider: this.normalizeNullable(input.identityProvider), isActive: input.isActive ?? true, @@ -216,6 +450,204 @@ export class RbacService implements IRbacService { return this.userRoleRepository.findMappings(); } + public async listOrganizations(): Promise { + return this.organizationRepository.findAll({ order: 'name' }); + } + + public async listMembers(actor?: AuthActor | null) { + this.assertAdmin(actor); + const organization = actor.organization; + return this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ); + } + + public async listInvitations(actor?: AuthActor | null) { + this.assertAdmin(actor); + const organization = actor.organization; + const invitations = await this.memberInvitationRepository.findMappings(); + return invitations.filter( + (invitation) => invitation.organizationId === organization.id, + ); + } + + public async inviteMember( + input: InviteMemberInput, + actor?: AuthActor | null, + ): Promise { + this.assertAdmin(actor); + const email = this.validateEmail(input.email); + const role = await this.getRoleOrThrow(input.roleId); + const organizationId = input.organizationId || actor.member.organizationId; + const organization = await this.organizationRepository.findOneBy({ + id: organizationId, + }); + if (!organization) throw new Error('Organization was not found.'); + + const pending = await this.memberInvitationRepository.findAllBy({ + organizationId, + email, + status: 'pending', + }); + if (pending.length) { + throw new Error(`An invitation for "${email}" is already pending.`); + } + + const existingUser = await this.userRepository.findOneBy({ email }); + if (existingUser) { + const existingMember = + await this.organizationMemberRepository.findActiveMappingByUserId( + existingUser.id, + ); + if (existingMember?.organizationId === organizationId) { + throw new Error(`"${email}" is already a member.`); + } + } + + const now = new Date().toISOString(); + return this.memberInvitationRepository.createOne({ + organizationId, + roleId: role.id, + email, + name: this.normalizeNullable(input.name), + token: this.generateToken(), + status: 'pending', + invitedByMemberId: actor.member.id, + expiresAt: new Date(Date.now() + INVITATION_TTL_MS).toISOString(), + createdAt: now, + updatedAt: now, + }); + } + + public async acceptInvitation( + input: AcceptInvitationInput, + ): Promise { + const invitation = await this.memberInvitationRepository.findMappingByToken( + input.token, + ); + if (!invitation || invitation.status !== 'pending') { + throw new Error('Invitation is invalid or has already been used.'); + } + if (new Date(invitation.expiresAt).getTime() < Date.now()) { + throw new Error('Invitation has expired.'); + } + + const name = this.validateRequiredText( + input.name || invitation.name || invitation.email, + 'Name', + ); + const passwordHash = await this.hashPassword(input.password); + const now = new Date().toISOString(); + const tx = await this.userRepository.transaction(); + + try { + let user = await this.userRepository.findOneBy( + { email: invitation.email }, + { tx }, + ); + if (user) { + user = await this.userRepository.updateOne( + user.id, + { + name, + passwordHash, + identityProvider: user.identityProvider || 'local', + isActive: true, + updatedAt: now, + }, + { tx }, + ); + } else { + user = await this.userRepository.createOne( + { + name, + email: invitation.email, + passwordHash, + identityProvider: 'local', + isActive: true, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + } + + const member = await this.organizationMemberRepository.createOne( + { + organizationId: invitation.organizationId, + userId: user.id, + roleId: invitation.roleId, + status: 'active', + invitedByMemberId: invitation.invitedByMemberId, + joinedAt: now, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + await this.memberInvitationRepository.updateOne( + invitation.id, + { status: 'accepted', acceptedAt: now, updatedAt: now }, + { tx }, + ); + const session = await this.createSession(user.id, member.id, tx); + await tx.commit(); + return { + user, + member: { + ...member, + user, + role: invitation.role, + organization: invitation.organization, + }, + role: invitation.role, + organization: invitation.organization, + session, + }; + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async updateMember( + input: UpdateMemberInput, + actor?: AuthActor | null, + ): Promise { + this.assertAdmin(actor); + const member = await this.getMemberOrThrow(input.id); + if (member.organizationId !== actor.member.organizationId) { + throw new Error('Member is outside of your organization.'); + } + const now = new Date().toISOString(); + if (input.name !== undefined && input.name !== null) { + await this.userRepository.updateOne(member.userId, { + name: this.validateRequiredText(input.name, 'Name'), + updatedAt: now, + }); + } + const data: any = { updatedAt: now }; + if (input.roleId !== undefined && input.roleId !== null) { + await this.getRoleOrThrow(input.roleId); + data.roleId = input.roleId; + } + if (input.status !== undefined && input.status !== null) { + data.status = this.validateMemberStatus(input.status); + } + await this.organizationMemberRepository.updateOne(member.id, data); + return this.getMemberOrThrow(member.id); + } + + public async updateMemberRole( + input: UpdateMemberRoleInput, + actor?: AuthActor | null, + ) { + return this.updateMember( + { id: input.memberId, roleId: input.roleId }, + actor, + ); + } + private async createUserRoleAssignments( userId: number, roleIds: number[], @@ -243,12 +675,51 @@ export class RbacService implements IRbacService { return role; } + private async getRoleByNameOrThrow(name: string): Promise { + const role = await this.roleRepository.findOneBy({ name }); + if (!role) throw new Error(`Role "${name}" was not found.`); + return role; + } + private async getUserOrThrow(id: number): Promise { const user = await this.userRepository.findOneBy({ id }); if (!user) throw new Error(`User ${id} was not found.`); return user; } + private async getMemberOrThrow( + id: number, + ): Promise { + const member = await this.organizationMemberRepository.findMappingById(id); + if (!member) throw new Error(`Member ${id} was not found.`); + return member; + } + + private assertAdmin(actor?: AuthActor | null): asserts actor is AuthActor { + if (!actor || actor.role.name !== ADMIN_ROLE_NAME) { + throw new Error('Admin role is required for this action.'); + } + } + + private async createSession( + userId: number, + organizationMemberId: number, + tx: Knex.Transaction, + ): Promise { + const now = new Date().toISOString(); + return this.authSessionRepository.createOne( + { + userId, + organizationMemberId, + token: this.generateToken(), + expiresAt: new Date(Date.now() + SESSION_TTL_MS).toISOString(), + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + } + private async assertUniqueRoleName(name: string, exceptId?: number) { const roles = await this.roleRepository.findAll(); const duplicate = roles.find( @@ -280,6 +751,22 @@ export class RbacService implements IRbacService { return normalized; } + private validateMemberStatus(status: string): string { + const normalized = this.validateRequiredText(status, 'Status'); + if (!['active', 'inactive', 'suspended'].includes(normalized)) { + throw new Error('Member status must be active, inactive, or suspended.'); + } + return normalized; + } + + private async hashPassword(password: string): Promise { + const normalized = this.validateRequiredText(password, 'Password'); + if (normalized.length < 8) { + throw new Error('Password must be at least 8 characters.'); + } + return bcrypt.hash(normalized, 12); + } + private validateRequiredText(value: string, label: string): string { const normalized = `${value || ''}`.trim(); if (!normalized) throw new Error(`${label} is required.`); @@ -294,4 +781,24 @@ export class RbacService implements IRbacService { private uniqueIds(ids: number[]): number[] { return Array.from(new Set((ids || []).filter(Boolean))); } + + private generateToken(): string { + return randomBytes(32).toString('hex'); + } + + private async uniqueOrganizationSlug(name: string): Promise { + const base = name + .toLowerCase() + .replace(/[^a-z0-9]+/g, '-') + .replace(/^-|-$/g, '') + .slice(0, 80); + const safeBase = base || 'organization'; + let slug = safeBase; + let index = 1; + while (await this.organizationRepository.findOneBy({ slug })) { + slug = `${safeBase}-${index}`; + index += 1; + } + return slug; + } } diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index f92562a6e0..e94338f8ae 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -23,6 +23,10 @@ import { IRoleRepository, IUserRepository, IUserRoleRepository, + IOrganizationRepository, + IOrganizationMemberRepository, + IMemberInvitationRepository, + IAuthSessionRepository, } from '@server/repositories'; import { IQueryService, @@ -34,6 +38,7 @@ import { IDashboardService, IInstructionService, IRbacService, + AuthActor, } from '@server/services'; import { ITelemetry } from '@server/telemetry/telemetry'; import { @@ -64,6 +69,7 @@ export interface IContext { sqlPairService: ISqlPairService; instructionService: IInstructionService; rbacService: IRbacService; + currentUser?: AuthActor | null; // repository projectRepository: IProjectRepository; @@ -84,6 +90,10 @@ export interface IContext { roleRepository: IRoleRepository; userRepository: IUserRepository; userRoleRepository: IUserRoleRepository; + organizationRepository: IOrganizationRepository; + organizationMemberRepository: IOrganizationMemberRepository; + memberInvitationRepository: IMemberInvitationRepository; + authSessionRepository: IAuthSessionRepository; // background trackers projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; diff --git a/wren-ui/src/apollo/server/utils/auth.ts b/wren-ui/src/apollo/server/utils/auth.ts new file mode 100644 index 0000000000..8a447b260b --- /dev/null +++ b/wren-ui/src/apollo/server/utils/auth.ts @@ -0,0 +1,54 @@ +import { NextApiRequest, NextApiResponse } from 'next'; + +export const AUTH_COOKIE_NAME = 'wren_auth_session'; + +export const getCookie = ( + req: Pick, + name: string, +): string | null => { + const cookieHeader = req.headers.cookie; + if (!cookieHeader) return null; + const cookies = cookieHeader.split(';').map((cookie) => cookie.trim()); + const cookie = cookies.find((item) => item.startsWith(`${name}=`)); + if (!cookie) return null; + return decodeURIComponent(cookie.slice(name.length + 1)); +}; + +export const setAuthCookie = ( + res: NextApiResponse, + token: string, + expiresAt: string, +) => { + res.setHeader( + 'Set-Cookie', + `${AUTH_COOKIE_NAME}=${encodeURIComponent( + token, + )}; Path=/; Expires=${new Date( + expiresAt, + ).toUTCString()}; HttpOnly; SameSite=Lax`, + ); +}; + +export const clearAuthCookie = (res: NextApiResponse) => { + res.setHeader( + 'Set-Cookie', + `${AUTH_COOKIE_NAME}=; Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT; HttpOnly; SameSite=Lax`, + ); +}; + +export const sanitizeAuthSession = (session: any) => { + if (!session) return session; + const sanitizeUser = (user: any) => { + if (!user) return user; + const { passwordHash: _passwordHash, ...rest } = user; + return rest; + }; + return { + user: sanitizeUser(session.user), + member: session.member + ? { ...session.member, user: sanitizeUser(session.member.user) } + : session.member, + organization: session.organization, + role: session.role, + }; +}; diff --git a/wren-ui/src/apollo/server/utils/index.ts b/wren-ui/src/apollo/server/utils/index.ts index f9628a9dc3..6256b1a11f 100644 --- a/wren-ui/src/apollo/server/utils/index.ts +++ b/wren-ui/src/apollo/server/utils/index.ts @@ -8,3 +8,4 @@ export * from './helper'; export * from './regex'; export * from './sseTypes'; export * from './sseUtils'; +export * from './auth'; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 18ef4da4d2..528222f1d5 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -22,6 +22,10 @@ import { RoleRepository, UserRepository, UserRoleRepository, + OrganizationRepository, + OrganizationMemberRepository, + MemberInvitationRepository, + AuthSessionRepository, } from '@server/repositories'; import { WrenEngineAdaptor, @@ -157,6 +161,10 @@ export const initComponents = () => { const roleRepository = new RoleRepository(knex); const userRepository = new UserRepository(knex); const userRoleRepository = new UserRoleRepository(knex); + const organizationRepository = new OrganizationRepository(knex); + const organizationMemberRepository = new OrganizationMemberRepository(knex); + const memberInvitationRepository = new MemberInvitationRepository(knex); + const authSessionRepository = new AuthSessionRepository(knex); // adaptors const wrenEngineAdaptor = new WrenEngineAdaptor({ @@ -266,6 +274,10 @@ export const initComponents = () => { roleRepository, userRepository, userRoleRepository, + organizationRepository, + organizationMemberRepository, + memberInvitationRepository, + authSessionRepository, }); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ @@ -303,6 +315,10 @@ export const initComponents = () => { roleRepository, userRepository, userRoleRepository, + organizationRepository, + organizationMemberRepository, + memberInvitationRepository, + authSessionRepository, // adaptors wrenEngineAdaptor, diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index 528dfdab5b..1dc177a302 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -1,9 +1,10 @@ import { useRouter } from 'next/router'; -import { Button, Layout, Space } from 'antd'; +import { Button, Dropdown, Layout, Menu, Space } from 'antd'; import styled from 'styled-components'; import LogoBar from '@/components/LogoBar'; import { Path } from '@/utils/enum'; import Deploy from '@/components/deploy/Deploy'; +import { useAuth } from '@/hooks/useAuth'; const { Header } = Layout; @@ -33,9 +34,17 @@ const StyledHeader = styled(Header)` export default function HeaderBar() { const router = useRouter(); + const auth = useAuth(); const { pathname } = router; const showNav = !pathname.startsWith(Path.Onboarding); const isModeling = pathname.startsWith(Path.Modeling); + const roleName = auth.role?.name; + const isAdmin = roleName === 'Admin'; + const isManager = roleName === 'Manager'; + const isAnalyst = roleName === 'Analyst'; + const canModel = isAdmin || isManager; + const canUseKnowledge = isAdmin || isManager || isAnalyst; + const canUseApi = isAdmin || isManager; return ( @@ -60,6 +69,7 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.Modeling)} onClick={() => router.push(Path.Modeling)} + style={{ display: canModel ? undefined : 'none' }} > Modeling @@ -68,6 +78,7 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.Knowledge)} onClick={() => router.push(Path.KnowledgeQuestionSQLPairs)} + style={{ display: canUseKnowledge ? undefined : 'none' }} > Knowledge @@ -76,6 +87,7 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.APIManagement)} onClick={() => router.push(Path.APIManagementHistory)} + style={{ display: canUseApi ? undefined : 'none' }} > API @@ -84,17 +96,35 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.Administration)} onClick={() => router.push(Path.AdministrationUsers)} + style={{ display: isAdmin ? undefined : 'none' }} > Admin )} - {isModeling && ( - - - - )} + + {isModeling && canModel && } + {auth.authenticated && ( + + + {auth.user?.email} - {roleName} + + auth.logout()}> + Sign out + + + } + trigger={['click']} + > + + + )} + ); diff --git a/wren-ui/src/components/pages/administration/types.tsx b/wren-ui/src/components/pages/administration/types.tsx index 896bf8cf40..548ff33662 100644 --- a/wren-ui/src/components/pages/administration/types.tsx +++ b/wren-ui/src/components/pages/administration/types.tsx @@ -21,6 +21,45 @@ export interface User { updatedAt: string; } +export interface Organization { + id: number; + name: string; + slug: string; + isActive: boolean; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationMember { + id: number; + organizationId: number; + userId: number; + roleId: number; + status: string; + joinedAt?: string | null; + user: User; + role: Role; + organization: Organization; + createdAt: string; + updatedAt: string; +} + +export interface MemberInvitation { + id: number; + organizationId: number; + roleId: number; + email: string; + name?: string | null; + token: string; + status: string; + expiresAt: string; + acceptedAt?: string | null; + role: Role; + organization: Organization; + createdAt: string; + updatedAt: string; +} + export interface UserRoleMapping { id: number; userId: number; diff --git a/wren-ui/src/components/sidebar/Administration.tsx b/wren-ui/src/components/sidebar/Administration.tsx index ecec7dd7ea..c95b034e63 100644 --- a/wren-ui/src/components/sidebar/Administration.tsx +++ b/wren-ui/src/components/sidebar/Administration.tsx @@ -33,7 +33,7 @@ export default function Administration() { { label: ( - User Management + Member Management ), icon: , diff --git a/wren-ui/src/hooks/useAuth.tsx b/wren-ui/src/hooks/useAuth.tsx new file mode 100644 index 0000000000..12bd963d9d --- /dev/null +++ b/wren-ui/src/hooks/useAuth.tsx @@ -0,0 +1,134 @@ +import { + createContext, + ReactNode, + useCallback, + useContext, + useEffect, + useMemo, + useState, +} from 'react'; +import { useRouter } from 'next/router'; +import PageLoading from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; + +type RoleName = 'Admin' | 'Manager' | 'Analyst' | 'Viewer'; + +type AuthState = { + loading: boolean; + authenticated: boolean; + bootstrapRequired: boolean; + user?: any; + member?: any; + organization?: any; + role?: { name: RoleName }; + refresh: () => Promise; + logout: () => Promise; + canAccessPath: (path: string) => boolean; +}; + +const AuthContext = createContext({ + loading: true, + authenticated: false, + bootstrapRequired: false, + refresh: async () => undefined, + logout: async () => undefined, + canAccessPath: () => true, +}); + +const PUBLIC_PATHS = [Path.Login, Path.AcceptInvitation, Path.Onboarding]; + +const ROLE_PATHS: Record = { + Admin: [ + Path.Home, + Path.Modeling, + Path.Knowledge, + Path.APIManagement, + Path.Administration, + ], + Manager: [Path.Home, Path.Modeling, Path.Knowledge, Path.APIManagement], + Analyst: [Path.Home, Path.Knowledge], + Viewer: [Path.Home], +}; + +const isPublicPath = (pathname: string) => + PUBLIC_PATHS.some((path) => pathname.startsWith(path)); + +export const AuthProvider = ({ children }: { children: ReactNode }) => { + const [state, setState] = useState< + Omit + >({ + loading: true, + authenticated: false, + bootstrapRequired: false, + }); + + const refresh = useCallback(async () => { + const [statusResponse, meResponse] = await Promise.all([ + fetch('/api/auth/status'), + fetch('/api/auth/me'), + ]); + const status = await statusResponse.json(); + const me = meResponse.ok ? await meResponse.json() : null; + setState({ + loading: false, + bootstrapRequired: Boolean(status.required), + authenticated: Boolean(me?.authenticated), + user: me?.user, + member: me?.member, + organization: me?.organization, + role: me?.role, + }); + }, []); + + const logout = useCallback(async () => { + await fetch('/api/auth/logout', { method: 'POST' }); + await refresh(); + }, [refresh]); + + useEffect(() => { + void refresh(); + }, [refresh]); + + const canAccessPath = useCallback( + (path: string) => { + if (isPublicPath(path)) return true; + if (!state.authenticated) return false; + const roleName = state.role?.name; + if (!roleName) return false; + return ROLE_PATHS[roleName].some((allowed) => path.startsWith(allowed)); + }, + [state.authenticated, state.role?.name], + ); + + const value = useMemo( + () => ({ ...state, refresh, logout, canAccessPath }), + [state, refresh, logout, canAccessPath], + ); + + return {children}; +}; + +export const AuthGate = ({ children }: { children: ReactNode }) => { + const auth = useAuth(); + const router = useRouter(); + + useEffect(() => { + if (auth.loading) return; + if (isPublicPath(router.pathname)) return; + if (!auth.authenticated) { + void router.replace(Path.Login); + return; + } + if (!auth.canAccessPath(router.pathname)) { + void router.replace(Path.Home); + } + }, [auth, router]); + + if (auth.loading) return ; + if (!isPublicPath(router.pathname) && !auth.authenticated) { + return ; + } + return <>{children}; +}; + +export const useAuth = () => useContext(AuthContext); diff --git a/wren-ui/src/pages/_app.tsx b/wren-ui/src/pages/_app.tsx index 0b3b3765ef..05ab7a5188 100644 --- a/wren-ui/src/pages/_app.tsx +++ b/wren-ui/src/pages/_app.tsx @@ -7,6 +7,7 @@ import { GlobalConfigProvider } from '@/hooks/useGlobalConfig'; import { PostHogProvider } from 'posthog-js/react'; import { ApolloProvider } from '@apollo/client'; import { defaultIndicator } from '@/components/PageLoading'; +import { AuthGate, AuthProvider } from '@/hooks/useAuth'; require('../styles/index.less'); @@ -21,11 +22,15 @@ function App({ Component, pageProps }: AppProps) { - -
- -
- + + + +
+ +
+ + + diff --git a/wren-ui/src/pages/accept-invitation.tsx b/wren-ui/src/pages/accept-invitation.tsx new file mode 100644 index 0000000000..2630bd9fc2 --- /dev/null +++ b/wren-ui/src/pages/accept-invitation.tsx @@ -0,0 +1,92 @@ +import { useState } from 'react'; +import { useRouter } from 'next/router'; +import { Button, Card, Form, Input, Typography, message } from 'antd'; +import styled from 'styled-components'; +import LogoBar from '@/components/LogoBar'; +import { useAuth } from '@/hooks/useAuth'; +import { Path } from '@/utils/enum'; + +const { Paragraph, Title } = Typography; + +const Layout = styled.div` + min-height: 100vh; + display: flex; + align-items: center; + justify-content: center; + background: var(--gray-2); +`; + +const Panel = styled(Card)` + width: 420px; + box-shadow: 0 12px 32px rgba(0, 0, 0, 0.08); +`; + +export default function AcceptInvitationPage() { + const router = useRouter(); + const auth = useAuth(); + const [loading, setLoading] = useState(false); + const token = `${router.query.token || ''}`; + + const submit = async (values: any) => { + setLoading(true); + try { + const response = await fetch('/api/auth/accept-invitation', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ ...values, token }), + }); + const payload = await response.json(); + if (!response.ok) throw new Error(payload.message); + await auth.refresh(); + await router.replace(Path.Home); + } catch (error) { + message.error((error as Error).message); + } finally { + setLoading(false); + } + }; + + return ( + + +
+ +
+ + Accept invitation + + + Create your member profile to access this Wren AI workspace. + +
+ + + + + + + + + + + ); +} diff --git a/wren-ui/src/pages/administration/assignments.tsx b/wren-ui/src/pages/administration/assignments.tsx index 252691c282..421988514d 100644 --- a/wren-ui/src/pages/administration/assignments.tsx +++ b/wren-ui/src/pages/administration/assignments.tsx @@ -4,7 +4,6 @@ import { Button, Form, Modal, - Popconfirm, Select, Table, TableColumnsType, @@ -13,200 +12,128 @@ import { } from 'antd'; import SafetyCertificateOutlined from '@ant-design/icons/SafetyCertificateOutlined'; import EditOutlined from '@ant-design/icons/EditOutlined'; -import DeleteOutlined from '@ant-design/icons/DeleteOutlined'; import SiderLayout from '@/components/layouts/SiderLayout'; import PageLayout from '@/components/layouts/PageLayout'; import { - ASSIGN_ROLE_TO_USER, LIST_USER_ROLE_MAPPINGS, - REMOVE_ROLE_FROM_USER, - UPDATE_USER_ROLES, + UPDATE_MEMBER_ROLE, } from '@/apollo/client/graphql/rbac'; import { + OrganizationMember, Role, RoleTags, - User, - UserRoleMapping, } from '@/components/pages/administration/types'; import { getAbsoluteTime } from '@/utils/time'; const { Text } = Typography; -type AssignmentModalState = { - visible: boolean; - user?: User; -}; - const AssignmentModal = ({ - users, + member, roles, - state, loading, onClose, onSubmit, }: { - users: User[]; + member?: OrganizationMember; roles: Role[]; - state: AssignmentModalState; loading: boolean; onClose: () => void; - onSubmit: (values: any, user?: User) => Promise; + onSubmit: (roleId: number, member: OrganizationMember) => Promise; }) => { const [form] = Form.useForm(); - const isUpdate = !!state.user; useEffect(() => { - if (!state.visible) return; - form.setFieldsValue({ - userId: state.user?.id, - roleId: undefined, - roleIds: state.user?.roles?.map((role) => role.id) || [], - }); - }, [form, state.visible, state.user]); + if (!member) return; + form.setFieldsValue({ roleId: member.roleId }); + }, [form, member]); const submit = async () => { + if (!member) return; const values = await form.validateFields(); - await onSubmit(values, state.user); + await onSubmit(values.roleId, member); form.resetFields(); onClose(); }; return ( form.resetFields()} >
+ +
+
{member?.user.name}
+ {member?.user.email} +
+ ({ - label: role.name, - value: role.id, - }))} - /> - - ) : ( - -
setEditingMember(undefined)} onSubmit={submitAssignment} /> diff --git a/wren-ui/src/pages/api/auth/accept-invitation.ts b/wren-ui/src/pages/api/auth/accept-invitation.ts new file mode 100644 index 0000000000..63d980a49d --- /dev/null +++ b/wren-ui/src/pages/api/auth/accept-invitation.ts @@ -0,0 +1,17 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { sanitizeAuthSession, setAuthCookie } from '@/apollo/server/utils'; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + if (req.method !== 'POST') return res.status(405).end(); + try { + const session = await components.rbacService.acceptInvitation(req.body); + setAuthCookie(res, session.session.token, session.session.expiresAt); + return res.status(200).json(sanitizeAuthSession(session)); + } catch (error) { + return res.status(400).json({ message: (error as Error).message }); + } +} diff --git a/wren-ui/src/pages/api/auth/bootstrap.ts b/wren-ui/src/pages/api/auth/bootstrap.ts new file mode 100644 index 0000000000..7ed11c3e8e --- /dev/null +++ b/wren-ui/src/pages/api/auth/bootstrap.ts @@ -0,0 +1,17 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { sanitizeAuthSession, setAuthCookie } from '@/apollo/server/utils'; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + if (req.method !== 'POST') return res.status(405).end(); + try { + const session = await components.rbacService.bootstrapAdmin(req.body); + setAuthCookie(res, session.session.token, session.session.expiresAt); + return res.status(200).json(sanitizeAuthSession(session)); + } catch (error) { + return res.status(400).json({ message: (error as Error).message }); + } +} diff --git a/wren-ui/src/pages/api/auth/login.ts b/wren-ui/src/pages/api/auth/login.ts new file mode 100644 index 0000000000..7fe64fbd2f --- /dev/null +++ b/wren-ui/src/pages/api/auth/login.ts @@ -0,0 +1,17 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { sanitizeAuthSession, setAuthCookie } from '@/apollo/server/utils'; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + if (req.method !== 'POST') return res.status(405).end(); + try { + const session = await components.rbacService.login(req.body); + setAuthCookie(res, session.session.token, session.session.expiresAt); + return res.status(200).json(sanitizeAuthSession(session)); + } catch (error) { + return res.status(401).json({ message: (error as Error).message }); + } +} diff --git a/wren-ui/src/pages/api/auth/logout.ts b/wren-ui/src/pages/api/auth/logout.ts new file mode 100644 index 0000000000..c5d6bbe328 --- /dev/null +++ b/wren-ui/src/pages/api/auth/logout.ts @@ -0,0 +1,18 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { + AUTH_COOKIE_NAME, + clearAuthCookie, + getCookie, +} from '@/apollo/server/utils'; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + if (req.method !== 'POST') return res.status(405).end(); + const token = getCookie(req, AUTH_COOKIE_NAME); + if (token) await components.rbacService.logout(token); + clearAuthCookie(res); + return res.status(200).json({ ok: true }); +} diff --git a/wren-ui/src/pages/api/auth/me.ts b/wren-ui/src/pages/api/auth/me.ts new file mode 100644 index 0000000000..cf304e037d --- /dev/null +++ b/wren-ui/src/pages/api/auth/me.ts @@ -0,0 +1,21 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { + AUTH_COOKIE_NAME, + getCookie, + sanitizeAuthSession, +} from '@/apollo/server/utils'; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + if (req.method !== 'GET') return res.status(405).end(); + const token = getCookie(req, AUTH_COOKIE_NAME); + const session = await components.rbacService.getSession(token); + if (!session) return res.status(401).json({ authenticated: false }); + return res.status(200).json({ + authenticated: true, + ...sanitizeAuthSession(session), + }); +} diff --git a/wren-ui/src/pages/api/auth/status.ts b/wren-ui/src/pages/api/auth/status.ts new file mode 100644 index 0000000000..344c459df1 --- /dev/null +++ b/wren-ui/src/pages/api/auth/status.ts @@ -0,0 +1,11 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + if (req.method !== 'GET') return res.status(405).end(); + const status = await components.rbacService.getBootstrapStatus(); + return res.status(200).json(status); +} diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index b31de21a75..ddaf00bd98 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -5,7 +5,7 @@ import { typeDefs } from '@server'; import resolvers from '@server/resolvers'; import { IContext } from '@server/types'; import { GraphQLError } from 'graphql'; -import { getLogger } from '@server/utils'; +import { AUTH_COOKIE_NAME, getCookie, getLogger } from '@server/utils'; import { getConfig } from '@server/config'; import { ModelService } from '@server/services/modelService'; import { @@ -50,6 +50,10 @@ const bootstrapServer = async () => { roleRepository, userRepository, userRoleRepository, + organizationRepository, + organizationMemberRepository, + memberInvitationRepository, + authSessionRepository, // adaptors wrenEngineAdaptor, ibisAdaptor, @@ -129,48 +133,57 @@ const bootstrapServer = async () => { return defaultApolloErrorHandler(error); }, introspection: process.env.NODE_ENV !== 'production', - context: (): IContext => ({ - config: serverConfig, - telemetry, - // adaptor - wrenEngineAdaptor, - ibisServerAdaptor: ibisAdaptor, - wrenAIAdaptor, - // services - projectService, - modelService, - mdlService, - deployService, - askingService, - queryService, - dashboardService, - sqlPairService, - instructionService, - rbacService, - // repository - projectRepository, - modelRepository, - modelColumnRepository, - modelNestedColumnRepository, - relationRepository, - viewRepository, - deployRepository: deployLogRepository, - schemaChangeRepository, - learningRepository, - dashboardRepository, - dashboardItemRepository, - sqlPairRepository, - instructionRepository, - apiHistoryRepository, - dashboardItemRefreshJobRepository, - roleRepository, - userRepository, - userRoleRepository, - // background trackers - projectRecommendQuestionBackgroundTracker, - threadRecommendQuestionBackgroundTracker, - dashboardCacheBackgroundTracker, - }), + context: async ({ req }): Promise => { + const token = getCookie(req, AUTH_COOKIE_NAME); + const currentUser = await rbacService.getSession(token); + return { + config: serverConfig, + telemetry, + // adaptor + wrenEngineAdaptor, + ibisServerAdaptor: ibisAdaptor, + wrenAIAdaptor, + // services + projectService, + modelService, + mdlService, + deployService, + askingService, + queryService, + dashboardService, + sqlPairService, + instructionService, + rbacService, + currentUser, + // repository + projectRepository, + modelRepository, + modelColumnRepository, + modelNestedColumnRepository, + relationRepository, + viewRepository, + deployRepository: deployLogRepository, + schemaChangeRepository, + learningRepository, + dashboardRepository, + dashboardItemRepository, + sqlPairRepository, + instructionRepository, + apiHistoryRepository, + dashboardItemRefreshJobRepository, + roleRepository, + userRepository, + userRoleRepository, + organizationRepository, + organizationMemberRepository, + memberInvitationRepository, + authSessionRepository, + // background trackers + projectRecommendQuestionBackgroundTracker, + threadRecommendQuestionBackgroundTracker, + dashboardCacheBackgroundTracker, + }; + }, }); await apolloServer.start(); return apolloServer; diff --git a/wren-ui/src/pages/login.tsx b/wren-ui/src/pages/login.tsx new file mode 100644 index 0000000000..2e8c47e4b2 --- /dev/null +++ b/wren-ui/src/pages/login.tsx @@ -0,0 +1,112 @@ +import { useState } from 'react'; +import { useRouter } from 'next/router'; +import { Button, Card, Form, Input, Typography, message } from 'antd'; +import styled from 'styled-components'; +import LogoBar from '@/components/LogoBar'; +import { useAuth } from '@/hooks/useAuth'; +import { Path } from '@/utils/enum'; + +const { Paragraph, Title } = Typography; + +const Layout = styled.div` + min-height: 100vh; + display: flex; + align-items: center; + justify-content: center; + background: var(--gray-2); +`; + +const Panel = styled(Card)` + width: 420px; + box-shadow: 0 12px 32px rgba(0, 0, 0, 0.08); +`; + +export default function LoginPage() { + const auth = useAuth(); + const router = useRouter(); + const [loading, setLoading] = useState(false); + const [form] = Form.useForm(); + + const submit = async (values: any) => { + setLoading(true); + const endpoint = auth.bootstrapRequired + ? '/api/auth/bootstrap' + : '/api/auth/login'; + try { + const response = await fetch(endpoint, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(values), + }); + const payload = await response.json(); + if (!response.ok) throw new Error(payload.message); + await auth.refresh(); + await router.replace(Path.Home); + } catch (error) { + message.error((error as Error).message); + } finally { + setLoading(false); + } + }; + + return ( + + +
+ +
+ + {auth.bootstrapRequired ? 'Create your Admin account' : 'Sign in'} + + + {auth.bootstrapRequired + ? 'Set up the first organization and Admin member.' + : 'Access your Wren AI workspace.'} + + + {auth.bootstrapRequired && ( + + + + )} + {auth.bootstrapRequired && ( + + + + )} + + + + + + + + + + + ); +} diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index 54c0eb91f0..588f71cc8c 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -1,5 +1,7 @@ export enum Path { Home = '/home', + Login = '/login', + AcceptInvitation = '/accept-invitation', HomeDashboard = '/home/dashboard', Thread = '/home/[id]', Modeling = '/modeling', From 2679fd3e357183bba171623260bdb187fd4d4330 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 3 Jun 2026 14:37:35 +0530 Subject: [PATCH 0086/1087] revert changes --- wren-ui/docs/rbac-architecture.md | 69 +-- ...0_create_organization_membership_tables.js | 174 ------ wren-ui/src/apollo/client/graphql/rbac.ts | 142 +---- wren-ui/src/apollo/server/models/rbac.ts | 38 -- .../server/repositories/rbacRepository.ts | 352 ------------ wren-ui/src/apollo/server/resolvers.ts | 8 - .../apollo/server/resolvers/rbacResolver.ts | 100 +--- wren-ui/src/apollo/server/schema.ts | 79 --- .../src/apollo/server/services/rbacService.ts | 515 +----------------- wren-ui/src/apollo/server/types/context.ts | 10 - wren-ui/src/apollo/server/utils/auth.ts | 54 -- wren-ui/src/apollo/server/utils/index.ts | 1 - wren-ui/src/common.ts | 16 - wren-ui/src/components/HeaderBar.tsx | 42 +- .../components/pages/administration/types.tsx | 39 -- .../src/components/sidebar/Administration.tsx | 2 +- wren-ui/src/hooks/useAuth.tsx | 134 ----- wren-ui/src/pages/_app.tsx | 15 +- wren-ui/src/pages/accept-invitation.tsx | 92 ---- .../src/pages/administration/assignments.tsx | 210 +++++-- .../src/pages/api/auth/accept-invitation.ts | 17 - wren-ui/src/pages/api/auth/bootstrap.ts | 17 - wren-ui/src/pages/api/auth/login.ts | 17 - wren-ui/src/pages/api/auth/logout.ts | 18 - wren-ui/src/pages/api/auth/me.ts | 21 - wren-ui/src/pages/api/auth/status.ts | 11 - wren-ui/src/pages/api/graphql.ts | 19 +- wren-ui/src/pages/login.tsx | 112 ---- wren-ui/src/utils/enum/path.ts | 2 - 29 files changed, 210 insertions(+), 2116 deletions(-) delete mode 100644 wren-ui/migrations/20250603000000_create_organization_membership_tables.js delete mode 100644 wren-ui/src/apollo/server/utils/auth.ts delete mode 100644 wren-ui/src/hooks/useAuth.tsx delete mode 100644 wren-ui/src/pages/accept-invitation.tsx delete mode 100644 wren-ui/src/pages/api/auth/accept-invitation.ts delete mode 100644 wren-ui/src/pages/api/auth/bootstrap.ts delete mode 100644 wren-ui/src/pages/api/auth/login.ts delete mode 100644 wren-ui/src/pages/api/auth/logout.ts delete mode 100644 wren-ui/src/pages/api/auth/me.ts delete mode 100644 wren-ui/src/pages/api/auth/status.ts delete mode 100644 wren-ui/src/pages/login.tsx diff --git a/wren-ui/docs/rbac-architecture.md b/wren-ui/docs/rbac-architecture.md index caa594f115..9f824c252b 100644 --- a/wren-ui/docs/rbac-architecture.md +++ b/wren-ui/docs/rbac-architecture.md @@ -1,18 +1,16 @@ # RBAC Foundation -This document describes the application-level organization and member RBAC foundation in Wren UI. It establishes durable organization, user, member, invitation, session, and role primitives for future governance work. +This document describes the application-level RBAC foundation in Wren UI. It does not enforce permissions yet; it establishes durable role, user, and user-role assignment primitives for future governance work. ## Scope Implemented: - Roles: `Admin`, `Manager`, `Analyst`, `Viewer`, plus custom roles. -- Organizations with active members. -- Users with local identity metadata and password hashes for local login. -- Organization-member role assignments. -- Member invitations and local auth sessions. -- GraphQL APIs for role, member, invitation, and assignment management. -- Administration UI for member management, role management, and role assignment. +- Users with local identity metadata. +- User-role assignments. +- GraphQL APIs for role, user, and assignment management. +- Administration UI for user management, role management, and role assignment. Deferred: @@ -45,39 +43,6 @@ Tables: - `user_id` - `role_id` - timestamps -- `organizations` - - `id` - - `name` - - `slug` - - external identity fields - - `is_active` - - timestamps -- `organization_members` - - `id` - - `organization_id` - - `user_id` - - `role_id` - - `status` - - `joined_at` - - timestamps -- `member_invitations` - - `id` - - `organization_id` - - `role_id` - - `email` - - `token` - - `status` - - `expires_at` - - timestamps -- `auth_sessions` - - `id` - - `user_id` - - `organization_member_id` - - `token` - - `expires_at` - - `revoked_at` - - timestamps - `external_id` and `identity_provider` are intentionally present now so Teams, LDAP, and Azure AD integrations can later attach external identities without replacing the RBAC tables. ## Backend Layers @@ -98,11 +63,6 @@ Queries: - `roles` - `users` - `userRoleMappings` -- `organizations` -- `organizationMembers` -- `memberInvitations` -- `currentSession` -- `bootstrapStatus` Mutations: @@ -113,16 +73,6 @@ Mutations: - `assignRoleToUser` - `updateUserRoles` - `removeRoleFromUser` -- `inviteMember` -- `updateMember` -- `updateMemberRole` - -Admin-only behavior: - -- Inviting members. -- Creating and editing roles. -- Editing members and updating member roles. -- Legacy user-role mutation paths are also guarded for Admin members. ## UI @@ -133,19 +83,12 @@ Navigation: Screens: -- `/administration/users` (Member Management) +- `/administration/users` - `/administration/roles` - `/administration/assignments` The UI uses the existing Next.js, Apollo Client, Ant Design, and `SiderLayout`/`PageLayout` patterns. -Authentication routes: - -- `/login` -- `/accept-invitation?token=...` - -The first Admin can bootstrap the first organization when no active Admin member exists. - ## Future Permission Model Future schema-level or table-level permissions should be added as separate tables referencing `roles.id`, for example: diff --git a/wren-ui/migrations/20250603000000_create_organization_membership_tables.js b/wren-ui/migrations/20250603000000_create_organization_membership_tables.js deleted file mode 100644 index ae917bd847..0000000000 --- a/wren-ui/migrations/20250603000000_create_organization_membership_tables.js +++ /dev/null @@ -1,174 +0,0 @@ -/** - * @param { import("knex").Knex } knex - * @returns { Promise } - */ -exports.up = async function (knex) { - const hasUsers = await knex.schema.hasTable('users'); - const hasRoles = await knex.schema.hasTable('roles'); - - if (hasUsers) { - const hasPasswordHash = await knex.schema.hasColumn( - 'users', - 'password_hash', - ); - const hasLastLoginAt = await knex.schema.hasColumn( - 'users', - 'last_login_at', - ); - await knex.schema.alterTable('users', (table) => { - if (!hasPasswordHash) { - table.string('password_hash', 255).nullable(); - } - if (!hasLastLoginAt) { - table.timestamp('last_login_at').nullable(); - } - }); - } - - const hasOrganizations = await knex.schema.hasTable('organizations'); - if (!hasOrganizations) { - await knex.schema.createTable('organizations', (table) => { - table.increments('id').primary(); - table.string('name', 160).notNullable().unique(); - table.string('slug', 180).notNullable().unique(); - table.string('external_id', 255).nullable().unique(); - table.string('identity_provider', 80).nullable(); - table.boolean('is_active').notNullable().defaultTo(true); - table.timestamps(true, true); - }); - } - - const hasOrganizationMembers = await knex.schema.hasTable( - 'organization_members', - ); - if (!hasOrganizationMembers) { - await knex.schema.createTable('organization_members', (table) => { - table.increments('id').primary(); - table.integer('organization_id').notNullable(); - table.integer('user_id').notNullable(); - table.integer('role_id').notNullable(); - table.string('status', 40).notNullable().defaultTo('active'); - table.integer('invited_by_member_id').nullable(); - table.timestamp('joined_at').nullable(); - table.timestamps(true, true); - - table - .foreign('organization_id') - .references('organizations.id') - .onDelete('CASCADE'); - table.foreign('user_id').references('users.id').onDelete('CASCADE'); - table.foreign('role_id').references('roles.id').onDelete('RESTRICT'); - table.unique(['organization_id', 'user_id']); - }); - } - - const hasMemberInvitations = await knex.schema.hasTable('member_invitations'); - if (!hasMemberInvitations) { - await knex.schema.createTable('member_invitations', (table) => { - table.increments('id').primary(); - table.integer('organization_id').notNullable(); - table.integer('role_id').notNullable(); - table.string('email', 320).notNullable(); - table.string('name', 160).nullable(); - table.string('token', 128).notNullable().unique(); - table.string('status', 40).notNullable().defaultTo('pending'); - table.integer('invited_by_member_id').nullable(); - table.timestamp('expires_at').notNullable(); - table.timestamp('accepted_at').nullable(); - table.timestamps(true, true); - - table - .foreign('organization_id') - .references('organizations.id') - .onDelete('CASCADE'); - table.foreign('role_id').references('roles.id').onDelete('RESTRICT'); - table.unique(['organization_id', 'email', 'status']); - }); - } - - const hasAuthSessions = await knex.schema.hasTable('auth_sessions'); - if (!hasAuthSessions) { - await knex.schema.createTable('auth_sessions', (table) => { - table.increments('id').primary(); - table.integer('user_id').notNullable(); - table.integer('organization_member_id').notNullable(); - table.string('token', 128).notNullable().unique(); - table.timestamp('expires_at').notNullable(); - table.timestamp('revoked_at').nullable(); - table.timestamps(true, true); - - table.foreign('user_id').references('users.id').onDelete('CASCADE'); - table - .foreign('organization_member_id') - .references('organization_members.id') - .onDelete('CASCADE'); - }); - } - - if (hasOrganizations || !hasUsers || !hasRoles) return; - - const now = new Date().toISOString(); - const [organization] = await knex('organizations') - .insert({ - name: 'Default organization', - slug: 'default', - is_active: true, - created_at: now, - updated_at: now, - }) - .returning('*'); - - const adminRole = await knex('roles').where({ name: 'Admin' }).first(); - const roles = await knex('roles'); - const users = await knex('users'); - const userRoles = await knex('user_roles'); - - const roleById = new Map(roles.map((role) => [role.id, role])); - const firstRoleByUserId = new Map(); - userRoles.forEach((mapping) => { - if (!firstRoleByUserId.has(mapping.user_id)) { - firstRoleByUserId.set(mapping.user_id, mapping.role_id); - } - }); - - for (const user of users) { - const roleId = firstRoleByUserId.get(user.id) || adminRole?.id; - if (!roleById.has(roleId)) continue; - await knex('organization_members').insert({ - organization_id: organization.id, - user_id: user.id, - role_id: roleId, - status: 'active', - joined_at: now, - created_at: now, - updated_at: now, - }); - } -}; - -/** - * @param { import("knex").Knex } knex - * @returns { Promise } - */ -exports.down = async function (knex) { - await knex.schema.dropTableIfExists('auth_sessions'); - await knex.schema.dropTableIfExists('member_invitations'); - await knex.schema.dropTableIfExists('organization_members'); - await knex.schema.dropTableIfExists('organizations'); - - const hasUsers = await knex.schema.hasTable('users'); - if (hasUsers) { - const hasPasswordHash = await knex.schema.hasColumn( - 'users', - 'password_hash', - ); - const hasLastLoginAt = await knex.schema.hasColumn( - 'users', - 'last_login_at', - ); - await knex.schema.alterTable('users', (table) => { - if (hasPasswordHash) table.dropColumn('password_hash'); - if (hasLastLoginAt) table.dropColumn('last_login_at'); - }); - } -}; diff --git a/wren-ui/src/apollo/client/graphql/rbac.ts b/wren-ui/src/apollo/client/graphql/rbac.ts index 9816fa8abd..6da6060285 100644 --- a/wren-ui/src/apollo/client/graphql/rbac.ts +++ b/wren-ui/src/apollo/client/graphql/rbac.ts @@ -23,68 +23,13 @@ export const USER_FIELDS = gql` } `; -export const ORGANIZATION_FIELDS = gql` - fragment OrganizationFields on Organization { - id - name - slug - isActive - createdAt - updatedAt - } -`; - -export const MEMBER_FIELDS = gql` - fragment MemberFields on OrganizationMember { - id - organizationId - userId - roleId - status - joinedAt - createdAt - updatedAt - user { - ...UserFields - } - role { - ...RoleFields - } - organization { - ...OrganizationFields - } - } -`; - -export const INVITATION_FIELDS = gql` - fragment InvitationFields on MemberInvitation { - id - organizationId - roleId - email - name - token - status - expiresAt - acceptedAt - createdAt - updatedAt - role { - ...RoleFields - } - organization { - ...OrganizationFields - } - } -`; - export const LIST_RBAC_USERS = gql` query RbacUsers { - organizationMembers { - ...MemberFields - } - memberInvitations { - ...InvitationFields + users { + ...UserFields + roles { + ...RoleFields + } } roles { ...RoleFields @@ -93,34 +38,48 @@ export const LIST_RBAC_USERS = gql` ${USER_FIELDS} ${ROLE_FIELDS} - ${ORGANIZATION_FIELDS} - ${MEMBER_FIELDS} - ${INVITATION_FIELDS} `; export const LIST_RBAC_ROLES = gql` query RbacRoles { roles { ...RoleFields + users { + ...UserFields + } } - organizationMembers { - ...MemberFields + users { + ...UserFields } } ${ROLE_FIELDS} ${USER_FIELDS} - ${ORGANIZATION_FIELDS} - ${MEMBER_FIELDS} `; export const LIST_USER_ROLE_MAPPINGS = gql` query UserRoleMappings { - organizationMembers { - ...MemberFields + userRoleMappings { + id + userId + roleId + createdAt + updatedAt + user { + ...UserFields + roles { + ...RoleFields + } + } + role { + ...RoleFields + } } - memberInvitations { - ...InvitationFields + users { + ...UserFields + roles { + ...RoleFields + } } roles { ...RoleFields @@ -129,9 +88,6 @@ export const LIST_USER_ROLE_MAPPINGS = gql` ${USER_FIELDS} ${ROLE_FIELDS} - ${ORGANIZATION_FIELDS} - ${MEMBER_FIELDS} - ${INVITATION_FIELDS} `; export const CREATE_ROLE = gql` @@ -211,41 +167,3 @@ export const REMOVE_ROLE_FROM_USER = gql` removeRoleFromUser(data: $data) } `; - -export const INVITE_MEMBER = gql` - mutation InviteMember($data: InviteMemberInput!) { - inviteMember(data: $data) { - ...InvitationFields - } - } - - ${ROLE_FIELDS} - ${ORGANIZATION_FIELDS} - ${INVITATION_FIELDS} -`; - -export const UPDATE_MEMBER = gql` - mutation UpdateMember($data: UpdateMemberInput!) { - updateMember(data: $data) { - ...MemberFields - } - } - - ${USER_FIELDS} - ${ROLE_FIELDS} - ${ORGANIZATION_FIELDS} - ${MEMBER_FIELDS} -`; - -export const UPDATE_MEMBER_ROLE = gql` - mutation UpdateMemberRole($data: UpdateMemberRoleInput!) { - updateMemberRole(data: $data) { - ...MemberFields - } - } - - ${USER_FIELDS} - ${ROLE_FIELDS} - ${ORGANIZATION_FIELDS} - ${MEMBER_FIELDS} -`; diff --git a/wren-ui/src/apollo/server/models/rbac.ts b/wren-ui/src/apollo/server/models/rbac.ts index e35defbc2f..52b1a98ce9 100644 --- a/wren-ui/src/apollo/server/models/rbac.ts +++ b/wren-ui/src/apollo/server/models/rbac.ts @@ -12,7 +12,6 @@ export interface UpdateRoleInput { export interface CreateUserInput { name: string; email: string; - password?: string | null; externalId?: string | null; identityProvider?: string | null; isActive?: boolean; @@ -37,40 +36,3 @@ export interface UpdateUserRolesInput { userId: number; roleIds: number[]; } - -export interface LoginInput { - email: string; - password: string; -} - -export interface BootstrapAdminInput { - organizationName: string; - name: string; - email: string; - password: string; -} - -export interface InviteMemberInput { - organizationId?: number | null; - email: string; - name?: string | null; - roleId: number; -} - -export interface AcceptInvitationInput { - token: string; - name?: string | null; - password: string; -} - -export interface UpdateMemberInput { - id: number; - name?: string | null; - roleId?: number | null; - status?: string | null; -} - -export interface UpdateMemberRoleInput { - memberId: number; - roleId: number; -} diff --git a/wren-ui/src/apollo/server/repositories/rbacRepository.ts b/wren-ui/src/apollo/server/repositories/rbacRepository.ts index c3c3bd6e3f..d742af6579 100644 --- a/wren-ui/src/apollo/server/repositories/rbacRepository.ts +++ b/wren-ui/src/apollo/server/repositories/rbacRepository.ts @@ -17,11 +17,9 @@ export interface RbacUser { id: number; name: string; email: string; - passwordHash?: string | null; externalId?: string | null; identityProvider?: string | null; isActive: boolean; - lastLoginAt?: string | null; createdAt: string; updatedAt: string; } @@ -39,67 +37,6 @@ export interface UserRoleMapping extends UserRole { role: Role; } -export interface Organization { - id: number; - name: string; - slug: string; - externalId?: string | null; - identityProvider?: string | null; - isActive: boolean; - createdAt: string; - updatedAt: string; -} - -export interface OrganizationMember { - id: number; - organizationId: number; - userId: number; - roleId: number; - status: string; - invitedByMemberId?: number | null; - joinedAt?: string | null; - createdAt: string; - updatedAt: string; -} - -export interface OrganizationMemberMapping extends OrganizationMember { - organization: Organization; - user: RbacUser; - role: Role; -} - -export interface MemberInvitation { - id: number; - organizationId: number; - roleId: number; - email: string; - name?: string | null; - token: string; - status: string; - invitedByMemberId?: number | null; - expiresAt: string; - acceptedAt?: string | null; - createdAt: string; - updatedAt: string; -} - -export interface MemberInvitationMapping extends MemberInvitation { - organization: Organization; - role: Role; - invitedBy?: OrganizationMemberMapping | null; -} - -export interface AuthSession { - id: number; - userId: number; - organizationMemberId: number; - token: string; - expiresAt: string; - revokedAt?: string | null; - createdAt: string; - updatedAt: string; -} - export interface IRoleRepository extends IBasicRepository {} export interface IUserRepository extends IBasicRepository {} @@ -116,46 +53,6 @@ export interface IUserRoleRepository extends IBasicRepository { ): Promise; } -export interface IOrganizationRepository - extends IBasicRepository {} - -export interface IOrganizationMemberRepository - extends IBasicRepository { - findMappings( - queryOptions?: IQueryOptions, - ): Promise; - findMappingsByOrganizationId( - organizationId: number, - queryOptions?: IQueryOptions, - ): Promise; - findMappingById( - id: number, - queryOptions?: IQueryOptions, - ): Promise; - findActiveMappingByUserId( - userId: number, - queryOptions?: IQueryOptions, - ): Promise; -} - -export interface IMemberInvitationRepository - extends IBasicRepository { - findMappings( - queryOptions?: IQueryOptions, - ): Promise; - findMappingByToken( - token: string, - queryOptions?: IQueryOptions, - ): Promise; -} - -export interface IAuthSessionRepository extends IBasicRepository { - findActiveByToken( - token: string, - queryOptions?: IQueryOptions, - ): Promise; -} - export class RoleRepository extends BaseRepository implements IRoleRepository @@ -261,252 +158,3 @@ export class UserRoleRepository })); } } - -export class OrganizationRepository - extends BaseRepository - implements IOrganizationRepository -{ - constructor(knexPg: Knex) { - super({ knexPg, tableName: 'organizations' }); - } -} - -export class OrganizationMemberRepository - extends BaseRepository - implements IOrganizationMemberRepository -{ - constructor(knexPg: Knex) { - super({ knexPg, tableName: 'organization_members' }); - } - - public async findMappings(queryOptions?: IQueryOptions) { - return this.queryMappings({}, queryOptions); - } - - public async findMappingsByOrganizationId( - organizationId: number, - queryOptions?: IQueryOptions, - ) { - return this.queryMappings({ organizationId }, queryOptions); - } - - public async findMappingById(id: number, queryOptions?: IQueryOptions) { - const [mapping] = await this.queryMappings({ id }, queryOptions); - return mapping || null; - } - - public async findActiveMappingByUserId( - userId: number, - queryOptions?: IQueryOptions, - ) { - const [mapping] = await this.queryMappings( - { userId, status: 'active' }, - queryOptions, - ); - return mapping || null; - } - - private async queryMappings( - filter: Partial, - queryOptions?: IQueryOptions, - ): Promise { - const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const query = executer('organization_members') - .select( - 'organization_members.*', - 'organizations.id as organization__id', - 'organizations.name as organization__name', - 'organizations.slug as organization__slug', - 'organizations.external_id as organization__external_id', - 'organizations.identity_provider as organization__identity_provider', - 'organizations.is_active as organization__is_active', - 'organizations.created_at as organization__created_at', - 'organizations.updated_at as organization__updated_at', - 'users.id as user__id', - 'users.name as user__name', - 'users.email as user__email', - 'users.password_hash as user__password_hash', - 'users.external_id as user__external_id', - 'users.identity_provider as user__identity_provider', - 'users.is_active as user__is_active', - 'users.last_login_at as user__last_login_at', - 'users.created_at as user__created_at', - 'users.updated_at as user__updated_at', - 'roles.id as role__id', - 'roles.name as role__name', - 'roles.description as role__description', - 'roles.created_at as role__created_at', - 'roles.updated_at as role__updated_at', - ) - .join( - 'organizations', - 'organization_members.organization_id', - 'organizations.id', - ) - .join('users', 'organization_members.user_id', 'users.id') - .join('roles', 'organization_members.role_id', 'roles.id') - .orderBy('users.email'); - - if (filter.id) query.where('organization_members.id', filter.id); - if (filter.organizationId) { - query.where( - 'organization_members.organization_id', - filter.organizationId, - ); - } - if (filter.userId) - query.where('organization_members.user_id', filter.userId); - if (filter.status) - query.where('organization_members.status', filter.status); - - const rows = await query; - return rows.map((row) => ({ - id: row.id, - organizationId: row.organization_id, - userId: row.user_id, - roleId: row.role_id, - status: row.status, - invitedByMemberId: row.invited_by_member_id, - joinedAt: row.joined_at, - createdAt: row.created_at, - updatedAt: row.updated_at, - organization: { - id: row.organization__id, - name: row.organization__name, - slug: row.organization__slug, - externalId: row.organization__external_id, - identityProvider: row.organization__identity_provider, - isActive: row.organization__is_active, - createdAt: row.organization__created_at, - updatedAt: row.organization__updated_at, - }, - user: { - id: row.user__id, - name: row.user__name, - email: row.user__email, - passwordHash: row.user__password_hash, - externalId: row.user__external_id, - identityProvider: row.user__identity_provider, - isActive: row.user__is_active, - lastLoginAt: row.user__last_login_at, - createdAt: row.user__created_at, - updatedAt: row.user__updated_at, - }, - role: { - id: row.role__id, - name: row.role__name, - description: row.role__description, - createdAt: row.role__created_at, - updatedAt: row.role__updated_at, - }, - })); - } -} - -export class MemberInvitationRepository - extends BaseRepository - implements IMemberInvitationRepository -{ - constructor(knexPg: Knex) { - super({ knexPg, tableName: 'member_invitations' }); - } - - public async findMappings(queryOptions?: IQueryOptions) { - return this.queryMappings({}, queryOptions); - } - - public async findMappingByToken(token: string, queryOptions?: IQueryOptions) { - const [mapping] = await this.queryMappings({ token }, queryOptions); - return mapping || null; - } - - private async queryMappings( - filter: Partial, - queryOptions?: IQueryOptions, - ): Promise { - const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const query = executer('member_invitations') - .select( - 'member_invitations.*', - 'organizations.id as organization__id', - 'organizations.name as organization__name', - 'organizations.slug as organization__slug', - 'organizations.external_id as organization__external_id', - 'organizations.identity_provider as organization__identity_provider', - 'organizations.is_active as organization__is_active', - 'organizations.created_at as organization__created_at', - 'organizations.updated_at as organization__updated_at', - 'roles.id as role__id', - 'roles.name as role__name', - 'roles.description as role__description', - 'roles.created_at as role__created_at', - 'roles.updated_at as role__updated_at', - ) - .join( - 'organizations', - 'member_invitations.organization_id', - 'organizations.id', - ) - .join('roles', 'member_invitations.role_id', 'roles.id') - .orderBy('member_invitations.created_at', 'desc'); - - if (filter.token) query.where('member_invitations.token', filter.token); - if (filter.organizationId) { - query.where('member_invitations.organization_id', filter.organizationId); - } - if (filter.status) query.where('member_invitations.status', filter.status); - - const rows = await query; - return rows.map((row) => ({ - id: row.id, - organizationId: row.organization_id, - roleId: row.role_id, - email: row.email, - name: row.name, - token: row.token, - status: row.status, - invitedByMemberId: row.invited_by_member_id, - expiresAt: row.expires_at, - acceptedAt: row.accepted_at, - createdAt: row.created_at, - updatedAt: row.updated_at, - organization: { - id: row.organization__id, - name: row.organization__name, - slug: row.organization__slug, - externalId: row.organization__external_id, - identityProvider: row.organization__identity_provider, - isActive: row.organization__is_active, - createdAt: row.organization__created_at, - updatedAt: row.organization__updated_at, - }, - role: { - id: row.role__id, - name: row.role__name, - description: row.role__description, - createdAt: row.role__created_at, - updatedAt: row.role__updated_at, - }, - invitedBy: null, - })); - } -} - -export class AuthSessionRepository - extends BaseRepository - implements IAuthSessionRepository -{ - constructor(knexPg: Knex) { - super({ knexPg, tableName: 'auth_sessions' }); - } - - public async findActiveByToken(token: string, queryOptions?: IQueryOptions) { - const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const result = await executer('auth_sessions') - .where({ token }) - .whereNull('revoked_at') - .where('expires_at', '>', new Date().toISOString()) - .limit(1); - return result?.[0] ? this.transformFromDBData(result[0]) : null; - } -} diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 7ab2a4ab6d..e5e6d5a309 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -79,11 +79,6 @@ const resolvers = { apiHistory: apiHistoryResolver.getApiHistory, // Administration / RBAC - bootstrapStatus: rbacResolver.bootstrapStatus, - currentSession: rbacResolver.currentSession, - organizations: rbacResolver.listOrganizations, - organizationMembers: rbacResolver.listMembers, - memberInvitations: rbacResolver.listInvitations, roles: rbacResolver.listRoles, users: rbacResolver.listUsers, userRoleMappings: rbacResolver.listUserRoleMappings, @@ -198,9 +193,6 @@ const resolvers = { assignRoleToUser: rbacResolver.assignRoleToUser, updateUserRoles: rbacResolver.updateUserRoles, removeRoleFromUser: rbacResolver.removeRoleFromUser, - inviteMember: rbacResolver.inviteMember, - updateMember: rbacResolver.updateMember, - updateMemberRole: rbacResolver.updateMemberRole, }, ThreadResponse: askingResolver.getThreadResponseNestedResolver(), DetailStep: askingResolver.getDetailStepNestedResolver(), diff --git a/wren-ui/src/apollo/server/resolvers/rbacResolver.ts b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts index 5bb9966bdc..26f4f761ab 100644 --- a/wren-ui/src/apollo/server/resolvers/rbacResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts @@ -2,9 +2,6 @@ import { IContext } from '@server/types'; import { CreateRoleInput, CreateUserInput, - InviteMemberInput, - UpdateMemberInput, - UpdateMemberRoleInput, UpdateRoleInput, UpdateUserInput, UpdateUserRolesInput, @@ -15,15 +12,8 @@ import { Role, UserRole, UserRoleMapping, - Organization, - OrganizationMemberMapping, - MemberInvitationMapping, } from '@server/repositories'; -import { - AuthSessionResult, - RbacUserWithRoles, - RoleWithUsers, -} from '@server/services'; +import { RbacUserWithRoles, RoleWithUsers } from '@server/services'; import { getLogger } from '@server/utils'; const logger = getLogger('RbacResolver'); @@ -34,11 +24,6 @@ export class RbacResolver { this.listRoles = this.listRoles.bind(this); this.listUsers = this.listUsers.bind(this); this.listUserRoleMappings = this.listUserRoleMappings.bind(this); - this.bootstrapStatus = this.bootstrapStatus.bind(this); - this.currentSession = this.currentSession.bind(this); - this.listOrganizations = this.listOrganizations.bind(this); - this.listMembers = this.listMembers.bind(this); - this.listInvitations = this.listInvitations.bind(this); this.createRole = this.createRole.bind(this); this.updateRole = this.updateRole.bind(this); this.createUser = this.createUser.bind(this); @@ -46,9 +31,6 @@ export class RbacResolver { this.assignRoleToUser = this.assignRoleToUser.bind(this); this.updateUserRoles = this.updateUserRoles.bind(this); this.removeRoleFromUser = this.removeRoleFromUser.bind(this); - this.inviteMember = this.inviteMember.bind(this); - this.updateMember = this.updateMember.bind(this); - this.updateMemberRole = this.updateMemberRole.bind(this); } public getRoleNestedResolver() { @@ -109,50 +91,12 @@ export class RbacResolver { return ctx.rbacService.getUserRoleMappings(); } - public async bootstrapStatus(_root: any, _args: any, ctx: IContext) { - return ctx.rbacService.getBootstrapStatus(); - } - - public async currentSession( - _root: any, - _args: any, - ctx: IContext, - ): Promise { - return ctx.currentUser - ? ({ ...ctx.currentUser } as AuthSessionResult) - : null; - } - - public async listOrganizations( - _root: any, - _args: any, - ctx: IContext, - ): Promise { - return ctx.rbacService.listOrganizations(); - } - - public async listMembers( - _root: any, - _args: any, - ctx: IContext, - ): Promise { - return ctx.rbacService.listMembers(ctx.currentUser); - } - - public async listInvitations( - _root: any, - _args: any, - ctx: IContext, - ): Promise { - return ctx.rbacService.listInvitations(ctx.currentUser); - } - public async createRole( _root: any, args: { data: CreateRoleInput }, ctx: IContext, ): Promise { - return ctx.rbacService.createRole(args.data, ctx.currentUser); + return ctx.rbacService.createRole(args.data); } public async updateRole( @@ -160,10 +104,7 @@ export class RbacResolver { args: { where: { id: number }; data: Omit }, ctx: IContext, ): Promise { - return ctx.rbacService.updateRole( - { id: args.where.id, ...args.data }, - ctx.currentUser, - ); + return ctx.rbacService.updateRole({ id: args.where.id, ...args.data }); } public async createUser( @@ -171,7 +112,6 @@ export class RbacResolver { args: { data: CreateUserInput }, ctx: IContext, ): Promise { - this.assertAdmin(ctx); return ctx.rbacService.createUser(args.data); } @@ -180,7 +120,6 @@ export class RbacResolver { args: { where: { id: number }; data: Omit }, ctx: IContext, ): Promise { - this.assertAdmin(ctx); return ctx.rbacService.updateUser({ id: args.where.id, ...args.data }); } @@ -189,7 +128,6 @@ export class RbacResolver { args: { data: UserRoleInput }, ctx: IContext, ): Promise { - this.assertAdmin(ctx); return ctx.rbacService.assignRoleToUser(args.data); } @@ -198,7 +136,6 @@ export class RbacResolver { args: { data: UpdateUserRolesInput }, ctx: IContext, ): Promise { - this.assertAdmin(ctx); return ctx.rbacService.updateUserRoles(args.data); } @@ -207,37 +144,6 @@ export class RbacResolver { args: { data: UserRoleInput }, ctx: IContext, ): Promise { - this.assertAdmin(ctx); return ctx.rbacService.removeRoleFromUser(args.data); } - - public async inviteMember( - _root: any, - args: { data: InviteMemberInput }, - ctx: IContext, - ) { - return ctx.rbacService.inviteMember(args.data, ctx.currentUser); - } - - public async updateMember( - _root: any, - args: { data: UpdateMemberInput }, - ctx: IContext, - ) { - return ctx.rbacService.updateMember(args.data, ctx.currentUser); - } - - public async updateMemberRole( - _root: any, - args: { data: UpdateMemberRoleInput }, - ctx: IContext, - ) { - return ctx.rbacService.updateMemberRole(args.data, ctx.currentUser); - } - - private assertAdmin(ctx: IContext) { - if (ctx.currentUser?.role.name !== 'Admin') { - throw new Error('Admin role is required for this action.'); - } - } } diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index f63517edb7..c5a4cbb63e 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -1126,7 +1126,6 @@ export const typeDefs = gql` externalId: String identityProvider: String isActive: Boolean! - lastLoginAt: String roles: [Role!]! createdAt: String! updatedAt: String! @@ -1150,56 +1149,6 @@ export const typeDefs = gql` updatedAt: String! } - type Organization { - id: Int! - name: String! - slug: String! - externalId: String - identityProvider: String - isActive: Boolean! - createdAt: String! - updatedAt: String! - } - - type OrganizationMember { - id: Int! - organizationId: Int! - userId: Int! - roleId: Int! - status: String! - invitedByMemberId: Int - joinedAt: String - organization: Organization! - user: User! - role: Role! - createdAt: String! - updatedAt: String! - } - - type MemberInvitation { - id: Int! - organizationId: Int! - roleId: Int! - email: String! - name: String - token: String! - status: String! - invitedByMemberId: Int - expiresAt: String! - acceptedAt: String - organization: Organization! - role: Role! - createdAt: String! - updatedAt: String! - } - - type CurrentSession { - user: User! - member: OrganizationMember! - organization: Organization! - role: Role! - } - input RoleWhereInput { id: Int! } @@ -1221,7 +1170,6 @@ export const typeDefs = gql` input CreateUserInput { name: String! email: String! - password: String externalId: String identityProvider: String isActive: Boolean @@ -1246,25 +1194,6 @@ export const typeDefs = gql` roleIds: [Int!]! } - input InviteMemberInput { - organizationId: Int - email: String! - name: String - roleId: Int! - } - - input UpdateMemberInput { - id: Int! - name: String - roleId: Int - status: String - } - - input UpdateMemberRoleInput { - memberId: Int! - roleId: Int! - } - # Query and Mutation type Query { # On Boarding Steps @@ -1324,11 +1253,6 @@ export const typeDefs = gql` ): ApiHistoryPaginatedResponse! # Administration / RBAC - bootstrapStatus: JSON! - currentSession: CurrentSession - organizations: [Organization!]! - organizationMembers: [OrganizationMember!]! - memberInvitations: [MemberInvitation!]! roles: [Role!]! users: [User!]! userRoleMappings: [UserRoleMapping!]! @@ -1486,8 +1410,5 @@ export const typeDefs = gql` assignRoleToUser(data: UserRoleInput!): UserRole! updateUserRoles(data: UpdateUserRolesInput!): User! removeRoleFromUser(data: UserRoleInput!): Boolean! - inviteMember(data: InviteMemberInput!): MemberInvitation! - updateMember(data: UpdateMemberInput!): OrganizationMember! - updateMemberRole(data: UpdateMemberRoleInput!): OrganizationMember! } `; diff --git a/wren-ui/src/apollo/server/services/rbacService.ts b/wren-ui/src/apollo/server/services/rbacService.ts index a716c2dfc6..76064df4b5 100644 --- a/wren-ui/src/apollo/server/services/rbacService.ts +++ b/wren-ui/src/apollo/server/services/rbacService.ts @@ -1,15 +1,7 @@ import { Knex } from 'knex'; -import bcrypt from 'bcryptjs'; -import { randomBytes } from 'crypto'; import { - AcceptInvitationInput, - BootstrapAdminInput, CreateRoleInput, CreateUserInput, - InviteMemberInput, - LoginInput, - UpdateMemberInput, - UpdateMemberRoleInput, UpdateRoleInput, UpdateUserInput, UpdateUserRolesInput, @@ -19,15 +11,6 @@ import { IRoleRepository, IUserRepository, IUserRoleRepository, - IOrganizationRepository, - IOrganizationMemberRepository, - IMemberInvitationRepository, - IAuthSessionRepository, - AuthSession, - MemberInvitation, - MemberInvitationMapping, - Organization, - OrganizationMemberMapping, RbacUser, Role, UserRole, @@ -43,14 +26,9 @@ export interface RoleWithUsers extends Role { } export interface IRbacService { - getBootstrapStatus(): Promise<{ required: boolean }>; - bootstrapAdmin(input: BootstrapAdminInput): Promise; - login(input: LoginInput): Promise; - logout(token: string): Promise; - getSession(token?: string | null): Promise; listRoles(): Promise; - createRole(input: CreateRoleInput, actor?: AuthActor | null): Promise; - updateRole(input: UpdateRoleInput, actor?: AuthActor | null): Promise; + createRole(input: CreateRoleInput): Promise; + updateRole(input: UpdateRoleInput): Promise; listUsers(): Promise; createUser(input: CreateUserInput): Promise; updateUser(input: UpdateUserInput): Promise; @@ -58,227 +36,27 @@ export interface IRbacService { updateUserRoles(input: UpdateUserRolesInput): Promise; removeRoleFromUser(input: UserRoleInput): Promise; getUserRoleMappings(): Promise; - listOrganizations(): Promise; - listMembers(actor?: AuthActor | null): Promise; - listInvitations(actor?: AuthActor | null): Promise; - inviteMember( - input: InviteMemberInput, - actor?: AuthActor | null, - ): Promise; - acceptInvitation(input: AcceptInvitationInput): Promise; - updateMember( - input: UpdateMemberInput, - actor?: AuthActor | null, - ): Promise; - updateMemberRole( - input: UpdateMemberRoleInput, - actor?: AuthActor | null, - ): Promise; } const EMAIL_PATTERN = /^[^\s@]+@[^\s@]+\.[^\s@]+$/; -const ADMIN_ROLE_NAME = 'Admin'; -const SESSION_TTL_MS = 1000 * 60 * 60 * 24 * 7; -const INVITATION_TTL_MS = 1000 * 60 * 60 * 24 * 7; - -export interface AuthActor { - user: RbacUser; - member: OrganizationMemberMapping; - organization: Organization; - role: Role; -} - -export interface AuthSessionResult extends AuthActor { - session: AuthSession; -} export class RbacService implements IRbacService { private readonly roleRepository: IRoleRepository; private readonly userRepository: IUserRepository; private readonly userRoleRepository: IUserRoleRepository; - private readonly organizationRepository: IOrganizationRepository; - private readonly organizationMemberRepository: IOrganizationMemberRepository; - private readonly memberInvitationRepository: IMemberInvitationRepository; - private readonly authSessionRepository: IAuthSessionRepository; constructor({ roleRepository, userRepository, userRoleRepository, - organizationRepository, - organizationMemberRepository, - memberInvitationRepository, - authSessionRepository, }: { roleRepository: IRoleRepository; userRepository: IUserRepository; userRoleRepository: IUserRoleRepository; - organizationRepository: IOrganizationRepository; - organizationMemberRepository: IOrganizationMemberRepository; - memberInvitationRepository: IMemberInvitationRepository; - authSessionRepository: IAuthSessionRepository; }) { this.roleRepository = roleRepository; this.userRepository = userRepository; this.userRoleRepository = userRoleRepository; - this.organizationRepository = organizationRepository; - this.organizationMemberRepository = organizationMemberRepository; - this.memberInvitationRepository = memberInvitationRepository; - this.authSessionRepository = authSessionRepository; - } - - public async getBootstrapStatus(): Promise<{ required: boolean }> { - const adminRole = await this.roleRepository.findOneBy({ - name: ADMIN_ROLE_NAME, - }); - if (!adminRole) return { required: true }; - - const activeMembers = - await this.organizationMemberRepository.findMappings(); - return { - required: !activeMembers.some( - (member) => - member.status === 'active' && member.role.name === ADMIN_ROLE_NAME, - ), - }; - } - - public async bootstrapAdmin( - input: BootstrapAdminInput, - ): Promise { - const status = await this.getBootstrapStatus(); - if (!status.required) { - throw new Error('An Admin member already exists.'); - } - - const name = this.validateRequiredText(input.name, 'Name'); - const email = this.validateEmail(input.email); - const passwordHash = await this.hashPassword(input.password); - const role = await this.getRoleByNameOrThrow(ADMIN_ROLE_NAME); - const organizationName = this.validateRequiredText( - input.organizationName, - 'Organization name', - ); - const now = new Date().toISOString(); - const tx = await this.userRepository.transaction(); - - try { - const organization = await this.organizationRepository.createOne( - { - name: organizationName, - slug: await this.uniqueOrganizationSlug(organizationName), - isActive: true, - createdAt: now, - updatedAt: now, - }, - { tx }, - ); - const user = await this.userRepository.createOne( - { - name, - email, - passwordHash, - identityProvider: 'local', - isActive: true, - createdAt: now, - updatedAt: now, - }, - { tx }, - ); - const member = await this.organizationMemberRepository.createOne( - { - organizationId: organization.id, - userId: user.id, - roleId: role.id, - status: 'active', - joinedAt: now, - createdAt: now, - updatedAt: now, - }, - { tx }, - ); - const session = await this.createSession(user.id, member.id, tx); - await tx.commit(); - return { - user, - member: { ...member, user, role, organization }, - role, - organization, - session, - }; - } catch (error) { - await tx.rollback(); - throw error; - } - } - - public async login(input: LoginInput): Promise { - const email = this.validateEmail(input.email); - const user = await this.userRepository.findOneBy({ email }); - if (!user?.passwordHash) throw new Error('Invalid email or password.'); - const passwordMatches = await bcrypt.compare( - input.password || '', - user.passwordHash, - ); - if (!passwordMatches) throw new Error('Invalid email or password.'); - if (!user.isActive) throw new Error('This user is inactive.'); - - const member = - await this.organizationMemberRepository.findActiveMappingByUserId( - user.id, - ); - if (!member) throw new Error('This user is not an active member.'); - - const now = new Date().toISOString(); - const tx = await this.userRepository.transaction(); - try { - await this.userRepository.updateOne( - user.id, - { lastLoginAt: now, updatedAt: now }, - { tx }, - ); - const session = await this.createSession(user.id, member.id, tx); - await tx.commit(); - return { - user: { ...user, lastLoginAt: now }, - member, - role: member.role, - organization: member.organization, - session, - }; - } catch (error) { - await tx.rollback(); - throw error; - } - } - - public async logout(token: string): Promise { - const session = await this.authSessionRepository.findOneBy({ token }); - if (!session) return true; - await this.authSessionRepository.updateOne(session.id, { - revokedAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - }); - return true; - } - - public async getSession(token?: string | null) { - if (!token) return null; - const session = await this.authSessionRepository.findActiveByToken(token); - if (!session) return null; - const member = await this.organizationMemberRepository.findMappingById( - session.organizationMemberId, - ); - if (!member || member.status !== 'active' || !member.user.isActive) { - return null; - } - return { - user: member.user, - member, - role: member.role, - organization: member.organization, - session, - }; } public async listRoles(): Promise { @@ -292,11 +70,7 @@ export class RbacService implements IRbacService { })); } - public async createRole( - input: CreateRoleInput, - actor?: AuthActor | null, - ): Promise { - this.assertAdmin(actor); + public async createRole(input: CreateRoleInput): Promise { const name = this.validateRoleName(input.name); await this.assertUniqueRoleName(name); const now = new Date().toISOString(); @@ -308,11 +82,7 @@ export class RbacService implements IRbacService { }); } - public async updateRole( - input: UpdateRoleInput, - actor?: AuthActor | null, - ): Promise { - this.assertAdmin(actor); + public async updateRole(input: UpdateRoleInput): Promise { const role = await this.getRoleOrThrow(input.id); const data: Partial = { updatedAt: new Date().toISOString() }; @@ -343,9 +113,6 @@ export class RbacService implements IRbacService { const name = this.validateRequiredText(input.name, 'User name'); const email = this.validateEmail(input.email); await this.assertUniqueUserEmail(email); - const passwordHash = input.password - ? await this.hashPassword(input.password) - : null; const now = new Date().toISOString(); const tx = await this.userRepository.transaction(); @@ -354,7 +121,6 @@ export class RbacService implements IRbacService { { name, email, - passwordHash, externalId: this.normalizeNullable(input.externalId), identityProvider: this.normalizeNullable(input.identityProvider), isActive: input.isActive ?? true, @@ -450,204 +216,6 @@ export class RbacService implements IRbacService { return this.userRoleRepository.findMappings(); } - public async listOrganizations(): Promise { - return this.organizationRepository.findAll({ order: 'name' }); - } - - public async listMembers(actor?: AuthActor | null) { - this.assertAdmin(actor); - const organization = actor.organization; - return this.organizationMemberRepository.findMappingsByOrganizationId( - organization.id, - ); - } - - public async listInvitations(actor?: AuthActor | null) { - this.assertAdmin(actor); - const organization = actor.organization; - const invitations = await this.memberInvitationRepository.findMappings(); - return invitations.filter( - (invitation) => invitation.organizationId === organization.id, - ); - } - - public async inviteMember( - input: InviteMemberInput, - actor?: AuthActor | null, - ): Promise { - this.assertAdmin(actor); - const email = this.validateEmail(input.email); - const role = await this.getRoleOrThrow(input.roleId); - const organizationId = input.organizationId || actor.member.organizationId; - const organization = await this.organizationRepository.findOneBy({ - id: organizationId, - }); - if (!organization) throw new Error('Organization was not found.'); - - const pending = await this.memberInvitationRepository.findAllBy({ - organizationId, - email, - status: 'pending', - }); - if (pending.length) { - throw new Error(`An invitation for "${email}" is already pending.`); - } - - const existingUser = await this.userRepository.findOneBy({ email }); - if (existingUser) { - const existingMember = - await this.organizationMemberRepository.findActiveMappingByUserId( - existingUser.id, - ); - if (existingMember?.organizationId === organizationId) { - throw new Error(`"${email}" is already a member.`); - } - } - - const now = new Date().toISOString(); - return this.memberInvitationRepository.createOne({ - organizationId, - roleId: role.id, - email, - name: this.normalizeNullable(input.name), - token: this.generateToken(), - status: 'pending', - invitedByMemberId: actor.member.id, - expiresAt: new Date(Date.now() + INVITATION_TTL_MS).toISOString(), - createdAt: now, - updatedAt: now, - }); - } - - public async acceptInvitation( - input: AcceptInvitationInput, - ): Promise { - const invitation = await this.memberInvitationRepository.findMappingByToken( - input.token, - ); - if (!invitation || invitation.status !== 'pending') { - throw new Error('Invitation is invalid or has already been used.'); - } - if (new Date(invitation.expiresAt).getTime() < Date.now()) { - throw new Error('Invitation has expired.'); - } - - const name = this.validateRequiredText( - input.name || invitation.name || invitation.email, - 'Name', - ); - const passwordHash = await this.hashPassword(input.password); - const now = new Date().toISOString(); - const tx = await this.userRepository.transaction(); - - try { - let user = await this.userRepository.findOneBy( - { email: invitation.email }, - { tx }, - ); - if (user) { - user = await this.userRepository.updateOne( - user.id, - { - name, - passwordHash, - identityProvider: user.identityProvider || 'local', - isActive: true, - updatedAt: now, - }, - { tx }, - ); - } else { - user = await this.userRepository.createOne( - { - name, - email: invitation.email, - passwordHash, - identityProvider: 'local', - isActive: true, - createdAt: now, - updatedAt: now, - }, - { tx }, - ); - } - - const member = await this.organizationMemberRepository.createOne( - { - organizationId: invitation.organizationId, - userId: user.id, - roleId: invitation.roleId, - status: 'active', - invitedByMemberId: invitation.invitedByMemberId, - joinedAt: now, - createdAt: now, - updatedAt: now, - }, - { tx }, - ); - await this.memberInvitationRepository.updateOne( - invitation.id, - { status: 'accepted', acceptedAt: now, updatedAt: now }, - { tx }, - ); - const session = await this.createSession(user.id, member.id, tx); - await tx.commit(); - return { - user, - member: { - ...member, - user, - role: invitation.role, - organization: invitation.organization, - }, - role: invitation.role, - organization: invitation.organization, - session, - }; - } catch (error) { - await tx.rollback(); - throw error; - } - } - - public async updateMember( - input: UpdateMemberInput, - actor?: AuthActor | null, - ): Promise { - this.assertAdmin(actor); - const member = await this.getMemberOrThrow(input.id); - if (member.organizationId !== actor.member.organizationId) { - throw new Error('Member is outside of your organization.'); - } - const now = new Date().toISOString(); - if (input.name !== undefined && input.name !== null) { - await this.userRepository.updateOne(member.userId, { - name: this.validateRequiredText(input.name, 'Name'), - updatedAt: now, - }); - } - const data: any = { updatedAt: now }; - if (input.roleId !== undefined && input.roleId !== null) { - await this.getRoleOrThrow(input.roleId); - data.roleId = input.roleId; - } - if (input.status !== undefined && input.status !== null) { - data.status = this.validateMemberStatus(input.status); - } - await this.organizationMemberRepository.updateOne(member.id, data); - return this.getMemberOrThrow(member.id); - } - - public async updateMemberRole( - input: UpdateMemberRoleInput, - actor?: AuthActor | null, - ) { - return this.updateMember( - { id: input.memberId, roleId: input.roleId }, - actor, - ); - } - private async createUserRoleAssignments( userId: number, roleIds: number[], @@ -675,51 +243,12 @@ export class RbacService implements IRbacService { return role; } - private async getRoleByNameOrThrow(name: string): Promise { - const role = await this.roleRepository.findOneBy({ name }); - if (!role) throw new Error(`Role "${name}" was not found.`); - return role; - } - private async getUserOrThrow(id: number): Promise { const user = await this.userRepository.findOneBy({ id }); if (!user) throw new Error(`User ${id} was not found.`); return user; } - private async getMemberOrThrow( - id: number, - ): Promise { - const member = await this.organizationMemberRepository.findMappingById(id); - if (!member) throw new Error(`Member ${id} was not found.`); - return member; - } - - private assertAdmin(actor?: AuthActor | null): asserts actor is AuthActor { - if (!actor || actor.role.name !== ADMIN_ROLE_NAME) { - throw new Error('Admin role is required for this action.'); - } - } - - private async createSession( - userId: number, - organizationMemberId: number, - tx: Knex.Transaction, - ): Promise { - const now = new Date().toISOString(); - return this.authSessionRepository.createOne( - { - userId, - organizationMemberId, - token: this.generateToken(), - expiresAt: new Date(Date.now() + SESSION_TTL_MS).toISOString(), - createdAt: now, - updatedAt: now, - }, - { tx }, - ); - } - private async assertUniqueRoleName(name: string, exceptId?: number) { const roles = await this.roleRepository.findAll(); const duplicate = roles.find( @@ -751,22 +280,6 @@ export class RbacService implements IRbacService { return normalized; } - private validateMemberStatus(status: string): string { - const normalized = this.validateRequiredText(status, 'Status'); - if (!['active', 'inactive', 'suspended'].includes(normalized)) { - throw new Error('Member status must be active, inactive, or suspended.'); - } - return normalized; - } - - private async hashPassword(password: string): Promise { - const normalized = this.validateRequiredText(password, 'Password'); - if (normalized.length < 8) { - throw new Error('Password must be at least 8 characters.'); - } - return bcrypt.hash(normalized, 12); - } - private validateRequiredText(value: string, label: string): string { const normalized = `${value || ''}`.trim(); if (!normalized) throw new Error(`${label} is required.`); @@ -781,24 +294,4 @@ export class RbacService implements IRbacService { private uniqueIds(ids: number[]): number[] { return Array.from(new Set((ids || []).filter(Boolean))); } - - private generateToken(): string { - return randomBytes(32).toString('hex'); - } - - private async uniqueOrganizationSlug(name: string): Promise { - const base = name - .toLowerCase() - .replace(/[^a-z0-9]+/g, '-') - .replace(/^-|-$/g, '') - .slice(0, 80); - const safeBase = base || 'organization'; - let slug = safeBase; - let index = 1; - while (await this.organizationRepository.findOneBy({ slug })) { - slug = `${safeBase}-${index}`; - index += 1; - } - return slug; - } } diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index e94338f8ae..f92562a6e0 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -23,10 +23,6 @@ import { IRoleRepository, IUserRepository, IUserRoleRepository, - IOrganizationRepository, - IOrganizationMemberRepository, - IMemberInvitationRepository, - IAuthSessionRepository, } from '@server/repositories'; import { IQueryService, @@ -38,7 +34,6 @@ import { IDashboardService, IInstructionService, IRbacService, - AuthActor, } from '@server/services'; import { ITelemetry } from '@server/telemetry/telemetry'; import { @@ -69,7 +64,6 @@ export interface IContext { sqlPairService: ISqlPairService; instructionService: IInstructionService; rbacService: IRbacService; - currentUser?: AuthActor | null; // repository projectRepository: IProjectRepository; @@ -90,10 +84,6 @@ export interface IContext { roleRepository: IRoleRepository; userRepository: IUserRepository; userRoleRepository: IUserRoleRepository; - organizationRepository: IOrganizationRepository; - organizationMemberRepository: IOrganizationMemberRepository; - memberInvitationRepository: IMemberInvitationRepository; - authSessionRepository: IAuthSessionRepository; // background trackers projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; diff --git a/wren-ui/src/apollo/server/utils/auth.ts b/wren-ui/src/apollo/server/utils/auth.ts deleted file mode 100644 index 8a447b260b..0000000000 --- a/wren-ui/src/apollo/server/utils/auth.ts +++ /dev/null @@ -1,54 +0,0 @@ -import { NextApiRequest, NextApiResponse } from 'next'; - -export const AUTH_COOKIE_NAME = 'wren_auth_session'; - -export const getCookie = ( - req: Pick, - name: string, -): string | null => { - const cookieHeader = req.headers.cookie; - if (!cookieHeader) return null; - const cookies = cookieHeader.split(';').map((cookie) => cookie.trim()); - const cookie = cookies.find((item) => item.startsWith(`${name}=`)); - if (!cookie) return null; - return decodeURIComponent(cookie.slice(name.length + 1)); -}; - -export const setAuthCookie = ( - res: NextApiResponse, - token: string, - expiresAt: string, -) => { - res.setHeader( - 'Set-Cookie', - `${AUTH_COOKIE_NAME}=${encodeURIComponent( - token, - )}; Path=/; Expires=${new Date( - expiresAt, - ).toUTCString()}; HttpOnly; SameSite=Lax`, - ); -}; - -export const clearAuthCookie = (res: NextApiResponse) => { - res.setHeader( - 'Set-Cookie', - `${AUTH_COOKIE_NAME}=; Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT; HttpOnly; SameSite=Lax`, - ); -}; - -export const sanitizeAuthSession = (session: any) => { - if (!session) return session; - const sanitizeUser = (user: any) => { - if (!user) return user; - const { passwordHash: _passwordHash, ...rest } = user; - return rest; - }; - return { - user: sanitizeUser(session.user), - member: session.member - ? { ...session.member, user: sanitizeUser(session.member.user) } - : session.member, - organization: session.organization, - role: session.role, - }; -}; diff --git a/wren-ui/src/apollo/server/utils/index.ts b/wren-ui/src/apollo/server/utils/index.ts index 6256b1a11f..f9628a9dc3 100644 --- a/wren-ui/src/apollo/server/utils/index.ts +++ b/wren-ui/src/apollo/server/utils/index.ts @@ -8,4 +8,3 @@ export * from './helper'; export * from './regex'; export * from './sseTypes'; export * from './sseUtils'; -export * from './auth'; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 528222f1d5..18ef4da4d2 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -22,10 +22,6 @@ import { RoleRepository, UserRepository, UserRoleRepository, - OrganizationRepository, - OrganizationMemberRepository, - MemberInvitationRepository, - AuthSessionRepository, } from '@server/repositories'; import { WrenEngineAdaptor, @@ -161,10 +157,6 @@ export const initComponents = () => { const roleRepository = new RoleRepository(knex); const userRepository = new UserRepository(knex); const userRoleRepository = new UserRoleRepository(knex); - const organizationRepository = new OrganizationRepository(knex); - const organizationMemberRepository = new OrganizationMemberRepository(knex); - const memberInvitationRepository = new MemberInvitationRepository(knex); - const authSessionRepository = new AuthSessionRepository(knex); // adaptors const wrenEngineAdaptor = new WrenEngineAdaptor({ @@ -274,10 +266,6 @@ export const initComponents = () => { roleRepository, userRepository, userRoleRepository, - organizationRepository, - organizationMemberRepository, - memberInvitationRepository, - authSessionRepository, }); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ @@ -315,10 +303,6 @@ export const initComponents = () => { roleRepository, userRepository, userRoleRepository, - organizationRepository, - organizationMemberRepository, - memberInvitationRepository, - authSessionRepository, // adaptors wrenEngineAdaptor, diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index 1dc177a302..528dfdab5b 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -1,10 +1,9 @@ import { useRouter } from 'next/router'; -import { Button, Dropdown, Layout, Menu, Space } from 'antd'; +import { Button, Layout, Space } from 'antd'; import styled from 'styled-components'; import LogoBar from '@/components/LogoBar'; import { Path } from '@/utils/enum'; import Deploy from '@/components/deploy/Deploy'; -import { useAuth } from '@/hooks/useAuth'; const { Header } = Layout; @@ -34,17 +33,9 @@ const StyledHeader = styled(Header)` export default function HeaderBar() { const router = useRouter(); - const auth = useAuth(); const { pathname } = router; const showNav = !pathname.startsWith(Path.Onboarding); const isModeling = pathname.startsWith(Path.Modeling); - const roleName = auth.role?.name; - const isAdmin = roleName === 'Admin'; - const isManager = roleName === 'Manager'; - const isAnalyst = roleName === 'Analyst'; - const canModel = isAdmin || isManager; - const canUseKnowledge = isAdmin || isManager || isAnalyst; - const canUseApi = isAdmin || isManager; return ( @@ -69,7 +60,6 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.Modeling)} onClick={() => router.push(Path.Modeling)} - style={{ display: canModel ? undefined : 'none' }} > Modeling @@ -78,7 +68,6 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.Knowledge)} onClick={() => router.push(Path.KnowledgeQuestionSQLPairs)} - style={{ display: canUseKnowledge ? undefined : 'none' }} > Knowledge @@ -87,7 +76,6 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.APIManagement)} onClick={() => router.push(Path.APIManagementHistory)} - style={{ display: canUseApi ? undefined : 'none' }} > API @@ -96,35 +84,17 @@ export default function HeaderBar() { size="small" $isHighlight={pathname.startsWith(Path.Administration)} onClick={() => router.push(Path.AdministrationUsers)} - style={{ display: isAdmin ? undefined : 'none' }} > Admin )} - - {isModeling && canModel && } - {auth.authenticated && ( - - - {auth.user?.email} - {roleName} - - auth.logout()}> - Sign out - - - } - trigger={['click']} - > - - - )} - + {isModeling && ( + + + + )} ); diff --git a/wren-ui/src/components/pages/administration/types.tsx b/wren-ui/src/components/pages/administration/types.tsx index 548ff33662..896bf8cf40 100644 --- a/wren-ui/src/components/pages/administration/types.tsx +++ b/wren-ui/src/components/pages/administration/types.tsx @@ -21,45 +21,6 @@ export interface User { updatedAt: string; } -export interface Organization { - id: number; - name: string; - slug: string; - isActive: boolean; - createdAt: string; - updatedAt: string; -} - -export interface OrganizationMember { - id: number; - organizationId: number; - userId: number; - roleId: number; - status: string; - joinedAt?: string | null; - user: User; - role: Role; - organization: Organization; - createdAt: string; - updatedAt: string; -} - -export interface MemberInvitation { - id: number; - organizationId: number; - roleId: number; - email: string; - name?: string | null; - token: string; - status: string; - expiresAt: string; - acceptedAt?: string | null; - role: Role; - organization: Organization; - createdAt: string; - updatedAt: string; -} - export interface UserRoleMapping { id: number; userId: number; diff --git a/wren-ui/src/components/sidebar/Administration.tsx b/wren-ui/src/components/sidebar/Administration.tsx index c95b034e63..ecec7dd7ea 100644 --- a/wren-ui/src/components/sidebar/Administration.tsx +++ b/wren-ui/src/components/sidebar/Administration.tsx @@ -33,7 +33,7 @@ export default function Administration() { { label: ( - Member Management + User Management ), icon: , diff --git a/wren-ui/src/hooks/useAuth.tsx b/wren-ui/src/hooks/useAuth.tsx deleted file mode 100644 index 12bd963d9d..0000000000 --- a/wren-ui/src/hooks/useAuth.tsx +++ /dev/null @@ -1,134 +0,0 @@ -import { - createContext, - ReactNode, - useCallback, - useContext, - useEffect, - useMemo, - useState, -} from 'react'; -import { useRouter } from 'next/router'; -import PageLoading from '@/components/PageLoading'; -import { Path } from '@/utils/enum'; - -type RoleName = 'Admin' | 'Manager' | 'Analyst' | 'Viewer'; - -type AuthState = { - loading: boolean; - authenticated: boolean; - bootstrapRequired: boolean; - user?: any; - member?: any; - organization?: any; - role?: { name: RoleName }; - refresh: () => Promise; - logout: () => Promise; - canAccessPath: (path: string) => boolean; -}; - -const AuthContext = createContext({ - loading: true, - authenticated: false, - bootstrapRequired: false, - refresh: async () => undefined, - logout: async () => undefined, - canAccessPath: () => true, -}); - -const PUBLIC_PATHS = [Path.Login, Path.AcceptInvitation, Path.Onboarding]; - -const ROLE_PATHS: Record = { - Admin: [ - Path.Home, - Path.Modeling, - Path.Knowledge, - Path.APIManagement, - Path.Administration, - ], - Manager: [Path.Home, Path.Modeling, Path.Knowledge, Path.APIManagement], - Analyst: [Path.Home, Path.Knowledge], - Viewer: [Path.Home], -}; - -const isPublicPath = (pathname: string) => - PUBLIC_PATHS.some((path) => pathname.startsWith(path)); - -export const AuthProvider = ({ children }: { children: ReactNode }) => { - const [state, setState] = useState< - Omit - >({ - loading: true, - authenticated: false, - bootstrapRequired: false, - }); - - const refresh = useCallback(async () => { - const [statusResponse, meResponse] = await Promise.all([ - fetch('/api/auth/status'), - fetch('/api/auth/me'), - ]); - const status = await statusResponse.json(); - const me = meResponse.ok ? await meResponse.json() : null; - setState({ - loading: false, - bootstrapRequired: Boolean(status.required), - authenticated: Boolean(me?.authenticated), - user: me?.user, - member: me?.member, - organization: me?.organization, - role: me?.role, - }); - }, []); - - const logout = useCallback(async () => { - await fetch('/api/auth/logout', { method: 'POST' }); - await refresh(); - }, [refresh]); - - useEffect(() => { - void refresh(); - }, [refresh]); - - const canAccessPath = useCallback( - (path: string) => { - if (isPublicPath(path)) return true; - if (!state.authenticated) return false; - const roleName = state.role?.name; - if (!roleName) return false; - return ROLE_PATHS[roleName].some((allowed) => path.startsWith(allowed)); - }, - [state.authenticated, state.role?.name], - ); - - const value = useMemo( - () => ({ ...state, refresh, logout, canAccessPath }), - [state, refresh, logout, canAccessPath], - ); - - return {children}; -}; - -export const AuthGate = ({ children }: { children: ReactNode }) => { - const auth = useAuth(); - const router = useRouter(); - - useEffect(() => { - if (auth.loading) return; - if (isPublicPath(router.pathname)) return; - if (!auth.authenticated) { - void router.replace(Path.Login); - return; - } - if (!auth.canAccessPath(router.pathname)) { - void router.replace(Path.Home); - } - }, [auth, router]); - - if (auth.loading) return ; - if (!isPublicPath(router.pathname) && !auth.authenticated) { - return ; - } - return <>{children}; -}; - -export const useAuth = () => useContext(AuthContext); diff --git a/wren-ui/src/pages/_app.tsx b/wren-ui/src/pages/_app.tsx index 05ab7a5188..0b3b3765ef 100644 --- a/wren-ui/src/pages/_app.tsx +++ b/wren-ui/src/pages/_app.tsx @@ -7,7 +7,6 @@ import { GlobalConfigProvider } from '@/hooks/useGlobalConfig'; import { PostHogProvider } from 'posthog-js/react'; import { ApolloProvider } from '@apollo/client'; import { defaultIndicator } from '@/components/PageLoading'; -import { AuthGate, AuthProvider } from '@/hooks/useAuth'; require('../styles/index.less'); @@ -22,15 +21,11 @@ function App({ Component, pageProps }: AppProps) { - - - -
- -
- - - + +
+ +
+ diff --git a/wren-ui/src/pages/accept-invitation.tsx b/wren-ui/src/pages/accept-invitation.tsx deleted file mode 100644 index 2630bd9fc2..0000000000 --- a/wren-ui/src/pages/accept-invitation.tsx +++ /dev/null @@ -1,92 +0,0 @@ -import { useState } from 'react'; -import { useRouter } from 'next/router'; -import { Button, Card, Form, Input, Typography, message } from 'antd'; -import styled from 'styled-components'; -import LogoBar from '@/components/LogoBar'; -import { useAuth } from '@/hooks/useAuth'; -import { Path } from '@/utils/enum'; - -const { Paragraph, Title } = Typography; - -const Layout = styled.div` - min-height: 100vh; - display: flex; - align-items: center; - justify-content: center; - background: var(--gray-2); -`; - -const Panel = styled(Card)` - width: 420px; - box-shadow: 0 12px 32px rgba(0, 0, 0, 0.08); -`; - -export default function AcceptInvitationPage() { - const router = useRouter(); - const auth = useAuth(); - const [loading, setLoading] = useState(false); - const token = `${router.query.token || ''}`; - - const submit = async (values: any) => { - setLoading(true); - try { - const response = await fetch('/api/auth/accept-invitation', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ ...values, token }), - }); - const payload = await response.json(); - if (!response.ok) throw new Error(payload.message); - await auth.refresh(); - await router.replace(Path.Home); - } catch (error) { - message.error((error as Error).message); - } finally { - setLoading(false); - } - }; - - return ( - - -
- -
- - Accept invitation - - - Create your member profile to access this Wren AI workspace. - -
- - - - - - - - - - - ); -} diff --git a/wren-ui/src/pages/administration/assignments.tsx b/wren-ui/src/pages/administration/assignments.tsx index 421988514d..252691c282 100644 --- a/wren-ui/src/pages/administration/assignments.tsx +++ b/wren-ui/src/pages/administration/assignments.tsx @@ -4,6 +4,7 @@ import { Button, Form, Modal, + Popconfirm, Select, Table, TableColumnsType, @@ -12,128 +13,200 @@ import { } from 'antd'; import SafetyCertificateOutlined from '@ant-design/icons/SafetyCertificateOutlined'; import EditOutlined from '@ant-design/icons/EditOutlined'; +import DeleteOutlined from '@ant-design/icons/DeleteOutlined'; import SiderLayout from '@/components/layouts/SiderLayout'; import PageLayout from '@/components/layouts/PageLayout'; import { + ASSIGN_ROLE_TO_USER, LIST_USER_ROLE_MAPPINGS, - UPDATE_MEMBER_ROLE, + REMOVE_ROLE_FROM_USER, + UPDATE_USER_ROLES, } from '@/apollo/client/graphql/rbac'; import { - OrganizationMember, Role, RoleTags, + User, + UserRoleMapping, } from '@/components/pages/administration/types'; import { getAbsoluteTime } from '@/utils/time'; const { Text } = Typography; +type AssignmentModalState = { + visible: boolean; + user?: User; +}; + const AssignmentModal = ({ - member, + users, roles, + state, loading, onClose, onSubmit, }: { - member?: OrganizationMember; + users: User[]; roles: Role[]; + state: AssignmentModalState; loading: boolean; onClose: () => void; - onSubmit: (roleId: number, member: OrganizationMember) => Promise; + onSubmit: (values: any, user?: User) => Promise; }) => { const [form] = Form.useForm(); + const isUpdate = !!state.user; useEffect(() => { - if (!member) return; - form.setFieldsValue({ roleId: member.roleId }); - }, [form, member]); + if (!state.visible) return; + form.setFieldsValue({ + userId: state.user?.id, + roleId: undefined, + roleIds: state.user?.roles?.map((role) => role.id) || [], + }); + }, [form, state.visible, state.user]); const submit = async () => { - if (!member) return; const values = await form.validateFields(); - await onSubmit(values.roleId, member); + await onSubmit(values, state.user); form.resetFields(); onClose(); }; return ( form.resetFields()} >
- -
-
{member?.user.name}
- {member?.user.email} -
- ({ + label: role.name, + value: role.id, + }))} + /> + + ) : ( + +
setEditingMember(undefined)} + state={modalState} + loading={assignRoleState.loading || updateUserRolesState.loading} + onClose={closeModal} onSubmit={submitAssignment} /> diff --git a/wren-ui/src/pages/api/auth/accept-invitation.ts b/wren-ui/src/pages/api/auth/accept-invitation.ts deleted file mode 100644 index 63d980a49d..0000000000 --- a/wren-ui/src/pages/api/auth/accept-invitation.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; -import { sanitizeAuthSession, setAuthCookie } from '@/apollo/server/utils'; - -export default async function handler( - req: NextApiRequest, - res: NextApiResponse, -) { - if (req.method !== 'POST') return res.status(405).end(); - try { - const session = await components.rbacService.acceptInvitation(req.body); - setAuthCookie(res, session.session.token, session.session.expiresAt); - return res.status(200).json(sanitizeAuthSession(session)); - } catch (error) { - return res.status(400).json({ message: (error as Error).message }); - } -} diff --git a/wren-ui/src/pages/api/auth/bootstrap.ts b/wren-ui/src/pages/api/auth/bootstrap.ts deleted file mode 100644 index 7ed11c3e8e..0000000000 --- a/wren-ui/src/pages/api/auth/bootstrap.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; -import { sanitizeAuthSession, setAuthCookie } from '@/apollo/server/utils'; - -export default async function handler( - req: NextApiRequest, - res: NextApiResponse, -) { - if (req.method !== 'POST') return res.status(405).end(); - try { - const session = await components.rbacService.bootstrapAdmin(req.body); - setAuthCookie(res, session.session.token, session.session.expiresAt); - return res.status(200).json(sanitizeAuthSession(session)); - } catch (error) { - return res.status(400).json({ message: (error as Error).message }); - } -} diff --git a/wren-ui/src/pages/api/auth/login.ts b/wren-ui/src/pages/api/auth/login.ts deleted file mode 100644 index 7fe64fbd2f..0000000000 --- a/wren-ui/src/pages/api/auth/login.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; -import { sanitizeAuthSession, setAuthCookie } from '@/apollo/server/utils'; - -export default async function handler( - req: NextApiRequest, - res: NextApiResponse, -) { - if (req.method !== 'POST') return res.status(405).end(); - try { - const session = await components.rbacService.login(req.body); - setAuthCookie(res, session.session.token, session.session.expiresAt); - return res.status(200).json(sanitizeAuthSession(session)); - } catch (error) { - return res.status(401).json({ message: (error as Error).message }); - } -} diff --git a/wren-ui/src/pages/api/auth/logout.ts b/wren-ui/src/pages/api/auth/logout.ts deleted file mode 100644 index c5d6bbe328..0000000000 --- a/wren-ui/src/pages/api/auth/logout.ts +++ /dev/null @@ -1,18 +0,0 @@ -import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; -import { - AUTH_COOKIE_NAME, - clearAuthCookie, - getCookie, -} from '@/apollo/server/utils'; - -export default async function handler( - req: NextApiRequest, - res: NextApiResponse, -) { - if (req.method !== 'POST') return res.status(405).end(); - const token = getCookie(req, AUTH_COOKIE_NAME); - if (token) await components.rbacService.logout(token); - clearAuthCookie(res); - return res.status(200).json({ ok: true }); -} diff --git a/wren-ui/src/pages/api/auth/me.ts b/wren-ui/src/pages/api/auth/me.ts deleted file mode 100644 index cf304e037d..0000000000 --- a/wren-ui/src/pages/api/auth/me.ts +++ /dev/null @@ -1,21 +0,0 @@ -import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; -import { - AUTH_COOKIE_NAME, - getCookie, - sanitizeAuthSession, -} from '@/apollo/server/utils'; - -export default async function handler( - req: NextApiRequest, - res: NextApiResponse, -) { - if (req.method !== 'GET') return res.status(405).end(); - const token = getCookie(req, AUTH_COOKIE_NAME); - const session = await components.rbacService.getSession(token); - if (!session) return res.status(401).json({ authenticated: false }); - return res.status(200).json({ - authenticated: true, - ...sanitizeAuthSession(session), - }); -} diff --git a/wren-ui/src/pages/api/auth/status.ts b/wren-ui/src/pages/api/auth/status.ts deleted file mode 100644 index 344c459df1..0000000000 --- a/wren-ui/src/pages/api/auth/status.ts +++ /dev/null @@ -1,11 +0,0 @@ -import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; - -export default async function handler( - req: NextApiRequest, - res: NextApiResponse, -) { - if (req.method !== 'GET') return res.status(405).end(); - const status = await components.rbacService.getBootstrapStatus(); - return res.status(200).json(status); -} diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index ddaf00bd98..569329e5c4 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -5,7 +5,7 @@ import { typeDefs } from '@server'; import resolvers from '@server/resolvers'; import { IContext } from '@server/types'; import { GraphQLError } from 'graphql'; -import { AUTH_COOKIE_NAME, getCookie, getLogger } from '@server/utils'; +import { getLogger } from '@server/utils'; import { getConfig } from '@server/config'; import { ModelService } from '@server/services/modelService'; import { @@ -50,10 +50,6 @@ const bootstrapServer = async () => { roleRepository, userRepository, userRoleRepository, - organizationRepository, - organizationMemberRepository, - memberInvitationRepository, - authSessionRepository, // adaptors wrenEngineAdaptor, ibisAdaptor, @@ -133,10 +129,7 @@ const bootstrapServer = async () => { return defaultApolloErrorHandler(error); }, introspection: process.env.NODE_ENV !== 'production', - context: async ({ req }): Promise => { - const token = getCookie(req, AUTH_COOKIE_NAME); - const currentUser = await rbacService.getSession(token); - return { + context: (): IContext => ({ config: serverConfig, telemetry, // adaptor @@ -154,7 +147,6 @@ const bootstrapServer = async () => { sqlPairService, instructionService, rbacService, - currentUser, // repository projectRepository, modelRepository, @@ -174,16 +166,11 @@ const bootstrapServer = async () => { roleRepository, userRepository, userRoleRepository, - organizationRepository, - organizationMemberRepository, - memberInvitationRepository, - authSessionRepository, // background trackers projectRecommendQuestionBackgroundTracker, threadRecommendQuestionBackgroundTracker, dashboardCacheBackgroundTracker, - }; - }, + }), }); await apolloServer.start(); return apolloServer; diff --git a/wren-ui/src/pages/login.tsx b/wren-ui/src/pages/login.tsx deleted file mode 100644 index 2e8c47e4b2..0000000000 --- a/wren-ui/src/pages/login.tsx +++ /dev/null @@ -1,112 +0,0 @@ -import { useState } from 'react'; -import { useRouter } from 'next/router'; -import { Button, Card, Form, Input, Typography, message } from 'antd'; -import styled from 'styled-components'; -import LogoBar from '@/components/LogoBar'; -import { useAuth } from '@/hooks/useAuth'; -import { Path } from '@/utils/enum'; - -const { Paragraph, Title } = Typography; - -const Layout = styled.div` - min-height: 100vh; - display: flex; - align-items: center; - justify-content: center; - background: var(--gray-2); -`; - -const Panel = styled(Card)` - width: 420px; - box-shadow: 0 12px 32px rgba(0, 0, 0, 0.08); -`; - -export default function LoginPage() { - const auth = useAuth(); - const router = useRouter(); - const [loading, setLoading] = useState(false); - const [form] = Form.useForm(); - - const submit = async (values: any) => { - setLoading(true); - const endpoint = auth.bootstrapRequired - ? '/api/auth/bootstrap' - : '/api/auth/login'; - try { - const response = await fetch(endpoint, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(values), - }); - const payload = await response.json(); - if (!response.ok) throw new Error(payload.message); - await auth.refresh(); - await router.replace(Path.Home); - } catch (error) { - message.error((error as Error).message); - } finally { - setLoading(false); - } - }; - - return ( - - -
- -
- - {auth.bootstrapRequired ? 'Create your Admin account' : 'Sign in'} - - - {auth.bootstrapRequired - ? 'Set up the first organization and Admin member.' - : 'Access your Wren AI workspace.'} - - - {auth.bootstrapRequired && ( - - - - )} - {auth.bootstrapRequired && ( - - - - )} - - - - - - - - - - - ); -} diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index 588f71cc8c..54c0eb91f0 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -1,7 +1,5 @@ export enum Path { Home = '/home', - Login = '/login', - AcceptInvitation = '/accept-invitation', HomeDashboard = '/home/dashboard', Thread = '/home/[id]', Modeling = '/modeling', From b15b12592aaf7da82dc6087453f4e59ab2dbb697 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 3 Jun 2026 15:20:36 +0530 Subject: [PATCH 0087/1087] db details --- wren-ui/package.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/wren-ui/package.json b/wren-ui/package.json index 9e5af55473..56508cec14 100644 --- a/wren-ui/package.json +++ b/wren-ui/package.json @@ -10,10 +10,10 @@ "test": "jest", "test:e2e": "npx playwright install chromium && npx playwright test", "check-types": "tsc --noEmit", - "migrate": "yarn knex migrate:latest", + "migrate": "knex migrate:latest", "migrate:sqlite-to-mssql": "node -e \"process.env.MIGRATE_SQLITE_TO_MSSQL='true'; require('./tools/knex.js')\"", - "rollback": "yarn knex migrate:rollback", - "generate-gql": "yarn graphql-codegen --config codegen.yaml" + "rollback": "knex migrate:rollback", + "generate-gql": "graphql-codegen --config codegen.yaml" }, "dependencies": { "@google-cloud/bigquery": "^6.0.3", From e5d6f38c6175cd8a359e22eab53c59127c0cba6c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 29 May 2026 15:48:40 +0530 Subject: [PATCH 0088/1087] Updated qdrant --- .../src/providers/document_store/qdrant.py | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/wren-ai-service/src/providers/document_store/qdrant.py b/wren-ai-service/src/providers/document_store/qdrant.py index 5c7894493d..facbbff8c1 100644 --- a/wren-ai-service/src/providers/document_store/qdrant.py +++ b/wren-ai-service/src/providers/document_store/qdrant.py @@ -1,3 +1,4 @@ +import inspect import logging import os from typing import Any, Dict, List, Optional @@ -178,6 +179,70 @@ def __init__( collection_name=index, field_name="project_id", field_schema="keyword" ) + def recreate_collection( + self, + collection_name: str, + distance, + embedding_dim: int, + on_disk: Optional[bool] = None, + use_sparse_embeddings: Optional[bool] = None, + sparse_idf: bool = False, + ): + if on_disk is None: + on_disk = self.on_disk + + if use_sparse_embeddings is None: + use_sparse_embeddings = self.use_sparse_embeddings + + vectors_config = rest.VectorParams( + size=embedding_dim, + on_disk=on_disk, + distance=distance, + ) + sparse_vectors_config = None + + if use_sparse_embeddings: + vectors_config = {DENSE_VECTORS_NAME: vectors_config} + sparse_vectors_config = { + SPARSE_VECTORS_NAME: rest.SparseVectorParams( + index=rest.SparseIndexParams(on_disk=on_disk), + modifier=rest.Modifier.IDF if sparse_idf else None, + ), + } + + if self.client.collection_exists(collection_name): + self.client.delete_collection(collection_name) + + create_collection_kwargs = { + "collection_name": collection_name, + "vectors_config": vectors_config, + "sparse_vectors_config": sparse_vectors_config + if use_sparse_embeddings + else None, + "shard_number": self.shard_number, + "replication_factor": self.replication_factor, + "write_consistency_factor": self.write_consistency_factor, + "on_disk_payload": self.on_disk_payload, + "hnsw_config": self.hnsw_config, + "optimizers_config": self.optimizers_config, + "wal_config": self.wal_config, + "quantization_config": self.quantization_config, + } + if self.init_from is not None: + create_collection_signature = inspect.signature( + self.client.create_collection + ) + if "init_from" in create_collection_signature.parameters: + create_collection_kwargs["init_from"] = self.init_from + else: + logger.warning( + "Ignoring init_from for collection %s because the installed " + "qdrant-client does not support that argument", + collection_name, + ) + + self.client.create_collection(**create_collection_kwargs) + async def _query_by_embedding( self, query_embedding: List[float], From e52d3c97a96f968703f6e2cb5c365260eaca61f8 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 2 Jun 2026 19:26:25 +0530 Subject: [PATCH 0089/1087] wren-engine restore --- wren-engine | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-engine b/wren-engine index 47ca29ebba..06fb43c3db 160000 --- a/wren-engine +++ b/wren-engine @@ -1 +1 @@ -Subproject commit 47ca29ebba291100ba5d70ce1790f9887eaed7a0 +Subproject commit 06fb43c3dbb6486c05b93e722d03c5aca2c0c5f4 From 24ae9edc56379181316c170d401f253f7ff10aab Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 16:22:24 +0530 Subject: [PATCH 0090/1087] Add organization management feature --- ...0250603000000_create_organization_table.js | 33 ++ wren-ui/openapi.yaml | 118 +++++++ .../src/apollo/client/graphql/__types__.ts | 4 + .../server/middlewares/organizationApi.ts | 32 ++ .../repositories/apiHistoryRepository.ts | 4 + .../src/apollo/server/repositories/index.ts | 1 - .../repositories/organizationRepository.ts | 45 +++ wren-ui/src/apollo/server/schema.ts | 4 + wren-ui/src/apollo/server/services/index.ts | 4 + .../server/services/organizationService.ts | 142 ++++++++ .../tests/organizationService.test.ts | 106 ++++++ wren-ui/src/apollo/server/types/context.ts | 16 + wren-ui/src/common.ts | 25 ++ wren-ui/src/components/HeaderBar.tsx | 10 +- .../src/components/OrganizationSwitcher.tsx | 323 ++++++++++++++++++ .../src/pages/api/v1/organizations/index.ts | 82 +++++ 16 files changed, 944 insertions(+), 5 deletions(-) create mode 100644 wren-ui/migrations/20250603000000_create_organization_table.js create mode 100644 wren-ui/src/apollo/server/middlewares/organizationApi.ts create mode 100644 wren-ui/src/apollo/server/repositories/organizationRepository.ts create mode 100644 wren-ui/src/apollo/server/services/organizationService.ts create mode 100644 wren-ui/src/apollo/server/services/tests/organizationService.test.ts create mode 100644 wren-ui/src/components/OrganizationSwitcher.tsx create mode 100644 wren-ui/src/pages/api/v1/organizations/index.ts diff --git a/wren-ui/migrations/20250603000000_create_organization_table.js b/wren-ui/migrations/20250603000000_create_organization_table.js new file mode 100644 index 0000000000..9f87776f3b --- /dev/null +++ b/wren-ui/migrations/20250603000000_create_organization_table.js @@ -0,0 +1,33 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = function (knex) { + return knex.schema.createTable('organization', (table) => { + table.increments('id').comment('ID'); + table.string('name').notNullable().comment('Organization display name'); + table + .string('identifier') + .notNullable() + .unique() + .comment('Organization identifier used in selectors and APIs'); + table + .text('description') + .nullable() + .comment('Optional organization description'); + table + .boolean('is_current') + .notNullable() + .defaultTo(false) + .comment('Whether the organization is currently selected'); + table.timestamps(true, true); + }); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = function (knex) { + return knex.schema.dropTable('organization'); +}; diff --git a/wren-ui/openapi.yaml b/wren-ui/openapi.yaml index 8424a08a53..6b0e5f3b11 100644 --- a/wren-ui/openapi.yaml +++ b/wren-ui/openapi.yaml @@ -288,6 +288,106 @@ paths: application/json: schema: $ref: '#/components/schemas/ErrorResponse' + /organizations: + get: + summary: List organizations + description: Returns all organizations stored in the local WrenAI workspace + responses: + '200': + description: Organization list + content: + application/json: + schema: + type: object + properties: + organizations: + type: array + items: + $ref: '#/components/schemas/Organization' + currentProjectName: + type: string + post: + summary: Create organization + description: Creates a new organization in the local WrenAI workspace + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - name + properties: + name: + type: string + identifier: + type: string + description: + type: string + responses: + '201': + description: Organization created + content: + application/json: + schema: + $ref: '#/components/schemas/Organization' + '400': + description: Invalid payload + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' + '409': + description: Duplicate organization identifier + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' + /organizations/current: + get: + summary: Get current organization + description: Returns the current organization, the organization list, and the current project label + responses: + '200': + description: Current organization payload + content: + application/json: + schema: + type: object + properties: + currentOrganization: + allOf: + - $ref: '#/components/schemas/Organization' + nullable: true + organizations: + type: array + items: + $ref: '#/components/schemas/Organization' + currentProjectName: + type: string + /organizations/{id}/select: + post: + summary: Select current organization + description: Marks the specified organization as the current active organization + parameters: + - name: id + in: path + required: true + schema: + type: integer + responses: + '200': + description: Selected organization + content: + application/json: + schema: + $ref: '#/components/schemas/Organization' + '404': + description: Organization not found + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' /stream_explanation: get: summary: Stream an explanation @@ -356,3 +456,21 @@ components: properties: type: object description: Additional column properties + Organization: + type: object + properties: + id: + type: integer + name: + type: string + identifier: + type: string + description: + type: string + nullable: true + isCurrent: + type: boolean + createdAt: + type: string + updatedAt: + type: string diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index a7f054a896..34d20f9895 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -78,6 +78,7 @@ export type ApiHistoryResponse = { export enum ApiType { ASK = 'ASK', + CREATE_ORGANIZATION = 'CREATE_ORGANIZATION', CREATE_INSTRUCTION = 'CREATE_INSTRUCTION', CREATE_SQL_PAIR = 'CREATE_SQL_PAIR', DELETE_INSTRUCTION = 'DELETE_INSTRUCTION', @@ -85,10 +86,13 @@ export enum ApiType { GENERATE_SQL = 'GENERATE_SQL', GENERATE_SUMMARY = 'GENERATE_SUMMARY', GENERATE_VEGA_CHART = 'GENERATE_VEGA_CHART', + GET_CURRENT_ORGANIZATION = 'GET_CURRENT_ORGANIZATION', GET_INSTRUCTIONS = 'GET_INSTRUCTIONS', GET_MODELS = 'GET_MODELS', + GET_ORGANIZATIONS = 'GET_ORGANIZATIONS', GET_SQL_PAIRS = 'GET_SQL_PAIRS', RUN_SQL = 'RUN_SQL', + SELECT_ORGANIZATION = 'SELECT_ORGANIZATION', STREAM_ASK = 'STREAM_ASK', STREAM_GENERATE_SQL = 'STREAM_GENERATE_SQL', UPDATE_INSTRUCTION = 'UPDATE_INSTRUCTION', diff --git a/wren-ui/src/apollo/server/middlewares/organizationApi.ts b/wren-ui/src/apollo/server/middlewares/organizationApi.ts new file mode 100644 index 0000000000..e87a0ab8df --- /dev/null +++ b/wren-ui/src/apollo/server/middlewares/organizationApi.ts @@ -0,0 +1,32 @@ +import { NextApiRequest } from 'next'; +import { ApiError } from '../utils/apiUtils'; +import { components } from '@/common'; + +const { projectService } = components; + +export const assertAllowedMethods = ( + req: NextApiRequest, + methods: string[], +) => { + if (!req.method || !methods.includes(req.method)) { + throw new ApiError('Method not allowed', 405); + } +}; + +export const parseOrganizationId = (value: string | string[] | undefined) => { + const rawValue = Array.isArray(value) ? value[0] : value; + const parsed = Number(rawValue); + if (!rawValue || !Number.isInteger(parsed) || parsed <= 0) { + throw new ApiError('Invalid organization id', 400); + } + return parsed; +}; + +export const getCurrentProjectName = async () => { + try { + const project = await projectService.getCurrentProject(); + return project?.displayName || 'Default Project'; + } catch { + return 'Default Project'; + } +}; diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 7de5e5c097..d941a9fb8f 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -25,6 +25,10 @@ export enum ApiType { GET_MODELS = 'GET_MODELS', STREAM_ASK = 'STREAM_ASK', STREAM_GENERATE_SQL = 'STREAM_GENERATE_SQL', + GET_ORGANIZATIONS = 'GET_ORGANIZATIONS', + GET_CURRENT_ORGANIZATION = 'GET_CURRENT_ORGANIZATION', + CREATE_ORGANIZATION = 'CREATE_ORGANIZATION', + SELECT_ORGANIZATION = 'SELECT_ORGANIZATION', } export interface ApiHistory { diff --git a/wren-ui/src/apollo/server/repositories/index.ts b/wren-ui/src/apollo/server/repositories/index.ts index 28e92cce08..0dd1cc5905 100644 --- a/wren-ui/src/apollo/server/repositories/index.ts +++ b/wren-ui/src/apollo/server/repositories/index.ts @@ -19,4 +19,3 @@ export * from './askingTaskRepository'; export * from './instructionRepository'; export * from './apiHistoryRepository'; export * from './dashboardItemRefreshJobRepository'; -export * from './rbacRepository'; diff --git a/wren-ui/src/apollo/server/repositories/organizationRepository.ts b/wren-ui/src/apollo/server/repositories/organizationRepository.ts new file mode 100644 index 0000000000..7059b5ad9e --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/organizationRepository.ts @@ -0,0 +1,45 @@ +import { Knex } from 'knex'; +import { BaseRepository, IBasicRepository, IQueryOptions } from './baseRepository'; + +export interface Organization { + id: number; + name: string; + identifier: string; + description?: string | null; + isCurrent: boolean; + createdAt?: string; + updatedAt?: string; +} + +export interface IOrganizationRepository extends IBasicRepository { + getCurrentOrganization: ( + queryOptions?: IQueryOptions, + ) => Promise; + setCurrentOrganization: ( + id: number, + queryOptions?: IQueryOptions, + ) => Promise; +} + +export class OrganizationRepository + extends BaseRepository + implements IOrganizationRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization' }); + } + + public async getCurrentOrganization(queryOptions?: IQueryOptions) { + return await this.findOneBy({ isCurrent: true }, queryOptions); + } + + public async setCurrentOrganization(id: number, queryOptions?: IQueryOptions) { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + await executer(this.tableName).update({ is_current: false }); + const [result] = await executer(this.tableName) + .where({ id }) + .update({ is_current: true }) + .returning('*'); + return this.transformFromDBData(result); + } +} diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index c5a4cbb63e..2c6fa18516 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -21,6 +21,10 @@ export const typeDefs = gql` GET_MODELS STREAM_ASK STREAM_GENERATE_SQL + GET_ORGANIZATIONS + GET_CURRENT_ORGANIZATION + CREATE_ORGANIZATION + SELECT_ORGANIZATION } input ApiHistoryFilterInput { diff --git a/wren-ui/src/apollo/server/services/index.ts b/wren-ui/src/apollo/server/services/index.ts index 1a7c8dadb0..e5b7f161cf 100644 --- a/wren-ui/src/apollo/server/services/index.ts +++ b/wren-ui/src/apollo/server/services/index.ts @@ -8,4 +8,8 @@ export * from './metadataService'; export * from './dashboardService'; export * from './askingTaskTracker'; export * from './instructionService'; +<<<<<<< HEAD export * from './rbacService'; +======= +export * from './organizationService'; +>>>>>>> 52a339fb5 (Add organization management feature) diff --git a/wren-ui/src/apollo/server/services/organizationService.ts b/wren-ui/src/apollo/server/services/organizationService.ts new file mode 100644 index 0000000000..293322b95e --- /dev/null +++ b/wren-ui/src/apollo/server/services/organizationService.ts @@ -0,0 +1,142 @@ +import { ApiError } from '../utils/apiUtils'; +import { + IOrganizationRepository, + Organization, +} from '../repositories/organizationRepository'; + +export interface CreateOrganizationData { + name: string; + identifier?: string; + description?: string; +} + +export interface IOrganizationService { + listOrganizations: () => Promise; + getCurrentOrganization: () => Promise; + createOrganization: (data: CreateOrganizationData) => Promise; + selectCurrentOrganization: (id: number) => Promise; +} + +const NAME_MAX_LENGTH = 64; +const IDENTIFIER_MAX_LENGTH = 64; +const DESCRIPTION_MAX_LENGTH = 255; +const IDENTIFIER_PATTERN = /^[a-z0-9]+(?:[_-][a-z0-9]+)*$/; + +export const normalizeOrganizationName = (value: string) => value.trim(); + +export const normalizeOrganizationIdentifier = (value: string) => + value + .trim() + .toLowerCase() + .replace(/[^a-z0-9]+/g, '_') + .replace(/_+/g, '_') + .replace(/^_+|_+$/g, ''); + +export const validateOrganizationPayload = ( + data: CreateOrganizationData, +): Required => { + const name = normalizeOrganizationName(data.name || ''); + const identifier = normalizeOrganizationIdentifier( + data.identifier || data.name || '', + ); + const description = (data.description || '').trim(); + + if (!name) { + throw new ApiError('Organization name is required', 400); + } + if (name.length > NAME_MAX_LENGTH) { + throw new ApiError( + `Organization name is too long (max ${NAME_MAX_LENGTH} characters)`, + 400, + ); + } + if (!identifier) { + throw new ApiError('Organization identifier is required', 400); + } + if (identifier.length > IDENTIFIER_MAX_LENGTH) { + throw new ApiError( + `Organization identifier is too long (max ${IDENTIFIER_MAX_LENGTH} characters)`, + 400, + ); + } + if (!IDENTIFIER_PATTERN.test(identifier)) { + throw new ApiError( + 'Organization identifier must contain only lowercase letters, numbers, underscores, or hyphens', + 400, + ); + } + if (description.length > DESCRIPTION_MAX_LENGTH) { + throw new ApiError( + `Organization description is too long (max ${DESCRIPTION_MAX_LENGTH} characters)`, + 400, + ); + } + + return { + name, + identifier, + description, + }; +}; + +export class OrganizationService implements IOrganizationService { + private organizationRepository: IOrganizationRepository; + + constructor({ + organizationRepository, + }: { + organizationRepository: IOrganizationRepository; + }) { + this.organizationRepository = organizationRepository; + } + + public async listOrganizations() { + return await this.organizationRepository.findAll({ order: 'id' }); + } + + public async getCurrentOrganization() { + return await this.organizationRepository.getCurrentOrganization(); + } + + public async createOrganization(data: CreateOrganizationData) { + const payload = validateOrganizationPayload(data); + const existingByIdentifier = await this.organizationRepository.findOneBy({ + identifier: payload.identifier, + }); + if (existingByIdentifier) { + throw new ApiError('Organization identifier already exists', 409); + } + + const tx = await this.organizationRepository.transaction(); + try { + const currentOrganization = + await this.organizationRepository.getCurrentOrganization({ tx }); + const created = await this.organizationRepository.createOne( + { + ...payload, + isCurrent: !currentOrganization, + }, + { tx }, + ); + const result = currentOrganization + ? created + : await this.organizationRepository.setCurrentOrganization(created.id, { + tx, + }); + await this.organizationRepository.commit(tx); + return result; + } catch (error) { + await this.organizationRepository.rollback(tx); + throw error; + } + } + + public async selectCurrentOrganization(id: number) { + const organization = await this.organizationRepository.findOneBy({ id }); + if (!organization) { + throw new ApiError('Organization not found', 404); + } + + return await this.organizationRepository.setCurrentOrganization(id); + } +} diff --git a/wren-ui/src/apollo/server/services/tests/organizationService.test.ts b/wren-ui/src/apollo/server/services/tests/organizationService.test.ts new file mode 100644 index 0000000000..41830424d9 --- /dev/null +++ b/wren-ui/src/apollo/server/services/tests/organizationService.test.ts @@ -0,0 +1,106 @@ +import { + OrganizationService, + normalizeOrganizationIdentifier, + validateOrganizationPayload, +} from '../organizationService'; + +describe('OrganizationService', () => { + const mockOrganizationRepository = () => ({ + findAll: jest.fn(), + findOneBy: jest.fn(), + createOne: jest.fn(), + updateOne: jest.fn(), + transaction: jest.fn(), + commit: jest.fn(), + rollback: jest.fn(), + getCurrentOrganization: jest.fn(), + setCurrentOrganization: jest.fn(), + }); + + describe('validateOrganizationPayload', () => { + it('should normalize valid payloads', () => { + expect( + validateOrganizationPayload({ + name: ' My Org ', + identifier: 'My Org', + description: ' Team workspace ', + }), + ).toEqual({ + name: 'My Org', + identifier: 'my_org', + description: 'Team workspace', + }); + }); + + it('should reject invalid identifiers', () => { + expect(() => + validateOrganizationPayload({ + name: 'Org', + identifier: '***', + }), + ).toThrow('Organization identifier is required'); + }); + }); + + describe('createOrganization', () => { + it('should make the first organization current', async () => { + const repository = mockOrganizationRepository(); + const tx = {}; + repository.transaction.mockResolvedValue(tx); + repository.getCurrentOrganization.mockResolvedValue(null); + repository.findOneBy.mockResolvedValue(null); + repository.createOne.mockResolvedValue({ + id: 1, + name: 'My Org', + identifier: 'my_org', + description: '', + isCurrent: true, + }); + repository.setCurrentOrganization.mockResolvedValue({ + id: 1, + name: 'My Org', + identifier: 'my_org', + description: '', + isCurrent: true, + }); + + const service = new OrganizationService({ + organizationRepository: repository as any, + }); + + const result = await service.createOrganization({ name: 'My Org' }); + + expect(repository.createOne).toHaveBeenCalledWith( + { + name: 'My Org', + identifier: 'my_org', + description: '', + isCurrent: true, + }, + { tx }, + ); + expect(repository.setCurrentOrganization).toHaveBeenCalledWith(1, { tx }); + expect(result.identifier).toBe('my_org'); + }); + + it('should reject duplicate identifiers', async () => { + const repository = mockOrganizationRepository(); + repository.findOneBy.mockResolvedValue({ id: 1, identifier: 'my_org' }); + const service = new OrganizationService({ + organizationRepository: repository as any, + }); + + await expect( + service.createOrganization({ name: 'My Org', identifier: 'my_org' }), + ).rejects.toThrow('Organization identifier already exists'); + }); + }); + + describe('normalizeOrganizationIdentifier', () => { + it('should collapse separators into underscores', () => { + expect(normalizeOrganizationIdentifier('My cool-org!!')).toBe( + 'my_cool_org', + ); + }); + }); +}); diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index f92562a6e0..ce7b0b5ba2 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -20,9 +20,13 @@ import { IInstructionRepository, IApiHistoryRepository, IDashboardItemRefreshJobRepository, +<<<<<<< HEAD IRoleRepository, IUserRepository, IUserRoleRepository, +======= + IOrganizationRepository, +>>>>>>> 52a339fb5 (Add organization management feature) } from '@server/repositories'; import { IQueryService, @@ -33,7 +37,11 @@ import { IProjectService, IDashboardService, IInstructionService, +<<<<<<< HEAD IRbacService, +======= + IOrganizationService, +>>>>>>> 52a339fb5 (Add organization management feature) } from '@server/services'; import { ITelemetry } from '@server/telemetry/telemetry'; import { @@ -63,7 +71,11 @@ export interface IContext { dashboardService: IDashboardService; sqlPairService: ISqlPairService; instructionService: IInstructionService; +<<<<<<< HEAD rbacService: IRbacService; +======= + organizationService: IOrganizationService; +>>>>>>> 52a339fb5 (Add organization management feature) // repository projectRepository: IProjectRepository; @@ -81,9 +93,13 @@ export interface IContext { instructionRepository: IInstructionRepository; apiHistoryRepository: IApiHistoryRepository; dashboardItemRefreshJobRepository: IDashboardItemRefreshJobRepository; +<<<<<<< HEAD roleRepository: IRoleRepository; userRepository: IUserRepository; userRoleRepository: IUserRoleRepository; +======= + organizationRepository: IOrganizationRepository; +>>>>>>> 52a339fb5 (Add organization management feature) // background trackers projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 18ef4da4d2..b7cd09b6ca 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -19,9 +19,13 @@ import { InstructionRepository, ApiHistoryRepository, DashboardItemRefreshJobRepository, +<<<<<<< HEAD RoleRepository, UserRepository, UserRoleRepository, +======= + OrganizationRepository, +>>>>>>> 52a339fb5 (Add organization management feature) } from '@server/repositories'; import { WrenEngineAdaptor, @@ -38,7 +42,11 @@ import { DashboardService, AskingTaskTracker, InstructionService, +<<<<<<< HEAD RbacService, +======= + OrganizationService, +>>>>>>> 52a339fb5 (Add organization management feature) } from '@server/services'; import { PostHogTelemetry } from './apollo/server/telemetry/telemetry'; import { @@ -154,9 +162,13 @@ export const initComponents = () => { const apiHistoryRepository = new ApiHistoryRepository(knex); const dashboardItemRefreshJobRepository = new DashboardItemRefreshJobRepository(knex); +<<<<<<< HEAD const roleRepository = new RoleRepository(knex); const userRepository = new UserRepository(knex); const userRoleRepository = new UserRoleRepository(knex); +======= + const organizationRepository = new OrganizationRepository(knex); +>>>>>>> 52a339fb5 (Add organization management feature) // adaptors const wrenEngineAdaptor = new WrenEngineAdaptor({ @@ -262,10 +274,15 @@ export const initComponents = () => { instructionRepository, wrenAIAdaptor, }); +<<<<<<< HEAD const rbacService = new RbacService({ roleRepository, userRepository, userRoleRepository, +======= + const organizationService = new OrganizationService({ + organizationRepository, +>>>>>>> 52a339fb5 (Add organization management feature) }); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ @@ -300,9 +317,13 @@ export const initComponents = () => { apiHistoryRepository, instructionRepository, dashboardItemRefreshJobRepository, +<<<<<<< HEAD roleRepository, userRepository, userRoleRepository, +======= + organizationRepository, +>>>>>>> 52a339fb5 (Add organization management feature) // adaptors wrenEngineAdaptor, @@ -319,7 +340,11 @@ export const initComponents = () => { dashboardService, sqlPairService, instructionService, +<<<<<<< HEAD rbacService, +======= + organizationService, +>>>>>>> 52a339fb5 (Add organization management feature) askingTaskTracker, // background trackers diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index 528dfdab5b..953440cd62 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -4,6 +4,7 @@ import styled from 'styled-components'; import LogoBar from '@/components/LogoBar'; import { Path } from '@/utils/enum'; import Deploy from '@/components/deploy/Deploy'; +import OrganizationSwitcher from '@/components/OrganizationSwitcher'; const { Header } = Layout; @@ -90,11 +91,12 @@ export default function HeaderBar() { )} - {isModeling && ( - + + + {isModeling && ( - - )} + )} + ); diff --git a/wren-ui/src/components/OrganizationSwitcher.tsx b/wren-ui/src/components/OrganizationSwitcher.tsx new file mode 100644 index 0000000000..e54ee19314 --- /dev/null +++ b/wren-ui/src/components/OrganizationSwitcher.tsx @@ -0,0 +1,323 @@ +import { useEffect, useMemo, useState } from 'react'; +import { + Button, + Dropdown, + Form, + Input, + Menu, + message, + Modal, + Space, + Spin, + Tooltip, + Typography, +} from 'antd'; +import styled from 'styled-components'; +import PlusOutlined from '@ant-design/icons/PlusOutlined'; +import DownOutlined from '@ant-design/icons/DownOutlined'; + +interface OrganizationRecord { + id: number; + name: string; + identifier: string; + description?: string | null; + isCurrent: boolean; +} + +interface OrganizationResponse { + organizations: OrganizationRecord[]; + currentOrganization: OrganizationRecord | null; + currentProjectName: string; +} + +const TriggerButton = styled(Button)` + display: flex; + align-items: center; + color: var(--gray-2); + border: none; + background: transparent; + padding: 0 8px; + + &:hover, + &:focus { + color: var(--gray-1); + background: rgba(255, 255, 255, 0.05); + } +`; + +const OrganizationBadge = styled.div` + width: 28px; + height: 28px; + border-radius: 6px; + background: #b7eb8f; + color: #274916; + display: flex; + align-items: center; + justify-content: center; + font-weight: 600; + text-transform: uppercase; +`; + +const Overlay = styled.div` + width: 280px; + background: var(--gray-1); + border-radius: 10px; + box-shadow: 0 10px 30px rgba(0, 0, 0, 0.16); + overflow: hidden; +`; + +const OverlayHeader = styled.div` + display: flex; + align-items: center; + justify-content: space-between; + padding: 12px 16px 8px; +`; + +const OverlayLabel = styled.div` + font-size: 12px; + letter-spacing: 0.02em; + color: var(--gray-6); + text-transform: uppercase; +`; + +const OverlayMenu = styled(Menu)` + border-right: none; + box-shadow: none; + + .ant-dropdown-menu-item, + .ant-dropdown-menu-submenu-title { + padding: 0; + } +`; + +const MenuRow = styled.button<{ $active?: boolean }>` + width: 100%; + display: flex; + align-items: center; + gap: 12px; + padding: 10px 16px; + border: none; + background: ${(props) => (props.$active ? 'var(--gray-3)' : 'transparent')}; + text-align: left; + cursor: pointer; + + &:hover { + background: var(--gray-3); + } +`; + +const getBadgeText = (name?: string) => + (name || 'O') + .trim() + .split(/\s+/) + .slice(0, 1) + .map((segment) => segment[0]) + .join('') + .toUpperCase(); + +export default function OrganizationSwitcher() { + const [form] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + const [visible, setVisible] = useState(false); + const [currentProjectName, setCurrentProjectName] = useState('Default Project'); + const [organizations, setOrganizations] = useState([]); + const [currentOrganization, setCurrentOrganization] = + useState(null); + + const loadOrganizations = async () => { + setLoading(true); + try { + const response = await fetch('/api/v1/organizations/current'); + const payload = (await response.json()) as OrganizationResponse & { + error?: string; + }; + if (!response.ok) { + throw new Error(payload.error || 'Failed to load organizations'); + } + setOrganizations(payload.organizations || []); + setCurrentOrganization(payload.currentOrganization || null); + setCurrentProjectName(payload.currentProjectName || 'Default Project'); + } catch (error: any) { + message.error(error.message || 'Failed to load organizations'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadOrganizations(); + }, []); + + const hasOrganizations = organizations.length > 0; + const currentLabel = currentOrganization?.identifier || 'Create organization'; + const projectLabel = useMemo( + () => currentProjectName || 'Default Project', + [currentProjectName], + ); + + const createOrganization = async () => { + try { + const values = await form.validateFields(); + setSaving(true); + const response = await fetch('/api/v1/organizations', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify(values), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to create organization'); + } + message.success('Organization created successfully.'); + setVisible(false); + form.resetFields(); + await loadOrganizations(); + } catch (error: any) { + if (error?.errorFields) { + return; + } + message.error(error.message || 'Failed to create organization'); + } finally { + setSaving(false); + } + }; + + const selectOrganization = async (organizationId: number) => { + try { + const response = await fetch( + `/api/v1/organizations/${organizationId}/select`, + { + method: 'POST', + }, + ); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to switch organization'); + } + message.success('Organization switched successfully.'); + await loadOrganizations(); + } catch (error: any) { + message.error(error.message || 'Failed to switch organization'); + } + }; + + const overlay = ( + + + Organizations + + + )} + + void createOrganization()} + onCancel={() => { + setVisible(false); + form.resetFields(); + }} + destroyOnClose + > +
+ + + + + + + + + + + + + ); +} diff --git a/wren-ui/src/pages/api/v1/organizations/index.ts b/wren-ui/src/pages/api/v1/organizations/index.ts new file mode 100644 index 0000000000..643ee2432d --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/index.ts @@ -0,0 +1,82 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectName, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ORGANIZATIONS'); +logger.level = 'debug'; + +const { organizationService } = components; + +const serializeOrganization = (organization) => ({ + id: organization.id, + name: organization.name, + identifier: organization.identifier, + description: organization.description, + isCurrent: Boolean(organization.isCurrent), + createdAt: organization.createdAt, + updatedAt: organization.updatedAt, +}); + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET', 'POST']); + + if (req.method === 'GET') { + const organizations = await organizationService.listOrganizations(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + organizations: organizations.map(serializeOrganization), + currentProjectName: await getCurrentProjectName(), + }, + projectId: 0, + apiType: ApiType.GET_ORGANIZATIONS, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + + const organization = await organizationService.createOrganization(req.body); + await respondWithSimple({ + res, + statusCode: 201, + responsePayload: serializeOrganization(organization), + projectId: 0, + apiType: ApiType.CREATE_ORGANIZATION, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: 0, + apiType: + req.method === 'GET' + ? ApiType.GET_ORGANIZATIONS + : ApiType.CREATE_ORGANIZATION, + requestPayload: req.method === 'GET' ? {} : req.body, + headers: req.headers as Record, + startTime, + logger, + }); + } +} From 5d78a952f97c22a66a314baa51a56436409682d2 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 15:35:43 +0530 Subject: [PATCH 0091/1087] Add organization management feature --- package-lock.json | 256 ++++++++++++++++++ package.json | 5 + .../pages/api/v1/organizations/[id]/select.ts | 61 +++++ .../src/pages/api/v1/organizations/current.ts | 73 +++++ 4 files changed, 395 insertions(+) create mode 100644 package-lock.json create mode 100644 package.json create mode 100644 wren-ui/src/pages/api/v1/organizations/[id]/select.ts create mode 100644 wren-ui/src/pages/api/v1/organizations/current.ts diff --git a/package-lock.json b/package-lock.json new file mode 100644 index 0000000000..f3e1aa42f9 --- /dev/null +++ b/package-lock.json @@ -0,0 +1,256 @@ +{ + "name": "WrenAI1", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "devDependencies": { + "@apollo/client": "^3.6.9" + } + }, + "node_modules/@apollo/client": { + "version": "3.6.9", + "resolved": "https://registry.npmjs.org/@apollo/client/-/client-3.6.9.tgz", + "integrity": "sha512-Y1yu8qa2YeaCUBVuw08x8NHenFi0sw2I3KCu7Kw9mDSu86HmmtHJkCAifKVrN2iPgDTW/BbP3EpSV8/EQCcxZA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@graphql-typed-document-node/core": "^3.1.1", + "@wry/context": "^0.6.0", + "@wry/equality": "^0.5.0", + "@wry/trie": "^0.3.0", + "graphql-tag": "^2.12.6", + "hoist-non-react-statics": "^3.3.2", + "optimism": "^0.16.1", + "prop-types": "^15.7.2", + "symbol-observable": "^4.0.0", + "ts-invariant": "^0.10.3", + "tslib": "^2.3.0", + "zen-observable-ts": "^1.2.5" + }, + "peerDependencies": { + "graphql": "^14.0.0 || ^15.0.0 || ^16.0.0", + "graphql-ws": "^5.5.5", + "react": "^16.8.0 || ^17.0.0 || ^18.0.0", + "subscriptions-transport-ws": "^0.9.0 || ^0.11.0" + }, + "peerDependenciesMeta": { + "graphql-ws": { + "optional": true + }, + "react": { + "optional": true + }, + "subscriptions-transport-ws": { + "optional": true + } + } + }, + "node_modules/@graphql-typed-document-node/core": { + "version": "3.2.0", + "resolved": "https://registry.npmjs.org/@graphql-typed-document-node/core/-/core-3.2.0.tgz", + "integrity": "sha512-mB9oAsNCm9aM3/SOv4YtBMqZbYj10R7dkq8byBqxGY/ncFwhf2oQzMV+LCRlWoDSEBJ3COiR1yeDvMtsoOsuFQ==", + "dev": true, + "license": "MIT", + "peerDependencies": { + "graphql": "^0.8.0 || ^0.9.0 || ^0.10.0 || ^0.11.0 || ^0.12.0 || ^0.13.0 || ^14.0.0 || ^15.0.0 || ^16.0.0 || ^17.0.0" + } + }, + "node_modules/@wry/context": { + "version": "0.6.1", + "resolved": "https://registry.npmjs.org/@wry/context/-/context-0.6.1.tgz", + "integrity": "sha512-LOmVnY1iTU2D8tv4Xf6MVMZZ+juIJ87Kt/plMijjN20NMAXGmH4u8bS1t0uT74cZ5gwpocYueV58YwyI8y+GKw==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/@wry/equality": { + "version": "0.5.7", + "resolved": "https://registry.npmjs.org/@wry/equality/-/equality-0.5.7.tgz", + "integrity": "sha512-BRFORjsTuQv5gxcXsuDXx6oGRhuVsEGwZy6LOzRRfgu+eSfxbhUQ9L9YtSEIuIjY/o7g3iWFjrc5eSY1GXP2Dw==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/@wry/trie": { + "version": "0.3.2", + "resolved": "https://registry.npmjs.org/@wry/trie/-/trie-0.3.2.tgz", + "integrity": "sha512-yRTyhWSls2OY/pYLfwff867r8ekooZ4UI+/gxot5Wj8EFwSf2rG+n+Mo/6LoLQm1TKA4GRj2+LCpbfS937dClQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/graphql": { + "version": "16.14.1", + "resolved": "https://registry.npmjs.org/graphql/-/graphql-16.14.1.tgz", + "integrity": "sha512-cQOsSMS/IrDz82PVyRDvf/Q1F/bRbBVjJlh+xYOkI1qw2bWRvWGiWc+m2O0d6l4Bt1fyY+8kzJ8JFWGJqNeDBg==", + "dev": true, + "license": "MIT", + "peer": true, + "engines": { + "node": "^12.22.0 || ^14.16.0 || ^16.0.0 || >=17.0.0" + } + }, + "node_modules/graphql-tag": { + "version": "2.12.6", + "resolved": "https://registry.npmjs.org/graphql-tag/-/graphql-tag-2.12.6.tgz", + "integrity": "sha512-FdSNcu2QQcWnM2VNvSCCDCVS5PpPqpzgFT8+GXzqJuoDd0CBncxCY278u4mhRO7tMgo2JjgJA5aZ+nWSQ/Z+xg==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.1.0" + }, + "engines": { + "node": ">=10" + }, + "peerDependencies": { + "graphql": "^0.9.0 || ^0.10.0 || ^0.11.0 || ^0.12.0 || ^0.13.0 || ^14.0.0 || ^15.0.0 || ^16.0.0" + } + }, + "node_modules/hoist-non-react-statics": { + "version": "3.3.2", + "resolved": "https://registry.npmjs.org/hoist-non-react-statics/-/hoist-non-react-statics-3.3.2.tgz", + "integrity": "sha512-/gGivxi8JPKWNm/W0jSmzcMPpfpPLc3dY/6GxhX2hQ9iGj3aDfklV4ET7NjKpSinLpJ5vafa9iiGIEZg10SfBw==", + "dev": true, + "license": "BSD-3-Clause", + "dependencies": { + "react-is": "^16.7.0" + } + }, + "node_modules/js-tokens": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-4.0.0.tgz", + "integrity": "sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/loose-envify": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/loose-envify/-/loose-envify-1.4.0.tgz", + "integrity": "sha512-lyuxPGr/Wfhrlem2CL/UcnUc1zcqKAImBDzukY7Y5F/yQiNdko6+fRLevlw1HgMySw7f611UIY408EtxRSoK3Q==", + "dev": true, + "license": "MIT", + "dependencies": { + "js-tokens": "^3.0.0 || ^4.0.0" + }, + "bin": { + "loose-envify": "cli.js" + } + }, + "node_modules/object-assign": { + "version": "4.1.1", + "resolved": "https://registry.npmjs.org/object-assign/-/object-assign-4.1.1.tgz", + "integrity": "sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/optimism": { + "version": "0.16.2", + "resolved": "https://registry.npmjs.org/optimism/-/optimism-0.16.2.tgz", + "integrity": "sha512-zWNbgWj+3vLEjZNIh/okkY2EUfX+vB9TJopzIZwT1xxaMqC5hRLLraePod4c5n4He08xuXNH+zhKFFCu390wiQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "@wry/context": "^0.7.0", + "@wry/trie": "^0.3.0" + } + }, + "node_modules/optimism/node_modules/@wry/context": { + "version": "0.7.4", + "resolved": "https://registry.npmjs.org/@wry/context/-/context-0.7.4.tgz", + "integrity": "sha512-jmT7Sb4ZQWI5iyu3lobQxICu2nC/vbUhP0vIdd6tHC9PTfenmRmuIFqktc6GH9cgi+ZHnsLWPvfSvc4DrYmKiQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/prop-types": { + "version": "15.8.1", + "resolved": "https://registry.npmjs.org/prop-types/-/prop-types-15.8.1.tgz", + "integrity": "sha512-oj87CgZICdulUohogVAR7AjlC0327U4el4L6eAvOqCeudMDVU0NThNaV+b9Df4dXgSP1gXMTnPdhfe/2qDH5cg==", + "dev": true, + "license": "MIT", + "dependencies": { + "loose-envify": "^1.4.0", + "object-assign": "^4.1.1", + "react-is": "^16.13.1" + } + }, + "node_modules/react-is": { + "version": "16.13.1", + "resolved": "https://registry.npmjs.org/react-is/-/react-is-16.13.1.tgz", + "integrity": "sha512-24e6ynE2H+OKt4kqsOvNd8kBpV65zoxbA4BVsEOB3ARVWQki/DHzaUoC5KuON/BiccDaCCTZBuOcfZs70kR8bQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/symbol-observable": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/symbol-observable/-/symbol-observable-4.0.0.tgz", + "integrity": "sha512-b19dMThMV4HVFynSAM1++gBHAbk2Tc/osgLIBZMKsyqh34jb2e8Os7T6ZW/Bt3pJFdBTd2JwAnAAEQV7rSNvcQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10" + } + }, + "node_modules/ts-invariant": { + "version": "0.10.3", + "resolved": "https://registry.npmjs.org/ts-invariant/-/ts-invariant-0.10.3.tgz", + "integrity": "sha512-uivwYcQaxAucv1CzRp2n/QdYPo4ILf9VXgH19zEIjFx2EJufV16P0JtJVpYHy89DItG6Kwj2oIUjrcK5au+4tQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.1.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/tslib": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", + "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", + "dev": true, + "license": "0BSD" + }, + "node_modules/zen-observable": { + "version": "0.8.15", + "resolved": "https://registry.npmjs.org/zen-observable/-/zen-observable-0.8.15.tgz", + "integrity": "sha512-PQ2PC7R9rslx84ndNBZB/Dkv8V8fZEpk83RLgXtYd0fwUgEjseMn1Dgajh2x6S8QbZAFa9p2qVCEuYZNgve0dQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/zen-observable-ts": { + "version": "1.2.5", + "resolved": "https://registry.npmjs.org/zen-observable-ts/-/zen-observable-ts-1.2.5.tgz", + "integrity": "sha512-QZWQekv6iB72Naeake9hS1KxHlotfRpe+WGNbNx5/ta+R3DNjVO2bswf63gXlWDcs+EMd7XY8HfVQyP1X6T4Zg==", + "dev": true, + "license": "MIT", + "dependencies": { + "zen-observable": "0.8.15" + } + } + } +} diff --git a/package.json b/package.json new file mode 100644 index 0000000000..10f0ca32dc --- /dev/null +++ b/package.json @@ -0,0 +1,5 @@ +{ + "devDependencies": { + "@apollo/client": "^3.6.9" + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts new file mode 100644 index 0000000000..f357d409db --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts @@ -0,0 +1,61 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + parseOrganizationId, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_SELECT_ORGANIZATION'); +logger.level = 'debug'; + +const { organizationService } = components; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['POST']); + const organizationId = parseOrganizationId(req.query.id); + const organization = + await organizationService.selectCurrentOrganization(organizationId); + + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + id: organization.id, + name: organization.name, + identifier: organization.identifier, + description: organization.description, + isCurrent: Boolean(organization.isCurrent), + createdAt: organization.createdAt, + updatedAt: organization.updatedAt, + }, + projectId: 0, + apiType: ApiType.SELECT_ORGANIZATION, + startTime, + requestPayload: { id: organizationId }, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: 0, + apiType: ApiType.SELECT_ORGANIZATION, + requestPayload: { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/current.ts b/wren-ui/src/pages/api/v1/organizations/current.ts new file mode 100644 index 0000000000..612d71ee62 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/current.ts @@ -0,0 +1,73 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectName, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_CURRENT_ORGANIZATION'); +logger.level = 'debug'; + +const { organizationService } = components; + +const serializeOrganization = (organization) => + organization + ? { + id: organization.id, + name: organization.name, + identifier: organization.identifier, + description: organization.description, + isCurrent: Boolean(organization.isCurrent), + createdAt: organization.createdAt, + updatedAt: organization.updatedAt, + } + : null; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET']); + const [currentOrganization, organizations, currentProjectName] = + await Promise.all([ + organizationService.getCurrentOrganization(), + organizationService.listOrganizations(), + getCurrentProjectName(), + ]); + + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + currentOrganization: serializeOrganization(currentOrganization), + organizations: organizations.map(serializeOrganization), + currentProjectName, + }, + projectId: 0, + apiType: ApiType.GET_CURRENT_ORGANIZATION, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: 0, + apiType: ApiType.GET_CURRENT_ORGANIZATION, + requestPayload: {}, + headers: req.headers as Record, + startTime, + logger, + }); + } +} From 043d304409b47c47dbf58eebd7e465d0e03a8655 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 19:35:36 +0530 Subject: [PATCH 0092/1087] APi Utils error handling --- wren-ui/src/apollo/server/utils/apiUtils.ts | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index a30b1ab372..248d0aa14f 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -12,7 +12,13 @@ import { TextBasedAnswerStatus, } from '@/apollo/server/models/adaptor'; -const { apiHistoryRepository } = components; +const getApiHistoryRepository = () => { + const repository = components?.apiHistoryRepository; + if (!repository) { + throw new Error('API history repository is not initialized'); + } + return repository; +}; export const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes @@ -195,7 +201,7 @@ export const respondWith = async ({ }) => { const durationMs = startTime ? Date.now() - startTime : undefined; const responseId = uuidv4(); - await apiHistoryRepository.createOne({ + await getApiHistoryRepository().createOne({ id: responseId, projectId, apiType, @@ -238,7 +244,7 @@ export const respondWithSimple = async ({ }) => { const durationMs = startTime ? Date.now() - startTime : undefined; const responseId = uuidv4(); - await apiHistoryRepository.createOne({ + await getApiHistoryRepository().createOne({ id: responseId, projectId, apiType, From 48b1e048e8ef52a53c51f9e79f66b5e3212b86da Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 19:40:15 +0530 Subject: [PATCH 0093/1087] Modified graphql --- wren-ui/src/pages/api/graphql.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index 569329e5c4..524d4d2808 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -13,7 +13,7 @@ import { GeneralErrorCodes, } from '@/apollo/server/utils/error'; import { TelemetryEvent } from '@/apollo/server/telemetry/telemetry'; -import { components } from '@/common'; +import { components, initComponents } from '@/common'; const serverConfig = getConfig(); const logger = getLogger('APOLLO'); @@ -28,6 +28,7 @@ export const config: PageConfig = { }; const bootstrapServer = async () => { + const componentGraph = components ?? initComponents(); const { telemetry, @@ -70,7 +71,7 @@ const bootstrapServer = async () => { projectRecommendQuestionBackgroundTracker, threadRecommendQuestionBackgroundTracker, dashboardCacheBackgroundTracker, - } = components; + } = componentGraph; const modelService = new ModelService({ projectService, From cfc643c5ca1f505970254e636a1a35572ce5bd86 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 19:47:42 +0530 Subject: [PATCH 0094/1087] Modified Components --- .../src/apollo/server/repositories/index.ts | 2 ++ wren-ui/src/apollo/server/services/index.ts | 3 --- wren-ui/src/apollo/server/types/context.ts | 12 ------------ wren-ui/src/common.ts | 18 ------------------ 4 files changed, 2 insertions(+), 33 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/index.ts b/wren-ui/src/apollo/server/repositories/index.ts index 0dd1cc5905..753e3f3c36 100644 --- a/wren-ui/src/apollo/server/repositories/index.ts +++ b/wren-ui/src/apollo/server/repositories/index.ts @@ -19,3 +19,5 @@ export * from './askingTaskRepository'; export * from './instructionRepository'; export * from './apiHistoryRepository'; export * from './dashboardItemRefreshJobRepository'; +export * from './rbacRepository'; +export * from './organizationRepository'; diff --git a/wren-ui/src/apollo/server/services/index.ts b/wren-ui/src/apollo/server/services/index.ts index e5b7f161cf..0d052c814a 100644 --- a/wren-ui/src/apollo/server/services/index.ts +++ b/wren-ui/src/apollo/server/services/index.ts @@ -8,8 +8,5 @@ export * from './metadataService'; export * from './dashboardService'; export * from './askingTaskTracker'; export * from './instructionService'; -<<<<<<< HEAD export * from './rbacService'; -======= export * from './organizationService'; ->>>>>>> 52a339fb5 (Add organization management feature) diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index ce7b0b5ba2..4aed639176 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -20,13 +20,10 @@ import { IInstructionRepository, IApiHistoryRepository, IDashboardItemRefreshJobRepository, -<<<<<<< HEAD IRoleRepository, IUserRepository, IUserRoleRepository, -======= IOrganizationRepository, ->>>>>>> 52a339fb5 (Add organization management feature) } from '@server/repositories'; import { IQueryService, @@ -37,11 +34,8 @@ import { IProjectService, IDashboardService, IInstructionService, -<<<<<<< HEAD IRbacService, -======= IOrganizationService, ->>>>>>> 52a339fb5 (Add organization management feature) } from '@server/services'; import { ITelemetry } from '@server/telemetry/telemetry'; import { @@ -71,11 +65,8 @@ export interface IContext { dashboardService: IDashboardService; sqlPairService: ISqlPairService; instructionService: IInstructionService; -<<<<<<< HEAD rbacService: IRbacService; -======= organizationService: IOrganizationService; ->>>>>>> 52a339fb5 (Add organization management feature) // repository projectRepository: IProjectRepository; @@ -93,13 +84,10 @@ export interface IContext { instructionRepository: IInstructionRepository; apiHistoryRepository: IApiHistoryRepository; dashboardItemRefreshJobRepository: IDashboardItemRefreshJobRepository; -<<<<<<< HEAD roleRepository: IRoleRepository; userRepository: IUserRepository; userRoleRepository: IUserRoleRepository; -======= organizationRepository: IOrganizationRepository; ->>>>>>> 52a339fb5 (Add organization management feature) // background trackers projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index b7cd09b6ca..e9faee4862 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -19,13 +19,10 @@ import { InstructionRepository, ApiHistoryRepository, DashboardItemRefreshJobRepository, -<<<<<<< HEAD RoleRepository, UserRepository, UserRoleRepository, -======= OrganizationRepository, ->>>>>>> 52a339fb5 (Add organization management feature) } from '@server/repositories'; import { WrenEngineAdaptor, @@ -42,11 +39,8 @@ import { DashboardService, AskingTaskTracker, InstructionService, -<<<<<<< HEAD RbacService, -======= OrganizationService, ->>>>>>> 52a339fb5 (Add organization management feature) } from '@server/services'; import { PostHogTelemetry } from './apollo/server/telemetry/telemetry'; import { @@ -162,13 +156,10 @@ export const initComponents = () => { const apiHistoryRepository = new ApiHistoryRepository(knex); const dashboardItemRefreshJobRepository = new DashboardItemRefreshJobRepository(knex); -<<<<<<< HEAD const roleRepository = new RoleRepository(knex); const userRepository = new UserRepository(knex); const userRoleRepository = new UserRoleRepository(knex); -======= const organizationRepository = new OrganizationRepository(knex); ->>>>>>> 52a339fb5 (Add organization management feature) // adaptors const wrenEngineAdaptor = new WrenEngineAdaptor({ @@ -274,15 +265,12 @@ export const initComponents = () => { instructionRepository, wrenAIAdaptor, }); -<<<<<<< HEAD const rbacService = new RbacService({ roleRepository, userRepository, userRoleRepository, -======= const organizationService = new OrganizationService({ organizationRepository, ->>>>>>> 52a339fb5 (Add organization management feature) }); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ @@ -317,13 +305,10 @@ export const initComponents = () => { apiHistoryRepository, instructionRepository, dashboardItemRefreshJobRepository, -<<<<<<< HEAD roleRepository, userRepository, userRoleRepository, -======= organizationRepository, ->>>>>>> 52a339fb5 (Add organization management feature) // adaptors wrenEngineAdaptor, @@ -340,11 +325,8 @@ export const initComponents = () => { dashboardService, sqlPairService, instructionService, -<<<<<<< HEAD rbacService, -======= organizationService, ->>>>>>> 52a339fb5 (Add organization management feature) askingTaskTracker, // background trackers From cb6626e8e1ab1bd6af87beae1f4fbe2157b73eec Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 19:58:40 +0530 Subject: [PATCH 0095/1087] Updated error handling --- wren-ui/src/common.ts | 1 + wren-ui/src/pages/api/graphql.ts | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index e9faee4862..a22dd08bdb 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -269,6 +269,7 @@ export const initComponents = () => { roleRepository, userRepository, userRoleRepository, + }); const organizationService = new OrganizationService({ organizationRepository, }); diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index 524d4d2808..bb5cd2e5f8 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -51,6 +51,7 @@ const bootstrapServer = async () => { roleRepository, userRepository, userRoleRepository, + organizationRepository, // adaptors wrenEngineAdaptor, ibisAdaptor, @@ -67,6 +68,7 @@ const bootstrapServer = async () => { instructionService, rbacService, + organizationService, // background trackers projectRecommendQuestionBackgroundTracker, threadRecommendQuestionBackgroundTracker, @@ -148,6 +150,7 @@ const bootstrapServer = async () => { sqlPairService, instructionService, rbacService, + organizationService, // repository projectRepository, modelRepository, @@ -167,6 +170,7 @@ const bootstrapServer = async () => { roleRepository, userRepository, userRoleRepository, + organizationRepository, // background trackers projectRecommendQuestionBackgroundTracker, threadRecommendQuestionBackgroundTracker, From 0d5b071a9cb130e1e1ceb196d1c28bc0e2737355 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 20:17:03 +0530 Subject: [PATCH 0096/1087] updated organization creation --- wren-ui/src/apollo/server/middlewares/organizationApi.ts | 8 ++++++-- wren-ui/src/apollo/server/utils/apiUtils.ts | 8 ++++++-- wren-ui/src/pages/api/v1/organizations/[id]/select.ts | 8 ++++++-- wren-ui/src/pages/api/v1/organizations/current.ts | 8 ++++++-- wren-ui/src/pages/api/v1/organizations/index.ts | 8 ++++++-- 5 files changed, 30 insertions(+), 10 deletions(-) diff --git a/wren-ui/src/apollo/server/middlewares/organizationApi.ts b/wren-ui/src/apollo/server/middlewares/organizationApi.ts index e87a0ab8df..f9ee394d91 100644 --- a/wren-ui/src/apollo/server/middlewares/organizationApi.ts +++ b/wren-ui/src/apollo/server/middlewares/organizationApi.ts @@ -1,8 +1,11 @@ import { NextApiRequest } from 'next'; import { ApiError } from '../utils/apiUtils'; -import { components } from '@/common'; -const { projectService } = components; +const getProjectService = () => { + const { components, initComponents } = require('@/common'); + const componentGraph = components ?? initComponents(); + return componentGraph.projectService; +}; export const assertAllowedMethods = ( req: NextApiRequest, @@ -24,6 +27,7 @@ export const parseOrganizationId = (value: string | string[] | undefined) => { export const getCurrentProjectName = async () => { try { + const projectService = getProjectService(); const project = await projectService.getCurrentProject(); return project?.displayName || 'Default Project'; } catch { diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index 248d0aa14f..6e0b87866f 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -2,7 +2,6 @@ import { NextApiResponse } from 'next'; import { v4 as uuidv4 } from 'uuid'; import { ApiType, ApiHistory } from '@server/repositories/apiHistoryRepository'; import * as Errors from '@server/utils/error'; -import { components } from '@/common'; import { AskResult, AskResultStatus, @@ -12,8 +11,13 @@ import { TextBasedAnswerStatus, } from '@/apollo/server/models/adaptor'; +const getComponentGraph = () => { + const { components, initComponents } = require('@/common'); + return components ?? initComponents(); +}; + const getApiHistoryRepository = () => { - const repository = components?.apiHistoryRepository; + const repository = getComponentGraph().apiHistoryRepository; if (!repository) { throw new Error('API history repository is not initialized'); } diff --git a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts index f357d409db..52838ac957 100644 --- a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts +++ b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts @@ -1,5 +1,4 @@ import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; import { ApiType } from '@server/repositories/apiHistoryRepository'; import { handleApiError, @@ -14,7 +13,11 @@ import { const logger = getLogger('API_SELECT_ORGANIZATION'); logger.level = 'debug'; -const { organizationService } = components; +const getOrganizationService = () => { + const { components, initComponents } = require('@/common'); + const componentGraph = components ?? initComponents(); + return componentGraph.organizationService; +}; export default async function handler( req: NextApiRequest, @@ -24,6 +27,7 @@ export default async function handler( try { assertAllowedMethods(req, ['POST']); + const organizationService = getOrganizationService(); const organizationId = parseOrganizationId(req.query.id); const organization = await organizationService.selectCurrentOrganization(organizationId); diff --git a/wren-ui/src/pages/api/v1/organizations/current.ts b/wren-ui/src/pages/api/v1/organizations/current.ts index 612d71ee62..3a95ba763f 100644 --- a/wren-ui/src/pages/api/v1/organizations/current.ts +++ b/wren-ui/src/pages/api/v1/organizations/current.ts @@ -1,5 +1,4 @@ import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; import { ApiType } from '@server/repositories/apiHistoryRepository'; import { handleApiError, @@ -14,7 +13,11 @@ import { const logger = getLogger('API_CURRENT_ORGANIZATION'); logger.level = 'debug'; -const { organizationService } = components; +const getOrganizationService = () => { + const { components, initComponents } = require('@/common'); + const componentGraph = components ?? initComponents(); + return componentGraph.organizationService; +}; const serializeOrganization = (organization) => organization @@ -37,6 +40,7 @@ export default async function handler( try { assertAllowedMethods(req, ['GET']); + const organizationService = getOrganizationService(); const [currentOrganization, organizations, currentProjectName] = await Promise.all([ organizationService.getCurrentOrganization(), diff --git a/wren-ui/src/pages/api/v1/organizations/index.ts b/wren-ui/src/pages/api/v1/organizations/index.ts index 643ee2432d..6ae798025d 100644 --- a/wren-ui/src/pages/api/v1/organizations/index.ts +++ b/wren-ui/src/pages/api/v1/organizations/index.ts @@ -1,5 +1,4 @@ import { NextApiRequest, NextApiResponse } from 'next'; -import { components } from '@/common'; import { ApiType } from '@server/repositories/apiHistoryRepository'; import { handleApiError, @@ -14,7 +13,11 @@ import { const logger = getLogger('API_ORGANIZATIONS'); logger.level = 'debug'; -const { organizationService } = components; +const getOrganizationService = () => { + const { components, initComponents } = require('@/common'); + const componentGraph = components ?? initComponents(); + return componentGraph.organizationService; +}; const serializeOrganization = (organization) => ({ id: organization.id, @@ -34,6 +37,7 @@ export default async function handler( try { assertAllowedMethods(req, ['GET', 'POST']); + const organizationService = getOrganizationService(); if (req.method === 'GET') { const organizations = await organizationService.listOrganizations(); From 1799918849ce60d413748b2ec77c51a26dc77ec1 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 20:25:36 +0530 Subject: [PATCH 0097/1087] Updated API handling --- .../server/middlewares/organizationApi.ts | 25 +++++++++++++------ wren-ui/src/apollo/server/utils/apiUtils.ts | 12 +++++++-- .../pages/api/v1/organizations/[id]/select.ts | 7 ++++-- .../src/pages/api/v1/organizations/current.ts | 7 ++++-- .../src/pages/api/v1/organizations/index.ts | 9 ++++--- 5 files changed, 44 insertions(+), 16 deletions(-) diff --git a/wren-ui/src/apollo/server/middlewares/organizationApi.ts b/wren-ui/src/apollo/server/middlewares/organizationApi.ts index f9ee394d91..b7ee5cb122 100644 --- a/wren-ui/src/apollo/server/middlewares/organizationApi.ts +++ b/wren-ui/src/apollo/server/middlewares/organizationApi.ts @@ -7,6 +7,22 @@ const getProjectService = () => { return componentGraph.projectService; }; +export const getCurrentProjectContext = async () => { + try { + const projectService = getProjectService(); + const project = await projectService.getCurrentProject(); + return { + id: project?.id ?? null, + displayName: project?.displayName || 'Default Project', + }; + } catch { + return { + id: null, + displayName: 'Default Project', + }; + } +}; + export const assertAllowedMethods = ( req: NextApiRequest, methods: string[], @@ -26,11 +42,6 @@ export const parseOrganizationId = (value: string | string[] | undefined) => { }; export const getCurrentProjectName = async () => { - try { - const projectService = getProjectService(); - const project = await projectService.getCurrentProject(); - return project?.displayName || 'Default Project'; - } catch { - return 'Default Project'; - } + const project = await getCurrentProjectContext(); + return project.displayName; }; diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index 6e0b87866f..c063c96d6a 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -24,6 +24,14 @@ const getApiHistoryRepository = () => { return repository; }; +const persistApiHistory = async (payload: Record) => { + try { + await getApiHistoryRepository().createOne(payload); + } catch (error) { + console.error('Failed to persist API history:', error); + } +}; + export const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes export const isAskResultFinished = (result: AskResult) => { @@ -205,7 +213,7 @@ export const respondWith = async ({ }) => { const durationMs = startTime ? Date.now() - startTime : undefined; const responseId = uuidv4(); - await getApiHistoryRepository().createOne({ + await persistApiHistory({ id: responseId, projectId, apiType, @@ -248,7 +256,7 @@ export const respondWithSimple = async ({ }) => { const durationMs = startTime ? Date.now() - startTime : undefined; const responseId = uuidv4(); - await getApiHistoryRepository().createOne({ + await persistApiHistory({ id: responseId, projectId, apiType, diff --git a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts index 52838ac957..7ae26c9669 100644 --- a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts +++ b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts @@ -7,6 +7,7 @@ import { import { getLogger } from '@server/utils'; import { assertAllowedMethods, + getCurrentProjectContext, parseOrganizationId, } from '@/apollo/server/middlewares/organizationApi'; @@ -28,6 +29,8 @@ export default async function handler( try { assertAllowedMethods(req, ['POST']); const organizationService = getOrganizationService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; const organizationId = parseOrganizationId(req.query.id); const organization = await organizationService.selectCurrentOrganization(organizationId); @@ -44,7 +47,7 @@ export default async function handler( createdAt: organization.createdAt, updatedAt: organization.updatedAt, }, - projectId: 0, + projectId, apiType: ApiType.SELECT_ORGANIZATION, startTime, requestPayload: { id: organizationId }, @@ -54,7 +57,7 @@ export default async function handler( await handleApiError({ error, res, - projectId: 0, + projectId: (await getCurrentProjectContext()).id ?? 0, apiType: ApiType.SELECT_ORGANIZATION, requestPayload: { id: req.query.id }, headers: req.headers as Record, diff --git a/wren-ui/src/pages/api/v1/organizations/current.ts b/wren-ui/src/pages/api/v1/organizations/current.ts index 3a95ba763f..93303349b8 100644 --- a/wren-ui/src/pages/api/v1/organizations/current.ts +++ b/wren-ui/src/pages/api/v1/organizations/current.ts @@ -7,6 +7,7 @@ import { import { getLogger } from '@server/utils'; import { assertAllowedMethods, + getCurrentProjectContext, getCurrentProjectName, } from '@/apollo/server/middlewares/organizationApi'; @@ -41,6 +42,8 @@ export default async function handler( try { assertAllowedMethods(req, ['GET']); const organizationService = getOrganizationService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; const [currentOrganization, organizations, currentProjectName] = await Promise.all([ organizationService.getCurrentOrganization(), @@ -56,7 +59,7 @@ export default async function handler( organizations: organizations.map(serializeOrganization), currentProjectName, }, - projectId: 0, + projectId, apiType: ApiType.GET_CURRENT_ORGANIZATION, startTime, requestPayload: {}, @@ -66,7 +69,7 @@ export default async function handler( await handleApiError({ error, res, - projectId: 0, + projectId: (await getCurrentProjectContext()).id ?? 0, apiType: ApiType.GET_CURRENT_ORGANIZATION, requestPayload: {}, headers: req.headers as Record, diff --git a/wren-ui/src/pages/api/v1/organizations/index.ts b/wren-ui/src/pages/api/v1/organizations/index.ts index 6ae798025d..f9bafdd358 100644 --- a/wren-ui/src/pages/api/v1/organizations/index.ts +++ b/wren-ui/src/pages/api/v1/organizations/index.ts @@ -7,6 +7,7 @@ import { import { getLogger } from '@server/utils'; import { assertAllowedMethods, + getCurrentProjectContext, getCurrentProjectName, } from '@/apollo/server/middlewares/organizationApi'; @@ -38,6 +39,8 @@ export default async function handler( try { assertAllowedMethods(req, ['GET', 'POST']); const organizationService = getOrganizationService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; if (req.method === 'GET') { const organizations = await organizationService.listOrganizations(); @@ -48,7 +51,7 @@ export default async function handler( organizations: organizations.map(serializeOrganization), currentProjectName: await getCurrentProjectName(), }, - projectId: 0, + projectId, apiType: ApiType.GET_ORGANIZATIONS, startTime, requestPayload: {}, @@ -62,7 +65,7 @@ export default async function handler( res, statusCode: 201, responsePayload: serializeOrganization(organization), - projectId: 0, + projectId, apiType: ApiType.CREATE_ORGANIZATION, startTime, requestPayload: req.body, @@ -72,7 +75,7 @@ export default async function handler( await handleApiError({ error, res, - projectId: 0, + projectId: (await getCurrentProjectContext()).id ?? 0, apiType: req.method === 'GET' ? ApiType.GET_ORGANIZATIONS From 8d750864b251f5bad10dda584bec561941aa40d0 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 20:32:47 +0530 Subject: [PATCH 0098/1087] Updated API handling --- wren-ui/src/apollo/server/middlewares/organizationApi.ts | 7 +++++-- wren-ui/src/apollo/server/utils/apiUtils.ts | 8 ++++++-- wren-ui/src/pages/api/v1/organizations/current.ts | 7 +++++-- wren-ui/src/pages/api/v1/organizations/index.ts | 7 +++++-- 4 files changed, 21 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/apollo/server/middlewares/organizationApi.ts b/wren-ui/src/apollo/server/middlewares/organizationApi.ts index b7ee5cb122..aa7f6e34ef 100644 --- a/wren-ui/src/apollo/server/middlewares/organizationApi.ts +++ b/wren-ui/src/apollo/server/middlewares/organizationApi.ts @@ -2,8 +2,11 @@ import { NextApiRequest } from 'next'; import { ApiError } from '../utils/apiUtils'; const getProjectService = () => { - const { components, initComponents } = require('@/common'); - const componentGraph = components ?? initComponents(); + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } return componentGraph.projectService; }; diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index c063c96d6a..8cd1169f01 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -12,8 +12,12 @@ import { } from '@/apollo/server/models/adaptor'; const getComponentGraph = () => { - const { components, initComponents } = require('@/common'); - return components ?? initComponents(); + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph; }; const getApiHistoryRepository = () => { diff --git a/wren-ui/src/pages/api/v1/organizations/current.ts b/wren-ui/src/pages/api/v1/organizations/current.ts index 93303349b8..5ed62df666 100644 --- a/wren-ui/src/pages/api/v1/organizations/current.ts +++ b/wren-ui/src/pages/api/v1/organizations/current.ts @@ -15,8 +15,11 @@ const logger = getLogger('API_CURRENT_ORGANIZATION'); logger.level = 'debug'; const getOrganizationService = () => { - const { components, initComponents } = require('@/common'); - const componentGraph = components ?? initComponents(); + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } return componentGraph.organizationService; }; diff --git a/wren-ui/src/pages/api/v1/organizations/index.ts b/wren-ui/src/pages/api/v1/organizations/index.ts index f9bafdd358..5973b70da1 100644 --- a/wren-ui/src/pages/api/v1/organizations/index.ts +++ b/wren-ui/src/pages/api/v1/organizations/index.ts @@ -15,8 +15,11 @@ const logger = getLogger('API_ORGANIZATIONS'); logger.level = 'debug'; const getOrganizationService = () => { - const { components, initComponents } = require('@/common'); - const componentGraph = components ?? initComponents(); + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } return componentGraph.organizationService; }; From b95f9a9a1153a6e3e3228d4bd44836b8a5b56f43 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 20:33:17 +0530 Subject: [PATCH 0099/1087] Updated API handling --- wren-ui/src/pages/api/v1/organizations/[id]/select.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts index 7ae26c9669..f97abe3d78 100644 --- a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts +++ b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts @@ -15,8 +15,11 @@ const logger = getLogger('API_SELECT_ORGANIZATION'); logger.level = 'debug'; const getOrganizationService = () => { - const { components, initComponents } = require('@/common'); - const componentGraph = components ?? initComponents(); + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } return componentGraph.organizationService; }; From 15564b91d4d998a57335e4c65e5e8ca0daeca6fe Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 3 Jun 2026 20:56:26 +0530 Subject: [PATCH 0100/1087] Add organization general settings page --- .../repositories/apiHistoryRepository.ts | 1 + .../server/services/organizationService.ts | 36 +++ .../src/components/OrganizationSwitcher.tsx | 22 ++ .../src/pages/api/v1/organizations/current.ts | 27 +- wren-ui/src/pages/organization/general.tsx | 269 ++++++++++++++++++ wren-ui/src/pages/organization/index.tsx | 14 + wren-ui/src/utils/enum/path.ts | 2 + 7 files changed, 368 insertions(+), 3 deletions(-) create mode 100644 wren-ui/src/pages/organization/general.tsx create mode 100644 wren-ui/src/pages/organization/index.tsx diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index d941a9fb8f..a07e7fbaa0 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -28,6 +28,7 @@ export enum ApiType { GET_ORGANIZATIONS = 'GET_ORGANIZATIONS', GET_CURRENT_ORGANIZATION = 'GET_CURRENT_ORGANIZATION', CREATE_ORGANIZATION = 'CREATE_ORGANIZATION', + UPDATE_CURRENT_ORGANIZATION = 'UPDATE_CURRENT_ORGANIZATION', SELECT_ORGANIZATION = 'SELECT_ORGANIZATION', } diff --git a/wren-ui/src/apollo/server/services/organizationService.ts b/wren-ui/src/apollo/server/services/organizationService.ts index 293322b95e..5c463b4e3e 100644 --- a/wren-ui/src/apollo/server/services/organizationService.ts +++ b/wren-ui/src/apollo/server/services/organizationService.ts @@ -10,11 +10,16 @@ export interface CreateOrganizationData { description?: string; } +export interface UpdateOrganizationData { + name: string; +} + export interface IOrganizationService { listOrganizations: () => Promise; getCurrentOrganization: () => Promise; createOrganization: (data: CreateOrganizationData) => Promise; selectCurrentOrganization: (id: number) => Promise; + updateCurrentOrganization: (data: UpdateOrganizationData) => Promise; } const NAME_MAX_LENGTH = 64; @@ -79,6 +84,24 @@ export const validateOrganizationPayload = ( }; }; +export const validateOrganizationUpdatePayload = ( + data: UpdateOrganizationData, +): UpdateOrganizationData => { + const name = normalizeOrganizationName(data.name || ''); + + if (!name) { + throw new ApiError('Organization name is required', 400); + } + if (name.length > NAME_MAX_LENGTH) { + throw new ApiError( + `Organization name is too long (max ${NAME_MAX_LENGTH} characters)`, + 400, + ); + } + + return { name }; +}; + export class OrganizationService implements IOrganizationService { private organizationRepository: IOrganizationRepository; @@ -139,4 +162,17 @@ export class OrganizationService implements IOrganizationService { return await this.organizationRepository.setCurrentOrganization(id); } + + public async updateCurrentOrganization(data: UpdateOrganizationData) { + const currentOrganization = + await this.organizationRepository.getCurrentOrganization(); + if (!currentOrganization) { + throw new ApiError('Current organization not found', 404); + } + + const payload = validateOrganizationUpdatePayload(data); + return await this.organizationRepository.updateOne(currentOrganization.id, { + name: payload.name, + }); + } } diff --git a/wren-ui/src/components/OrganizationSwitcher.tsx b/wren-ui/src/components/OrganizationSwitcher.tsx index e54ee19314..5e7e62b53f 100644 --- a/wren-ui/src/components/OrganizationSwitcher.tsx +++ b/wren-ui/src/components/OrganizationSwitcher.tsx @@ -15,6 +15,9 @@ import { import styled from 'styled-components'; import PlusOutlined from '@ant-design/icons/PlusOutlined'; import DownOutlined from '@ant-design/icons/DownOutlined'; +import SettingOutlined from '@ant-design/icons/SettingOutlined'; +import { useRouter } from 'next/router'; +import { Path } from '@/utils/enum'; interface OrganizationRecord { id: number; @@ -90,6 +93,11 @@ const OverlayMenu = styled(Menu)` } `; +const OverlayFooter = styled.div` + border-top: 1px solid var(--gray-4); + padding: 8px; +`; + const MenuRow = styled.button<{ $active?: boolean }>` width: 100%; display: flex; @@ -116,6 +124,7 @@ const getBadgeText = (name?: string) => .toUpperCase(); export default function OrganizationSwitcher() { + const router = useRouter(); const [form] = Form.useForm(); const [loading, setLoading] = useState(true); const [saving, setSaving] = useState(false); @@ -175,6 +184,7 @@ export default function OrganizationSwitcher() { setVisible(false); form.resetFields(); await loadOrganizations(); + await router.push(Path.OrganizationGeneral); } catch (error: any) { if (error?.errorFields) { return; @@ -238,6 +248,18 @@ export default function OrganizationSwitcher() { ))} + {currentOrganization && ( + + + + )} ); diff --git a/wren-ui/src/pages/api/v1/organizations/current.ts b/wren-ui/src/pages/api/v1/organizations/current.ts index 5ed62df666..cf8a39a103 100644 --- a/wren-ui/src/pages/api/v1/organizations/current.ts +++ b/wren-ui/src/pages/api/v1/organizations/current.ts @@ -43,10 +43,28 @@ export default async function handler( const startTime = Date.now(); try { - assertAllowedMethods(req, ['GET']); + assertAllowedMethods(req, ['GET', 'PUT']); const organizationService = getOrganizationService(); const projectContext = await getCurrentProjectContext(); const projectId = projectContext.id ?? 0; + + if (req.method === 'PUT') { + const organization = await organizationService.updateCurrentOrganization( + req.body, + ); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: serializeOrganization(organization), + projectId, + apiType: ApiType.UPDATE_CURRENT_ORGANIZATION, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + return; + } + const [currentOrganization, organizations, currentProjectName] = await Promise.all([ organizationService.getCurrentOrganization(), @@ -73,8 +91,11 @@ export default async function handler( error, res, projectId: (await getCurrentProjectContext()).id ?? 0, - apiType: ApiType.GET_CURRENT_ORGANIZATION, - requestPayload: {}, + apiType: + req.method === 'PUT' + ? ApiType.UPDATE_CURRENT_ORGANIZATION + : ApiType.GET_CURRENT_ORGANIZATION, + requestPayload: req.method === 'PUT' ? req.body : {}, headers: req.headers as Record, startTime, logger, diff --git a/wren-ui/src/pages/organization/general.tsx b/wren-ui/src/pages/organization/general.tsx new file mode 100644 index 0000000000..e8415f2766 --- /dev/null +++ b/wren-ui/src/pages/organization/general.tsx @@ -0,0 +1,269 @@ +import { useEffect, useMemo, useState } from 'react'; +import Link from 'next/link'; +import { + Button, + Form, + Input, + Layout, + Typography, + message, +} from 'antd'; +import styled from 'styled-components'; +import SimpleLayout from '@/components/layouts/SimpleLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; + +const { Sider, Content } = Layout; + +interface OrganizationRecord { + id: number; + name: string; + identifier: string; + description?: string | null; + isCurrent: boolean; +} + +interface OrganizationResponse { + organizations: OrganizationRecord[]; + currentOrganization: OrganizationRecord | null; + currentProjectName: string; + error?: string; +} + +const linkStyle = { color: 'inherit', transition: 'none' }; + +const StyledSider = styled(Sider)` + height: calc(100vh - 48px); + background: var(--gray-2); + border-right: 1px solid var(--gray-4); + overflow-y: auto; +`; + +const StyledContent = styled(Content)` + height: calc(100vh - 48px); + overflow-y: auto; + background: white; +`; + +const SidebarSection = styled.div` + padding: 14px 16px 8px; + font-size: 12px; + font-weight: 700; + color: var(--gray-7); +`; + +const SidebarItem = styled.div<{ $active?: boolean; $disabled?: boolean }>` + padding: 6px 20px; + color: ${(props) => + props.$active + ? 'var(--gray-10)' + : props.$disabled + ? 'var(--gray-7)' + : 'var(--gray-8)'}; + background: ${(props) => (props.$active ? 'var(--gray-4)' : 'transparent')}; + font-weight: ${(props) => (props.$active ? 600 : 400)}; +`; + +const PageBody = styled.div` + padding: 24px 48px; +`; + +const SettingsCard = styled.div` + margin-top: 16px; + border: 1px solid var(--gray-4); + border-radius: 4px; + padding: 20px 28px 28px; + background: white; +`; + +const InlineRow = styled.div` + display: flex; + align-items: center; + gap: 16px; + max-width: 860px; +`; + +const LabelCell = styled.div` + width: 160px; + text-align: right; + color: var(--gray-8); + flex-shrink: 0; +`; + +const FieldCell = styled.div` + flex: 1; +`; + +const Actions = styled.div` + margin-left: 176px; + display: flex; + gap: 8px; +`; + +const PlaceholderItem = ({ children }: { children: React.ReactNode }) => ( + {children} +); + +export default function OrganizationGeneralPage() { + const [form] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + const [hasOrganization, setHasOrganization] = useState(false); + const [initialName, setInitialName] = useState(''); + + const loadOrganization = async () => { + setLoading(true); + try { + const response = await fetch('/api/v1/organizations/current'); + const payload = (await response.json()) as OrganizationResponse; + if (!response.ok) { + throw new Error(payload.error || 'Failed to load organization'); + } + + const organization = payload.currentOrganization; + setHasOrganization(Boolean(organization)); + const organizationName = organization?.name || ''; + setInitialName(organizationName); + form.setFieldsValue({ name: organizationName }); + } catch (error: any) { + message.error(error.message || 'Failed to load organization'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadOrganization(); + }, []); + + const currentName = Form.useWatch('name', form); + const hasChanges = useMemo( + () => (currentName || '').trim() !== initialName, + [currentName, initialName], + ); + + const resetChanges = () => { + form.setFieldsValue({ name: initialName }); + }; + + const saveChanges = async () => { + try { + const values = await form.validateFields(); + setSaving(true); + const response = await fetch('/api/v1/organizations/current', { + method: 'PUT', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify(values), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to update organization'); + } + + const updatedName = payload.name || values.name; + setInitialName(updatedName); + form.setFieldsValue({ name: updatedName }); + message.success('Organization updated successfully.'); + } catch (error: any) { + if (error?.errorFields) return; + message.error(error.message || 'Failed to update organization'); + } finally { + setSaving(false); + } + }; + + return ( + + + + Project + General + Access control + Data source + Danger zone + + Organization + + + General + + + Members + Billing + Danger zone + + User + Profile + Danger zone + + + + + + General + + + Organization name + + + + {!hasOrganization && ( + + Create an organization from the header before editing + organization settings. + + )} + + Organization name + +
+ + + + + + + + + + + + + + + + + ); +} diff --git a/wren-ui/src/pages/organization/index.tsx b/wren-ui/src/pages/organization/index.tsx new file mode 100644 index 0000000000..da6625737f --- /dev/null +++ b/wren-ui/src/pages/organization/index.tsx @@ -0,0 +1,14 @@ +import { useEffect } from 'react'; +import { useRouter } from 'next/router'; +import PageLoading from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; + +export default function OrganizationIndex() { + const router = useRouter(); + + useEffect(() => { + router.replace(Path.OrganizationGeneral); + }, [router]); + + return ; +} diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index 54c0eb91f0..4b42349b79 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -16,4 +16,6 @@ export enum Path { AdministrationUsers = '/administration/users', AdministrationRoles = '/administration/roles', AdministrationAssignments = '/administration/assignments', + Organization = '/organization', + OrganizationGeneral = '/organization/general', } From e53da02eb0a0bc0f1569017682d88a80bcd09318 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 4 Jun 2026 01:36:23 +0530 Subject: [PATCH 0101/1087] Implemented general org settings --- wren-ui/src/components/HeaderBar.tsx | 33 +++++++++++++++++++++------- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index 953440cd62..9437d03503 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -32,6 +32,27 @@ const StyledHeader = styled(Header)` padding: 10px 16px; `; +const HeaderLeft = styled.div` + display: flex; + align-items: center; + gap: 48px; +`; + +const HeaderCenter = styled.div` + display: flex; + align-items: center; + justify-content: center; + flex: 1; + min-width: 0; +`; + +const HeaderRight = styled.div` + display: flex; + align-items: center; + justify-content: flex-end; + min-width: 120px; +`; + export default function HeaderBar() { const router = useRouter(); const { pathname } = router; @@ -44,7 +65,7 @@ export default function HeaderBar() { className="d-flex justify-space-between align-center" style={{ marginTop: -2 }} > - + {showNav && ( @@ -90,13 +111,9 @@ export default function HeaderBar() { )} - - - - {isModeling && ( - - )} - + + {showNav && } + {isModeling && } ); From e8f7098cb77a32049f9c414287be2d4aa69adb29 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 4 Jun 2026 01:48:46 +0530 Subject: [PATCH 0102/1087] Implemented general org settings --- wren-ui/src/components/sidebar/index.tsx | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/wren-ui/src/components/sidebar/index.tsx b/wren-ui/src/components/sidebar/index.tsx index 0ed182995b..7850355774 100644 --- a/wren-ui/src/components/sidebar/index.tsx +++ b/wren-ui/src/components/sidebar/index.tsx @@ -80,11 +80,10 @@ const DynamicSidebar = ( }; export default function Sidebar(props: Props) { - const { onOpenSettings } = props; const router = useRouter(); const onSettingsClick = (event) => { - onOpenSettings && onOpenSettings(); + router.push(Path.OrganizationGeneral); event.target.blur(); }; From 7a48cb9837bb7712c2792717ff899702e3777e35 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 4 Jun 2026 02:16:55 +0530 Subject: [PATCH 0103/1087] Add organization members module --- ...00000_create_organization_member_tables.js | 44 ++ .../repositories/apiHistoryRepository.ts | 4 + .../src/apollo/server/repositories/index.ts | 1 + .../organizationMemberRepository.ts | 175 +++++++ wren-ui/src/apollo/server/services/index.ts | 1 + .../services/organizationMemberService.ts | 308 +++++++++++++ wren-ui/src/common.ts | 16 + .../organization/SettingsLayout.tsx | 105 +++++ .../api/v1/organizations/members/[id].ts | 83 ++++ .../api/v1/organizations/members/index.ts | 79 ++++ wren-ui/src/pages/organization/general.tsx | 207 +++------ wren-ui/src/pages/organization/members.tsx | 427 ++++++++++++++++++ wren-ui/src/utils/enum/path.ts | 1 + 13 files changed, 1307 insertions(+), 144 deletions(-) create mode 100644 wren-ui/migrations/20250604000000_create_organization_member_tables.js create mode 100644 wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts create mode 100644 wren-ui/src/apollo/server/services/organizationMemberService.ts create mode 100644 wren-ui/src/components/organization/SettingsLayout.tsx create mode 100644 wren-ui/src/pages/api/v1/organizations/members/[id].ts create mode 100644 wren-ui/src/pages/api/v1/organizations/members/index.ts create mode 100644 wren-ui/src/pages/organization/members.tsx diff --git a/wren-ui/migrations/20250604000000_create_organization_member_tables.js b/wren-ui/migrations/20250604000000_create_organization_member_tables.js new file mode 100644 index 0000000000..8361f107f0 --- /dev/null +++ b/wren-ui/migrations/20250604000000_create_organization_member_tables.js @@ -0,0 +1,44 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + await knex.schema.createTable('organization_members', (table) => { + table.increments('id').primary(); + table.integer('organization_id').notNullable(); + table.integer('user_id').notNullable(); + table.string('organization_role', 80).notNullable(); + table.timestamps(true, true); + + table + .foreign('organization_id') + .references('organization.id') + .onDelete('CASCADE'); + table.foreign('user_id').references('users.id').onDelete('CASCADE'); + table.unique(['organization_id', 'user_id']); + }); + + await knex.schema.createTable('organization_member_projects', (table) => { + table.increments('id').primary(); + table.integer('organization_member_id').notNullable(); + table.integer('project_id').notNullable(); + table.string('permission', 80).notNullable(); + table.timestamps(true, true); + + table + .foreign('organization_member_id') + .references('organization_members.id') + .onDelete('CASCADE'); + table.foreign('project_id').references('project.id').onDelete('CASCADE'); + table.unique(['organization_member_id', 'project_id']); + }); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex.schema.dropTableIfExists('organization_member_projects'); + await knex.schema.dropTableIfExists('organization_members'); +}; diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index a07e7fbaa0..d7862963f4 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -30,6 +30,10 @@ export enum ApiType { CREATE_ORGANIZATION = 'CREATE_ORGANIZATION', UPDATE_CURRENT_ORGANIZATION = 'UPDATE_CURRENT_ORGANIZATION', SELECT_ORGANIZATION = 'SELECT_ORGANIZATION', + GET_ORGANIZATION_MEMBERS = 'GET_ORGANIZATION_MEMBERS', + INVITE_ORGANIZATION_MEMBER = 'INVITE_ORGANIZATION_MEMBER', + UPDATE_ORGANIZATION_MEMBER = 'UPDATE_ORGANIZATION_MEMBER', + REMOVE_ORGANIZATION_MEMBER = 'REMOVE_ORGANIZATION_MEMBER', } export interface ApiHistory { diff --git a/wren-ui/src/apollo/server/repositories/index.ts b/wren-ui/src/apollo/server/repositories/index.ts index 753e3f3c36..51af16549d 100644 --- a/wren-ui/src/apollo/server/repositories/index.ts +++ b/wren-ui/src/apollo/server/repositories/index.ts @@ -21,3 +21,4 @@ export * from './apiHistoryRepository'; export * from './dashboardItemRefreshJobRepository'; export * from './rbacRepository'; export * from './organizationRepository'; +export * from './organizationMemberRepository'; diff --git a/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts b/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts new file mode 100644 index 0000000000..c5cef36f33 --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts @@ -0,0 +1,175 @@ +import { Knex } from 'knex'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; +import { Project } from './projectRepository'; +import { DataSourceName } from '@server/types'; +import { RbacUser } from './rbacRepository'; + +export interface OrganizationMember { + id: number; + organizationId: number; + userId: number; + organizationRole: string; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationMemberProject { + id: number; + organizationMemberId: number; + projectId: number; + permission: string; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationMemberProjectAssignment { + id: number; + projectId: number; + permission: string; + project: Project; +} + +export interface OrganizationMemberMapping extends OrganizationMember { + user: RbacUser; + projects: OrganizationMemberProjectAssignment[]; +} + +export interface IOrganizationMemberRepository + extends IBasicRepository { + findMappingsByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export interface IOrganizationMemberProjectRepository + extends IBasicRepository { + findByOrganizationMemberId( + organizationMemberId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export class OrganizationMemberRepository + extends BaseRepository + implements IOrganizationMemberRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization_members' }); + } + + public async findMappingsByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const rows = await executer('organization_members') + .select( + 'organization_members.*', + 'users.id as user__id', + 'users.name as user__name', + 'users.email as user__email', + 'users.external_id as user__external_id', + 'users.identity_provider as user__identity_provider', + 'users.is_active as user__is_active', + 'users.created_at as user__created_at', + 'users.updated_at as user__updated_at', + ) + .join('users', 'organization_members.user_id', 'users.id') + .where('organization_members.organization_id', organizationId) + .orderBy('users.email'); + + return rows.map((row) => ({ + id: row.id, + organizationId: row.organization_id, + userId: row.user_id, + organizationRole: row.organization_role, + createdAt: row.created_at, + updatedAt: row.updated_at, + user: { + id: row.user__id, + name: row.user__name, + email: row.user__email, + externalId: row.user__external_id, + identityProvider: row.user__identity_provider, + isActive: row.user__is_active, + createdAt: row.user__created_at, + updatedAt: row.user__updated_at, + }, + projects: [], + })); + } +} + +export class OrganizationMemberProjectRepository + extends BaseRepository + implements IOrganizationMemberProjectRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization_member_projects' }); + } + + public async findByOrganizationMemberId( + organizationMemberId: number, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const rows = await executer('organization_member_projects') + .select( + 'organization_member_projects.*', + 'project.id as project__id', + 'project.display_name as project__display_name', + 'project.type as project__type', + 'project.version as project__version', + 'project.catalog as project__catalog', + 'project.schema as project__schema', + 'project.sample_dataset as project__sample_dataset', + 'project.connection_info as project__connection_info', + 'project.language as project__language', + 'project.query_id as project__query_id', + 'project.questions as project__questions', + 'project.questions_status as project__questions_status', + 'project.questions_error as project__questions_error', + ) + .join('project', 'organization_member_projects.project_id', 'project.id') + .where( + 'organization_member_projects.organization_member_id', + organizationMemberId, + ) + .orderBy('project.id'); + + return rows.map((row) => ({ + id: row.id, + projectId: row.project_id, + permission: row.permission, + project: { + id: row.project__id, + displayName: row.project__display_name, + type: DataSourceName[row.project__type], + version: row.project__version, + catalog: row.project__catalog, + schema: row.project__schema, + sampleDataset: row.project__sample_dataset, + connectionInfo: + typeof row.project__connection_info === 'string' + ? JSON.parse(row.project__connection_info || '{}') + : row.project__connection_info, + language: row.project__language, + queryId: row.project__query_id, + questions: + typeof row.project__questions === 'string' + ? JSON.parse(row.project__questions || '[]') + : row.project__questions, + questionsStatus: row.project__questions_status, + questionsError: + typeof row.project__questions_error === 'string' + ? JSON.parse(row.project__questions_error || '{}') + : row.project__questions_error, + }, + })); + } +} diff --git a/wren-ui/src/apollo/server/services/index.ts b/wren-ui/src/apollo/server/services/index.ts index 0d052c814a..8d341f923e 100644 --- a/wren-ui/src/apollo/server/services/index.ts +++ b/wren-ui/src/apollo/server/services/index.ts @@ -10,3 +10,4 @@ export * from './askingTaskTracker'; export * from './instructionService'; export * from './rbacService'; export * from './organizationService'; +export * from './organizationMemberService'; diff --git a/wren-ui/src/apollo/server/services/organizationMemberService.ts b/wren-ui/src/apollo/server/services/organizationMemberService.ts new file mode 100644 index 0000000000..1938da1d94 --- /dev/null +++ b/wren-ui/src/apollo/server/services/organizationMemberService.ts @@ -0,0 +1,308 @@ +import { ApiError } from '../utils/apiUtils'; +import { + IOrganizationMemberProjectRepository, + IOrganizationMemberRepository, + OrganizationMemberMapping, +} from '../repositories/organizationMemberRepository'; +import { IOrganizationRepository } from '../repositories/organizationRepository'; +import { IProjectRepository } from '../repositories/projectRepository'; +import { IUserRepository, RbacUser } from '../repositories/rbacRepository'; + +export const ORGANIZATION_MEMBER_ROLES = ['Admin', 'Member'] as const; +export const PROJECT_PERMISSION_ROLES = ['Owner', 'Editor', 'Viewer'] as const; + +export type OrganizationMemberRole = + (typeof ORGANIZATION_MEMBER_ROLES)[number]; +export type ProjectPermissionRole = (typeof PROJECT_PERMISSION_ROLES)[number]; + +export interface MemberProjectInput { + projectId: number; + permission: ProjectPermissionRole; +} + +export interface InviteOrganizationMemberInput { + email: string; + organizationRole: OrganizationMemberRole; + projects: MemberProjectInput[]; +} + +export interface UpdateOrganizationMemberInput { + organizationRole: OrganizationMemberRole; +} + +export interface OrganizationMemberSummary { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationMemberRole; + projects: Array<{ + projectId: number; + displayName: string; + permission: ProjectPermissionRole; + }>; +} + +export interface IOrganizationMemberService { + listCurrentOrganizationMembers(): Promise<{ + members: OrganizationMemberSummary[]; + projects: Array<{ id: number; displayName: string }>; + }>; + inviteMember( + input: InviteOrganizationMemberInput, + ): Promise; + updateMember( + id: number, + input: UpdateOrganizationMemberInput, + ): Promise; + removeMember(id: number): Promise; +} + +const EMAIL_PATTERN = /^[^\s@]+@[^\s@]+\.[^\s@]+$/; + +export class OrganizationMemberService implements IOrganizationMemberService { + constructor( + private readonly organizationRepository: IOrganizationRepository, + private readonly organizationMemberRepository: IOrganizationMemberRepository, + private readonly organizationMemberProjectRepository: IOrganizationMemberProjectRepository, + private readonly userRepository: IUserRepository, + private readonly projectRepository: IProjectRepository, + ) {} + + public async listCurrentOrganizationMembers() { + const organization = await this.getCurrentOrganizationOrThrow(); + const [members, projects] = await Promise.all([ + this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ), + this.projectRepository.findAll({ order: 'id' }), + ]); + + const hydratedMembers = await Promise.all( + members.map((member) => this.serializeMember(member)), + ); + + return { + members: hydratedMembers, + projects: projects.map((project) => ({ + id: project.id, + displayName: project.displayName, + })), + }; + } + + public async inviteMember(input: InviteOrganizationMemberInput) { + const organization = await this.getCurrentOrganizationOrThrow(); + const payload = await this.validateInvitePayload(input); + const existingUser = await this.userRepository.findOneBy({ + email: payload.email, + }); + const user = existingUser + ? existingUser + : await this.createUserFromInvite(payload.email); + + const existingMembership = + await this.organizationMemberRepository.findOneBy({ + organizationId: organization.id, + userId: user.id, + }); + if (existingMembership) { + throw new ApiError('Member already exists in this organization', 409); + } + + const tx = await this.organizationMemberRepository.transaction(); + try { + const now = new Date().toISOString(); + const member = await this.organizationMemberRepository.createOne( + { + organizationId: organization.id, + userId: user.id, + organizationRole: payload.organizationRole, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + + if (payload.projects.length) { + await this.organizationMemberProjectRepository.createMany( + payload.projects.map((project) => ({ + organizationMemberId: member.id, + projectId: project.projectId, + permission: project.permission, + createdAt: now, + updatedAt: now, + })), + { tx }, + ); + } + + await tx.commit(); + const mapping = await this.getMemberMappingOrThrow(member.id); + return this.serializeMember(mapping); + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async updateMember( + id: number, + input: UpdateOrganizationMemberInput, + ): Promise { + const member = await this.organizationMemberRepository.findOneBy({ id }); + if (!member) { + throw new ApiError('Member not found', 404); + } + + const organization = await this.getCurrentOrganizationOrThrow(); + if (member.organizationId !== organization.id) { + throw new ApiError('Member not found', 404); + } + + const role = this.validateOrganizationRole(input.organizationRole); + await this.organizationMemberRepository.updateOne(member.id, { + organizationRole: role, + updatedAt: new Date().toISOString(), + }); + + const mapping = await this.getMemberMappingOrThrow(member.id); + return this.serializeMember(mapping); + } + + public async removeMember(id: number): Promise { + const member = await this.organizationMemberRepository.findOneBy({ id }); + if (!member) return true; + + const organization = await this.getCurrentOrganizationOrThrow(); + if (member.organizationId !== organization.id) { + throw new ApiError('Member not found', 404); + } + + await this.organizationMemberRepository.deleteOne(member.id); + return true; + } + + private async getCurrentOrganizationOrThrow() { + const organization = + await this.organizationRepository.getCurrentOrganization(); + if (!organization) { + throw new ApiError('Current organization not found', 404); + } + return organization; + } + + private async createUserFromInvite(email: string): Promise { + const now = new Date().toISOString(); + const localPart = email.split('@')[0] || 'user'; + return this.userRepository.createOne({ + name: localPart, + email, + isActive: true, + createdAt: now, + updatedAt: now, + }); + } + + private async validateInvitePayload(input: InviteOrganizationMemberInput) { + const email = `${input.email || ''}`.trim().toLowerCase(); + if (!EMAIL_PATTERN.test(email)) { + throw new ApiError('A valid email address is required', 400); + } + + const organizationRole = this.validateOrganizationRole( + input.organizationRole, + ); + const allProjects = await this.projectRepository.findAll({ order: 'id' }); + const projectMap = new Map(allProjects.map((project) => [project.id, project])); + + let projects = (input.projects || []).map((project) => ({ + projectId: Number(project.projectId), + permission: this.validateProjectPermission(project.permission), + })); + + if (organizationRole === 'Admin') { + projects = allProjects.map((project) => ({ + projectId: project.id, + permission: 'Owner' as ProjectPermissionRole, + })); + } else if (!projects.length) { + throw new ApiError( + 'Select at least one project for organization members', + 400, + ); + } + + for (const project of projects) { + if (!projectMap.has(project.projectId)) { + throw new ApiError(`Project ${project.projectId} not found`, 400); + } + } + + const dedupedProjects = Array.from( + new Map(projects.map((project) => [project.projectId, project])).values(), + ); + + return { + email, + organizationRole, + projects: dedupedProjects, + }; + } + + private validateOrganizationRole(role: string): OrganizationMemberRole { + const normalized = `${role || ''}`.trim(); + if ( + !ORGANIZATION_MEMBER_ROLES.includes( + normalized as OrganizationMemberRole, + ) + ) { + throw new ApiError('Invalid organization role', 400); + } + return normalized as OrganizationMemberRole; + } + + private validateProjectPermission(role: string): ProjectPermissionRole { + const normalized = `${role || ''}`.trim(); + if ( + !PROJECT_PERMISSION_ROLES.includes(normalized as ProjectPermissionRole) + ) { + throw new ApiError('Invalid project permission', 400); + } + return normalized as ProjectPermissionRole; + } + + private async getMemberMappingOrThrow(id: number) { + const organization = await this.getCurrentOrganizationOrThrow(); + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ); + const member = members.find((item) => item.id === id); + if (!member) { + throw new ApiError('Member not found', 404); + } + return member; + } + + private async serializeMember( + member: OrganizationMemberMapping, + ): Promise { + const projects = + await this.organizationMemberProjectRepository.findByOrganizationMemberId( + member.id, + ); + return { + id: member.id, + userId: member.userId, + name: member.user.name, + email: member.user.email, + organizationRole: member.organizationRole as OrganizationMemberRole, + projects: projects.map((project) => ({ + projectId: project.projectId, + displayName: project.project.displayName, + permission: project.permission as ProjectPermissionRole, + })), + }; + } +} diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index a22dd08bdb..57033442db 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -23,6 +23,8 @@ import { UserRepository, UserRoleRepository, OrganizationRepository, + OrganizationMemberRepository, + OrganizationMemberProjectRepository, } from '@server/repositories'; import { WrenEngineAdaptor, @@ -41,6 +43,7 @@ import { InstructionService, RbacService, OrganizationService, + OrganizationMemberService, } from '@server/services'; import { PostHogTelemetry } from './apollo/server/telemetry/telemetry'; import { @@ -160,6 +163,9 @@ export const initComponents = () => { const userRepository = new UserRepository(knex); const userRoleRepository = new UserRoleRepository(knex); const organizationRepository = new OrganizationRepository(knex); + const organizationMemberRepository = new OrganizationMemberRepository(knex); + const organizationMemberProjectRepository = + new OrganizationMemberProjectRepository(knex); // adaptors const wrenEngineAdaptor = new WrenEngineAdaptor({ @@ -273,6 +279,13 @@ export const initComponents = () => { const organizationService = new OrganizationService({ organizationRepository, }); + const organizationMemberService = new OrganizationMemberService( + organizationRepository, + organizationMemberRepository, + organizationMemberProjectRepository, + userRepository, + projectRepository, + ); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ dashboardRepository, @@ -310,6 +323,8 @@ export const initComponents = () => { userRepository, userRoleRepository, organizationRepository, + organizationMemberRepository, + organizationMemberProjectRepository, // adaptors wrenEngineAdaptor, @@ -328,6 +343,7 @@ export const initComponents = () => { instructionService, rbacService, organizationService, + organizationMemberService, askingTaskTracker, // background trackers diff --git a/wren-ui/src/components/organization/SettingsLayout.tsx b/wren-ui/src/components/organization/SettingsLayout.tsx new file mode 100644 index 0000000000..a04f6c49e2 --- /dev/null +++ b/wren-ui/src/components/organization/SettingsLayout.tsx @@ -0,0 +1,105 @@ +import { ReactNode } from 'react'; +import Link from 'next/link'; +import { Layout, Typography } from 'antd'; +import styled from 'styled-components'; +import SimpleLayout from '@/components/layouts/SimpleLayout'; +import { Path } from '@/utils/enum'; + +const { Sider, Content } = Layout; + +const linkStyle = { color: 'inherit', transition: 'none' }; + +const StyledSider = styled(Sider)` + height: calc(100vh - 48px); + background: var(--gray-2); + border-right: 1px solid var(--gray-4); + overflow-y: auto; +`; + +const StyledContent = styled(Content)` + height: calc(100vh - 48px); + overflow-y: auto; + background: white; +`; + +const SidebarSection = styled.div` + padding: 14px 16px 8px; + font-size: 12px; + font-weight: 700; + color: var(--gray-7); +`; + +const SidebarItem = styled.div<{ $active?: boolean; $disabled?: boolean }>` + padding: 6px 20px; + color: ${(props) => + props.$active + ? 'var(--gray-10)' + : props.$disabled + ? 'var(--gray-7)' + : 'var(--gray-8)'}; + background: ${(props) => (props.$active ? 'var(--gray-4)' : 'transparent')}; + font-weight: ${(props) => (props.$active ? 600 : 400)}; +`; + +const PageBody = styled.div` + padding: 24px 48px; +`; + +const PlaceholderItem = ({ children }: { children: ReactNode }) => ( + {children} +); + +export default function OrganizationSettingsLayout({ + section, + title, + titleExtra, + children, +}: { + section: 'general' | 'members'; + title: ReactNode; + titleExtra?: ReactNode; + children: ReactNode; +}) { + return ( + + + + Project + General + Access control + Data source + Danger zone + + Organization + + + General + + + + + Members + + + Billing + Danger zone + + User + Profile + Danger zone + + + +
+ + {title} + + {titleExtra} +
+ {children} + + + + + ); +} diff --git a/wren-ui/src/pages/api/v1/organizations/members/[id].ts b/wren-ui/src/pages/api/v1/organizations/members/[id].ts new file mode 100644 index 0000000000..88c6d290c8 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/members/[id].ts @@ -0,0 +1,83 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, + parseOrganizationId, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ORGANIZATION_MEMBER'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['PATCH', 'DELETE']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + const memberId = parseOrganizationId(req.query.id); + + if (req.method === 'PATCH') { + const member = await organizationMemberService.updateMember( + memberId, + req.body, + ); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: member, + projectId, + apiType: ApiType.UPDATE_ORGANIZATION_MEMBER, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + return; + } + + await organizationMemberService.removeMember(memberId); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.REMOVE_ORGANIZATION_MEMBER, + startTime, + requestPayload: { id: memberId }, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'PATCH' + ? ApiType.UPDATE_ORGANIZATION_MEMBER + : ApiType.REMOVE_ORGANIZATION_MEMBER, + requestPayload: req.method === 'PATCH' ? req.body : { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/members/index.ts b/wren-ui/src/pages/api/v1/organizations/members/index.ts new file mode 100644 index 0000000000..f343b8079e --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/members/index.ts @@ -0,0 +1,79 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ORGANIZATION_MEMBERS'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET', 'POST']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + + if (req.method === 'GET') { + const payload = + await organizationMemberService.listCurrentOrganizationMembers(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: payload, + projectId, + apiType: ApiType.GET_ORGANIZATION_MEMBERS, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + + const member = await organizationMemberService.inviteMember(req.body); + await respondWithSimple({ + res, + statusCode: 201, + responsePayload: member, + projectId, + apiType: ApiType.INVITE_ORGANIZATION_MEMBER, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'GET' + ? ApiType.GET_ORGANIZATION_MEMBERS + : ApiType.INVITE_ORGANIZATION_MEMBER, + requestPayload: req.method === 'GET' ? {} : req.body, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/organization/general.tsx b/wren-ui/src/pages/organization/general.tsx index e8415f2766..59b656956d 100644 --- a/wren-ui/src/pages/organization/general.tsx +++ b/wren-ui/src/pages/organization/general.tsx @@ -1,19 +1,8 @@ import { useEffect, useMemo, useState } from 'react'; -import Link from 'next/link'; -import { - Button, - Form, - Input, - Layout, - Typography, - message, -} from 'antd'; +import { Button, Form, Input, Typography, message } from 'antd'; import styled from 'styled-components'; -import SimpleLayout from '@/components/layouts/SimpleLayout'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; import { LoadingWrapper } from '@/components/PageLoading'; -import { Path } from '@/utils/enum'; - -const { Sider, Content } = Layout; interface OrganizationRecord { id: number; @@ -30,44 +19,6 @@ interface OrganizationResponse { error?: string; } -const linkStyle = { color: 'inherit', transition: 'none' }; - -const StyledSider = styled(Sider)` - height: calc(100vh - 48px); - background: var(--gray-2); - border-right: 1px solid var(--gray-4); - overflow-y: auto; -`; - -const StyledContent = styled(Content)` - height: calc(100vh - 48px); - overflow-y: auto; - background: white; -`; - -const SidebarSection = styled.div` - padding: 14px 16px 8px; - font-size: 12px; - font-weight: 700; - color: var(--gray-7); -`; - -const SidebarItem = styled.div<{ $active?: boolean; $disabled?: boolean }>` - padding: 6px 20px; - color: ${(props) => - props.$active - ? 'var(--gray-10)' - : props.$disabled - ? 'var(--gray-7)' - : 'var(--gray-8)'}; - background: ${(props) => (props.$active ? 'var(--gray-4)' : 'transparent')}; - font-weight: ${(props) => (props.$active ? 600 : 400)}; -`; - -const PageBody = styled.div` - padding: 24px 48px; -`; - const SettingsCard = styled.div` margin-top: 16px; border: 1px solid var(--gray-4); @@ -100,10 +51,6 @@ const Actions = styled.div` gap: 8px; `; -const PlaceholderItem = ({ children }: { children: React.ReactNode }) => ( - {children} -); - export default function OrganizationGeneralPage() { const [form] = Form.useForm(); const [loading, setLoading] = useState(true); @@ -175,95 +122,67 @@ export default function OrganizationGeneralPage() { }; return ( - - - - Project - General - Access control - Data source - Danger zone - - Organization - - - General - - - Members - Billing - Danger zone - - User - Profile - Danger zone - - - - - - General - - - Organization name - - - - {!hasOrganization && ( - - Create an organization from the header before editing - organization settings. - - )} - - Organization name - -
- - - - - - - - - - - - - - - - + + + + + + + + + + + + + ); } diff --git a/wren-ui/src/pages/organization/members.tsx b/wren-ui/src/pages/organization/members.tsx new file mode 100644 index 0000000000..04e390c79f --- /dev/null +++ b/wren-ui/src/pages/organization/members.tsx @@ -0,0 +1,427 @@ +import { useEffect, useMemo, useState } from 'react'; +import { + Avatar, + Button, + Checkbox, + Form, + Input, + Modal, + Select, + Table, + TableColumnsType, + Typography, + message, +} from 'antd'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; + +type OrganizationRole = 'Admin' | 'Member'; +type ProjectPermissionRole = 'Owner' | 'Editor' | 'Viewer'; + +interface ProjectOption { + id: number; + displayName: string; +} + +interface MemberProject { + projectId: number; + displayName: string; + permission: ProjectPermissionRole; +} + +interface MemberRecord { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationRole; + projects: MemberProject[]; +} + +interface MembersResponse { + members: MemberRecord[]; + projects: ProjectOption[]; + error?: string; +} + +const ROLE_OPTIONS = ['Admin', 'Member']; +const PROJECT_PERMISSION_OPTIONS = ['Owner', 'Editor', 'Viewer']; + +const MembersCard = styled.div` + margin-top: 16px; + border: 1px solid var(--gray-4); + border-radius: 4px; + background: white; + overflow: hidden; +`; + +const ProjectsBox = styled.div` + border: 1px solid var(--gray-4); + border-radius: 4px; + background: white; + overflow: hidden; +`; + +const ProjectsToolbar = styled.div` + padding: 12px 12px 0; +`; + +const ProjectRow = styled.div` + display: grid; + grid-template-columns: 28px 1fr 140px; + gap: 12px; + align-items: center; + padding: 10px 12px; + border-top: 1px solid var(--gray-4); +`; + +const ProjectHeader = styled(ProjectRow)` + background: var(--gray-3); + font-weight: 600; +`; + +const NameCell = styled.div` + display: flex; + align-items: center; + gap: 12px; +`; + +const MemberName = styled.div` + display: flex; + flex-direction: column; +`; + +const getInitials = (name: string) => + (name || 'U') + .trim() + .split(/\s+/) + .slice(0, 2) + .map((part) => part[0]) + .join('') + .toUpperCase(); + +const buildProjectSelectionMap = ( + projects: ProjectOption[], + organizationRole: OrganizationRole, +) => { + if (organizationRole === 'Admin') { + return projects.reduce>((acc, project) => { + acc[project.id] = 'Owner'; + return acc; + }, {}); + } + return {}; +}; + +export default function OrganizationMembersPage() { + const [form] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + const [updatingMemberId, setUpdatingMemberId] = useState(null); + const [visible, setVisible] = useState(false); + const [members, setMembers] = useState([]); + const [projects, setProjects] = useState([]); + const [projectSelections, setProjectSelections] = useState< + Record + >({}); + const organizationRole = Form.useWatch( + 'organizationRole', + form, + ) as OrganizationRole | undefined; + const projectSearch = Form.useWatch('projectSearch', form) as string | undefined; + + const loadMembers = async () => { + setLoading(true); + try { + const response = await fetch('/api/v1/organizations/members'); + const payload = (await response.json()) as MembersResponse; + if (!response.ok) { + throw new Error(payload.error || 'Failed to load members'); + } + setMembers(payload.members || []); + setProjects(payload.projects || []); + } catch (error: any) { + message.error(error.message || 'Failed to load members'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadMembers(); + }, []); + + useEffect(() => { + if (!visible) return; + if (organizationRole === 'Admin') { + setProjectSelections(buildProjectSelectionMap(projects, organizationRole)); + } + }, [organizationRole, visible, projects]); + + const openInviteModal = () => { + form.resetFields(); + form.setFieldsValue({ + organizationRole: 'Admin', + projectSearch: '', + }); + setProjectSelections(buildProjectSelectionMap(projects, 'Admin')); + setVisible(true); + }; + + const filteredProjects = useMemo(() => { + const search = `${projectSearch || ''}`.trim().toLowerCase(); + return projects.filter((project) => + !search + ? true + : project.displayName.toLowerCase().includes(search), + ); + }, [projects, projectSearch]); + + const toggleProject = (projectId: number, checked: boolean) => { + setProjectSelections((prev) => { + const next = { ...prev }; + if (checked) { + next[projectId] = next[projectId] || 'Owner'; + } else { + delete next[projectId]; + } + return next; + }); + }; + + const updateProjectPermission = ( + projectId: number, + permission: ProjectPermissionRole, + ) => { + setProjectSelections((prev) => ({ + ...prev, + [projectId]: permission, + })); + }; + + const inviteMember = async () => { + try { + const values = await form.validateFields(); + const response = await fetch('/api/v1/organizations/members', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + email: values.email, + organizationRole: values.organizationRole, + projects: Object.entries(projectSelections).map( + ([projectId, permission]) => ({ + projectId: Number(projectId), + permission, + }), + ), + }), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to invite member'); + } + message.success('Member invited successfully.'); + setVisible(false); + form.resetFields(); + setProjectSelections({}); + await loadMembers(); + } catch (error: any) { + if (error?.errorFields) return; + message.error(error.message || 'Failed to invite member'); + } finally { + setSaving(false); + } + }; + + const updateMemberRole = async ( + memberId: number, + organizationRole: OrganizationRole, + ) => { + try { + setUpdatingMemberId(memberId); + const response = await fetch(`/api/v1/organizations/members/${memberId}`, { + method: 'PATCH', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ organizationRole }), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to update member'); + } + setMembers((prev) => + prev.map((member) => + member.id === memberId + ? { ...member, organizationRole: payload.organizationRole } + : member, + ), + ); + message.success('Member role updated successfully.'); + } catch (error: any) { + message.error(error.message || 'Failed to update member'); + } finally { + setUpdatingMemberId(null); + } + }; + + const columns: TableColumnsType = [ + { + title: 'Name', + dataIndex: 'name', + render: (_value, record: MemberRecord) => ( + + + {getInitials(record.name)} + + + {record.name} + + {record.email} + + + + ), + }, + { + title: 'Role', + dataIndex: 'organizationRole', + width: 220, + render: (value: OrganizationRole, record: MemberRecord) => ( +
+ + + + + { + setVisible(false); + setProjectSelections({}); + form.resetFields(); + }} + onOk={() => { + setSaving(true); + void inviteMember(); + }} + confirmLoading={saving} + okText="Invite" + destroyOnClose + width={720} + > +
+ + + + + + + + + + +
+
Project name
+
Permission
+ + + {filteredProjects.map((project) => { + const checked = Boolean(projectSelections[project.id]); + const permission = projectSelections[project.id] || 'Owner'; + return ( + + + toggleProject(project.id, event.target.checked) + } + /> +
{project.displayName}
+ ({ label: role, value: role }))} - onChange={(nextValue) => - void updateMemberRole(record.id, nextValue as OrganizationRole) - } - /> - ), - }, - ]; - return (
- - Organization members - +

Organization members

-
+ +
Name
+
Role
+ + {members.map((member) => ( + + + {getInitials(member.name)} + + {member.name} + {member.email} + + + + void updateMemberRole( + member.id, + event.target.value as OrganizationRole, + ) + } + > + {ROLE_OPTIONS.map((role) => ( + + ))} + + + ))} @@ -369,20 +410,21 @@ export default function OrganizationMembersPage() { name="organizationRole" rules={[{ required: true, message: 'Organization role is required' }]} > - @@ -407,20 +449,22 @@ export default function OrganizationMembersPage() { } />
{project.displayName}
- + + + + + Email : + + + + + + + + + + + + + + + Account Verification + + + To ensure security, account verification is handled by the configured + identity provider or invitation flow for your organization. + + + + Change password + + + + Password : + + + + You will receive a confirmation link via email when password + reset delivery is configured. + + + + + + + ); +} diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index a174186e26..4256c0a980 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -21,4 +21,5 @@ export enum Path { OrganizationMembers = '/organization/members', OrganizationDangerZone = '/organization/danger-zone', ProjectCreate = '/projects/create', + UserProfile = '/user/profile', } From 07c276252390eb409a39a1a255155bc2dea3e360 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 13:53:08 +0530 Subject: [PATCH 0114/1087] Add database --- wren-ui/tools/knex.js | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/wren-ui/tools/knex.js b/wren-ui/tools/knex.js index 27beb1b0ec..4c74e1f9b5 100644 --- a/wren-ui/tools/knex.js +++ b/wren-ui/tools/knex.js @@ -8,6 +8,14 @@ const SQLITE_FILE = process.env.SQLITE_FILE; // export SQLITE_FILE=./db.sqlite3 const APP_TABLE_ORDER = [ 'project', + 'roles', + 'users', + 'user_roles', + 'organization', + 'organization_members', + 'organization_member_projects', + 'organization_invitations', + 'organization_invitation_projects', 'model', 'model_column', 'model_nested_column', From 98a0116b6c52e6841e8ffa133e223b076d949520 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 15:07:13 +0530 Subject: [PATCH 0115/1087] Set API history timestamps before insert --- .../repositories/apiHistoryRepository.ts | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index bbae711ef2..175bbdf7ce 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -5,7 +5,11 @@ import { mapValues, snakeCase, } from 'lodash'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { Knex } from 'knex'; export enum ApiType { @@ -91,6 +95,20 @@ export class ApiHistoryRepository super({ knexPg, tableName: 'api_history' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany(data.map(this.withTimestamps), queryOptions); + } + /** * Count API history records with filtering */ @@ -184,6 +202,15 @@ export class ApiHistoryRepository return formattedData; }; + private withTimestamps = (data: Partial): Partial => { + const now = new Date().toISOString(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + protected override transformToDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); From 7b7bc2049473afaf9fc62304150c88a8eaa126bc Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 15:35:01 +0530 Subject: [PATCH 0116/1087] Skip invalid calculated fields in MDL build --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 104 ++++++++++++++------ 1 file changed, 73 insertions(+), 31 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 4f7a87671a..21ebb8e261 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -241,6 +241,12 @@ export class MDLBuilder implements IMDLBuilder { const relatedModel = this.relatedModels.find( (model: any) => model.id === column.modelId, ); + if (!relatedModel) { + logger.debug( + `Build MDL Column Error: can not find related model, modelId "${column.modelId}", columnId: "${column.id}"`, + ); + return; + } const model = this.manifest.models.find( (model: any) => model.name === relatedModel.referenceName, ); @@ -251,6 +257,12 @@ export class MDLBuilder implements IMDLBuilder { return; } const expression = this.getColumnExpression(column, model); + if (expression === null) { + logger.debug( + `Build MDL Column Error: invalid calculated field metadata, modelId "${column.modelId}", columnId: "${column.id}"`, + ); + return; + } const columnValue = { name: column.referenceName, type: column.type, @@ -283,6 +295,12 @@ export class MDLBuilder implements IMDLBuilder { return; } const expression = this.getColumnExpression(calculatedField, model); + if (expression === null) { + logger.debug( + `Can not add calculated field "${calculatedField.referenceName}" because its metadata is invalid`, + ); + return; + } const columnValue = { name: calculatedField.referenceName, type: calculatedField.type, @@ -379,7 +397,7 @@ export class MDLBuilder implements IMDLBuilder { protected getColumnExpression( column: ModelColumn, currentModel?: Partial, - ): string { + ): string | null { if (!column.isCalculated) { // columns existed in the data source. // Provide original column name in expression to MDL if referenceName has converted. @@ -389,40 +407,52 @@ export class MDLBuilder implements IMDLBuilder { return ''; } // calculated field - const lineage = JSON.parse(column.lineage) as number[]; + const lineage = this.parseLineage(column.lineage); + if (isEmpty(lineage) || !column.aggregation) { + return null; + } // lineage = [relationId1, relationId2, ..., columnId] - const fieldExpression = Object.entries(lineage).reduce( - (acc, [index, id]) => { - const isLast = parseInt(index) == lineage.length - 1; - if (isLast) { - // id is columnId - const columnReferenceName = this.relatedColumns.find( - (relatedColumn) => relatedColumn.id === id, - )?.referenceName; - acc.push(`\"${columnReferenceName}\"`); + const fieldExpression = lineage.reduce((acc, id, index) => { + const isLast = index === lineage.length - 1; + if (isLast) { + // id is columnId + const columnReferenceName = this.relatedColumns.find( + (relatedColumn) => relatedColumn.id === id, + )?.referenceName; + if (!columnReferenceName) { return acc; } - // id is relationId - const usedRelation = this.relatedRelations.find( - (relatedRelation) => relatedRelation.id === id, - ); - const relationColumnName = currentModel!.columns.find( - (c) => c.relationship === usedRelation.name, - ).name; - // move to next model - const nextModelName = - currentModel.name === usedRelation.fromModelName - ? usedRelation.toModelName - : usedRelation.fromModelName; - const nextModel = this.manifest.models.find( - (model) => model.name === nextModelName, - ); - currentModel = nextModel; - acc.push(relationColumnName); + acc.push(`\"${columnReferenceName}\"`); return acc; - }, - [], - ); + } + // id is relationId + const usedRelation = this.relatedRelations.find( + (relatedRelation) => relatedRelation.id === id, + ); + if (!usedRelation || !currentModel?.columns) { + return acc; + } + const relationColumnName = currentModel.columns.find( + (c) => c.relationship === usedRelation.name, + )?.name; + if (!relationColumnName) { + return acc; + } + // move to next model + const nextModelName = + currentModel.name === usedRelation.fromModelName + ? usedRelation.toModelName + : usedRelation.fromModelName; + const nextModel = this.manifest.models.find( + (model) => model.name === nextModelName, + ); + currentModel = nextModel; + acc.push(relationColumnName); + return acc; + }, []); + if (fieldExpression.length !== lineage.length) { + return null; + } return `${column.aggregation}(${fieldExpression.join('.')})`; } @@ -447,6 +477,18 @@ export class MDLBuilder implements IMDLBuilder { table: modelProps.table, }; } + private parseLineage(lineage?: string): number[] { + if (!lineage) { + return []; + } + try { + const parsedLineage = JSON.parse(lineage); + return Array.isArray(parsedLineage) ? parsedLineage : []; + } catch (error) { + logger.debug(`Can not parse calculated field lineage "${lineage}"`); + return []; + } + } private postProcessManifest() { if (this.useRustWrenEngine()) { // 1. remove all the key that the value is null From 53d10d64f7e8b727f88c2e2c08305ce5a8337e75 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 16:10:10 +0530 Subject: [PATCH 0117/1087] Normalize MSSQL boolean metadata values --- .../apollo/server/repositories/baseRepository.ts | 13 +++++++++++++ .../server/repositories/dashboardRepository.ts | 14 +++++++++++++- .../server/repositories/instructionRepository.ts | 9 ++++++++- .../server/repositories/metricsRepository.ts | 14 +++++++++++++- .../server/repositories/modelColumnRepository.ts | 11 +++++++++++ .../apollo/server/repositories/modelRepository.ts | 10 ++++++++++ .../repositories/organizationMemberRepository.ts | 3 ++- .../server/repositories/organizationRepository.ts | 15 ++++++++++++++- .../server/repositories/projectRepository.ts | 9 ++++++++- .../apollo/server/repositories/rbacRepository.ts | 11 ++++++++++- .../apollo/server/repositories/viewRepository.ts | 14 +++++++++++++- 11 files changed, 115 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index f76fe27654..b67008cd67 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -41,6 +41,19 @@ export interface IBasicRepository { ) => Promise; } +export const coerceBoolean = (value: unknown): boolean => { + if (typeof value === 'boolean') { + return value; + } + if (typeof value === 'number') { + return value === 1; + } + if (typeof value === 'string') { + return ['1', 'true'].includes(value.toLowerCase()); + } + return Boolean(value); +}; + export class BaseRepository implements IBasicRepository { protected knex: Knex; protected tableName: string; diff --git a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts index d88fe35c68..b73a2bb889 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, +} from './baseRepository'; import { ScheduleFrequencyEnum } from '@server/models/dashboard'; export interface Dashboard { @@ -22,4 +26,12 @@ export class DashboardRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'dashboard' }); } + + protected override transformFromDBData = (data: any): Dashboard => { + const dashboard = super.transformFromDBData(data) as Dashboard; + return { + ...dashboard, + cacheEnabled: coerceBoolean(dashboard.cacheEnabled), + }; + }; } diff --git a/wren-ui/src/apollo/server/repositories/instructionRepository.ts b/wren-ui/src/apollo/server/repositories/instructionRepository.ts index d3cef2cbc2..05d12d4eeb 100644 --- a/wren-ui/src/apollo/server/repositories/instructionRepository.ts +++ b/wren-ui/src/apollo/server/repositories/instructionRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -43,6 +47,9 @@ export class InstructionRepository return value; } } + if (key === 'isDefault') { + return coerceBoolean(value); + } return value; }); return transformData as Instruction; diff --git a/wren-ui/src/apollo/server/repositories/metricsRepository.ts b/wren-ui/src/apollo/server/repositories/metricsRepository.ts index 57bec8f920..81ad196a37 100644 --- a/wren-ui/src/apollo/server/repositories/metricsRepository.ts +++ b/wren-ui/src/apollo/server/repositories/metricsRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, +} from './baseRepository'; export interface Metric { id: number; // ID @@ -24,4 +28,12 @@ export class MetricRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'metric' }); } + + protected override transformFromDBData = (data: any): Metric => { + const metric = super.transformFromDBData(data) as Metric; + return { + ...metric, + cached: coerceBoolean(metric.cached), + }; + }; } diff --git a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts index df4198e214..bdc5bdb869 100644 --- a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts @@ -3,6 +3,7 @@ import { BaseRepository, IBasicRepository, IQueryOptions, + coerceBoolean, } from './baseRepository'; export interface ModelColumn { @@ -55,6 +56,16 @@ export class ModelColumnRepository super({ knexPg, tableName: 'model_column' }); } + protected override transformFromDBData = (data: any): ModelColumn => { + const column = super.transformFromDBData(data) as ModelColumn; + return { + ...column, + isCalculated: coerceBoolean(column.isCalculated), + notNull: coerceBoolean(column.notNull), + isPk: coerceBoolean(column.isPk), + }; + }; + public async findColumnsByModelIds(modelIds, queryOptions?: IQueryOptions) { if (queryOptions && queryOptions.tx) { const { tx } = queryOptions; diff --git a/wren-ui/src/apollo/server/repositories/modelRepository.ts b/wren-ui/src/apollo/server/repositories/modelRepository.ts index 57b92c9983..7c5a8ec4dd 100644 --- a/wren-ui/src/apollo/server/repositories/modelRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelRepository.ts @@ -3,6 +3,7 @@ import { BaseRepository, IBasicRepository, IQueryOptions, + coerceBoolean, } from './baseRepository'; export interface Model { @@ -32,6 +33,15 @@ export class ModelRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'model' }); } + + protected override transformFromDBData = (data: any): Model => { + const model = super.transformFromDBData(data) as Model; + return { + ...model, + cached: coerceBoolean(model.cached), + }; + }; + public async findAllByIds(ids: number[]) { const res = await this.knex(this.tableName).whereIn('id', ids); return res.map((r) => this.transformFromDBData(r)); diff --git a/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts b/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts index c5cef36f33..4506c3df63 100644 --- a/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts +++ b/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts @@ -3,6 +3,7 @@ import { BaseRepository, IBasicRepository, IQueryOptions, + coerceBoolean, } from './baseRepository'; import { Project } from './projectRepository'; import { DataSourceName } from '@server/types'; @@ -96,7 +97,7 @@ export class OrganizationMemberRepository email: row.user__email, externalId: row.user__external_id, identityProvider: row.user__identity_provider, - isActive: row.user__is_active, + isActive: coerceBoolean(row.user__is_active), createdAt: row.user__created_at, updatedAt: row.user__updated_at, }, diff --git a/wren-ui/src/apollo/server/repositories/organizationRepository.ts b/wren-ui/src/apollo/server/repositories/organizationRepository.ts index 7059b5ad9e..cf50e2584e 100644 --- a/wren-ui/src/apollo/server/repositories/organizationRepository.ts +++ b/wren-ui/src/apollo/server/repositories/organizationRepository.ts @@ -1,5 +1,10 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository, IQueryOptions } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, + coerceBoolean, +} from './baseRepository'; export interface Organization { id: number; @@ -29,6 +34,14 @@ export class OrganizationRepository super({ knexPg, tableName: 'organization' }); } + protected override transformFromDBData = (data: any): Organization => { + const organization = super.transformFromDBData(data) as Organization; + return { + ...organization, + isCurrent: coerceBoolean(organization.isCurrent), + }; + }; + public async getCurrentOrganization(queryOptions?: IQueryOptions) { return await this.findOneBy({ isCurrent: true }, queryOptions); } diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 557049ae26..8555704f49 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -273,6 +277,9 @@ export class ProjectRepository const camelCaseData = mapKeys(formattedData, (_value, key) => camelCase(key), ); + if (Object.prototype.hasOwnProperty.call(camelCaseData, 'isCurrent')) { + camelCaseData.isCurrent = coerceBoolean(camelCaseData.isCurrent); + } return camelCaseData as Project; }; diff --git a/wren-ui/src/apollo/server/repositories/rbacRepository.ts b/wren-ui/src/apollo/server/repositories/rbacRepository.ts index d742af6579..53d01ba988 100644 --- a/wren-ui/src/apollo/server/repositories/rbacRepository.ts +++ b/wren-ui/src/apollo/server/repositories/rbacRepository.ts @@ -3,6 +3,7 @@ import { BaseRepository, IBasicRepository, IQueryOptions, + coerceBoolean, } from './baseRepository'; export interface Role { @@ -69,6 +70,14 @@ export class UserRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'users' }); } + + protected override transformFromDBData = (data: any): RbacUser => { + const user = super.transformFromDBData(data) as RbacUser; + return { + ...user, + isActive: coerceBoolean(user.isActive), + }; + }; } export class UserRoleRepository @@ -144,7 +153,7 @@ export class UserRoleRepository email: row.user__email, externalId: row.user__external_id, identityProvider: row.user__identity_provider, - isActive: row.user__is_active, + isActive: coerceBoolean(row.user__is_active), createdAt: row.user__created_at, updatedAt: row.user__updated_at, }, diff --git a/wren-ui/src/apollo/server/repositories/viewRepository.ts b/wren-ui/src/apollo/server/repositories/viewRepository.ts index 39ed245c35..bbe0072dfb 100644 --- a/wren-ui/src/apollo/server/repositories/viewRepository.ts +++ b/wren-ui/src/apollo/server/repositories/viewRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, +} from './baseRepository'; export interface View { id: number; // ID @@ -20,4 +24,12 @@ export class ViewRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'view' }); } + + protected override transformFromDBData = (data: any): View => { + const view = super.transformFromDBData(data) as View; + return { + ...view, + cached: coerceBoolean(view.cached), + }; + }; } From 4640b1219e851a7e8487050dc1f26d97b5a7ef0e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 16:45:58 +0530 Subject: [PATCH 0118/1087] Fix MSSQL boolean repository transforms --- wren-ui/src/apollo/server/repositories/baseRepository.ts | 7 +++++-- .../src/apollo/server/repositories/dashboardRepository.ts | 2 +- .../src/apollo/server/repositories/metricsRepository.ts | 2 +- .../apollo/server/repositories/modelColumnRepository.ts | 2 +- wren-ui/src/apollo/server/repositories/modelRepository.ts | 2 +- .../apollo/server/repositories/organizationRepository.ts | 2 +- wren-ui/src/apollo/server/repositories/rbacRepository.ts | 2 +- wren-ui/src/apollo/server/repositories/viewRepository.ts | 2 +- 8 files changed, 12 insertions(+), 9 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index b67008cd67..f7c211206b 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -188,11 +188,14 @@ export class BaseRepository implements IBasicRepository { return mapKeys(data, (_value, key) => snakeCase(key)); }; - protected transformFromDBData = (data: any): T => { + protected defaultTransformFromDBData(data: any): T { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); } const camelCaseData = mapKeys(data, (_value, key) => camelCase(key)); return camelCaseData as T; - }; + } + + protected transformFromDBData = (data: any): T => + this.defaultTransformFromDBData(data); } diff --git a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts index b73a2bb889..b43cea8cdc 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts @@ -28,7 +28,7 @@ export class DashboardRepository } protected override transformFromDBData = (data: any): Dashboard => { - const dashboard = super.transformFromDBData(data) as Dashboard; + const dashboard = this.defaultTransformFromDBData(data) as Dashboard; return { ...dashboard, cacheEnabled: coerceBoolean(dashboard.cacheEnabled), diff --git a/wren-ui/src/apollo/server/repositories/metricsRepository.ts b/wren-ui/src/apollo/server/repositories/metricsRepository.ts index 81ad196a37..411ae9f45c 100644 --- a/wren-ui/src/apollo/server/repositories/metricsRepository.ts +++ b/wren-ui/src/apollo/server/repositories/metricsRepository.ts @@ -30,7 +30,7 @@ export class MetricRepository } protected override transformFromDBData = (data: any): Metric => { - const metric = super.transformFromDBData(data) as Metric; + const metric = this.defaultTransformFromDBData(data) as Metric; return { ...metric, cached: coerceBoolean(metric.cached), diff --git a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts index bdc5bdb869..c63df2daa9 100644 --- a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts @@ -57,7 +57,7 @@ export class ModelColumnRepository } protected override transformFromDBData = (data: any): ModelColumn => { - const column = super.transformFromDBData(data) as ModelColumn; + const column = this.defaultTransformFromDBData(data) as ModelColumn; return { ...column, isCalculated: coerceBoolean(column.isCalculated), diff --git a/wren-ui/src/apollo/server/repositories/modelRepository.ts b/wren-ui/src/apollo/server/repositories/modelRepository.ts index 7c5a8ec4dd..31823ffb4f 100644 --- a/wren-ui/src/apollo/server/repositories/modelRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelRepository.ts @@ -35,7 +35,7 @@ export class ModelRepository } protected override transformFromDBData = (data: any): Model => { - const model = super.transformFromDBData(data) as Model; + const model = this.defaultTransformFromDBData(data) as Model; return { ...model, cached: coerceBoolean(model.cached), diff --git a/wren-ui/src/apollo/server/repositories/organizationRepository.ts b/wren-ui/src/apollo/server/repositories/organizationRepository.ts index cf50e2584e..7f35ac49e4 100644 --- a/wren-ui/src/apollo/server/repositories/organizationRepository.ts +++ b/wren-ui/src/apollo/server/repositories/organizationRepository.ts @@ -35,7 +35,7 @@ export class OrganizationRepository } protected override transformFromDBData = (data: any): Organization => { - const organization = super.transformFromDBData(data) as Organization; + const organization = this.defaultTransformFromDBData(data) as Organization; return { ...organization, isCurrent: coerceBoolean(organization.isCurrent), diff --git a/wren-ui/src/apollo/server/repositories/rbacRepository.ts b/wren-ui/src/apollo/server/repositories/rbacRepository.ts index 53d01ba988..d6c2dbf55a 100644 --- a/wren-ui/src/apollo/server/repositories/rbacRepository.ts +++ b/wren-ui/src/apollo/server/repositories/rbacRepository.ts @@ -72,7 +72,7 @@ export class UserRepository } protected override transformFromDBData = (data: any): RbacUser => { - const user = super.transformFromDBData(data) as RbacUser; + const user = this.defaultTransformFromDBData(data) as RbacUser; return { ...user, isActive: coerceBoolean(user.isActive), diff --git a/wren-ui/src/apollo/server/repositories/viewRepository.ts b/wren-ui/src/apollo/server/repositories/viewRepository.ts index bbe0072dfb..ff125c76ab 100644 --- a/wren-ui/src/apollo/server/repositories/viewRepository.ts +++ b/wren-ui/src/apollo/server/repositories/viewRepository.ts @@ -26,7 +26,7 @@ export class ViewRepository } protected override transformFromDBData = (data: any): View => { - const view = super.transformFromDBData(data) as View; + const view = this.defaultTransformFromDBData(data) as View; return { ...view, cached: coerceBoolean(view.cached), From 259d9a1b0f1713b9df3c3300babeea466debcf9b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 18:05:15 +0530 Subject: [PATCH 0119/1087] Fix SQLite to MSSQL identity insert detection --- wren-ui/tools/knex.js | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/wren-ui/tools/knex.js b/wren-ui/tools/knex.js index 4c74e1f9b5..12c3ba8a51 100644 --- a/wren-ui/tools/knex.js +++ b/wren-ui/tools/knex.js @@ -161,6 +161,16 @@ const getTargetColumns = async (targetDb, tableName) => { return rows.map((row) => row.COLUMN_NAME); }; +const getTargetIdentityColumns = async (targetDb, tableName) => { + const rows = await targetDb('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ TABLE_SCHEMA: 'dbo', TABLE_NAME: tableName }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ); + return rows.map((row) => row.COLUMN_NAME); +}; + const parseCount = (row) => Number(row.count || row.Count || row[''] || 0); const getTableCount = async (db, tableName) => { @@ -246,10 +256,12 @@ const copyTable = async (sourceDb, targetDb, tableName) => { 1, Math.floor(1800 / Math.max(commonColumns.length, 1)), ); - const hasId = commonColumns.includes('id'); + const identityColumns = await getTargetIdentityColumns(targetDb, tableName); + const hasIdentityId = + commonColumns.includes('id') && identityColumns.includes('id'); - if (hasId) { - await targetDb.raw(`SET IDENTITY_INSERT [${tableName}] ON`); + if (hasIdentityId) { + await targetDb.raw(`SET IDENTITY_INSERT [dbo].[${tableName}] ON`); } try { @@ -257,8 +269,8 @@ const copyTable = async (sourceDb, targetDb, tableName) => { await targetDb(tableName).insert(rows.slice(index, index + chunkSize)); } } finally { - if (hasId) { - await targetDb.raw(`SET IDENTITY_INSERT [${tableName}] OFF`); + if (hasIdentityId) { + await targetDb.raw(`SET IDENTITY_INSERT [dbo].[${tableName}] OFF`); } } From 33329c68617dddc16db929250841c849a879ce6e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 18:46:55 +0530 Subject: [PATCH 0120/1087] Restore original settings modal behavior --- wren-ui/src/components/sidebar/index.tsx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/components/sidebar/index.tsx b/wren-ui/src/components/sidebar/index.tsx index 7850355774..0ed182995b 100644 --- a/wren-ui/src/components/sidebar/index.tsx +++ b/wren-ui/src/components/sidebar/index.tsx @@ -80,10 +80,11 @@ const DynamicSidebar = ( }; export default function Sidebar(props: Props) { + const { onOpenSettings } = props; const router = useRouter(); const onSettingsClick = (event) => { - router.push(Path.OrganizationGeneral); + onOpenSettings && onOpenSettings(); event.target.blur(); }; From cded0b0a2c2e684d6951c9a9afde515e76d782b8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 19:06:30 +0530 Subject: [PATCH 0121/1087] Revert sidebar settings modal change --- wren-ui/src/components/sidebar/index.tsx | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/wren-ui/src/components/sidebar/index.tsx b/wren-ui/src/components/sidebar/index.tsx index 0ed182995b..7850355774 100644 --- a/wren-ui/src/components/sidebar/index.tsx +++ b/wren-ui/src/components/sidebar/index.tsx @@ -80,11 +80,10 @@ const DynamicSidebar = ( }; export default function Sidebar(props: Props) { - const { onOpenSettings } = props; const router = useRouter(); const onSettingsClick = (event) => { - onOpenSettings && onOpenSettings(); + router.push(Path.OrganizationGeneral); event.target.blur(); }; From 53b98ceae63adbffdaae93a84754b0ceafcbbe06 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 5 Jun 2026 20:26:10 +0530 Subject: [PATCH 0122/1087] Update organization danger zone copy --- wren-ui/src/pages/organization/danger-zone.tsx | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/pages/organization/danger-zone.tsx b/wren-ui/src/pages/organization/danger-zone.tsx index 9e22236426..70262861ab 100644 --- a/wren-ui/src/pages/organization/danger-zone.tsx +++ b/wren-ui/src/pages/organization/danger-zone.tsx @@ -195,20 +195,25 @@ export default function OrganizationDangerZonePage() { Delete organization - Please be aware that deleting the organization will permanently - delete organization membership and associations. This cannot be - undone. + Only organization admins can delete the organization. This + action can not be reversed; be careful when performing this + action. + + Once the organization is deleted, its projects will no longer be + available to all the organization's users. + , + , + ]} + onCancel={() => { + setClassicProjectModalVisible(false); + classicProjectForm.resetFields(); + }} + destroyOnClose + > +
+ + + + + ); } diff --git a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx index a95d73c48c..5810fd3967 100644 --- a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx +++ b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx @@ -1,9 +1,11 @@ import Image from 'next/image'; import Link from 'next/link'; +import { useRouter } from 'next/router'; import { Alert, Typography, Form, Row, Col, Button } from 'antd'; import styled from 'styled-components'; import { DATA_SOURCES } from '@/utils/enum/dataSources'; import { getDataSource, getPostgresErrorMessage } from './utils'; +import { useEffect } from 'react'; const StyledForm = styled(Form)` border: 1px var(--gray-4) solid; @@ -26,8 +28,15 @@ interface Props { export default function ConnectDataSource(props: Props) { const { connectError, dataSource, submitting, onNext, onBack } = props; const [form] = Form.useForm(); + const router = useRouter(); const current = getDataSource(dataSource); + useEffect(() => { + if (typeof router.query.projectName === 'string') { + form.setFieldValue('displayName', router.query.projectName); + } + }, [form, router.query.projectName]); + const submit = () => { form .validateFields() From 99fed9ce1d5a3441f1a1359e5ee0b52d8a6eea15 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 6 Jun 2026 21:14:37 +0530 Subject: [PATCH 0129/1087] Add project general settings page --- .../src/apollo/client/graphql/__types__.ts | 1 + .../server/resolvers/projectResolver.ts | 14 +- wren-ui/src/apollo/server/schema.ts | 1 + .../organization/SettingsLayout.tsx | 15 +- wren-ui/src/components/sidebar/index.tsx | 2 +- wren-ui/src/pages/project/general.tsx | 279 ++++++++++++++++++ wren-ui/src/utils/enum/path.ts | 1 + 7 files changed, 305 insertions(+), 8 deletions(-) create mode 100644 wren-ui/src/pages/project/general.tsx diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 0d6a320ee0..d122814dcc 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -1557,6 +1557,7 @@ export type UpdateColumnMetadataInput = { }; export type UpdateCurrentProjectInput = { + displayName?: InputMaybe; language: ProjectLanguage; }; diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index c58a0090bf..4460506a70 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -98,14 +98,20 @@ export class ProjectResolver { public async updateCurrentProject( _root: any, - arg: { data: { language: string } }, + arg: { data: { language: string; displayName?: string } }, ctx: IContext, ) { - const { language } = arg.data; + const { language, displayName } = arg.data; const project = await ctx.projectService.getCurrentProject(); - await ctx.projectRepository.updateOne(project.id, { + const changes: Record = { language, - }); + }; + + if (typeof displayName === 'string' && trim(displayName)) { + changes.displayName = trim(displayName); + } + + await ctx.projectRepository.updateOne(project.id, changes); // only generating for user's data source if (project.sampleDataset === null) { diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 4280e88e01..15f4fcae2d 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -883,6 +883,7 @@ export const typeDefs = gql` input UpdateCurrentProjectInput { language: ProjectLanguage! + displayName: String } type Settings { diff --git a/wren-ui/src/components/organization/SettingsLayout.tsx b/wren-ui/src/components/organization/SettingsLayout.tsx index 3e325c7d82..6ad11fd19e 100644 --- a/wren-ui/src/components/organization/SettingsLayout.tsx +++ b/wren-ui/src/components/organization/SettingsLayout.tsx @@ -55,7 +55,12 @@ export default function OrganizationSettingsLayout({ titleExtra, children, }: { - section: 'general' | 'members' | 'danger-zone' | 'user-profile'; + section: + | 'project-general' + | 'general' + | 'members' + | 'danger-zone' + | 'user-profile'; title: ReactNode; titleExtra?: ReactNode; children: ReactNode; @@ -65,9 +70,13 @@ export default function OrganizationSettingsLayout({ Project - General + + + General + + Access control - Data source + Data connection Danger zone Organization diff --git a/wren-ui/src/components/sidebar/index.tsx b/wren-ui/src/components/sidebar/index.tsx index 7850355774..58c44c50e7 100644 --- a/wren-ui/src/components/sidebar/index.tsx +++ b/wren-ui/src/components/sidebar/index.tsx @@ -83,7 +83,7 @@ export default function Sidebar(props: Props) { const router = useRouter(); const onSettingsClick = (event) => { - router.push(Path.OrganizationGeneral); + router.push(Path.ProjectGeneral); event.target.blur(); }; diff --git a/wren-ui/src/pages/project/general.tsx b/wren-ui/src/pages/project/general.tsx new file mode 100644 index 0000000000..e9869c3f04 --- /dev/null +++ b/wren-ui/src/pages/project/general.tsx @@ -0,0 +1,279 @@ +import { useEffect, useMemo, useState } from 'react'; +import { Button, Form, Input, Select, Typography, message } from 'antd'; +import CopyOutlined from '@ant-design/icons/CopyOutlined'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; +import { + useGetSettingsQuery, + useUpdateCurrentProjectMutation, +} from '@/apollo/client/graphql/settings.generated'; +import { + ProjectLanguage, + WorkspaceProjectType, +} from '@/apollo/client/graphql/__types__'; +import { getLanguageText } from '@/utils/language'; + +interface CurrentProjectRecord { + id: number; + displayName: string; + projectType: WorkspaceProjectType; + isCurrent: boolean; + hasDataSource: boolean; + type?: string | null; + createdAt?: string; + updatedAt?: string; +} + +interface CurrentProjectResponse { + currentProject: CurrentProjectRecord | null; + projects: CurrentProjectRecord[]; + error?: string; +} + +const SettingsCard = styled.div` + margin-top: 16px; + border: 1px solid var(--gray-4); + border-radius: 4px; + padding: 20px 28px 28px; + background: white; +`; + +const DetailsTitle = styled(Typography.Title)` + && { + margin-top: 20px; + margin-bottom: 0; + color: var(--gray-8); + } +`; + +const InlineRow = styled.div` + display: flex; + align-items: flex-start; + gap: 16px; + max-width: 860px; + margin-bottom: 20px; +`; + +const LabelCell = styled.div` + width: 160px; + text-align: right; + color: var(--gray-8); + flex-shrink: 0; + padding-top: 8px; +`; + +const FieldCell = styled.div` + flex: 1; +`; + +const Actions = styled.div` + margin-left: 176px; + display: flex; + gap: 8px; +`; + +const HelperText = styled(Typography.Text)` + display: block; + margin-top: 6px; + color: var(--gray-6); +`; + +const languageOptions = Object.keys(ProjectLanguage).map((key) => ({ + label: getLanguageText(key as ProjectLanguage), + value: key, +})); + +export default function ProjectGeneralPage() { + const [form] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [project, setProject] = useState(null); + const { data: settingsData } = useGetSettingsQuery(); + const [initialValues, setInitialValues] = useState({ + displayName: '', + language: ProjectLanguage.EN, + }); + + const [updateCurrentProject, { loading: saving }] = + useUpdateCurrentProjectMutation({ + onError: (error) => + message.error(error.message || 'Failed to update project'), + onCompleted: () => { + message.success('Project updated successfully.'); + }, + }); + + const loadProject = async () => { + setLoading(true); + try { + const currentProjectResponse = await fetch('/api/v1/projects/current'); + const currentProjectPayload = + (await currentProjectResponse.json()) as CurrentProjectResponse; + if (!currentProjectResponse.ok) { + throw new Error( + currentProjectPayload.error || 'Failed to load current project', + ); + } + + const currentProject = currentProjectPayload.currentProject; + if (!currentProject) { + throw new Error('No current project found'); + } + + const values = { + displayName: currentProject.displayName || '', + language: settingsData?.settings?.language || ProjectLanguage.EN, + }; + + setProject(currentProject); + setInitialValues(values); + form.setFieldsValue(values); + } catch (error: any) { + message.error(error.message || 'Failed to load project settings'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadProject(); + }, [settingsData?.settings?.language]); + + const currentDisplayName = Form.useWatch('displayName', form); + const currentLanguage = Form.useWatch('language', form); + const hasChanges = useMemo( + () => + (currentDisplayName || '').trim() !== initialValues.displayName || + currentLanguage !== initialValues.language, + [currentDisplayName, currentLanguage, initialValues], + ); + + const resetChanges = () => { + form.setFieldsValue(initialValues); + }; + + const copyProjectId = async () => { + if (!project?.id) return; + try { + await navigator.clipboard.writeText(String(project.id)); + message.success('Project ID copied.'); + } catch { + message.error('Failed to copy project ID'); + } + }; + + const saveChanges = async () => { + try { + const values = await form.validateFields(); + await updateCurrentProject({ + variables: { + data: { + displayName: values.displayName.trim(), + language: values.language, + }, + }, + }); + const updatedValues = { + displayName: values.displayName.trim(), + language: values.language, + }; + setInitialValues(updatedValues); + setProject((current) => + current + ? { + ...current, + displayName: updatedValues.displayName, + } + : current, + ); + form.setFieldsValue(updatedValues); + } catch (error: any) { + if (error?.errorFields) return; + message.error(error.message || 'Failed to update project'); + } + }; + + return ( + + +
+ Details + +
+ + Project ID: + + } + onClick={() => void copyProjectId()} + /> + } + /> + + + + + Project name: + + + + + + + + + Project language: + + + + + + + member.id === selectedMemberId) + ?.name || '' + : '' + } + readOnly + /> + + + + + + + + {filteredAvailableMembers.length ? ( + filteredAvailableMembers.map((member) => ( + form.setFieldsValue({ organizationMemberId: member.id })} + > + {getInitials(member.name)} + + {member.name} + {member.email} + + + )) + ) : ( + + No organization members are available to add. + + )} + + + + + {PERMISSION_OPTIONS.map((permission) => ( + + ))} + + + + + + ); +} diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index edc8c82f75..9e76697188 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -21,6 +21,7 @@ export enum Path { OrganizationMembers = '/organization/members', OrganizationDangerZone = '/organization/danger-zone', ProjectGeneral = '/project/general', + ProjectAccessControl = '/project/access-control', ProjectCreate = '/projects/create', UserProfile = '/user/profile', } From b80bb5843e2e4fa39301b0f5475e1942512334ab Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 8 Jun 2026 14:19:11 +0530 Subject: [PATCH 0131/1087] Add profile menu to header --- wren-ui/src/components/HeaderBar.tsx | 65 +++++++++++++++++++++++++++- 1 file changed, 63 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index 9437d03503..026d9890fd 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -1,5 +1,6 @@ +import { useEffect, useMemo, useState } from 'react'; import { useRouter } from 'next/router'; -import { Button, Layout, Space } from 'antd'; +import { Avatar, Button, Dropdown, Layout, Menu, Space } from 'antd'; import styled from 'styled-components'; import LogoBar from '@/components/LogoBar'; import { Path } from '@/utils/enum'; @@ -50,14 +51,63 @@ const HeaderRight = styled.div` display: flex; align-items: center; justify-content: flex-end; + gap: 12px; min-width: 120px; `; +const UserAvatar = styled(Avatar)` + cursor: pointer; + background: var(--geekblue-6); + color: var(--gray-1); + font-weight: 600; +`; + +interface CurrentUserProfile { + name?: string; + email?: string; +} + export default function HeaderBar() { const router = useRouter(); const { pathname } = router; const showNav = !pathname.startsWith(Path.Onboarding); const isModeling = pathname.startsWith(Path.Modeling); + const [currentUser, setCurrentUser] = useState( + null, + ); + + useEffect(() => { + if (!showNav) return; + + const loadCurrentUser = async () => { + try { + const response = await fetch('/api/v1/users/current'); + if (!response.ok) return; + const payload = (await response.json()) as CurrentUserProfile; + setCurrentUser(payload); + } catch { + setCurrentUser(null); + } + }; + + loadCurrentUser(); + }, [showNav]); + + const userInitials = useMemo(() => { + const displayName = currentUser?.name || currentUser?.email || 'User'; + const parts = displayName.trim().split(/\s+/).filter(Boolean); + if (!parts.length) return 'U'; + if (parts.length === 1) return parts[0].slice(0, 2).toUpperCase(); + return `${parts[0][0]}${parts[parts.length - 1][0]}`.toUpperCase(); + }, [currentUser]); + + const userMenu = ( + + router.push(Path.UserProfile)}> + {currentUser?.email || 'Profile'} + + + ); return ( @@ -113,7 +163,18 @@ export default function HeaderBar() { )} {showNav && } - {isModeling && } + + {isModeling && } + {showNav && ( + + {userInitials} + + )} +
); From 70b12bb1437f8dccc9c4882e571d5347f9f2b1ef Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 8 Jun 2026 14:24:20 +0530 Subject: [PATCH 0132/1087] Fix header profile menu avatar render --- wren-ui/src/components/HeaderBar.tsx | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index 026d9890fd..dfacb3fd57 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -1,6 +1,6 @@ import { useEffect, useMemo, useState } from 'react'; import { useRouter } from 'next/router'; -import { Avatar, Button, Dropdown, Layout, Menu, Space } from 'antd'; +import { Button, Dropdown, Layout, Menu, Space } from 'antd'; import styled from 'styled-components'; import LogoBar from '@/components/LogoBar'; import { Path } from '@/utils/enum'; @@ -55,11 +55,18 @@ const HeaderRight = styled.div` min-width: 120px; `; -const UserAvatar = styled(Avatar)` +const UserAvatar = styled.div` + width: 32px; + height: 32px; + border-radius: 50%; + display: flex; + align-items: center; + justify-content: center; cursor: pointer; background: var(--geekblue-6); color: var(--gray-1); font-weight: 600; + line-height: 1; `; interface CurrentUserProfile { @@ -171,7 +178,7 @@ export default function HeaderBar() { trigger={['click']} placement="bottomRight" > - {userInitials} + {userInitials} )} From 1c174e5e25f2bd3ecd31021b4fd308fb00323de7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 8 Jun 2026 15:11:15 +0530 Subject: [PATCH 0133/1087] Add user account danger zone --- .../src/apollo/client/graphql/__types__.ts | 1 + .../repositories/apiHistoryRepository.ts | 2 + wren-ui/src/apollo/server/schema.ts | 1 + .../services/organizationMemberService.ts | 40 ++++++++ .../organization/SettingsLayout.tsx | 9 +- wren-ui/src/pages/api/v1/users/current.ts | 25 ++++- wren-ui/src/pages/user/danger-zone.tsx | 96 +++++++++++++++++++ wren-ui/src/utils/enum/path.ts | 1 + 8 files changed, 169 insertions(+), 6 deletions(-) create mode 100644 wren-ui/src/pages/user/danger-zone.tsx diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 2f79034152..48960a07db 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -84,6 +84,7 @@ export enum ApiType { CREATE_INSTRUCTION = 'CREATE_INSTRUCTION', CREATE_SQL_PAIR = 'CREATE_SQL_PAIR', DELETE_CURRENT_ORGANIZATION = 'DELETE_CURRENT_ORGANIZATION', + DELETE_CURRENT_USER = 'DELETE_CURRENT_USER', DELETE_INSTRUCTION = 'DELETE_INSTRUCTION', DELETE_SQL_PAIR = 'DELETE_SQL_PAIR', GENERATE_SQL = 'GENERATE_SQL', diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 2ef536fd41..46364bbeff 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -47,6 +47,7 @@ export enum ApiType { SELECT_PROJECT = 'SELECT_PROJECT', GET_CURRENT_USER = 'GET_CURRENT_USER', UPDATE_CURRENT_USER = 'UPDATE_CURRENT_USER', + DELETE_CURRENT_USER = 'DELETE_CURRENT_USER', GET_PROJECT_ACCESS = 'GET_PROJECT_ACCESS', ADD_PROJECT_MEMBER = 'ADD_PROJECT_MEMBER', UPDATE_PROJECT_MEMBER = 'UPDATE_PROJECT_MEMBER', @@ -72,6 +73,7 @@ export const INTERNAL_API_HISTORY_TYPES = [ ApiType.SELECT_PROJECT, ApiType.GET_CURRENT_USER, ApiType.UPDATE_CURRENT_USER, + ApiType.DELETE_CURRENT_USER, ApiType.GET_PROJECT_ACCESS, ApiType.ADD_PROJECT_MEMBER, ApiType.UPDATE_PROJECT_MEMBER, diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index e334dfc4a0..62177089a8 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -39,6 +39,7 @@ export const typeDefs = gql` SELECT_PROJECT GET_CURRENT_USER UPDATE_CURRENT_USER + DELETE_CURRENT_USER GET_PROJECT_ACCESS ADD_PROJECT_MEMBER UPDATE_PROJECT_MEMBER diff --git a/wren-ui/src/apollo/server/services/organizationMemberService.ts b/wren-ui/src/apollo/server/services/organizationMemberService.ts index 6313fa680e..cc1524844d 100644 --- a/wren-ui/src/apollo/server/services/organizationMemberService.ts +++ b/wren-ui/src/apollo/server/services/organizationMemberService.ts @@ -135,6 +135,7 @@ export interface IOrganizationMemberService { updateCurrentUserProfile( input: UpdateCurrentUserProfileInput, ): Promise; + deleteCurrentUserAccount(): Promise; listCurrentProjectAccess(): Promise<{ members: ProjectAccessMemberSummary[]; availableMembers: ProjectAccessAvailableMember[]; @@ -418,6 +419,18 @@ export class OrganizationMemberService implements IOrganizationMemberService { return this.serializeCurrentUserProfile(user); } + public async deleteCurrentUserAccount(): Promise { + const organization = await this.getCurrentOrganizationOrThrow(); + const currentUserId = await this.getCurrentUserId(organization.id); + if (!currentUserId) { + throw new ApiError('Current user not found', 404); + } + + await this.assertCurrentUserCanDeleteAccount(currentUserId); + await this.userRepository.deleteOne(currentUserId); + return true; + } + public async listCurrentProjectAccess() { const organization = await this.getCurrentOrganizationOrThrow(); const project = await this.projectRepository.getCurrentProject(); @@ -935,6 +948,33 @@ export class OrganizationMemberService implements IOrganizationMemberService { } } + private async assertCurrentUserCanDeleteAccount(userId: number) { + const organizations = await this.organizationRepository.findAll({ + order: 'id', + }); + + for (const organization of organizations) { + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ); + const currentMember = members.find((member) => member.userId === userId); + if (!currentMember || currentMember.organizationRole !== 'Admin') { + continue; + } + + const adminCount = members.filter( + (member) => member.organizationRole === 'Admin', + ).length; + if (adminCount <= 1) { + throw new ApiError( + 'If you are the last Organization admin, you cannot delete your account. Assign another Organization admin or delete the organization first.', + 400, + ); + } + } + } + private async getFallbackProjects() { try { const currentProject = await this.projectRepository.getCurrentProject(); diff --git a/wren-ui/src/components/organization/SettingsLayout.tsx b/wren-ui/src/components/organization/SettingsLayout.tsx index fb0b263160..adbd33e992 100644 --- a/wren-ui/src/components/organization/SettingsLayout.tsx +++ b/wren-ui/src/components/organization/SettingsLayout.tsx @@ -61,7 +61,8 @@ export default function OrganizationSettingsLayout({ | 'general' | 'members' | 'danger-zone' - | 'user-profile'; + | 'user-profile' + | 'user-danger-zone'; title: ReactNode; titleExtra?: ReactNode; children: ReactNode; @@ -108,7 +109,11 @@ export default function OrganizationSettingsLayout({ Profile - Danger zone + + + Danger zone + + diff --git a/wren-ui/src/pages/api/v1/users/current.ts b/wren-ui/src/pages/api/v1/users/current.ts index 812e0aeaa3..6e9346d2d0 100644 --- a/wren-ui/src/pages/api/v1/users/current.ts +++ b/wren-ui/src/pages/api/v1/users/current.ts @@ -29,11 +29,26 @@ export default async function handler( const startTime = Date.now(); try { - assertAllowedMethods(req, ['GET', 'PUT']); + assertAllowedMethods(req, ['GET', 'PUT', 'DELETE']); const organizationMemberService = getOrganizationMemberService(); const projectContext = await getCurrentProjectContext(); const projectId = projectContext.id ?? 0; + if (req.method === 'DELETE') { + await organizationMemberService.deleteCurrentUserAccount(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.DELETE_CURRENT_USER, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + if (req.method === 'PUT') { const user = await organizationMemberService.updateCurrentUserProfile( req.body, @@ -68,9 +83,11 @@ export default async function handler( res, projectId: (await getCurrentProjectContext()).id ?? 0, apiType: - req.method === 'PUT' - ? ApiType.UPDATE_CURRENT_USER - : ApiType.GET_CURRENT_USER, + req.method === 'DELETE' + ? ApiType.DELETE_CURRENT_USER + : req.method === 'PUT' + ? ApiType.UPDATE_CURRENT_USER + : ApiType.GET_CURRENT_USER, requestPayload: req.method === 'PUT' ? req.body : {}, headers: req.headers as Record, startTime, diff --git a/wren-ui/src/pages/user/danger-zone.tsx b/wren-ui/src/pages/user/danger-zone.tsx new file mode 100644 index 0000000000..433e20bb27 --- /dev/null +++ b/wren-ui/src/pages/user/danger-zone.tsx @@ -0,0 +1,96 @@ +import { Button, Modal, Typography, message } from 'antd'; +import { useRouter } from 'next/router'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { Path } from '@/utils/enum'; + +const IntroText = styled(Typography.Text)` + display: block; + margin-top: 12px; + color: var(--gray-7); +`; + +const DangerPanel = styled.div` + margin-top: 24px; + border: 1px solid var(--red-5); + border-radius: 4px; + background: white; +`; + +const DangerRow = styled.div` + display: flex; + align-items: center; + justify-content: space-between; + gap: 24px; + padding: 20px 24px; +`; + +const DangerCopy = styled.div` + min-width: 0; +`; + +const DangerTitle = styled.div` + color: var(--gray-8); + font-weight: 600; + margin-bottom: 6px; +`; + +const DangerDescription = styled(Typography.Text)` + color: var(--gray-6); +`; + +export default function UserDangerZonePage() { + const router = useRouter(); + + const deleteAccount = async () => { + try { + const response = await fetch('/api/v1/users/current', { + method: 'DELETE', + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to delete account'); + } + message.success('Account deleted successfully.'); + await router.push(Path.OrganizationGeneral); + } catch (error: any) { + message.error(error.message || 'Failed to delete account'); + } + }; + + return ( + + In the Danger Zone section, you can delete your account. + + + + + Delete account + + Please be aware that deleting the account will permanently delete + all data and associations, it cannot be undone. + + + + + + + ); +} diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index 9e76697188..18c8d36b1f 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -24,4 +24,5 @@ export enum Path { ProjectAccessControl = '/project/access-control', ProjectCreate = '/projects/create', UserProfile = '/user/profile', + UserDangerZone = '/user/danger-zone', } From 64205f95ccf331dd5c433a3346308556ecab43bd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 8 Jun 2026 15:31:00 +0530 Subject: [PATCH 0134/1087] Add project danger zone page --- .../organization/SettingsLayout.tsx | 7 +- wren-ui/src/pages/project/danger-zone.tsx | 201 ++++++++++++++++++ wren-ui/src/utils/enum/path.ts | 1 + 3 files changed, 208 insertions(+), 1 deletion(-) create mode 100644 wren-ui/src/pages/project/danger-zone.tsx diff --git a/wren-ui/src/components/organization/SettingsLayout.tsx b/wren-ui/src/components/organization/SettingsLayout.tsx index adbd33e992..0d889caa1e 100644 --- a/wren-ui/src/components/organization/SettingsLayout.tsx +++ b/wren-ui/src/components/organization/SettingsLayout.tsx @@ -58,6 +58,7 @@ export default function OrganizationSettingsLayout({ section: | 'project-general' | 'project-access-control' + | 'project-danger-zone' | 'general' | 'members' | 'danger-zone' @@ -83,7 +84,11 @@ export default function OrganizationSettingsLayout({ Data source - Danger zone + + + Danger zone + + Organization diff --git a/wren-ui/src/pages/project/danger-zone.tsx b/wren-ui/src/pages/project/danger-zone.tsx new file mode 100644 index 0000000000..5ad315d2de --- /dev/null +++ b/wren-ui/src/pages/project/danger-zone.tsx @@ -0,0 +1,201 @@ +import { useEffect, useMemo, useState } from 'react'; +import { Button, Modal, Typography, message } from 'antd'; +import { useRouter } from 'next/router'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; +import { useResetCurrentProjectMutation } from '@/apollo/client/graphql/settings.generated'; + +type OrganizationRole = 'Admin' | 'Member'; + +interface ProjectAccessMember { + userId: number; + organizationRole: OrganizationRole; + isCurrentUser: boolean; +} + +interface ProjectAccessResponse { + members: ProjectAccessMember[]; + currentUserId: number | null; + error?: string; +} + +const IntroText = styled(Typography.Text)` + display: block; + margin-top: 12px; + color: var(--gray-7); +`; + +const WarningText = styled(Typography.Text)` + display: block; + margin-top: 12px; + color: var(--red-6); +`; + +const DangerPanel = styled.div` + margin-top: 24px; + border: 1px solid var(--red-5); + border-radius: 4px; + background: white; +`; + +const DangerRow = styled.div` + display: flex; + align-items: center; + justify-content: space-between; + gap: 24px; + padding: 20px 24px; + + & + & { + border-top: 1px solid var(--gray-4); + } +`; + +const DangerCopy = styled.div` + min-width: 0; +`; + +const DangerTitle = styled.div` + color: var(--gray-8); + font-weight: 600; + margin-bottom: 6px; +`; + +const DangerDescription = styled(Typography.Text)` + color: var(--gray-6); +`; + +const readJsonResponse = async (response: Response): Promise => { + const text = await response.text(); + return (text ? JSON.parse(text) : {}) as T; +}; + +export default function ProjectDangerZonePage() { + const router = useRouter(); + const [loading, setLoading] = useState(true); + const [access, setAccess] = useState(null); + const [resetCurrentProject, { loading: resetting, client }] = + useResetCurrentProjectMutation({ + onError: (error) => + message.error(error.message || 'Failed to reset project'), + }); + + const loadProjectAccess = async () => { + setLoading(true); + try { + const response = await fetch('/api/v1/projects/access/current', { + cache: 'no-store', + headers: { + 'Cache-Control': 'no-cache', + }, + }); + const payload = await readJsonResponse(response); + if (!response.ok) { + throw new Error(payload.error || 'Failed to load project access'); + } + setAccess(payload); + } catch (error: any) { + message.error(error.message || 'Failed to load project access'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadProjectAccess(); + }, []); + + const currentMember = useMemo( + () => + access?.members.find( + (member) => + member.isCurrentUser || member.userId === access.currentUserId, + ) || null, + [access], + ); + const canDeleteProject = currentMember?.organizationRole === 'Admin'; + + const resetProject = async () => { + await resetCurrentProject(); + await client.clearStore(); + message.success('Project reset successfully.'); + await router.push(Path.OnboardingConnection); + }; + + const deleteProject = async () => { + await resetCurrentProject(); + await client.clearStore(); + message.success('Project deleted successfully.'); + await router.push(Path.OrganizationGeneral); + }; + + return ( + + + Use Danger Zone for destructive project actions. + Actions in this section cannot be undone. + + + + + Reset project + + Resetting a project removes its current settings and records, + including data connection information, Modeling page + information, and Home page threads. + + + + + + + + Delete project + + Only organization admins can delete a project. Deleting a + project permanently removes access to its ask records for all + owners and members. + + + + + + + + ); +} diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index 18c8d36b1f..2bf5289333 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -22,6 +22,7 @@ export enum Path { OrganizationDangerZone = '/organization/danger-zone', ProjectGeneral = '/project/general', ProjectAccessControl = '/project/access-control', + ProjectDangerZone = '/project/danger-zone', ProjectCreate = '/projects/create', UserProfile = '/user/profile', UserDangerZone = '/user/danger-zone', From 12529695d20d902e2440af154d91aa020f2dd42c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 8 Jun 2026 20:02:12 +0530 Subject: [PATCH 0135/1087] Hide admin and billing navigation items --- wren-ui/src/components/HeaderBar.tsx | 8 -------- wren-ui/src/components/organization/SettingsLayout.tsx | 1 - 2 files changed, 9 deletions(-) diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index dfacb3fd57..53e248612d 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -158,14 +158,6 @@ export default function HeaderBar() { > API - router.push(Path.AdministrationUsers)} - > - Admin - )} diff --git a/wren-ui/src/components/organization/SettingsLayout.tsx b/wren-ui/src/components/organization/SettingsLayout.tsx index 0d889caa1e..323b6f7b8d 100644 --- a/wren-ui/src/components/organization/SettingsLayout.tsx +++ b/wren-ui/src/components/organization/SettingsLayout.tsx @@ -101,7 +101,6 @@ export default function OrganizationSettingsLayout({ Members - Billing Danger zone From 433fe9bbc30a5606759b3e6ce6c8bbafb90a0c7f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 8 Jun 2026 21:02:42 +0530 Subject: [PATCH 0136/1087] Route schema questions to text-to-sql instead of user guide --- wren-ai-service/src/web/v1/services/ask.py | 56 +++++++++++++++++++++- 1 file changed, 54 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 682d53c600..94fca0320f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -184,6 +184,49 @@ def _is_data_analysis_query(self, query: str) -> bool: } return any(term in normalized for term in analysis_terms) + def _is_schema_grounded_query( + self, query: str, db_schemas: Optional[list[str]] = None + ) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + explicit_schema_terms = ( + "table", + "column", + "schema", + "dataset", + "dbo.", + "select ", + " from ", + " join ", + " where ", + " group by ", + " order by ", + ) + if any(term in normalized for term in explicit_schema_terms): + return True + + identifier_tokens = re.findall(r"[a-zA-Z_][a-zA-Z0-9_\.]*", normalized) + if any("." in token for token in identifier_tokens): + return True + + for schema in db_schemas or []: + schema_text = schema.lower() + table_matches = re.findall( + r"create\s+table\s+([a-zA-Z0-9_\.\"]+)", schema_text + ) + column_matches = re.findall(r"\n\s*\"?([a-zA-Z_][a-zA-Z0-9_]*)\"?\s+", schema_text) + candidates = { + token.strip('"') + for token in table_matches + column_matches + if token and len(token.strip('"')) > 2 + } + if any(candidate in normalized for candidate in candidates): + return True + + return False + def _get_unqueryable_metric_message( self, query: str, table_ddls: list[str] ) -> str | None: @@ -408,16 +451,25 @@ async def ask( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") + retrieved_db_schemas = intent_classification_result.get( + "db_schemas" + ) or [] is_original_analytics_query = self._is_data_analysis_query( original_user_query ) + is_schema_grounded_query = self._is_schema_grounded_query( + original_user_query, retrieved_db_schemas + ) or self._is_schema_grounded_query( + rephrased_question or "", retrieved_db_schemas + ) - if intent in {"GENERAL", "MISLEADING_QUERY"} and ( + if intent in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"} and ( is_original_analytics_query + or is_schema_grounded_query or self._is_data_analysis_query(rephrased_question or "") ): logger.info( - "Overriding intent %s to TEXT_TO_SQL for analytics query: %s", + "Overriding intent %s to TEXT_TO_SQL for schema/data query: %s", intent, user_query, ) From 44df955dbffbbfbe2422e0978b396ae981714b04 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 8 Jun 2026 21:17:26 +0530 Subject: [PATCH 0137/1087] Hide administration sidebar module --- wren-ui/src/components/sidebar/index.tsx | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/wren-ui/src/components/sidebar/index.tsx b/wren-ui/src/components/sidebar/index.tsx index 58c44c50e7..b076a18e94 100644 --- a/wren-ui/src/components/sidebar/index.tsx +++ b/wren-ui/src/components/sidebar/index.tsx @@ -9,7 +9,6 @@ import Home, { Props as HomeSidebarProps } from './Home'; import Modeling, { Props as ModelingSidebarProps } from './Modeling'; import Knowledge from './Knowledge'; import APIManagement from './APIManagement'; -import Administration from './Administration'; import LearningSection from '@/components/learning'; const Layout = styled.div` @@ -69,10 +68,6 @@ const DynamicSidebar = ( return ; } - if (pathname.startsWith(Path.Administration)) { - return ; - } - return null; }; @@ -83,7 +78,7 @@ export default function Sidebar(props: Props) { const router = useRouter(); const onSettingsClick = (event) => { - router.push(Path.ProjectGeneral); + router.push(Path.OrganizationGeneral); event.target.blur(); }; From 5e6acb98d895ab81846e954b06a36c1fff872ff0 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 8 Jun 2026 21:38:07 +0530 Subject: [PATCH 0138/1087] Improve PCB analytics retrieval and SQL grounding --- .../src/pipelines/generation/utils/sql.py | 6 ++++++ .../pipelines/retrieval/db_schema_retrieval.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c78a4f1d7d..89efd1c885 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1321,6 +1321,12 @@ async def _classify_generation_result( - For throughput trends across manufacturing/business units, use "dbo_DebugEntries"."BusinessUnit" as the unit dimension and a real debug-entry timestamp such as "dbo_DebugEntries"."DateIn" or "dbo_DebugEntries"."FailedAt" for the trend bucket. Do not use "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH", or invented manufacturing/date fields. - For top/common PCB failure questions, prefer grouping by "dbo_failure_patterns"."name" or "dbo_failure_patterns"."category" and counting "dbo_DebugEntries"."DebugEntryId" after joining "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". - If a useful aggregate already exists in "dbo_failure_patterns" such as "occurrences", it can be used directly for top failure pattern questions without joining event rows. + - For requests such as "show top 10 most common PCB failures", "bar chart of failures by category", or "count of repairs grouped by failure category", generate SQL first. Do not answer with general charting guidance. Return the categorical failure field plus a count metric that can drive a bar chart. + - For failure-category charts, prefer one of these patterns depending on schema availability: + 1. `GROUP BY "dbo_failure_patterns"."category"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` + 2. `GROUP BY "dbo_failure_patterns"."name"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` + 3. `GROUP BY "dbo_repair_logs"."failure_code"` and `COUNT(*)` + - For chart-oriented questions, ensure the final SELECT contains only the chart-ready dimension and metric columns. Avoid prose-like outputs or helper columns. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 2c6a6f322d..6a9ce54be8 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -125,11 +125,24 @@ def _build_view_ddl(content: dict) -> str: def expand_business_terms_for_retrieval(query: str) -> str: normalized = (query or "").lower() pcb_terms = { + "assembly", + "bar chart", + "business unit", + "category", + "common", + "contributor", + "count by", "pcb", "repair", "debug", "turnaround", "failure", + "failure code", + "failure category", + "failure pattern", + "top 10", + "top ten", + "most common", "resolved", "trend", "volume", @@ -151,6 +164,9 @@ def expand_business_terms_for_retrieval(query: str) -> str: "repair trends repair volume repair counts debug entries debug fixes", "average debug hours turnaround time resolved entries failure category failure code", "monthly trend quarter grouped by month bar chart line chart", + "top common pcb failures top 10 failures most common failure categories", + "failure patterns category occurrences debugentryid failuresys material workorder serialnumber", + "dbo_DebugEntries dbo_failure_patterns dbo_repair_logs created_at failedat datein dateout", ] ) From 485b7814d2888c9cf8dfa5c67bd8f56e6f9d567e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 8 Jun 2026 22:11:22 +0530 Subject: [PATCH 0139/1087] Improve SQL routing for chart and failure analytics prompts --- wren-ai-service/src/web/v1/services/ask.py | 68 +++++++++++++++++++--- 1 file changed, 61 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 94fca0320f..bc12f81fc4 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -184,6 +184,50 @@ def _is_data_analysis_query(self, query: str) -> bool: } return any(term in normalized for term in analysis_terms) + def _rewrite_query_for_text_to_sql(self, query: str) -> str: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return query + + guidance: list[str] = [] + + if any( + term in normalized + for term in ("chart", "bar chart", "line chart", "pie chart", "graph") + ): + guidance.append( + "Return SQL only for the aggregated dataset required to build the requested chart." + ) + + if any( + term in normalized + for term in ( + "failure category", + "failure categories", + "common failure", + "common failures", + "failure code", + "top 10", + "most common", + ) + ): + guidance.append( + "Use an exposed failure category, failure name, or failure code field from the schema and return that dimension with a count metric." + ) + + if any( + term in normalized + for term in ("monthly", "last 12 months", "last month", "trend", "volume") + ): + guidance.append( + "Use a real timestamp column from the schema and aggregate results by calendar month when a monthly trend is requested." + ) + + if not guidance: + return query + + return f"{query}\n\nSQL generation guidance:\n- " + "\n- ".join(guidance) + def _is_schema_grounded_query( self, query: str, db_schemas: Optional[list[str]] = None ) -> bool: @@ -356,6 +400,7 @@ async def ask( try: user_query = ask_request.query + sql_user_query = user_query # ask status can be understanding, searching, generating, finished, failed, stopped # we will need to handle business logic for each status @@ -487,6 +532,12 @@ async def ask( elif rephrased_question: user_query = rephrased_question + sql_user_query = ( + self._rewrite_query_for_text_to_sql(user_query) + if self._is_data_analysis_query(user_query) + else user_query + ) + if intent == "MISLEADING_QUERY": general_result = await self._run_with_timeout( "Misleading assistance", @@ -595,10 +646,13 @@ async def ask( retrieval_result = await self._run_with_timeout( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( - query=user_query, + query=sql_user_query, histories=histories, project_id=ask_request.project_id, - enable_column_pruning=enable_column_pruning, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), ), ) _retrieval_result = retrieval_result.get( @@ -681,7 +735,7 @@ async def ask( self._pipelines[ "followup_sql_generation_reasoning" ].run( - query=user_query, + query=sql_user_query, contexts=table_ddls, histories=histories, sql_samples=sql_samples, @@ -704,7 +758,7 @@ async def ask( await self._run_with_timeout( "SQL generation reasoning", self._pipelines["sql_generation_reasoning"].run( - query=user_query, + query=sql_user_query, contexts=table_ddls, sql_samples=sql_samples, instructions=instructions, @@ -774,7 +828,7 @@ async def ask( text_to_sql_generation_results = await self._run_with_timeout( "Follow-up SQL generation", self._pipelines["followup_sql_generation"].run( - query=user_query, + query=sql_user_query, contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, histories=histories, @@ -794,7 +848,7 @@ async def ask( text_to_sql_generation_results = await self._run_with_timeout( "SQL generation", self._pipelines["sql_generation"].run( - query=user_query, + query=sql_user_query, contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, @@ -888,7 +942,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - query=user_query, + query=sql_user_query, ), ) From b99f9a71127a3bd4da632acc291e82ed0eec300e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 8 Jun 2026 22:58:49 +0530 Subject: [PATCH 0140/1087] Add heuristic SQL fallback for PCB analytics prompts --- wren-ai-service/src/web/v1/services/ask.py | 146 +++++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bc12f81fc4..a2b2bf1f79 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -228,6 +228,120 @@ def _rewrite_query_for_text_to_sql(self, query: str) -> str: return f"{query}\n\nSQL generation guidance:\n- " + "\n- ".join(guidance) + def _schema_contains(self, table_ddls: list[str], pattern: str) -> bool: + schema_text = "\n".join(table_ddls or []) + return bool(re.search(pattern, schema_text, flags=re.IGNORECASE)) + + def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: + if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): + return max(1, min(int(match.group(1)), 100)) + return default_value + + def _build_heuristic_text_to_sql_fallback( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + schema_text = "\n".join(table_ddls or []) + wants_chart = any( + term in normalized for term in ("chart", "bar chart", "line chart", "graph") + ) + wants_failure_counts = any( + term in normalized + for term in ( + "failure", + "failure category", + "failure code", + "common pcb failures", + "common failures", + "most common", + "top 10", + "top ten", + ) + ) + wants_monthly_repairs = ( + "repair" in normalized + and any( + term in normalized + for term in ("monthly", "last 12 months", "trend", "volume") + ) + ) + + if wants_failure_counts and wants_chart: + top_n = self._extract_requested_top_n(query) + has_debug_entries = self._schema_contains( + table_ddls, r"\bCREATE\s+TABLE\s+dbo_DebugEntries\b" + ) + has_failure_patterns = self._schema_contains( + table_ddls, r"\bCREATE\s+TABLE\s+dbo_failure_patterns\b" + ) + has_failure_sys = self._schema_contains(table_ddls, r"\bFailureSys\b") + has_debug_entry_id = self._schema_contains(table_ddls, r"\bDebugEntryId\b") + has_pattern_id = self._schema_contains(table_ddls, r"\bid\b") + has_pattern_category = self._schema_contains(table_ddls, r"\bcategory\b") + has_pattern_name = self._schema_contains(table_ddls, r"\bname\b") + + if ( + has_debug_entries + and has_failure_patterns + and has_failure_sys + and has_debug_entry_id + and has_pattern_id + ): + dimension_column = ( + "category" + if ("category" in normalized and has_pattern_category) + else ("name" if has_pattern_name else "category") + ) + if dimension_column == "category" and not has_pattern_category: + dimension_column = "name" + + return ( + f'SELECT TOP {top_n} ' + f'"dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' + f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' + f'FROM "dbo_DebugEntries" ' + f'JOIN "dbo_failure_patterns" ' + f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' + f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' + f'ORDER BY "repair_count" DESC' + ) + + has_repair_logs = self._schema_contains( + table_ddls, r"\bCREATE\s+TABLE\s+dbo_repair_logs\b" + ) + has_failure_code = self._schema_contains(table_ddls, r"\bfailure_code\b") + if has_repair_logs and has_failure_code: + return ( + f'SELECT TOP {top_n} ' + f'"dbo_repair_logs"."failure_code" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_repair_logs" ' + f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' + f'GROUP BY "dbo_repair_logs"."failure_code" ' + f'ORDER BY "repair_count" DESC' + ) + + if wants_monthly_repairs and self._schema_contains( + table_ddls, r"\bCREATE\s+TABLE\s+dbo_repair_logs\b" + ) and self._schema_contains(table_ddls, r"\bcreated_at\b"): + return ( + 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."created_at" >= DATEADD(month, -12, GETDATE()) ' + 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + ) + + return None + def _is_schema_grounded_query( self, query: str, db_schemas: Optional[list[str]] = None ) -> bool: @@ -983,6 +1097,38 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls + ): + logger.info( + "Using heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + api_results = [ + AskResult( + **{ + "sql": heuristic_sql, + "type": "llm", + } + ) + ] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( From 8e26eb3a24653ddf032fe26d2750c647cb2c8571 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 8 Jun 2026 23:36:49 +0530 Subject: [PATCH 0141/1087] Broaden heuristic fallback for PCB analytics queries --- wren-ai-service/src/web/v1/services/ask.py | 111 ++++++++++++++++++--- 1 file changed, 96 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 a2b2bf1f79..5198a42c6d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -228,8 +228,15 @@ def _rewrite_query_for_text_to_sql(self, query: str) -> str: return f"{query}\n\nSQL generation guidance:\n- " + "\n- ".join(guidance) - def _schema_contains(self, table_ddls: list[str], pattern: str) -> bool: + def _schema_contains( + self, + table_ddls: list[str], + pattern: str, + table_names: Optional[list[str]] = None, + ) -> bool: schema_text = "\n".join(table_ddls or []) + if table_names: + schema_text += "\n" + "\n".join(table_names) return bool(re.search(pattern, schema_text, flags=re.IGNORECASE)) def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: @@ -238,13 +245,15 @@ def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: return default_value def _build_heuristic_text_to_sql_fallback( - self, query: str, table_ddls: list[str] + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, ) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None - schema_text = "\n".join(table_ddls or []) wants_chart = any( term in normalized for term in ("chart", "bar chart", "line chart", "graph") ) @@ -272,16 +281,26 @@ def _build_heuristic_text_to_sql_fallback( if wants_failure_counts and wants_chart: top_n = self._extract_requested_top_n(query) has_debug_entries = self._schema_contains( - table_ddls, r"\bCREATE\s+TABLE\s+dbo_DebugEntries\b" + table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names ) has_failure_patterns = self._schema_contains( - table_ddls, r"\bCREATE\s+TABLE\s+dbo_failure_patterns\b" + table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names + ) + has_failure_sys = self._schema_contains( + table_ddls, r"\bFailureSys\b", table_names=table_names + ) + has_debug_entry_id = self._schema_contains( + table_ddls, r"\bDebugEntryId\b", table_names=table_names + ) + has_pattern_id = self._schema_contains( + table_ddls, r"\bid\b", table_names=table_names + ) + has_pattern_category = self._schema_contains( + table_ddls, r"\bcategory\b", table_names=table_names + ) + has_pattern_name = self._schema_contains( + table_ddls, r"\bname\b", table_names=table_names ) - has_failure_sys = self._schema_contains(table_ddls, r"\bFailureSys\b") - has_debug_entry_id = self._schema_contains(table_ddls, r"\bDebugEntryId\b") - has_pattern_id = self._schema_contains(table_ddls, r"\bid\b") - has_pattern_category = self._schema_contains(table_ddls, r"\bcategory\b") - has_pattern_name = self._schema_contains(table_ddls, r"\bname\b") if ( has_debug_entries @@ -311,9 +330,11 @@ def _build_heuristic_text_to_sql_fallback( ) has_repair_logs = self._schema_contains( - table_ddls, r"\bCREATE\s+TABLE\s+dbo_repair_logs\b" + table_ddls, r"\bdbo_repair_logs\b", table_names=table_names + ) + has_failure_code = self._schema_contains( + table_ddls, r"\bfailure_code\b", table_names=table_names ) - has_failure_code = self._schema_contains(table_ddls, r"\bfailure_code\b") if has_repair_logs and has_failure_code: return ( f'SELECT TOP {top_n} ' @@ -326,8 +347,37 @@ def _build_heuristic_text_to_sql_fallback( ) if wants_monthly_repairs and self._schema_contains( - table_ddls, r"\bCREATE\s+TABLE\s+dbo_repair_logs\b" - ) and self._schema_contains(table_ddls, r"\bcreated_at\b"): + table_ddls, r"\bdbo_repair_logs\b", table_names=table_names + ) and self._schema_contains( + table_ddls, r"\bcreated_at\b", table_names=table_names + ): + return ( + 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."created_at" >= DATEADD(month, -12, GETDATE()) ' + 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + ) + + if wants_failure_counts and wants_chart: + top_n = self._extract_requested_top_n(query) + return ( + f'SELECT TOP {top_n} ' + f'"dbo_failure_patterns"."category" AS "failure_category", ' + f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' + f'FROM "dbo_DebugEntries" ' + f'JOIN "dbo_failure_patterns" ' + f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' + f'WHERE "dbo_failure_patterns"."category" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."category" ' + f'ORDER BY "repair_count" DESC' + ) + + if wants_monthly_repairs: return ( 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' @@ -808,6 +858,37 @@ async def ask( return results if not documents: + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", + query_id, + user_query, + ) + api_results = [ + AskResult( + **{ + "sql": heuristic_sql, + "type": "llm", + } + ) + ] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( @@ -1098,7 +1179,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" else: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls + user_query, table_ddls, table_names=table_names ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", From 31480bde4c85f452f1888dc62044e8d79c523192 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 8 Jun 2026 23:51:18 +0530 Subject: [PATCH 0142/1087] Set asking task timestamps before insert --- .../repositories/askingTaskRepository.ts | 46 ++++++++++++++++++- 1 file changed, 45 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts index 942f445ee1..ff98b7502a 100644 --- a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts +++ b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -44,6 +48,35 @@ export class AskingTaskRepository return this.findOneBy({ queryId }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany(data.map(this.withTimestamps), queryOptions); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + protected override transformFromDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); @@ -74,4 +107,15 @@ export class AskingTaskRepository }); return mapKeys(transformedData, (_value, key) => snakeCase(key)); }; + + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; } From 150582413d1bf10017590a5c92d561dada2c0f63 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 00:23:07 +0530 Subject: [PATCH 0143/1087] Generate asking task ids for MSSQL migration tables --- .../repositories/askingTaskRepository.ts | 87 ++++++++++++++++++- 1 file changed, 85 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts index ff98b7502a..6d7e495018 100644 --- a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts +++ b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts @@ -39,6 +39,7 @@ export class AskingTaskRepository implements IAskingTaskRepository { private readonly jsonbColumns = ['detail']; + private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'asking_task' }); @@ -52,14 +53,20 @@ export class AskingTaskRepository data: Partial, queryOptions?: IQueryOptions, ): Promise { - return super.createOne(this.withTimestamps(data), queryOptions); + return super.createOne( + await this.withMssqlId(this.withTimestamps(data), queryOptions), + queryOptions, + ); } public override async createMany( data: Partial[], queryOptions?: IQueryOptions, ): Promise { - return super.createMany(data.map(this.withTimestamps), queryOptions); + return super.createMany( + await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), + queryOptions, + ); } public override async updateOne( @@ -118,4 +125,80 @@ export class AskingTaskRepository updatedAt: data.updatedAt ?? now, }; }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; } From 8f0f9b2aac31794195a15eca4f9df9ddb99e07d5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 00:36:10 +0530 Subject: [PATCH 0144/1087] Generate thread response ids for MSSQL migration tables --- .../repositories/threadResponseRepository.ts | 116 ++++++++++++++++-- 1 file changed, 107 insertions(+), 9 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts index bc490cd1b3..d68757e51b 100644 --- a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts @@ -4,7 +4,7 @@ import { IBasicRepository, IQueryOptions, } from './baseRepository'; -import { camelCase, isPlainObject, mapKeys, mapValues, snakeCase } from 'lodash'; +import { camelCase, isPlainObject, mapKeys, mapValues } from 'lodash'; import { AskResultStatus } from '@server/models/adaptor'; export interface DetailStep { @@ -73,6 +73,8 @@ export interface ThreadResponse { breakdownDetail?: ThreadResponseBreakdownDetail; // Thread response breakdown detail chartDetail?: ThreadResponseChartDetail; // Thread response chart detail adjustment?: ThreadResponseAdjustment; // Thread response adjustment + createdAt?: Date; + updatedAt?: Date; } export interface IThreadResponseRepository @@ -93,11 +95,32 @@ export class ThreadResponseRepository 'chartDetail', 'adjustment', ]; + private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'thread_response' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne( + await this.withMssqlId(this.withTimestamps(data), queryOptions), + queryOptions, + ); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany( + await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), + queryOptions, + ); + } + public async getResponsesWithThread(threadId: number, limit?: number) { const query = this.knex(this.tableName) .select('thread_response.*') @@ -168,6 +191,7 @@ export class ThreadResponseRepository ? JSON.stringify(data.chartDetail) : undefined, adjustment: data.adjustment ? JSON.stringify(data.adjustment) : undefined, + updatedAt: new Date(), }; const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const [result] = await executer(this.tableName) @@ -196,16 +220,90 @@ export class ThreadResponseRepository return formattedData; }; - protected override transformToDBData = (data: Partial) => { - if (!isPlainObject(data)) { - throw new Error('Unexpected dbdata'); + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; } - const transformedData = mapValues(data, (value, key) => { - if (this.jsonbColumns.includes(key)) { - return value ? JSON.stringify(value) : value; + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; } - return value; + return { + ...item, + id: nextId++, + }; }); - return mapKeys(transformedData, (_value, key) => snakeCase(key)); }; } From e80669404a64b4f67e007d43317ca3ee29820d47 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 01:01:21 +0530 Subject: [PATCH 0145/1087] Generate thread ids for MSSQL migration tables --- .../server/repositories/threadRepository.ts | 129 +++++++++++++++++- 1 file changed, 128 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/repositories/threadRepository.ts b/wren-ui/src/apollo/server/repositories/threadRepository.ts index 65720ba5fb..d3fc858ae4 100644 --- a/wren-ui/src/apollo/server/repositories/threadRepository.ts +++ b/wren-ui/src/apollo/server/repositories/threadRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -24,6 +28,8 @@ export interface Thread { questions?: ThreadRecommendationQuestionResult[]; // Recommended questions questionsStatus?: string; // Status of the recommended questions questionsError?: object; // Error of the recommended questions + createdAt?: Date; + updatedAt?: Date; } export interface IThreadRepository extends IBasicRepository { @@ -35,11 +41,47 @@ export class ThreadRepository implements IThreadRepository { private readonly jsonbColumns = ['questions', 'questionsError']; + private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'thread' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne( + await this.withMssqlId(this.withTimestamps(data), queryOptions), + queryOptions, + ); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany( + await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), + queryOptions, + ); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + public async listAllTimeDescOrder(projectId: number): Promise { const threads = await this.knex(this.tableName) .where(this.transformToDBData({ projectId })) @@ -78,4 +120,89 @@ export class ThreadRepository }); return mapKeys(transformedData, (_value, key) => snakeCase(key)); }; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; } From cecc8ab2149d5e968716122f8d5e5737dbaaacc5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 12:05:45 +0530 Subject: [PATCH 0146/1087] Guard MSSQL repair cost asks without cost columns --- wren-ai-service/src/web/v1/services/ask.py | 72 +++++++++++++ .../services/test_ask_unqueryable_metrics.py | 102 ++++++++++++++++++ 2 files changed, 174 insertions(+) create mode 100644 wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 5198a42c6d..3dfe881996 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -169,6 +169,7 @@ def _is_data_analysis_query(self, query: str) -> bool: "chart", "compare", "count", + "cost", "debug", "failure", "group", @@ -239,6 +240,43 @@ def _schema_contains( schema_text += "\n" + "\n".join(table_names) return bool(re.search(pattern, schema_text, flags=re.IGNORECASE)) + def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: + column_names: list[str] = [] + non_column_prefixes = ( + "create ", + "constraint ", + "foreign ", + "primary ", + "unique ", + "index ", + ")", + "/*", + "--", + ) + + for ddl in table_ddls: + for line in ddl.splitlines(): + stripped = line.strip().rstrip(",") + if not stripped: + continue + if stripped.lower().startswith(non_column_prefixes): + continue + + column_match = re.match( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_]*))\s+", + stripped, + ) + if not column_match: + continue + + column_name = next( + value for value in column_match.groupdict().values() if value + ) + column_names.append(column_name.lower()) + + return column_names + def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): return max(1, min(int(match.group(1)), 100)) @@ -440,10 +478,44 @@ def _get_unqueryable_metric_message( ) -> str | None: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) normalized_schema = re.sub(r"\s+", " ", " ".join(table_ddls).lower()) + schema_column_names = self._extract_schema_column_names(table_ddls) if not normalized_query: return None + repair_cost_terms = ( + "repair cost", + "repair_cost", + "repaircost", + "cost", + "cost impact", + "cost_impact", + ) + if any(term in normalized_query for term in repair_cost_terms): + cost_field_patterns = ( + r"\brepair[_ ]?cost\b", + r"\bcost[_ ]?impact\b", + r"\bcost[_ ]?amount\b", + r"\btotal[_ ]?cost\b", + r"\bunit[_ ]?cost\b", + r"\bcost\b", + r"\bamount\b", + ) + has_cost_field = any( + re.search(pattern, column_name) + for pattern in cost_field_patterns + for column_name in schema_column_names + ) + + if not has_cost_field: + return ( + "The schema does not expose repair cost as a queryable " + "column. The MSSQL Wren/Ibis runtime cannot extract cost " + "from generic JSON/text fields such as data. Add repair " + "cost as a first-class column or calculated field, then " + "ask again." + ) + first_pass_yield_terms = ( "first pass yield", "first-pass yield", diff --git a/wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py b/wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py new file mode 100644 index 0000000000..e13af38fc6 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py @@ -0,0 +1,102 @@ +from src.web.v1.services.ask import AskService + + +def test_first_pass_yield_requires_queryable_attempt_fields(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + org_id VARCHAR, + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP, + data JSON + ); + """ + ] + + message = service._get_unqueryable_metric_message( + "Show First Pass Yield percentage trend over time.", + table_ddls, + ) + + assert message + assert "first-pass yield" in message.lower() + assert "first-class columns" in message + + +def test_first_pass_yield_guard_allows_queryable_attempt_fields(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + created_at TIMESTAMP, + attempt_number INTEGER, + pass_fail VARCHAR + ); + """ + ] + + assert ( + service._get_unqueryable_metric_message( + "Show FPY trend over time.", + table_ddls, + ) + is None + ) + + +def test_repair_cost_requires_queryable_cost_field(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + org_id VARCHAR, + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP, + data JSON + ); + """ + ] + + message = service._get_unqueryable_metric_message( + "Create a line chart comparing repair cost and turnaround time.", + table_ddls, + ) + + assert message + assert "repair cost" in message.lower() + assert "first-class column" in message + assert "JSON/text" in message + + +def test_repair_cost_guard_allows_queryable_cost_field(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + repair_cost DOUBLE, + created_at TIMESTAMP, + updated_at TIMESTAMP + ); + """ + ] + + assert ( + service._get_unqueryable_metric_message( + "Create a line chart comparing repair cost and turnaround time.", + table_ddls, + ) + is None + ) From 16c7bde64eb14691870d6f651b6cc65bf4684480 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 14:49:22 +0530 Subject: [PATCH 0147/1087] Route manufacturing throughput asks to debug entries --- wren-ai-service/src/web/v1/services/ask.py | 108 ++++++++++++++++++ .../test_ask_heuristic_text_to_sql.py | 56 +++++++++ 2 files changed, 164 insertions(+) create mode 100644 wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3dfe881996..893d3c41e2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -282,6 +282,78 @@ def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: return max(1, min(int(match.group(1)), 100)) return default_value + def _build_manufacturing_throughput_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_throughput = "throughput" in normalized or ( + "repair" in normalized and "volume" in normalized + ) + wants_unit_breakdown = any( + term in normalized + for term in ( + "manufacturing unit", + "manufacturing units", + "business unit", + "business units", + "different unit", + "different units", + ) + ) + + if not (wants_throughput and wants_unit_breakdown): + return None + + has_debug_entries = self._schema_contains( + table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names + ) + has_business_unit = self._schema_contains( + table_ddls, r"\bBusinessUnit\b", table_names=table_names + ) + if not (has_debug_entries and has_business_unit): + return None + + timestamp_column = None + for candidate in ("DateIn", "FailedAt"): + if self._schema_contains( + table_ddls, rf"\b{candidate}\b", table_names=table_names + ): + timestamp_column = candidate + break + + if timestamp_column and any( + term in normalized for term in ("trend", "monthly", "over time") + ): + timestamp_expression = f'"dbo_DebugEntries"."{timestamp_column}"' + return ( + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' + f'DATEPART(YEAR, {timestamp_expression}) AS "year", ' + f'DATEPART(MONTH, {timestamp_expression}) AS "month", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries" ' + 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' + f'AND {timestamp_expression} IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' + f'DATEPART(YEAR, {timestamp_expression}), ' + f'DATEPART(MONTH, {timestamp_expression}) ' + 'ORDER BY "unit_name" ASC, "year" ASC, "month" ASC' + ) + + return ( + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries" ' + 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries"."BusinessUnit" ' + 'ORDER BY "throughput" DESC' + ) + def _build_heuristic_text_to_sql_fallback( self, query: str, @@ -292,6 +364,11 @@ def _build_heuristic_text_to_sql_fallback( if not normalized: return None + if throughput_sql := self._build_manufacturing_throughput_sql( + query, table_ddls, table_names=table_names + ): + return throughput_sql + wants_chart = any( term in normalized for term in ("chart", "bar chart", "line chart", "graph") ) @@ -929,6 +1006,37 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + if heuristic_sql := self._build_manufacturing_throughput_sql( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using manufacturing throughput fallback for query_id %s: %s", + query_id, + user_query, + ) + api_results = [ + AskResult( + **{ + "sql": heuristic_sql, + "type": "llm", + } + ) + ] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py new file mode 100644 index 0000000000..673aa65fc1 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -0,0 +1,56 @@ +from src.web.v1.services.ask import AskService + + +def test_manufacturing_throughput_trend_uses_debug_entry_business_unit(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_DebugEntries ( + DebugEntryId VARCHAR, + BusinessUnit VARCHAR, + DateIn TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_batch_records ( + id VARCHAR, + board_model VARCHAR, + production_date TIMESTAMP + ); + """, + ] + + sql = service._build_manufacturing_throughput_sql( + "Show throughput trends across different manufacturing units.", + table_ddls, + table_names=["dbo_DebugEntries", "dbo_batch_records"], + ) + + assert sql + assert '"dbo_DebugEntries"."BusinessUnit"' in sql + assert '"dbo_DebugEntries"."DateIn"' in sql + assert "dbo_batch_records" not in sql + assert 'COUNT(*) AS "throughput"' in sql + assert "DATEPART(MONTH" in sql + + +def test_manufacturing_throughput_fallback_requires_business_unit_column(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_batch_records ( + id VARCHAR, + board_model VARCHAR, + production_date TIMESTAMP + ); + """ + ] + + assert ( + service._build_manufacturing_throughput_sql( + "Show throughput trends across different manufacturing units.", + table_ddls, + table_names=["dbo_batch_records"], + ) + is None + ) From 61f33b0dc477e9eea10384041f5aede853da0824 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 15:23:27 +0530 Subject: [PATCH 0148/1087] Route repair failure count asks to failure code --- wren-ai-service/src/web/v1/services/ask.py | 124 ++++++++++++++++++ .../test_ask_heuristic_text_to_sql.py | 56 ++++++++ 2 files changed, 180 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 893d3c41e2..683cc53d46 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -354,6 +354,94 @@ def _build_manufacturing_throughput_sql( 'ORDER BY "throughput" DESC' ) + def _build_repair_failure_count_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_failure_counts = ( + "failure" in normalized + and any( + term in normalized + for term in ("count", "counts", "category", "code", "grouped") + ) + and any(term in normalized for term in ("repair", "bar chart", "chart")) + ) + if not wants_failure_counts: + return None + + top_n = self._extract_requested_top_n(query) + has_repair_logs = self._schema_contains( + table_ddls, r"\bdbo_repair_logs\b", table_names=table_names + ) + has_failure_code = self._schema_contains( + table_ddls, r"\bfailure_code\b", table_names=table_names + ) + if has_repair_logs and has_failure_code: + return ( + f'SELECT TOP {top_n} ' + f'"dbo_repair_logs"."failure_code" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_repair_logs" ' + f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' + f'GROUP BY "dbo_repair_logs"."failure_code" ' + f'ORDER BY "repair_count" DESC' + ) + + has_debug_entries = self._schema_contains( + table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names + ) + has_failure_patterns = self._schema_contains( + table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names + ) + has_failure_sys = self._schema_contains( + table_ddls, r"\bFailureSys\b", table_names=table_names + ) + has_debug_entry_id = self._schema_contains( + table_ddls, r"\bDebugEntryId\b", table_names=table_names + ) + has_pattern_id = self._schema_contains( + table_ddls, r"\bid\b", table_names=table_names + ) + has_pattern_category = self._schema_contains( + table_ddls, r"\bcategory\b", table_names=table_names + ) + has_pattern_name = self._schema_contains( + table_ddls, r"\bname\b", table_names=table_names + ) + + if ( + has_debug_entries + and has_failure_patterns + and has_failure_sys + and has_debug_entry_id + and has_pattern_id + and (has_pattern_category or has_pattern_name) + ): + dimension_column = ( + "category" + if ("category" in normalized and has_pattern_category) + else ("name" if has_pattern_name else "category") + ) + return ( + f'SELECT TOP {top_n} ' + f'"dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' + f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' + f'FROM "dbo_DebugEntries" ' + f'JOIN "dbo_failure_patterns" ' + f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' + f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' + f'ORDER BY "repair_count" DESC' + ) + + return None + def _build_heuristic_text_to_sql_fallback( self, query: str, @@ -369,6 +457,11 @@ def _build_heuristic_text_to_sql_fallback( ): return throughput_sql + if repair_failure_count_sql := self._build_repair_failure_count_sql( + query, table_ddls, table_names=table_names + ): + return repair_failure_count_sql + wants_chart = any( term in normalized for term in ("chart", "bar chart", "line chart", "graph") ) @@ -1037,6 +1130,37 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + if heuristic_sql := self._build_repair_failure_count_sql( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using repair failure-count fallback for query_id %s: %s", + query_id, + user_query, + ) + api_results = [ + AskResult( + **{ + "sql": heuristic_sql, + "type": "llm", + } + ) + ] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 673aa65fc1..1d6b41e781 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -54,3 +54,59 @@ def test_manufacturing_throughput_fallback_requires_business_unit_column(): ) is None ) + + +def test_repair_failure_count_uses_repair_log_failure_code(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + created_at TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_reports ( + id VARCHAR, + name VARCHAR + ); + """, + ] + + sql = service._build_repair_failure_count_sql( + "Create a bar chart of repair counts grouped by failure category.", + table_ddls, + table_names=["dbo_repair_logs", "dbo_reports"], + ) + + assert sql + assert '"dbo_repair_logs"."failure_code" AS "failure_category"' in sql + assert 'COUNT(*) AS "repair_count"' in sql + assert "Failure Category" not in sql + assert "dbo_reports" not in sql + + +def test_repair_failure_count_requires_schema_backed_failure_dimension(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + status VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + assert ( + service._build_repair_failure_count_sql( + "Create a bar chart of repair counts grouped by failure category.", + table_ddls, + table_names=["dbo_repair_logs"], + ) + is None + ) From 9f8ed641664f450df84bd211cd19dab5175def38 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 16:57:03 +0530 Subject: [PATCH 0149/1087] Start answer generation when answer detail is missing --- .../src/components/pages/home/promptThread/AnswerResult.tsx | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx index 19cad3e1d2..697016d966 100644 --- a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx +++ b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx @@ -177,6 +177,8 @@ const AdjustmentInformation = (props: { }; const isNeedGenerateAnswer = (answerDetail: ThreadResponseAnswerDetail) => { + if (!answerDetail) return true; + const isFinished = getAnswerIsFinished(answerDetail?.status); // it means the background task has not started yet, but answer is pending for generating const isProcessing = [ @@ -184,7 +186,7 @@ const isNeedGenerateAnswer = (answerDetail: ThreadResponseAnswerDetail) => { ThreadResponseAnswerStatus.PREPROCESSING, ThreadResponseAnswerStatus.FETCHING_DATA, ].includes(answerDetail?.status); - return answerDetail?.queryId === null && !isFinished && !isProcessing; + return !answerDetail?.queryId && !isFinished && !isProcessing; }; const isAnswerGenerationInProgress = ( From 45e8c823259565da581209be07d9e1df90290c00 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 17:46:39 +0530 Subject: [PATCH 0150/1087] Bypass intent classification for common PCB failure counts --- wren-ai-service/src/web/v1/services/ask.py | 95 ++++++++++++++++++- .../test_ask_heuristic_text_to_sql.py | 41 ++++++++ 2 files changed, 135 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 683cc53d46..383d3846af 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -368,7 +368,16 @@ def _build_repair_failure_count_sql( "failure" in normalized and any( term in normalized - for term in ("count", "counts", "category", "code", "grouped") + for term in ( + "count", + "counts", + "category", + "code", + "grouped", + "common", + "most common", + "top", + ) ) and any(term in normalized for term in ("repair", "bar chart", "chart")) ) @@ -442,6 +451,28 @@ def _build_repair_failure_count_sql( return None + def _is_direct_heuristic_sql_query(self, query: str) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + asks_manufacturing_throughput = ( + "throughput" in normalized + and any(term in normalized for term in ("manufacturing", "unit", "units")) + ) + asks_failure_counts = ( + "failure" in normalized + and any( + term in normalized + for term in ("count", "counts", "common", "most common", "top") + ) + and any( + term in normalized + for term in ("pcb", "repair", "bar chart", "chart", "category") + ) + ) + return asks_manufacturing_throughput or asks_failure_counts + def _build_heuristic_text_to_sql_fallback( self, query: str, @@ -832,6 +863,68 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results + if self._is_direct_heuristic_sql_query(user_query): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents = _retrieval_result.get("retrieval_results", []) + table_names = [ + document.get("table_name") for document in documents + ] + table_ddls = [ + document.get("table_ddl", "") or "" for document in documents + ] + logger.info( + "Retrieved tables for direct heuristic query_id %s: %s", + query_id, + table_names, + ) + + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using direct heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + api_results = [ + AskResult( + **{ + "sql": heuristic_sql, + "type": "llm", + } + ) + ] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + historical_question = await self._run_with_timeout( "Historical question retrieval", self._pipelines["historical_question"].run( diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 1d6b41e781..6429c4f143 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -89,6 +89,47 @@ def test_repair_failure_count_uses_repair_log_failure_code(): assert "dbo_reports" not in sql +def test_common_pcb_failures_uses_repair_log_failure_code(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + created_at TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_DebugEntries ( + DebugEntryId VARCHAR, + FailureSys VARCHAR + ); + """, + ] + + sql = service._build_repair_failure_count_sql( + "Show top 10 most common PCB failures in a bar chart.", + table_ddls, + table_names=["dbo_repair_logs", "dbo_DebugEntries"], + ) + + assert sql + assert sql.startswith('SELECT TOP 10 ') + assert '"dbo_repair_logs"."failure_code" AS "failure_category"' in sql + assert 'COUNT(*) AS "repair_count"' in sql + assert "FailureSys" not in sql + + +def test_common_pcb_failures_uses_direct_heuristic_route(): + service = AskService(pipelines={}) + + assert service._is_direct_heuristic_sql_query( + "Show top 10 most common PCB failures in a bar chart." + ) + + def test_repair_failure_count_requires_schema_backed_failure_dimension(): service = AskService(pipelines={}) table_ddls = [ From fd93e009370439e9d93a4f90b77091a4023c92d9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 19:12:03 +0530 Subject: [PATCH 0151/1087] Generate project ids for MSSQL app database --- .../server/repositories/projectRepository.ts | 102 +++++++++++++++++- 1 file changed, 100 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 8555704f49..5a37ac6f59 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -2,6 +2,7 @@ import { Knex } from 'knex'; import { BaseRepository, IBasicRepository, + IQueryOptions, coerceBoolean, } from './baseRepository'; import { @@ -205,11 +206,32 @@ export class ProjectRepository implements IProjectRepository { private jsonTypeColumns = ['questions', 'questions_error', 'connection_info']; + private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'project' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne( + await this.withMssqlId(data, queryOptions), + queryOptions, + ); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany( + await this.withMssqlIds(data, queryOptions), + queryOptions, + ); + } + public async getCurrentProject() { const currentProject = await this.findCurrentProject(); if (currentProject) { @@ -283,8 +305,8 @@ export class ProjectRepository return camelCaseData as Project; }; - public override transformToDBData: (data: Project) => any = ( - data: Project, + public override transformToDBData: (data: Partial) => any = ( + data: Partial, ) => { if (!isPlainObject(data)) { throw new Error('Unexpected db data'); @@ -298,4 +320,80 @@ export class ProjectRepository }); return formattedData; }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; } From b0ff168b8c1495a0f8ccf47c3f01409ce01f2c0e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 9 Jun 2026 19:37:55 +0530 Subject: [PATCH 0152/1087] Set project timestamps for MSSQL inserts --- .../server/repositories/projectRepository.ts | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 5a37ac6f59..ce83fb3496 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -192,6 +192,8 @@ export interface Project { questionsError?: object; projectType?: WorkspaceProjectType; isCurrent?: boolean; + createdAt?: Date; + updatedAt?: Date; } export interface IProjectRepository extends IBasicRepository { @@ -217,7 +219,7 @@ export class ProjectRepository queryOptions?: IQueryOptions, ): Promise { return super.createOne( - await this.withMssqlId(data, queryOptions), + await this.withMssqlId(this.withTimestamps(data), queryOptions), queryOptions, ); } @@ -227,7 +229,22 @@ export class ProjectRepository queryOptions?: IQueryOptions, ): Promise { return super.createMany( - await this.withMssqlIds(data, queryOptions), + await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), + queryOptions, + ); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, queryOptions, ); } @@ -324,6 +341,15 @@ export class ProjectRepository private isMssql = () => String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + private hasIdentityId = async (): Promise => { if (!this.isMssql()) { return true; From 6ea736ea4d90c9f2945bb2cb5579a41eec9d5e1a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 9 Jun 2026 18:19:49 +0530 Subject: [PATCH 0153/1087] Handle manual MSSQL ids for migrated tables --- .../server/repositories/baseRepository.ts | 84 +++++++++++++++++-- 1 file changed, 79 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index f7c211206b..7d81514e9b 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -57,6 +57,7 @@ export const coerceBoolean = (value: unknown): boolean => { export class BaseRepository implements IBasicRepository { protected knex: Knex; protected tableName: string; + private hasIdColumnCache: boolean | null = null; constructor({ knexPg, tableName }: { knexPg: Knex; tableName: string }) { this.knex = knexPg; @@ -118,23 +119,25 @@ export class BaseRepository implements IBasicRepository { public async createOne(data: Partial, queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const insertValue = await this.prepareInsertData(data, executer); const [result] = await executer(this.tableName) - .insert(this.transformToDBData(data)) + .insert(insertValue) .returning('*'); return this.transformFromDBData(result); } public async createMany(data: Partial[], queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const preparedData = await this.prepareInsertManyData(data, executer); const batchSize = 100; - const batchCount = Math.ceil(data.length / batchSize); + const batchCount = Math.ceil(preparedData.length / batchSize); const result = []; for (let i = 0; i < batchCount; i++) { const start = i * batchSize; - const end = Math.min((i + 1) * batchSize, data.length); - const batchValues = data.slice(start, end); + const end = Math.min((i + 1) * batchSize, preparedData.length); + const batchValues = preparedData.slice(start, end); const chunk = await executer(this.tableName) - .insert(batchValues.map(this.transformToDBData)) + .insert(batchValues) .returning('*'); result.push(...chunk); } @@ -198,4 +201,75 @@ export class BaseRepository implements IBasicRepository { protected transformFromDBData = (data: any): T => this.defaultTransformFromDBData(data); + + private isMssql(executer: Knex | Knex.Transaction) { + return executer.client.config.client === 'mssql'; + } + + private async hasIdColumn(executer: Knex | Knex.Transaction) { + if (this.hasIdColumnCache !== null) { + return this.hasIdColumnCache; + } + + const hasIdColumn = await executer.schema.hasColumn(this.tableName, 'id'); + this.hasIdColumnCache = hasIdColumn; + return hasIdColumn; + } + + private async getNextId(executer: Knex | Knex.Transaction) { + const [row] = await executer(this.tableName).max<{ maxId: number | null }>( + 'id as maxId', + ); + return (row?.maxId || 0) + 1; + } + + private async prepareInsertData( + data: Partial, + executer: Knex | Knex.Transaction, + ) { + const dbData = this.transformToDBData(data); + if (!this.isMssql(executer)) { + return dbData; + } + + if (!(await this.hasIdColumn(executer)) || dbData.id !== undefined) { + return dbData; + } + + return { + ...dbData, + id: await this.getNextId(executer), + }; + } + + private async prepareInsertManyData( + data: Partial[], + executer: Knex | Knex.Transaction, + ) { + const dbData = data.map((item) => this.transformToDBData(item)); + if (!this.isMssql(executer) || !(await this.hasIdColumn(executer))) { + return dbData; + } + + const missingIdIndexes = dbData.reduce((acc, item, index) => { + if (item.id === undefined) { + acc.push(index); + } + return acc; + }, []); + + if (!missingIdIndexes.length) { + return dbData; + } + + let nextId = await this.getNextId(executer); + for (const index of missingIdIndexes) { + dbData[index] = { + ...dbData[index], + id: nextId++, + }; + } + + return dbData; + } } From acdf0ee697604b8fcabb5d211656a5f7eb567950 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 9 Jun 2026 21:54:35 +0530 Subject: [PATCH 0154/1087] Prefer grounded failure pattern fields for PCB failure charts --- wren-ai-service/src/web/v1/services/ask.py | 57 +++++++++++++++++++--- 1 file changed, 50 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 383d3846af..e4f8a797b6 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -519,6 +519,12 @@ def _build_heuristic_text_to_sql_fallback( if wants_failure_counts and wants_chart: top_n = self._extract_requested_top_n(query) + has_pattern_failure_sys = self._schema_contains( + table_ddls, r"\bFailuresys\b", table_names=table_names + ) + has_pattern_occurrences = self._schema_contains( + table_ddls, r"\boccurrences\b", table_names=table_names + ) has_debug_entries = self._schema_contains( table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names ) @@ -541,6 +547,28 @@ def _build_heuristic_text_to_sql_fallback( table_ddls, r"\bname\b", table_names=table_names ) + if has_failure_patterns and has_pattern_failure_sys and has_pattern_occurrences: + return ( + f'SELECT TOP {top_n} ' + f'"dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'"dbo_failure_patterns"."occurrences" AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'AND "dbo_failure_patterns"."occurrences" IS NOT NULL ' + f'ORDER BY "dbo_failure_patterns"."occurrences" DESC' + ) + + if has_failure_patterns and has_pattern_failure_sys: + return ( + f'SELECT TOP {top_n} ' + f'"dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."Failuresys" ' + f'ORDER BY "repair_count" DESC' + ) + if ( has_debug_entries and has_failure_patterns @@ -606,13 +634,11 @@ def _build_heuristic_text_to_sql_fallback( top_n = self._extract_requested_top_n(query) return ( f'SELECT TOP {top_n} ' - f'"dbo_failure_patterns"."category" AS "failure_category", ' - f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' - f'FROM "dbo_DebugEntries" ' - f'JOIN "dbo_failure_patterns" ' - f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' - f'WHERE "dbo_failure_patterns"."category" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."category" ' + f'"dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."Failuresys" ' f'ORDER BY "repair_count" DESC' ) @@ -1304,6 +1330,23 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using heuristic text-to-sql shortcut after retrieval for query_id %s: %s", + query_id, + user_query, + ) + api_results = [ + AskResult( + **{ + "sql": heuristic_sql, + "type": "llm", + } + ) + ] + if ( not self._is_stopped(query_id, self._ask_results) and not api_results From 76151da9399ffb75d1d6d076e37e0b0ced23b309 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 9 Jun 2026 23:19:48 +0530 Subject: [PATCH 0155/1087] Ground analytics prompts against full deployed project schema --- .../retrieval/db_schema_retrieval.py | 80 ++++++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 79 ------------------ 2 files changed, 78 insertions(+), 81 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6a9ce54be8..ea5c2bb3e2 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -171,6 +171,52 @@ def expand_business_terms_for_retrieval(query: str) -> str: ) +def _is_project_wide_analysis_query(query: str) -> bool: + normalized = (query or "").lower() + if not normalized: + return False + + analysis_terms = { + "average", + "avg", + "bar chart", + "chart", + "compare", + "count", + "counts", + "group by", + "grouped", + "line chart", + "monthly", + "most common", + "pie chart", + "quarter", + "recommend", + "recommended", + "show", + "top", + "trend", + "volume", + } + return any(term in normalized for term in analysis_terms) + + +def _dedupe_documents(documents: list[Document]) -> list[Document]: + deduped: list[Document] = [] + seen: set[tuple[str, str, str]] = set() + for document in documents: + key = ( + str(document.meta.get("name", "")), + str(document.meta.get("type", "")), + document.content, + ) + if key in seen: + continue + seen.add(key) + deduped.append(document) + return deduped + + @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: @@ -275,7 +321,22 @@ async def dbschema_retrieval( results = await dbschema_retriever.run(query_embedding=[], filters=filters) if results.get("documents") or not project_id: - return results["documents"] + documents = results["documents"] + if project_id and _is_project_wide_analysis_query(query): + all_project_results = await dbschema_retriever.run( + query_embedding=[], + filters={ + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": project_id}, + ], + }, + ) + documents = _dedupe_documents( + documents + all_project_results.get("documents", []) + ) + return documents fallback_filters = { "operator": "AND", @@ -287,7 +348,22 @@ async def dbschema_retrieval( fallback_results = await dbschema_retriever.run( query_embedding=[], filters=fallback_filters ) - return fallback_results["documents"] + documents = fallback_results["documents"] + if project_id and _is_project_wide_analysis_query(query): + all_project_results = await dbschema_retriever.run( + query_embedding=[], + filters={ + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": project_id}, + ], + }, + ) + documents = _dedupe_documents( + documents + all_project_results.get("documents", []) + ) + return documents filters = { "operator": "AND", diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e4f8a797b6..3330d80a6a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1218,68 +1218,6 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if heuristic_sql := self._build_manufacturing_throughput_sql( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using manufacturing throughput fallback for query_id %s: %s", - query_id, - user_query, - ) - api_results = [ - AskResult( - **{ - "sql": heuristic_sql, - "type": "llm", - } - ) - ] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - if heuristic_sql := self._build_repair_failure_count_sql( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using repair failure-count fallback for query_id %s: %s", - query_id, - user_query, - ) - api_results = [ - AskResult( - **{ - "sql": heuristic_sql, - "type": "llm", - } - ) - ] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names @@ -1330,23 +1268,6 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql shortcut after retrieval for query_id %s: %s", - query_id, - user_query, - ) - api_results = [ - AskResult( - **{ - "sql": heuristic_sql, - "type": "llm", - } - ) - ] - if ( not self._is_stopped(query_id, self._ask_results) and not api_results From c25c3a51d88f1f830851373488771cdd68692212 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 9 Jun 2026 23:57:50 +0530 Subject: [PATCH 0156/1087] Stabilize grouped bar chart schema fallback --- .../src/pipelines/generation/utils/chart.py | 56 ++++++++++++++++++- 1 file changed, 53 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index aee5d66d13..f972e90095 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -57,7 +57,7 @@ def _normalize_chart_schema_fields(chart_schema: dict, columns: list[str]) -> di normalized = deepcopy(chart_schema) encoding = normalized.get("encoding", {}) - for key in ("x", "y", "color", "xOffset", "theta"): + for key in ("x", "y", "x2", "y2", "color", "xOffset", "theta"): axis = encoding.get(key) if isinstance(axis, dict) and axis.get("field"): axis["field"] = _match_column_name(axis["field"], columns) @@ -207,7 +207,7 @@ def _is_schema_compatible_with_sample_data( columns = set(sample_data[0].keys()) encoding = chart_schema.get("encoding", {}) - for key in ("x", "y", "color", "xOffset", "theta"): + for key in ("x", "y", "x2", "y2", "color", "xOffset", "theta"): axis = encoding.get(key) if isinstance(axis, dict) and axis.get("field") and axis["field"] not in columns: return False @@ -221,6 +221,51 @@ def _is_schema_compatible_with_sample_data( return True +def _needs_deterministic_bar_fallback( + chart_schema: dict, + chart_type: str, + sample_data: list[dict], +) -> bool: + if chart_type not in {"bar", "grouped_bar", "stacked_bar"}: + return False + if not sample_data: + return False + + encoding = chart_schema.get("encoding", {}) if chart_schema else {} + + # Reject range-style bar encodings for simple grouped-count datasets. + for key in ("x2", "y2"): + axis = encoding.get(key) + if isinstance(axis, dict) and axis.get("field"): + return True + + for axis_name in ("x", "y", "color", "xOffset", "theta"): + axis = encoding.get(axis_name) + if not isinstance(axis, dict): + continue + field = axis.get("field", "") + if isinstance(field, str) and ( + field.endswith("_start") or field.endswith("_end") + ): + return True + + inferred = _infer_column_types(sample_data) + quantitative = inferred["quantitative"] + nominal = inferred["nominal"] + + # For the common case "category + count", prefer a deterministic bar spec + # if the model did not produce a usable quantitative y axis. + if len(quantitative) == 1 and len(nominal) >= 1: + y_axis = encoding.get("y") + x_axis = encoding.get("x") + if not isinstance(y_axis, dict) or y_axis.get("field") not in quantitative: + return True + if not isinstance(x_axis, dict) or x_axis.get("field") not in nominal: + return True + + return False + + chart_generation_instructions = """ ### INSTRUCTIONS ### @@ -515,7 +560,12 @@ def run( chart_schema, list(sample_data[0].keys()) if sample_data else [] ) - if not _is_schema_compatible_with_sample_data(chart_schema, sample_data): + if ( + not _is_schema_compatible_with_sample_data(chart_schema, sample_data) + or _needs_deterministic_bar_fallback( + chart_schema, chart_type or "", sample_data + ) + ): chart_schema = _build_fallback_chart_schema( query, chart_type or "bar", sample_data ) From 5bce5637c37d5a73eb6be9f13431da04c05db87b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 10 Jun 2026 00:05:16 +0530 Subject: [PATCH 0157/1087] Use Wren parser compatible limits for repair fallbacks --- wren-ai-service/src/web/v1/services/ask.py | 42 +++++++++---------- .../test_ask_heuristic_text_to_sql.py | 4 +- 2 files changed, 24 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3330d80a6a..3d65b33809 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -393,13 +393,13 @@ def _build_repair_failure_count_sql( ) if has_repair_logs and has_failure_code: return ( - f'SELECT TOP {top_n} ' - f'"dbo_repair_logs"."failure_code" AS "failure_category", ' + f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' f'COUNT(*) AS "repair_count" ' f'FROM "dbo_repair_logs" ' f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' ) has_debug_entries = self._schema_contains( @@ -438,15 +438,15 @@ def _build_repair_failure_count_sql( else ("name" if has_pattern_name else "category") ) return ( - f'SELECT TOP {top_n} ' - f'"dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' + f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' f'FROM "dbo_DebugEntries" ' f'JOIN "dbo_failure_patterns" ' f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' - f'ORDER BY "repair_count" DESC' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' ) return None @@ -549,24 +549,24 @@ def _build_heuristic_text_to_sql_fallback( if has_failure_patterns and has_pattern_failure_sys and has_pattern_occurrences: return ( - f'SELECT TOP {top_n} ' - f'"dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' f'"dbo_failure_patterns"."occurrences" AS "repair_count" ' f'FROM "dbo_failure_patterns" ' f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' f'AND "dbo_failure_patterns"."occurrences" IS NOT NULL ' - f'ORDER BY "dbo_failure_patterns"."occurrences" DESC' + f'ORDER BY "dbo_failure_patterns"."occurrences" DESC ' + f'LIMIT {top_n}' ) if has_failure_patterns and has_pattern_failure_sys: return ( - f'SELECT TOP {top_n} ' - f'"dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' f'COUNT(*) AS "repair_count" ' f'FROM "dbo_failure_patterns" ' f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' ) if ( @@ -585,15 +585,15 @@ def _build_heuristic_text_to_sql_fallback( dimension_column = "name" return ( - f'SELECT TOP {top_n} ' - f'"dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' + f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' f'FROM "dbo_DebugEntries" ' f'JOIN "dbo_failure_patterns" ' f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' - f'ORDER BY "repair_count" DESC' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' ) has_repair_logs = self._schema_contains( @@ -604,13 +604,13 @@ def _build_heuristic_text_to_sql_fallback( ) if has_repair_logs and has_failure_code: return ( - f'SELECT TOP {top_n} ' - f'"dbo_repair_logs"."failure_code" AS "failure_category", ' + f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' f'COUNT(*) AS "repair_count" ' f'FROM "dbo_repair_logs" ' f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' ) if wants_monthly_repairs and self._schema_contains( @@ -633,13 +633,13 @@ def _build_heuristic_text_to_sql_fallback( if wants_failure_counts and wants_chart: top_n = self._extract_requested_top_n(query) return ( - f'SELECT TOP {top_n} ' - f'"dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' f'COUNT(*) AS "repair_count" ' f'FROM "dbo_failure_patterns" ' f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' ) if wants_monthly_repairs: diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 6429c4f143..87e2fa2f09 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -116,10 +116,12 @@ def test_common_pcb_failures_uses_repair_log_failure_code(): ) assert sql - assert sql.startswith('SELECT TOP 10 ') + assert sql.startswith('SELECT "dbo_repair_logs"."failure_code"') assert '"dbo_repair_logs"."failure_code" AS "failure_category"' in sql assert 'COUNT(*) AS "repair_count"' in sql assert "FailureSys" not in sql + assert "TOP 10" not in sql + assert sql.endswith("LIMIT 10") def test_common_pcb_failures_uses_direct_heuristic_route(): From 48bb90183aacd1f0a8ecc1bb4eeec3f00fc18be7 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 10 Jun 2026 18:07:52 +0530 Subject: [PATCH 0158/1087] Guard MDL builder against null metadata --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 30 +++++++++++++-------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 21ebb8e261..bebef38b10 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -103,7 +103,7 @@ export class MDLBuilder implements IMDLBuilder { return; } this.manifest.models = this.models.map((model: Model) => { - const properties = model.properties ? JSON.parse(model.properties) : {}; + const properties = this.parseProperties(model.properties); // put displayName in properties if (model.displayName) { properties.displayName = model.displayName; @@ -136,7 +136,7 @@ export class MDLBuilder implements IMDLBuilder { return; } this.manifest.views = this.views.map((view: View) => { - const properties = JSON.parse(view.properties) || {}; + const properties = this.parseProperties(view.properties); // filter out properties that are not null or undefined // and are in the list of properties that are allowed @@ -193,9 +193,7 @@ export class MDLBuilder implements IMDLBuilder { if (!model.columns) { model.columns = []; } - const properties = column.properties - ? JSON.parse(column.properties) - : {}; + const properties = this.parseProperties(column.properties); // put displayName in properties if (column.displayName) { properties.displayName = column.displayName; @@ -269,7 +267,7 @@ export class MDLBuilder implements IMDLBuilder { isCalculated: true, expression, notNull: column.notNull ? true : false, - properties: JSON.parse(column.properties), + properties: this.parseProperties(column.properties), }; model.columns.push(columnValue); }); @@ -307,7 +305,7 @@ export class MDLBuilder implements IMDLBuilder { isCalculated: true, expression, notNull: calculatedField.notNull ? true : false, - properties: JSON.parse(calculatedField.properties), + properties: this.parseProperties(calculatedField.properties), }; model.columns.push(columnValue); } @@ -335,9 +333,7 @@ export class MDLBuilder implements IMDLBuilder { relation: name, }); - const properties = relation.properties - ? JSON.parse(relation.properties) - : {}; + const properties = this.parseProperties(relation.properties); return { name: name, @@ -466,7 +462,7 @@ export class MDLBuilder implements IMDLBuilder { private buildTableReference(model: Model): TableReference | null { const modelProps = model.properties && typeof model.properties === 'string' - ? JSON.parse(model.properties) + ? this.parseProperties(model.properties) : {}; if (!modelProps.table) { return null; @@ -489,6 +485,18 @@ export class MDLBuilder implements IMDLBuilder { return []; } } + private parseProperties(properties?: string | null): Record { + if (!properties) { + return {}; + } + try { + const parsed = JSON.parse(properties); + return parsed && typeof parsed === 'object' ? parsed : {}; + } catch (error) { + logger.debug(`Can not parse properties "${properties}"`); + return {}; + } + } private postProcessManifest() { if (this.useRustWrenEngine()) { // 1. remove all the key that the value is null From ae32946df7bc635c4cbfa090c62a2d9403028f0e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 10 Jun 2026 19:04:34 +0530 Subject: [PATCH 0159/1087] Fail safe model sync status checks --- .../apollo/server/resolvers/modelResolver.ts | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 59ad3cd7bb..4afff6351e 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -200,19 +200,24 @@ export class ModelResolver { } public async checkModelSync(_root: any, _args: any, ctx: IContext) { - const { id } = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest, id); - const lastDeploy = await ctx.deployService.getLastDeployment(id); - const lastDeployHash = lastDeploy?.hash; - const inProgressDeployment = - await ctx.deployService.getInProgressDeployment(id); - if (inProgressDeployment) { - return { status: SyncStatusEnum.IN_PROGRESS }; + try { + const { id } = await ctx.projectService.getCurrentProject(); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const currentHash = ctx.deployService.createMDLHash(manifest, id); + const lastDeploy = await ctx.deployService.getLastDeployment(id); + const lastDeployHash = lastDeploy?.hash; + const inProgressDeployment = + await ctx.deployService.getInProgressDeployment(id); + if (inProgressDeployment) { + return { status: SyncStatusEnum.IN_PROGRESS }; + } + return currentHash == lastDeployHash + ? { status: SyncStatusEnum.SYNCRONIZED } + : { status: SyncStatusEnum.UNSYNCRONIZED }; + } catch (err: any) { + logger.error(`checkModelSync failed: ${err.message}`, err); + return { status: SyncStatusEnum.UNSYNCRONIZED }; } - return currentHash == lastDeployHash - ? { status: SyncStatusEnum.SYNCRONIZED } - : { status: SyncStatusEnum.UNSYNCRONIZED }; } public async deploy( From 12fbc3ce86f1bf10130875a83c8183418707c809 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 10 Jun 2026 21:40:36 +0530 Subject: [PATCH 0160/1087] Guard diagram resolver against null metadata --- .../server/resolvers/diagramResolver.ts | 40 +++++++++++++++---- 1 file changed, 32 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index f57b6980cb..fcbc0e29d4 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -133,7 +133,7 @@ export class DiagramResolver { } private transformModel(model: Model): DiagramModel { - const properties = JSON.parse(model.properties); + const properties = this.parseProperties(model.properties); return { id: uuidv4(), modelId: model.id, @@ -155,7 +155,7 @@ export class DiagramResolver { column: ModelColumn, nestedColumns: ModelNestedColumn[], ): DiagramModelField { - const properties = JSON.parse(column.properties); + const properties = this.parseProperties(column.properties); return { id: uuidv4(), columnId: column.id, @@ -186,8 +186,8 @@ export class DiagramResolver { column: ModelColumn, columnsMDL: ColumnMDL[], ): DiagramModelField { - const properties = JSON.parse(column.properties); - const lineage = JSON.parse(column.lineage); + const properties = this.parseProperties(column.properties); + const lineage = this.parseLineage(column.lineage); const columnMDL = columnsMDL.find( ({ name }) => name === column.referenceName, ); @@ -222,9 +222,7 @@ export class DiagramResolver { const displayName = models.find( (model) => model.referenceName === referenceName, )?.displayName; - const properties = relation.properties - ? JSON.parse(relation.properties) - : null; + const properties = this.parseProperties(relation.properties); return { id: uuidv4(), relationId: relation.id, @@ -249,7 +247,7 @@ export class DiagramResolver { } private transformView(view: View): DiagramView { - const properties = JSON.parse(view.properties); + const properties = this.parseProperties(view.properties); const fields = (properties?.columns || []).map((column: any) => ({ id: uuidv4(), nodeType: NodeType.FIELD, @@ -270,4 +268,30 @@ export class DiagramResolver { description: properties?.description, }; } + + private parseProperties(properties?: string | null): Record { + if (!properties) { + return {}; + } + try { + const parsed = JSON.parse(properties); + return parsed && typeof parsed === 'object' ? parsed : {}; + } catch (error) { + logger.debug(`Can not parse properties "${properties}"`); + return {}; + } + } + + private parseLineage(lineage?: string | null): number[] { + if (!lineage) { + return []; + } + try { + const parsed = JSON.parse(lineage); + return Array.isArray(parsed) ? parsed : []; + } catch (error) { + logger.debug(`Can not parse lineage "${lineage}"`); + return []; + } + } } From 7c6e578087f9ebc806ddf2b55b5dcf8e5a307efc Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 11 Jun 2026 19:04:36 +0530 Subject: [PATCH 0161/1087] Skip broken calculated fields during MDL build --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 134 +++++++++++--------- 1 file changed, 74 insertions(+), 60 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index bebef38b10..9ba7ce8f5f 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -235,41 +235,48 @@ export class MDLBuilder implements IMDLBuilder { this.columns .filter(({ isCalculated }) => isCalculated) .forEach((column: ModelColumn) => { - // validate manifest.model exist - const relatedModel = this.relatedModels.find( - (model: any) => model.id === column.modelId, - ); - if (!relatedModel) { - logger.debug( - `Build MDL Column Error: can not find related model, modelId "${column.modelId}", columnId: "${column.id}"`, + try { + // validate manifest.model exist + const relatedModel = this.relatedModels.find( + (model: any) => model.id === column.modelId, ); - return; - } - const model = this.manifest.models.find( - (model: any) => model.name === relatedModel.referenceName, - ); - if (!model) { - logger.debug( - `Build MDL Column Error: can not find model, modelId "${column.modelId}", columnId: "${column.id}"`, + if (!relatedModel) { + logger.debug( + `Build MDL Column Error: can not find related model, modelId "${column.modelId}", columnId: "${column.id}"`, + ); + return; + } + const model = this.manifest.models.find( + (model: any) => model.name === relatedModel.referenceName, ); - return; - } - const expression = this.getColumnExpression(column, model); - if (expression === null) { - logger.debug( - `Build MDL Column Error: invalid calculated field metadata, modelId "${column.modelId}", columnId: "${column.id}"`, + if (!model) { + logger.debug( + `Build MDL Column Error: can not find model, modelId "${column.modelId}", columnId: "${column.id}"`, + ); + return; + } + const expression = this.getColumnExpression(column, model); + if (expression === null) { + logger.debug( + `Build MDL Column Error: invalid calculated field metadata, modelId "${column.modelId}", columnId: "${column.id}"`, + ); + return; + } + const columnValue = { + name: column.referenceName, + type: column.type, + isCalculated: true, + expression, + notNull: column.notNull ? true : false, + properties: this.parseProperties(column.properties), + }; + model.columns.push(columnValue); + } catch (error: any) { + logger.error( + `Build MDL Column Error: failed to add calculated field, modelId "${column.modelId}", columnId: "${column.id}", message: ${error.message}`, + error, ); - return; } - const columnValue = { - name: column.referenceName, - type: column.type, - isCalculated: true, - expression, - notNull: column.notNull ? true : false, - properties: this.parseProperties(column.properties), - }; - model.columns.push(columnValue); }); } @@ -277,37 +284,44 @@ export class MDLBuilder implements IMDLBuilder { modelName: string, calculatedField: ModelColumn, ) { - const model = this.manifest.models.find( - (model: any) => model.name === modelName, - ); - if (!model) { - logger.debug(`Can not find model "${modelName}" to add calculated field`); - return; - } - // if calculated field is already in the model, skip - if ( - model.columns.find( - (column: any) => column.name === calculatedField.referenceName, - ) - ) { - return; - } - const expression = this.getColumnExpression(calculatedField, model); - if (expression === null) { - logger.debug( - `Can not add calculated field "${calculatedField.referenceName}" because its metadata is invalid`, + try { + const model = this.manifest.models.find( + (model: any) => model.name === modelName, + ); + if (!model) { + logger.debug(`Can not find model "${modelName}" to add calculated field`); + return; + } + // if calculated field is already in the model, skip + if ( + model.columns.find( + (column: any) => column.name === calculatedField.referenceName, + ) + ) { + return; + } + const expression = this.getColumnExpression(calculatedField, model); + if (expression === null) { + logger.debug( + `Can not add calculated field "${calculatedField.referenceName}" because its metadata is invalid`, + ); + return; + } + const columnValue = { + name: calculatedField.referenceName, + type: calculatedField.type, + isCalculated: true, + expression, + notNull: calculatedField.notNull ? true : false, + properties: this.parseProperties(calculatedField.properties), + }; + model.columns.push(columnValue); + } catch (error: any) { + logger.error( + `Can not add calculated field "${calculatedField.referenceName}" because building it failed: ${error.message}`, + error, ); - return; } - const columnValue = { - name: calculatedField.referenceName, - type: calculatedField.type, - isCalculated: true, - expression, - notNull: calculatedField.notNull ? true : false, - properties: this.parseProperties(calculatedField.properties), - }; - model.columns.push(columnValue); } public addRelation(): void { From 7f418180cb0789bb8c3d5915fd42c2b2bac5aa4b Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 11 Jun 2026 19:19:15 +0530 Subject: [PATCH 0162/1087] Skip invalid calculated fields in diagram resolver --- .../apollo/server/resolvers/diagramResolver.ts | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index fcbc0e29d4..25fb2e9c2d 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -113,9 +113,13 @@ export class DiagramResolver { } if (column.isCalculated) { - transformedModel.calculatedFields.push( - this.transformCalculatedField(column, modelMDL.columns), + const transformedCalculatedField = this.transformCalculatedField( + column, + modelMDL?.columns || [], ); + if (transformedCalculatedField) { + transformedModel.calculatedFields.push(transformedCalculatedField); + } } else { const nestedColumns = modelNestedColumns.filter( (nestedColumn) => nestedColumn.columnId === column.id, @@ -185,12 +189,18 @@ export class DiagramResolver { private transformCalculatedField( column: ModelColumn, columnsMDL: ColumnMDL[], - ): DiagramModelField { + ): DiagramModelField | null { const properties = this.parseProperties(column.properties); const lineage = this.parseLineage(column.lineage); const columnMDL = columnsMDL.find( ({ name }) => name === column.referenceName, ); + if (!columnMDL) { + logger.debug( + `Skip diagram calculated field "${column.referenceName}" because it is missing from built MDL`, + ); + return null; + } return { id: uuidv4(), columnId: column.id, From 1c59cccafe86110b7589ece9c62842dba343f7c7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 12 Jun 2026 15:36:21 +0530 Subject: [PATCH 0163/1087] Validate ask pipeline SQL responses --- wren-ai-service/src/web/v1/services/ask.py | 276 ++++++++++++------ .../test_ask_heuristic_text_to_sql.py | 30 ++ 2 files changed, 218 insertions(+), 88 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3d65b33809..f8815b61c1 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -810,6 +810,84 @@ def _extract_pipeline_reply(self, result: dict, key: str) -> str: return "" + def _extract_retrieval_documents(self, retrieval_result: dict) -> list[dict]: + construct_result = retrieval_result.get("construct_retrieval_results", {}) + documents = construct_result.get("retrieval_results", []) + if not isinstance(documents, list): + logger.warning("Schema retrieval returned invalid document payload") + return [] + + valid_documents = [] + for document in documents: + if not isinstance(document, dict): + logger.warning("Ignoring malformed retrieval document: %s", document) + continue + if not document.get("table_name") and not document.get("table_ddl"): + logger.warning("Ignoring retrieval document without table metadata") + continue + valid_documents.append(document) + + return valid_documents + + def _extract_retrieval_metadata( + self, retrieval_result: dict + ) -> tuple[list[dict], list[str], list[str]]: + documents = self._extract_retrieval_documents(retrieval_result) + table_names = [ + table_name + for document in documents + if isinstance(table_name := document.get("table_name"), str) + and table_name.strip() + ] + table_ddls = [ + table_ddl + for document in documents + if isinstance(table_ddl := document.get("table_ddl"), str) + and table_ddl.strip() + ] + return documents, table_names, table_ddls + + def _is_valid_select_sql(self, sql: Optional[str]) -> bool: + if not isinstance(sql, str): + return False + + normalized = re.sub(r"\s+", " ", sql.strip()) + if not normalized: + return False + + return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) + + def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: + if not self._is_valid_select_sql(sql): + return None + return AskResult(sql=sql.strip(), type="llm") + + def _build_failed_text_to_sql_response( + self, + trace_id: Optional[str], + message: str, + *, + rephrased_question: Optional[str] = None, + intent_reasoning: Optional[str] = None, + retrieved_tables: Optional[list[str]] = None, + sql_generation_reasoning: Optional[str] = None, + invalid_sql: Optional[str] = None, + is_followup: bool = False, + code: Literal["NO_RELEVANT_DATA", "NO_RELEVANT_SQL", "OTHERS"] = "NO_RELEVANT_SQL", + ) -> AskResultResponse: + return AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError(code=code, message=message), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=retrieved_tables, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + trace_id=trace_id, + is_followup=is_followup, + ) + @observe(name="Ask Question") @trace_metadata async def ask( @@ -832,6 +910,18 @@ async def ask( if not query_id: raise ValueError("query_id is required for ask service execution") + user_query = (ask_request.query or "").strip() + if not user_query: + self._ask_results[query_id] = self._build_failed_text_to_sql_response( + trace_id, + "Question is required", + code="OTHERS", + ) + results["metadata"]["error_type"] = "OTHERS" + results["metadata"]["error_message"] = "Question is required" + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + logger.info(f"Ask pipeline started for query_id: {query_id}") histories = ask_request.histories[: self._max_histories][ ::-1 @@ -862,7 +952,6 @@ async def ask( sql_knowledge = None try: - user_query = ask_request.query sql_user_query = user_query # ask status can be understanding, searching, generating, finished, failed, stopped @@ -905,16 +994,9 @@ async def ask( enable_column_pruning=False, ), ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) ) - documents = _retrieval_result.get("retrieval_results", []) - table_names = [ - document.get("table_name") for document in documents - ] - table_ddls = [ - document.get("table_ddl", "") or "" for document in documents - ] logger.info( "Retrieved tables for direct heuristic query_id %s: %s", query_id, @@ -929,27 +1011,25 @@ async def ask( query_id, user_query, ) - api_results = [ - AskResult( - **{ - "sql": heuristic_sql, - "type": "llm", - } - ) - ] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results + if ask_result := self._build_ask_result_from_sql( + heuristic_sql + ): + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback did not produce a valid SELECT statement." historical_question = await self._run_with_timeout( "Historical question retrieval", @@ -964,17 +1044,27 @@ async def ask( "formatted_output", {} ).get("documents", [])[:1] - if historical_question_result: - api_results = [ + valid_historical_results = [] + for result in historical_question_result: + sql_statement = result.get("statement") + if not self._is_valid_select_sql(sql_statement): + logger.warning( + "Ignoring historical question without valid SQL for query_id %s", + query_id, + ) + continue + valid_historical_results.append( AskResult( **{ - "sql": result.get("statement"), + "sql": sql_statement.strip(), "type": "view" if result.get("viewId") else "llm", "viewId": result.get("viewId"), } ) - for result in historical_question_result - ] + ) + + if valid_historical_results: + api_results = valid_historical_results sql_generation_reasoning = "" else: original_user_query = user_query @@ -1183,11 +1273,9 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - documents = _retrieval_result.get("retrieval_results", []) - table_names = [document.get("table_name") for document in documents] - table_ddls = [ - document.get("table_ddl", "") or "" for document in documents - ] + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) @@ -1227,14 +1315,27 @@ async def ask( query_id, user_query, ) - api_results = [ - AskResult( - **{ - "sql": heuristic_sql, - "type": "llm", - } - ) - ] + ask_result = self._build_ask_result_from_sql(heuristic_sql) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback did not produce a valid SELECT statement." + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_failed_text_to_sql_response( + trace_id, + error_message, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + invalid_sql=invalid_sql, + is_followup=True if histories else False, + ) + ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = error_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + api_results = [ask_result] if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="finished", @@ -1423,14 +1524,15 @@ async def ask( if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - api_results = [ - AskResult( - **{ - "sql": sql_valid_result.get("sql"), - "type": "llm", - } + if ask_result := self._build_ask_result_from_sql( + sql_valid_result.get("sql") + ): + api_results = [ask_result] + else: + invalid_sql = sql_valid_result.get("sql") + error_message = ( + "SQL generation did not produce a valid SELECT statement." ) - ] elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: @@ -1505,15 +1607,15 @@ async def ask( if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - api_results = [ - AskResult( - **{ - "sql": valid_generation_result.get("sql"), - "type": "llm", - } - ) - ] - break + if ask_result := self._build_ask_result_from_sql( + valid_generation_result.get("sql") + ): + api_results = [ask_result] + break + invalid_sql = valid_generation_result.get("sql") + error_message = ( + "SQL correction did not produce a valid SELECT statement." + ) failed_dry_run_result = sql_correction_results["post_process"][ "invalid_generation_result" @@ -1547,29 +1649,27 @@ async def ask( query_id, user_query, ) - api_results = [ - AskResult( - **{ - "sql": heuristic_sql, - "type": "llm", - } - ) - ] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results + ask_result = self._build_ask_result_from_sql(heuristic_sql) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback did not produce a valid SELECT statement." + else: + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 87e2fa2f09..d3bf2f834c 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -153,3 +153,33 @@ def test_repair_failure_count_requires_schema_backed_failure_dimension(): ) is None ) + + +def test_ask_result_validation_requires_select_sql(): + service = AskService(pipelines={}) + + assert service._build_ask_result_from_sql("SELECT 1") + assert service._build_ask_result_from_sql("WITH rows AS (SELECT 1) SELECT * FROM rows") + assert service._build_ask_result_from_sql("") is None + assert service._build_ask_result_from_sql("DELETE FROM dbo_repair_logs") is None + assert service._build_ask_result_from_sql(None) is None + + +def test_retrieval_metadata_ignores_malformed_documents(): + service = AskService(pipelines={}) + + documents, table_names, table_ddls = service._extract_retrieval_metadata( + { + "construct_retrieval_results": { + "retrieval_results": [ + {"table_name": "dbo_repair_logs", "table_ddl": "CREATE TABLE dbo_repair_logs (id varchar)"}, + {}, + "bad-document", + ] + } + } + ) + + assert len(documents) == 1 + assert table_names == ["dbo_repair_logs"] + assert table_ddls == ["CREATE TABLE dbo_repair_logs (id varchar)"] From d88641bc3b8867c7df22cac4b25afe1bee2808ec Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 12 Jun 2026 16:07:08 +0530 Subject: [PATCH 0164/1087] Summarize invalid calculated field logs in MDL builder --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 67 ++++++++++++++++----- 1 file changed, 53 insertions(+), 14 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 9ba7ce8f5f..05cd3bf220 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -41,6 +41,11 @@ export interface IMDLBuilder { // responsible to generate a valid manifest json export class MDLBuilder implements IMDLBuilder { private manifest: Manifest; + private invalidCalculatedFields: Array<{ + modelId: number; + columnId: number; + reason: string; + }> = []; private project: Project; private readonly models: Model[]; @@ -84,6 +89,7 @@ export class MDLBuilder implements IMDLBuilder { } public build(): Manifest { + this.invalidCalculatedFields = []; this.addProject(); this.addModel(); this.addNormalField(); @@ -91,6 +97,7 @@ export class MDLBuilder implements IMDLBuilder { this.addCalculatedField(); this.addView(); this.postProcessManifest(); + this.logInvalidCalculatedFieldSummary(); return this.getManifest(); } @@ -241,8 +248,10 @@ export class MDLBuilder implements IMDLBuilder { (model: any) => model.id === column.modelId, ); if (!relatedModel) { - logger.debug( - `Build MDL Column Error: can not find related model, modelId "${column.modelId}", columnId: "${column.id}"`, + this.recordInvalidCalculatedField( + column.modelId, + column.id, + 'can not find related model', ); return; } @@ -250,15 +259,19 @@ export class MDLBuilder implements IMDLBuilder { (model: any) => model.name === relatedModel.referenceName, ); if (!model) { - logger.debug( - `Build MDL Column Error: can not find model, modelId "${column.modelId}", columnId: "${column.id}"`, + this.recordInvalidCalculatedField( + column.modelId, + column.id, + 'can not find model', ); return; } const expression = this.getColumnExpression(column, model); if (expression === null) { - logger.debug( - `Build MDL Column Error: invalid calculated field metadata, modelId "${column.modelId}", columnId: "${column.id}"`, + this.recordInvalidCalculatedField( + column.modelId, + column.id, + 'invalid calculated field metadata', ); return; } @@ -272,9 +285,10 @@ export class MDLBuilder implements IMDLBuilder { }; model.columns.push(columnValue); } catch (error: any) { - logger.error( - `Build MDL Column Error: failed to add calculated field, modelId "${column.modelId}", columnId: "${column.id}", message: ${error.message}`, - error, + this.recordInvalidCalculatedField( + column.modelId, + column.id, + `failed to add calculated field: ${error.message}`, ); } }); @@ -302,8 +316,10 @@ export class MDLBuilder implements IMDLBuilder { } const expression = this.getColumnExpression(calculatedField, model); if (expression === null) { - logger.debug( - `Can not add calculated field "${calculatedField.referenceName}" because its metadata is invalid`, + this.recordInvalidCalculatedField( + calculatedField.modelId, + calculatedField.id, + `insert skipped because metadata is invalid for "${calculatedField.referenceName}"`, ); return; } @@ -317,9 +333,10 @@ export class MDLBuilder implements IMDLBuilder { }; model.columns.push(columnValue); } catch (error: any) { - logger.error( - `Can not add calculated field "${calculatedField.referenceName}" because building it failed: ${error.message}`, - error, + this.recordInvalidCalculatedField( + calculatedField.modelId, + calculatedField.id, + `insert failed for "${calculatedField.referenceName}": ${error.message}`, ); } } @@ -511,6 +528,28 @@ export class MDLBuilder implements IMDLBuilder { return {}; } } + private recordInvalidCalculatedField( + modelId: number, + columnId: number, + reason: string, + ) { + this.invalidCalculatedFields.push({ modelId, columnId, reason }); + } + private logInvalidCalculatedFieldSummary() { + if (this.invalidCalculatedFields.length === 0) { + return; + } + const preview = this.invalidCalculatedFields + .slice(0, 10) + .map( + ({ modelId, columnId, reason }) => + `modelId="${modelId}", columnId="${columnId}", reason="${reason}"`, + ) + .join('; '); + logger.warn( + `Skipped ${this.invalidCalculatedFields.length} invalid calculated field(s) while building MDL. ${preview}${this.invalidCalculatedFields.length > 10 ? '; ...' : ''}`, + ); + } private postProcessManifest() { if (this.useRustWrenEngine()) { // 1. remove all the key that the value is null From 899afa3d9c316da3607f821452151843b23ee33f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 12 Jun 2026 16:48:15 +0530 Subject: [PATCH 0165/1087] Ground repair failure SQL generation --- .../src/pipelines/generation/utils/sql.py | 14 +-- wren-ai-service/src/web/v1/services/ask.py | 94 ++++++++++++++++--- .../pipelines/generation/test_sql_utils.py | 5 + .../test_ask_heuristic_text_to_sql.py | 47 ++++++++++ 4 files changed, 141 insertions(+), 19 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 89efd1c885..e65128234d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1294,8 +1294,9 @@ async def _classify_generation_result( - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. - Never invent foreign key columns or relationship fields such as "FailurePatternID", "FailurePatternId", "TicketID", or "
ID" unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. - Never invent time bucket columns such as "MONTH", "YEAR", "DAY", "month", "year", or "date" unless that exact column appears in the DATABASE SCHEMA. For monthly, yearly, or daily trends, apply a supported date/time bucket function from SQL FUNCTIONS to a real timestamp column from the selected table. +- Every generated SQL query must be grounded only in the connected datasource metadata, deployed semantic model definitions, relationships, and DATABASE SCHEMA shown in the prompt. Do not use table names, column names, join paths, JSON keys, or business dimensions that are not explicitly present in that context. - For synced repair-log schemas, if "dbo_repair_logs" contains "created_at" and the user asks for monthly repair volume or repair trends, count repair rows and bucket "dbo_repair_logs"."created_at". Do not select, group by, or order by "dbo_repair_logs"."MONTH" or bare "MONTH" unless the schema explicitly contains that column. -- For repair counts grouped by failure category in synced repair-log schemas, use "dbo_repair_logs"."failure_code" when that column appears in the schema. Do not invent "failure_category" unless it appears in the DATABASE SCHEMA. +- For repair counts grouped by failure category, prefer the richest explicit category field exposed by the connected datasource. If the schema includes "dbo_DebugEntries", "dbo_DebugFixLogs", and "dbo_DebugFixes", group by "dbo_DebugFixes"."Description" after joining "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" and "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id". Otherwise use "dbo_repair_logs"."failure_code" only when that column appears in the schema. Do not invent "failure_category" unless it appears in the DATABASE SCHEMA. - For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. - For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. - Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. @@ -1312,20 +1313,21 @@ async def _classify_generation_result( - If a requested metric such as debug hours, risk score, repair cost, or turnaround time is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. - Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. - For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. -- For repair counts grouped by failure category, use explicit exposed fields such as "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. +- For repair counts grouped by failure category, use explicit exposed fields and schema-backed joins only. Prefer "dbo_DebugFixes"."Description" joined through "dbo_DebugFixLogs" when "dbo_DebugEntries"."DebugEntryId", "dbo_DebugFixLogs"."DebugEntryId", "dbo_DebugFixLogs"."FixId", and "dbo_DebugFixes"."Id" are present. Otherwise use "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. - For PCB/debug-entry failure charts, do not join "dbo_DebugEntries"."DebugEntryId" to "dbo_failure_patterns"."id"; those fields have incompatible types. If both "dbo_DebugEntries"."FailureSys" and "dbo_failure_patterns"."id" exist, join "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". - For PCB synced database questions: - Use "dbo_DebugEntries" for debug/PCB event records when the schema contains it. - Use columns such as "Material", "WorkOrder", "SerialNumber", "FailedAt", "DateIn", "DateOut", "Hours", "Priority", "Actions", "Notes", and "FailureSys" only when they appear in the schema. - Use "dbo_failure_patterns" for failure names, categories, severity, trend, occurrence counts, daily pattern summaries, and cost impact when those columns appear in the schema. - For throughput trends across manufacturing/business units, use "dbo_DebugEntries"."BusinessUnit" as the unit dimension and a real debug-entry timestamp such as "dbo_DebugEntries"."DateIn" or "dbo_DebugEntries"."FailedAt" for the trend bucket. Do not use "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH", or invented manufacturing/date fields. - - For top/common PCB failure questions, prefer grouping by "dbo_failure_patterns"."name" or "dbo_failure_patterns"."category" and counting "dbo_DebugEntries"."DebugEntryId" after joining "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". + - For top/common PCB failure questions, first prefer grouping by "dbo_DebugFixes"."Description" and counting rows through the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join when those tables and join columns are in the schema. Otherwise prefer grouping by "dbo_failure_patterns"."name" or "dbo_failure_patterns"."category" and counting "dbo_DebugEntries"."DebugEntryId" after joining "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". - If a useful aggregate already exists in "dbo_failure_patterns" such as "occurrences", it can be used directly for top failure pattern questions without joining event rows. - For requests such as "show top 10 most common PCB failures", "bar chart of failures by category", or "count of repairs grouped by failure category", generate SQL first. Do not answer with general charting guidance. Return the categorical failure field plus a count metric that can drive a bar chart. - For failure-category charts, prefer one of these patterns depending on schema availability: - 1. `GROUP BY "dbo_failure_patterns"."category"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` - 2. `GROUP BY "dbo_failure_patterns"."name"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` - 3. `GROUP BY "dbo_repair_logs"."failure_code"` and `COUNT(*)` + 1. `GROUP BY "dbo_DebugFixes"."Description"` and `COUNT(*)` using the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join + 2. `GROUP BY "dbo_failure_patterns"."category"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` + 3. `GROUP BY "dbo_failure_patterns"."name"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` + 4. `GROUP BY "dbo_repair_logs"."failure_code"` and `COUNT(*)` - For chart-oriented questions, ensure the final SELECT contains only the chart-ready dimension and metric columns. Avoid prose-like outputs or helper columns. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f8815b61c1..c71a191835 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -240,6 +240,24 @@ def _schema_contains( schema_text += "\n" + "\n".join(table_names) return bool(re.search(pattern, schema_text, flags=re.IGNORECASE)) + def _schema_has_table_column( + self, + table_ddls: list[str], + table_name: str, + column_name: str, + table_names: Optional[list[str]] = None, + ) -> bool: + table_pattern = rf"\b{re.escape(table_name)}\b" + column_pattern = rf"\b{re.escape(column_name)}\b" + + for ddl in table_ddls or []: + if re.search(table_pattern, ddl, flags=re.IGNORECASE) and re.search( + column_pattern, ddl, flags=re.IGNORECASE + ): + return True + + return False + def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: column_names: list[str] = [] non_column_prefixes = ( @@ -385,21 +403,54 @@ def _build_repair_failure_count_sql( return None top_n = self._extract_requested_top_n(query) - has_repair_logs = self._schema_contains( - table_ddls, r"\bdbo_repair_logs\b", table_names=table_names - ) - has_failure_code = self._schema_contains( - table_ddls, r"\bfailure_code\b", table_names=table_names + + has_debug_fix_route = all( + ( + self._schema_has_table_column( + table_ddls, + "dbo_DebugEntries", + "DebugEntryId", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixLogs", + "DebugEntryId", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixLogs", + "FixId", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixes", + "Id", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixes", + "Description", + table_names=table_names, + ), + ) ) - if has_repair_logs and has_failure_code: + if has_debug_fix_route: return ( - f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_repair_logs" ' - f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' - f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' + 'SELECT "dbo_DebugFixes"."Description" AS "failure_category", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_DebugEntries" ' + 'JOIN "dbo_DebugFixLogs" ' + 'ON "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" ' + 'JOIN "dbo_DebugFixes" ' + 'ON "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id" ' + 'WHERE "dbo_DebugFixes"."Description" IS NOT NULL ' + 'GROUP BY "dbo_DebugFixes"."Description" ' + 'ORDER BY "repair_count" DESC ' + f"LIMIT {top_n}" ) has_debug_entries = self._schema_contains( @@ -449,6 +500,23 @@ def _build_repair_failure_count_sql( f'LIMIT {top_n}' ) + has_repair_logs = self._schema_has_table_column( + table_ddls, + "dbo_repair_logs", + "failure_code", + table_names=table_names, + ) + if has_repair_logs: + return ( + f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_repair_logs" ' + f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' + f'GROUP BY "dbo_repair_logs"."failure_code" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + return None def _is_direct_heuristic_sql_query(self, query: str) -> bool: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index e1801778af..7bcee2f7f1 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -77,6 +77,9 @@ def test_get_text_to_sql_rules_adds_mssql_specific_constraints(): assert "Resolve relative time phrases" in rules assert "Do not include helper ranking columns" in rules assert "prefer SELECT TOP (N)" in rules + assert "connected datasource metadata" in rules + assert '"dbo_DebugFixes"."Description"' in rules + assert '"dbo_DebugFixLogs"."FixId"' in rules assert "FailurePatternID" in rules assert "failure_code" in rules assert "CURRENT_DATE - INTERVAL '1 month'" not in rules @@ -117,6 +120,8 @@ def test_get_sql_generation_system_prompt_uses_data_source_specific_rules(): assert "The target database is MSSQL." in prompt assert "DATEPART(YEAR, )" in prompt + assert "deployed semantic model definitions" in prompt + assert '"dbo_DebugFixes"."Description"' in prompt def test_normalize_generation_result_sql_rewrites_common_mssql_time_patterns(): diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index d3bf2f834c..8a4612c556 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -89,6 +89,53 @@ def test_repair_failure_count_uses_repair_log_failure_code(): assert "dbo_reports" not in sql +def test_repair_failure_count_prefers_debug_fix_description_when_available(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_DebugEntries ( + DebugEntryId VARCHAR, + FailureSys VARCHAR + ); + """, + """ + CREATE TABLE dbo_DebugFixLogs ( + DebugEntryId VARCHAR, + FixId VARCHAR + ); + """, + """ + CREATE TABLE dbo_DebugFixes ( + Id VARCHAR, + Description VARCHAR + ); + """, + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + failure_code VARCHAR + ); + """, + ] + + sql = service._build_repair_failure_count_sql( + "Create a bar chart of repair counts grouped by failure category.", + table_ddls, + table_names=[ + "dbo_DebugEntries", + "dbo_DebugFixLogs", + "dbo_DebugFixes", + "dbo_repair_logs", + ], + ) + + assert sql + assert '"dbo_DebugFixes"."Description" AS "failure_category"' in sql + assert '"dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id"' in sql + assert '"dbo_repair_logs"."failure_code"' not in sql + assert "FailurePatternID" not in sql + + def test_common_pcb_failures_uses_repair_log_failure_code(): service = AskService(pipelines={}) table_ddls = [ From b13cf6426595c934a2622f11b45bd41eb121f8db Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 12 Jun 2026 16:59:46 +0530 Subject: [PATCH 0166/1087] Guard diagram resolver against missing expressions --- wren-ui/src/apollo/server/resolvers/diagramResolver.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index 25fb2e9c2d..dab3331b83 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -201,6 +201,12 @@ export class DiagramResolver { ); return null; } + if (columnMDL.expression == null) { + logger.debug( + `Skip diagram calculated field "${column.referenceName}" because its MDL expression is missing`, + ); + return null; + } return { id: uuidv4(), columnId: column.id, From edda57f82b367b8c371d3dd94567b9a14b028815 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 12 Jun 2026 19:23:02 +0530 Subject: [PATCH 0167/1087] Ground repair SLA chart SQL --- .../src/pipelines/generation/utils/sql.py | 2 + wren-ai-service/src/web/v1/services/ask.py | 50 ++++++++++++++++++- .../pipelines/generation/test_sql_utils.py | 3 ++ .../test_ask_heuristic_text_to_sql.py | 40 +++++++++++++++ 4 files changed, 94 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index e65128234d..5cf077d6e2 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1297,6 +1297,7 @@ async def _classify_generation_result( - Every generated SQL query must be grounded only in the connected datasource metadata, deployed semantic model definitions, relationships, and DATABASE SCHEMA shown in the prompt. Do not use table names, column names, join paths, JSON keys, or business dimensions that are not explicitly present in that context. - For synced repair-log schemas, if "dbo_repair_logs" contains "created_at" and the user asks for monthly repair volume or repair trends, count repair rows and bucket "dbo_repair_logs"."created_at". Do not select, group by, or order by "dbo_repair_logs"."MONTH" or bare "MONTH" unless the schema explicitly contains that column. - For repair counts grouped by failure category, prefer the richest explicit category field exposed by the connected datasource. If the schema includes "dbo_DebugEntries", "dbo_DebugFixLogs", and "dbo_DebugFixes", group by "dbo_DebugFixes"."Description" after joining "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" and "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id". Otherwise use "dbo_repair_logs"."failure_code" only when that column appears in the schema. Do not invent "failure_category" unless it appears in the DATABASE SCHEMA. +- For repair SLA compliance dashboard/chart requests, do not invent "DAY", "MONTH", "turnaround_time", "sla_due_at", or due-date fields. If the schema only exposes "dbo_repair_logs"."status" and no explicit SLA/duration/deadline column, return a status distribution using "dbo_repair_logs"."status" and COUNT(*) so the UI can render a grounded chart. - For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. - For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. - Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. @@ -1313,6 +1314,7 @@ async def _classify_generation_result( - If a requested metric such as debug hours, risk score, repair cost, or turnaround time is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. - Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. - For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. +- For repair SLA compliance charts on "dbo_repair_logs", use "dbo_repair_logs"."status" as the compliance/status dimension when no explicit SLA, due-date, duration, or turnaround column appears in the DATABASE SCHEMA. Never use invented "DAY", "MONTH", or "turnaround_time" fields for SLA compliance. - For repair counts grouped by failure category, use explicit exposed fields and schema-backed joins only. Prefer "dbo_DebugFixes"."Description" joined through "dbo_DebugFixLogs" when "dbo_DebugEntries"."DebugEntryId", "dbo_DebugFixLogs"."DebugEntryId", "dbo_DebugFixLogs"."FixId", and "dbo_DebugFixes"."Id" are present. Otherwise use "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. - For PCB/debug-entry failure charts, do not join "dbo_DebugEntries"."DebugEntryId" to "dbo_failure_patterns"."id"; those fields have incompatible types. If both "dbo_DebugEntries"."FailureSys" and "dbo_failure_patterns"."id" exist, join "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". - For PCB synced database questions: diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c71a191835..927298858f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -519,6 +519,41 @@ def _build_repair_failure_count_sql( return None + def _build_repair_sla_compliance_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_sla = "sla" in normalized and any( + term in normalized + for term in ("compliance", "dashboard", "chart", "repair", "repairs") + ) + if not wants_sla: + return None + + has_repair_status = self._schema_has_table_column( + table_ddls, + "dbo_repair_logs", + "status", + table_names=table_names, + ) + if has_repair_status: + return ( + 'SELECT "dbo_repair_logs"."status" AS "sla_status", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."status" IS NOT NULL ' + 'GROUP BY "dbo_repair_logs"."status" ' + 'ORDER BY "repair_count" DESC' + ) + + return None + def _is_direct_heuristic_sql_query(self, query: str) -> bool: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -539,7 +574,15 @@ def _is_direct_heuristic_sql_query(self, query: str) -> bool: for term in ("pcb", "repair", "bar chart", "chart", "category") ) ) - return asks_manufacturing_throughput or asks_failure_counts + asks_sla_compliance = "sla" in normalized and any( + term in normalized + for term in ("compliance", "dashboard", "chart", "repair", "repairs") + ) + return ( + asks_manufacturing_throughput + or asks_failure_counts + or asks_sla_compliance + ) def _build_heuristic_text_to_sql_fallback( self, @@ -561,6 +604,11 @@ def _build_heuristic_text_to_sql_fallback( ): return repair_failure_count_sql + if repair_sla_sql := self._build_repair_sla_compliance_sql( + query, table_ddls, table_names=table_names + ): + return repair_sla_sql + wants_chart = any( term in normalized for term in ("chart", "bar chart", "line chart", "graph") ) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 7bcee2f7f1..eec70b3e4e 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -80,6 +80,8 @@ def test_get_text_to_sql_rules_adds_mssql_specific_constraints(): assert "connected datasource metadata" in rules assert '"dbo_DebugFixes"."Description"' in rules assert '"dbo_DebugFixLogs"."FixId"' in rules + assert "repair SLA compliance" in rules + assert '"dbo_repair_logs"."status"' in rules assert "FailurePatternID" in rules assert "failure_code" in rules assert "CURRENT_DATE - INTERVAL '1 month'" not in rules @@ -122,6 +124,7 @@ def test_get_sql_generation_system_prompt_uses_data_source_specific_rules(): assert "DATEPART(YEAR, )" in prompt assert "deployed semantic model definitions" in prompt assert '"dbo_DebugFixes"."Description"' in prompt + assert "repair SLA compliance" in prompt def test_normalize_generation_result_sql_rewrites_common_mssql_time_patterns(): diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 8a4612c556..b67ba370b8 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -179,6 +179,46 @@ def test_common_pcb_failures_uses_direct_heuristic_route(): ) +def test_repair_sla_compliance_uses_status_when_no_sla_duration_field(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP, + data VARCHAR + ); + """ + ] + + sql = service._build_repair_sla_compliance_sql( + "Generate a dashboard chart for repair SLA compliance.", + table_ddls, + table_names=["dbo_repair_logs"], + ) + + assert sql + assert '"dbo_repair_logs"."status" AS "sla_status"' in sql + assert 'COUNT(*) AS "repair_count"' in sql + assert '"dbo_repair_logs"."turnaround_time"' not in sql + assert '"DAY"' not in sql + assert '"MONTH"' not in sql + assert "DATEDIFF" not in sql.upper() + + +def test_repair_sla_compliance_uses_direct_heuristic_route(): + service = AskService(pipelines={}) + + assert service._is_direct_heuristic_sql_query( + "Generate a dashboard chart for repair SLA compliance." + ) + + def test_repair_failure_count_requires_schema_backed_failure_dimension(): service = AskService(pipelines={}) table_ddls = [ From 596724fd1f1507a878b787fe8e5abfd30ffc2f2f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 12 Jun 2026 19:28:31 +0530 Subject: [PATCH 0168/1087] Normalize current project booleans in API --- wren-ui/src/pages/api/v1/projects/current.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/pages/api/v1/projects/current.ts b/wren-ui/src/pages/api/v1/projects/current.ts index a88a7ea380..30e4e422c4 100644 --- a/wren-ui/src/pages/api/v1/projects/current.ts +++ b/wren-ui/src/pages/api/v1/projects/current.ts @@ -1,5 +1,6 @@ import { NextApiRequest, NextApiResponse } from 'next'; import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { coerceBoolean } from '@server/repositories/baseRepository'; import { ApiError, handleApiError, @@ -25,7 +26,7 @@ const serializeProject = (project) => id: project.id, displayName: project.displayName, projectType: project.projectType || 'CLASSIC', - isCurrent: Boolean(project.isCurrent), + isCurrent: coerceBoolean(project.isCurrent), hasDataSource: Boolean(project.type), type: project.type || null, createdAt: project.createdAt, @@ -46,7 +47,8 @@ export default async function handler( const projectService = getProjectService(); const projects = await projectService.listProjects(); - const currentProject = projects.find((project) => project.isCurrent) || null; + const currentProject = + projects.find((project) => coerceBoolean(project.isCurrent)) || null; await respondWithSimple({ res, From 698f7b2e0f23a52f984f9968704801c8eaf7a8b8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 12 Jun 2026 19:56:34 +0530 Subject: [PATCH 0169/1087] Ground monthly repair volume SQL --- wren-ai-service/src/web/v1/services/ask.py | 84 ++++++++++++------- .../test_ask_heuristic_text_to_sql.py | 41 +++++++++ 2 files changed, 95 insertions(+), 30 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 927298858f..313775c377 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -554,6 +554,47 @@ def _build_repair_sla_compliance_sql( return None + def _build_monthly_repair_volume_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_monthly_repairs = ( + "repair" in normalized + and any( + term in normalized + for term in ("monthly", "last 12 months", "trend", "volume") + ) + ) + if not wants_monthly_repairs: + return None + + has_repair_created_at = self._schema_has_table_column( + table_ddls, + "dbo_repair_logs", + "created_at", + table_names=table_names, + ) + if has_repair_created_at: + return ( + 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + ) + + return None + def _is_direct_heuristic_sql_query(self, query: str) -> bool: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -578,10 +619,18 @@ def _is_direct_heuristic_sql_query(self, query: str) -> bool: term in normalized for term in ("compliance", "dashboard", "chart", "repair", "repairs") ) + asks_monthly_repairs = ( + "repair" in normalized + and any( + term in normalized + for term in ("monthly", "last 12 months", "trend", "volume") + ) + ) return ( asks_manufacturing_throughput or asks_failure_counts or asks_sla_compliance + or asks_monthly_repairs ) def _build_heuristic_text_to_sql_fallback( @@ -609,6 +658,11 @@ def _build_heuristic_text_to_sql_fallback( ): return repair_sla_sql + if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( + query, table_ddls, table_names=table_names + ): + return monthly_repair_volume_sql + wants_chart = any( term in normalized for term in ("chart", "bar chart", "line chart", "graph") ) @@ -729,23 +783,6 @@ def _build_heuristic_text_to_sql_fallback( f'LIMIT {top_n}' ) - if wants_monthly_repairs and self._schema_contains( - table_ddls, r"\bdbo_repair_logs\b", table_names=table_names - ) and self._schema_contains( - table_ddls, r"\bcreated_at\b", table_names=table_names - ): - return ( - 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."created_at" >= DATEADD(month, -12, GETDATE()) ' - 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' - ) - if wants_failure_counts and wants_chart: top_n = self._extract_requested_top_n(query) return ( @@ -758,19 +795,6 @@ def _build_heuristic_text_to_sql_fallback( f'LIMIT {top_n}' ) - if wants_monthly_repairs: - return ( - 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."created_at" >= DATEADD(month, -12, GETDATE()) ' - 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' - ) - return None def _is_schema_grounded_query( diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index b67ba370b8..a11aa59c8e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -219,6 +219,47 @@ def test_repair_sla_compliance_uses_direct_heuristic_route(): ) +def test_monthly_repair_volume_uses_created_at_bucket(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP, + data VARCHAR + ); + """ + ] + + sql = service._build_monthly_repair_volume_sql( + "Generate a line chart showing monthly repair volume for the last 12 months.", + table_ddls, + table_names=["dbo_repair_logs"], + ) + + assert sql + assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year"' in sql + assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' in sql + assert 'COUNT(*) AS "repair_count"' in sql + assert '"dbo_repair_logs"."MONTH"' not in sql + assert '"MONTH"' not in sql + assert "DATEADD" not in sql.upper() + assert "GETDATE" not in sql.upper() + + +def test_monthly_repair_volume_uses_direct_heuristic_route(): + service = AskService(pipelines={}) + + assert service._is_direct_heuristic_sql_query( + "Generate a line chart showing monthly repair volume for the last 12 months." + ) + + def test_repair_failure_count_requires_schema_backed_failure_dimension(): service = AskService(pipelines={}) table_ddls = [ From 1e248c21e473692e7a49e876a31f54831afe42f4 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 12 Jun 2026 20:05:20 +0530 Subject: [PATCH 0170/1087] Retry project inserts with manual MSSQL ids --- .../server/repositories/projectRepository.ts | 72 ++++++++++++++++--- 1 file changed, 62 insertions(+), 10 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index ce83fb3496..85149c9d49 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -218,20 +218,44 @@ export class ProjectRepository data: Partial, queryOptions?: IQueryOptions, ): Promise { - return super.createOne( - await this.withMssqlId(this.withTimestamps(data), queryOptions), - queryOptions, - ); + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } } public override async createMany( data: Partial[], queryOptions?: IQueryOptions, ): Promise { - return super.createMany( - await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), - queryOptions, - ); + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } } public override async updateOne( @@ -376,11 +400,12 @@ export class ProjectRepository private withMssqlId = async ( data: Partial, queryOptions?: IQueryOptions, + forceManualId = false, ): Promise> => { if ( (data.id !== undefined && data.id !== null) || !this.isMssql() || - (await this.hasIdentityId()) + (!forceManualId && (await this.hasIdentityId())) ) { return data; } @@ -398,11 +423,12 @@ export class ProjectRepository private withMssqlIds = async ( data: Partial[], queryOptions?: IQueryOptions, + forceManualId = false, ): Promise[]> => { if ( data.every((item) => item.id !== undefined && item.id !== null) || !this.isMssql() || - (await this.hasIdentityId()) + (!forceManualId && (await this.hasIdentityId())) ) { return data; } @@ -422,4 +448,30 @@ export class ProjectRepository }; }); }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } From f19f382fbfa68d542e9f9eddb23f4f6f1fd8128d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 12 Jun 2026 20:30:37 +0530 Subject: [PATCH 0171/1087] Update project access API routes --- .../src/pages/api/v1/projects/access/[id].ts | 18 ++++++++++++++---- .../pages/api/v1/projects/access/current.ts | 18 ++++++++++++++---- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/pages/api/v1/projects/access/[id].ts b/wren-ui/src/pages/api/v1/projects/access/[id].ts index d55f477225..ce3244638e 100644 --- a/wren-ui/src/pages/api/v1/projects/access/[id].ts +++ b/wren-ui/src/pages/api/v1/projects/access/[id].ts @@ -15,10 +15,20 @@ const logger = getLogger('API_PROJECT_ACCESS_MEMBER'); logger.level = 'debug'; const getOrganizationMemberService = () => { - const { components } = require('@/common'); - const componentGraph = components ?? globalThis.__wrenComponents; - if (!componentGraph) { - throw new Error('Components are not initialized'); + const { components, initComponents } = require('@/common'); + let componentGraph = components ?? globalThis.__wrenComponents; + + if ( + !componentGraph?.organizationMemberService || + typeof componentGraph.organizationMemberService.listCurrentProjectAccess !== + 'function' + ) { + componentGraph = initComponents(); + globalThis.__wrenComponents = componentGraph; + } + + if (!componentGraph?.organizationMemberService) { + throw new Error('Organization member service is not initialized'); } return componentGraph.organizationMemberService; }; diff --git a/wren-ui/src/pages/api/v1/projects/access/current.ts b/wren-ui/src/pages/api/v1/projects/access/current.ts index 990e3900a0..ba5915b875 100644 --- a/wren-ui/src/pages/api/v1/projects/access/current.ts +++ b/wren-ui/src/pages/api/v1/projects/access/current.ts @@ -14,10 +14,20 @@ const logger = getLogger('API_PROJECT_ACCESS'); logger.level = 'debug'; const getOrganizationMemberService = () => { - const { components } = require('@/common'); - const componentGraph = components ?? globalThis.__wrenComponents; - if (!componentGraph) { - throw new Error('Components are not initialized'); + const { components, initComponents } = require('@/common'); + let componentGraph = components ?? globalThis.__wrenComponents; + + if ( + !componentGraph?.organizationMemberService || + typeof componentGraph.organizationMemberService.listCurrentProjectAccess !== + 'function' + ) { + componentGraph = initComponents(); + globalThis.__wrenComponents = componentGraph; + } + + if (!componentGraph?.organizationMemberService) { + throw new Error('Organization member service is not initialized'); } return componentGraph.organizationMemberService; }; From 5144f79a21640871bebfdc3df86bbaffcc4d745c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 12 Jun 2026 21:08:42 +0530 Subject: [PATCH 0172/1087] Normalize MSSQL generated SQL syntax --- .../src/pipelines/generation/utils/sql.py | 100 +++++++++++++----- .../pipelines/generation/test_sql_utils.py | 48 +++++++++ .../apollo/server/utils/mssqlSqlNormalizer.ts | 39 ++++++- 3 files changed, 162 insertions(+), 25 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5cf077d6e2..80c7e80daf 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -772,33 +772,36 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: rf'(?:"{bucket}"|\[{bucket}\]|{bucket})', re.IGNORECASE, ) - select_identifier_pattern = re.compile( - rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', - re.IGNORECASE, - ) - select_bare_identifier_pattern = re.compile( - rf"(?P\bSELECT\s+|,\s*){bucket}(?P\s*(?:,|\bFROM\b))", - re.IGNORECASE, - ) - select_qualified_identifier_pattern = re.compile( - rf'(?P\bSELECT\s+|,\s*){qualified_bucket_pattern.pattern}(?P\s*(?:,|\bFROM\b))', - re.IGNORECASE, + select_pattern = re.compile( + r"\bSELECT\b(?P.*?)(?=\bFROM\b)", + re.IGNORECASE | re.DOTALL, ) - def replace_select_identifier(match: re.Match[str]) -> str: - prefix = match.group("prefix") - suffix = match.group("suffix") - return f'{prefix}{expression} AS "{bucket}"{suffix}' + def replace_select(match: re.Match[str]) -> str: + body = match.group("body") + items = _split_top_level_select_items(body) + if not items: + return match.group(0) - rewritten = select_qualified_identifier_pattern.sub( - replace_select_identifier, rewritten - ) - rewritten = select_identifier_pattern.sub( - replace_select_identifier, rewritten - ) - rewritten = select_bare_identifier_pattern.sub( - replace_select_identifier, rewritten - ) + rebuilt: list[str] = [] + changed = False + select_identifier_pattern = re.compile( + rf"^(?:\"{bucket}\"|\[{bucket}\]|{bucket}|{qualified_bucket_pattern.pattern})$", + re.IGNORECASE, + ) + for item in items: + if select_identifier_pattern.fullmatch(item.strip()): + rebuilt.append(f'{expression} AS "{bucket}"') + changed = True + else: + rebuilt.append(item) + + if not changed: + return match.group(0) + + return "SELECT " + ", ".join(rebuilt) + " " + + rewritten = select_pattern.sub(replace_select, rewritten) clause_pattern = re.compile( r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", @@ -837,6 +840,50 @@ def replace_clause(match: re.Match[str]) -> str: return clause_pattern.sub(replace_clause, rewritten) +def _rewrite_mssql_limit_clause(sql: str) -> str: + limit_match = re.search(r"\s+LIMIT\s+(\d+)\s*;?\s*$", sql, flags=re.IGNORECASE) + if not limit_match: + return sql + + limit = limit_match.group(1) + without_limit = sql[: limit_match.start()].rstrip() + if re.search( + r"\bSELECT\s+(?:DISTINCT\s+)?TOP\s+(?:\(\s*)?\d+", + without_limit, + flags=re.IGNORECASE, + ): + return without_limit + + if re.match(r"\s*SELECT\s+DISTINCT\b", without_limit, flags=re.IGNORECASE): + return re.sub( + r"\bSELECT\s+DISTINCT\b", + f"SELECT DISTINCT TOP {limit}", + without_limit, + count=1, + flags=re.IGNORECASE, + ) + + if re.match(r"\s*SELECT\b", without_limit, flags=re.IGNORECASE): + return re.sub( + r"\bSELECT\b", + f"SELECT TOP {limit}", + without_limit, + count=1, + flags=re.IGNORECASE, + ) + + return without_limit + + +def _unwrap_simple_mssql_where_parentheses(sql: str) -> str: + return re.sub( + r"\bWHERE\s*\(\s*([^()]+?)\s*\)(?=\s*(?:GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|$))", + r"WHERE \1", + sql, + flags=re.IGNORECASE | re.DOTALL, + ) + + def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", @@ -994,6 +1041,8 @@ def _references_known_hallucination_prone_schema(sql: str) -> bool: def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _replace_relative_current_date_calls(sql, now) + normalized = _unwrap_simple_mssql_where_parentheses(normalized) + normalized = _rewrite_mssql_limit_clause(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) normalized = _rewrite_mssql_invented_repair_relationship_identifiers(normalized) @@ -1018,6 +1067,8 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = re.sub( r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE ) + normalized = _unwrap_simple_mssql_where_parentheses(normalized) + normalized = _rewrite_mssql_limit_clause(normalized) normalized = re.sub( r"CAST\(\s*('(?:[^']|'')*')\s+AS\s+DATETIME(?:2|OFFSET)\s*\)", r"\1", @@ -1044,6 +1095,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) + normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) elif _references_known_hallucination_prone_schema(normalized): normalized = _rewrite_known_schema_hallucinations(normalized, datetime.now()) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index eec70b3e4e..110c5db6cb 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -724,3 +724,51 @@ def test_normalize_generation_result_sql_rewrites_timestamp_subtraction_for_mssq 'DATEDIFF(\'second\', "created_at", "updated_at") AS "turnaround_seconds"' in normalized ) + + +def test_normalize_generation_result_sql_rewrites_mssql_time_buckets_and_ordering(): + sql = """ + SELECT + YEAR, + MONTH, + COUNT(*) AS repair_count + FROM dbo_repair_logs + GROUP BY YEAR, MONTH + ORDER BY YEAR ASC NULLS LAST, MONTH ASC NULLS LAST + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "NULLS LAST" not in normalized + assert "SELECT YEAR" not in normalized + assert "GROUP BY YEAR" not in normalized + assert "ORDER BY YEAR" not in normalized + assert ( + 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year"' + in normalized + ) + assert ( + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' + in normalized + ) + assert 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at")' in normalized + assert 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC' in normalized + + +def test_normalize_generation_result_sql_rewrites_mssql_limit_and_where_parentheses(): + sql = """ + SELECT model_id, COUNT(*) AS ticket_count + FROM dbo_tickets + WHERE (source = 'AI') + GROUP BY model_id + ORDER BY ticket_count DESC NULLS LAST + LIMIT 1 + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert normalized.startswith("SELECT TOP 1") + assert "WHERE (source = 'AI')" not in normalized + assert "WHERE source = 'AI'" in normalized + assert "NULLS LAST" not in normalized + assert "LIMIT 1" not in normalized diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 6f12946a97..06bd655953 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -150,6 +150,41 @@ const replaceInventedTimeBuckets = (sql: string): string => { return sql; }; +const rewriteMssqlLimitClause = (sql: string): string => { + const limitMatch = sql.match(/\s+LIMIT\s+(\d+)\s*;?\s*$/i); + if (!limitMatch || limitMatch.index === undefined) { + return sql; + } + + const limit = limitMatch[1]; + const withoutLimit = sql.slice(0, limitMatch.index).trimEnd(); + if (/\bSELECT\s+(?:DISTINCT\s+)?TOP\s+(?:\(\s*)?\d+/i.test(withoutLimit)) { + return withoutLimit; + } + + if (/^\s*SELECT\s+DISTINCT\b/i.test(withoutLimit)) { + return withoutLimit.replace(/\bSELECT\s+DISTINCT\b/i, `SELECT DISTINCT TOP ${limit}`); + } + + if (/^\s*SELECT\b/i.test(withoutLimit)) { + return withoutLimit.replace(/\bSELECT\b/i, `SELECT TOP ${limit}`); + } + + return withoutLimit; +}; + +const unwrapSimpleMssqlWhereParentheses = (sql: string): string => + sql.replace( + /\bWHERE\s*\(\s*([^()]+?)\s*\)(?=\s*(?:GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|$))/gis, + 'WHERE $1', + ); + +const normalizeMssqlGeneratedSqlSyntax = (sql: string): string => { + sql = sql.replace(/\s+NULLS\s+(?:LAST|FIRST)\b/gi, ''); + sql = unwrapSimpleMssqlWhereParentheses(sql); + return rewriteMssqlLimitClause(sql); +}; + const replaceBadFailurePatternJoins = (sql: string): string => { if ( !/\bdbo_DebugEntries\b/i.test(sql) || @@ -294,6 +329,7 @@ export const normalizeMssqlGeneratedSqlFields = ( } sql = sql.replace(/\\"/g, '"'); + sql = normalizeMssqlGeneratedSqlSyntax(sql); sql = replaceRelativeCurrentDateCalls(sql); sql = replaceInventedDateFields(sql); sql = replaceRepairLogThroughputShape(sql); @@ -379,5 +415,6 @@ export const normalizeMssqlSqlForIbis = ( dataSource: DataSourceName, ): string => { sql = normalizeMssqlGeneratedSqlFields(sql, dataSource); - return rewriteMssqlDatepartAliasReferences(sql, dataSource); + sql = rewriteMssqlDatepartAliasReferences(sql, dataSource); + return normalizeMssqlGeneratedSqlFields(sql, dataSource); }; From 4dfd85f6eab859e51d9f0d70c47107dd560be978 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 12 Jun 2026 23:18:28 +0530 Subject: [PATCH 0173/1087] Fix project current flag coercion --- .../server/repositories/projectRepository.ts | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 85149c9d49..8d1213c133 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -3,7 +3,6 @@ import { BaseRepository, IBasicRepository, IQueryOptions, - coerceBoolean, } from './baseRepository'; import { camelCase, @@ -174,6 +173,19 @@ export enum WorkspaceProjectType { CLASSIC = 'CLASSIC', } +const coerceProjectBoolean = (value: unknown): boolean => { + if (typeof value === 'boolean') { + return value; + } + if (typeof value === 'number') { + return value === 1; + } + if (typeof value === 'string') { + return ['1', 'true'].includes(value.toLowerCase()); + } + return Boolean(value); +}; + export interface Project { id: number; // ID type: DataSourceName; // Project datasource type. ex: bigquery, mysql, postgresql, mongodb, etc @@ -341,7 +353,7 @@ export class ProjectRepository camelCase(key), ); if (Object.prototype.hasOwnProperty.call(camelCaseData, 'isCurrent')) { - camelCaseData.isCurrent = coerceBoolean(camelCaseData.isCurrent); + camelCaseData.isCurrent = coerceProjectBoolean(camelCaseData.isCurrent); } return camelCaseData as Project; }; From 020c61ea0356490d964be0759866add9bde91cda Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 13 Jun 2026 01:42:41 +0530 Subject: [PATCH 0174/1087] Fix project boolean normalization in project repository --- .../server/repositories/projectRepository.ts | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 8d1213c133..cb142b1201 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -173,19 +173,6 @@ export enum WorkspaceProjectType { CLASSIC = 'CLASSIC', } -const coerceProjectBoolean = (value: unknown): boolean => { - if (typeof value === 'boolean') { - return value; - } - if (typeof value === 'number') { - return value === 1; - } - if (typeof value === 'string') { - return ['1', 'true'].includes(value.toLowerCase()); - } - return Boolean(value); -}; - export interface Project { id: number; // ID type: DataSourceName; // Project datasource type. ex: bigquery, mysql, postgresql, mongodb, etc @@ -353,7 +340,7 @@ export class ProjectRepository camelCase(key), ); if (Object.prototype.hasOwnProperty.call(camelCaseData, 'isCurrent')) { - camelCaseData.isCurrent = coerceProjectBoolean(camelCaseData.isCurrent); + camelCaseData.isCurrent = this.normalizeProjectBoolean(camelCaseData.isCurrent); } return camelCaseData as Project; }; @@ -374,6 +361,19 @@ export class ProjectRepository return formattedData; }; + private normalizeProjectBoolean(value: unknown): boolean { + if (typeof value === 'boolean') { + return value; + } + if (typeof value === 'number') { + return value === 1; + } + if (typeof value === 'string') { + return ['1', 'true'].includes(value.toLowerCase()); + } + return Boolean(value); + } + private isMssql = () => String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; @@ -487,3 +487,4 @@ export class ProjectRepository return message.includes("Cannot insert the value NULL into column 'id'"); }; } + From bab26f48c207a3e9dd1beb231483e2918e8dbdf4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 12:02:08 +0530 Subject: [PATCH 0175/1087] Fix MSSQL normalization for knowledge article fields --- .../src/pipelines/generation/utils/sql.py | 102 +++++++++++++++++- .../pipelines/generation/test_sql_utils.py | 76 ++++++++++++- .../apollo/server/utils/mssqlSqlNormalizer.ts | 96 ++++++++++++++++- .../utils/tests/mssqlSqlNormalizer.test.ts | 59 ++++++++++ 4 files changed, 330 insertions(+), 3 deletions(-) create mode 100644 wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 80c7e80daf..3e0e873339 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -514,6 +514,11 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: return f'{quoted_table_name}."DateIn"' if "report" in normalized_table_name: return f'{quoted_table_name}."generated_at"' + if any( + token in normalized_table_name + for token in ("knowledge", "kb_article", "kb_articles", "article") + ): + return f'{quoted_table_name}."created_at"' if any( token in normalized_table_name for token in ("repair", "ticket", "event", "log") @@ -735,6 +740,71 @@ def _rewrite_mssql_invented_report_fields(sql: str) -> str: return rewritten +def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: + if not re.search( + r"\b(?:dbo_knowledge_articles|dbo_kb_articles)\b", sql, flags=re.IGNORECASE + ): + return sql + + rewritten = sql + table_replacements = { + "dbo_knowledge_articles": { + "effectiveness_score": '"helpful"', + "created_by": '"author"', + "created_by_user": '"author"', + "created_by_user_id": '"author"', + "author_id": '"author"', + }, + "dbo_kb_articles": { + "created_by": '"created_by_user_id"', + "created_by_user": '"created_by_user_id"', + "author": '"created_by_user_id"', + "author_id": '"created_by_user_id"', + }, + } + + for table_name, field_replacements in table_replacements.items(): + article_table = ( + rf'(?:"{table_name}"|\[{table_name}\]|{table_name})' + ) + for invented_field, replacement_field in field_replacements.items(): + rewritten = re.sub( + rf"(?P
{article_table})\s*\.\s*(?:\"{invented_field}\"|\[{invented_field}\]|\b{invented_field}\b)", + rf"\g
.{replacement_field}", + rewritten, + flags=re.IGNORECASE, + ) + + if re.search(r"\bdbo_knowledge_articles\b", rewritten, flags=re.IGNORECASE): + unqualified_replacements = table_replacements["dbo_knowledge_articles"] + elif re.search(r"\bdbo_kb_articles\b", rewritten, flags=re.IGNORECASE): + unqualified_replacements = table_replacements["dbo_kb_articles"] + else: + unqualified_replacements = {} + + for invented_field, replacement_field in unqualified_replacements.items(): + rewritten = re.sub( + rf'(? bool: if re.search(r"(?:->>|->)", sql): return True @@ -1032,7 +1102,7 @@ def replace_clause(match: re.Match[str]) -> str: def _references_known_hallucination_prone_schema(sql: str) -> bool: return bool( re.search( - r"\b(?:dbo_repair_logs|dbo_DebugEntries|dbo_reports)\b", + r"\b(?:dbo_repair_logs|dbo_DebugEntries|dbo_reports|dbo_knowledge_articles|dbo_kb_articles)\b", sql, flags=re.IGNORECASE, ) @@ -1051,6 +1121,7 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_invented_failure_category(normalized) normalized = _rewrite_mssql_invented_report_fields(normalized) + normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) @@ -1058,6 +1129,29 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: return normalized +def _rewrite_mssql_limit_clause(sql: str) -> str: + limit_match = re.search(r"\s+LIMIT\s+(\d+)\s*;?\s*$", sql, flags=re.IGNORECASE) + if not limit_match: + return sql + + limit = limit_match.group(1) + without_limit = sql[: limit_match.start()].rstrip() + if re.search( + r"\bSELECT\s+(?:DISTINCT\s+)?TOP\s*\(?\s*\d+\s*\)?", + without_limit, + flags=re.IGNORECASE, + ): + return without_limit + + return re.sub( + r"\bSELECT\s+(DISTINCT\s+)?", + lambda match: f"{match.group(0)}TOP {limit} ", + without_limit, + count=1, + flags=re.IGNORECASE, + ) + + def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = sql normalized_data_source = normalize_data_source(data_source) @@ -1090,12 +1184,14 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_invented_failure_category(normalized) normalized = _rewrite_mssql_invented_report_fields(normalized) + normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) + normalized = _rewrite_mssql_limit_clause(normalized) elif _references_known_hallucination_prone_schema(normalized): normalized = _rewrite_known_schema_hallucinations(normalized, datetime.now()) @@ -1383,6 +1479,10 @@ async def _classify_generation_result( 3. `GROUP BY "dbo_failure_patterns"."name"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` 4. `GROUP BY "dbo_repair_logs"."failure_code"` and `COUNT(*)` - For chart-oriented questions, ensure the final SELECT contains only the chart-ready dimension and metric columns. Avoid prose-like outputs or helper columns. +- For knowledge article tables: + - Use "created_at" for year/month trend buckets. Do not select, group by, or order by invented "YEAR" or "MONTH" columns. + - In "dbo_knowledge_articles", use "helpful" and "views" for effectiveness-style questions, and use "author" for creator/author groupings. Do not invent "effectiveness_score" or "created_by". + - In "dbo_kb_articles", use "created_by_user_id" for creator groupings. Do not invent "created_by" or "author" unless those exact columns appear in the schema. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 110c5db6cb..370d68fcb4 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -725,7 +725,6 @@ def test_normalize_generation_result_sql_rewrites_timestamp_subtraction_for_mssq in normalized ) - def test_normalize_generation_result_sql_rewrites_mssql_time_buckets_and_ordering(): sql = """ SELECT @@ -772,3 +771,78 @@ def test_normalize_generation_result_sql_rewrites_mssql_limit_and_where_parenthe assert "WHERE source = 'AI'" in normalized assert "NULLS LAST" not in normalized assert "LIMIT 1" not in normalized + + +def test_normalize_generation_result_sql_removes_mssql_limit_when_top_exists(): + sql = """ + SELECT TOP 5 + id + FROM dbo_tickets + ORDER BY id DESC + LIMIT 1 + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "LIMIT" not in normalized + assert "SELECT TOP 5 id" in normalized + + +def test_normalize_generation_result_sql_rewrites_knowledge_article_time_buckets_for_mssql(): + sql = """ + SELECT + "YEAR", + COUNT("dbo_knowledge_articles"."id") AS "article_count" + FROM "dbo_knowledge_articles" + GROUP BY "YEAR" + ORDER BY "YEAR" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert 'SELECT "YEAR"' not in normalized + assert 'GROUP BY "YEAR"' not in normalized + assert 'ORDER BY "YEAR"' not in normalized + assert ( + 'DATEPART(YEAR, "dbo_knowledge_articles"."created_at") AS "year"' + in normalized + ) + assert ( + 'GROUP BY DATEPART(YEAR, "dbo_knowledge_articles"."created_at")' + in normalized + ) + + +def test_normalize_generation_result_sql_rewrites_knowledge_article_hallucinated_fields_for_mssql(): + sql = """ + SELECT + AVG("dbo_knowledge_articles"."effectiveness_score") AS "avg_effectiveness", + "dbo_knowledge_articles"."created_by" AS "created_by" + FROM "dbo_knowledge_articles" + GROUP BY "dbo_knowledge_articles"."created_by" + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "effectiveness_score" not in normalized + assert '"dbo_knowledge_articles"."created_by"' not in normalized + assert 'AVG("dbo_knowledge_articles"."helpful") AS "avg_effectiveness"' in normalized + assert '"dbo_knowledge_articles"."author" AS "author"' in normalized + assert 'GROUP BY "dbo_knowledge_articles"."author"' in normalized + + +def test_normalize_generation_result_sql_rewrites_kb_article_created_by_for_mssql(): + sql = """ + SELECT + created_by, + COUNT(*) AS article_count + FROM dbo_kb_articles + GROUP BY created_by + ORDER BY article_count DESC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "created_by," not in normalized + assert "GROUP BY created_by" not in normalized + assert '"created_by_user_id"' in normalized diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 06bd655953..b68e431da2 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -59,6 +59,12 @@ const inferMssqlTimestampExpression = (sql: string): string => { if (tableName.toLowerCase() === 'dbo_reports') { return `"${tableName}"."generated_at"`; } + if ( + tableName.toLowerCase() === 'dbo_knowledge_articles' || + tableName.toLowerCase() === 'dbo_kb_articles' + ) { + return `"${tableName}"."created_at"`; + } if (tableName.toLowerCase() === 'dbo_repair_logs') { return `"${tableName}"."created_at"`; } @@ -320,6 +326,93 @@ const replaceInventedReportFields = (sql: string): string => { return sql; }; +const replaceInventedKnowledgeArticleFields = (sql: string): string => { + if (!/\b(?:dbo_knowledge_articles|dbo_kb_articles)\b/i.test(sql)) { + return sql; + } + + const replacementsByTable: Record> = { + dbo_knowledge_articles: { + effectiveness_score: '"helpful"', + created_by: '"author"', + created_by_user: '"author"', + created_by_user_id: '"author"', + author_id: '"author"', + }, + dbo_kb_articles: { + created_by: '"created_by_user_id"', + created_by_user: '"created_by_user_id"', + author: '"created_by_user_id"', + author_id: '"created_by_user_id"', + }, + }; + + Object.entries(replacementsByTable).forEach(([tableName, replacements]) => { + const tablePattern = String.raw`(?:"${tableName}"|\[${tableName}\]|${tableName})`; + Object.entries(replacements).forEach(([inventedField, replacementField]) => { + const escapedField = escapeRegex(inventedField); + sql = sql.replace( + new RegExp( + String.raw`(${tablePattern})\s*\.\s*(?:"${escapedField}"|\[${escapedField}\]|\b${escapedField}\b)`, + 'gi', + ), + `$1.${replacementField}`, + ); + }); + }); + + const activeTable = /\bdbo_knowledge_articles\b/i.test(sql) + ? 'dbo_knowledge_articles' + : /\bdbo_kb_articles\b/i.test(sql) + ? 'dbo_kb_articles' + : null; + + if (activeTable) { + Object.entries(replacementsByTable[activeTable]).forEach( + ([inventedField, replacementField]) => { + const escapedField = escapeRegex(inventedField); + sql = sql.replace( + new RegExp(String.raw`(? { + const limitMatch = sql.match(/\s+LIMIT\s+(\d+)\s*;?\s*$/i); + if (!limitMatch || limitMatch.index === undefined) { + return sql; + } + + const limit = limitMatch[1]; + const withoutLimit = sql.slice(0, limitMatch.index).trimEnd(); + if (/\bSELECT\s+(?:DISTINCT\s+)?TOP\s*\(?\s*\d+\s*\)?/i.test(withoutLimit)) { + return withoutLimit; + } + + return withoutLimit.replace( + /\bSELECT\s+(DISTINCT\s+)?/i, + (match) => `${match}TOP ${limit} `, + ); +}; + +const normalizeMssqlGeneratedSqlSyntax = (sql: string): string => { + sql = sql.replace(/\s+NULLS\s+(?:LAST|FIRST)\b/gi, ''); + return rewriteMssqlLimitClause(sql); +}; + export const normalizeMssqlGeneratedSqlFields = ( sql: string, dataSource: DataSourceName, @@ -336,9 +429,10 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = replacePcbThroughputFields(sql); sql = replaceInventedFailureCategory(sql); sql = replaceInventedReportFields(sql); + sql = replaceInventedKnowledgeArticleFields(sql); sql = replaceInventedTimeBuckets(sql); sql = replaceBadFailurePatternJoins(sql); - return sql; + return normalizeMssqlGeneratedSqlSyntax(sql); }; export const rewriteMssqlDatepartAliasReferences = ( diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts new file mode 100644 index 0000000000..070e2ba8eb --- /dev/null +++ b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts @@ -0,0 +1,59 @@ +import { DataSourceName } from '../../types'; +import { normalizeMssqlSqlForIbis } from '../mssqlSqlNormalizer'; + +describe('mssqlSqlNormalizer', () => { + it('rewrites knowledge article time buckets', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT "YEAR", COUNT("dbo_knowledge_articles"."id") AS "article_count" + FROM "dbo_knowledge_articles" + GROUP BY "YEAR" + ORDER BY "YEAR" ASC + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('SELECT "YEAR"'); + expect(normalized).not.toContain('GROUP BY "YEAR"'); + expect(normalized).toContain( + 'DATEPART(YEAR, "dbo_knowledge_articles"."created_at") AS "year"', + ); + }); + + it('rewrites hallucinated knowledge article fields', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT + AVG("dbo_knowledge_articles"."effectiveness_score") AS "avg_effectiveness", + "dbo_knowledge_articles"."created_by" AS "created_by" + FROM "dbo_knowledge_articles" + GROUP BY "dbo_knowledge_articles"."created_by" + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('effectiveness_score'); + expect(normalized).not.toContain('"dbo_knowledge_articles"."created_by"'); + expect(normalized).toContain( + 'AVG("dbo_knowledge_articles"."helpful") AS "avg_effectiveness"', + ); + expect(normalized).toContain('"dbo_knowledge_articles"."author" AS "author"'); + expect(normalized).toContain('GROUP BY "dbo_knowledge_articles"."author"'); + }); + + it('rewrites hallucinated kb article creator fields', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT created_by, COUNT(*) AS article_count + FROM dbo_kb_articles + GROUP BY created_by + ORDER BY article_count DESC + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('created_by,'); + expect(normalized).not.toContain('GROUP BY created_by'); + expect(normalized).toContain('"created_by_user_id"'); + }); +}); From 5ab8d3eb70fe1c904b6873a9041774cbbaaf4b42 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 12:25:32 +0530 Subject: [PATCH 0176/1087] Remove duplicate MSSQL normalizer helpers --- .../apollo/server/utils/mssqlSqlNormalizer.ts | 23 ------------------- 1 file changed, 23 deletions(-) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index b68e431da2..018e389742 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -390,29 +390,6 @@ const replaceInventedKnowledgeArticleFields = (sql: string): string => { return sql; }; -const rewriteMssqlLimitClause = (sql: string): string => { - const limitMatch = sql.match(/\s+LIMIT\s+(\d+)\s*;?\s*$/i); - if (!limitMatch || limitMatch.index === undefined) { - return sql; - } - - const limit = limitMatch[1]; - const withoutLimit = sql.slice(0, limitMatch.index).trimEnd(); - if (/\bSELECT\s+(?:DISTINCT\s+)?TOP\s*\(?\s*\d+\s*\)?/i.test(withoutLimit)) { - return withoutLimit; - } - - return withoutLimit.replace( - /\bSELECT\s+(DISTINCT\s+)?/i, - (match) => `${match}TOP ${limit} `, - ); -}; - -const normalizeMssqlGeneratedSqlSyntax = (sql: string): string => { - sql = sql.replace(/\s+NULLS\s+(?:LAST|FIRST)\b/gi, ''); - return rewriteMssqlLimitClause(sql); -}; - export const normalizeMssqlGeneratedSqlFields = ( sql: string, dataSource: DataSourceName, From b41cf3dc9c36c28ed908eb9b2e6bfc770e51c42c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 12:36:34 +0530 Subject: [PATCH 0177/1087] Rewrite aliased MSSQL time bucket fields --- .../src/pipelines/generation/utils/sql.py | 11 +++-- .../pipelines/generation/test_sql_utils.py | 48 +++++++++++++++++++ .../apollo/server/utils/mssqlSqlNormalizer.ts | 11 ++++- .../utils/tests/mssqlSqlNormalizer.test.ts | 48 +++++++++++++++++++ 4 files changed, 112 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 3e0e873339..1ec8c2dc7f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -855,13 +855,16 @@ def replace_select(match: re.Match[str]) -> str: rebuilt: list[str] = [] changed = False - select_identifier_pattern = re.compile( - rf"^(?:\"{bucket}\"|\[{bucket}\]|{bucket}|{qualified_bucket_pattern.pattern})$", + bucket_select_item_pattern = re.compile( + rf"^(?P(?:\"{bucket}\"|\[{bucket}\]|{bucket}|{qualified_bucket_pattern.pattern}))" + rf"(?:\s+(?:AS\s+)?(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*))?$", re.IGNORECASE, ) for item in items: - if select_identifier_pattern.fullmatch(item.strip()): - rebuilt.append(f'{expression} AS "{bucket}"') + item_match = bucket_select_item_pattern.fullmatch(item.strip()) + if item_match: + alias = item_match.group("alias") or f'"{bucket}"' + rebuilt.append(f"{expression} AS {alias}") changed = True else: rebuilt.append(item) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 370d68fcb4..cd59992342 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -754,6 +754,54 @@ def test_normalize_generation_result_sql_rewrites_mssql_time_buckets_and_orderin assert 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC' in normalized +def test_normalize_generation_result_sql_rewrites_aliased_mssql_time_buckets(): + sql = """ + SELECT + YEAR AS year, + MONTH AS month, + COUNT(*) AS repair_count + FROM dbo_repair_logs + GROUP BY YEAR, MONTH + ORDER BY YEAR ASC, MONTH ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "YEAR AS year" not in normalized + assert "MONTH AS month" not in normalized + assert ( + 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS year' + in normalized + ) + assert ( + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS month' + in normalized + ) + assert 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at")' in normalized + assert 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC' in normalized + + +def test_normalize_generation_result_sql_rewrites_debug_entry_quoted_year_alias(): + sql = """ + SELECT + "YEAR" AS "YEAR", + "dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit", + COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" + FROM "dbo_DebugEntries" + GROUP BY "YEAR", "dbo_DebugEntries"."BusinessUnit" + ORDER BY "YEAR" ASC + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert '"YEAR" AS "YEAR"' not in normalized + assert 'GROUP BY "YEAR"' not in normalized + assert 'ORDER BY "YEAR"' not in normalized + assert 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "YEAR"' in normalized + assert 'GROUP BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn")' in normalized + assert 'ORDER BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC' in normalized + + def test_normalize_generation_result_sql_rewrites_mssql_limit_and_where_parentheses(): sql = """ SELECT model_id, COUNT(*) AS ticket_count diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 018e389742..7d58589c82 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -114,10 +114,17 @@ const replaceInventedTimeBuckets = (sql: string): string => { const alias = bucket.toLowerCase(); body = body.replace( new RegExp( - String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"?${bucket}"?|(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})|\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b|"${bucket}"|\[${bucket}\]|\b${bucket}\b)(?=\s*(?:,|$))`, + String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"?${bucket}"?|(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})|\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b|"${bucket}"|\[${bucket}\]|\b${bucket}\b)(?:\s+(?:AS\s+)?(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*)))?(?=\s*(?:,|$))`, 'gi', ), - `$1 ${expression} AS "${alias}"`, + (_match, prefix, quotedAlias, bracketAlias, bareAlias) => { + const selectedAlias = quotedAlias + ? `"${quotedAlias}"` + : bracketAlias + ? `[${bracketAlias}]` + : bareAlias || `"${alias}"`; + return `${prefix} ${expression} AS ${selectedAlias}`; + }, ); }); return `SELECT${body}`; diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts index 070e2ba8eb..ab43d6bf12 100644 --- a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts +++ b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts @@ -2,6 +2,54 @@ import { DataSourceName } from '../../types'; import { normalizeMssqlSqlForIbis } from '../mssqlSqlNormalizer'; describe('mssqlSqlNormalizer', () => { + it('rewrites aliased repair log time buckets', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT YEAR AS year, MONTH AS month, COUNT(*) AS repair_count + FROM dbo_repair_logs + GROUP BY YEAR, MONTH + ORDER BY YEAR ASC, MONTH ASC + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('YEAR AS year'); + expect(normalized).not.toContain('MONTH AS month'); + expect(normalized).toContain( + 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS year', + ); + expect(normalized).toContain( + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS month', + ); + expect(normalized).toContain( + 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at")', + ); + }); + + it('rewrites quoted debug entry year aliases', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT + "YEAR" AS "YEAR", + "dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit", + COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" + FROM "dbo_DebugEntries" + GROUP BY "YEAR", "dbo_DebugEntries"."BusinessUnit" + ORDER BY "YEAR" ASC + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('"YEAR" AS "YEAR"'); + expect(normalized).not.toContain('GROUP BY "YEAR"'); + expect(normalized).toContain( + 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "YEAR"', + ); + expect(normalized).toContain( + 'GROUP BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn")', + ); + }); + it('rewrites knowledge article time buckets', () => { const normalized = normalizeMssqlSqlForIbis( ` From 1ddace635deab3f32eaef8c58976296a3bd3a5e1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 13:23:39 +0530 Subject: [PATCH 0178/1087] Use planner-safe MSSQL date buckets --- .../src/pipelines/generation/utils/sql.py | 114 +++++++++++------- .../apollo/server/utils/mssqlSqlNormalizer.ts | 43 ++++++- 2 files changed, 109 insertions(+), 48 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1ec8c2dc7f..ae528246eb 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -248,14 +248,14 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", lambda m: ( - f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" + f"(DATEPART('YEAR', {m.group(1)}) * 100 + DATEPART('MONTH', {m.group(1)}))" ), sql, flags=re.IGNORECASE, ) sql = re.sub( rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", sql, flags=re.IGNORECASE, ) @@ -270,124 +270,124 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"DATEPART('YEAR', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"DATEPART('MONTH', {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"DATEPART('DAY', {m.group(1)})", ), ] @@ -422,9 +422,9 @@ def _rewrite_mssql_to_date_buckets(sql: str) -> str: def make_day_bucket(expression: str) -> str: timestamp_expression = _qualify_mssql_temporal_expression(expression, sql) return ( - f"(DATEPART(YEAR, {timestamp_expression}) * 10000 + " - f"DATEPART(MONTH, {timestamp_expression}) * 100 + " - f"DATEPART(DAY, {timestamp_expression}))" + f"(DATEPART('YEAR', {timestamp_expression}) * 10000 + " + f"DATEPART('MONTH', {timestamp_expression}) * 100 + " + f"DATEPART('DAY', {timestamp_expression}))" ) rewritten = re.sub( @@ -671,21 +671,44 @@ def _rewrite_mssql_repair_log_turnaround_trend_shape(sql: str) -> str: return sql if not re.search(r"\bavg_turnaround_time\b|\bturnaround\b", sql, flags=re.IGNORECASE): return sql - if not re.search(r"\bMONTH\b|DATEPART\(\s*MONTH", sql, flags=re.IGNORECASE): + if not re.search(r"\bMONTH\b|DATEPART\(\s*'?\s*MONTH", sql, flags=re.IGNORECASE): return sql return ( - 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' + "SELECT DATEPART('YEAR', \"dbo_repair_logs\".\"created_at\") AS \"year\", " + "DATEPART('MONTH', \"dbo_repair_logs\".\"created_at\") AS \"month\", " 'AVG(DATEDIFF(\'second\', "dbo_repair_logs"."created_at", ' '"dbo_repair_logs"."updated_at")) AS "avg_turnaround_seconds" ' 'FROM "dbo_repair_logs" ' 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' 'AND "dbo_repair_logs"."updated_at" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + "GROUP BY DATEPART('YEAR', \"dbo_repair_logs\".\"created_at\"), " + "DATEPART('MONTH', \"dbo_repair_logs\".\"created_at\") " + "ORDER BY DATEPART('YEAR', \"dbo_repair_logs\".\"created_at\") ASC, " + "DATEPART('MONTH', \"dbo_repair_logs\".\"created_at\") ASC" + ) + + +def _rewrite_mssql_ticket_cycle_turnaround_shape(sql: str) -> str: + if not re.search(r"\bdbo_ticket_cycles\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\bturnaround_time\b|\bavg_turnaround_time\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\bMONTH\b|DATEPART\(\s*'?\s*MONTH", sql, flags=re.IGNORECASE): + return sql + + return ( + 'SELECT DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") AS "year", ' + 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") AS "month", ' + 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", ' + '"dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds" ' + 'FROM "dbo_ticket_cycles" ' + 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL ' + 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL ' + 'GROUP BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at"), ' + 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") ' + 'ORDER BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") ASC, ' + 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") ASC' ) @@ -756,6 +779,9 @@ def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: "author_id": '"author"', }, "dbo_kb_articles": { + "category": '"category"', + "section": '"category"', + "article_section": '"category"', "created_by": '"created_by_user_id"', "created_by_user": '"created_by_user_id"', "author": '"created_by_user_id"', @@ -830,9 +856,9 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: return sql bucket_expressions = { - "year": f"DATEPART(YEAR, {timestamp_expression})", - "month": f"DATEPART(MONTH, {timestamp_expression})", - "day": f"DATEPART(DAY, {timestamp_expression})", + "year": f"DATEPART('YEAR', {timestamp_expression})", + "month": f"DATEPART('MONTH', {timestamp_expression})", + "day": f"DATEPART('DAY', {timestamp_expression})", } rewritten = sql @@ -903,7 +929,7 @@ def replace_clause(match: re.Match[str]) -> str: flags=re.IGNORECASE, ) body = re.sub( - rf"(? str: def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( - r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", + r"\b(DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", re.IGNORECASE, ) aliases: dict[str, str] = {} for match in datepart_alias_pattern.finditer(sql): expression = match.group(1) - alias = match.group(4) or match.group(5) or match.group(6) + alias = match.group(3) or match.group(4) or match.group(5) aliases[alias.lower()] = expression if not aliases: @@ -1105,7 +1131,7 @@ def replace_clause(match: re.Match[str]) -> str: def _references_known_hallucination_prone_schema(sql: str) -> bool: return bool( re.search( - r"\b(?:dbo_repair_logs|dbo_DebugEntries|dbo_reports|dbo_knowledge_articles|dbo_kb_articles)\b", + r"\b(?:dbo_repair_logs|dbo_ticket_cycles|dbo_DebugEntries|dbo_reports|dbo_knowledge_articles|dbo_kb_articles)\b", sql, flags=re.IGNORECASE, ) @@ -1120,6 +1146,7 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _rewrite_mssql_invented_date_identifiers(normalized) normalized = _rewrite_mssql_invented_repair_relationship_identifiers(normalized) normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) + normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_invented_failure_category(normalized) @@ -1183,6 +1210,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized ) normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) + normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_invented_failure_category(normalized) @@ -1458,7 +1486,7 @@ async def _classify_generation_result( _MSSQL_TEXT_TO_SQL_RULES = """ ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. -- Prefer native T-SQL date bucket syntax such as DATEPART(YEAR, "created_at") and DATEPART(MONTH, "created_at"). +- Prefer native T-SQL date bucket syntax such as DATEPART('YEAR', "created_at") and DATEPART('MONTH', "created_at"). - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. - DO NOT use JSON extraction functions or operators such as JSON_VALUE, JSON_QUERY, JSON_EXTRACT, JSON_EXTRACT_SCALAR, JSON_EXTRACT_ARRAY, json_value, json_extract, ->, or ->>. The MSSQL Wren/Ibis runtime does not support them. - If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. @@ -1490,11 +1518,11 @@ async def _classify_generation_result( - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - - DATEPART(YEAR, ) AS "year" - - DATEPART(MONTH, ) AS "month" + - DATEPART('YEAR', ) AS "year" + - DATEPART('MONTH', ) AS "month" Then GROUP BY and ORDER BY the same year/month expressions. - Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the DATEPART(...) expression instead. -- For year bucketing, prefer DATEPART(YEAR, ). +- For year bucketing, prefer DATEPART('YEAR', ). - For top/bottom N questions in MSSQL, prefer SELECT TOP (N) with ORDER BY over DENSE_RANK/ROW_NUMBER when the user did not explicitly request ranks. - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' @@ -1781,7 +1809,7 @@ def get_metric_instructions( #### MSSQL Metric Notes #### - Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. - Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. -- For month trend metrics, prefer DATEPART(YEAR, ) and DATEPART(MONTH, ) as separate grouped columns. +- For month trend metrics, prefer DATEPART('YEAR', ) and DATEPART('MONTH', ) as separate grouped columns. """ return instructions diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 7d58589c82..787d1286e9 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -100,12 +100,18 @@ const replaceInventedDateFields = (sql: string): string => { return sql; }; +const quoteMssqlDatepartUnits = (sql: string): string => + sql.replace( + /\bDATEPART\(\s*'?\s*(YEAR|MONTH|DAY)\s*'?\s*,/gi, + (_match, part) => `DATEPART('${String(part).toUpperCase()}',`, + ); + const replaceInventedTimeBuckets = (sql: string): string => { const timestampExpression = inferMssqlTimestampExpression(sql); const bucketExpressions: Record = { - YEAR: `DATEPART(YEAR, ${timestampExpression})`, - MONTH: `DATEPART(MONTH, ${timestampExpression})`, - DAY: `DATEPART(DAY, ${timestampExpression})`, + YEAR: `DATEPART('YEAR', ${timestampExpression})`, + MONTH: `DATEPART('MONTH', ${timestampExpression})`, + DAY: `DATEPART('DAY', ${timestampExpression})`, }; sql = sql.replace(/\bSELECT\b(?.*?)(?=\bFROM\b)/is, (match, _body, _offset, _source, groups) => { @@ -153,7 +159,7 @@ const replaceInventedTimeBuckets = (sql: string): string => { body = body.replace(new RegExp(String.raw`"${bucket}"`, 'gi'), expression); body = body.replace(new RegExp(String.raw`\[${bucket}\]`, 'gi'), expression); body = body.replace( - new RegExp(String.raw`(? { ].join(' '); }; +const replaceTicketCycleTurnaroundShape = (sql: string): string => { + if ( + !/\bdbo_ticket_cycles\b/i.test(sql) || + !/\b(?:turnaround_time|avg_turnaround_time)\b/i.test(sql) || + !/\bMONTH\b|DATEPART\(\s*'?\s*MONTH/i.test(sql) + ) { + return sql; + } + + return [ + 'SELECT DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") AS "year",', + 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") AS "month",', + 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", "dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds"', + 'FROM "dbo_ticket_cycles"', + 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL', + 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL', + 'GROUP BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at"), DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at")', + 'ORDER BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") ASC, DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") ASC', + ].join(' '); +}; + const replaceInventedFailureCategory = (sql: string): string => { if (!/\bdbo_repair_logs\b/i.test(sql) || !/\bfailure_category\b/i.test(sql)) { return sql; @@ -347,6 +374,9 @@ const replaceInventedKnowledgeArticleFields = (sql: string): string => { author_id: '"author"', }, dbo_kb_articles: { + category: '"category"', + section: '"category"', + article_section: '"category"', created_by: '"created_by_user_id"', created_by_user: '"created_by_user_id"', author: '"created_by_user_id"', @@ -407,15 +437,18 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = sql.replace(/\\"/g, '"'); sql = normalizeMssqlGeneratedSqlSyntax(sql); + sql = quoteMssqlDatepartUnits(sql); sql = replaceRelativeCurrentDateCalls(sql); sql = replaceInventedDateFields(sql); sql = replaceRepairLogThroughputShape(sql); + sql = replaceTicketCycleTurnaroundShape(sql); sql = replacePcbThroughputFields(sql); sql = replaceInventedFailureCategory(sql); sql = replaceInventedReportFields(sql); sql = replaceInventedKnowledgeArticleFields(sql); sql = replaceInventedTimeBuckets(sql); sql = replaceBadFailurePatternJoins(sql); + sql = quoteMssqlDatepartUnits(sql); return normalizeMssqlGeneratedSqlSyntax(sql); }; @@ -434,7 +467,7 @@ export const rewriteMssqlDatepartAliasReferences = ( String.raw`(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))`; const aliasPatterns = [ new RegExp( - String.raw`\b(DATEPART\(\s*(?:YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, + String.raw`\b(DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, 'gi', ), new RegExp( From 9d414b8a52d5e143e6f7a220f11bad3f17ae7511 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 13:43:29 +0530 Subject: [PATCH 0179/1087] Fix fallback charts for categorical results --- .../src/pipelines/generation/utils/chart.py | 34 ++++++++++++++----- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index f972e90095..15ac720604 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -132,27 +132,36 @@ def axis(field: str, field_type: str) -> dict: base["timeUnit"] = "yearmonth" return base + def count_axis() -> dict: + return { + "aggregate": "count", + "type": "quantitative", + "title": "Count", + } + if chart_type == "pie": color_field = nominal[0] if nominal else columns[0] - theta_field = quantitative[0] if quantitative else ( - columns[1] if len(columns) > 1 else columns[0] + theta_encoding = ( + axis(quantitative[0], "quantitative") if quantitative else count_axis() ) return { "title": title, "mark": {"type": "arc"}, "encoding": { - "theta": axis(theta_field, "quantitative"), + "theta": theta_encoding, "color": axis(color_field, "nominal"), }, } if chart_type in {"line", "area", "multi_line"}: - y_field = quantitative[0] if quantitative else columns[-1] + y_encoding = ( + axis(quantitative[0], "quantitative") if quantitative else count_axis() + ) if {"year", "month"}.issubset({c.lower() for c in columns}): month_field = next(c for c in columns if c.lower() == "month") encoding = { "x": axis(month_field, "ordinal"), - "y": axis(y_field, "quantitative"), + "y": y_encoding, } years = [c for c in columns if c.lower() == "year"] if years: @@ -170,16 +179,18 @@ def axis(field: str, field_type: str) -> dict: "mark": {"type": "area" if chart_type == "area" else "line"}, "encoding": { "x": axis(x_field, x_type), - "y": axis(y_field, "quantitative"), + "y": y_encoding, }, } x_field = nominal[0] if nominal else (temporal[0] if temporal else columns[0]) - y_field = quantitative[0] if quantitative else (columns[1] if len(columns) > 1 else columns[0]) x_type = "nominal" if x_field in nominal else ("temporal" if x_field in temporal else "ordinal") + y_encoding = ( + axis(quantitative[0], "quantitative") if quantitative else count_axis() + ) encoding = { "x": axis(x_field, x_type), - "y": axis(y_field, "quantitative"), + "y": y_encoding, } if chart_type == "grouped_bar" and len(nominal) > 1: encoding["xOffset"] = axis(nominal[1], "nominal") @@ -263,6 +274,13 @@ def _needs_deterministic_bar_fallback( if not isinstance(x_axis, dict) or x_axis.get("field") not in nominal: return True + if len(quantitative) == 0 and len(nominal) >= 1: + y_axis = encoding.get("y") + if isinstance(y_axis, dict) and y_axis.get("field") in nominal: + return True + if not isinstance(y_axis, dict) or y_axis.get("aggregate") != "count": + return True + return False From 217b601e9abe6db9c05d4eebb4e40b15b3e05f28 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 14:18:13 +0530 Subject: [PATCH 0180/1087] Use scalar MSSQL date bucket functions --- .../src/pipelines/generation/utils/sql.py | 98 ++++++++++--------- .../apollo/server/utils/mssqlSqlNormalizer.ts | 25 ++--- 2 files changed, 63 insertions(+), 60 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ae528246eb..c9801f5274 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -248,14 +248,14 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", lambda m: ( - f"(DATEPART('YEAR', {m.group(1)}) * 100 + DATEPART('MONTH', {m.group(1)}))" + f"(YEAR({m.group(1)}) * 100 + MONTH({m.group(1)}))" ), sql, flags=re.IGNORECASE, ) sql = re.sub( rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", sql, flags=re.IGNORECASE, ) @@ -270,124 +270,124 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"MONTH({m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DAY({m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"MONTH({m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DAY({m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"MONTH({m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"MONTH({m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"MONTH({m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DAY({m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"MONTH({m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DAY({m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('YEAR', {m.group(1)})", + lambda m: f"YEAR({m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('MONTH', {m.group(1)})", + lambda m: f"MONTH({m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART('DAY', {m.group(1)})", + lambda m: f"DAY({m.group(1)})", ), ] @@ -422,9 +422,9 @@ def _rewrite_mssql_to_date_buckets(sql: str) -> str: def make_day_bucket(expression: str) -> str: timestamp_expression = _qualify_mssql_temporal_expression(expression, sql) return ( - f"(DATEPART('YEAR', {timestamp_expression}) * 10000 + " - f"DATEPART('MONTH', {timestamp_expression}) * 100 + " - f"DATEPART('DAY', {timestamp_expression}))" + f"(YEAR({timestamp_expression}) * 10000 + " + f"MONTH({timestamp_expression}) * 100 + " + f"DAY({timestamp_expression}))" ) rewritten = re.sub( @@ -675,17 +675,17 @@ def _rewrite_mssql_repair_log_turnaround_trend_shape(sql: str) -> str: return sql return ( - "SELECT DATEPART('YEAR', \"dbo_repair_logs\".\"created_at\") AS \"year\", " - "DATEPART('MONTH', \"dbo_repair_logs\".\"created_at\") AS \"month\", " + "SELECT YEAR(\"dbo_repair_logs\".\"created_at\") AS \"year\", " + "MONTH(\"dbo_repair_logs\".\"created_at\") AS \"month\", " 'AVG(DATEDIFF(\'second\', "dbo_repair_logs"."created_at", ' '"dbo_repair_logs"."updated_at")) AS "avg_turnaround_seconds" ' 'FROM "dbo_repair_logs" ' 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' 'AND "dbo_repair_logs"."updated_at" IS NOT NULL ' - "GROUP BY DATEPART('YEAR', \"dbo_repair_logs\".\"created_at\"), " - "DATEPART('MONTH', \"dbo_repair_logs\".\"created_at\") " - "ORDER BY DATEPART('YEAR', \"dbo_repair_logs\".\"created_at\") ASC, " - "DATEPART('MONTH', \"dbo_repair_logs\".\"created_at\") ASC" + "GROUP BY YEAR(\"dbo_repair_logs\".\"created_at\"), " + "MONTH(\"dbo_repair_logs\".\"created_at\") " + "ORDER BY YEAR(\"dbo_repair_logs\".\"created_at\") ASC, " + "MONTH(\"dbo_repair_logs\".\"created_at\") ASC" ) @@ -698,17 +698,17 @@ def _rewrite_mssql_ticket_cycle_turnaround_shape(sql: str) -> str: return sql return ( - 'SELECT DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") AS "year", ' - 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") AS "month", ' + 'SELECT YEAR("dbo_ticket_cycles"."created_at") AS "year", ' + 'MONTH("dbo_ticket_cycles"."created_at") AS "month", ' 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", ' '"dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds" ' 'FROM "dbo_ticket_cycles" ' 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL ' 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL ' - 'GROUP BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at"), ' - 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") ' - 'ORDER BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") ASC, ' - 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") ASC' + 'GROUP BY YEAR("dbo_ticket_cycles"."created_at"), ' + 'MONTH("dbo_ticket_cycles"."created_at") ' + 'ORDER BY YEAR("dbo_ticket_cycles"."created_at") ASC, ' + 'MONTH("dbo_ticket_cycles"."created_at") ASC' ) @@ -856,9 +856,9 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: return sql bucket_expressions = { - "year": f"DATEPART('YEAR', {timestamp_expression})", - "month": f"DATEPART('MONTH', {timestamp_expression})", - "day": f"DATEPART('DAY', {timestamp_expression})", + "year": f"YEAR({timestamp_expression})", + "month": f"MONTH({timestamp_expression})", + "day": f"DAY({timestamp_expression})", } rewritten = sql @@ -985,14 +985,14 @@ def _unwrap_simple_mssql_where_parentheses(sql: str) -> str: def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( - r"\b(DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", + r"\b((?:DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)|(?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\)))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", re.IGNORECASE, ) aliases: dict[str, str] = {} for match in datepart_alias_pattern.finditer(sql): expression = match.group(1) - alias = match.group(3) or match.group(4) or match.group(5) + alias = match.group(4) or match.group(5) or match.group(6) aliases[alias.lower()] = expression if not aliases: @@ -1074,7 +1074,9 @@ def _rewrite_mssql_temporal_bucket_alias_references(sql: str) -> str: aliases: dict[str, str] = {} for item in _split_top_level_select_items(select_match.group("body")): - if not re.search(r"\bDATEPART\s*\(", item, flags=re.IGNORECASE): + if not re.search( + r"\b(?:DATEPART|YEAR|MONTH|DAY)\s*\(", item, flags=re.IGNORECASE + ): continue alias_match = re.search( @@ -1486,7 +1488,7 @@ async def _classify_generation_result( _MSSQL_TEXT_TO_SQL_RULES = """ ### MSSQL-SPECIFIC RULES ### - The target database is MSSQL. -- Prefer native T-SQL date bucket syntax such as DATEPART('YEAR', "created_at") and DATEPART('MONTH', "created_at"). +- Prefer native T-SQL date bucket syntax such as YEAR("created_at") and MONTH("created_at"). - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. - DO NOT use JSON extraction functions or operators such as JSON_VALUE, JSON_QUERY, JSON_EXTRACT, JSON_EXTRACT_SCALAR, JSON_EXTRACT_ARRAY, json_value, json_extract, ->, or ->>. The MSSQL Wren/Ibis runtime does not support them. - If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. @@ -1518,11 +1520,11 @@ async def _classify_generation_result( - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - - DATEPART('YEAR', ) AS "year" - - DATEPART('MONTH', ) AS "month" + - YEAR() AS "year" + - MONTH() AS "month" Then GROUP BY and ORDER BY the same year/month expressions. -- Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the DATEPART(...) expression instead. -- For year bucketing, prefer DATEPART('YEAR', ). +- Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the YEAR(...) or MONTH(...) expression instead. +- For year bucketing, prefer YEAR(). - For top/bottom N questions in MSSQL, prefer SELECT TOP (N) with ORDER BY over DENSE_RANK/ROW_NUMBER when the user did not explicitly request ranks. - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' @@ -1809,7 +1811,7 @@ def get_metric_instructions( #### MSSQL Metric Notes #### - Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. - Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. -- For month trend metrics, prefer DATEPART('YEAR', ) and DATEPART('MONTH', ) as separate grouped columns. +- For month trend metrics, prefer YEAR() and MONTH() as separate grouped columns. """ return instructions diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 787d1286e9..cf8d27741f 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -100,18 +100,19 @@ const replaceInventedDateFields = (sql: string): string => { return sql; }; -const quoteMssqlDatepartUnits = (sql: string): string => +const rewriteMssqlDatepartFunctions = (sql: string): string => sql.replace( - /\bDATEPART\(\s*'?\s*(YEAR|MONTH|DAY)\s*'?\s*,/gi, - (_match, part) => `DATEPART('${String(part).toUpperCase()}',`, + /\bDATEPART\(\s*'?\s*(YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)/gi, + (_match, part, expression) => + `${String(part).toUpperCase()}(${String(expression).trim()})`, ); const replaceInventedTimeBuckets = (sql: string): string => { const timestampExpression = inferMssqlTimestampExpression(sql); const bucketExpressions: Record = { - YEAR: `DATEPART('YEAR', ${timestampExpression})`, - MONTH: `DATEPART('MONTH', ${timestampExpression})`, - DAY: `DATEPART('DAY', ${timestampExpression})`, + YEAR: `YEAR(${timestampExpression})`, + MONTH: `MONTH(${timestampExpression})`, + DAY: `DAY(${timestampExpression})`, }; sql = sql.replace(/\bSELECT\b(?.*?)(?=\bFROM\b)/is, (match, _body, _offset, _source, groups) => { @@ -300,14 +301,14 @@ const replaceTicketCycleTurnaroundShape = (sql: string): string => { } return [ - 'SELECT DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") AS "year",', - 'DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") AS "month",', + 'SELECT YEAR("dbo_ticket_cycles"."created_at") AS "year",', + 'MONTH("dbo_ticket_cycles"."created_at") AS "month",', 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", "dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds"', 'FROM "dbo_ticket_cycles"', 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL', 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL', - 'GROUP BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at"), DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at")', - 'ORDER BY DATEPART(\'YEAR\', "dbo_ticket_cycles"."created_at") ASC, DATEPART(\'MONTH\', "dbo_ticket_cycles"."created_at") ASC', + 'GROUP BY YEAR("dbo_ticket_cycles"."created_at"), MONTH("dbo_ticket_cycles"."created_at")', + 'ORDER BY YEAR("dbo_ticket_cycles"."created_at") ASC, MONTH("dbo_ticket_cycles"."created_at") ASC', ].join(' '); }; @@ -437,7 +438,7 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = sql.replace(/\\"/g, '"'); sql = normalizeMssqlGeneratedSqlSyntax(sql); - sql = quoteMssqlDatepartUnits(sql); + sql = rewriteMssqlDatepartFunctions(sql); sql = replaceRelativeCurrentDateCalls(sql); sql = replaceInventedDateFields(sql); sql = replaceRepairLogThroughputShape(sql); @@ -448,7 +449,7 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = replaceInventedKnowledgeArticleFields(sql); sql = replaceInventedTimeBuckets(sql); sql = replaceBadFailurePatternJoins(sql); - sql = quoteMssqlDatepartUnits(sql); + sql = rewriteMssqlDatepartFunctions(sql); return normalizeMssqlGeneratedSqlSyntax(sql); }; From 25c4aedcadf1c7d8e77d637a4b4f09e9b6e3c036 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 14:46:53 +0530 Subject: [PATCH 0181/1087] Normalize MSSQL buckets with extract expressions --- .../src/pipelines/generation/utils/sql.py | 112 +++++++++--------- .../src/pipelines/sql_normalizer.py | 98 ++++++++++----- .../apollo/server/utils/mssqlSqlNormalizer.ts | 23 ++-- 3 files changed, 142 insertions(+), 91 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c9801f5274..01a5efbec5 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -248,14 +248,14 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", lambda m: ( - f"(YEAR({m.group(1)}) * 100 + MONTH({m.group(1)}))" + f"(EXTRACT(YEAR FROM {m.group(1)}) * 100 + EXTRACT(MONTH FROM {m.group(1)}))" ), sql, flags=re.IGNORECASE, ) sql = re.sub( rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", sql, flags=re.IGNORECASE, ) @@ -270,124 +270,124 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"MONTH({m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DAY({m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"MONTH({m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DAY({m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"MONTH({m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"MONTH({m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"MONTH({m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DAY({m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"MONTH({m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DAY({m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"YEAR({m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"MONTH({m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DAY({m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ] @@ -422,9 +422,9 @@ def _rewrite_mssql_to_date_buckets(sql: str) -> str: def make_day_bucket(expression: str) -> str: timestamp_expression = _qualify_mssql_temporal_expression(expression, sql) return ( - f"(YEAR({timestamp_expression}) * 10000 + " - f"MONTH({timestamp_expression}) * 100 + " - f"DAY({timestamp_expression}))" + f"(EXTRACT(YEAR FROM {timestamp_expression}) * 10000 + " + f"EXTRACT(MONTH FROM {timestamp_expression}) * 100 + " + f"EXTRACT(DAY FROM {timestamp_expression}))" ) rewritten = re.sub( @@ -675,17 +675,17 @@ def _rewrite_mssql_repair_log_turnaround_trend_shape(sql: str) -> str: return sql return ( - "SELECT YEAR(\"dbo_repair_logs\".\"created_at\") AS \"year\", " - "MONTH(\"dbo_repair_logs\".\"created_at\") AS \"month\", " + "SELECT EXTRACT(YEAR FROM \"dbo_repair_logs\".\"created_at\") AS \"year\", " + "EXTRACT(MONTH FROM \"dbo_repair_logs\".\"created_at\") AS \"month\", " 'AVG(DATEDIFF(\'second\', "dbo_repair_logs"."created_at", ' '"dbo_repair_logs"."updated_at")) AS "avg_turnaround_seconds" ' 'FROM "dbo_repair_logs" ' 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' 'AND "dbo_repair_logs"."updated_at" IS NOT NULL ' - "GROUP BY YEAR(\"dbo_repair_logs\".\"created_at\"), " - "MONTH(\"dbo_repair_logs\".\"created_at\") " - "ORDER BY YEAR(\"dbo_repair_logs\".\"created_at\") ASC, " - "MONTH(\"dbo_repair_logs\".\"created_at\") ASC" + "GROUP BY EXTRACT(YEAR FROM \"dbo_repair_logs\".\"created_at\"), " + "EXTRACT(MONTH FROM \"dbo_repair_logs\".\"created_at\") " + "ORDER BY EXTRACT(YEAR FROM \"dbo_repair_logs\".\"created_at\") ASC, " + "EXTRACT(MONTH FROM \"dbo_repair_logs\".\"created_at\") ASC" ) @@ -698,33 +698,39 @@ def _rewrite_mssql_ticket_cycle_turnaround_shape(sql: str) -> str: return sql return ( - 'SELECT YEAR("dbo_ticket_cycles"."created_at") AS "year", ' - 'MONTH("dbo_ticket_cycles"."created_at") AS "month", ' + 'SELECT EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") AS "year", ' + 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") AS "month", ' 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", ' '"dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds" ' 'FROM "dbo_ticket_cycles" ' 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL ' 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL ' - 'GROUP BY YEAR("dbo_ticket_cycles"."created_at"), ' - 'MONTH("dbo_ticket_cycles"."created_at") ' - 'ORDER BY YEAR("dbo_ticket_cycles"."created_at") ASC, ' - 'MONTH("dbo_ticket_cycles"."created_at") ASC' + 'GROUP BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at"), ' + 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") ' + 'ORDER BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") ASC, ' + 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") ASC' ) def _rewrite_mssql_invented_failure_category(sql: str) -> str: if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): return sql - if not re.search(r"\bfailure_category\b", sql, flags=re.IGNORECASE): + if not re.search(r"\bfailure[_\s]+category\b", sql, flags=re.IGNORECASE): return sql failure_code_expression = '"dbo_repair_logs"."failure_code"' rewritten = re.sub( - r'(?P\bSELECT\s+|,\s*)(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)(?P\s*(?:,|\bFROM\b))', + r'(?P\bSELECT\s+|,\s*)(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)(?P\s*(?:,|\bFROM\b))', rf'\g{failure_code_expression} AS "failure_category"\g', sql, flags=re.IGNORECASE, ) + rewritten = re.sub( + r"\bAS\s+failure\s+category\b", + 'AS "failure_category"', + rewritten, + flags=re.IGNORECASE, + ) clause_pattern = re.compile( r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", @@ -733,7 +739,7 @@ def _rewrite_mssql_invented_failure_category(sql: str) -> str: def replace_clause(match: re.Match[str]) -> str: body = re.sub( - r'(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)', + r'(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)', failure_code_expression, match.group("body"), flags=re.IGNORECASE, @@ -856,9 +862,9 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: return sql bucket_expressions = { - "year": f"YEAR({timestamp_expression})", - "month": f"MONTH({timestamp_expression})", - "day": f"DAY({timestamp_expression})", + "year": f"EXTRACT(YEAR FROM {timestamp_expression})", + "month": f"EXTRACT(MONTH FROM {timestamp_expression})", + "day": f"EXTRACT(DAY FROM {timestamp_expression})", } rewritten = sql @@ -985,14 +991,14 @@ def _unwrap_simple_mssql_where_parentheses(sql: str) -> str: def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( - r"\b((?:DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)|(?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\)))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", + r"\b((?:DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)|(?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\)|EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\)))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", re.IGNORECASE, ) aliases: dict[str, str] = {} for match in datepart_alias_pattern.finditer(sql): expression = match.group(1) - alias = match.group(4) or match.group(5) or match.group(6) + alias = match.group(5) or match.group(6) or match.group(7) aliases[alias.lower()] = expression if not aliases: @@ -1022,7 +1028,7 @@ def replace_clause(match: re.Match) -> str: flags=re.IGNORECASE, ) body = re.sub( - rf"(? str: aliases: dict[str, str] = {} for item in _split_top_level_select_items(select_match.group("body")): if not re.search( - r"\b(?:DATEPART|YEAR|MONTH|DAY)\s*\(", item, flags=re.IGNORECASE + r"\b(?:DATEPART|YEAR|MONTH|DAY|EXTRACT)\s*\(", item, flags=re.IGNORECASE ): continue @@ -1118,7 +1124,7 @@ def replace_clause(match: re.Match[str]) -> str: flags=re.IGNORECASE, ) body = re.sub( - rf"(?, or ->>. The MSSQL Wren/Ibis runtime does not support them. - If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. @@ -1520,11 +1526,11 @@ async def _classify_generation_result( - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. - For month bucketing, prefer separate year/month fields: - - YEAR() AS "year" - - MONTH() AS "month" + - EXTRACT(YEAR FROM ) AS "year" + - EXTRACT(MONTH FROM ) AS "month" Then GROUP BY and ORDER BY the same year/month expressions. -- Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the YEAR(...) or MONTH(...) expression instead. -- For year bucketing, prefer YEAR(). +- Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the EXTRACT(...) expression instead. +- For year bucketing, prefer EXTRACT(YEAR FROM ). - For top/bottom N questions in MSSQL, prefer SELECT TOP (N) with ORDER BY over DENSE_RANK/ROW_NUMBER when the user did not explicitly request ranks. - For filtering a specific year such as 2025, prefer a closed-open range: - >= '2025-01-01 00:00:00' @@ -1811,7 +1817,7 @@ def get_metric_instructions( #### MSSQL Metric Notes #### - Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. - Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. -- For month trend metrics, prefer YEAR() and MONTH() as separate grouped columns. +- For month trend metrics, prefer EXTRACT(YEAR FROM ) and EXTRACT(MONTH FROM ) as separate grouped columns. """ return instructions diff --git a/wren-ai-service/src/pipelines/sql_normalizer.py b/wren-ai-service/src/pipelines/sql_normalizer.py index b8071e268c..2adfe7398e 100644 --- a/wren-ai-service/src/pipelines/sql_normalizer.py +++ b/wren-ai-service/src/pipelines/sql_normalizer.py @@ -100,14 +100,14 @@ def _rewrite_mssql_bucket_functions(sql: str) -> str: sql = re.sub( rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", lambda m: ( - f"(DATEPART(YEAR, {m.group(1)}) * 100 + DATEPART(MONTH, {m.group(1)}))" + f"(EXTRACT(YEAR FROM {m.group(1)}) * 100 + EXTRACT(MONTH FROM {m.group(1)}))" ), sql, flags=re.IGNORECASE, ) sql = re.sub( rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", sql, flags=re.IGNORECASE, ) @@ -122,124 +122,124 @@ def _rewrite_temporal_bucket_functions(sql: str) -> str: rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(YEAR, {m.group(1)})", + lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(MONTH, {m.group(1)})", + lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", ), ( re.compile( rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", re.IGNORECASE, ), - lambda m: f"DATEPART(DAY, {m.group(1)})", + lambda m: f"EXTRACT(DAY FROM {m.group(1)})", ), ] @@ -345,9 +345,9 @@ def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: return sql bucket_expressions = { - "year": f"DATEPART(YEAR, {timestamp_expression})", - "month": f"DATEPART(MONTH, {timestamp_expression})", - "day": f"DATEPART(DAY, {timestamp_expression})", + "year": f"EXTRACT(YEAR FROM {timestamp_expression})", + "month": f"EXTRACT(MONTH FROM {timestamp_expression})", + "day": f"EXTRACT(DAY FROM {timestamp_expression})", } rewritten = sql @@ -391,16 +391,53 @@ def replace_clause(match: re.Match[str]) -> str: return clause_pattern.sub(replace_clause, rewritten) +def _rewrite_mssql_invented_failure_category(sql: str) -> str: + if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): + return sql + if not re.search(r"\bfailure[_\s]+category\b", sql, flags=re.IGNORECASE): + return sql + + failure_code_expression = '"dbo_repair_logs"."failure_code"' + rewritten = re.sub( + r'(?P\bSELECT\s+|,\s*)(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)(?P\s*(?:,|\bFROM\b))', + rf'\g{failure_code_expression} AS "failure_category"\g', + sql, + flags=re.IGNORECASE, + ) + rewritten = re.sub( + r"\bAS\s+failure\s+category\b", + 'AS "failure_category"', + rewritten, + flags=re.IGNORECASE, + ) + + clause_pattern = re.compile( + r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + + def replace_clause(match: re.Match[str]) -> str: + body = re.sub( + r'(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)', + failure_code_expression, + match.group("body"), + flags=re.IGNORECASE, + ) + return f"{match.group(1)}{body}" + + return clause_pattern.sub(replace_clause, rewritten) + + def _rewrite_mssql_datepart_alias_references(sql: str) -> str: datepart_alias_pattern = re.compile( - r"\b(DATEPART\(\s*(YEAR|MONTH|DAY)\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", + r"\b((?:DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)|(?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\)|EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\)))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", re.IGNORECASE, ) aliases: dict[str, str] = {} for match in datepart_alias_pattern.finditer(sql): expression = match.group(1) - alias = match.group(4) or match.group(5) or match.group(6) + alias = match.group(5) or match.group(6) or match.group(7) aliases[alias.lower()] = expression if not aliases: @@ -430,7 +467,7 @@ def replace_clause(match: re.Match) -> str: flags=re.IGNORECASE, ) body = re.sub( - rf"(? flags=re.IGNORECASE, ) normalized = _replace_relative_getdate_calls(normalized, now) + normalized = re.sub( + r"\bAS\s+failure\s+category\b", + 'AS "failure_category"', + normalized, + flags=re.IGNORECASE, + ) normalized = _rewrite_mssql_to_unixtime(normalized) normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_invented_failure_category(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index cf8d27741f..1366592138 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -104,15 +104,15 @@ const rewriteMssqlDatepartFunctions = (sql: string): string => sql.replace( /\bDATEPART\(\s*'?\s*(YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)/gi, (_match, part, expression) => - `${String(part).toUpperCase()}(${String(expression).trim()})`, + `EXTRACT(${String(part).toUpperCase()} FROM ${String(expression).trim()})`, ); const replaceInventedTimeBuckets = (sql: string): string => { const timestampExpression = inferMssqlTimestampExpression(sql); const bucketExpressions: Record = { - YEAR: `YEAR(${timestampExpression})`, - MONTH: `MONTH(${timestampExpression})`, - DAY: `DAY(${timestampExpression})`, + YEAR: `EXTRACT(YEAR FROM ${timestampExpression})`, + MONTH: `EXTRACT(MONTH FROM ${timestampExpression})`, + DAY: `EXTRACT(DAY FROM ${timestampExpression})`, }; sql = sql.replace(/\bSELECT\b(?.*?)(?=\bFROM\b)/is, (match, _body, _offset, _source, groups) => { @@ -301,19 +301,19 @@ const replaceTicketCycleTurnaroundShape = (sql: string): string => { } return [ - 'SELECT YEAR("dbo_ticket_cycles"."created_at") AS "year",', - 'MONTH("dbo_ticket_cycles"."created_at") AS "month",', + 'SELECT EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") AS "year",', + 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") AS "month",', 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", "dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds"', 'FROM "dbo_ticket_cycles"', 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL', 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL', - 'GROUP BY YEAR("dbo_ticket_cycles"."created_at"), MONTH("dbo_ticket_cycles"."created_at")', - 'ORDER BY YEAR("dbo_ticket_cycles"."created_at") ASC, MONTH("dbo_ticket_cycles"."created_at") ASC', + 'GROUP BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at"), EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at")', + 'ORDER BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") ASC, EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") ASC', ].join(' '); }; const replaceInventedFailureCategory = (sql: string): string => { - if (!/\bdbo_repair_logs\b/i.test(sql) || !/\bfailure_category\b/i.test(sql)) { + if (!/\bdbo_repair_logs\b/i.test(sql) || !/\bfailure[_\s]+category\b/i.test(sql)) { return sql; } @@ -323,18 +323,19 @@ const replaceInventedFailureCategory = (sql: string): string => { (match, _body, _offset, _source, groups) => { let body = groups?.body || ''; body = body.replace( - /(^|,)\s*(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)(?=\s*(?:,|$))/gi, + /(^|,)\s*(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)(?=\s*(?:,|$))/gi, `$1 ${failureCodeExpression} AS "failure_category"`, ); return `SELECT${body}`; }, ); + sql = sql.replace(/\bAS\s+failure\s+category\b/gi, 'AS "failure_category"'); const clausePattern = /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { const body = (groups?.body || '').replace( - /(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure_category"|\[failure_category\]|failure_category)/gi, + /(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)/gi, failureCodeExpression, ); return `${clause}${body}`; From 982d2c911dfe1e27c0e6309fbf4850eaf5aeb041 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 15:08:41 +0530 Subject: [PATCH 0182/1087] Support legacy dashboard item titles --- .../repositories/dashboardItemRepository.ts | 85 ++++++++++++++++++- 1 file changed, 83 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index 12959f0167..ad91bfc4e0 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -40,6 +44,7 @@ export interface DashboardItem { layout: DashboardItemLayout; detail: DashboardItemDetail; displayName?: string; + title?: string; } export interface IDashboardItemRepository @@ -50,11 +55,35 @@ export class DashboardItemRepository implements IDashboardItemRepository { private readonly jsonbColumns = ['layout', 'detail']; + private hasTitleColumnCache: boolean | null = null; + private hasDisplayNameColumnCache: boolean | null = null; constructor(knexPg: Knex) { super({ knexPg, tableName: 'dashboard_item' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return await super.createOne( + await this.normalizeWriteData(data, queryOptions), + queryOptions, + ); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return await super.updateOne( + id, + await this.normalizeWriteData(data, queryOptions), + queryOptions, + ); + } + protected override transformFromDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); @@ -70,7 +99,10 @@ export class DashboardItemRepository } return value; }); - return transformData as DashboardItem; + return { + ...transformData, + displayName: transformData.displayName || transformData.title, + } as DashboardItem; }; protected override transformToDBData = (data: any) => { @@ -86,4 +118,53 @@ export class DashboardItemRepository }); return mapKeys(transformedData, (_value, key) => snakeCase(key)); }; + + private async normalizeWriteData( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise> { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [hasTitleColumn, hasDisplayNameColumn] = await Promise.all([ + this.hasColumn('title', executer), + this.hasColumn('display_name', executer), + ]); + const normalizedData: Partial = { ...data }; + const displayName = + typeof data.displayName === 'string' ? data.displayName.trim() : ''; + const chartTitle = + typeof data.detail?.chartSchema?.title === 'string' + ? data.detail.chartSchema.title.trim() + : ''; + const title = displayName || chartTitle || 'Untitled dashboard item'; + + if (hasTitleColumn && !normalizedData.title) { + normalizedData.title = title; + } + if (!hasDisplayNameColumn) { + delete normalizedData.displayName; + } + + return normalizedData; + } + + private async hasColumn(column: string, executer: Knex | Knex.Transaction) { + if (column === 'title' && this.hasTitleColumnCache !== null) { + return this.hasTitleColumnCache; + } + if ( + column === 'display_name' && + this.hasDisplayNameColumnCache !== null + ) { + return this.hasDisplayNameColumnCache; + } + + const result = await executer.schema.hasColumn(this.tableName, column); + if (column === 'title') { + this.hasTitleColumnCache = result; + } + if (column === 'display_name') { + this.hasDisplayNameColumnCache = result; + } + return result; + } } From 61c44a223125c3d62a20283b451fedca709da588 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 15:18:01 +0530 Subject: [PATCH 0183/1087] Generate dashboard item ids for legacy MSSQL --- .../repositories/dashboardItemRepository.ts | 47 +++++++++++++++++-- 1 file changed, 42 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index ad91bfc4e0..de5a57f10c 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -55,6 +55,7 @@ export class DashboardItemRepository implements IDashboardItemRepository { private readonly jsonbColumns = ['layout', 'detail']; + private hasIdColumnCache: boolean | null = null; private hasTitleColumnCache: boolean | null = null; private hasDisplayNameColumnCache: boolean | null = null; @@ -67,7 +68,7 @@ export class DashboardItemRepository queryOptions?: IQueryOptions, ): Promise { return await super.createOne( - await this.normalizeWriteData(data, queryOptions), + await this.normalizeWriteData(data, queryOptions, true), queryOptions, ); } @@ -122,12 +123,15 @@ export class DashboardItemRepository private async normalizeWriteData( data: Partial, queryOptions?: IQueryOptions, + includeGeneratedId = false, ): Promise> { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [hasTitleColumn, hasDisplayNameColumn] = await Promise.all([ - this.hasColumn('title', executer), - this.hasColumn('display_name', executer), - ]); + const [hasIdColumn, hasTitleColumn, hasDisplayNameColumn] = + await Promise.all([ + this.hasColumn('id', executer), + this.hasColumn('title', executer), + this.hasColumn('display_name', executer), + ]); const normalizedData: Partial = { ...data }; const displayName = typeof data.displayName === 'string' ? data.displayName.trim() : ''; @@ -143,11 +147,22 @@ export class DashboardItemRepository if (!hasDisplayNameColumn) { delete normalizedData.displayName; } + if ( + includeGeneratedId && + hasIdColumn && + normalizedData.id === undefined && + this.isMssqlLike(executer) + ) { + normalizedData.id = await this.getNextId(executer); + } return normalizedData; } private async hasColumn(column: string, executer: Knex | Knex.Transaction) { + if (column === 'id' && this.hasIdColumnCache !== null) { + return this.hasIdColumnCache; + } if (column === 'title' && this.hasTitleColumnCache !== null) { return this.hasTitleColumnCache; } @@ -159,6 +174,9 @@ export class DashboardItemRepository } const result = await executer.schema.hasColumn(this.tableName, column); + if (column === 'id') { + this.hasIdColumnCache = result; + } if (column === 'title') { this.hasTitleColumnCache = result; } @@ -167,4 +185,23 @@ export class DashboardItemRepository } return result; } + + private isMssqlLike(executer: Knex | Knex.Transaction) { + const clientName = String(executer.client.config.client || '').toLowerCase(); + const dialect = String((executer.client as any).dialect || '').toLowerCase(); + const driverName = String( + (executer.client as any).driverName || '', + ).toLowerCase(); + + return [clientName, dialect, driverName].some((value) => + value.includes('mssql'), + ); + } + + private async getNextId(executer: Knex | Knex.Transaction) { + const [row] = await executer(this.tableName).max<{ maxId: number | null }>( + 'id as maxId', + ); + return (row?.maxId || 0) + 1; + } } From 832d0c0f1c7c58e521b46822f17d37e7e9cd0410 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 15:32:54 +0530 Subject: [PATCH 0184/1087] Set dashboard item timestamps for legacy MSSQL --- .../repositories/dashboardItemRepository.ts | 23 ++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index de5a57f10c..bb9f695ecd 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -45,6 +45,8 @@ export interface DashboardItem { detail: DashboardItemDetail; displayName?: string; title?: string; + createdAt?: Date; + updatedAt?: Date; } export interface IDashboardItemRepository @@ -58,6 +60,7 @@ export class DashboardItemRepository private hasIdColumnCache: boolean | null = null; private hasTitleColumnCache: boolean | null = null; private hasDisplayNameColumnCache: boolean | null = null; + private columnCache = new Map(); constructor(knexPg: Knex) { super({ knexPg, tableName: 'dashboard_item' }); @@ -126,11 +129,19 @@ export class DashboardItemRepository includeGeneratedId = false, ): Promise> { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [hasIdColumn, hasTitleColumn, hasDisplayNameColumn] = + const [ + hasIdColumn, + hasTitleColumn, + hasDisplayNameColumn, + hasCreatedAtColumn, + hasUpdatedAtColumn, + ] = await Promise.all([ this.hasColumn('id', executer), this.hasColumn('title', executer), this.hasColumn('display_name', executer), + this.hasColumn('created_at', executer), + this.hasColumn('updated_at', executer), ]); const normalizedData: Partial = { ...data }; const displayName = @@ -155,11 +166,20 @@ export class DashboardItemRepository ) { normalizedData.id = await this.getNextId(executer); } + if (includeGeneratedId && hasCreatedAtColumn && !normalizedData.createdAt) { + normalizedData.createdAt = new Date(); + } + if (includeGeneratedId && hasUpdatedAtColumn && !normalizedData.updatedAt) { + normalizedData.updatedAt = normalizedData.createdAt || new Date(); + } return normalizedData; } private async hasColumn(column: string, executer: Knex | Knex.Transaction) { + if (this.columnCache.has(column)) { + return this.columnCache.get(column); + } if (column === 'id' && this.hasIdColumnCache !== null) { return this.hasIdColumnCache; } @@ -174,6 +194,7 @@ export class DashboardItemRepository } const result = await executer.schema.hasColumn(this.tableName, column); + this.columnCache.set(column, result); if (column === 'id') { this.hasIdColumnCache = result; } From 16e2fd8497a676bc86ab4244e27180e37c535cd3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 16:30:37 +0530 Subject: [PATCH 0185/1087] Isolate AI context by active project --- .../pipelines/generation/sql_generation.py | 16 ++++ .../retrieval/db_schema_retrieval.py | 18 ---- wren-ai-service/src/web/v1/services/ask.py | 8 +- .../apollo/server/services/askingService.ts | 83 +++++++++++++------ .../src/apollo/server/services/mdlService.ts | 6 ++ .../apollo/server/services/projectService.ts | 32 +++---- .../src/components/OrganizationSwitcher.tsx | 4 +- wren-ui/src/hooks/useHomeSidebar.tsx | 3 +- .../useRecommendedQuestionsInstruction.tsx | 3 + wren-ui/src/pages/home/[id].tsx | 3 +- wren-ui/src/pages/home/index.tsx | 6 +- 11 files changed, 117 insertions(+), 65 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index efc6d93fa9..6f847edcec 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -83,11 +83,13 @@ ### QUESTION ### User's Question: {{ query }} +{% if has_pcb_context %} ### PCB ANALYTICS TERM MAPPING ### If the user asks about PCB repair trends, repair volume, repair counts, debug hours, turnaround time, resolved entries, or failure category, map those business terms to the closest explicit table and column names in DATABASE SCHEMA and VALID TABLE NAMES. Do not answer with general guidance when a SQL aggregation, comparison, trend, or chart is requested. +{% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### @@ -114,10 +116,24 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: + schema_context = "\n".join(documents or []).lower() + has_pcb_context = any( + term in schema_context + for term in ( + "dbo_debugentries", + "debugentryid", + "failure_patterns", + "repair_logs", + "failedat", + "failuresys", + "workorder", + ) + ) _prompt = prompt_builder.run( query=query, data_source=data_source, documents=documents, + has_pcb_context=has_pcb_context, valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ea5c2bb3e2..c7271c04f2 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -125,13 +125,7 @@ def _build_view_ddl(content: dict) -> str: def expand_business_terms_for_retrieval(query: str) -> str: normalized = (query or "").lower() pcb_terms = { - "assembly", - "bar chart", "business unit", - "category", - "common", - "contributor", - "count by", "pcb", "repair", "debug", @@ -140,19 +134,7 @@ def expand_business_terms_for_retrieval(query: str) -> str: "failure code", "failure category", "failure pattern", - "top 10", - "top ten", - "most common", "resolved", - "trend", - "volume", - "count", - "counts", - "average", - "avg", - "chart", - "month", - "monthly", } if not any(term in normalized for term in pcb_terms): return query diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 313775c377..7b4bae8ee7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -930,11 +930,11 @@ async def _run_with_timeout(self, label: str, coroutine): def _build_greeting_response(self, query: str) -> str: return ( - f"Hi. I can help with questions about your PCB database and Wren AI.\n\n" + f"Hi. I can help with questions about your active datasource and Wren AI.\n\n" f"Try a data question like:\n" - f"- Show repair trends for the last 12 months\n" - f"- Compare average debug hours by product family\n" - f"- Which failure codes occur most often?\n\n" + f"- Show monthly trends for the last 12 months\n" + f"- Compare totals by category\n" + f"- Which records occur most often?\n\n" f"If you want, ask a database question directly instead of `{query}`." ) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 9e0860dd5c..0e0652aa6d 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -536,10 +536,7 @@ export class AskingService implements IAskingService { public async getThreadRecommendationQuestions( threadId: number, ): Promise { - const thread = await this.threadRepository.findOneBy({ id: threadId }); - if (!thread) { - throw new Error(`Thread ${threadId} not found`); - } + const thread = await this.ensureThreadInCurrentProject(threadId); // handle not started const res: ThreadRecommendQuestionResult = { @@ -580,10 +577,7 @@ export class AskingService implements IAskingService { private async doGenerateThreadRecommendationQuestions( threadId: number, ): Promise { - const thread = await this.threadRepository.findOneBy({ id: threadId }); - if (!thread) { - throw new Error(`Thread ${threadId} not found`); - } + const thread = await this.ensureThreadInCurrentProject(threadId); if (this.threadRecommendQuestionBackgroundTracker.isExist(thread)) { logger.debug( @@ -668,8 +662,14 @@ export class AskingService implements IAskingService { threadResponseId?: number, ): Promise { const { threadId, language } = payload; - const projectId = - payload.projectId ?? (await this.projectService.getCurrentProject()).id; + const currentProject = await this.projectService.getCurrentProject(); + const projectId = payload.projectId ?? currentProject.id; + if (projectId !== currentProject.id) { + throw new Error(`Project ${projectId} is not the active project`); + } + if (threadId) { + await this.ensureThreadInCurrentProject(threadId); + } const deployId = await this.getDeployId(); // if it's a follow-up question, then the input will have a threadId @@ -704,6 +704,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); // get the original question and ask again const question = threadResponse.question; @@ -800,12 +801,14 @@ export class AskingService implements IAskingService { throw new Error('Update thread input is empty'); } + await this.ensureThreadInCurrentProject(threadId); return this.threadRepository.updateOne(threadId, { summary: input.summary, }); } public async deleteThread(threadId: number): Promise { + await this.ensureThreadInCurrentProject(threadId); await this.threadRepository.deleteOne(threadId); } @@ -813,13 +816,7 @@ export class AskingService implements IAskingService { input: AskingDetailTaskInput, threadId: number, ): Promise { - const thread = await this.threadRepository.findOneBy({ - id: threadId, - }); - - if (!thread) { - throw new Error(`Thread ${threadId} not found`); - } + const thread = await this.ensureThreadInCurrentProject(threadId); const threadResponse = await this.threadResponseRepository.createOne({ threadId: thread.id, @@ -852,6 +849,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${responseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); return await this.threadResponseRepository.updateOne(responseId, { sql: data.sql, @@ -870,6 +868,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); // 1. create a task on AI service to generate the detail const response = await this.wrenAIAdaptor.generateAskDetail({ @@ -906,6 +905,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); if (isAnswerGenerationInProgress(threadResponse.answerDetail?.status)) { logger.debug( @@ -941,6 +941,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); if (isChartGenerationInProgress(threadResponse.chartDetail?.status)) { logger.debug( @@ -985,6 +986,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); if ( isChartGenerationInProgress(threadResponse.chartDetail?.status) && @@ -1024,11 +1026,19 @@ export class AskingService implements IAskingService { } public async getResponsesWithThread(threadId: number) { + await this.ensureThreadInCurrentProject(threadId); return this.threadResponseRepository.getResponsesWithThread(threadId); } public async getResponse(responseId: number) { - return this.threadResponseRepository.findOneBy({ id: responseId }); + const response = await this.threadResponseRepository.findOneBy({ + id: responseId, + }); + if (!response) { + return null; + } + await this.ensureThreadInCurrentProject(response.threadId); + return response; } public async previewData(responseId: number, limit?: number) { @@ -1104,14 +1114,18 @@ export class AskingService implements IAskingService { public async createInstantRecommendedQuestions( input: InstantRecommendedQuestionsInput, ): Promise { - const key = JSON.stringify(input.previousQuestions || []); + const project = await this.projectService.getCurrentProject(); + const key = JSON.stringify({ + projectId: project.id, + previousQuestions: input.previousQuestions || [], + }); const existingJob = this.instantRecommendationJobs.get(key); if (existingJob) { logger.debug('instant recommended questions are already being requested'); return existingJob; } - const job = this.doCreateInstantRecommendedQuestions(input); + const job = this.doCreateInstantRecommendedQuestions(input, project); this.instantRecommendationJobs.set(key, job); try { return await job; @@ -1122,15 +1136,19 @@ export class AskingService implements IAskingService { private async doCreateInstantRecommendedQuestions( input: InstantRecommendedQuestionsInput, + project?: Project, ): Promise { - const project = await this.projectService.getCurrentProject(); - const { manifest } = await this.deployService.getLastDeployment(project.id); + const currentProject = + project ?? (await this.projectService.getCurrentProject()); + const { manifest } = await this.deployService.getLastDeployment( + currentProject.id, + ); const response = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, - projectId: project.id.toString(), + projectId: currentProject.id.toString(), previousQuestions: input.previousQuestions, - ...this.getThreadRecommendationQuestionsConfig(project), + ...this.getThreadRecommendationQuestionsConfig(currentProject), }); return { id: response.queryId }; } @@ -1194,6 +1212,7 @@ export class AskingService implements IAskingService { if (!response) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(response.threadId); return await this.threadResponseRepository.createOne({ sql: input.sql, @@ -1221,6 +1240,7 @@ export class AskingService implements IAskingService { if (!originalThreadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(originalThreadResponse.threadId); const { createdThreadResponse } = await this.adjustmentBackgroundTracker.createAdjustmentTask({ @@ -1252,6 +1272,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); const { queryId } = await this.adjustmentBackgroundTracker.rerunAdjustmentTask({ @@ -1287,6 +1308,7 @@ export class AskingService implements IAskingService { if (!threadId) { return []; } + await this.ensureThreadInCurrentProject(threadId); let responses = await this.threadResponseRepository.getResponsesWithThread( threadId, 10, @@ -1314,4 +1336,17 @@ export class AskingService implements IAskingService { }, }; } + + private async ensureThreadInCurrentProject(threadId: number): Promise { + const [thread, project] = await Promise.all([ + this.threadRepository.findOneBy({ id: threadId }), + this.projectService.getCurrentProject(), + ]); + + if (!thread || thread.projectId !== project.id) { + throw new Error(`Thread ${threadId} not found in current project`); + } + + return thread; + } } diff --git a/wren-ui/src/apollo/server/services/mdlService.ts b/wren-ui/src/apollo/server/services/mdlService.ts index c60199aad0..5d65604a5b 100644 --- a/wren-ui/src/apollo/server/services/mdlService.ts +++ b/wren-ui/src/apollo/server/services/mdlService.ts @@ -3,6 +3,7 @@ import { IModelNestedColumnRepository, IModelColumnRepository, IModelRepository, + Project, IProjectRepository, IRelationRepository, IViewRepository, @@ -15,6 +16,7 @@ export interface MakeCurrentModelMDLResult { } export interface IMDLService { makeCurrentModelMDL(): Promise; + makeModelMDL(project: Project): Promise; } export class MDLService implements IMDLService { @@ -50,6 +52,10 @@ export class MDLService implements IMDLService { public async makeCurrentModelMDL() { const project = await this.projectRepository.getCurrentProject(); + return this.makeModelMDL(project); + } + + public async makeModelMDL(project: Project) { const projectId = project.id; const models = await this.modelRepository.findAllBy({ projectId }); const modelIds = models.map((m) => m.id); diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 65ce3d09cd..4621689dea 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -95,7 +95,7 @@ export class ProjectService implements IProjectService { private mdlService: IMDLService; private wrenAIAdaptor: IWrenAIAdaptor; private projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; - private projectRecommendationJob: Promise | null = null; + private projectRecommendationJobs = new Map>(); constructor({ projectRepository, metadataService, @@ -147,28 +147,32 @@ export class ProjectService implements IProjectService { } public async generateProjectRecommendationQuestions(): Promise { - if (this.projectRecommendationJob) { + const project = await this.getCurrentProject(); + if (!project) { + throw new Error(`Project not found`); + } + + const existingJob = this.projectRecommendationJobs.get(project.id); + if (existingJob) { logger.debug( - 'project recommended questions are already being requested, reusing in-flight job', + `project "${project.id}" recommended questions are already being requested, reusing in-flight job`, ); - return this.projectRecommendationJob; + return existingJob; } - this.projectRecommendationJob = - this.doGenerateProjectRecommendationQuestions(); + const job = this.doGenerateProjectRecommendationQuestions(project); + this.projectRecommendationJobs.set(project.id, job); try { - return await this.projectRecommendationJob; + return await job; } finally { - this.projectRecommendationJob = null; + this.projectRecommendationJobs.delete(project.id); } } - private async doGenerateProjectRecommendationQuestions(): Promise { - const project = await this.getCurrentProject(); - if (!project) { - throw new Error(`Project not found`); - } - const { manifest } = await this.mdlService.makeCurrentModelMDL(); + private async doGenerateProjectRecommendationQuestions( + project: Project, + ): Promise { + const { manifest } = await this.mdlService.makeModelMDL(project); const recommendQuestionResult = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, diff --git a/wren-ui/src/components/OrganizationSwitcher.tsx b/wren-ui/src/components/OrganizationSwitcher.tsx index 5d74580fcf..a4c04d952c 100644 --- a/wren-ui/src/components/OrganizationSwitcher.tsx +++ b/wren-ui/src/components/OrganizationSwitcher.tsx @@ -20,6 +20,7 @@ import SearchOutlined from '@ant-design/icons/SearchOutlined'; import { useRouter } from 'next/router'; import { Path } from '@/utils/enum'; import { WorkspaceProjectType } from '@/apollo/client/graphql/__types__'; +import apolloClient from '@/apollo/client'; interface OrganizationRecord { id: number; @@ -300,8 +301,9 @@ export default function OrganizationSwitcher() { throw new Error(payload.error || 'Failed to switch project'); } message.success('Project switched successfully.'); + await apolloClient.clearStore(); await loadOrganizations(); - await router.push(Path.Home); + await router.replace(Path.Home); } catch (error: any) { message.error(error.message || 'Failed to switch project'); } diff --git a/wren-ui/src/hooks/useHomeSidebar.tsx b/wren-ui/src/hooks/useHomeSidebar.tsx index 8611af4795..c914fa6837 100644 --- a/wren-ui/src/hooks/useHomeSidebar.tsx +++ b/wren-ui/src/hooks/useHomeSidebar.tsx @@ -10,7 +10,8 @@ import { export default function useHomeSidebar() { const router = useRouter(); const { data, refetch } = useThreadsQuery({ - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); const [updateThread] = useUpdateThreadMutation({ onError: (error) => console.error(error), diff --git a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx index e3c2b85c5c..543ebcf62f 100644 --- a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx +++ b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx @@ -45,6 +45,8 @@ export default function useRecommendedQuestionsInstruction() { const [fetchRecommendationQuestions, recommendationQuestionsResult] = useGetProjectRecommendationQuestionsLazyQuery({ + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', pollInterval: 2000, }); @@ -79,6 +81,7 @@ export default function useRecommendedQuestionsInstruction() { }; fetchRecommendationQuestionsData(); + return () => recommendationQuestionsResult.stopPolling(); }, []); useEffect(() => { diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index c3b3cb48ee..7ce5543de0 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -98,7 +98,8 @@ export default function HomeThread() { const { data, updateQuery: updateThreadQuery } = useThreadQuery({ variables: { threadId }, - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', skip: threadId === null, onError: () => router.push(Path.Home), }); diff --git a/wren-ui/src/pages/home/index.tsx b/wren-ui/src/pages/home/index.tsx index afff9b5abd..14e33ccde9 100644 --- a/wren-ui/src/pages/home/index.tsx +++ b/wren-ui/src/pages/home/index.tsx @@ -94,14 +94,16 @@ export default function Home() { const askPrompt = useAskPrompt(); const { data: suggestedQuestionsData } = useSuggestedQuestionsQuery({ - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); const [createThread, { loading: threadCreating }] = useCreateThreadMutation({ onError: (error) => console.error(error), onCompleted: () => homeSidebar.refetch(), }); const [preloadThread] = useThreadLazyQuery({ - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); const { data: settingsResult } = useGetSettingsQuery(); From ec6525c38aec61f954fe2428f3d39f9fa7be2d56 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 17:02:12 +0530 Subject: [PATCH 0186/1087] Use source table fallback for Rust MDL models --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 7 ++- .../apollo/server/mdl/test/mdlBuilder.test.ts | 46 +++++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 05cd3bf220..95774bed98 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -495,13 +495,16 @@ export class MDLBuilder implements IMDLBuilder { model.properties && typeof model.properties === 'string' ? this.parseProperties(model.properties) : {}; - if (!modelProps.table) { + const table = + modelProps.table || + (this.useRustWrenEngine() ? model.sourceTableName : null); + if (!table) { return null; } return { catalog: modelProps.catalog || null, schema: modelProps.schema || null, - table: modelProps.table, + table, }; } private parseLineage(lineage?: string): number[] { diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 4c7976d665..7c8a40c686 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -539,6 +539,52 @@ describe('MDLBuilder', () => { expect(manifest.views).toEqual(expectedViews); }); + it('should use source table name as rust engine table reference fallback.', () => { + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'my project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Search Queries', + sourceTableName: 'dbo.search_queries', + referenceName: 'dbo_search_queries', + refSql: 'SELECT * FROM dbo.search_queries', + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: null, + table: 'dbo.search_queries', + }); + expect(manifest.models[0].refSql).toBeUndefined(); + }); + it('should return correct expression in calculated field.', () => { const models = [ // customer model From 11871b1198c16ebfe4d5b37ef35c415fe8771446 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 17:18:06 +0530 Subject: [PATCH 0187/1087] Use full project schema for broad analysis queries --- .../src/pipelines/retrieval/db_schema_retrieval.py | 14 ++++++++++++++ .../retrieval/test_db_schema_retrieval.py | 11 +++++++++++ 2 files changed, 25 insertions(+) create mode 100644 wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index c7271c04f2..212cd2f0b2 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -162,20 +162,34 @@ def _is_project_wide_analysis_query(query: str) -> bool: "average", "avg", "bar chart", + "breakdown", "chart", + "completed", "compare", "count", "counts", + "distribution", "group by", "grouped", + "highest", "line chart", + "lowest", + "maximum", + "minimum", "monthly", "most common", + "number of", "pie chart", "quarter", + "rank", + "ranking", "recommend", "recommended", "show", + "status", + "sum", + "total", + "totals", "top", "trend", "volume", diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py new file mode 100644 index 0000000000..e0b0f918e1 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -0,0 +1,11 @@ +from src.pipelines.retrieval.db_schema_retrieval import _is_project_wide_analysis_query + + +def test_project_wide_analysis_query_includes_broad_ranking_questions(): + assert _is_project_wide_analysis_query( + "Which projects have the highest number of completed questions?" + ) + + +def test_project_wide_analysis_query_ignores_empty_query(): + assert not _is_project_wide_analysis_query("") From 621a4bd53c41d2a7c73d1853f6bdc93a5d60f641 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 18:00:58 +0530 Subject: [PATCH 0188/1087] Resolve MSSQL source table references in Rust MDL --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 48 ++++++++++++++++-- .../apollo/server/mdl/test/mdlBuilder.test.ts | 50 ++++++++++++++++++- 2 files changed, 91 insertions(+), 7 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 95774bed98..ed07f22e73 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -495,18 +495,56 @@ export class MDLBuilder implements IMDLBuilder { model.properties && typeof model.properties === 'string' ? this.parseProperties(model.properties) : {}; - const table = - modelProps.table || - (this.useRustWrenEngine() ? model.sourceTableName : null); + const fallbackTableReference = this.buildFallbackTableReference(model); + const table = modelProps.table || fallbackTableReference?.table; if (!table) { return null; } return { - catalog: modelProps.catalog || null, - schema: modelProps.schema || null, + catalog: modelProps.catalog || fallbackTableReference?.catalog || null, + schema: modelProps.schema || fallbackTableReference?.schema || null, table, }; } + + private buildFallbackTableReference(model: Model): TableReference | null { + if (!this.useRustWrenEngine() || !model.sourceTableName) { + return null; + } + + if (this.project.type !== DataSourceName.MSSQL) { + return { + catalog: null, + schema: null, + table: model.sourceTableName, + }; + } + + const sourceTableName = model.sourceTableName.trim(); + const dotQualifiedMatch = sourceTableName.match(/^([^.]+)\.([^.]+)$/); + if (dotQualifiedMatch) { + return { + catalog: null, + schema: dotQualifiedMatch[1], + table: dotQualifiedMatch[2], + }; + } + + const underscoreQualifiedMatch = sourceTableName.match(/^(dbo)_(.+)$/i); + if (underscoreQualifiedMatch) { + return { + catalog: null, + schema: underscoreQualifiedMatch[1], + table: underscoreQualifiedMatch[2], + }; + } + + return { + catalog: null, + schema: null, + table: sourceTableName, + }; + } private parseLineage(lineage?: string): number[] { if (!lineage) { return []; diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 7c8a40c686..a382a883e1 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -579,8 +579,54 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: null, - schema: null, - table: 'dbo.search_queries', + schema: 'dbo', + table: 'search_queries', + }); + expect(manifest.models[0].refSql).toBeUndefined(); + }); + + it('should split mssql schema-normalized source table names in rust engine table reference fallback.', () => { + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'my project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Tickets', + sourceTableName: 'dbo_tickets', + referenceName: 'dbo_tickets', + refSql: 'SELECT * FROM dbo.tickets', + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'dbo', + table: 'tickets', }); expect(manifest.models[0].refSql).toBeUndefined(); }); From 73fabb1f0fecabbe7246fd9c33bb62dc473be76a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 18:29:19 +0530 Subject: [PATCH 0189/1087] Normalize MSSQL dbo-prefixed table references --- .../apollo/server/utils/mssqlSqlNormalizer.ts | 25 ++++++++++++++++ .../utils/tests/mssqlSqlNormalizer.test.ts | 30 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 1366592138..8bb66b3b90 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -205,6 +205,30 @@ const normalizeMssqlGeneratedSqlSyntax = (sql: string): string => { return rewriteMssqlLimitClause(sql); }; +const rewriteSchemaNormalizedDboTables = (sql: string): string => { + const tableRef = + String.raw`(?:"(dbo_[A-Za-z0-9_]+)"|\[(dbo_[A-Za-z0-9_]+)\]|\b(dbo_[A-Za-z0-9_]+)\b)`; + const tableClause = new RegExp( + String.raw`\b(FROM|JOIN)\s+${tableRef}(?=\s*(?:WHERE|JOIN|LEFT|RIGHT|INNER|FULL|CROSS|GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|ON|,|\)|$))`, + 'gi', + ); + + return sql.replace( + tableClause, + (_match, clause, quotedName, bracketName, bareName) => { + const modelName = quotedName || bracketName || bareName; + const tableName = modelName.replace(/^dbo_/i, ''); + const schemaQualifiedTable = + quotedName || bracketName + ? `"dbo"."${tableName}"` + : `dbo.${tableName}`; + const alias = + quotedName || bracketName ? ` AS "${modelName}"` : ` AS ${modelName}`; + return `${clause} ${schemaQualifiedTable}${alias}`; + }, + ); +}; + const replaceBadFailurePatternJoins = (sql: string): string => { if ( !/\bdbo_DebugEntries\b/i.test(sql) || @@ -451,6 +475,7 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = replaceInventedTimeBuckets(sql); sql = replaceBadFailurePatternJoins(sql); sql = rewriteMssqlDatepartFunctions(sql); + sql = rewriteSchemaNormalizedDboTables(sql); return normalizeMssqlGeneratedSqlSyntax(sql); }; diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts index ab43d6bf12..9280c2f15a 100644 --- a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts +++ b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts @@ -104,4 +104,34 @@ describe('mssqlSqlNormalizer', () => { expect(normalized).not.toContain('GROUP BY created_by'); expect(normalized).toContain('"created_by_user_id"'); }); + + it('rewrites quoted dbo-prefixed table names to schema-qualified tables with aliases', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT "dbo_search_queries"."org_id", COUNT(*) AS completed_questions + FROM "dbo_search_queries" + WHERE "dbo_search_queries"."result_count" > 0 + GROUP BY "dbo_search_queries"."org_id" + `, + DataSourceName.MSSQL, + ); + + expect(normalized).toContain('FROM "dbo"."search_queries" AS "dbo_search_queries"'); + expect(normalized).toContain('"dbo_search_queries"."org_id"'); + expect(normalized).not.toContain('FROM "dbo_search_queries"'); + }); + + it('rewrites unquoted dbo-prefixed table names to schema-qualified tables with aliases', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT status, COUNT(*) AS count_of_questions + FROM dbo_tickets + GROUP BY status + `, + DataSourceName.MSSQL, + ); + + expect(normalized).toContain('FROM dbo.tickets AS dbo_tickets'); + expect(normalized).not.toContain('FROM dbo_tickets'); + }); }); From 8023d2f2d32a2b682bf538862ded4303e1d52ee2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 19:05:00 +0530 Subject: [PATCH 0190/1087] Quote MSSQL dbo model references for Ibis --- .../apollo/server/utils/mssqlSqlNormalizer.ts | 43 +++++++++---------- .../utils/tests/mssqlSqlNormalizer.test.ts | 28 +++++++++--- 2 files changed, 44 insertions(+), 27 deletions(-) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 8bb66b3b90..56c3ed974b 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -205,28 +205,27 @@ const normalizeMssqlGeneratedSqlSyntax = (sql: string): string => { return rewriteMssqlLimitClause(sql); }; -const rewriteSchemaNormalizedDboTables = (sql: string): string => { - const tableRef = - String.raw`(?:"(dbo_[A-Za-z0-9_]+)"|\[(dbo_[A-Za-z0-9_]+)\]|\b(dbo_[A-Za-z0-9_]+)\b)`; - const tableClause = new RegExp( - String.raw`\b(FROM|JOIN)\s+${tableRef}(?=\s*(?:WHERE|JOIN|LEFT|RIGHT|INNER|FULL|CROSS|GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|ON|,|\)|$))`, - 'gi', - ); +const quoteMssqlDboModelReferences = (sql: string): string => { + sql = sql.replace(/\bdbo\.([A-Za-z0-9_]+)\b/g, '"dbo_$1"'); - return sql.replace( - tableClause, - (_match, clause, quotedName, bracketName, bareName) => { - const modelName = quotedName || bracketName || bareName; - const tableName = modelName.replace(/^dbo_/i, ''); - const schemaQualifiedTable = - quotedName || bracketName - ? `"dbo"."${tableName}"` - : `dbo.${tableName}`; - const alias = - quotedName || bracketName ? ` AS "${modelName}"` : ` AS ${modelName}`; - return `${clause} ${schemaQualifiedTable}${alias}`; - }, - ); + const quotedModels = new Set(); + sql.replace(/"dbo_[A-Za-z0-9_]+"/gi, (match) => { + quotedModels.add(match.slice(1, -1)); + return match; + }); + + const bareDboModel = /\bdbo_[A-Za-z0-9_]+\b/g; + return sql.replace(bareDboModel, (modelName, offset, source) => { + if (source[offset - 1] === '"' || source[offset + modelName.length] === '"') { + return modelName; + } + + if (!quotedModels.has(modelName)) { + quotedModels.add(modelName); + } + + return `"${modelName}"`; + }); }; const replaceBadFailurePatternJoins = (sql: string): string => { @@ -475,7 +474,7 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = replaceInventedTimeBuckets(sql); sql = replaceBadFailurePatternJoins(sql); sql = rewriteMssqlDatepartFunctions(sql); - sql = rewriteSchemaNormalizedDboTables(sql); + sql = quoteMssqlDboModelReferences(sql); return normalizeMssqlGeneratedSqlSyntax(sql); }; diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts index 9280c2f15a..c958e47cb5 100644 --- a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts +++ b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts @@ -105,7 +105,7 @@ describe('mssqlSqlNormalizer', () => { expect(normalized).toContain('"created_by_user_id"'); }); - it('rewrites quoted dbo-prefixed table names to schema-qualified tables with aliases', () => { + it('keeps quoted dbo-prefixed model names for ibis model resolution', () => { const normalized = normalizeMssqlSqlForIbis( ` SELECT "dbo_search_queries"."org_id", COUNT(*) AS completed_questions @@ -116,12 +116,12 @@ describe('mssqlSqlNormalizer', () => { DataSourceName.MSSQL, ); - expect(normalized).toContain('FROM "dbo"."search_queries" AS "dbo_search_queries"'); + expect(normalized).toContain('FROM "dbo_search_queries"'); expect(normalized).toContain('"dbo_search_queries"."org_id"'); - expect(normalized).not.toContain('FROM "dbo_search_queries"'); + expect(normalized).not.toContain('FROM dbo_search_queries'); }); - it('rewrites unquoted dbo-prefixed table names to schema-qualified tables with aliases', () => { + it('quotes unquoted dbo-prefixed model names for ibis model resolution', () => { const normalized = normalizeMssqlSqlForIbis( ` SELECT status, COUNT(*) AS count_of_questions @@ -131,7 +131,25 @@ describe('mssqlSqlNormalizer', () => { DataSourceName.MSSQL, ); - expect(normalized).toContain('FROM dbo.tickets AS dbo_tickets'); + expect(normalized).toContain('FROM "dbo_tickets"'); expect(normalized).not.toContain('FROM dbo_tickets'); }); + + it('normalizes schema-qualified dbo references back to model names', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT dbo.search_queries.org_id, COUNT(*) AS completed_questions + FROM dbo.search_queries + INNER JOIN dbo.organizations ON dbo.search_queries.org_id = dbo.organizations.id + GROUP BY dbo.organizations.name + `, + DataSourceName.MSSQL, + ); + + expect(normalized).toContain('FROM "dbo_search_queries"'); + expect(normalized).toContain('INNER JOIN "dbo_organizations"'); + expect(normalized).toContain('"dbo_search_queries".org_id'); + expect(normalized).not.toContain('dbo.search_queries'); + expect(normalized).not.toContain('dbo.organizations'); + }); }); From 7abdbcfcb938bec0913baeb64cf0bc996430ab56 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 19:34:10 +0530 Subject: [PATCH 0191/1087] Quote dbo model references across project types --- .../src/apollo/server/utils/mssqlSqlNormalizer.ts | 4 +++- .../server/utils/tests/mssqlSqlNormalizer.test.ts | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 56c3ed974b..6a0b7ec195 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -456,11 +456,13 @@ export const normalizeMssqlGeneratedSqlFields = ( sql: string, dataSource: DataSourceName, ): string => { + sql = sql.replace(/\\"/g, '"'); + sql = quoteMssqlDboModelReferences(sql); + if (dataSource !== DataSourceName.MSSQL) { return sql; } - sql = sql.replace(/\\"/g, '"'); sql = normalizeMssqlGeneratedSqlSyntax(sql); sql = rewriteMssqlDatepartFunctions(sql); sql = replaceRelativeCurrentDateCalls(sql); diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts index c958e47cb5..61238daa54 100644 --- a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts +++ b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts @@ -135,6 +135,20 @@ describe('mssqlSqlNormalizer', () => { expect(normalized).not.toContain('FROM dbo_tickets'); }); + it('quotes dbo-prefixed model names for non-MSSQL project contexts', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT org_id, COUNT(*) AS num_questions + FROM dbo_search_queries + GROUP BY org_id + `, + DataSourceName.POSTGRES, + ); + + expect(normalized).toContain('FROM "dbo_search_queries"'); + expect(normalized).not.toContain('FROM dbo_search_queries'); + }); + it('normalizes schema-qualified dbo references back to model names', () => { const normalized = normalizeMssqlSqlForIbis( ` From bd5e965fdc83437cd8cc58c48a9cccab45f90ea0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 20:00:00 +0530 Subject: [PATCH 0192/1087] Collapse qualified dbo model references --- .../src/apollo/server/utils/mssqlSqlNormalizer.ts | 12 ++++++++++++ .../server/utils/tests/mssqlSqlNormalizer.test.ts | 15 +++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 6a0b7ec195..655f0f293f 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -206,6 +206,18 @@ const normalizeMssqlGeneratedSqlSyntax = (sql: string): string => { }; const quoteMssqlDboModelReferences = (sql: string): string => { + sql = sql.replace( + /(?:"[^"]+"\.){1,2}"(dbo_[A-Za-z0-9_]+)"/gi, + '"$1"', + ); + sql = sql.replace( + /\b[A-Za-z_][A-Za-z0-9_]*\.[A-Za-z_][A-Za-z0-9_]*\.(dbo_[A-Za-z0-9_]+)\b/g, + '"$1"', + ); + sql = sql.replace( + /\[[^\]]+\]\.\[[^\]]+\]\.\[(dbo_[A-Za-z0-9_]+)\]/gi, + '"$1"', + ); sql = sql.replace(/\bdbo\.([A-Za-z0-9_]+)\b/g, '"dbo_$1"'); const quotedModels = new Set(); diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts index 61238daa54..76d9100825 100644 --- a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts +++ b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts @@ -149,6 +149,21 @@ describe('mssqlSqlNormalizer', () => { expect(normalized).not.toContain('FROM dbo_search_queries'); }); + it('collapses fully qualified dbo-prefixed model names for switched projects', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT wrenai.public.dbo_search_queries.org_id, COUNT(*) AS num_questions + FROM wrenai.public.dbo_search_queries + GROUP BY wrenai.public.dbo_search_queries.org_id + `, + DataSourceName.POSTGRES, + ); + + expect(normalized).toContain('FROM "dbo_search_queries"'); + expect(normalized).toContain('"dbo_search_queries".org_id'); + expect(normalized).not.toContain('wrenai.public.dbo_search_queries'); + }); + it('normalizes schema-qualified dbo references back to model names', () => { const normalized = normalizeMssqlSqlForIbis( ` From 13c8e99288022df1b71f1b224a19501a2d09fe50 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 21:28:30 +0530 Subject: [PATCH 0193/1087] Resolve dbo-prefixed table references for all sources --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 16 +++---- .../apollo/server/mdl/test/mdlBuilder.test.ts | 46 +++++++++++++++++++ 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index ed07f22e73..a5dbdd6018 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -512,14 +512,6 @@ export class MDLBuilder implements IMDLBuilder { return null; } - if (this.project.type !== DataSourceName.MSSQL) { - return { - catalog: null, - schema: null, - table: model.sourceTableName, - }; - } - const sourceTableName = model.sourceTableName.trim(); const dotQualifiedMatch = sourceTableName.match(/^([^.]+)\.([^.]+)$/); if (dotQualifiedMatch) { @@ -539,6 +531,14 @@ export class MDLBuilder implements IMDLBuilder { }; } + if (this.project.type !== DataSourceName.MSSQL) { + return { + catalog: null, + schema: null, + table: sourceTableName, + }; + } + return { catalog: null, schema: null, diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index a382a883e1..6101b7b778 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -631,6 +631,52 @@ describe('MDLBuilder', () => { expect(manifest.models[0].refSql).toBeUndefined(); }); + it('should split dbo-prefixed source table names for non-mssql projects.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Search Queries', + sourceTableName: 'dbo_search_queries', + referenceName: 'dbo_search_queries', + refSql: 'SELECT * FROM dbo.search_queries', + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'dbo', + table: 'search_queries', + }); + expect(manifest.models[0].refSql).toBeUndefined(); + }); + it('should return correct expression in calculated field.', () => { const models = [ // customer model From f50cf187367de60cbc7256f1bbe4de4a7e712301 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 21:53:35 +0530 Subject: [PATCH 0194/1087] Continue SQL flow when intent classification times out --- wren-ai-service/src/web/v1/services/ask.py | 38 +++++++++++++++------- 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7b4bae8ee7..a0f0805ee2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1233,19 +1233,33 @@ async def ask( ) if self._allow_intent_classification: - intent_classification_result = ( - await self._run_with_timeout( - "Intent classification", - self._pipelines["intent_classification"].run( - query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, - ), + try: + intent_classification_result = ( + await self._run_with_timeout( + "Intent classification", + self._pipelines["intent_classification"].run( + query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, + ), + ) + ).get("post_process", {}) + except TimeoutError as exc: + logger.warning( + "Intent classification timed out; continuing with TEXT_TO_SQL. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, ) - ).get("post_process", {}) + intent_classification_result = { + "intent": "TEXT_TO_SQL", + "rephrased_question": user_query, + "reasoning": "Intent classification timed out; using SQL generation.", + "db_schemas": [], + } intent = intent_classification_result.get("intent") rephrased_question = intent_classification_result.get( "rephrased_question" From 97ebb904d892104b2ded9f340f3184b50b5a4a11 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 22:47:24 +0530 Subject: [PATCH 0195/1087] Normalize explicit dbo table references in manifest --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 63 ++++++++++++++----- .../apollo/server/mdl/test/mdlBuilder.test.ts | 49 +++++++++++++++ 2 files changed, 96 insertions(+), 16 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index a5dbdd6018..da961b5377 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -495,14 +495,29 @@ export class MDLBuilder implements IMDLBuilder { model.properties && typeof model.properties === 'string' ? this.parseProperties(model.properties) : {}; + const propertyTableReference = + typeof modelProps.table === 'string' + ? this.buildTableReferenceFromTableName(modelProps.table) + : null; const fallbackTableReference = this.buildFallbackTableReference(model); - const table = modelProps.table || fallbackTableReference?.table; + const table = + propertyTableReference?.table || + modelProps.table || + fallbackTableReference?.table; if (!table) { return null; } return { - catalog: modelProps.catalog || fallbackTableReference?.catalog || null, - schema: modelProps.schema || fallbackTableReference?.schema || null, + catalog: + propertyTableReference?.catalog || + modelProps.catalog || + fallbackTableReference?.catalog || + null, + schema: + propertyTableReference?.schema || + modelProps.schema || + fallbackTableReference?.schema || + null, table, }; } @@ -513,6 +528,34 @@ export class MDLBuilder implements IMDLBuilder { } const sourceTableName = model.sourceTableName.trim(); + const normalizedTableReference = + this.buildTableReferenceFromTableName(sourceTableName); + if (normalizedTableReference) { + return normalizedTableReference; + } + + return { + catalog: null, + schema: null, + table: sourceTableName, + }; + } + + private buildTableReferenceFromTableName( + tableName: string, + ): TableReference | null { + const sourceTableName = tableName.trim(); + const catalogQualifiedMatch = sourceTableName.match( + /^([^.]+)\.([^.]+)\.([^.]+)$/, + ); + if (catalogQualifiedMatch) { + return { + catalog: catalogQualifiedMatch[1], + schema: catalogQualifiedMatch[2], + table: catalogQualifiedMatch[3], + }; + } + const dotQualifiedMatch = sourceTableName.match(/^([^.]+)\.([^.]+)$/); if (dotQualifiedMatch) { return { @@ -531,19 +574,7 @@ export class MDLBuilder implements IMDLBuilder { }; } - if (this.project.type !== DataSourceName.MSSQL) { - return { - catalog: null, - schema: null, - table: sourceTableName, - }; - } - - return { - catalog: null, - schema: null, - table: sourceTableName, - }; + return null; } private parseLineage(lineage?: string): number[] { if (!lineage) { diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 6101b7b778..c1c8af88c3 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -677,6 +677,55 @@ describe('MDLBuilder', () => { expect(manifest.models[0].refSql).toBeUndefined(); }); + it('should split dbo-prefixed property table names before project schema fallback.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Search Queries', + sourceTableName: 'dbo_search_queries', + referenceName: 'dbo_search_queries', + refSql: 'SELECT * FROM dbo.search_queries', + cached: false, + refreshTime: null, + properties: JSON.stringify({ + schema: 'public', + table: 'dbo_search_queries', + }), + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'dbo', + table: 'search_queries', + }); + expect(manifest.models[0].refSql).toBeUndefined(); + }); + it('should return correct expression in calculated field.', () => { const models = [ // customer model From dd177f8af72c09a05a1774c409c3edc98f53b2d8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 23:45:02 +0530 Subject: [PATCH 0196/1087] Strip dbo model prefixes for non-MSSQL manifests --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 30 +++++++++- .../apollo/server/mdl/test/mdlBuilder.test.ts | 56 +++++++++++++++++-- 2 files changed, 79 insertions(+), 7 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index da961b5377..d3980b3dee 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -549,10 +549,13 @@ export class MDLBuilder implements IMDLBuilder { /^([^.]+)\.([^.]+)\.([^.]+)$/, ); if (catalogQualifiedMatch) { + const normalizedTableName = this.normalizeDboPrefixedTableName( + catalogQualifiedMatch[3], + ); return { catalog: catalogQualifiedMatch[1], - schema: catalogQualifiedMatch[2], - table: catalogQualifiedMatch[3], + schema: normalizedTableName.schema || catalogQualifiedMatch[2], + table: normalizedTableName.table, }; } @@ -569,13 +572,34 @@ export class MDLBuilder implements IMDLBuilder { if (underscoreQualifiedMatch) { return { catalog: null, - schema: underscoreQualifiedMatch[1], + schema: + this.project.type === DataSourceName.MSSQL + ? underscoreQualifiedMatch[1] + : null, table: underscoreQualifiedMatch[2], }; } return null; } + + private normalizeDboPrefixedTableName(tableName: string): { + schema: string | null; + table: string; + } { + const underscoreQualifiedMatch = tableName.match(/^(dbo)_(.+)$/i); + if (!underscoreQualifiedMatch) { + return { schema: null, table: tableName }; + } + + return { + schema: + this.project.type === DataSourceName.MSSQL + ? underscoreQualifiedMatch[1] + : null, + table: underscoreQualifiedMatch[2], + }; + } private parseLineage(lineage?: string): number[] { if (!lineage) { return []; diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index c1c8af88c3..97c1c45b0f 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -631,7 +631,7 @@ describe('MDLBuilder', () => { expect(manifest.models[0].refSql).toBeUndefined(); }); - it('should split dbo-prefixed source table names for non-mssql projects.', () => { + it('should strip dbo-prefixed source table names for non-mssql projects.', () => { const project = { id: 1, type: DataSourceName.POSTGRES, @@ -671,13 +671,13 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: null, - schema: 'dbo', + schema: null, table: 'search_queries', }); expect(manifest.models[0].refSql).toBeUndefined(); }); - it('should split dbo-prefixed property table names before project schema fallback.', () => { + it('should strip dbo-prefixed property table names before project schema fallback.', () => { const project = { id: 1, type: DataSourceName.POSTGRES, @@ -720,7 +720,55 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: null, - schema: 'dbo', + schema: 'public', + table: 'search_queries', + }); + expect(manifest.models[0].refSql).toBeUndefined(); + }); + + it('should strip dbo-prefixed table names from catalog-qualified non-mssql table references.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Search Queries', + sourceTableName: 'wrenai.public.dbo_search_queries', + referenceName: 'dbo_search_queries', + refSql: 'SELECT * FROM wrenai.public.dbo_search_queries', + cached: false, + refreshTime: null, + properties: JSON.stringify({ + table: 'wrenai.public.dbo_search_queries', + }), + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: 'wrenai', + schema: 'public', table: 'search_queries', }); expect(manifest.models[0].refSql).toBeUndefined(); From 39d91b29e8c43410bf43acdf264f392aa9e2ff64 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 15 Jun 2026 23:51:25 +0530 Subject: [PATCH 0197/1087] Normalize deployed manifests during query execution --- .../apollo/server/services/queryService.ts | 71 +++++++++++- .../services/tests/queryService.test.ts | 109 ++++++++++++++++++ 2 files changed, 177 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 269da4cbba..54b7e13746 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -1,5 +1,5 @@ import { DataSourceName } from '@server/types'; -import { Manifest } from '@server/mdl/type'; +import { Manifest, TableReference } from '@server/mdl/type'; import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; import { SupportedDataSource, @@ -107,6 +107,69 @@ const normalizePreviewSqlForIbis = ( }; }; +const normalizeDeployedManifestForDatasource = ( + manifest: Manifest, + project: Project, +): Manifest => { + if (project.type === DataSourceName.MSSQL || !manifest?.models?.length) { + return manifest; + } + + const fallbackSchema = manifest.schema || project.schema || null; + + return { + ...manifest, + models: manifest.models.map((model) => { + const tableReference = normalizeTableReference( + model.tableReference, + fallbackSchema, + ); + + if (!tableReference) { + return model; + } + + return { + ...model, + tableReference, + }; + }), + }; +}; + +const normalizeTableReference = ( + tableReference: TableReference | undefined, + fallbackSchema: string | null, +): TableReference | undefined => { + if (!tableReference?.table) { + return tableReference; + } + + const normalizedTableName = normalizeDboPrefixedTableName( + tableReference.table, + ); + const shouldReplaceDboSchema = + tableReference.schema?.toLowerCase() === 'dbo' && fallbackSchema; + + if ( + normalizedTableName === tableReference.table && + !shouldReplaceDboSchema + ) { + return tableReference; + } + + return { + ...tableReference, + schema: shouldReplaceDboSchema ? fallbackSchema : tableReference.schema, + table: normalizedTableName, + }; +}; + +const normalizeDboPrefixedTableName = (tableName: string): string => { + const match = tableName.match(/^dbo_(.+)$/i); + return match ? match[1] : tableName; +}; + export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; @@ -132,12 +195,13 @@ export class QueryService implements IQueryService { ): Promise { const { project, - manifest: mdl, + manifest: rawMdl, limit, dryRun, refresh, cacheEnabled, } = options; + const mdl = normalizeDeployedManifestForDatasource(rawMdl, project); const { type: dataSource, connectionInfo } = project; if (this.useEngine(dataSource)) { if (dryRun) { @@ -193,11 +257,12 @@ export class QueryService implements IQueryService { parameters: Record, ): Promise { const { type: dataSource, connectionInfo } = project; + const mdl = normalizeDeployedManifestForDatasource(manifest, project); const res = await this.ibisAdaptor.validate( dataSource, rule, connectionInfo, - manifest, + mdl, parameters, ); return res; diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 011b56176b..6a3c308096 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -54,6 +54,115 @@ describe('QueryService', () => { }); }); + it('should normalize deployed dbo-prefixed table references for non-mssql previews', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview('SELECT * FROM "dbo_search_queries"', { + project: { + type: DataSourceName.POSTGRES, + connectionInfo: {}, + schema: 'public', + }, + manifest: { + schema: 'public', + models: [ + { + name: 'dbo_search_queries', + tableReference: { + catalog: 'wrenai', + schema: 'public', + table: 'dbo_search_queries', + }, + }, + { + name: 'dbo_tickets', + tableReference: { + catalog: 'wrenai', + schema: 'dbo', + table: 'tickets', + }, + }, + ], + }, + dryRun: true, + }); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT * FROM "dbo_search_queries"', + expect.objectContaining({ + mdl: expect.objectContaining({ + models: [ + expect.objectContaining({ + name: 'dbo_search_queries', + tableReference: { + catalog: 'wrenai', + schema: 'public', + table: 'search_queries', + }, + }), + expect.objectContaining({ + name: 'dbo_tickets', + tableReference: { + catalog: 'wrenai', + schema: 'public', + table: 'tickets', + }, + }), + ], + }), + }), + ); + }); + + it('should preserve deployed dbo-prefixed table references for mssql previews', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview('SELECT * FROM "dbo_search_queries"', { + project: { + type: DataSourceName.MSSQL, + connectionInfo: {}, + schema: 'public', + }, + manifest: { + schema: 'public', + models: [ + { + name: 'dbo_search_queries', + tableReference: { + catalog: null, + schema: 'dbo', + table: 'search_queries', + }, + }, + ], + }, + dryRun: true, + }); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT * FROM "dbo_search_queries"', + expect.objectContaining({ + mdl: expect.objectContaining({ + models: [ + expect.objectContaining({ + tableReference: { + catalog: null, + schema: 'dbo', + table: 'search_queries', + }, + }), + ], + }), + }), + ); + }); + it('should send event when previewing via ibis dry run fails', async () => { mockIbisAdaptor.dryRun.mockRejectedValue({ message: 'Error message', From dc44a737c7d9cfcbd819850ff573be6b7d241cab Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 00:05:43 +0530 Subject: [PATCH 0198/1087] Repair stale dbo refSql during query execution --- .../apollo/server/services/queryService.ts | 100 ++++++++++++++++-- .../services/tests/queryService.test.ts | 37 +++++++ 2 files changed, 127 insertions(+), 10 deletions(-) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 54b7e13746..128a256d26 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -115,24 +115,47 @@ const normalizeDeployedManifestForDatasource = ( return manifest; } + const fallbackCatalog = manifest.catalog || project.catalog || null; const fallbackSchema = manifest.schema || project.schema || null; return { ...manifest, models: manifest.models.map((model) => { - const tableReference = normalizeTableReference( + const tableReferenceResult = normalizeTableReference( model.tableReference, fallbackSchema, ); + const synthesizedTableReference = + tableReferenceResult.tableReference || + buildTableReferenceFromDboModelName( + model.name, + fallbackCatalog, + fallbackSchema, + ) || + buildTableReferenceFromDboRefSql( + model.refSql, + fallbackCatalog, + fallbackSchema, + ); - if (!tableReference) { + if (!synthesizedTableReference) { return model; } - return { + const normalizedModel = { ...model, - tableReference, + tableReference: synthesizedTableReference, }; + + if ( + tableReferenceResult.changed || + isDboPrefixedModelName(model.name) || + containsDboPhysicalReference(model.refSql) + ) { + delete normalizedModel.refSql; + } + + return normalizedModel; }), }; }; @@ -140,9 +163,9 @@ const normalizeDeployedManifestForDatasource = ( const normalizeTableReference = ( tableReference: TableReference | undefined, fallbackSchema: string | null, -): TableReference | undefined => { +): { tableReference?: TableReference; changed: boolean } => { if (!tableReference?.table) { - return tableReference; + return { tableReference, changed: false }; } const normalizedTableName = normalizeDboPrefixedTableName( @@ -155,13 +178,16 @@ const normalizeTableReference = ( normalizedTableName === tableReference.table && !shouldReplaceDboSchema ) { - return tableReference; + return { tableReference, changed: false }; } return { - ...tableReference, - schema: shouldReplaceDboSchema ? fallbackSchema : tableReference.schema, - table: normalizedTableName, + tableReference: { + ...tableReference, + schema: shouldReplaceDboSchema ? fallbackSchema : tableReference.schema, + table: normalizedTableName, + }, + changed: true, }; }; @@ -170,6 +196,60 @@ const normalizeDboPrefixedTableName = (tableName: string): string => { return match ? match[1] : tableName; }; +const buildTableReferenceFromDboModelName = ( + modelName: string | undefined, + fallbackCatalog: string | null, + fallbackSchema: string | null, +): TableReference | undefined => { + if (!modelName || !isDboPrefixedModelName(modelName)) { + return undefined; + } + + return { + catalog: fallbackCatalog, + schema: fallbackSchema, + table: normalizeDboPrefixedTableName(modelName), + }; +}; + +const buildTableReferenceFromDboRefSql = ( + refSql: string | undefined, + fallbackCatalog: string | null, + fallbackSchema: string | null, +): TableReference | undefined => { + if (!refSql || !containsDboPhysicalReference(refSql)) { + return undefined; + } + + const tableName = + extractDboPrefixedTableName(refSql) || extractDboSchemaTableName(refSql); + if (!tableName) { + return undefined; + } + + return { + catalog: fallbackCatalog, + schema: fallbackSchema, + table: normalizeDboPrefixedTableName(tableName), + }; +}; + +const isDboPrefixedModelName = (modelName: string | undefined): boolean => + !!modelName && /^dbo_.+/i.test(modelName); + +const containsDboPhysicalReference = (sql: string | undefined): boolean => + !!sql && /(?:^|[.\s"])(?:dbo_[\w]+|dbo\.[\w"]+)/i.test(sql); + +const extractDboPrefixedTableName = (sql: string): string | undefined => { + const match = sql.match(/\bdbo_([A-Za-z0-9_]+)\b/i); + return match ? `dbo_${match[1]}` : undefined; +}; + +const extractDboSchemaTableName = (sql: string): string | undefined => { + const match = sql.match(/\bdbo\.("?)([A-Za-z0-9_]+)\1/i); + return match ? match[2] : undefined; +}; + export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 6a3c308096..c5b024d538 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -163,6 +163,43 @@ describe('QueryService', () => { ); }); + it('should repair old non-mssql deployments that still use dbo refSql', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview('SELECT * FROM "dbo_search_queries"', { + project: { + type: DataSourceName.POSTGRES, + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + }, + manifest: { + catalog: 'wrenai', + schema: 'public', + models: [ + { + name: 'dbo_search_queries', + refSql: 'SELECT * FROM wrenai.public.dbo_search_queries', + }, + ], + }, + dryRun: true, + }); + + const dryRunOptions = mockIbisAdaptor.dryRun.mock.calls[0][1]; + expect(dryRunOptions.mdl.models[0]).toEqual({ + name: 'dbo_search_queries', + tableReference: { + catalog: 'wrenai', + schema: 'public', + table: 'search_queries', + }, + }); + }); + it('should send event when previewing via ibis dry run fails', async () => { mockIbisAdaptor.dryRun.mockRejectedValue({ message: 'Error message', From 3d08d68d0e7158179304c2ec5b05a1ab8590c880 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 00:29:32 +0530 Subject: [PATCH 0199/1087] Normalize non-MSSQL generated SQL before Ibis preview --- .../apollo/server/services/queryService.ts | 48 ++++++++++++++++-- .../services/tests/queryService.test.ts | 50 +++++++++++++++++++ 2 files changed, 95 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 128a256d26..e1b8ebc123 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -82,12 +82,15 @@ const normalizePreviewSqlForIbis = ( dataSource: DataSourceName, limit?: number, ): { sql: string; limit?: number } => { + sql = normalizeMssqlSqlForIbis(sql, dataSource); + if (dataSource !== DataSourceName.MSSQL) { - return { sql, limit }; + return { + sql: normalizeNonMssqlGeneratedSqlSyntax(sql, dataSource), + limit, + }; } - sql = normalizeMssqlSqlForIbis(sql, dataSource); - const topMatch = sql.match(/^\s*SELECT\s+(DISTINCT\s+)?TOP\s*\(?\s*(\d+)\s*\)?\s+/i); if (!topMatch) { return { sql, limit }; @@ -160,6 +163,45 @@ const normalizeDeployedManifestForDatasource = ( }; }; +const normalizeNonMssqlGeneratedSqlSyntax = ( + sql: string, + dataSource: DataSourceName, +): string => { + if (dataSource === DataSourceName.MSSQL) { + return sql; + } + + return rewriteGeneratedDateDiff(sql); +}; + +const rewriteGeneratedDateDiff = (sql: string): string => { + const dateDiffPattern = + /\bdate_?diff\s*\(\s*'?([A-Za-z]+)'?\s*,\s*([^,()]+(?:\([^)]*\))?[^,()]*)\s*,\s*([^()]+(?:\([^)]*\))?[^()]*)\)/gi; + + return sql.replace( + dateDiffPattern, + (_match, unit: string, startExpression: string, endExpression: string) => { + const normalizedUnit = unit.toLowerCase(); + const start = startExpression.trim(); + const end = endExpression.trim(); + + if (['day', 'dd', 'd'].includes(normalizedUnit)) { + return `EXTRACT(DAY FROM (${end} - ${start}))`; + } + + if (['month', 'mm', 'm'].includes(normalizedUnit)) { + return `((EXTRACT(YEAR FROM ${end}) - EXTRACT(YEAR FROM ${start})) * 12 + (EXTRACT(MONTH FROM ${end}) - EXTRACT(MONTH FROM ${start})))`; + } + + if (['year', 'yy', 'yyyy'].includes(normalizedUnit)) { + return `(EXTRACT(YEAR FROM ${end}) - EXTRACT(YEAR FROM ${start}))`; + } + + return `EXTRACT(DAY FROM (${end} - ${start}))`; + }, + ); +}; + const normalizeTableReference = ( tableReference: TableReference | undefined, fallbackSchema: string | null, diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index c5b024d538..81441d4794 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -200,6 +200,56 @@ describe('QueryService', () => { }); }); + it('should normalize non-mssql dbo model names before previewing with ibis', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT created_at FROM dbo_search_queries ORDER BY created_at', + { + project: { + type: DataSourceName.POSTGRES, + connectionInfo: {}, + schema: 'public', + }, + manifest: {}, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT created_at FROM "dbo_search_queries" ORDER BY created_at', + expect.any(Object), + ); + }); + + it('should normalize generated datediff calls for non-mssql previews', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + `SELECT DATEDIFF('day', "dbo_tickets"."created_at", CURRENT_DATE) AS ticket_age FROM "dbo_tickets"`, + { + project: { + type: DataSourceName.POSTGRES, + connectionInfo: {}, + schema: 'public', + }, + manifest: {}, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT EXTRACT(DAY FROM (CURRENT_DATE - "dbo_tickets"."created_at")) AS ticket_age FROM "dbo_tickets"', + expect.any(Object), + ); + }); + it('should send event when previewing via ibis dry run fails', async () => { mockIbisAdaptor.dryRun.mockRejectedValue({ message: 'Error message', From cc5d288b44f2920532c0a373c9bb2ccaf56ee16f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 00:59:14 +0530 Subject: [PATCH 0200/1087] Keep retrieval scoped to active project --- wren-ai-service/src/pipelines/common.py | 4 - .../historical_question_retrieval.py | 7 - .../src/pipelines/retrieval/instructions.py | 24 --- .../retrieval/sql_pairs_retrieval.py | 7 - .../retrieval/test_project_scope_isolation.py | 155 ++++++++++++++++++ 5 files changed, 155 insertions(+), 42 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 825e264e98..940fa66eb7 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -88,10 +88,6 @@ async def retrieve_metadata(project_id: str, retriever) -> dict[str, Any]: result = await retriever.run(query_embedding=[], filters=filters) documents = result["documents"] - if not documents and project_id: - result = await retriever.run(query_embedding=[], filters=None) - documents = result["documents"] - # only one document for a project, thus we can return the first one if documents: doc = documents[0] diff --git a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py index 77911ef659..68f1ef158c 100644 --- a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py @@ -52,8 +52,6 @@ async def count_documents( ) count = await view_questions_store.count_documents(filters=filters) - if count == 0 and project_id: - count = await view_questions_store.count_documents(filters=None) return count @@ -87,11 +85,6 @@ async def retrieval( query_embedding=embedding.get("embedding"), filters=filters, ) - if not view_question_res.get("documents") and project_id: - view_question_res = await view_questions_retriever.run( - query_embedding=embedding.get("embedding"), - filters=None, - ) return dict(documents=view_question_res.get("documents")) return {} diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 22e688159e..86c17e93de 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -70,8 +70,6 @@ async def count_documents( else None ) document_count = await store.count_documents(filters=filters) - if document_count == 0 and project_id: - document_count = await store.count_documents(filters=None) return document_count @@ -104,17 +102,6 @@ async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: query_embedding=embedding.get("embedding"), filters=filters, ) - if not res.get("documents") and project_id: - fallback_filters = { - "operator": "AND", - "conditions": [ - {"field": "is_default", "operator": "==", "value": False}, - ], - } - res = await retriever.run( - query_embedding=embedding.get("embedding"), - filters=fallback_filters, - ) return dict(documents=res.get("documents")) @@ -169,17 +156,6 @@ async def default_instructions( query_embedding=None, filters=filters, ) - if not _res.get("documents") and project_id: - fallback_filters = { - "operator": "AND", - "conditions": [ - {"field": "is_default", "operator": "==", "value": True}, - ], - } - _res = await retriever.run( - query_embedding=None, - filters=fallback_filters, - ) res = scope_filter.run( documents=_res.get("documents"), diff --git a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py index c8d2997c17..3fe44f32eb 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py @@ -49,8 +49,6 @@ async def count_documents( else None ) document_count = await store.count_documents(filters=filters) - if document_count == 0 and project_id: - document_count = await store.count_documents(filters=None) return document_count @@ -80,11 +78,6 @@ async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: query_embedding=embedding.get("embedding"), filters=filters, ) - if not res.get("documents") and project_id: - res = await retriever.run( - query_embedding=embedding.get("embedding"), - filters=None, - ) return dict(documents=res.get("documents")) return {} diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py new file mode 100644 index 0000000000..d058274eff --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py @@ -0,0 +1,155 @@ +import pytest + +from src.pipelines.common import retrieve_metadata +from src.pipelines.retrieval import historical_question_retrieval, instructions +from src.pipelines.retrieval import sql_pairs_retrieval + + +PROJECT_FILTER = { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-a"}, + ], +} + + +class StoreSpy: + def __init__(self, count=0): + self.count = count + self.filters = [] + + async def count_documents(self, filters=None): + self.filters.append(filters) + return self.count + + +class RetrieverSpy: + def __init__(self, documents=None): + self.documents = documents or [] + self.calls = [] + + async def run(self, query_embedding=None, filters=None): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + return {"documents": self.documents} + + +@pytest.mark.asyncio +async def test_metadata_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await retrieve_metadata("project-a", retriever) + + assert result == {} + assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_sql_pairs_count_stays_project_scoped_when_project_has_no_documents(): + store = StoreSpy(count=0) + + count = await sql_pairs_retrieval.count_documents(store, project_id="project-a") + + assert count == 0 + assert store.filters == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_sql_pairs_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await sql_pairs_retrieval.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_historical_question_count_stays_project_scoped_when_project_has_no_documents(): + store = StoreSpy(count=0) + + count = await historical_question_retrieval.count_documents( + store, + project_id="project-a", + ) + + assert count == 0 + assert store.filters == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_historical_question_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await historical_question_retrieval.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + view_questions_retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_instruction_count_stays_project_scoped_when_project_has_no_documents(): + store = StoreSpy(count=0) + + count = await instructions.count_documents(store, project_id="project-a") + + assert count == 0 + assert store.filters == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_instruction_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await instructions.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": False}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + ], + } + ] + + +@pytest.mark.asyncio +async def test_default_instructions_do_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await instructions.default_instructions( + count_documents=1, + retriever=retriever, + project_id="project-a", + scope_filter=instructions.ScopeFilter(), + scope="sql", + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": True}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + ], + } + ] From fc4780cc59e57d4d675eb0fac0d44f04813255dd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 13:29:40 +0530 Subject: [PATCH 0201/1087] Handle dashboard ids on MSSQL app database --- .../repositories/dashboardRepository.ts | 177 ++++++++++++++++++ 1 file changed, 177 insertions(+) diff --git a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts index b43cea8cdc..503036b7c5 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts @@ -3,6 +3,7 @@ import { BaseRepository, IBasicRepository, coerceBoolean, + IQueryOptions, } from './baseRepository'; import { ScheduleFrequencyEnum } from '@server/models/dashboard'; @@ -15,6 +16,8 @@ export interface Dashboard { scheduleTimezone: string | null; // e.g. 'America/New_York', 'Asia/Taipei' scheduleCron: string | null; // cron expression string nextScheduledAt: Date | null; // Next scheduled run timestamp + createdAt?: Date; + updatedAt?: Date; } export interface IDashboardRepository extends IBasicRepository {} @@ -23,10 +26,71 @@ export class DashboardRepository extends BaseRepository implements IDashboardRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'dashboard' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + protected override transformFromDBData = (data: any): Dashboard => { const dashboard = this.defaultTransformFromDBData(data) as Dashboard; return { @@ -34,4 +98,117 @@ export class DashboardRepository cacheEnabled: coerceBoolean(dashboard.cacheEnabled), }; }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } From e3ed919e4a9a7b2b45af63aa41f24a25ce8a57d1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 16:47:48 +0530 Subject: [PATCH 0202/1087] Add timestamps for model writes on MSSQL --- .../server/repositories/modelRepository.ts | 176 ++++++++++++++++++ 1 file changed, 176 insertions(+) diff --git a/wren-ui/src/apollo/server/repositories/modelRepository.ts b/wren-ui/src/apollo/server/repositories/modelRepository.ts index 31823ffb4f..dd44274364 100644 --- a/wren-ui/src/apollo/server/repositories/modelRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelRepository.ts @@ -16,6 +16,8 @@ export interface Model { cached: boolean; // Model is cached or not refreshTime: string | null; // Contain a number followed by a time unit (ns, us, ms, s, m, h, d). For example, "2h" properties: string | null; // Model properties, a json string, the description and displayName should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IModelRepository extends IBasicRepository { @@ -30,10 +32,71 @@ export class ModelRepository extends BaseRepository implements IModelRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'model' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + protected override transformFromDBData = (data: any): Model => { const model = this.defaultTransformFromDBData(data) as Model; return { @@ -57,4 +120,117 @@ export class ModelRepository .delete(); return await builder; } + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } From bb64159fce5e61e40a7615639fe308c90e091a0c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 17:04:36 +0530 Subject: [PATCH 0203/1087] Add timestamps for model column writes on MSSQL --- .../repositories/modelColumnRepository.ts | 178 +++++++++++++++++ .../modelNestedColumnRepository.ts | 184 +++++++++++++++++- 2 files changed, 361 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts index c63df2daa9..78eb5ace83 100644 --- a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts @@ -20,6 +20,8 @@ export interface ModelColumn { notNull: boolean; // Is not null isPk: boolean; // Is primary key of the table properties?: string; // Column properties, a json string, the description and displayName should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IModelColumnRepository extends IBasicRepository { @@ -52,10 +54,71 @@ export class ModelColumnRepository extends BaseRepository implements IModelColumnRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'model_column' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + protected override transformFromDBData = (data: any): ModelColumn => { const column = this.defaultTransformFromDBData(data) as ModelColumn; return { @@ -139,4 +202,119 @@ export class ModelColumnRepository .whereIn('id', columnIds) .delete(); } + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } diff --git a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts index 5438ff47ec..a6b4e5bdee 100644 --- a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -18,6 +22,8 @@ export interface ModelNestedColumn { sourceColumnName: string; // The nested column name in the datasource type: string; // Data type, refer to the nested column type in the datasource properties?: Record; // Nested column properties, a json string, the description should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IModelNestedColumnRepository @@ -30,10 +36,71 @@ export class ModelNestedColumnRepository extends BaseRepository implements IModelNestedColumnRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'model_nested_column' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + public findNestedColumnsByModelIds = async (modelIds: number[]) => { const result = await this.knex(this.tableName) .select('*') @@ -79,4 +146,119 @@ export class ModelNestedColumnRepository }) as ModelNestedColumn; return formattedData; }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } From 39a1648f4382f2d9c1162de79247258a59974c96 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 17:50:38 +0530 Subject: [PATCH 0204/1087] Handle stale setup relations on MSSQL --- .../repositories/relationshipRepository.ts | 79 +++++++++++++++++++ .../apollo/server/services/modelService.ts | 38 ++++++--- 2 files changed, 104 insertions(+), 13 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts index 9aba1b6bc7..e33f033091 100644 --- a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts +++ b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts @@ -15,6 +15,8 @@ export interface Relation { fromColumnId: number; // from column id, "{fromColumn} {joinType} {toColumn}" toColumnId: number; // to column id, "{fromColumn} {joinType} {toColumn}" properties: string | null; // Model properties, a json string, the description should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface ExtraRelationInfo { @@ -62,10 +64,50 @@ export class RelationRepository extends BaseRepository implements IRelationRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'relation' }); } + public async createOne(data: Partial): Promise { + if (!this.isMssql() || (await this.hasIdentityId())) { + return super.createOne(this.withTimestamps(data)); + } + + const id = await this.nextRelationId(); + return super.createOne(this.withTimestamps({ id, ...data })); + } + + public async createMany(data: Partial[]): Promise { + if (data.length === 0) { + return []; + } + + const dataWithTimestamps = data.map((relation) => + this.withTimestamps(relation), + ); + + if (!this.isMssql() || (await this.hasIdentityId())) { + return super.createMany(dataWithTimestamps); + } + + const startId = await this.nextRelationId(); + return super.createMany( + dataWithTimestamps.map((relation, index) => ({ + id: startId + index, + ...relation, + })), + ); + } + + public async updateOne(id: number, data: Partial) { + return super.updateOne(id, { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }); + } + public async findRelationsBy( { columnIds, modelIds }, queryOptions?: IQueryOptions, @@ -220,4 +262,41 @@ export class RelationRepository const result = await query; return result.map((r) => this.transformFromDBData(r)) as RelationInfo[]; } + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private async hasIdentityId() { + this.hasIdentityIdPromise ??= this.knex('sys.columns as c') + .join('sys.tables as t', 'c.object_id', 't.object_id') + .where('t.name', this.tableName) + .where('c.name', 'id') + .select( + this.knex.raw( + 'COLUMNPROPERTY(c.object_id, c.name, ?) as isIdentity', + ['IsIdentity'], + ), + ) + .first() + .then((row) => Number(row?.isIdentity ?? 0) === 1); + + return this.hasIdentityIdPromise; + } + + private async nextRelationId() { + const row = await this.knex(this.tableName) + .max<{ maxId?: number | string }>('id as maxId') + .first(); + + return Number(row?.maxId ?? 0) + 1; + } } diff --git a/wren-ui/src/apollo/server/services/modelService.ts b/wren-ui/src/apollo/server/services/modelService.ts index 5299ade75c..c0b8bb5083 100644 --- a/wren-ui/src/apollo/server/services/modelService.ts +++ b/wren-ui/src/apollo/server/services/modelService.ts @@ -353,32 +353,44 @@ export class ModelService implements IModelService { .flat(); const columns = await this.modelColumnRepository.findColumnsByIds(columnIds); - const relationValues = relations.map((relation) => { + const relationValues = relations.flatMap((relation) => { const fromColumn = columns.find( (column) => column.id === relation.fromColumnId, ); if (!fromColumn) { - throw new Error(`Column not found, column Id ${relation.fromColumnId}`); + logger.warn( + `Skip relation because column ${relation.fromColumnId} was not found`, + ); + return []; } const toColumn = columns.find( (column) => column.id === relation.toColumnId, ); if (!toColumn) { - throw new Error(`Column not found, column Id ${relation.toColumnId}`); + logger.warn( + `Skip relation because column ${relation.toColumnId} was not found`, + ); + return []; } const relationName = this.generateRelationName(relation, models, columns); - return { - projectId: id, - name: relationName, - fromColumnId: relation.fromColumnId, - toColumnId: relation.toColumnId, - joinType: relation.type, - properties: relation.description - ? JSON.stringify({ description: relation.description }) - : null, - } as Partial; + return [ + { + projectId: id, + name: relationName, + fromColumnId: relation.fromColumnId, + toColumnId: relation.toColumnId, + joinType: relation.type, + properties: relation.description + ? JSON.stringify({ description: relation.description }) + : null, + } as Partial, + ]; }); + if (isEmpty(relationValues)) { + return []; + } + const savedRelations = await this.relationRepository.createMany(relationValues); From ae971947333ee4fb4d8efbfb2238a977f704472c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 16 Jun 2026 20:49:52 +0530 Subject: [PATCH 0205/1087] Add MSSQL timestamps for deploy logs --- .../repositories/deployLogRepository.ts | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts index 6e2e0471fc..ab7dc75baa 100644 --- a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts +++ b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts @@ -9,6 +9,8 @@ export interface Deploy { hash: string; status: string; // Deploy status error: string; // Error message + createdAt?: Date; + updatedAt?: Date; } export enum DeployStatusEnum { @@ -26,10 +28,28 @@ export class DeployLogRepository extends BaseRepository implements IDeployLogRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'deploy_log' }); } + public async createOne(data: Partial): Promise { + if (!this.isMssql() || (await this.hasIdentityId())) { + return super.createOne(this.withTimestamps(data)); + } + + const id = await this.nextDeployLogId(); + return super.createOne(this.withTimestamps({ id, ...data })); + } + + public async updateOne(id: number, data: Partial) { + return super.updateOne(id, { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }); + } + public async findLastProjectDeployLog(projectId: number) { const res = await this.knex .select('*') @@ -71,4 +91,41 @@ export class DeployLogRepository }); return formattedData as Deploy; }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private async hasIdentityId() { + this.hasIdentityIdPromise ??= this.knex('sys.columns as c') + .join('sys.tables as t', 'c.object_id', 't.object_id') + .where('t.name', this.tableName) + .where('c.name', 'id') + .select( + this.knex.raw( + 'COLUMNPROPERTY(c.object_id, c.name, ?) as isIdentity', + ['IsIdentity'], + ), + ) + .first() + .then((row) => Number(row?.isIdentity ?? 0) === 1); + + return this.hasIdentityIdPromise; + } + + private async nextDeployLogId() { + const row = await this.knex(this.tableName) + .max<{ maxId?: number | string }>('id as maxId') + .first(); + + return Number(row?.maxId ?? 0) + 1; + } } From fad918808ac009bb925733c5c935f3761469380c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 18 Jun 2026 22:13:50 +0530 Subject: [PATCH 0206/1087] Batch repository operations for large schemas --- .../server/repositories/baseRepository.ts | 37 ++++++- .../repositories/modelColumnRepository.ts | 103 ++++++++++++------ .../modelNestedColumnRepository.ts | 37 ++++++- .../server/repositories/modelRepository.ts | 35 +++++- 4 files changed, 164 insertions(+), 48 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index 7d81514e9b..e374cdb6b3 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -129,7 +129,11 @@ export class BaseRepository implements IBasicRepository { public async createMany(data: Partial[], queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const preparedData = await this.prepareInsertManyData(data, executer); - const batchSize = 100; + if (preparedData.length === 0) { + return []; + } + + const batchSize = this.getCreateManyBatchSize(preparedData); const batchCount = Math.ceil(preparedData.length / batchSize); const result = []; for (let i = 0; i < batchCount; i++) { @@ -202,6 +206,37 @@ export class BaseRepository implements IBasicRepository { protected transformFromDBData = (data: any): T => this.defaultTransformFromDBData(data); + protected getMssqlWhereInBatchSize() { + return 2000; + } + + protected getCreateManyBatchSize(insertValues: any[]) { + const defaultBatchSize = 100; + if (insertValues.length === 0) { + return defaultBatchSize; + } + + const client = String(this.knex.client.config.client || '').toLowerCase(); + if (client !== 'mssql') { + return defaultBatchSize; + } + + const parameterLimit = 2100; + const safetyMargin = 100; + const columnCount = Math.max( + ...insertValues.map((value) => Object.keys(value).length), + 1, + ); + + return Math.max( + 1, + Math.min( + defaultBatchSize, + Math.floor((parameterLimit - safetyMargin) / columnCount), + ), + ); + } + private isMssql(executer: Knex | Knex.Transaction) { return executer.client.config.client === 'mssql'; } diff --git a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts index 78eb5ace83..89edb96434 100644 --- a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts @@ -130,28 +130,16 @@ export class ModelColumnRepository }; public async findColumnsByModelIds(modelIds, queryOptions?: IQueryOptions) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - const result = await tx(this.tableName) - .whereIn('model_id', modelIds) - .select('*'); - return result.map((r) => this.transformFromDBData(r)); - } - const result = await this.knex('model_column') - .whereIn('model_id', modelIds) - .select('*'); + const result = await this.findByColumnIn( + 'model_id', + modelIds, + queryOptions, + ); return result.map((r) => this.transformFromDBData(r)); } public async findColumnsByIds(ids: number[], queryOptions?: IQueryOptions) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - const result = await tx(this.tableName).whereIn('id', ids).select('*'); - return result.map((r) => this.transformFromDBData(r)); - } - const result = await this.knex('model_column') - .whereIn('id', ids) - .select('*'); + const result = await this.findByColumnIn('id', ids, queryOptions); return result.map((r) => this.transformFromDBData(r)); } @@ -159,14 +147,7 @@ export class ModelColumnRepository modelIds: number[], queryOptions?: IQueryOptions, ) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - await tx(this.tableName).whereIn('model_id', modelIds).delete(); - return; - } - await this.knex('model_column') - .whereIn('model_id', modelIds) - .delete(); + await this.deleteByColumnIn('model_id', modelIds, queryOptions); } public async resetModelPrimaryKey(modelId: number) { @@ -185,22 +166,74 @@ export class ModelColumnRepository sourceColumnNames: string[], queryOptions?: IQueryOptions, ): Promise { - const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer(this.tableName) - .where(this.transformToDBData({ modelId })) - .whereIn('source_column_name', sourceColumnNames) - .delete(); - return await builder; + return await this.deleteByColumnIn( + 'source_column_name', + sourceColumnNames, + queryOptions, + this.transformToDBData({ modelId }), + ); } public async deleteAllByColumnIds( columnIds: number[], queryOptions?: IQueryOptions, ): Promise { + await this.deleteByColumnIn('id', columnIds, queryOptions); + } + + private async findByColumnIn( + columnName: string, + values: Array, + queryOptions?: IQueryOptions, + ) { + if (values.length === 0) { + return []; + } + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - await executer(this.tableName) - .whereIn('id', columnIds) - .delete(); + const rows = []; + for (const batch of this.toWhereInBatches(values)) { + const result = await executer(this.tableName) + .whereIn(columnName, batch) + .select('*'); + rows.push(...result); + } + + return rows; + } + + private async deleteByColumnIn( + columnName: string, + values: Array, + queryOptions?: IQueryOptions, + extraWhere?: Record, + ) { + if (values.length === 0) { + return 0; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + let deleted = 0; + for (const batch of this.toWhereInBatches(values)) { + const query = executer(this.tableName).whereIn(columnName, batch); + if (extraWhere) { + query.where(extraWhere); + } + deleted += await query.delete(); + } + + return deleted; + } + + private toWhereInBatches(values: TValue[]) { + const batchSize = this.isMssql() + ? this.getMssqlWhereInBatchSize() + : Math.max(values.length, 1); + const batches: TValue[][] = []; + for (let index = 0; index < values.length; index += batchSize) { + batches.push(values.slice(index, index + batchSize)); + } + return batches; } private isMssql = () => diff --git a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts index a6b4e5bdee..fec5189f87 100644 --- a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts @@ -102,19 +102,44 @@ export class ModelNestedColumnRepository } public findNestedColumnsByModelIds = async (modelIds: number[]) => { - const result = await this.knex(this.tableName) - .select('*') - .whereIn('model_id', modelIds); + const result = await this.findByColumnIn('model_id', modelIds); return result.map((r) => this.transformFromDBData(r)); }; public findNestedColumnsByIds = async (ids: number[]) => { - const result = await this.knex(this.tableName) - .select('*') - .whereIn('id', ids); + const result = await this.findByColumnIn('id', ids); return result.map((r) => this.transformFromDBData(r)); }; + private async findByColumnIn( + columnName: string, + values: Array, + ) { + if (values.length === 0) { + return []; + } + + const rows = []; + for (const batch of this.toWhereInBatches(values)) { + const result = await this.knex(this.tableName) + .select('*') + .whereIn(columnName, batch); + rows.push(...result); + } + return rows; + } + + private toWhereInBatches(values: TValue[]) { + const batchSize = this.isMssql() + ? this.getMssqlWhereInBatchSize() + : Math.max(values.length, 1); + const batches: TValue[][] = []; + for (let index = 0; index < values.length; index += batchSize) { + batches.push(values.slice(index, index + batchSize)); + } + return batches; + } + protected override transformToDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); diff --git a/wren-ui/src/apollo/server/repositories/modelRepository.ts b/wren-ui/src/apollo/server/repositories/modelRepository.ts index dd44274364..6ef5dacd04 100644 --- a/wren-ui/src/apollo/server/repositories/modelRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelRepository.ts @@ -106,8 +106,17 @@ export class ModelRepository }; public async findAllByIds(ids: number[]) { - const res = await this.knex(this.tableName).whereIn('id', ids); - return res.map((r) => this.transformFromDBData(r)); + if (ids.length === 0) { + return []; + } + + const rows = []; + for (const batch of this.toWhereInBatches(ids)) { + const res = await this.knex(this.tableName).whereIn('id', batch); + rows.push(...res); + } + + return rows.map((r) => this.transformFromDBData(r)); } public async deleteAllBySourceTableNames( @@ -115,10 +124,24 @@ export class ModelRepository queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer(this.tableName) - .whereIn('source_table_name', sourceTableNames) - .delete(); - return await builder; + let deleted = 0; + for (const batch of this.toWhereInBatches(sourceTableNames)) { + deleted += await executer(this.tableName) + .whereIn('source_table_name', batch) + .delete(); + } + return deleted; + } + + private toWhereInBatches(values: TValue[]) { + const batchSize = this.isMssql() + ? this.getMssqlWhereInBatchSize() + : Math.max(values.length, 1); + const batches: TValue[][] = []; + for (let index = 0; index < values.length; index += batchSize) { + batches.push(values.slice(index, index + batchSize)); + } + return batches; } private isMssql = () => From ef10329bdab6c8eca07745bac0600a9fedd81fa9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 18 Jun 2026 22:43:09 +0530 Subject: [PATCH 0207/1087] Chunk relationship queries for large schemas --- .../server/repositories/baseRepository.ts | 23 ++- .../repositories/modelColumnRepository.ts | 11 - .../modelNestedColumnRepository.ts | 11 - .../server/repositories/modelRepository.ts | 11 - .../repositories/relationshipRepository.ts | 191 ++++++++++-------- 5 files changed, 128 insertions(+), 119 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index e374cdb6b3..a286555250 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -173,8 +173,14 @@ export class BaseRepository implements IBasicRepository { queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer.from(this.tableName).whereIn('id', ids).delete(); - return await builder; + let deleted = 0; + for (const batch of this.toWhereInBatches(ids)) { + deleted += await executer + .from(this.tableName) + .whereIn('id', batch) + .delete(); + } + return deleted; } public deleteAllBy = async ( @@ -210,6 +216,19 @@ export class BaseRepository implements IBasicRepository { return 2000; } + protected toWhereInBatches(values: TValue[]) { + const client = String(this.knex.client.config.client || '').toLowerCase(); + const batchSize = + client === 'mssql' + ? this.getMssqlWhereInBatchSize() + : Math.max(values.length, 1); + const batches: TValue[][] = []; + for (let index = 0; index < values.length; index += batchSize) { + batches.push(values.slice(index, index + batchSize)); + } + return batches; + } + protected getCreateManyBatchSize(insertValues: any[]) { const defaultBatchSize = 100; if (insertValues.length === 0) { diff --git a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts index 89edb96434..43f6a89f97 100644 --- a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts @@ -225,17 +225,6 @@ export class ModelColumnRepository return deleted; } - private toWhereInBatches(values: TValue[]) { - const batchSize = this.isMssql() - ? this.getMssqlWhereInBatchSize() - : Math.max(values.length, 1); - const batches: TValue[][] = []; - for (let index = 0; index < values.length; index += batchSize) { - batches.push(values.slice(index, index + batchSize)); - } - return batches; - } - private isMssql = () => String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; diff --git a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts index fec5189f87..b46952b5ab 100644 --- a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts @@ -129,17 +129,6 @@ export class ModelNestedColumnRepository return rows; } - private toWhereInBatches(values: TValue[]) { - const batchSize = this.isMssql() - ? this.getMssqlWhereInBatchSize() - : Math.max(values.length, 1); - const batches: TValue[][] = []; - for (let index = 0; index < values.length; index += batchSize) { - batches.push(values.slice(index, index + batchSize)); - } - return batches; - } - protected override transformToDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); diff --git a/wren-ui/src/apollo/server/repositories/modelRepository.ts b/wren-ui/src/apollo/server/repositories/modelRepository.ts index 6ef5dacd04..bbd1bd93d1 100644 --- a/wren-ui/src/apollo/server/repositories/modelRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelRepository.ts @@ -133,17 +133,6 @@ export class ModelRepository return deleted; } - private toWhereInBatches(values: TValue[]) { - const batchSize = this.isMssql() - ? this.getMssqlWhereInBatchSize() - : Math.max(values.length, 1); - const batches: TValue[][] = []; - for (let index = 0; index < values.length; index += batchSize) { - batches.push(values.slice(index, index + batchSize)); - } - return batches; - } - private isMssql = () => String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; diff --git a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts index e33f033091..309b9a46db 100644 --- a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts +++ b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts @@ -112,41 +112,41 @@ export class RelationRepository { columnIds, modelIds }, queryOptions?: IQueryOptions, ) { - let executer = this.knex; - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - executer = tx; - } - // select the leftModel name and rightModel name along with relation - const builder = executer(this.tableName) - .join( - 'model_column AS fmc', - `${this.tableName}.from_column_id`, - '=', - 'fmc.id', - ) - .join( - 'model_column AS tmc', - `${this.tableName}.to_column_id`, - '=', - 'tmc.id', + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const selectRows = (builder) => + builder.select( + `${this.tableName}.*`, + 'fmc.model_id AS fromModelId', + 'tmc.model_id AS toModelId', ); + const rows = []; if (columnIds && columnIds.length > 0) { - builder - .whereIn(`${this.tableName}.from_column_id`, columnIds) - .orWhereIn(`${this.tableName}.to_column_id`, columnIds); + for (const batch of this.toWhereInBatches(columnIds)) { + const result = await selectRows(this.relationJoinBuilder(executer)).where( + (builder) => + builder + .whereIn(`${this.tableName}.from_column_id`, batch) + .orWhereIn(`${this.tableName}.to_column_id`, batch), + ); + rows.push(...result); + } + return rows.map((r) => this.transformFromDBData(r)); } + if (modelIds && modelIds.length > 0) { - builder - .whereIn('fmc.model_id', modelIds) - .orWhereIn('tmc.model_id', modelIds); + for (const batch of this.toWhereInBatches(modelIds)) { + const result = await selectRows(this.relationJoinBuilder(executer)).where( + (builder) => + builder + .whereIn('fmc.model_id', batch) + .orWhereIn('tmc.model_id', batch), + ); + rows.push(...result); + } + return rows.map((r) => this.transformFromDBData(r)); } - const result = await builder.select( - `${this.tableName}.*`, - 'fmc.model_id AS fromModelId', - 'tmc.model_id AS toModelId', - ); + const result = await selectRows(this.relationJoinBuilder(executer)); return result.map((r) => this.transformFromDBData(r)); } @@ -157,9 +157,11 @@ export class RelationRepository executer = tx; } - const result = await executer(this.tableName) - .whereIn('id', ids) - .select('*'); + const result = []; + for (const batch of this.toWhereInBatches(ids)) { + const rows = await executer(this.tableName).whereIn('id', batch).select('*'); + result.push(...rows); + } return result.map((r) => this.transformFromDBData(r)); } @@ -167,70 +169,69 @@ export class RelationRepository columnIds: number[], queryOptions?: IQueryOptions, ) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - await tx(this.tableName) - .whereIn('from_column_id', columnIds) - .orWhereIn('to_column_id', columnIds) + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + for (const batch of this.toWhereInBatches(columnIds)) { + await executer(this.tableName) + .where((builder) => + builder + .whereIn('from_column_id', batch) + .orWhereIn('to_column_id', batch), + ) .delete(); - return; } - await this.knex(this.tableName) - .whereIn('from_column_id', columnIds) - .orWhereIn('to_column_id', columnIds) - .delete(); } public async findRelationInfoBy(filter, queryOptions) { const { projectId, columnIds, modelIds } = filter; - let executer = this.knex; - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - executer = tx; - } - // select the leftModel name and rightModel name along with relation - const builder = executer(this.tableName) - .join( - 'model_column AS fmc', - `${this.tableName}.from_column_id`, - '=', - 'fmc.id', - ) - .join( - 'model_column AS tmc', - `${this.tableName}.to_column_id`, - '=', - 'tmc.id', - ) - .join('model AS fm', 'fmc.model_id', '=', 'fm.id') - .join('model AS tm', 'tmc.model_id', '=', 'tm.id'); + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const selectRows = (builder) => + builder.select( + `${this.tableName}.*`, + 'fm.id AS fromModelId', + 'fm.reference_name AS fromModelName', + 'fm.display_name AS fromModelDisplayName', + 'tm.id AS toModelId', + 'tm.reference_name AS toModelName', + 'tm.display_name AS toModelDisplayName', + 'fmc.reference_name AS fromColumnName', + 'fmc.display_name AS fromColumnDisplayName', + 'tmc.reference_name AS toColumnName', + 'tmc.display_name AS toColumnDisplayName', + ); if (projectId) { + const builder = this.relationInfoJoinBuilder(executer); builder.where(`${this.tableName}.project_id`, projectId); - } else if (columnIds && columnIds.length > 0) { - builder - .whereIn(`${this.tableName}.from_column_id`, columnIds) - .orWhereIn(`${this.tableName}.to_column_id`, columnIds); + const result = await selectRows(builder); + return result.map((r) => this.transformFromDBData(r)) as RelationInfo[]; + } + + const rows = []; + if (columnIds && columnIds.length > 0) { + for (const batch of this.toWhereInBatches(columnIds)) { + const result = await selectRows(this.relationInfoJoinBuilder(executer)).where( + (builder) => + builder + .whereIn(`${this.tableName}.from_column_id`, batch) + .orWhereIn(`${this.tableName}.to_column_id`, batch), + ); + rows.push(...result); + } } else if (modelIds && modelIds.length > 0) { - builder - .whereIn('fmc.model_id', modelIds) - .orWhereIn('tmc.model_id', modelIds); + for (const batch of this.toWhereInBatches(modelIds)) { + const result = await selectRows(this.relationInfoJoinBuilder(executer)).where( + (builder) => + builder + .whereIn('fmc.model_id', batch) + .orWhereIn('tmc.model_id', batch), + ); + rows.push(...result); + } + } else { + rows.push(...(await selectRows(this.relationInfoJoinBuilder(executer)))); } - const result = await builder.select( - `${this.tableName}.*`, - 'fm.id AS fromModelId', - 'fm.reference_name AS fromModelName', - 'fm.display_name AS fromModelDisplayName', - 'tm.id AS toModelId', - 'tm.reference_name AS toModelName', - 'tm.display_name AS toModelDisplayName', - 'fmc.reference_name AS fromColumnName', - 'fmc.display_name AS fromColumnDisplayName', - 'tmc.reference_name AS toColumnName', - 'tmc.display_name AS toColumnDisplayName', - ); - return result.map((r) => this.transformFromDBData(r)) as RelationInfo[]; + return rows.map((r) => this.transformFromDBData(r)) as RelationInfo[]; } public async findExistedRelationBetweenModels(relation: RelationData) { @@ -266,6 +267,28 @@ export class RelationRepository private isMssql = () => String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + private relationJoinBuilder(executer: Knex | Knex.Transaction) { + return executer(this.tableName) + .join( + 'model_column AS fmc', + `${this.tableName}.from_column_id`, + '=', + 'fmc.id', + ) + .join( + 'model_column AS tmc', + `${this.tableName}.to_column_id`, + '=', + 'tmc.id', + ); + } + + private relationInfoJoinBuilder(executer: Knex | Knex.Transaction) { + return this.relationJoinBuilder(executer) + .join('model AS fm', 'fmc.model_id', '=', 'fm.id') + .join('model AS tm', 'tmc.model_id', '=', 'tm.id'); + } + private withTimestamps = (data: Partial): Partial => { const now = new Date(); return { From 2d1e9eadfaa1308ca62298d99c1f33662bebc3c9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 18 Jun 2026 23:19:30 +0530 Subject: [PATCH 0208/1087] Keep relation queries under MSSQL parameter limit --- .../repositories/relationshipRepository.ts | 26 +++++++++++++++---- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts index 309b9a46db..d02aa4f6fc 100644 --- a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts +++ b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts @@ -121,7 +121,7 @@ export class RelationRepository ); const rows = []; if (columnIds && columnIds.length > 0) { - for (const batch of this.toWhereInBatches(columnIds)) { + for (const batch of this.toDualWhereInBatches(columnIds)) { const result = await selectRows(this.relationJoinBuilder(executer)).where( (builder) => builder @@ -134,7 +134,7 @@ export class RelationRepository } if (modelIds && modelIds.length > 0) { - for (const batch of this.toWhereInBatches(modelIds)) { + for (const batch of this.toDualWhereInBatches(modelIds)) { const result = await selectRows(this.relationJoinBuilder(executer)).where( (builder) => builder @@ -170,7 +170,7 @@ export class RelationRepository queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - for (const batch of this.toWhereInBatches(columnIds)) { + for (const batch of this.toDualWhereInBatches(columnIds)) { await executer(this.tableName) .where((builder) => builder @@ -208,7 +208,7 @@ export class RelationRepository const rows = []; if (columnIds && columnIds.length > 0) { - for (const batch of this.toWhereInBatches(columnIds)) { + for (const batch of this.toDualWhereInBatches(columnIds)) { const result = await selectRows(this.relationInfoJoinBuilder(executer)).where( (builder) => builder @@ -218,7 +218,7 @@ export class RelationRepository rows.push(...result); } } else if (modelIds && modelIds.length > 0) { - for (const batch of this.toWhereInBatches(modelIds)) { + for (const batch of this.toDualWhereInBatches(modelIds)) { const result = await selectRows(this.relationInfoJoinBuilder(executer)).where( (builder) => builder @@ -267,6 +267,22 @@ export class RelationRepository private isMssql = () => String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + private toDualWhereInBatches(values: TValue[]) { + if (!this.isMssql()) { + return this.toWhereInBatches(values); + } + + const batchSize = Math.max( + Math.floor(this.getMssqlWhereInBatchSize() / 2), + 1, + ); + const batches: TValue[][] = []; + for (let index = 0; index < values.length; index += batchSize) { + batches.push(values.slice(index, index + batchSize)); + } + return batches; + } + private relationJoinBuilder(executer: Knex | Knex.Transaction) { return executer(this.tableName) .join( From 0facdba9a9f1a301bd545262941777e114b5aa53 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 19 Jun 2026 00:33:23 +0530 Subject: [PATCH 0209/1087] Deduplicate model column references in MDL --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 117 +++++++++++++----- .../apollo/server/resolvers/modelResolver.ts | 18 ++- .../server/resolvers/projectResolver.ts | 30 +++-- wren-ui/src/apollo/server/utils/model.ts | 27 ++++ 4 files changed, 147 insertions(+), 45 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index d3980b3dee..6836a251c7 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -8,6 +8,7 @@ import { View, } from '../repositories'; import { + ColumnMDL, Manifest, ModelMDL, TableReference, @@ -16,6 +17,7 @@ import { import { getLogger } from '@server/utils'; import { getConfig } from '@server/config'; import { DataSourceName } from '../types'; +import { getUniqueReferenceName } from '../utils/model'; const logger = getLogger('MDLBuilder'); logger.level = 'debug'; @@ -61,6 +63,8 @@ export class MDLBuilder implements IMDLBuilder { private readonly relatedColumns: ModelColumn[]; // eslint-disable-next-line @typescript-eslint/no-unused-vars private readonly relatedRelations: RelationInfo[]; + private readonly columnNameAliases = new Map(); + private readonly manifestColumnNamesByModel = new Map>(); constructor(builderOptions: MDLBuilderBuildFromOptions) { const { @@ -191,11 +195,6 @@ export class MDLBuilder implements IMDLBuilder { (model: any) => model.name === modelRefName, ); - // modify model primary key - if (column.isPk) { - model.primaryKey = column.referenceName; - } - // add column into model if (!model.columns) { model.columns = []; @@ -221,9 +220,15 @@ export class MDLBuilder implements IMDLBuilder { } }, {}); } - const expression = this.getColumnExpression(column, model); + const columnName = this.getManifestColumnName(column, model); + // modify model primary key + if (column.isPk) { + model.primaryKey = columnName; + } + + const expression = this.getColumnExpression(column, model, columnName); model.columns.push({ - name: column.referenceName, + name: columnName, type: column.type, isCalculated: column.isCalculated ? true : false, notNull: column.notNull ? true : false, @@ -266,7 +271,8 @@ export class MDLBuilder implements IMDLBuilder { ); return; } - const expression = this.getColumnExpression(column, model); + const columnName = this.getManifestColumnName(column, model); + const expression = this.getColumnExpression(column, model, columnName); if (expression === null) { this.recordInvalidCalculatedField( column.modelId, @@ -276,7 +282,7 @@ export class MDLBuilder implements IMDLBuilder { return; } const columnValue = { - name: column.referenceName, + name: columnName, type: column.type, isCalculated: true, expression, @@ -306,15 +312,12 @@ export class MDLBuilder implements IMDLBuilder { logger.debug(`Can not find model "${modelName}" to add calculated field`); return; } - // if calculated field is already in the model, skip - if ( - model.columns.find( - (column: any) => column.name === calculatedField.referenceName, - ) - ) { - return; - } - const expression = this.getColumnExpression(calculatedField, model); + const columnName = this.getManifestColumnName(calculatedField, model); + const expression = this.getColumnExpression( + calculatedField, + model, + columnName, + ); if (expression === null) { this.recordInvalidCalculatedField( calculatedField.modelId, @@ -324,7 +327,7 @@ export class MDLBuilder implements IMDLBuilder { return; } const columnValue = { - name: calculatedField.referenceName, + name: columnName, type: calculatedField.type, isCalculated: true, expression, @@ -349,18 +352,22 @@ export class MDLBuilder implements IMDLBuilder { joinType, fromModelName, fromColumnName, + fromColumnId, toModelName, toColumnName, + toColumnId, } = relation; const condition = this.getRelationCondition(relation); this.addRelationColumn(fromModelName, { modelReferenceName: toModelName, - columnReferenceName: toColumnName, + columnReferenceName: + this.columnNameAliases.get(toColumnId) || toColumnName, relation: name, }); this.addRelationColumn(toModelName, { modelReferenceName: fromModelName, - columnReferenceName: fromColumnName, + columnReferenceName: + this.columnNameAliases.get(fromColumnId) || fromColumnName, relation: name, }); @@ -405,13 +412,18 @@ export class MDLBuilder implements IMDLBuilder { model.columns = []; } // check if the modelReferenceName is already in the model column - const modelNameDuplicated = model.columns.find( - (column: any) => column.name === columnData.modelReferenceName, + const modelColumnNames = this.getManifestColumnNames(model); + const modelNameDuplicated = modelColumnNames.has( + columnData.modelReferenceName.toLowerCase(), ); - const column = { - name: modelNameDuplicated + const columnName = getUniqueReferenceName( + modelNameDuplicated ? `${columnData.modelReferenceName}_${columnData.columnReferenceName}` : columnData.modelReferenceName, + modelColumnNames, + ); + const column = { + name: columnName, type: columnData.modelReferenceName, properties: null, relationship: columnData.relation, @@ -424,11 +436,12 @@ export class MDLBuilder implements IMDLBuilder { protected getColumnExpression( column: ModelColumn, currentModel?: Partial, + columnReferenceName = column.referenceName, ): string | null { if (!column.isCalculated) { // columns existed in the data source. // Provide original column name in expression to MDL if referenceName has converted. - if (column.sourceColumnName !== column.referenceName) { + if (column.sourceColumnName !== columnReferenceName) { return `"${column.sourceColumnName}"`; } return ''; @@ -443,9 +456,13 @@ export class MDLBuilder implements IMDLBuilder { const isLast = index === lineage.length - 1; if (isLast) { // id is columnId - const columnReferenceName = this.relatedColumns.find( + const relatedColumn = this.relatedColumns.find( (relatedColumn) => relatedColumn.id === id, - )?.referenceName; + ); + const columnReferenceName = relatedColumn + ? this.columnNameAliases.get(relatedColumn.id) || + relatedColumn.referenceName + : null; if (!columnReferenceName) { return acc; } @@ -485,9 +502,19 @@ export class MDLBuilder implements IMDLBuilder { protected getRelationCondition(relation: RelationInfo): string { //TODO phase2: implement the expression for relation condition - const { fromColumnName, toColumnName, fromModelName, toModelName } = - relation; - return `"${fromModelName}".${fromColumnName} = "${toModelName}".${toColumnName}`; + const { + fromColumnId, + fromColumnName, + toColumnId, + toColumnName, + fromModelName, + toModelName, + } = relation; + const fromColumnReferenceName = + this.columnNameAliases.get(fromColumnId) || fromColumnName; + const toColumnReferenceName = + this.columnNameAliases.get(toColumnId) || toColumnName; + return `"${fromModelName}".${fromColumnReferenceName} = "${toModelName}".${toColumnReferenceName}`; } private buildTableReference(model: Model): TableReference | null { @@ -646,6 +673,34 @@ export class MDLBuilder implements IMDLBuilder { `Skipped ${this.invalidCalculatedFields.length} invalid calculated field(s) while building MDL. ${preview}${this.invalidCalculatedFields.length > 10 ? '; ...' : ''}`, ); } + + private getManifestColumnName( + column: ModelColumn, + model: Partial, + ): string { + if (this.columnNameAliases.has(column.id)) { + return this.columnNameAliases.get(column.id)!; + } + const columnName = getUniqueReferenceName( + column.referenceName, + this.getManifestColumnNames(model), + ); + this.columnNameAliases.set(column.id, columnName); + return columnName; + } + + private getManifestColumnNames(model: Partial): Set { + const modelName = model.name || ''; + if (!this.manifestColumnNamesByModel.has(modelName)) { + const existingColumns = (model.columns || []) as ColumnMDL[]; + this.manifestColumnNamesByModel.set( + modelName, + new Set(existingColumns.map((column) => column.name.toLowerCase())), + ); + } + return this.manifestColumnNamesByModel.get(modelName)!; + } + private postProcessManifest() { if (this.useRustWrenEngine()) { // 1. remove all the key that the value is null diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 4afff6351e..76d9abc504 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -13,7 +13,7 @@ import { RelationData, UpdateRelationData, } from '../types'; -import { getLogger, transformInvalidColumnName } from '@server/utils'; +import { getLogger, transformUniqueInvalidColumnName } from '@server/utils'; import { DeployResponse } from '../services/deployService'; import { safeFormatSQL } from '@server/utils/sqlFormat'; import { isEmpty, isNil } from 'lodash'; @@ -396,13 +396,17 @@ export class ModelResolver { const compactColumns = dataSourceTable.columns.filter((c) => fields.includes(c.name), ); + const usedReferenceNames = new Set(); const columnValues = compactColumns.map( (column) => ({ modelId: model.id, isCalculated: false, displayName: column.name, - referenceName: transformInvalidColumnName(column.name), + referenceName: transformUniqueInvalidColumnName( + column.name, + usedReferenceNames, + ), sourceColumnName: column.name, type: column.type || 'string', notNull: column.notNull || false, @@ -491,6 +495,11 @@ export class ModelResolver { // create columns if (toCreateColumns.length) { + const usedReferenceNames = new Set( + existingColumns + .filter(({ id }) => !toDeleteColumnIds.includes(id)) + .map(({ referenceName }) => referenceName.toLowerCase()), + ); const compactColumns = sourceTableColumns.filter((sourceColumn) => toCreateColumns.includes(sourceColumn.name), ); @@ -500,7 +509,10 @@ export class ModelResolver { isCalculated: false, displayName: column.name, sourceColumnName: column.name, - referenceName: transformInvalidColumnName(column.name), + referenceName: transformUniqueInvalidColumnName( + column.name, + usedReferenceNames, + ), type: column.type || 'string', notNull: column.notNull, isPk: primaryKey === column.name, diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 4460506a70..6ca33be5cc 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -12,7 +12,7 @@ import { trim, getLogger, replaceInvalidReferenceName, - transformInvalidColumnName, + transformUniqueInvalidColumnName, handleNestedColumns, } from '@server/utils'; import { @@ -733,13 +733,17 @@ export class ProjectResolver { const compactColumns = table.columns; const primaryKey = table.primaryKey; const model = models.find((m) => m.sourceTableName === table.name); + const usedReferenceNames = new Set(); return compactColumns.map( (column) => ({ modelId: model.id, isCalculated: false, displayName: column.name, - referenceName: transformInvalidColumnName(column.name), + referenceName: transformUniqueInvalidColumnName( + column.name, + usedReferenceNames, + ), sourceColumnName: column.name, type: column.type || 'string', notNull: column.notNull || false, @@ -753,15 +757,19 @@ export class ProjectResolver { const columns = await ctx.modelColumnRepository.createMany(columnValues); // create nested columns - const compactColumns = selectedTables.flatMap((table) => table.columns); - const nestedColumnValues = compactColumns.flatMap((compactColumn) => { - const column = columns.find( - (c) => c.sourceColumnName === compactColumn.name, - ); - return handleNestedColumns(compactColumn, { - modelId: column.modelId, - columnId: column.id, - sourceColumnName: column.sourceColumnName, + const nestedColumnValues = selectedTables.flatMap((table) => { + const model = models.find((m) => m.sourceTableName === table.name); + const tableColumns = columns.filter((c) => c.modelId === model.id); + return table.columns.flatMap((compactColumn) => { + const column = tableColumns.find( + (c) => c.sourceColumnName === compactColumn.name, + ); + if (!column) return []; + return handleNestedColumns(compactColumn, { + modelId: column.modelId, + columnId: column.id, + sourceColumnName: column.sourceColumnName, + }); }); }); await ctx.modelNestedColumnRepository.createMany(nestedColumnValues); diff --git a/wren-ui/src/apollo/server/utils/model.ts b/wren-ui/src/apollo/server/utils/model.ts index 05e79cb2a3..22c5542f6b 100644 --- a/wren-ui/src/apollo/server/utils/model.ts +++ b/wren-ui/src/apollo/server/utils/model.ts @@ -22,6 +22,33 @@ export function transformInvalidColumnName(columnName: string) { return referenceName; } +export function getUniqueReferenceName( + referenceName: string, + usedReferenceNames: Set, +) { + const baseName = referenceName || 'column'; + let uniqueName = baseName; + let suffix = 2; + + while (usedReferenceNames.has(uniqueName.toLowerCase())) { + uniqueName = `${baseName}_${suffix}`; + suffix += 1; + } + + usedReferenceNames.add(uniqueName.toLowerCase()); + return uniqueName; +} + +export function transformUniqueInvalidColumnName( + columnName: string, + usedReferenceNames: Set, +) { + return getUniqueReferenceName( + transformInvalidColumnName(columnName), + usedReferenceNames, + ); +} + export function replaceInvalidReferenceName(referenceName: string) { // replace dot with underscore return referenceName.replace(/\./g, '_'); From f80a9239af220557d75bc69379a6930cb7f91485 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 19 Jun 2026 13:57:44 +0530 Subject: [PATCH 0210/1087] Harden datasource metadata retrieval --- .../retrieval/db_schema_retrieval.py | 71 +---------- .../web/v1/services/semantics_preparation.py | 120 +++++++++++++++++- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 31 +++++ 3 files changed, 154 insertions(+), 68 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 212cd2f0b2..33ad534b63 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -250,18 +250,7 @@ async def table_retrieval( query_embedding=embedding.get("embedding"), filters=base_filters, ) - if result.get("documents") or not project_id: - return result - fallback_filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, - ], - } - return await table_retriever.run( - query_embedding=embedding.get("embedding"), - filters=fallback_filters, - ) + return result else: base_filters["conditions"].append( {"field": "name", "operator": "in", "value": tables} @@ -271,19 +260,7 @@ async def table_retrieval( query_embedding=[], filters=base_filters, ) - if result.get("documents") or not project_id: - return result - fallback_filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, - {"field": "name", "operator": "in", "value": tables}, - ], - } - return await table_retriever.run( - query_embedding=[], - filters=fallback_filters, - ) + return result @observe(capture_input=False) @@ -316,35 +293,7 @@ async def dbschema_retrieval( ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) - if results.get("documents") or not project_id: - documents = results["documents"] - if project_id and _is_project_wide_analysis_query(query): - all_project_results = await dbschema_retriever.run( - query_embedding=[], - filters={ - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": project_id}, - ], - }, - ) - documents = _dedupe_documents( - documents + all_project_results.get("documents", []) - ) - return documents - - fallback_filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } - fallback_results = await dbschema_retriever.run( - query_embedding=[], filters=fallback_filters - ) - documents = fallback_results["documents"] + documents = results["documents"] if project_id and _is_project_wide_analysis_query(query): all_project_results = await dbschema_retriever.run( query_embedding=[], @@ -377,19 +326,7 @@ async def dbschema_retrieval( project_id, ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) - if results.get("documents") or not project_id: - return results.get("documents", []) - - fallback_results = await dbschema_retriever.run( - query_embedding=[], - filters={ - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - }, - ) - return fallback_results.get("documents", []) + return results.get("documents", []) @observe() 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 2ff6215cbe..502e943dee 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -1,9 +1,10 @@ import asyncio import logging -from typing import Dict, Literal, Optional +from typing import Any, Dict, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe +import orjson from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline @@ -55,6 +56,117 @@ def __init__( str, SemanticsPreparationStatusResponse ] = TTLCache(maxsize=maxsize, ttl=ttl) + def _parse_mdl(self, mdl: str) -> dict[str, Any]: + parsed = orjson.loads(mdl) + parsed.setdefault("models", []) + parsed.setdefault("views", []) + parsed.setdefault("metrics", []) + parsed.setdefault("relationships", []) + return parsed + + def _validate_mdl_integrity(self, mdl: dict[str, Any]) -> None: + model_names = set() + for model in mdl["models"]: + model_name = model.get("name") + if not model_name: + raise ValueError("MDL contains a model without a name") + + normalized_model_name = model_name.lower() + if normalized_model_name in model_names: + raise ValueError(f'MDL contains duplicate model name "{model_name}"') + model_names.add(normalized_model_name) + + column_names = set() + for column in model.get("columns", []): + column_name = column.get("name") + if not column_name: + raise ValueError( + f'MDL model "{model_name}" contains a column without a name' + ) + + normalized_column_name = column_name.lower() + if normalized_column_name in column_names: + raise ValueError( + f'MDL model "{model_name}" contains duplicate column name "{column_name}"' + ) + column_names.add(normalized_column_name) + + for relationship in mdl["relationships"]: + for model_name in relationship.get("models", []): + if not model_name: + raise ValueError( + f'MDL relationship "{relationship.get("name", "")}" references an empty model name' + ) + if model_name.lower() not in model_names: + raise ValueError( + f'MDL relationship "{relationship.get("name", "")}" references missing model "{model_name}"' + ) + + def _project_filter( + self, project_id: Optional[str], *conditions: dict[str, Any] + ) -> dict[str, Any] | None: + all_conditions = list(conditions) + if project_id: + all_conditions.append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + if not all_conditions: + return None + return {"operator": "AND", "conditions": all_conditions} + + async def _count_indexed_documents( + self, + pipeline_name: str, + project_id: Optional[str], + *conditions: dict[str, Any], + ) -> int: + pipeline = self._pipelines[pipeline_name] + writer = pipeline._components["writer"] + return await writer.document_store.count_documents( + filters=self._project_filter(project_id, *conditions) + ) + + async def _validate_index_integrity( + self, mdl: dict[str, Any], project_id: Optional[str] + ) -> None: + resource_count = ( + len(mdl["models"]) + len(mdl["views"]) + len(mdl["metrics"]) + ) + expected_schema_documents = len(mdl["views"]) + len(mdl["metrics"]) + for model in mdl["models"]: + expected_schema_documents += 1 + if model.get("columns") or mdl["relationships"]: + expected_schema_documents += 1 + + schema_count, table_description_count, project_meta_count = await asyncio.gather( + self._count_indexed_documents( + "db_schema", + project_id, + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ), + self._count_indexed_documents( + "table_description", + project_id, + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + ), + self._count_indexed_documents("project_meta", project_id), + ) + + if schema_count < expected_schema_documents: + raise ValueError( + "Incomplete DB schema index: " + f"expected at least {expected_schema_documents} documents, found {schema_count}" + ) + + if table_description_count < resource_count: + raise ValueError( + "Incomplete table-description index: " + f"expected at least {resource_count} documents, found {table_description_count}" + ) + + if project_meta_count < 1: + raise ValueError("Project metadata was not indexed") + @observe(name="Prepare Semantics") @trace_metadata async def prepare_semantics( @@ -71,6 +183,8 @@ async def prepare_semantics( } try: + mdl = self._parse_mdl(prepare_semantics_request.mdl) + self._validate_mdl_integrity(mdl) logger.info(f"MDL: {prepare_semantics_request.mdl}") input = { @@ -90,6 +204,10 @@ async def prepare_semantics( ] await asyncio.gather(*tasks) + await self._validate_index_integrity( + mdl, + prepare_semantics_request.project_id, + ) self._prepare_semantics_statuses[ prepare_semantics_request.mdl_hash diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 6836a251c7..0b3aeef753 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -65,6 +65,10 @@ export class MDLBuilder implements IMDLBuilder { private readonly relatedRelations: RelationInfo[]; private readonly columnNameAliases = new Map(); private readonly manifestColumnNamesByModel = new Map>(); + private readonly manifestColumnNameBySourceByModel = new Map< + string, + Map + >(); constructor(builderOptions: MDLBuilderBuildFromOptions) { const { @@ -220,7 +224,24 @@ export class MDLBuilder implements IMDLBuilder { } }, {}); } + const sourceColumnName = column.sourceColumnName || column.referenceName; + const sourceColumnNames = this.getManifestSourceColumnNameMap(model); + const existingColumnName = sourceColumnNames.get( + sourceColumnName.toLowerCase(), + ); + if (existingColumnName) { + this.columnNameAliases.set(column.id, existingColumnName); + if (column.isPk) { + model.primaryKey = existingColumnName; + } + logger.debug( + `Skipping duplicate source column "${sourceColumnName}" for model "${model.name}". Reusing manifest column "${existingColumnName}".`, + ); + return; + } + const columnName = this.getManifestColumnName(column, model); + sourceColumnNames.set(sourceColumnName.toLowerCase(), columnName); // modify model primary key if (column.isPk) { model.primaryKey = columnName; @@ -701,6 +722,16 @@ export class MDLBuilder implements IMDLBuilder { return this.manifestColumnNamesByModel.get(modelName)!; } + private getManifestSourceColumnNameMap( + model: Partial, + ): Map { + const modelName = model.name || ''; + if (!this.manifestColumnNameBySourceByModel.has(modelName)) { + this.manifestColumnNameBySourceByModel.set(modelName, new Map()); + } + return this.manifestColumnNameBySourceByModel.get(modelName)!; + } + private postProcessManifest() { if (this.useRustWrenEngine()) { // 1. remove all the key that the value is null From aeabe24699ff6a0983209b041a207178161f1aff Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 19 Jun 2026 14:28:52 +0530 Subject: [PATCH 0211/1087] Normalize duplicate SQL projections --- .../src/pipelines/generation/utils/sql.py | 164 ++++++++++++++++++ 1 file changed, 164 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 01a5efbec5..bc1eacee42 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1069,6 +1069,168 @@ def _split_top_level_select_items(select_body: str) -> list[str]: return items +def _is_word_at(sql: str, index: int, word: str) -> bool: + end = index + len(word) + if sql[index:end].upper() != word: + return False + before = sql[index - 1] if index > 0 else "" + after = sql[end] if end < len(sql) else "" + return not (before.isalnum() or before == "_") and not ( + after.isalnum() or after == "_" + ) + + +def _find_select_list_spans(sql: str) -> list[tuple[int, int]]: + spans: list[tuple[int, int]] = [] + depth = 0 + in_single_quote = False + in_double_quote = False + in_bracket = False + index = 0 + + while index < len(sql): + char = sql[index] + if char == "'" and not in_double_quote and not in_bracket: + in_single_quote = not in_single_quote + elif char == '"' and not in_single_quote and not in_bracket: + in_double_quote = not in_double_quote + elif char == "[" and not in_single_quote and not in_double_quote: + in_bracket = True + elif char == "]" and in_bracket: + in_bracket = False + elif not in_single_quote and not in_double_quote and not in_bracket: + if char == "(": + depth += 1 + elif char == ")" and depth > 0: + depth -= 1 + elif _is_word_at(sql, index, "SELECT"): + select_depth = depth + select_body_start = index + len("SELECT") + cursor = select_body_start + cursor_depth = depth + cursor_in_single_quote = False + cursor_in_double_quote = False + cursor_in_bracket = False + + while cursor < len(sql): + cursor_char = sql[cursor] + if ( + cursor_char == "'" + and not cursor_in_double_quote + and not cursor_in_bracket + ): + cursor_in_single_quote = not cursor_in_single_quote + elif ( + cursor_char == '"' + and not cursor_in_single_quote + and not cursor_in_bracket + ): + cursor_in_double_quote = not cursor_in_double_quote + elif ( + cursor_char == "[" + and not cursor_in_single_quote + and not cursor_in_double_quote + ): + cursor_in_bracket = True + elif cursor_char == "]" and cursor_in_bracket: + cursor_in_bracket = False + elif ( + not cursor_in_single_quote + and not cursor_in_double_quote + and not cursor_in_bracket + ): + if cursor_char == "(": + cursor_depth += 1 + elif cursor_char == ")" and cursor_depth > 0: + cursor_depth -= 1 + elif cursor_depth == select_depth and _is_word_at( + sql, cursor, "FROM" + ): + spans.append((select_body_start, cursor)) + break + cursor += 1 + index = cursor + index += 1 + + return spans + + +def _strip_projection_alias(item: str) -> str: + alias_match = re.search( + r"\s+(?:AS\s+)?(?:\"[^\"]+\"|\[[^\]]+\]|`[^`]+`|[A-Za-z_][A-Za-z0-9_]*)\s*$", + item, + flags=re.IGNORECASE, + ) + if not alias_match: + return item.strip() + + expression = item[: alias_match.start()].strip() + if not expression: + return item.strip() + return expression + + +def _simple_projection_key(item: str) -> str | None: + cleaned = re.sub(r"^\s*DISTINCT\s+", "", item, flags=re.IGNORECASE).strip() + cleaned = _strip_projection_alias(cleaned) + identifier_pattern = ( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_]*|\*))" + ) + identifier_match_pattern = ( + r'(?:"[^"]+"|\[[^\]]+\]|`[^`]+`|[A-Za-z_][A-Za-z0-9_]*|\*)' + ) + qualified_identifier_pattern = rf"^\s*{identifier_match_pattern}(?:\s*\.\s*{identifier_match_pattern})*\s*$" + if not re.match(qualified_identifier_pattern, cleaned): + return None + + parts = re.findall(identifier_pattern, cleaned) + if not parts: + return None + last_part = next(value for value in parts[-1] if value) + return last_part.lower() + + +def _dedupe_duplicate_simple_select_items(sql: str) -> str: + spans = _find_select_list_spans(sql) + if not spans: + return sql + + normalized = sql + for start, end in reversed(spans): + body = normalized[start:end] + items = _split_top_level_select_items(body) + if len(items) < 2: + continue + + seen_simple_projection_keys: set[str] = set() + deduped_items: list[str] = [] + changed = False + for item in items: + key = _simple_projection_key(item) + if key and key in seen_simple_projection_keys: + changed = True + logger.debug( + 'Removing duplicate simple projection "%s" from generated SQL', + item, + ) + continue + if key: + seen_simple_projection_keys.add(key) + deduped_items.append(item) + + if changed: + normalized = ( + normalized[:start] + + " " + + ", ".join(deduped_items) + + " " + + normalized[end:] + ) + + return normalized + + def _rewrite_mssql_temporal_bucket_alias_references(sql: str) -> str: select_match = re.search( r"\bSELECT\b(?P.*?)(?=\bFROM\b)", @@ -1234,6 +1396,8 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> elif _references_known_hallucination_prone_schema(normalized): normalized = _rewrite_known_schema_hallucinations(normalized, datetime.now()) + normalized = _dedupe_duplicate_simple_select_items(normalized) + return re.sub(r"\s+", " ", normalized).strip() From 5ad8ff756bffdca99e94a5bc53fdc673af9cb0e0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 19 Jun 2026 14:46:04 +0530 Subject: [PATCH 0212/1087] Use deduped projections for duplicate source columns --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 87 ++++++++++++++++++++- 1 file changed, 85 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 0b3aeef753..bc166e7f55 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -124,8 +124,7 @@ export class MDLBuilder implements IMDLBuilder { properties.displayName = model.displayName; } const tableReference = this.buildTableReference(model); - - return { + const modelMdl = { name: model.referenceName, columns: [], tableReference, @@ -143,6 +142,23 @@ export class MDLBuilder implements IMDLBuilder { }, primaryKey: '', // will be modified in addColumn } as ModelMDL; + + if (tableReference && this.hasDuplicateSourceColumns(model.id)) { + const refSql = this.buildDedupedTableReferenceSql( + model.id, + modelMdl, + tableReference, + ); + if (refSql) { + logger.debug( + `Using deduped explicit projection for model "${model.referenceName}" because its source table contains duplicate column names.`, + ); + modelMdl.tableReference = null; + modelMdl.refSql = refSql; + } + } + + return modelMdl; }); } @@ -648,6 +664,73 @@ export class MDLBuilder implements IMDLBuilder { table: underscoreQualifiedMatch[2], }; } + private hasDuplicateSourceColumns(modelId: number): boolean { + const sourceColumnNames = new Set(); + for (const column of this.columns.filter( + ({ isCalculated, modelId: columnModelId }) => + !isCalculated && columnModelId === modelId, + )) { + const sourceColumnName = ( + column.sourceColumnName || column.referenceName + ).toLowerCase(); + if (sourceColumnNames.has(sourceColumnName)) { + return true; + } + sourceColumnNames.add(sourceColumnName); + } + return false; + } + private buildDedupedTableReferenceSql( + modelId: number, + model: Partial, + tableReference: TableReference, + ): string | null { + const sourceColumnNames = new Map(); + const projections: string[] = []; + + this.columns + .filter( + ({ isCalculated, modelId: columnModelId }) => + !isCalculated && columnModelId === modelId, + ) + .forEach((column) => { + const sourceColumnName = column.sourceColumnName || column.referenceName; + const normalizedSourceColumnName = sourceColumnName.toLowerCase(); + const existingColumnName = sourceColumnNames.get( + normalizedSourceColumnName, + ); + + if (existingColumnName) { + this.columnNameAliases.set(column.id, existingColumnName); + return; + } + + const columnName = this.getManifestColumnName(column, model); + sourceColumnNames.set(normalizedSourceColumnName, columnName); + const sourceExpression = this.quoteSqlIdentifier(sourceColumnName); + projections.push( + sourceColumnName === columnName + ? sourceExpression + : `${sourceExpression} AS ${this.quoteSqlIdentifier(columnName)}`, + ); + }); + + if (!projections.length) { + return null; + } + + const tableParts = [ + tableReference.catalog, + tableReference.schema, + tableReference.table, + ].filter((part): part is string => Boolean(part)); + return `SELECT ${projections.join(', ')} FROM ${tableParts + .map((part) => this.quoteSqlIdentifier(part)) + .join('.')}`; + } + private quoteSqlIdentifier(identifier: string): string { + return `"${identifier.replace(/"/g, '""')}"`; + } private parseLineage(lineage?: string): number[] { if (!lineage) { return []; From 8108ea25f1edf30de8a90a1fa41f06acf8a4845b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 03:01:27 +0530 Subject: [PATCH 0213/1087] Validate generated SQL table references --- .../generation/followup_sql_generation.py | 2 + .../pipelines/generation/sql_correction.py | 2 + .../pipelines/generation/sql_generation.py | 2 + .../pipelines/generation/sql_regeneration.py | 3 + .../src/pipelines/generation/utils/sql.py | 91 +++++++++++++++++++ 5 files changed, 100 insertions(+) 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 55e5ddbd83..243d32db67 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -166,6 +166,7 @@ async def generate_sql_in_followup( async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, + documents: list[str], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -177,6 +178,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + valid_table_names=construct_valid_table_names(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 223e50c100..9d15472e9d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -155,6 +155,7 @@ async def generate_sql_correction( async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, + documents: List[Document], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -166,6 +167,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + valid_table_names=construct_valid_table_names(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 6f847edcec..d3eb8749df 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -180,6 +180,7 @@ async def generate_sql( async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, + documents: list[str], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -193,6 +194,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + valid_table_names=construct_valid_table_names(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 7cf5787bc5..7d7e76d843 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -172,6 +173,7 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, + documents: list[str], data_source: str, project_id: str | None = None, ) -> dict: @@ -179,6 +181,7 @@ async def post_process( regenerate_sql.get("replies"), project_id=project_id, data_source=data_source, + valid_table_names=construct_valid_table_names(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index bc1eacee42..7b93ca7fba 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1418,6 +1418,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + valid_table_names: list[str] | None = None, ) -> dict: try: cleaned_generation_result = extract_sql_generation_result(replies[0]) @@ -1438,6 +1439,29 @@ async def run( }, } + invalid_table_references = find_invalid_table_references( + cleaned_generation_result, + valid_table_names or [], + ) + if invalid_table_references: + valid_table_list = ", ".join(valid_table_names or []) + invalid_table_list = ", ".join(invalid_table_references) + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_VALIDATION", + "error": ( + "Generated SQL references table(s) not present in the " + f"active datasource metadata: {invalid_table_list}. " + "Use only these valid table names exactly as shown: " + f"{valid_table_list}" + ), + "correlation_id": "", + }, + } + if normalize_data_source( data_source ) == "MSSQL" and contains_unsupported_mssql_json_access( @@ -2104,6 +2128,73 @@ def construct_valid_table_names(documents: list[Any] | None = None) -> list[str] return sorted(set(table_names)) +_SQL_IDENTIFIER_PATTERN = ( + r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)' +) +_SQL_TABLE_REFERENCE_PATTERN = re.compile( + rf"\b(?:FROM|JOIN)\s+" + rf"(?P
{_SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SQL_IDENTIFIER_PATTERN})*)", + flags=re.IGNORECASE, +) +_SQL_CTE_PATTERN = re.compile( + rf"(?:\bWITH\b|,)\s*(?P{_SQL_IDENTIFIER_PATTERN})\s+AS\s*\(", + flags=re.IGNORECASE, +) + + +def _normalize_sql_identifier(identifier: str) -> str: + identifier = identifier.strip() + if ( + (identifier.startswith('"') and identifier.endswith('"')) + or (identifier.startswith("`") and identifier.endswith("`")) + or (identifier.startswith("[") and identifier.endswith("]")) + ): + return identifier[1:-1] + return identifier + + +def _split_table_reference(table_reference: str) -> list[str]: + return [ + _normalize_sql_identifier(part) + for part in re.split(r"\s*\.\s*", table_reference.strip()) + if part.strip() + ] + + +def extract_sql_table_references(sql: str) -> list[str]: + references = [] + for match in _SQL_TABLE_REFERENCE_PATTERN.finditer(sql): + table_reference = match.group("table") + if table_reference.startswith("("): + continue + references.append(".".join(_split_table_reference(table_reference))) + return references + + +def extract_cte_names(sql: str) -> set[str]: + return { + _normalize_sql_identifier(match.group("cte")).lower() + for match in _SQL_CTE_PATTERN.finditer(sql) + } + + +def find_invalid_table_references(sql: str, valid_table_names: list[str]) -> list[str]: + if not valid_table_names: + return [] + + valid_tables = {table_name.lower() for table_name in valid_table_names} + cte_names = extract_cte_names(sql) + invalid_references = [] + + for table_reference in extract_sql_table_references(sql): + normalized_reference = table_reference.lower() + if normalized_reference in valid_tables or normalized_reference in cte_names: + continue + invalid_references.append(table_reference) + + return sorted(set(invalid_references)) + + def construct_ask_history_messages( histories: list[Any] | list[dict], ) -> list[ChatMessage]: From 190be5d98d815b3aaadee1999280322beb8852fd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 03:18:38 +0530 Subject: [PATCH 0214/1087] Validate generated SQL column references --- .../generation/followup_sql_generation.py | 2 + .../pipelines/generation/sql_correction.py | 2 + .../pipelines/generation/sql_generation.py | 2 + .../pipelines/generation/sql_regeneration.py | 2 + .../src/pipelines/generation/utils/sql.py | 175 +++++++++++++++++- 5 files changed, 181 insertions(+), 2 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 243d32db67..710d5d5dfe 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -15,6 +15,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, @@ -179,6 +180,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 9d15472e9d..a8e3174b60 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_columns, construct_valid_table_names, get_sql_generation_model_kwargs, get_text_to_sql_rules, @@ -168,6 +169,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index d3eb8749df..1f9ba2b81e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, @@ -195,6 +196,7 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 7d7e76d843..8b3eaafcb1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, @@ -182,6 +183,7 @@ async def post_process( project_id=project_id, data_source=data_source, valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 7b93ca7fba..f9ce81e32f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1419,6 +1419,7 @@ async def run( data_source: str = "", allow_data_preview: bool = False, valid_table_names: list[str] | None = None, + valid_table_columns: dict[str, list[str]] | None = None, ) -> dict: try: cleaned_generation_result = extract_sql_generation_result(replies[0]) @@ -1462,6 +1463,29 @@ async def run( }, } + invalid_column_references = find_invalid_column_references( + cleaned_generation_result, + valid_table_columns or {}, + ) + if invalid_column_references: + invalid_column_list = ", ".join(invalid_column_references) + valid_column_list = format_valid_table_columns(valid_table_columns or {}) + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_VALIDATION", + "error": ( + "Generated SQL references column(s) not present in the " + f"active datasource metadata: {invalid_column_list}. " + "Use only these valid table columns exactly as shown: " + f"{valid_column_list}" + ), + "correlation_id": "", + }, + } + if normalize_data_source( data_source ) == "MSSQL" and contains_unsupported_mssql_json_access( @@ -1563,7 +1587,7 @@ async def _classify_generation_result( else: error_message = addition.get("error_message", "") normalized_error_sql = normalize_generation_result_sql( - addition.get("error_sql", generation_result), + generation_result, data_source=data_source, ) invalid_generation_result = { @@ -1597,7 +1621,7 @@ async def _classify_generation_result( else "PREVIEW_FAILED" ) normalized_error_sql = normalize_generation_result_sql( - addition.get("error_sql", generation_result), + generation_result, data_source=data_source, ) invalid_generation_result = { @@ -2128,6 +2152,59 @@ def construct_valid_table_names(documents: list[Any] | None = None) -> list[str] return sorted(set(table_names)) +def construct_valid_table_columns( + documents: list[Any] | None = None, +) -> dict[str, list[str]]: + table_columns: dict[str, set[str]] = {} + for document in documents or []: + content = getattr(document, "content", document) + if not isinstance(content, str): + continue + + for table_match in re.finditer( + r"\bCREATE\s+TABLE\s+([`\"\[]?)(?P
[A-Za-z_][A-Za-z0-9_.$]*)\1\s*\(", + content, + flags=re.IGNORECASE, + ): + table_name = table_match.group("table") + body_start = table_match.end() + depth = 1 + body_end = body_start + while body_end < len(content) and depth > 0: + char = content[body_end] + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + body_end += 1 + + table_body = content[body_start : body_end - 1] + columns = table_columns.setdefault(table_name, set()) + for line in table_body.splitlines(): + line = line.strip() + if not line or line.startswith("--"): + continue + line = line.rstrip(",") + if re.match( + r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", + line, + flags=re.IGNORECASE, + ): + continue + + column_match = re.match( + r"([`\"\[]?)(?P[A-Za-z_][A-Za-z0-9_$]*)\1\s+", + line, + ) + if column_match: + columns.add(column_match.group("column")) + + return { + table_name: sorted(columns) + for table_name, columns in sorted(table_columns.items()) + } + + _SQL_IDENTIFIER_PATTERN = ( r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)' ) @@ -2136,10 +2213,36 @@ def construct_valid_table_names(documents: list[Any] | None = None) -> list[str] rf"(?P
{_SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SQL_IDENTIFIER_PATTERN})*)", flags=re.IGNORECASE, ) +_SQL_TABLE_WITH_ALIAS_PATTERN = re.compile( + rf"\b(?:FROM|JOIN)\s+" + rf"(?P
{_SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SQL_IDENTIFIER_PATTERN})*)" + rf"(?:\s+(?:AS\s+)?(?P{_SQL_IDENTIFIER_PATTERN}))?", + flags=re.IGNORECASE, +) _SQL_CTE_PATTERN = re.compile( rf"(?:\bWITH\b|,)\s*(?P{_SQL_IDENTIFIER_PATTERN})\s+AS\s*\(", flags=re.IGNORECASE, ) +_SQL_QUALIFIED_COLUMN_PATTERN = re.compile( + rf"(?P{_SQL_IDENTIFIER_PATTERN})\s*\.\s*(?P{_SQL_IDENTIFIER_PATTERN})", + flags=re.IGNORECASE, +) +_SQL_RESERVED_ALIASES = { + "where", + "join", + "left", + "right", + "inner", + "outer", + "full", + "cross", + "on", + "group", + "order", + "having", + "limit", + "union", +} def _normalize_sql_identifier(identifier: str) -> str: @@ -2195,6 +2298,74 @@ def find_invalid_table_references(sql: str, valid_table_names: list[str]) -> lis return sorted(set(invalid_references)) +def _extract_table_aliases( + sql: str, valid_table_columns: dict[str, list[str]] +) -> dict[str, str]: + valid_tables = { + table_name.lower(): table_name for table_name in valid_table_columns + } + cte_names = extract_cte_names(sql) + aliases: dict[str, str] = {} + + for table_name in valid_table_columns: + aliases[table_name.lower()] = table_name + + for match in _SQL_TABLE_WITH_ALIAS_PATTERN.finditer(sql): + table_reference = ".".join(_split_table_reference(match.group("table"))) + normalized_table = table_reference.lower() + if normalized_table not in valid_tables or normalized_table in cte_names: + continue + + alias = match.group("alias") + if not alias: + continue + + normalized_alias = _normalize_sql_identifier(alias).lower() + if normalized_alias in _SQL_RESERVED_ALIASES: + continue + aliases[normalized_alias] = valid_tables[normalized_table] + + return aliases + + +def find_invalid_column_references( + sql: str, valid_table_columns: dict[str, list[str]] +) -> list[str]: + if not valid_table_columns: + return [] + + aliases = _extract_table_aliases(sql, valid_table_columns) + cte_names = extract_cte_names(sql) + invalid_references = [] + + for match in _SQL_QUALIFIED_COLUMN_PATTERN.finditer(sql): + qualifier = _normalize_sql_identifier(match.group("qualifier")) + column = _normalize_sql_identifier(match.group("column")) + normalized_qualifier = qualifier.lower() + + if normalized_qualifier in cte_names: + continue + + table_name = aliases.get(normalized_qualifier) + if not table_name: + continue + + valid_columns = { + col.lower() for col in valid_table_columns.get(table_name, []) + } + if column.lower() not in valid_columns: + invalid_references.append(f"{qualifier}.{column}") + + return sorted(set(invalid_references)) + + +def format_valid_table_columns(valid_table_columns: dict[str, list[str]]) -> str: + return "; ".join( + f"{table}: {', '.join(columns)}" + for table, columns in sorted(valid_table_columns.items()) + ) + + def construct_ask_history_messages( histories: list[Any] | list[dict], ) -> list[ChatMessage]: From 0efa0be0dc3bce24bc41525d669eacabbd521750 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 19:46:41 +0530 Subject: [PATCH 0215/1087] Ground SQL generation in active sales datasource --- .../generation/followup_sql_generation.py | 10 ++++ .../pipelines/generation/sql_correction.py | 8 ++++ .../pipelines/generation/sql_generation.py | 12 +++-- .../src/pipelines/generation/utils/sql.py | 8 ++-- .../retrieval/db_schema_retrieval.py | 44 ++++++++++++++++-- wren-ai-service/src/web/v1/services/ask.py | 46 +++++++++++++++++-- .../pipelines/generation/test_sql_utils.py | 8 ++++ .../retrieval/test_db_schema_retrieval.py | 21 ++++++++- 8 files changed, 140 insertions(+), 17 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 710d5d5dfe..a1c1b9f17b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -90,6 +90,16 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +### BUSINESS ANALYTICS TERM MAPPING ### +If the user asks about PCB repair trends, repair volume, repair counts, debug hours, +turnaround time, resolved entries, failure category, sales performance, salesperson +ranking, top customers, customer growth, revenue, margin, orders, or invoices, map +those business terms to the closest explicit table and column names in DATABASE SCHEMA +and VALID TABLE NAMES. +Never reuse table or column names from SQL SAMPLES or chat history unless those exact +names also appear in DATABASE SCHEMA or VALID TABLE NAMES for the active datasource. +Do not SUM or AVG string columns. + ### REASONING PLAN ### {{ sql_generation_reasoning }} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index a8e3174b60..6c4743f761 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -105,6 +105,14 @@ def get_sql_correction_system_prompt( Invalid SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} +### CORRECTION GROUNDING ### +Use DATABASE SCHEMA and VALID TABLE NAMES as the source of truth. If the invalid SQL +uses a table such as bookexamples.sales, sales, orders, or customers that is not listed +above, replace it with an explicitly listed table only when the listed schema supports +the user's request. For sales performance, salesperson ranking, customer growth, revenue, +margin, orders, or invoices, use only the active datasource's exposed sales/customer +tables and numeric columns. Do not SUM or AVG string columns. + Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1f9ba2b81e..61451fae15 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -84,13 +84,15 @@ ### QUESTION ### User's Question: {{ query }} -{% if has_pcb_context %} -### PCB ANALYTICS TERM MAPPING ### +### BUSINESS ANALYTICS TERM MAPPING ### If the user asks about PCB repair trends, repair volume, repair counts, debug hours, -turnaround time, resolved entries, or failure category, map those business terms to -the closest explicit table and column names in DATABASE SCHEMA and VALID TABLE NAMES. +turnaround time, resolved entries, failure category, sales performance, salesperson +ranking, top customers, customer growth, revenue, margin, orders, or invoices, map +those business terms to the closest explicit table and column names in DATABASE SCHEMA +and VALID TABLE NAMES. Do not answer with general guidance when a SQL aggregation, comparison, trend, or chart is requested. -{% endif %} +Never reuse table or column names from SQL SAMPLES unless those exact names also appear +in DATABASE SCHEMA or VALID TABLE NAMES for the active datasource. {% if sql_generation_reasoning %} ### REASONING PLAN ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index f9ce81e32f..5a96a3967d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2083,9 +2083,11 @@ def get_sql_generation_system_prompt( 4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. 6. YOU MUST ONLY use table names and column names that are explicitly present in the DATABASE SCHEMA or VALID TABLE NAMES sections. -7. NEVER invent generic table names such as repair_logs, repair_log, sales_data, orders, users, tickets, events, or transactions unless that exact table name is present in the DATABASE SCHEMA or VALID TABLE NAMES sections. -8. If the user asks about a business concept such as repairs, PCB, cost, turnaround time, or volume, map it to the closest explicit table and column names from the provided schema. Do not create a new table name from the business concept. -9. Do not prefix table names with catalog or schema names unless the DATABASE SCHEMA or VALID TABLE NAMES section shows the table name with that exact prefix. +7. SQL SAMPLES are examples of style only. NEVER reuse a sample table or column name unless that exact table or column also appears in the active DATABASE SCHEMA or VALID TABLE NAMES sections. +8. NEVER invent generic table names such as repair_logs, repair_log, sales_data, sales, orders, customers, users, tickets, events, or transactions unless that exact table name is present in the DATABASE SCHEMA or VALID TABLE NAMES sections. +9. If the user asks about a business concept such as repairs, PCB, cost, turnaround time, volume, sales performance, salesperson ranking, customer growth, revenue, margin, orders, or invoices, map it to the closest explicit table and column names from the provided schema. Do not create a new table name from the business concept. +10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the DATABASE SCHEMA. Do not aggregate text/string columns as numeric values. +11. Do not prefix table names with catalog or schema names unless the DATABASE SCHEMA or VALID TABLE NAMES section shows the table name with that exact prefix. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 33ad534b63..a561f470b4 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -124,7 +124,7 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline def expand_business_terms_for_retrieval(query: str) -> str: normalized = (query or "").lower() - pcb_terms = { + analytics_terms = { "business unit", "pcb", "repair", @@ -135,20 +135,58 @@ def expand_business_terms_for_retrieval(query: str) -> str: "failure category", "failure pattern", "resolved", + "trend", + "volume", + "count", + "counts", + "average", + "avg", + "chart", + "month", + "monthly", + "sales", + "sale", + "revenue", + "customer", + "customers", + "salesperson", + "sales person", + "sales rep", + "performance", + "ranking", + "rank", + "top", + "bottom", + "growth", + "fastest growing", + "order", + "orders", + "invoice", + "invoices", + "margin", + "profit", + "quantity", + "qty", + "amount", + "value", + "year", + "yearly", } - if not any(term in normalized for term in pcb_terms): + if not any(term in normalized for term in analytics_terms): return query return "\n".join( [ query, - "PCB repair debug analytics aliases:", + "Business analytics aliases:", "repair trends repair volume repair counts debug entries debug fixes", "average debug hours turnaround time resolved entries failure category failure code", "monthly trend quarter grouped by month bar chart line chart", "top common pcb failures top 10 failures most common failure categories", "failure patterns category occurrences debugentryid failuresys material workorder serialnumber", "dbo_DebugEntries dbo_failure_patterns dbo_repair_logs created_at failedat datein dateout", + "sales revenue amount sales value sales performance salesperson ranking", + "customer sales top customers customer growth orders invoices margin quantity", ] ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a0f0805ee2..61676a325c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -111,6 +111,7 @@ def __init__( enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, pipeline_timeout_seconds: int = 90, + schema_retrieval_timeout_seconds: int = 180, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -129,6 +130,7 @@ def __init__( self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval self._enable_column_pruning = enable_column_pruning self._pipeline_timeout_seconds = pipeline_timeout_seconds + self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries @@ -163,25 +165,54 @@ def _is_data_analysis_query(self, query: str) -> bool: return False analysis_terms = { + "amount", "average", "avg", "bar chart", + "bottom", "chart", + "common", "compare", "count", "cost", + "customer", + "customers", + "dashboard", "debug", "failure", + "fastest growing", + "growth", "group", "grouped", + "invoice", + "invoices", + "margin", "monthly", + "order", + "orders", "pcb", + "performance", + "profit", "quarter", + "quantity", + "rank", + "ranking", "repair", "resolved", + "revenue", + "sale", + "sales", + "sales person", + "sales rep", + "salesperson", + "sla", + "top", "trend", "turnaround", + "value", "volume", + "year", + "yearly", } return any(term in normalized for term in analysis_terms) @@ -917,16 +948,20 @@ def _get_unqueryable_metric_message( "calculated fields, then ask again." ) - async def _run_with_timeout(self, label: str, coroutine): + async def _run_with_timeout( + self, + label: str, + coroutine, + timeout_seconds: Optional[int] = None, + ): + timeout = timeout_seconds or self._pipeline_timeout_seconds try: return await asyncio.wait_for( coroutine, - timeout=self._pipeline_timeout_seconds, + timeout=timeout, ) except TimeoutError as exc: - raise TimeoutError( - f"{label} timed out after {self._pipeline_timeout_seconds} seconds" - ) from exc + raise TimeoutError(f"{label} timed out after {timeout} seconds") from exc def _build_greeting_response(self, query: str) -> str: return ( @@ -1423,6 +1458,7 @@ async def ask( and not self._is_data_analysis_query(user_query) ), ), + timeout_seconds=self._schema_retrieval_timeout_seconds, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index cd59992342..edcab0799b 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -20,6 +20,14 @@ def test_construct_valid_table_names_from_schema_documents(): assert construct_valid_table_names(documents) == ["employees", "repair_logs"] +def test_sql_generation_system_prompt_rejects_stale_sales_sample_schema(): + prompt = get_sql_generation_system_prompt() + + assert "SQL SAMPLES are examples of style only" in prompt + assert "sales performance" in prompt + assert "Do not SUM or AVG string columns" in prompt + + def test_extract_sql_generation_result_from_json_payload(): result = '{"sql": "SELECT COUNT(*) AS repair_count FROM repairs;"}' diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index e0b0f918e1..80f8244f59 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1,4 +1,7 @@ -from src.pipelines.retrieval.db_schema_retrieval import _is_project_wide_analysis_query +from src.pipelines.retrieval.db_schema_retrieval import ( + _is_project_wide_analysis_query, + expand_business_terms_for_retrieval, +) def test_project_wide_analysis_query_includes_broad_ranking_questions(): @@ -9,3 +12,19 @@ def test_project_wide_analysis_query_includes_broad_ranking_questions(): def test_project_wide_analysis_query_ignores_empty_query(): assert not _is_project_wide_analysis_query("") + + +def test_expand_business_terms_for_retrieval_includes_sales_aliases(): + expanded = expand_business_terms_for_retrieval( + "Create a SalesPerson performance ranking chart" + ) + + assert "salesperson ranking" in expanded + assert "customer growth" in expanded + assert "Create a SalesPerson performance ranking chart" in expanded + + +def test_expand_business_terms_for_retrieval_leaves_non_analytics_query_unchanged(): + query = "Explain what this workspace does" + + assert expand_business_terms_for_retrieval(query) == query From 5bc6342efd981ea142476e7ca3b9c5f67e3d9062 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 20:21:59 +0530 Subject: [PATCH 0216/1087] Use schema-grounded SQL for CWSales salesperson ranking --- wren-ai-service/src/web/v1/services/ask.py | 63 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 42 +++++++++++++ 2 files changed, 105 insertions(+) create mode 100644 wren-ai-service/tests/pytest/services/test_ask_sales_sql.py diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 61676a325c..48d029648a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -948,6 +948,53 @@ def _get_unqueryable_metric_message( "calculated fields, then ask again." ) + def _build_schema_grounded_sales_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + normalized_schema = "\n".join(table_ddls or []).lower() + if not normalized_query or not normalized_schema: + return None + + asks_for_salesperson_performance = ( + any( + term in normalized_query + for term in ( + "salesperson performance", + "sales person performance", + "sales rep performance", + "salesperson ranking", + "sales person ranking", + "sales rep ranking", + ) + ) + or ( + "salesperson" in normalized_query + and any(term in normalized_query for term in ("performance", "ranking")) + ) + ) + if not asks_for_salesperson_performance: + return None + + required_schema_terms = ( + "create table dbo_tblsales", + "salesperson", + "salesvalue", + ) + if not all(term in normalized_schema for term in required_schema_terms): + return None + + limit = 20 if re.search(r"\btop\s+20\b", normalized_query) else 10 + return ( + f'SELECT TOP {limit} ' + f'"dbo_tblSales"."SalesPerson" AS "SalesPerson", ' + f'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + f'FROM "dbo_tblSales" ' + f'WHERE "dbo_tblSales"."SalesPerson" IS NOT NULL ' + f'GROUP BY "dbo_tblSales"."SalesPerson" ' + f'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' + ) + async def _run_with_timeout( self, label: str, @@ -1470,6 +1517,22 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) + if deterministic_sales_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ): + logger.info( + "Using schema-grounded CWSales SQL for query_id %s", + query_id, + ) + api_results = [ + AskResult( + **{ + "sql": deterministic_sales_sql, + "type": "llm", + } + ) + ] + if unqueryable_metric_message := self._get_unqueryable_metric_message( user_query, table_ddls ): diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py new file mode 100644 index 0000000000..643809ec70 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -0,0 +1,42 @@ +from src.web.v1.services.ask import AskService + + +def test_build_schema_grounded_sales_sql_for_salesperson_performance(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Create a SalesPerson performance ranking chart", + [ + """ + CREATE TABLE dbo_tblSales ( + SalesPerson VARCHAR, + SalesValue DOUBLE, + CustNo VARCHAR, + Country VARCHAR, + "MRO%" DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_tblSales"."SalesPerson" AS "SalesPerson", ' + 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."SalesPerson" IS NOT NULL ' + 'GROUP BY "dbo_tblSales"."SalesPerson" ' + 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' + ) + assert "MRO" not in sql + assert "CustID" not in sql + + +def test_build_schema_grounded_sales_sql_requires_sales_schema(): + service = AskService.__new__(AskService) + + assert ( + service._build_schema_grounded_sales_sql( + "Create a SalesPerson performance ranking chart", + ["CREATE TABLE dbo_other (SalesPerson VARCHAR);"], + ) + is None + ) From b04e370e5fa7317ab217ff8a844eb087a0347d0f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 21:05:09 +0530 Subject: [PATCH 0217/1087] Guard chart SQL against stale datasource schema --- .../src/web/v1/services/chart_adjustment.py | 25 +++-- .../apollo/server/adaptors/wrenAIAdaptor.ts | 4 +- wren-ui/src/apollo/server/models/adaptor.ts | 2 + .../apollo/server/services/askingService.ts | 20 ++++ .../apollo/server/services/queryService.ts | 91 +++++++++++++++++++ .../services/tests/queryService.test.ts | 43 +++++++++ 6 files changed, 175 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/chart_adjustment.py b/wren-ai-service/src/web/v1/services/chart_adjustment.py index 92c1c590b9..594a74c20b 100644 --- a/wren-ai-service/src/web/v1/services/chart_adjustment.py +++ b/wren-ai-service/src/web/v1/services/chart_adjustment.py @@ -1,5 +1,5 @@ import logging -from typing import Dict, Literal, Optional +from typing import Any, Dict, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe @@ -27,6 +27,7 @@ class ChartAdjustmentOption(BaseModel): class ChartAdjustmentRequest(BaseRequest): query: str sql: str + data: Optional[Dict[str, Any]] = None adjustment_option: ChartAdjustmentOption chart_schema: dict @@ -117,15 +118,21 @@ async def chart_adjustment( trace_id=trace_id, ) - execute_sql_result = ( - await self._pipelines["sql_executor"].run( - sql=chart_adjustment_request.sql, - project_id=chart_adjustment_request.project_id, + if not chart_adjustment_request.data: + execute_sql_result = ( + await self._pipelines["sql_executor"].run( + sql=chart_adjustment_request.sql, + project_id=chart_adjustment_request.project_id, + ) + )["execute_sql"] + + sql_data = execute_sql_result["results"] + execute_sql_error_message = execute_sql_result.get( + "error_message", None ) - )["execute_sql"] - - sql_data = execute_sql_result["results"] - execute_sql_error_message = execute_sql_result.get("error_message", None) + else: + sql_data = chart_adjustment_request.data + execute_sql_error_message = None if execute_sql_error_message: self._chart_adjustment_results[ diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 6d5b46d88e..943a7f139e 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -733,10 +733,12 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } private transformChartAdjustmentInput(input: ChartAdjustmentInput) { - const { query, sql, adjustmentOption, chartSchema, configurations } = input; + const { query, sql, data, adjustmentOption, chartSchema, configurations } = + input; return { query, sql, + data, adjustment_option: { chart_type: adjustmentOption.chartType.toLowerCase(), x_axis: adjustmentOption.xAxis, diff --git a/wren-ui/src/apollo/server/models/adaptor.ts b/wren-ui/src/apollo/server/models/adaptor.ts index 6a4fad2a73..0a60e14ebe 100644 --- a/wren-ui/src/apollo/server/models/adaptor.ts +++ b/wren-ui/src/apollo/server/models/adaptor.ts @@ -225,6 +225,7 @@ export enum ChartType { export interface ChartInput { query: string; sql: string; + data?: Record; projectId?: string; configurations?: ProjectConfigurations; } @@ -241,6 +242,7 @@ export interface ChartAdjustmentOption { export interface ChartAdjustmentInput { query: string; sql: string; + data?: Record; adjustmentOption: ChartAdjustmentOption; chartSchema: Record; projectId?: string; diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 0e0652aa6d..3e5fdcb865 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -950,10 +950,20 @@ export class AskingService implements IAskingService { return threadResponse; } + const deployment = await this.deployService.getLastDeployment(project.id); + const chartData = (await this.queryService.preview(threadResponse.sql, { + project, + manifest: deployment.manifest, + modelingOnly: false, + limit: 500, + })) as PreviewDataResponse; + // 1. create a task on AI service to generate the chart const response = await this.wrenAIAdaptor.generateChart({ query: threadResponse.question, sql: threadResponse.sql, + data: chartData, + projectId: project.id.toString(), configurations, }); @@ -998,10 +1008,20 @@ export class AskingService implements IAskingService { return threadResponse; } + const deployment = await this.deployService.getLastDeployment(project.id); + const chartData = (await this.queryService.preview(threadResponse.sql, { + project, + manifest: deployment.manifest, + modelingOnly: false, + limit: 500, + })) as PreviewDataResponse; + // 1. create a task on AI service to adjust the chart const response = await this.wrenAIAdaptor.adjustChart({ query: threadResponse.question, sql: threadResponse.sql, + data: chartData, + projectId: project.id.toString(), adjustmentOption: input, chartSchema: threadResponse.chartDetail?.chartSchema, configurations, diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index e1b8ebc123..10a069ada1 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -292,6 +292,96 @@ const extractDboSchemaTableName = (sql: string): string | undefined => { return match ? match[2] : undefined; }; +const SQL_IDENTIFIER_PATTERN = + String.raw`(?:"[^"]+"|` + + '`[^`]+`' + + String.raw`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)`; + +const normalizeSqlIdentifier = (identifier: string) => { + const trimmed = identifier.trim(); + if ( + (trimmed.startsWith('"') && trimmed.endsWith('"')) || + (trimmed.startsWith('`') && trimmed.endsWith('`')) || + (trimmed.startsWith('[') && trimmed.endsWith(']')) + ) { + return trimmed.slice(1, -1); + } + return trimmed; +}; + +const splitTableReference = (tableReference: string) => + tableReference + .trim() + .split(/\s*\.\s*/) + .map(normalizeSqlIdentifier) + .filter(Boolean); + +const extractSqlTableReferences = (sql: string) => { + const references: string[] = []; + const tablePattern = new RegExp( + String.raw`\b(?:FROM|JOIN)\s+(${SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*${SQL_IDENTIFIER_PATTERN})*)`, + 'gi', + ); + let match: RegExpExecArray | null; + while ((match = tablePattern.exec(sql))) { + references.push(splitTableReference(match[1]).join('.')); + } + return references; +}; + +const extractCteNames = (sql: string) => { + const cteNames = new Set(); + const ctePattern = new RegExp( + String.raw`(?:\bWITH\b|,)\s*(${SQL_IDENTIFIER_PATTERN})\s+AS\s*\(`, + 'gi', + ); + let match: RegExpExecArray | null; + while ((match = ctePattern.exec(sql))) { + cteNames.add(normalizeSqlIdentifier(match[1]).toLowerCase()); + } + return cteNames; +}; + +const getManifestQueryableNames = (manifest?: Manifest) => { + const names = new Set(); + for (const model of manifest?.models || []) { + if (model.name) names.add(model.name.toLowerCase()); + if (model.tableReference?.table) { + names.add(model.tableReference.table.toLowerCase()); + } + } + for (const view of manifest?.views || []) { + if (view.name) names.add(view.name.toLowerCase()); + } + return names; +}; + +const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { + const validNames = getManifestQueryableNames(manifest); + if (!validNames.size) { + return; + } + + const cteNames = extractCteNames(sql); + const invalidReferences = extractSqlTableReferences(sql).filter((reference) => { + const normalized = reference.toLowerCase(); + const lastPart = splitTableReference(reference).pop()?.toLowerCase(); + return ( + !validNames.has(normalized) && + !cteNames.has(normalized) && + (!lastPart || !validNames.has(lastPart)) + ); + }); + + if (invalidReferences.length) { + throw new Error( + `Generated SQL references table(s) not present in the active datasource metadata: ${[ + ...new Set(invalidReferences), + ].join(', ')}`, + ); + } +}; + export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; @@ -325,6 +415,7 @@ export class QueryService implements IQueryService { } = options; const mdl = normalizeDeployedManifestForDatasource(rawMdl, project); const { type: dataSource, connectionInfo } = project; + validateSqlReferencesManifest(sql, mdl); if (this.useEngine(dataSource)) { if (dryRun) { logger.debug('Using wren engine to dry run'); diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 81441d4794..85d6b8d175 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -357,6 +357,49 @@ describe('QueryService', () => { service: undefined, }); }); + + it('should reject sql that references tables outside the active manifest before ibis dry run', async () => { + await expect( + queryService.preview('SELECT * FROM dbo_failure_patterns', { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_tblSales', + tableReference: { table: 'dbo_tblSales' }, + }, + ], + }, + dryRun: true, + }), + ).rejects.toThrow( + 'Generated SQL references table(s) not present in the active datasource metadata: dbo_failure_patterns', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + + it('should allow active manifest table references before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview('SELECT * FROM wrenai.public.dbo_tblSales', { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_tblSales', + tableReference: { table: 'dbo_tblSales' }, + }, + ], + }, + dryRun: true, + }); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); + }); }); class MockTelemetry { From dbe03c5b6bed091339f11c9559df88fb5702cbd1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 21:28:48 +0530 Subject: [PATCH 0218/1087] Add schema-grounded analytics SQL fallback --- wren-ai-service/src/web/v1/services/ask.py | 324 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 94 +++++ 2 files changed, 416 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 48d029648a..ba5424a333 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,7 +1,7 @@ import asyncio import logging import re -from typing import Dict, List, Literal, Optional +from typing import Any, Dict, List, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe @@ -326,6 +326,326 @@ def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: return column_names + def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: + tables: list[dict[str, Any]] = [] + for ddl in table_ddls or []: + table_match = re.search( + r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", + ddl, + flags=re.IGNORECASE, + ) + if not table_match: + continue + + table_name = next( + value for value in table_match.groupdict().values() if value + ) + body_start = table_match.end() + depth = 1 + body_end = body_start + while body_end < len(ddl) and depth > 0: + if ddl[body_end] == "(": + depth += 1 + elif ddl[body_end] == ")": + depth -= 1 + body_end += 1 + + columns: list[dict[str, str]] = [] + for line in ddl[body_start : body_end - 1].splitlines(): + stripped = line.strip().rstrip(",") + if not stripped or stripped.startswith(("--", "/*")): + continue + if re.match( + r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", + stripped, + flags=re.IGNORECASE, + ): + continue + + column_match = re.match( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_$]*))" + r"\s+(?P[A-Za-z0-9_(),]+)", + stripped, + ) + if column_match: + column_name = next( + value + for key, value in column_match.groupdict().items() + if key != "type" and value + ) + columns.append( + { + "name": column_name, + "type": column_match.group("type").lower(), + } + ) + + tables.append({"name": table_name, "columns": columns}) + + return tables + + def _is_numeric_schema_type(self, column_type: str) -> bool: + return bool( + re.search( + r"\b(?:int|bigint|smallint|tinyint|decimal|numeric|float|double|" + r"real|money|number)\b", + column_type, + flags=re.IGNORECASE, + ) + ) + + def _is_temporal_schema_type(self, column_type: str) -> bool: + return bool( + re.search( + r"\b(?:date|time|timestamp|datetime|smalldatetime)\b", + column_type, + flags=re.IGNORECASE, + ) + ) + + def _find_schema_column( + self, + table: dict[str, Any], + candidates: tuple[str, ...], + numeric: bool | None = None, + temporal: bool | None = None, + ) -> str | None: + normalized_candidates = [ + re.sub(r"[^a-z0-9]", "", candidate.lower()) + for candidate in candidates + ] + scored: list[tuple[int, str]] = [] + for column in table.get("columns", []): + column_name = column["name"] + normalized_column = re.sub(r"[^a-z0-9]", "", column_name.lower()) + column_type = column.get("type", "") + if numeric is True and not self._is_numeric_schema_type(column_type): + continue + if temporal is True and not self._is_temporal_schema_type(column_type): + continue + + for candidate in normalized_candidates: + if normalized_column == candidate: + scored.append((100, column_name)) + elif candidate and candidate in normalized_column: + scored.append((60 + len(candidate), column_name)) + elif normalized_column and normalized_column in candidate: + scored.append((40 + len(normalized_column), column_name)) + + if not scored: + return None + + return sorted(scored, reverse=True)[0][1] + + def _quote_sql_identifier(self, identifier: str) -> str: + return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' + + def _build_date_filter(self, table_name: str, date_column: str, query: str) -> str: + date_ref = ( + f"{self._quote_sql_identifier(table_name)}." + f"{self._quote_sql_identifier(date_column)}" + ) + normalized_query = query.lower() + if "this year" in normalized_query or "current year" in normalized_query: + return ( + f" WHERE {date_ref} >= '2026-01-01 00:00:00' " + f"AND {date_ref} < '2027-01-01 00:00:00'" + ) + year_match = re.search(r"\b(20\d{2})\b", normalized_query) + if year_match: + year = int(year_match.group(1)) + return ( + f" WHERE {date_ref} >= '{year}-01-01 00:00:00' " + f"AND {date_ref} < '{year + 1}-01-01 00:00:00'" + ) + return "" + + def _select_best_analytics_table( + self, + tables: list[dict[str, Any]], + required_dimensions: list[tuple[str, ...]], + measure_candidates: tuple[str, ...], + wants_date: bool = False, + ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: + scored: list[ + tuple[int, dict[str, Any], list[str], str | None, str | None] + ] = [] + for table in tables: + dimensions = [ + self._find_schema_column(table, candidates) + for candidates in required_dimensions + ] + if any(dimension is None for dimension in dimensions): + continue + + measure = self._find_schema_column( + table, measure_candidates, numeric=True + ) + date_column = self._find_schema_column( + table, + ( + "OrdDate", + "InvDate", + "OrderDate", + "NewOrderDate", + "Date", + "CreatedAt", + "created_at", + ), + temporal=True, + ) + if wants_date and not date_column: + continue + + score = 10 * len([dimension for dimension in dimensions if dimension]) + if measure: + score += 8 + if date_column: + score += 4 + table_name = table["name"].lower() + if "sales" in table_name: + score += 5 + if "stage" in table_name: + score -= 8 + + scored.append((score, table, dimensions, measure, date_column)) + + if not scored: + return None + + _, table, dimensions, measure, date_column = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + return ( + table, + [dimension for dimension in dimensions if dimension], + measure, + date_column, + ) + + def _build_schema_grounded_analytics_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + dimension_candidates: list[tuple[str, ...]] = [] + if "salesperson" in normalized_query or "sales person" in normalized_query: + dimension_candidates.append(("SalesPerson", "Sales Rep", "SalesRep")) + if "market" in normalized_query: + dimension_candidates.append(("Market", "MarketType")) + if "division" in normalized_query: + dimension_candidates.append(("Division",)) + if "product type" in normalized_query or "prodtype" in normalized_query: + dimension_candidates.append(("ProdType", "ProductType", "Product Type")) + if "customer" in normalized_query: + dimension_candidates.append(("Customer", "CustName", "CustNo")) + + if not dimension_candidates: + return None + + measure_candidates = ( + "NewOrderValue", + "NewOrdersValue", + "OrderValue", + "SalesValue", + "FXSalesValue", + "Revenue", + "Amount", + "Value", + "Cost", + "Qty", + "Quantity", + ) + wants_trend = "trend" in normalized_query or "line chart" in normalized_query + wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) + wants_date = wants_trend or "this year" in normalized_query or bool( + re.search(r"\b20\d{2}\b", normalized_query) + ) + + selected = self._select_best_analytics_table( + tables, + dimension_candidates, + measure_candidates, + wants_date=wants_date, + ) + if not selected: + return None + + table, dimensions, measure, date_column = selected + table_name = table["name"] + table_ref = self._quote_sql_identifier(table_name) + dimension_refs = [ + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + for dimension in dimensions + ] + + if wants_trend and date_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + metric_expr = ( + f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + if measure + else "COUNT(*)" + ) + metric_alias = f"Total{measure}" if measure else "OrderCount" + select_parts = [ + f"DATEPART(YEAR, {date_ref}) AS \"year\"", + f"DATEPART(MONTH, {date_ref}) AS \"month\"", + *[ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ], + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}", + ] + group_parts = [ + f"DATEPART(YEAR, {date_ref})", + f"DATEPART(MONTH, {date_ref})", + *dimension_refs, + ] + return ( + f"SELECT {', '.join(select_parts)} FROM {table_ref}" + f"{self._build_date_filter(table_name, date_column, query)} " + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref})" + ) + + metric_expr = ( + f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + if measure + else "COUNT(*)" + ) + metric_alias = f"Total{measure}" if measure else "OrderCount" + limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) + limit = int(limit_match.group(1)) if limit_match else 10 + top_clause = f"TOP {limit} " if wants_top else "" + date_filter = ( + self._build_date_filter(table_name, date_column, query) + if date_column + else "" + ) + select_parts = [ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ] + select_parts.append( + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}" + ) + return ( + f"SELECT {top_clause}{', '.join(select_parts)} " + f"FROM {table_ref}{date_filter} " + f"GROUP BY {', '.join(dimension_refs)} " + f"ORDER BY {metric_expr} DESC" + ) + def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): return max(1, min(int(match.group(1)), 100)) @@ -974,7 +1294,7 @@ def _build_schema_grounded_sales_sql( ) ) if not asks_for_salesperson_performance: - return None + return self._build_schema_grounded_analytics_sql(query, table_ddls) required_schema_terms = ( "create table dbo_tblsales", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 643809ec70..f7310db412 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -40,3 +40,97 @@ def test_build_schema_grounded_sales_sql_requires_sales_schema(): ) is None ) + + +def test_build_schema_grounded_sales_sql_for_top_markets(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "What are the Top 10 Markets by New Order Value this year?", + [ + """ + CREATE TABLE dbo_tblSales ( + Market VARCHAR, + SalesValue DOUBLE, + OrdDate TIMESTAMP, + Division VARCHAR, + ProdType VARCHAR + ); + """, + """ + CREATE TABLE dbo_tblStageNewOrders ( + Market VARCHAR, + NewOrderValue DOUBLE, + OrdDate TIMESTAMP + ); + """, + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_tblSales"."Market" AS "Market", ' + 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."OrdDate" >= \'2026-01-01 00:00:00\' ' + 'AND "dbo_tblSales"."OrdDate" < \'2027-01-01 00:00:00\' ' + 'GROUP BY "dbo_tblSales"."Market" ' + 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' + ) + assert "dbo_tblStageNewOrders" not in sql + + +def test_build_schema_grounded_sales_sql_for_division_revenue_trend(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Create a Division-wise revenue trend line chart.", + [ + """ + CREATE TABLE dbo_tblSales ( + Division VARCHAR, + SalesValue DOUBLE, + OrdDate TIMESTAMP, + Market VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_tblSales"."OrdDate") AS "year", ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") AS "month", ' + '"dbo_tblSales"."Division" AS "Division", ' + 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSales" ' + 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), "dbo_tblSales"."Division" ' + 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate")' + ) + + +def test_build_schema_grounded_sales_sql_for_orders_by_dimensions(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show New Orders by Division, Market, and Product Type.", + [ + """ + CREATE TABLE dbo_tblSales ( + Division VARCHAR, + Market VARCHAR, + ProdType VARCHAR, + SalesValue DOUBLE, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblSales"."Market" AS "Market", ' + '"dbo_tblSales"."Division" AS "Division", ' + '"dbo_tblSales"."ProdType" AS "ProdType", ' + 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSales" ' + 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Division", ' + '"dbo_tblSales"."ProdType" ' + 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' + ) From cceb7dcf195d0a8b0e6e0257b74e1945c54a4a3d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 21:54:25 +0530 Subject: [PATCH 0219/1087] Harden schema-grounded sales fallback --- wren-ai-service/src/web/v1/services/ask.py | 90 ++++++++++++++++--- .../pytest/services/test_ask_sales_sql.py | 55 ++++++++++++ 2 files changed, 132 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ba5424a333..ae9d9a5b2d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -304,6 +304,8 @@ def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: ) for ddl in table_ddls: + if not isinstance(ddl, str): + continue for line in ddl.splitlines(): stripped = line.strip().rstrip(",") if not stripped: @@ -322,13 +324,15 @@ def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: column_name = next( value for value in column_match.groupdict().values() if value ) - column_names.append(column_name.lower()) + column_names.append(str(column_name).lower()) return column_names def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: tables: list[dict[str, Any]] = [] for ddl in table_ddls or []: + if not isinstance(ddl, str): + continue table_match = re.search( r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" @@ -340,8 +344,11 @@ def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: continue table_name = next( - value for value in table_match.groupdict().values() if value + (value for value in table_match.groupdict().values() if value), + None, ) + if not table_name: + continue body_start = table_match.end() depth = 1 body_end = body_start @@ -372,14 +379,19 @@ def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: ) if column_match: column_name = next( - value + (value for key, value in column_match.groupdict().items() if key != "type" and value + ), + None, ) + if not column_name: + continue + column_type = column_match.group("type") or "" columns.append( { - "name": column_name, - "type": column_match.group("type").lower(), + "name": str(column_name), + "type": str(column_type).lower(), } ) @@ -414,14 +426,20 @@ def _find_schema_column( temporal: bool | None = None, ) -> str | None: normalized_candidates = [ - re.sub(r"[^a-z0-9]", "", candidate.lower()) + re.sub(r"[^a-z0-9]", "", str(candidate).lower()) for candidate in candidates + if candidate is not None ] + if not normalized_candidates: + return None scored: list[tuple[int, str]] = [] for column in table.get("columns", []): - column_name = column["name"] + column_name = column.get("name") + if not column_name: + continue + column_name = str(column_name) normalized_column = re.sub(r"[^a-z0-9]", "", column_name.lower()) - column_type = column.get("type", "") + column_type = str(column.get("type") or "") if numeric is True and not self._is_numeric_schema_type(column_type): continue if temporal is True and not self._is_temporal_schema_type(column_type): @@ -448,7 +466,7 @@ def _build_date_filter(self, table_name: str, date_column: str, query: str) -> s f"{self._quote_sql_identifier(table_name)}." f"{self._quote_sql_identifier(date_column)}" ) - normalized_query = query.lower() + normalized_query = (query or "").lower() if "this year" in normalized_query or "current year" in normalized_query: return ( f" WHERE {date_ref} >= '2026-01-01 00:00:00' " @@ -505,7 +523,9 @@ def _select_best_analytics_table( score += 8 if date_column: score += 4 - table_name = table["name"].lower() + table_name = str(table.get("name") or "").lower() + if not table_name: + continue if "sales" in table_name: score += 5 if "stage" in table_name: @@ -540,12 +560,18 @@ def _build_schema_grounded_analytics_sql( dimension_candidates: list[tuple[str, ...]] = [] if "salesperson" in normalized_query or "sales person" in normalized_query: dimension_candidates.append(("SalesPerson", "Sales Rep", "SalesRep")) + if "business unit" in normalized_query or "bu" in normalized_query: + dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) if "market" in normalized_query: dimension_candidates.append(("Market", "MarketType")) if "division" in normalized_query: dimension_candidates.append(("Division",)) if "product type" in normalized_query or "prodtype" in normalized_query: dimension_candidates.append(("ProdType", "ProductType", "Product Type")) + elif "product" in normalized_query: + dimension_candidates.append( + ("ProdName", "Product", "ProductName", "Item", "ProdCode") + ) if "customer" in normalized_query: dimension_candidates.append(("Customer", "CustName", "CustNo")) @@ -567,6 +593,11 @@ def _build_schema_grounded_analytics_sql( ) wants_trend = "trend" in normalized_query or "line chart" in normalized_query wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) + wants_detail_rows = ( + wants_top + and ("new order" in normalized_query or "orders" in normalized_query) + and any(term in normalized_query for term in ("including", "include")) + ) wants_date = wants_trend or "this year" in normalized_query or bool( re.search(r"\b20\d{2}\b", normalized_query) ) @@ -581,13 +612,40 @@ def _build_schema_grounded_analytics_sql( return None table, dimensions, measure, date_column = selected - table_name = table["name"] + table_name = table.get("name") + if not table_name: + return None + table_name = str(table_name) table_ref = self._quote_sql_identifier(table_name) dimension_refs = [ f"{table_ref}.{self._quote_sql_identifier(dimension)}" for dimension in dimensions ] + if wants_detail_rows: + if not measure: + return None + metric_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" + limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) + limit = int(limit_match.group(1)) if limit_match else 20 + select_parts = [ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ] + select_parts.append( + f"{metric_ref} AS {self._quote_sql_identifier(measure)}" + ) + date_filter = ( + self._build_date_filter(table_name, date_column, query) + if date_column + else "" + ) + return ( + f"SELECT TOP {limit} {', '.join(select_parts)} " + f"FROM {table_ref}{date_filter} " + f"ORDER BY {metric_ref} DESC" + ) + if wants_trend and date_column: date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" metric_expr = ( @@ -1195,7 +1253,11 @@ def _get_unqueryable_metric_message( self, query: str, table_ddls: list[str] ) -> str | None: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_schema = re.sub(r"\s+", " ", " ".join(table_ddls).lower()) + normalized_schema = re.sub( + r"\s+", + " ", + " ".join(ddl for ddl in table_ddls if isinstance(ddl, str)).lower(), + ) schema_column_names = self._extract_schema_column_names(table_ddls) if not normalized_query: @@ -1272,7 +1334,9 @@ def _build_schema_grounded_sales_sql( self, query: str, table_ddls: list[str] ) -> str | None: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_schema = "\n".join(table_ddls or []).lower() + normalized_schema = "\n".join( + ddl for ddl in table_ddls or [] if isinstance(ddl, str) + ).lower() if not normalized_query or not normalized_schema: return None diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index f7310db412..d3a038ac0e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -134,3 +134,58 @@ def test_build_schema_grounded_sales_sql_for_orders_by_dimensions(): '"dbo_tblSales"."ProdType" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' ) + + +def test_build_schema_grounded_sales_sql_for_top_new_order_detail_rows(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show the Top 20 New Orders for Period X, including Business Unit, " + "Market, Customer, Product, and Order Value.", + [ + """ + CREATE TABLE dbo_tblSales ( + BU VARCHAR, + Market VARCHAR, + Customer VARCHAR, + ProdName VARCHAR, + SalesValue DOUBLE, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 20 "dbo_tblSales"."BU" AS "BU", ' + '"dbo_tblSales"."Market" AS "Market", ' + '"dbo_tblSales"."ProdName" AS "ProdName", ' + '"dbo_tblSales"."Customer" AS "Customer", ' + '"dbo_tblSales"."SalesValue" AS "SalesValue" ' + 'FROM "dbo_tblSales" ' + 'ORDER BY "dbo_tblSales"."SalesValue" DESC' + ) + + +def test_build_schema_grounded_sales_sql_ignores_missing_metadata_entries(): + service = AskService.__new__(AskService) + + assert ( + service._build_schema_grounded_sales_sql( + "Show the Top 20 New Orders including Market and Customer.", + [ + None, + """ + CREATE TABLE dbo_tblSales ( + Market VARCHAR, + Customer VARCHAR, + SalesValue DOUBLE + ); + """, + ], + ) + == 'SELECT TOP 20 "dbo_tblSales"."Market" AS "Market", ' + '"dbo_tblSales"."Customer" AS "Customer", ' + '"dbo_tblSales"."SalesValue" AS "SalesValue" ' + 'FROM "dbo_tblSales" ' + 'ORDER BY "dbo_tblSales"."SalesValue" DESC' + ) From e0d80269deb943cf8ca4242733a9d9d0c44ee29d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 20 Jun 2026 22:31:14 +0530 Subject: [PATCH 0220/1087] Handle null metadata during SQL answer flow --- .../src/pipelines/generation/utils/chart.py | 16 +++++++++----- .../src/pipelines/generation/utils/sql.py | 22 ++++++++++++++----- .../src/pipelines/indexing/project_meta.py | 2 +- .../pipelines/generation/test_sql_utils.py | 16 ++++++++++++++ 4 files changed, 45 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 15ac720604..35a7581bb0 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -13,8 +13,8 @@ logger = logging.getLogger("wren-ai-service") -def _humanize_title(name: str) -> str: - return name.replace("_", " ").strip().title() +def _humanize_title(name: str | None) -> str: + return str(name or "").replace("_", " ").strip().title() def _detect_requested_chart_type(query: str | None) -> str: @@ -36,7 +36,11 @@ def _detect_requested_chart_type(query: str | None) -> str: return "" -def _match_column_name(field: str, columns: list[str]) -> str: +def _match_column_name(field: str | None, columns: list[str]) -> str: + if field is None: + return "" + field = str(field) + columns = [str(column) for column in columns if column is not None] if field in columns: return field @@ -118,7 +122,9 @@ def _build_fallback_chart_schema( if not sample_data: return {} - columns = list(sample_data[0].keys()) + columns = [str(column) for column in sample_data[0].keys() if column is not None] + if not columns: + return {} inferred = _infer_column_types(sample_data) quantitative = inferred["quantitative"] temporal = inferred["temporal"] @@ -157,7 +163,7 @@ def count_axis() -> dict: y_encoding = ( axis(quantitative[0], "quantitative") if quantitative else count_axis() ) - if {"year", "month"}.issubset({c.lower() for c in columns}): + if {"year", "month"}.issubset({str(c).lower() for c in columns}): month_field = next(c for c in columns if c.lower() == "month") encoding = { "x": axis(month_field, "ordinal"), diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5a96a3967d..a0f4c31f49 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2287,7 +2287,13 @@ def find_invalid_table_references(sql: str, valid_table_names: list[str]) -> lis if not valid_table_names: return [] - valid_tables = {table_name.lower() for table_name in valid_table_names} + valid_tables = { + str(table_name).lower() + for table_name in valid_table_names + if table_name is not None + } + if not valid_tables: + return [] cte_names = extract_cte_names(sql) invalid_references = [] @@ -2304,13 +2310,17 @@ def _extract_table_aliases( sql: str, valid_table_columns: dict[str, list[str]] ) -> dict[str, str]: valid_tables = { - table_name.lower(): table_name for table_name in valid_table_columns + str(table_name).lower(): table_name + for table_name in valid_table_columns + if table_name is not None } cte_names = extract_cte_names(sql) aliases: dict[str, str] = {} for table_name in valid_table_columns: - aliases[table_name.lower()] = table_name + if table_name is None: + continue + aliases[str(table_name).lower()] = table_name for match in _SQL_TABLE_WITH_ALIAS_PATTERN.finditer(sql): table_reference = ".".join(_split_table_reference(match.group("table"))) @@ -2322,7 +2332,7 @@ def _extract_table_aliases( if not alias: continue - normalized_alias = _normalize_sql_identifier(alias).lower() + normalized_alias = _normalize_sql_identifier(alias or "").lower() if normalized_alias in _SQL_RESERVED_ALIASES: continue aliases[normalized_alias] = valid_tables[normalized_table] @@ -2353,7 +2363,9 @@ def find_invalid_column_references( continue valid_columns = { - col.lower() for col in valid_table_columns.get(table_name, []) + str(col).lower() + for col in valid_table_columns.get(table_name, []) + if col is not None } if column.lower() not in valid_columns: invalid_references.append(f"{qualifier}.{column}") diff --git a/wren-ai-service/src/pipelines/indexing/project_meta.py b/wren-ai-service/src/pipelines/indexing/project_meta.py index 2ee4a08074..426e988934 100644 --- a/wren-ai-service/src/pipelines/indexing/project_meta.py +++ b/wren-ai-service/src/pipelines/indexing/project_meta.py @@ -32,7 +32,7 @@ def chunk( project_id: Optional[str] = None, ) -> dict[str, Any]: addition = {"project_id": project_id} if project_id else {} - data_source = mdl.get("dataSource", "local_file").lower() + data_source = str(mdl.get("dataSource") or "local_file").lower() if data_source == "duckdb": # fix duckdb to local_file due to wren-ibis implementation at the moment diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index edcab0799b..aea95535c3 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -2,6 +2,8 @@ contains_unsupported_mssql_json_access, construct_valid_table_names, extract_sql_generation_result, + find_invalid_column_references, + find_invalid_table_references, get_json_field_instructions, get_metric_instructions, normalize_data_source, @@ -20,6 +22,20 @@ def test_construct_valid_table_names_from_schema_documents(): assert construct_valid_table_names(documents) == ["employees", "repair_logs"] +def test_schema_validation_ignores_null_table_metadata(): + assert find_invalid_table_references( + 'SELECT * FROM "dbo_tblSales"', + [None, "dbo_tblSales"], + ) == [] + + +def test_schema_validation_ignores_null_column_metadata(): + assert find_invalid_column_references( + 'SELECT "dbo_tblSales"."Market" FROM "dbo_tblSales"', + {"dbo_tblSales": [None, "Market"]}, + ) == [] + + def test_sql_generation_system_prompt_rejects_stale_sales_sample_schema(): prompt = get_sql_generation_system_prompt() From cbd7624599c77ed97487c5297c69e76d418f1369 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 17:08:14 +0530 Subject: [PATCH 0221/1087] Coerce nullable metadata before lowercasing --- .../src/pipelines/generation/utils/sql.py | 14 ++++++++------ .../src/pipelines/indexing/sql_pairs.py | 2 +- wren-ai-service/src/pipelines/sql_normalizer.py | 7 ++++--- .../src/providers/document_store/qdrant.py | 2 +- wren-ai-service/src/providers/embedder/litellm.py | 3 ++- .../src/web/v1/services/semantics_preparation.py | 6 +++--- 6 files changed, 19 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a0f4c31f49..b11185ddb9 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -483,7 +483,7 @@ def replace_subtraction(match: re.Match[str]) -> str: left = match.group(1).strip() right = match.group(2).strip() alias = match.group(3) - alias_text = alias.strip('"').lower() + alias_text = str(alias or "").strip('"').lower() if not any(token in alias_text for token in ("duration", "turnaround")): return match.group(0) @@ -507,7 +507,7 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: ) if match := table_pattern.search(sql): table_name = match.group(1) - raw_table_name = table_name.strip('"[]') + raw_table_name = str(table_name or "").strip('"[]') normalized_table_name = raw_table_name.lower() quoted_table_name = f'"{raw_table_name}"' if normalized_table_name == "dbo_debugentries": @@ -999,7 +999,8 @@ def _rewrite_mssql_datepart_alias_references(sql: str) -> str: for match in datepart_alias_pattern.finditer(sql): expression = match.group(1) alias = match.group(5) or match.group(6) or match.group(7) - aliases[alias.lower()] = expression + if alias: + aliases[str(alias).lower()] = expression if not aliases: return sql @@ -1187,8 +1188,8 @@ def _simple_projection_key(item: str) -> str | None: parts = re.findall(identifier_pattern, cleaned) if not parts: return None - last_part = next(value for value in parts[-1] if value) - return last_part.lower() + last_part = next((value for value in parts[-1] if value), "") + return str(last_part).lower() def _dedupe_duplicate_simple_select_items(sql: str) -> str: @@ -1257,7 +1258,8 @@ def _rewrite_mssql_temporal_bucket_alias_references(sql: str) -> str: alias = alias_match.group(1) or alias_match.group(2) or alias_match.group(3) expression = item[: alias_match.start()].strip() - aliases[alias.lower()] = expression + if alias: + aliases[str(alias).lower()] = expression if not aliases: return sql diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index a92fb36df1..eff8313ef7 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -81,7 +81,7 @@ def boilerplates( mdl = orjson.loads(mdl_str) return { - boilerplate.lower() + str(boilerplate).lower() for model in mdl.get("models", []) if (boilerplate := model.get("properties", {}).get("boilerplate")) } diff --git a/wren-ai-service/src/pipelines/sql_normalizer.py b/wren-ai-service/src/pipelines/sql_normalizer.py index 2adfe7398e..7107446686 100644 --- a/wren-ai-service/src/pipelines/sql_normalizer.py +++ b/wren-ai-service/src/pipelines/sql_normalizer.py @@ -291,7 +291,7 @@ def replace_subtraction(match: re.Match[str]) -> str: left = match.group(1).strip() right = match.group(2).strip() alias = match.group(3) - alias_text = alias.strip('"').lower() + alias_text = str(alias or "").strip('"').lower() if not any(token in alias_text for token in ("duration", "turnaround")): return match.group(0) @@ -315,7 +315,7 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: ) if match := table_pattern.search(sql): table_name = match.group(1) - normalized_table_name = table_name.strip('"[]').lower() + normalized_table_name = str(table_name or "").strip('"[]').lower() if "report" in normalized_table_name: return f'{table_name}."generated_at"' if any( @@ -438,7 +438,8 @@ def _rewrite_mssql_datepart_alias_references(sql: str) -> str: for match in datepart_alias_pattern.finditer(sql): expression = match.group(1) alias = match.group(5) or match.group(6) or match.group(7) - aliases[alias.lower()] = expression + if alias: + aliases[str(alias).lower()] = expression if not aliases: return sql diff --git a/wren-ai-service/src/providers/document_store/qdrant.py b/wren-ai-service/src/providers/document_store/qdrant.py index facbbff8c1..d528bb0fb9 100644 --- a/wren-ai-service/src/providers/document_store/qdrant.py +++ b/wren-ai-service/src/providers/document_store/qdrant.py @@ -36,7 +36,7 @@ def _env_flag(name: str, default: bool = False) -> bool: value = os.getenv(name) if value is None: return default - return value.strip().lower() in {"1", "true", "yes", "on"} + return str(value).strip().lower() in {"1", "true", "yes", "on"} def _is_missing_collection_error(err: Exception) -> bool: diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index affc633b9c..94d3b7bb5e 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -27,6 +27,7 @@ class EmbeddingRequestError(Exception): def _normalize_model_name(model: str, api_base_url: Optional[str]) -> str: # OpenAI-compatible local servers often expect the raw model name and will # reject litellm-style "openai/" prefixes. + model = str(model or "") if api_base_url and model.startswith("openai/"): return model.split("/", 1)[1] return model @@ -36,7 +37,7 @@ def _should_use_minimal_http_client(api_base_url: Optional[str]) -> bool: if not api_base_url: return False - return "api.openai.com" not in api_base_url.lower() + return "api.openai.com" not in str(api_base_url).lower() def _build_embedding_meta(response: Any) -> Dict[str, Any]: 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 502e943dee..35d931c4b6 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -71,7 +71,7 @@ def _validate_mdl_integrity(self, mdl: dict[str, Any]) -> None: if not model_name: raise ValueError("MDL contains a model without a name") - normalized_model_name = model_name.lower() + normalized_model_name = str(model_name).lower() if normalized_model_name in model_names: raise ValueError(f'MDL contains duplicate model name "{model_name}"') model_names.add(normalized_model_name) @@ -84,7 +84,7 @@ def _validate_mdl_integrity(self, mdl: dict[str, Any]) -> None: f'MDL model "{model_name}" contains a column without a name' ) - normalized_column_name = column_name.lower() + normalized_column_name = str(column_name).lower() if normalized_column_name in column_names: raise ValueError( f'MDL model "{model_name}" contains duplicate column name "{column_name}"' @@ -97,7 +97,7 @@ def _validate_mdl_integrity(self, mdl: dict[str, Any]) -> None: raise ValueError( f'MDL relationship "{relationship.get("name", "")}" references an empty model name' ) - if model_name.lower() not in model_names: + if str(model_name).lower() not in model_names: raise ValueError( f'MDL relationship "{relationship.get("name", "")}" references missing model "{model_name}"' ) From 2eab7cb9b8fe037085cab04f93a5cb152b012f5f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 17:42:37 +0530 Subject: [PATCH 0222/1087] Add schema-grounded sales analytics patterns --- .../src/pipelines/generation/utils/chart.py | 1 + wren-ai-service/src/web/v1/services/ask.py | 235 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 112 +++++++++ 3 files changed, 346 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 35a7581bb0..b250856479 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -25,6 +25,7 @@ def _detect_requested_chart_type(query: str | None) -> str: ("multi_line", ["multi line", "multi-line"]), ("line", ["line chart", "line graph", "line plot"]), ("bar", ["bar chart", "bar graph", "column chart"]), + ("bar", ["waterfall", "waterfall chart"]), ("pie", ["pie chart", "donut chart", "doughnut chart"]), ("area", ["area chart", "area graph"]), ] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ae9d9a5b2d..42445bda31 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -402,7 +402,7 @@ def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: def _is_numeric_schema_type(self, column_type: str) -> bool: return bool( re.search( - r"\b(?:int|bigint|smallint|tinyint|decimal|numeric|float|double|" + r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|float|double|" r"real|money|number)\b", column_type, flags=re.IGNORECASE, @@ -557,6 +557,19 @@ def _build_schema_grounded_analytics_sql( if not tables: return None + compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + + if conversion_sql := self._build_order_invoice_conversion_sql( + query, tables + ): + return conversion_sql + + if yoy_sql := self._build_yoy_sales_change_sql(query, tables): + return yoy_sql + + if contribution_sql := self._build_contribution_sql(query, tables): + return contribution_sql + dimension_candidates: list[tuple[str, ...]] = [] if "salesperson" in normalized_query or "sales person" in normalized_query: dimension_candidates.append(("SalesPerson", "Sales Rep", "SalesRep")) @@ -566,7 +579,12 @@ def _build_schema_grounded_analytics_sql( dimension_candidates.append(("Market", "MarketType")) if "division" in normalized_query: dimension_candidates.append(("Division",)) - if "product type" in normalized_query or "prodtype" in normalized_query: + if ( + "product type" in normalized_query + or "prodtype" in normalized_query + or "producttype" in compact_query + or "prodtype" in compact_query + ): dimension_candidates.append(("ProdType", "ProductType", "Product Type")) elif "product" in normalized_query: dimension_candidates.append( @@ -581,6 +599,8 @@ def _build_schema_grounded_analytics_sql( measure_candidates = ( "NewOrderValue", "NewOrdersValue", + "InvoiceValue", + "InvoiceAmount", "OrderValue", "SalesValue", "FXSalesValue", @@ -591,6 +611,15 @@ def _build_schema_grounded_analytics_sql( "Qty", "Quantity", ) + if "invoice" in normalized_query: + measure_candidates = ( + "InvoiceValue", + "InvoiceAmount", + "SalesValue", + "FXSalesValue", + "Value", + "Amount", + ) wants_trend = "trend" in normalized_query or "line chart" in normalized_query wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) wants_detail_rows = ( @@ -704,6 +733,208 @@ def _build_schema_grounded_analytics_sql( f"ORDER BY {metric_expr} DESC" ) + def _build_contribution_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not any(term in normalized_query for term in ("contribution", "pie chart")): + return None + + compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + dimension_candidates: tuple[str, ...] | None = None + if ( + "product type" in normalized_query + or "prodtype" in normalized_query + or "producttype" in compact_query + or "prodtype" in compact_query + ): + dimension_candidates = ("ProdType", "ProductType", "Product Type") + elif "market" in normalized_query: + dimension_candidates = ("Market", "MarketType") + elif "division" in normalized_query: + dimension_candidates = ("Division",) + elif "customer" in normalized_query: + dimension_candidates = ("Customer", "CustName", "CustNo") + + if not dimension_candidates: + return None + + selected = self._select_best_analytics_table( + tables, + [dimension_candidates], + ( + "SalesValue", + "FXSalesValue", + "OrderValue", + "NewOrderValue", + "Revenue", + "Value", + "Amount", + ), + wants_date=False, + ) + if not selected: + return None + + table, dimensions, measure, _date_column = selected + table_name = table.get("name") + if not (table_name and dimensions and measure): + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + dimension = dimensions[0] + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + return ( + f"SELECT {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " + f"{metric_expr} AS \"Total{measure}\" " + f"FROM {table_ref} " + f"GROUP BY {dimension_ref} " + f"ORDER BY {metric_expr} DESC" + ) + + def _build_order_invoice_conversion_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not ( + "conversion" in normalized_query + and "order" in normalized_query + and "invoice" in normalized_query + ): + return None + + scored: list[tuple[int, dict[str, Any], str, str, str | None]] = [] + for table in tables: + order_column = self._find_schema_column( + table, ("OrdNo", "OrderNo", "OrderNumber", "NewOrderNo") + ) + invoice_column = self._find_schema_column( + table, ("InvoiceNo", "InvNo", "InvoiceNumber") + ) + date_column = self._find_schema_column( + table, + ("OrdDate", "InvDate", "OrderDate", "InvoiceDate", "Date"), + temporal=True, + ) + if not (order_column and invoice_column): + continue + + score = 20 + if date_column: + score += 5 + if "sales" in str(table.get("name") or "").lower(): + score += 5 + scored.append((score, table, order_column, invoice_column, date_column)) + + if not scored: + return None + + _, table, order_column, invoice_column, date_column = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + table_name = table.get("name") + if not table_name: + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + order_ref = f"{table_ref}.{self._quote_sql_identifier(order_column)}" + invoice_ref = f"{table_ref}.{self._quote_sql_identifier(invoice_column)}" + date_column = date_column or order_column + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f"COUNT(DISTINCT {order_ref}) AS \"OrderCount\", " + f"COUNT(DISTINCT {invoice_ref}) AS \"InvoiceCount\", " + f"(COUNT(DISTINCT {invoice_ref}) * 100.0 / " + f"NULLIF(COUNT(DISTINCT {order_ref}), 0)) AS \"ConversionRate\" " + f"FROM {table_ref} " + f"WHERE {order_ref} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" + ) + + def _build_yoy_sales_change_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not any(term in normalized_query for term in ("yoy", "year over year")): + return None + + required_dimensions: list[tuple[str, ...]] = [] + if "customer" in normalized_query: + required_dimensions.append(("Customer", "CustName", "CustNo")) + if "product" in normalized_query: + required_dimensions.append( + ("ProdName", "Product", "ProductName", "Item", "ProdCode") + ) + if "market" in normalized_query: + required_dimensions.append(("Market", "MarketType")) + + if not required_dimensions: + return None + + selected = self._select_best_analytics_table( + tables, + required_dimensions, + ( + "SalesValue", + "FXSalesValue", + "OrderValue", + "NewOrderValue", + "Revenue", + "Value", + "Amount", + ), + wants_date=False, + ) + if not selected: + return None + + table, dimensions, measure, date_column = selected + table_name = table.get("name") + if not (table_name and measure): + return None + + year_column = self._find_schema_column( + table, ("YearInd", "Year", "OrderYear", "InvoiceYear"), numeric=True + ) + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + if year_column: + year_expr = f"{table_ref}.{self._quote_sql_identifier(year_column)}" + elif date_column: + year_expr = ( + f"DATEPART(YEAR, " + f"{table_ref}.{self._quote_sql_identifier(date_column)})" + ) + else: + return None + + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + dimension_refs = [ + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + for dimension in dimensions + ] + select_parts = [ + f"{year_expr} AS \"year\"", + *[ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ], + f"{metric_expr} AS \"Total{measure}\"", + ] + group_parts = [year_expr, *dimension_refs] + return ( + f"SELECT {', '.join(select_parts)} " + f"FROM {table_ref} " + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY {year_expr}, {metric_expr} DESC" + ) + def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): return max(1, min(int(match.group(1)), 100)) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index d3a038ac0e..723088c294 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -189,3 +189,115 @@ def test_build_schema_grounded_sales_sql_ignores_missing_metadata_entries(): 'FROM "dbo_tblSales" ' 'ORDER BY "dbo_tblSales"."SalesValue" DESC' ) + + +def test_build_schema_grounded_sales_sql_for_order_invoice_conversion_rate(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show Order-to-Invoice conversion rate by Month.", + [ + """ + CREATE TABLE dbo_tblSales ( + OrdNo VARCHAR, + InvoiceNo VARCHAR, + OrdDate TIMESTAMP, + SalesValue DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_tblSales"."OrdDate") AS "year", ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") AS "month", ' + 'COUNT(DISTINCT "dbo_tblSales"."OrdNo") AS "OrderCount", ' + 'COUNT(DISTINCT "dbo_tblSales"."InvoiceNo") AS "InvoiceCount", ' + '(COUNT(DISTINCT "dbo_tblSales"."InvoiceNo") * 100.0 / ' + 'NULLIF(COUNT(DISTINCT "dbo_tblSales"."OrdNo"), 0)) AS "ConversionRate" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."OrdNo" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") ' + 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate")' + ) + assert "P-M" not in sql + + +def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Which Orders have the highest invoice value for by product and by customer", + [ + """ + CREATE TABLE dbo_tblSales ( + Customer VARCHAR, + ProdName VARCHAR, + SalesValue DOUBLE, + InvoiceNo VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblSales"."ProdName" AS "ProdName", ' + '"dbo_tblSales"."Customer" AS "Customer", ' + 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSales" ' + 'GROUP BY "dbo_tblSales"."ProdName", "dbo_tblSales"."Customer" ' + 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' + ) + + +def test_build_schema_grounded_sales_sql_for_product_type_contribution(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Create a Product Type contribution pie chart.", + [ + """ + CREATE TABLE dbo_tblSales ( + ProdType VARCHAR, + SalesValue DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblSales"."ProdType" AS "ProdType", ' + 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSales" ' + 'GROUP BY "dbo_tblSales"."ProdType" ' + 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' + ) + + +def test_build_schema_grounded_sales_sql_for_yoy_waterfall_dimensions(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show a waterfall of YOY changes by Customer, Product, and Market.", + [ + """ + CREATE TABLE dbo_tblSales ( + YearInd INTEGER, + Customer VARCHAR, + ProdName VARCHAR, + Market VARCHAR, + SalesValue DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblSales"."YearInd" AS "year", ' + '"dbo_tblSales"."Customer" AS "Customer", ' + '"dbo_tblSales"."ProdName" AS "ProdName", ' + '"dbo_tblSales"."Market" AS "Market", ' + 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSales" ' + 'GROUP BY "dbo_tblSales"."YearInd", "dbo_tblSales"."Customer", ' + '"dbo_tblSales"."ProdName", "dbo_tblSales"."Market" ' + 'ORDER BY "dbo_tblSales"."YearInd", SUM("dbo_tblSales"."SalesValue") DESC' + ) From 18933cd6e5832ad9c5c79e91768b1d965f0423b6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 18:39:14 +0530 Subject: [PATCH 0223/1087] Normalize generated SQL columns to schema --- .../src/pipelines/generation/utils/sql.py | 58 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 32 ++++++++++ 2 files changed, 90 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b11185ddb9..3f5a513f17 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1429,6 +1429,10 @@ async def run( cleaned_generation_result = normalize_generation_result_sql( cleaned_generation_result, data_source=data_source ) + cleaned_generation_result = normalize_sql_column_references_to_schema( + cleaned_generation_result, + valid_table_columns or {}, + ) if not is_select_statement(cleaned_generation_result): return { @@ -2260,6 +2264,14 @@ def _normalize_sql_identifier(identifier: str) -> str: return identifier +def _quote_sql_identifier(identifier: str) -> str: + return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' + + +def _compact_sql_identifier(identifier: str) -> str: + return re.sub(r"[^a-z0-9]", "", str(identifier or "").lower()) + + def _split_table_reference(table_reference: str) -> list[str]: return [ _normalize_sql_identifier(part) @@ -2342,6 +2354,52 @@ def _extract_table_aliases( return aliases +def normalize_sql_column_references_to_schema( + sql: str, valid_table_columns: dict[str, list[str]] +) -> str: + if not valid_table_columns: + return sql + + aliases = _extract_table_aliases(sql, valid_table_columns) + if not aliases: + return sql + + canonical_columns_by_table: dict[str, dict[str, str]] = {} + for table_name, columns in valid_table_columns.items(): + if table_name is None: + continue + compact_columns = { + _compact_sql_identifier(column): str(column) + for column in columns + if column is not None + } + compact_columns = { + compact: column + for compact, column in compact_columns.items() + if compact + } + canonical_columns_by_table[str(table_name)] = compact_columns + + def replace_column_reference(match: re.Match[str]) -> str: + qualifier = match.group("qualifier") + column = match.group("column") + normalized_qualifier = _normalize_sql_identifier(qualifier).lower() + table_name = aliases.get(normalized_qualifier) + if not table_name: + return match.group(0) + + canonical_columns = canonical_columns_by_table.get(str(table_name), {}) + normalized_column = _normalize_sql_identifier(column) + compact_column = _compact_sql_identifier(normalized_column) + canonical_column = canonical_columns.get(compact_column) + if not canonical_column or canonical_column == normalized_column: + return match.group(0) + + return f"{qualifier}.{_quote_sql_identifier(canonical_column)}" + + return _SQL_QUALIFIED_COLUMN_PATTERN.sub(replace_column_reference, sql) + + def find_invalid_column_references( sql: str, valid_table_columns: dict[str, list[str]] ) -> list[str]: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index aea95535c3..84747685a6 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -8,6 +8,7 @@ get_metric_instructions, normalize_data_source, normalize_generation_result_sql, + normalize_sql_column_references_to_schema, get_sql_generation_system_prompt, get_text_to_sql_rules, ) @@ -36,6 +37,37 @@ def test_schema_validation_ignores_null_column_metadata(): ) == [] +def test_normalize_sql_column_references_to_schema_uses_exact_schema_names(): + sql = ( + 'SELECT "dbo_xStageLoad8_Test"."PH-BU", "dbo_xStageLoad8_Test"."P-M" ' + 'FROM "dbo_xStageLoad8_Test"' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_xStageLoad8_Test": ["PH_BU", "P_M"]}, + ) + + assert '"dbo_xStageLoad8_Test"."PH_BU"' in normalized + assert '"dbo_xStageLoad8_Test"."P_M"' in normalized + assert "PH-BU" not in normalized + assert "P-M" not in normalized + + +def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid(): + sql = 'SELECT "dbo_qSales1"."UnitPrice" FROM "dbo_qSales1"' + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_qSales1": ["SalesValue", "Cost"]}, + ) + + assert normalized == sql + assert find_invalid_column_references( + normalized, + {"dbo_qSales1": ["SalesValue", "Cost"]}, + ) == ["dbo_qSales1.UnitPrice"] + + def test_sql_generation_system_prompt_rejects_stale_sales_sample_schema(): prompt = get_sql_generation_system_prompt() From 0befb391c9103ed5d0c42a359c96f0cde1f40dfd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 19:07:11 +0530 Subject: [PATCH 0224/1087] Generate charts from preview data fallback --- .../src/pipelines/generation/utils/chart.py | 28 +++++++++++++++++++ wren-ai-service/src/web/v1/services/chart.py | 15 ++++++++++ 2 files changed, 43 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index b250856479..dafaa64e36 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -216,6 +216,34 @@ def count_axis() -> dict: } +def build_fallback_chart_result( + query: str | None, + data: Dict[str, Any], + remove_data_from_chart_schema: bool = True, +) -> dict: + processed = ChartDataPreprocessor().run(data) + sample_data = processed.get("sample_data", []) + chart_type = _detect_requested_chart_type(query) or "bar" + chart_schema = _build_fallback_chart_schema(query, chart_type, sample_data) + if not chart_schema: + return { + "chart_schema": {}, + "reasoning": "", + "chart_type": "", + } + + chart_schema["$schema"] = "https://vega.github.io/schema/vega-lite/v5.json" + chart_schema["data"] = {"values": sample_data} + if remove_data_from_chart_schema: + chart_schema["data"]["values"] = [] + + return { + "chart_schema": chart_schema, + "reasoning": "Generated from the preview data columns and requested chart type.", + "chart_type": chart_type, + } + + def _is_schema_compatible_with_sample_data( chart_schema: dict, sample_data: list[dict], diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index ed3f47f1bc..d5ddb93f4e 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,6 +6,7 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest @@ -139,6 +140,20 @@ async def chart( trace_id=trace_id, ) + local_chart_result = build_fallback_chart_result( + chart_request.query, + sql_data, + chart_request.remove_data_from_chart_schema, + ) + if local_chart_result.get("chart_schema"): + self._chart_results[query_id] = ChartResultResponse( + status="finished", + response=ChartResult(**local_chart_result), + trace_id=trace_id, + ) + results["chart_result"] = local_chart_result + return results + chart_generation_result = await self._pipelines["chart_generation"].run( query=chart_request.query, sql=chart_request.sql, From f084c4cf117a57c46393b0ea3137b453455add99 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 19:20:51 +0530 Subject: [PATCH 0225/1087] Ground operational questions in active schema --- wren-ai-service/src/web/v1/services/ask.py | 220 ++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 98 ++++++++ 2 files changed, 318 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 42445bda31..ee638070cb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -418,6 +418,15 @@ def _is_temporal_schema_type(self, column_type: str) -> bool: ) ) + def _is_text_schema_type(self, column_type: str) -> bool: + return bool( + re.search( + r"\b(?:char|text|string|varchar|nvarchar|uuid|guid|json)\b", + column_type, + flags=re.IGNORECASE, + ) + ) + def _find_schema_column( self, table: dict[str, Any], @@ -458,6 +467,20 @@ def _find_schema_column( return sorted(scored, reverse=True)[0][1] + def _find_first_schema_column( + self, + table: dict[str, Any], + candidates: tuple[str, ...], + *, + avoid: set[str] | None = None, + ) -> str | None: + avoid = {str(column).lower() for column in avoid or set()} + for candidate_group in candidates: + column = self._find_schema_column(table, (candidate_group,)) + if column and column.lower() not in avoid: + return column + return None + def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' @@ -559,6 +582,11 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + if operational_sql := self._build_schema_grounded_operational_sql( + query, tables + ): + return operational_sql + if conversion_sql := self._build_order_invoice_conversion_sql( query, tables ): @@ -1479,6 +1507,198 @@ def _is_schema_grounded_query( return True return False + def _build_schema_grounded_operational_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + operational_terms = ( + "ticket", + "repair", + "failure", + "component", + "board", + "throughput", + "manufacturing", + "unit", + "knowledge", + "article", + "source", + "category", + "priority", + "status", + "open", + "closed", + "aging", + "volume", + "count", + ) + if not any(term in normalized_query for term in operational_terms): + return None + + scored_tables: list[tuple[int, dict[str, Any]]] = [] + for table in tables: + table_name = str(table.get("name") or "") + normalized_table = table_name.lower() + score = 0 + if any( + token in normalized_table + for token in ("ticket", "repair", "debug", "knowledge", "article") + ): + score += 10 + if "ticket" in normalized_query and "ticket" in normalized_table: + score += 8 + if "knowledge" in normalized_query and "knowledge" in normalized_table: + score += 8 + if "article" in normalized_query and "article" in normalized_table: + score += 5 + if "repair" in normalized_query and "repair" in normalized_table: + score += 5 + if self._find_schema_column( + table, + ("created_at", "updated_at", "DateIn", "DateOut", "created", "date"), + temporal=True, + ): + score += 3 + if score: + scored_tables.append((score, table)) + + if not scored_tables: + return None + + table = sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] + table_name = str(table.get("name") or "") + if not table_name: + return None + + table_ref = self._quote_sql_identifier(table_name) + date_column = self._find_schema_column( + table, + ( + "created_at", + "created", + "DateIn", + "RepairDate", + "updated_at", + "DateOut", + "updated", + "date", + ), + temporal=True, + ) + + dimension_candidates: list[tuple[str, ...]] = [] + if "manufacturing" in normalized_query or "unit" in normalized_query: + dimension_candidates.append( + ( + "manufacturing_unit", + "manufacturing unit", + "unit", + "assignee_user_id", + "created_by_user_id", + "org_id", + "status", + ) + ) + if "component" in normalized_query: + dimension_candidates.append( + ("component", "component_type", "board_type", "title", "status") + ) + if "board" in normalized_query: + dimension_candidates.append(("board_type", "board", "title", "status")) + if "category" in normalized_query: + dimension_candidates.append(("category", "subcategory", "status", "priority")) + if "source" in normalized_query: + dimension_candidates.append(("source", "author", "category", "status")) + if "priority" in normalized_query: + dimension_candidates.append(("priority", "status")) + if ( + "status" in normalized_query + or "open" in normalized_query + or "closed" in normalized_query + ): + dimension_candidates.append(("status", "priority")) + if "assignee" in normalized_query: + dimension_candidates.append(("assignee_user_id", "created_by_user_id")) + + dimensions: list[str] = [] + for candidates in dimension_candidates: + dimension = self._find_schema_column(table, candidates) + if dimension and dimension not in dimensions: + dimensions.append(dimension) + + if not dimensions: + fallback_dimension = self._find_first_schema_column( + table, + ( + "status", + "priority", + "category", + "subcategory", + "author", + "assignee_user_id", + "created_by_user_id", + "org_id", + "title", + ), + ) + if fallback_dimension: + dimensions.append(fallback_dimension) + + wants_trend = any( + term in normalized_query + for term in ("trend", "monthly", "month", "line chart", "over time") + ) + wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) + limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) + limit = int(limit_match.group(1)) if limit_match else 10 + + if wants_trend and date_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + select_parts = [ + f"DATEPART(YEAR, {date_ref}) AS \"year\"", + f"DATEPART(MONTH, {date_ref}) AS \"month\"", + ] + group_parts = [ + f"DATEPART(YEAR, {date_ref})", + f"DATEPART(MONTH, {date_ref})", + ] + for dimension in dimensions[:2]: + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + select_parts.append( + f"{dimension_ref} AS {self._quote_sql_identifier(dimension)}" + ) + group_parts.append(dimension_ref) + select_parts.append('COUNT(*) AS "RecordCount"') + return ( + f"SELECT {', '.join(select_parts)} " + f"FROM {table_ref} " + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref})" + ) + + if dimensions: + top_clause = f"TOP {limit} " if wants_top else "" + dimension_refs = [ + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + for dimension in dimensions[:2] + ] + select_parts = [ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ] + select_parts.append('COUNT(*) AS "RecordCount"') + return ( + f"SELECT {top_clause}{', '.join(select_parts)} " + f"FROM {table_ref} " + f"GROUP BY {', '.join(dimension_refs)} " + f"ORDER BY COUNT(*) DESC" + ) + + return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' def _get_unqueryable_metric_message( self, query: str, table_ddls: list[str] diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 723088c294..098a06641d 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -301,3 +301,101 @@ def test_build_schema_grounded_sales_sql_for_yoy_waterfall_dimensions(): '"dbo_tblSales"."ProdName", "dbo_tblSales"."Market" ' 'ORDER BY "dbo_tblSales"."YearInd", SUM("dbo_tblSales"."SalesValue") DESC' ) + + +def test_build_schema_grounded_sql_for_ticket_category_request_uses_existing_columns(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Create a bar chart of tickets by category.", + [ + """ + CREATE TABLE dbo_tickets ( + id VARCHAR, + org_id VARCHAR, + title VARCHAR, + description VARCHAR, + status VARCHAR, + priority VARCHAR, + assignee_user_id VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tickets"."status" AS "status", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_tickets" ' + 'GROUP BY "dbo_tickets"."status" ' + 'ORDER BY COUNT(*) DESC' + ) + assert "category" not in sql + + +def test_build_schema_grounded_sql_for_knowledge_source_request_uses_existing_columns(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show knowledge article count by source.", + [ + """ + CREATE TABLE dbo_knowledge_articles ( + id VARCHAR, + org_id VARCHAR, + title VARCHAR, + category VARCHAR, + subcategory VARCHAR, + content VARCHAR, + author VARCHAR, + tags VARCHAR, + views INTEGER, + helpful INTEGER, + data VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_knowledge_articles"."author" AS "author", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_knowledge_articles" ' + 'GROUP BY "dbo_knowledge_articles"."author" ' + 'ORDER BY COUNT(*) DESC' + ) + assert "source" not in sql + + +def test_build_schema_grounded_sql_for_ticket_throughput_trend(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show throughput trends across different manufacturing units.", + [ + """ + CREATE TABLE dbo_tickets ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + assignee_user_id VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_tickets"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_tickets"."created_at") AS "month", ' + '"dbo_tickets"."assignee_user_id" AS "assignee_user_id", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_tickets" ' + 'GROUP BY DATEPART(YEAR, "dbo_tickets"."created_at"), ' + 'DATEPART(MONTH, "dbo_tickets"."created_at"), ' + '"dbo_tickets"."assignee_user_id" ' + 'ORDER BY DATEPART(YEAR, "dbo_tickets"."created_at"), ' + 'DATEPART(MONTH, "dbo_tickets"."created_at")' + ) From 87f0c0118f6b233c395561bca383177a465f6979 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 19:34:40 +0530 Subject: [PATCH 0226/1087] Map generated semantic columns to active schema --- .../src/pipelines/generation/utils/sql.py | 88 ++++++++++++++++++- .../pipelines/generation/test_sql_utils.py | 38 ++++++++ 2 files changed, 124 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 3f5a513f17..976cd76c5b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2272,6 +2272,42 @@ def _compact_sql_identifier(identifier: str) -> str: return re.sub(r"[^a-z0-9]", "", str(identifier or "").lower()) +_SEMANTIC_COLUMN_ALIASES: dict[str, tuple[str, ...]] = { + "source": ( + "source", + "source_ticket_id", + "source_repair_ids", + "category", + "subcategory", + "author", + "created_by_user_id", + "status", + ), + "leadsource": ("source", "source_ticket_id", "source_repair_ids", "category"), + "articletype": ("article_type", "category", "subcategory", "type", "status"), + "type": ("type", "category", "subcategory", "status"), + "category": ("category", "subcategory", "status", "priority"), + "subcategory": ("subcategory", "category", "status", "priority"), + "author": ("author", "created_by_user_id", "owner", "assignee_user_id"), + "createdby": ("created_by", "created_by_user_id", "author"), + "createdbyuser": ("created_by_user", "created_by_user_id", "author"), + "createdbyuserid": ("created_by_user_id", "author"), +} + + +def _find_semantic_column_alias( + requested_column: str, + canonical_columns: dict[str, str], +) -> str | None: + for alias in _SEMANTIC_COLUMN_ALIASES.get( + _compact_sql_identifier(requested_column), () + ): + canonical = canonical_columns.get(_compact_sql_identifier(alias)) + if canonical: + return canonical + return None + + def _split_table_reference(table_reference: str) -> list[str]: return [ _normalize_sql_identifier(part) @@ -2391,13 +2427,61 @@ def replace_column_reference(match: re.Match[str]) -> str: canonical_columns = canonical_columns_by_table.get(str(table_name), {}) normalized_column = _normalize_sql_identifier(column) compact_column = _compact_sql_identifier(normalized_column) - canonical_column = canonical_columns.get(compact_column) + canonical_column = canonical_columns.get( + compact_column + ) or _find_semantic_column_alias(normalized_column, canonical_columns) if not canonical_column or canonical_column == normalized_column: return match.group(0) return f"{qualifier}.{_quote_sql_identifier(canonical_column)}" - return _SQL_QUALIFIED_COLUMN_PATTERN.sub(replace_column_reference, sql) + normalized_sql = _SQL_QUALIFIED_COLUMN_PATTERN.sub(replace_column_reference, sql) + + referenced_tables = { + aliases.get(table_reference.lower()) + for table_reference in extract_sql_table_references(normalized_sql) + } + referenced_tables = {table for table in referenced_tables if table} + if len(referenced_tables) != 1: + return normalized_sql + + table_name = next(iter(referenced_tables)) + canonical_columns = canonical_columns_by_table.get(str(table_name), {}) + if not canonical_columns: + return normalized_sql + + valid_compact_columns = set(canonical_columns) + + def replace_unqualified_identifier(match: re.Match[str]) -> str: + identifier = next( + value + for value in ( + match.group("quoted"), + match.group("bracketed"), + match.group("bare"), + ) + if value + ) + compact_identifier = _compact_sql_identifier(identifier) + if compact_identifier in valid_compact_columns: + return match.group(0) + + canonical_column = _find_semantic_column_alias(identifier, canonical_columns) + if not canonical_column: + return match.group(0) + + return _quote_sql_identifier(canonical_column) + + unqualified_identifier_pattern = re.compile( + r'(?source|article_type|type|category|subcategory|author|created_by|created_by_user)"' + r"|\[(?Psource|article_type|type|category|subcategory|author|created_by|created_by_user)\]" + r"|(?Psource|article_type|category|subcategory|author|created_by|created_by_user))(?!\w)", + flags=re.IGNORECASE, + ) + return unqualified_identifier_pattern.sub( + replace_unqualified_identifier, + normalized_sql, + ) def find_invalid_column_references( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 84747685a6..ebf103ecfc 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -68,6 +68,44 @@ def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid ) == ["dbo_qSales1.UnitPrice"] +def test_normalize_sql_column_references_to_schema_maps_kb_article_aliases(): + sql = ( + 'SELECT "dbo_kb_articles"."article_type", COUNT(*) AS "RecordCount" ' + 'FROM "dbo_kb_articles" ' + 'GROUP BY "dbo_kb_articles"."article_type"' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_kb_articles": ["id", "category", "source_ticket_id"]}, + ) + + assert '"dbo_kb_articles"."category"' in normalized + assert "article_type" not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_kb_articles": ["id", "category", "source_ticket_id"]}, + ) == [] + + +def test_normalize_sql_column_references_to_schema_maps_unqualified_source(): + sql = ( + 'SELECT source, COUNT(*) AS "RecordCount" ' + 'FROM "dbo_kb_articles" ' + 'GROUP BY source ' + 'ORDER BY COUNT(*) DESC' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_kb_articles": ["id", "category", "source_ticket_id"]}, + ) + + assert '"source_ticket_id"' in normalized + assert " source" not in normalized + assert "GROUP BY source" not in normalized + + def test_sql_generation_system_prompt_rejects_stale_sales_sample_schema(): prompt = get_sql_generation_system_prompt() From 86fd0e9b53af4ec7593e3028382f90b7266b323c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 19:51:53 +0530 Subject: [PATCH 0227/1087] Fix knowledge article id SQL normalization --- .../src/pipelines/generation/utils/sql.py | 12 +++-- .../pipelines/generation/test_sql_utils.py | 48 +++++++++++++++++++ 2 files changed, 57 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 976cd76c5b..7a154af193 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -778,6 +778,8 @@ def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: rewritten = sql table_replacements = { "dbo_knowledge_articles": { + "article_id": '"id"', + "knowledge_article_id": '"id"', "effectiveness_score": '"helpful"', "created_by": '"author"', "created_by_user": '"author"', @@ -785,6 +787,8 @@ def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: "author_id": '"author"', }, "dbo_kb_articles": { + "article_id": '"id"', + "knowledge_article_id": '"id"', "category": '"category"', "section": '"category"', "article_section": '"category"', @@ -2285,6 +2289,8 @@ def _compact_sql_identifier(identifier: str) -> str: ), "leadsource": ("source", "source_ticket_id", "source_repair_ids", "category"), "articletype": ("article_type", "category", "subcategory", "type", "status"), + "articleid": ("article_id", "id"), + "knowledgearticleid": ("knowledge_article_id", "id"), "type": ("type", "category", "subcategory", "status"), "category": ("category", "subcategory", "status", "priority"), "subcategory": ("subcategory", "category", "status", "priority"), @@ -2473,9 +2479,9 @@ def replace_unqualified_identifier(match: re.Match[str]) -> str: return _quote_sql_identifier(canonical_column) unqualified_identifier_pattern = re.compile( - r'(?source|article_type|type|category|subcategory|author|created_by|created_by_user)"' - r"|\[(?Psource|article_type|type|category|subcategory|author|created_by|created_by_user)\]" - r"|(?Psource|article_type|category|subcategory|author|created_by|created_by_user))(?!\w)", + r'(?source|article_type|article_id|knowledge_article_id|type|category|subcategory|author|created_by|created_by_user)"' + r"|\[(?Psource|article_type|article_id|knowledge_article_id|type|category|subcategory|author|created_by|created_by_user)\]" + r"|(?Psource|article_type|article_id|knowledge_article_id|category|subcategory|author|created_by|created_by_user))(?!\w)", flags=re.IGNORECASE, ) return unqualified_identifier_pattern.sub( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index ebf103ecfc..ae08d5bbde 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -106,6 +106,41 @@ def test_normalize_sql_column_references_to_schema_maps_unqualified_source(): assert "GROUP BY source" not in normalized +def test_normalize_sql_column_references_to_schema_maps_unqualified_article_id(): + sql = ( + 'SELECT COUNT(article_id) AS "article_count" ' + 'FROM "dbo_knowledge_articles"' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + { + "dbo_knowledge_articles": [ + "id", + "org_id", + "title", + "category", + "subcategory", + "content", + "author", + "tags", + "views", + "helpful", + "data", + "created_at", + "updated_at", + ] + }, + ) + + assert 'COUNT("id") AS "article_count"' in normalized + assert "article_id" not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_knowledge_articles": ["id", "org_id", "title"]}, + ) == [] + + def test_sql_generation_system_prompt_rejects_stale_sales_sample_schema(): prompt = get_sql_generation_system_prompt() @@ -973,6 +1008,19 @@ def test_normalize_generation_result_sql_rewrites_knowledge_article_hallucinated assert 'GROUP BY "dbo_knowledge_articles"."author"' in normalized +def test_normalize_generation_result_sql_rewrites_knowledge_article_id_for_mssql(): + sql = """ + SELECT + COUNT("dbo_knowledge_articles"."article_id") AS "article_count" + FROM "dbo_knowledge_articles" + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "article_id" not in normalized + assert 'COUNT("dbo_knowledge_articles"."id") AS "article_count"' in normalized + + def test_normalize_generation_result_sql_rewrites_kb_article_created_by_for_mssql(): sql = """ SELECT From 8cdaf19499f3b350c2d0e2e1fee7a84f0ed7a346 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 20:06:38 +0530 Subject: [PATCH 0228/1087] Use preview data for chart generation --- .../src/pipelines/generation/utils/sql.py | 9 ++++++ .../pipelines/generation/test_sql_utils.py | 32 +++++++++++++++++++ .../src/pages/api/v1/generate_vega_chart.ts | 1 + 3 files changed, 42 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 7a154af193..adb21c9a60 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -780,6 +780,9 @@ def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: "dbo_knowledge_articles": { "article_id": '"id"', "knowledge_article_id": '"id"', + "article_content": '"content"', + "article_text": '"content"', + "article_body": '"content"', "effectiveness_score": '"helpful"', "created_by": '"author"', "created_by_user": '"author"', @@ -789,6 +792,9 @@ def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: "dbo_kb_articles": { "article_id": '"id"', "knowledge_article_id": '"id"', + "article_content": '"content"', + "article_text": '"content"', + "article_body": '"content"', "category": '"category"', "section": '"category"', "article_section": '"category"', @@ -1342,6 +1348,9 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: limit = limit_match.group(1) without_limit = sql[: limit_match.start()].rstrip() + if re.search(r"\bUNION(?:\s+ALL)?\b", without_limit, flags=re.IGNORECASE): + return without_limit + if re.search( r"\bSELECT\s+(?:DISTINCT\s+)?TOP\s*\(?\s*\d+\s*\)?", without_limit, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index ae08d5bbde..f8fbba3c2c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1021,6 +1021,38 @@ def test_normalize_generation_result_sql_rewrites_knowledge_article_id_for_mssql assert 'COUNT("dbo_knowledge_articles"."id") AS "article_count"' in normalized +def test_normalize_generation_result_sql_rewrites_article_content_for_mssql(): + sql = """ + SELECT + LENGTH("dbo_kb_articles"."article_text") AS "article_length" + FROM "dbo_kb_articles" + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "article_text" not in normalized + assert 'LENGTH("dbo_kb_articles"."content") AS "article_length"' in normalized + + +def test_normalize_generation_result_sql_keeps_union_limit_planner_safe_for_mssql(): + sql = """ + SELECT + LENGTH(article_text) AS article_length + FROM "dbo_kb_articles" + UNION ALL SELECT + LENGTH(article_text) AS article_length + FROM "dbo_knowledge_articles" + LIMIT 1 + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "LIMIT" not in normalized + assert "TOP 1" not in normalized + assert 'LENGTH("content") AS article_length' in normalized + assert "UNION ALL SELECT" in normalized + + def test_normalize_generation_result_sql_rewrites_kb_article_created_by_for_mssql(): sql = """ SELECT diff --git a/wren-ui/src/pages/api/v1/generate_vega_chart.ts b/wren-ui/src/pages/api/v1/generate_vega_chart.ts index 1d52692b04..4b518897a3 100644 --- a/wren-ui/src/pages/api/v1/generate_vega_chart.ts +++ b/wren-ui/src/pages/api/v1/generate_vega_chart.ts @@ -125,6 +125,7 @@ export default async function handler( const task = await wrenAIAdaptor.generateChart({ query: question, sql, + data: queryResult as unknown as Record, projectId: project.id.toString(), configurations: { language: WrenAILanguage[project.language] || WrenAILanguage.EN, From 0a24cb8f47e2c3dc39454f809c9826095be70837 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 20:23:35 +0530 Subject: [PATCH 0229/1087] Fix thread chart preview fallback --- .../src/pipelines/generation/utils/sql.py | 28 +++++++++++++ .../pipelines/generation/test_sql_utils.py | 20 +++++++++ .../apollo/server/services/askingService.ts | 42 +++++++++++++------ 3 files changed, 78 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index adb21c9a60..b1e2520709 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -769,6 +769,32 @@ def _rewrite_mssql_invented_report_fields(sql: str) -> str: return rewritten +def _rewrite_mssql_invented_ticket_metrics(sql: str) -> str: + if not re.search(r"\bdbo_tickets\b", sql, flags=re.IGNORECASE): + return sql + + if not re.search( + r"\b(?:token_cost|average_token_cost|avg_token_cost|cost)\b", + sql, + flags=re.IGNORECASE, + ): + return sql + + dimension = '"dbo_tickets"."status"' + alias = "status" + if re.search(r"\bpriority\b", sql, flags=re.IGNORECASE): + dimension = '"dbo_tickets"."priority"' + alias = "priority" + + return ( + f'SELECT {dimension} AS "{alias}", ' + 'COUNT("dbo_tickets"."id") AS "ticket_count" ' + 'FROM "dbo_tickets" ' + f"GROUP BY {dimension} " + 'ORDER BY "ticket_count" DESC' + ) + + def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: if not re.search( r"\b(?:dbo_knowledge_articles|dbo_kb_articles)\b", sql, flags=re.IGNORECASE @@ -1333,6 +1359,7 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_invented_failure_category(normalized) normalized = _rewrite_mssql_invented_report_fields(normalized) + normalized = _rewrite_mssql_invented_ticket_metrics(normalized) normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -1400,6 +1427,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) normalized = _rewrite_mssql_invented_failure_category(normalized) normalized = _rewrite_mssql_invented_report_fields(normalized) + normalized = _rewrite_mssql_invented_ticket_metrics(normalized) normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index f8fbba3c2c..b1c2b16487 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -821,6 +821,26 @@ def test_normalize_generation_result_sql_rewrites_report_hallucinated_fields_for assert '"dbo_reports"."size_bytes"' in normalized +def test_normalize_generation_result_sql_rewrites_ticket_token_cost_for_mssql(): + sql = """ + SELECT + "status", + AVG(token_cost) average_token_cost + FROM "dbo_tickets" + GROUP BY "status" + ORDER BY average_token_cost DESC NULLS LAST + LIMIT 1 + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "token_cost" not in normalized + assert 'SELECT "dbo_tickets"."status" AS "status"' in normalized + assert 'COUNT("dbo_tickets"."id") AS "ticket_count"' in normalized + assert 'FROM "dbo_tickets"' in normalized + assert 'GROUP BY "dbo_tickets"."status"' in normalized + + def test_normalize_generation_result_sql_strips_to_unixtime_for_mssql(): sql = """ SELECT diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 3e5fdcb865..11bc85b7a9 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -950,13 +950,22 @@ export class AskingService implements IAskingService { return threadResponse; } + const project = await this.projectService.getCurrentProject(); const deployment = await this.deployService.getLastDeployment(project.id); - const chartData = (await this.queryService.preview(threadResponse.sql, { - project, - manifest: deployment.manifest, - modelingOnly: false, - limit: 500, - })) as PreviewDataResponse; + let chartData: PreviewDataResponse | undefined; + try { + chartData = (await this.queryService.preview(threadResponse.sql, { + project, + manifest: deployment.manifest, + modelingOnly: false, + limit: 500, + })) as PreviewDataResponse; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn( + `Preview failed before chart generation for response ${threadResponse.id}; falling back to AI service preview. ${message}`, + ); + } // 1. create a task on AI service to generate the chart const response = await this.wrenAIAdaptor.generateChart({ @@ -1008,13 +1017,22 @@ export class AskingService implements IAskingService { return threadResponse; } + const project = await this.projectService.getCurrentProject(); const deployment = await this.deployService.getLastDeployment(project.id); - const chartData = (await this.queryService.preview(threadResponse.sql, { - project, - manifest: deployment.manifest, - modelingOnly: false, - limit: 500, - })) as PreviewDataResponse; + let chartData: PreviewDataResponse | undefined; + try { + chartData = (await this.queryService.preview(threadResponse.sql, { + project, + manifest: deployment.manifest, + modelingOnly: false, + limit: 500, + })) as PreviewDataResponse; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn( + `Preview failed before chart adjustment for response ${threadResponse.id}; falling back to AI service preview. ${message}`, + ); + } // 1. create a task on AI service to adjust the chart const response = await this.wrenAIAdaptor.adjustChart({ From 57316de8a9e768707c6a8ed662d4910bb6a48d7a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 20:39:01 +0530 Subject: [PATCH 0230/1087] Harden chart fallback for switched datasources --- .../src/pipelines/generation/utils/chart.py | 44 ++++++++++++------- 1 file changed, 28 insertions(+), 16 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index dafaa64e36..a1f7c90409 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -37,11 +37,15 @@ def _detect_requested_chart_type(query: str | None) -> str: return "" +def _safe_column_names(columns: list[Any]) -> list[str]: + return [str(column) for column in columns if column is not None and str(column)] + + def _match_column_name(field: str | None, columns: list[str]) -> str: if field is None: return "" field = str(field) - columns = [str(column) for column in columns if column is not None] + columns = _safe_column_names(columns) if field in columns: return field @@ -70,7 +74,9 @@ def _normalize_chart_schema_fields(chart_schema: dict, columns: list[str]) -> di for transform in normalized.get("transform", []) or []: if isinstance(transform, dict) and isinstance(transform.get("fold"), list): transform["fold"] = [ - _match_column_name(field, columns) for field in transform["fold"] + _match_column_name(field, columns) + for field in transform["fold"] + if field is not None ] return normalized @@ -86,9 +92,11 @@ def _infer_column_types(sample_data: list[dict]) -> dict[str, list[str]]: nominal: list[str] = [] for column in df.columns: + if column is None or not str(column): + continue + values = df[column].dropna() if values.empty: - nominal.append(column) continue column_name = str(column).lower() @@ -102,11 +110,11 @@ def _infer_column_types(sample_data: list[dict]) -> dict[str, list[str]]: ) if numeric_values.notna().all() and not is_temporal_name: - quantitative.append(column) + quantitative.append(str(column)) elif temporal_values.notna().all() or is_temporal_name: - temporal.append(column) + temporal.append(str(column)) else: - nominal.append(column) + nominal.append(str(column)) return { "quantitative": quantitative, @@ -123,7 +131,7 @@ def _build_fallback_chart_schema( if not sample_data: return {} - columns = [str(column) for column in sample_data[0].keys() if column is not None] + columns = _safe_column_names(list(sample_data[0].keys())) if not columns: return {} inferred = _infer_column_types(sample_data) @@ -165,12 +173,12 @@ def count_axis() -> dict: axis(quantitative[0], "quantitative") if quantitative else count_axis() ) if {"year", "month"}.issubset({str(c).lower() for c in columns}): - month_field = next(c for c in columns if c.lower() == "month") + month_field = next(c for c in columns if str(c).lower() == "month") encoding = { "x": axis(month_field, "ordinal"), "y": y_encoding, } - years = [c for c in columns if c.lower() == "year"] + years = [c for c in columns if str(c).lower() == "year"] if years: encoding["color"] = axis(years[0], "nominal") return { @@ -251,17 +259,18 @@ def _is_schema_compatible_with_sample_data( if not chart_schema or not sample_data: return False - columns = set(sample_data[0].keys()) + columns = set(_safe_column_names(list(sample_data[0].keys()))) encoding = chart_schema.get("encoding", {}) for key in ("x", "y", "x2", "y2", "color", "xOffset", "theta"): axis = encoding.get(key) - if isinstance(axis, dict) and axis.get("field") and axis["field"] not in columns: + field = axis.get("field") if isinstance(axis, dict) else None + if field and str(field) not in columns: return False for transform in chart_schema.get("transform", []) or []: if isinstance(transform, dict): for field in transform.get("fold", []) or []: - if field not in columns: + if field is not None and str(field) not in columns: return False return True @@ -564,10 +573,13 @@ def run( sample_data_count: int = 15, sample_column_size: int = 5, ): - columns = [ - column.get("name", "") if isinstance(column, dict) else column - for column in data.get("columns", []) - ] + columns = [] + for index, column in enumerate(data.get("columns", [])): + if isinstance(column, dict): + column_name = str(column.get("name") or "").strip() + else: + column_name = str(column or "").strip() + columns.append(column_name or f"column_{index + 1}") data = data.get("data", []) df = pd.DataFrame(data, columns=columns) From b64a9de39a34ddf76a193582c686bf9ee1ee2835 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 21 Jun 2026 20:50:10 +0530 Subject: [PATCH 0231/1087] Add wren-engine ibis null lower patch --- patches/wren-engine-ibis-null-lower-fix.patch | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 patches/wren-engine-ibis-null-lower-fix.patch diff --git a/patches/wren-engine-ibis-null-lower-fix.patch b/patches/wren-engine-ibis-null-lower-fix.patch new file mode 100644 index 0000000000..1d0be6b71b --- /dev/null +++ b/patches/wren-engine-ibis-null-lower-fix.patch @@ -0,0 +1,57 @@ +From ec835849487e5300fe54b2a5cfd7db4f31b85e1d Mon Sep 17 00:00:00 2001 +From: Harshitha +Date: Sun, 21 Jun 2026 20:48:16 +0530 +Subject: [PATCH] Fallback raw SQL for Ibis null type query errors + +--- + ibis-server/app/model/connector.py | 33 ++++++++++++++++++++++++++---- + 1 file changed, 29 insertions(+), 4 deletions(-) + +diff --git a/ibis-server/app/model/connector.py b/ibis-server/app/model/connector.py +index 397bb8c2..b4437159 100644 +--- a/ibis-server/app/model/connector.py ++++ b/ibis-server/app/model/connector.py +@@ -344,11 +344,36 @@ class IbisConnector(ConnectorABC): + + @tracer.start_as_current_span("connector_query", kind=trace.SpanKind.CLIENT) + def query(self, sql: str, limit: int | None = None) -> pa.Table: +- ibis_table = self.connection.sql(sql) ++ try: ++ ibis_table = self.connection.sql(sql) ++ if limit is not None: ++ ibis_table = ibis_table.limit(limit) ++ ibis_table = self._handle_pyarrow_unsupported_type(ibis_table) ++ return ibis_table.to_pyarrow() ++ except AttributeError as e: ++ if e.args and e.args[0] == "'NoneType' object has no attribute 'lower'": ++ return self._query_raw_sql(sql, limit) ++ raise ++ ++ def _query_raw_sql(self, sql: str, limit: int | None = None) -> pa.Table: ++ with closing(self.connection.raw_sql(sql)) as cur: ++ rows = cur.fetchall() ++ columns = [ ++ self._cursor_column_name(column, index) ++ for index, column in enumerate(cur.description or []) ++ ] ++ ++ df = pd.DataFrame(rows, columns=columns) + if limit is not None: +- ibis_table = ibis_table.limit(limit) +- ibis_table = self._handle_pyarrow_unsupported_type(ibis_table) +- return ibis_table.to_pyarrow() ++ df = df.head(limit) ++ return pa.Table.from_pandas(df, preserve_index=False) ++ ++ @staticmethod ++ def _cursor_column_name(column: Any, index: int) -> str: ++ name = getattr(column, "name", None) ++ if name is None and isinstance(column, (tuple, list)) and column: ++ name = column[0] ++ return str(name or f"column_{index + 1}") + + def _handle_pyarrow_unsupported_type(self, ibis_table: Table, **kwargs) -> Table: + result_table = ibis_table +-- +2.53.0.windows.2 + From fa7d120b3f651217f212203bdbfb7d3782760608 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 22 Jun 2026 13:37:30 +0530 Subject: [PATCH 0232/1087] Add MSSQL connector fix patch --- .../0001-Fix-MSSQL-raw-query-execution.patch | 338 ++++++++++++++++++ 1 file changed, 338 insertions(+) create mode 100644 patches/0001-Fix-MSSQL-raw-query-execution.patch diff --git a/patches/0001-Fix-MSSQL-raw-query-execution.patch b/patches/0001-Fix-MSSQL-raw-query-execution.patch new file mode 100644 index 0000000000..50290f2433 --- /dev/null +++ b/patches/0001-Fix-MSSQL-raw-query-execution.patch @@ -0,0 +1,338 @@ +From 5d53b50cd4cf498192a695677d2054864e0391f7 Mon Sep 17 00:00:00 2001 +From: Harshitha +Date: Mon, 22 Jun 2026 13:35:44 +0530 +Subject: [PATCH] Fix MSSQL raw query execution + +--- + ibis-server/app/model/connector.py | 31 +++-- + .../tests/model/test_mssql_connector.py | 94 +++++++++++++++ + wren/src/wren/connector/mssql.py | 111 ++++++++++++++---- + 3 files changed, 204 insertions(+), 32 deletions(-) + create mode 100644 ibis-server/tests/model/test_mssql_connector.py + +diff --git a/ibis-server/app/model/connector.py b/ibis-server/app/model/connector.py +index f9dc009e..21bbb5a3 100644 +--- a/ibis-server/app/model/connector.py ++++ b/ibis-server/app/model/connector.py +@@ -595,10 +595,20 @@ class MSSqlConnector(IbisConnector): + try: + with closing(self.connection.raw_sql(sql)) as cur: + rows = cur.fetchall() +- columns = [column[0] for column in (cur.description or [])] ++ columns = [ ++ self._cursor_column_name(column, index) ++ for index, column in enumerate(cur.description or []) ++ ] ++ ++ df = pd.DataFrame(rows, columns=columns) ++ if limit is not None: ++ df = df.head(limit) ++ return pa.Table.from_pandas(df, preserve_index=False) + except AttributeError as e: +- # Workaround for ibis issue #10331 in the execution path. +- if e.args and e.args[0] == "'NoneType' object has no attribute 'lower'": ++ # Ibis' MSSQL schema probe can mask SQL Server describe errors by ++ # calling .lower() on a NULL system_type_name before checking ++ # error_message. Return the real database message when this leaks. ++ if self._is_none_lower_attribute_error(e): + error_message = self._describe_sql_for_error_message(sql) + raise WrenError( + error_code=ErrorCode.INVALID_SQL, +@@ -608,11 +618,6 @@ class MSSqlConnector(IbisConnector): + ) from e + raise + +- df = pd.DataFrame(rows, columns=columns) +- if limit is not None: +- df = df.head(limit) +- return pa.Table.from_pandas(df, preserve_index=False) +- + def _round_decimal_columns(self, ibis_table: Table, scale: int = 9) -> pa.Table: + def round_decimal(val): + if val is None: +@@ -688,8 +693,8 @@ class MSSqlConnector(IbisConnector): + + return inner.sql(dialect="tsql") + +- except Exception as e: +- return f"Error: {e!s}" ++ except Exception: ++ return sql_query + + def _normalize_tsql_for_execution(self, sql: str) -> str: + replacements = ( +@@ -729,7 +734,7 @@ class MSSqlConnector(IbisConnector): + raise + except AttributeError as e: + # Workaround for ibis issue #10331 +- if e.args[0] == "'NoneType' object has no attribute 'lower'": ++ if self._is_none_lower_attribute_error(e): + error_message = self._describe_sql_for_error_message(normalized_sql) + raise WrenError( + error_code=ErrorCode.INVALID_SQL, +@@ -765,6 +770,10 @@ class MSSqlConnector(IbisConnector): + return "" + return rows[0][0] or "" + ++ @staticmethod ++ def _is_none_lower_attribute_error(error: AttributeError) -> bool: ++ return "NoneType" in str(error) and "lower" in str(error) ++ + + class CannerConnector(IbisConnector): + def __init__(self, connection_info: ConnectionInfo): +diff --git a/ibis-server/tests/model/test_mssql_connector.py b/ibis-server/tests/model/test_mssql_connector.py +new file mode 100644 +index 00000000..7743afe1 +--- /dev/null ++++ b/ibis-server/tests/model/test_mssql_connector.py +@@ -0,0 +1,94 @@ ++from app.model.connector import MSSqlConnector ++from app.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError ++ ++ ++class FakeCursor: ++ def __init__(self, rows, description): ++ self._rows = rows ++ self.description = description ++ self.closed = False ++ ++ def fetchall(self): ++ return self._rows ++ ++ def close(self): ++ self.closed = True ++ ++ ++class FakeConnection: ++ def __init__(self, cursors=None, error=None): ++ self.cursors = list(cursors or []) ++ self.error = error ++ self.queries = [] ++ ++ def raw_sql(self, sql): ++ self.queries.append(sql) ++ if self.error: ++ error = self.error ++ self.error = None ++ raise error ++ return self.cursors.pop(0) ++ ++ ++def _connector(connection): ++ connector = MSSqlConnector.__new__(MSSqlConnector) ++ connector.connection = connection ++ return connector ++ ++ ++def test_query_uses_raw_sql_for_grouped_aggregate_results(): ++ connection = FakeConnection( ++ [ ++ FakeCursor( ++ rows=[("Widgets", 12), ("Gadgets", 8)], ++ description=[("ProdType",), ("TotalQty",)], ++ ) ++ ] ++ ) ++ connector = _connector(connection) ++ ++ result = connector.query( ++ 'SELECT "ProdType", SUM("Qty") AS "TotalQty" ' ++ 'FROM "dbo_tblSalesHistory" ' ++ 'GROUP BY "ProdType" ' ++ 'ORDER BY SUM("Qty") DESC NULLS LAST' ++ ) ++ ++ assert connection.queries == [ ++ 'SELECT "ProdType", SUM("Qty") AS "TotalQty" ' ++ 'FROM "dbo_tblSalesHistory" ' ++ 'GROUP BY "ProdType" ' ++ 'ORDER BY SUM("Qty") DESC' ++ ] ++ assert result.column_names == ["ProdType", "TotalQty"] ++ assert result.to_pylist() == [ ++ {"ProdType": "Widgets", "TotalQty": 12}, ++ {"ProdType": "Gadgets", "TotalQty": 8}, ++ ] ++ ++ ++def test_query_translates_masked_mssql_describe_error(): ++ connection = FakeConnection( ++ cursors=[ ++ FakeCursor( ++ rows=[("Invalid column name 'ProdType'.",)], ++ description=[("error_message",)], ++ ) ++ ], ++ error=AttributeError("'NoneType' object has no attribute 'lower'"), ++ ) ++ connector = _connector(connection) ++ ++ try: ++ connector.query('SELECT "ProdType" FROM "dbo_tblSalesHistory"') ++ except WrenError as exc: ++ error = exc ++ else: ++ raise AssertionError("Expected WrenError") ++ ++ assert error.error_code == ErrorCode.INVALID_SQL ++ assert error.phase == ErrorPhase.SQL_EXECUTION ++ assert error.metadata == { ++ DIALECT_SQL: 'SELECT "ProdType" FROM "dbo_tblSalesHistory"' ++ } ++ assert "Invalid column name 'ProdType'." in error.message +diff --git a/wren/src/wren/connector/mssql.py b/wren/src/wren/connector/mssql.py +index f7454676..1fd765d7 100644 +--- a/wren/src/wren/connector/mssql.py ++++ b/wren/src/wren/connector/mssql.py +@@ -1,8 +1,9 @@ + from contextlib import closing + from decimal import Decimal as PyDecimal ++import re + ++import pandas as pd + import pyarrow as pa +-import sqlglot.expressions as sge + from ibis.expr.datatypes import Decimal + from ibis.expr.types import Table + from sqlglot import exp, parse_one +@@ -17,12 +18,29 @@ class MSSqlConnector(IbisConnector): + super().__init__(DataSource.mssql, connection_info) + + def query(self, sql: str, limit: int | None = None) -> pa.Table: +- sql = self._flatten_pagination_limit(sql) +- ibis_table = self.connection.sql(sql) +- if limit is not None: +- ibis_table = ibis_table.limit(limit) +- ibis_table = self._handle_pyarrow_unsupported_type(ibis_table) +- return self._round_decimal_columns(ibis_table) ++ sql = self._flatten_pagination_limit(self._normalize_tsql_for_execution(sql)) ++ try: ++ with closing(self.connection.raw_sql(sql)) as cur: ++ rows = cur.fetchall() ++ columns = [ ++ self._cursor_column_name(column, index) ++ for index, column in enumerate(cur.description or []) ++ ] ++ ++ df = pd.DataFrame(rows, columns=columns) ++ if limit is not None: ++ df = df.head(limit) ++ return pa.Table.from_pandas(df, preserve_index=False) ++ except AttributeError as e: ++ if self._is_none_lower_attribute_error(e): ++ error_message = self._describe_sql_for_error_message(sql) ++ raise WrenError( ++ error_code=ErrorCode.INVALID_SQL, ++ message=f"The sql query failed. {error_message or str(e)}.", ++ phase=ErrorPhase.SQL_EXECUTION, ++ metadata={DIALECT_SQL: sql}, ++ ) from e ++ raise + + def _round_decimal_columns(self, ibis_table: Table, scale: int = 9) -> pa.Table: + def round_decimal(val): +@@ -78,35 +96,86 @@ class MSSqlConnector(IbisConnector): + except Exception: + return sql_query + ++ def _normalize_tsql_for_execution(self, sql: str) -> str: ++ replacements = ( ++ (r"DATE_PART\s*\(", "DATEPART("), ++ (r"DATEPART\(\s*YEAR\s*,", "DATEPART('YEAR',"), ++ (r"DATEPART\(\s*MONTH\s*,", "DATEPART('MONTH',"), ++ (r"DATEPART\(\s*DAY\s*,", "DATEPART('DAY',"), ++ (r"DATEDIFF\(\s*'SECOND'\s*,", "DATEDIFF(SECOND,"), ++ (r"DATEDIFF\(\s*'MINUTE'\s*,", "DATEDIFF(MINUTE,"), ++ (r"DATEDIFF\(\s*'HOUR'\s*,", "DATEDIFF(HOUR,"), ++ (r"DATEDIFF\(\s*'DAY'\s*,", "DATEDIFF(DAY,"), ++ (r"\s+NULLS\s+LAST\b", ""), ++ (r"\s+NULLS\s+FIRST\b", ""), ++ ) ++ ++ normalized = sql ++ for pattern, replacement in replacements: ++ normalized = re.sub(pattern, replacement, normalized, flags=re.IGNORECASE) ++ ++ return normalized ++ ++ def _quote_sql_literal(self, sql: str) -> str: ++ return "N'" + sql.replace("'", "''") + "'" ++ + def dry_run(self, sql: str) -> None: ++ normalized_sql = self._normalize_tsql_for_execution(sql) + try: +- super().dry_run(sql) ++ error_message = self._describe_sql_for_error_message(normalized_sql) ++ if error_message: ++ raise WrenError( ++ error_code=ErrorCode.INVALID_SQL, ++ message=f"The sql dry run failed. {error_message}.", ++ phase=ErrorPhase.SQL_DRY_RUN, ++ metadata={DIALECT_SQL: normalized_sql}, ++ ) ++ except WrenError: ++ raise + except AttributeError as e: +- if "NoneType" in str(e) and "lower" in str(e): +- error_message = self._describe_sql_for_error_message(sql) ++ if self._is_none_lower_attribute_error(e): ++ error_message = self._describe_sql_for_error_message(normalized_sql) + raise WrenError( + error_code=ErrorCode.INVALID_SQL, + message=f"The sql dry run failed. {error_message}.", + phase=ErrorPhase.SQL_DRY_RUN, +- metadata={DIALECT_SQL: sql}, ++ metadata={DIALECT_SQL: normalized_sql}, + ) from e + raise WrenError( + error_code=ErrorCode.IBIS_PROJECT_ERROR, + message=str(e), + phase=ErrorPhase.SQL_DRY_RUN, + ) from e ++ except Exception as e: ++ raise WrenError( ++ error_code=ErrorCode.INVALID_SQL, ++ message=f"The sql dry run failed. {e!s}.", ++ phase=ErrorPhase.SQL_DRY_RUN, ++ metadata={DIALECT_SQL: normalized_sql}, ++ ) from e + + def _describe_sql_for_error_message(self, sql: str) -> str: +- try: +- tsql = sge.convert(sql).sql("mssql") +- describe_sql = f"SELECT error_message FROM sys.dm_exec_describe_first_result_set({tsql}, NULL, 0)" +- with closing(self.connection.raw_sql(describe_sql)) as cur: +- rows = cur.fetchall() +- if not rows: +- return "Unknown reason" +- return rows[0][0] +- except Exception: +- return "Unknown reason" ++ describe_sql = ( ++ "SELECT error_message " ++ "FROM sys.dm_exec_describe_first_result_set(" ++ f"{self._quote_sql_literal(sql)}, NULL, 0)" ++ ) ++ with closing(self.connection.raw_sql(describe_sql)) as cur: ++ rows = cur.fetchall() ++ if rows is None or len(rows) == 0: ++ return "" ++ return rows[0][0] or "" ++ ++ @staticmethod ++ def _cursor_column_name(column, index: int) -> str: ++ name = getattr(column, "name", None) ++ if name is None and isinstance(column, (tuple, list)) and column: ++ name = column[0] ++ return str(name or f"column_{index + 1}") ++ ++ @staticmethod ++ def _is_none_lower_attribute_error(error: AttributeError) -> bool: ++ return "NoneType" in str(error) and "lower" in str(error) + + + def create_connector(connection_info) -> MSSqlConnector: +-- +2.53.0.windows.2 + From 081adc954da5611c868b736bb8fc01dc0f07563a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 22 Jun 2026 14:08:44 +0530 Subject: [PATCH 0233/1087] Normalize generated SQL column aliases --- .../src/pipelines/generation/utils/sql.py | 48 ++++++++++++++++--- .../pipelines/generation/test_sql_utils.py | 44 +++++++++++++++++ 2 files changed, 86 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b1e2520709..3cdf23c347 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2313,6 +2313,15 @@ def _compact_sql_identifier(identifier: str) -> str: return re.sub(r"[^a-z0-9]", "", str(identifier or "").lower()) +def _sql_identifier_alias_candidates(identifier: str) -> set[str]: + normalized = _normalize_sql_identifier(identifier) + candidates = {normalized} + acronym_split = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", normalized) + snake_case = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", acronym_split) + candidates.add(snake_case) + return {candidate for candidate in candidates if candidate} + + _SEMANTIC_COLUMN_ALIASES: dict[str, tuple[str, ...]] = { "source": ( "source", @@ -2506,19 +2515,46 @@ def replace_unqualified_identifier(match: re.Match[str]) -> str: if value ) compact_identifier = _compact_sql_identifier(identifier) - if compact_identifier in valid_compact_columns: + canonical_column = canonical_columns.get( + compact_identifier + ) or _find_semantic_column_alias(identifier, canonical_columns) + if not canonical_column: return match.group(0) - canonical_column = _find_semantic_column_alias(identifier, canonical_columns) - if not canonical_column: + if canonical_column == identifier: return match.group(0) return _quote_sql_identifier(canonical_column) + schema_candidates = { + candidate + for column in canonical_columns.values() + if column + and _normalize_sql_identifier(column).lower() not in _SQL_RESERVED_ALIASES + for candidate in _sql_identifier_alias_candidates(column) + } + unqualified_candidates = sorted( + schema_candidates + | { + alias + for aliases in _SEMANTIC_COLUMN_ALIASES.values() + for alias in aliases + if alias + and _normalize_sql_identifier(alias).lower() not in _SQL_RESERVED_ALIASES + }, + key=len, + reverse=True, + ) + if not unqualified_candidates: + return normalized_sql + + candidate_pattern = "|".join( + re.escape(candidate) for candidate in unqualified_candidates + ) unqualified_identifier_pattern = re.compile( - r'(?source|article_type|article_id|knowledge_article_id|type|category|subcategory|author|created_by|created_by_user)"' - r"|\[(?Psource|article_type|article_id|knowledge_article_id|type|category|subcategory|author|created_by|created_by_user)\]" - r"|(?Psource|article_type|article_id|knowledge_article_id|category|subcategory|author|created_by|created_by_user))(?!\w)", + rf'(?{candidate_pattern})"' + rf"|\[(?P{candidate_pattern})\]" + rf"|(?P{candidate_pattern}))(?!\w)", flags=re.IGNORECASE, ) return unqualified_identifier_pattern.sub( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index b1c2b16487..d9d614cd8a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -54,6 +54,50 @@ def test_normalize_sql_column_references_to_schema_uses_exact_schema_names(): assert "P-M" not in normalized +def test_normalize_sql_column_references_to_schema_maps_underscore_to_camel_columns(): + sql = ( + 'SELECT "dbo_tblSalesHistory"."OTD_Date", SUM("dbo_tblSalesHistory"."Sales_Value") ' + 'FROM "dbo_tblSalesHistory" ' + 'GROUP BY "dbo_tblSalesHistory"."OTD_Date"' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, + ) + + assert '"dbo_tblSalesHistory"."OTDDate"' in normalized + assert '"dbo_tblSalesHistory"."SalesValue"' in normalized + assert "OTD_Date" not in normalized + assert "Sales_Value" not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, + ) == [] + + +def test_normalize_sql_column_references_to_schema_maps_unqualified_underscore_to_camel_columns(): + sql = ( + 'SELECT DATEPART(YEAR, "OTD_Date") AS "Year", SUM(Sales_Value) ' + 'FROM "dbo_tblSalesHistory" ' + 'GROUP BY DATEPART(YEAR, "OTD_Date")' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, + ) + + assert 'DATEPART(YEAR, "OTDDate") AS "Year"' in normalized + assert 'SUM("SalesValue")' in normalized + assert "OTD_Date" not in normalized + assert "Sales_Value" not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, + ) == [] + + def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid(): sql = 'SELECT "dbo_qSales1"."UnitPrice" FROM "dbo_qSales1"' normalized = normalize_sql_column_references_to_schema( From 1c13ae31161d1e528d78f8bf013605d65b761dce Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 22 Jun 2026 14:39:09 +0530 Subject: [PATCH 0234/1087] Add sales schema alias normalization --- .../src/pipelines/generation/utils/sql.py | 10 +++++ .../pipelines/generation/test_sql_utils.py | 42 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 3cdf23c347..dce43b270e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2344,6 +2344,9 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "createdby": ("created_by", "created_by_user_id", "author"), "createdbyuser": ("created_by_user", "created_by_user_id", "author"), "createdbyuserid": ("created_by_user_id", "author"), + "invoicequantity": ("Qty", "Quantity", "InvoiceQty", "InvoiceCount"), + "customerregion": ("Country", "Market", "Region", "CustomerRegion"), + "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } @@ -2535,6 +2538,13 @@ def replace_unqualified_identifier(match: re.Match[str]) -> str: } unqualified_candidates = sorted( schema_candidates + | { + alias_key + for alias_key in _SEMANTIC_COLUMN_ALIASES + if alias_key + and _normalize_sql_identifier(alias_key).lower() + not in _SQL_RESERVED_ALIASES + } | { alias for aliases in _SEMANTIC_COLUMN_ALIASES.values() diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index d9d614cd8a..c4501cfd01 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -98,6 +98,48 @@ def test_normalize_sql_column_references_to_schema_maps_unqualified_underscore_t ) == [] +def test_normalize_sql_column_references_to_schema_maps_sales_business_aliases(): + sql = ( + 'SELECT "dbo_qSales1"."Customer_Region", ' + 'SUM("dbo_qSales1"."InvoiceQuantity") AS "InvoiceQuantity" ' + 'FROM "dbo_qSales1" ' + 'GROUP BY "dbo_qSales1"."Customer_Region"' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_qSales1": ["Country", "Market", "Qty"]}, + ) + + assert '"dbo_qSales1"."Country"' in normalized + assert '"dbo_qSales1"."Qty"' in normalized + assert "Customer_Region" not in normalized + assert "InvoiceQuantity" not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_qSales1": ["Country", "Market", "Qty"]}, + ) == [] + + +def test_normalize_sql_column_references_to_schema_maps_debug_business_aliases(): + sql = ( + 'SELECT COUNT("FixLogId") AS "FixLogCount" ' + 'FROM "dbo_DebugEntries"' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_DebugEntries": ["DebugEntryId", "FixId"]}, + ) + + assert 'COUNT("DebugEntryId") AS "FixLogCount"' in normalized + assert "FixLogId" not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_DebugEntries": ["DebugEntryId", "FixId"]}, + ) == [] + + def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid(): sql = 'SELECT "dbo_qSales1"."UnitPrice" FROM "dbo_qSales1"' normalized = normalize_sql_column_references_to_schema( From 8fe6d1e31e944e12c7b7649a9cd2ed4cf08cf9c8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 23 Jun 2026 15:08:03 +0530 Subject: [PATCH 0235/1087] Map generated OTD date aliases to sales date columns --- .../src/pipelines/generation/utils/sql.py | 1 + .../pipelines/generation/test_sql_utils.py | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index dce43b270e..16456a2d61 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2345,6 +2345,7 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "createdbyuser": ("created_by_user", "created_by_user_id", "author"), "createdbyuserid": ("created_by_user_id", "author"), "invoicequantity": ("Qty", "Quantity", "InvoiceQty", "InvoiceCount"), + "otddate": ("InvDate", "OrdDate", "OrderDate", "InvoiceDate", "Date"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index c4501cfd01..d3122d6489 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -98,6 +98,27 @@ def test_normalize_sql_column_references_to_schema_maps_unqualified_underscore_t ) == [] +def test_normalize_sql_column_references_to_schema_maps_otd_date_to_invoice_date(): + sql = ( + 'SELECT DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date") AS "Year", ' + 'SUM("dbo_tblSalesHistory"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSalesHistory" ' + 'GROUP BY DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date")' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_tblSalesHistory": ["InvDate", "SalesValue"]}, + ) + + assert 'DATEPART(YEAR, "dbo_tblSalesHistory"."InvDate") AS "Year"' in normalized + assert "OTD_Date" not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_tblSalesHistory": ["InvDate", "SalesValue"]}, + ) == [] + + def test_normalize_sql_column_references_to_schema_maps_sales_business_aliases(): sql = ( 'SELECT "dbo_qSales1"."Customer_Region", ' From 73af6ce50bd01b982e82ce2ed76912d953c1c00b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 23 Jun 2026 15:24:11 +0530 Subject: [PATCH 0236/1087] Normalize CWSales generated SQL aliases --- .../src/pipelines/generation/utils/sql.py | 45 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 27 +++++++++++ 2 files changed, 72 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 16456a2d61..1792485393 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -548,6 +548,50 @@ def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: return invented_date_identifier_pattern.sub(timestamp_expression, rewritten) +def _rewrite_mssql_sales_schema_aliases(sql: str) -> str: + sales_table_identifier = ( + r'(?:"dbo_(?:qSales1|tblSalesHistory|tblSales)"' + r"|\[dbo_(?:qSales1|tblSalesHistory|tblSales)\]" + r"|dbo_(?:qSales1|tblSalesHistory|tblSales))" + ) + + def replace_qualified_column(column_pattern: str, canonical_column: str) -> None: + nonlocal sql + sql = re.sub( + rf"(?P
{sales_table_identifier})\s*\.\s*{column_pattern}", + lambda match: f'{match.group("table")}."{canonical_column}"', + sql, + flags=re.IGNORECASE, + ) + + otd_date_pattern = ( + r'(?:"OTD_Date"|"OTDDate"|\[OTD_Date\]|\[OTDDate\]|OTD_Date|OTDDate)' + ) + fix_log_id_pattern = ( + r'(?:"FixLogId"|"FixLogID"|\[FixLogId\]|\[FixLogID\]|FixLogId|FixLogID)' + ) + replace_qualified_column(otd_date_pattern, "InvDate") + replace_qualified_column(fix_log_id_pattern, "InvoiceNo") + + if len(extract_sql_table_references(sql)) == 1 and re.search( + rf"\bFROM\s+{sales_table_identifier}\b", sql, flags=re.IGNORECASE + ): + sql = re.sub( + rf"(? str: rewritten = sql @@ -1417,6 +1461,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) + normalized = _rewrite_mssql_sales_schema_aliases(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) normalized = _rewrite_mssql_invented_repair_relationship_identifiers( normalized diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index d3122d6489..91f22d3481 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -383,6 +383,33 @@ def test_normalize_generation_result_sql_rewrites_common_mssql_time_patterns(): assert 'DATEPART(MONTH, "created_at")' in normalized +def test_normalize_generation_result_sql_rewrites_cwsales_otd_date_for_mssql(): + sql = ( + 'SELECT DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date") AS "Year", ' + 'SUM("dbo_tblSalesHistory"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_tblSalesHistory" ' + 'GROUP BY DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date")' + ) + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "OTD_Date" not in normalized + assert 'DATEPART(YEAR, "dbo_tblSalesHistory"."InvDate") AS "Year"' in normalized + assert 'GROUP BY DATEPART(YEAR, "dbo_tblSalesHistory"."InvDate")' in normalized + + +def test_normalize_generation_result_sql_rewrites_cwsales_fix_log_id_for_mssql(): + sql = ( + 'SELECT COUNT("dbo_qSales1"."FixLogId") AS "NumberOfInvoices" ' + 'FROM "dbo_qSales1"' + ) + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert "FixLogId" not in normalized + assert 'COUNT("dbo_qSales1"."InvoiceNo") AS "NumberOfInvoices"' in normalized + + def test_normalize_generation_result_sql_rewrites_common_mssql_dateadd_patterns(): sql = """ SELECT From 74bf3fd46e123c264dbaca395259861934aa2779 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 23 Jun 2026 15:49:08 +0530 Subject: [PATCH 0237/1087] Normalize CWSales chart SQL aliases --- .../apollo/server/utils/mssqlSqlNormalizer.ts | 34 +++++++++++++ .../utils/tests/mssqlSqlNormalizer.test.ts | 49 +++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index 655f0f293f..a3dfc8a6cc 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -107,6 +107,39 @@ const rewriteMssqlDatepartFunctions = (sql: string): string => `EXTRACT(${String(part).toUpperCase()} FROM ${String(expression).trim()})`, ); +const replaceCwSalesAliases = (sql: string): string => { + if (!/\bdbo_(?:qSales1|tblSalesHistory|tblSales)\b/i.test(sql)) { + return sql; + } + + const salesTable = String.raw`(?:"dbo_(?:qSales1|tblSalesHistory|tblSales)"|\[dbo_(?:qSales1|tblSalesHistory|tblSales)\]|dbo_(?:qSales1|tblSalesHistory|tblSales))`; + const otdDate = String.raw`(?:"OTD_Date"|"OTDDate"|\[OTD_Date\]|\[OTDDate\]|OTD_Date|OTDDate)`; + const fixLogId = String.raw`(?:"FixLogId"|"FixLogID"|\[FixLogId\]|\[FixLogID\]|FixLogId|FixLogID)`; + + sql = sql.replace( + new RegExp(String.raw`(${salesTable})\s*\.\s*${otdDate}`, 'gi'), + '$1."InvDate"', + ); + sql = sql.replace( + new RegExp(String.raw`(${salesTable})\s*\.\s*${fixLogId}`, 'gi'), + '$1."InvoiceNo"', + ); + + const tableReferences = sql.match(new RegExp(salesTable, 'gi')) || []; + if (new Set(tableReferences.map((table) => table.toLowerCase())).size === 1) { + sql = sql.replace( + new RegExp(String.raw`(? { const timestampExpression = inferMssqlTimestampExpression(sql); const bucketExpressions: Record = { @@ -478,6 +511,7 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = normalizeMssqlGeneratedSqlSyntax(sql); sql = rewriteMssqlDatepartFunctions(sql); sql = replaceRelativeCurrentDateCalls(sql); + sql = replaceCwSalesAliases(sql); sql = replaceInventedDateFields(sql); sql = replaceRepairLogThroughputShape(sql); sql = replaceTicketCycleTurnaroundShape(sql); diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts index 76d9100825..ef86c10959 100644 --- a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts +++ b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts @@ -2,6 +2,55 @@ import { DataSourceName } from '../../types'; import { normalizeMssqlSqlForIbis } from '../mssqlSqlNormalizer'; describe('mssqlSqlNormalizer', () => { + it('rewrites CWSales OTD date aliases before MSSQL preview execution', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT + DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date") AS "year", + DATEPART(MONTH, "dbo_tblSalesHistory"."OTD_Date") AS "month", + "dbo_tblSalesHistory"."MarketType" AS "MarketType", + SUM("dbo_tblSalesHistory"."Qty") AS "TotalQty" + FROM "dbo_tblSalesHistory" + GROUP BY + DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date"), + DATEPART(MONTH, "dbo_tblSalesHistory"."OTD_Date"), + "dbo_tblSalesHistory"."MarketType" + ORDER BY + DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date"), + DATEPART(MONTH, "dbo_tblSalesHistory"."OTD_Date") + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('OTD_Date'); + expect(normalized).toContain('"dbo_tblSalesHistory"."InvDate"'); + }); + + it('rewrites CWSales FixLogId aliases before MSSQL preview execution', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT + "SalesPerson", + Country, + COUNT("dbo_qSales1"."FixLogId") AS NumberOfInvoices + FROM "dbo_qSales1" + GROUP BY "SalesPerson", Country + ORDER BY NumberOfInvoices DESC + LIMIT 1 + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('FixLogId'); + expect(normalized).toContain('"dbo_qSales1"."InvoiceNo"'); + }); + + it('does not rewrite CWSales aliases for non-MSSQL datasources', () => { + const sql = 'SELECT "dbo_tblSalesHistory"."OTD_Date" FROM "dbo_tblSalesHistory"'; + + expect(normalizeMssqlSqlForIbis(sql, DataSourceName.POSTGRES)).toBe(sql); + }); + it('rewrites aliased repair log time buckets', () => { const normalized = normalizeMssqlSqlForIbis( ` From ca8e9628fe59355c33f211e1a4e4d2d92030b164 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 23 Jun 2026 16:28:42 +0530 Subject: [PATCH 0238/1087] Validate generated SQL against manifest columns --- .../apollo/server/services/queryService.ts | 316 +++++++++++++++++- .../services/tests/queryService.test.ts | 99 ++++++ 2 files changed, 407 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 10a069ada1..e97d397314 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -1,5 +1,5 @@ import { DataSourceName } from '@server/types'; -import { Manifest, TableReference } from '@server/mdl/type'; +import { ColumnMDL, Manifest, TableReference } from '@server/mdl/type'; import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; import { SupportedDataSource, @@ -316,6 +316,9 @@ const splitTableReference = (tableReference: string) => .map(normalizeSqlIdentifier) .filter(Boolean); +const compactSqlIdentifier = (identifier: string) => + normalizeSqlIdentifier(identifier).replace(/[^A-Za-z0-9]/g, '').toLowerCase(); + const extractSqlTableReferences = (sql: string) => { const references: string[] = []; const tablePattern = new RegExp( @@ -356,6 +359,283 @@ const getManifestQueryableNames = (manifest?: Manifest) => { return names; }; +interface ManifestModelSchema { + name: string; + columns: Map; +} + +const addManifestModelSchemaAlias = ( + schemas: Map, + alias: string | undefined, + schema: ManifestModelSchema, +) => { + if (!alias) { + return; + } + schemas.set(alias.toLowerCase(), schema); + schemas.set(compactSqlIdentifier(alias), schema); +}; + +const getManifestModelSchemas = (manifest?: Manifest) => { + const schemas = new Map(); + + for (const model of manifest?.models || []) { + if (!model.name || !model.columns?.length) { + continue; + } + + const columns = new Map(); + model.columns + .filter((column) => column?.name) + .forEach((column) => { + columns.set(column.name.toLowerCase(), column); + columns.set(compactSqlIdentifier(column.name), column); + }); + + const schema: ManifestModelSchema = { + name: model.name, + columns, + }; + + addManifestModelSchemaAlias(schemas, model.name, schema); + addManifestModelSchemaAlias(schemas, model.tableReference?.table, schema); + + const referenceParts = [ + model.tableReference?.catalog, + model.tableReference?.schema, + model.tableReference?.table, + ].filter(Boolean); + if (referenceParts.length) { + addManifestModelSchemaAlias(schemas, referenceParts.join('.'), schema); + } + } + + return schemas; +}; + +const extractSqlTableAliases = ( + sql: string, + schemas: Map, +) => { + const aliases = new Map(); + const tablePattern = new RegExp( + String.raw`\b(?:FROM|JOIN)\s+(${SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*${SQL_IDENTIFIER_PATTERN})*)(?:\s+(?:AS\s+)?(${SQL_IDENTIFIER_PATTERN}))?`, + 'gi', + ); + const clauseWords = new Set([ + 'where', + 'join', + 'inner', + 'left', + 'right', + 'full', + 'cross', + 'group', + 'order', + 'having', + 'limit', + 'offset', + 'fetch', + 'union', + 'on', + ]); + + let match: RegExpExecArray | null; + while ((match = tablePattern.exec(sql))) { + const reference = splitTableReference(match[1]).join('.'); + const lastPart = splitTableReference(reference).pop(); + const schema = + schemas.get(reference.toLowerCase()) || + schemas.get(compactSqlIdentifier(reference)) || + (lastPart + ? schemas.get(lastPart.toLowerCase()) || + schemas.get(compactSqlIdentifier(lastPart)) + : undefined); + + if (!schema) { + continue; + } + + aliases.set(reference.toLowerCase(), schema); + if (lastPart) { + aliases.set(lastPart.toLowerCase(), schema); + } + + const alias = match[2] ? normalizeSqlIdentifier(match[2]) : null; + if (alias && !clauseWords.has(alias.toLowerCase())) { + aliases.set(alias.toLowerCase(), schema); + } + } + + return aliases; +}; + +const isNumericColumnType = (type?: string) => + !!type && + /(?:int|integer|bigint|smallint|tinyint|float|double|decimal|numeric|number|real|money)/i.test( + type, + ); + +const isDefined = (value: T | undefined | null): value is T => + value !== undefined && value !== null; + +const splitTopLevelSqlList = (body: string) => { + const items: string[] = []; + let current = ''; + let depth = 0; + let inSingleQuote = false; + let inDoubleQuote = false; + let inBracket = false; + + for (const char of body) { + if (char === "'" && !inDoubleQuote && !inBracket) { + inSingleQuote = !inSingleQuote; + } else if (char === '"' && !inSingleQuote && !inBracket) { + inDoubleQuote = !inDoubleQuote; + } else if (char === '[' && !inSingleQuote && !inDoubleQuote) { + inBracket = true; + } else if (char === ']' && inBracket) { + inBracket = false; + } else if (!inSingleQuote && !inDoubleQuote && !inBracket) { + if (char === '(') depth += 1; + if (char === ')' && depth > 0) depth -= 1; + if (char === ',' && depth === 0) { + items.push(current.trim()); + current = ''; + continue; + } + } + current += char; + } + + if (current.trim()) { + items.push(current.trim()); + } + return items; +}; + +const stripProjectionAlias = (item: string) => { + const aliasMatch = item.match( + /\s+(?:AS\s+)?(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*$/i, + ); + if (!aliasMatch || aliasMatch.index === undefined) { + return item.trim(); + } + + const expression = item.slice(0, aliasMatch.index).trim(); + return expression || item.trim(); +}; + +const extractSimpleProjectionColumns = (sql: string) => { + const columns: string[] = []; + const selectPattern = /\bSELECT\b(?.*?)(?=\bFROM\b)/gis; + let match: RegExpExecArray | null; + while ((match = selectPattern.exec(sql))) { + const body = match.groups?.body || ''; + splitTopLevelSqlList(body).forEach((item) => { + const expression = stripProjectionAlias( + item.replace(/^\s*DISTINCT\s+/i, ''), + ); + const identifierMatch = expression.match( + new RegExp(String.raw`^${SQL_IDENTIFIER_PATTERN}$`, 'i'), + ); + if (identifierMatch && normalizeSqlIdentifier(expression) !== '*') { + columns.push(normalizeSqlIdentifier(expression)); + } + }); + } + return columns; +}; + +const findSqlReferenceValidationErrors = ( + sql: string, + manifest?: Manifest, +) => { + const schemas = getManifestModelSchemas(manifest); + if (!schemas.size) { + return []; + } + + const aliases = extractSqlTableAliases(sql, schemas); + const errors: string[] = []; + const qualifiedColumnPattern = new RegExp( + String.raw`(${SQL_IDENTIFIER_PATTERN})\s*\.\s*(${SQL_IDENTIFIER_PATTERN})`, + 'gi', + ); + + let match: RegExpExecArray | null; + while ((match = qualifiedColumnPattern.exec(sql))) { + const qualifier = normalizeSqlIdentifier(match[1]); + const column = normalizeSqlIdentifier(match[2]); + const schema = aliases.get(qualifier.toLowerCase()); + + if (!schema || column === '*') { + continue; + } + + if ( + !schema.columns.has(column.toLowerCase()) && + !schema.columns.has(compactSqlIdentifier(column)) + ) { + errors.push(`${qualifier}.${column}`); + } + } + + const aggregatePattern = new RegExp( + String.raw`\b(AVG|SUM)\s*\(\s*(?:DISTINCT\s+)?(${SQL_IDENTIFIER_PATTERN})(?:\s*\.\s*(${SQL_IDENTIFIER_PATTERN}))?\s*\)`, + 'gi', + ); + while ((match = aggregatePattern.exec(sql))) { + const functionName = match[1].toUpperCase(); + const qualifier = match[3] ? normalizeSqlIdentifier(match[2]) : null; + const column = normalizeSqlIdentifier(match[3] || match[2]); + if (column === '*') { + continue; + } + + const candidateSchemas = qualifier + ? [aliases.get(qualifier.toLowerCase())].filter(isDefined) + : [...new Set(aliases.values())]; + const matchingColumns = candidateSchemas + .map( + (schema) => + schema?.columns.get(column.toLowerCase()) || + schema?.columns.get(compactSqlIdentifier(column)), + ) + .filter(isDefined); + + if (!matchingColumns.length) { + errors.push(qualifier ? `${qualifier}.${column}` : column); + continue; + } + + if ( + matchingColumns.some( + (columnSchema) => !isNumericColumnType(columnSchema.type), + ) + ) { + errors.push( + `${functionName}(${qualifier ? `${qualifier}.` : ''}${column}) uses a non-numeric column`, + ); + } + } + + const activeSchemas = [...new Set(aliases.values())]; + if (activeSchemas.length === 1) { + const [schema] = activeSchemas; + extractSimpleProjectionColumns(sql).forEach((column) => { + if ( + !schema.columns.has(column.toLowerCase()) && + !schema.columns.has(compactSqlIdentifier(column)) + ) { + errors.push(column); + } + }); + } + + return [...new Set(errors)]; +}; + const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { const validNames = getManifestQueryableNames(manifest); if (!validNames.size) { @@ -380,6 +660,16 @@ const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { ].join(', ')}`, ); } + + const invalidColumnReferences = findSqlReferenceValidationErrors( + sql, + manifest, + ); + if (invalidColumnReferences.length) { + throw new Error( + `Generated SQL references column(s) or expressions not valid for the active datasource metadata: ${invalidColumnReferences.join(', ')}`, + ); + } }; export class QueryService implements IQueryService { @@ -415,32 +705,42 @@ export class QueryService implements IQueryService { } = options; const mdl = normalizeDeployedManifestForDatasource(rawMdl, project); const { type: dataSource, connectionInfo } = project; - validateSqlReferencesManifest(sql, mdl); + const normalizedPreview = normalizePreviewSqlForIbis(sql, dataSource, limit); + validateSqlReferencesManifest(normalizedPreview.sql, mdl); if (this.useEngine(dataSource)) { if (dryRun) { logger.debug('Using wren engine to dry run'); - await this.wrenEngineAdaptor.dryRun(sql, { + await this.wrenEngineAdaptor.dryRun(normalizedPreview.sql, { manifest: mdl, - limit, + limit: normalizedPreview.limit, }); return true; } else { logger.debug('Using wren engine to preview'); - const data = await this.wrenEngineAdaptor.previewData(sql, mdl, limit); + const data = await this.wrenEngineAdaptor.previewData( + normalizedPreview.sql, + mdl, + normalizedPreview.limit, + ); return data as PreviewDataResponse; } } else { this.checkDataSourceIsSupported(dataSource); logger.debug('Use ibis adaptor to preview'); if (dryRun) { - return await this.ibisDryRun(sql, dataSource, connectionInfo, mdl); + return await this.ibisDryRun( + normalizedPreview.sql, + dataSource, + connectionInfo, + mdl, + ); } else { return await this.ibisQuery( - sql, + normalizedPreview.sql, dataSource, connectionInfo, mdl, - limit, + normalizedPreview.limit, refresh, cacheEnabled, ); diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 85d6b8d175..0f6873cca8 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -379,6 +379,105 @@ describe('QueryService', () => { expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); }); + it('should reject sql that references columns outside the active manifest before ibis dry run', async () => { + await expect( + queryService.preview('SELECT "orders"."OTD_Date" FROM "orders"', { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'orders', + tableReference: { table: 'orders' }, + columns: [ + { name: 'order_date', type: 'timestamp', isCalculated: false }, + { name: 'quantity', type: 'integer', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }), + ).rejects.toThrow( + 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: orders.OTD_Date', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + + it('should reject unqualified projected columns outside a single active manifest table', async () => { + await expect( + queryService.preview('SELECT tools_required FROM "knowledge_articles"', { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'knowledge_articles', + tableReference: { table: 'knowledge_articles' }, + columns: [ + { name: 'id', type: 'integer', isCalculated: false }, + { name: 'content', type: 'string', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }), + ).rejects.toThrow( + 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: tools_required', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + + it('should reject numeric aggregates on non-numeric manifest columns before ibis planning', async () => { + await expect( + queryService.preview('SELECT AVG("orders"."quantity") FROM "orders"', { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'orders', + tableReference: { table: 'orders' }, + columns: [ + { name: 'quantity', type: 'string', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }), + ).rejects.toThrow( + 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: AVG(orders.quantity) uses a non-numeric column', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + + it('should allow numeric aggregates on numeric manifest columns before ibis planning', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview('SELECT AVG("orders"."quantity") FROM "orders"', { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'orders', + tableReference: { table: 'orders' }, + columns: [ + { name: 'quantity', type: 'integer', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); + }); + it('should allow active manifest table references before ibis dry run', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ correlationId: '123', From 26254d576a33a76842300f7215c4986c31e40392 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 23 Jun 2026 16:41:25 +0530 Subject: [PATCH 0239/1087] Stop chart generation raw SQL fallback --- wren-ui/src/apollo/server/services/askingService.ts | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 11bc85b7a9..793299e4f5 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -952,7 +952,7 @@ export class AskingService implements IAskingService { const project = await this.projectService.getCurrentProject(); const deployment = await this.deployService.getLastDeployment(project.id); - let chartData: PreviewDataResponse | undefined; + let chartData: PreviewDataResponse; try { chartData = (await this.queryService.preview(threadResponse.sql, { project, @@ -963,8 +963,9 @@ export class AskingService implements IAskingService { } catch (error) { const message = error instanceof Error ? error.message : String(error); logger.warn( - `Preview failed before chart generation for response ${threadResponse.id}; falling back to AI service preview. ${message}`, + `Preview failed before chart generation for response ${threadResponse.id}. ${message}`, ); + throw error; } // 1. create a task on AI service to generate the chart @@ -1019,7 +1020,7 @@ export class AskingService implements IAskingService { const project = await this.projectService.getCurrentProject(); const deployment = await this.deployService.getLastDeployment(project.id); - let chartData: PreviewDataResponse | undefined; + let chartData: PreviewDataResponse; try { chartData = (await this.queryService.preview(threadResponse.sql, { project, @@ -1030,8 +1031,9 @@ export class AskingService implements IAskingService { } catch (error) { const message = error instanceof Error ? error.message : String(error); logger.warn( - `Preview failed before chart adjustment for response ${threadResponse.id}; falling back to AI service preview. ${message}`, + `Preview failed before chart adjustment for response ${threadResponse.id}. ${message}`, ); + throw error; } // 1. create a task on AI service to adjust the chart From 46d5280cf4e51ff2159a9311e07580f80694e7cb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 16:45:08 +0530 Subject: [PATCH 0240/1087] Fix datasource metadata validation context --- .../src/pipelines/generation/utils/sql.py | 58 +++++++++++++++---- .../pipelines/generation/test_sql_utils.py | 48 +++++++++++++++ .../dashboardCacheBackgroundTracker.ts | 4 +- .../textBasedAnswerBackgroundTracker.ts | 20 ++++++- .../apollo/server/services/askingService.ts | 28 ++++++--- .../apollo/server/services/queryService.ts | 27 ++++++++- .../services/tests/queryService.test.ts | 22 +++++++ 7 files changed, 184 insertions(+), 23 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1792485393..b0bc9b7cd8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2217,6 +2217,14 @@ def get_sql_generation_model_kwargs(llm_provider: LLMProvider) -> dict: return SQL_GENERATION_MODEL_KWARGS +_SCHEMA_IDENTIFIER_PATTERN = ( + r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)' +) +_SCHEMA_TABLE_REFERENCE_PATTERN = ( + rf"{_SCHEMA_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SCHEMA_IDENTIFIER_PATTERN})*" +) + + def construct_instructions( instructions: list[dict] | None = None, ): @@ -2230,20 +2238,25 @@ def construct_instructions( def construct_valid_table_names(documents: list[Any] | None = None) -> list[str]: - table_names = [] + table_names: set[str] = set() for document in documents or []: content = getattr(document, "content", document) if not isinstance(content, str): continue for match in re.finditer( - r"\bCREATE\s+TABLE\s+([`\"\[]?)([A-Za-z_][A-Za-z0-9_.$]*)\1", + rf"\bCREATE\s+TABLE\s+(?P
{_SCHEMA_TABLE_REFERENCE_PATTERN})", content, flags=re.IGNORECASE, ): - table_names.append(match.group(2)) + for table_name in _table_reference_suffixes(match.group("table")): + table_names.add(table_name) - return sorted(set(table_names)) + for table_reference in extract_sql_table_references(content): + for table_name in _table_reference_suffixes(table_reference): + table_names.add(table_name) + + return sorted(table_names) def construct_valid_table_columns( @@ -2256,11 +2269,11 @@ def construct_valid_table_columns( continue for table_match in re.finditer( - r"\bCREATE\s+TABLE\s+([`\"\[]?)(?P
[A-Za-z_][A-Za-z0-9_.$]*)\1\s*\(", + rf"\bCREATE\s+TABLE\s+(?P
{_SCHEMA_TABLE_REFERENCE_PATTERN})\s*\(", content, flags=re.IGNORECASE, ): - table_name = table_match.group("table") + table_names = _table_reference_suffixes(table_match.group("table")) body_start = table_match.end() depth = 1 body_end = body_start @@ -2273,7 +2286,10 @@ def construct_valid_table_columns( body_end += 1 table_body = content[body_start : body_end - 1] - columns = table_columns.setdefault(table_name, set()) + columns_by_table = [ + table_columns.setdefault(table_name, set()) + for table_name in table_names + ] for line in table_body.splitlines(): line = line.strip() if not line or line.startswith("--"): @@ -2291,7 +2307,8 @@ def construct_valid_table_columns( line, ) if column_match: - columns.add(column_match.group("column")) + for columns in columns_by_table: + columns.add(column_match.group("column")) return { table_name: sorted(columns) @@ -2417,6 +2434,11 @@ def _split_table_reference(table_reference: str) -> list[str]: ] +def _table_reference_suffixes(table_reference: str) -> list[str]: + parts = _split_table_reference(table_reference) + return [".".join(parts[index:]) for index in range(len(parts))] + + def extract_sql_table_references(sql: str) -> list[str]: references = [] for match in _SQL_TABLE_REFERENCE_PATTERN.finditer(sql): @@ -2450,7 +2472,15 @@ def find_invalid_table_references(sql: str, valid_table_names: list[str]) -> lis for table_reference in extract_sql_table_references(sql): normalized_reference = table_reference.lower() - if normalized_reference in valid_tables or normalized_reference in cte_names: + reference_suffixes = { + suffix.lower() for suffix in _table_reference_suffixes(table_reference) + } + if ( + normalized_reference in valid_tables + or normalized_reference in cte_names + or reference_suffixes.intersection(valid_tables) + or reference_suffixes.intersection(cte_names) + ): continue invalid_references.append(table_reference) @@ -2476,7 +2506,13 @@ def _extract_table_aliases( for match in _SQL_TABLE_WITH_ALIAS_PATTERN.finditer(sql): table_reference = ".".join(_split_table_reference(match.group("table"))) normalized_table = table_reference.lower() - if normalized_table not in valid_tables or normalized_table in cte_names: + matched_table = valid_tables.get(normalized_table) + if not matched_table: + for suffix in _table_reference_suffixes(table_reference): + matched_table = valid_tables.get(suffix.lower()) + if matched_table: + break + if not matched_table or normalized_table in cte_names: continue alias = match.group("alias") @@ -2486,7 +2522,7 @@ def _extract_table_aliases( normalized_alias = _normalize_sql_identifier(alias or "").lower() if normalized_alias in _SQL_RESERVED_ALIASES: continue - aliases[normalized_alias] = valid_tables[normalized_table] + aliases[normalized_alias] = matched_table return aliases diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 91f22d3481..6f1637f4bf 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,5 +1,6 @@ from src.pipelines.generation.utils.sql import ( contains_unsupported_mssql_json_access, + construct_valid_table_columns, construct_valid_table_names, extract_sql_generation_result, find_invalid_column_references, @@ -23,6 +24,53 @@ def test_construct_valid_table_names_from_schema_documents(): assert construct_valid_table_names(documents) == ["employees", "repair_logs"] +def test_construct_valid_table_names_includes_ref_sql_source_tables(): + documents = [ + ''' + CREATE TABLE repair_logs ("id" INTEGER); + refSql: SELECT "created_at", "warning_signals" FROM "wrenai"."public"."dbo_repair_logs" + ''', + ] + + assert construct_valid_table_names(documents) == [ + "dbo_repair_logs", + "public.dbo_repair_logs", + "repair_logs", + "wrenai.public.dbo_repair_logs", + ] + + +def test_schema_validation_allows_qualified_suffix_table_references(): + assert find_invalid_table_references( + 'SELECT * FROM "wrenai"."public"."dbo_repair_logs"', + ["dbo_repair_logs"], + ) == [] + + +def test_column_validation_allows_qualified_suffix_table_references(): + sql = ( + 'SELECT "wrenai"."public"."dbo_repair_logs"."warning_signals" ' + 'FROM "wrenai"."public"."dbo_repair_logs"' + ) + + assert find_invalid_column_references( + sql, + {"dbo_repair_logs": ["warning_signals"]}, + ) == [] + + +def test_construct_valid_table_columns_adds_qualified_suffix_tables(): + documents = [ + 'CREATE TABLE "wrenai"."public"."dbo_repair_logs" ("warning_signals" INTEGER);', + ] + + assert construct_valid_table_columns(documents) == { + "dbo_repair_logs": ["warning_signals"], + "public.dbo_repair_logs": ["warning_signals"], + "wrenai.public.dbo_repair_logs": ["warning_signals"], + } + + def test_schema_validation_ignores_null_table_metadata(): assert find_invalid_table_references( 'SELECT * FROM "dbo_tblSales"', diff --git a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts index bad4c69b13..3f30fa5b86 100644 --- a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts @@ -114,7 +114,9 @@ export class DashboardCacheBackgroundTracker { }); // Get project and deployment info - const project = await this.projectService.getCurrentProject(); + const project = await this.projectService.getProjectById( + dashboard.projectId, + ); const deployment = await this.deployService.getLastDeployment(project.id); const mdl = deployment.manifest; const hash = uuidv4(); diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index 79b769e18c..6904573abc 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -4,7 +4,11 @@ import { TextBasedAnswerResult, TextBasedAnswerStatus, } from '../models/adaptor'; -import { ThreadResponse, IThreadResponseRepository } from '../repositories'; +import { + ThreadResponse, + IThreadRepository, + IThreadResponseRepository, +} from '../repositories'; import { IProjectService, IDeployService, @@ -24,6 +28,7 @@ export class TextBasedAnswerBackgroundTracker { private tasks: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; + private threadRepository: IThreadRepository; private threadResponseRepository: IThreadResponseRepository; private projectService: IProjectService; private deployService: IDeployService; @@ -32,18 +37,21 @@ export class TextBasedAnswerBackgroundTracker { constructor({ wrenAIAdaptor, + threadRepository, threadResponseRepository, projectService, deployService, queryService, }: { wrenAIAdaptor: IWrenAIAdaptor; + threadRepository: IThreadRepository; threadResponseRepository: IThreadResponseRepository; projectService: IProjectService; deployService: IDeployService; queryService: IQueryService; }) { this.wrenAIAdaptor = wrenAIAdaptor; + this.threadRepository = threadRepository; this.threadResponseRepository = threadResponseRepository; this.projectService = projectService; this.deployService = deployService; @@ -79,7 +87,15 @@ export class TextBasedAnswerBackgroundTracker { }); threadResponse.answerDetail = fetchingDetail; - const project = await this.projectService.getCurrentProject(); + const thread = await this.threadRepository.findOneBy({ + id: threadResponse.threadId, + }); + if (!thread) { + throw new Error(`Thread ${threadResponse.threadId} not found`); + } + const project = await this.projectService.getProjectById( + thread.projectId, + ); const deployment = await this.deployService.getLastDeployment( project.id, ); diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 793299e4f5..2ae1e66b99 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -488,6 +488,7 @@ export class AskingService implements IAskingService { this.textBasedAnswerBackgroundTracker = new TextBasedAnswerBackgroundTracker({ wrenAIAdaptor, + threadRepository, threadResponseRepository, projectService, deployService, @@ -586,8 +587,8 @@ export class AskingService implements IAskingService { return; } - const project = await this.projectService.getCurrentProject(); - const { manifest } = await this.mdlService.makeCurrentModelMDL(); + const project = await this.projectService.getProjectById(thread.projectId); + const { manifest } = await this.mdlService.makeModelMDL(project); const threadResponses = await this.threadResponseRepository.findAllBy({ threadId, @@ -950,7 +951,7 @@ export class AskingService implements IAskingService { return threadResponse; } - const project = await this.projectService.getCurrentProject(); + const project = await this.getProjectForThreadResponse(threadResponse); const deployment = await this.deployService.getLastDeployment(project.id); let chartData: PreviewDataResponse; try { @@ -1018,7 +1019,7 @@ export class AskingService implements IAskingService { return threadResponse; } - const project = await this.projectService.getCurrentProject(); + const project = await this.getProjectForThreadResponse(threadResponse); const deployment = await this.deployService.getLastDeployment(project.id); let chartData: PreviewDataResponse; try { @@ -1086,7 +1087,7 @@ export class AskingService implements IAskingService { if (!response) { throw new Error(`Thread response ${responseId} not found`); } - const project = await this.projectService.getCurrentProject(); + const project = await this.getProjectForThreadResponse(response); const deployment = await this.deployService.getLastDeployment(project.id); const mdl = deployment.manifest; const eventName = TelemetryEvent.HOME_PREVIEW_ANSWER; @@ -1126,7 +1127,7 @@ export class AskingService implements IAskingService { if (!response) { throw new Error(`Thread response ${responseId} not found`); } - const project = await this.projectService.getCurrentProject(); + const project = await this.getProjectForThreadResponse(response); const deployment = await this.deployService.getLastDeployment(project.id); const mdl = deployment.manifest; const steps = response?.breakdownDetail?.steps; @@ -1236,12 +1237,23 @@ export class AskingService implements IAskingService { return updatedResponse; } - private async getDeployId() { - const { id } = await this.projectService.getCurrentProject(); + private async getDeployId(projectId?: number) { + const id = projectId ?? (await this.projectService.getCurrentProject()).id; const lastDeploy = await this.deployService.getLastDeployment(id); return lastDeploy.hash; } + private async getProjectForThreadResponse(threadResponse: ThreadResponse) { + const thread = await this.threadRepository.findOneBy({ + id: threadResponse.threadId, + }); + if (!thread) { + throw new Error(`Thread ${threadResponse.threadId} not found`); + } + + return this.projectService.getProjectById(thread.projectId); + } + public async adjustThreadResponseWithSQL( threadResponseId: number, input: AdjustmentSqlInput, diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index e97d397314..c3adfd5e64 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -332,6 +332,22 @@ const extractSqlTableReferences = (sql: string) => { return references; }; +const addTableReferenceName = ( + names: Set, + parts: Array, +) => { + const normalizedParts = parts + .filter((part): part is string => Boolean(part)) + .map((part) => part.toLowerCase()); + if (!normalizedParts.length) { + return; + } + + for (let index = 0; index < normalizedParts.length; index += 1) { + names.add(normalizedParts.slice(index).join('.')); + } +}; + const extractCteNames = (sql: string) => { const cteNames = new Set(); const ctePattern = new RegExp( @@ -350,7 +366,16 @@ const getManifestQueryableNames = (manifest?: Manifest) => { for (const model of manifest?.models || []) { if (model.name) names.add(model.name.toLowerCase()); if (model.tableReference?.table) { - names.add(model.tableReference.table.toLowerCase()); + addTableReferenceName(names, [ + model.tableReference.catalog, + model.tableReference.schema, + model.tableReference.table, + ]); + } + if (model.refSql) { + for (const reference of extractSqlTableReferences(model.refSql)) { + addTableReferenceName(names, splitTableReference(reference)); + } } } for (const view of manifest?.views || []) { diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 0f6873cca8..475c836274 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -499,6 +499,28 @@ describe('QueryService', () => { expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); }); + + it('should allow source tables referenced by active manifest refSql before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview('SELECT * FROM dbo_repair_logs', { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'repair_logs', + refSql: 'SELECT created_at, ticket_id FROM dbo_repair_logs', + }, + ], + }, + dryRun: true, + }); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); + }); }); class MockTelemetry { From dc4e27b6ea8ae6d0cd59054fd6b8672a5e8ccf77 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 16:58:44 +0530 Subject: [PATCH 0241/1087] Fix thread response context and SQL identifier quoting --- .../src/pipelines/generation/utils/sql.py | 21 ++++++++++++++- .../pipelines/generation/test_sql_utils.py | 15 +++++++++++ .../apollo/server/services/askingService.ts | 27 ++++++++++--------- 3 files changed, 49 insertions(+), 14 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b0bc9b7cd8..09400312bc 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1438,8 +1438,27 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: ) +def _normalize_identifier_quote_syntax(sql: str) -> str: + normalized = re.sub( + r"`([^`]+)`", + lambda match: _quote_sql_identifier(match.group(1)), + sql, + ) + normalized = re.sub( + r"\[([^\]]+)\]", + lambda match: _quote_sql_identifier(match.group(1)), + normalized, + ) + normalized = re.sub( + r'"{2,}([A-Za-z_][A-Za-z0-9_$]*)"{2,}', + lambda match: _quote_sql_identifier(match.group(1)), + normalized, + ) + return normalized + + def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: - normalized = sql + normalized = _normalize_identifier_quote_syntax(sql) normalized_data_source = normalize_data_source(data_source) if normalized_data_source == "MSSQL": diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6f1637f4bf..6f6659d1ac 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -71,6 +71,21 @@ def test_construct_valid_table_columns_adds_qualified_suffix_tables(): } +def test_normalize_generation_result_sql_standardizes_identifier_quotes(): + sql = ( + 'SELECT COUNT(*) AS `num_tags`, SUM(`tokenCost`) AS `popularity` ' + 'FROM `dbo_kb_articles` WHERE (""""category"""" = \'Ticket Sourcing\')' + ) + + normalized = normalize_generation_result_sql(sql, data_source="mssql") + + assert "`" not in normalized + assert '""""category""""' not in normalized + assert '"num_tags"' in normalized + assert '"tokenCost"' in normalized + assert '"category" = \'Ticket Sourcing\'' in normalized + + def test_schema_validation_ignores_null_table_metadata(): assert find_invalid_table_references( 'SELECT * FROM "dbo_tblSales"', diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 2ae1e66b99..4b92ecbfdf 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -664,14 +664,22 @@ export class AskingService implements IAskingService { ): Promise { const { threadId, language } = payload; const currentProject = await this.projectService.getCurrentProject(); - const projectId = payload.projectId ?? currentProject.id; - if (projectId !== currentProject.id) { - throw new Error(`Project ${projectId} is not the active project`); - } + let projectId = payload.projectId ?? currentProject.id; if (threadId) { - await this.ensureThreadInCurrentProject(threadId); + const thread = await this.threadRepository.findOneBy({ id: threadId }); + if (!thread) { + throw new Error(`Thread ${threadId} not found`); + } + if (payload.projectId && payload.projectId !== thread.projectId) { + throw new Error( + `Thread ${threadId} does not belong to project ${payload.projectId}`, + ); + } + projectId = thread.projectId; + } else if (projectId !== currentProject.id) { + throw new Error(`Project ${projectId} is not the active project`); } - const deployId = await this.getDeployId(); + const deployId = await this.getDeployId(projectId); // if it's a follow-up question, then the input will have a threadId // then use the threadId to get the sql and get the steps of last thread response @@ -705,7 +713,6 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } - await this.ensureThreadInCurrentProject(threadResponse.threadId); // get the original question and ask again const question = threadResponse.question; @@ -869,7 +876,6 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } - await this.ensureThreadInCurrentProject(threadResponse.threadId); // 1. create a task on AI service to generate the detail const response = await this.wrenAIAdaptor.generateAskDetail({ @@ -906,7 +912,6 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } - await this.ensureThreadInCurrentProject(threadResponse.threadId); if (isAnswerGenerationInProgress(threadResponse.answerDetail?.status)) { logger.debug( @@ -942,7 +947,6 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } - await this.ensureThreadInCurrentProject(threadResponse.threadId); if (isChartGenerationInProgress(threadResponse.chartDetail?.status)) { logger.debug( @@ -1007,7 +1011,6 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } - await this.ensureThreadInCurrentProject(threadResponse.threadId); if ( isChartGenerationInProgress(threadResponse.chartDetail?.status) && @@ -1078,7 +1081,6 @@ export class AskingService implements IAskingService { if (!response) { return null; } - await this.ensureThreadInCurrentProject(response.threadId); return response; } @@ -1360,7 +1362,6 @@ export class AskingService implements IAskingService { if (!threadId) { return []; } - await this.ensureThreadInCurrentProject(threadId); let responses = await this.threadResponseRepository.getResponsesWithThread( threadId, 10, From ea86b472b61432540c2246e1cc431f86807e9a21 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 17:06:58 +0530 Subject: [PATCH 0242/1087] Parse semantic columns for SQL validation --- .../src/pipelines/generation/utils/sql.py | 122 ++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 80 ++++++++++++ 2 files changed, 202 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 09400312bc..b9472d1c71 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2242,6 +2242,27 @@ def get_sql_generation_model_kwargs(llm_provider: LLMProvider) -> dict: _SCHEMA_TABLE_REFERENCE_PATTERN = ( rf"{_SCHEMA_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SCHEMA_IDENTIFIER_PATTERN})*" ) +_SEMANTIC_TABLE_NAME_KEYS = ( + "name", + "referenceName", + "sourceTableName", + "tableName", + "table", +) +_SEMANTIC_COLUMN_CONTAINER_KEYS = ( + "columns", + "fields", + "calculatedFields", + "dimensions", + "measures", +) +_SEMANTIC_COLUMN_NAME_KEYS = ( + "name", + "referenceName", + "sourceColumnName", + "columnName", + "fieldName", +) def construct_instructions( @@ -2256,6 +2277,99 @@ def construct_instructions( return _instructions +def _parse_semantic_metadata_content(content: str) -> Any | None: + content = content.strip() + if not content: + return None + + if content.startswith("```"): + content = re.sub(r"^```(?:json|mdl)?\s*", "", content, flags=re.IGNORECASE) + content = re.sub(r"\s*```$", "", content) + + candidates = [content] + object_start = content.find("{") + object_end = content.rfind("}") + if object_start >= 0 and object_end > object_start: + candidates.append(content[object_start : object_end + 1]) + array_start = content.find("[") + array_end = content.rfind("]") + if array_start >= 0 and array_end > array_start: + candidates.append(content[array_start : array_end + 1]) + + for candidate in candidates: + try: + return orjson.loads(candidate) + except orjson.JSONDecodeError: + continue + + return None + + +def _semantic_name_values(metadata: dict[str, Any], keys: tuple[str, ...]) -> list[str]: + values = [] + for key in keys: + value = metadata.get(key) + if isinstance(value, str) and value.strip(): + values.append(value.strip()) + return values + + +def _semantic_column_names(metadata: Any) -> set[str]: + columns: set[str] = set() + if isinstance(metadata, dict): + for column_name in _semantic_name_values(metadata, _SEMANTIC_COLUMN_NAME_KEYS): + columns.add(column_name) + for key in _SEMANTIC_COLUMN_CONTAINER_KEYS: + value = metadata.get(key) + if value is not None: + columns.update(_semantic_column_names(value)) + elif isinstance(metadata, list): + for item in metadata: + columns.update(_semantic_column_names(item)) + + return columns + + +def _construct_semantic_table_columns(content: str) -> dict[str, set[str]]: + parsed_content = _parse_semantic_metadata_content(content) + if parsed_content is None: + return {} + + table_columns: dict[str, set[str]] = {} + + def collect(metadata: Any) -> None: + if isinstance(metadata, list): + for item in metadata: + collect(item) + return + + if not isinstance(metadata, dict): + return + + column_containers = [ + metadata.get(key) + for key in _SEMANTIC_COLUMN_CONTAINER_KEYS + if metadata.get(key) is not None + ] + if column_containers: + columns = set() + for container in column_containers: + columns.update(_semantic_column_names(container)) + + if columns: + for table_reference in _semantic_name_values( + metadata, _SEMANTIC_TABLE_NAME_KEYS + ): + for table_name in _table_reference_suffixes(table_reference): + table_columns.setdefault(table_name, set()).update(columns) + + for value in metadata.values(): + collect(value) + + collect(parsed_content) + return table_columns + + def construct_valid_table_names(documents: list[Any] | None = None) -> list[str]: table_names: set[str] = set() for document in documents or []: @@ -2263,6 +2377,9 @@ def construct_valid_table_names(documents: list[Any] | None = None) -> list[str] if not isinstance(content, str): continue + for table_name in _construct_semantic_table_columns(content): + table_names.add(table_name) + for match in re.finditer( rf"\bCREATE\s+TABLE\s+(?P
{_SCHEMA_TABLE_REFERENCE_PATTERN})", content, @@ -2287,6 +2404,11 @@ def construct_valid_table_columns( if not isinstance(content, str): continue + for table_name, columns in _construct_semantic_table_columns( + content + ).items(): + table_columns.setdefault(table_name, set()).update(columns) + for table_match in re.finditer( rf"\bCREATE\s+TABLE\s+(?P
{_SCHEMA_TABLE_REFERENCE_PATTERN})\s*\(", content, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6f6659d1ac..894ea1fcd2 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -71,6 +71,86 @@ def test_construct_valid_table_columns_adds_qualified_suffix_tables(): } +def test_construct_valid_table_columns_from_semantic_metadata_document(): + documents = [ + """ + { + "models": [ + { + "name": "dbo_new_orders", + "referenceName": "sales.dbo_new_orders", + "columns": [ + {"name": "business"}, + {"name": "market"}, + {"name": "customer_name"}, + {"name": "product_name"}, + {"name": "order_value"} + ], + "calculatedFields": [ + {"name": "order_month"} + ] + } + ] + } + """, + ] + + assert construct_valid_table_names(documents) == [ + "dbo_new_orders", + "sales.dbo_new_orders", + ] + assert construct_valid_table_columns(documents) == { + "dbo_new_orders": [ + "business", + "customer_name", + "market", + "order_month", + "order_value", + "product_name", + ], + "sales.dbo_new_orders": [ + "business", + "customer_name", + "market", + "order_month", + "order_value", + "product_name", + ], + } + + +def test_column_validation_uses_semantic_metadata_document_columns(): + documents = [ + """ + { + "models": [ + { + "name": "dbo_new_orders", + "columns": [ + {"name": "customer_name"}, + {"name": "product_name"}, + {"name": "order_value"} + ] + } + ] + } + """, + ] + valid_table_columns = construct_valid_table_columns(documents) + + assert find_invalid_column_references( + 'SELECT "dbo_new_orders"."customer_name", ' + '"dbo_new_orders"."product_name", ' + '"dbo_new_orders"."order_value" ' + 'FROM "dbo_new_orders"', + valid_table_columns, + ) == [] + assert find_invalid_column_references( + 'SELECT "dbo_new_orders"."missing_value" FROM "dbo_new_orders"', + valid_table_columns, + ) == ["dbo_new_orders.missing_value"] + + def test_normalize_generation_result_sql_standardizes_identifier_quotes(): sql = ( 'SELECT COUNT(*) AS `num_tags`, SUM(`tokenCost`) AS `popularity` ' From 5543e2c7ec941adda4dafa9bc9fdf39611fcfeeb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 17:23:36 +0530 Subject: [PATCH 0243/1087] Include datasource identity in deployment sync --- .../apollo/server/resolvers/modelResolver.ts | 10 ++--- .../server/resolvers/projectResolver.ts | 2 +- .../apollo/server/services/deployService.ts | 38 ++++++++++++++----- .../services/tests/deployService.test.ts | 35 ++++++++++++++++- 4 files changed, 69 insertions(+), 16 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 76d9abc504..dfa49c9078 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -201,13 +201,13 @@ export class ModelResolver { public async checkModelSync(_root: any, _args: any, ctx: IContext) { try { - const { id } = await ctx.projectService.getCurrentProject(); + const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest, id); - const lastDeploy = await ctx.deployService.getLastDeployment(id); + const currentHash = ctx.deployService.createMDLHash(manifest, project); + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = - await ctx.deployService.getInProgressDeployment(id); + await ctx.deployService.getInProgressDeployment(project.id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } @@ -236,7 +236,7 @@ export class ModelResolver { const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy( manifest, - project.id, + project, args.force, ); diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 6ca33be5cc..99bbb9b467 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -640,7 +640,7 @@ export class ProjectResolver { private async deploy(ctx: IContext) { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const deployRes = await ctx.deployService.deploy(manifest, project.id); + const deployRes = await ctx.deployService.deploy(manifest, project); // Recommendation generation depends on a successful deployment because // question validation calls previewSql against the deployed manifest. diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 3c751d3156..7a330bf03e 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -5,6 +5,7 @@ import { DeployStatusEnum, IDeployLogRepository, } from '../repositories/deployLogRepository'; +import { Project } from '../repositories/projectRepository'; import { Manifest } from '../mdl/type'; import { createHash } from 'node:crypto'; import { getLogger } from '@server/utils'; @@ -29,12 +30,12 @@ export interface MDLSyncResponse { export interface IDeployService { deploy( manifest: Manifest, - projectId: number, + project: Project | number, force?: boolean, ): Promise; getLastDeployment(projectId: number): Promise; getInProgressDeployment(projectId: number): Promise; - createMDLHash(manifest: Manifest, projectId: number): string; + createMDLHash(manifest: Manifest, project: Project | number): string; getMDLByHash(hash: string): Promise; deleteAllByProjectId(projectId: number): Promise; } @@ -58,7 +59,7 @@ export class DeployService implements IDeployService { this.telemetry = telemetry; } - public async getLastDeployment(projectId) { + public async getLastDeployment(projectId: number) { const lastDeploy = await this.deployLogRepository.findLastProjectDeployLog(projectId); if (!lastDeploy) { @@ -67,17 +68,22 @@ export class DeployService implements IDeployService { return lastDeploy; } - public async getInProgressDeployment(projectId) { + public async getInProgressDeployment(projectId: number) { return await this.deployLogRepository.findInProgressProjectDeployLog( projectId, ); } - public async deploy(manifest, projectId, force = false) { + public async deploy( + manifest: Manifest, + project: Project | number, + force = false, + ) { const eventName = TelemetryEvent.MODELING_DEPLOY_MDL; + const projectId = typeof project === 'number' ? project : project.id; try { // generate hash of manifest - const hash = this.createMDLHash(manifest, projectId); + const hash = this.createMDLHash(manifest, project); logger.debug(`Deploying model, hash: ${hash}`); if (!force) { @@ -139,9 +145,23 @@ export class DeployService implements IDeployService { } } - public createMDLHash(manifest: Manifest, projectId: number) { - const manifestStr = JSON.stringify(manifest); - const content = `${projectId} ${manifestStr}`; + public createMDLHash(manifest: Manifest, project: Project | number) { + const projectFingerprint = + typeof project === 'number' + ? { id: project } + : { + id: project.id, + type: project.type, + version: project.version, + catalog: project.catalog, + schema: project.schema, + sampleDataset: project.sampleDataset, + connectionInfo: project.connectionInfo, + }; + const content = JSON.stringify({ + project: projectFingerprint, + manifest, + }); const hash = createHash('sha1').update(content).digest('hex'); return hash; } diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 14f922de85..31688702e2 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -10,7 +10,7 @@ describe('DeployService', () => { beforeEach(() => { mockTelemetry = { sendEvent: jest.fn() }; - mockWrenAIAdaptor = { deploy: jest.fn() }; + mockWrenAIAdaptor = { deploy: jest.fn(), delete: jest.fn() }; mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), createOne: jest.fn(), @@ -72,5 +72,38 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); + it('should include datasource identity in deployment hash', () => { + const manifest = { models: [{ name: 'orders', columns: [] }] }; + const project = { + id: 1, + type: 'mssql', + version: '16', + catalog: 'catalog', + schema: 'dbo', + sampleDataset: null, + connectionInfo: { + host: 'db-a', + port: 1433, + database: 'sales_a', + }, + }; + + const sameMdlDifferentDatabaseHash = deployService.createMDLHash( + manifest, + { + ...project, + connectionInfo: { + ...project.connectionInfo, + host: 'db-b', + database: 'sales_b', + }, + }, + ); + + expect(deployService.createMDLHash(manifest, project)).not.toEqual( + sameMdlDifferentDatabaseHash, + ); + }); + // Add more tests here to cover other scenarios and error handling }); From 35dcfd718032f45dac3fa7909e52e78df4a4d431 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 17:39:22 +0530 Subject: [PATCH 0244/1087] Ground asking pipelines in active datasource schema --- .../src/pipelines/generation/utils/sql.py | 33 +++++++++++ .../retrieval/db_schema_retrieval.py | 10 ++-- wren-ai-service/src/web/v1/services/ask.py | 56 ++++++++++++++++--- .../pipelines/generation/test_sql_utils.py | 21 +++++++ .../apollo/server/services/queryService.ts | 9 ++- .../services/tests/queryService.test.ts | 28 ++++++++++ 6 files changed, 144 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b9472d1c71..99bd1f6a91 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2826,6 +2826,39 @@ def find_invalid_column_references( if column.lower() not in valid_columns: invalid_references.append(f"{qualifier}.{column}") + referenced_tables = { + aliases.get(table_reference.lower()) + for table_reference in extract_sql_table_references(sql) + } + referenced_tables = {table for table in referenced_tables if table} + if len(referenced_tables) == 1: + table_name = next(iter(referenced_tables)) + valid_columns = { + str(col).lower() + for col in valid_table_columns.get(table_name, []) + if col is not None + } + valid_compact_columns = { + _compact_sql_identifier(col) + for col in valid_table_columns.get(table_name, []) + if col is not None + } + for start, end in _find_select_list_spans(sql): + for item in _split_top_level_select_items(sql[start:end]): + expression = _strip_projection_alias( + re.sub(r"^\s*DISTINCT\s+", "", item, flags=re.IGNORECASE) + ) + if not re.fullmatch(_SQL_IDENTIFIER_PATTERN, expression.strip()): + continue + column = _normalize_sql_identifier(expression) + if column == "*": + continue + if ( + column.lower() not in valid_columns + and _compact_sql_identifier(column) not in valid_compact_columns + ): + invalid_references.append(column) + return sorted(set(invalid_references)) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index a561f470b4..939f7d1934 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -179,12 +179,12 @@ def expand_business_terms_for_retrieval(query: str) -> str: [ query, "Business analytics aliases:", - "repair trends repair volume repair counts debug entries debug fixes", - "average debug hours turnaround time resolved entries failure category failure code", + "throughput trend volume count counts average total ranking top bottom grouped distribution", + "business unit manufacturing unit department location site plant team region category status", + "repair trends repair volume repair counts debug entries debug fixes failure code", "monthly trend quarter grouped by month bar chart line chart", - "top common pcb failures top 10 failures most common failure categories", - "failure patterns category occurrences debugentryid failuresys material workorder serialnumber", - "dbo_DebugEntries dbo_failure_patterns dbo_repair_logs created_at failedat datein dateout", + "top common failures most common failure categories", + "failure patterns category occurrences material workorder serial number", "sales revenue amount sales value sales performance salesperson ranking", "customer sales top customers customer growth orders invoices margin quantity", ] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ee638070cb..493cc01a6b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -8,6 +8,12 @@ from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + find_invalid_column_references, + find_invalid_table_references, +) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -1919,6 +1925,35 @@ def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: return None return AskResult(sql=sql.strip(), type="llm") + def _build_validated_ask_result_from_sql( + self, + sql: Optional[str], + table_ddls: list[str], + ) -> Optional[AskResult]: + ask_result = self._build_ask_result_from_sql(sql) + if not ask_result: + return None + + invalid_tables = find_invalid_table_references( + ask_result.sql, + construct_valid_table_names(table_ddls), + ) + invalid_columns = find_invalid_column_references( + ask_result.sql, + construct_valid_table_columns(table_ddls), + ) + if invalid_tables or invalid_columns: + logger.warning( + "Ignoring heuristic SQL because it is not valid for active schema. " + "invalid_tables=%s invalid_columns=%s sql=%s", + invalid_tables, + invalid_columns, + ask_result.sql, + ) + return None + + return ask_result + def _build_failed_text_to_sql_response( self, trace_id: Optional[str], @@ -2068,8 +2103,9 @@ async def ask( query_id, user_query, ) - if ask_result := self._build_ask_result_from_sql( - heuristic_sql + if ask_result := self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, ): api_results = [ask_result] if not self._is_stopped(query_id, self._ask_results): @@ -2086,7 +2122,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback did not produce a valid SELECT statement." + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." historical_question = await self._run_with_timeout( "Historical question retrieval", @@ -2403,10 +2439,13 @@ async def ask( query_id, user_query, ) - ask_result = self._build_ask_result_from_sql(heuristic_sql) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + ) if not ask_result: invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback did not produce a valid SELECT statement." + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( self._build_failed_text_to_sql_response( @@ -2737,10 +2776,13 @@ async def ask( query_id, user_query, ) - ask_result = self._build_ask_result_from_sql(heuristic_sql) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + ) if not ask_result: invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback did not produce a valid SELECT statement." + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." else: api_results = [ask_result] if not self._is_stopped(query_id, self._ask_results): diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 894ea1fcd2..6d37b4bc71 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -59,6 +59,27 @@ def test_column_validation_allows_qualified_suffix_table_references(): ) == [] +def test_column_validation_rejects_invalid_unqualified_projection_for_single_table(): + assert find_invalid_column_references( + 'SELECT warning_signals FROM "dbo_repair_logs"', + {"dbo_repair_logs": ["id", "created_at", "status"]}, + ) == ["warning_signals"] + + +def test_column_validation_rejects_invalid_unqualified_projection_alias(): + assert find_invalid_column_references( + 'SELECT categories AS category_count FROM "dbo_repair_logs"', + {"dbo_repair_logs": ["id", "created_at", "status"]}, + ) == ["categories"] + + +def test_column_validation_allows_valid_unqualified_projection_for_single_table(): + assert find_invalid_column_references( + 'SELECT status AS repair_status FROM "dbo_repair_logs"', + {"dbo_repair_logs": ["id", "created_at", "status"]}, + ) == [] + + def test_construct_valid_table_columns_adds_qualified_suffix_tables(): documents = [ 'CREATE TABLE "wrenai"."public"."dbo_repair_logs" ("warning_signals" INTEGER);', diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index c3adfd5e64..df46ea7c28 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -670,10 +670,17 @@ const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { const cteNames = extractCteNames(sql); const invalidReferences = extractSqlTableReferences(sql).filter((reference) => { const normalized = reference.toLowerCase(); - const lastPart = splitTableReference(reference).pop()?.toLowerCase(); + const parts = splitTableReference(reference); + const lastPart = parts[parts.length - 1]?.toLowerCase(); + const firstPart = parts[0]?.toLowerCase(); + const suffixes = parts.map((_, index) => + parts.slice(index).join('.').toLowerCase(), + ); return ( !validNames.has(normalized) && !cteNames.has(normalized) && + !suffixes.some((suffix) => validNames.has(suffix)) && + !(parts.length === 2 && firstPart && validNames.has(firstPart)) && (!lastPart || !validNames.has(lastPart)) ); }); diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 475c836274..e51b2bd27d 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -404,6 +404,34 @@ describe('QueryService', () => { expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); }); + it('should not treat model column expressions as missing table references', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT DATEPART(WEEK, "dbo_DebugEntries"."DateIn") AS "week" FROM "dbo_DebugEntries"', + { + project: { type: DataSourceName.MSSQL, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_DebugEntries', + tableReference: { table: 'dbo_DebugEntries' }, + columns: [ + { name: 'DateIn', type: 'timestamp', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalled(); + }); + it('should reject unqualified projected columns outside a single active manifest table', async () => { await expect( queryService.preview('SELECT tools_required FROM "knowledge_articles"', { From f5aba9375f6b4b05f26ad5b33a0e28c5fd6afd7e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 17:54:40 +0530 Subject: [PATCH 0245/1087] Persist streamed text answers reliably --- .../src/pipelines/generation/sql_answer.py | 3 ++ .../apollo/server/services/askingService.ts | 9 +++-- .../pages/api/ask_task/streaming_answer.ts | 34 ++++++++++++++----- 3 files changed, 34 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 81289081b5..948cff0291 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -32,6 +32,9 @@ 6. Answer must be in the same language user specified. 7. Do not include ```markdown or ``` in the answer. 8. If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. +9. Always produce a narrative answer. Never return an empty response. +10. If the user asks for a chart or trend, still summarize the result in words and mention the chart-ready fields. +11. If the data contains only raw rows or a single column, summarize what those rows show, mention the visible date/category range when possible, and state that the result table contains the detailed rows. ### OUTPUT FORMAT diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 4b92ecbfdf..d71eca9964 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1221,8 +1221,11 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${responseId} not found`); } - if (response.answerDetail?.status === status) { - return; + if ( + response.answerDetail?.status === status && + (!content || response.answerDetail?.content === content) + ) { + return response; } const updatedResponse = await this.threadResponseRepository.updateOne( @@ -1231,7 +1234,7 @@ export class AskingService implements IAskingService { answerDetail: { ...response.answerDetail, status, - content, + content: content ?? response.answerDetail?.content, }, }, ); diff --git a/wren-ui/src/pages/api/ask_task/streaming_answer.ts b/wren-ui/src/pages/api/ask_task/streaming_answer.ts index c878bbae94..eea12a95a6 100644 --- a/wren-ui/src/pages/api/ask_task/streaming_answer.ts +++ b/wren-ui/src/pages/api/ask_task/streaming_answer.ts @@ -29,6 +29,24 @@ class ContentMap { const contentMap = new ContentMap(); +const parseSSEMessages = (chunk: Buffer): string[] => { + return chunk + .toString('utf-8') + .split(/\r?\n/) + .filter((line) => line.startsWith('data: ')) + .map((line) => line.slice('data: '.length).trim()) + .filter(Boolean) + .flatMap((payload) => { + try { + const eventData = JSON.parse(payload); + return eventData?.message ? [String(eventData.message)] : []; + } catch (error) { + console.error(`Failed to parse streaming answer payload: ${payload}`); + return []; + } + }); +}; + export default async function handler( req: NextApiRequest, res: NextApiResponse, @@ -69,22 +87,17 @@ export default async function handler( } const stream = await wrenAIAdaptor.streamTextBasedAnswer(queryId); + let streamEnded = false; stream.on('data', (chunk) => { - // pass the chunk directly to the client - const chunkString = chunk.toString('utf-8'); - let message = ''; - const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); - if (match && match[1]) { - message = match[1]; - } else { - console.log(`not able to match: ${chunkString}`); + for (const message of parseSSEMessages(chunk)) { + contentMap.appendContent(queryId, message); } - contentMap.appendContent(queryId, message); res.write(chunk); }); stream.on('end', () => { + streamEnded = true; res.write(`data: ${JSON.stringify({ done: true })}\n\n`); res.end(); askingService @@ -122,6 +135,9 @@ export default async function handler( // destroy the stream if the client closes the connection req.on('close', () => { + if (streamEnded) { + return; + } stream.destroy(); askingService .changeThreadResponseAnswerDetailStatus( From 784e441ec3e58376d2e509bf67ba2a093f85cf19 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 19:31:39 +0530 Subject: [PATCH 0246/1087] Regenerate blank finished answers --- .../pages/home/promptThread/AnswerResult.tsx | 6 ++++++ wren-ui/src/pages/api/ask_task/streaming_answer.ts | 11 ++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx index 697016d966..9c751867f0 100644 --- a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx +++ b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx @@ -179,6 +179,11 @@ const AdjustmentInformation = (props: { const isNeedGenerateAnswer = (answerDetail: ThreadResponseAnswerDetail) => { if (!answerDetail) return true; + const isFinishedWithoutContent = + answerDetail?.status === ThreadResponseAnswerStatus.FINISHED && + !answerDetail?.content?.trim(); + if (isFinishedWithoutContent) return true; + const isFinished = getAnswerIsFinished(answerDetail?.status); // it means the background task has not started yet, but answer is pending for generating const isProcessing = [ @@ -289,6 +294,7 @@ export default function AnswerResult(props: Props) { askingTask?.status, adjustmentTask?.status, answerDetail?.status, + answerDetail?.content, ]); useEffect(() => { diff --git a/wren-ui/src/pages/api/ask_task/streaming_answer.ts b/wren-ui/src/pages/api/ask_task/streaming_answer.ts index eea12a95a6..6a5f604159 100644 --- a/wren-ui/src/pages/api/ask_task/streaming_answer.ts +++ b/wren-ui/src/pages/api/ask_task/streaming_answer.ts @@ -47,6 +47,13 @@ const parseSSEMessages = (chunk: Buffer): string[] => { }); }; +const buildFallbackAnswer = (question: string) => + [ + `I found results for: **${question}**.`, + '', + 'The result table below contains the data returned from the active datasource. Use the visible fields and rows to review the detailed records, and switch to the chart tab when a visualization is available.', + ].join('\n'); + export default async function handler( req: NextApiRequest, res: NextApiResponse, @@ -98,13 +105,15 @@ export default async function handler( stream.on('end', () => { streamEnded = true; + const finalContent = + contentMap.getContent(queryId)?.trim() || buildFallbackAnswer(response.question); res.write(`data: ${JSON.stringify({ done: true })}\n\n`); res.end(); askingService .changeThreadResponseAnswerDetailStatus( Number(responseId), ThreadResponseAnswerStatus.FINISHED, - contentMap.getContent(queryId), + finalContent, ) .then(() => { console.log( From b409fe8f7f6e145638f456a6a40cf9891af85ee9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 19:54:50 +0530 Subject: [PATCH 0247/1087] Guard original SQL datasource option --- .../home/promptThread/ViewSQLTabContent.tsx | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx index d52c103f4d..135f8e2e5f 100644 --- a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx @@ -77,6 +77,9 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { const { hasNativeSQL, dataSourceType } = nativeSQLResult; const showNativeSQL = hasNativeSQL; + const dataSourceOption = dataSourceType + ? DATA_SOURCE_OPTIONS[dataSourceType] + : undefined; const sqls = nativeSQLResult.nativeSQLMode && nativeSQLResult.loading === false @@ -132,15 +135,17 @@ export default function ViewSQLTabContent(props: AnswerResultProps) {
{nativeSQLResult.nativeSQLMode ? ( <> - + {dataSourceOption?.logo && ( + + )} - {DATA_SOURCE_OPTIONS[dataSourceType].label} + {dataSourceOption?.label || 'Original SQL'} ) : ( From ad16b02b233c880e920e4c6fa03c6985fd5cdf6e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 24 Jun 2026 21:53:40 +0530 Subject: [PATCH 0248/1087] Make fallback charts follow question intent --- .../src/pipelines/generation/utils/chart.py | 226 ++++++++++++++++-- 1 file changed, 201 insertions(+), 25 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index a1f7c90409..8e3778a4d8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -14,7 +14,10 @@ def _humanize_title(name: str | None) -> str: - return str(name or "").replace("_", " ").strip().title() + text = str(name or "") + text = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", " ", text) + text = re.sub(r"[_\s]+", " ", text) + return text.strip().title() def _detect_requested_chart_type(query: str | None) -> str: @@ -41,6 +44,153 @@ def _safe_column_names(columns: list[Any]) -> list[str]: return [str(column) for column in columns if column is not None and str(column)] +def _normalize_identifier(value: str | None) -> str: + return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) + + +def _identifier_tokens(value: str | None) -> set[str]: + text = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", " ", str(value or "")) + return { + token + for token in re.split(r"[^a-zA-Z0-9]+", text.lower()) + if len(token) > 1 + } + + +def _query_relevant_columns(query: str | None, columns: list[str]) -> list[str]: + normalized_query = str(query or "").lower() + compact_query = _normalize_identifier(normalized_query) + query_tokens = _identifier_tokens(normalized_query) + + scored_columns: list[tuple[int, int, str]] = [] + for index, column in enumerate(columns): + tokens = _identifier_tokens(column) + compact_column = _normalize_identifier(column) + score = 0 + + if compact_column and compact_column in compact_query: + score += 100 + + matched_tokens = tokens.intersection(query_tokens) + score += len(matched_tokens) * 20 + + # Prefer multi-word business dimensions that the user explicitly asks + # for, e.g. "Business Unit" matching BusinessUnit. + if tokens and tokens.issubset(query_tokens): + score += 40 + + if score: + scored_columns.append((-score, index, column)) + + return [column for _, _, column in sorted(scored_columns)] + + +def _select_measure_column(query: str | None, quantitative: list[str]) -> str | None: + if not quantitative: + return None + + relevant = _query_relevant_columns(query, quantitative) + if relevant: + return relevant[0] + + metric_keywords = ( + "count", + "total", + "sum", + "amount", + "value", + "volume", + "order", + "sales", + "revenue", + "quantity", + "workload", + "throughput", + "failure", + "repair", + ) + for column in quantitative: + normalized = str(column).lower() + if any(keyword in normalized for keyword in metric_keywords): + return column + + return quantitative[0] + + +def _count_axis_title(query: str | None) -> str: + normalized = str(query or "").lower() + if "new order" in normalized: + return "New Orders Count" + if "order" in normalized: + return "Order Count" + if "repair" in normalized: + return "Repair Count" + if "failure" in normalized: + return "Failure Count" + if "ticket" in normalized: + return "Ticket Count" + return "Count" + + +def _wants_time_axis(query: str | None, chart_type: str) -> bool: + normalized = str(query or "").lower() + time_pattern = ( + r"\b(trend|over time|timeline|monthly|weekly|daily|yearly|" + r"by month|by week|by day|by year)\b" + ) + return chart_type in {"line", "area", "multi_line"} or bool( + re.search(time_pattern, normalized) + ) + + +def _select_dimension_columns( + query: str | None, + chart_type: str, + nominal: list[str], + temporal: list[str], + columns: list[str], +) -> list[str]: + relevant = _query_relevant_columns(query, columns) + relevant_nominal = [column for column in relevant if column in nominal] + relevant_temporal = [column for column in relevant if column in temporal] + + if _wants_time_axis(query, chart_type): + ordered = relevant_temporal + [ + column for column in temporal if column not in relevant_temporal + ] + ordered += relevant_nominal + [ + column for column in nominal if column not in relevant_nominal + ] + return ordered + + ordered = relevant_nominal + [ + column for column in nominal if column not in relevant_nominal + ] + ordered += relevant_temporal + [ + column for column in temporal if column not in relevant_temporal + ] + return ordered + + +def _refine_chart_type_for_columns( + query: str | None, + chart_type: str, + sample_data: list[dict], +) -> str: + if chart_type != "bar" or not sample_data: + return chart_type + + columns = _safe_column_names(list(sample_data[0].keys())) + inferred = _infer_column_types(sample_data) + dimensions = _select_dimension_columns( + query, chart_type, inferred["nominal"], inferred["temporal"], columns + ) + if len([column for column in dimensions if column in inferred["nominal"]]) > 1: + return "grouped_bar" + + return chart_type + + def _match_column_name(field: str | None, columns: list[str]) -> str: if field is None: return "" @@ -104,14 +254,24 @@ def _infer_column_types(sample_data: list[dict]) -> dict[str, list[str]]: values.astype(str).str.replace(",", "", regex=False), errors="coerce", ) - temporal_values = pd.to_datetime(values, errors="coerce") is_temporal_name = bool( re.search(r"(date|time|month|year|day|created|updated)", column_name) ) + string_values = values.astype(str) + looks_temporal = string_values.str.match( + r"^\d{4}[-/]\d{1,2}([-/]\d{1,2})?" + ).all() + temporal_values = ( + pd.to_datetime(values, errors="coerce") + if is_temporal_name or looks_temporal + else None + ) if numeric_values.notna().all() and not is_temporal_name: quantitative.append(str(column)) - elif temporal_values.notna().all() or is_temporal_name: + elif is_temporal_name or ( + temporal_values is not None and temporal_values.notna().all() + ): temporal.append(str(column)) else: nominal.append(str(column)) @@ -138,6 +298,10 @@ def _build_fallback_chart_schema( quantitative = inferred["quantitative"] temporal = inferred["temporal"] nominal = inferred["nominal"] + dimensions = _select_dimension_columns( + query, chart_type, nominal, temporal, columns + ) + measure = _select_measure_column(query, quantitative) title = _humanize_title(query or "Chart") @@ -151,14 +315,12 @@ def count_axis() -> dict: return { "aggregate": "count", "type": "quantitative", - "title": "Count", + "title": _count_axis_title(query), } if chart_type == "pie": - color_field = nominal[0] if nominal else columns[0] - theta_encoding = ( - axis(quantitative[0], "quantitative") if quantitative else count_axis() - ) + color_field = dimensions[0] if dimensions else columns[0] + theta_encoding = axis(measure, "quantitative") if measure else count_axis() return { "title": title, "mark": {"type": "arc"}, @@ -169,9 +331,7 @@ def count_axis() -> dict: } if chart_type in {"line", "area", "multi_line"}: - y_encoding = ( - axis(quantitative[0], "quantitative") if quantitative else count_axis() - ) + y_encoding = axis(measure, "quantitative") if measure else count_axis() if {"year", "month"}.issubset({str(c).lower() for c in columns}): month_field = next(c for c in columns if str(c).lower() == "month") encoding = { @@ -187,31 +347,46 @@ def count_axis() -> dict: "encoding": encoding, } - x_field = temporal[0] if temporal else (nominal[0] if nominal else columns[0]) + x_field = dimensions[0] if dimensions else columns[0] x_type = "temporal" if x_field in temporal else "ordinal" + encoding = { + "x": axis(x_field, x_type), + "y": y_encoding, + } + series_field = next( + (column for column in dimensions[1:] if column in nominal), + None, + ) + if series_field: + encoding["color"] = axis(series_field, "nominal") return { "title": title, "mark": {"type": "area" if chart_type == "area" else "line"}, - "encoding": { - "x": axis(x_field, x_type), - "y": y_encoding, - }, + "encoding": encoding, } - x_field = nominal[0] if nominal else (temporal[0] if temporal else columns[0]) - x_type = "nominal" if x_field in nominal else ("temporal" if x_field in temporal else "ordinal") - y_encoding = ( - axis(quantitative[0], "quantitative") if quantitative else count_axis() + x_field = dimensions[0] if dimensions else columns[0] + x_type = ( + "nominal" + if x_field in nominal + else ("temporal" if x_field in temporal else "ordinal") ) + y_encoding = axis(measure, "quantitative") if measure else count_axis() encoding = { "x": axis(x_field, x_type), "y": y_encoding, } - if chart_type == "grouped_bar" and len(nominal) > 1: - encoding["xOffset"] = axis(nominal[1], "nominal") - encoding["color"] = axis(nominal[1], "nominal") - elif nominal: - encoding["color"] = axis(nominal[0], "nominal") + comparison_field = next( + (column for column in dimensions[1:] if column in nominal), + None, + ) + if comparison_field: + encoding["color"] = axis(comparison_field, "nominal") + nominal_dimension_count = len([c for c in dimensions if c in nominal]) + if chart_type == "grouped_bar" or nominal_dimension_count > 1: + encoding["xOffset"] = axis(comparison_field, "nominal") + elif x_field in nominal: + encoding["color"] = axis(x_field, "nominal") mark = {"type": "bar"} if chart_type == "stacked_bar": @@ -232,6 +407,7 @@ def build_fallback_chart_result( processed = ChartDataPreprocessor().run(data) sample_data = processed.get("sample_data", []) chart_type = _detect_requested_chart_type(query) or "bar" + chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) chart_schema = _build_fallback_chart_schema(query, chart_type, sample_data) if not chart_schema: return { From ba5fd7e121b1e1c2812c5c2e19f4d0b62b8371c1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 00:26:17 +0530 Subject: [PATCH 0249/1087] Prevent stale deployment status from blocking sync --- .../repositories/deployLogRepository.ts | 11 +++ .../apollo/server/services/deployService.ts | 38 +++++++++-- .../services/tests/deployService.test.ts | 67 ++++++++++++++++++- 3 files changed, 111 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts index ab7dc75baa..f374758668 100644 --- a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts +++ b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts @@ -21,6 +21,7 @@ export enum DeployStatusEnum { export interface IDeployLogRepository extends IBasicRepository { findLastProjectDeployLog(projectId: number): Promise; + findLatestProjectDeployLog(projectId: number): Promise; findInProgressProjectDeployLog(projectId: number): Promise; } @@ -62,6 +63,16 @@ export class DeployLogRepository return (res && this.transformFromDBData(res)) || null; } + public async findLatestProjectDeployLog(projectId: number) { + const res = await this.knex + .select('*') + .from(this.tableName) + .where(this.transformToDBData({ projectId })) + .orderBy('created_at', 'desc') + .first(); + return (res && this.transformFromDBData(res)) || null; + } + public async findInProgressProjectDeployLog(projectId: number) { const res = await this.knex .select('*') diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 7a330bf03e..604df21332 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -18,6 +18,8 @@ import { const logger = getLogger('DeployService'); logger.level = 'debug'; +const STALE_DEPLOYMENT_MS = 10 * 60 * 1000; + export interface DeployResponse { status: DeployStatusEnum; error?: string; @@ -69,9 +71,26 @@ export class DeployService implements IDeployService { } public async getInProgressDeployment(projectId: number) { - return await this.deployLogRepository.findInProgressProjectDeployLog( - projectId, - ); + const latestDeploy = + await this.deployLogRepository.findLatestProjectDeployLog(projectId); + if (latestDeploy?.status !== DeployStatusEnum.IN_PROGRESS) { + return null; + } + + const updatedAt = latestDeploy.updatedAt || latestDeploy.createdAt; + const updatedAtTime = updatedAt ? new Date(updatedAt).getTime() : 0; + const isStale = + updatedAtTime > 0 && Date.now() - updatedAtTime > STALE_DEPLOYMENT_MS; + + if (isStale) { + await this.deployLogRepository.updateOne(latestDeploy.id, { + status: DeployStatusEnum.FAILED, + error: 'Deployment timed out before completion.', + }); + return null; + } + + return latestDeploy; } public async deploy( @@ -81,6 +100,7 @@ export class DeployService implements IDeployService { ) { const eventName = TelemetryEvent.MODELING_DEPLOY_MDL; const projectId = typeof project === 'number' ? project : project.id; + let deploy: Deploy | null = null; try { // generate hash of manifest const hash = this.createMDLHash(manifest, project); @@ -101,7 +121,7 @@ export class DeployService implements IDeployService { projectId, status: DeployStatusEnum.IN_PROGRESS, } as Deploy; - const deploy = await this.deployLogRepository.createOne(deployData); + deploy = await this.deployLogRepository.createOne(deployData); // deploy to AI-service const { status: aiStatus, error: aiError } = @@ -135,6 +155,16 @@ export class DeployService implements IDeployService { return { status, error: aiError }; } catch (err: any) { logger.error(`Error deploying model: ${err.message}`); + if (deploy?.id) { + try { + await this.deployLogRepository.updateOne(deploy.id, { + status: DeployStatusEnum.FAILED, + error: err.message, + }); + } catch (updateErr: any) { + logger.error(`Error marking deployment failed: ${updateErr.message}`); + } + } this.telemetry.sendEvent( eventName, { mdl: manifest, error: err.message }, diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 31688702e2..9ed22e9d53 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -13,6 +13,7 @@ describe('DeployService', () => { mockWrenAIAdaptor = { deploy: jest.fn(), delete: jest.fn() }; mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), + findLatestProjectDeployLog: jest.fn(), createOne: jest.fn(), updateOne: jest.fn(), }; @@ -105,5 +106,69 @@ describe('DeployService', () => { ); }); - // Add more tests here to cover other scenarios and error handling + it('should not report in-progress when the latest deployment is successful', async () => { + mockDeployLogRepository.findLatestProjectDeployLog.mockResolvedValue({ + id: 2, + status: DeployStatusEnum.SUCCESS, + }); + + const deployment = await deployService.getInProgressDeployment(1); + + expect(deployment).toBeNull(); + expect( + mockDeployLogRepository.findLatestProjectDeployLog, + ).toHaveBeenCalledWith(1); + expect(mockDeployLogRepository.updateOne).not.toHaveBeenCalled(); + }); + + it('should mark stale in-progress deployments failed', async () => { + const oldDate = new Date(Date.now() - 11 * 60 * 1000); + mockDeployLogRepository.findLatestProjectDeployLog.mockResolvedValue({ + id: 3, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: oldDate, + }); + + const deployment = await deployService.getInProgressDeployment(1); + + expect(deployment).toBeNull(); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(3, { + status: DeployStatusEnum.FAILED, + error: 'Deployment timed out before completion.', + }); + }); + + it('should keep recent in-progress deployments active', async () => { + const recentDate = new Date(); + const inProgressDeployment = { + id: 4, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: recentDate, + }; + mockDeployLogRepository.findLatestProjectDeployLog.mockResolvedValue( + inProgressDeployment, + ); + + const deployment = await deployService.getInProgressDeployment(1); + + expect(deployment).toBe(inProgressDeployment); + expect(mockDeployLogRepository.updateOne).not.toHaveBeenCalled(); + }); + + it('should mark created deployment failed when deployment throws', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockRejectedValue(new Error('network error')); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.FAILED); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.FAILED, + error: 'network error', + }); + }); }); From 2f0de7c9ccc806714fcfd56add899ce04cfabbf6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 00:35:01 +0530 Subject: [PATCH 0250/1087] Avoid redeploy prompts after project switches --- .../apollo/server/resolvers/modelResolver.ts | 3 +- .../apollo/server/services/deployService.ts | 19 +++- .../services/tests/deployService.test.ts | 92 +++++++++++++++++-- 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index dfa49c9078..25dfe5b627 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -225,13 +225,14 @@ export class ModelResolver { args: { force: boolean }, ctx: IContext, ): Promise { - const project = await ctx.projectService.getCurrentProject(); + let project = await ctx.projectService.getCurrentProject(); if (!project.version && project.type !== DataSourceName.DUCKDB) { const version = await ctx.projectService.getProjectDataSourceVersion(project); await ctx.projectService.updateProject(project.id, { version, }); + project = { ...project, version }; } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy( diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 604df21332..484d27bb43 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -182,13 +182,11 @@ export class DeployService implements IDeployService { : { id: project.id, type: project.type, - version: project.version, catalog: project.catalog, schema: project.schema, sampleDataset: project.sampleDataset, - connectionInfo: project.connectionInfo, }; - const content = JSON.stringify({ + const content = this.stableStringify({ project: projectFingerprint, manifest, }); @@ -196,6 +194,21 @@ export class DeployService implements IDeployService { return hash; } + private stableStringify(value: any): string { + if (Array.isArray(value)) { + return `[${value.map((item) => this.stableStringify(item)).join(',')}]`; + } + + if (value && typeof value === 'object') { + return `{${Object.keys(value) + .sort() + .map((key) => `${JSON.stringify(key)}:${this.stableStringify(value[key])}`) + .join(',')}}`; + } + + return JSON.stringify(value); + } + public async getMDLByHash(hash: string) { const deploy = await this.deployLogRepository.findOneBy({ hash }); if (!deploy) { diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 9ed22e9d53..99ac7997d1 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -73,7 +73,7 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); - it('should include datasource identity in deployment hash', () => { + it('should include project schema identity in deployment hash', () => { const manifest = { models: [{ name: 'orders', columns: [] }] }; const project = { id: 1, @@ -89,20 +89,92 @@ describe('DeployService', () => { }, }; - const sameMdlDifferentDatabaseHash = deployService.createMDLHash( - manifest, - { + expect(deployService.createMDLHash(manifest, project)).not.toEqual( + deployService.createMDLHash(manifest, { + ...project, + id: 2, + }), + ); + expect(deployService.createMDLHash(manifest, project)).not.toEqual( + deployService.createMDLHash(manifest, { ...project, - connectionInfo: { - ...project.connectionInfo, - host: 'db-b', - database: 'sales_b', + schema: 'analytics', + }), + ); + }); + + it('should create the same deployment hash for equivalent project objects', () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], }, + ], + }; + const project = { + id: 1, + type: 'mssql', + version: '16', + catalog: 'catalog', + schema: 'dbo', + sampleDataset: null, + connectionInfo: { + host: 'db-a', + port: 1433, + database: 'sales_a', + }, + }; + + const reorderedProject = { + ...project, + connectionInfo: { + database: 'sales_a', + port: 1433, + host: 'db-a', }, + }; + const reorderedManifest = { + models: [ + { + columns: [{ name: 'id' }, { name: 'amount' }], + name: 'orders', + }, + ], + }; + + expect(deployService.createMDLHash(manifest, project)).toEqual( + deployService.createMDLHash(reorderedManifest, reorderedProject), ); + }); - expect(deployService.createMDLHash(manifest, project)).not.toEqual( - sameMdlDifferentDatabaseHash, + it('should not change deployment hash for connection or version refreshes', () => { + const manifest = { models: [{ name: 'orders', columns: [] }] }; + const project = { + id: 1, + type: 'mssql', + version: null, + catalog: 'catalog', + schema: 'dbo', + sampleDataset: null, + connectionInfo: { + host: 'db-a', + port: 1433, + database: 'sales_a', + }, + }; + const refreshedProject = { + ...project, + version: '16', + connectionInfo: { + ...project.connectionInfo, + host: 'db-b', + database: 'sales_b', + }, + }; + + expect(deployService.createMDLHash(manifest, project)).toEqual( + deployService.createMDLHash(manifest, refreshedProject), ); }); From 8ad8f4817aaba276f827f113320ef519f649b8cc Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 00:47:22 +0530 Subject: [PATCH 0251/1087] Make deployment sync check repository compatible --- .../apollo/server/services/deployService.ts | 18 ++++++++++++++-- .../services/tests/deployService.test.ts | 21 +++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 484d27bb43..5a335c8add 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -71,8 +71,7 @@ export class DeployService implements IDeployService { } public async getInProgressDeployment(projectId: number) { - const latestDeploy = - await this.deployLogRepository.findLatestProjectDeployLog(projectId); + const latestDeploy = await this.findLatestOrInProgressDeployment(projectId); if (latestDeploy?.status !== DeployStatusEnum.IN_PROGRESS) { return null; } @@ -93,6 +92,21 @@ export class DeployService implements IDeployService { return latestDeploy; } + private async findLatestOrInProgressDeployment(projectId: number) { + const repository = this.deployLogRepository as IDeployLogRepository & { + findLatestProjectDeployLog?: (projectId: number) => Promise; + }; + + if (typeof repository.findLatestProjectDeployLog === 'function') { + return await repository.findLatestProjectDeployLog(projectId); + } + + logger.warn( + 'DeployLogRepository.findLatestProjectDeployLog is unavailable; falling back to in-progress deployment lookup.', + ); + return await repository.findInProgressProjectDeployLog(projectId); + } + public async deploy( manifest: Manifest, project: Project | number, diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 99ac7997d1..52dfb9290f 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -14,6 +14,7 @@ describe('DeployService', () => { mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), findLatestProjectDeployLog: jest.fn(), + findInProgressProjectDeployLog: jest.fn(), createOne: jest.fn(), updateOne: jest.fn(), }; @@ -227,6 +228,26 @@ describe('DeployService', () => { expect(mockDeployLogRepository.updateOne).not.toHaveBeenCalled(); }); + it('should fall back when latest deployment lookup is unavailable', async () => { + const recentDate = new Date(); + const inProgressDeployment = { + id: 5, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: recentDate, + }; + delete mockDeployLogRepository.findLatestProjectDeployLog; + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue( + inProgressDeployment, + ); + + const deployment = await deployService.getInProgressDeployment(1); + + expect(deployment).toBe(inProgressDeployment); + expect( + mockDeployLogRepository.findInProgressProjectDeployLog, + ).toHaveBeenCalledWith(1); + }); + it('should mark created deployment failed when deployment throws', async () => { const manifest = { key: 'value' }; const projectId = 1; From 380134431343790dc438fa655fd7f57d0d3f32c5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 01:00:07 +0530 Subject: [PATCH 0252/1087] Clear stuck deployments before redeploy --- .../apollo/server/services/deployService.ts | 65 +++++++++++----- .../services/tests/deployService.test.ts | 74 ++++++++++++------- 2 files changed, 91 insertions(+), 48 deletions(-) diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 5a335c8add..6c43232203 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -71,40 +71,55 @@ export class DeployService implements IDeployService { } public async getInProgressDeployment(projectId: number) { - const latestDeploy = await this.findLatestOrInProgressDeployment(projectId); - if (latestDeploy?.status !== DeployStatusEnum.IN_PROGRESS) { + const inProgressDeploy = + await this.deployLogRepository.findInProgressProjectDeployLog(projectId); + if (!inProgressDeploy) { return null; } - const updatedAt = latestDeploy.updatedAt || latestDeploy.createdAt; + const lastSuccessfulDeploy = + await this.deployLogRepository.findLastProjectDeployLog(projectId); + const successTime = lastSuccessfulDeploy + ? this.getDeployTime(lastSuccessfulDeploy) + : 0; + const inProgressTime = this.getDeployTime(inProgressDeploy); + if ( + lastSuccessfulDeploy && + successTime >= inProgressTime + ) { + await this.markDeploymentFailed( + inProgressDeploy, + 'Deployment was superseded by a successful deployment.', + ); + return null; + } + + const updatedAt = inProgressDeploy.updatedAt || inProgressDeploy.createdAt; const updatedAtTime = updatedAt ? new Date(updatedAt).getTime() : 0; const isStale = updatedAtTime > 0 && Date.now() - updatedAtTime > STALE_DEPLOYMENT_MS; if (isStale) { - await this.deployLogRepository.updateOne(latestDeploy.id, { - status: DeployStatusEnum.FAILED, - error: 'Deployment timed out before completion.', - }); + await this.markDeploymentFailed( + inProgressDeploy, + 'Deployment timed out before completion.', + ); return null; } - return latestDeploy; + return inProgressDeploy; } - private async findLatestOrInProgressDeployment(projectId: number) { - const repository = this.deployLogRepository as IDeployLogRepository & { - findLatestProjectDeployLog?: (projectId: number) => Promise; - }; - - if (typeof repository.findLatestProjectDeployLog === 'function') { - return await repository.findLatestProjectDeployLog(projectId); - } + private getDeployTime(deploy: Deploy) { + const deployTime = deploy.updatedAt || deploy.createdAt; + return deployTime ? new Date(deployTime).getTime() : 0; + } - logger.warn( - 'DeployLogRepository.findLatestProjectDeployLog is unavailable; falling back to in-progress deployment lookup.', - ); - return await repository.findInProgressProjectDeployLog(projectId); + private async markDeploymentFailed(deploy: Deploy, error: string) { + await this.deployLogRepository.updateOne(deploy.id, { + status: DeployStatusEnum.FAILED, + error, + }); } public async deploy( @@ -129,6 +144,16 @@ export class DeployService implements IDeployService { return { status: DeployStatusEnum.SUCCESS }; } } + + const previousInProgressDeploy = + await this.deployLogRepository.findInProgressProjectDeployLog(projectId); + if (previousInProgressDeploy) { + await this.markDeploymentFailed( + previousInProgressDeploy, + 'Deployment was superseded by a new deployment.', + ); + } + const deployData = { manifest, hash, diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 52dfb9290f..73ef57a081 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -179,28 +179,37 @@ describe('DeployService', () => { ); }); - it('should not report in-progress when the latest deployment is successful', async () => { - mockDeployLogRepository.findLatestProjectDeployLog.mockResolvedValue({ + it('should not report in-progress when a successful deployment supersedes it', async () => { + const oldDate = new Date(Date.now() - 60 * 1000); + const newDate = new Date(); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ + id: 1, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: oldDate, + }); + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ id: 2, status: DeployStatusEnum.SUCCESS, + updatedAt: newDate, }); const deployment = await deployService.getInProgressDeployment(1); expect(deployment).toBeNull(); - expect( - mockDeployLogRepository.findLatestProjectDeployLog, - ).toHaveBeenCalledWith(1); - expect(mockDeployLogRepository.updateOne).not.toHaveBeenCalled(); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(1, { + status: DeployStatusEnum.FAILED, + error: 'Deployment was superseded by a successful deployment.', + }); }); it('should mark stale in-progress deployments failed', async () => { const oldDate = new Date(Date.now() - 11 * 60 * 1000); - mockDeployLogRepository.findLatestProjectDeployLog.mockResolvedValue({ + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ id: 3, status: DeployStatusEnum.IN_PROGRESS, updatedAt: oldDate, }); + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); const deployment = await deployService.getInProgressDeployment(1); @@ -218,34 +227,15 @@ describe('DeployService', () => { status: DeployStatusEnum.IN_PROGRESS, updatedAt: recentDate, }; - mockDeployLogRepository.findLatestProjectDeployLog.mockResolvedValue( - inProgressDeployment, - ); - - const deployment = await deployService.getInProgressDeployment(1); - - expect(deployment).toBe(inProgressDeployment); - expect(mockDeployLogRepository.updateOne).not.toHaveBeenCalled(); - }); - - it('should fall back when latest deployment lookup is unavailable', async () => { - const recentDate = new Date(); - const inProgressDeployment = { - id: 5, - status: DeployStatusEnum.IN_PROGRESS, - updatedAt: recentDate, - }; - delete mockDeployLogRepository.findLatestProjectDeployLog; mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue( inProgressDeployment, ); + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); const deployment = await deployService.getInProgressDeployment(1); expect(deployment).toBe(inProgressDeployment); - expect( - mockDeployLogRepository.findInProgressProjectDeployLog, - ).toHaveBeenCalledWith(1); + expect(mockDeployLogRepository.updateOne).not.toHaveBeenCalled(); }); it('should mark created deployment failed when deployment throws', async () => { @@ -264,4 +254,32 @@ describe('DeployService', () => { error: 'network error', }); }); + + it('should clear previous in-progress deployment before creating a new one', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ + id: 122, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: new Date(), + }); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.SUCCESS); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(122, { + status: DeployStatusEnum.FAILED, + error: 'Deployment was superseded by a new deployment.', + }); + expect(mockDeployLogRepository.createOne).toHaveBeenCalledWith( + expect.objectContaining({ + projectId, + status: DeployStatusEnum.IN_PROGRESS, + }), + ); + }); }); From e7daacab9a7f7e74a335539f716876c8b4b6bdd6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 01:12:58 +0530 Subject: [PATCH 0253/1087] Stabilize deployment hash for manifest ordering --- .../apollo/server/services/deployService.ts | 6 ++- .../services/tests/deployService.test.ts | 48 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 6c43232203..8966d74d7e 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -235,7 +235,11 @@ export class DeployService implements IDeployService { private stableStringify(value: any): string { if (Array.isArray(value)) { - return `[${value.map((item) => this.stableStringify(item)).join(',')}]`; + const serializedItems = value.map((item) => this.stableStringify(item)); + if (value.every((item) => item && typeof item === 'object')) { + serializedItems.sort(); + } + return `[${serializedItems.join(',')}]`; } if (value && typeof value === 'object') { diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 73ef57a081..30294fd3d8 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -149,6 +149,54 @@ describe('DeployService', () => { ); }); + it('should create the same deployment hash for equivalent manifest ordering', () => { + const project = { + id: 1, + type: 'mssql', + version: '16', + catalog: 'catalog', + schema: 'dbo', + sampleDataset: null, + connectionInfo: {}, + }; + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + { + name: 'customers', + columns: [{ name: 'name' }, { name: 'id' }], + }, + ], + relationships: [ + { name: 'orders_customer', from: 'orders', to: 'customers' }, + { name: 'customers_region', from: 'customers', to: 'regions' }, + ], + }; + const reorderedManifest = { + relationships: [ + { to: 'regions', from: 'customers', name: 'customers_region' }, + { to: 'customers', from: 'orders', name: 'orders_customer' }, + ], + models: [ + { + columns: [{ name: 'id' }, { name: 'name' }], + name: 'customers', + }, + { + columns: [{ name: 'amount' }, { name: 'id' }], + name: 'orders', + }, + ], + }; + + expect(deployService.createMDLHash(manifest, project)).toEqual( + deployService.createMDLHash(reorderedManifest, project), + ); + }); + it('should not change deployment hash for connection or version refreshes', () => { const manifest = { models: [{ name: 'orders', columns: [] }] }; const project = { From c8cbf87f5d5b4be381df307cfe996e7b2a06ae9b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 01:25:18 +0530 Subject: [PATCH 0254/1087] Revert datasource deployment sync changes --- .../apollo/server/resolvers/modelResolver.ts | 13 +- .../server/resolvers/projectResolver.ts | 2 +- .../apollo/server/services/deployService.ts | 132 +-------- .../services/tests/deployService.test.ts | 261 +----------------- 4 files changed, 22 insertions(+), 386 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 25dfe5b627..76d9abc504 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -201,13 +201,13 @@ export class ModelResolver { public async checkModelSync(_root: any, _args: any, ctx: IContext) { try { - const project = await ctx.projectService.getCurrentProject(); + const { id } = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest, project); - const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + const currentHash = ctx.deployService.createMDLHash(manifest, id); + const lastDeploy = await ctx.deployService.getLastDeployment(id); const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = - await ctx.deployService.getInProgressDeployment(project.id); + await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } @@ -225,19 +225,18 @@ export class ModelResolver { args: { force: boolean }, ctx: IContext, ): Promise { - let project = await ctx.projectService.getCurrentProject(); + const project = await ctx.projectService.getCurrentProject(); if (!project.version && project.type !== DataSourceName.DUCKDB) { const version = await ctx.projectService.getProjectDataSourceVersion(project); await ctx.projectService.updateProject(project.id, { version, }); - project = { ...project, version }; } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy( manifest, - project, + project.id, args.force, ); diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 99bbb9b467..6ca33be5cc 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -640,7 +640,7 @@ export class ProjectResolver { private async deploy(ctx: IContext) { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const deployRes = await ctx.deployService.deploy(manifest, project); + const deployRes = await ctx.deployService.deploy(manifest, project.id); // Recommendation generation depends on a successful deployment because // question validation calls previewSql against the deployed manifest. diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 8966d74d7e..3c751d3156 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -5,7 +5,6 @@ import { DeployStatusEnum, IDeployLogRepository, } from '../repositories/deployLogRepository'; -import { Project } from '../repositories/projectRepository'; import { Manifest } from '../mdl/type'; import { createHash } from 'node:crypto'; import { getLogger } from '@server/utils'; @@ -18,8 +17,6 @@ import { const logger = getLogger('DeployService'); logger.level = 'debug'; -const STALE_DEPLOYMENT_MS = 10 * 60 * 1000; - export interface DeployResponse { status: DeployStatusEnum; error?: string; @@ -32,12 +29,12 @@ export interface MDLSyncResponse { export interface IDeployService { deploy( manifest: Manifest, - project: Project | number, + projectId: number, force?: boolean, ): Promise; getLastDeployment(projectId: number): Promise; getInProgressDeployment(projectId: number): Promise; - createMDLHash(manifest: Manifest, project: Project | number): string; + createMDLHash(manifest: Manifest, projectId: number): string; getMDLByHash(hash: string): Promise; deleteAllByProjectId(projectId: number): Promise; } @@ -61,7 +58,7 @@ export class DeployService implements IDeployService { this.telemetry = telemetry; } - public async getLastDeployment(projectId: number) { + public async getLastDeployment(projectId) { const lastDeploy = await this.deployLogRepository.findLastProjectDeployLog(projectId); if (!lastDeploy) { @@ -70,69 +67,17 @@ export class DeployService implements IDeployService { return lastDeploy; } - public async getInProgressDeployment(projectId: number) { - const inProgressDeploy = - await this.deployLogRepository.findInProgressProjectDeployLog(projectId); - if (!inProgressDeploy) { - return null; - } - - const lastSuccessfulDeploy = - await this.deployLogRepository.findLastProjectDeployLog(projectId); - const successTime = lastSuccessfulDeploy - ? this.getDeployTime(lastSuccessfulDeploy) - : 0; - const inProgressTime = this.getDeployTime(inProgressDeploy); - if ( - lastSuccessfulDeploy && - successTime >= inProgressTime - ) { - await this.markDeploymentFailed( - inProgressDeploy, - 'Deployment was superseded by a successful deployment.', - ); - return null; - } - - const updatedAt = inProgressDeploy.updatedAt || inProgressDeploy.createdAt; - const updatedAtTime = updatedAt ? new Date(updatedAt).getTime() : 0; - const isStale = - updatedAtTime > 0 && Date.now() - updatedAtTime > STALE_DEPLOYMENT_MS; - - if (isStale) { - await this.markDeploymentFailed( - inProgressDeploy, - 'Deployment timed out before completion.', - ); - return null; - } - - return inProgressDeploy; - } - - private getDeployTime(deploy: Deploy) { - const deployTime = deploy.updatedAt || deploy.createdAt; - return deployTime ? new Date(deployTime).getTime() : 0; + public async getInProgressDeployment(projectId) { + return await this.deployLogRepository.findInProgressProjectDeployLog( + projectId, + ); } - private async markDeploymentFailed(deploy: Deploy, error: string) { - await this.deployLogRepository.updateOne(deploy.id, { - status: DeployStatusEnum.FAILED, - error, - }); - } - - public async deploy( - manifest: Manifest, - project: Project | number, - force = false, - ) { + public async deploy(manifest, projectId, force = false) { const eventName = TelemetryEvent.MODELING_DEPLOY_MDL; - const projectId = typeof project === 'number' ? project : project.id; - let deploy: Deploy | null = null; try { // generate hash of manifest - const hash = this.createMDLHash(manifest, project); + const hash = this.createMDLHash(manifest, projectId); logger.debug(`Deploying model, hash: ${hash}`); if (!force) { @@ -144,23 +89,13 @@ export class DeployService implements IDeployService { return { status: DeployStatusEnum.SUCCESS }; } } - - const previousInProgressDeploy = - await this.deployLogRepository.findInProgressProjectDeployLog(projectId); - if (previousInProgressDeploy) { - await this.markDeploymentFailed( - previousInProgressDeploy, - 'Deployment was superseded by a new deployment.', - ); - } - const deployData = { manifest, hash, projectId, status: DeployStatusEnum.IN_PROGRESS, } as Deploy; - deploy = await this.deployLogRepository.createOne(deployData); + const deploy = await this.deployLogRepository.createOne(deployData); // deploy to AI-service const { status: aiStatus, error: aiError } = @@ -194,16 +129,6 @@ export class DeployService implements IDeployService { return { status, error: aiError }; } catch (err: any) { logger.error(`Error deploying model: ${err.message}`); - if (deploy?.id) { - try { - await this.deployLogRepository.updateOne(deploy.id, { - status: DeployStatusEnum.FAILED, - error: err.message, - }); - } catch (updateErr: any) { - logger.error(`Error marking deployment failed: ${updateErr.message}`); - } - } this.telemetry.sendEvent( eventName, { mdl: manifest, error: err.message }, @@ -214,44 +139,13 @@ export class DeployService implements IDeployService { } } - public createMDLHash(manifest: Manifest, project: Project | number) { - const projectFingerprint = - typeof project === 'number' - ? { id: project } - : { - id: project.id, - type: project.type, - catalog: project.catalog, - schema: project.schema, - sampleDataset: project.sampleDataset, - }; - const content = this.stableStringify({ - project: projectFingerprint, - manifest, - }); + public createMDLHash(manifest: Manifest, projectId: number) { + const manifestStr = JSON.stringify(manifest); + const content = `${projectId} ${manifestStr}`; const hash = createHash('sha1').update(content).digest('hex'); return hash; } - private stableStringify(value: any): string { - if (Array.isArray(value)) { - const serializedItems = value.map((item) => this.stableStringify(item)); - if (value.every((item) => item && typeof item === 'object')) { - serializedItems.sort(); - } - return `[${serializedItems.join(',')}]`; - } - - if (value && typeof value === 'object') { - return `{${Object.keys(value) - .sort() - .map((key) => `${JSON.stringify(key)}:${this.stableStringify(value[key])}`) - .join(',')}}`; - } - - return JSON.stringify(value); - } - public async getMDLByHash(hash: string) { const deploy = await this.deployLogRepository.findOneBy({ hash }); if (!deploy) { diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 30294fd3d8..14f922de85 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -10,11 +10,9 @@ describe('DeployService', () => { beforeEach(() => { mockTelemetry = { sendEvent: jest.fn() }; - mockWrenAIAdaptor = { deploy: jest.fn(), delete: jest.fn() }; + mockWrenAIAdaptor = { deploy: jest.fn() }; mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), - findLatestProjectDeployLog: jest.fn(), - findInProgressProjectDeployLog: jest.fn(), createOne: jest.fn(), updateOne: jest.fn(), }; @@ -74,260 +72,5 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); - it('should include project schema identity in deployment hash', () => { - const manifest = { models: [{ name: 'orders', columns: [] }] }; - const project = { - id: 1, - type: 'mssql', - version: '16', - catalog: 'catalog', - schema: 'dbo', - sampleDataset: null, - connectionInfo: { - host: 'db-a', - port: 1433, - database: 'sales_a', - }, - }; - - expect(deployService.createMDLHash(manifest, project)).not.toEqual( - deployService.createMDLHash(manifest, { - ...project, - id: 2, - }), - ); - expect(deployService.createMDLHash(manifest, project)).not.toEqual( - deployService.createMDLHash(manifest, { - ...project, - schema: 'analytics', - }), - ); - }); - - it('should create the same deployment hash for equivalent project objects', () => { - const manifest = { - models: [ - { - name: 'orders', - columns: [{ name: 'id' }, { name: 'amount' }], - }, - ], - }; - const project = { - id: 1, - type: 'mssql', - version: '16', - catalog: 'catalog', - schema: 'dbo', - sampleDataset: null, - connectionInfo: { - host: 'db-a', - port: 1433, - database: 'sales_a', - }, - }; - - const reorderedProject = { - ...project, - connectionInfo: { - database: 'sales_a', - port: 1433, - host: 'db-a', - }, - }; - const reorderedManifest = { - models: [ - { - columns: [{ name: 'id' }, { name: 'amount' }], - name: 'orders', - }, - ], - }; - - expect(deployService.createMDLHash(manifest, project)).toEqual( - deployService.createMDLHash(reorderedManifest, reorderedProject), - ); - }); - - it('should create the same deployment hash for equivalent manifest ordering', () => { - const project = { - id: 1, - type: 'mssql', - version: '16', - catalog: 'catalog', - schema: 'dbo', - sampleDataset: null, - connectionInfo: {}, - }; - const manifest = { - models: [ - { - name: 'orders', - columns: [{ name: 'id' }, { name: 'amount' }], - }, - { - name: 'customers', - columns: [{ name: 'name' }, { name: 'id' }], - }, - ], - relationships: [ - { name: 'orders_customer', from: 'orders', to: 'customers' }, - { name: 'customers_region', from: 'customers', to: 'regions' }, - ], - }; - const reorderedManifest = { - relationships: [ - { to: 'regions', from: 'customers', name: 'customers_region' }, - { to: 'customers', from: 'orders', name: 'orders_customer' }, - ], - models: [ - { - columns: [{ name: 'id' }, { name: 'name' }], - name: 'customers', - }, - { - columns: [{ name: 'amount' }, { name: 'id' }], - name: 'orders', - }, - ], - }; - - expect(deployService.createMDLHash(manifest, project)).toEqual( - deployService.createMDLHash(reorderedManifest, project), - ); - }); - - it('should not change deployment hash for connection or version refreshes', () => { - const manifest = { models: [{ name: 'orders', columns: [] }] }; - const project = { - id: 1, - type: 'mssql', - version: null, - catalog: 'catalog', - schema: 'dbo', - sampleDataset: null, - connectionInfo: { - host: 'db-a', - port: 1433, - database: 'sales_a', - }, - }; - const refreshedProject = { - ...project, - version: '16', - connectionInfo: { - ...project.connectionInfo, - host: 'db-b', - database: 'sales_b', - }, - }; - - expect(deployService.createMDLHash(manifest, project)).toEqual( - deployService.createMDLHash(manifest, refreshedProject), - ); - }); - - it('should not report in-progress when a successful deployment supersedes it', async () => { - const oldDate = new Date(Date.now() - 60 * 1000); - const newDate = new Date(); - mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ - id: 1, - status: DeployStatusEnum.IN_PROGRESS, - updatedAt: oldDate, - }); - mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ - id: 2, - status: DeployStatusEnum.SUCCESS, - updatedAt: newDate, - }); - - const deployment = await deployService.getInProgressDeployment(1); - - expect(deployment).toBeNull(); - expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(1, { - status: DeployStatusEnum.FAILED, - error: 'Deployment was superseded by a successful deployment.', - }); - }); - - it('should mark stale in-progress deployments failed', async () => { - const oldDate = new Date(Date.now() - 11 * 60 * 1000); - mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ - id: 3, - status: DeployStatusEnum.IN_PROGRESS, - updatedAt: oldDate, - }); - mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); - - const deployment = await deployService.getInProgressDeployment(1); - - expect(deployment).toBeNull(); - expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(3, { - status: DeployStatusEnum.FAILED, - error: 'Deployment timed out before completion.', - }); - }); - - it('should keep recent in-progress deployments active', async () => { - const recentDate = new Date(); - const inProgressDeployment = { - id: 4, - status: DeployStatusEnum.IN_PROGRESS, - updatedAt: recentDate, - }; - mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue( - inProgressDeployment, - ); - mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); - - const deployment = await deployService.getInProgressDeployment(1); - - expect(deployment).toBe(inProgressDeployment); - expect(mockDeployLogRepository.updateOne).not.toHaveBeenCalled(); - }); - - it('should mark created deployment failed when deployment throws', async () => { - const manifest = { key: 'value' }; - const projectId = 1; - - mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); - mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); - mockWrenAIAdaptor.deploy.mockRejectedValue(new Error('network error')); - - const response = await deployService.deploy(manifest, projectId); - - expect(response.status).toEqual(DeployStatusEnum.FAILED); - expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { - status: DeployStatusEnum.FAILED, - error: 'network error', - }); - }); - - it('should clear previous in-progress deployment before creating a new one', async () => { - const manifest = { key: 'value' }; - const projectId = 1; - - mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); - mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ - id: 122, - status: DeployStatusEnum.IN_PROGRESS, - updatedAt: new Date(), - }); - mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); - mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); - - const response = await deployService.deploy(manifest, projectId); - - expect(response.status).toEqual(DeployStatusEnum.SUCCESS); - expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(122, { - status: DeployStatusEnum.FAILED, - error: 'Deployment was superseded by a new deployment.', - }); - expect(mockDeployLogRepository.createOne).toHaveBeenCalledWith( - expect.objectContaining({ - projectId, - status: DeployStatusEnum.IN_PROGRESS, - }), - ); - }); + // Add more tests here to cover other scenarios and error handling }); From c9bf93d0e23cf700596d1582b2f7b86510693136 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 01:32:38 +0530 Subject: [PATCH 0255/1087] Clear stuck deployment rows --- .../apollo/server/services/deployService.ts | 45 +++++++++++++- .../services/tests/deployService.test.ts | 62 ++++++++++++++++++- 2 files changed, 104 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 3c751d3156..c4ab29868b 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -17,6 +17,8 @@ import { const logger = getLogger('DeployService'); logger.level = 'debug'; +const STALE_DEPLOYMENT_MS = 10 * 60 * 1000; + export interface DeployResponse { status: DeployStatusEnum; error?: string; @@ -68,13 +70,29 @@ export class DeployService implements IDeployService { } public async getInProgressDeployment(projectId) { - return await this.deployLogRepository.findInProgressProjectDeployLog( + const inProgressDeploy = await this.deployLogRepository.findInProgressProjectDeployLog( projectId, ); + if (!inProgressDeploy) { + return null; + } + + const updatedAt = inProgressDeploy.updatedAt || inProgressDeploy.createdAt; + const updatedAtTime = updatedAt ? new Date(updatedAt).getTime() : 0; + if (updatedAtTime > 0 && Date.now() - updatedAtTime > STALE_DEPLOYMENT_MS) { + await this.markDeploymentFailed( + inProgressDeploy, + 'Deployment timed out before completion.', + ); + return null; + } + + return inProgressDeploy; } public async deploy(manifest, projectId, force = false) { const eventName = TelemetryEvent.MODELING_DEPLOY_MDL; + let deploy: Deploy | null = null; try { // generate hash of manifest const hash = this.createMDLHash(manifest, projectId); @@ -89,13 +107,22 @@ export class DeployService implements IDeployService { return { status: DeployStatusEnum.SUCCESS }; } } + const previousInProgressDeploy = + await this.deployLogRepository.findInProgressProjectDeployLog(projectId); + if (previousInProgressDeploy) { + await this.markDeploymentFailed( + previousInProgressDeploy, + 'Deployment was superseded by a new deployment.', + ); + } + const deployData = { manifest, hash, projectId, status: DeployStatusEnum.IN_PROGRESS, } as Deploy; - const deploy = await this.deployLogRepository.createOne(deployData); + deploy = await this.deployLogRepository.createOne(deployData); // deploy to AI-service const { status: aiStatus, error: aiError } = @@ -129,6 +156,13 @@ export class DeployService implements IDeployService { return { status, error: aiError }; } catch (err: any) { logger.error(`Error deploying model: ${err.message}`); + if (deploy?.id) { + try { + await this.markDeploymentFailed(deploy, err.message); + } catch (updateErr: any) { + logger.error(`Error marking deployment failed: ${updateErr.message}`); + } + } this.telemetry.sendEvent( eventName, { mdl: manifest, error: err.message }, @@ -139,6 +173,13 @@ export class DeployService implements IDeployService { } } + private async markDeploymentFailed(deploy: Deploy, error: string) { + await this.deployLogRepository.updateOne(deploy.id, { + status: DeployStatusEnum.FAILED, + error, + }); + } + public createMDLHash(manifest: Manifest, projectId: number) { const manifestStr = JSON.stringify(manifest); const content = `${projectId} ${manifestStr}`; diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 14f922de85..cdf539089c 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -13,6 +13,7 @@ describe('DeployService', () => { mockWrenAIAdaptor = { deploy: jest.fn() }; mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), + findInProgressProjectDeployLog: jest.fn(), createOne: jest.fn(), updateOne: jest.fn(), }; @@ -72,5 +73,64 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); - // Add more tests here to cover other scenarios and error handling + it('should clear stale in-progress deployments', async () => { + const oldDate = new Date(Date.now() - 11 * 60 * 1000); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ + id: 122, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: oldDate, + }); + + const deployment = await deployService.getInProgressDeployment(1); + + expect(deployment).toBeNull(); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(122, { + status: DeployStatusEnum.FAILED, + error: 'Deployment timed out before completion.', + }); + }); + + it('should clear previous in-progress deployment before creating a new one', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ + id: 122, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: new Date(), + }); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.SUCCESS); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(122, { + status: DeployStatusEnum.FAILED, + error: 'Deployment was superseded by a new deployment.', + }); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.SUCCESS, + error: undefined, + }); + }); + + it('should mark created deployment failed when deployment throws', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockRejectedValue(new Error('network error')); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.FAILED); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.FAILED, + error: 'network error', + }); + }); }); From 335003ff0b65f34324f9bd257683d91093090b15 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 01:36:27 +0530 Subject: [PATCH 0256/1087] Recognize already deployed manifests after hash changes --- .../apollo/server/resolvers/modelResolver.ts | 4 +- .../apollo/server/services/deployService.ts | 43 ++++++++++++++++ .../services/tests/deployService.test.ts | 50 +++++++++++++++++++ 3 files changed, 94 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 76d9abc504..54e04a8dfb 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -203,15 +203,13 @@ export class ModelResolver { try { const { id } = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest, id); const lastDeploy = await ctx.deployService.getLastDeployment(id); - const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - return currentHash == lastDeployHash + return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index c4ab29868b..a70ad92181 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -37,6 +37,11 @@ export interface IDeployService { getLastDeployment(projectId: number): Promise; getInProgressDeployment(projectId: number): Promise; createMDLHash(manifest: Manifest, projectId: number): string; + isSameDeployment( + manifest: Manifest, + projectId: number, + deployment?: Deploy | null, + ): boolean; getMDLByHash(hash: string): Promise; deleteAllByProjectId(projectId: number): Promise; } @@ -187,6 +192,44 @@ export class DeployService implements IDeployService { return hash; } + public isSameDeployment( + manifest: Manifest, + projectId: number, + deployment?: Deploy | null, + ) { + if (!deployment) { + return false; + } + + if (deployment.hash === this.createMDLHash(manifest, projectId)) { + return true; + } + + return ( + this.canonicalStringify(deployment.manifest) === + this.canonicalStringify(manifest) + ); + } + + private canonicalStringify(value: any): string { + if (Array.isArray(value)) { + const serializedItems = value.map((item) => this.canonicalStringify(item)); + if (value.every((item) => item && typeof item === 'object')) { + serializedItems.sort(); + } + return `[${serializedItems.join(',')}]`; + } + + if (value && typeof value === 'object') { + return `{${Object.keys(value) + .sort() + .map((key) => `${JSON.stringify(key)}:${this.canonicalStringify(value[key])}`) + .join(',')}}`; + } + + return JSON.stringify(value); + } + public async getMDLByHash(hash: string) { const deploy = await this.deployLogRepository.findOneBy({ hash }); if (!deploy) { diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index cdf539089c..586494b102 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -73,6 +73,56 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); + it('should treat equivalent deployed manifests as the same deployment', () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + { + name: 'customers', + columns: [{ name: 'id' }, { name: 'name' }], + }, + ], + }; + const reorderedManifest = { + models: [ + { + columns: [{ name: 'name' }, { name: 'id' }], + name: 'customers', + }, + { + columns: [{ name: 'amount' }, { name: 'id' }], + name: 'orders', + }, + ], + }; + + expect( + deployService.isSameDeployment(manifest, 1, { + hash: 'different-hash-version', + manifest: reorderedManifest, + }), + ).toBe(true); + }); + + it('should not treat changed deployed manifests as the same deployment', () => { + const manifest = { + models: [{ name: 'orders', columns: [{ name: 'id' }] }], + }; + const changedManifest = { + models: [{ name: 'orders', columns: [{ name: 'id' }, { name: 'amount' }] }], + }; + + expect( + deployService.isSameDeployment(manifest, 1, { + hash: 'different-hash-version', + manifest: changedManifest, + }), + ).toBe(false); + }); + it('should clear stale in-progress deployments', async () => { const oldDate = new Date(Date.now() - 11 * 60 * 1000); mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ From 18aaa029dc0de322c5c751a58fee46ef430c9005 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 16:24:37 +0530 Subject: [PATCH 0257/1087] Fix chart generation result alignment --- .../src/pipelines/generation/utils/chart.py | 147 ++++++++++++++++-- wren-ai-service/src/web/v1/services/chart.py | 15 -- .../generation/test_chart_generation_utils.py | 76 +++++++++ 3 files changed, 212 insertions(+), 26 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 8e3778a4d8..8c05a69f2e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -283,6 +283,70 @@ def _infer_column_types(sample_data: list[dict]) -> dict[str, list[str]]: } +def _chart_mark_type(chart_schema: dict) -> str: + mark = chart_schema.get("mark", {}) + if isinstance(mark, str): + return mark + if isinstance(mark, dict): + return str(mark.get("type") or "") + return "" + + +def _chart_type_from_schema(chart_schema: dict, default: str = "") -> str: + mark_type = _chart_mark_type(chart_schema) + encoding = chart_schema.get("encoding", {}) + + if mark_type == "arc": + return "pie" + if mark_type == "area": + return "area" + if mark_type == "line": + return "multi_line" if chart_schema.get("transform") else "line" + if mark_type == "bar": + if encoding.get("xOffset"): + return "grouped_bar" + if isinstance(encoding.get("y"), dict) and encoding["y"].get("stack"): + return "stacked_bar" + return "bar" + + return default + + +def _is_quantitative_encoding(axis: Any) -> bool: + return isinstance(axis, dict) and axis.get("type") == "quantitative" + + +def _is_categorical_encoding(axis: Any) -> bool: + return isinstance(axis, dict) and axis.get("type") in { + "nominal", + "ordinal", + "temporal", + } + + +def _fallback_chart_type( + requested_chart_type: str, + quantitative: list[str], + temporal: list[str], + nominal: list[str], +) -> str: + if not quantitative: + return "" + + chart_type = requested_chart_type or "bar" + + if chart_type == "pie": + return "pie" if nominal else "" + + if chart_type in {"line", "area", "multi_line"}: + return chart_type if temporal or nominal else "" + + if chart_type in {"grouped_bar", "stacked_bar"}: + return chart_type if len(nominal) > 1 else "" + + return "bar" if nominal or temporal else "" + + def _build_fallback_chart_schema( query: str | None, chart_type: str, @@ -298,6 +362,10 @@ def _build_fallback_chart_schema( quantitative = inferred["quantitative"] temporal = inferred["temporal"] nominal = inferred["nominal"] + chart_type = _fallback_chart_type(chart_type, quantitative, temporal, nominal) + if not chart_type: + return {} + dimensions = _select_dimension_columns( query, chart_type, nominal, temporal, columns ) @@ -415,6 +483,7 @@ def build_fallback_chart_result( "reasoning": "", "chart_type": "", } + chart_type = _chart_type_from_schema(chart_schema, chart_type) chart_schema["$schema"] = "https://vega.github.io/schema/vega-lite/v5.json" chart_schema["data"] = {"values": sample_data} @@ -423,7 +492,7 @@ def build_fallback_chart_result( return { "chart_schema": chart_schema, - "reasoning": "Generated from the preview data columns and requested chart type.", + "reasoning": "Generated from the SQL result columns and requested chart type.", "chart_type": chart_type, } @@ -436,6 +505,11 @@ def _is_schema_compatible_with_sample_data( return False columns = set(_safe_column_names(list(sample_data[0].keys()))) + inferred = _infer_column_types(sample_data) + quantitative = set(inferred["quantitative"]) + temporal = set(inferred["temporal"]) + nominal = set(inferred["nominal"]) + categorical = nominal | temporal encoding = chart_schema.get("encoding", {}) for key in ("x", "y", "x2", "y2", "color", "xOffset", "theta"): axis = encoding.get(key) @@ -443,12 +517,43 @@ def _is_schema_compatible_with_sample_data( if field and str(field) not in columns: return False + if ( + field + and _is_quantitative_encoding(axis) + and str(field) not in quantitative + and axis.get("aggregate") != "count" + ): + return False + + if field and key in {"color", "xOffset"} and str(field) not in categorical: + return False + for transform in chart_schema.get("transform", []) or []: if isinstance(transform, dict): for field in transform.get("fold", []) or []: if field is not None and str(field) not in columns: return False + mark_type = _chart_mark_type(chart_schema) + x_axis = encoding.get("x") + y_axis = encoding.get("y") + theta_axis = encoding.get("theta") + has_quantitative_measure = any( + _is_quantitative_encoding(axis) + or (isinstance(axis, dict) and axis.get("aggregate") == "count") + for axis in (x_axis, y_axis, theta_axis) + ) + + if mark_type == "arc": + return _is_quantitative_encoding(theta_axis) and _is_categorical_encoding( + encoding.get("color") + ) + + if mark_type in {"bar", "line", "area"}: + return has_quantitative_measure and ( + _is_categorical_encoding(x_axis) or _is_categorical_encoding(y_axis) + ) + return True @@ -763,10 +868,7 @@ def run( col: list(df[col].unique())[:sample_column_size] for col in df.columns } - if len(df) > sample_data_count: - sample_data = df.sample(n=sample_data_count).to_dict(orient="records") - else: - sample_data = df.to_dict(orient="records") + sample_data = df.head(sample_data_count).to_dict(orient="records") return { "sample_data": sample_data, @@ -810,6 +912,16 @@ def run( chart_schema = _build_fallback_chart_schema( query, chart_type or "bar", sample_data ) + chart_type = _chart_type_from_schema(chart_schema, chart_type) + + if not chart_schema: + return { + "results": { + "chart_schema": {}, + "reasoning": reasoning, + "chart_type": "", + } + } chart_schema[ "$schema" @@ -829,13 +941,18 @@ def run( } } + fallback_schema = _build_fallback_chart_schema( + query, chart_type or "bar", sample_data + ) + fallback_chart_type = _chart_type_from_schema(fallback_schema, chart_type) + if not fallback_schema: + fallback_chart_type = "" + return { "results": { - "chart_schema": _build_fallback_chart_schema( - query, chart_type or "bar", sample_data - ), + "chart_schema": fallback_schema, "reasoning": reasoning, - "chart_type": chart_type, + "chart_type": fallback_chart_type, } } except ValidationError as e: @@ -850,7 +967,11 @@ def run( "results": { "chart_schema": fallback_schema, "reasoning": "", - "chart_type": _detect_requested_chart_type(query) or "", + "chart_type": _chart_type_from_schema( + fallback_schema, _detect_requested_chart_type(query) or "" + ) + if fallback_schema + else "", } } except Exception as e: @@ -864,7 +985,11 @@ def run( "results": { "chart_schema": fallback_schema, "reasoning": "", - "chart_type": fallback_chart_type, + "chart_type": _chart_type_from_schema( + fallback_schema, fallback_chart_type + ) + if fallback_schema + else "", } } diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index d5ddb93f4e..ed3f47f1bc 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,7 +6,6 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest @@ -140,20 +139,6 @@ async def chart( trace_id=trace_id, ) - local_chart_result = build_fallback_chart_result( - chart_request.query, - sql_data, - chart_request.remove_data_from_chart_schema, - ) - if local_chart_result.get("chart_schema"): - self._chart_results[query_id] = ChartResultResponse( - status="finished", - response=ChartResult(**local_chart_result), - trace_id=trace_id, - ) - results["chart_result"] = local_chart_result - return results - chart_generation_result = await self._pipelines["chart_generation"].run( query=chart_request.query, sql=chart_request.sql, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py new file mode 100644 index 0000000000..57322badc4 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -0,0 +1,76 @@ +import pytest + +from src.pipelines.generation.utils.chart import ( + ChartDataPreprocessor, + build_fallback_chart_result, +) +from src.web.v1.services.chart import ChartRequest, ChartService + + +def test_chart_preprocessor_uses_deterministic_sql_result_order(): + data = { + "columns": [{"name": "Market"}, {"name": "Revenue"}], + "data": [[f"Market {index}", index] for index in range(20)], + } + + result = ChartDataPreprocessor().run(data, sample_data_count=3) + + assert result["sample_data"] == [ + {"Market": "Market 0", "Revenue": 0}, + {"Market": "Market 1", "Revenue": 1}, + {"Market": "Market 2", "Revenue": 2}, + ] + + +def test_fallback_chart_requires_a_real_quantitative_measure(): + result = build_fallback_chart_result( + "Create a chart comparing completed repairs across engineers.", + { + "columns": [{"name": "Status"}], + "data": [["completed"], ["completed"], ["in-progress"]], + }, + ) + + assert result == {"chart_schema": {}, "reasoning": "", "chart_type": ""} + + +@pytest.mark.asyncio +async def test_chart_service_does_not_short_circuit_to_generic_fallback(): + class FakeChartGenerationPipeline: + async def run(self, **kwargs): + assert kwargs["data"]["columns"] == [ + {"name": "Market"}, + {"name": "Revenue"}, + ] + return { + "post_process": { + "results": { + "chart_schema": { + "mark": {"type": "bar"}, + "encoding": { + "x": {"field": "Market", "type": "nominal"}, + "y": {"field": "Revenue", "type": "quantitative"}, + }, + }, + "reasoning": "Generated from the executed SQL result.", + "chart_type": "bar", + } + } + } + + service = ChartService({"chart_generation": FakeChartGenerationPipeline()}) + request = ChartRequest( + query_id="chart-task", + query="Compare customer performance across markets.", + sql="SELECT Market, SUM(Revenue) AS Revenue FROM Sales GROUP BY Market", + data={ + "columns": [{"name": "Market"}, {"name": "Revenue"}], + "data": [["North", 100], ["South", 200]], + }, + ) + + result = await service.chart(request) + + assert result["chart_result"]["reasoning"] == ( + "Generated from the executed SQL result." + ) From 7ab08a565f40b32a09e197fa9b63c82083c65f1b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 25 Jun 2026 17:55:19 +0530 Subject: [PATCH 0258/1087] Require SQL metrics for chart measures --- .../src/pipelines/generation/utils/chart.py | 27 +++++++------------ .../generation/test_chart_generation_utils.py | 14 ++++++++++ 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 8c05a69f2e..d111565788 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -370,6 +370,8 @@ def _build_fallback_chart_schema( query, chart_type, nominal, temporal, columns ) measure = _select_measure_column(query, quantitative) + if not measure: + return {} title = _humanize_title(query or "Chart") @@ -379,27 +381,19 @@ def axis(field: str, field_type: str) -> dict: base["timeUnit"] = "yearmonth" return base - def count_axis() -> dict: - return { - "aggregate": "count", - "type": "quantitative", - "title": _count_axis_title(query), - } - if chart_type == "pie": color_field = dimensions[0] if dimensions else columns[0] - theta_encoding = axis(measure, "quantitative") if measure else count_axis() return { "title": title, "mark": {"type": "arc"}, "encoding": { - "theta": theta_encoding, + "theta": axis(measure, "quantitative"), "color": axis(color_field, "nominal"), }, } if chart_type in {"line", "area", "multi_line"}: - y_encoding = axis(measure, "quantitative") if measure else count_axis() + y_encoding = axis(measure, "quantitative") if {"year", "month"}.issubset({str(c).lower() for c in columns}): month_field = next(c for c in columns if str(c).lower() == "month") encoding = { @@ -439,7 +433,7 @@ def count_axis() -> dict: if x_field in nominal else ("temporal" if x_field in temporal else "ordinal") ) - y_encoding = axis(measure, "quantitative") if measure else count_axis() + y_encoding = axis(measure, "quantitative") encoding = { "x": axis(x_field, x_type), "y": y_encoding, @@ -513,6 +507,8 @@ def _is_schema_compatible_with_sample_data( encoding = chart_schema.get("encoding", {}) for key in ("x", "y", "x2", "y2", "color", "xOffset", "theta"): axis = encoding.get(key) + if isinstance(axis, dict) and axis.get("aggregate"): + return False field = axis.get("field") if isinstance(axis, dict) else None if field and str(field) not in columns: return False @@ -521,7 +517,6 @@ def _is_schema_compatible_with_sample_data( field and _is_quantitative_encoding(axis) and str(field) not in quantitative - and axis.get("aggregate") != "count" ): return False @@ -540,7 +535,6 @@ def _is_schema_compatible_with_sample_data( theta_axis = encoding.get("theta") has_quantitative_measure = any( _is_quantitative_encoding(axis) - or (isinstance(axis, dict) and axis.get("aggregate") == "count") for axis in (x_axis, y_axis, theta_axis) ) @@ -600,11 +594,7 @@ def _needs_deterministic_bar_fallback( return True if len(quantitative) == 0 and len(nominal) >= 1: - y_axis = encoding.get("y") - if isinstance(y_axis, dict) and y_axis.get("field") in nominal: - return True - if not isinstance(y_axis, dict) or y_axis.get("aggregate") != "count": - return True + return True return False @@ -636,6 +626,7 @@ def _needs_deterministic_bar_fallback( - Default time unit is "yearmonth". - For each axis, generate the corresponding human-readable title based on the language provided by the user. - Make sure all of the fields(x, y, xOffset, color, etc.) in the encoding section of the chart schema are present in the column names of the data. +- Do not use Vega-Lite aggregate count or calculate new measures in the chart schema. The SQL must return the metric column, and the chart must encode that returned metric field. ### GUIDELINES TO PLOT CHART ### diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index 57322badc4..495716e4d6 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -2,6 +2,7 @@ from src.pipelines.generation.utils.chart import ( ChartDataPreprocessor, + _is_schema_compatible_with_sample_data, build_fallback_chart_result, ) from src.web.v1.services.chart import ChartRequest, ChartService @@ -34,6 +35,19 @@ def test_fallback_chart_requires_a_real_quantitative_measure(): assert result == {"chart_schema": {}, "reasoning": "", "chart_type": ""} +def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): + assert not _is_schema_compatible_with_sample_data( + { + "mark": {"type": "bar"}, + "encoding": { + "x": {"field": "Inv Date", "type": "temporal"}, + "y": {"aggregate": "count", "type": "quantitative"}, + }, + }, + [{"Inv Date": "2026-01-01"}, {"Inv Date": "2026-07-01"}], + ) + + @pytest.mark.asyncio async def test_chart_service_does_not_short_circuit_to_generic_fallback(): class FakeChartGenerationPipeline: From bd20eceb377b5b106a39801f61d6bf8ca47c9efb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 26 Jun 2026 14:07:45 +0530 Subject: [PATCH 0259/1087] Fix AI service timeout settings --- wren-ai-service/src/config.py | 2 ++ wren-ai-service/src/globals.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 20637032a4..46f53cb250 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -44,6 +44,8 @@ class Settings(BaseSettings): allow_sql_knowledge_retrieval: bool = Field(default=True) max_histories: int = Field(default=5) max_sql_correction_retries: int = Field(default=3) + pipeline_timeout_seconds: int = Field(default=90) + schema_retrieval_timeout_seconds: int = Field(default=600) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 9343344616..d5cf1bb6fb 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -162,6 +162,8 @@ def create_service_container( max_histories=settings.max_histories, enable_column_pruning=settings.enable_column_pruning, max_sql_correction_retries=settings.max_sql_correction_retries, + pipeline_timeout_seconds=settings.pipeline_timeout_seconds, + schema_retrieval_timeout_seconds=settings.schema_retrieval_timeout_seconds, **query_cache, ), ask_feedback_service=services.AskFeedbackService( From 4065a76794eb1f0e865fd2c1f28724a7837698c6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 20:44:10 +0530 Subject: [PATCH 0260/1087] Fix scoped ask SQL generation and empty results --- .../pipelines/generation/chart_adjustment.py | 2 +- .../pipelines/generation/chart_generation.py | 2 +- .../generation/intent_classification.py | 2 +- .../src/pipelines/generation/sql_answer.py | 2 +- .../pipelines/generation/sql_generation.py | 2 +- .../retrieval/db_schema_retrieval.py | 2 +- .../retrieval/preprocess_sql_data.py | 2 +- wren-ai-service/src/web/v1/routers/ask.py | 30 ++++++++++++-- wren-ai-service/src/web/v1/services/ask.py | 18 ++++++++- .../web/v1/services/semantics_preparation.py | 10 +++++ .../adjustmentBackgroundTracker.ts | 4 +- .../src/apollo/server/backgrounds/chart.ts | 8 +--- .../server/backgrounds/recommend-question.ts | 10 +---- .../textBasedAnswerBackgroundTracker.ts | 19 +++++++++ .../apollo/server/services/askingService.ts | 40 ++++++++++++++++--- .../server/services/askingTaskTracker.ts | 3 +- .../apollo/server/services/queryService.ts | 40 +++++++++++++++++++ wren-ui/src/pages/api/v1/ask.ts | 22 ++++++++++ 18 files changed, 182 insertions(+), 36 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/chart_adjustment.py b/wren-ai-service/src/pipelines/generation/chart_adjustment.py index 2e4e7ba8a3..341d246c79 100644 --- a/wren-ai-service/src/pipelines/generation/chart_adjustment.py +++ b/wren-ai-service/src/pipelines/generation/chart_adjustment.py @@ -185,7 +185,7 @@ async def run( data: dict, language: str, ) -> dict: - logger.info("Chart Adjustment pipeline is running...") + logger.debug("Chart Adjustment pipeline is running...") return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/chart_generation.py b/wren-ai-service/src/pipelines/generation/chart_generation.py index 1c22161977..f06fae2800 100644 --- a/wren-ai-service/src/pipelines/generation/chart_generation.py +++ b/wren-ai-service/src/pipelines/generation/chart_generation.py @@ -161,7 +161,7 @@ async def run( remove_data_from_chart_schema: bool = True, custom_instruction: Optional[str] = None, ) -> dict: - logger.info("Chart Generation pipeline is running...") + logger.debug("Chart Generation pipeline is running...") return await self._pipe.execute( ["post_process"], inputs={ diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..bbb5feca38 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -384,7 +384,7 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), ): - logger.info("Intent Classification pipeline is running...") + logger.debug("Intent Classification pipeline is running...") return await self._pipe.execute( ["post_process"], inputs={ diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 948cff0291..42b067b7a6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -160,7 +160,7 @@ async def run( query_id: Optional[str] = None, custom_instruction: Optional[str] = None, ) -> dict: - logger.info("Sql_Answer Generation pipeline is running...") + logger.debug("Sql_Answer Generation pipeline is running...") return await self._pipe.execute( ["generate_answer"], inputs={ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 61451fae15..67e9eb025e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -251,7 +251,7 @@ async def run( allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, ): - logger.info("SQL Generation pipeline is running...") + logger.debug("SQL Generation pipeline is running...") metadata = await retrieve_metadata(project_id or "", self._retriever) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 939f7d1934..e19a445843 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -668,7 +668,7 @@ async def run( histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, ): - logger.info("Ask Retrieval pipeline is running...") + logger.debug("Ask Retrieval pipeline is running...") return await self._pipe.execute( ["construct_retrieval_results"], inputs={ diff --git a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py index e6dadd32d0..e2b0e125de 100644 --- a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py +++ b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py @@ -100,7 +100,7 @@ def run( self, sql_data: Dict, ): - logger.info("Preprocess SQL Data pipeline is running...") + logger.debug("Preprocess SQL Data pipeline is running...") return self._pipe.execute( ["preprocess"], inputs={ diff --git a/wren-ai-service/src/web/v1/routers/ask.py b/wren-ai-service/src/web/v1/routers/ask.py index 8e63f711d1..ae017b5614 100644 --- a/wren-ai-service/src/web/v1/routers/ask.py +++ b/wren-ai-service/src/web/v1/routers/ask.py @@ -12,6 +12,7 @@ get_service_metadata, ) from src.web.v1.services.ask import ( + AskError, AskRequest, AskResponse, AskResultRequest, @@ -32,9 +33,6 @@ async def ask( query_id = str(uuid.uuid4()) ask_request.query_id = query_id ask_service = service_container.ask_service - ask_service._ask_results[query_id] = AskResultResponse( - status="understanding", - ) if ask_service._is_greeting_query(ask_request.query): ask_service._general_streaming_results[query_id] = ( @@ -46,6 +44,32 @@ async def ask( ) return AskResponse(query_id=query_id) + if not ask_request.project_id and ask_request.mdl_hash: + ask_request.project_id = ( + service_container.semantics_preparation_service.get_project_id( + ask_request.mdl_hash + ) + ) + + if not ask_request.project_id: + ask_service._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="OTHERS", + message=( + "project_id is required for /v1/asks. Include the project_id " + "used for semantics preparation, or run semantics preparation " + "again so project_id can be inferred from mdl_hash." + ), + ), + ) + return AskResponse(query_id=query_id) + + ask_service._ask_results[query_id] = AskResultResponse( + status="understanding", + ) + task = asyncio.create_task( ask_service.ask( ask_request, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 493cc01a6b..aa383baaf1 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2070,6 +2070,22 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results + if not ask_request.project_id: + error_message = ( + "project_id is required for scoped schema retrieval. " + "Pass the active project_id with the ask request." + ) + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError(code="OTHERS", message=error_message), + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["metadata"]["error_type"] = "OTHERS" + results["metadata"]["error_message"] = error_message + return results + if self._is_direct_heuristic_sql_query(user_query): self._ask_results[query_id] = AskResultResponse( status="searching", @@ -2384,7 +2400,7 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - logger.info( + logger.debug( "Retrieved tables for query_id %s: %s", query_id, table_names ) 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 35d931c4b6..4ef95e4e1b 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -55,6 +55,7 @@ def __init__( self._prepare_semantics_statuses: Dict[ str, SemanticsPreparationStatusResponse ] = TTLCache(maxsize=maxsize, ttl=ttl) + self._mdl_hash_project_ids: Dict[str, str] = TTLCache(maxsize=maxsize, ttl=ttl) def _parse_mdl(self, mdl: str) -> dict[str, Any]: parsed = orjson.loads(mdl) @@ -214,6 +215,10 @@ async def prepare_semantics( ] = SemanticsPreparationStatusResponse( status="finished", ) + if prepare_semantics_request.project_id: + self._mdl_hash_project_ids[prepare_semantics_request.mdl_hash] = str( + prepare_semantics_request.project_id + ) except Exception as e: logger.exception(f"Failed to prepare semantics: {e}") @@ -232,6 +237,11 @@ async def prepare_semantics( return results + def get_project_id(self, mdl_hash: Optional[str]) -> Optional[str]: + if not mdl_hash: + return None + return self._mdl_hash_project_ids.get(mdl_hash) + def get_prepare_semantics_status( self, prepare_semantics_status_request: SemanticsPreparationStatusRequest ) -> SemanticsPreparationStatusResponse: diff --git a/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts index 4aafbdf053..6fb125458f 100644 --- a/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts @@ -16,7 +16,7 @@ import { TelemetryEvent, WrenService } from '../telemetry/telemetry'; import { PostHogTelemetry } from '../telemetry/telemetry'; const logger = getLogger('AdjustmentTaskTracker'); -logger.level = 'debug'; +logger.level = 'info'; interface TrackedTask { queryId: string; @@ -332,8 +332,6 @@ export class AdjustmentBackgroundTaskTracker // Mark the job as running this.runningJobs.add(queryId); - // Poll for updates - logger.info(`Polling for updates for task ${queryId}`); const result = await this.wrenAIAdaptor.getAskFeedbackResult(queryId); task.lastPolled = now; diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index 35c9c22b79..448c6323b0 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -12,7 +12,7 @@ import { } from '@server/telemetry/telemetry'; const logger = getLogger('ChartBackgroundTracker'); -logger.level = 'debug'; +logger.level = 'info'; const isFinalized = (status: ChartStatus) => { return ( @@ -96,9 +96,6 @@ export class ChartBackgroundTracker { // check if status change if (!statusChanged) { // mark the job as finished - logger.debug( - `Job ${threadResponse.id} chart status not changed, finished`, - ); this.runningJobs.delete(threadResponse.id); return; } @@ -276,9 +273,6 @@ export class ChartAdjustmentBackgroundTracker { // check if status change if (!statusChanged) { // mark the job as finished - logger.debug( - `Job ${threadResponse.id} chart status not changed, finished`, - ); this.runningJobs.delete(threadResponse.id); return; } diff --git a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts index 4831425afa..c6cb50fbab 100644 --- a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts +++ b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts @@ -47,7 +47,7 @@ export class ProjectRecommendQuestionBackgroundTracker { projectRepository: IProjectRepository; }) { this.logger = getLogger('PRQ Background Tracker'); - this.logger.level = 'debug'; + this.logger.level = 'info'; this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.projectRepository = projectRepository; @@ -96,9 +96,6 @@ export class ProjectRecommendQuestionBackgroundTracker { // check if status change if (!changed) { // mark the job as finished - this.logger.debug( - `${loggerPrefix}job ${this.taskKey(project)} status not changed, returning question count: ${result.response?.questions.length || 0}`, - ); this.runningJobs.delete(this.taskKey(project)); return; } @@ -257,7 +254,7 @@ export class ThreadRecommendQuestionBackgroundTracker { threadRepository: IThreadRepository; }) { this.logger = getLogger('TRQ Background Tracker'); - this.logger.level = 'debug'; + this.logger.level = 'info'; this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadRepository = threadRepository; @@ -305,9 +302,6 @@ export class ThreadRecommendQuestionBackgroundTracker { // check if status change if (!changed) { // mark the job as finished - this.logger.debug( - `${loggerPrefix}job ${this.taskKey(thread)} status not changed, returning question count: ${result.response?.questions.length || 0}`, - ); this.runningJobs.delete(this.taskKey(thread)); return; } diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index 6904573abc..3c5c59afac 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -15,6 +15,7 @@ import { IQueryService, ThreadResponseAnswerStatus, PreviewDataResponse, + isPreviewDataEmpty, } from '../services'; import { getLogger } from '@server/utils'; @@ -22,6 +23,8 @@ const logger = getLogger('TextBasedAnswerBackgroundTracker'); logger.level = 'debug'; const ANSWER_PREVIEW_LIMIT = 200; +const EMPTY_RESULT_ANSWER = + 'The SQL query ran successfully, but it returned no rows for the current filters and datasource. Review the SQL or broaden the question if you expected matching records.'; export class TextBasedAnswerBackgroundTracker { // tasks is a kv pair of task id and thread response @@ -122,6 +125,22 @@ export class TextBasedAnswerBackgroundTracker { throw error; } + if (isPreviewDataEmpty(data)) { + const finishedDetail = { + ...threadResponse.answerDetail, + status: ThreadResponseAnswerStatus.FINISHED, + content: EMPTY_RESULT_ANSWER, + numRowsUsedInLLM: 0, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: finishedDetail, + }); + threadResponse.answerDetail = finishedDetail; + delete this.tasks[threadResponse.id]; + this.runningJobs.delete(threadResponse.id); + return; + } + const response = await this.wrenAIAdaptor.createTextBasedAnswer({ query: threadResponse.question, sql: threadResponse.sql, diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index d71eca9964..201e0358c1 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -32,7 +32,11 @@ import { IViewRepository, Project, } from '../repositories'; -import { IQueryService, PreviewDataResponse } from './queryService'; +import { + IQueryService, + PreviewDataResponse, + isPreviewDataEmpty, +} from './queryService'; import { IMDLService } from './mdlService'; import { ThreadRecommendQuestionBackgroundTracker, @@ -44,11 +48,12 @@ import { import { getConfig } from '@server/config'; import { TextBasedAnswerBackgroundTracker } from '../backgrounds/textBasedAnswerBackgroundTracker'; import { IAskingTaskTracker, TrackedAskingResult } from './askingTaskTracker'; +import * as Errors from '@server/utils/error'; const config = getConfig(); const logger = getLogger('AskingService'); -logger.level = 'debug'; +logger.level = 'info'; // const QUERY_ID_PLACEHOLDER = '0'; @@ -121,6 +126,13 @@ const isChartGenerationInProgress = (status?: ChartStatus | string | null) => status || '', ); +const emptyResultChartError = { + code: Errors.GeneralErrorCodes.NO_CHART, + message: + 'The SQL query ran successfully, but it returned no rows to visualize.', + shortMessage: 'No chart data', +}; + // adjustment input export interface AdjustmentReasoningInput { tables: string[]; @@ -349,9 +361,6 @@ class BreakdownBackgroundTracker { // check if status change if (breakdownDetail.status === result.status) { // mark the job as finished - logger.debug( - `Job ${threadResponse.id} status not changed, finished`, - ); this.runningJobs.delete(threadResponse.id); return; } @@ -965,6 +974,16 @@ export class AskingService implements IAskingService { modelingOnly: false, limit: 500, })) as PreviewDataResponse; + + if (isPreviewDataEmpty(chartData)) { + return await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: { + status: ChartStatus.FAILED, + error: emptyResultChartError, + chartSchema: {}, + }, + }); + } } catch (error) { const message = error instanceof Error ? error.message : String(error); logger.warn( @@ -1032,6 +1051,17 @@ export class AskingService implements IAskingService { modelingOnly: false, limit: 500, })) as PreviewDataResponse; + + if (isPreviewDataEmpty(chartData)) { + return await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: { + status: ChartStatus.FAILED, + error: emptyResultChartError, + chartSchema: {}, + adjustment: true, + }, + }); + } } catch (error) { const message = error instanceof Error ? error.message : String(error); logger.warn( diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index 04c2cf8707..ac962b0e8e 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -16,7 +16,7 @@ import { IWrenAIAdaptor } from '../adaptors'; import * as Errors from '@server/utils/error'; const logger = getLogger('AskingTaskTracker'); -logger.level = 'debug'; +logger.level = 'info'; interface TrackedTask { queryId: string; @@ -320,7 +320,6 @@ export class AskingTaskTracker implements IAskingTaskTracker { this.runningJobs.add(queryId); // Poll for updates - logger.debug(`Polling for updates for task ${queryId}`); const result = await this.wrenAIAdaptor.getAskResult(queryId); task.lastPolled = now; const resultChanged = this.isResultChanged(task.result, result); diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index df46ea7c28..cdd29a805b 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -361,6 +361,43 @@ const extractCteNames = (sql: string) => { return cteNames; }; +const extractSqlTableAliases = (sql: string) => { + const aliases = new Map(); + const tablePattern = new RegExp( + String.raw`\b(?:FROM|JOIN)\s+(${SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*${SQL_IDENTIFIER_PATTERN})*)(?:\s+(?:AS\s+)?(${SQL_IDENTIFIER_PATTERN}))?`, + 'gi', + ); + let match: RegExpExecArray | null; + while ((match = tablePattern.exec(sql))) { + const tableReference = splitTableReference(match[1]).join('.').toLowerCase(); + const alias = match[2] ? normalizeSqlIdentifier(match[2]).toLowerCase() : ''; + + if ( + alias && + ![ + 'on', + 'where', + 'join', + 'left', + 'right', + 'inner', + 'outer', + 'full', + 'cross', + ].includes(alias) + ) { + aliases.set(alias, tableReference); + } + aliases.set(tableReference, tableReference); + + const lastPart = splitTableReference(match[1]).pop()?.toLowerCase(); + if (lastPart) { + aliases.set(lastPart, tableReference); + } + } + return aliases; +}; + const getManifestQueryableNames = (manifest?: Manifest) => { const names = new Set(); for (const model of manifest?.models || []) { @@ -704,6 +741,9 @@ const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { } }; +export const isPreviewDataEmpty = (data?: Partial | null) => + !data || !Array.isArray(data.data) || data.data.length === 0; + export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index d2e98d0eb2..839415cb32 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -22,9 +22,12 @@ import { WrenAIError, } from '@/apollo/server/models/adaptor'; import { getLogger } from '@server/utils'; +import { isPreviewDataEmpty } from '@/apollo/server/services'; const logger = getLogger('API_ASK'); logger.level = 'debug'; +const EMPTY_RESULT_SUMMARY = + 'The SQL query ran successfully, but it returned no rows for the current filters and datasource. Review the SQL or broaden the question if you expected matching records.'; const { apiHistoryRepository, @@ -207,6 +210,25 @@ export default async function handler( ); } + if (isPreviewDataEmpty(sqlData as any)) { + await respondWith({ + res, + statusCode: 200, + responsePayload: { + sql, + summary: EMPTY_RESULT_SUMMARY, + threadId: newThreadId, + }, + projectId: project.id, + apiType: ApiType.ASK, + startTime, + requestPayload: req.body, + threadId: newThreadId, + headers: req.headers as Record, + }); + return; + } + // Step 3: Generate summary using text-based answer const textBasedAnswerInput: TextBasedAnswerInput = { query: question, From cf68186d6bfb6e11074efe5bc2068b639c622894 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 20:54:14 +0530 Subject: [PATCH 0261/1087] Fix AskHistory circular imports --- .../pipelines/generation/data_assistance.py | 8 +- .../generation/followup_sql_generation.py | 8 +- .../followup_sql_generation_reasoning.py | 8 +- .../generation/intent_classification.py | 8 +- .../generation/misleading_assistance.py | 8 +- .../retrieval/db_schema_retrieval.py | 96 ++----------------- 6 files changed, 39 insertions(+), 97 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index 51b91197f9..b6d6912ded 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import asyncio import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -12,7 +14,9 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") 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 a1c1b9f17b..dce78f488e 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import logging import sys -from typing import Any +from typing import TYPE_CHECKING, Any from hamilton import base from hamilton.async_driver import AsyncDriver @@ -26,7 +28,9 @@ from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 42b28c5b8f..c5fa6a78ce 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import asyncio import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -17,7 +19,9 @@ ) from src.utils import trace_cost from src.web.v1.services import Configuration -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index bbb5feca38..df9078da23 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import ast import logging import sys -from typing import Any, Literal, Optional +from typing import TYPE_CHECKING, Any, Literal, Optional import orjson from hamilton import base @@ -17,7 +19,9 @@ from src.pipelines.generation.utils.sql import construct_instructions from src.utils import trace_cost from src.web.v1.services import Configuration -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/misleading_assistance.py b/wren-ai-service/src/pipelines/generation/misleading_assistance.py index a35738ecf5..373d073cb6 100644 --- a/wren-ai-service/src/pipelines/generation/misleading_assistance.py +++ b/wren-ai-service/src/pipelines/generation/misleading_assistance.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import asyncio import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -12,7 +14,9 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e19a445843..5b962dff93 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import ast import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional import orjson import tiktoken @@ -21,7 +23,9 @@ normalize_data_type, ) from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -125,15 +129,11 @@ def _build_view_ddl(content: dict) -> str: def expand_business_terms_for_retrieval(query: str) -> str: normalized = (query or "").lower() analytics_terms = { - "business unit", "pcb", "repair", "debug", "turnaround", "failure", - "failure code", - "failure category", - "failure pattern", "resolved", "trend", "volume", @@ -179,78 +179,15 @@ def expand_business_terms_for_retrieval(query: str) -> str: [ query, "Business analytics aliases:", - "throughput trend volume count counts average total ranking top bottom grouped distribution", - "business unit manufacturing unit department location site plant team region category status", - "repair trends repair volume repair counts debug entries debug fixes failure code", + "repair trends repair volume repair counts debug entries debug fixes", + "average debug hours turnaround time resolved entries failure category failure code", "monthly trend quarter grouped by month bar chart line chart", - "top common failures most common failure categories", - "failure patterns category occurrences material workorder serial number", "sales revenue amount sales value sales performance salesperson ranking", "customer sales top customers customer growth orders invoices margin quantity", ] ) -def _is_project_wide_analysis_query(query: str) -> bool: - normalized = (query or "").lower() - if not normalized: - return False - - analysis_terms = { - "average", - "avg", - "bar chart", - "breakdown", - "chart", - "completed", - "compare", - "count", - "counts", - "distribution", - "group by", - "grouped", - "highest", - "line chart", - "lowest", - "maximum", - "minimum", - "monthly", - "most common", - "number of", - "pie chart", - "quarter", - "rank", - "ranking", - "recommend", - "recommended", - "show", - "status", - "sum", - "total", - "totals", - "top", - "trend", - "volume", - } - return any(term in normalized for term in analysis_terms) - - -def _dedupe_documents(documents: list[Document]) -> list[Document]: - deduped: list[Document] = [] - seen: set[tuple[str, str, str]] = set() - for document in documents: - key = ( - str(document.meta.get("name", "")), - str(document.meta.get("type", "")), - document.content, - ) - if key in seen: - continue - seen.add(key) - deduped.append(document) - return deduped - - @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: @@ -331,22 +268,7 @@ async def dbschema_retrieval( ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) - documents = results["documents"] - if project_id and _is_project_wide_analysis_query(query): - all_project_results = await dbschema_retriever.run( - query_embedding=[], - filters={ - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": project_id}, - ], - }, - ) - documents = _dedupe_documents( - documents + all_project_results.get("documents", []) - ) - return documents + return results["documents"] filters = { "operator": "AND", From c57b0b190303dc4cc709f77b93df75b9be50e411 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 20:58:56 +0530 Subject: [PATCH 0262/1087] Fix runtime AskHistory type hints --- wren-ai-service/src/pipelines/generation/data_assistance.py | 2 ++ .../src/pipelines/generation/followup_sql_generation.py | 2 ++ .../pipelines/generation/followup_sql_generation_reasoning.py | 2 ++ .../src/pipelines/generation/intent_classification.py | 2 ++ .../src/pipelines/generation/misleading_assistance.py | 2 ++ wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py | 2 ++ 6 files changed, 12 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index b6d6912ded..b5598774f0 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -17,6 +17,8 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") 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 dce78f488e..f120dd027d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -31,6 +31,8 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index c5fa6a78ce..d30b0ba62b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -22,6 +22,8 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index df9078da23..6e46206b6a 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -22,6 +22,8 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/misleading_assistance.py b/wren-ai-service/src/pipelines/generation/misleading_assistance.py index 373d073cb6..a84b2f016e 100644 --- a/wren-ai-service/src/pipelines/generation/misleading_assistance.py +++ b/wren-ai-service/src/pipelines/generation/misleading_assistance.py @@ -17,6 +17,8 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 5b962dff93..f972ff707b 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -26,6 +26,8 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") From 57a32c96efa63f8ef0766014eb7f7dab7221a4f1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 21:05:36 +0530 Subject: [PATCH 0263/1087] Revert "Fix runtime AskHistory type hints" This reverts commit c57b0b190303dc4cc709f77b93df75b9be50e411. --- wren-ai-service/src/pipelines/generation/data_assistance.py | 2 -- .../src/pipelines/generation/followup_sql_generation.py | 2 -- .../pipelines/generation/followup_sql_generation_reasoning.py | 2 -- .../src/pipelines/generation/intent_classification.py | 2 -- .../src/pipelines/generation/misleading_assistance.py | 2 -- wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py | 2 -- 6 files changed, 12 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index b5598774f0..b6d6912ded 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -17,8 +17,6 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any logger = logging.getLogger("wren-ai-service") 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 f120dd027d..dce78f488e 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -31,8 +31,6 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index d30b0ba62b..c5fa6a78ce 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -22,8 +22,6 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 6e46206b6a..df9078da23 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -22,8 +22,6 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/misleading_assistance.py b/wren-ai-service/src/pipelines/generation/misleading_assistance.py index a84b2f016e..373d073cb6 100644 --- a/wren-ai-service/src/pipelines/generation/misleading_assistance.py +++ b/wren-ai-service/src/pipelines/generation/misleading_assistance.py @@ -17,8 +17,6 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index f972ff707b..5b962dff93 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -26,8 +26,6 @@ if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any logger = logging.getLogger("wren-ai-service") From 8dcc7b41fc4cb078925516b6a7153fe23fc367eb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 21:05:36 +0530 Subject: [PATCH 0264/1087] Revert "Fix AskHistory circular imports" This reverts commit cf68186d6bfb6e11074efe5bc2068b639c622894. --- .../pipelines/generation/data_assistance.py | 8 +- .../generation/followup_sql_generation.py | 8 +- .../followup_sql_generation_reasoning.py | 8 +- .../generation/intent_classification.py | 8 +- .../generation/misleading_assistance.py | 8 +- .../retrieval/db_schema_retrieval.py | 96 +++++++++++++++++-- 6 files changed, 97 insertions(+), 39 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index b6d6912ded..51b91197f9 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -1,9 +1,7 @@ -from __future__ import annotations - import asyncio import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -14,9 +12,7 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") 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 dce78f488e..a1c1b9f17b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -1,8 +1,6 @@ -from __future__ import annotations - import logging import sys -from typing import TYPE_CHECKING, Any +from typing import Any from hamilton import base from hamilton.async_driver import AsyncDriver @@ -28,9 +26,7 @@ from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index c5fa6a78ce..42b28c5b8f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -1,9 +1,7 @@ -from __future__ import annotations - import asyncio import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -19,9 +17,7 @@ ) from src.utils import trace_cost from src.web.v1.services import Configuration - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index df9078da23..bbb5feca38 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,9 +1,7 @@ -from __future__ import annotations - import ast import logging import sys -from typing import TYPE_CHECKING, Any, Literal, Optional +from typing import Any, Literal, Optional import orjson from hamilton import base @@ -19,9 +17,7 @@ from src.pipelines.generation.utils.sql import construct_instructions from src.utils import trace_cost from src.web.v1.services import Configuration - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/misleading_assistance.py b/wren-ai-service/src/pipelines/generation/misleading_assistance.py index 373d073cb6..a35738ecf5 100644 --- a/wren-ai-service/src/pipelines/generation/misleading_assistance.py +++ b/wren-ai-service/src/pipelines/generation/misleading_assistance.py @@ -1,9 +1,7 @@ -from __future__ import annotations - import asyncio import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -14,9 +12,7 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 5b962dff93..e19a445843 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,9 +1,7 @@ -from __future__ import annotations - import ast import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional import orjson import tiktoken @@ -23,9 +21,7 @@ normalize_data_type, ) from src.utils import trace_cost - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -129,11 +125,15 @@ def _build_view_ddl(content: dict) -> str: def expand_business_terms_for_retrieval(query: str) -> str: normalized = (query or "").lower() analytics_terms = { + "business unit", "pcb", "repair", "debug", "turnaround", "failure", + "failure code", + "failure category", + "failure pattern", "resolved", "trend", "volume", @@ -179,15 +179,78 @@ def expand_business_terms_for_retrieval(query: str) -> str: [ query, "Business analytics aliases:", - "repair trends repair volume repair counts debug entries debug fixes", - "average debug hours turnaround time resolved entries failure category failure code", + "throughput trend volume count counts average total ranking top bottom grouped distribution", + "business unit manufacturing unit department location site plant team region category status", + "repair trends repair volume repair counts debug entries debug fixes failure code", "monthly trend quarter grouped by month bar chart line chart", + "top common failures most common failure categories", + "failure patterns category occurrences material workorder serial number", "sales revenue amount sales value sales performance salesperson ranking", "customer sales top customers customer growth orders invoices margin quantity", ] ) +def _is_project_wide_analysis_query(query: str) -> bool: + normalized = (query or "").lower() + if not normalized: + return False + + analysis_terms = { + "average", + "avg", + "bar chart", + "breakdown", + "chart", + "completed", + "compare", + "count", + "counts", + "distribution", + "group by", + "grouped", + "highest", + "line chart", + "lowest", + "maximum", + "minimum", + "monthly", + "most common", + "number of", + "pie chart", + "quarter", + "rank", + "ranking", + "recommend", + "recommended", + "show", + "status", + "sum", + "total", + "totals", + "top", + "trend", + "volume", + } + return any(term in normalized for term in analysis_terms) + + +def _dedupe_documents(documents: list[Document]) -> list[Document]: + deduped: list[Document] = [] + seen: set[tuple[str, str, str]] = set() + for document in documents: + key = ( + str(document.meta.get("name", "")), + str(document.meta.get("type", "")), + document.content, + ) + if key in seen: + continue + seen.add(key) + deduped.append(document) + return deduped + + @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: @@ -268,7 +331,22 @@ async def dbschema_retrieval( ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] + documents = results["documents"] + if project_id and _is_project_wide_analysis_query(query): + all_project_results = await dbschema_retriever.run( + query_embedding=[], + filters={ + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": project_id}, + ], + }, + ) + documents = _dedupe_documents( + documents + all_project_results.get("documents", []) + ) + return documents filters = { "operator": "AND", From 956084b9a21e5d53922f109b588490ad171f0c70 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 21:05:37 +0530 Subject: [PATCH 0265/1087] Revert "Fix scoped ask SQL generation and empty results" This reverts commit 4065a76794eb1f0e865fd2c1f28724a7837698c6. --- .../pipelines/generation/chart_adjustment.py | 2 +- .../pipelines/generation/chart_generation.py | 2 +- .../generation/intent_classification.py | 2 +- .../src/pipelines/generation/sql_answer.py | 2 +- .../pipelines/generation/sql_generation.py | 2 +- .../retrieval/db_schema_retrieval.py | 2 +- .../retrieval/preprocess_sql_data.py | 2 +- wren-ai-service/src/web/v1/routers/ask.py | 30 ++------------ wren-ai-service/src/web/v1/services/ask.py | 18 +-------- .../web/v1/services/semantics_preparation.py | 10 ----- .../adjustmentBackgroundTracker.ts | 4 +- .../src/apollo/server/backgrounds/chart.ts | 8 +++- .../server/backgrounds/recommend-question.ts | 10 ++++- .../textBasedAnswerBackgroundTracker.ts | 19 --------- .../apollo/server/services/askingService.ts | 40 +++---------------- .../server/services/askingTaskTracker.ts | 3 +- .../apollo/server/services/queryService.ts | 40 ------------------- wren-ui/src/pages/api/v1/ask.ts | 22 ---------- 18 files changed, 36 insertions(+), 182 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/chart_adjustment.py b/wren-ai-service/src/pipelines/generation/chart_adjustment.py index 341d246c79..2e4e7ba8a3 100644 --- a/wren-ai-service/src/pipelines/generation/chart_adjustment.py +++ b/wren-ai-service/src/pipelines/generation/chart_adjustment.py @@ -185,7 +185,7 @@ async def run( data: dict, language: str, ) -> dict: - logger.debug("Chart Adjustment pipeline is running...") + logger.info("Chart Adjustment pipeline is running...") return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/chart_generation.py b/wren-ai-service/src/pipelines/generation/chart_generation.py index f06fae2800..1c22161977 100644 --- a/wren-ai-service/src/pipelines/generation/chart_generation.py +++ b/wren-ai-service/src/pipelines/generation/chart_generation.py @@ -161,7 +161,7 @@ async def run( remove_data_from_chart_schema: bool = True, custom_instruction: Optional[str] = None, ) -> dict: - logger.debug("Chart Generation pipeline is running...") + logger.info("Chart Generation pipeline is running...") return await self._pipe.execute( ["post_process"], inputs={ diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index bbb5feca38..4d6cd313cd 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -384,7 +384,7 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), ): - logger.debug("Intent Classification pipeline is running...") + logger.info("Intent Classification pipeline is running...") return await self._pipe.execute( ["post_process"], inputs={ diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 42b067b7a6..948cff0291 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -160,7 +160,7 @@ async def run( query_id: Optional[str] = None, custom_instruction: Optional[str] = None, ) -> dict: - logger.debug("Sql_Answer Generation pipeline is running...") + logger.info("Sql_Answer Generation pipeline is running...") return await self._pipe.execute( ["generate_answer"], inputs={ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 67e9eb025e..61451fae15 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -251,7 +251,7 @@ async def run( allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, ): - logger.debug("SQL Generation pipeline is running...") + logger.info("SQL Generation pipeline is running...") metadata = await retrieve_metadata(project_id or "", self._retriever) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e19a445843..939f7d1934 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -668,7 +668,7 @@ async def run( histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, ): - logger.debug("Ask Retrieval pipeline is running...") + logger.info("Ask Retrieval pipeline is running...") return await self._pipe.execute( ["construct_retrieval_results"], inputs={ diff --git a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py index e2b0e125de..e6dadd32d0 100644 --- a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py +++ b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py @@ -100,7 +100,7 @@ def run( self, sql_data: Dict, ): - logger.debug("Preprocess SQL Data pipeline is running...") + logger.info("Preprocess SQL Data pipeline is running...") return self._pipe.execute( ["preprocess"], inputs={ diff --git a/wren-ai-service/src/web/v1/routers/ask.py b/wren-ai-service/src/web/v1/routers/ask.py index ae017b5614..8e63f711d1 100644 --- a/wren-ai-service/src/web/v1/routers/ask.py +++ b/wren-ai-service/src/web/v1/routers/ask.py @@ -12,7 +12,6 @@ get_service_metadata, ) from src.web.v1.services.ask import ( - AskError, AskRequest, AskResponse, AskResultRequest, @@ -33,6 +32,9 @@ async def ask( query_id = str(uuid.uuid4()) ask_request.query_id = query_id ask_service = service_container.ask_service + ask_service._ask_results[query_id] = AskResultResponse( + status="understanding", + ) if ask_service._is_greeting_query(ask_request.query): ask_service._general_streaming_results[query_id] = ( @@ -44,32 +46,6 @@ async def ask( ) return AskResponse(query_id=query_id) - if not ask_request.project_id and ask_request.mdl_hash: - ask_request.project_id = ( - service_container.semantics_preparation_service.get_project_id( - ask_request.mdl_hash - ) - ) - - if not ask_request.project_id: - ask_service._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="OTHERS", - message=( - "project_id is required for /v1/asks. Include the project_id " - "used for semantics preparation, or run semantics preparation " - "again so project_id can be inferred from mdl_hash." - ), - ), - ) - return AskResponse(query_id=query_id) - - ask_service._ask_results[query_id] = AskResultResponse( - status="understanding", - ) - task = asyncio.create_task( ask_service.ask( ask_request, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index aa383baaf1..493cc01a6b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2070,22 +2070,6 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results - if not ask_request.project_id: - error_message = ( - "project_id is required for scoped schema retrieval. " - "Pass the active project_id with the ask request." - ) - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError(code="OTHERS", message=error_message), - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["metadata"]["error_type"] = "OTHERS" - results["metadata"]["error_message"] = error_message - return results - if self._is_direct_heuristic_sql_query(user_query): self._ask_results[query_id] = AskResultResponse( status="searching", @@ -2400,7 +2384,7 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - logger.debug( + logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) 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 4ef95e4e1b..35d931c4b6 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -55,7 +55,6 @@ def __init__( self._prepare_semantics_statuses: Dict[ str, SemanticsPreparationStatusResponse ] = TTLCache(maxsize=maxsize, ttl=ttl) - self._mdl_hash_project_ids: Dict[str, str] = TTLCache(maxsize=maxsize, ttl=ttl) def _parse_mdl(self, mdl: str) -> dict[str, Any]: parsed = orjson.loads(mdl) @@ -215,10 +214,6 @@ async def prepare_semantics( ] = SemanticsPreparationStatusResponse( status="finished", ) - if prepare_semantics_request.project_id: - self._mdl_hash_project_ids[prepare_semantics_request.mdl_hash] = str( - prepare_semantics_request.project_id - ) except Exception as e: logger.exception(f"Failed to prepare semantics: {e}") @@ -237,11 +232,6 @@ async def prepare_semantics( return results - def get_project_id(self, mdl_hash: Optional[str]) -> Optional[str]: - if not mdl_hash: - return None - return self._mdl_hash_project_ids.get(mdl_hash) - def get_prepare_semantics_status( self, prepare_semantics_status_request: SemanticsPreparationStatusRequest ) -> SemanticsPreparationStatusResponse: diff --git a/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts index 6fb125458f..4aafbdf053 100644 --- a/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/adjustmentBackgroundTracker.ts @@ -16,7 +16,7 @@ import { TelemetryEvent, WrenService } from '../telemetry/telemetry'; import { PostHogTelemetry } from '../telemetry/telemetry'; const logger = getLogger('AdjustmentTaskTracker'); -logger.level = 'info'; +logger.level = 'debug'; interface TrackedTask { queryId: string; @@ -332,6 +332,8 @@ export class AdjustmentBackgroundTaskTracker // Mark the job as running this.runningJobs.add(queryId); + // Poll for updates + logger.info(`Polling for updates for task ${queryId}`); const result = await this.wrenAIAdaptor.getAskFeedbackResult(queryId); task.lastPolled = now; diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index 448c6323b0..35c9c22b79 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -12,7 +12,7 @@ import { } from '@server/telemetry/telemetry'; const logger = getLogger('ChartBackgroundTracker'); -logger.level = 'info'; +logger.level = 'debug'; const isFinalized = (status: ChartStatus) => { return ( @@ -96,6 +96,9 @@ export class ChartBackgroundTracker { // check if status change if (!statusChanged) { // mark the job as finished + logger.debug( + `Job ${threadResponse.id} chart status not changed, finished`, + ); this.runningJobs.delete(threadResponse.id); return; } @@ -273,6 +276,9 @@ export class ChartAdjustmentBackgroundTracker { // check if status change if (!statusChanged) { // mark the job as finished + logger.debug( + `Job ${threadResponse.id} chart status not changed, finished`, + ); this.runningJobs.delete(threadResponse.id); return; } diff --git a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts index c6cb50fbab..4831425afa 100644 --- a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts +++ b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts @@ -47,7 +47,7 @@ export class ProjectRecommendQuestionBackgroundTracker { projectRepository: IProjectRepository; }) { this.logger = getLogger('PRQ Background Tracker'); - this.logger.level = 'info'; + this.logger.level = 'debug'; this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.projectRepository = projectRepository; @@ -96,6 +96,9 @@ export class ProjectRecommendQuestionBackgroundTracker { // check if status change if (!changed) { // mark the job as finished + this.logger.debug( + `${loggerPrefix}job ${this.taskKey(project)} status not changed, returning question count: ${result.response?.questions.length || 0}`, + ); this.runningJobs.delete(this.taskKey(project)); return; } @@ -254,7 +257,7 @@ export class ThreadRecommendQuestionBackgroundTracker { threadRepository: IThreadRepository; }) { this.logger = getLogger('TRQ Background Tracker'); - this.logger.level = 'info'; + this.logger.level = 'debug'; this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadRepository = threadRepository; @@ -302,6 +305,9 @@ export class ThreadRecommendQuestionBackgroundTracker { // check if status change if (!changed) { // mark the job as finished + this.logger.debug( + `${loggerPrefix}job ${this.taskKey(thread)} status not changed, returning question count: ${result.response?.questions.length || 0}`, + ); this.runningJobs.delete(this.taskKey(thread)); return; } diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index 3c5c59afac..6904573abc 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -15,7 +15,6 @@ import { IQueryService, ThreadResponseAnswerStatus, PreviewDataResponse, - isPreviewDataEmpty, } from '../services'; import { getLogger } from '@server/utils'; @@ -23,8 +22,6 @@ const logger = getLogger('TextBasedAnswerBackgroundTracker'); logger.level = 'debug'; const ANSWER_PREVIEW_LIMIT = 200; -const EMPTY_RESULT_ANSWER = - 'The SQL query ran successfully, but it returned no rows for the current filters and datasource. Review the SQL or broaden the question if you expected matching records.'; export class TextBasedAnswerBackgroundTracker { // tasks is a kv pair of task id and thread response @@ -125,22 +122,6 @@ export class TextBasedAnswerBackgroundTracker { throw error; } - if (isPreviewDataEmpty(data)) { - const finishedDetail = { - ...threadResponse.answerDetail, - status: ThreadResponseAnswerStatus.FINISHED, - content: EMPTY_RESULT_ANSWER, - numRowsUsedInLLM: 0, - }; - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: finishedDetail, - }); - threadResponse.answerDetail = finishedDetail; - delete this.tasks[threadResponse.id]; - this.runningJobs.delete(threadResponse.id); - return; - } - const response = await this.wrenAIAdaptor.createTextBasedAnswer({ query: threadResponse.question, sql: threadResponse.sql, diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 201e0358c1..d71eca9964 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -32,11 +32,7 @@ import { IViewRepository, Project, } from '../repositories'; -import { - IQueryService, - PreviewDataResponse, - isPreviewDataEmpty, -} from './queryService'; +import { IQueryService, PreviewDataResponse } from './queryService'; import { IMDLService } from './mdlService'; import { ThreadRecommendQuestionBackgroundTracker, @@ -48,12 +44,11 @@ import { import { getConfig } from '@server/config'; import { TextBasedAnswerBackgroundTracker } from '../backgrounds/textBasedAnswerBackgroundTracker'; import { IAskingTaskTracker, TrackedAskingResult } from './askingTaskTracker'; -import * as Errors from '@server/utils/error'; const config = getConfig(); const logger = getLogger('AskingService'); -logger.level = 'info'; +logger.level = 'debug'; // const QUERY_ID_PLACEHOLDER = '0'; @@ -126,13 +121,6 @@ const isChartGenerationInProgress = (status?: ChartStatus | string | null) => status || '', ); -const emptyResultChartError = { - code: Errors.GeneralErrorCodes.NO_CHART, - message: - 'The SQL query ran successfully, but it returned no rows to visualize.', - shortMessage: 'No chart data', -}; - // adjustment input export interface AdjustmentReasoningInput { tables: string[]; @@ -361,6 +349,9 @@ class BreakdownBackgroundTracker { // check if status change if (breakdownDetail.status === result.status) { // mark the job as finished + logger.debug( + `Job ${threadResponse.id} status not changed, finished`, + ); this.runningJobs.delete(threadResponse.id); return; } @@ -974,16 +965,6 @@ export class AskingService implements IAskingService { modelingOnly: false, limit: 500, })) as PreviewDataResponse; - - if (isPreviewDataEmpty(chartData)) { - return await this.threadResponseRepository.updateOne(threadResponse.id, { - chartDetail: { - status: ChartStatus.FAILED, - error: emptyResultChartError, - chartSchema: {}, - }, - }); - } } catch (error) { const message = error instanceof Error ? error.message : String(error); logger.warn( @@ -1051,17 +1032,6 @@ export class AskingService implements IAskingService { modelingOnly: false, limit: 500, })) as PreviewDataResponse; - - if (isPreviewDataEmpty(chartData)) { - return await this.threadResponseRepository.updateOne(threadResponse.id, { - chartDetail: { - status: ChartStatus.FAILED, - error: emptyResultChartError, - chartSchema: {}, - adjustment: true, - }, - }); - } } catch (error) { const message = error instanceof Error ? error.message : String(error); logger.warn( diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index ac962b0e8e..04c2cf8707 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -16,7 +16,7 @@ import { IWrenAIAdaptor } from '../adaptors'; import * as Errors from '@server/utils/error'; const logger = getLogger('AskingTaskTracker'); -logger.level = 'info'; +logger.level = 'debug'; interface TrackedTask { queryId: string; @@ -320,6 +320,7 @@ export class AskingTaskTracker implements IAskingTaskTracker { this.runningJobs.add(queryId); // Poll for updates + logger.debug(`Polling for updates for task ${queryId}`); const result = await this.wrenAIAdaptor.getAskResult(queryId); task.lastPolled = now; const resultChanged = this.isResultChanged(task.result, result); diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index cdd29a805b..df46ea7c28 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -361,43 +361,6 @@ const extractCteNames = (sql: string) => { return cteNames; }; -const extractSqlTableAliases = (sql: string) => { - const aliases = new Map(); - const tablePattern = new RegExp( - String.raw`\b(?:FROM|JOIN)\s+(${SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*${SQL_IDENTIFIER_PATTERN})*)(?:\s+(?:AS\s+)?(${SQL_IDENTIFIER_PATTERN}))?`, - 'gi', - ); - let match: RegExpExecArray | null; - while ((match = tablePattern.exec(sql))) { - const tableReference = splitTableReference(match[1]).join('.').toLowerCase(); - const alias = match[2] ? normalizeSqlIdentifier(match[2]).toLowerCase() : ''; - - if ( - alias && - ![ - 'on', - 'where', - 'join', - 'left', - 'right', - 'inner', - 'outer', - 'full', - 'cross', - ].includes(alias) - ) { - aliases.set(alias, tableReference); - } - aliases.set(tableReference, tableReference); - - const lastPart = splitTableReference(match[1]).pop()?.toLowerCase(); - if (lastPart) { - aliases.set(lastPart, tableReference); - } - } - return aliases; -}; - const getManifestQueryableNames = (manifest?: Manifest) => { const names = new Set(); for (const model of manifest?.models || []) { @@ -741,9 +704,6 @@ const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { } }; -export const isPreviewDataEmpty = (data?: Partial | null) => - !data || !Array.isArray(data.data) || data.data.length === 0; - export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index 839415cb32..d2e98d0eb2 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -22,12 +22,9 @@ import { WrenAIError, } from '@/apollo/server/models/adaptor'; import { getLogger } from '@server/utils'; -import { isPreviewDataEmpty } from '@/apollo/server/services'; const logger = getLogger('API_ASK'); logger.level = 'debug'; -const EMPTY_RESULT_SUMMARY = - 'The SQL query ran successfully, but it returned no rows for the current filters and datasource. Review the SQL or broaden the question if you expected matching records.'; const { apiHistoryRepository, @@ -210,25 +207,6 @@ export default async function handler( ); } - if (isPreviewDataEmpty(sqlData as any)) { - await respondWith({ - res, - statusCode: 200, - responsePayload: { - sql, - summary: EMPTY_RESULT_SUMMARY, - threadId: newThreadId, - }, - projectId: project.id, - apiType: ApiType.ASK, - startTime, - requestPayload: req.body, - threadId: newThreadId, - headers: req.headers as Record, - }); - return; - } - // Step 3: Generate summary using text-based answer const textBasedAnswerInput: TextBasedAnswerInput = { query: question, From 93796b30e24d0145fe3d61cb944bd857d4996f22 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 22:18:54 +0530 Subject: [PATCH 0266/1087] Fix AskHistory runtime circular imports --- .../src/pipelines/generation/data_assistance.py | 7 +++++-- .../src/pipelines/generation/followup_sql_generation.py | 7 +++++-- .../generation/followup_sql_generation_reasoning.py | 7 +++++-- .../src/pipelines/generation/intent_classification.py | 7 +++++-- .../src/pipelines/generation/misleading_assistance.py | 7 +++++-- .../src/pipelines/retrieval/db_schema_retrieval.py | 7 +++++-- 6 files changed, 30 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index 51b91197f9..09ea192ad0 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -1,7 +1,7 @@ import asyncio import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -12,7 +12,10 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") 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 a1c1b9f17b..a49281d04c 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -1,6 +1,6 @@ import logging import sys -from typing import Any +from typing import TYPE_CHECKING, Any from hamilton import base from hamilton.async_driver import AsyncDriver @@ -26,7 +26,10 @@ from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 42b28c5b8f..ebd3923a5d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -1,7 +1,7 @@ import asyncio import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -17,7 +17,10 @@ ) from src.utils import trace_cost from src.web.v1.services import Configuration -from src.web.v1.services.ask import AskHistory +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..8fe3ada814 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,7 +1,7 @@ import ast import logging import sys -from typing import Any, Literal, Optional +from typing import TYPE_CHECKING, Any, Literal, Optional import orjson from hamilton import base @@ -17,7 +17,10 @@ from src.pipelines.generation.utils.sql import construct_instructions from src.utils import trace_cost from src.web.v1.services import Configuration -from src.web.v1.services.ask import AskHistory +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/misleading_assistance.py b/wren-ai-service/src/pipelines/generation/misleading_assistance.py index a35738ecf5..3053ab53e0 100644 --- a/wren-ai-service/src/pipelines/generation/misleading_assistance.py +++ b/wren-ai-service/src/pipelines/generation/misleading_assistance.py @@ -1,7 +1,7 @@ import asyncio import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -12,7 +12,10 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 939f7d1934..177448fd6f 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,7 @@ import ast import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional import orjson import tiktoken @@ -21,7 +21,10 @@ normalize_data_type, ) from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") From 59961842b5bb304548997cbab3147d14f04466c9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 23:09:50 +0530 Subject: [PATCH 0267/1087] Reduce noisy recommendation polling logs --- wren-ui/src/apollo/server/utils/logger.ts | 48 +++++++++- .../useRecommendedQuestionsInstruction.tsx | 93 +++++++++++++++++-- 2 files changed, 132 insertions(+), 9 deletions(-) diff --git a/wren-ui/src/apollo/server/utils/logger.ts b/wren-ui/src/apollo/server/utils/logger.ts index f19e429efb..7218acf7ee 100644 --- a/wren-ui/src/apollo/server/utils/logger.ts +++ b/wren-ui/src/apollo/server/utils/logger.ts @@ -1 +1,47 @@ -export { getLogger } from 'log4js'; +import { + getLogger as getLog4jsLogger, + type Logger, + type LoggingEvent, +} from 'log4js'; + +const DEFAULT_LOG_LEVEL = process.env.LOG_LEVEL || 'info'; +const DEBUG_ENABLED = DEFAULT_LOG_LEVEL.toLowerCase() === 'debug'; +const wrappedLoggers = new WeakMap(); + +const normalizeLevel = (level: unknown) => { + const requestedLevel = + typeof level === 'string' ? level.toLowerCase() : String(level); + + if (requestedLevel === 'debug' && !DEBUG_ENABLED) { + return DEFAULT_LOG_LEVEL; + } + + return level; +}; + +export const getLogger = (category?: string): Logger => { + const logger = getLog4jsLogger(category); + const cachedLogger = wrappedLoggers.get(logger); + + if (cachedLogger) { + return cachedLogger; + } + + logger.level = DEFAULT_LOG_LEVEL; + + const wrappedLogger = new Proxy(logger, { + set(target, property, value, receiver) { + if (property === 'level') { + target.level = normalizeLevel(value) as string; + return true; + } + + return Reflect.set(target, property, value, receiver); + }, + }); + + wrappedLoggers.set(logger, wrappedLogger); + return wrappedLogger; +}; + +export type { Logger, LoggingEvent }; diff --git a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx index 543ebcf62f..a99a645b20 100644 --- a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx +++ b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx @@ -1,4 +1,4 @@ -import { useMemo, useState, useEffect } from 'react'; +import { useCallback, useMemo, useRef, useState, useEffect } from 'react'; import { groupBy, orderBy, flatMap } from 'lodash'; import { message } from 'antd'; import Icon from '@/import/icon'; @@ -20,6 +20,9 @@ export interface GroupedQuestion { sql: string; } +const RECOMMENDATION_POLL_INTERVAL_MS = 2000; +const RECOMMENDATION_POLL_MAX_INTERVAL_MS = 10000; + const getGroupedQuestions = ( questions: ResultQuestion[], ): GroupedQuestion[] => { @@ -42,12 +45,16 @@ export default function useRecommendedQuestionsInstruction() { const [recommendedQuestions, setRecommendedQuestions] = useState< GroupedQuestion[] >([]); + const pollingRef = useRef | null>(null); + const pollingSessionRef = useRef(0); + const pollingRequestRef = useRef | null>(null); + const pollingDelayRef = useRef(RECOMMENDATION_POLL_INTERVAL_MS); + const lastFingerprintRef = useRef(null); const [fetchRecommendationQuestions, recommendationQuestionsResult] = useGetProjectRecommendationQuestionsLazyQuery({ fetchPolicy: 'network-only', nextFetchPolicy: 'network-only', - pollInterval: 2000, }); // Handle errors via try/catch blocks rather than onError callback @@ -61,6 +68,47 @@ export default function useRecommendedQuestionsInstruction() { [recommendationQuestionsResult.data], ); + const stopPolling = useCallback(() => { + pollingSessionRef.current += 1; + if (pollingRef.current) { + clearTimeout(pollingRef.current); + pollingRef.current = null; + } + pollingDelayRef.current = RECOMMENDATION_POLL_INTERVAL_MS; + }, []); + + const startPolling = useCallback(async () => { + if (pollingRequestRef.current || pollingRef.current) { + return; + } + + stopPolling(); + const pollingSessionId = pollingSessionRef.current; + + const run = async () => { + if (pollingSessionRef.current !== pollingSessionId) return; + if (pollingRequestRef.current) { + await pollingRequestRef.current; + if (pollingSessionRef.current !== pollingSessionId) return; + } + + try { + const request = fetchRecommendationQuestions().then(() => undefined); + pollingRequestRef.current = request; + await request; + } catch (error) { + console.error(error); + } finally { + pollingRequestRef.current = null; + if (pollingSessionRef.current === pollingSessionId) { + pollingRef.current = setTimeout(run, pollingDelayRef.current); + } + } + }; + + await run(); + }, [fetchRecommendationQuestions, stopPolling]); + useEffect(() => { const fetchRecommendationQuestionsData = async () => { const result = await fetchRecommendationQuestions(); @@ -69,7 +117,6 @@ export default function useRecommendedQuestionsInstruction() { return; } - // for existing projects that do not have to generate recommended questions yet if (isRecommendedFinished(data.status)) { if (data.questions.length > 0) { // for regenerate then leave and go back to the home page @@ -77,12 +124,15 @@ export default function useRecommendedQuestionsInstruction() { setShowRecommendedQuestionsPromptMode(true); } + } else { + setGenerating(true); + await startPolling(); } }; fetchRecommendationQuestionsData(); - return () => recommendationQuestionsResult.stopPolling(); - }, []); + return () => stopPolling(); + }, [fetchRecommendationQuestions, startPolling, stopPolling]); useEffect(() => { if (!recommendedQuestionsTask) { @@ -90,7 +140,7 @@ export default function useRecommendedQuestionsInstruction() { } if (isRecommendedFinished(recommendedQuestionsTask?.status)) { - recommendationQuestionsResult.stopPolling(); + stopPolling(); if (recommendedQuestionsTask.questions.length === 0) { isRegenerate && setShowRetry(true); @@ -116,14 +166,41 @@ export default function useRecommendedQuestionsInstruction() { setGenerating(false); } - }, [recommendedQuestionsTask]); + }, [ + isRegenerate, + recommendedQuestionsTask, + showRecommendedQuestionsPromptMode, + stopPolling, + ]); + + useEffect(() => { + const fingerprint = JSON.stringify({ + status: recommendedQuestionsTask?.status || null, + count: recommendedQuestionsTask?.questions?.length || 0, + errorCode: recommendedQuestionsTask?.error?.code || null, + }); + + if (lastFingerprintRef.current === fingerprint) { + pollingDelayRef.current = Math.min( + pollingDelayRef.current * 2, + RECOMMENDATION_POLL_MAX_INTERVAL_MS, + ); + } else { + pollingDelayRef.current = RECOMMENDATION_POLL_INTERVAL_MS; + lastFingerprintRef.current = fingerprint; + } + }, [ + recommendedQuestionsTask?.status, + recommendedQuestionsTask?.questions?.length, + recommendedQuestionsTask?.error?.code, + ]); const onGetRecommendationQuestions = async () => { setGenerating(true); setIsRegenerate(true); try { await generateProjectRecommendationQuestions(); - fetchRecommendationQuestions(); + await startPolling(); } catch (error) { console.error(error); } From 3d5c45edff1d314f3a8e3e1863f61114eacdd0d8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 23:19:14 +0530 Subject: [PATCH 0268/1087] Normalize invented last update date columns --- .../src/pipelines/generation/utils/sql.py | 33 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 24 ++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 99bd1f6a91..711858e036 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2547,6 +2547,39 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "createdby": ("created_by", "created_by_user_id", "author"), "createdbyuser": ("created_by_user", "created_by_user_id", "author"), "createdbyuserid": ("created_by_user_id", "author"), + "lastupdatedate": ( + "last_update_date", + "last_updated_at", + "updated_at", + "DateOut", + "DateIn", + "FailedAt", + "ModifiedAt", + "CreatedAt", + "created_at", + ), + "lastupdate": ( + "last_update_date", + "last_updated_at", + "updated_at", + "DateOut", + "DateIn", + "FailedAt", + "ModifiedAt", + "CreatedAt", + "created_at", + ), + "updateddate": ( + "updated_at", + "last_update_date", + "last_updated_at", + "DateOut", + "DateIn", + "FailedAt", + "ModifiedAt", + "CreatedAt", + "created_at", + ), "invoicequantity": ("Qty", "Quantity", "InvoiceQty", "InvoiceCount"), "otddate": ("InvDate", "OrdDate", "OrderDate", "InvoiceDate", "Date"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6d37b4bc71..52790d0e1b 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -325,6 +325,30 @@ def test_normalize_sql_column_references_to_schema_maps_debug_business_aliases() ) == [] +def test_normalize_sql_column_references_to_schema_maps_last_update_date_alias(): + sql = ( + 'SELECT DATEPART(YEAR, last_update_date) AS "year", ' + 'DATEPART(MONTH, last_update_date) AS "month", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries" ' + 'GROUP BY DATEPART(YEAR, last_update_date), ' + 'DATEPART(MONTH, last_update_date)' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_DebugEntries": ["DebugEntryId", "BusinessUnit", "DateIn", "FailedAt"]}, + ) + + assert "last_update_date" not in normalized + assert 'DATEPART(YEAR, "DateIn") AS "year"' in normalized + assert 'DATEPART(MONTH, "DateIn") AS "month"' in normalized + assert find_invalid_column_references( + normalized, + {"dbo_DebugEntries": ["DebugEntryId", "BusinessUnit", "DateIn", "FailedAt"]}, + ) == [] + + def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid(): sql = 'SELECT "dbo_qSales1"."UnitPrice" FROM "dbo_qSales1"' normalized = normalize_sql_column_references_to_schema( From 9d9fcc4213f700311c306874ec9bcc2784c90c7c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 29 Jun 2026 23:49:03 +0530 Subject: [PATCH 0269/1087] Normalize stale MSSQL date aliases --- .../apollo/server/utils/mssqlSqlNormalizer.ts | 13 ++++++++++ .../utils/tests/mssqlSqlNormalizer.test.ts | 25 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index a3dfc8a6cc..b2c7a1fdb4 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -83,6 +83,13 @@ const replaceInventedDateFields = (sql: string): string => { 'Repair_Date', 'EventDate', 'event_date', + 'last_update_date', + 'last_updated_date', + 'last_updated_at', + 'lastUpdateDate', + 'lastUpdatedDate', + 'updated_date', + 'updated_at', 'Date', 'date', ]; @@ -95,6 +102,12 @@ const replaceInventedDateFields = (sql: string): string => { ); sql = sql.replace(new RegExp(String.raw`"${escaped}"`, 'gi'), timestampExpression); sql = sql.replace(new RegExp(String.raw`\[${escaped}\]`, 'gi'), timestampExpression); + if (field.toLowerCase() !== 'date') { + sql = sql.replace( + new RegExp(String.raw`(? { ); }); + it('rewrites stale last_update_date references for debug entry trends', () => { + const normalized = normalizeMssqlSqlForIbis( + ` + SELECT + DATEPART(YEAR, last_update_date) AS "year", + DATEPART(MONTH, last_update_date) AS "month", + "dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit", + COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" + FROM "dbo_DebugEntries" + GROUP BY + DATEPART(YEAR, last_update_date), + DATEPART(MONTH, last_update_date), + "dbo_DebugEntries"."BusinessUnit" + ORDER BY + DATEPART(YEAR, last_update_date), + DATEPART(MONTH, last_update_date) + `, + DataSourceName.MSSQL, + ); + + expect(normalized).not.toContain('last_update_date'); + expect(normalized).toContain('DATEPART(YEAR, "dbo_DebugEntries"."DateIn")'); + expect(normalized).toContain('DATEPART(MONTH, "dbo_DebugEntries"."DateIn")'); + }); + it('rewrites hallucinated knowledge article fields', () => { const normalized = normalizeMssqlSqlForIbis( ` From f2b9f3e62f46dc3b65035992be06af3e12f4de33 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 00:37:31 +0530 Subject: [PATCH 0270/1087] Bypass retrieval for explicit grouped table columns --- wren-ai-service/src/web/v1/services/ask.py | 103 ++++++++++++++++-- .../pytest/services/test_ask_sales_sql.py | 15 +++ 2 files changed, 106 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 493cc01a6b..37a33ca6dc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -490,6 +490,60 @@ def _find_first_schema_column( def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' + def _extract_explicit_table_column_reference( + self, query: str + ) -> tuple[str, str] | None: + normalized_query = query or "" + reference_match = re.search( + r"\b(?P[A-Za-z_][A-Za-z0-9_]*)[._]" + r"(?P
[A-Za-z_][A-Za-z0-9_]*)[._]" + r"(?P[A-Za-z_][A-Za-z0-9_]*)\b", + normalized_query, + ) + if not reference_match: + return None + + schema = reference_match.group("schema") + table = reference_match.group("table") + column = reference_match.group("column") + table_name = f"{schema}_{table}" + return table_name, column + + def _build_explicit_group_count_sql(self, query: str) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + if not any( + term in normalized_query + for term in ( + "group by", + "grouped by", + "by ", + "pie chart", + "donut chart", + "bar chart", + "count", + "counts", + ) + ): + return None + + explicit_reference = self._extract_explicit_table_column_reference(query) + if not explicit_reference: + return None + + table_name, column = explicit_reference + table_ref = self._quote_sql_identifier(table_name) + column_ref = f"{table_ref}.{self._quote_sql_identifier(column)}" + return ( + f"SELECT {column_ref} AS {self._quote_sql_identifier(column)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"GROUP BY {column_ref} " + f"ORDER BY COUNT(*) DESC" + ) + def _build_date_filter(self, table_name: str, date_column: str, query: str) -> str: date_ref = ( f"{self._quote_sql_identifier(table_name)}." @@ -2124,18 +2178,43 @@ async def ask( invalid_sql = heuristic_sql error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - historical_question = await self._run_with_timeout( - "Historical question retrieval", - self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ), - ) + if explicit_group_count_sql := self._build_explicit_group_count_sql( + user_query + ): + table_column_reference = ( + self._extract_explicit_table_column_reference(user_query) + ) + table_names = ( + [table_column_reference[0]] if table_column_reference else [] + ) + api_results = [ + AskResult( + **{ + "sql": explicit_group_count_sql, + "type": "llm", + } + ) + ] + rephrased_question = user_query + logger.info( + "Using explicit table-column grouped count SQL for query_id %s", + query_id, + ) + + historical_question_result = [] + if not api_results: + historical_question = await self._run_with_timeout( + "Historical question retrieval", + self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ), + ) - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] valid_historical_results = [] for result in historical_question_result: @@ -2159,7 +2238,7 @@ async def ask( if valid_historical_results: api_results = valid_historical_results sql_generation_reasoning = "" - else: + elif not api_results: original_user_query = user_query # Run both pipeline operations concurrently sql_samples_task, instructions_task = await self._run_with_timeout( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 098a06641d..48614e05a7 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -334,6 +334,21 @@ def test_build_schema_grounded_sql_for_ticket_category_request_uses_existing_col assert "category" not in sql +def test_build_explicit_group_count_sql_for_schema_table_column_reference(): + service = AskService.__new__(AskService) + sql = service._build_explicit_group_count_sql( + "Show a pie chart grouped by dbo.tickets.status." + ) + + assert sql == ( + 'SELECT "dbo_tickets"."status" AS "status", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_tickets" ' + 'GROUP BY "dbo_tickets"."status" ' + 'ORDER BY COUNT(*) DESC' + ) + + def test_build_schema_grounded_sql_for_knowledge_source_request_uses_existing_columns(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 603210438efc6de8ee88b368948a94be7f5bb0d9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 01:35:28 +0530 Subject: [PATCH 0271/1087] Fix ask service circular import --- wren-ai-service/src/web/v1/services/ask.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 37a33ca6dc..7e5ef1d7ff 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -8,12 +8,6 @@ from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - find_invalid_column_references, - find_invalid_table_references, -) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -1988,6 +1982,13 @@ def _build_validated_ask_result_from_sql( if not ask_result: return None + from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + find_invalid_column_references, + find_invalid_table_references, + ) + invalid_tables = find_invalid_table_references( ask_result.sql, construct_valid_table_names(table_ddls), From 2fda10623e44c036522d38666a2fc4cdd24590dc Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 01:55:37 +0530 Subject: [PATCH 0272/1087] Generate audit log activity SQL from schema --- wren-ai-service/src/web/v1/services/ask.py | 99 ++++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 41 ++++++++ 2 files changed, 138 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7e5ef1d7ff..cfa1ea5a6c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1094,6 +1094,83 @@ def _build_manufacturing_throughput_sql( 'ORDER BY "throughput" DESC' ) + def _build_audit_log_activity_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + if not ( + "audit" in normalized + and "log" in normalized + and any(term in normalized for term in ("activity", "over time", "trend")) + ): + return None + + table_name = "dbo_audit_log" + timestamp_column = "created_at" + if not self._schema_has_table_column( + table_ddls, + table_name, + timestamp_column, + table_names=table_names, + ): + return None + + dimension_column = None + condition_candidates = ( + "is_name_condition", + "name", + "action", + "entity_type", + ) + activity_candidates = ( + "action", + "entity_type", + "actor_name", + "actor_user_id", + "name", + ) + candidates = ( + condition_candidates + if "condition" in normalized + else activity_candidates + ) + for candidate in candidates: + if self._schema_has_table_column( + table_ddls, + table_name, + candidate, + table_names=table_names, + ): + dimension_column = candidate + break + + if not dimension_column: + return None + + timestamp_expression = f'"{table_name}"."{timestamp_column}"' + dimension_expression = f'"{table_name}"."{dimension_column}"' + return ( + f"SELECT DATEPART(YEAR, {timestamp_expression}) AS \"year\", " + f"DATEPART(MONTH, {timestamp_expression}) AS \"month\", " + f"{dimension_expression} AS \"{dimension_column}\", " + f'COUNT(*) AS "activity_count" ' + f'FROM "{table_name}" ' + f"WHERE {timestamp_expression} IS NOT NULL " + f"AND {dimension_expression} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {timestamp_expression}), " + f"DATEPART(MONTH, {timestamp_expression}), " + f"{dimension_expression} " + f"ORDER BY DATEPART(YEAR, {timestamp_expression}), " + f"DATEPART(MONTH, {timestamp_expression}), " + f'"activity_count" DESC' + ) + def _build_repair_failure_count_sql( self, query: str, @@ -2468,8 +2545,26 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if deterministic_sales_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls + if audit_log_activity_sql := self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using schema-grounded audit log activity SQL for query_id %s", + query_id, + ) + api_results = [ + AskResult( + **{ + "sql": audit_log_activity_sql, + "type": "llm", + } + ) + ] + + if not api_results and ( + deterministic_sales_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ) ): logger.info( "Using schema-grounded CWSales SQL for query_id %s", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 48614e05a7..fc1e52fb2f 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -414,3 +414,44 @@ def test_build_schema_grounded_sql_for_ticket_throughput_trend(): 'ORDER BY DATEPART(YEAR, "dbo_tickets"."created_at"), ' 'DATEPART(MONTH, "dbo_tickets"."created_at")' ) + + +def test_build_audit_log_activity_sql_uses_existing_condition_columns(): + service = AskService.__new__(AskService) + sql = service._build_audit_log_activity_sql( + "Show audit log activity by condition name over time.", + [ + """ + CREATE TABLE dbo_audit_log ( + id VARCHAR, + action VARCHAR, + actor_name VARCHAR, + actor_user_id VARCHAR, + after_state VARCHAR, + before_state VARCHAR, + created_at TIMESTAMP, + entity_type VARCHAR, + is_name_condition BOOLEAN, + name VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_audit_log"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_audit_log"."created_at") AS "month", ' + '"dbo_audit_log"."is_name_condition" AS "is_name_condition", ' + 'COUNT(*) AS "activity_count" ' + 'FROM "dbo_audit_log" ' + 'WHERE "dbo_audit_log"."created_at" IS NOT NULL ' + 'AND "dbo_audit_log"."is_name_condition" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_audit_log"."created_at"), ' + 'DATEPART(MONTH, "dbo_audit_log"."created_at"), ' + '"dbo_audit_log"."is_name_condition" ' + 'ORDER BY DATEPART(YEAR, "dbo_audit_log"."created_at"), ' + 'DATEPART(MONTH, "dbo_audit_log"."created_at"), ' + '"activity_count" DESC' + ) + assert "condition_name" not in sql + assert "timestamp" not in sql From 5c87b544f6ffb4e005c2e892829ab266fbf48596 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 02:08:20 +0530 Subject: [PATCH 0273/1087] Keep ask validation local to service --- wren-ai-service/src/web/v1/services/ask.py | 78 ++++++++++++++++--- .../pytest/services/test_ask_sales_sql.py | 19 +++++ 2 files changed, 86 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cfa1ea5a6c..d3eb4d8ff0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2059,21 +2059,77 @@ def _build_validated_ask_result_from_sql( if not ask_result: return None - from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - find_invalid_column_references, - find_invalid_table_references, + schema_tables = self._parse_schema_tables(table_ddls) + valid_tables = { + str(table.get("name") or "").lower(): table + for table in schema_tables + if table.get("name") + } + valid_table_suffixes = { + table_name.split(".")[-1].lower(): table + for table_name, table in valid_tables.items() + } + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, ) + referenced_tables = [ + next(value for value in match.groupdict().values() if value) + for match in table_reference_pattern.finditer(ask_result.sql) + ] + invalid_tables = [ + table + for table in referenced_tables + if table.lower() not in valid_tables + and table.lower().split(".")[-1] not in valid_table_suffixes + ] - invalid_tables = find_invalid_table_references( - ask_result.sql, - construct_valid_table_names(table_ddls), + columns_by_table = { + table_name: { + str(column.get("name") or "").lower() + for column in table.get("columns", []) + if column.get("name") + } + for table_name, table in valid_tables.items() + } + columns_by_table.update( + { + table_name.split(".")[-1].lower(): columns + for table_name, columns in columns_by_table.items() + } ) - invalid_columns = find_invalid_column_references( - ask_result.sql, - construct_valid_table_columns(table_ddls), + + qualified_column_pattern = re.compile( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_$]*))", + flags=re.IGNORECASE, ) + invalid_columns = [] + for match in qualified_column_pattern.finditer(ask_result.sql): + table_reference = ( + match.group("table_quoted") + or match.group("table_bracketed") + or match.group("table_bare") + or "" + ) + column_reference = ( + match.group("column_quoted") + or match.group("column_bracketed") + or match.group("column_bare") + or "" + ) + table_key = table_reference.lower() + column_key = column_reference.lower() + table_columns = columns_by_table.get(table_key) or columns_by_table.get( + table_key.split(".")[-1] + ) + if table_columns is not None and column_key not in table_columns: + invalid_columns.append(f"{table_reference}.{column_reference}") + if invalid_tables or invalid_columns: logger.warning( "Ignoring heuristic SQL because it is not valid for active schema. " diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index fc1e52fb2f..eab00038f2 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -455,3 +455,22 @@ def test_build_audit_log_activity_sql_uses_existing_condition_columns(): ) assert "condition_name" not in sql assert "timestamp" not in sql + + +def test_build_validated_ask_result_from_sql_uses_local_schema_validation(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + 'SELECT "dbo_tblSales"."SalesPerson" FROM "dbo_tblSales"', + [ + """ + CREATE TABLE dbo_tblSales ( + SalesPerson VARCHAR, + SalesValue INTEGER + ); + """ + ], + ) + + assert result is not None + assert result.sql == 'SELECT "dbo_tblSales"."SalesPerson" FROM "dbo_tblSales"' From 93ee7a1113b9844f3cc7c00d6d096eafb8d8baf7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 02:43:45 +0530 Subject: [PATCH 0274/1087] Add direct Orders sales SQL fast path --- wren-ai-service/src/web/v1/services/ask.py | 125 ++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 50 +++++++ 2 files changed, 175 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d3eb4d8ff0..2764fd9877 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -484,6 +484,103 @@ def _find_first_schema_column( def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' + def _build_direct_orders_sales_sql(self, query: str) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + is_sales_or_orders_query = any( + term in normalized + for term in ( + "sales", + "sale", + "order", + "orders", + "new order", + "new orders", + "market", + "salesperson", + "sales person", + ) + ) + if not is_sales_or_orders_query: + return None + + table_ref = '"dbo_tblSales"' + limit = self._extract_requested_top_n(query, default_value=10) + + if ( + ("salesperson" in normalized or "sales person" in normalized) + and ("order count" in normalized or "orders" in normalized or "count" in normalized) + ): + return ( + f'SELECT TOP {limit} {table_ref}."SalesPerson" AS "SalesPerson", ' + f'COUNT(*) AS "OrderCount" ' + f"FROM {table_ref} " + f'WHERE {table_ref}."SalesPerson" IS NOT NULL ' + f'GROUP BY {table_ref}."SalesPerson" ' + f"ORDER BY COUNT(*) DESC" + ) + + if "top" in normalized and "new order" in normalized: + date_filter = "" + if re.search(r"\b2026[\s-]*q1\b", normalized): + date_filter = ( + f'WHERE {table_ref}."OrdDate" >= \'2026-01-01 00:00:00\' ' + f'AND {table_ref}."OrdDate" < \'2026-04-01 00:00:00\' ' + ) + return ( + f'SELECT TOP {limit} {table_ref}."BU" AS "BU", ' + f'{table_ref}."Market" AS "Market", ' + f'{table_ref}."Customer" AS "Customer", ' + f'{table_ref}."ProdName" AS "ProdName", ' + f'{table_ref}."SalesValue" AS "SalesValue" ' + f"FROM {table_ref} " + f"{date_filter}" + f'ORDER BY {table_ref}."SalesValue" DESC' + ) + + if ( + "market" in normalized + and "growth" in normalized + and any(term in normalized for term in ("last year", "previous year")) + ): + return ( + f'SELECT {table_ref}."Market" AS "Market", ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "CurrentPeriodSales", ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "PreviousPeriodSales", ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) - ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "SalesGrowth" ' + f"FROM {table_ref} " + f'WHERE {table_ref}."Market" IS NOT NULL ' + f'GROUP BY {table_ref}."Market" ' + f'ORDER BY "SalesGrowth" DESC' + ) + + if ( + "distribution" in normalized + and "sales" in normalized + and ("market" in normalized or "by market" in normalized) + ): + return ( + f'SELECT {table_ref}."Market" AS "Market", ' + f'SUM({table_ref}."SalesValue") AS "TotalSalesValue" ' + f"FROM {table_ref} " + f'WHERE {table_ref}."Market" IS NOT NULL ' + f'GROUP BY {table_ref}."Market" ' + f'ORDER BY SUM({table_ref}."SalesValue") DESC' + ) + + return None + def _extract_explicit_table_column_reference( self, query: str ) -> tuple[str, str] | None: @@ -2258,6 +2355,34 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results + if direct_orders_sales_sql := self._build_direct_orders_sales_sql( + user_query + ): + api_results = [ + AskResult( + **{ + "sql": direct_orders_sales_sql, + "type": "llm", + } + ) + ] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=["dbo_tblSales"], + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + logger.info( + "Using direct Orders/Sales SQL for query_id %s", + query_id, + ) + return results + if self._is_direct_heuristic_sql_query(user_query): self._ask_results[query_id] = AskResultResponse( status="searching", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index eab00038f2..29e947dff8 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -30,6 +30,56 @@ def test_build_schema_grounded_sales_sql_for_salesperson_performance(): assert "CustID" not in sql +def test_build_direct_orders_sales_sql_for_salesperson_order_count(): + service = AskService.__new__(AskService) + sql = service._build_direct_orders_sales_sql( + "Create a bar chart of top 10 SalesPerson by order count" + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_tblSales"."SalesPerson" AS "SalesPerson", ' + 'COUNT(*) AS "OrderCount" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."SalesPerson" IS NOT NULL ' + 'GROUP BY "dbo_tblSales"."SalesPerson" ' + 'ORDER BY COUNT(*) DESC' + ) + + +def test_build_direct_orders_sales_sql_for_top_new_orders_q1(): + service = AskService.__new__(AskService) + sql = service._build_direct_orders_sales_sql( + "Show the top 20 new orders for period 2026-Q1" + ) + + assert sql == ( + 'SELECT TOP 20 "dbo_tblSales"."BU" AS "BU", ' + '"dbo_tblSales"."Market" AS "Market", ' + '"dbo_tblSales"."Customer" AS "Customer", ' + '"dbo_tblSales"."ProdName" AS "ProdName", ' + '"dbo_tblSales"."SalesValue" AS "SalesValue" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."OrdDate" >= \'2026-01-01 00:00:00\' ' + 'AND "dbo_tblSales"."OrdDate" < \'2026-04-01 00:00:00\' ' + 'ORDER BY "dbo_tblSales"."SalesValue" DESC' + ) + + +def test_build_direct_orders_sales_sql_for_market_growth_comparison(): + service = AskService.__new__(AskService) + sql = service._build_direct_orders_sales_sql( + "Which markets had the highest growth in the first 6 months of this year compared with the same period last year?" + ) + + assert sql is not None + assert '"dbo_tblSales"."Market" AS "Market"' in sql + assert '"CurrentPeriodSales"' in sql + assert '"PreviousPeriodSales"' in sql + assert '"SalesGrowth"' in sql + assert "2026-01-01" in sql + assert "2025-01-01" in sql + + def test_build_schema_grounded_sales_sql_requires_sales_schema(): service = AskService.__new__(AskService) From 0825093d6fa3d617af83624c57f8126b925857eb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 03:07:46 +0530 Subject: [PATCH 0275/1087] Stop completed UI polling loops --- .../pages/home/promptThread/ChartAnswer.tsx | 6 ++-- wren-ui/src/hooks/useAdjustAnswer.tsx | 19 ++++++++--- wren-ui/src/hooks/useAskPrompt.tsx | 34 ++++++++++++++----- .../useRecommendedQuestionsInstruction.tsx | 17 +++++++--- wren-ui/src/pages/home/[id].tsx | 34 +++++++++++++++---- 5 files changed, 85 insertions(+), 25 deletions(-) diff --git a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx index 851d7f2ae4..9ec2b9c894 100644 --- a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx @@ -100,12 +100,14 @@ export default function ChartAnswer(props: AnswerResultProps) { }, }); - // initial trigger when render + // Fetch preview data after the chart task has a terminal schema/result. useEffect(() => { + if (!getIsChartFinished(status)) return; + previewData({ variables: { where: { responseId: threadResponse.id } }, }); - }, []); + }, [previewData, status, threadResponse.id]); const chartSpec = useMemo(() => { if ( diff --git a/wren-ui/src/hooks/useAdjustAnswer.tsx b/wren-ui/src/hooks/useAdjustAnswer.tsx index bec1fc39a4..7a960785f0 100644 --- a/wren-ui/src/hooks/useAdjustAnswer.tsx +++ b/wren-ui/src/hooks/useAdjustAnswer.tsx @@ -127,17 +127,28 @@ export default function useAdjustAnswer(threadId?: number) { if (threadResponsePollingSessionRef.current !== pollingSessionId) return; } + let shouldContinuePolling = true; try { const request = fetchThreadResponse({ variables: { responseId }, - }).then(() => undefined); - threadResponsePollingRequestRef.current = request; - await request; + }); + threadResponsePollingRequestRef.current = request.then( + () => undefined, + ); + const result = await request; + const task = result.data?.threadResponse?.adjustmentTask; + if (!task || getIsFinished(task.status)) { + shouldContinuePolling = false; + stopThreadResponsePolling(); + } } catch (error) { console.error(error); } finally { threadResponsePollingRequestRef.current = null; - if (threadResponsePollingSessionRef.current === pollingSessionId) { + if ( + shouldContinuePolling && + threadResponsePollingSessionRef.current === pollingSessionId + ) { threadResponsePollingRef.current = setTimeout( run, threadResponsePollingDelayRef.current, diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index a524ee7b3c..b34e0a66f9 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -267,17 +267,26 @@ export default function useAskPrompt(threadId?: number) { if (askingTaskPollingSessionRef.current !== pollingSessionId) return; } + let shouldContinuePolling = true; try { const request = fetchAskingTask({ variables: { taskId }, - }).then(() => undefined); - askingTaskPollingRequestRef.current = request; - await request; + }); + askingTaskPollingRequestRef.current = request.then(() => undefined); + const result = await request; + const task = result.data?.askingTask; + if (!task || getIsFinished(task.status)) { + shouldContinuePolling = false; + stopAskingTaskPolling(); + } } catch (error) { console.error(error); } finally { askingTaskPollingRequestRef.current = null; - if (askingTaskPollingSessionRef.current === pollingSessionId) { + if ( + shouldContinuePolling && + askingTaskPollingSessionRef.current === pollingSessionId + ) { askingTaskPollingRef.current = setTimeout( run, askingTaskPollingDelayRef.current, @@ -312,17 +321,26 @@ export default function useAskPrompt(threadId?: number) { if (recommendedPollingSessionRef.current !== pollingSessionId) return; } + let shouldContinuePolling = true; try { const request = fetchInstantRecommendedQuestions({ variables: { taskId }, - }).then(() => undefined); - recommendedPollingRequestRef.current = request; - await request; + }); + recommendedPollingRequestRef.current = request.then(() => undefined); + const result = await request; + const task = result.data?.instantRecommendedQuestions; + if (!task || isRecommendedFinished(task.status)) { + shouldContinuePolling = false; + stopRecommendedPolling(); + } } catch (error) { console.error(error); } finally { recommendedPollingRequestRef.current = null; - if (recommendedPollingSessionRef.current === pollingSessionId) { + if ( + shouldContinuePolling && + recommendedPollingSessionRef.current === pollingSessionId + ) { recommendedPollingRef.current = setTimeout( run, recommendedPollingDelayRef.current, diff --git a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx index a99a645b20..243843006d 100644 --- a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx +++ b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx @@ -92,15 +92,24 @@ export default function useRecommendedQuestionsInstruction() { if (pollingSessionRef.current !== pollingSessionId) return; } + let shouldContinuePolling = true; try { - const request = fetchRecommendationQuestions().then(() => undefined); - pollingRequestRef.current = request; - await request; + const request = fetchRecommendationQuestions(); + pollingRequestRef.current = request.then(() => undefined); + const result = await request; + const task = result.data?.getProjectRecommendationQuestions; + if (!task || isRecommendedFinished(task.status)) { + shouldContinuePolling = false; + stopPolling(); + } } catch (error) { console.error(error); } finally { pollingRequestRef.current = null; - if (pollingSessionRef.current === pollingSessionId) { + if ( + shouldContinuePolling && + pollingSessionRef.current === pollingSessionId + ) { pollingRef.current = setTimeout(run, pollingDelayRef.current); } } diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index 7ce5543de0..9cc0a8176e 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -244,17 +244,28 @@ export default function HomeThread() { } } + let shouldContinuePolling = true; try { const request = fetchThreadResponse({ variables: { responseId }, - }).then(() => undefined); - threadResponsePollingRequestRef.current = request; - await request; + }); + threadResponsePollingRequestRef.current = request.then( + () => undefined, + ); + const result = await request; + if (getThreadResponseIsFinished(result.data?.threadResponse)) { + shouldContinuePolling = false; + stopThreadResponsePolling(); + setShowRecommendedQuestions(true); + } } catch (error) { console.error(error); } finally { threadResponsePollingRequestRef.current = null; - if (threadResponsePollingSessionRef.current === pollingSessionId) { + if ( + shouldContinuePolling && + threadResponsePollingSessionRef.current === pollingSessionId + ) { threadResponsePollingRef.current = setTimeout( run, threadResponsePollingDelayRef.current, @@ -309,17 +320,26 @@ export default function HomeThread() { } } + let shouldContinuePolling = true; try { const request = fetchThreadRecommendationQuestions({ variables: { threadId: nextThreadId }, - }).then(() => undefined); - threadRecommendationPollingRequestRef.current = request; - await request; + }); + threadRecommendationPollingRequestRef.current = request.then( + () => undefined, + ); + const result = await request; + const task = result.data?.getThreadRecommendationQuestions; + if (!task || isRecommendedFinished(task.status)) { + shouldContinuePolling = false; + stopThreadRecommendationPolling(); + } } catch (error) { console.error(error); } finally { threadRecommendationPollingRequestRef.current = null; if ( + shouldContinuePolling && threadRecommendationPollingSessionRef.current === pollingSessionId ) { threadRecommendationPollingRef.current = setTimeout( From 0878b775b75535556562c9f24d419c92e77a86b9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 03:22:52 +0530 Subject: [PATCH 0276/1087] Generate monthly order count SQL from date columns --- wren-ai-service/src/web/v1/services/ask.py | 58 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 27 +++++++++ 2 files changed, 85 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2764fd9877..17714fed15 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -749,6 +749,64 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql + wants_monthly_count = ( + "monthly" in normalized_query + and any(term in normalized_query for term in ("count", "volume")) + and any(term in normalized_query for term in ("order", "orders")) + ) + if wants_monthly_count: + date_candidates = ( + ("InvDate", "InvoiceDate", "Invoice Date") + if "invdate" in compact_query or "invoice" in normalized_query + else ( + "OrdDate", + "OrderDate", + "NewOrderDate", + "InvDate", + "InvoiceDate", + "Date", + ) + ) + scored_tables: list[tuple[int, dict[str, Any], str]] = [] + for table in tables: + date_column = self._find_schema_column( + table, date_candidates, temporal=True + ) + if not date_column: + continue + table_name = str(table.get("name") or "") + score = 10 + if "sales" in table_name.lower() or "order" in table_name.lower(): + score += 5 + if self._find_schema_column( + table, ("OrdNo", "OrderNo", "OrderId", "InvoiceNo") + ): + score += 3 + scored_tables.append((score, table, date_column)) + + if scored_tables: + _, table, date_column = sorted( + scored_tables, key=lambda item: item[0], reverse=True + )[0] + table_name = table.get("name") + if table_name and date_column: + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + date_ref = ( + f"{table_ref}.{self._quote_sql_identifier(date_column)}" + ) + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "OrderCount" ' + f"FROM {table_ref}" + f"{self._build_date_filter(table_name, date_column, query)} " + f"GROUP BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref})" + ) + dimension_candidates: list[tuple[str, ...]] = [] if "salesperson" in normalized_query or "sales person" in normalized_query: dimension_candidates.append(("SalesPerson", "Sales Rep", "SalesRep")) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 29e947dff8..03f070016e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -274,6 +274,33 @@ def test_build_schema_grounded_sales_sql_for_order_invoice_conversion_rate(): assert "P-M" not in sql +def test_build_schema_grounded_sales_sql_for_monthly_order_count_by_invdate(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show monthly order count by InvDate.", + [ + """ + CREATE TABLE dbo_tblSales ( + OrdNo VARCHAR, + InvDate TIMESTAMP, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_tblSales"."InvDate") AS "year", ' + 'DATEPART(MONTH, "dbo_tblSales"."InvDate") AS "month", ' + 'COUNT(*) AS "OrderCount" ' + 'FROM "dbo_tblSales" ' + 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."InvDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."InvDate") ' + 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."InvDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."InvDate")' + ) + + def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From fc3fe74a5fdc5b05871e8557c3a437c46eec5eb7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 10:52:58 +0530 Subject: [PATCH 0277/1087] Handle explicit table preview questions --- wren-ai-service/src/web/v1/services/ask.py | 118 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 26 ++++ 2 files changed, 142 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 17714fed15..c3b947e86c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -484,6 +484,71 @@ def _find_first_schema_column( def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' + def _build_explicit_table_preview_sql( + self, query: str, table_ddls: list[str] + ) -> tuple[str, str] | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) + if not normalized_query: + return None + + if not re.search( + r"\b(?:first|top|sample|preview|show|list)\b", + normalized_query, + flags=re.IGNORECASE, + ): + return None + if not re.search( + r"\b(?:rows?|records?|data)\b", normalized_query, flags=re.IGNORECASE + ): + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + normalized_query_key = re.sub(r"[^a-z0-9]", "", normalized_query.lower()) + scored_tables: list[tuple[int, str]] = [] + for table in tables: + table_name = table.get("name") + if not table_name: + continue + table_name = str(table_name) + normalized_table = re.sub(r"[^a-z0-9]", "", table_name.lower()) + if not normalized_table: + continue + if normalized_table in normalized_query_key: + scored_tables.append((100 + len(normalized_table), table_name)) + continue + + table_without_schema = re.split(r"[.$]", table_name)[-1] + normalized_short_name = re.sub( + r"[^a-z0-9]", "", table_without_schema.lower() + ) + if normalized_short_name and normalized_short_name in normalized_query_key: + scored_tables.append((80 + len(normalized_short_name), table_name)) + + if not scored_tables: + return None + + _, table_name = sorted(scored_tables, reverse=True)[0] + limit = self._extract_requested_top_n(query, default_value=10) + return ( + f"SELECT TOP {limit} * FROM {self._quote_sql_identifier(table_name)}", + table_name, + ) + + def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: + table_names: list[str] = [] + for match in re.finditer( + r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if table_name and table_name not in table_names: + table_names.append(table_name) + return table_names + def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -2780,12 +2845,61 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + if not documents: + explicit_table_names = self._extract_explicit_table_names_from_query( + user_query + ) + if explicit_table_names: + logger.info( + "Retrying schema retrieval for explicit tables query_id %s: %s", + query_id, + explicit_table_names, + ) + retrieval_result = await self._run_with_timeout( + "Explicit table schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=histories, + enable_column_pruning=enable_column_pruning, + ), + timeout_seconds=self._schema_retrieval_timeout_seconds, + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if audit_log_activity_sql := self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names + if explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ): + explicit_sql, explicit_table_name = explicit_table_preview + logger.info( + "Using explicit table preview SQL for query_id %s and table %s", + query_id, + explicit_table_name, + ) + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + api_results = [ + AskResult( + **{ + "sql": explicit_sql, + "type": "llm", + } + ) + ] + + if not api_results and ( + audit_log_activity_sql := self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ) ): logger.info( "Using schema-grounded audit log activity SQL for query_id %s", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 03f070016e..869b0e2734 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -92,6 +92,32 @@ def test_build_schema_grounded_sales_sql_requires_sales_schema(): ) +def test_build_explicit_table_preview_sql_for_named_table(): + service = AskService.__new__(AskService) + result = service._build_explicit_table_preview_sql( + "Show the first 10 rows from tblNewOrders", + [ + """ + CREATE TABLE tblNewOrders ( + OrdNo VARCHAR, + Customer VARCHAR, + InvDate TIMESTAMP + ); + """ + ], + ) + + assert result == ('SELECT TOP 10 * FROM "tblNewOrders"', "tblNewOrders") + + +def test_extract_explicit_table_names_from_query(): + service = AskService.__new__(AskService) + + assert service._extract_explicit_table_names_from_query( + "Show the first 10 rows from tblNewOrders" + ) == ["tblNewOrders"] + + def test_build_schema_grounded_sales_sql_for_top_markets(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From a0703420f272e33d515256837601237a20b4f821 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 11:16:57 +0530 Subject: [PATCH 0278/1087] Fast path explicit table preview asks --- wren-ai-service/src/web/v1/services/ask.py | 69 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 21 +++++- 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c3b947e86c..282cb0c897 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1240,6 +1240,12 @@ def _build_yoy_sales_change_sql( def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): return max(1, min(int(match.group(1)), 100)) + if match := re.search( + r"\b(?:first|limit)\s+(\d+)\b", query or "", flags=re.IGNORECASE + ): + return max(1, min(int(match.group(1)), 100)) + if match := re.search(r"\b(\d+)\s+rows?\b", query or "", flags=re.IGNORECASE): + return max(1, min(int(match.group(1)), 100)) return default_value def _build_manufacturing_throughput_sql( @@ -2506,6 +2512,69 @@ async def ask( ) return results + explicit_table_names = self._extract_explicit_table_names_from_query( + user_query + ) + if explicit_table_names: + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=user_query, + intent_reasoning="Explicit table name detected; retrieving that deployed schema directly.", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Explicit table schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=histories, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + logger.info( + "Retrieved explicit tables for query_id %s: %s", + query_id, + table_names, + ) + + if explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ): + explicit_sql, explicit_table_name = explicit_table_preview + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + api_results = [ + AskResult( + **{ + "sql": explicit_sql, + "type": "llm", + } + ) + ] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table preview request matched deployed schema.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if self._is_direct_heuristic_sql_query(user_query): self._ask_results[query_id] = AskResultResponse( status="searching", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 869b0e2734..e0b3046ce1 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -95,7 +95,7 @@ def test_build_schema_grounded_sales_sql_requires_sales_schema(): def test_build_explicit_table_preview_sql_for_named_table(): service = AskService.__new__(AskService) result = service._build_explicit_table_preview_sql( - "Show the first 10 rows from tblNewOrders", + "Show the first 5 rows from tblNewOrders", [ """ CREATE TABLE tblNewOrders ( @@ -107,7 +107,24 @@ def test_build_explicit_table_preview_sql_for_named_table(): ], ) - assert result == ('SELECT TOP 10 * FROM "tblNewOrders"', "tblNewOrders") + assert result == ('SELECT TOP 5 * FROM "tblNewOrders"', "tblNewOrders") + + +def test_build_explicit_table_preview_sql_for_show_data_prompt(): + service = AskService.__new__(AskService) + result = service._build_explicit_table_preview_sql( + "Show data from CustomerMaster", + [ + """ + CREATE TABLE CustomerMaster ( + CustomerId VARCHAR, + CustomerName VARCHAR + ); + """ + ], + ) + + assert result == ('SELECT TOP 10 * FROM "CustomerMaster"', "CustomerMaster") def test_extract_explicit_table_names_from_query(): From 71bafc2684d0795acedccbe76a6c7a2d72d9836b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 11:31:36 +0530 Subject: [PATCH 0279/1087] Use direct schema lookup for explicit table asks --- .../retrieval/db_schema_retrieval.py | 28 +++++++++++++++---- wren-ai-service/src/web/v1/services/ask.py | 23 +++++++++++++++ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 177448fd6f..45d8b20e38 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -255,7 +255,16 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: @observe(capture_input=False, capture_output=False) -async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: +async def embedding( + query: str, + embedder: Any, + histories: list[AskHistory], + tables: Optional[list[str]] = None, +) -> dict: + if tables: + logger.info("Skipping embedding retrieval for explicit tables: %s", tables) + return {} + if query: if histories: previous_query_summaries = [history.question for history in histories] @@ -306,13 +315,20 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - query: str, table_retrieval: dict, project_id: str, dbschema_retriever: Any + query: str, + table_retrieval: dict, + project_id: str, + dbschema_retriever: Any, + tables: Optional[list[str]] = None, ) -> list[Document]: - tables = table_retrieval.get("documents", []) table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) + if tables: + table_names.extend(tables) + else: + retrieved_tables = table_retrieval.get("documents", []) + for table in retrieved_tables: + content = ast.literal_eval(table.content) + table_names.append(content["name"]) table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 282cb0c897..2866075581 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2536,6 +2536,7 @@ async def ask( timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, + 20, ), ) documents, table_names, table_ddls = ( @@ -2575,6 +2576,28 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + error_message = ( + "The requested table was not found in the deployed schema: " + + ", ".join(explicit_table_names) + ) + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_DATA", + message=error_message, + ), + rephrased_question=user_query, + intent_reasoning="Explicit table preview request did not match any deployed schema table.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = error_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if self._is_direct_heuristic_sql_query(user_query): self._ask_results[query_id] = AskResultResponse( status="searching", From 5ee0cf4a297be1808c30189ae52983a92a98b000 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 11:50:48 +0530 Subject: [PATCH 0280/1087] Bypass table retriever for explicit schema lookup --- .../retrieval/db_schema_retrieval.py | 14 ++---- wren-ai-service/src/web/v1/services/ask.py | 50 ++++++++++++++----- 2 files changed, 42 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 45d8b20e38..59eca37846 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -301,16 +301,12 @@ async def table_retrieval( filters=base_filters, ) return result - else: - base_filters["conditions"].append( - {"field": "name", "operator": "in", "value": tables} - ) - result = await table_retriever.run( - query_embedding=[], - filters=base_filters, - ) - return result + if tables: + logger.info("Skipping table-description retrieval for explicit tables: %s", tables) + return {"documents": []} + + return {"documents": []} @observe(capture_input=False) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2866075581..0685988a03 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2918,19 +2918,40 @@ async def ask( is_followup=True if histories else False, ) - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - histories=histories, - project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) + try: + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), ), - ), - timeout_seconds=self._schema_retrieval_timeout_seconds, - ) + timeout_seconds=self._schema_retrieval_timeout_seconds, + ) + except TimeoutError as error: + logger.warning( + "Schema retrieval timed out; falling back to deployed schemas. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + error, + ) + retrieval_result = await self._run_with_timeout( + "Deployed schema fallback retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 30, + ), + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) @@ -2956,7 +2977,10 @@ async def ask( histories=histories, enable_column_pruning=enable_column_pruning, ), - timeout_seconds=self._schema_retrieval_timeout_seconds, + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 20, + ), ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} From cda9661e69155c69acffab945baecceb9b27ef60 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 12:04:10 +0530 Subject: [PATCH 0281/1087] Skip intent classification for analytics asks --- wren-ai-service/src/web/v1/services/ask.py | 32 ++++++++++++++++++++-- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0685988a03..5944a9aa8d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -175,6 +175,8 @@ def _is_data_analysis_query(self, query: str) -> bool: "compare", "count", "cost", + "claim", + "claims", "customer", "customers", "dashboard", @@ -193,10 +195,16 @@ def _is_data_analysis_query(self, query: str) -> bool: "pcb", "performance", "profit", + "product", + "products", + "product type", + "product types", "quarter", "quantity", "rank", "ranking", + "region", + "regions", "repair", "resolved", "revenue", @@ -878,7 +886,9 @@ def _build_schema_grounded_analytics_sql( if "business unit" in normalized_query or "bu" in normalized_query: dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) if "market" in normalized_query: - dimension_candidates.append(("Market", "MarketType")) + dimension_candidates.append(("Market", "MarketType", "Region")) + if "region" in normalized_query: + dimension_candidates.append(("Region", "Market", "Area", "Territory")) if "division" in normalized_query: dimension_candidates.append(("Division",)) if ( @@ -2676,7 +2686,23 @@ async def ask( ) historical_question_result = [] - if not api_results: + should_skip_pre_sql_retrieval = self._is_data_analysis_query( + user_query + ) + if should_skip_pre_sql_retrieval: + rephrased_question = user_query + intent_reasoning = ( + "Detected a deployed-data analytics question; skipping " + "intent classification and using SQL generation." + ) + sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + logger.info( + "Skipping pre-SQL retrieval for analytics query_id %s: %s", + query_id, + user_query, + ) + + if not api_results and not should_skip_pre_sql_retrieval: historical_question = await self._run_with_timeout( "Historical question retrieval", self._pipelines["historical_question"].run( @@ -2712,7 +2738,7 @@ async def ask( if valid_historical_results: api_results = valid_historical_results sql_generation_reasoning = "" - elif not api_results: + elif not api_results and not should_skip_pre_sql_retrieval: original_user_query = user_query # Run both pipeline operations concurrently sql_samples_task, instructions_task = await self._run_with_timeout( From 9e5b66bd267f5e7f9e01c194cfdfa73d5e6482f2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 13:20:09 +0530 Subject: [PATCH 0282/1087] Count orders for customer order trends --- wren-ai-service/src/web/v1/services/ask.py | 38 +++++++++++++++---- .../pytest/services/test_ask_sales_sql.py | 29 ++++++++++++++ 2 files changed, 60 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 5944a9aa8d..3fc3369949 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -932,7 +932,19 @@ def _build_schema_grounded_analytics_sql( "Value", "Amount", ) - wants_trend = "trend" in normalized_query or "line chart" in normalized_query + wants_trend = ( + "trend" in normalized_query + or "line chart" in normalized_query + or "over time" in normalized_query + or "last 12 months" in normalized_query + ) + wants_order_count_metric = ( + any(term in normalized_query for term in ("order", "orders", "new order", "new orders")) + and not any( + term in normalized_query + for term in ("value", "amount", "revenue", "sales", "cost", "margin") + ) + ) wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) wants_detail_rows = ( wants_top @@ -989,12 +1001,24 @@ def _build_schema_grounded_analytics_sql( if wants_trend and date_column: date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - metric_expr = ( - f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - if measure - else "COUNT(*)" - ) - metric_alias = f"Total{measure}" if measure else "OrderCount" + if wants_order_count_metric: + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), + ) + metric_expr = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" + if order_column + else "COUNT(*)" + ) + metric_alias = "OrderCount" + else: + metric_expr = ( + f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + if measure + else "COUNT(*)" + ) + metric_alias = f"Total{measure}" if measure else "RecordCount" select_parts = [ f"DATEPART(YEAR, {date_ref}) AS \"year\"", f"DATEPART(MONTH, {date_ref}) AS \"month\"", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index e0b3046ce1..79f44dfd00 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -344,6 +344,35 @@ def test_build_schema_grounded_sales_sql_for_monthly_order_count_by_invdate(): ) +def test_build_schema_grounded_sales_sql_counts_new_orders_by_customer_over_time(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show new orders by CustName over the last 12 months using dbo.XStageNewOrders OrdDate and CustName.", + [ + """ + CREATE TABLE dbo_XStageNewOrders ( + OrdNo VARCHAR, + OrdDate TIMESTAMP, + CustName VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate") AS "year", ' + 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate") AS "month", ' + '"dbo_XStageNewOrders"."CustName" AS "CustName", ' + 'COUNT(DISTINCT "dbo_XStageNewOrders"."OrdNo") AS "OrderCount" ' + 'FROM "dbo_XStageNewOrders" ' + 'GROUP BY DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate"), ' + '"dbo_XStageNewOrders"."CustName" ' + 'ORDER BY DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate")' + ) + + def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From c803c22791325cb55646ed9bba8ceaf34cd29ba6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 13:33:29 +0530 Subject: [PATCH 0283/1087] Select count metrics for explicit order trends --- wren-ai-service/src/web/v1/services/ask.py | 29 +++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3fc3369949..a0536c2afb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -734,6 +734,7 @@ def _select_best_analytics_table( required_dimensions: list[tuple[str, ...]], measure_candidates: tuple[str, ...], wants_date: bool = False, + allow_count_metric: bool = False, ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] @@ -749,6 +750,8 @@ def _select_best_analytics_table( measure = self._find_schema_column( table, measure_candidates, numeric=True ) + if not measure and not allow_count_metric: + continue date_column = self._find_schema_column( table, ( @@ -768,6 +771,8 @@ def _select_best_analytics_table( score = 10 * len([dimension for dimension in dimensions if dimension]) if measure: score += 8 + elif allow_count_metric: + score += 2 if date_column: score += 4 table_name = str(table.get("name") or "").lower() @@ -902,7 +907,11 @@ def _build_schema_grounded_analytics_sql( dimension_candidates.append( ("ProdName", "Product", "ProductName", "Item", "ProdCode") ) - if "customer" in normalized_query: + if ( + "customer" in normalized_query + or "custname" in compact_query + or "custno" in compact_query + ): dimension_candidates.append(("Customer", "CustName", "CustNo")) if not dimension_candidates: @@ -951,8 +960,21 @@ def _build_schema_grounded_analytics_sql( and ("new order" in normalized_query or "orders" in normalized_query) and any(term in normalized_query for term in ("including", "include")) ) - wants_date = wants_trend or "this year" in normalized_query or bool( - re.search(r"\b20\d{2}\b", normalized_query) + mentions_date_column = any( + column_name in compact_query + for column_name in ( + "orddate", + "invdate", + "orderdate", + "invoicedate", + "createdat", + ) + ) + wants_date = ( + wants_trend + or mentions_date_column + or "this year" in normalized_query + or bool(re.search(r"\b20\d{2}\b", normalized_query)) ) selected = self._select_best_analytics_table( @@ -960,6 +982,7 @@ def _build_schema_grounded_analytics_sql( dimension_candidates, measure_candidates, wants_date=wants_date, + allow_count_metric=wants_order_count_metric, ) if not selected: return None From 874c968f2c309412aaf77bb880e87519ca0a766f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 13:47:40 +0530 Subject: [PATCH 0284/1087] Prune SQL generation context after schema retrieval --- wren-ai-service/src/web/v1/services/ask.py | 323 ++++++++++++++---- .../pytest/services/test_ask_sales_sql.py | 52 +++ 2 files changed, 316 insertions(+), 59 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a0536c2afb..8b5f39ca4b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -555,6 +555,18 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_name = match.group(1).strip(".,;:()[]{}") if table_name and table_name not in table_names: table_names.append(table_name) + for match in re.finditer( + r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if ( + table_name + and ("." in table_name or "_" in table_name) + and table_name not in table_names + ): + table_names.append(table_name) return table_names def _build_direct_orders_sales_sql(self, query: str) -> str | None: @@ -1725,6 +1737,11 @@ def _build_heuristic_text_to_sql_fallback( if not normalized: return None + if schema_grounded_sql := self._build_schema_grounded_sales_sql( + query, table_ddls + ): + return schema_grounded_sql + if throughput_sql := self._build_manufacturing_throughput_sql( query, table_ddls, table_names=table_names ): @@ -2318,6 +2335,128 @@ def _extract_retrieval_metadata( ] return documents, table_names, table_ddls + def _normalize_schema_token(self, value: str) -> str: + return re.sub(r"[^a-z0-9]", "", (value or "").lower()) + + def _query_schema_terms(self, query: str) -> set[str]: + normalized_query = (query or "").lower() + stop_words = { + "about", + "against", + "from", + "give", + "group", + "grouped", + "list", + "month", + "monthly", + "over", + "rows", + "show", + "table", + "tables", + "the", + "this", + "using", + "what", + "which", + "with", + "year", + } + terms = { + self._normalize_schema_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", normalized_query) + if len(token) > 2 and token not in stop_words + } + return {term for term in terms if term} + + def _prune_sql_generation_context( + self, + query: str, + documents: list[dict], + table_names: list[str], + table_ddls: list[str], + *, + max_tables: int = 8, + ) -> tuple[list[dict], list[str], list[str]]: + if len(table_ddls) <= max_tables: + return documents, table_names, table_ddls + + parsed_tables = self._parse_schema_tables(table_ddls) + if not parsed_tables: + return documents, table_names, table_ddls[:max_tables] + + query_key = self._normalize_schema_token(query) + query_terms = self._query_schema_terms(query) + explicit_tables = { + self._normalize_schema_token(table_name) + for table_name in self._extract_explicit_table_names_from_query(query) + } + + scored: list[tuple[int, int]] = [] + for index, table in enumerate(parsed_tables): + table_name = str(table.get("name") or "") + normalized_table = self._normalize_schema_token(table_name) + normalized_short_table = self._normalize_schema_token( + re.split(r"[.$]", table_name)[-1] + ) + column_terms = { + self._normalize_schema_token(str(column.get("name") or "")) + for column in table.get("columns", []) + if column.get("name") + } + + score = 0 + if normalized_table in explicit_tables or normalized_short_table in explicit_tables: + score += 1000 + if normalized_table and normalized_table in query_key: + score += 500 + if normalized_short_table and normalized_short_table in query_key: + score += 450 + for term in query_terms: + if not term: + continue + if term == normalized_table or term == normalized_short_table: + score += 80 + elif term in normalized_table or term in normalized_short_table: + score += 40 + for column_term in column_terms: + if term == column_term: + score += 60 + elif term in column_term or column_term in term: + score += 25 + + if score > 0: + scored.append((score, index)) + + if not scored: + return documents, table_names, table_ddls[:max_tables] + + selected_indexes = [ + index + for _, index in sorted(scored, key=lambda item: item[0], reverse=True)[ + :max_tables + ] + ] + selected_indexes = sorted(selected_indexes) + pruned_documents = [ + documents[index] for index in selected_indexes if index < len(documents) + ] + pruned_table_names = [ + table_names[index] for index in selected_indexes if index < len(table_names) + ] + pruned_table_ddls = [ + table_ddls[index] for index in selected_indexes if index < len(table_ddls) + ] + + logger.info( + "Pruned SQL generation context from %s to %s tables for query: %s", + len(table_ddls), + len(pruned_table_ddls), + query, + ) + return pruned_documents, pruned_table_names, pruned_table_ddls + def _is_valid_select_sql(self, sql: Optional[str]) -> bool: if not isinstance(sql, str): return False @@ -2495,7 +2634,10 @@ async def ask( sql_samples = [] instructions = [] api_results = [] + documents = [] table_names = [] + table_ddls = [] + _retrieval_result = {} error_message = None invalid_sql = None allow_sql_generation_reasoning = ( @@ -2599,6 +2741,9 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) logger.info( "Retrieved explicit tables for query_id %s: %s", query_id, @@ -2633,27 +2778,61 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - error_message = ( - "The requested table was not found in the deployed schema: " - + ", ".join(explicit_table_names) - ) - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_DATA", - message=error_message, - ), - rephrased_question=user_query, - intent_reasoning="Explicit table preview request did not match any deployed schema table.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, + if documents and ( + deterministic_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ) + ): + api_results = [ + AskResult( + **{ + "sql": deterministic_sql, + "type": "llm", + } + ) + ] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + if not documents: + error_message = ( + "The requested table was not found in the deployed schema: " + + ", ".join(explicit_table_names) + ) + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_DATA", + message=error_message, + ), + rephrased_question=user_query, + intent_reasoning="Explicit table request did not match any deployed schema table.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = error_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + rephrased_question = user_query + intent_reasoning = ( + "Explicit table request matched deployed schema; generating SQL against retrieved schema." ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = error_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results + sql_user_query = self._rewrite_query_for_text_to_sql(user_query) if self._is_direct_heuristic_sql_query(user_query): self._ask_results[query_id] = AskResultResponse( @@ -2981,7 +3160,11 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - if not self._is_stopped(query_id, self._ask_results) and not api_results: + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and not documents + ): self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -3213,6 +3396,14 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + if documents and not api_results: + documents, table_names, table_ddls = self._prune_sql_generation_context( + sql_user_query, + documents, + table_names, + table_ddls, + ) + if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -3325,45 +3516,59 @@ async def ask( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - if histories: - text_to_sql_generation_results = await self._run_with_timeout( - "Follow-up SQL generation", - self._pipelines["followup_sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), - ) - else: - text_to_sql_generation_results = await self._run_with_timeout( - "SQL generation", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), + try: + if histories: + text_to_sql_generation_results = await self._run_with_timeout( + "Follow-up SQL generation", + self._pipelines["followup_sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), + ) + else: + text_to_sql_generation_results = await self._run_with_timeout( + "SQL generation", + self._pipelines["sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), + ) + except TimeoutError as generation_timeout: + logger.warning( + "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + query_id, + generation_timeout, ) + text_to_sql_generation_results = { + "post_process": { + "valid_generation_result": None, + "invalid_generation_result": None, + } + } + error_message = str(generation_timeout) if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 79f44dfd00..8c4a26f966 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -135,6 +135,58 @@ def test_extract_explicit_table_names_from_query(): ) == ["tblNewOrders"] +def test_extract_explicit_table_names_from_using_clause(): + service = AskService.__new__(AskService) + + assert service._extract_explicit_table_names_from_query( + "Show new orders by CustName using dbo.XStageNewOrders OrdDate and CustName" + ) == ["dbo.XStageNewOrders"] + assert service._extract_explicit_table_names_from_query( + "Show new orders using OrdDate and CustName" + ) == [] + + +def test_prune_sql_generation_context_prefers_referenced_table_and_columns(): + service = AskService.__new__(AskService) + table_ddls = [ + """ + CREATE TABLE dbo_Customers ( + CustomerId VARCHAR, + CustomerName VARCHAR + ); + """, + """ + CREATE TABLE dbo_Products ( + ProductId VARCHAR, + ProductName VARCHAR + ); + """, + """ + CREATE TABLE dbo_XStageNewOrders ( + OrdNo VARCHAR, + OrdDate TIMESTAMP, + CustName VARCHAR + ); + """, + ] + documents = [ + {"table_name": "dbo_Customers", "table_ddl": table_ddls[0]}, + {"table_name": "dbo_Products", "table_ddl": table_ddls[1]}, + {"table_name": "dbo_XStageNewOrders", "table_ddl": table_ddls[2]}, + ] + + _, table_names, pruned_ddls = service._prune_sql_generation_context( + "Show new orders by CustName using dbo.XStageNewOrders OrdDate and CustName", + documents, + [document["table_name"] for document in documents], + table_ddls, + max_tables=1, + ) + + assert table_names == ["dbo_XStageNewOrders"] + assert pruned_ddls == [table_ddls[2]] + + def test_build_schema_grounded_sales_sql_for_top_markets(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From fbd712168b169cca044fbf7e607c0535c99c67c4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 13:51:05 +0530 Subject: [PATCH 0285/1087] Aggregate categorical fields for status questions --- wren-ai-service/src/web/v1/services/ask.py | 126 ++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 25 ++++ 2 files changed, 151 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8b5f39ca4b..e7add61a4d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -839,6 +839,11 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + wants_monthly_count = ( "monthly" in normalized_query and any(term in normalized_query for term in ("count", "volume")) @@ -1165,6 +1170,127 @@ def _build_contribution_sql( f"ORDER BY {metric_expr} DESC" ) + def _build_generic_categorical_count_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + if re.search( + r"\b(?:first|top|sample|preview|show|list)\b.*\b(?:rows?|records?|data)\b", + normalized_query, + ): + return None + + wants_categorical_summary = any( + term in normalized_query + for term in ( + "bar chart", + "by ", + "chart", + "count", + "distribution", + "donut chart", + "frequency", + "group by", + "grouped by", + "pie chart", + "status", + "type", + "category", + ) + ) + if not wants_categorical_summary: + return None + + query_key = self._normalize_schema_token(query) + query_terms = self._query_schema_terms(query) + scored: list[tuple[int, dict[str, Any], str]] = [] + low_value_column_patterns = ( + "id", + "no", + "number", + "date", + "time", + "description", + "comment", + "note", + "remark", + ) + + for table in tables: + table_name = str(table.get("name") or "") + normalized_table = self._normalize_schema_token(table_name) + normalized_short_table = self._normalize_schema_token( + re.split(r"[.$]", table_name)[-1] + ) + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if not column_name or not self._is_text_schema_type(column_type): + continue + normalized_column = self._normalize_schema_token(column_name) + if not normalized_column: + continue + + score = 0 + if normalized_table and normalized_table in query_key: + score += 120 + if normalized_short_table and normalized_short_table in query_key: + score += 100 + if normalized_column and normalized_column in query_key: + score += 180 + for term in query_terms: + if term == normalized_column: + score += 100 + elif term in normalized_column or normalized_column in term: + score += 45 + if term == normalized_table or term == normalized_short_table: + score += 40 + elif term in normalized_table or term in normalized_short_table: + score += 20 + if "status" in normalized_query and "status" in normalized_column: + score += 90 + if "category" in normalized_query and "category" in normalized_column: + score += 80 + if "type" in normalized_query and "type" in normalized_column: + score += 70 + if any(pattern == normalized_column for pattern in low_value_column_patterns): + score -= 100 + elif any( + normalized_column.endswith(pattern) + for pattern in low_value_column_patterns + ): + score -= 35 + + if score > 0: + scored.append((score, table, column_name)) + + if not scored: + return None + + _, table, dimension = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + table_name = table.get("name") + if not table_name: + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + top_n = self._extract_requested_top_n(query, default_value=0) + top_clause = f"TOP {top_n} " if top_n else "" + return ( + f"SELECT {top_clause}{dimension_ref} AS {self._quote_sql_identifier(dimension)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f"ORDER BY COUNT(*) DESC" + ) + def _build_order_invoice_conversion_sql( self, query: str, tables: list[dict[str, Any]] ) -> str | None: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 8c4a26f966..f97e7ba25a 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -474,6 +474,31 @@ def test_build_schema_grounded_sales_sql_for_product_type_contribution(): ) +def test_build_schema_grounded_sql_counts_categorical_status_values(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show refund status distribution.", + [ + """ + CREATE TABLE dbo_ytblRefund ( + Refund_Id VARCHAR, + Refund_Status VARCHAR, + CustomerName VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_ytblRefund"."Refund_Status" AS "Refund_Status", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ytblRefund" ' + 'WHERE "dbo_ytblRefund"."Refund_Status" IS NOT NULL ' + 'GROUP BY "dbo_ytblRefund"."Refund_Status" ' + 'ORDER BY COUNT(*) DESC' + ) + + def test_build_schema_grounded_sales_sql_for_yoy_waterfall_dimensions(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From e975cec0ef43d7db95984c01dc2216e9261f2c79 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 16:31:49 +0530 Subject: [PATCH 0286/1087] Use fast SQL generation for standalone data questions --- wren-ai-service/src/web/v1/services/ask.py | 82 ++++++++++++++----- .../pytest/services/test_ask_sales_sql.py | 11 +++ 2 files changed, 71 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e7add61a4d..f1dc6dd6b1 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -224,6 +224,22 @@ def _is_data_analysis_query(self, query: str) -> bool: } return any(term in normalized for term in analysis_terms) + def _needs_conversation_context(self, query: str) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + return any( + re.search(pattern, normalized) + for pattern in ( + r"\b(previous|last|above|earlier)\s+(query|question|answer|result|sql|chart)\b", + r"\b(same|that|those|them|it)\s+(table|query|question|result|chart|sql|period|filter)\b", + r"\b(use|using|based on|compare with|compared with)\s+(that|previous|last|above|earlier)\b", + r"\bwhat about\b", + r"\bhow about\b", + ) + ) + def _rewrite_query_for_text_to_sql(self, query: str) -> str: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -3530,6 +3546,19 @@ async def ask( table_ddls, ) + sql_generation_histories = histories + if self._is_data_analysis_query( + sql_user_query + ) and not self._needs_conversation_context(sql_user_query): + sql_generation_histories = [] + allow_sql_generation_reasoning = False + allow_sql_knowledge_retrieval = False + max_sql_correction_retries = min(max_sql_correction_retries, 1) + logger.info( + "Using fast standalone SQL generation path for query_id %s", + query_id, + ) + if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -3545,7 +3574,7 @@ async def ask( is_followup=True if histories else False, ) - if histories: + if sql_generation_histories: try: sql_generation_reasoning = ( await self._run_with_timeout( @@ -3555,7 +3584,7 @@ async def ask( ].run( query=sql_user_query, contexts=table_ddls, - histories=histories, + histories=sql_generation_histories, sql_samples=sql_samples, instructions=instructions, configuration=ask_request.configurations, @@ -3616,25 +3645,34 @@ async def ask( is_followup=True if histories else False, ) - sql_functions, sql_knowledge = await self._run_with_timeout( - "SQL helper retrieval", - asyncio.gather( - ( - self._pipelines["sql_functions_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_functions_retrieval - else _return_value([]) - ), - ( - self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_knowledge_retrieval - else _return_value(None) + try: + sql_functions, sql_knowledge = await self._run_with_timeout( + "SQL helper retrieval", + asyncio.gather( + ( + self._pipelines["sql_functions_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_functions_retrieval + else _return_value([]) + ), + ( + self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_knowledge_retrieval + else _return_value(None) + ), ), - ), - ) + timeout_seconds=min(self._pipeline_timeout_seconds, 10), + ) + except TimeoutError as helper_timeout: + logger.warning( + "SQL helper retrieval timed out for query_id %s; continuing with schema only: %s", + query_id, + helper_timeout, + ) + sql_functions, sql_knowledge = [], None has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -3643,14 +3681,14 @@ async def ask( has_json_field = _retrieval_result.get("has_json_field", False) try: - if histories: + if sql_generation_histories: text_to_sql_generation_results = await self._run_with_timeout( "Follow-up SQL generation", self._pipelines["followup_sql_generation"].run( query=sql_user_query, contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, - histories=histories, + histories=sql_generation_histories, project_id=ask_request.project_id, sql_samples=sql_samples, instructions=instructions, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index f97e7ba25a..86add879ae 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -146,6 +146,17 @@ def test_extract_explicit_table_names_from_using_clause(): ) == [] +def test_needs_conversation_context_only_for_true_followups(): + service = AskService.__new__(AskService) + + assert not service._needs_conversation_context( + "Show refund status distribution by Refund_Status" + ) + assert service._needs_conversation_context( + "What about the same period from the previous result?" + ) + + def test_prune_sql_generation_context_prefers_referenced_table_and_columns(): service = AskService.__new__(AskService) table_ddls = [ From 069e7474e8aad795ceda85f9267af6c256c82756 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 17:09:12 +0530 Subject: [PATCH 0287/1087] Filter null dimensions in grouped SQL --- wren-ai-service/src/web/v1/services/ask.py | 106 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 49 ++++++++ 2 files changed, 152 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f1dc6dd6b1..497447beb3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -756,6 +756,17 @@ def _build_date_filter(self, table_name: str, date_column: str, query: str) -> s ) return "" + def _append_not_null_filters( + self, where_clause: str, column_refs: list[str] + ) -> str: + conditions = [f"{column_ref} IS NOT NULL" for column_ref in column_refs] + if not conditions: + return where_clause + + if where_clause.strip(): + return f"{where_clause.rstrip()} AND {' AND '.join(conditions)} " + return f" WHERE {' AND '.join(conditions)} " + def _select_best_analytics_table( self, tables: list[dict[str, Any]], @@ -1051,10 +1062,98 @@ def _build_schema_grounded_analytics_sql( ) return ( f"SELECT TOP {limit} {', '.join(select_parts)} " - f"FROM {table_ref}{date_filter} " + f"FROM {table_ref}" + f"{self._append_not_null_filters(date_filter, dimension_refs)} " f"ORDER BY {metric_ref} DESC" ) + wants_top_per_group = ( + len(dimensions) >= 2 + and any(term in normalized_query for term in ("highest", "top", "most")) + and any(term in normalized_query for term in ("each", "per ")) + ) + if wants_top_per_group: + partition_dimension = None + rank_dimension = None + if "market" in normalized_query: + partition_dimension = self._find_schema_column( + table, ("Market", "MarketType", "Region") + ) + if "region" in normalized_query and not partition_dimension: + partition_dimension = self._find_schema_column( + table, ("Region", "Market", "Area", "Territory") + ) + if "customer" in normalized_query: + rank_dimension = self._find_schema_column( + table, ("Customer", "CustName", "CustNo") + ) + if not partition_dimension: + partition_dimension = dimensions[0] + if not rank_dimension: + rank_dimension = next( + ( + dimension + for dimension in dimensions + if dimension != partition_dimension + ), + None, + ) + + if partition_dimension and rank_dimension: + partition_ref = ( + f"{table_ref}.{self._quote_sql_identifier(partition_dimension)}" + ) + rank_ref = f"{table_ref}.{self._quote_sql_identifier(rank_dimension)}" + if wants_order_count_metric: + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), + ) + metric_expr = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" + if order_column + else "COUNT(*)" + ) + metric_alias = "OrderCount" + else: + metric_expr = ( + f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + if measure + else "COUNT(*)" + ) + metric_alias = f"Total{measure}" if measure else "RecordCount" + where_clause = self._append_not_null_filters( + ( + self._build_date_filter(table_name, date_column, query) + if date_column + else "" + ), + [partition_ref, rank_ref], + ) + return ( + "WITH grouped_results AS (" + f"SELECT {partition_ref} AS {self._quote_sql_identifier(partition_dimension)}, " + f"{rank_ref} AS {self._quote_sql_identifier(rank_dimension)}, " + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " + f"FROM {table_ref}" + f"{where_clause} " + f"GROUP BY {partition_ref}, {rank_ref}" + "), ranked_results AS (" + f"SELECT {self._quote_sql_identifier(partition_dimension)}, " + f"{self._quote_sql_identifier(rank_dimension)}, " + f"{self._quote_sql_identifier(metric_alias)}, " + f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(partition_dimension)} " + f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC) AS \"rank\" " + "FROM grouped_results" + ") " + f"SELECT {self._quote_sql_identifier(partition_dimension)}, " + f"{self._quote_sql_identifier(rank_dimension)}, " + f"{self._quote_sql_identifier(metric_alias)} " + "FROM ranked_results " + "WHERE \"rank\" = 1 " + f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC" + ) + if wants_trend and date_column: date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" if wants_order_count_metric: @@ -1091,7 +1190,7 @@ def _build_schema_grounded_analytics_sql( ] return ( f"SELECT {', '.join(select_parts)} FROM {table_ref}" - f"{self._build_date_filter(table_name, date_column, query)} " + f"{self._append_not_null_filters(self._build_date_filter(table_name, date_column, query), dimension_refs)} " f"GROUP BY {', '.join(group_parts)} " f"ORDER BY DATEPART(YEAR, {date_ref}), " f"DATEPART(MONTH, {date_ref})" @@ -1120,7 +1219,8 @@ def _build_schema_grounded_analytics_sql( ) return ( f"SELECT {top_clause}{', '.join(select_parts)} " - f"FROM {table_ref}{date_filter} " + f"FROM {table_ref}" + f"{self._append_not_null_filters(date_filter, dimension_refs)} " f"GROUP BY {', '.join(dimension_refs)} " f"ORDER BY {metric_expr} DESC" ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 86add879ae..2f43719ce3 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -228,6 +228,7 @@ def test_build_schema_grounded_sales_sql_for_top_markets(): 'FROM "dbo_tblSales" ' 'WHERE "dbo_tblSales"."OrdDate" >= \'2026-01-01 00:00:00\' ' 'AND "dbo_tblSales"."OrdDate" < \'2027-01-01 00:00:00\' ' + 'AND "dbo_tblSales"."Market" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' ) @@ -256,6 +257,7 @@ def test_build_schema_grounded_sales_sql_for_division_revenue_trend(): '"dbo_tblSales"."Division" AS "Division", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Division" IS NOT NULL ' 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), "dbo_tblSales"."Division" ' 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' @@ -286,6 +288,9 @@ def test_build_schema_grounded_sales_sql_for_orders_by_dimensions(): '"dbo_tblSales"."ProdType" AS "ProdType", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."Division" IS NOT NULL ' + 'AND "dbo_tblSales"."ProdType" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Division", ' '"dbo_tblSales"."ProdType" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' @@ -318,6 +323,10 @@ def test_build_schema_grounded_sales_sql_for_top_new_order_detail_rows(): '"dbo_tblSales"."Customer" AS "Customer", ' '"dbo_tblSales"."SalesValue" AS "SalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."BU" IS NOT NULL ' + 'AND "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."ProdName" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'ORDER BY "dbo_tblSales"."SalesValue" DESC' ) @@ -343,6 +352,8 @@ def test_build_schema_grounded_sales_sql_ignores_missing_metadata_entries(): '"dbo_tblSales"."Customer" AS "Customer", ' '"dbo_tblSales"."SalesValue" AS "SalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'ORDER BY "dbo_tblSales"."SalesValue" DESC' ) @@ -428,6 +439,7 @@ def test_build_schema_grounded_sales_sql_counts_new_orders_by_customer_over_time '"dbo_XStageNewOrders"."CustName" AS "CustName", ' 'COUNT(DISTINCT "dbo_XStageNewOrders"."OrdNo") AS "OrderCount" ' 'FROM "dbo_XStageNewOrders" ' + 'WHERE "dbo_XStageNewOrders"."CustName" IS NOT NULL ' 'GROUP BY DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate"), ' 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate"), ' '"dbo_XStageNewOrders"."CustName" ' @@ -457,11 +469,48 @@ def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): '"dbo_tblSales"."Customer" AS "Customer", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."ProdName" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."ProdName", "dbo_tblSales"."Customer" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' ) +def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Which Customers have the highest New Orders in each Market?", + [ + """ + CREATE TABLE dbo_tblSales ( + Market VARCHAR, + Customer VARCHAR, + OrdNo VARCHAR, + SalesValue DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'WITH grouped_results AS (SELECT "dbo_tblSales"."Market" AS "Market", ' + '"dbo_tblSales"."Customer" AS "Customer", ' + 'COUNT(DISTINCT "dbo_tblSales"."OrdNo") AS "OrderCount" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' + 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Customer"), ' + 'ranked_results AS (SELECT "Market", "Customer", "OrderCount", ' + 'ROW_NUMBER() OVER (PARTITION BY "Market" ' + 'ORDER BY "OrderCount" DESC) AS "rank" ' + 'FROM grouped_results) ' + 'SELECT "Market", "Customer", "OrderCount" ' + 'FROM ranked_results ' + 'WHERE "rank" = 1 ' + 'ORDER BY "OrderCount" DESC' + ) + + def test_build_schema_grounded_sales_sql_for_product_type_contribution(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 34303ecfda4e253b13f44baa2b787f9fdedafe5f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 17:35:00 +0530 Subject: [PATCH 0288/1087] Use deterministic chart fallback and reliable pinning --- .../src/pipelines/generation/utils/chart.py | 11 ++++ wren-ai-service/src/web/v1/services/chart.py | 15 +++++ .../generation/test_chart_generation_utils.py | 45 +++++++------ .../server/resolvers/dashboardResolver.ts | 66 ++++++++++++++----- .../pages/home/promptThread/ChartAnswer.tsx | 18 ++++- 5 files changed, 117 insertions(+), 38 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index d111565788..faf97df01f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -169,6 +169,17 @@ def _select_dimension_columns( ordered += relevant_temporal + [ column for column in temporal if column not in relevant_temporal ] + if len(ordered) > 1: + normalized_query = str(query or "").lower() + for column in list(ordered): + tokens = _identifier_tokens(column) + if any( + re.search(rf"\b(?:each|per)\s+{re.escape(token)}s?\b", normalized_query) + for token in tokens + ): + ordered.remove(column) + ordered.insert(1, column) + break return ordered diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index ed3f47f1bc..3392cad921 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,6 +6,7 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest @@ -139,6 +140,20 @@ async def chart( trace_id=trace_id, ) + deterministic_chart_result = build_fallback_chart_result( + chart_request.query, + sql_data, + chart_request.remove_data_from_chart_schema, + ) + if deterministic_chart_result.get("chart_schema"): + self._chart_results[query_id] = ChartResultResponse( + status="finished", + response=ChartResult(**deterministic_chart_result), + trace_id=trace_id, + ) + results["chart_result"] = deterministic_chart_result + return results + chart_generation_result = await self._pipelines["chart_generation"].run( query=chart_request.query, sql=chart_request.sql, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index 495716e4d6..a35a2b538c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -35,6 +35,26 @@ def test_fallback_chart_requires_a_real_quantitative_measure(): assert result == {"chart_schema": {}, "reasoning": "", "chart_type": ""} +def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): + result = build_fallback_chart_result( + "Which Customers have the highest New Orders in each Market?", + { + "columns": [ + {"name": "Market"}, + {"name": "Customer"}, + {"name": "OrderCount"}, + ], + "data": [["North", "Acme", 10], ["South", "Globex", 8]], + }, + ) + + assert result["chart_type"] == "grouped_bar" + assert result["chart_schema"]["encoding"]["x"]["field"] == "Customer" + assert result["chart_schema"]["encoding"]["y"]["field"] == "OrderCount" + assert result["chart_schema"]["encoding"]["color"]["field"] == "Market" + assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Market" + + def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): assert not _is_schema_compatible_with_sample_data( { @@ -49,28 +69,10 @@ def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): @pytest.mark.asyncio -async def test_chart_service_does_not_short_circuit_to_generic_fallback(): +async def test_chart_service_returns_deterministic_chart_without_llm_wait(): class FakeChartGenerationPipeline: async def run(self, **kwargs): - assert kwargs["data"]["columns"] == [ - {"name": "Market"}, - {"name": "Revenue"}, - ] - return { - "post_process": { - "results": { - "chart_schema": { - "mark": {"type": "bar"}, - "encoding": { - "x": {"field": "Market", "type": "nominal"}, - "y": {"field": "Revenue", "type": "quantitative"}, - }, - }, - "reasoning": "Generated from the executed SQL result.", - "chart_type": "bar", - } - } - } + raise AssertionError("chart LLM pipeline should not be called") service = ChartService({"chart_generation": FakeChartGenerationPipeline()}) request = ChartRequest( @@ -86,5 +88,6 @@ async def run(self, **kwargs): result = await service.chart(request) assert result["chart_result"]["reasoning"] == ( - "Generated from the executed SQL result." + "Generated from the SQL result columns and requested chart type." ) + assert result["chart_result"]["chart_type"] == "bar" diff --git a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts index 8411b5409c..addfa8cb8c 100644 --- a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts @@ -77,14 +77,15 @@ export class DashboardResolver { args: { data: { itemType: DashboardItemType; responseId: number } }, ctx: IContext, ): Promise { - const { responseId, itemType } = args.data; + const { responseId } = args.data; + const itemType = this.normalizeDashboardItemType(args.data.itemType); const dashboard = await ctx.dashboardService.getCurrentDashboard(); const response = await ctx.askingService.getResponse(responseId); if (!response) { throw new Error(`Thread response not found. responseId: ${responseId}`); } - if (!Object.keys(ChartType).includes(itemType)) { + if (!itemType) { throw new Error(`Chart type not supported. responseId: ${responseId}`); } if (!response.chartDetail?.chartSchema) { @@ -93,24 +94,59 @@ export class DashboardResolver { ); } - // query with cache enabled - const project = await ctx.projectService.getCurrentProject(); - const deployment = await ctx.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; - await ctx.queryService.preview(response.sql, { - project, - manifest: mdl, - limit: DEFAULT_PREVIEW_LIMIT, - cacheEnabled: true, - refresh: true, - }); - - return await ctx.dashboardService.createDashboardItem({ + const dashboardItem = await ctx.dashboardService.createDashboardItem({ dashboardId: dashboard.id, type: itemType, sql: response.sql, chartSchema: response.chartDetail?.chartSchema, }); + + // Warm dashboard cache after persisting the item. Cache warm-up failures should + // not prevent a valid chart from being pinned. + const project = await ctx.projectService.getCurrentProject(); + const deployment = await ctx.deployService.getLastDeployment(project.id); + const mdl = deployment.manifest; + try { + await ctx.queryService.preview(response.sql, { + project, + manifest: mdl, + limit: DEFAULT_PREVIEW_LIMIT, + cacheEnabled: true, + refresh: true, + }); + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : String(error); + logger.warn( + `Dashboard item ${dashboardItem.id} was pinned but cache warm-up failed: ${errorMessage}`, + ); + } + + return dashboardItem; + } + + private normalizeDashboardItemType( + itemType: DashboardItemType | ChartType | string, + ): DashboardItemType | null { + const rawValue = String(itemType || ''); + const normalizedKey = rawValue.toUpperCase() as keyof typeof DashboardItemType; + const normalizedValue = DashboardItemType[normalizedKey]; + if (normalizedValue) { + return normalizedValue; + } + + const chartTypeValue = Object.values(ChartType).find( + (value) => value === rawValue, + ); + if (!chartTypeValue) { + return null; + } + + return ( + DashboardItemType[ + chartTypeValue.toUpperCase() as keyof typeof DashboardItemType + ] || null + ); } public async updateDashboardItem( diff --git a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx index 9ec2b9c894..6a4dbb5683 100644 --- a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx @@ -77,6 +77,15 @@ const getDynamicProperties = (chartType: ChartType) => { return propertiesMap[chartType] || BasicProperties; }; +const chartTypeToDashboardItemType = ( + chartType: ChartType, +): DashboardItemType | null => { + const normalized = String( + chartType || '', + ).toUpperCase() as keyof typeof DashboardItemType; + return DashboardItemType[normalized] || null; +}; + export default function ChartAnswer(props: AnswerResultProps) { const { onGenerateChartAnswer, onAdjustChartAnswer } = usePromptThreadStore(); const { threadResponse } = props; @@ -181,6 +190,12 @@ export default function ChartAnswer(props: AnswerResultProps) { }; const onPin = () => { + const dashboardItemType = chartTypeToDashboardItemType(chartType as ChartType); + if (!dashboardItemType) { + message.error('Chart type is not supported for dashboard pinning.'); + return; + } + Modal.confirm({ title: 'Are you sure you want to pin this chart to the dashboard?', okText: 'Save', @@ -188,8 +203,7 @@ export default function ChartAnswer(props: AnswerResultProps) { await createDashboardItem({ variables: { data: { - // DashboardItemType is compatible with ChartType - itemType: chartType as unknown as DashboardItemType, + itemType: dashboardItemType, responseId: threadResponse.id, }, }, From b66fe3a80e0b447e8b87c9c6079c45d0ea096014 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 17:51:39 +0530 Subject: [PATCH 0289/1087] Keep sparse grouped dimensions in SQL results --- wren-ai-service/src/web/v1/services/ask.py | 12 +++--- .../pytest/services/test_ask_sales_sql.py | 42 +++++++++++-------- 2 files changed, 30 insertions(+), 24 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 497447beb3..e164a0a219 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -759,13 +759,11 @@ def _build_date_filter(self, table_name: str, date_column: str, query: str) -> s def _append_not_null_filters( self, where_clause: str, column_refs: list[str] ) -> str: - conditions = [f"{column_ref} IS NOT NULL" for column_ref in column_refs] - if not conditions: - return where_clause - - if where_clause.strip(): - return f"{where_clause.rstrip()} AND {' AND '.join(conditions)} " - return f" WHERE {' AND '.join(conditions)} " + # Grouping dimensions can be sparsely populated in deployed customer + # models. Adding implicit IS NOT NULL filters can turn valid aggregate + # questions into empty result sets, so keep only caller-provided filters + # such as date ranges. + return where_clause def _select_best_analytics_table( self, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2f43719ce3..ed489de6c6 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -228,13 +228,37 @@ def test_build_schema_grounded_sales_sql_for_top_markets(): 'FROM "dbo_tblSales" ' 'WHERE "dbo_tblSales"."OrdDate" >= \'2026-01-01 00:00:00\' ' 'AND "dbo_tblSales"."OrdDate" < \'2027-01-01 00:00:00\' ' - 'AND "dbo_tblSales"."Market" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' ) assert "dbo_tblStageNewOrders" not in sql +def test_build_schema_grounded_sales_sql_keeps_sparse_grouped_dimensions(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show new orders by Market.", + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Market VARCHAR, + OrderAmount DOUBLE, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tnoStageNewOrders"."Market" AS "Market", ' + 'SUM("dbo_tnoStageNewOrders"."OrderAmount") AS "TotalOrderAmount" ' + 'FROM "dbo_tnoStageNewOrders" ' + 'GROUP BY "dbo_tnoStageNewOrders"."Market" ' + 'ORDER BY SUM("dbo_tnoStageNewOrders"."OrderAmount") DESC' + ) + assert '"dbo_tnoStageNewOrders"."Market" IS NOT NULL' not in sql + + def test_build_schema_grounded_sales_sql_for_division_revenue_trend(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( @@ -257,7 +281,6 @@ def test_build_schema_grounded_sales_sql_for_division_revenue_trend(): '"dbo_tblSales"."Division" AS "Division", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Division" IS NOT NULL ' 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), "dbo_tblSales"."Division" ' 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' @@ -288,9 +311,6 @@ def test_build_schema_grounded_sales_sql_for_orders_by_dimensions(): '"dbo_tblSales"."ProdType" AS "ProdType", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."Division" IS NOT NULL ' - 'AND "dbo_tblSales"."ProdType" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Division", ' '"dbo_tblSales"."ProdType" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' @@ -323,10 +343,6 @@ def test_build_schema_grounded_sales_sql_for_top_new_order_detail_rows(): '"dbo_tblSales"."Customer" AS "Customer", ' '"dbo_tblSales"."SalesValue" AS "SalesValue" ' 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."BU" IS NOT NULL ' - 'AND "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."ProdName" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'ORDER BY "dbo_tblSales"."SalesValue" DESC' ) @@ -352,8 +368,6 @@ def test_build_schema_grounded_sales_sql_ignores_missing_metadata_entries(): '"dbo_tblSales"."Customer" AS "Customer", ' '"dbo_tblSales"."SalesValue" AS "SalesValue" ' 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'ORDER BY "dbo_tblSales"."SalesValue" DESC' ) @@ -439,7 +453,6 @@ def test_build_schema_grounded_sales_sql_counts_new_orders_by_customer_over_time '"dbo_XStageNewOrders"."CustName" AS "CustName", ' 'COUNT(DISTINCT "dbo_XStageNewOrders"."OrdNo") AS "OrderCount" ' 'FROM "dbo_XStageNewOrders" ' - 'WHERE "dbo_XStageNewOrders"."CustName" IS NOT NULL ' 'GROUP BY DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate"), ' 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate"), ' '"dbo_XStageNewOrders"."CustName" ' @@ -469,8 +482,6 @@ def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): '"dbo_tblSales"."Customer" AS "Customer", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."ProdName" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."ProdName", "dbo_tblSales"."Customer" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' ) @@ -497,8 +508,6 @@ def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): '"dbo_tblSales"."Customer" AS "Customer", ' 'COUNT(DISTINCT "dbo_tblSales"."OrdNo") AS "OrderCount" ' 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Customer"), ' 'ranked_results AS (SELECT "Market", "Customer", "OrderCount", ' 'ROW_NUMBER() OVER (PARTITION BY "Market" ' @@ -553,7 +562,6 @@ def test_build_schema_grounded_sql_counts_categorical_status_values(): 'SELECT "dbo_ytblRefund"."Refund_Status" AS "Refund_Status", ' 'COUNT(*) AS "RecordCount" ' 'FROM "dbo_ytblRefund" ' - 'WHERE "dbo_ytblRefund"."Refund_Status" IS NOT NULL ' 'GROUP BY "dbo_ytblRefund"."Refund_Status" ' 'ORDER BY COUNT(*) DESC' ) From 39356428d83f7d9b1bb081459fa4519595a39ad3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 18:18:45 +0530 Subject: [PATCH 0290/1087] Ground asks in active metadata and fix dashboard pinning --- wren-ai-service/src/web/v1/services/ask.py | 28 ------------- .../repositories/dashboardItemRepository.ts | 39 ++++++++++++++++++- .../server/services/dashboardService.ts | 11 +++++- .../services/tests/dashboardService.test.ts | 35 +++++++++++++++++ 4 files changed, 82 insertions(+), 31 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e164a0a219..69937315a2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2923,34 +2923,6 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results - if direct_orders_sales_sql := self._build_direct_orders_sales_sql( - user_query - ): - api_results = [ - AskResult( - **{ - "sql": direct_orders_sales_sql, - "type": "llm", - } - ) - ] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=["dbo_tblSales"], - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - logger.info( - "Using direct Orders/Sales SQL for query_id %s", - query_id, - ) - return results - explicit_table_names = self._extract_explicit_table_names_from_query( user_query ) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index bb9f695ecd..f507f52976 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -18,6 +18,7 @@ export enum DashboardItemType { BAR = 'BAR', GROUPED_BAR = 'GROUPED_BAR', LINE = 'LINE', + MULTI_LINE = 'MULTI_LINE', PIE = 'PIE', STACKED_BAR = 'STACKED_BAR', // other types @@ -96,15 +97,16 @@ export class DashboardItemRepository const transformData = mapValues(camelCaseData, (value, key) => { if (this.jsonbColumns.includes(key)) { if (typeof value === 'string') { - return value ? JSON.parse(value) : value; + return this.parseJsonColumn(value, key); } else { - return value; + return isPlainObject(value) ? value : this.defaultJsonColumnValue(key); } } return value; }); return { ...transformData, + type: this.normalizeItemType(transformData.type), displayName: transformData.displayName || transformData.title, } as DashboardItem; }; @@ -144,6 +146,9 @@ export class DashboardItemRepository this.hasColumn('updated_at', executer), ]); const normalizedData: Partial = { ...data }; + if (normalizedData.type) { + normalizedData.type = this.normalizeItemType(normalizedData.type); + } const displayName = typeof data.displayName === 'string' ? data.displayName.trim() : ''; const chartTitle = @@ -176,6 +181,36 @@ export class DashboardItemRepository return normalizedData; } + private parseJsonColumn(value: string, key: string) { + if (!value) { + return this.defaultJsonColumnValue(key); + } + try { + const parsed = JSON.parse(value); + return isPlainObject(parsed) ? parsed : this.defaultJsonColumnValue(key); + } catch { + return this.defaultJsonColumnValue(key); + } + } + + private defaultJsonColumnValue(key: string) { + if (key === 'layout') { + return { x: 0, y: 0, w: 3, h: 2 }; + } + if (key === 'detail') { + return { sql: '', chartSchema: null }; + } + return {}; + } + + private normalizeItemType(type: unknown): DashboardItemType { + const normalized = String(type || '').toUpperCase(); + return ( + DashboardItemType[normalized as keyof typeof DashboardItemType] || + DashboardItemType.BAR + ); + } + private async hasColumn(column: string, executer: Knex | Knex.Transaction) { if (this.columnCache.has(column)) { return this.columnCache.get(column); diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 59818eef9f..4f48ede13a 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -143,7 +143,10 @@ export class DashboardService implements IDashboardService { const dashboard = await this.dashboardRepository.findOneBy({ projectId: project.id, }); - return { ...dashboard }; + if (dashboard) { + return dashboard; + } + return await this.initDashboard(); } public async getDashboardItem( @@ -169,6 +172,12 @@ export class DashboardService implements IDashboardService { public async createDashboardItem( input: CreateDashboardItemInput, ): Promise { + if (!input.dashboardId) { + throw new Error('Dashboard id is required.'); + } + if (!input.sql) { + throw new Error('Dashboard item SQL is required.'); + } const layout = await this.calculateNewLayout(input.dashboardId); return await this.dashboardItemRepository.createOne({ dashboardId: input.dashboardId, diff --git a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts index 55814d5610..e47a967fd4 100644 --- a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts @@ -85,6 +85,41 @@ describe('DashboardService', () => { }); }); + describe('getCurrentDashboard', () => { + it('should return the existing dashboard for the current project', async () => { + const project = { id: 7 }; + const dashboard = { id: 11, projectId: 7, name: 'Dashboard' }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + + await expect(dashboardService.getCurrentDashboard()).resolves.toBe( + dashboard, + ); + expect(mockDashboardRepository.findOneBy).toHaveBeenCalledWith({ + projectId: 7, + }); + expect(mockDashboardRepository.createOne).not.toHaveBeenCalled(); + }); + + it('should initialize a dashboard when the current project has none', async () => { + const project = { id: 7 }; + const dashboard = { id: 12, projectId: 7, name: 'Dashboard' }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null); + mockDashboardRepository.createOne.mockResolvedValue(dashboard); + + await expect(dashboardService.getCurrentDashboard()).resolves.toBe( + dashboard, + ); + expect(mockDashboardRepository.createOne).toHaveBeenCalledWith({ + name: 'Dashboard', + projectId: 7, + }); + }); + }); + describe('generateCronExpression', () => { it('should generate correct cron expression for daily schedule', () => { const schedule = { From 957d527d72305ea5ce7f8873b82fd8666cb5414f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 18:55:24 +0530 Subject: [PATCH 0291/1087] Revert "Ground asks in active metadata and fix dashboard pinning" This reverts commit 39356428d83f7d9b1bb081459fa4519595a39ad3. --- wren-ai-service/src/web/v1/services/ask.py | 28 +++++++++++++ .../repositories/dashboardItemRepository.ts | 39 +------------------ .../server/services/dashboardService.ts | 11 +----- .../services/tests/dashboardService.test.ts | 35 ----------------- 4 files changed, 31 insertions(+), 82 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 69937315a2..e164a0a219 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2923,6 +2923,34 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results + if direct_orders_sales_sql := self._build_direct_orders_sales_sql( + user_query + ): + api_results = [ + AskResult( + **{ + "sql": direct_orders_sales_sql, + "type": "llm", + } + ) + ] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=["dbo_tblSales"], + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + logger.info( + "Using direct Orders/Sales SQL for query_id %s", + query_id, + ) + return results + explicit_table_names = self._extract_explicit_table_names_from_query( user_query ) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index f507f52976..bb9f695ecd 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -18,7 +18,6 @@ export enum DashboardItemType { BAR = 'BAR', GROUPED_BAR = 'GROUPED_BAR', LINE = 'LINE', - MULTI_LINE = 'MULTI_LINE', PIE = 'PIE', STACKED_BAR = 'STACKED_BAR', // other types @@ -97,16 +96,15 @@ export class DashboardItemRepository const transformData = mapValues(camelCaseData, (value, key) => { if (this.jsonbColumns.includes(key)) { if (typeof value === 'string') { - return this.parseJsonColumn(value, key); + return value ? JSON.parse(value) : value; } else { - return isPlainObject(value) ? value : this.defaultJsonColumnValue(key); + return value; } } return value; }); return { ...transformData, - type: this.normalizeItemType(transformData.type), displayName: transformData.displayName || transformData.title, } as DashboardItem; }; @@ -146,9 +144,6 @@ export class DashboardItemRepository this.hasColumn('updated_at', executer), ]); const normalizedData: Partial = { ...data }; - if (normalizedData.type) { - normalizedData.type = this.normalizeItemType(normalizedData.type); - } const displayName = typeof data.displayName === 'string' ? data.displayName.trim() : ''; const chartTitle = @@ -181,36 +176,6 @@ export class DashboardItemRepository return normalizedData; } - private parseJsonColumn(value: string, key: string) { - if (!value) { - return this.defaultJsonColumnValue(key); - } - try { - const parsed = JSON.parse(value); - return isPlainObject(parsed) ? parsed : this.defaultJsonColumnValue(key); - } catch { - return this.defaultJsonColumnValue(key); - } - } - - private defaultJsonColumnValue(key: string) { - if (key === 'layout') { - return { x: 0, y: 0, w: 3, h: 2 }; - } - if (key === 'detail') { - return { sql: '', chartSchema: null }; - } - return {}; - } - - private normalizeItemType(type: unknown): DashboardItemType { - const normalized = String(type || '').toUpperCase(); - return ( - DashboardItemType[normalized as keyof typeof DashboardItemType] || - DashboardItemType.BAR - ); - } - private async hasColumn(column: string, executer: Knex | Knex.Transaction) { if (this.columnCache.has(column)) { return this.columnCache.get(column); diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 4f48ede13a..59818eef9f 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -143,10 +143,7 @@ export class DashboardService implements IDashboardService { const dashboard = await this.dashboardRepository.findOneBy({ projectId: project.id, }); - if (dashboard) { - return dashboard; - } - return await this.initDashboard(); + return { ...dashboard }; } public async getDashboardItem( @@ -172,12 +169,6 @@ export class DashboardService implements IDashboardService { public async createDashboardItem( input: CreateDashboardItemInput, ): Promise { - if (!input.dashboardId) { - throw new Error('Dashboard id is required.'); - } - if (!input.sql) { - throw new Error('Dashboard item SQL is required.'); - } const layout = await this.calculateNewLayout(input.dashboardId); return await this.dashboardItemRepository.createOne({ dashboardId: input.dashboardId, diff --git a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts index e47a967fd4..55814d5610 100644 --- a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts @@ -85,41 +85,6 @@ describe('DashboardService', () => { }); }); - describe('getCurrentDashboard', () => { - it('should return the existing dashboard for the current project', async () => { - const project = { id: 7 }; - const dashboard = { id: 11, projectId: 7, name: 'Dashboard' }; - mockProjectService.getCurrentProject.mockResolvedValue(project); - mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); - - await expect(dashboardService.getCurrentDashboard()).resolves.toBe( - dashboard, - ); - expect(mockDashboardRepository.findOneBy).toHaveBeenCalledWith({ - projectId: 7, - }); - expect(mockDashboardRepository.createOne).not.toHaveBeenCalled(); - }); - - it('should initialize a dashboard when the current project has none', async () => { - const project = { id: 7 }; - const dashboard = { id: 12, projectId: 7, name: 'Dashboard' }; - mockProjectService.getCurrentProject.mockResolvedValue(project); - mockDashboardRepository.findOneBy - .mockResolvedValueOnce(null) - .mockResolvedValueOnce(null); - mockDashboardRepository.createOne.mockResolvedValue(dashboard); - - await expect(dashboardService.getCurrentDashboard()).resolves.toBe( - dashboard, - ); - expect(mockDashboardRepository.createOne).toHaveBeenCalledWith({ - name: 'Dashboard', - projectId: 7, - }); - }); - }); - describe('generateCronExpression', () => { it('should generate correct cron expression for daily schedule', () => { const schedule = { From 24fc6f6f16bc4293a6e852fa830423cb4c9d689f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 18:55:25 +0530 Subject: [PATCH 0292/1087] Revert "Keep sparse grouped dimensions in SQL results" This reverts commit b66fe3a80e0b447e8b87c9c6079c45d0ea096014. --- wren-ai-service/src/web/v1/services/ask.py | 12 +++--- .../pytest/services/test_ask_sales_sql.py | 42 ++++++++----------- 2 files changed, 24 insertions(+), 30 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e164a0a219..497447beb3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -759,11 +759,13 @@ def _build_date_filter(self, table_name: str, date_column: str, query: str) -> s def _append_not_null_filters( self, where_clause: str, column_refs: list[str] ) -> str: - # Grouping dimensions can be sparsely populated in deployed customer - # models. Adding implicit IS NOT NULL filters can turn valid aggregate - # questions into empty result sets, so keep only caller-provided filters - # such as date ranges. - return where_clause + conditions = [f"{column_ref} IS NOT NULL" for column_ref in column_refs] + if not conditions: + return where_clause + + if where_clause.strip(): + return f"{where_clause.rstrip()} AND {' AND '.join(conditions)} " + return f" WHERE {' AND '.join(conditions)} " def _select_best_analytics_table( self, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index ed489de6c6..2f43719ce3 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -228,37 +228,13 @@ def test_build_schema_grounded_sales_sql_for_top_markets(): 'FROM "dbo_tblSales" ' 'WHERE "dbo_tblSales"."OrdDate" >= \'2026-01-01 00:00:00\' ' 'AND "dbo_tblSales"."OrdDate" < \'2027-01-01 00:00:00\' ' + 'AND "dbo_tblSales"."Market" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' ) assert "dbo_tblStageNewOrders" not in sql -def test_build_schema_grounded_sales_sql_keeps_sparse_grouped_dimensions(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show new orders by Market.", - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Market VARCHAR, - OrderAmount DOUBLE, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tnoStageNewOrders"."Market" AS "Market", ' - 'SUM("dbo_tnoStageNewOrders"."OrderAmount") AS "TotalOrderAmount" ' - 'FROM "dbo_tnoStageNewOrders" ' - 'GROUP BY "dbo_tnoStageNewOrders"."Market" ' - 'ORDER BY SUM("dbo_tnoStageNewOrders"."OrderAmount") DESC' - ) - assert '"dbo_tnoStageNewOrders"."Market" IS NOT NULL' not in sql - - def test_build_schema_grounded_sales_sql_for_division_revenue_trend(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( @@ -281,6 +257,7 @@ def test_build_schema_grounded_sales_sql_for_division_revenue_trend(): '"dbo_tblSales"."Division" AS "Division", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Division" IS NOT NULL ' 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), "dbo_tblSales"."Division" ' 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' @@ -311,6 +288,9 @@ def test_build_schema_grounded_sales_sql_for_orders_by_dimensions(): '"dbo_tblSales"."ProdType" AS "ProdType", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."Division" IS NOT NULL ' + 'AND "dbo_tblSales"."ProdType" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Division", ' '"dbo_tblSales"."ProdType" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' @@ -343,6 +323,10 @@ def test_build_schema_grounded_sales_sql_for_top_new_order_detail_rows(): '"dbo_tblSales"."Customer" AS "Customer", ' '"dbo_tblSales"."SalesValue" AS "SalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."BU" IS NOT NULL ' + 'AND "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."ProdName" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'ORDER BY "dbo_tblSales"."SalesValue" DESC' ) @@ -368,6 +352,8 @@ def test_build_schema_grounded_sales_sql_ignores_missing_metadata_entries(): '"dbo_tblSales"."Customer" AS "Customer", ' '"dbo_tblSales"."SalesValue" AS "SalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'ORDER BY "dbo_tblSales"."SalesValue" DESC' ) @@ -453,6 +439,7 @@ def test_build_schema_grounded_sales_sql_counts_new_orders_by_customer_over_time '"dbo_XStageNewOrders"."CustName" AS "CustName", ' 'COUNT(DISTINCT "dbo_XStageNewOrders"."OrdNo") AS "OrderCount" ' 'FROM "dbo_XStageNewOrders" ' + 'WHERE "dbo_XStageNewOrders"."CustName" IS NOT NULL ' 'GROUP BY DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate"), ' 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate"), ' '"dbo_XStageNewOrders"."CustName" ' @@ -482,6 +469,8 @@ def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): '"dbo_tblSales"."Customer" AS "Customer", ' 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."ProdName" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."ProdName", "dbo_tblSales"."Customer" ' 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' ) @@ -508,6 +497,8 @@ def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): '"dbo_tblSales"."Customer" AS "Customer", ' 'COUNT(DISTINCT "dbo_tblSales"."OrdNo") AS "OrderCount" ' 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."Customer" IS NOT NULL ' 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Customer"), ' 'ranked_results AS (SELECT "Market", "Customer", "OrderCount", ' 'ROW_NUMBER() OVER (PARTITION BY "Market" ' @@ -562,6 +553,7 @@ def test_build_schema_grounded_sql_counts_categorical_status_values(): 'SELECT "dbo_ytblRefund"."Refund_Status" AS "Refund_Status", ' 'COUNT(*) AS "RecordCount" ' 'FROM "dbo_ytblRefund" ' + 'WHERE "dbo_ytblRefund"."Refund_Status" IS NOT NULL ' 'GROUP BY "dbo_ytblRefund"."Refund_Status" ' 'ORDER BY COUNT(*) DESC' ) From 4105cc2b7e9e6b444ded50bad79c12bbd15c54db Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 19:03:28 +0530 Subject: [PATCH 0293/1087] Prevent stuck answer and chart background jobs --- .../src/apollo/server/backgrounds/chart.ts | 231 ++++++++++-------- .../textBasedAnswerBackgroundTracker.ts | 195 ++++++++------- 2 files changed, 236 insertions(+), 190 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index 35c9c22b79..d1211c8651 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -72,61 +72,76 @@ export class ChartBackgroundTracker { // mark the job as running this.runningJobs.add(threadResponse.id); - // get the chart detail - const chartDetail = threadResponse.chartDetail; - - // get the latest result from AI service - const result = await this.wrenAIAdaptor.getChartResult( - chartDetail.queryId, - ); - - const statusChanged = chartDetail.status !== result.status; - this.scheduleNextPoll( - threadResponse.id, - result.status, - statusChanged, - ); - - if (isFinalized(result.status) && !statusChanged) { - this.finalizeTask(threadResponse, result); - this.runningJobs.delete(threadResponse.id); - return; - } + try { + // get the chart detail + const chartDetail = threadResponse.chartDetail; - // check if status change - if (!statusChanged) { - // mark the job as finished + // get the latest result from AI service + const result = await this.wrenAIAdaptor.getChartResult( + chartDetail.queryId, + ); + + const statusChanged = chartDetail.status !== result.status; + this.scheduleNextPoll( + threadResponse.id, + result.status, + statusChanged, + ); + + if (isFinalized(result.status) && !statusChanged) { + this.finalizeTask(threadResponse, result); + return; + } + + // check if status change + if (!statusChanged) { + // mark the job as finished + logger.debug( + `Job ${threadResponse.id} chart status not changed, finished`, + ); + return; + } + + // update database + const updatedChartDetail = { + queryId: chartDetail.queryId, + status: result?.status, + error: result?.error, + description: result?.response?.reasoning, + chartType: result?.response?.chartType?.toUpperCase() || null, + chartSchema: result?.response?.chartSchema, + }; logger.debug( - `Job ${threadResponse.id} chart status not changed, finished`, + `Job ${threadResponse.id} chart status changed, updating`, ); + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: updatedChartDetail, + }); + threadResponse.chartDetail = updatedChartDetail; + + // remove the task from tracker if it is finalized + if (isFinalized(result.status)) { + this.finalizeTask(threadResponse, result); + } + } catch (error) { + logger.error(`Chart job ${threadResponse.id} failed: ${error}`); + const failedChartDetail = { + ...threadResponse.chartDetail, + status: ChartStatus.FAILED, + error: error?.extensions || error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: failedChartDetail, + }); + threadResponse.chartDetail = failedChartDetail; + this.finalizeTask(threadResponse, { + status: ChartStatus.FAILED, + error, + }); + throw error; + } finally { this.runningJobs.delete(threadResponse.id); - return; } - - // update database - const updatedChartDetail = { - queryId: chartDetail.queryId, - status: result?.status, - error: result?.error, - description: result?.response?.reasoning, - chartType: result?.response?.chartType?.toUpperCase() || null, - chartSchema: result?.response?.chartSchema, - }; - logger.debug( - `Job ${threadResponse.id} chart status changed, updating`, - ); - await this.threadResponseRepository.updateOne(threadResponse.id, { - chartDetail: updatedChartDetail, - }); - threadResponse.chartDetail = updatedChartDetail; - - // remove the task from tracker if it is finalized - if (isFinalized(result.status)) { - this.finalizeTask(threadResponse, result); - } - - // mark the job as finished - this.runningJobs.delete(threadResponse.id); }, ); @@ -252,62 +267,80 @@ export class ChartAdjustmentBackgroundTracker { // mark the job as running this.runningJobs.add(threadResponse.id); - // get the chart detail - const chartDetail = threadResponse.chartDetail; + try { + // get the chart detail + const chartDetail = threadResponse.chartDetail; - // get the latest result from AI service - const result = await this.wrenAIAdaptor.getChartAdjustmentResult( - chartDetail.queryId, - ); - - const statusChanged = chartDetail.status !== result.status; - this.scheduleNextPoll( - threadResponse.id, - result.status, - statusChanged, - ); + // get the latest result from AI service + const result = await this.wrenAIAdaptor.getChartAdjustmentResult( + chartDetail.queryId, + ); - if (isFinalized(result.status) && !statusChanged) { - this.finalizeTask(threadResponse, result); - this.runningJobs.delete(threadResponse.id); - return; - } + const statusChanged = chartDetail.status !== result.status; + this.scheduleNextPoll( + threadResponse.id, + result.status, + statusChanged, + ); - // check if status change - if (!statusChanged) { - // mark the job as finished + if (isFinalized(result.status) && !statusChanged) { + this.finalizeTask(threadResponse, result); + return; + } + + // check if status change + if (!statusChanged) { + // mark the job as finished + logger.debug( + `Job ${threadResponse.id} chart status not changed, finished`, + ); + return; + } + + // update database + const updatedChartDetail = { + queryId: chartDetail.queryId, + status: result?.status, + error: result?.error, + description: result?.response?.reasoning, + chartType: result?.response?.chartType?.toUpperCase() || null, + chartSchema: result?.response?.chartSchema, + adjustment: true, + }; logger.debug( - `Job ${threadResponse.id} chart status not changed, finished`, + `Job ${threadResponse.id} chart status changed, updating`, ); + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: updatedChartDetail, + }); + threadResponse.chartDetail = updatedChartDetail; + + // remove the task from tracker if it is finalized + if (isFinalized(result.status)) { + this.finalizeTask(threadResponse, result); + } + } catch (error) { + logger.error( + `Chart adjustment job ${threadResponse.id} failed: ${error}`, + ); + const failedChartDetail = { + ...threadResponse.chartDetail, + status: ChartStatus.FAILED, + error: error?.extensions || error, + adjustment: true, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: failedChartDetail, + }); + threadResponse.chartDetail = failedChartDetail; + this.finalizeTask(threadResponse, { + status: ChartStatus.FAILED, + error, + }); + throw error; + } finally { this.runningJobs.delete(threadResponse.id); - return; } - - // update database - const updatedChartDetail = { - queryId: chartDetail.queryId, - status: result?.status, - error: result?.error, - description: result?.response?.reasoning, - chartType: result?.response?.chartType?.toUpperCase() || null, - chartSchema: result?.response?.chartSchema, - adjustment: true, - }; - logger.debug( - `Job ${threadResponse.id} chart status changed, updating`, - ); - await this.threadResponseRepository.updateOne(threadResponse.id, { - chartDetail: updatedChartDetail, - }); - threadResponse.chartDetail = updatedChartDetail; - - // remove the task from tracker if it is finalized - if (isFinalized(result.status)) { - this.finalizeTask(threadResponse, result); - } - - // mark the job as finished - this.runningJobs.delete(threadResponse.id); }, ); diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index 6904573abc..f271039633 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -72,111 +72,124 @@ export class TextBasedAnswerBackgroundTracker { } this.runningJobs.add(threadResponse.id); - const answerDetail = threadResponse.answerDetail; - - if ( - !answerDetail.queryId && - answerDetail.status !== ThreadResponseAnswerStatus.FETCHING_DATA - ) { - const fetchingDetail = { - ...answerDetail, - status: ThreadResponseAnswerStatus.FETCHING_DATA, - }; - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: fetchingDetail, - }); - threadResponse.answerDetail = fetchingDetail; - - const thread = await this.threadRepository.findOneBy({ - id: threadResponse.threadId, - }); - if (!thread) { - throw new Error(`Thread ${threadResponse.threadId} not found`); - } - const project = await this.projectService.getProjectById( - thread.projectId, - ); - const deployment = await this.deployService.getLastDeployment( - project.id, - ); - const mdl = deployment.manifest; - let data: PreviewDataResponse; - try { - data = (await this.queryService.preview(threadResponse.sql, { - project, - manifest: mdl, - modelingOnly: false, - limit: ANSWER_PREVIEW_LIMIT, - })) as PreviewDataResponse; - } catch (error) { - logger.error(`Error when query sql data: ${error}`); - const failedDetail = { - ...threadResponse.answerDetail, - status: ThreadResponseAnswerStatus.FAILED, - error: error?.extensions || error, + try { + const answerDetail = threadResponse.answerDetail; + + if ( + !answerDetail.queryId && + answerDetail.status !== ThreadResponseAnswerStatus.FETCHING_DATA + ) { + const fetchingDetail = { + ...answerDetail, + status: ThreadResponseAnswerStatus.FETCHING_DATA, }; await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: failedDetail, + answerDetail: fetchingDetail, }); - threadResponse.answerDetail = failedDetail; - throw error; - } - - const response = await this.wrenAIAdaptor.createTextBasedAnswer({ - query: threadResponse.question, - sql: threadResponse.sql, - sqlData: data, - threadId: threadResponse.threadId.toString(), - configurations: { - language: WrenAILanguage[project.language] || WrenAILanguage.EN, - }, - }); + threadResponse.answerDetail = fetchingDetail; - const preprocessingDetail = { - ...threadResponse.answerDetail, - queryId: response.queryId, - status: ThreadResponseAnswerStatus.PREPROCESSING, - }; - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: preprocessingDetail, - }); - threadResponse.answerDetail = preprocessingDetail; - this.runningJobs.delete(threadResponse.id); - return; - } - - if ( - answerDetail.queryId && - answerDetail.status === ThreadResponseAnswerStatus.PREPROCESSING - ) { - const result: TextBasedAnswerResult = - await this.wrenAIAdaptor.getTextBasedAnswerResult( - answerDetail.queryId, + const thread = await this.threadRepository.findOneBy({ + id: threadResponse.threadId, + }); + if (!thread) { + throw new Error(`Thread ${threadResponse.threadId} not found`); + } + const project = await this.projectService.getProjectById( + thread.projectId, ); + const deployment = await this.deployService.getLastDeployment( + project.id, + ); + const mdl = deployment.manifest; + let data: PreviewDataResponse; + try { + data = (await this.queryService.preview(threadResponse.sql, { + project, + manifest: mdl, + modelingOnly: false, + limit: ANSWER_PREVIEW_LIMIT, + })) as PreviewDataResponse; + } catch (error) { + logger.error(`Error when query sql data: ${error}`); + const failedDetail = { + ...threadResponse.answerDetail, + status: ThreadResponseAnswerStatus.FAILED, + error: error?.extensions || error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: failedDetail, + }); + threadResponse.answerDetail = failedDetail; + delete this.tasks[threadResponse.id]; + throw error; + } + + const response = await this.wrenAIAdaptor.createTextBasedAnswer({ + query: threadResponse.question, + sql: threadResponse.sql, + sqlData: data, + threadId: threadResponse.threadId.toString(), + configurations: { + language: WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }); - if (result.status === TextBasedAnswerStatus.PREPROCESSING) { - this.runningJobs.delete(threadResponse.id); + const preprocessingDetail = { + ...threadResponse.answerDetail, + queryId: response.queryId, + status: ThreadResponseAnswerStatus.PREPROCESSING, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: preprocessingDetail, + }); + threadResponse.answerDetail = preprocessingDetail; return; } - const updatedAnswerDetail = { - queryId: answerDetail.queryId, - status: - result.status === TextBasedAnswerStatus.SUCCEEDED - ? ThreadResponseAnswerStatus.STREAMING - : ThreadResponseAnswerStatus.FAILED, - numRowsUsedInLLM: result.numRowsUsedInLLM, - error: result.error, + if ( + answerDetail.queryId && + answerDetail.status === ThreadResponseAnswerStatus.PREPROCESSING + ) { + const result: TextBasedAnswerResult = + await this.wrenAIAdaptor.getTextBasedAnswerResult( + answerDetail.queryId, + ); + + if (result.status === TextBasedAnswerStatus.PREPROCESSING) { + return; + } + + const updatedAnswerDetail = { + queryId: answerDetail.queryId, + status: + result.status === TextBasedAnswerStatus.SUCCEEDED + ? ThreadResponseAnswerStatus.STREAMING + : ThreadResponseAnswerStatus.FAILED, + numRowsUsedInLLM: result.numRowsUsedInLLM, + error: result.error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: updatedAnswerDetail, + }); + threadResponse.answerDetail = updatedAnswerDetail; + delete this.tasks[threadResponse.id]; + } + } catch (error) { + logger.error(`Answer job ${threadResponse.id} failed: ${error}`); + const failedDetail = { + ...threadResponse.answerDetail, + status: ThreadResponseAnswerStatus.FAILED, + error: error?.extensions || error, }; await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: updatedAnswerDetail, + answerDetail: failedDetail, }); - threadResponse.answerDetail = updatedAnswerDetail; + threadResponse.answerDetail = failedDetail; delete this.tasks[threadResponse.id]; + throw error; + } finally { + this.runningJobs.delete(threadResponse.id); } - - // Mark the job as finished - this.runningJobs.delete(threadResponse.id); }, ); From 76db0fd59552bf40c46f290b080b4ea252ce9498 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 19:06:39 +0530 Subject: [PATCH 0294/1087] Guard thread response cache updates --- wren-ui/src/pages/home/[id].tsx | 32 +++++++++++++++++++++++--------- 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index 9cc0a8176e..566fee4fc3 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -107,7 +107,13 @@ export default function HomeThread() { onError: (error) => console.error(error), onCompleted(next) { const nextResponse = next.createThreadResponse; + if (!nextResponse) { + return; + } updateThreadQuery((prev) => { + if (!prev?.thread?.responses) { + return prev; + } return { ...prev, thread: { @@ -133,15 +139,23 @@ export default function HomeThread() { nextFetchPolicy: 'network-only', onCompleted(next) { const nextResponse = next.threadResponse; - updateThreadQuery((prev) => ({ - ...prev, - thread: { - ...prev.thread, - responses: prev.thread.responses.map((response) => - response.id === nextResponse.id ? nextResponse : response, - ), - }, - })); + if (!nextResponse) { + return; + } + updateThreadQuery((prev) => { + if (!prev?.thread?.responses) { + return prev; + } + return { + ...prev, + thread: { + ...prev.thread, + responses: prev.thread.responses.map((response) => + response.id === nextResponse.id ? nextResponse : response, + ), + }, + }; + }); }, }); From 907b36ee87f6bd4557ec3b1cd02ae9ed3b0907f3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 30 Jun 2026 19:19:05 +0530 Subject: [PATCH 0295/1087] Return fast answers from SQL preview data --- .../textBasedAnswerBackgroundTracker.ts | 93 ++++++++++++++++++- .../apollo/server/services/askingService.ts | 2 +- 2 files changed, 93 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index f271039633..b913497f76 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -21,7 +21,79 @@ import { getLogger } from '@server/utils'; const logger = getLogger('TextBasedAnswerBackgroundTracker'); logger.level = 'debug'; -const ANSWER_PREVIEW_LIMIT = 200; +const ANSWER_PREVIEW_LIMIT = 50; + +const formatValue = (value: unknown) => { + if (value === null || value === undefined || value === '') { + return '(blank)'; + } + if (typeof value === 'number') { + return Number.isInteger(value) ? value.toString() : value.toLocaleString(); + } + return String(value); +}; + +const buildFastAnswer = ( + question: string, + data: PreviewDataResponse, +): string | null => { + const rows = data?.data || []; + const columns = data?.columns || []; + if (!columns.length) { + return null; + } + if (!rows.length) { + return 'No rows were returned for this question.'; + } + + const columnNames = columns.map((column) => column.name); + const rowCount = rows.length; + const sampleRows = rows.slice(0, Math.min(rowCount, 10)); + const hasMetricColumn = columns.some((column) => + /count|total|sum|amount|value|revenue|sales|qty|quantity|rate|percent/i.test( + column.name, + ), + ); + const questionPrefix = question ? `For "${question}", ` : ''; + + if (columns.length === 1) { + const values = sampleRows.map((row) => formatValue(row[0])).join(', '); + return `${questionPrefix}the query returned ${rowCount} row${ + rowCount === 1 ? '' : 's' + }. Values: ${values}.`; + } + + if (columns.length === 2 && hasMetricColumn) { + const [labelColumn, metricColumn] = columnNames; + const lines = sampleRows.map( + (row, index) => + `${index + 1}. ${formatValue(row[0])}: ${formatValue(row[1])}`, + ); + return [ + `${questionPrefix}the top ${sampleRows.length} results by ${metricColumn} are:`, + ...lines, + `Columns used: ${labelColumn}, ${metricColumn}.`, + ].join('\n'); + } + + const preview = sampleRows + .map((row, index) => { + const values = columnNames + .map((columnName, columnIndex) => { + return `${columnName}: ${formatValue(row[columnIndex])}`; + }) + .join(', '); + return `${index + 1}. ${values}`; + }) + .join('\n'); + + return [ + `${questionPrefix}the query returned ${rowCount} row${ + rowCount === 1 ? '' : 's' + }. Showing the first ${sampleRows.length}:`, + preview, + ].join('\n'); +}; export class TextBasedAnswerBackgroundTracker { // tasks is a kv pair of task id and thread response @@ -124,6 +196,25 @@ export class TextBasedAnswerBackgroundTracker { throw error; } + const fastAnswer = buildFastAnswer(threadResponse.question, data); + if (fastAnswer) { + const finishedDetail = { + ...threadResponse.answerDetail, + status: ThreadResponseAnswerStatus.FINISHED, + content: fastAnswer, + numRowsUsedInLLM: Math.min( + data?.data?.length || 0, + ANSWER_PREVIEW_LIMIT, + ), + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: finishedDetail, + }); + threadResponse.answerDetail = finishedDetail; + delete this.tasks[threadResponse.id]; + return; + } + const response = await this.wrenAIAdaptor.createTextBasedAnswer({ query: threadResponse.question, sql: threadResponse.sql, diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index d71eca9964..e227f955ff 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -963,7 +963,7 @@ export class AskingService implements IAskingService { project, manifest: deployment.manifest, modelingOnly: false, - limit: 500, + limit: 100, })) as PreviewDataResponse; } catch (error) { const message = error instanceof Error ? error.message : String(error); From 6c9dd13b3d5514710f2698e7857c231c71b58df5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 01:39:51 +0530 Subject: [PATCH 0296/1087] Ground recommendations in active metadata --- wren-ai-service/src/web/v1/services/ask.py | 28 ---- .../apollo/server/services/askingService.ts | 16 ++ .../apollo/server/services/projectService.ts | 15 ++ .../server/utils/recommendationQuestions.ts | 148 ++++++++++++++++++ 4 files changed, 179 insertions(+), 28 deletions(-) create mode 100644 wren-ui/src/apollo/server/utils/recommendationQuestions.ts diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 497447beb3..1e09e1f372 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2925,34 +2925,6 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results - if direct_orders_sales_sql := self._build_direct_orders_sales_sql( - user_query - ): - api_results = [ - AskResult( - **{ - "sql": direct_orders_sales_sql, - "type": "llm", - } - ) - ] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=["dbo_tblSales"], - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - logger.info( - "Using direct Orders/Sales SQL for query_id %s", - query_id, - ) - return results - explicit_table_names = self._extract_explicit_table_names_from_query( user_query ) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index e227f955ff..076bc9432a 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -42,6 +42,7 @@ import { TrackedAdjustmentResult, } from '../backgrounds'; import { getConfig } from '@server/config'; +import { buildFastRecommendationQuestions } from '@server/utils/recommendationQuestions'; import { TextBasedAnswerBackgroundTracker } from '../backgrounds/textBasedAnswerBackgroundTracker'; import { IAskingTaskTracker, TrackedAskingResult } from './askingTaskTracker'; @@ -598,6 +599,21 @@ export class AskingService implements IAskingService { .sort((a, b) => b.id - a.id) .slice(0, 5); const questions = slicedThreadResponses.map(({ question }) => question); + const fastQuestions = buildFastRecommendationQuestions( + manifest, + this.getThreadRecommendationQuestionsConfig(project).maxQuestions, + questions, + ); + if (fastQuestions.length) { + await this.threadRepository.updateOne(threadId, { + queryId: `fast-thread-${threadId}-${Date.now()}`, + questionsStatus: RecommendationQuestionStatus.FINISHED, + questions: fastQuestions, + questionsError: null, + }); + return; + } + const recommendQuestionData: RecommendationQuestionsInput = { manifest, projectId: project.id.toString(), diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 4621689dea..95ee000b5f 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -27,6 +27,7 @@ import { IMDLService } from './mdlService'; import { ProjectRecommendQuestionBackgroundTracker } from '../backgrounds'; import { ITelemetry } from '../telemetry/telemetry'; import { getConfig } from '../config'; +import { buildFastRecommendationQuestions } from '@server/utils/recommendationQuestions'; const config = getConfig(); @@ -173,6 +174,20 @@ export class ProjectService implements IProjectService { project: Project, ): Promise { const { manifest } = await this.mdlService.makeModelMDL(project); + const fastQuestions = buildFastRecommendationQuestions( + manifest, + this.getProjectRecommendationQuestionsConfig(project).maxQuestions, + ); + if (fastQuestions.length) { + await this.projectRepository.updateOne(project.id, { + queryId: `fast-project-${project.id}-${Date.now()}`, + questionsStatus: RecommendationQuestionStatus.FINISHED, + questions: fastQuestions, + questionsError: null, + }); + return; + } + const recommendQuestionResult = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, diff --git a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts new file mode 100644 index 0000000000..7de04883b9 --- /dev/null +++ b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts @@ -0,0 +1,148 @@ +import { RecommendationQuestion } from '@server/models/adaptor'; +import { Manifest, ModelMDL, ColumnMDL } from '@server/mdl/type'; + +const MAX_MODEL_COUNT = 6; +type RecommendationModel = Partial & { + name: string; + columns?: Partial[]; +}; + +const quoteIdentifier = (identifier: unknown) => { + return `"${String(identifier || '').replace(/"/g, '""')}"`; +}; + +const isMetricColumn = (column: Partial) => { + const name = String(column.name || '').toLowerCase(); + const type = String(column.type || '').toLowerCase(); + return ( + /int|float|double|decimal|numeric|number|real|money/.test(type) || + /amount|value|total|count|qty|quantity|price|cost|revenue|sales|margin|rate/.test( + name, + ) + ); +}; + +const isDateColumn = (column: Partial) => { + const name = String(column.name || '').toLowerCase(); + const type = String(column.type || '').toLowerCase(); + return ( + /date|time|timestamp/.test(type) || /date|time|created|updated/.test(name) + ); +}; + +const isDimensionColumn = (column: Partial) => { + const name = String(column.name || '').toLowerCase(); + if ( + !column.name || + column.isCalculated || + isMetricColumn(column) || + isDateColumn(column) + ) { + return false; + } + return !/(^id$|_id$|uuid|guid|password|token|secret|json|payload)/.test(name); +}; + +const displayName = (model: RecommendationModel) => { + return model.properties?.displayName || model.name; +}; + +const firstUsableModels = (manifest: Manifest) => { + return (manifest.models || []) + .filter((model): model is RecommendationModel => Boolean(model?.name)) + .filter((model) => (model.columns || []).some((column) => column.name)) + .slice(0, MAX_MODEL_COUNT); +}; + +export const buildFastRecommendationQuestions = ( + manifest: Manifest, + maxQuestions = 5, + previousQuestions: string[] = [], +): RecommendationQuestion[] => { + const seen = new Set( + previousQuestions.map((question) => question.trim().toLowerCase()), + ); + const questions: RecommendationQuestion[] = []; + const addQuestion = (question: RecommendationQuestion) => { + if (questions.length >= maxQuestions) { + return; + } + const key = question.question.trim().toLowerCase(); + if (seen.has(key)) { + return; + } + seen.add(key); + questions.push(question); + }; + + for (const model of firstUsableModels(manifest)) { + const modelRef = quoteIdentifier(model.name); + const columns = model.columns || []; + const dimensions = columns.filter(isDimensionColumn); + const metrics = columns.filter(isMetricColumn); + const dates = columns.filter(isDateColumn); + const label = displayName(model); + + if (dimensions[0]) { + const column = dimensions[0]; + const columnRef = `${modelRef}.${quoteIdentifier(column.name)}`; + addQuestion({ + category: label, + question: `What is the distribution of ${column.name} in ${label}?`, + sql: + `SELECT ${columnRef} AS ${quoteIdentifier(column.name)}, ` + + `COUNT(*) AS "RecordCount" FROM ${modelRef} ` + + `GROUP BY ${columnRef} ORDER BY COUNT(*) DESC`, + }); + } + + if (dimensions[0] && metrics[0]) { + const dimension = dimensions[0]; + const metric = metrics[0]; + const dimensionRef = `${modelRef}.${quoteIdentifier(dimension.name)}`; + const metricRef = `${modelRef}.${quoteIdentifier(metric.name)}`; + addQuestion({ + category: label, + question: `Which ${dimension.name} values have the highest ${metric.name} in ${label}?`, + sql: + `SELECT ${dimensionRef} AS ${quoteIdentifier(dimension.name)}, ` + + `SUM(${metricRef}) AS ${quoteIdentifier(`Total${metric.name}`)} ` + + `FROM ${modelRef} GROUP BY ${dimensionRef} ` + + `ORDER BY SUM(${metricRef}) DESC`, + }); + } + + if (dates[0]) { + const date = dates[0]; + const dateRef = `${modelRef}.${quoteIdentifier(date.name)}`; + addQuestion({ + category: label, + question: `Show monthly record count by ${date.name} in ${label}.`, + sql: + `SELECT DATEPART(YEAR, ${dateRef}) AS "year", ` + + `DATEPART(MONTH, ${dateRef}) AS "month", ` + + `COUNT(*) AS "RecordCount" FROM ${modelRef} ` + + `GROUP BY DATEPART(YEAR, ${dateRef}), DATEPART(MONTH, ${dateRef}) ` + + `ORDER BY DATEPART(YEAR, ${dateRef}), DATEPART(MONTH, ${dateRef})`, + }); + } + + const previewColumns = columns + .filter((column) => column.name) + .slice(0, 8) + .map( + (column) => + `${modelRef}.${quoteIdentifier(column.name)} AS ${quoteIdentifier( + column.name, + )}`, + ); + + addQuestion({ + category: label, + question: `Show the first 10 rows from ${label}.`, + sql: `SELECT TOP 10 ${previewColumns.join(', ')} FROM ${modelRef}`, + }); + } + + return questions; +}; From 250959c54559fa9f6dc6386f7c296e436f47c57b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 01:54:46 +0530 Subject: [PATCH 0297/1087] Use generic active datasource SQL grounding --- .../src/pipelines/generation/utils/sql.py | 50 +-- wren-ai-service/src/web/v1/services/ask.py | 305 +----------------- 2 files changed, 9 insertions(+), 346 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 711858e036..04c6f1b8de 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1395,16 +1395,6 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _unwrap_simple_mssql_where_parentheses(normalized) normalized = _rewrite_mssql_limit_clause(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) - normalized = _rewrite_mssql_invented_date_identifiers(normalized) - normalized = _rewrite_mssql_invented_repair_relationship_identifiers(normalized) - normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) - normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) - normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) - normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) - normalized = _rewrite_mssql_invented_failure_category(normalized) - normalized = _rewrite_mssql_invented_report_fields(normalized) - normalized = _rewrite_mssql_invented_ticket_metrics(normalized) - normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) @@ -1480,19 +1470,6 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) - normalized = _rewrite_mssql_sales_schema_aliases(normalized) - normalized = _rewrite_mssql_invented_date_identifiers(normalized) - normalized = _rewrite_mssql_invented_repair_relationship_identifiers( - normalized - ) - normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) - normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) - normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) - normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) - normalized = _rewrite_mssql_invented_failure_category(normalized) - normalized = _rewrite_mssql_invented_report_fields(normalized) - normalized = _rewrite_mssql_invented_ticket_metrics(normalized) - normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -1805,9 +1782,6 @@ async def _classify_generation_result( - Never invent foreign key columns or relationship fields such as "FailurePatternID", "FailurePatternId", "TicketID", or "
ID" unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. - Never invent time bucket columns such as "MONTH", "YEAR", "DAY", "month", "year", or "date" unless that exact column appears in the DATABASE SCHEMA. For monthly, yearly, or daily trends, apply a supported date/time bucket function from SQL FUNCTIONS to a real timestamp column from the selected table. - Every generated SQL query must be grounded only in the connected datasource metadata, deployed semantic model definitions, relationships, and DATABASE SCHEMA shown in the prompt. Do not use table names, column names, join paths, JSON keys, or business dimensions that are not explicitly present in that context. -- For synced repair-log schemas, if "dbo_repair_logs" contains "created_at" and the user asks for monthly repair volume or repair trends, count repair rows and bucket "dbo_repair_logs"."created_at". Do not select, group by, or order by "dbo_repair_logs"."MONTH" or bare "MONTH" unless the schema explicitly contains that column. -- For repair counts grouped by failure category, prefer the richest explicit category field exposed by the connected datasource. If the schema includes "dbo_DebugEntries", "dbo_DebugFixLogs", and "dbo_DebugFixes", group by "dbo_DebugFixes"."Description" after joining "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" and "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id". Otherwise use "dbo_repair_logs"."failure_code" only when that column appears in the schema. Do not invent "failure_category" unless it appears in the DATABASE SCHEMA. -- For repair SLA compliance dashboard/chart requests, do not invent "DAY", "MONTH", "turnaround_time", "sla_due_at", or due-date fields. If the schema only exposes "dbo_repair_logs"."status" and no explicit SLA/duration/deadline column, return a status distribution using "dbo_repair_logs"."status" and COUNT(*) so the UI can render a grounded chart. - For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. - For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. - Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. @@ -1823,28 +1797,8 @@ async def _classify_generation_result( - If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. - If a requested metric such as debug hours, risk score, repair cost, or turnaround time is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. - Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. -- For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. -- For repair SLA compliance charts on "dbo_repair_logs", use "dbo_repair_logs"."status" as the compliance/status dimension when no explicit SLA, due-date, duration, or turnaround column appears in the DATABASE SCHEMA. Never use invented "DAY", "MONTH", or "turnaround_time" fields for SLA compliance. -- For repair counts grouped by failure category, use explicit exposed fields and schema-backed joins only. Prefer "dbo_DebugFixes"."Description" joined through "dbo_DebugFixLogs" when "dbo_DebugEntries"."DebugEntryId", "dbo_DebugFixLogs"."DebugEntryId", "dbo_DebugFixLogs"."FixId", and "dbo_DebugFixes"."Id" are present. Otherwise use "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. -- For PCB/debug-entry failure charts, do not join "dbo_DebugEntries"."DebugEntryId" to "dbo_failure_patterns"."id"; those fields have incompatible types. If both "dbo_DebugEntries"."FailureSys" and "dbo_failure_patterns"."id" exist, join "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". -- For PCB synced database questions: - - Use "dbo_DebugEntries" for debug/PCB event records when the schema contains it. - - Use columns such as "Material", "WorkOrder", "SerialNumber", "FailedAt", "DateIn", "DateOut", "Hours", "Priority", "Actions", "Notes", and "FailureSys" only when they appear in the schema. - - Use "dbo_failure_patterns" for failure names, categories, severity, trend, occurrence counts, daily pattern summaries, and cost impact when those columns appear in the schema. - - For throughput trends across manufacturing/business units, use "dbo_DebugEntries"."BusinessUnit" as the unit dimension and a real debug-entry timestamp such as "dbo_DebugEntries"."DateIn" or "dbo_DebugEntries"."FailedAt" for the trend bucket. Do not use "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH", or invented manufacturing/date fields. - - For top/common PCB failure questions, first prefer grouping by "dbo_DebugFixes"."Description" and counting rows through the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join when those tables and join columns are in the schema. Otherwise prefer grouping by "dbo_failure_patterns"."name" or "dbo_failure_patterns"."category" and counting "dbo_DebugEntries"."DebugEntryId" after joining "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". - - If a useful aggregate already exists in "dbo_failure_patterns" such as "occurrences", it can be used directly for top failure pattern questions without joining event rows. - - For requests such as "show top 10 most common PCB failures", "bar chart of failures by category", or "count of repairs grouped by failure category", generate SQL first. Do not answer with general charting guidance. Return the categorical failure field plus a count metric that can drive a bar chart. - - For failure-category charts, prefer one of these patterns depending on schema availability: - 1. `GROUP BY "dbo_DebugFixes"."Description"` and `COUNT(*)` using the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join - 2. `GROUP BY "dbo_failure_patterns"."category"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` - 3. `GROUP BY "dbo_failure_patterns"."name"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` - 4. `GROUP BY "dbo_repair_logs"."failure_code"` and `COUNT(*)` - - For chart-oriented questions, ensure the final SELECT contains only the chart-ready dimension and metric columns. Avoid prose-like outputs or helper columns. -- For knowledge article tables: - - Use "created_at" for year/month trend buckets. Do not select, group by, or order by invented "YEAR" or "MONTH" columns. - - In "dbo_knowledge_articles", use "helpful" and "views" for effectiveness-style questions, and use "author" for creator/author groupings. Do not invent "effectiveness_score" or "created_by". - - In "dbo_kb_articles", use "created_by_user_id" for creator groupings. Do not invent "created_by" or "author" unless those exact columns appear in the schema. +- For any trend, volume, SLA, duration, growth, or time-series question, use only explicit timestamp/date/duration columns shown in the DATABASE SCHEMA or semantic model. If no suitable field exists, do not invent one. +- For categorical charts, choose explicit category/status/type/name fields from the selected schema and return chart-ready dimension and metric columns. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 1e09e1f372..fa59d26de1 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1932,42 +1932,7 @@ def _build_monthly_repair_volume_sql( return None def _is_direct_heuristic_sql_query(self, query: str) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - asks_manufacturing_throughput = ( - "throughput" in normalized - and any(term in normalized for term in ("manufacturing", "unit", "units")) - ) - asks_failure_counts = ( - "failure" in normalized - and any( - term in normalized - for term in ("count", "counts", "common", "most common", "top") - ) - and any( - term in normalized - for term in ("pcb", "repair", "bar chart", "chart", "category") - ) - ) - asks_sla_compliance = "sla" in normalized and any( - term in normalized - for term in ("compliance", "dashboard", "chart", "repair", "repairs") - ) - asks_monthly_repairs = ( - "repair" in normalized - and any( - term in normalized - for term in ("monthly", "last 12 months", "trend", "volume") - ) - ) - return ( - asks_manufacturing_throughput - or asks_failure_counts - or asks_sla_compliance - or asks_monthly_repairs - ) + return False def _build_heuristic_text_to_sql_fallback( self, @@ -1975,168 +1940,7 @@ def _build_heuristic_text_to_sql_fallback( table_ddls: list[str], table_names: Optional[list[str]] = None, ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - if schema_grounded_sql := self._build_schema_grounded_sales_sql( - query, table_ddls - ): - return schema_grounded_sql - - if throughput_sql := self._build_manufacturing_throughput_sql( - query, table_ddls, table_names=table_names - ): - return throughput_sql - - if repair_failure_count_sql := self._build_repair_failure_count_sql( - query, table_ddls, table_names=table_names - ): - return repair_failure_count_sql - - if repair_sla_sql := self._build_repair_sla_compliance_sql( - query, table_ddls, table_names=table_names - ): - return repair_sla_sql - - if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( - query, table_ddls, table_names=table_names - ): - return monthly_repair_volume_sql - - wants_chart = any( - term in normalized for term in ("chart", "bar chart", "line chart", "graph") - ) - wants_failure_counts = any( - term in normalized - for term in ( - "failure", - "failure category", - "failure code", - "common pcb failures", - "common failures", - "most common", - "top 10", - "top ten", - ) - ) - wants_monthly_repairs = ( - "repair" in normalized - and any( - term in normalized - for term in ("monthly", "last 12 months", "trend", "volume") - ) - ) - - if wants_failure_counts and wants_chart: - top_n = self._extract_requested_top_n(query) - has_pattern_failure_sys = self._schema_contains( - table_ddls, r"\bFailuresys\b", table_names=table_names - ) - has_pattern_occurrences = self._schema_contains( - table_ddls, r"\boccurrences\b", table_names=table_names - ) - has_debug_entries = self._schema_contains( - table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names - ) - has_failure_patterns = self._schema_contains( - table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names - ) - has_failure_sys = self._schema_contains( - table_ddls, r"\bFailureSys\b", table_names=table_names - ) - has_debug_entry_id = self._schema_contains( - table_ddls, r"\bDebugEntryId\b", table_names=table_names - ) - has_pattern_id = self._schema_contains( - table_ddls, r"\bid\b", table_names=table_names - ) - has_pattern_category = self._schema_contains( - table_ddls, r"\bcategory\b", table_names=table_names - ) - has_pattern_name = self._schema_contains( - table_ddls, r"\bname\b", table_names=table_names - ) - - if has_failure_patterns and has_pattern_failure_sys and has_pattern_occurrences: - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'"dbo_failure_patterns"."occurrences" AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'AND "dbo_failure_patterns"."occurrences" IS NOT NULL ' - f'ORDER BY "dbo_failure_patterns"."occurrences" DESC ' - f'LIMIT {top_n}' - ) - - if has_failure_patterns and has_pattern_failure_sys: - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - if ( - has_debug_entries - and has_failure_patterns - and has_failure_sys - and has_debug_entry_id - and has_pattern_id - ): - dimension_column = ( - "category" - if ("category" in normalized and has_pattern_category) - else ("name" if has_pattern_name else "category") - ) - if dimension_column == "category" and not has_pattern_category: - dimension_column = "name" - - return ( - f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' - f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' - f'FROM "dbo_DebugEntries" ' - f'JOIN "dbo_failure_patterns" ' - f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' - f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - has_repair_logs = self._schema_contains( - table_ddls, r"\bdbo_repair_logs\b", table_names=table_names - ) - has_failure_code = self._schema_contains( - table_ddls, r"\bfailure_code\b", table_names=table_names - ) - if has_repair_logs and has_failure_code: - return ( - f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_repair_logs" ' - f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' - f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - if wants_failure_counts and wants_chart: - top_n = self._extract_requested_top_n(query) - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - return None + return self._build_schema_grounded_analytics_sql(query, table_ddls) def _is_schema_grounded_query( self, query: str, db_schemas: Optional[list[str]] = None @@ -2993,7 +2797,7 @@ async def ask( return results if documents and ( - deterministic_sql := self._build_schema_grounded_sales_sql( + deterministic_sql := self._build_schema_grounded_analytics_sql( user_query, table_ddls ) ): @@ -3010,7 +2814,7 @@ async def ask( type="TEXT_TO_SQL", response=api_results, rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", + intent_reasoning="Explicit table request matched deployed schema and generated SQL from active metadata.", retrieved_tables=table_names, trace_id=trace_id, is_followup=True if histories else False, @@ -3048,83 +2852,6 @@ async def ask( ) sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - if self._is_direct_heuristic_sql_query(user_query): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - histories=histories, - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - logger.info( - "Retrieved tables for direct heuristic query_id %s: %s", - query_id, - table_names, - ) - - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using direct heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - if ask_result := self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - ): - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - - if explicit_group_count_sql := self._build_explicit_group_count_sql( - user_query - ): - table_column_reference = ( - self._extract_explicit_table_column_reference(user_query) - ) - table_names = ( - [table_column_reference[0]] if table_column_reference else [] - ) - api_results = [ - AskResult( - **{ - "sql": explicit_group_count_sql, - "type": "llm", - } - ) - ] - rephrased_question = user_query - logger.info( - "Using explicit table-column grouped count SQL for query_id %s", - query_id, - ) - historical_question_result = [] should_skip_pre_sql_retrieval = self._is_data_analysis_query( user_query @@ -3483,36 +3210,18 @@ async def ask( ] if not api_results and ( - audit_log_activity_sql := self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ) - ): - logger.info( - "Using schema-grounded audit log activity SQL for query_id %s", - query_id, - ) - api_results = [ - AskResult( - **{ - "sql": audit_log_activity_sql, - "type": "llm", - } - ) - ] - - if not api_results and ( - deterministic_sales_sql := self._build_schema_grounded_sales_sql( + deterministic_sql := self._build_schema_grounded_analytics_sql( user_query, table_ddls ) ): logger.info( - "Using schema-grounded CWSales SQL for query_id %s", + "Using generic schema-grounded SQL for query_id %s", query_id, ) api_results = [ AskResult( **{ - "sql": deterministic_sales_sql, + "sql": deterministic_sql, "type": "llm", } ) From 62d79e2ce54a79c0ae2f2e004b54d521f473f516 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 01:58:10 +0530 Subject: [PATCH 0298/1087] Revert "Use generic active datasource SQL grounding" This reverts commit 250959c54559fa9f6dc6386f7c296e436f47c57b. --- .../src/pipelines/generation/utils/sql.py | 50 ++- wren-ai-service/src/web/v1/services/ask.py | 305 +++++++++++++++++- 2 files changed, 346 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 04c6f1b8de..711858e036 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1395,6 +1395,16 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _unwrap_simple_mssql_where_parentheses(normalized) normalized = _rewrite_mssql_limit_clause(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) + normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_invented_repair_relationship_identifiers(normalized) + normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) + normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) + normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) + normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) + normalized = _rewrite_mssql_invented_failure_category(normalized) + normalized = _rewrite_mssql_invented_report_fields(normalized) + normalized = _rewrite_mssql_invented_ticket_metrics(normalized) + normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) @@ -1470,6 +1480,19 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) + normalized = _rewrite_mssql_sales_schema_aliases(normalized) + normalized = _rewrite_mssql_invented_date_identifiers(normalized) + normalized = _rewrite_mssql_invented_repair_relationship_identifiers( + normalized + ) + normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) + normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) + normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) + normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) + normalized = _rewrite_mssql_invented_failure_category(normalized) + normalized = _rewrite_mssql_invented_report_fields(normalized) + normalized = _rewrite_mssql_invented_ticket_metrics(normalized) + normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -1782,6 +1805,9 @@ async def _classify_generation_result( - Never invent foreign key columns or relationship fields such as "FailurePatternID", "FailurePatternId", "TicketID", or "
ID" unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. - Never invent time bucket columns such as "MONTH", "YEAR", "DAY", "month", "year", or "date" unless that exact column appears in the DATABASE SCHEMA. For monthly, yearly, or daily trends, apply a supported date/time bucket function from SQL FUNCTIONS to a real timestamp column from the selected table. - Every generated SQL query must be grounded only in the connected datasource metadata, deployed semantic model definitions, relationships, and DATABASE SCHEMA shown in the prompt. Do not use table names, column names, join paths, JSON keys, or business dimensions that are not explicitly present in that context. +- For synced repair-log schemas, if "dbo_repair_logs" contains "created_at" and the user asks for monthly repair volume or repair trends, count repair rows and bucket "dbo_repair_logs"."created_at". Do not select, group by, or order by "dbo_repair_logs"."MONTH" or bare "MONTH" unless the schema explicitly contains that column. +- For repair counts grouped by failure category, prefer the richest explicit category field exposed by the connected datasource. If the schema includes "dbo_DebugEntries", "dbo_DebugFixLogs", and "dbo_DebugFixes", group by "dbo_DebugFixes"."Description" after joining "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" and "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id". Otherwise use "dbo_repair_logs"."failure_code" only when that column appears in the schema. Do not invent "failure_category" unless it appears in the DATABASE SCHEMA. +- For repair SLA compliance dashboard/chart requests, do not invent "DAY", "MONTH", "turnaround_time", "sla_due_at", or due-date fields. If the schema only exposes "dbo_repair_logs"."status" and no explicit SLA/duration/deadline column, return a status distribution using "dbo_repair_logs"."status" and COUNT(*) so the UI can render a grounded chart. - For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. - For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. - Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. @@ -1797,8 +1823,28 @@ async def _classify_generation_result( - If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. - If a requested metric such as debug hours, risk score, repair cost, or turnaround time is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. - Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. -- For any trend, volume, SLA, duration, growth, or time-series question, use only explicit timestamp/date/duration columns shown in the DATABASE SCHEMA or semantic model. If no suitable field exists, do not invent one. -- For categorical charts, choose explicit category/status/type/name fields from the selected schema and return chart-ready dimension and metric columns. +- For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. +- For repair SLA compliance charts on "dbo_repair_logs", use "dbo_repair_logs"."status" as the compliance/status dimension when no explicit SLA, due-date, duration, or turnaround column appears in the DATABASE SCHEMA. Never use invented "DAY", "MONTH", or "turnaround_time" fields for SLA compliance. +- For repair counts grouped by failure category, use explicit exposed fields and schema-backed joins only. Prefer "dbo_DebugFixes"."Description" joined through "dbo_DebugFixLogs" when "dbo_DebugEntries"."DebugEntryId", "dbo_DebugFixLogs"."DebugEntryId", "dbo_DebugFixLogs"."FixId", and "dbo_DebugFixes"."Id" are present. Otherwise use "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. +- For PCB/debug-entry failure charts, do not join "dbo_DebugEntries"."DebugEntryId" to "dbo_failure_patterns"."id"; those fields have incompatible types. If both "dbo_DebugEntries"."FailureSys" and "dbo_failure_patterns"."id" exist, join "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". +- For PCB synced database questions: + - Use "dbo_DebugEntries" for debug/PCB event records when the schema contains it. + - Use columns such as "Material", "WorkOrder", "SerialNumber", "FailedAt", "DateIn", "DateOut", "Hours", "Priority", "Actions", "Notes", and "FailureSys" only when they appear in the schema. + - Use "dbo_failure_patterns" for failure names, categories, severity, trend, occurrence counts, daily pattern summaries, and cost impact when those columns appear in the schema. + - For throughput trends across manufacturing/business units, use "dbo_DebugEntries"."BusinessUnit" as the unit dimension and a real debug-entry timestamp such as "dbo_DebugEntries"."DateIn" or "dbo_DebugEntries"."FailedAt" for the trend bucket. Do not use "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH", or invented manufacturing/date fields. + - For top/common PCB failure questions, first prefer grouping by "dbo_DebugFixes"."Description" and counting rows through the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join when those tables and join columns are in the schema. Otherwise prefer grouping by "dbo_failure_patterns"."name" or "dbo_failure_patterns"."category" and counting "dbo_DebugEntries"."DebugEntryId" after joining "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". + - If a useful aggregate already exists in "dbo_failure_patterns" such as "occurrences", it can be used directly for top failure pattern questions without joining event rows. + - For requests such as "show top 10 most common PCB failures", "bar chart of failures by category", or "count of repairs grouped by failure category", generate SQL first. Do not answer with general charting guidance. Return the categorical failure field plus a count metric that can drive a bar chart. + - For failure-category charts, prefer one of these patterns depending on schema availability: + 1. `GROUP BY "dbo_DebugFixes"."Description"` and `COUNT(*)` using the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join + 2. `GROUP BY "dbo_failure_patterns"."category"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` + 3. `GROUP BY "dbo_failure_patterns"."name"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` + 4. `GROUP BY "dbo_repair_logs"."failure_code"` and `COUNT(*)` + - For chart-oriented questions, ensure the final SELECT contains only the chart-ready dimension and metric columns. Avoid prose-like outputs or helper columns. +- For knowledge article tables: + - Use "created_at" for year/month trend buckets. Do not select, group by, or order by invented "YEAR" or "MONTH" columns. + - In "dbo_knowledge_articles", use "helpful" and "views" for effectiveness-style questions, and use "author" for creator/author groupings. Do not invent "effectiveness_score" or "created_by". + - In "dbo_kb_articles", use "created_by_user_id" for creator groupings. Do not invent "created_by" or "author" unless those exact columns appear in the schema. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index fa59d26de1..1e09e1f372 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1932,7 +1932,42 @@ def _build_monthly_repair_volume_sql( return None def _is_direct_heuristic_sql_query(self, query: str) -> bool: - return False + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + asks_manufacturing_throughput = ( + "throughput" in normalized + and any(term in normalized for term in ("manufacturing", "unit", "units")) + ) + asks_failure_counts = ( + "failure" in normalized + and any( + term in normalized + for term in ("count", "counts", "common", "most common", "top") + ) + and any( + term in normalized + for term in ("pcb", "repair", "bar chart", "chart", "category") + ) + ) + asks_sla_compliance = "sla" in normalized and any( + term in normalized + for term in ("compliance", "dashboard", "chart", "repair", "repairs") + ) + asks_monthly_repairs = ( + "repair" in normalized + and any( + term in normalized + for term in ("monthly", "last 12 months", "trend", "volume") + ) + ) + return ( + asks_manufacturing_throughput + or asks_failure_counts + or asks_sla_compliance + or asks_monthly_repairs + ) def _build_heuristic_text_to_sql_fallback( self, @@ -1940,7 +1975,168 @@ def _build_heuristic_text_to_sql_fallback( table_ddls: list[str], table_names: Optional[list[str]] = None, ) -> str | None: - return self._build_schema_grounded_analytics_sql(query, table_ddls) + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + if schema_grounded_sql := self._build_schema_grounded_sales_sql( + query, table_ddls + ): + return schema_grounded_sql + + if throughput_sql := self._build_manufacturing_throughput_sql( + query, table_ddls, table_names=table_names + ): + return throughput_sql + + if repair_failure_count_sql := self._build_repair_failure_count_sql( + query, table_ddls, table_names=table_names + ): + return repair_failure_count_sql + + if repair_sla_sql := self._build_repair_sla_compliance_sql( + query, table_ddls, table_names=table_names + ): + return repair_sla_sql + + if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( + query, table_ddls, table_names=table_names + ): + return monthly_repair_volume_sql + + wants_chart = any( + term in normalized for term in ("chart", "bar chart", "line chart", "graph") + ) + wants_failure_counts = any( + term in normalized + for term in ( + "failure", + "failure category", + "failure code", + "common pcb failures", + "common failures", + "most common", + "top 10", + "top ten", + ) + ) + wants_monthly_repairs = ( + "repair" in normalized + and any( + term in normalized + for term in ("monthly", "last 12 months", "trend", "volume") + ) + ) + + if wants_failure_counts and wants_chart: + top_n = self._extract_requested_top_n(query) + has_pattern_failure_sys = self._schema_contains( + table_ddls, r"\bFailuresys\b", table_names=table_names + ) + has_pattern_occurrences = self._schema_contains( + table_ddls, r"\boccurrences\b", table_names=table_names + ) + has_debug_entries = self._schema_contains( + table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names + ) + has_failure_patterns = self._schema_contains( + table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names + ) + has_failure_sys = self._schema_contains( + table_ddls, r"\bFailureSys\b", table_names=table_names + ) + has_debug_entry_id = self._schema_contains( + table_ddls, r"\bDebugEntryId\b", table_names=table_names + ) + has_pattern_id = self._schema_contains( + table_ddls, r"\bid\b", table_names=table_names + ) + has_pattern_category = self._schema_contains( + table_ddls, r"\bcategory\b", table_names=table_names + ) + has_pattern_name = self._schema_contains( + table_ddls, r"\bname\b", table_names=table_names + ) + + if has_failure_patterns and has_pattern_failure_sys and has_pattern_occurrences: + return ( + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'"dbo_failure_patterns"."occurrences" AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'AND "dbo_failure_patterns"."occurrences" IS NOT NULL ' + f'ORDER BY "dbo_failure_patterns"."occurrences" DESC ' + f'LIMIT {top_n}' + ) + + if has_failure_patterns and has_pattern_failure_sys: + return ( + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."Failuresys" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + if ( + has_debug_entries + and has_failure_patterns + and has_failure_sys + and has_debug_entry_id + and has_pattern_id + ): + dimension_column = ( + "category" + if ("category" in normalized and has_pattern_category) + else ("name" if has_pattern_name else "category") + ) + if dimension_column == "category" and not has_pattern_category: + dimension_column = "name" + + return ( + f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' + f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' + f'FROM "dbo_DebugEntries" ' + f'JOIN "dbo_failure_patterns" ' + f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' + f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + has_repair_logs = self._schema_contains( + table_ddls, r"\bdbo_repair_logs\b", table_names=table_names + ) + has_failure_code = self._schema_contains( + table_ddls, r"\bfailure_code\b", table_names=table_names + ) + if has_repair_logs and has_failure_code: + return ( + f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_repair_logs" ' + f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' + f'GROUP BY "dbo_repair_logs"."failure_code" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + if wants_failure_counts and wants_chart: + top_n = self._extract_requested_top_n(query) + return ( + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."Failuresys" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + return None def _is_schema_grounded_query( self, query: str, db_schemas: Optional[list[str]] = None @@ -2797,7 +2993,7 @@ async def ask( return results if documents and ( - deterministic_sql := self._build_schema_grounded_analytics_sql( + deterministic_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) ): @@ -2814,7 +3010,7 @@ async def ask( type="TEXT_TO_SQL", response=api_results, rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL from active metadata.", + intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", retrieved_tables=table_names, trace_id=trace_id, is_followup=True if histories else False, @@ -2852,6 +3048,83 @@ async def ask( ) sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + if self._is_direct_heuristic_sql_query(user_query): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + logger.info( + "Retrieved tables for direct heuristic query_id %s: %s", + query_id, + table_names, + ) + + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using direct heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + if ask_result := self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + ): + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + + if explicit_group_count_sql := self._build_explicit_group_count_sql( + user_query + ): + table_column_reference = ( + self._extract_explicit_table_column_reference(user_query) + ) + table_names = ( + [table_column_reference[0]] if table_column_reference else [] + ) + api_results = [ + AskResult( + **{ + "sql": explicit_group_count_sql, + "type": "llm", + } + ) + ] + rephrased_question = user_query + logger.info( + "Using explicit table-column grouped count SQL for query_id %s", + query_id, + ) + historical_question_result = [] should_skip_pre_sql_retrieval = self._is_data_analysis_query( user_query @@ -3210,18 +3483,36 @@ async def ask( ] if not api_results and ( - deterministic_sql := self._build_schema_grounded_analytics_sql( + audit_log_activity_sql := self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ) + ): + logger.info( + "Using schema-grounded audit log activity SQL for query_id %s", + query_id, + ) + api_results = [ + AskResult( + **{ + "sql": audit_log_activity_sql, + "type": "llm", + } + ) + ] + + if not api_results and ( + deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) ): logger.info( - "Using generic schema-grounded SQL for query_id %s", + "Using schema-grounded CWSales SQL for query_id %s", query_id, ) api_results = [ AskResult( **{ - "sql": deterministic_sql, + "sql": deterministic_sales_sql, "type": "llm", } ) From ac8b8d5eaec49e061e92070e71affa8306cb22e5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 14:42:32 +0530 Subject: [PATCH 0299/1087] Use count charts for categorical SQL results --- .../src/pipelines/generation/utils/chart.py | 24 +++++++++++-------- .../generation/test_chart_generation_utils.py | 10 ++++++-- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index faf97df01f..a2596887c8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -341,21 +341,18 @@ def _fallback_chart_type( temporal: list[str], nominal: list[str], ) -> str: - if not quantitative: - return "" - chart_type = requested_chart_type or "bar" if chart_type == "pie": return "pie" if nominal else "" if chart_type in {"line", "area", "multi_line"}: - return chart_type if temporal or nominal else "" + return chart_type if quantitative and (temporal or nominal) else "" if chart_type in {"grouped_bar", "stacked_bar"}: - return chart_type if len(nominal) > 1 else "" + return chart_type if quantitative and len(nominal) > 1 else "" - return "bar" if nominal or temporal else "" + return "bar" if nominal or temporal or quantitative else "" def _build_fallback_chart_schema( @@ -381,8 +378,6 @@ def _build_fallback_chart_schema( query, chart_type, nominal, temporal, columns ) measure = _select_measure_column(query, quantitative) - if not measure: - return {} title = _humanize_title(query or "Chart") @@ -392,18 +387,27 @@ def axis(field: str, field_type: str) -> dict: base["timeUnit"] = "yearmonth" return base + count_axis = { + "aggregate": "count", + "type": "quantitative", + "title": _count_axis_title(query), + } + if chart_type == "pie": color_field = dimensions[0] if dimensions else columns[0] + theta_axis = axis(measure, "quantitative") if measure else count_axis return { "title": title, "mark": {"type": "arc"}, "encoding": { - "theta": axis(measure, "quantitative"), + "theta": theta_axis, "color": axis(color_field, "nominal"), }, } if chart_type in {"line", "area", "multi_line"}: + if not measure: + return {} y_encoding = axis(measure, "quantitative") if {"year", "month"}.issubset({str(c).lower() for c in columns}): month_field = next(c for c in columns if str(c).lower() == "month") @@ -444,7 +448,7 @@ def axis(field: str, field_type: str) -> dict: if x_field in nominal else ("temporal" if x_field in temporal else "ordinal") ) - y_encoding = axis(measure, "quantitative") + y_encoding = axis(measure, "quantitative") if measure else count_axis encoding = { "x": axis(x_field, x_type), "y": y_encoding, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index a35a2b538c..d95099cf3d 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -23,7 +23,7 @@ def test_chart_preprocessor_uses_deterministic_sql_result_order(): ] -def test_fallback_chart_requires_a_real_quantitative_measure(): +def test_fallback_chart_counts_categorical_only_results(): result = build_fallback_chart_result( "Create a chart comparing completed repairs across engineers.", { @@ -32,7 +32,13 @@ def test_fallback_chart_requires_a_real_quantitative_measure(): }, ) - assert result == {"chart_schema": {}, "reasoning": "", "chart_type": ""} + assert result["chart_type"] == "bar" + assert result["chart_schema"]["encoding"]["x"]["field"] == "Status" + assert result["chart_schema"]["encoding"]["y"] == { + "aggregate": "count", + "type": "quantitative", + "title": "Count", + } def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): From 3bc111d144d791dd192c21fb6dd1b9e516d31873 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 15:10:45 +0530 Subject: [PATCH 0300/1087] Make dashboard chart pinning reliable --- wren-ui/src/components/chart/index.tsx | 1 + .../pages/home/promptThread/ChartAnswer.tsx | 19 ++++++++++++------- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/wren-ui/src/components/chart/index.tsx b/wren-ui/src/components/chart/index.tsx index dfc81a7c7e..955ca891cd 100644 --- a/wren-ui/src/components/chart/index.tsx +++ b/wren-ui/src/components/chart/index.tsx @@ -125,6 +125,7 @@ export default function Chart(props: VegaLiteProps) { }; const getChartContent = () => { + if (!values) return null; if (values.length === 0) return
No available data
; if (parsedError) { diff --git a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx index 6a4dbb5683..614873a87d 100644 --- a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx @@ -11,7 +11,11 @@ import LineProperties from '@/components/chart/properties/LineProperties'; import StackedBarProperties from '@/components/chart/properties/StackedBarProperties'; import GroupedBarProperties from '@/components/chart/properties/GroupedBarProperties'; import { Props as AnswerResultProps } from '@/components/pages/home/promptThread/AnswerResult'; -import { ChartTaskStatus, ChartType } from '@/apollo/client/graphql/__types__'; +import { + ChartTaskStatus, + ChartType, + DashboardItemType, +} from '@/apollo/client/graphql/__types__'; import { usePreviewDataMutation } from '@/apollo/client/graphql/home.generated'; import { isEmpty, isEqual } from 'lodash'; import { @@ -19,7 +23,6 @@ import { getChartSpecOptionValues, } from '@/components/chart/handler'; import { useCreateDashboardItemMutation } from '@/apollo/client/graphql/dashboard.generated'; -import { DashboardItemType } from '@/apollo/server/repositories'; import usePromptThreadStore from './store'; const Chart = dynamic(() => import('@/components/chart'), { @@ -78,11 +81,11 @@ const getDynamicProperties = (chartType: ChartType) => { }; const chartTypeToDashboardItemType = ( - chartType: ChartType, + chartType: ChartType | string, ): DashboardItemType | null => { - const normalized = String( - chartType || '', - ).toUpperCase() as keyof typeof DashboardItemType; + const normalized = String(chartType || '') + .replace(/-/g, '_') + .toUpperCase() as keyof typeof DashboardItemType; return DashboardItemType[normalized] || null; }; @@ -190,7 +193,9 @@ export default function ChartAnswer(props: AnswerResultProps) { }; const onPin = () => { - const dashboardItemType = chartTypeToDashboardItemType(chartType as ChartType); + const dashboardItemType = chartTypeToDashboardItemType( + chartType || chartDetail?.chartType, + ); if (!dashboardItemType) { message.error('Chart type is not supported for dashboard pinning.'); return; From 9606924badf346697127614ddd19431442a493de Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 15:27:31 +0530 Subject: [PATCH 0301/1087] Fix dashboard item creation on MSSQL --- .../server/repositories/baseRepository.ts | 36 +++++++++++++++++-- .../repositories/dashboardItemRepository.ts | 27 +++++++++++++- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index a286555250..3ebbfcdf77 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -58,6 +58,7 @@ export class BaseRepository implements IBasicRepository { protected knex: Knex; protected tableName: string; private hasIdColumnCache: boolean | null = null; + private hasIdentityIdPromise?: Promise; constructor({ knexPg, tableName }: { knexPg: Knex; tableName: string }) { this.knex = knexPg; @@ -277,6 +278,29 @@ export class BaseRepository implements IBasicRepository { return (row?.maxId || 0) + 1; } + private async hasIdentityId(executer: Knex | Knex.Transaction) { + if (!this.isMssql(executer)) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = executer('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + } + private async prepareInsertData( data: Partial, executer: Knex | Knex.Transaction, @@ -286,7 +310,11 @@ export class BaseRepository implements IBasicRepository { return dbData; } - if (!(await this.hasIdColumn(executer)) || dbData.id !== undefined) { + if ( + !(await this.hasIdColumn(executer)) || + dbData.id !== undefined || + (await this.hasIdentityId(executer)) + ) { return dbData; } @@ -301,7 +329,11 @@ export class BaseRepository implements IBasicRepository { executer: Knex | Knex.Transaction, ) { const dbData = data.map((item) => this.transformToDBData(item)); - if (!this.isMssql(executer) || !(await this.hasIdColumn(executer))) { + if ( + !this.isMssql(executer) || + !(await this.hasIdColumn(executer)) || + (await this.hasIdentityId(executer)) + ) { return dbData; } diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index bb9f695ecd..564264a94c 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -58,6 +58,7 @@ export class DashboardItemRepository { private readonly jsonbColumns = ['layout', 'detail']; private hasIdColumnCache: boolean | null = null; + private hasIdentityIdPromise?: Promise; private hasTitleColumnCache: boolean | null = null; private hasDisplayNameColumnCache: boolean | null = null; private columnCache = new Map(); @@ -162,7 +163,8 @@ export class DashboardItemRepository includeGeneratedId && hasIdColumn && normalizedData.id === undefined && - this.isMssqlLike(executer) + this.isMssqlLike(executer) && + !(await this.hasIdentityId(executer)) ) { normalizedData.id = await this.getNextId(executer); } @@ -219,6 +221,29 @@ export class DashboardItemRepository ); } + private async hasIdentityId(executer: Knex | Knex.Transaction) { + if (!this.isMssqlLike(executer)) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = executer('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + } + private async getNextId(executer: Knex | Knex.Transaction) { const [row] = await executer(this.tableName).max<{ maxId: number | null }>( 'id as maxId', From f1fc58b184098315e99005251a4e2ad8848bea59 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 15:50:59 +0530 Subject: [PATCH 0302/1087] Route basic datasource questions to SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 25 +++++++++++++++++++ .../test_ask_heuristic_text_to_sql.py | 16 ++++++++++++ 2 files changed, 41 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 1e09e1f372..60bac804be 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -164,6 +164,24 @@ def _is_data_analysis_query(self, query: str) -> bool: if not normalized: return False + app_help_patterns = ( + r"\bhow\s+(do|can)\s+i\s+(use|setup|configure|connect|deploy)\b", + r"\bwhat\s+is\s+wren\b", + r"\bwren\s+(ai|sql|docs|documentation|guide|api)\b", + r"\bhelp\s+(me\s+)?(use|setup|configure|connect)\b", + ) + if any(re.search(pattern, normalized) for pattern in app_help_patterns): + return False + + datasource_question_patterns = ( + r"\b(show|list|display|give|tell|find|get|fetch)\b.+\b(by|from|for|where|with|top|bottom|first|last|count|sum|total|average|avg|rating|status|category|type|month|year|date)\b", + r"\b(how many|what are|what is|which|who)\b.+\b(count|sum|total|average|avg|top|bottom|highest|lowest|distribution|trend|pattern|correlation|rating|status|category|type|month|year|date)\b", + r"\b(group|grouped|breakdown|distribution|trend|correlation|pattern|rank|ranking|summarize|summary)\b", + r"\b(first|last|top|bottom)\s+\d+\b", + ) + if any(re.search(pattern, normalized) for pattern in datasource_question_patterns): + return True + analysis_terms = { "amount", "average", @@ -183,6 +201,7 @@ def _is_data_analysis_query(self, query: str) -> bool: "debug", "failure", "fastest growing", + "feedback", "growth", "group", "grouped", @@ -203,6 +222,8 @@ def _is_data_analysis_query(self, query: str) -> bool: "quantity", "rank", "ranking", + "rating", + "ratings", "region", "regions", "repair", @@ -213,7 +234,11 @@ def _is_data_analysis_query(self, query: str) -> bool: "sales person", "sales rep", "salesperson", + "score", + "scores", "sla", + "status", + "summary", "top", "trend", "turnaround", diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index a11aa59c8e..a7b475e729 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -1,6 +1,22 @@ from src.web.v1.services.ask import AskService +def test_basic_datasource_questions_are_detected_as_data_analysis(): + service = AskService(pipelines={}) + + assert service._is_data_analysis_query("Show copilot feedback by rating") + assert service._is_data_analysis_query("List first 10 rows from CustomerMaster") + assert service._is_data_analysis_query("What is the distribution of status?") + assert service._is_data_analysis_query("Summarize records by category") + + +def test_wren_help_questions_do_not_force_text_to_sql(): + service = AskService(pipelines={}) + + assert not service._is_data_analysis_query("How do I use Wren AI?") + assert not service._is_data_analysis_query("Show me how to configure Wren SQL") + + def test_manufacturing_throughput_trend_uses_debug_entry_business_unit(): service = AskService(pipelines={}) table_ddls = [ From 212539bd9aea0d0e62ecc99c90598e818a371ba8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 15:57:00 +0530 Subject: [PATCH 0303/1087] Revert "Route basic datasource questions to SQL generation" This reverts commit f1fc58b184098315e99005251a4e2ad8848bea59. --- wren-ai-service/src/web/v1/services/ask.py | 25 ------------------- .../test_ask_heuristic_text_to_sql.py | 16 ------------ 2 files changed, 41 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 60bac804be..1e09e1f372 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -164,24 +164,6 @@ def _is_data_analysis_query(self, query: str) -> bool: if not normalized: return False - app_help_patterns = ( - r"\bhow\s+(do|can)\s+i\s+(use|setup|configure|connect|deploy)\b", - r"\bwhat\s+is\s+wren\b", - r"\bwren\s+(ai|sql|docs|documentation|guide|api)\b", - r"\bhelp\s+(me\s+)?(use|setup|configure|connect)\b", - ) - if any(re.search(pattern, normalized) for pattern in app_help_patterns): - return False - - datasource_question_patterns = ( - r"\b(show|list|display|give|tell|find|get|fetch)\b.+\b(by|from|for|where|with|top|bottom|first|last|count|sum|total|average|avg|rating|status|category|type|month|year|date)\b", - r"\b(how many|what are|what is|which|who)\b.+\b(count|sum|total|average|avg|top|bottom|highest|lowest|distribution|trend|pattern|correlation|rating|status|category|type|month|year|date)\b", - r"\b(group|grouped|breakdown|distribution|trend|correlation|pattern|rank|ranking|summarize|summary)\b", - r"\b(first|last|top|bottom)\s+\d+\b", - ) - if any(re.search(pattern, normalized) for pattern in datasource_question_patterns): - return True - analysis_terms = { "amount", "average", @@ -201,7 +183,6 @@ def _is_data_analysis_query(self, query: str) -> bool: "debug", "failure", "fastest growing", - "feedback", "growth", "group", "grouped", @@ -222,8 +203,6 @@ def _is_data_analysis_query(self, query: str) -> bool: "quantity", "rank", "ranking", - "rating", - "ratings", "region", "regions", "repair", @@ -234,11 +213,7 @@ def _is_data_analysis_query(self, query: str) -> bool: "sales person", "sales rep", "salesperson", - "score", - "scores", "sla", - "status", - "summary", "top", "trend", "turnaround", diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index a7b475e729..a11aa59c8e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -1,22 +1,6 @@ from src.web.v1.services.ask import AskService -def test_basic_datasource_questions_are_detected_as_data_analysis(): - service = AskService(pipelines={}) - - assert service._is_data_analysis_query("Show copilot feedback by rating") - assert service._is_data_analysis_query("List first 10 rows from CustomerMaster") - assert service._is_data_analysis_query("What is the distribution of status?") - assert service._is_data_analysis_query("Summarize records by category") - - -def test_wren_help_questions_do_not_force_text_to_sql(): - service = AskService(pipelines={}) - - assert not service._is_data_analysis_query("How do I use Wren AI?") - assert not service._is_data_analysis_query("Show me how to configure Wren SQL") - - def test_manufacturing_throughput_trend_uses_debug_entry_business_unit(): service = AskService(pipelines={}) table_ddls = [ From f8875850019435422d639693a577efa6f182e63a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 16:07:39 +0530 Subject: [PATCH 0304/1087] Bound understanding stage before SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 84 ++++++++++++++-------- 1 file changed, 53 insertions(+), 31 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 1e09e1f372..1deb96bb0b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2897,6 +2897,7 @@ async def ask( use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None + understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) try: sql_user_query = user_query @@ -3143,18 +3144,27 @@ async def ask( ) if not api_results and not should_skip_pre_sql_retrieval: - historical_question = await self._run_with_timeout( - "Historical question retrieval", - self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ), - ) + try: + historical_question = await self._run_with_timeout( + "Historical question retrieval", + self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ), + timeout_seconds=min(understanding_timeout_seconds, 10), + ) - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] + except TimeoutError as exc: + logger.warning( + "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) valid_historical_results = [] for result in historical_question_result: @@ -3181,28 +3191,39 @@ async def ask( elif not api_results and not should_skip_pre_sql_retrieval: original_user_query = user_query # Run both pipeline operations concurrently - sql_samples_task, instructions_task = await self._run_with_timeout( - "SQL pair and instruction retrieval", - asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", + try: + sql_samples_task, instructions_task = await self._run_with_timeout( + "SQL pair and instruction retrieval", + asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), ), - ), - ) + timeout_seconds=understanding_timeout_seconds, + ) - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] - ) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] + ) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) + except TimeoutError as exc: + logger.warning( + "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) + sql_samples = [] + instructions = [] if self._allow_intent_classification: try: @@ -3217,6 +3238,7 @@ async def ask( project_id=ask_request.project_id, configuration=ask_request.configurations, ), + timeout_seconds=understanding_timeout_seconds, ) ).get("post_process", {}) except TimeoutError as exc: From 34af647ad47f129683791b83811b42f4eeb44924 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 16:39:26 +0530 Subject: [PATCH 0305/1087] Show datasource search during pre SQL work --- wren-ai-service/src/web/v1/services/ask.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 1deb96bb0b..cae28b0bd3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -3144,6 +3144,16 @@ async def ask( ) if not api_results and not should_skip_pre_sql_retrieval: + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + try: historical_question = await self._run_with_timeout( "Historical question retrieval", From 4820726154c492a1bf38ea5bd996185729f3a64d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 17:06:14 +0530 Subject: [PATCH 0306/1087] Handle basic metadata questions without SQL --- wren-ai-service/src/web/v1/services/ask.py | 216 ++++++++++++++++++ .../test_ask_heuristic_text_to_sql.py | 86 +++++++ 2 files changed, 302 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cae28b0bd3..d441e7b43f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2577,6 +2577,167 @@ def _extract_retrieval_metadata( ] return documents, table_names, table_ddls + def _is_visualization_request(self, query: str) -> bool: + normalized = (query or "").lower() + return bool( + re.search( + r"\b(?:chart|graph|plot|visuali[sz]e|dashboard|bar|line|pie|donut|" + r"scatter|histogram|heatmap|trend|trends|distribution)\b", + normalized, + ) + ) + + def _get_metadata_question_kind(self, query: str) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").lower()).strip() + if not normalized: + return None + + if self._is_visualization_request(normalized): + return None + + explicit_column_patterns = ( + r"\b(?:what|which|list|show|display|give|describe)\b.*\b(?:columns?|fields?)\b", + r"\b(?:columns?|fields?)\b.*\b(?:available|present|there|exist|schema|metadata)\b", + ) + if any(re.search(pattern, normalized) for pattern in explicit_column_patterns): + return "columns" + + table_patterns = ( + r"\b(?:what|which|list|show|display|give)\b.*\b(?:tables?|models?)\b", + r"\b(?:tables?|models?)\b.*\b(?:available|present|there|exist|in this datasource|in the datasource)\b", + r"\b(?:datasource|database|semantic layer|semantic model)\b.*\b(?:tables?|models?)\b", + ) + if any(re.search(pattern, normalized) for pattern in table_patterns): + return "tables" + + schema_patterns = ( + r"\b(?:what|show|display|describe|list)\b.*\b(?:schema|metadata)\b", + r"\b(?:schema|metadata)\b.*\b(?:of|for|in)\b", + ) + if any(re.search(pattern, normalized) for pattern in schema_patterns): + return "columns" + + return None + + def _find_metadata_table_matches( + self, query: str, tables: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + query_key = self._normalize_schema_token(query) + if not query_key: + return [] + + matches: list[tuple[int, dict[str, Any]]] = [] + for table in tables: + table_name = str(table.get("name") or "") + if not table_name: + continue + short_name = re.split(r"[.$]", table_name)[-1] + normalized_name = self._normalize_schema_token(table_name) + normalized_short_name = self._normalize_schema_token(short_name) + + score = 0 + if normalized_name and normalized_name in query_key: + score = 100 + len(normalized_name) + elif normalized_short_name and normalized_short_name in query_key: + score = 80 + len(normalized_short_name) + + if score: + matches.append((score, table)) + + return [ + table + for _, table in sorted(matches, key=lambda item: item[0], reverse=True) + ] + + def _format_metadata_table_list( + self, tables: list[dict[str, Any]], *, max_tables: int = 120 + ) -> str: + if not tables: + return "I couldn't find any deployed tables in the active datasource metadata." + + sorted_tables = sorted( + {str(table.get("name")) for table in tables if table.get("name")}, + key=str.lower, + ) + shown_tables = sorted_tables[:max_tables] + lines = [ + f"The active datasource has {len(sorted_tables)} deployed table" + f"{'' if len(sorted_tables) == 1 else 's'}:" + ] + lines.extend(f"- {table_name}" for table_name in shown_tables) + if len(sorted_tables) > max_tables: + lines.append( + f"- ...and {len(sorted_tables) - max_tables} more tables." + ) + return "\n".join(lines) + + def _format_metadata_columns( + self, + query: str, + tables: list[dict[str, Any]], + *, + max_tables: int = 25, + max_columns_per_table: int = 60, + ) -> str: + if not tables: + return "I couldn't find any deployed columns in the active datasource metadata." + + matched_tables = self._find_metadata_table_matches(query, tables) + selected_tables = matched_tables or sorted( + tables, key=lambda table: str(table.get("name") or "").lower() + ) + selected_tables = selected_tables[:max_tables] + + heading = ( + "Columns available in the matched deployed table" + if matched_tables and len(selected_tables) == 1 + else "Columns available in the active datasource metadata" + ) + lines = [f"{heading}:"] + for table in selected_tables: + table_name = str(table.get("name") or "unknown_table") + columns = [ + column + for column in table.get("columns", []) + if isinstance(column, dict) and column.get("name") + ] + if not columns: + lines.append(f"- {table_name}: no columns found") + continue + + column_parts = [] + for column in columns[:max_columns_per_table]: + column_name = str(column.get("name")) + column_type = str(column.get("type") or "").upper() + column_parts.append( + f"{column_name} ({column_type})" if column_type else column_name + ) + if len(columns) > max_columns_per_table: + column_parts.append( + f"...and {len(columns) - max_columns_per_table} more" + ) + lines.append(f"- {table_name}: {', '.join(column_parts)}") + + if len(tables) > max_tables and not matched_tables: + lines.append(f"- ...and {len(tables) - max_tables} more tables.") + + return "\n".join(lines) + + def _build_metadata_response( + self, query: str, table_ddls: list[str], table_names: list[str] + ) -> str: + kind = self._get_metadata_question_kind(query) + parsed_tables = self._parse_schema_tables(table_ddls) + + if not parsed_tables and table_names: + parsed_tables = [ + {"name": table_name, "columns": []} for table_name in table_names + ] + + if kind == "columns": + return self._format_metadata_columns(query, parsed_tables) + return self._format_metadata_table_list(parsed_tables) + def _normalize_schema_token(self, value: str) -> str: return re.sub(r"[^a-z0-9]", "", (value or "").lower()) @@ -2926,6 +3087,61 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results + metadata_question_kind = self._get_metadata_question_kind(user_query) + if metadata_question_kind: + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="GENERAL", + rephrased_question=user_query, + intent_reasoning=( + "Basic datasource metadata question detected; " + "retrieving deployed schema metadata directly." + ), + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + retrieval_result = await self._run_with_timeout( + "Metadata schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + metadata_answer = self._build_metadata_response( + user_query, table_ddls, table_names + ) + self._general_streaming_results[query_id] = metadata_answer + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=user_query, + intent_reasoning=( + "Answered from active datasource deployed metadata " + "without SQL generation." + ), + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + results["metadata"]["type"] = "GENERAL" + results["metadata"]["metadata_question_kind"] = ( + metadata_question_kind + ) + results["metadata"]["retrieved_table_count"] = len(documents) + return results + explicit_table_names = self._extract_explicit_table_names_from_query( user_query ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index a11aa59c8e..c8748be59e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -1,6 +1,92 @@ from src.web.v1.services.ask import AskService +def test_metadata_table_question_is_not_sql_or_chart_intent(): + service = AskService(pipelines={}) + + assert service._get_metadata_question_kind( + "What tables are there in this datasource?" + ) == "tables" + assert service._get_metadata_question_kind( + "List the available models in the semantic layer" + ) == "tables" + assert service._get_metadata_question_kind( + "Create a bar chart of orders by table category" + ) is None + + +def test_metadata_column_question_is_not_sql_or_chart_intent(): + service = AskService(pipelines={}) + + assert service._get_metadata_question_kind( + "What columns are available in dbo_orders?" + ) == "columns" + assert service._get_metadata_question_kind( + "Show fields in the CustomerMaster table" + ) == "columns" + assert service._get_metadata_question_kind( + "Show schema for CustomerMaster" + ) == "columns" + assert service._get_metadata_question_kind( + "Show a line chart of monthly order count by customer field" + ) is None + + +def test_metadata_table_answer_lists_deployed_tables(): + service = AskService(pipelines={}) + answer = service._build_metadata_response( + "What tables are there in this datasource?", + [ + """ + CREATE TABLE dbo_orders ( + OrderId INT, + CustomerName VARCHAR + ); + """, + """ + CREATE TABLE dbo_customers ( + CustomerId INT, + Region VARCHAR + ); + """, + ], + [], + ) + + assert "active datasource has 2 deployed tables" in answer + assert "- dbo_orders" in answer + assert "- dbo_customers" in answer + + +def test_metadata_column_answer_lists_matching_table_columns(): + service = AskService(pipelines={}) + answer = service._build_metadata_response( + "What columns are available in dbo_orders?", + [ + """ + CREATE TABLE dbo_orders ( + OrderId INT, + CustomerName VARCHAR, + OrderDate TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_customers ( + CustomerId INT, + Region VARCHAR + ); + """, + ], + [], + ) + + assert "dbo_orders" in answer + assert "OrderId (INT)" in answer + assert "CustomerName (VARCHAR)" in answer + assert "OrderDate (TIMESTAMP)" in answer + assert "dbo_customers" not in answer + + def test_manufacturing_throughput_trend_uses_debug_entry_business_unit(): service = AskService(pipelines={}) table_ddls = [ From 887de8a4a0512688c437be38066bb1a7df1a435b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 17:33:36 +0530 Subject: [PATCH 0307/1087] Make metadata answers intent specific --- wren-ai-service/src/web/v1/services/ask.py | 190 +++++++++++++++++- .../test_ask_heuristic_text_to_sql.py | 102 +++++++++- 2 files changed, 284 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d441e7b43f..960ffc2608 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2595,6 +2595,37 @@ def _get_metadata_question_kind(self, query: str) -> str | None: if self._is_visualization_request(normalized): return None + if re.search(r"\b(?:row|rows|record|records)\s+count\b", normalized): + return None + + relationship_patterns = ( + r"\b(?:relationships?|relations?|joins?|foreign keys?|primary keys?)\b", + r"\b(?:how|what|which|show|list|describe)\b.*\b(?:tables?|models?)\b.*\b(?:connected|related|joined)\b", + ) + if any(re.search(pattern, normalized) for pattern in relationship_patterns): + return "relationships" + + table_count_patterns = ( + r"\b(?:how many|count|number of)\b.*\b(?:tables?|models?)\b", + r"\b(?:tables?|models?)\b.*\b(?:count|number)\b", + ) + if any(re.search(pattern, normalized) for pattern in table_count_patterns): + return "table_count" + + column_count_patterns = ( + r"\b(?:how many|count|number of)\b.*\b(?:columns?|fields?)\b", + r"\b(?:columns?|fields?)\b.*\b(?:count|number)\b", + ) + if any(re.search(pattern, normalized) for pattern in column_count_patterns): + return "column_count" + + schema_patterns = ( + r"\b(?:what|show|display|describe|list)\b.*\b(?:schema|metadata)\b", + r"\b(?:schema|metadata)\b.*\b(?:of|for|in)\b", + ) + if any(re.search(pattern, normalized) for pattern in schema_patterns): + return "schema" + explicit_column_patterns = ( r"\b(?:what|which|list|show|display|give|describe)\b.*\b(?:columns?|fields?)\b", r"\b(?:columns?|fields?)\b.*\b(?:available|present|there|exist|schema|metadata)\b", @@ -2610,13 +2641,6 @@ def _get_metadata_question_kind(self, query: str) -> str | None: if any(re.search(pattern, normalized) for pattern in table_patterns): return "tables" - schema_patterns = ( - r"\b(?:what|show|display|describe|list)\b.*\b(?:schema|metadata)\b", - r"\b(?:schema|metadata)\b.*\b(?:of|for|in)\b", - ) - if any(re.search(pattern, normalized) for pattern in schema_patterns): - return "columns" - return None def _find_metadata_table_matches( @@ -2723,6 +2747,150 @@ def _format_metadata_columns( return "\n".join(lines) + def _extract_metadata_relationships(self, table_ddls: list[str]) -> list[str]: + relationships: list[str] = [] + for ddl in table_ddls or []: + if not isinstance(ddl, str): + continue + table_match = re.search( + r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", + ddl, + flags=re.IGNORECASE, + ) + if not table_match: + continue + source_table = next( + (value for value in table_match.groupdict().values() if value), + "unknown_table", + ) + + for relationship_match in re.finditer( + r"FOREIGN\s+KEY\s*\((?P[^)]+)\)\s+REFERENCES\s+" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" + r"\s*\((?P[^)]+)\)", + ddl, + flags=re.IGNORECASE, + ): + target_table = next( + ( + value + for key, value in relationship_match.groupdict().items() + if key + in { + "quoted", + "bracketed", + "backticked", + "bare", + } + and value + ), + "unknown_table", + ) + source_columns = relationship_match.group("source_columns") + target_columns = relationship_match.group("target_columns") + relationships.append( + f"{source_table}({source_columns}) -> " + f"{target_table}({target_columns})" + ) + + return sorted(set(relationships), key=str.lower) + + def _format_metadata_relationships(self, table_ddls: list[str]) -> str: + relationships = self._extract_metadata_relationships(table_ddls) + if not relationships: + return ( + "I couldn't find explicit relationships or foreign keys in the " + "active datasource metadata." + ) + + lines = [ + f"The active datasource metadata has {len(relationships)} " + f"relationship{'' if len(relationships) == 1 else 's'}:" + ] + lines.extend(f"- {relationship}" for relationship in relationships[:120]) + if len(relationships) > 120: + lines.append(f"- ...and {len(relationships) - 120} more relationships.") + return "\n".join(lines) + + def _format_metadata_schema( + self, query: str, tables: list[dict[str, Any]], table_ddls: list[str] + ) -> str: + matched_tables = self._find_metadata_table_matches(query, tables) + selected_tables = matched_tables or sorted( + tables, key=lambda table: str(table.get("name") or "").lower() + ) + selected_tables = selected_tables[:20] + if not selected_tables: + return "I couldn't find schema details in the active datasource metadata." + + lines = ["Schema details from the active datasource metadata:"] + for table in selected_tables: + table_name = str(table.get("name") or "unknown_table") + columns = [ + column + for column in table.get("columns", []) + if isinstance(column, dict) and column.get("name") + ] + lines.append(f"- {table_name}") + if columns: + column_parts = [] + for column in columns[:60]: + column_name = str(column.get("name")) + column_type = str(column.get("type") or "").upper() + column_parts.append( + f"{column_name} ({column_type})" + if column_type + else column_name + ) + if len(columns) > 60: + column_parts.append(f"...and {len(columns) - 60} more") + lines.append(f" Columns: {', '.join(column_parts)}") + else: + lines.append(" Columns: no columns found") + + relationships = self._extract_metadata_relationships(table_ddls) + if relationships: + lines.append("Relationships:") + lines.extend(f"- {relationship}" for relationship in relationships[:40]) + if len(relationships) > 40: + lines.append(f"- ...and {len(relationships) - 40} more relationships.") + + return "\n".join(lines) + + def _format_metadata_table_count(self, tables: list[dict[str, Any]]) -> str: + table_names = {str(table.get("name")) for table in tables if table.get("name")} + return ( + f"The active datasource has {len(table_names)} deployed table" + f"{'' if len(table_names) == 1 else 's'}." + ) + + def _format_metadata_column_count( + self, query: str, tables: list[dict[str, Any]] + ) -> str: + matched_tables = self._find_metadata_table_matches(query, tables) + selected_tables = matched_tables or tables + total_columns = sum( + len( + [ + column + for column in table.get("columns", []) + if isinstance(column, dict) and column.get("name") + ] + ) + for table in selected_tables + ) + if matched_tables and len(selected_tables) == 1: + table_name = str(selected_tables[0].get("name") or "the matched table") + return f"{table_name} has {total_columns} deployed columns." + return ( + f"The active datasource metadata has {total_columns} deployed columns " + f"across {len(selected_tables)} table" + f"{'' if len(selected_tables) == 1 else 's'}." + ) + def _build_metadata_response( self, query: str, table_ddls: list[str], table_names: list[str] ) -> str: @@ -2734,6 +2902,14 @@ def _build_metadata_response( {"name": table_name, "columns": []} for table_name in table_names ] + if kind == "schema": + return self._format_metadata_schema(query, parsed_tables, table_ddls) + if kind == "relationships": + return self._format_metadata_relationships(table_ddls) + if kind == "table_count": + return self._format_metadata_table_count(parsed_tables) + if kind == "column_count": + return self._format_metadata_column_count(query, parsed_tables) if kind == "columns": return self._format_metadata_columns(query, parsed_tables) return self._format_metadata_table_list(parsed_tables) diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index c8748be59e..ee114780cc 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -26,10 +26,27 @@ def test_metadata_column_question_is_not_sql_or_chart_intent(): ) == "columns" assert service._get_metadata_question_kind( "Show schema for CustomerMaster" - ) == "columns" + ) == "schema" assert service._get_metadata_question_kind( "Show a line chart of monthly order count by customer field" ) is None + assert service._get_metadata_question_kind( + "What is the row count for dbo_orders?" + ) is None + + +def test_metadata_relationship_and_count_questions_have_specific_intents(): + service = AskService(pipelines={}) + + assert service._get_metadata_question_kind( + "What relationships exist between tables?" + ) == "relationships" + assert service._get_metadata_question_kind( + "How many tables are in this datasource?" + ) == "table_count" + assert service._get_metadata_question_kind( + "How many columns are in dbo_orders?" + ) == "column_count" def test_metadata_table_answer_lists_deployed_tables(): @@ -87,6 +104,89 @@ def test_metadata_column_answer_lists_matching_table_columns(): assert "dbo_customers" not in answer +def test_metadata_schema_answer_includes_columns_and_relationships(): + service = AskService(pipelines={}) + answer = service._build_metadata_response( + "Show schema for dbo_orders", + [ + """ + CREATE TABLE dbo_orders ( + OrderId INT, + CustomerId INT, + CONSTRAINT fk_customer FOREIGN KEY (CustomerId) + REFERENCES dbo_customers(CustomerId) + ); + """, + """ + CREATE TABLE dbo_customers ( + CustomerId INT, + Region VARCHAR + ); + """, + ], + [], + ) + + assert "Schema details from the active datasource metadata" in answer + assert "- dbo_orders" in answer + assert "OrderId (INT)" in answer + assert "Relationships:" in answer + assert "dbo_orders(CustomerId) -> dbo_customers(CustomerId)" in answer + + +def test_metadata_relationship_answer_lists_foreign_keys(): + service = AskService(pipelines={}) + answer = service._build_metadata_response( + "What relationships exist between tables?", + [ + """ + CREATE TABLE dbo_orders ( + OrderId INT, + CustomerId INT, + FOREIGN KEY (CustomerId) REFERENCES dbo_customers(CustomerId) + ); + """, + ], + [], + ) + + assert "active datasource metadata has 1 relationship" in answer + assert "dbo_orders(CustomerId) -> dbo_customers(CustomerId)" in answer + + +def test_metadata_count_answers_are_intent_specific(): + service = AskService(pipelines={}) + table_ddls = [ + """ + CREATE TABLE dbo_orders ( + OrderId INT, + CustomerId INT + ); + """, + """ + CREATE TABLE dbo_customers ( + CustomerId INT, + Region VARCHAR, + Segment VARCHAR + ); + """, + ] + + table_count = service._build_metadata_response( + "How many tables are in this datasource?", + table_ddls, + [], + ) + column_count = service._build_metadata_response( + "How many columns are in dbo_customers?", + table_ddls, + [], + ) + + assert table_count == "The active datasource has 2 deployed tables." + assert column_count == "dbo_customers has 3 deployed columns." + + def test_manufacturing_throughput_trend_uses_debug_entry_business_unit(): service = AskService(pipelines={}) table_ddls = [ From 504d6fe404cd72e34172204c4173cf5a47e5fdcf Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 17:50:37 +0530 Subject: [PATCH 0308/1087] Improve generic AI response latency --- .../src/apollo/server/backgrounds/chart.ts | 6 ++-- .../server/backgrounds/recommend-question.ts | 2 +- .../textBasedAnswerBackgroundTracker.ts | 2 +- .../apollo/server/services/askingService.ts | 29 +++++++++++++++++++ .../server/services/askingTaskTracker.ts | 6 ++-- wren-ui/src/hooks/useAskPrompt.tsx | 4 +-- wren-ui/src/pages/home/[id].tsx | 4 +-- 7 files changed, 41 insertions(+), 12 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index d1211c8651..d4a2600b58 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -22,7 +22,7 @@ const isFinalized = (status: ChartStatus) => { ); }; -const MIN_POLL_DELAY = 2000; +const MIN_POLL_DELAY = 1000; const MAX_POLL_DELAY = 10000; export class ChartBackgroundTracker { @@ -48,7 +48,7 @@ export class ChartBackgroundTracker { this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadResponseRepository = threadResponseRepository; - this.intervalTime = 2000; + this.intervalTime = 1000; this.start(); } @@ -243,7 +243,7 @@ export class ChartAdjustmentBackgroundTracker { this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadResponseRepository = threadResponseRepository; - this.intervalTime = 2000; + this.intervalTime = 1000; this.start(); } diff --git a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts index 4831425afa..6575fa83c8 100644 --- a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts +++ b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts @@ -20,7 +20,7 @@ const isFinalized = (status: RecommendationQuestionStatus) => { ].includes(status); }; -const MIN_POLL_DELAY = 2000; +const MIN_POLL_DELAY = 1000; const MAX_POLL_DELAY = 10000; export class ProjectRecommendQuestionBackgroundTracker { diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index b913497f76..fc4ac6a910 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -128,7 +128,7 @@ export class TextBasedAnswerBackgroundTracker { this.projectService = projectService; this.deployService = deployService; this.queryService = queryService; - this.intervalTime = 2000; + this.intervalTime = 1000; this.start(); } diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 076bc9432a..58b123b21a 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -439,6 +439,10 @@ export class AskingService implements IAskingService { private askingTaskRepository: IAskingTaskRepository; private adjustmentBackgroundTracker: AdjustmentBackgroundTaskTracker; private instantRecommendationJobs = new Map>(); + private instantRecommendationResults = new Map< + string, + RecommendationQuestionsResult + >(); private threadRecommendationJobs = new Map>(); private initialized = false; @@ -1203,6 +1207,26 @@ export class AskingService implements IAskingService { currentProject.id, ); + const fastQuestions = buildFastRecommendationQuestions( + manifest, + this.getThreadRecommendationQuestionsConfig(currentProject).maxQuestions, + input.previousQuestions || [], + ); + if (fastQuestions.length) { + const queryId = `fast-instant-${currentProject.id}-${Date.now()}`; + this.instantRecommendationResults.set(queryId, { + status: RecommendationQuestionStatus.FINISHED, + type: null, + response: { questions: fastQuestions }, + error: null, + }); + setTimeout( + () => this.instantRecommendationResults.delete(queryId), + 5 * 60 * 1000, + ); + return { id: queryId }; + } + const response = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, projectId: currentProject.id.toString(), @@ -1215,6 +1239,11 @@ export class AskingService implements IAskingService { public async getInstantRecommendedQuestions( queryId: string, ): Promise { + const localResult = this.instantRecommendationResults.get(queryId); + if (localResult) { + return localResult; + } + const response = await this.wrenAIAdaptor.getRecommendationQuestionsResult(queryId); return response; diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index 04c2cf8707..8dec6fe67b 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -59,8 +59,8 @@ export interface IAskingTaskTracker { } export class AskingTaskTracker implements IAskingTaskTracker { - private readonly minPollDelay = 10000; - private readonly maxPollDelay = 60000; + private readonly minPollDelay = 1000; + private readonly maxPollDelay = 10000; private wrenAIAdaptor: IWrenAIAdaptor; private askingTaskRepository: IAskingTaskRepository; private trackedTasks: Map = new Map(); @@ -598,6 +598,6 @@ export class AskingTaskTracker implements IAskingTaskTracker { return this.minPollDelay; } - return 3000; + return 1500; } } diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index b34e0a66f9..9270483192 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -30,8 +30,8 @@ export interface AskPromptData { recommendedQuestions?: RecommendedQuestionsTask; } -const ASKING_TASK_POLL_INTERVAL_MS = 2000; -const RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS = 2000; +const ASKING_TASK_POLL_INTERVAL_MS = 1000; +const RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS = 1000; const ASKING_TASK_POLL_MAX_INTERVAL_MS = 10000; const RECOMMENDED_QUESTIONS_POLL_MAX_INTERVAL_MS = 10000; diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index 566fee4fc3..7b7efacc81 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -70,8 +70,8 @@ const getThreadResponseIsFinished = (threadResponse: ThreadResponse) => { return isAnswerFinished !== false && isChartFinished !== false; }; -const THREAD_RESPONSE_POLL_INTERVAL_MS = 2000; -const THREAD_RECOMMENDATION_POLL_INTERVAL_MS = 2000; +const THREAD_RESPONSE_POLL_INTERVAL_MS = 1000; +const THREAD_RECOMMENDATION_POLL_INTERVAL_MS = 1000; const THREAD_RESPONSE_POLL_MAX_INTERVAL_MS = 10000; const THREAD_RECOMMENDATION_POLL_MAX_INTERVAL_MS = 10000; From b93febf61929ae3d2f1d7435e380c6d841fe9b2e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 18:10:46 +0530 Subject: [PATCH 0309/1087] Fix recommendation helper import resolution --- wren-ui/src/apollo/server/services/askingService.ts | 2 +- wren-ui/src/apollo/server/services/projectService.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 58b123b21a..dcf1043b17 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -42,7 +42,7 @@ import { TrackedAdjustmentResult, } from '../backgrounds'; import { getConfig } from '@server/config'; -import { buildFastRecommendationQuestions } from '@server/utils/recommendationQuestions'; +import { buildFastRecommendationQuestions } from '../utils/recommendationQuestions'; import { TextBasedAnswerBackgroundTracker } from '../backgrounds/textBasedAnswerBackgroundTracker'; import { IAskingTaskTracker, TrackedAskingResult } from './askingTaskTracker'; diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 95ee000b5f..d6770044a7 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -27,7 +27,7 @@ import { IMDLService } from './mdlService'; import { ProjectRecommendQuestionBackgroundTracker } from '../backgrounds'; import { ITelemetry } from '../telemetry/telemetry'; import { getConfig } from '../config'; -import { buildFastRecommendationQuestions } from '@server/utils/recommendationQuestions'; +import { buildFastRecommendationQuestions } from '../utils/recommendationQuestions'; const config = getConfig(); From 3a0ed4aa56e5b8fa94e4ad989bd672bf27243c1d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 18:32:47 +0530 Subject: [PATCH 0310/1087] Validate recommended questions against datasource --- .../textBasedAnswerBackgroundTracker.ts | 5 ++ .../apollo/server/services/askingService.ts | 49 +++++++++++++++++-- .../apollo/server/services/projectService.ts | 43 +++++++++++++++- .../server/utils/recommendationQuestions.ts | 9 +++- wren-ui/src/common.ts | 1 + 5 files changed, 101 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index fc4ac6a910..eabaf4fd85 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -58,6 +58,11 @@ const buildFastAnswer = ( if (columns.length === 1) { const values = sampleRows.map((row) => formatValue(row[0])).join(', '); + const columnName = columnNames[0]; + const isCountColumn = /count|recordcount|rowcount/i.test(columnName); + if (isCountColumn && Number(sampleRows[0]?.[0]) === 0) { + return `${questionPrefix}the active datasource returned 0 matching records.`; + } return `${questionPrefix}the query returned ${rowCount} row${ rowCount === 1 ? '' : 's' }. Values: ${values}.`; diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index dcf1043b17..8b19eda094 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -603,10 +603,15 @@ export class AskingService implements IAskingService { .sort((a, b) => b.id - a.id) .slice(0, 5); const questions = slicedThreadResponses.map(({ question }) => question); - const fastQuestions = buildFastRecommendationQuestions( + const fastQuestions = await this.filterExecutableRecommendationQuestions( + buildFastRecommendationQuestions( + manifest, + this.getThreadRecommendationQuestionsConfig(project).maxQuestions, + questions, + ), + project, manifest, this.getThreadRecommendationQuestionsConfig(project).maxQuestions, - questions, ); if (fastQuestions.length) { await this.threadRepository.updateOne(threadId, { @@ -1207,10 +1212,15 @@ export class AskingService implements IAskingService { currentProject.id, ); - const fastQuestions = buildFastRecommendationQuestions( + const fastQuestions = await this.filterExecutableRecommendationQuestions( + buildFastRecommendationQuestions( + manifest, + this.getThreadRecommendationQuestionsConfig(currentProject).maxQuestions, + input.previousQuestions || [], + ), + currentProject, manifest, this.getThreadRecommendationQuestionsConfig(currentProject).maxQuestions, - input.previousQuestions || [], ); if (fastQuestions.length) { const queryId = `fast-instant-${currentProject.id}-${Date.now()}`; @@ -1249,6 +1259,37 @@ export class AskingService implements IAskingService { return response; } + private async filterExecutableRecommendationQuestions( + questions: RecommendationQuestion[], + project: Project, + manifest: any, + maxQuestions: number, + ): Promise { + const validQuestions: RecommendationQuestion[] = []; + for (const question of questions) { + try { + const result = (await this.queryService.preview(question.sql, { + project, + manifest, + modelingOnly: false, + limit: 1, + cacheEnabled: false, + })) as PreviewDataResponse; + if (result?.data?.length) { + validQuestions.push(question); + if (validQuestions.length >= maxQuestions) { + break; + } + } + } catch (error) { + logger.warn( + `Skipping recommended question because SQL preview failed: ${question.question}. ${error}`, + ); + } + } + return validQuestions; + } + public async deleteAllByProjectId(projectId: number): Promise { // delete all threads await this.threadRepository.deleteAllBy({ projectId }); diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index d6770044a7..67479ae932 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -28,6 +28,7 @@ import { ProjectRecommendQuestionBackgroundTracker } from '../backgrounds'; import { ITelemetry } from '../telemetry/telemetry'; import { getConfig } from '../config'; import { buildFastRecommendationQuestions } from '../utils/recommendationQuestions'; +import { IQueryService, PreviewDataResponse } from './queryService'; const config = getConfig(); @@ -94,6 +95,7 @@ export class ProjectService implements IProjectService { private projectRepository: IProjectRepository; private metadataService: IDataSourceMetadataService; private mdlService: IMDLService; + private queryService: IQueryService; private wrenAIAdaptor: IWrenAIAdaptor; private projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; private projectRecommendationJobs = new Map>(); @@ -101,6 +103,7 @@ export class ProjectService implements IProjectService { projectRepository, metadataService, mdlService, + queryService, wrenAIAdaptor, telemetry, projectRecommendQuestionBackgroundTracker, @@ -108,6 +111,7 @@ export class ProjectService implements IProjectService { projectRepository: IProjectRepository; metadataService: IDataSourceMetadataService; mdlService: IMDLService; + queryService: IQueryService; wrenAIAdaptor: IWrenAIAdaptor; telemetry: ITelemetry; projectRecommendQuestionBackgroundTracker?: ProjectRecommendQuestionBackgroundTracker; @@ -115,6 +119,7 @@ export class ProjectService implements IProjectService { this.projectRepository = projectRepository; this.metadataService = metadataService; this.mdlService = mdlService; + this.queryService = queryService; this.wrenAIAdaptor = wrenAIAdaptor; this.projectRecommendQuestionBackgroundTracker = projectRecommendQuestionBackgroundTracker ?? @@ -174,7 +179,12 @@ export class ProjectService implements IProjectService { project: Project, ): Promise { const { manifest } = await this.mdlService.makeModelMDL(project); - const fastQuestions = buildFastRecommendationQuestions( + const fastQuestions = await this.filterExecutableRecommendationQuestions( + buildFastRecommendationQuestions( + manifest, + this.getProjectRecommendationQuestionsConfig(project).maxQuestions, + ), + project, manifest, this.getProjectRecommendationQuestionsConfig(project).maxQuestions, ); @@ -331,6 +341,37 @@ export class ProjectService implements IProjectService { ); } + private async filterExecutableRecommendationQuestions( + questions: RecommendationQuestion[], + project: Project, + manifest: any, + maxQuestions: number, + ): Promise { + const validQuestions: RecommendationQuestion[] = []; + for (const question of questions) { + try { + const result = (await this.queryService.preview(question.sql, { + project, + manifest, + modelingOnly: false, + limit: 1, + cacheEnabled: false, + })) as PreviewDataResponse; + if (result?.data?.length) { + validQuestions.push(question); + if (validQuestions.length >= maxQuestions) { + break; + } + } + } catch (error) { + logger.warn( + `Skipping project recommended question because SQL preview failed: ${question.question}. ${error}`, + ); + } + } + return validQuestions; + } + private getProjectRecommendationQuestionsConfig(project: Project) { return { maxCategories: config.projectRecommendationQuestionMaxCategories, diff --git a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts index 7de04883b9..691db14ec7 100644 --- a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts +++ b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts @@ -59,12 +59,13 @@ export const buildFastRecommendationQuestions = ( maxQuestions = 5, previousQuestions: string[] = [], ): RecommendationQuestion[] => { + const candidateLimit = Math.max(maxQuestions * 3, maxQuestions); const seen = new Set( previousQuestions.map((question) => question.trim().toLowerCase()), ); const questions: RecommendationQuestion[] = []; const addQuestion = (question: RecommendationQuestion) => { - if (questions.length >= maxQuestions) { + if (questions.length >= candidateLimit) { return; } const key = question.question.trim().toLowerCase(); @@ -83,6 +84,12 @@ export const buildFastRecommendationQuestions = ( const dates = columns.filter(isDateColumn); const label = displayName(model); + addQuestion({ + category: label, + question: `How many records are in ${label}?`, + sql: `SELECT COUNT(*) AS "RecordCount" FROM ${modelRef}`, + }); + if (dimensions[0]) { const column = dimensions[0]; const columnRef = `${modelRef}.${quoteIdentifier(column.name)}`; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 7075c0d863..e9be30d55b 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -237,6 +237,7 @@ export const initComponents = () => { projectRepository, metadataService, mdlService, + queryService, wrenAIAdaptor, telemetry, projectRecommendQuestionBackgroundTracker, From 7cedfb489b51d0180682b78a528662071e863b1c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 18:57:29 +0530 Subject: [PATCH 0311/1087] Prevent stale ask response reuse --- .../repositories/threadResponseRepository.ts | 33 ++++++++++++------- .../apollo/server/services/askingService.ts | 13 +++++++- .../server/services/askingTaskTracker.ts | 7 ++++ 3 files changed, 41 insertions(+), 12 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts index d68757e51b..6cd9896668 100644 --- a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts @@ -169,10 +169,10 @@ export class ThreadResponseRepository data: Partial<{ status: AskResultStatus; sql: string; - viewId: number; - answerDetail: ThreadResponseAnswerDetail; - breakdownDetail: ThreadResponseBreakdownDetail; - chartDetail: ThreadResponseChartDetail; + viewId: number | null; + answerDetail: ThreadResponseAnswerDetail | null; + breakdownDetail: ThreadResponseBreakdownDetail | null; + chartDetail: ThreadResponseChartDetail | null; adjustment: ThreadResponseAdjustment; }>, queryOptions?: IQueryOptions, @@ -180,15 +180,26 @@ export class ThreadResponseRepository const transformedData = { status: data.status ? data.status : undefined, sql: data.sql ? data.sql : undefined, - viewId: data.viewId ? data.viewId : undefined, - answerDetail: data.answerDetail - ? JSON.stringify(data.answerDetail) + viewId: Object.prototype.hasOwnProperty.call(data, 'viewId') + ? data.viewId : undefined, - breakdownDetail: data.breakdownDetail - ? JSON.stringify(data.breakdownDetail) + answerDetail: Object.prototype.hasOwnProperty.call(data, 'answerDetail') + ? data.answerDetail + ? JSON.stringify(data.answerDetail) + : null : undefined, - chartDetail: data.chartDetail - ? JSON.stringify(data.chartDetail) + breakdownDetail: Object.prototype.hasOwnProperty.call( + data, + 'breakdownDetail', + ) + ? data.breakdownDetail + ? JSON.stringify(data.breakdownDetail) + : null + : undefined, + chartDetail: Object.prototype.hasOwnProperty.call(data, 'chartDetail') + ? data.chartDetail + ? JSON.stringify(data.chartDetail) + : null : undefined, adjustment: data.adjustment ? JSON.stringify(data.adjustment) : undefined, updatedAt: new Date(), diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 8b19eda094..ec966ae5f5 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -122,6 +122,13 @@ const isChartGenerationInProgress = (status?: ChartStatus | string | null) => status || '', ); +const isContextualFollowUpQuestion = (question?: string) => { + if (!question) return false; + return /\b(previous|last|above|earlier|same|that|those|them|it|this)\b/i.test( + question, + ); +}; + // adjustment input export interface AdjustmentReasoningInput { tables: string[]; @@ -709,7 +716,7 @@ export class AskingService implements IAskingService { // if it's a follow-up question, then the input will have a threadId // then use the threadId to get the sql and get the steps of last thread response // construct it into AskHistory and pass to ask - const histories = threadId + const histories = threadId && isContextualFollowUpQuestion(input.question) ? await this.getAskingHistory(threadId, threadResponseId) : null; const response = await this.askingTaskTracker.createAskingTask({ @@ -886,6 +893,10 @@ export class AskingService implements IAskingService { return await this.threadResponseRepository.updateOne(responseId, { sql: data.sql, + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, }); } diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index 8dec6fe67b..f81f87f5e5 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -448,10 +448,17 @@ export class AskingTaskTracker implements IAskingTaskTracker { await this.threadResponseRepository.updateOne(task.threadResponseId, { sql: view.statement, viewId: response.viewId, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, }); } else { await this.threadResponseRepository.updateOne(task.threadResponseId, { sql: response?.sql, + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, }); } } From 809d74c08bc58e1136b9e3d15f3b4f3020f7d635 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 19:17:35 +0530 Subject: [PATCH 0312/1087] Respect explicit chart intent --- .../src/pipelines/generation/utils/chart.py | 93 +++++++++++++++++-- .../generation/test_chart_generation_utils.py | 48 ++++++++++ 2 files changed, 134 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index a2596887c8..0b43b0ab51 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -31,15 +31,68 @@ def _detect_requested_chart_type(query: str | None) -> str: ("bar", ["waterfall", "waterfall chart"]), ("pie", ["pie chart", "donut chart", "doughnut chart"]), ("area", ["area chart", "area graph"]), + ("scatter", ["scatter chart", "scatter plot", "scatter graph"]), ] for chart_type, patterns in checks: if any(pattern in normalized for pattern in patterns): return chart_type - if re.search(r"\b(chart|graph|plot|visuali[sz](?:e|ation)?)\b", normalized): + return "" + + +def _closest_supported_chart_type( + requested_chart_type: str, + quantitative: list[str], + temporal: list[str], + nominal: list[str], +) -> str: + if requested_chart_type != "scatter": + return requested_chart_type + + if temporal and quantitative: + return "line" + if nominal and quantitative: return "bar" return "" +def _chart_reasoning( + requested_chart_type: str, + chart_type: str, + chart_schema: dict, + existing_reasoning: str = "", +) -> str: + if not chart_schema: + return existing_reasoning + if requested_chart_type == "scatter" and chart_type: + return ( + "The user requested a scatter chart, but scatter charts are not " + "supported by the current chart type contract. Generated the closest " + f"supported visualization ({chart_type.replace('_', ' ')} chart) " + "from the SQL result columns." + ) + if requested_chart_type and requested_chart_type != chart_type: + return ( + f"The requested {requested_chart_type.replace('_', ' ')} chart was " + "not suitable for the returned SQL result shape. Generated the " + f"closest meaningful {chart_type.replace('_', ' ')} chart from the " + "SQL result columns." + ) + if existing_reasoning: + return existing_reasoning + return "Generated from the SQL result columns and requested chart type." + + +def _chart_type_matches_request( + requested_chart_type: str, + actual_chart_type: str, +) -> bool: + if not requested_chart_type: + return True + if requested_chart_type == "scatter": + return actual_chart_type in {"line", "bar"} + return requested_chart_type == actual_chart_type + + def _safe_column_names(columns: list[Any]) -> list[str]: return [str(column) for column in columns if column is not None and str(column)] @@ -341,7 +394,9 @@ def _fallback_chart_type( temporal: list[str], nominal: list[str], ) -> str: - chart_type = requested_chart_type or "bar" + chart_type = _closest_supported_chart_type( + requested_chart_type or "bar", quantitative, temporal, nominal + ) if chart_type == "pie": return "pie" if nominal else "" @@ -483,8 +538,10 @@ def build_fallback_chart_result( ) -> dict: processed = ChartDataPreprocessor().run(data) sample_data = processed.get("sample_data", []) - chart_type = _detect_requested_chart_type(query) or "bar" - chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) + requested_chart_type = _detect_requested_chart_type(query) + chart_type = requested_chart_type or "bar" + if not requested_chart_type: + chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) chart_schema = _build_fallback_chart_schema(query, chart_type, sample_data) if not chart_schema: return { @@ -501,7 +558,11 @@ def build_fallback_chart_result( return { "chart_schema": chart_schema, - "reasoning": "Generated from the SQL result columns and requested chart type.", + "reasoning": _chart_reasoning( + requested_chart_type, + chart_type, + chart_schema, + ), "chart_type": chart_type, } @@ -908,9 +969,15 @@ def run( chart_schema = _normalize_chart_schema_fields( chart_schema, list(sample_data[0].keys()) if sample_data else [] ) + actual_chart_type = _chart_type_from_schema( + chart_schema, generation_result.get("chart_type", "") + ) if ( not _is_schema_compatible_with_sample_data(chart_schema, sample_data) + or not _chart_type_matches_request( + requested_chart_type, actual_chart_type + ) or _needs_deterministic_bar_fallback( chart_schema, chart_type or "", sample_data ) @@ -919,6 +986,8 @@ def run( query, chart_type or "bar", sample_data ) chart_type = _chart_type_from_schema(chart_schema, chart_type) + else: + chart_type = actual_chart_type if not chart_schema: return { @@ -942,7 +1011,12 @@ def run( return { "results": { "chart_schema": chart_schema, - "reasoning": reasoning, + "reasoning": _chart_reasoning( + requested_chart_type, + chart_type, + chart_schema, + reasoning, + ), "chart_type": chart_type, } } @@ -957,7 +1031,12 @@ def run( return { "results": { "chart_schema": fallback_schema, - "reasoning": reasoning, + "reasoning": _chart_reasoning( + requested_chart_type, + fallback_chart_type, + fallback_schema, + reasoning, + ), "chart_type": fallback_chart_type, } } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index d95099cf3d..0790fc3131 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -61,6 +61,54 @@ def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Market" +def test_explicit_bar_chart_is_not_refined_to_grouped_bar(): + result = build_fallback_chart_result( + "Create a bar chart of new orders by Customer in each Market.", + { + "columns": [ + {"name": "Market"}, + {"name": "Customer"}, + {"name": "OrderCount"}, + ], + "data": [["North", "Acme", 10], ["South", "Globex", 8]], + }, + ) + + assert result["chart_type"] == "bar" + assert result["chart_schema"]["mark"]["type"] == "bar" + assert "xOffset" not in result["chart_schema"]["encoding"] + + +def test_explicit_stacked_bar_chart_is_respected(): + result = build_fallback_chart_result( + "Create a stacked bar chart of sales by Market and Division.", + { + "columns": [ + {"name": "Market"}, + {"name": "Division"}, + {"name": "Sales"}, + ], + "data": [["North", "A", 10], ["North", "B", 8]], + }, + ) + + assert result["chart_type"] == "stacked_bar" + assert result["chart_schema"]["encoding"]["y"]["stack"] == "zero" + + +def test_scatter_request_uses_closest_supported_chart_with_reasoning(): + result = build_fallback_chart_result( + "Create a scatter plot of revenue by market.", + { + "columns": [{"name": "Market"}, {"name": "Revenue"}], + "data": [["North", 100], ["South", 200]], + }, + ) + + assert result["chart_type"] == "bar" + assert "scatter charts are not supported" in result["reasoning"] + + def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): assert not _is_schema_compatible_with_sample_data( { From 52f65a6fa50ef16c2bb7191c6e02a10ae63396de Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 19:30:03 +0530 Subject: [PATCH 0313/1087] Revert "Respect explicit chart intent" This reverts commit 809d74c08bc58e1136b9e3d15f3b4f3020f7d635. --- .../src/pipelines/generation/utils/chart.py | 93 ++----------------- .../generation/test_chart_generation_utils.py | 48 ---------- 2 files changed, 7 insertions(+), 134 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 0b43b0ab51..a2596887c8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -31,68 +31,15 @@ def _detect_requested_chart_type(query: str | None) -> str: ("bar", ["waterfall", "waterfall chart"]), ("pie", ["pie chart", "donut chart", "doughnut chart"]), ("area", ["area chart", "area graph"]), - ("scatter", ["scatter chart", "scatter plot", "scatter graph"]), ] for chart_type, patterns in checks: if any(pattern in normalized for pattern in patterns): return chart_type - return "" - - -def _closest_supported_chart_type( - requested_chart_type: str, - quantitative: list[str], - temporal: list[str], - nominal: list[str], -) -> str: - if requested_chart_type != "scatter": - return requested_chart_type - - if temporal and quantitative: - return "line" - if nominal and quantitative: + if re.search(r"\b(chart|graph|plot|visuali[sz](?:e|ation)?)\b", normalized): return "bar" return "" -def _chart_reasoning( - requested_chart_type: str, - chart_type: str, - chart_schema: dict, - existing_reasoning: str = "", -) -> str: - if not chart_schema: - return existing_reasoning - if requested_chart_type == "scatter" and chart_type: - return ( - "The user requested a scatter chart, but scatter charts are not " - "supported by the current chart type contract. Generated the closest " - f"supported visualization ({chart_type.replace('_', ' ')} chart) " - "from the SQL result columns." - ) - if requested_chart_type and requested_chart_type != chart_type: - return ( - f"The requested {requested_chart_type.replace('_', ' ')} chart was " - "not suitable for the returned SQL result shape. Generated the " - f"closest meaningful {chart_type.replace('_', ' ')} chart from the " - "SQL result columns." - ) - if existing_reasoning: - return existing_reasoning - return "Generated from the SQL result columns and requested chart type." - - -def _chart_type_matches_request( - requested_chart_type: str, - actual_chart_type: str, -) -> bool: - if not requested_chart_type: - return True - if requested_chart_type == "scatter": - return actual_chart_type in {"line", "bar"} - return requested_chart_type == actual_chart_type - - def _safe_column_names(columns: list[Any]) -> list[str]: return [str(column) for column in columns if column is not None and str(column)] @@ -394,9 +341,7 @@ def _fallback_chart_type( temporal: list[str], nominal: list[str], ) -> str: - chart_type = _closest_supported_chart_type( - requested_chart_type or "bar", quantitative, temporal, nominal - ) + chart_type = requested_chart_type or "bar" if chart_type == "pie": return "pie" if nominal else "" @@ -538,10 +483,8 @@ def build_fallback_chart_result( ) -> dict: processed = ChartDataPreprocessor().run(data) sample_data = processed.get("sample_data", []) - requested_chart_type = _detect_requested_chart_type(query) - chart_type = requested_chart_type or "bar" - if not requested_chart_type: - chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) + chart_type = _detect_requested_chart_type(query) or "bar" + chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) chart_schema = _build_fallback_chart_schema(query, chart_type, sample_data) if not chart_schema: return { @@ -558,11 +501,7 @@ def build_fallback_chart_result( return { "chart_schema": chart_schema, - "reasoning": _chart_reasoning( - requested_chart_type, - chart_type, - chart_schema, - ), + "reasoning": "Generated from the SQL result columns and requested chart type.", "chart_type": chart_type, } @@ -969,15 +908,9 @@ def run( chart_schema = _normalize_chart_schema_fields( chart_schema, list(sample_data[0].keys()) if sample_data else [] ) - actual_chart_type = _chart_type_from_schema( - chart_schema, generation_result.get("chart_type", "") - ) if ( not _is_schema_compatible_with_sample_data(chart_schema, sample_data) - or not _chart_type_matches_request( - requested_chart_type, actual_chart_type - ) or _needs_deterministic_bar_fallback( chart_schema, chart_type or "", sample_data ) @@ -986,8 +919,6 @@ def run( query, chart_type or "bar", sample_data ) chart_type = _chart_type_from_schema(chart_schema, chart_type) - else: - chart_type = actual_chart_type if not chart_schema: return { @@ -1011,12 +942,7 @@ def run( return { "results": { "chart_schema": chart_schema, - "reasoning": _chart_reasoning( - requested_chart_type, - chart_type, - chart_schema, - reasoning, - ), + "reasoning": reasoning, "chart_type": chart_type, } } @@ -1031,12 +957,7 @@ def run( return { "results": { "chart_schema": fallback_schema, - "reasoning": _chart_reasoning( - requested_chart_type, - fallback_chart_type, - fallback_schema, - reasoning, - ), + "reasoning": reasoning, "chart_type": fallback_chart_type, } } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index 0790fc3131..d95099cf3d 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -61,54 +61,6 @@ def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Market" -def test_explicit_bar_chart_is_not_refined_to_grouped_bar(): - result = build_fallback_chart_result( - "Create a bar chart of new orders by Customer in each Market.", - { - "columns": [ - {"name": "Market"}, - {"name": "Customer"}, - {"name": "OrderCount"}, - ], - "data": [["North", "Acme", 10], ["South", "Globex", 8]], - }, - ) - - assert result["chart_type"] == "bar" - assert result["chart_schema"]["mark"]["type"] == "bar" - assert "xOffset" not in result["chart_schema"]["encoding"] - - -def test_explicit_stacked_bar_chart_is_respected(): - result = build_fallback_chart_result( - "Create a stacked bar chart of sales by Market and Division.", - { - "columns": [ - {"name": "Market"}, - {"name": "Division"}, - {"name": "Sales"}, - ], - "data": [["North", "A", 10], ["North", "B", 8]], - }, - ) - - assert result["chart_type"] == "stacked_bar" - assert result["chart_schema"]["encoding"]["y"]["stack"] == "zero" - - -def test_scatter_request_uses_closest_supported_chart_with_reasoning(): - result = build_fallback_chart_result( - "Create a scatter plot of revenue by market.", - { - "columns": [{"name": "Market"}, {"name": "Revenue"}], - "data": [["North", 100], ["South", 200]], - }, - ) - - assert result["chart_type"] == "bar" - assert "scatter charts are not supported" in result["reasoning"] - - def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): assert not _is_schema_compatible_with_sample_data( { From e10b28f5d2b04611acc2c19398317c44789dc21a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 23:25:09 +0530 Subject: [PATCH 0314/1087] Avoid historical SQL reuse for independent asks --- wren-ai-service/src/web/v1/services/ask.py | 15 +++++++++++- .../test_ask_heuristic_text_to_sql.py | 24 ++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 960ffc2608..71bb725adf 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -240,6 +240,13 @@ def _needs_conversation_context(self, query: str) -> bool: ) ) + def _should_reuse_historical_question_sql( + self, + query: str, + histories: list[AskHistory] | None, + ) -> bool: + return bool(histories) and self._needs_conversation_context(query) + def _rewrite_query_for_text_to_sql(self, query: str) -> str: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -3535,7 +3542,13 @@ async def ask( user_query, ) - if not api_results and not should_skip_pre_sql_retrieval: + if ( + not api_results + and not should_skip_pre_sql_retrieval + and self._should_reuse_historical_question_sql( + user_query, histories + ) + ): if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="searching", diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index ee114780cc..ee3cdbc5cb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -1,4 +1,26 @@ -from src.web.v1.services.ask import AskService +from src.web.v1.services.ask import AskHistory, AskService + + +def test_independent_question_does_not_reuse_historical_sql(): + service = AskService(pipelines={}) + + assert not service._should_reuse_historical_question_sql( + "Show monthly order count by market.", + [], + ) + assert not service._should_reuse_historical_question_sql( + "Show monthly order count by market.", + [AskHistory(question="previous", sql="SELECT 1")], + ) + + +def test_contextual_followup_can_reuse_historical_sql(): + service = AskService(pipelines={}) + + assert service._should_reuse_historical_question_sql( + "Use the same table and show it by month.", + [AskHistory(question="previous", sql="SELECT 1")], + ) def test_metadata_table_question_is_not_sql_or_chart_intent(): From 92f8860207b4de069fde2eab2ea15c3f30718e1a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 1 Jul 2026 23:45:37 +0530 Subject: [PATCH 0315/1087] Bypass cached previews for ask responses --- .../server/backgrounds/textBasedAnswerBackgroundTracker.ts | 1 + wren-ui/src/apollo/server/services/askingService.ts | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index eabaf4fd85..083e9c579f 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -185,6 +185,7 @@ export class TextBasedAnswerBackgroundTracker { manifest: mdl, modelingOnly: false, limit: ANSWER_PREVIEW_LIMIT, + cacheEnabled: false, })) as PreviewDataResponse; } catch (error) { logger.error(`Error when query sql data: ${error}`); diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index ec966ae5f5..3a0bb5cf63 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1000,6 +1000,7 @@ export class AskingService implements IAskingService { manifest: deployment.manifest, modelingOnly: false, limit: 100, + cacheEnabled: false, })) as PreviewDataResponse; } catch (error) { const message = error instanceof Error ? error.message : String(error); @@ -1067,6 +1068,7 @@ export class AskingService implements IAskingService { manifest: deployment.manifest, modelingOnly: false, limit: 500, + cacheEnabled: false, })) as PreviewDataResponse; } catch (error) { const message = error instanceof Error ? error.message : String(error); @@ -1134,6 +1136,7 @@ export class AskingService implements IAskingService { project, manifest: mdl, limit, + cacheEnabled: false, })) as PreviewDataResponse; this.telemetry.sendEvent(eventName, { sql: response.sql }); return data; @@ -1176,6 +1179,7 @@ export class AskingService implements IAskingService { project, manifest: mdl, limit, + cacheEnabled: false, })) as PreviewDataResponse; this.telemetry.sendEvent(eventName, { sql }); return data; From 36f236b56c65b85c91b08f72768f0c97d6d942b5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 00:17:08 +0530 Subject: [PATCH 0316/1087] Reject low-signal SQL for analytical asks --- wren-ai-service/src/web/v1/services/ask.py | 163 ++++++++++++++++-- .../test_ask_heuristic_text_to_sql.py | 94 ++++++++++ 2 files changed, 244 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 71bb725adf..14d60a79ef 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2218,6 +2218,19 @@ def _build_schema_grounded_operational_sql( if not any(term in normalized_query for term in operational_terms): return None + asks_status_dimension = any( + term in normalized_query + for term in ( + "status", + "open", + "closed", + "completed", + "in-progress", + "in progress", + "active", + ) + ) + scored_tables: list[tuple[int, dict[str, Any]]] = [] for table in tables: table_name = str(table.get("name") or "") @@ -2270,16 +2283,25 @@ def _build_schema_grounded_operational_sql( ) dimension_candidates: list[tuple[str, ...]] = [] - if "manufacturing" in normalized_query or "unit" in normalized_query: + asks_unit_dimension = ( + "manufacturing" in normalized_query + or "business unit" in normalized_query + or "business units" in normalized_query + or re.search(r"\bunit(?:s)?\b", normalized_query) is not None + ) + if asks_unit_dimension: dimension_candidates.append( ( "manufacturing_unit", "manufacturing unit", + "business_unit", + "business unit", + "BusinessUnit", + "BU", "unit", "assignee_user_id", "created_by_user_id", "org_id", - "status", ) ) if "component" in normalized_query: @@ -2294,11 +2316,7 @@ def _build_schema_grounded_operational_sql( dimension_candidates.append(("source", "author", "category", "status")) if "priority" in normalized_query: dimension_candidates.append(("priority", "status")) - if ( - "status" in normalized_query - or "open" in normalized_query - or "closed" in normalized_query - ): + if asks_status_dimension: dimension_candidates.append(("status", "priority")) if "assignee" in normalized_query: dimension_candidates.append(("assignee_user_id", "created_by_user_id")) @@ -2309,11 +2327,13 @@ def _build_schema_grounded_operational_sql( if dimension and dimension not in dimensions: dimensions.append(dimension) + if asks_unit_dimension and not dimensions: + return None + if not dimensions: fallback_dimension = self._find_first_schema_column( table, ( - "status", "priority", "category", "subcategory", @@ -2327,6 +2347,11 @@ def _build_schema_grounded_operational_sql( if fallback_dimension: dimensions.append(fallback_dimension) + if not dimensions and asks_status_dimension: + fallback_dimension = self._find_schema_column(table, ("status",)) + if fallback_dimension: + dimensions.append(fallback_dimension) + wants_trend = any( term in normalized_query for term in ("trend", "monthly", "month", "line chart", "over time") @@ -2335,6 +2360,9 @@ def _build_schema_grounded_operational_sql( limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) limit = int(limit_match.group(1)) if limit_match else 10 + if wants_trend and not date_column: + return None + if wants_trend and date_column: date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" select_parts = [ @@ -3053,6 +3081,98 @@ def _is_valid_select_sql(self, sql: Optional[str]) -> bool: return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) + def _extract_top_level_select_expressions(self, sql: str) -> list[str]: + normalized = re.sub(r"\s+", " ", (sql or "").strip()) + select_match = re.search(r"\bSELECT\b", normalized, flags=re.IGNORECASE) + if not select_match: + return [] + + start = select_match.end() + from_match_start: int | None = None + depth = 0 + index = start + while index < len(normalized): + char = normalized[index] + if char == "(": + depth += 1 + elif char == ")" and depth > 0: + depth -= 1 + elif ( + depth == 0 + and normalized[index : index + 6].upper() == " FROM " + ): + from_match_start = index + break + index += 1 + + if from_match_start is None: + return [] + + select_clause = normalized[start:from_match_start].strip() + select_clause = re.sub( + r"^TOP\s+\d+\s+", + "", + select_clause, + flags=re.IGNORECASE, + ) + expressions: list[str] = [] + depth = 0 + expression_start = 0 + for index, char in enumerate(select_clause): + if char == "(": + depth += 1 + elif char == ")" and depth > 0: + depth -= 1 + elif char == "," and depth == 0: + expression = select_clause[expression_start:index].strip() + if expression: + expressions.append(expression) + expression_start = index + 1 + + last_expression = select_clause[expression_start:].strip() + if last_expression: + expressions.append(last_expression) + return expressions + + def _is_low_signal_sql_for_question(self, query: str, sql: str) -> bool: + if not self._is_data_analysis_query(query): + return False + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not any( + term in normalized_query + for term in ( + "chart", + "trend", + "monthly", + "rate", + "revenue", + "sales", + "volume", + "count", + "top", + "highest", + "underperforming", + "distribution", + "group", + "grouped", + " by ", + ) + ): + return False + + normalized_sql = re.sub(r"\s+", " ", (sql or "").strip()) + if re.search( + r"\b(?:COUNT|SUM|AVG|MIN|MAX|DATEPART|DATETRUNC|DATE_TRUNC|" + r"ROW_NUMBER|RANK|DENSE_RANK)\s*\(", + normalized_sql, + flags=re.IGNORECASE, + ): + return False + + expressions = self._extract_top_level_select_expressions(normalized_sql) + return len(expressions) <= 1 + def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: if not self._is_valid_select_sql(sql): return None @@ -4246,12 +4366,19 @@ async def ask( if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - if ask_result := self._build_ask_result_from_sql( - sql_valid_result.get("sql") + generated_sql = sql_valid_result.get("sql") + if generated_sql and self._is_low_signal_sql_for_question( + sql_user_query, generated_sql ): + invalid_sql = generated_sql + error_message = ( + "SQL generation produced a low-signal result shape for " + "the requested analysis." + ) + elif ask_result := self._build_ask_result_from_sql(generated_sql): api_results = [ask_result] else: - invalid_sql = sql_valid_result.get("sql") + invalid_sql = generated_sql error_message = ( "SQL generation did not produce a valid SELECT statement." ) @@ -4329,12 +4456,22 @@ async def ask( if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: + generated_sql = valid_generation_result.get("sql") + if generated_sql and self._is_low_signal_sql_for_question( + sql_user_query, generated_sql + ): + invalid_sql = generated_sql + error_message = ( + "SQL correction produced a low-signal result " + "shape for the requested analysis." + ) + break if ask_result := self._build_ask_result_from_sql( - valid_generation_result.get("sql") + generated_sql ): api_results = [ask_result] break - invalid_sql = valid_generation_result.get("sql") + invalid_sql = generated_sql error_message = ( "SQL correction did not produce a valid SELECT statement." ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index ee3cdbc5cb..551ae1dc41 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -468,6 +468,74 @@ def test_monthly_repair_volume_uses_direct_heuristic_route(): ) +def test_operational_fallback_does_not_use_status_for_unit_question_without_unit_column(): + service = AskService(pipelines={}) + tables = service._parse_schema_tables( + [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + created_at TIMESTAMP + ); + """ + ] + ) + + assert ( + service._build_schema_grounded_operational_sql( + "Which business units are underperforming?", + tables, + ) + is None + ) + + +def test_operational_fallback_requires_date_for_trend_question(): + service = AskService(pipelines={}) + tables = service._parse_schema_tables( + [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR + ); + """ + ] + ) + + assert ( + service._build_schema_grounded_operational_sql( + "Generate a line chart showing monthly repair volume for the last 12 months.", + tables, + ) + is None + ) + + +def test_operational_fallback_allows_status_when_status_is_requested(): + service = AskService(pipelines={}) + tables = service._parse_schema_tables( + [ + """ + CREATE TABLE dbo_ticket_cycles ( + id VARCHAR, + status VARCHAR + ); + """ + ] + ) + + sql = service._build_schema_grounded_operational_sql( + "Show a pie chart of open, closed, and in-progress repair tickets.", + tables, + ) + + assert sql + assert '"dbo_ticket_cycles"."status" AS "status"' in sql + assert 'COUNT(*) AS "RecordCount"' in sql + + def test_repair_failure_count_requires_schema_backed_failure_dimension(): service = AskService(pipelines={}) table_ddls = [ @@ -501,6 +569,32 @@ def test_ask_result_validation_requires_select_sql(): assert service._build_ask_result_from_sql(None) is None +def test_low_signal_sql_is_rejected_for_analytical_questions(): + service = AskService(pipelines={}) + + assert service._is_low_signal_sql_for_question( + "Generate a line chart showing monthly repair volume for the last 12 months.", + 'SELECT "dbo_ticket_cycles"."status" AS "status" FROM "dbo_ticket_cycles"', + ) + assert service._is_low_signal_sql_for_question( + "Show Top 20 Sales Accounts by revenue.", + 'SELECT "dbo_refunds"."Refund_Status" AS "Refund_Status" FROM "dbo_refunds"', + ) + + +def test_aggregate_sql_is_allowed_for_analytical_questions(): + service = AskService(pipelines={}) + + assert not service._is_low_signal_sql_for_question( + "Show a pie chart of open, closed, and in-progress repair tickets.", + ( + 'SELECT "dbo_ticket_cycles"."status" AS "status", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ticket_cycles" GROUP BY "dbo_ticket_cycles"."status"' + ), + ) + + def test_retrieval_metadata_ignores_malformed_documents(): service = AskService(pipelines={}) From 52e35b5223e7a9ddc7d27f19aa3d064f0287ab43 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 00:28:18 +0530 Subject: [PATCH 0317/1087] Revert "Reject low-signal SQL for analytical asks" This reverts commit 36f236b56c65b85c91b08f72768f0c97d6d942b5. --- wren-ai-service/src/web/v1/services/ask.py | 163 ++---------------- .../test_ask_heuristic_text_to_sql.py | 94 ---------- 2 files changed, 13 insertions(+), 244 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 14d60a79ef..71bb725adf 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2218,19 +2218,6 @@ def _build_schema_grounded_operational_sql( if not any(term in normalized_query for term in operational_terms): return None - asks_status_dimension = any( - term in normalized_query - for term in ( - "status", - "open", - "closed", - "completed", - "in-progress", - "in progress", - "active", - ) - ) - scored_tables: list[tuple[int, dict[str, Any]]] = [] for table in tables: table_name = str(table.get("name") or "") @@ -2283,25 +2270,16 @@ def _build_schema_grounded_operational_sql( ) dimension_candidates: list[tuple[str, ...]] = [] - asks_unit_dimension = ( - "manufacturing" in normalized_query - or "business unit" in normalized_query - or "business units" in normalized_query - or re.search(r"\bunit(?:s)?\b", normalized_query) is not None - ) - if asks_unit_dimension: + if "manufacturing" in normalized_query or "unit" in normalized_query: dimension_candidates.append( ( "manufacturing_unit", "manufacturing unit", - "business_unit", - "business unit", - "BusinessUnit", - "BU", "unit", "assignee_user_id", "created_by_user_id", "org_id", + "status", ) ) if "component" in normalized_query: @@ -2316,7 +2294,11 @@ def _build_schema_grounded_operational_sql( dimension_candidates.append(("source", "author", "category", "status")) if "priority" in normalized_query: dimension_candidates.append(("priority", "status")) - if asks_status_dimension: + if ( + "status" in normalized_query + or "open" in normalized_query + or "closed" in normalized_query + ): dimension_candidates.append(("status", "priority")) if "assignee" in normalized_query: dimension_candidates.append(("assignee_user_id", "created_by_user_id")) @@ -2327,13 +2309,11 @@ def _build_schema_grounded_operational_sql( if dimension and dimension not in dimensions: dimensions.append(dimension) - if asks_unit_dimension and not dimensions: - return None - if not dimensions: fallback_dimension = self._find_first_schema_column( table, ( + "status", "priority", "category", "subcategory", @@ -2347,11 +2327,6 @@ def _build_schema_grounded_operational_sql( if fallback_dimension: dimensions.append(fallback_dimension) - if not dimensions and asks_status_dimension: - fallback_dimension = self._find_schema_column(table, ("status",)) - if fallback_dimension: - dimensions.append(fallback_dimension) - wants_trend = any( term in normalized_query for term in ("trend", "monthly", "month", "line chart", "over time") @@ -2360,9 +2335,6 @@ def _build_schema_grounded_operational_sql( limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) limit = int(limit_match.group(1)) if limit_match else 10 - if wants_trend and not date_column: - return None - if wants_trend and date_column: date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" select_parts = [ @@ -3081,98 +3053,6 @@ def _is_valid_select_sql(self, sql: Optional[str]) -> bool: return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) - def _extract_top_level_select_expressions(self, sql: str) -> list[str]: - normalized = re.sub(r"\s+", " ", (sql or "").strip()) - select_match = re.search(r"\bSELECT\b", normalized, flags=re.IGNORECASE) - if not select_match: - return [] - - start = select_match.end() - from_match_start: int | None = None - depth = 0 - index = start - while index < len(normalized): - char = normalized[index] - if char == "(": - depth += 1 - elif char == ")" and depth > 0: - depth -= 1 - elif ( - depth == 0 - and normalized[index : index + 6].upper() == " FROM " - ): - from_match_start = index - break - index += 1 - - if from_match_start is None: - return [] - - select_clause = normalized[start:from_match_start].strip() - select_clause = re.sub( - r"^TOP\s+\d+\s+", - "", - select_clause, - flags=re.IGNORECASE, - ) - expressions: list[str] = [] - depth = 0 - expression_start = 0 - for index, char in enumerate(select_clause): - if char == "(": - depth += 1 - elif char == ")" and depth > 0: - depth -= 1 - elif char == "," and depth == 0: - expression = select_clause[expression_start:index].strip() - if expression: - expressions.append(expression) - expression_start = index + 1 - - last_expression = select_clause[expression_start:].strip() - if last_expression: - expressions.append(last_expression) - return expressions - - def _is_low_signal_sql_for_question(self, query: str, sql: str) -> bool: - if not self._is_data_analysis_query(query): - return False - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not any( - term in normalized_query - for term in ( - "chart", - "trend", - "monthly", - "rate", - "revenue", - "sales", - "volume", - "count", - "top", - "highest", - "underperforming", - "distribution", - "group", - "grouped", - " by ", - ) - ): - return False - - normalized_sql = re.sub(r"\s+", " ", (sql or "").strip()) - if re.search( - r"\b(?:COUNT|SUM|AVG|MIN|MAX|DATEPART|DATETRUNC|DATE_TRUNC|" - r"ROW_NUMBER|RANK|DENSE_RANK)\s*\(", - normalized_sql, - flags=re.IGNORECASE, - ): - return False - - expressions = self._extract_top_level_select_expressions(normalized_sql) - return len(expressions) <= 1 - def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: if not self._is_valid_select_sql(sql): return None @@ -4366,19 +4246,12 @@ async def ask( if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - generated_sql = sql_valid_result.get("sql") - if generated_sql and self._is_low_signal_sql_for_question( - sql_user_query, generated_sql + if ask_result := self._build_ask_result_from_sql( + sql_valid_result.get("sql") ): - invalid_sql = generated_sql - error_message = ( - "SQL generation produced a low-signal result shape for " - "the requested analysis." - ) - elif ask_result := self._build_ask_result_from_sql(generated_sql): api_results = [ask_result] else: - invalid_sql = generated_sql + invalid_sql = sql_valid_result.get("sql") error_message = ( "SQL generation did not produce a valid SELECT statement." ) @@ -4456,22 +4329,12 @@ async def ask( if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - generated_sql = valid_generation_result.get("sql") - if generated_sql and self._is_low_signal_sql_for_question( - sql_user_query, generated_sql - ): - invalid_sql = generated_sql - error_message = ( - "SQL correction produced a low-signal result " - "shape for the requested analysis." - ) - break if ask_result := self._build_ask_result_from_sql( - generated_sql + valid_generation_result.get("sql") ): api_results = [ask_result] break - invalid_sql = generated_sql + invalid_sql = valid_generation_result.get("sql") error_message = ( "SQL correction did not produce a valid SELECT statement." ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 551ae1dc41..ee3cdbc5cb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -468,74 +468,6 @@ def test_monthly_repair_volume_uses_direct_heuristic_route(): ) -def test_operational_fallback_does_not_use_status_for_unit_question_without_unit_column(): - service = AskService(pipelines={}) - tables = service._parse_schema_tables( - [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - created_at TIMESTAMP - ); - """ - ] - ) - - assert ( - service._build_schema_grounded_operational_sql( - "Which business units are underperforming?", - tables, - ) - is None - ) - - -def test_operational_fallback_requires_date_for_trend_question(): - service = AskService(pipelines={}) - tables = service._parse_schema_tables( - [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR - ); - """ - ] - ) - - assert ( - service._build_schema_grounded_operational_sql( - "Generate a line chart showing monthly repair volume for the last 12 months.", - tables, - ) - is None - ) - - -def test_operational_fallback_allows_status_when_status_is_requested(): - service = AskService(pipelines={}) - tables = service._parse_schema_tables( - [ - """ - CREATE TABLE dbo_ticket_cycles ( - id VARCHAR, - status VARCHAR - ); - """ - ] - ) - - sql = service._build_schema_grounded_operational_sql( - "Show a pie chart of open, closed, and in-progress repair tickets.", - tables, - ) - - assert sql - assert '"dbo_ticket_cycles"."status" AS "status"' in sql - assert 'COUNT(*) AS "RecordCount"' in sql - - def test_repair_failure_count_requires_schema_backed_failure_dimension(): service = AskService(pipelines={}) table_ddls = [ @@ -569,32 +501,6 @@ def test_ask_result_validation_requires_select_sql(): assert service._build_ask_result_from_sql(None) is None -def test_low_signal_sql_is_rejected_for_analytical_questions(): - service = AskService(pipelines={}) - - assert service._is_low_signal_sql_for_question( - "Generate a line chart showing monthly repair volume for the last 12 months.", - 'SELECT "dbo_ticket_cycles"."status" AS "status" FROM "dbo_ticket_cycles"', - ) - assert service._is_low_signal_sql_for_question( - "Show Top 20 Sales Accounts by revenue.", - 'SELECT "dbo_refunds"."Refund_Status" AS "Refund_Status" FROM "dbo_refunds"', - ) - - -def test_aggregate_sql_is_allowed_for_analytical_questions(): - service = AskService(pipelines={}) - - assert not service._is_low_signal_sql_for_question( - "Show a pie chart of open, closed, and in-progress repair tickets.", - ( - 'SELECT "dbo_ticket_cycles"."status" AS "status", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ticket_cycles" GROUP BY "dbo_ticket_cycles"."status"' - ), - ) - - def test_retrieval_metadata_ignores_malformed_documents(): service = AskService(pipelines={}) From 8583a9f77bc0fefd48f30aa5b622fe45207f4112 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 00:52:19 +0530 Subject: [PATCH 0318/1087] Respect explicit chart requests --- .../src/pipelines/generation/utils/chart.py | 95 +++++++++++++++++-- .../generation/test_chart_generation_utils.py | 48 ++++++++++ 2 files changed, 136 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index a2596887c8..7390f8fac5 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -31,15 +31,70 @@ def _detect_requested_chart_type(query: str | None) -> str: ("bar", ["waterfall", "waterfall chart"]), ("pie", ["pie chart", "donut chart", "doughnut chart"]), ("area", ["area chart", "area graph"]), + ("scatter", ["scatter chart", "scatter plot", "scatter graph"]), ] for chart_type, patterns in checks: if any(pattern in normalized for pattern in patterns): return chart_type - if re.search(r"\b(chart|graph|plot|visuali[sz](?:e|ation)?)\b", normalized): + return "" + + +def _closest_supported_chart_type( + requested_chart_type: str, + quantitative: list[str], + temporal: list[str], + nominal: list[str], +) -> str: + if requested_chart_type != "scatter": + return requested_chart_type + + if len(quantitative) >= 2: + return "bar" + if temporal and quantitative: + return "line" + if nominal and quantitative: return "bar" return "" +def _chart_reasoning( + requested_chart_type: str, + chart_type: str, + chart_schema: dict, + existing_reasoning: str = "", +) -> str: + if not chart_schema: + return existing_reasoning + if requested_chart_type == "scatter" and chart_type: + return ( + "The user requested a scatter chart, but scatter charts are not " + "supported by the current chart type contract. Generated the closest " + f"supported visualization ({chart_type.replace('_', ' ')} chart) " + "from the SQL result columns." + ) + if requested_chart_type and requested_chart_type != chart_type: + return ( + f"The requested {requested_chart_type.replace('_', ' ')} chart was " + "not suitable for the returned SQL result shape. Generated the " + f"closest meaningful {chart_type.replace('_', ' ')} chart from the " + "SQL result columns." + ) + if existing_reasoning: + return existing_reasoning + return "Generated from the SQL result columns and requested chart type." + + +def _chart_type_matches_request( + requested_chart_type: str, + actual_chart_type: str, +) -> bool: + if not requested_chart_type: + return True + if requested_chart_type == "scatter": + return actual_chart_type in {"line", "bar"} + return requested_chart_type == actual_chart_type + + def _safe_column_names(columns: list[Any]) -> list[str]: return [str(column) for column in columns if column is not None and str(column)] @@ -341,7 +396,9 @@ def _fallback_chart_type( temporal: list[str], nominal: list[str], ) -> str: - chart_type = requested_chart_type or "bar" + chart_type = _closest_supported_chart_type( + requested_chart_type or "bar", quantitative, temporal, nominal + ) if chart_type == "pie": return "pie" if nominal else "" @@ -483,8 +540,10 @@ def build_fallback_chart_result( ) -> dict: processed = ChartDataPreprocessor().run(data) sample_data = processed.get("sample_data", []) - chart_type = _detect_requested_chart_type(query) or "bar" - chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) + requested_chart_type = _detect_requested_chart_type(query) + chart_type = requested_chart_type or "bar" + if not requested_chart_type: + chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) chart_schema = _build_fallback_chart_schema(query, chart_type, sample_data) if not chart_schema: return { @@ -501,7 +560,11 @@ def build_fallback_chart_result( return { "chart_schema": chart_schema, - "reasoning": "Generated from the SQL result columns and requested chart type.", + "reasoning": _chart_reasoning( + requested_chart_type, + chart_type, + chart_schema, + ), "chart_type": chart_type, } @@ -908,9 +971,15 @@ def run( chart_schema = _normalize_chart_schema_fields( chart_schema, list(sample_data[0].keys()) if sample_data else [] ) + actual_chart_type = _chart_type_from_schema( + chart_schema, generation_result.get("chart_type", "") + ) if ( not _is_schema_compatible_with_sample_data(chart_schema, sample_data) + or not _chart_type_matches_request( + requested_chart_type, actual_chart_type + ) or _needs_deterministic_bar_fallback( chart_schema, chart_type or "", sample_data ) @@ -919,6 +988,8 @@ def run( query, chart_type or "bar", sample_data ) chart_type = _chart_type_from_schema(chart_schema, chart_type) + else: + chart_type = actual_chart_type if not chart_schema: return { @@ -942,7 +1013,12 @@ def run( return { "results": { "chart_schema": chart_schema, - "reasoning": reasoning, + "reasoning": _chart_reasoning( + requested_chart_type, + chart_type, + chart_schema, + reasoning, + ), "chart_type": chart_type, } } @@ -957,7 +1033,12 @@ def run( return { "results": { "chart_schema": fallback_schema, - "reasoning": reasoning, + "reasoning": _chart_reasoning( + requested_chart_type, + fallback_chart_type, + fallback_schema, + reasoning, + ), "chart_type": fallback_chart_type, } } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index d95099cf3d..0790fc3131 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -61,6 +61,54 @@ def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Market" +def test_explicit_bar_chart_is_not_refined_to_grouped_bar(): + result = build_fallback_chart_result( + "Create a bar chart of new orders by Customer in each Market.", + { + "columns": [ + {"name": "Market"}, + {"name": "Customer"}, + {"name": "OrderCount"}, + ], + "data": [["North", "Acme", 10], ["South", "Globex", 8]], + }, + ) + + assert result["chart_type"] == "bar" + assert result["chart_schema"]["mark"]["type"] == "bar" + assert "xOffset" not in result["chart_schema"]["encoding"] + + +def test_explicit_stacked_bar_chart_is_respected(): + result = build_fallback_chart_result( + "Create a stacked bar chart of sales by Market and Division.", + { + "columns": [ + {"name": "Market"}, + {"name": "Division"}, + {"name": "Sales"}, + ], + "data": [["North", "A", 10], ["North", "B", 8]], + }, + ) + + assert result["chart_type"] == "stacked_bar" + assert result["chart_schema"]["encoding"]["y"]["stack"] == "zero" + + +def test_scatter_request_uses_closest_supported_chart_with_reasoning(): + result = build_fallback_chart_result( + "Create a scatter plot of revenue by market.", + { + "columns": [{"name": "Market"}, {"name": "Revenue"}], + "data": [["North", 100], ["South", 200]], + }, + ) + + assert result["chart_type"] == "bar" + assert "scatter charts are not supported" in result["reasoning"] + + def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): assert not _is_schema_compatible_with_sample_data( { From 9542d1935f6b8aa5d6e6645debec70a3f5cc78ba Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 01:09:11 +0530 Subject: [PATCH 0319/1087] Enforce requested chart types --- .../src/pipelines/generation/utils/chart.py | 42 ++++++- wren-ai-service/src/web/v1/services/chart.py | 4 +- .../generation/test_chart_generation_utils.py | 105 ++++++++++++++++++ 3 files changed, 142 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 7390f8fac5..28a5c592be 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -404,10 +404,10 @@ def _fallback_chart_type( return "pie" if nominal else "" if chart_type in {"line", "area", "multi_line"}: - return chart_type if quantitative and (temporal or nominal) else "" + return chart_type if temporal or nominal or quantitative else "" if chart_type in {"grouped_bar", "stacked_bar"}: - return chart_type if quantitative and len(nominal) > 1 else "" + return chart_type if len(nominal) > 1 else "" return "bar" if nominal or temporal or quantitative else "" @@ -463,9 +463,37 @@ def axis(field: str, field_type: str) -> dict: } if chart_type in {"line", "area", "multi_line"}: - if not measure: - return {} - y_encoding = axis(measure, "quantitative") + y_encoding = axis(measure, "quantitative") if measure else count_axis + if chart_type == "multi_line" and len(quantitative) > 1: + x_field = ( + dimensions[0] + if dimensions + else (temporal[0] if temporal else nominal[0] if nominal else columns[0]) + ) + x_type = "temporal" if x_field in temporal else "ordinal" + return { + "title": title, + "mark": {"type": "line"}, + "transform": [ + { + "fold": quantitative, + "as": ["Metric", "Value"], + } + ], + "encoding": { + "x": axis(x_field, x_type), + "y": { + "field": "Value", + "type": "quantitative", + "title": "Value", + }, + "color": { + "field": "Metric", + "type": "nominal", + "title": "Metric", + }, + }, + } if {"year", "month"}.issubset({str(c).lower() for c in columns}): month_field = next(c for c in columns if str(c).lower() == "month") encoding = { @@ -517,7 +545,9 @@ def axis(field: str, field_type: str) -> dict: if comparison_field: encoding["color"] = axis(comparison_field, "nominal") nominal_dimension_count = len([c for c in dimensions if c in nominal]) - if chart_type == "grouped_bar" or nominal_dimension_count > 1: + if chart_type == "grouped_bar" or ( + not chart_type and nominal_dimension_count > 1 + ): encoding["xOffset"] = axis(comparison_field, "nominal") elif x_field in nominal: encoding["color"] = axis(x_field, "nominal") diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index 3392cad921..f63ad9caa1 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -164,9 +164,7 @@ async def chart( ) chart_result = chart_generation_result["post_process"]["results"] - if not chart_result.get("chart_schema", {}) and not chart_result.get( - "reasoning", "" - ): + if not chart_result.get("chart_schema", {}): self._chart_results[query_id] = ChartResultResponse( status="failed", error=ChartError( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index 0790fc3131..7bb3f70ada 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -79,6 +79,62 @@ def test_explicit_bar_chart_is_not_refined_to_grouped_bar(): assert "xOffset" not in result["chart_schema"]["encoding"] +def test_explicit_line_chart_is_respected(): + result = build_fallback_chart_result( + "Generate a line chart of monthly revenue.", + { + "columns": [{"name": "Month"}, {"name": "Revenue"}], + "data": [["2026-01", 100], ["2026-02", 150]], + }, + ) + + assert result["chart_type"] == "line" + assert result["chart_schema"]["mark"]["type"] == "line" + + +def test_explicit_area_chart_is_respected(): + result = build_fallback_chart_result( + "Generate an area chart of monthly revenue.", + { + "columns": [{"name": "Month"}, {"name": "Revenue"}], + "data": [["2026-01", 100], ["2026-02", 150]], + }, + ) + + assert result["chart_type"] == "area" + assert result["chart_schema"]["mark"]["type"] == "area" + + +def test_explicit_pie_chart_is_respected(): + result = build_fallback_chart_result( + "Show a pie chart of ticket status.", + { + "columns": [{"name": "Status"}, {"name": "RecordCount"}], + "data": [["open", 10], ["closed", 20]], + }, + ) + + assert result["chart_type"] == "pie" + assert result["chart_schema"]["mark"]["type"] == "arc" + + +def test_explicit_grouped_bar_chart_is_respected(): + result = build_fallback_chart_result( + "Create a grouped bar chart of sales by Market and Division.", + { + "columns": [ + {"name": "Market"}, + {"name": "Division"}, + {"name": "Sales"}, + ], + "data": [["North", "A", 10], ["North", "B", 8]], + }, + ) + + assert result["chart_type"] == "grouped_bar" + assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Division" + + def test_explicit_stacked_bar_chart_is_respected(): result = build_fallback_chart_result( "Create a stacked bar chart of sales by Market and Division.", @@ -96,6 +152,24 @@ def test_explicit_stacked_bar_chart_is_respected(): assert result["chart_schema"]["encoding"]["y"]["stack"] == "zero" +def test_explicit_multi_line_chart_is_respected(): + result = build_fallback_chart_result( + "Create a multi-line chart of revenue and cost by month.", + { + "columns": [ + {"name": "Month"}, + {"name": "Revenue"}, + {"name": "Cost"}, + ], + "data": [["2026-01", 100, 60], ["2026-02", 150, 70]], + }, + ) + + assert result["chart_type"] == "multi_line" + assert result["chart_schema"]["mark"]["type"] == "line" + assert result["chart_schema"]["transform"][0]["fold"] == ["Revenue", "Cost"] + + def test_scatter_request_uses_closest_supported_chart_with_reasoning(): result = build_fallback_chart_result( "Create a scatter plot of revenue by market.", @@ -145,3 +219,34 @@ async def run(self, **kwargs): "Generated from the SQL result columns and requested chart type." ) assert result["chart_result"]["chart_type"] == "bar" + + +@pytest.mark.asyncio +async def test_chart_service_rejects_text_only_chart_generation_result(): + class FakeChartGenerationPipeline: + async def run(self, **kwargs): + return { + "post_process": { + "results": { + "chart_schema": {}, + "reasoning": "The data is not suitable for this chart.", + "chart_type": "", + } + } + } + + service = ChartService({"chart_generation": FakeChartGenerationPipeline()}) + request = ChartRequest( + query_id="chart-task", + query="Generate a line chart.", + sql="SELECT NULL AS empty_value", + data={ + "columns": [{"name": "empty_value"}], + "data": [[None]], + }, + ) + + result = await service.chart(request) + + assert result["metadata"]["error_type"] == "NO_CHART" + assert result["chart_result"] == {} From 9aa9ec43d775510610663a0d0793c978b5fc3189 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 01:19:38 +0530 Subject: [PATCH 0320/1087] Revert "Enforce requested chart types" This reverts commit 9542d1935f6b8aa5d6e6645debec70a3f5cc78ba. --- .../src/pipelines/generation/utils/chart.py | 42 +------ wren-ai-service/src/web/v1/services/chart.py | 4 +- .../generation/test_chart_generation_utils.py | 105 ------------------ 3 files changed, 9 insertions(+), 142 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 28a5c592be..7390f8fac5 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -404,10 +404,10 @@ def _fallback_chart_type( return "pie" if nominal else "" if chart_type in {"line", "area", "multi_line"}: - return chart_type if temporal or nominal or quantitative else "" + return chart_type if quantitative and (temporal or nominal) else "" if chart_type in {"grouped_bar", "stacked_bar"}: - return chart_type if len(nominal) > 1 else "" + return chart_type if quantitative and len(nominal) > 1 else "" return "bar" if nominal or temporal or quantitative else "" @@ -463,37 +463,9 @@ def axis(field: str, field_type: str) -> dict: } if chart_type in {"line", "area", "multi_line"}: - y_encoding = axis(measure, "quantitative") if measure else count_axis - if chart_type == "multi_line" and len(quantitative) > 1: - x_field = ( - dimensions[0] - if dimensions - else (temporal[0] if temporal else nominal[0] if nominal else columns[0]) - ) - x_type = "temporal" if x_field in temporal else "ordinal" - return { - "title": title, - "mark": {"type": "line"}, - "transform": [ - { - "fold": quantitative, - "as": ["Metric", "Value"], - } - ], - "encoding": { - "x": axis(x_field, x_type), - "y": { - "field": "Value", - "type": "quantitative", - "title": "Value", - }, - "color": { - "field": "Metric", - "type": "nominal", - "title": "Metric", - }, - }, - } + if not measure: + return {} + y_encoding = axis(measure, "quantitative") if {"year", "month"}.issubset({str(c).lower() for c in columns}): month_field = next(c for c in columns if str(c).lower() == "month") encoding = { @@ -545,9 +517,7 @@ def axis(field: str, field_type: str) -> dict: if comparison_field: encoding["color"] = axis(comparison_field, "nominal") nominal_dimension_count = len([c for c in dimensions if c in nominal]) - if chart_type == "grouped_bar" or ( - not chart_type and nominal_dimension_count > 1 - ): + if chart_type == "grouped_bar" or nominal_dimension_count > 1: encoding["xOffset"] = axis(comparison_field, "nominal") elif x_field in nominal: encoding["color"] = axis(x_field, "nominal") diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index f63ad9caa1..3392cad921 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -164,7 +164,9 @@ async def chart( ) chart_result = chart_generation_result["post_process"]["results"] - if not chart_result.get("chart_schema", {}): + if not chart_result.get("chart_schema", {}) and not chart_result.get( + "reasoning", "" + ): self._chart_results[query_id] = ChartResultResponse( status="failed", error=ChartError( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index 7bb3f70ada..0790fc3131 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -79,62 +79,6 @@ def test_explicit_bar_chart_is_not_refined_to_grouped_bar(): assert "xOffset" not in result["chart_schema"]["encoding"] -def test_explicit_line_chart_is_respected(): - result = build_fallback_chart_result( - "Generate a line chart of monthly revenue.", - { - "columns": [{"name": "Month"}, {"name": "Revenue"}], - "data": [["2026-01", 100], ["2026-02", 150]], - }, - ) - - assert result["chart_type"] == "line" - assert result["chart_schema"]["mark"]["type"] == "line" - - -def test_explicit_area_chart_is_respected(): - result = build_fallback_chart_result( - "Generate an area chart of monthly revenue.", - { - "columns": [{"name": "Month"}, {"name": "Revenue"}], - "data": [["2026-01", 100], ["2026-02", 150]], - }, - ) - - assert result["chart_type"] == "area" - assert result["chart_schema"]["mark"]["type"] == "area" - - -def test_explicit_pie_chart_is_respected(): - result = build_fallback_chart_result( - "Show a pie chart of ticket status.", - { - "columns": [{"name": "Status"}, {"name": "RecordCount"}], - "data": [["open", 10], ["closed", 20]], - }, - ) - - assert result["chart_type"] == "pie" - assert result["chart_schema"]["mark"]["type"] == "arc" - - -def test_explicit_grouped_bar_chart_is_respected(): - result = build_fallback_chart_result( - "Create a grouped bar chart of sales by Market and Division.", - { - "columns": [ - {"name": "Market"}, - {"name": "Division"}, - {"name": "Sales"}, - ], - "data": [["North", "A", 10], ["North", "B", 8]], - }, - ) - - assert result["chart_type"] == "grouped_bar" - assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Division" - - def test_explicit_stacked_bar_chart_is_respected(): result = build_fallback_chart_result( "Create a stacked bar chart of sales by Market and Division.", @@ -152,24 +96,6 @@ def test_explicit_stacked_bar_chart_is_respected(): assert result["chart_schema"]["encoding"]["y"]["stack"] == "zero" -def test_explicit_multi_line_chart_is_respected(): - result = build_fallback_chart_result( - "Create a multi-line chart of revenue and cost by month.", - { - "columns": [ - {"name": "Month"}, - {"name": "Revenue"}, - {"name": "Cost"}, - ], - "data": [["2026-01", 100, 60], ["2026-02", 150, 70]], - }, - ) - - assert result["chart_type"] == "multi_line" - assert result["chart_schema"]["mark"]["type"] == "line" - assert result["chart_schema"]["transform"][0]["fold"] == ["Revenue", "Cost"] - - def test_scatter_request_uses_closest_supported_chart_with_reasoning(): result = build_fallback_chart_result( "Create a scatter plot of revenue by market.", @@ -219,34 +145,3 @@ async def run(self, **kwargs): "Generated from the SQL result columns and requested chart type." ) assert result["chart_result"]["chart_type"] == "bar" - - -@pytest.mark.asyncio -async def test_chart_service_rejects_text_only_chart_generation_result(): - class FakeChartGenerationPipeline: - async def run(self, **kwargs): - return { - "post_process": { - "results": { - "chart_schema": {}, - "reasoning": "The data is not suitable for this chart.", - "chart_type": "", - } - } - } - - service = ChartService({"chart_generation": FakeChartGenerationPipeline()}) - request = ChartRequest( - query_id="chart-task", - query="Generate a line chart.", - sql="SELECT NULL AS empty_value", - data={ - "columns": [{"name": "empty_value"}], - "data": [[None]], - }, - ) - - result = await service.chart(request) - - assert result["metadata"]["error_type"] == "NO_CHART" - assert result["chart_result"] == {} From 0419d19bff9a0f15ae7b995fcabdb0963d4a5b84 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 01:19:38 +0530 Subject: [PATCH 0321/1087] Revert "Respect explicit chart requests" This reverts commit 8583a9f77bc0fefd48f30aa5b622fe45207f4112. --- .../src/pipelines/generation/utils/chart.py | 95 ++----------------- .../generation/test_chart_generation_utils.py | 48 ---------- 2 files changed, 7 insertions(+), 136 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 7390f8fac5..a2596887c8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -31,70 +31,15 @@ def _detect_requested_chart_type(query: str | None) -> str: ("bar", ["waterfall", "waterfall chart"]), ("pie", ["pie chart", "donut chart", "doughnut chart"]), ("area", ["area chart", "area graph"]), - ("scatter", ["scatter chart", "scatter plot", "scatter graph"]), ] for chart_type, patterns in checks: if any(pattern in normalized for pattern in patterns): return chart_type - return "" - - -def _closest_supported_chart_type( - requested_chart_type: str, - quantitative: list[str], - temporal: list[str], - nominal: list[str], -) -> str: - if requested_chart_type != "scatter": - return requested_chart_type - - if len(quantitative) >= 2: - return "bar" - if temporal and quantitative: - return "line" - if nominal and quantitative: + if re.search(r"\b(chart|graph|plot|visuali[sz](?:e|ation)?)\b", normalized): return "bar" return "" -def _chart_reasoning( - requested_chart_type: str, - chart_type: str, - chart_schema: dict, - existing_reasoning: str = "", -) -> str: - if not chart_schema: - return existing_reasoning - if requested_chart_type == "scatter" and chart_type: - return ( - "The user requested a scatter chart, but scatter charts are not " - "supported by the current chart type contract. Generated the closest " - f"supported visualization ({chart_type.replace('_', ' ')} chart) " - "from the SQL result columns." - ) - if requested_chart_type and requested_chart_type != chart_type: - return ( - f"The requested {requested_chart_type.replace('_', ' ')} chart was " - "not suitable for the returned SQL result shape. Generated the " - f"closest meaningful {chart_type.replace('_', ' ')} chart from the " - "SQL result columns." - ) - if existing_reasoning: - return existing_reasoning - return "Generated from the SQL result columns and requested chart type." - - -def _chart_type_matches_request( - requested_chart_type: str, - actual_chart_type: str, -) -> bool: - if not requested_chart_type: - return True - if requested_chart_type == "scatter": - return actual_chart_type in {"line", "bar"} - return requested_chart_type == actual_chart_type - - def _safe_column_names(columns: list[Any]) -> list[str]: return [str(column) for column in columns if column is not None and str(column)] @@ -396,9 +341,7 @@ def _fallback_chart_type( temporal: list[str], nominal: list[str], ) -> str: - chart_type = _closest_supported_chart_type( - requested_chart_type or "bar", quantitative, temporal, nominal - ) + chart_type = requested_chart_type or "bar" if chart_type == "pie": return "pie" if nominal else "" @@ -540,10 +483,8 @@ def build_fallback_chart_result( ) -> dict: processed = ChartDataPreprocessor().run(data) sample_data = processed.get("sample_data", []) - requested_chart_type = _detect_requested_chart_type(query) - chart_type = requested_chart_type or "bar" - if not requested_chart_type: - chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) + chart_type = _detect_requested_chart_type(query) or "bar" + chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) chart_schema = _build_fallback_chart_schema(query, chart_type, sample_data) if not chart_schema: return { @@ -560,11 +501,7 @@ def build_fallback_chart_result( return { "chart_schema": chart_schema, - "reasoning": _chart_reasoning( - requested_chart_type, - chart_type, - chart_schema, - ), + "reasoning": "Generated from the SQL result columns and requested chart type.", "chart_type": chart_type, } @@ -971,15 +908,9 @@ def run( chart_schema = _normalize_chart_schema_fields( chart_schema, list(sample_data[0].keys()) if sample_data else [] ) - actual_chart_type = _chart_type_from_schema( - chart_schema, generation_result.get("chart_type", "") - ) if ( not _is_schema_compatible_with_sample_data(chart_schema, sample_data) - or not _chart_type_matches_request( - requested_chart_type, actual_chart_type - ) or _needs_deterministic_bar_fallback( chart_schema, chart_type or "", sample_data ) @@ -988,8 +919,6 @@ def run( query, chart_type or "bar", sample_data ) chart_type = _chart_type_from_schema(chart_schema, chart_type) - else: - chart_type = actual_chart_type if not chart_schema: return { @@ -1013,12 +942,7 @@ def run( return { "results": { "chart_schema": chart_schema, - "reasoning": _chart_reasoning( - requested_chart_type, - chart_type, - chart_schema, - reasoning, - ), + "reasoning": reasoning, "chart_type": chart_type, } } @@ -1033,12 +957,7 @@ def run( return { "results": { "chart_schema": fallback_schema, - "reasoning": _chart_reasoning( - requested_chart_type, - fallback_chart_type, - fallback_schema, - reasoning, - ), + "reasoning": reasoning, "chart_type": fallback_chart_type, } } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index 0790fc3131..d95099cf3d 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -61,54 +61,6 @@ def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Market" -def test_explicit_bar_chart_is_not_refined_to_grouped_bar(): - result = build_fallback_chart_result( - "Create a bar chart of new orders by Customer in each Market.", - { - "columns": [ - {"name": "Market"}, - {"name": "Customer"}, - {"name": "OrderCount"}, - ], - "data": [["North", "Acme", 10], ["South", "Globex", 8]], - }, - ) - - assert result["chart_type"] == "bar" - assert result["chart_schema"]["mark"]["type"] == "bar" - assert "xOffset" not in result["chart_schema"]["encoding"] - - -def test_explicit_stacked_bar_chart_is_respected(): - result = build_fallback_chart_result( - "Create a stacked bar chart of sales by Market and Division.", - { - "columns": [ - {"name": "Market"}, - {"name": "Division"}, - {"name": "Sales"}, - ], - "data": [["North", "A", 10], ["North", "B", 8]], - }, - ) - - assert result["chart_type"] == "stacked_bar" - assert result["chart_schema"]["encoding"]["y"]["stack"] == "zero" - - -def test_scatter_request_uses_closest_supported_chart_with_reasoning(): - result = build_fallback_chart_result( - "Create a scatter plot of revenue by market.", - { - "columns": [{"name": "Market"}, {"name": "Revenue"}], - "data": [["North", 100], ["South", 200]], - }, - ) - - assert result["chart_type"] == "bar" - assert "scatter charts are not supported" in result["reasoning"] - - def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): assert not _is_schema_compatible_with_sample_data( { From 91ad25fb42fadcbde67dedeb46e7f84bfba1f6e8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 01:29:44 +0530 Subject: [PATCH 0322/1087] Select valid chart axes from results --- .../src/pipelines/generation/utils/chart.py | 84 ++++++++++++++++--- .../generation/test_chart_generation_utils.py | 53 ++++++++++++ 2 files changed, 124 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index a2596887c8..673e8be576 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -85,11 +85,40 @@ def _query_relevant_columns(query: str | None, columns: list[str]) -> list[str]: return [column for _, _, column in sorted(scored_columns)] +def _query_grouping_columns(query: str | None, columns: list[str]) -> set[str]: + normalized_query = str(query or "").lower() + grouping_terms = set() + for match in re.finditer( + r"\b(?:by|per|each|grouped by|group by)\s+([a-zA-Z0-9_ ]+)", + normalized_query, + ): + phrase = match.group(1) + phrase = re.split( + r"\b(?:and|with|over|for|where|order|sort|top|last|using)\b", + phrase, + maxsplit=1, + )[0] + grouping_terms.update(_identifier_tokens(phrase)) + + grouped_columns = set() + for column in columns: + tokens = _identifier_tokens(column) + if tokens and tokens.intersection(grouping_terms): + grouped_columns.add(column) + + return grouped_columns + + def _select_measure_column(query: str | None, quantitative: list[str]) -> str | None: if not quantitative: return None - relevant = _query_relevant_columns(query, quantitative) + grouping_columns = _query_grouping_columns(query, quantitative) + measure_candidates = [ + column for column in quantitative if column not in grouping_columns + ] or quantitative + + relevant = _query_relevant_columns(query, measure_candidates) if relevant: return relevant[0] @@ -109,12 +138,41 @@ def _select_measure_column(query: str | None, quantitative: list[str]) -> str | "failure", "repair", ) - for column in quantitative: + for column in measure_candidates: normalized = str(column).lower() if any(keyword in normalized for keyword in metric_keywords): return column - return quantitative[0] + return measure_candidates[0] + + +def _select_axis_columns( + query: str | None, + chart_type: str, + quantitative: list[str], + temporal: list[str], + nominal: list[str], + columns: list[str], +) -> tuple[str | None, str | None, list[str]]: + dimensions = _select_dimension_columns( + query, chart_type, nominal, temporal, columns + ) + measure = _select_measure_column(query, quantitative) + + if dimensions: + return dimensions[0], measure, dimensions + + # Numeric-only SQL results are still chartable, but x/y must never point to + # the same column. Use the first non-measure numeric column as the category + # axis and the selected measure as y. + if measure and len(quantitative) > 1: + x_field = next( + (column for column in quantitative if column != measure), + quantitative[0], + ) + return x_field, measure, [x_field] + + return (columns[0] if columns else None), measure, dimensions def _count_axis_title(query: str | None) -> str: @@ -347,10 +405,10 @@ def _fallback_chart_type( return "pie" if nominal else "" if chart_type in {"line", "area", "multi_line"}: - return chart_type if quantitative and (temporal or nominal) else "" + return chart_type if temporal or nominal or quantitative else "" if chart_type in {"grouped_bar", "stacked_bar"}: - return chart_type if quantitative and len(nominal) > 1 else "" + return chart_type if len(nominal) > 1 else "" return "bar" if nominal or temporal or quantitative else "" @@ -374,10 +432,9 @@ def _build_fallback_chart_schema( if not chart_type: return {} - dimensions = _select_dimension_columns( - query, chart_type, nominal, temporal, columns + x_field, measure, dimensions = _select_axis_columns( + query, chart_type, quantitative, temporal, nominal, columns ) - measure = _select_measure_column(query, quantitative) title = _humanize_title(query or "Chart") @@ -394,7 +451,7 @@ def axis(field: str, field_type: str) -> dict: } if chart_type == "pie": - color_field = dimensions[0] if dimensions else columns[0] + color_field = x_field or columns[0] theta_axis = axis(measure, "quantitative") if measure else count_axis return { "title": title, @@ -407,8 +464,9 @@ def axis(field: str, field_type: str) -> dict: if chart_type in {"line", "area", "multi_line"}: if not measure: - return {} - y_encoding = axis(measure, "quantitative") + y_encoding = count_axis + else: + y_encoding = axis(measure, "quantitative") if {"year", "month"}.issubset({str(c).lower() for c in columns}): month_field = next(c for c in columns if str(c).lower() == "month") encoding = { @@ -424,7 +482,7 @@ def axis(field: str, field_type: str) -> dict: "encoding": encoding, } - x_field = dimensions[0] if dimensions else columns[0] + x_field = x_field or columns[0] x_type = "temporal" if x_field in temporal else "ordinal" encoding = { "x": axis(x_field, x_type), @@ -442,7 +500,7 @@ def axis(field: str, field_type: str) -> dict: "encoding": encoding, } - x_field = dimensions[0] if dimensions else columns[0] + x_field = x_field or columns[0] x_type = ( "nominal" if x_field in nominal diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py index d95099cf3d..4f80816d21 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -61,6 +61,59 @@ def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Market" +def test_fallback_chart_selects_question_relevant_x_and_numeric_y(): + result = build_fallback_chart_result( + "Show Top 20 Sales Accounts by revenue.", + { + "columns": [ + {"name": "SalesAccount"}, + {"name": "Refund_Status"}, + {"name": "Revenue"}, + ], + "data": [["Acme", "TRANSMITTED", 100], ["Globex", "TRANSMITTED", 200]], + }, + ) + + assert result["chart_schema"]["encoding"]["x"]["field"] == "SalesAccount" + assert result["chart_schema"]["encoding"]["y"]["field"] == "Revenue" + assert result["chart_schema"]["encoding"]["y"]["type"] == "quantitative" + + +def test_fallback_chart_never_uses_same_numeric_field_for_x_and_y(): + result = build_fallback_chart_result( + "Show revenue by discount.", + { + "columns": [{"name": "Discount"}, {"name": "Revenue"}], + "data": [[5, 100], [10, 200]], + }, + ) + + assert result["chart_schema"]["encoding"]["x"]["field"] == "Discount" + assert result["chart_schema"]["encoding"]["y"]["field"] == "Revenue" + assert ( + result["chart_schema"]["encoding"]["x"]["field"] + != result["chart_schema"]["encoding"]["y"]["field"] + ) + + +def test_line_chart_uses_count_y_when_no_numeric_metric_exists(): + result = build_fallback_chart_result( + "Generate a line chart showing monthly repair volume.", + { + "columns": [{"name": "Month"}, {"name": "Status"}], + "data": [["2026-01", "completed"], ["2026-02", "in-progress"]], + }, + ) + + assert result["chart_type"] == "line" + assert result["chart_schema"]["encoding"]["x"]["field"] == "Month" + assert result["chart_schema"]["encoding"]["y"] == { + "aggregate": "count", + "type": "quantitative", + "title": "Repair Count", + } + + def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): assert not _is_schema_compatible_with_sample_data( { From b267ee1d87178965de9479c5e2de83c61e2a5b8c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 14:45:38 +0530 Subject: [PATCH 0323/1087] Add multi-dashboard support --- .../src/apollo/client/graphql/dashboard.ts | 64 ++++++- wren-ui/src/apollo/server/resolvers.ts | 4 + .../server/resolvers/dashboardResolver.ts | 62 ++++++- wren-ui/src/apollo/server/schema.ts | 45 +++-- .../server/services/dashboardService.ts | 99 +++++++++- wren-ui/src/components/modals/DeleteModal.tsx | 7 + .../pages/home/promptThread/ChartAnswer.tsx | 174 ++++++++++++------ wren-ui/src/components/sidebar/Home.tsx | 112 ++++++----- .../components/sidebar/home/DashboardTree.tsx | 128 +++++++++++++ .../src/components/sidebar/home/TreeTitle.tsx | 12 +- wren-ui/src/hooks/useHomeSidebar.tsx | 117 +++++++++++- wren-ui/src/pages/home/dashboard.tsx | 105 ++++++++--- 12 files changed, 754 insertions(+), 175 deletions(-) create mode 100644 wren-ui/src/components/sidebar/home/DashboardTree.tsx diff --git a/wren-ui/src/apollo/client/graphql/dashboard.ts b/wren-ui/src/apollo/client/graphql/dashboard.ts index 1c79838cc7..23d5bb8f7f 100644 --- a/wren-ui/src/apollo/client/graphql/dashboard.ts +++ b/wren-ui/src/apollo/client/graphql/dashboard.ts @@ -77,8 +77,11 @@ export const PREVIEW_ITEM_SQL = gql` `; export const SET_DASHBOARD_SCHEDULE = gql` - mutation SetDashboardSchedule($data: SetDashboardScheduleInput!) { - setDashboardSchedule(data: $data) { + mutation SetDashboardSchedule( + $where: DashboardWhereInput! + $data: SetDashboardScheduleInput! + ) { + setDashboardSchedule(where: $where, data: $data) { id projectId name @@ -91,9 +94,19 @@ export const SET_DASHBOARD_SCHEDULE = gql` } `; +export const DASHBOARDS = gql` + query Dashboards { + dashboards { + id + projectId + name + } + } +`; + export const DASHBOARD = gql` - query Dashboard { - dashboard { + query Dashboard($where: DashboardWhereInput) { + dashboard(where: $where) { id name description @@ -114,3 +127,46 @@ export const DASHBOARD = gql` } ${COMMON_DASHBOARD_ITEM} `; + +export const CREATE_DASHBOARD = gql` + mutation CreateDashboard($data: CreateDashboardInput) { + createDashboard(data: $data) { + id + projectId + name + cacheEnabled + scheduleFrequency + scheduleTimezone + scheduleCron + nextScheduledAt + } + } +`; + +export const UPDATE_DASHBOARD = gql` + mutation UpdateDashboard( + $where: DashboardWhereInput! + $data: UpdateDashboardInput! + ) { + updateDashboard(where: $where, data: $data) { + id + projectId + name + cacheEnabled + scheduleFrequency + scheduleTimezone + scheduleCron + nextScheduledAt + } + } +`; + +export const DELETE_DASHBOARD = gql` + mutation DeleteDashboard($where: DashboardWhereInput!) { + deleteDashboard(where: $where) { + id + projectId + name + } + } +`; diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index e5e6d5a309..96169f0e52 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -67,6 +67,7 @@ const resolvers = { projectResolver.getProjectRecommendationQuestions, // Dashboard + dashboards: dashboardResolver.getDashboards, dashboardItems: dashboardResolver.getDashboardItems, dashboard: dashboardResolver.getDashboard, @@ -167,6 +168,9 @@ const resolvers = { askingResolver.generateProjectRecommendationQuestions, // Dashboard + createDashboard: dashboardResolver.createDashboard, + updateDashboard: dashboardResolver.updateDashboard, + deleteDashboard: dashboardResolver.deleteDashboard, updateDashboardItemLayouts: dashboardResolver.updateDashboardItemLayouts, createDashboardItem: dashboardResolver.createDashboardItem, updateDashboardItem: dashboardResolver.updateDashboardItem, diff --git a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts index addfa8cb8c..eb9eaf2a3a 100644 --- a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts @@ -22,8 +22,12 @@ logger.level = 'debug'; export class DashboardResolver { constructor() { + this.getDashboards = this.getDashboards.bind(this); this.getDashboard = this.getDashboard.bind(this); this.getDashboardItems = this.getDashboardItems.bind(this); + this.createDashboard = this.createDashboard.bind(this); + this.updateDashboard = this.updateDashboard.bind(this); + this.deleteDashboard = this.deleteDashboard.bind(this); this.createDashboardItem = this.createDashboardItem.bind(this); this.updateDashboardItem = this.updateDashboardItem.bind(this); this.deleteDashboardItem = this.deleteDashboardItem.bind(this); @@ -33,10 +37,18 @@ export class DashboardResolver { this.setDashboardSchedule = this.setDashboardSchedule.bind(this); } - public async getDashboard( + public async getDashboards( _root: any, _args: any, ctx: IContext, + ): Promise { + return await ctx.dashboardService.getDashboards(); + } + + public async getDashboard( + _root: any, + args: { where?: { id: number } }, + ctx: IContext, ): Promise< Omit & { schedule: DashboardSchedule; @@ -44,7 +56,7 @@ export class DashboardResolver { nextScheduledAt: string | null; } > { - const dashboard = await ctx.dashboardService.getCurrentDashboard(); + const dashboard = await ctx.dashboardService.getDashboard(args.where?.id); if (!dashboard) { throw new Error('Dashboard not found.'); } @@ -65,21 +77,51 @@ export class DashboardResolver { _args: any, ctx: IContext, ): Promise { - const dashboard = await ctx.dashboardService.getCurrentDashboard(); + const dashboard = await ctx.dashboardService.getDashboard(); if (!dashboard) { throw new Error('Dashboard not found.'); } return await ctx.dashboardService.getDashboardItems(dashboard.id); } + public async createDashboard( + _root: any, + args: { data?: { name?: string } }, + ctx: IContext, + ): Promise { + return await ctx.dashboardService.createDashboard(args.data || {}); + } + + public async updateDashboard( + _root: any, + args: { where: { id: number }; data: { name: string } }, + ctx: IContext, + ): Promise { + return await ctx.dashboardService.updateDashboard(args.where.id, args.data); + } + + public async deleteDashboard( + _root: any, + args: { where: { id: number } }, + ctx: IContext, + ): Promise { + return await ctx.dashboardService.deleteDashboard(args.where.id); + } + public async createDashboardItem( _root: any, - args: { data: { itemType: DashboardItemType; responseId: number } }, + args: { + data: { + itemType: DashboardItemType; + responseId: number; + dashboardId?: number; + }; + }, ctx: IContext, ): Promise { const { responseId } = args.data; const itemType = this.normalizeDashboardItemType(args.data.itemType); - const dashboard = await ctx.dashboardService.getCurrentDashboard(); + const dashboard = await ctx.dashboardService.getDashboard(args.data.dashboardId); const response = await ctx.askingService.getResponse(responseId); if (!response) { @@ -196,7 +238,9 @@ export class DashboardResolver { const { itemId, limit, refresh } = args.data; try { const item = await ctx.dashboardService.getDashboardItem(itemId); - const { cacheEnabled } = await ctx.dashboardService.getCurrentDashboard(); + const { cacheEnabled } = await ctx.dashboardService.getDashboard( + item.dashboardId, + ); const project = await ctx.projectService.getCurrentProject(); const deployment = await ctx.deployService.getLastDeployment(project.id); const mdl = deployment.manifest; @@ -230,17 +274,17 @@ export class DashboardResolver { public async setDashboardSchedule( _root: any, - args: { data: SetDashboardCacheData }, + args: { where: { id: number }; data: SetDashboardCacheData }, ctx: IContext, ): Promise { try { - const dashboard = await ctx.dashboardService.getCurrentDashboard(); + const dashboard = await ctx.dashboardService.getDashboard(args.where.id); if (!dashboard) { throw new Error('Dashboard not found.'); } return await ctx.dashboardService.setDashboardSchedule( - dashboard.id, + args.where.id, args.data, ); } catch (error) { diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 62177089a8..9e7f0bcdb7 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -975,12 +975,25 @@ export const typeDefs = gql` input CreateDashboardItemInput { itemType: DashboardItemType! responseId: Int! + dashboardId: Int } input UpdateDashboardItemInput { displayName: String! } + input DashboardWhereInput { + id: Int! + } + + input CreateDashboardInput { + name: String + } + + input UpdateDashboardInput { + name: String! + } + input ItemLayoutInput { itemId: Int! x: Int! @@ -1277,9 +1290,10 @@ export const typeDefs = gql` getProjectRecommendationQuestions: RecommendedQuestionsTask! instantRecommendedQuestions(taskId: String): RecommendedQuestionsTask! - # Dashboard - dashboardItems: [DashboardItem!]! - dashboard: DetailedDashboard! + # Dashboard + dashboards: [Dashboard!]! + dashboardItems: [DashboardItem!]! + dashboard(where: DashboardWhereInput): DetailedDashboard! # SQL Pairs sqlPairs: [SqlPair]! @@ -1412,18 +1426,27 @@ export const typeDefs = gql` data: InstantRecommendedQuestionsInput! ): Task! - # Dashboard - updateDashboardItemLayouts( - data: UpdateDashboardItemLayoutsInput! - ): [DashboardItem!]! - createDashboardItem(data: CreateDashboardItemInput!): DashboardItem! + # Dashboard + createDashboard(data: CreateDashboardInput): Dashboard! + updateDashboard( + where: DashboardWhereInput! + data: UpdateDashboardInput! + ): Dashboard! + deleteDashboard(where: DashboardWhereInput!): Dashboard + updateDashboardItemLayouts( + data: UpdateDashboardItemLayoutsInput! + ): [DashboardItem!]! + createDashboardItem(data: CreateDashboardItemInput!): DashboardItem! updateDashboardItem( where: DashboardItemWhereInput! data: UpdateDashboardItemInput! ): DashboardItem! - deleteDashboardItem(where: DashboardItemWhereInput!): Boolean! - previewItemSQL(data: PreviewItemSQLInput!): PreviewItemResponse! - setDashboardSchedule(data: SetDashboardScheduleInput!): Dashboard! + deleteDashboardItem(where: DashboardItemWhereInput!): Boolean! + previewItemSQL(data: PreviewItemSQLInput!): PreviewItemResponse! + setDashboardSchedule( + where: DashboardWhereInput! + data: SetDashboardScheduleInput! + ): Dashboard! # SQL Pairs createSqlPair(data: CreateSqlPairInput!): SqlPair! diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 59818eef9f..28811fde24 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -32,15 +32,30 @@ export interface UpdateDashboardItemInput { displayName: string; } +export interface CreateDashboardInput { + name?: string; +} + +export interface UpdateDashboardInput { + name: string; +} + export type UpdateDashboardItemLayouts = (DashboardItemLayout & { itemId: number; })[]; export interface IDashboardService { initDashboard(): Promise; - getCurrentDashboard(): Promise; + getDashboard(dashboardId?: number): Promise; + getDashboards(): Promise; getDashboardItem(dashboardItemId: number): Promise; getDashboardItems(dashboardId: number): Promise; + createDashboard(input: CreateDashboardInput): Promise; + updateDashboard( + dashboardId: number, + input: UpdateDashboardInput, + ): Promise; + deleteDashboard(dashboardId: number): Promise; createDashboardItem(input: CreateDashboardItemInput): Promise; updateDashboardItem( dashboardItemId: number, @@ -131,21 +146,78 @@ export class DashboardService implements IDashboardService { projectId: project.id, }); if (existingDashboard) return existingDashboard; - // only support one dashboard for oss return await this.dashboardRepository.createOne({ name: 'Dashboard', projectId: project.id, }); } - public async getCurrentDashboard(): Promise { + public async getDashboards(): Promise { + const project = await this.projectService.getCurrentProject(); + const dashboards = await this.dashboardRepository.findAllBy({ + projectId: project.id, + }); + if (dashboards.length > 0) { + return dashboards.sort((left, right) => left.id - right.id); + } + + return [await this.initDashboard()]; + } + + public async getDashboard(dashboardId?: number): Promise { + if (!dashboardId) { + const [dashboard] = await this.getDashboards(); + return { ...dashboard }; + } + const project = await this.projectService.getCurrentProject(); const dashboard = await this.dashboardRepository.findOneBy({ + id: dashboardId, projectId: project.id, }); + if (!dashboard) { + throw new Error(`Dashboard with id ${dashboardId} not found.`); + } return { ...dashboard }; } + public async createDashboard(input: CreateDashboardInput): Promise { + const project = await this.projectService.getCurrentProject(); + const dashboards = await this.getDashboards(); + const defaultName = this.getNextDashboardName(dashboards); + const name = this.normalizeDashboardName(input.name || defaultName); + + return await this.dashboardRepository.createOne({ + name, + projectId: project.id, + }); + } + + public async updateDashboard( + dashboardId: number, + input: UpdateDashboardInput, + ): Promise { + await this.getDashboard(dashboardId); + return await this.dashboardRepository.updateOne(dashboardId, { + name: this.normalizeDashboardName(input.name), + }); + } + + public async deleteDashboard(dashboardId: number): Promise { + const dashboards = await this.getDashboards(); + if (dashboards.length <= 1) { + throw new Error('At least one dashboard is required.'); + } + + const dashboard = dashboards.find((item) => item.id === dashboardId); + if (!dashboard) { + throw new Error(`Dashboard with id ${dashboardId} not found.`); + } + + await this.dashboardRepository.deleteOne(dashboardId); + return dashboards.find((item) => item.id !== dashboardId) || null; + } + public async getDashboardItem( dashboardItemId: number, ): Promise { @@ -253,6 +325,27 @@ export class DashboardService implements IDashboardService { return { x, y, w: 3, h: 2 }; } + private normalizeDashboardName(name: string): string { + const normalized = name?.trim(); + if (!normalized) { + throw new Error('Dashboard name is required.'); + } + return normalized; + } + + private getNextDashboardName(dashboards: Dashboard[]): string { + if (dashboards.length === 0) { + return 'Dashboard'; + } + + let index = dashboards.length + 1; + const existingNames = new Set(dashboards.map((dashboard) => dashboard.name)); + while (existingNames.has(`Dashboard ${index}`)) { + index += 1; + } + return `Dashboard ${index}`; + } + protected toUTC(schedule: DashboardSchedule): DashboardSchedule { // If no timezone is specified or it's a custom schedule, return as is if ( diff --git a/wren-ui/src/components/modals/DeleteModal.tsx b/wren-ui/src/components/modals/DeleteModal.tsx index da18ec5f09..13a06fbcfe 100644 --- a/wren-ui/src/components/modals/DeleteModal.tsx +++ b/wren-ui/src/components/modals/DeleteModal.tsx @@ -101,6 +101,13 @@ export const DeleteDashboardItemModal = makeDeleteModal(DefaultDeleteButton, { 'This will be permanently deleted, please confirm you want to delete it.', }); +export const DeleteDashboardModal = makeDeleteModal(DefaultDeleteButton, { + icon: , + itemName: 'dashboard', + content: + 'This will permanently delete this dashboard and its pinned charts, please confirm you want to delete it.', +}); + export const DeleteQuestionSQLPairModal = makeDeleteModal(DefaultDeleteButton, { icon: , itemName: 'question-SQL pair', diff --git a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx index 614873a87d..8bc532f03f 100644 --- a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx @@ -2,7 +2,8 @@ import clsx from 'clsx'; import dynamic from 'next/dynamic'; import styled from 'styled-components'; import { useEffect, useMemo, useState } from 'react'; -import { Alert, Form, Button, Skeleton, Modal, message } from 'antd'; +import { useMutation, useQuery } from '@apollo/client'; +import { Alert, Form, Button, Skeleton, Modal, Select, message } from 'antd'; import { attachLoading } from '@/utils/helper'; import ReloadOutlined from '@ant-design/icons/ReloadOutlined'; import BasicProperties from '@/components/chart/properties/BasicProperties'; @@ -22,7 +23,7 @@ import { getChartSpecFieldTitleMap, getChartSpecOptionValues, } from '@/components/chart/handler'; -import { useCreateDashboardItemMutation } from '@/apollo/client/graphql/dashboard.generated'; +import { CREATE_DASHBOARD_ITEM, DASHBOARDS } from '@/apollo/client/graphql/dashboard'; import usePromptThreadStore from './store'; const Chart = dynamic(() => import('@/components/chart'), { @@ -95,23 +96,36 @@ export default function ChartAnswer(props: AnswerResultProps) { const [regenerating, setRegenerating] = useState(false); const [isEditMode, setIsEditMode] = useState(false); const [newValues, setNewValues] = useState(null); + const [isPinModalOpen, setIsPinModalOpen] = useState(false); + const [selectedDashboardId, setSelectedDashboardId] = useState( + null, + ); const [form] = Form.useForm(); const chartType = Form.useWatch('chartType', form); const { chartDetail } = threadResponse; const { error, status, adjustment } = chartDetail || {}; + const { data: dashboardsResult } = useQuery(DASHBOARDS); + const dashboards = dashboardsResult?.dashboards || []; const [previewData, previewDataResult] = usePreviewDataMutation({ onError: (error) => console.error(error), }); - const [createDashboardItem] = useCreateDashboardItemMutation({ + const [createDashboardItem] = useMutation(CREATE_DASHBOARD_ITEM, { onError: (error) => console.error(error), onCompleted: () => { message.success('Successfully pinned chart to dashboard.'); + setIsPinModalOpen(false); }, }); + useEffect(() => { + if (!selectedDashboardId && dashboards.length > 0) { + setSelectedDashboardId(dashboards[0].id); + } + }, [dashboards, selectedDashboardId]); + // Fetch preview data after the chart task has a terminal schema/result. useEffect(() => { if (!getIsChartFinished(status)) return; @@ -201,18 +215,31 @@ export default function ChartAnswer(props: AnswerResultProps) { return; } - Modal.confirm({ - title: 'Are you sure you want to pin this chart to the dashboard?', - okText: 'Save', - onOk: async () => - await createDashboardItem({ - variables: { - data: { - itemType: dashboardItemType, - responseId: threadResponse.id, - }, - }, - }), + if (dashboards.length === 0) { + message.error('No dashboards are available.'); + return; + } + + setSelectedDashboardId((prev) => prev || dashboards[0].id); + setIsPinModalOpen(true); + }; + + const onConfirmPin = async () => { + const dashboardItemType = chartTypeToDashboardItemType( + chartType || chartDetail?.chartType, + ); + if (!dashboardItemType || !selectedDashboardId) { + return; + } + + await createDashboardItem({ + variables: { + data: { + itemType: dashboardItemType, + responseId: threadResponse.id, + dashboardId: selectedDashboardId, + }, + }, }); }; @@ -263,53 +290,78 @@ export default function ChartAnswer(props: AnswerResultProps) {
{chartDetail?.description} {chartSpec ? ( - - -
-
-
- -
- {isAdjusted && ( -
- - + <> + + + +
+
+
- )} -
- - - - + {isAdjusted && ( +
+ + +
+ )} +
+ + + + + setIsPinModalOpen(false)} + > +
+ Choose which dashboard should receive this chart. +
+ + +
setSelectedDashboardId(value)} - options={dashboards.map((dashboard) => ({ - value: dashboard.id, - label: dashboard.name, - }))} - /> - - + form={form} + initialValues={chartOptionValues} + onFieldsChange={onFormChange} + > +
+
+ +
+ {isAdjusted && ( +
+ + +
+ )} +
+ + + + ) : ( chartRegenerateBtn )} diff --git a/wren-ui/src/components/sidebar/Home.tsx b/wren-ui/src/components/sidebar/Home.tsx index c01ba3bd23..824d435f80 100644 --- a/wren-ui/src/components/sidebar/Home.tsx +++ b/wren-ui/src/components/sidebar/Home.tsx @@ -1,68 +1,58 @@ +import clsx from 'clsx'; import { useEffect } from 'react'; import { useRouter } from 'next/router'; import { useParams } from 'next/navigation'; +import styled from 'styled-components'; import { Path } from '@/utils/enum'; -import { useSidebarTreeState } from './SidebarTree'; -import DashboardTree, { DashboardData } from './home/DashboardTree'; +import FundViewOutlined from '@ant-design/icons/FundViewOutlined'; +import SidebarTree, { + StyledTreeNodeLink, + useSidebarTreeState, +} from './SidebarTree'; import ThreadTree, { ThreadData } from './home/ThreadTree'; export interface Props { data: { - dashboards: DashboardData[]; threads: ThreadData[]; }; - onSelectThread: (selectKeys: string[]) => void; - onSelectDashboard: (selectKeys: string[]) => void; - onDeleteThread: (id: string) => Promise; - onRenameThread: (id: string, newName: string) => Promise; - onCreateDashboard: () => Promise; - onDeleteDashboard: (id: string) => Promise; - onRenameDashboard: (id: string, newName: string) => Promise; + onSelect: (selectKeys) => void; + onDelete: (id: string) => Promise; + onRename: (id: string, newName: string) => Promise; } +export const StyledSidebarTree = styled(SidebarTree)` + .adm-treeNode { + &.adm-treeNode__thread { + padding: 0px 16px 0px 4px !important; + + .ant-tree-title { + flex-grow: 1; + display: inline-flex; + align-items: center; + span:first-child, + .adm-treeTitle__title { + flex-grow: 1; + } + } + } + } +`; + export default function Home(props: Props) { - const { - data, - onSelectThread, - onSelectDashboard, - onRenameThread, - onDeleteThread: deleteThread, - onCreateDashboard, - onDeleteDashboard, - onRenameDashboard, - } = props; + const { data, onSelect, onRename, onDelete } = props; const router = useRouter(); const params = useParams<{ id: string }>(); - const { threads, dashboards } = data; + const { threads } = data; - const { - treeSelectedKeys: threadSelectedKeys, - setTreeSelectedKeys: setThreadSelectedKeys, - } = useSidebarTreeState(); - const { - treeSelectedKeys: dashboardSelectedKeys, - setTreeSelectedKeys: setDashboardSelectedKeys, - } = useSidebarTreeState(); + const { treeSelectedKeys, setTreeSelectedKeys } = useSidebarTreeState(); useEffect(() => { - params?.id && setThreadSelectedKeys([params.id] as string[]); + params?.id && setTreeSelectedKeys([params.id] as string[]); }, [params?.id]); - useEffect(() => { - const selectedDashboardId = router.query.dashboardId; - if (router.pathname !== Path.HomeDashboard) { - setDashboardSelectedKeys([]); - return; - } - - if (typeof selectedDashboardId === 'string') { - setDashboardSelectedKeys([selectedDashboardId]); - } - }, [router.pathname, router.query.dashboardId]); - - const handleDeleteThread = async (threadId: string) => { + const onDeleteThread = async (threadId: string) => { try { - await deleteThread(threadId); + await onDelete(threadId); if (params?.id == threadId) { router.push(Path.Home); } @@ -75,33 +65,27 @@ export default function Home(props: Props) { // prevent deselected if (selectedKeys.length === 0) return; - setThreadSelectedKeys(selectedKeys); - onSelectThread(selectedKeys as string[]); - }; - - const onDashboardSelect = (selectedKeys: React.Key[], _info: any) => { - if (selectedKeys.length === 0) return; - - setDashboardSelectedKeys(selectedKeys); - onSelectDashboard(selectedKeys as string[]); + setTreeSelectedKeys(selectedKeys); + onSelect(selectedKeys); }; return ( <> - + + + Dashboard + ); diff --git a/wren-ui/src/components/sidebar/home/DashboardTree.tsx b/wren-ui/src/components/sidebar/home/DashboardTree.tsx deleted file mode 100644 index d34e26212c..0000000000 --- a/wren-ui/src/components/sidebar/home/DashboardTree.tsx +++ /dev/null @@ -1,128 +0,0 @@ -import { useEffect, useState } from 'react'; -import styled from 'styled-components'; -import { DataNode } from 'antd/lib/tree'; -import PlusOutlined from '@ant-design/icons/PlusOutlined'; -import SidebarTree, { - sidebarCommonStyle, -} from '@/components/sidebar/SidebarTree'; -import { - createTreeGroupNode, - GroupActionButton, -} from '@/components/sidebar/utils'; -import TreeTitle from './TreeTitle'; -import { DeleteDashboardModal } from '@/components/modals/DeleteModal'; - -const StyledSidebarTree = styled(SidebarTree)` - ${sidebarCommonStyle} - - .adm-treeNode { - &.adm-treeNode__dashboard { - padding: 0px 16px 0px 4px !important; - - .ant-tree-title { - flex-grow: 1; - display: inline-flex; - align-items: center; - span:first-child, - .adm-treeTitle__title { - flex-grow: 1; - } - } - } - } -`; - -export interface DashboardData { - id: string; - name: string; -} - -interface Props { - dashboards: DashboardData[]; - selectedKeys: React.Key[]; - onSelect: (selectKeys: React.Key[], info: any) => void; - onRename: (id: string, newName: string) => Promise; - onDeleteDashboard: (id: string) => Promise; - onCreateDashboard: () => Promise; -} - -export default function DashboardTree(props: Props) { - const { - dashboards = [], - selectedKeys, - onSelect, - onRename, - onDeleteDashboard, - onCreateDashboard, - } = props; - - const getDashboardGroupNode = createTreeGroupNode({ - groupName: 'Dashboards', - groupKey: 'dashboards', - appendSlot: ( - - Beta - - ), - actions: [ - { - key: 'new-dashboard', - render: () => ( - } - onClick={() => onCreateDashboard()} - > - New - - ), - }, - ], - }); - - const [tree, setTree] = useState(getDashboardGroupNode()); - - useEffect(() => { - setTree(() => - getDashboardGroupNode({ - quotaUsage: dashboards.length, - children: dashboards.map((dashboard) => { - const nodeKey = dashboard.id; - - return { - className: 'adm-treeNode adm-treeNode__dashboard', - id: nodeKey, - isLeaf: true, - key: nodeKey, - title: ( - ( - - )} - /> - ), - }; - }), - }), - ); - }, [dashboards, onCreateDashboard, onDeleteDashboard, onRename]); - - return ( - - ); -} diff --git a/wren-ui/src/components/sidebar/home/TreeTitle.tsx b/wren-ui/src/components/sidebar/home/TreeTitle.tsx index d0c9e6c76b..5ac77ea9f5 100644 --- a/wren-ui/src/components/sidebar/home/TreeTitle.tsx +++ b/wren-ui/src/components/sidebar/home/TreeTitle.tsx @@ -23,11 +23,10 @@ interface TreeTitleProps { title: string; onDelete?: (id: string) => void; onRename?: (id: string, newName: string) => void; - renderDeleteModal?: (props: { onConfirm: () => void }) => React.ReactNode; } export default function TreeTitle(props: TreeTitleProps) { - const { id, onDelete, onRename, renderDeleteModal } = props; + const { id, onDelete, onRename } = props; const [title, setTitle] = useState(props.title); const [isEditing, setIsEditing] = useState(false); @@ -46,11 +45,6 @@ export default function TreeTitle(props: TreeTitleProps) { onDelete && onDelete(id); }; - const deleteModal = - renderDeleteModal?.({ - onConfirm: () => onDeleteData(id), - }) || onDeleteData(id)} />; - return isEditing ? ( onDeleteData(id)} /> + ), key: MENU_ITEM_KEYS.DELETE, onClick: ({ domEvent }) => { domEvent.stopPropagation(); diff --git a/wren-ui/src/hooks/useHomeSidebar.tsx b/wren-ui/src/hooks/useHomeSidebar.tsx index d46e807e76..c914fa6837 100644 --- a/wren-ui/src/hooks/useHomeSidebar.tsx +++ b/wren-ui/src/hooks/useHomeSidebar.tsx @@ -1,19 +1,11 @@ import { useMemo } from 'react'; import { useRouter } from 'next/router'; -import { useMutation, useQuery } from '@apollo/client'; -import { message } from 'antd'; import { Path } from '@/utils/enum'; import { useDeleteThreadMutation, useThreadsQuery, useUpdateThreadMutation, } from '@/apollo/client/graphql/home.generated'; -import { - CREATE_DASHBOARD, - DASHBOARDS, - DELETE_DASHBOARD, - UPDATE_DASHBOARD, -} from '@/apollo/client/graphql/dashboard'; export default function useHomeSidebar() { const router = useRouter(); @@ -21,13 +13,6 @@ export default function useHomeSidebar() { fetchPolicy: 'network-only', nextFetchPolicy: 'network-only', }); - const { data: dashboardData, refetch: refetchDashboards } = useQuery( - DASHBOARDS, - { - fetchPolicy: 'network-only', - nextFetchPolicy: 'network-only', - }, - ); const [updateThread] = useUpdateThreadMutation({ onError: (error) => console.error(error), }); @@ -44,115 +29,27 @@ export default function useHomeSidebar() { [data], ); - const dashboards = useMemo( - () => - (dashboardData?.dashboards || []).map((dashboard) => ({ - id: dashboard.id.toString(), - name: dashboard.name, - })), - [dashboardData], - ); - - const [createDashboard] = useMutation(CREATE_DASHBOARD, { - onError: (error) => { - console.error(error); - message.error('Failed to create dashboard.'); - }, - }); - const [updateDashboard] = useMutation(UPDATE_DASHBOARD, { - onError: (error) => { - console.error(error); - message.error('Failed to rename dashboard.'); - }, - }); - const [deleteDashboard] = useMutation(DELETE_DASHBOARD, { - onError: (error) => { - console.error(error); - message.error('Failed to delete dashboard.'); - }, - }); - - const onSelectThread = (selectKeys: string[]) => { + const onSelect = (selectKeys: string[]) => { router.push(`${Path.Home}/${selectKeys[0]}`); }; - const onSelectDashboard = (selectKeys: string[]) => { - router.push({ - pathname: Path.HomeDashboard, - query: { dashboardId: selectKeys[0] }, - }); - }; - - const onRenameThread = async (id: string, newName: string) => { + const onRename = async (id: string, newName: string) => { await updateThread({ variables: { where: { id: Number(id) }, data: { summary: newName } }, }); refetch(); }; - const onDeleteThread = async (id: string) => { + const onDelete = async (id) => { await deleteThread({ variables: { where: { id: Number(id) } } }); refetch(); }; - const onCreateDashboard = async () => { - const result = await createDashboard({ - variables: { data: {} }, - }); - await refetchDashboards(); - - const dashboardId = result.data?.createDashboard?.id; - if (dashboardId) { - message.success('Successfully created dashboard.'); - router.push({ - pathname: Path.HomeDashboard, - query: { dashboardId }, - }); - } - }; - - const onRenameDashboard = async (id: string, newName: string) => { - await updateDashboard({ - variables: { where: { id: Number(id) }, data: { name: newName } }, - }); - refetchDashboards(); - message.success('Successfully renamed dashboard.'); - }; - - const onDeleteDashboard = async (id: string) => { - const result = await deleteDashboard({ - variables: { where: { id: Number(id) } }, - }); - const nextDashboardId = result.data?.deleteDashboard?.id; - await refetchDashboards(); - message.success('Successfully deleted dashboard.'); - - const currentDashboardId = - typeof router.query.dashboardId === 'string' - ? Number(router.query.dashboardId) - : undefined; - - if (router.pathname === Path.HomeDashboard && currentDashboardId === Number(id)) { - if (nextDashboardId) { - router.push({ - pathname: Path.HomeDashboard, - query: { dashboardId: nextDashboardId }, - }); - } else { - router.push(Path.Home); - } - } - }; - return { - data: { dashboards, threads }, - onSelectThread, - onSelectDashboard, - onRenameThread, - onDeleteThread, - onCreateDashboard, - onRenameDashboard, - onDeleteDashboard, + data: { threads }, + onSelect, + onRename, + onDelete, refetch, }; } diff --git a/wren-ui/src/pages/home/dashboard.tsx b/wren-ui/src/pages/home/dashboard.tsx index b4414c5723..eda1cc9ee2 100644 --- a/wren-ui/src/pages/home/dashboard.tsx +++ b/wren-ui/src/pages/home/dashboard.tsx @@ -1,5 +1,4 @@ -import { useEffect, useMemo, useRef } from 'react'; -import { useMutation, useQuery } from '@apollo/client'; +import { useMemo, useRef } from 'react'; import { message } from 'antd'; import { Path } from '@/utils/enum'; import { useRouter } from 'next/router'; @@ -14,12 +13,11 @@ import CacheSettingsDrawer, { Schedule, } from '@/components/pages/home/dashboardGrid/CacheSettingsDrawer'; import { - DASHBOARD, - DASHBOARDS, - DELETE_DASHBOARD_ITEM, - SET_DASHBOARD_SCHEDULE, - UPDATE_DASHBOARD_ITEM_LAYOUTS, -} from '@/apollo/client/graphql/dashboard'; + useDashboardQuery, + useDeleteDashboardItemMutation, + useUpdateDashboardItemLayoutsMutation, + useSetDashboardScheduleMutation, +} from '@/apollo/client/graphql/dashboard.generated'; import { useGetSettingsQuery } from '@/apollo/client/graphql/settings.generated'; import { DataSource, @@ -36,10 +34,6 @@ const isSupportCachedSettings = (dataSource: DataSource) => { export default function Dashboard() { const router = useRouter(); - const selectedDashboardId = - typeof router.query.dashboardId === 'string' - ? Number(router.query.dashboardId) - : undefined; const dashboardGridRef = useRef<{ onRefreshAll: () => void }>(null); const homeSidebar = useHomeSidebar(); const cacheSettingsDrawer = useDrawerAction(); @@ -50,10 +44,11 @@ export default function Dashboard() { [settings?.dataSource], ); - const { data, loading, client } = useQuery(DASHBOARD, { - variables: selectedDashboardId - ? { where: { id: selectedDashboardId } } - : undefined, + const { + data, + loading, + updateQuery: updateDashboardQuery, + } = useDashboardQuery({ fetchPolicy: 'cache-and-network', onError: () => { message.error('Failed to fetch dashboard items.'); @@ -65,31 +60,20 @@ export default function Dashboard() { [data?.dashboard?.items], ); - const [setDashboardSchedule] = useMutation(SET_DASHBOARD_SCHEDULE, { - refetchQueries: [ - { - query: DASHBOARD, - variables: selectedDashboardId - ? { where: { id: selectedDashboardId } } - : undefined, - }, - { query: DASHBOARDS }, - ], + const [setDashboardSchedule] = useSetDashboardScheduleMutation({ + refetchQueries: ['Dashboard'], onCompleted: () => { message.success('Successfully updated dashboard schedule.'); }, onError: (error) => console.error(error), }); - const [updateDashboardItemLayouts] = useMutation( - UPDATE_DASHBOARD_ITEM_LAYOUTS, - { - onError: () => { - message.error('Failed to update dashboard item layouts.'); - }, + const [updateDashboardItemLayouts] = useUpdateDashboardItemLayoutsMutation({ + onError: () => { + message.error('Failed to update dashboard item layouts.'); }, - ); - const [deleteDashboardItem] = useMutation(DELETE_DASHBOARD_ITEM, { + }); + const [deleteDashboardItem] = useDeleteDashboardItemMutation({ onError: (error) => console.error(error), onCompleted: (_, query) => { message.success('Successfully deleted dashboard item.'); @@ -98,25 +82,15 @@ export default function Dashboard() { }); const onRemoveDashboardItemFromQueryCache = (id: number) => { - client.cache.updateQuery( - { - query: DASHBOARD, - variables: selectedDashboardId - ? { where: { id: selectedDashboardId } } - : undefined, - }, - (prev: any) => { - if (!prev?.dashboard) return prev; - return { - ...prev, - dashboard: { - ...prev.dashboard, - items: - prev?.dashboard?.items?.filter((item) => item.id !== id) || [], - }, - }; - }, - ); + updateDashboardQuery((prev) => { + return { + ...prev, + dashboard: { + ...prev.dashboard, + items: prev?.dashboard?.items?.filter((item) => item.id !== id) || [], + }, + }; + }); }; const onUpdateChange = async (layouts: ItemLayoutInput[]) => { @@ -129,23 +103,6 @@ export default function Dashboard() { await deleteDashboardItem({ variables: { where: { id } } }); }; - const dashboardId = data?.dashboard?.id; - - useEffect(() => { - if (!dashboardId || selectedDashboardId || router.pathname !== Path.HomeDashboard) { - return; - } - - router.replace( - { - pathname: Path.HomeDashboard, - query: { dashboardId }, - }, - undefined, - { shallow: true }, - ); - }, [dashboardId, router.pathname, selectedDashboardId]); - return ( @@ -178,13 +135,7 @@ export default function Dashboard() { {...cacheSettingsDrawer.state} onClose={cacheSettingsDrawer.closeDrawer} onSubmit={async (values) => { - if (!dashboardId) return; - await setDashboardSchedule({ - variables: { - where: { id: dashboardId }, - data: values, - }, - }); + await setDashboardSchedule({ variables: { data: values } }); }} /> )} From effb64591dfdbbabfc422eee086d5895b732ba6b Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 15:24:33 +0530 Subject: [PATCH 0325/1087] Harden dashboard chart pinning --- .../repositories/dashboardItemRepository.ts | 1 + .../server/resolvers/dashboardResolver.ts | 203 ++++++++++++++---- .../server/services/dashboardService.ts | 2 + .../pages/home/promptThread/ChartAnswer.tsx | 7 +- 4 files changed, 168 insertions(+), 45 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index 564264a94c..0971685238 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -35,6 +35,7 @@ export interface DashboardItemLayout { export interface DashboardItemDetail { sql: string; chartSchema?: Record; + previewDataSnapshot?: Record[]; } export interface DashboardItem { diff --git a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts index addfa8cb8c..81cb27402a 100644 --- a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts @@ -16,9 +16,12 @@ import { DashboardSchedule, PreviewItemResponse, } from '@server/models/dashboard'; +import { Manifest } from '@server/mdl/type'; const logger = getLogger('DashboardResolver'); logger.level = 'debug'; +const DASHBOARD_SNAPSHOT_TIMEOUT_MS = 5000; +const DASHBOARD_PREVIEW_TIMEOUT_MS = 15000; export class DashboardResolver { constructor() { @@ -93,34 +96,25 @@ export class DashboardResolver { `Chart schema not found in thread response. responseId: ${responseId}`, ); } + if (!response.sql) { + throw new Error(`Chart SQL not found in thread response. responseId: ${responseId}`); + } + + const previewDataSnapshot = await this.capturePreviewSnapshot( + ctx, + response.sql, + ); const dashboardItem = await ctx.dashboardService.createDashboardItem({ dashboardId: dashboard.id, type: itemType, sql: response.sql, chartSchema: response.chartDetail?.chartSchema, + previewDataSnapshot, }); - // Warm dashboard cache after persisting the item. Cache warm-up failures should - // not prevent a valid chart from being pinned. - const project = await ctx.projectService.getCurrentProject(); - const deployment = await ctx.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; - try { - await ctx.queryService.preview(response.sql, { - project, - manifest: mdl, - limit: DEFAULT_PREVIEW_LIMIT, - cacheEnabled: true, - refresh: true, - }); - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : String(error); - logger.warn( - `Dashboard item ${dashboardItem.id} was pinned but cache warm-up failed: ${errorMessage}`, - ); - } + // Warm dashboard cache asynchronously so datasource latency does not block pinning. + void this.warmDashboardCache(ctx, response.sql, dashboardItem.id); return dashboardItem; } @@ -198,30 +192,40 @@ export class DashboardResolver { const item = await ctx.dashboardService.getDashboardItem(itemId); const { cacheEnabled } = await ctx.dashboardService.getCurrentDashboard(); const project = await ctx.projectService.getCurrentProject(); - const deployment = await ctx.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; - const data = (await ctx.queryService.preview(item.detail.sql, { - project, - manifest: mdl, - limit: limit || DEFAULT_PREVIEW_LIMIT, - cacheEnabled, - refresh: refresh || false, - })) as PreviewDataResponse; - - // handle data to [{ column1: value1, column2: value2, ... }] - const values = data.data.map((val) => { - return data.columns.reduce((acc, col, index) => { - acc[col.name] = val[index]; - return acc; - }, {}); - }); - return { - cacheHit: data.cacheHit || false, - cacheCreatedAt: data.cacheCreatedAt || null, - cacheOverrodeAt: data.cacheOverrodeAt || null, - override: data.override || false, - data: values, - } as PreviewItemResponse; + const manifest = await this.getPreviewManifest(ctx, project.id); + + try { + const data = (await this.withTimeout( + ctx.queryService.preview(item.detail.sql, { + project, + manifest, + limit: limit || DEFAULT_PREVIEW_LIMIT, + cacheEnabled, + refresh: refresh || false, + }), + DASHBOARD_PREVIEW_TIMEOUT_MS, + `Dashboard preview timed out for item ${itemId}`, + )) as PreviewDataResponse; + + return this.formatPreviewItemResponse(data); + } catch (error) { + const snapshot = item.detail.previewDataSnapshot; + if (snapshot?.length) { + logger.warn( + `Using stored dashboard preview snapshot for item ${itemId}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return { + cacheHit: false, + cacheCreatedAt: null, + cacheOverrodeAt: null, + override: false, + data: snapshot, + } as PreviewItemResponse; + } + throw error; + } } catch (error) { logger.error(`Error previewing SQL item ${itemId}: ${error}`); throw error; @@ -248,4 +252,115 @@ export class DashboardResolver { throw error; } } + + private async capturePreviewSnapshot( + ctx: IContext, + sql: string, + ): Promise[] | undefined> { + try { + const project = await ctx.projectService.getCurrentProject(); + const manifest = await this.getPreviewManifest(ctx, project.id); + const data = (await this.withTimeout( + ctx.queryService.preview(sql, { + project, + manifest, + limit: DEFAULT_PREVIEW_LIMIT, + cacheEnabled: false, + refresh: false, + }), + DASHBOARD_SNAPSHOT_TIMEOUT_MS, + 'Dashboard snapshot preview timed out', + )) as PreviewDataResponse; + + return this.formatPreviewItemResponse(data).data; + } catch (error) { + logger.warn( + `Failed to capture dashboard preview snapshot: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return undefined; + } + } + + private async warmDashboardCache( + ctx: IContext, + sql: string, + dashboardItemId: number, + ): Promise { + try { + const project = await ctx.projectService.getCurrentProject(); + const manifest = await this.getPreviewManifest(ctx, project.id); + await this.withTimeout( + ctx.queryService.preview(sql, { + project, + manifest, + limit: DEFAULT_PREVIEW_LIMIT, + cacheEnabled: true, + refresh: true, + }), + DASHBOARD_PREVIEW_TIMEOUT_MS, + `Dashboard cache warm-up timed out for item ${dashboardItemId}`, + ); + } catch (error) { + logger.warn( + `Dashboard item ${dashboardItemId} was pinned but cache warm-up failed: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } + } + + private async getPreviewManifest( + ctx: IContext, + projectId: number, + ): Promise { + const deployment = await ctx.deployService.getLastDeployment(projectId); + if (deployment?.manifest) { + return deployment.manifest as Manifest; + } + + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + return manifest; + } + + private formatPreviewItemResponse(data: PreviewDataResponse): PreviewItemResponse { + const values = data.data.map((val) => { + return data.columns.reduce((acc, col, index) => { + acc[col.name] = val[index]; + return acc; + }, {}); + }); + + return { + cacheHit: data.cacheHit || false, + cacheCreatedAt: data.cacheCreatedAt || null, + cacheOverrodeAt: data.cacheOverrodeAt || null, + override: data.override || false, + data: values, + } as PreviewItemResponse; + } + + private async withTimeout( + promise: Promise, + timeoutMs: number, + timeoutMessage: string, + ): Promise { + let timeoutHandle: ReturnType | null = null; + try { + return await Promise.race([ + promise, + new Promise((_resolve, reject) => { + timeoutHandle = setTimeout( + () => reject(new Error(timeoutMessage)), + timeoutMs, + ); + }), + ]); + } finally { + if (timeoutHandle) { + clearTimeout(timeoutHandle); + } + } + } } diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 59818eef9f..25fee67332 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -26,6 +26,7 @@ export interface CreateDashboardItemInput { type: DashboardItemType; sql: string; chartSchema: DashboardItemDetail['chartSchema']; + previewDataSnapshot?: DashboardItemDetail['previewDataSnapshot']; } export interface UpdateDashboardItemInput { @@ -176,6 +177,7 @@ export class DashboardService implements IDashboardService { detail: { sql: input.sql, chartSchema: input.chartSchema, + previewDataSnapshot: input.previewDataSnapshot, }, layout, }); diff --git a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx index 614873a87d..e5ef14bb01 100644 --- a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx @@ -106,7 +106,12 @@ export default function ChartAnswer(props: AnswerResultProps) { }); const [createDashboardItem] = useCreateDashboardItemMutation({ - onError: (error) => console.error(error), + onError: (error) => { + console.error(error); + message.error( + error.message || 'Failed to pin chart to dashboard. Please try again.', + ); + }, onCompleted: () => { message.success('Successfully pinned chart to dashboard.'); }, From 4b304cd77ca943f2139c2110a343103c46921a86 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 15:31:31 +0530 Subject: [PATCH 0326/1087] Fix MSSQL dashboard item ID generation --- wren-ui/src/apollo/server/repositories/baseRepository.ts | 8 ++++---- .../apollo/server/repositories/dashboardItemRepository.ts | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index 3ebbfcdf77..d47517de37 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -272,10 +272,10 @@ export class BaseRepository implements IBasicRepository { } private async getNextId(executer: Knex | Knex.Transaction) { - const [row] = await executer(this.tableName).max<{ maxId: number | null }>( - 'id as maxId', - ); - return (row?.maxId || 0) + 1; + const [row] = await executer(this.tableName).max<{ + maxId: number | string | null; + }>('id as maxId'); + return Number(row?.maxId || 0) + 1; } private async hasIdentityId(executer: Knex | Knex.Transaction) { diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index 0971685238..d0b333166e 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -246,9 +246,9 @@ export class DashboardItemRepository } private async getNextId(executer: Knex | Knex.Transaction) { - const [row] = await executer(this.tableName).max<{ maxId: number | null }>( - 'id as maxId', - ); - return (row?.maxId || 0) + 1; + const [row] = await executer(this.tableName).max<{ + maxId: number | string | null; + }>('id as maxId'); + return Number(row?.maxId || 0) + 1; } } From d0c05071d7dd2b7134f9b1c045c25dc347fe7c8c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 15:43:19 +0530 Subject: [PATCH 0327/1087] Fix MSSQL dashboard item insert strategy --- .../repositories/dashboardItemRepository.ts | 45 ++++++++++++++++--- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index d0b333166e..4bedbaa8bb 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -72,10 +72,19 @@ export class DashboardItemRepository data: Partial, queryOptions?: IQueryOptions, ): Promise { - return await super.createOne( - await this.normalizeWriteData(data, queryOptions, true), - queryOptions, - ); + const normalized = await this.normalizeWriteData(data, queryOptions, true); + try { + return await super.createOne(normalized, queryOptions); + } catch (error) { + if (!this.shouldRetryManualId(error, normalized)) { + throw error; + } + + return await super.createOne( + await this.normalizeWriteData(data, queryOptions, true, true), + queryOptions, + ); + } } public override async updateOne( @@ -129,6 +138,7 @@ export class DashboardItemRepository data: Partial, queryOptions?: IQueryOptions, includeGeneratedId = false, + forceManualId = false, ): Promise> { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const [ @@ -165,7 +175,7 @@ export class DashboardItemRepository hasIdColumn && normalizedData.id === undefined && this.isMssqlLike(executer) && - !(await this.hasIdentityId(executer)) + (forceManualId || !(await this.hasIdentityId(executer))) ) { normalizedData.id = await this.getNextId(executer); } @@ -251,4 +261,29 @@ export class DashboardItemRepository }>('id as maxId'); return Number(row?.maxId || 0) + 1; } + + private shouldRetryManualId( + error: unknown, + data: Partial, + ): boolean { + if (!this.isMssqlLike(this.knex)) { + return false; + } + + if (data.id !== undefined && data.id !== null) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return ( + message.includes("Cannot insert the value NULL into column 'id'") || + message.includes("Cannot insert explicit value for identity column") + ); + } } From 94a58bdc256f541514ec279eb0e948a030e88507 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 15:50:50 +0530 Subject: [PATCH 0328/1087] Stop forcing dashboard item IDs on first MSSQL insert --- .../src/apollo/server/repositories/dashboardItemRepository.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index 4bedbaa8bb..1fe6bdb396 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -175,7 +175,7 @@ export class DashboardItemRepository hasIdColumn && normalizedData.id === undefined && this.isMssqlLike(executer) && - (forceManualId || !(await this.hasIdentityId(executer))) + forceManualId ) { normalizedData.id = await this.getNextId(executer); } From 7dfcfb32b010561893ea4e5e62a4e8a24aeb251e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 16:17:13 +0530 Subject: [PATCH 0329/1087] Normalize MSSQL BIGINT repository bindings --- .../server/repositories/baseRepository.ts | 169 +++++++++++++++--- .../repositories/dashboardItemRepository.ts | 66 +------ .../repositories/dashboardRepository.ts | 4 +- 3 files changed, 147 insertions(+), 92 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index d47517de37..af323841a7 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -80,7 +80,7 @@ export class BaseRepository implements IBasicRepository { public async findOneBy(filter: Partial, queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const query = executer(this.tableName).where( - this.transformToDBData(filter), + this.normalizeMssqlBindings(this.transformToDBData(filter)), ); if (queryOptions?.limit) { query.limit(queryOptions.limit); @@ -96,7 +96,7 @@ export class BaseRepository implements IBasicRepository { // format filter keys to snake_case const query = executer(this.tableName).where( - this.transformToDBData(filter), + this.normalizeMssqlBindings(this.transformToDBData(filter)), ); if (queryOptions?.order) { query.orderBy(queryOptions.order); @@ -120,34 +120,47 @@ export class BaseRepository implements IBasicRepository { public async createOne(data: Partial, queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const insertValue = await this.prepareInsertData(data, executer); - const [result] = await executer(this.tableName) - .insert(insertValue) - .returning('*'); - return this.transformFromDBData(result); + try { + const insertValue = await this.prepareInsertData(data, executer); + const [result] = await executer(this.tableName) + .insert(this.normalizeMssqlBindings(insertValue)) + .returning('*'); + return this.transformFromDBData(result); + } catch (error) { + if (!this.shouldRetryManualId(error, data, executer)) { + throw error; + } + + const insertValue = await this.prepareInsertData(data, executer, true); + const [result] = await executer(this.tableName) + .insert(this.normalizeMssqlBindings(insertValue)) + .returning('*'); + return this.transformFromDBData(result); + } } public async createMany(data: Partial[], queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const preparedData = await this.prepareInsertManyData(data, executer); + let preparedData: any[]; + try { + preparedData = await this.prepareInsertManyData(data, executer); + } catch (error) { + throw error; + } if (preparedData.length === 0) { return []; } - const batchSize = this.getCreateManyBatchSize(preparedData); - const batchCount = Math.ceil(preparedData.length / batchSize); - const result = []; - for (let i = 0; i < batchCount; i++) { - const start = i * batchSize; - const end = Math.min((i + 1) * batchSize, preparedData.length); - const batchValues = preparedData.slice(start, end); - const chunk = await executer(this.tableName) - .insert(batchValues) - .returning('*'); - result.push(...chunk); - } + try { + return await this.insertMany(executer, preparedData); + } catch (error) { + if (!this.shouldRetryManualId(error, data, executer)) { + throw error; + } - return result.map((data) => this.transformFromDBData(data)); + preparedData = await this.prepareInsertManyData(data, executer, true); + return await this.insertMany(executer, preparedData); + } } public async updateOne( @@ -156,16 +169,21 @@ export class BaseRepository implements IBasicRepository { queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const normalizedId = this.normalizeMssqlBindings({ id }).id; const [result] = await executer(this.tableName) - .where({ id }) - .update(this.transformToDBData(data)) + .where({ id: normalizedId }) + .update(this.normalizeMssqlBindings(this.transformToDBData(data))) .returning('*'); return this.transformFromDBData(result); } public async deleteOne(id: string, queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer.from(this.tableName).where({ id }).delete(); + const normalizedId = this.normalizeMssqlBindings({ id }).id; + const builder = executer + .from(this.tableName) + .where({ id: normalizedId }) + .delete(); return await builder; } @@ -174,8 +192,9 @@ export class BaseRepository implements IBasicRepository { queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const normalizedIds = this.normalizeMssqlBindings(ids); let deleted = 0; - for (const batch of this.toWhereInBatches(ids)) { + for (const batch of this.toWhereInBatches(normalizedIds)) { deleted += await executer .from(this.tableName) .whereIn('id', batch) @@ -190,7 +209,7 @@ export class BaseRepository implements IBasicRepository { ) => { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const builder = executer(this.tableName) - .where(this.transformToDBData(where)) + .where(this.normalizeMssqlBindings(this.transformToDBData(where))) .delete(); return await builder; }; @@ -304,6 +323,7 @@ export class BaseRepository implements IBasicRepository { private async prepareInsertData( data: Partial, executer: Knex | Knex.Transaction, + forceManualId = false, ) { const dbData = this.transformToDBData(data); if (!this.isMssql(executer)) { @@ -313,7 +333,7 @@ export class BaseRepository implements IBasicRepository { if ( !(await this.hasIdColumn(executer)) || dbData.id !== undefined || - (await this.hasIdentityId(executer)) + !forceManualId ) { return dbData; } @@ -327,12 +347,13 @@ export class BaseRepository implements IBasicRepository { private async prepareInsertManyData( data: Partial[], executer: Knex | Knex.Transaction, + forceManualId = false, ) { const dbData = data.map((item) => this.transformToDBData(item)); if ( !this.isMssql(executer) || !(await this.hasIdColumn(executer)) || - (await this.hasIdentityId(executer)) + !forceManualId ) { return dbData; } @@ -358,4 +379,96 @@ export class BaseRepository implements IBasicRepository { return dbData; } + + private async insertMany( + executer: Knex | Knex.Transaction, + preparedData: any[], + ) { + const batchSize = this.getCreateManyBatchSize(preparedData); + const batchCount = Math.ceil(preparedData.length / batchSize); + const result = []; + for (let i = 0; i < batchCount; i++) { + const start = i * batchSize; + const end = Math.min((i + 1) * batchSize, preparedData.length); + const batchValues = preparedData.slice(start, end); + const chunk = await executer(this.tableName) + .insert(this.normalizeMssqlBindings(batchValues)) + .returning('*'); + result.push(...chunk); + } + + return result.map((data) => this.transformFromDBData(data)); + } + + private normalizeMssqlBindings(value: any): any { + if (!this.isMssql(this.knex)) { + return value; + } + + if (Array.isArray(value)) { + return value.map((item) => this.normalizeMssqlBindings(item)); + } + + if (typeof value === 'bigint') { + return value.toString(); + } + + if (!isPlainObject(value)) { + if ( + typeof value === 'number' && + Number.isInteger(value) && + !Number.isSafeInteger(value) + ) { + return String(value); + } + return value; + } + + return Object.fromEntries( + Object.entries(value).map(([key, entryValue]) => { + if ( + typeof entryValue === 'number' && + Number.isInteger(entryValue) && + (!Number.isSafeInteger(entryValue) || + key === 'id' || + key.endsWith('_id')) + ) { + return [key, String(entryValue)]; + } + + if (typeof entryValue === 'bigint') { + return [key, entryValue.toString()]; + } + + return [key, this.normalizeMssqlBindings(entryValue)]; + }), + ); + } + + private shouldRetryManualId( + error: unknown, + data: Partial | Partial[], + executer: Knex | Knex.Transaction, + ) { + if (!this.isMssql(executer)) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item: any) => item?.id === undefined || item?.id === null) + : (data as any)?.id === undefined || (data as any)?.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + } } diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index 1fe6bdb396..b3694ed3e2 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -59,7 +59,6 @@ export class DashboardItemRepository { private readonly jsonbColumns = ['layout', 'detail']; private hasIdColumnCache: boolean | null = null; - private hasIdentityIdPromise?: Promise; private hasTitleColumnCache: boolean | null = null; private hasDisplayNameColumnCache: boolean | null = null; private columnCache = new Map(); @@ -72,19 +71,10 @@ export class DashboardItemRepository data: Partial, queryOptions?: IQueryOptions, ): Promise { - const normalized = await this.normalizeWriteData(data, queryOptions, true); - try { - return await super.createOne(normalized, queryOptions); - } catch (error) { - if (!this.shouldRetryManualId(error, normalized)) { - throw error; - } - - return await super.createOne( - await this.normalizeWriteData(data, queryOptions, true, true), - queryOptions, - ); - } + return await super.createOne( + await this.normalizeWriteData(data, queryOptions, true), + queryOptions, + ); } public override async updateOne( @@ -232,58 +222,10 @@ export class DashboardItemRepository ); } - private async hasIdentityId(executer: Knex | Knex.Transaction) { - if (!this.isMssqlLike(executer)) { - return true; - } - - if (!this.hasIdentityIdPromise) { - this.hasIdentityIdPromise = executer('INFORMATION_SCHEMA.COLUMNS') - .select('COLUMN_NAME') - .where({ - TABLE_SCHEMA: 'dbo', - TABLE_NAME: this.tableName, - COLUMN_NAME: 'id', - }) - .whereRaw( - "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", - ) - .first() - .then(Boolean); - } - - return this.hasIdentityIdPromise; - } - private async getNextId(executer: Knex | Knex.Transaction) { const [row] = await executer(this.tableName).max<{ maxId: number | string | null; }>('id as maxId'); return Number(row?.maxId || 0) + 1; } - - private shouldRetryManualId( - error: unknown, - data: Partial, - ): boolean { - if (!this.isMssqlLike(this.knex)) { - return false; - } - - if (data.id !== undefined && data.id !== null) { - return false; - } - - const message = - error instanceof Error - ? error.message - : typeof error === 'string' - ? error - : ''; - - return ( - message.includes("Cannot insert the value NULL into column 'id'") || - message.includes("Cannot insert explicit value for identity column") - ); - } } diff --git a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts index 503036b7c5..92a9b865ef 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts @@ -142,7 +142,7 @@ export class DashboardRepository if ( (data.id !== undefined && data.id !== null) || !this.isMssql() || - (!forceManualId && (await this.hasIdentityId())) + !forceManualId ) { return data; } @@ -165,7 +165,7 @@ export class DashboardRepository if ( data.every((item) => item.id !== undefined && item.id !== null) || !this.isMssql() || - (!forceManualId && (await this.hasIdentityId())) + !forceManualId ) { return data; } From ef6521b900e7fce38bad61f2b050ba31b5de510f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 16:38:40 +0530 Subject: [PATCH 0330/1087] Handle dashboard bigint IDs end-to-end --- .../src/apollo/client/graphql/__types__.ts | 19 +++++------ .../client/graphql/dashboard.generated.ts | 16 +++++----- .../dashboardCacheBackgroundTracker.ts | 2 +- wren-ui/src/apollo/server/models/dashboard.ts | 4 +-- .../server/repositories/baseRepository.ts | 20 +++++++++--- .../repositories/dashboardItemRepository.ts | 11 ++++--- .../repositories/dashboardRepository.ts | 30 +++++++++++++---- wren-ui/src/apollo/server/resolvers.ts | 3 +- .../server/resolvers/dashboardResolver.ts | 10 +++--- wren-ui/src/apollo/server/scalars.ts | 32 +++++++++++++++++++ wren-ui/src/apollo/server/schema.ts | 19 +++++------ .../server/services/dashboardService.ts | 28 ++++++++-------- .../pages/home/dashboardGrid/index.tsx | 10 +++--- wren-ui/src/pages/home/dashboard.tsx | 4 +-- 14 files changed, 139 insertions(+), 69 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 48960a07db..72135acbd4 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -11,6 +11,7 @@ export type Scalars = { Boolean: boolean; Int: number; Float: number; + BigIntString: string; DialectSQL: any; JSON: any; }; @@ -288,10 +289,10 @@ export type CustomFieldInput = { export type Dashboard = { __typename?: 'Dashboard'; cacheEnabled: Scalars['Boolean']; - id: Scalars['Int']; + id: Scalars['BigIntString']; name: Scalars['String']; nextScheduledAt?: Maybe; - projectId: Scalars['Int']; + projectId: Scalars['BigIntString']; scheduleCron?: Maybe; scheduleFrequency?: Maybe; scheduleTimezone?: Maybe; @@ -299,10 +300,10 @@ export type Dashboard = { export type DashboardItem = { __typename?: 'DashboardItem'; - dashboardId: Scalars['Int']; + dashboardId: Scalars['BigIntString']; detail: DashboardItemDetail; displayName?: Maybe; - id: Scalars['Int']; + id: Scalars['BigIntString']; layout: DashboardItemLayout; type: DashboardItemType; }; @@ -334,7 +335,7 @@ export enum DashboardItemType { } export type DashboardItemWhereInput = { - id: Scalars['Int']; + id: Scalars['BigIntString']; }; export type DashboardSchedule = { @@ -381,7 +382,7 @@ export enum DatabricksConnectionType { } export type DeleteDashboardItemInput = { - itemId: Scalars['Int']; + itemId: Scalars['BigIntString']; }; export type DetailStep = { @@ -436,7 +437,7 @@ export type DetailedDashboard = { __typename?: 'DetailedDashboard'; cacheEnabled: Scalars['Boolean']; description?: Maybe; - id: Scalars['Int']; + id: Scalars['BigIntString']; items: Array; name: Scalars['String']; nextScheduledAt?: Maybe; @@ -664,7 +665,7 @@ export type InstructionWhereInput = { export type ItemLayoutInput = { h: Scalars['Int']; - itemId: Scalars['Int']; + itemId: Scalars['BigIntString']; w: Scalars['Int']; x: Scalars['Int']; y: Scalars['Int']; @@ -1137,7 +1138,7 @@ export type PreviewItemResponse = { }; export type PreviewItemSqlInput = { - itemId: Scalars['Int']; + itemId: Scalars['BigIntString']; limit?: InputMaybe; refresh?: InputMaybe; }; diff --git a/wren-ui/src/apollo/client/graphql/dashboard.generated.ts b/wren-ui/src/apollo/client/graphql/dashboard.generated.ts index 1014a11ee8..67bde922c3 100644 --- a/wren-ui/src/apollo/client/graphql/dashboard.generated.ts +++ b/wren-ui/src/apollo/client/graphql/dashboard.generated.ts @@ -3,19 +3,19 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type CommonDashboardItemFragment = { __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }; +export type CommonDashboardItemFragment = { __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }; export type DashboardItemsQueryVariables = Types.Exact<{ [key: string]: never; }>; -export type DashboardItemsQuery = { __typename?: 'Query', dashboardItems: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; +export type DashboardItemsQuery = { __typename?: 'Query', dashboardItems: Array<{ __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; export type CreateDashboardItemMutationVariables = Types.Exact<{ data: Types.CreateDashboardItemInput; }>; -export type CreateDashboardItemMutation = { __typename?: 'Mutation', createDashboardItem: { __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; +export type CreateDashboardItemMutation = { __typename?: 'Mutation', createDashboardItem: { __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; export type UpdateDashboardItemMutationVariables = Types.Exact<{ where: Types.DashboardItemWhereInput; @@ -23,14 +23,14 @@ export type UpdateDashboardItemMutationVariables = Types.Exact<{ }>; -export type UpdateDashboardItemMutation = { __typename?: 'Mutation', updateDashboardItem: { __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; +export type UpdateDashboardItemMutation = { __typename?: 'Mutation', updateDashboardItem: { __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; export type UpdateDashboardItemLayoutsMutationVariables = Types.Exact<{ data: Types.UpdateDashboardItemLayoutsInput; }>; -export type UpdateDashboardItemLayoutsMutation = { __typename?: 'Mutation', updateDashboardItemLayouts: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; +export type UpdateDashboardItemLayoutsMutation = { __typename?: 'Mutation', updateDashboardItemLayouts: Array<{ __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; export type DeleteDashboardItemMutationVariables = Types.Exact<{ where: Types.DashboardItemWhereInput; @@ -51,12 +51,12 @@ export type SetDashboardScheduleMutationVariables = Types.Exact<{ }>; -export type SetDashboardScheduleMutation = { __typename?: 'Mutation', setDashboardSchedule: { __typename?: 'Dashboard', id: number, projectId: number, name: string, cacheEnabled: boolean, scheduleFrequency?: Types.ScheduleFrequencyEnum | null, scheduleTimezone?: string | null, scheduleCron?: string | null, nextScheduledAt?: string | null } }; +export type SetDashboardScheduleMutation = { __typename?: 'Mutation', setDashboardSchedule: { __typename?: 'Dashboard', id: string, projectId: string, name: string, cacheEnabled: boolean, scheduleFrequency?: Types.ScheduleFrequencyEnum | null, scheduleTimezone?: string | null, scheduleCron?: string | null, nextScheduledAt?: string | null } }; export type DashboardQueryVariables = Types.Exact<{ [key: string]: never; }>; -export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: number, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule?: { __typename?: 'DashboardSchedule', frequency?: Types.ScheduleFrequencyEnum | null, hour?: number | null, minute?: number | null, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null } | null, items: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } }; +export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: string, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule?: { __typename?: 'DashboardSchedule', frequency?: Types.ScheduleFrequencyEnum | null, hour?: number | null, minute?: number | null, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null } | null, items: Array<{ __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } }; export const CommonDashboardItemFragmentDoc = gql` fragment CommonDashboardItem on DashboardItem { @@ -366,4 +366,4 @@ export function useDashboardLazyQuery(baseOptions?: Apollo.LazyQueryHookOptions< } export type DashboardQueryHookResult = ReturnType; export type DashboardLazyQueryHookResult = ReturnType; -export type DashboardQueryResult = Apollo.QueryResult; \ No newline at end of file +export type DashboardQueryResult = Apollo.QueryResult; diff --git a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts index 3f30fa5b86..472897ec07 100644 --- a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts @@ -24,7 +24,7 @@ export class DashboardCacheBackgroundTracker { private projectService: IProjectService; private deployService: IDeployService; private queryService: IQueryService; - private runningJobs = new Set(); + private runningJobs = new Set(); private intervalId?: NodeJS.Timeout; constructor({ diff --git a/wren-ui/src/apollo/server/models/dashboard.ts b/wren-ui/src/apollo/server/models/dashboard.ts index d4c7e9d57b..f7da453868 100644 --- a/wren-ui/src/apollo/server/models/dashboard.ts +++ b/wren-ui/src/apollo/server/models/dashboard.ts @@ -40,8 +40,8 @@ export interface SetDashboardCacheData { } export interface DetailedDashboard { - id: number; - projectId: number; + id: string | number; + projectId: string | number; name: string; cacheEnabled: boolean; scheduleFrequency: ScheduleFrequencyEnum | null; diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index af323841a7..ab3be09621 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -290,11 +290,23 @@ export class BaseRepository implements IBasicRepository { return hasIdColumn; } + private serializeIdValue(value: number | string | bigint) { + const bigintValue = BigInt(value); + return bigintValue <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(bigintValue) + : bigintValue.toString(); + } + + private toNextIdValue(maxId: number | string | bigint | null | undefined) { + const nextId = BigInt(maxId ?? 0) + 1n; + return this.serializeIdValue(nextId); + } + private async getNextId(executer: Knex | Knex.Transaction) { const [row] = await executer(this.tableName).max<{ - maxId: number | string | null; + maxId: number | string | bigint | null; }>('id as maxId'); - return Number(row?.maxId || 0) + 1; + return this.toNextIdValue(row?.maxId); } private async hasIdentityId(executer: Knex | Knex.Transaction) { @@ -369,11 +381,11 @@ export class BaseRepository implements IBasicRepository { return dbData; } - let nextId = await this.getNextId(executer); + let nextId = BigInt((await this.getNextId(executer)) ?? 0); for (const index of missingIdIndexes) { dbData[index] = { ...dbData[index], - id: nextId++, + id: this.serializeIdValue(nextId++), }; } diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index b3694ed3e2..fb9a5b9a73 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -39,8 +39,8 @@ export interface DashboardItemDetail { } export interface DashboardItem { - id: number; - dashboardId: number; + id: string | number; + dashboardId: string | number; type: DashboardItemType; layout: DashboardItemLayout; detail: DashboardItemDetail; @@ -224,8 +224,11 @@ export class DashboardItemRepository private async getNextId(executer: Knex | Knex.Transaction) { const [row] = await executer(this.tableName).max<{ - maxId: number | string | null; + maxId: number | string | bigint | null; }>('id as maxId'); - return Number(row?.maxId || 0) + 1; + const nextId = BigInt(row?.maxId ?? 0) + 1n; + return nextId <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(nextId) + : nextId.toString(); } } diff --git a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts index 92a9b865ef..86c0abaf45 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts @@ -8,8 +8,8 @@ import { import { ScheduleFrequencyEnum } from '@server/models/dashboard'; export interface Dashboard { - id: number; - projectId: number; + id: string | number; + projectId: string | number; name: string; cacheEnabled: boolean; scheduleFrequency: ScheduleFrequencyEnum | null; @@ -148,12 +148,14 @@ export class DashboardRepository } const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ maxId: 'id', }); return { ...data, - id: Number(row?.maxId || 0) + 1, + id: this.toNextIdValue(row?.maxId), }; }; @@ -171,21 +173,35 @@ export class DashboardRepository } const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ maxId: 'id', }); - let nextId = Number(row?.maxId || 0) + 1; + let nextId = BigInt(row?.maxId ?? 0) + 1n; return data.map((item) => { if (item.id !== undefined && item.id !== null) { return item; } return { ...item, - id: nextId++, + id: this.serializeIdValue(nextId++), }; }); }; + private toNextIdValue = (maxId: number | string | bigint | null | undefined) => { + const nextId = BigInt(maxId ?? 0) + 1n; + return this.serializeIdValue(nextId); + }; + + private serializeIdValue = (value: number | string | bigint) => { + const bigintValue = BigInt(value); + return bigintValue <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(bigintValue) + : bigintValue.toString(); + }; + private shouldRetryManualId = ( error: unknown, data: Partial | Partial[], diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index e5e6d5a309..8381952531 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -10,7 +10,7 @@ import { InstructionResolver } from './resolvers/instructionResolver'; import { ApiHistoryResolver } from './resolvers/apiHistoryResolver'; import { RbacResolver } from './resolvers/rbacResolver'; import { convertColumnType } from '@server/utils'; -import { DialectSQLScalar } from './scalars'; +import { BigIntStringScalar, DialectSQLScalar } from './scalars'; const projectResolver = new ProjectResolver(); const modelResolver = new ModelResolver(); @@ -25,6 +25,7 @@ const rbacResolver = new RbacResolver(); const resolvers = { JSON: GraphQLJSON, DialectSQL: DialectSQLScalar, + BigIntString: BigIntStringScalar, Query: { listDataSourceTables: projectResolver.listDataSourceTables, autoGenerateRelation: projectResolver.autoGenerateRelation, diff --git a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts index 81cb27402a..460c803bff 100644 --- a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts @@ -145,7 +145,7 @@ export class DashboardResolver { public async updateDashboardItem( _root: any, - args: { where: { id: number }; data: { displayName: string } }, + args: { where: { id: string | number }; data: { displayName: string } }, ctx: IContext, ): Promise { const { id } = args.where; @@ -159,7 +159,7 @@ export class DashboardResolver { public async deleteDashboardItem( _root: any, - args: { where: { id: number } }, + args: { where: { id: string | number } }, ctx: IContext, ): Promise { const { id } = args.where; @@ -184,7 +184,9 @@ export class DashboardResolver { public async previewItemSQL( _root: any, - args: { data: { itemId: number; limit?: number; refresh?: boolean } }, + args: { + data: { itemId: string | number; limit?: number; refresh?: boolean }; + }, ctx: IContext, ): Promise { const { itemId, limit, refresh } = args.data; @@ -286,7 +288,7 @@ export class DashboardResolver { private async warmDashboardCache( ctx: IContext, sql: string, - dashboardItemId: number, + dashboardItemId: string | number, ): Promise { try { const project = await ctx.projectService.getCurrentProject(); diff --git a/wren-ui/src/apollo/server/scalars.ts b/wren-ui/src/apollo/server/scalars.ts index d137671730..3c8be92bdc 100644 --- a/wren-ui/src/apollo/server/scalars.ts +++ b/wren-ui/src/apollo/server/scalars.ts @@ -23,3 +23,35 @@ export const DialectSQLScalar = new GraphQLScalarType({ return ast.value as DialectSQL; }, }); + +export const BigIntStringScalar = new GraphQLScalarType({ + name: 'BigIntString', + description: + 'A bigint-compatible scalar serialized as a string to preserve precision across GraphQL and MSSQL.', + serialize(value: unknown): string { + if ( + typeof value !== 'string' && + typeof value !== 'number' && + typeof value !== 'bigint' + ) { + throw new Error('BigIntString must be a string, number, or bigint'); + } + return String(value); + }, + parseValue(value: unknown): string { + if ( + typeof value !== 'string' && + typeof value !== 'number' && + typeof value !== 'bigint' + ) { + throw new Error('BigIntString must be a string, number, or bigint'); + } + return String(value); + }, + parseLiteral(ast: any): string { + if (ast.kind !== 'StringValue' && ast.kind !== 'IntValue') { + throw new Error('BigIntString must be a string or int literal'); + } + return ast.value; + }, +}); diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 62177089a8..3caeed2d25 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -3,6 +3,7 @@ import { gql } from 'apollo-server-micro'; export const typeDefs = gql` scalar JSON scalar DialectSQL + scalar BigIntString enum ApiType { GENERATE_SQL @@ -969,7 +970,7 @@ export const typeDefs = gql` } input DashboardItemWhereInput { - id: Int! + id: BigIntString! } input CreateDashboardItemInput { @@ -982,7 +983,7 @@ export const typeDefs = gql` } input ItemLayoutInput { - itemId: Int! + itemId: BigIntString! x: Int! y: Int! w: Int! @@ -994,11 +995,11 @@ export const typeDefs = gql` } input DeleteDashboardItemInput { - itemId: Int! + itemId: BigIntString! } input PreviewItemSQLInput { - itemId: Int! + itemId: BigIntString! limit: Int refresh: Boolean = false } @@ -1064,8 +1065,8 @@ export const typeDefs = gql` } type DashboardItem { - id: Int! - dashboardId: Int! + id: BigIntString! + dashboardId: BigIntString! type: DashboardItemType! layout: DashboardItemLayout! detail: DashboardItemDetail! @@ -1073,8 +1074,8 @@ export const typeDefs = gql` } type Dashboard { - id: Int! - projectId: Int! + id: BigIntString! + projectId: BigIntString! name: String! cacheEnabled: Boolean! scheduleFrequency: ScheduleFrequencyEnum @@ -1084,7 +1085,7 @@ export const typeDefs = gql` } type DetailedDashboard { - id: Int! + id: BigIntString! name: String! description: String cacheEnabled: Boolean! diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 25fee67332..d26ba42440 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -22,7 +22,7 @@ const logger = getLogger('DashboardService'); logger.level = 'debug'; export interface CreateDashboardItemInput { - dashboardId: number; + dashboardId: string | number; type: DashboardItemType; sql: string; chartSchema: DashboardItemDetail['chartSchema']; @@ -34,25 +34,25 @@ export interface UpdateDashboardItemInput { } export type UpdateDashboardItemLayouts = (DashboardItemLayout & { - itemId: number; + itemId: string | number; })[]; export interface IDashboardService { initDashboard(): Promise; getCurrentDashboard(): Promise; - getDashboardItem(dashboardItemId: number): Promise; - getDashboardItems(dashboardId: number): Promise; + getDashboardItem(dashboardItemId: string | number): Promise; + getDashboardItems(dashboardId: string | number): Promise; createDashboardItem(input: CreateDashboardItemInput): Promise; updateDashboardItem( - dashboardItemId: number, + dashboardItemId: string | number, input: UpdateDashboardItemInput, ): Promise; - deleteDashboardItem(dashboardItemId: number): Promise; + deleteDashboardItem(dashboardItemId: string | number): Promise; updateDashboardItemLayouts( layouts: UpdateDashboardItemLayouts, ): Promise; setDashboardSchedule( - dashboardId: number, + dashboardId: string | number, data: SetDashboardCacheData, ): Promise; parseCronExpression(dashboard: Dashboard): DashboardSchedule; @@ -78,7 +78,7 @@ export class DashboardService implements IDashboardService { } public async setDashboardSchedule( - dashboardId: number, + dashboardId: string | number, data: SetDashboardCacheData, ): Promise { try { @@ -148,7 +148,7 @@ export class DashboardService implements IDashboardService { } public async getDashboardItem( - dashboardItemId: number, + dashboardItemId: string | number, ): Promise { const item = await this.dashboardItemRepository.findOneBy({ id: dashboardItemId, @@ -160,7 +160,7 @@ export class DashboardService implements IDashboardService { } public async getDashboardItems( - dashboardId: number, + dashboardId: string | number, ): Promise { return await this.dashboardItemRepository.findAllBy({ dashboardId, @@ -184,7 +184,7 @@ export class DashboardService implements IDashboardService { } public async updateDashboardItem( - dashboardItemId: number, + dashboardItemId: string | number, input: UpdateDashboardItemInput, ): Promise { return await this.dashboardItemRepository.updateOne(dashboardItemId, { @@ -226,13 +226,15 @@ export class DashboardService implements IDashboardService { return updatedItems; } - public async deleteDashboardItem(dashboardItemId: number): Promise { + public async deleteDashboardItem( + dashboardItemId: string | number, + ): Promise { await this.dashboardItemRepository.deleteOne(dashboardItemId); return true; } private async calculateNewLayout( - dashboardId: number, + dashboardId: string | number, ): Promise { const dashboardItems = await this.dashboardItemRepository.findAllBy({ dashboardId, diff --git a/wren-ui/src/components/pages/home/dashboardGrid/index.tsx b/wren-ui/src/components/pages/home/dashboardGrid/index.tsx index e27b4c6b00..b320be8342 100644 --- a/wren-ui/src/components/pages/home/dashboardGrid/index.tsx +++ b/wren-ui/src/components/pages/home/dashboardGrid/index.tsx @@ -129,7 +129,7 @@ const getLayoutToGrid = (item: DashboardItem) => { const getLayoutToUpdateItem = (layout: Layout) => { return { - itemId: Number(layout.i), + itemId: layout.i, x: layout.x, y: layout.y, w: layout.w, @@ -141,7 +141,7 @@ interface Props { items: DashboardItem[]; isSupportCached: boolean; onUpdateChange: (layouts: ItemLayoutInput[]) => void; - onDelete: (id: number) => Promise; + onDelete: (id: string) => Promise; } const DashboardGrid = forwardRef( @@ -240,7 +240,7 @@ const DashboardGrid = forwardRef( export default DashboardGrid; -const PinnedItemTitle = (props: { id: number; title: string }) => { +const PinnedItemTitle = (props: { id: string; title: string }) => { const { title } = props; const [form] = Form.useForm(); @@ -248,7 +248,7 @@ const PinnedItemTitle = (props: { id: number; title: string }) => { onError: (error) => console.error(error), }); - const handleSave = (dashboardItemId: number, values: { title: string }) => { + const handleSave = (dashboardItemId: string, values: { title: string }) => { if (values.title === title) return; updateDashboardItem({ variables: { @@ -280,7 +280,7 @@ const PinnedItem = forwardRef( props: { item: DashboardItem; isSupportCached: boolean; - onDelete: (id: number) => Promise; + onDelete: (id: string) => Promise; }, ref: React.RefObject<{ onRefresh: () => void }>, ) => { diff --git a/wren-ui/src/pages/home/dashboard.tsx b/wren-ui/src/pages/home/dashboard.tsx index eda1cc9ee2..b6cc935c45 100644 --- a/wren-ui/src/pages/home/dashboard.tsx +++ b/wren-ui/src/pages/home/dashboard.tsx @@ -81,7 +81,7 @@ export default function Dashboard() { }, }); - const onRemoveDashboardItemFromQueryCache = (id: number) => { + const onRemoveDashboardItemFromQueryCache = (id: string) => { updateDashboardQuery((prev) => { return { ...prev, @@ -99,7 +99,7 @@ export default function Dashboard() { } }; - const onDelete = async (id: number) => { + const onDelete = async (id: string) => { await deleteDashboardItem({ variables: { where: { id } } }); }; From 9f96d08d8e264db73747541cf1cc8c90999d5e85 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 16:59:39 +0530 Subject: [PATCH 0331/1087] Auto-initialize dashboards for current projects --- .../server/services/dashboardService.ts | 9 ++++- .../services/tests/dashboardService.test.ts | 35 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index d26ba42440..4366d9a150 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -144,7 +144,14 @@ export class DashboardService implements IDashboardService { const dashboard = await this.dashboardRepository.findOneBy({ projectId: project.id, }); - return { ...dashboard }; + if (dashboard) { + return dashboard; + } + + logger.debug( + `Dashboard not found for project ${project.id}; initializing a default dashboard.`, + ); + return await this.initDashboard(); } public async getDashboardItem( diff --git a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts index 55814d5610..7ca88841b2 100644 --- a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts @@ -85,6 +85,41 @@ describe('DashboardService', () => { }); }); + describe('dashboard initialization', () => { + it('should return the existing dashboard for the current project', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + + await expect(dashboardService.getCurrentDashboard()).resolves.toEqual( + dashboard, + ); + expect(mockDashboardRepository.createOne).not.toHaveBeenCalled(); + }); + + it('should initialize a dashboard when the current project has none', async () => { + const project = { id: 42 }; + const createdDashboard = { + id: '1002', + projectId: 42, + name: 'Dashboard', + }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValueOnce(null); + mockDashboardRepository.findOneBy.mockResolvedValueOnce(null); + mockDashboardRepository.createOne.mockResolvedValue(createdDashboard); + + await expect(dashboardService.getCurrentDashboard()).resolves.toEqual( + createdDashboard, + ); + expect(mockDashboardRepository.createOne).toHaveBeenCalledWith({ + name: 'Dashboard', + projectId: project.id, + }); + }); + }); + describe('generateCronExpression', () => { it('should generate correct cron expression for daily schedule', () => { const schedule = { From 05606f94640d1d2d55d4d8ba96bcfc23e4dcac25 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 17:26:42 +0530 Subject: [PATCH 0332/1087] Scope dashboard APIs to current project context --- .../server/repositories/projectRepository.ts | 6 +- .../apollo/server/services/askingService.ts | 1 + .../server/services/dashboardService.ts | 11 +++- .../apollo/server/services/projectService.ts | 26 +++++++++ .../services/tests/dashboardService.test.ts | 57 +++++++++++++++++++ 5 files changed, 98 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index cb142b1201..bfb6de9f91 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -280,12 +280,14 @@ export class ProjectRepository const projects = await this.findAll({ order: 'id', - limit: 1, }); if (!projects.length) { throw new Error('No project found'); } - return projects[0]; + if (projects.length === 1) { + return await this.setCurrentProject(projects[0].id); + } + throw new Error('No current project selected'); } public async listProjects() { diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 3a0bb5cf63..01eb9e68e5 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1119,6 +1119,7 @@ export class AskingService implements IAskingService { if (!response) { return null; } + await this.ensureThreadInCurrentProject(response.threadId); return response; } diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 4366d9a150..ca5356ca7a 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -157,10 +157,11 @@ export class DashboardService implements IDashboardService { public async getDashboardItem( dashboardItemId: string | number, ): Promise { + const dashboard = await this.getCurrentDashboard(); const item = await this.dashboardItemRepository.findOneBy({ id: dashboardItemId, }); - if (!item) { + if (!item || String(item.dashboardId) !== String(dashboard.id)) { throw new Error('Dashboard item not found.'); } return item; @@ -202,10 +203,18 @@ export class DashboardService implements IDashboardService { public async updateDashboardItemLayouts( layouts: UpdateDashboardItemLayouts, ): Promise { + const dashboard = await this.getCurrentDashboard(); + const dashboardItems = await this.dashboardItemRepository.findAllBy({ + dashboardId: dashboard.id, + }); + const dashboardItemIds = new Set( + dashboardItems.map((item) => String(item.id)), + ); const updatedItems: DashboardItem[] = []; const isValidLayouts = layouts.every( (layout) => layout.itemId && + dashboardItemIds.has(String(layout.itemId)) && layout.x >= 0 && layout.y >= 0 && layout.w > 0 && diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 67479ae932..040b30c3c6 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -326,7 +326,33 @@ export class ProjectService implements IProjectService { } public async deleteProject(projectId: number): Promise { + const [projectToDelete, currentProject, remainingProjects] = + await Promise.all([ + this.projectRepository.findOneBy({ id: projectId }), + this.projectRepository.findCurrentProject(), + this.projectRepository.listProjects(), + ]); + await this.projectRepository.deleteOne(projectId); + + if (currentProject?.id !== projectId) { + return; + } + + const nextProject = remainingProjects + .filter((project) => project.id !== projectId) + .sort((a, b) => b.id - a.id)[0]; + + if (nextProject) { + await this.projectRepository.setCurrentProject(nextProject.id); + return; + } + + if (projectToDelete) { + logger.debug( + `Deleted the last current project ${projectToDelete.id}; no active project remains.`, + ); + } } public getGeneralConnectionInfo(project) { diff --git a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts index 7ca88841b2..11e09af2c6 100644 --- a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts @@ -120,6 +120,63 @@ describe('DashboardService', () => { }); }); + describe('dashboard scoping', () => { + it('should only return dashboard items from the current dashboard', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + const dashboardItem = { + id: '2001', + dashboardId: '1001', + layout: { x: 0, y: 0, w: 3, h: 2 }, + }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + mockDashboardItemRepository.findOneBy.mockResolvedValue(dashboardItem); + + await expect(dashboardService.getDashboardItem('2001')).resolves.toEqual( + dashboardItem, + ); + }); + + it('should reject dashboard items from another dashboard context', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + const dashboardItem = { + id: '2001', + dashboardId: '9999', + layout: { x: 0, y: 0, w: 3, h: 2 }, + }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + mockDashboardItemRepository.findOneBy.mockResolvedValue(dashboardItem); + + await expect(dashboardService.getDashboardItem('2001')).rejects.toThrow( + 'Dashboard item not found.', + ); + }); + + it('should reject layout updates for items outside the current dashboard', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + const dashboardItems = [ + { + id: '2001', + dashboardId: '1001', + layout: { x: 0, y: 0, w: 3, h: 2 }, + }, + ]; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + mockDashboardItemRepository.findAllBy.mockResolvedValue(dashboardItems); + + await expect( + dashboardService.updateDashboardItemLayouts([ + { itemId: '9999', x: 0, y: 0, w: 3, h: 2 }, + ]), + ).rejects.toThrow('Invalid layouts boundaries.'); + }); + }); + describe('generateCronExpression', () => { it('should generate correct cron expression for daily schedule', () => { const schedule = { From d55cc348a0d06ac41b16cf0b5c2b5ab9a95f51ab Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 17:42:23 +0530 Subject: [PATCH 0333/1087] Align current project API with dashboard context --- wren-ui/src/pages/api/v1/projects/current.ts | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/pages/api/v1/projects/current.ts b/wren-ui/src/pages/api/v1/projects/current.ts index 30e4e422c4..fe680e0a19 100644 --- a/wren-ui/src/pages/api/v1/projects/current.ts +++ b/wren-ui/src/pages/api/v1/projects/current.ts @@ -46,18 +46,26 @@ export default async function handler( } const projectService = getProjectService(); + let currentProject = null; + try { + currentProject = await projectService.getCurrentProject(); + } catch { + currentProject = null; + } const projects = await projectService.listProjects(); - const currentProject = - projects.find((project) => coerceBoolean(project.isCurrent)) || null; + const serializedCurrentProject = + currentProject && projects.some((project) => project.id === currentProject.id) + ? currentProject + : projects.find((project) => coerceBoolean(project.isCurrent)) || null; await respondWithSimple({ res, statusCode: 200, responsePayload: { - currentProject: serializeProject(currentProject), + currentProject: serializeProject(serializedCurrentProject), projects: projects.map(serializeProject), }, - projectId: currentProject?.id ?? 0, + projectId: serializedCurrentProject?.id ?? 0, apiType: ApiType.GET_CURRENT_PROJECT, startTime, requestPayload: {}, From 0497bdaaa3032c65b804207e4ee8d36c25d6e252 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 18:37:30 +0530 Subject: [PATCH 0334/1087] Fix dashboard project context switching --- .../server/repositories/projectRepository.ts | 59 ++++++++++++++++--- .../server/resolvers/dashboardResolver.ts | 20 ++++++- .../server/services/dashboardService.ts | 24 +++++++- .../apollo/server/services/projectService.ts | 38 ++++++++---- .../src/components/OrganizationSwitcher.tsx | 4 +- .../src/pages/api/v1/projects/[id]/select.ts | 11 ++-- wren-ui/src/pages/api/v1/projects/current.ts | 33 +++++++++-- wren-ui/src/pages/project/general.tsx | 2 +- 8 files changed, 157 insertions(+), 34 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index bfb6de9f91..5e48177ba1 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -17,6 +17,10 @@ import { IbisRedshiftConnectionType, IbisDatabricksConnectionType, } from '@server/adaptors/ibisAdaptor'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('ProjectRepository'); +logger.level = 'debug'; export interface BIG_QUERY_CONNECTION_INFO { projectId: string; @@ -199,7 +203,7 @@ export interface IProjectRepository extends IBasicRepository { getCurrentProject: () => Promise; listProjects: () => Promise; findCurrentProject: () => Promise; - setCurrentProject: (projectId: number) => Promise; + setCurrentProject: (projectId: string | number) => Promise; } export class ProjectRepository @@ -275,16 +279,29 @@ export class ProjectRepository public async getCurrentProject() { const currentProject = await this.findCurrentProject(); if (currentProject) { + logger.debug( + `Resolved current project ${String(currentProject.id)} (${currentProject.type || 'unknown'})`, + ); return currentProject; } const projects = await this.findAll({ order: 'id', }); + logger.warn( + `Current project marker missing. Available projects: ${projects + .map((project) => String(project.id)) + .join(', ') || 'none'}`, + ); if (!projects.length) { throw new Error('No project found'); } if (projects.length === 1) { + logger.warn( + `Repairing missing current project marker by selecting the only project ${String( + projects[0].id, + )}`, + ); return await this.setCurrentProject(projects[0].id); } throw new Error('No current project selected'); @@ -302,11 +319,14 @@ export class ProjectRepository } as Partial); } - public async setCurrentProject(projectId: number) { + public async setCurrentProject(projectId: string | number) { const tx = await this.transaction(); try { + logger.debug(`Selecting current project ${String(projectId)}`); await tx(this.tableName).update({ is_current: false }); - await tx(this.tableName).where({ id: projectId }).update({ is_current: true }); + await tx(this.tableName) + .where({ id: this.normalizeProjectId(projectId) }) + .update({ is_current: true }); const project = await this.findOneBy({ id: projectId } as Partial, { tx, }); @@ -314,6 +334,9 @@ export class ProjectRepository throw new Error(`Project ${projectId} not found`); } await this.commit(tx); + logger.debug( + `Selected current project ${String(project.id)} (${project.type || 'unknown'})`, + ); return project; } catch (error) { await this.rollback(tx); @@ -425,12 +448,14 @@ export class ProjectRepository } const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ maxId: 'id', }); return { ...data, - id: Number(row?.maxId || 0) + 1, + id: this.toNextIdValue(row?.maxId), }; }; @@ -448,21 +473,39 @@ export class ProjectRepository } const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ maxId: 'id', }); - let nextId = Number(row?.maxId || 0) + 1; + let nextId = BigInt(row?.maxId ?? 0) + 1n; return data.map((item) => { if (item.id !== undefined && item.id !== null) { return item; } return { ...item, - id: nextId++, + id: this.serializeIdValue(nextId++), }; }); }; + private serializeIdValue(value: number | string | bigint) { + const bigintValue = BigInt(value); + return bigintValue <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(bigintValue) + : bigintValue.toString(); + } + + private toNextIdValue(maxId: number | string | bigint | null | undefined) { + const nextId = BigInt(maxId ?? 0) + 1n; + return this.serializeIdValue(nextId); + } + + private normalizeProjectId(projectId: string | number) { + return typeof projectId === 'string' ? projectId : projectId; + } + private shouldRetryManualId = ( error: unknown, data: Partial | Partial[], diff --git a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts index 460c803bff..2f4f88db80 100644 --- a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts @@ -51,6 +51,10 @@ export class DashboardResolver { if (!dashboard) { throw new Error('Dashboard not found.'); } + const project = await ctx.projectService.getCurrentProject(); + logger.debug( + `Resolving dashboard ${String(dashboard.id)} for project ${String(project.id)} (${project.type || 'unknown'})`, + ); const schedule = ctx.dashboardService.parseCronExpression(dashboard); const items = await ctx.dashboardService.getDashboardItems(dashboard.id); return { @@ -72,6 +76,10 @@ export class DashboardResolver { if (!dashboard) { throw new Error('Dashboard not found.'); } + const project = await ctx.projectService.getCurrentProject(); + logger.debug( + `Resolving dashboard items for dashboard ${String(dashboard.id)} in project ${String(project.id)} (${project.type || 'unknown'})`, + ); return await ctx.dashboardService.getDashboardItems(dashboard.id); } @@ -84,6 +92,11 @@ export class DashboardResolver { const itemType = this.normalizeDashboardItemType(args.data.itemType); const dashboard = await ctx.dashboardService.getCurrentDashboard(); const response = await ctx.askingService.getResponse(responseId); + const project = await ctx.projectService.getCurrentProject(); + + logger.debug( + `Pinning response ${responseId} into dashboard ${String(dashboard.id)} for project ${String(project.id)} (${project.type || 'unknown'})`, + ); if (!response) { throw new Error(`Thread response not found. responseId: ${responseId}`); @@ -194,6 +207,9 @@ export class DashboardResolver { const item = await ctx.dashboardService.getDashboardItem(itemId); const { cacheEnabled } = await ctx.dashboardService.getCurrentDashboard(); const project = await ctx.projectService.getCurrentProject(); + logger.debug( + `Previewing dashboard item ${String(itemId)} for project ${String(project.id)} (${project.type || 'unknown'})`, + ); const manifest = await this.getPreviewManifest(ctx, project.id); try { @@ -315,13 +331,15 @@ export class DashboardResolver { private async getPreviewManifest( ctx: IContext, - projectId: number, + projectId: string | number, ): Promise { const deployment = await ctx.deployService.getLastDeployment(projectId); if (deployment?.manifest) { + logger.debug(`Using deployed manifest for project ${projectId}`); return deployment.manifest as Manifest; } + logger.debug(`Using current model manifest fallback for project ${projectId}`); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); return manifest; } diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index ca5356ca7a..190e168549 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -128,6 +128,9 @@ export class DashboardService implements IDashboardService { public async initDashboard(): Promise { const project = await this.projectService.getCurrentProject(); + logger.debug( + `Initializing dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, + ); const existingDashboard = await this.dashboardRepository.findOneBy({ projectId: project.id, }); @@ -141,10 +144,16 @@ export class DashboardService implements IDashboardService { public async getCurrentDashboard(): Promise { const project = await this.projectService.getCurrentProject(); + logger.debug( + `Loading current dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, + ); const dashboard = await this.dashboardRepository.findOneBy({ projectId: project.id, }); if (dashboard) { + logger.debug( + `Resolved dashboard ${String(dashboard.id)} for project ${String(project.id)}`, + ); return dashboard; } @@ -170,16 +179,23 @@ export class DashboardService implements IDashboardService { public async getDashboardItems( dashboardId: string | number, ): Promise { - return await this.dashboardItemRepository.findAllBy({ + const items = await this.dashboardItemRepository.findAllBy({ dashboardId, }); + logger.debug( + `Loaded ${items.length} dashboard item(s) for dashboard ${String(dashboardId)}`, + ); + return items; } public async createDashboardItem( input: CreateDashboardItemInput, ): Promise { const layout = await this.calculateNewLayout(input.dashboardId); - return await this.dashboardItemRepository.createOne({ + logger.debug( + `Creating dashboard item for dashboard ${String(input.dashboardId)} with type ${input.type}`, + ); + const dashboardItem = await this.dashboardItemRepository.createOne({ dashboardId: input.dashboardId, type: input.type, detail: { @@ -189,6 +205,10 @@ export class DashboardService implements IDashboardService { }, layout, }); + logger.debug( + `Created dashboard item ${String(dashboardItem.id)} for dashboard ${String(input.dashboardId)}`, + ); + return dashboardItem; } public async updateDashboardItem( diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index 040b30c3c6..f27ca78255 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -59,7 +59,7 @@ export type ProjectRecommendationQuestionsResult = { export interface IProjectService { createProject: (projectData: ProjectData) => Promise; updateProject: ( - projectId: number, + projectId: string | number, projectData: Partial, ) => Promise; getGeneralConnectionInfo: (project: Project) => Record; @@ -78,8 +78,8 @@ export interface IProjectService { getCurrentProject: () => Promise; listProjects: () => Promise; - selectProject: (projectId: number) => Promise; - getProjectById: (projectId: number) => Promise; + selectProject: (projectId: string | number) => Promise; + getProjectById: (projectId: string | number) => Promise; writeCredentialFile: ( credentials: JSON, persistCredentialDir: string, @@ -98,7 +98,7 @@ export class ProjectService implements IProjectService { private queryService: IQueryService; private wrenAIAdaptor: IWrenAIAdaptor; private projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; - private projectRecommendationJobs = new Map>(); + private projectRecommendationJobs = new Map>(); constructor({ projectRepository, metadataService, @@ -134,7 +134,7 @@ export class ProjectService implements IProjectService { this.projectRecommendQuestionBackgroundTracker.stop(); } public async updateProject( - projectId: number, + projectId: string | number, projectData: Partial, ): Promise { return await this.projectRepository.updateOne(projectId, projectData); @@ -158,7 +158,8 @@ export class ProjectService implements IProjectService { throw new Error(`Project not found`); } - const existingJob = this.projectRecommendationJobs.get(project.id); + const projectJobKey = String(project.id); + const existingJob = this.projectRecommendationJobs.get(projectJobKey); if (existingJob) { logger.debug( `project "${project.id}" recommended questions are already being requested, reusing in-flight job`, @@ -167,11 +168,11 @@ export class ProjectService implements IProjectService { } const job = this.doGenerateProjectRecommendationQuestions(project); - this.projectRecommendationJobs.set(project.id, job); + this.projectRecommendationJobs.set(projectJobKey, job); try { return await job; } finally { - this.projectRecommendationJobs.delete(project.id); + this.projectRecommendationJobs.delete(projectJobKey); } } @@ -251,11 +252,12 @@ export class ProjectService implements IProjectService { return await this.projectRepository.listProjects(); } - public async selectProject(projectId: number) { + public async selectProject(projectId: string | number) { + logger.debug(`Selecting active project ${String(projectId)}`); return await this.projectRepository.setCurrentProject(projectId); } - public async getProjectById(projectId: number) { + public async getProjectById(projectId: string | number) { return await this.projectRepository.findOneBy({ id: projectId }); } @@ -340,8 +342,8 @@ export class ProjectService implements IProjectService { } const nextProject = remainingProjects - .filter((project) => project.id !== projectId) - .sort((a, b) => b.id - a.id)[0]; + .filter((project) => String(project.id) !== String(projectId)) + .sort((a, b) => this.compareProjectIdsDescending(a.id, b.id))[0]; if (nextProject) { await this.projectRepository.setCurrentProject(nextProject.id); @@ -408,4 +410,16 @@ export class ProjectService implements IProjectService { }, }; } + + private compareProjectIdsDescending( + left: string | number, + right: string | number, + ) { + const leftId = BigInt(left); + const rightId = BigInt(right); + if (leftId === rightId) { + return 0; + } + return leftId > rightId ? -1 : 1; + } } diff --git a/wren-ui/src/components/OrganizationSwitcher.tsx b/wren-ui/src/components/OrganizationSwitcher.tsx index a4c04d952c..03c9340c4a 100644 --- a/wren-ui/src/components/OrganizationSwitcher.tsx +++ b/wren-ui/src/components/OrganizationSwitcher.tsx @@ -37,7 +37,7 @@ interface OrganizationResponse { } interface ProjectRecord { - id: number; + id: string; displayName: string; projectType: WorkspaceProjectType; isCurrent: boolean; @@ -291,7 +291,7 @@ export default function OrganizationSwitcher() { } }; - const selectProject = async (projectId: number) => { + const selectProject = async (projectId: string) => { try { const response = await fetch(`/api/v1/projects/${projectId}/select`, { method: 'POST', diff --git a/wren-ui/src/pages/api/v1/projects/[id]/select.ts b/wren-ui/src/pages/api/v1/projects/[id]/select.ts index 3f538cc8e2..7aef051872 100644 --- a/wren-ui/src/pages/api/v1/projects/[id]/select.ts +++ b/wren-ui/src/pages/api/v1/projects/[id]/select.ts @@ -21,11 +21,10 @@ const getProjectService = () => { const parseProjectId = (value: string | string[] | undefined) => { const rawValue = Array.isArray(value) ? value[0] : value; - const parsed = Number(rawValue); - if (!rawValue || !Number.isInteger(parsed) || parsed <= 0) { + if (!rawValue || !/^\d+$/.test(rawValue)) { throw new ApiError('Invalid project id', 400); } - return parsed; + return rawValue; }; export default async function handler( @@ -41,7 +40,11 @@ export default async function handler( const projectId = parseProjectId(req.query.id); const projectService = getProjectService(); + logger.debug(`API select project request for ${projectId}`); const project = await projectService.selectProject(projectId); + logger.debug( + `API selected project ${String(project.id)} (${project.type || 'unknown'})`, + ); await respondWithSimple({ res, @@ -53,7 +56,7 @@ export default async function handler( projectType: project.projectType || 'CLASSIC', }, }, - projectId, + projectId: Number.isSafeInteger(Number(projectId)) ? Number(projectId) : 0, apiType: ApiType.SELECT_PROJECT, startTime, requestPayload: {}, diff --git a/wren-ui/src/pages/api/v1/projects/current.ts b/wren-ui/src/pages/api/v1/projects/current.ts index fe680e0a19..852f028440 100644 --- a/wren-ui/src/pages/api/v1/projects/current.ts +++ b/wren-ui/src/pages/api/v1/projects/current.ts @@ -23,7 +23,7 @@ const getProjectService = () => { const serializeProject = (project) => project ? { - id: project.id, + id: String(project.id), displayName: project.displayName, projectType: project.projectType || 'CLASSIC', isCurrent: coerceBoolean(project.isCurrent), @@ -49,14 +49,37 @@ export default async function handler( let currentProject = null; try { currentProject = await projectService.getCurrentProject(); - } catch { + logger.debug( + `Resolved current project ${String(currentProject?.id)} (${currentProject?.type || 'unknown'})`, + ); + } catch (error) { + logger.warn( + `Failed to resolve current project through project service: ${ + error instanceof Error ? error.message : String(error) + }`, + ); currentProject = null; } const projects = await projectService.listProjects(); + logger.debug( + `Loaded ${projects.length} project(s) for current project API: ${projects + .map((project) => `${String(project.id)}:${project.type || 'unknown'}:${coerceBoolean(project.isCurrent)}`) + .join(', ') || 'none'}`, + ); const serializedCurrentProject = - currentProject && projects.some((project) => project.id === currentProject.id) + currentProject && + projects.some( + (project) => String(project.id) === String(currentProject.id), + ) ? currentProject : projects.find((project) => coerceBoolean(project.isCurrent)) || null; + logger.debug( + `Current project API returning ${ + serializedCurrentProject + ? `${String(serializedCurrentProject.id)} (${serializedCurrentProject.type || 'unknown'})` + : 'no current project' + }`, + ); await respondWithSimple({ res, @@ -65,7 +88,9 @@ export default async function handler( currentProject: serializeProject(serializedCurrentProject), projects: projects.map(serializeProject), }, - projectId: serializedCurrentProject?.id ?? 0, + projectId: Number.isSafeInteger(Number(serializedCurrentProject?.id)) + ? Number(serializedCurrentProject?.id) + : 0, apiType: ApiType.GET_CURRENT_PROJECT, startTime, requestPayload: {}, diff --git a/wren-ui/src/pages/project/general.tsx b/wren-ui/src/pages/project/general.tsx index e9869c3f04..4f9fd7b09b 100644 --- a/wren-ui/src/pages/project/general.tsx +++ b/wren-ui/src/pages/project/general.tsx @@ -15,7 +15,7 @@ import { import { getLanguageText } from '@/utils/language'; interface CurrentProjectRecord { - id: number; + id: string; displayName: string; projectType: WorkspaceProjectType; isCurrent: boolean; From d1f6859dd98b13461eb2973c0337dc954ca75b1d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 19:36:29 +0530 Subject: [PATCH 0335/1087] Fix dashboard fetch context after project switch --- .../server/services/dashboardService.ts | 201 ++++++++++-------- .../services/tests/dashboardService.test.ts | 29 +++ wren-ui/src/pages/home/dashboard.tsx | 12 +- 3 files changed, 151 insertions(+), 91 deletions(-) diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 190e168549..b29a70d781 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -63,6 +63,37 @@ export class DashboardService implements IDashboardService { private dashboardItemRepository: IDashboardItemRepository; private dashboardRepository: IDashboardRepository; + private async getCurrentDashboardContext(): Promise<{ + project: Awaited>; + dashboard: Dashboard; + }> { + const project = await this.projectService.getCurrentProject(); + logger.debug( + `Loading current dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, + ); + const existingDashboard = await this.dashboardRepository.findOneBy({ + projectId: project.id, + }); + if (existingDashboard) { + logger.debug( + `Resolved dashboard ${String(existingDashboard.id)} for project ${String(project.id)}`, + ); + return { project, dashboard: existingDashboard }; + } + + logger.debug( + `Dashboard not found for project ${project.id}; initializing a default dashboard.`, + ); + const dashboard = await this.dashboardRepository.createOne({ + name: 'Dashboard', + projectId: project.id, + }); + logger.debug( + `Initialized dashboard ${String(dashboard.id)} for project ${String(project.id)}`, + ); + return { project, dashboard }; + } + constructor({ projectService, dashboardItemRepository, @@ -86,15 +117,12 @@ export class DashboardService implements IDashboardService { // Validate input this.validateScheduleInput(data); - // Check if dashboard exists - const dashboard = await this.dashboardRepository.findOneBy({ - id: dashboardId, - }); - if (!dashboard) { + const { dashboard } = await this.getCurrentDashboardContext(); + if (String(dashboard.id) !== String(dashboardId)) { throw new Error(`Dashboard with id ${dashboardId} not found`); } if (!cacheEnabled) { - return await this.dashboardRepository.updateOne(dashboardId, { + return await this.dashboardRepository.updateOne(dashboard.id, { cacheEnabled: false, scheduleFrequency: null, scheduleTimezone: null, @@ -113,7 +141,7 @@ export class DashboardService implements IDashboardService { } // Update dashboard with new schedule - return await this.dashboardRepository.updateOne(dashboardId, { + return await this.dashboardRepository.updateOne(dashboard.id, { cacheEnabled, scheduleFrequency: schedule.frequency, scheduleTimezone: schedule.timezone, @@ -127,40 +155,13 @@ export class DashboardService implements IDashboardService { } public async initDashboard(): Promise { - const project = await this.projectService.getCurrentProject(); - logger.debug( - `Initializing dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, - ); - const existingDashboard = await this.dashboardRepository.findOneBy({ - projectId: project.id, - }); - if (existingDashboard) return existingDashboard; - // only support one dashboard for oss - return await this.dashboardRepository.createOne({ - name: 'Dashboard', - projectId: project.id, - }); + const { dashboard } = await this.getCurrentDashboardContext(); + return dashboard; } public async getCurrentDashboard(): Promise { - const project = await this.projectService.getCurrentProject(); - logger.debug( - `Loading current dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, - ); - const dashboard = await this.dashboardRepository.findOneBy({ - projectId: project.id, - }); - if (dashboard) { - logger.debug( - `Resolved dashboard ${String(dashboard.id)} for project ${String(project.id)}`, - ); - return dashboard; - } - - logger.debug( - `Dashboard not found for project ${project.id}; initializing a default dashboard.`, - ); - return await this.initDashboard(); + const { dashboard } = await this.getCurrentDashboardContext(); + return dashboard; } public async getDashboardItem( @@ -179,11 +180,15 @@ export class DashboardService implements IDashboardService { public async getDashboardItems( dashboardId: string | number, ): Promise { + const { dashboard } = await this.getCurrentDashboardContext(); + if (String(dashboard.id) !== String(dashboardId)) { + throw new Error('Dashboard not found.'); + } const items = await this.dashboardItemRepository.findAllBy({ - dashboardId, + dashboardId: dashboard.id, }); logger.debug( - `Loaded ${items.length} dashboard item(s) for dashboard ${String(dashboardId)}`, + `Loaded ${items.length} dashboard item(s) for dashboard ${String(dashboard.id)}`, ); return items; } @@ -191,12 +196,16 @@ export class DashboardService implements IDashboardService { public async createDashboardItem( input: CreateDashboardItemInput, ): Promise { - const layout = await this.calculateNewLayout(input.dashboardId); + const { dashboard } = await this.getCurrentDashboardContext(); + if (String(dashboard.id) !== String(input.dashboardId)) { + throw new Error('Dashboard not found.'); + } + const layout = await this.calculateNewLayout(dashboard.id); logger.debug( - `Creating dashboard item for dashboard ${String(input.dashboardId)} with type ${input.type}`, + `Creating dashboard item for dashboard ${String(dashboard.id)} with type ${input.type}`, ); const dashboardItem = await this.dashboardItemRepository.createOne({ - dashboardId: input.dashboardId, + dashboardId: dashboard.id, type: input.type, detail: { sql: input.sql, @@ -206,7 +215,7 @@ export class DashboardService implements IDashboardService { layout, }); logger.debug( - `Created dashboard item ${String(dashboardItem.id)} for dashboard ${String(input.dashboardId)}`, + `Created dashboard item ${String(dashboardItem.id)} for dashboard ${String(dashboard.id)}`, ); return dashboardItem; } @@ -506,58 +515,74 @@ export class DashboardService implements IDashboardService { } public parseCronExpression(dashboard: Dashboard): DashboardSchedule { - if (!dashboard.scheduleCron) { - return { - frequency: dashboard.scheduleFrequency, - hour: 0, - minute: 0, - day: null, - timezone: dashboard.scheduleTimezone || '', - cron: '', - } as DashboardSchedule; - } - switch (dashboard.scheduleFrequency) { - case ScheduleFrequencyEnum.CUSTOM: + try { + if (!dashboard.scheduleCron) { return { - frequency: ScheduleFrequencyEnum.CUSTOM, + frequency: dashboard.scheduleFrequency, hour: 0, minute: 0, day: null, timezone: dashboard.scheduleTimezone || '', - cron: dashboard.scheduleCron, - }; - case ScheduleFrequencyEnum.DAILY: - case ScheduleFrequencyEnum.WEEKLY: { - const parts = dashboard.scheduleCron.split(' '); - if (parts.length !== 5) { - throw new Error('Invalid cron expression format'); - } - const [minute, hour, , , dayOfWeek] = parts; - return this.toTimezone({ - frequency: dashboard.scheduleFrequency, - hour: parseInt(hour, 10), - minute: parseInt(minute, 10), - day: - dashboard.scheduleFrequency === ScheduleFrequencyEnum.WEEKLY - ? (dayOfWeek as CacheScheduleDayEnum) - : null, - timezone: dashboard.scheduleTimezone || '', - cron: null, + cron: '', } as DashboardSchedule); } - case ScheduleFrequencyEnum.NEVER: { - return { - frequency: ScheduleFrequencyEnum.NEVER, - hour: null, - minute: null, - day: null, - timezone: dashboard.scheduleTimezone || '', - cron: null, - } as DashboardSchedule; - } - default: { - throw new Error('Invalid schedule frequency'); + switch (dashboard.scheduleFrequency) { + case ScheduleFrequencyEnum.CUSTOM: + return { + frequency: ScheduleFrequencyEnum.CUSTOM, + hour: 0, + minute: 0, + day: null, + timezone: dashboard.scheduleTimezone || '', + cron: dashboard.scheduleCron, + }; + case ScheduleFrequencyEnum.DAILY: + case ScheduleFrequencyEnum.WEEKLY: { + const parts = dashboard.scheduleCron.split(' '); + if (parts.length !== 5) { + throw new Error('Invalid cron expression format'); + } + const [minute, hour, , , dayOfWeek] = parts; + return this.toTimezone({ + frequency: dashboard.scheduleFrequency, + hour: parseInt(hour, 10), + minute: parseInt(minute, 10), + day: + dashboard.scheduleFrequency === ScheduleFrequencyEnum.WEEKLY + ? (dayOfWeek as CacheScheduleDayEnum) + : null, + timezone: dashboard.scheduleTimezone || '', + cron: null, + } as DashboardSchedule); + } + case ScheduleFrequencyEnum.NEVER: { + return { + frequency: ScheduleFrequencyEnum.NEVER, + hour: null, + minute: null, + day: null, + timezone: dashboard.scheduleTimezone || '', + cron: null, + } as DashboardSchedule; + } + default: { + throw new Error('Invalid schedule frequency'); + } } + } catch (error) { + logger.warn( + `Failed to parse dashboard schedule for dashboard ${String(dashboard.id)}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return { + frequency: ScheduleFrequencyEnum.NEVER, + hour: null, + minute: null, + day: null, + timezone: dashboard.scheduleTimezone || '', + cron: null, + } as DashboardSchedule; } } } diff --git a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts index 11e09af2c6..0799756b7c 100644 --- a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts @@ -175,6 +175,35 @@ describe('DashboardService', () => { ]), ).rejects.toThrow('Invalid layouts boundaries.'); }); + + it('should reject fetching items for a dashboard outside the current context', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + + await expect( + dashboardService.getDashboardItems('9999'), + ).rejects.toThrow('Dashboard not found.'); + expect(mockDashboardItemRepository.findAllBy).not.toHaveBeenCalled(); + }); + + it('should reject creating items for a dashboard outside the current context', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + + await expect( + dashboardService.createDashboardItem({ + dashboardId: '9999', + type: 'BAR' as any, + sql: 'select 1', + chartSchema: { title: 'Test' }, + }), + ).rejects.toThrow('Dashboard not found.'); + expect(mockDashboardItemRepository.createOne).not.toHaveBeenCalled(); + }); }); describe('generateCronExpression', () => { diff --git a/wren-ui/src/pages/home/dashboard.tsx b/wren-ui/src/pages/home/dashboard.tsx index b6cc935c45..e7c1a1946f 100644 --- a/wren-ui/src/pages/home/dashboard.tsx +++ b/wren-ui/src/pages/home/dashboard.tsx @@ -24,6 +24,7 @@ import { DataSourceName, ItemLayoutInput, } from '@/apollo/client/graphql/__types__'; +import { parseGraphQLError } from '@/utils/errorHandler'; const isSupportCachedSettings = (dataSource: DataSource) => { // DuckDB not supported, sample dataset as well @@ -49,9 +50,14 @@ export default function Dashboard() { loading, updateQuery: updateDashboardQuery, } = useDashboardQuery({ - fetchPolicy: 'cache-and-network', - onError: () => { - message.error('Failed to fetch dashboard items.'); + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', + onError: (error) => { + console.error('Dashboard query failed', error); + const parsedError = parseGraphQLError(error); + message.error( + parsedError?.message || 'Failed to fetch dashboard items.', + ); router.push(Path.Home); }, }); From 1353a37757486c28c7c2d57df41ce12b19c086f3 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 2 Jul 2026 22:37:12 +0530 Subject: [PATCH 0336/1087] Revert "Fix dashboard fetch context after project switch" This reverts commit d1f6859dd98b13461eb2973c0337dc954ca75b1d. --- .../server/services/dashboardService.ts | 201 ++++++++---------- .../services/tests/dashboardService.test.ts | 29 --- wren-ui/src/pages/home/dashboard.tsx | 12 +- 3 files changed, 91 insertions(+), 151 deletions(-) diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index b29a70d781..190e168549 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -63,37 +63,6 @@ export class DashboardService implements IDashboardService { private dashboardItemRepository: IDashboardItemRepository; private dashboardRepository: IDashboardRepository; - private async getCurrentDashboardContext(): Promise<{ - project: Awaited>; - dashboard: Dashboard; - }> { - const project = await this.projectService.getCurrentProject(); - logger.debug( - `Loading current dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, - ); - const existingDashboard = await this.dashboardRepository.findOneBy({ - projectId: project.id, - }); - if (existingDashboard) { - logger.debug( - `Resolved dashboard ${String(existingDashboard.id)} for project ${String(project.id)}`, - ); - return { project, dashboard: existingDashboard }; - } - - logger.debug( - `Dashboard not found for project ${project.id}; initializing a default dashboard.`, - ); - const dashboard = await this.dashboardRepository.createOne({ - name: 'Dashboard', - projectId: project.id, - }); - logger.debug( - `Initialized dashboard ${String(dashboard.id)} for project ${String(project.id)}`, - ); - return { project, dashboard }; - } - constructor({ projectService, dashboardItemRepository, @@ -117,12 +86,15 @@ export class DashboardService implements IDashboardService { // Validate input this.validateScheduleInput(data); - const { dashboard } = await this.getCurrentDashboardContext(); - if (String(dashboard.id) !== String(dashboardId)) { + // Check if dashboard exists + const dashboard = await this.dashboardRepository.findOneBy({ + id: dashboardId, + }); + if (!dashboard) { throw new Error(`Dashboard with id ${dashboardId} not found`); } if (!cacheEnabled) { - return await this.dashboardRepository.updateOne(dashboard.id, { + return await this.dashboardRepository.updateOne(dashboardId, { cacheEnabled: false, scheduleFrequency: null, scheduleTimezone: null, @@ -141,7 +113,7 @@ export class DashboardService implements IDashboardService { } // Update dashboard with new schedule - return await this.dashboardRepository.updateOne(dashboard.id, { + return await this.dashboardRepository.updateOne(dashboardId, { cacheEnabled, scheduleFrequency: schedule.frequency, scheduleTimezone: schedule.timezone, @@ -155,13 +127,40 @@ export class DashboardService implements IDashboardService { } public async initDashboard(): Promise { - const { dashboard } = await this.getCurrentDashboardContext(); - return dashboard; + const project = await this.projectService.getCurrentProject(); + logger.debug( + `Initializing dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, + ); + const existingDashboard = await this.dashboardRepository.findOneBy({ + projectId: project.id, + }); + if (existingDashboard) return existingDashboard; + // only support one dashboard for oss + return await this.dashboardRepository.createOne({ + name: 'Dashboard', + projectId: project.id, + }); } public async getCurrentDashboard(): Promise { - const { dashboard } = await this.getCurrentDashboardContext(); - return dashboard; + const project = await this.projectService.getCurrentProject(); + logger.debug( + `Loading current dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, + ); + const dashboard = await this.dashboardRepository.findOneBy({ + projectId: project.id, + }); + if (dashboard) { + logger.debug( + `Resolved dashboard ${String(dashboard.id)} for project ${String(project.id)}`, + ); + return dashboard; + } + + logger.debug( + `Dashboard not found for project ${project.id}; initializing a default dashboard.`, + ); + return await this.initDashboard(); } public async getDashboardItem( @@ -180,15 +179,11 @@ export class DashboardService implements IDashboardService { public async getDashboardItems( dashboardId: string | number, ): Promise { - const { dashboard } = await this.getCurrentDashboardContext(); - if (String(dashboard.id) !== String(dashboardId)) { - throw new Error('Dashboard not found.'); - } const items = await this.dashboardItemRepository.findAllBy({ - dashboardId: dashboard.id, + dashboardId, }); logger.debug( - `Loaded ${items.length} dashboard item(s) for dashboard ${String(dashboard.id)}`, + `Loaded ${items.length} dashboard item(s) for dashboard ${String(dashboardId)}`, ); return items; } @@ -196,16 +191,12 @@ export class DashboardService implements IDashboardService { public async createDashboardItem( input: CreateDashboardItemInput, ): Promise { - const { dashboard } = await this.getCurrentDashboardContext(); - if (String(dashboard.id) !== String(input.dashboardId)) { - throw new Error('Dashboard not found.'); - } - const layout = await this.calculateNewLayout(dashboard.id); + const layout = await this.calculateNewLayout(input.dashboardId); logger.debug( - `Creating dashboard item for dashboard ${String(dashboard.id)} with type ${input.type}`, + `Creating dashboard item for dashboard ${String(input.dashboardId)} with type ${input.type}`, ); const dashboardItem = await this.dashboardItemRepository.createOne({ - dashboardId: dashboard.id, + dashboardId: input.dashboardId, type: input.type, detail: { sql: input.sql, @@ -215,7 +206,7 @@ export class DashboardService implements IDashboardService { layout, }); logger.debug( - `Created dashboard item ${String(dashboardItem.id)} for dashboard ${String(dashboard.id)}`, + `Created dashboard item ${String(dashboardItem.id)} for dashboard ${String(input.dashboardId)}`, ); return dashboardItem; } @@ -515,74 +506,58 @@ export class DashboardService implements IDashboardService { } public parseCronExpression(dashboard: Dashboard): DashboardSchedule { - try { - if (!dashboard.scheduleCron) { + if (!dashboard.scheduleCron) { + return { + frequency: dashboard.scheduleFrequency, + hour: 0, + minute: 0, + day: null, + timezone: dashboard.scheduleTimezone || '', + cron: '', + } as DashboardSchedule; + } + switch (dashboard.scheduleFrequency) { + case ScheduleFrequencyEnum.CUSTOM: return { - frequency: dashboard.scheduleFrequency, + frequency: ScheduleFrequencyEnum.CUSTOM, hour: 0, minute: 0, day: null, timezone: dashboard.scheduleTimezone || '', - cron: '', + cron: dashboard.scheduleCron, + }; + case ScheduleFrequencyEnum.DAILY: + case ScheduleFrequencyEnum.WEEKLY: { + const parts = dashboard.scheduleCron.split(' '); + if (parts.length !== 5) { + throw new Error('Invalid cron expression format'); + } + const [minute, hour, , , dayOfWeek] = parts; + return this.toTimezone({ + frequency: dashboard.scheduleFrequency, + hour: parseInt(hour, 10), + minute: parseInt(minute, 10), + day: + dashboard.scheduleFrequency === ScheduleFrequencyEnum.WEEKLY + ? (dayOfWeek as CacheScheduleDayEnum) + : null, + timezone: dashboard.scheduleTimezone || '', + cron: null, } as DashboardSchedule); } - switch (dashboard.scheduleFrequency) { - case ScheduleFrequencyEnum.CUSTOM: - return { - frequency: ScheduleFrequencyEnum.CUSTOM, - hour: 0, - minute: 0, - day: null, - timezone: dashboard.scheduleTimezone || '', - cron: dashboard.scheduleCron, - }; - case ScheduleFrequencyEnum.DAILY: - case ScheduleFrequencyEnum.WEEKLY: { - const parts = dashboard.scheduleCron.split(' '); - if (parts.length !== 5) { - throw new Error('Invalid cron expression format'); - } - const [minute, hour, , , dayOfWeek] = parts; - return this.toTimezone({ - frequency: dashboard.scheduleFrequency, - hour: parseInt(hour, 10), - minute: parseInt(minute, 10), - day: - dashboard.scheduleFrequency === ScheduleFrequencyEnum.WEEKLY - ? (dayOfWeek as CacheScheduleDayEnum) - : null, - timezone: dashboard.scheduleTimezone || '', - cron: null, - } as DashboardSchedule); - } - case ScheduleFrequencyEnum.NEVER: { - return { - frequency: ScheduleFrequencyEnum.NEVER, - hour: null, - minute: null, - day: null, - timezone: dashboard.scheduleTimezone || '', - cron: null, - } as DashboardSchedule; - } - default: { - throw new Error('Invalid schedule frequency'); - } + case ScheduleFrequencyEnum.NEVER: { + return { + frequency: ScheduleFrequencyEnum.NEVER, + hour: null, + minute: null, + day: null, + timezone: dashboard.scheduleTimezone || '', + cron: null, + } as DashboardSchedule; + } + default: { + throw new Error('Invalid schedule frequency'); } - } catch (error) { - logger.warn( - `Failed to parse dashboard schedule for dashboard ${String(dashboard.id)}: ${ - error instanceof Error ? error.message : String(error) - }`, - ); - return { - frequency: ScheduleFrequencyEnum.NEVER, - hour: null, - minute: null, - day: null, - timezone: dashboard.scheduleTimezone || '', - cron: null, - } as DashboardSchedule; } } } diff --git a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts index 0799756b7c..11e09af2c6 100644 --- a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts @@ -175,35 +175,6 @@ describe('DashboardService', () => { ]), ).rejects.toThrow('Invalid layouts boundaries.'); }); - - it('should reject fetching items for a dashboard outside the current context', async () => { - const project = { id: 42 }; - const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; - mockProjectService.getCurrentProject.mockResolvedValue(project); - mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); - - await expect( - dashboardService.getDashboardItems('9999'), - ).rejects.toThrow('Dashboard not found.'); - expect(mockDashboardItemRepository.findAllBy).not.toHaveBeenCalled(); - }); - - it('should reject creating items for a dashboard outside the current context', async () => { - const project = { id: 42 }; - const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; - mockProjectService.getCurrentProject.mockResolvedValue(project); - mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); - - await expect( - dashboardService.createDashboardItem({ - dashboardId: '9999', - type: 'BAR' as any, - sql: 'select 1', - chartSchema: { title: 'Test' }, - }), - ).rejects.toThrow('Dashboard not found.'); - expect(mockDashboardItemRepository.createOne).not.toHaveBeenCalled(); - }); }); describe('generateCronExpression', () => { diff --git a/wren-ui/src/pages/home/dashboard.tsx b/wren-ui/src/pages/home/dashboard.tsx index e7c1a1946f..b6cc935c45 100644 --- a/wren-ui/src/pages/home/dashboard.tsx +++ b/wren-ui/src/pages/home/dashboard.tsx @@ -24,7 +24,6 @@ import { DataSourceName, ItemLayoutInput, } from '@/apollo/client/graphql/__types__'; -import { parseGraphQLError } from '@/utils/errorHandler'; const isSupportCachedSettings = (dataSource: DataSource) => { // DuckDB not supported, sample dataset as well @@ -50,14 +49,9 @@ export default function Dashboard() { loading, updateQuery: updateDashboardQuery, } = useDashboardQuery({ - fetchPolicy: 'network-only', - nextFetchPolicy: 'network-only', - onError: (error) => { - console.error('Dashboard query failed', error); - const parsedError = parseGraphQLError(error); - message.error( - parsedError?.message || 'Failed to fetch dashboard items.', - ); + fetchPolicy: 'cache-and-network', + onError: () => { + message.error('Failed to fetch dashboard items.'); router.push(Path.Home); }, }); From 0ffa2fc22f293d7971de65a351f8cad65378e27b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 2 Jul 2026 23:50:29 +0530 Subject: [PATCH 0337/1087] Complete selected schemas before SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 71 ++++++++++++++++++ .../test_ask_heuristic_text_to_sql.py | 75 +++++++++++++++++++ 2 files changed, 146 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 71bb725adf..6475730d46 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2584,6 +2584,63 @@ def _extract_retrieval_metadata( ] return documents, table_names, table_ddls + async def _complete_sql_generation_context( + self, + *, + query: str, + project_id: Optional[str], + documents: list[dict], + table_names: list[str], + table_ddls: list[str], + ) -> tuple[list[dict], list[str], list[str], dict]: + if not table_names or "db_schema_retrieval" not in self._pipelines: + return documents, table_names, table_ddls, {} + + selected_table_names = list(dict.fromkeys(table_names)) + try: + retrieval_result = await self._run_with_timeout( + "Complete selected schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=query, + tables=selected_table_names, + project_id=project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min(self._schema_retrieval_timeout_seconds, 30), + ) + except Exception as error: + logger.warning( + "Complete selected schema retrieval failed; using existing retrieval context. project_id=%s tables=%s error=%s", + project_id, + selected_table_names, + error, + ) + return documents, table_names, table_ddls, {} + + complete_documents, complete_table_names, complete_table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if not complete_documents: + logger.warning( + "Complete selected schema retrieval returned no documents; using existing retrieval context. project_id=%s tables=%s", + project_id, + selected_table_names, + ) + return documents, table_names, table_ddls, {} + + logger.info( + "Completed SQL generation context with full schemas for project_id %s tables=%s", + project_id, + complete_table_names, + ) + return ( + complete_documents, + complete_table_names, + complete_table_ddls, + retrieval_result.get("construct_retrieval_results", {}), + ) + def _is_visualization_request(self, query: str) -> bool: normalized = (query or "").lower() return bool( @@ -4054,6 +4111,20 @@ async def ask( table_names, table_ddls, ) + ( + documents, + table_names, + table_ddls, + completed_retrieval_result, + ) = await self._complete_sql_generation_context( + query=sql_user_query, + project_id=ask_request.project_id, + documents=documents, + table_names=table_names, + table_ddls=table_ddls, + ) + if completed_retrieval_result: + _retrieval_result = completed_retrieval_result sql_generation_histories = histories if self._is_data_analysis_query( diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index ee3cdbc5cb..4ef95c24b1 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -1,6 +1,33 @@ +import asyncio + from src.web.v1.services.ask import AskHistory, AskService +class _FakeSchemaRetrievalPipeline: + def __init__(self): + self.calls = [] + + async def run(self, **kwargs): + self.calls.append(kwargs) + return { + "construct_retrieval_results": { + "retrieval_results": [ + { + "table_name": "dbo_orders", + "table_ddl": ( + "CREATE TABLE dbo_orders (" + "OrderId INT, CustomerId INT, OrderDate TIMESTAMP" + ")" + ), + } + ], + "has_calculated_field": False, + "has_metric": False, + "has_json_field": False, + } + } + + def test_independent_question_does_not_reuse_historical_sql(): service = AskService(pipelines={}) @@ -519,3 +546,51 @@ def test_retrieval_metadata_ignores_malformed_documents(): assert len(documents) == 1 assert table_names == ["dbo_repair_logs"] assert table_ddls == ["CREATE TABLE dbo_repair_logs (id varchar)"] + + +def test_complete_sql_generation_context_refetches_full_selected_schema(): + pipeline = _FakeSchemaRetrievalPipeline() + service = AskService( + pipelines={"db_schema_retrieval": pipeline}, + schema_retrieval_timeout_seconds=180, + ) + + documents, table_names, table_ddls, retrieval_result = asyncio.run( + service._complete_sql_generation_context( + query="Show monthly orders by customer.", + project_id="project-1", + documents=[ + { + "table_name": "dbo_orders", + "table_ddl": "CREATE TABLE dbo_orders (OrderId INT)", + } + ], + table_names=["dbo_orders"], + table_ddls=["CREATE TABLE dbo_orders (OrderId INT)"], + ) + ) + + assert table_names == ["dbo_orders"] + assert table_ddls == [ + "CREATE TABLE dbo_orders (OrderId INT, CustomerId INT, OrderDate TIMESTAMP)" + ] + assert documents == [ + { + "table_name": "dbo_orders", + "table_ddl": ( + "CREATE TABLE dbo_orders (" + "OrderId INT, CustomerId INT, OrderDate TIMESTAMP" + ")" + ), + } + ] + assert retrieval_result["retrieval_results"] == documents + assert pipeline.calls == [ + { + "query": "Show monthly orders by customer.", + "tables": ["dbo_orders"], + "project_id": "project-1", + "histories": [], + "enable_column_pruning": False, + } + ] From 7f38b532762bc698c6b347036595affe3dbe4beb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 15:25:45 +0530 Subject: [PATCH 0338/1087] Honor explicit table references in SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 68 +++++++++++++++++++ .../test_ask_heuristic_text_to_sql.py | 29 ++++++++ 2 files changed, 97 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6475730d46..2764e7bb5a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -592,6 +592,66 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names.append(table_name) return table_names + def _explicit_table_name_tokens(self, query: str) -> set[str]: + tokens: set[str] = set() + for table_name in self._extract_explicit_table_names_from_query(query): + table_name = str(table_name or "").strip() + if not table_name: + continue + + variants = { + table_name, + table_name.replace(".", "_"), + table_name.replace("_", "."), + re.split(r"[.$]", table_name)[-1], + } + tokens.update( + self._normalize_schema_token(variant) + for variant in variants + if variant + ) + return {token for token in tokens if token} + + def _filter_context_to_explicit_tables( + self, + query: str, + documents: list[dict], + table_names: list[str], + table_ddls: list[str], + ) -> tuple[list[dict], list[str], list[str]]: + explicit_tokens = self._explicit_table_name_tokens(query) + if not explicit_tokens: + return documents, table_names, table_ddls + + selected_indexes: list[int] = [] + for index, table_name in enumerate(table_names): + table_name = str(table_name or "") + table_tokens = { + self._normalize_schema_token(table_name), + self._normalize_schema_token(table_name.replace("_", ".")), + self._normalize_schema_token(table_name.replace(".", "_")), + self._normalize_schema_token(re.split(r"[.$]", table_name)[-1]), + } + if explicit_tokens.intersection(table_tokens): + selected_indexes.append(index) + + if not selected_indexes: + return documents, table_names, table_ddls + + logger.info( + "Restricting SQL generation context to explicit tables for query: %s", + query, + ) + return ( + [documents[index] for index in selected_indexes if index < len(documents)], + [ + table_names[index] + for index in selected_indexes + if index < len(table_names) + ], + [table_ddls[index] for index in selected_indexes if index < len(table_ddls)], + ) + def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -4111,6 +4171,14 @@ async def ask( table_names, table_ddls, ) + documents, table_names, table_ddls = ( + self._filter_context_to_explicit_tables( + sql_user_query, + documents, + table_names, + table_ddls, + ) + ) ( documents, table_names, diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 4ef95c24b1..9a00e0e9ed 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -548,6 +548,35 @@ def test_retrieval_metadata_ignores_malformed_documents(): assert table_ddls == ["CREATE TABLE dbo_repair_logs (id varchar)"] +def test_explicit_table_filter_matches_dot_and_underscore_variants(): + service = AskService(pipelines={}) + documents = [ + { + "table_name": "dbo_failure_patterns", + "table_ddl": "CREATE TABLE dbo_failure_patterns (name varchar)", + }, + { + "table_name": "dbo_DebugEntries_Staging", + "table_ddl": "CREATE TABLE dbo_DebugEntries_Staging (Priority varchar)", + }, + ] + + filtered_documents, filtered_table_names, filtered_table_ddls = ( + service._filter_context_to_explicit_tables( + "Which name values have the highest occurrences in dbo.failure_patterns?", + documents, + [document["table_name"] for document in documents], + [document["table_ddl"] for document in documents], + ) + ) + + assert filtered_documents == [documents[0]] + assert filtered_table_names == ["dbo_failure_patterns"] + assert filtered_table_ddls == [ + "CREATE TABLE dbo_failure_patterns (name varchar)" + ] + + def test_complete_sql_generation_context_refetches_full_selected_schema(): pipeline = _FakeSchemaRetrievalPipeline() service = AskService( From b247ab1b52f38432ad7e1d2f3e8bc40e117f9006 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 16:07:28 +0530 Subject: [PATCH 0339/1087] Prioritize explicit tables before SQL pruning --- wren-ai-service/src/web/v1/services/ask.py | 25 +++++++++++-------- .../test_ask_heuristic_text_to_sql.py | 23 +++++++++++++++++ 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2764e7bb5a..7d5c2b6a5a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -3091,10 +3091,7 @@ def _prune_sql_generation_context( query_key = self._normalize_schema_token(query) query_terms = self._query_schema_terms(query) - explicit_tables = { - self._normalize_schema_token(table_name) - for table_name in self._extract_explicit_table_names_from_query(query) - } + explicit_tables = self._explicit_table_name_tokens(query) scored: list[tuple[int, int]] = [] for index, table in enumerate(parsed_tables): @@ -3110,7 +3107,13 @@ def _prune_sql_generation_context( } score = 0 - if normalized_table in explicit_tables or normalized_short_table in explicit_tables: + explicit_table_tokens = { + normalized_table, + normalized_short_table, + self._normalize_schema_token(table_name.replace("_", ".")), + self._normalize_schema_token(table_name.replace(".", "_")), + } + if explicit_tables.intersection(explicit_table_tokens): score += 1000 if normalized_table and normalized_table in query_key: score += 500 @@ -4165,12 +4168,6 @@ async def ask( return results if documents and not api_results: - documents, table_names, table_ddls = self._prune_sql_generation_context( - sql_user_query, - documents, - table_names, - table_ddls, - ) documents, table_names, table_ddls = ( self._filter_context_to_explicit_tables( sql_user_query, @@ -4179,6 +4176,12 @@ async def ask( table_ddls, ) ) + documents, table_names, table_ddls = self._prune_sql_generation_context( + sql_user_query, + documents, + table_names, + table_ddls, + ) ( documents, table_names, diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 9a00e0e9ed..edcfd2b7e7 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -577,6 +577,29 @@ def test_explicit_table_filter_matches_dot_and_underscore_variants(): ] +def test_prune_sql_generation_context_prioritizes_explicit_table_variant(): + service = AskService(pipelines={}) + table_ddls = [ + "CREATE TABLE dbo_DebugEntries_Staging (Priority varchar)", + "CREATE TABLE dbo_failure_patterns (name varchar, occurrences int)", + ] + documents = [ + {"table_name": "dbo_DebugEntries_Staging", "table_ddl": table_ddls[0]}, + {"table_name": "dbo_failure_patterns", "table_ddl": table_ddls[1]}, + ] + + _, table_names, pruned_ddls = service._prune_sql_generation_context( + "Which name values have the highest occurrences in dbo.failure_patterns?", + documents, + [document["table_name"] for document in documents], + table_ddls, + max_tables=1, + ) + + assert table_names == ["dbo_failure_patterns"] + assert pruned_ddls == [table_ddls[1]] + + def test_complete_sql_generation_context_refetches_full_selected_schema(): pipeline = _FakeSchemaRetrievalPipeline() service = AskService( From 2a2c7039a79904d1b53cdae2c1b08f46fe7edb65 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 16:43:49 +0530 Subject: [PATCH 0340/1087] Resolve explicit metadata objects generically --- wren-ai-service/src/web/v1/services/ask.py | 80 ++++++++++++++----- .../test_ask_heuristic_text_to_sql.py | 49 ++++++++++++ 2 files changed, 110 insertions(+), 19 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7d5c2b6a5a..12cf0396b0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -569,28 +569,28 @@ def _build_explicit_table_preview_sql( ) def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: - table_names: list[str] = [] + object_names: list[str] = [] for match in re.finditer( - r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\b(?:from|table|model|view|metric)\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): - table_name = match.group(1).strip(".,;:()[]{}") - if table_name and table_name not in table_names: - table_names.append(table_name) + object_name = match.group(1).strip(".,;:()[]{}") + if object_name and object_name not in object_names: + object_names.append(object_name) for match in re.finditer( - r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\b(?:using|in|on|for|against)\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): - table_name = match.group(1).strip(".,;:()[]{}") + object_name = match.group(1).strip(".,;:()[]{}") if ( - table_name - and ("." in table_name or "_" in table_name) - and table_name not in table_names + object_name + and ("." in object_name or "_" in object_name) + and object_name not in object_names ): - table_names.append(table_name) - return table_names + object_names.append(object_name) + return object_names def _explicit_table_name_tokens(self, query: str) -> set[str]: tokens: set[str] = set() @@ -612,6 +612,40 @@ def _explicit_table_name_tokens(self, query: str) -> set[str]: ) return {token for token in tokens if token} + def _metadata_object_name_tokens(self, object_name: str) -> set[str]: + object_name = str(object_name or "").strip() + if not object_name: + return set() + + variants = { + object_name, + object_name.replace("_", "."), + object_name.replace(".", "_"), + re.split(r"[.$]", object_name)[-1], + } + return { + token + for variant in variants + if (token := self._normalize_schema_token(variant)) + } + + def _referenced_relationship_table_tokens( + self, table_ddls: list[str], selected_indexes: list[int] + ) -> set[str]: + selected_ddls = [ + table_ddls[index] for index in selected_indexes if index < len(table_ddls) + ] + relationship_tokens: set[str] = set() + for relationship in self._extract_metadata_relationships(selected_ddls): + for table_name in re.findall( + r"([A-Za-z_][A-Za-z0-9_.$]*)\s*\(", + relationship, + ): + relationship_tokens.update( + self._metadata_object_name_tokens(table_name) + ) + return relationship_tokens + def _filter_context_to_explicit_tables( self, query: str, @@ -626,20 +660,28 @@ def _filter_context_to_explicit_tables( selected_indexes: list[int] = [] for index, table_name in enumerate(table_names): table_name = str(table_name or "") - table_tokens = { - self._normalize_schema_token(table_name), - self._normalize_schema_token(table_name.replace("_", ".")), - self._normalize_schema_token(table_name.replace(".", "_")), - self._normalize_schema_token(re.split(r"[.$]", table_name)[-1]), - } + table_tokens = self._metadata_object_name_tokens(table_name) if explicit_tokens.intersection(table_tokens): selected_indexes.append(index) if not selected_indexes: return documents, table_names, table_ddls + relationship_tokens = self._referenced_relationship_table_tokens( + table_ddls, selected_indexes + ) + if relationship_tokens: + selected_index_set = set(selected_indexes) + for index, table_name in enumerate(table_names): + if index in selected_index_set: + continue + if relationship_tokens.intersection( + self._metadata_object_name_tokens(str(table_name or "")) + ): + selected_indexes.append(index) + logger.info( - "Restricting SQL generation context to explicit tables for query: %s", + "Restricting SQL generation context to explicit metadata objects for query: %s", query, ) return ( diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index edcfd2b7e7..1ae7350aeb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -550,6 +550,9 @@ def test_retrieval_metadata_ignores_malformed_documents(): def test_explicit_table_filter_matches_dot_and_underscore_variants(): service = AskService(pipelines={}) + assert service._extract_explicit_table_names_from_query( + "Which name values have the highest occurrences in dbo.failure_patterns?" + ) == ["dbo.failure_patterns"] documents = [ { "table_name": "dbo_failure_patterns", @@ -577,6 +580,52 @@ def test_explicit_table_filter_matches_dot_and_underscore_variants(): ] +def test_explicit_metadata_object_extraction_avoids_plain_language_prepositions(): + service = AskService(pipelines={}) + + assert service._extract_explicit_table_names_from_query( + "Show total sales in each market" + ) == [] + assert service._extract_explicit_table_names_from_query( + "Show the first 10 rows from CustomerMaster" + ) == ["CustomerMaster"] + assert service._extract_explicit_table_names_from_query( + "Create a chart on dbo.ticket_cycles by status" + ) == ["dbo.ticket_cycles"] + + +def test_explicit_table_filter_includes_relationship_dependencies(): + service = AskService(pipelines={}) + documents = [ + { + "table_name": "dbo_orders", + "table_ddl": ( + "CREATE TABLE dbo_orders (" + "OrderId INT, CustomerId INT, " + "FOREIGN KEY (CustomerId) REFERENCES dbo_customers(CustomerId)" + ")" + ), + }, + { + "table_name": "dbo_customers", + "table_ddl": "CREATE TABLE dbo_customers (CustomerId INT, Region varchar)", + }, + { + "table_name": "dbo_inventory", + "table_ddl": "CREATE TABLE dbo_inventory (Sku varchar)", + }, + ] + + _, filtered_table_names, _ = service._filter_context_to_explicit_tables( + "Show orders from dbo.orders by customer region", + documents, + [document["table_name"] for document in documents], + [document["table_ddl"] for document in documents], + ) + + assert filtered_table_names == ["dbo_orders", "dbo_customers"] + + def test_prune_sql_generation_context_prioritizes_explicit_table_variant(): service = AskService(pipelines={}) table_ddls = [ From a7f3de2022847628509af94e7b2f290e64644a8d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 18:03:16 +0530 Subject: [PATCH 0341/1087] Prevent standalone asks from reusing thread SQL history --- wren-ai-service/src/web/v1/services/ask.py | 26 ++++++++++---- .../test_ask_heuristic_text_to_sql.py | 35 +++++++++++++++++-- 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 12cf0396b0..f2940b3ad2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -247,6 +247,15 @@ def _should_reuse_historical_question_sql( ) -> bool: return bool(histories) and self._needs_conversation_context(query) + def _sql_generation_histories_for_query( + self, + query: str, + histories: list[AskHistory] | None, + ) -> list[AskHistory]: + if not histories or not self._needs_conversation_context(query): + return [] + return histories + def _rewrite_query_for_text_to_sql(self, query: str) -> str: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -4239,7 +4248,10 @@ async def ask( if completed_retrieval_result: _retrieval_result = completed_retrieval_result - sql_generation_histories = histories + sql_generation_histories = self._sql_generation_histories_for_query( + sql_user_query, + histories, + ) if self._is_data_analysis_query( sql_user_query ) and not self._needs_conversation_context(sql_user_query): @@ -4264,7 +4276,7 @@ async def ask( intent_reasoning=intent_reasoning, retrieved_tables=table_names, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) if sql_generation_histories: @@ -4323,7 +4335,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) if not self._is_stopped(query_id, self._ask_results) and not api_results: @@ -4335,7 +4347,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) try: @@ -4467,7 +4479,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) if allow_sql_diagnosis: @@ -4542,7 +4554,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" @@ -4574,7 +4586,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 1ae7350aeb..e79451aeb0 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -30,6 +30,7 @@ async def run(self, **kwargs): def test_independent_question_does_not_reuse_historical_sql(): service = AskService(pipelines={}) + histories = [AskHistory(question="previous", sql="SELECT 1")] assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", @@ -37,17 +38,47 @@ def test_independent_question_does_not_reuse_historical_sql(): ) assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", - [AskHistory(question="previous", sql="SELECT 1")], + histories, ) + assert service._sql_generation_histories_for_query( + "Show monthly order count by market.", + histories, + ) == [] def test_contextual_followup_can_reuse_historical_sql(): service = AskService(pipelines={}) + histories = [AskHistory(question="previous", sql="SELECT 1")] assert service._should_reuse_historical_question_sql( "Use the same table and show it by month.", - [AskHistory(question="previous", sql="SELECT 1")], + histories, ) + assert service._sql_generation_histories_for_query( + "Use the same table and show it by month.", + histories, + ) == histories + + +def test_independent_question_with_thread_history_uses_standalone_sql_generation(): + service = AskService(pipelines={}) + + assert service._sql_generation_histories_for_query( + "Which name values have the highest occurrences in dbo.failure_patterns?", + [AskHistory(question="Show repair ticket status", sql="SELECT status FROM x")], + ) == [] + + +def test_explicit_followup_question_can_use_history(): + service = AskService(pipelines={}) + histories = [ + AskHistory(question="Show orders by market", sql="SELECT Market FROM x") + ] + + assert service._sql_generation_histories_for_query( + "Use that same table and show it by month.", + histories, + ) == histories def test_metadata_table_question_is_not_sql_or_chart_intent(): From 61f43c8cd5bab98ecf8936b528bbe29fcf64c52a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 18:27:01 +0530 Subject: [PATCH 0342/1087] Use active metadata consistently across datasource asks --- wren-ai-service/src/web/v1/services/ask.py | 18 +++++------ .../test_ask_heuristic_text_to_sql.py | 32 +++++++++++++++++++ 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f2940b3ad2..705dc5b242 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -932,13 +932,8 @@ def _select_best_analytics_table( score += 2 if date_column: score += 4 - table_name = str(table.get("name") or "").lower() - if not table_name: + if not str(table.get("name") or ""): continue - if "sales" in table_name: - score += 5 - if "stage" in table_name: - score -= 8 scored.append((score, table, dimensions, measure, date_column)) @@ -3982,6 +3977,11 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) + retrieval_histories = self._sql_generation_histories_for_query( + sql_user_query, + histories, + ) + if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -3993,7 +3993,7 @@ async def ask( rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if retrieval_histories else False, ) try: @@ -4001,7 +4001,7 @@ async def ask( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( query=sql_user_query, - histories=histories, + histories=retrieval_histories, project_id=ask_request.project_id, enable_column_pruning=( enable_column_pruning @@ -4052,7 +4052,7 @@ async def ask( query=user_query, tables=explicit_table_names, project_id=ask_request.project_id, - histories=histories, + histories=retrieval_histories, enable_column_pruning=enable_column_pruning, ), timeout_seconds=min( diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index e79451aeb0..6395b5630b 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -81,6 +81,38 @@ def test_explicit_followup_question_can_use_history(): ) == histories +def test_generic_table_selection_does_not_prefer_sales_named_tables(): + service = AskService(pipelines={}) + tables = [ + { + "name": "dbo_agent_events", + "columns": [ + {"name": "Region", "type": "varchar"}, + {"name": "Amount", "type": "decimal"}, + ], + }, + { + "name": "dbo_sales_archive", + "columns": [ + {"name": "Region", "type": "varchar"}, + {"name": "Amount", "type": "decimal"}, + ], + }, + ] + + selected = service._select_best_analytics_table( + tables, + required_dimensions=[("Region",)], + measure_candidates=("Amount",), + ) + + assert selected is not None + table, dimensions, measure, _ = selected + assert table["name"] == "dbo_agent_events" + assert dimensions == ["Region"] + assert measure == "Amount" + + def test_metadata_table_question_is_not_sql_or_chart_intent(): service = AskService(pipelines={}) From f6cc7ce29850cea666eba133798e4e92fabcfede Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 18:53:06 +0530 Subject: [PATCH 0343/1087] Force AI metadata sync on manual deploy --- .../src/apollo/client/graphql/deploy.generated.ts | 10 ++++++---- wren-ui/src/apollo/client/graphql/deploy.ts | 4 ++-- wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts | 12 ++++++++---- wren-ui/src/components/deploy/Deploy.tsx | 2 +- 4 files changed, 17 insertions(+), 11 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 83d2456a17..77b43cc981 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,7 +3,9 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; +export type DeployMutationVariables = Types.Exact<{ + force?: Types.InputMaybe; +}>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -15,8 +17,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -77,4 +79,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file +export type DeployStatusQueryResult = Apollo.QueryResult; diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 75fe308d2f..0967269ebd 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 943a7f139e..141c3b3015 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -776,11 +776,15 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { private async waitDeployFinished(deployId: string): Promise { let deploySuccess = false; - // timeout after 30 seconds - for (let waitTime = 1; waitTime <= 7; waitTime++) { + const maxAttempts = 90; + const pollingIntervalMs = 2000; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { try { const status = await this.getDeployStatus(deployId); - logger.debug(`Wren AI: Deploy status: ${status}`); + logger.debug( + `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, + ); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -795,7 +799,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); + await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); } return deploySuccess; } diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index 14560b42db..8d75bdc6a7 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -66,7 +66,7 @@ export default function Deploy() { const syncStatus = data?.modelSync.status; const onDeploy = () => { - deployMutation(); + deployMutation({ variables: { force: true } }); startPolling(1000); }; From 9ee698bc8eafc2946c16f3d4142461d5b7dd2520 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:03:23 +0530 Subject: [PATCH 0344/1087] Revert "Force AI metadata sync on manual deploy" This reverts commit f6cc7ce29850cea666eba133798e4e92fabcfede. --- .../src/apollo/client/graphql/deploy.generated.ts | 10 ++++------ wren-ui/src/apollo/client/graphql/deploy.ts | 4 ++-- wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts | 12 ++++-------- wren-ui/src/components/deploy/Deploy.tsx | 2 +- 4 files changed, 11 insertions(+), 17 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 77b43cc981..83d2456a17 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,9 +3,7 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ - force?: Types.InputMaybe; -}>; +export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -17,8 +15,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy($force: Boolean) { - deploy(force: $force) + mutation Deploy { + deploy } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -79,4 +77,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; +export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 0967269ebd..75fe308d2f 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy($force: Boolean) { - deploy(force: $force) + mutation Deploy { + deploy } `; diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 141c3b3015..943a7f139e 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -776,15 +776,11 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { private async waitDeployFinished(deployId: string): Promise { let deploySuccess = false; - const maxAttempts = 90; - const pollingIntervalMs = 2000; - - for (let attempt = 1; attempt <= maxAttempts; attempt++) { + // timeout after 30 seconds + for (let waitTime = 1; waitTime <= 7; waitTime++) { try { const status = await this.getDeployStatus(deployId); - logger.debug( - `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, - ); + logger.debug(`Wren AI: Deploy status: ${status}`); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -799,7 +795,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); + await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); } return deploySuccess; } diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index 8d75bdc6a7..14560b42db 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -66,7 +66,7 @@ export default function Deploy() { const syncStatus = data?.modelSync.status; const onDeploy = () => { - deployMutation({ variables: { force: true } }); + deployMutation(); startPolling(1000); }; From 505d45ede47df33017b902e1463e52422b8fc60d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:09:09 +0530 Subject: [PATCH 0345/1087] Mark schema changes as unsynced before deploy --- .../apollo/server/resolvers/modelResolver.ts | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 54e04a8dfb..b2a3a22731 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -28,6 +28,7 @@ import { } from '../utils/model'; import { CompactTable, PreviewDataResponse } from '@server/services'; import { TelemetryEvent } from '../telemetry/telemetry'; +import DataSourceSchemaDetector from '@server/managers/dataSourceSchemaDetector'; const logger = getLogger('ModelResolver'); logger.level = 'debug'; @@ -209,6 +210,9 @@ export class ModelResolver { if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } + if (await this.hasUnresolvedSchemaChange(ctx, id)) { + return { status: SyncStatusEnum.UNSYNCRONIZED }; + } return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; @@ -241,11 +245,24 @@ export class ModelResolver { // Recommendation generation depends on a successful deployment because // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + await schemaDetector.detectSchemaChange(); await ctx.projectService.generateProjectRecommendationQuestions(); } return deployRes; } + private async hasUnresolvedSchemaChange(ctx: IContext, projectId: number) { + const lastSchemaChange = + await ctx.schemaChangeRepository.findLastSchemaChange(projectId); + return Object.values(lastSchemaChange?.resolve || {}).some( + (resolved) => resolved === false, + ); + } + public async getMDL(_root: any, args: { hash: string }, ctx: IContext) { const mdl = await ctx.deployService.getMDLByHash(args.hash); return { From 289669b93a2e9b47e5e956adacb0fc2f8372e541 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:17:54 +0530 Subject: [PATCH 0346/1087] Force deploy when schema changes are unresolved --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index b2a3a22731..c13e10fa45 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -236,10 +236,14 @@ export class ModelResolver { }); } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const hasUnresolvedSchemaChange = await this.hasUnresolvedSchemaChange( + ctx, + project.id, + ); const deployRes = await ctx.deployService.deploy( manifest, project.id, - args.force, + args.force || hasUnresolvedSchemaChange, ); // Recommendation generation depends on a successful deployment because From 5cfc8a4787d3f50a296019a59082464877ffc25c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:28:10 +0530 Subject: [PATCH 0347/1087] Revert "Force deploy when schema changes are unresolved" This reverts commit 289669b93a2e9b47e5e956adacb0fc2f8372e541. --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index c13e10fa45..b2a3a22731 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -236,14 +236,10 @@ export class ModelResolver { }); } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const hasUnresolvedSchemaChange = await this.hasUnresolvedSchemaChange( - ctx, - project.id, - ); const deployRes = await ctx.deployService.deploy( manifest, project.id, - args.force || hasUnresolvedSchemaChange, + args.force, ); // Recommendation generation depends on a successful deployment because From f5504d7f8a52b341fe7a79a3cc801447cf2e8f22 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:28:10 +0530 Subject: [PATCH 0348/1087] Revert "Mark schema changes as unsynced before deploy" This reverts commit 505d45ede47df33017b902e1463e52422b8fc60d. --- .../apollo/server/resolvers/modelResolver.ts | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index b2a3a22731..54e04a8dfb 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -28,7 +28,6 @@ import { } from '../utils/model'; import { CompactTable, PreviewDataResponse } from '@server/services'; import { TelemetryEvent } from '../telemetry/telemetry'; -import DataSourceSchemaDetector from '@server/managers/dataSourceSchemaDetector'; const logger = getLogger('ModelResolver'); logger.level = 'debug'; @@ -210,9 +209,6 @@ export class ModelResolver { if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - if (await this.hasUnresolvedSchemaChange(ctx, id)) { - return { status: SyncStatusEnum.UNSYNCRONIZED }; - } return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; @@ -245,24 +241,11 @@ export class ModelResolver { // Recommendation generation depends on a successful deployment because // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { - const schemaDetector = new DataSourceSchemaDetector({ - ctx, - projectId: project.id, - }); - await schemaDetector.detectSchemaChange(); await ctx.projectService.generateProjectRecommendationQuestions(); } return deployRes; } - private async hasUnresolvedSchemaChange(ctx: IContext, projectId: number) { - const lastSchemaChange = - await ctx.schemaChangeRepository.findLastSchemaChange(projectId); - return Object.values(lastSchemaChange?.resolve || {}).some( - (resolved) => resolved === false, - ); - } - public async getMDL(_root: any, args: { hash: string }, ctx: IContext) { const mdl = await ctx.deployService.getMDLByHash(args.hash); return { From bacc51653e87ba5993e12f532b14a89838a72ac9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:28:10 +0530 Subject: [PATCH 0349/1087] Reapply "Force AI metadata sync on manual deploy" This reverts commit 9ee698bc8eafc2946c16f3d4142461d5b7dd2520. --- .../src/apollo/client/graphql/deploy.generated.ts | 10 ++++++---- wren-ui/src/apollo/client/graphql/deploy.ts | 4 ++-- wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts | 12 ++++++++---- wren-ui/src/components/deploy/Deploy.tsx | 2 +- 4 files changed, 17 insertions(+), 11 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 83d2456a17..77b43cc981 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,7 +3,9 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; +export type DeployMutationVariables = Types.Exact<{ + force?: Types.InputMaybe; +}>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -15,8 +17,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -77,4 +79,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file +export type DeployStatusQueryResult = Apollo.QueryResult; diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 75fe308d2f..0967269ebd 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 943a7f139e..141c3b3015 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -776,11 +776,15 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { private async waitDeployFinished(deployId: string): Promise { let deploySuccess = false; - // timeout after 30 seconds - for (let waitTime = 1; waitTime <= 7; waitTime++) { + const maxAttempts = 90; + const pollingIntervalMs = 2000; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { try { const status = await this.getDeployStatus(deployId); - logger.debug(`Wren AI: Deploy status: ${status}`); + logger.debug( + `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, + ); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -795,7 +799,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); + await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); } return deploySuccess; } diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index 14560b42db..8d75bdc6a7 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -66,7 +66,7 @@ export default function Deploy() { const syncStatus = data?.modelSync.status; const onDeploy = () => { - deployMutation(); + deployMutation({ variables: { force: true } }); startPolling(1000); }; From 4ec4e0bbef7d2761f0bd3df4f7647baef8700e83 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:36:34 +0530 Subject: [PATCH 0350/1087] Revert "Reapply "Force AI metadata sync on manual deploy"" This reverts commit bacc51653e87ba5993e12f532b14a89838a72ac9. --- .../src/apollo/client/graphql/deploy.generated.ts | 10 ++++------ wren-ui/src/apollo/client/graphql/deploy.ts | 4 ++-- wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts | 12 ++++-------- wren-ui/src/components/deploy/Deploy.tsx | 2 +- 4 files changed, 11 insertions(+), 17 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 77b43cc981..83d2456a17 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,9 +3,7 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ - force?: Types.InputMaybe; -}>; +export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -17,8 +15,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy($force: Boolean) { - deploy(force: $force) + mutation Deploy { + deploy } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -79,4 +77,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; +export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 0967269ebd..75fe308d2f 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy($force: Boolean) { - deploy(force: $force) + mutation Deploy { + deploy } `; diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 141c3b3015..943a7f139e 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -776,15 +776,11 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { private async waitDeployFinished(deployId: string): Promise { let deploySuccess = false; - const maxAttempts = 90; - const pollingIntervalMs = 2000; - - for (let attempt = 1; attempt <= maxAttempts; attempt++) { + // timeout after 30 seconds + for (let waitTime = 1; waitTime <= 7; waitTime++) { try { const status = await this.getDeployStatus(deployId); - logger.debug( - `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, - ); + logger.debug(`Wren AI: Deploy status: ${status}`); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -799,7 +795,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); + await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); } return deploySuccess; } diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index 8d75bdc6a7..14560b42db 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -66,7 +66,7 @@ export default function Deploy() { const syncStatus = data?.modelSync.status; const onDeploy = () => { - deployMutation({ variables: { force: true } }); + deployMutation(); startPolling(1000); }; From 064ca44f0443020bcaae1139c2194a25ecb9d009 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:45:47 +0530 Subject: [PATCH 0351/1087] Revert "Use active metadata consistently across datasource asks" This reverts commit 61f43c8cd5bab98ecf8936b528bbe29fcf64c52a. --- wren-ai-service/src/web/v1/services/ask.py | 18 +++++------ .../test_ask_heuristic_text_to_sql.py | 32 ------------------- 2 files changed, 9 insertions(+), 41 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 705dc5b242..f2940b3ad2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -932,8 +932,13 @@ def _select_best_analytics_table( score += 2 if date_column: score += 4 - if not str(table.get("name") or ""): + table_name = str(table.get("name") or "").lower() + if not table_name: continue + if "sales" in table_name: + score += 5 + if "stage" in table_name: + score -= 8 scored.append((score, table, dimensions, measure, date_column)) @@ -3977,11 +3982,6 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - retrieval_histories = self._sql_generation_histories_for_query( - sql_user_query, - histories, - ) - if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -3993,7 +3993,7 @@ async def ask( rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, - is_followup=True if retrieval_histories else False, + is_followup=True if histories else False, ) try: @@ -4001,7 +4001,7 @@ async def ask( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( query=sql_user_query, - histories=retrieval_histories, + histories=histories, project_id=ask_request.project_id, enable_column_pruning=( enable_column_pruning @@ -4052,7 +4052,7 @@ async def ask( query=user_query, tables=explicit_table_names, project_id=ask_request.project_id, - histories=retrieval_histories, + histories=histories, enable_column_pruning=enable_column_pruning, ), timeout_seconds=min( diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 6395b5630b..e79451aeb0 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -81,38 +81,6 @@ def test_explicit_followup_question_can_use_history(): ) == histories -def test_generic_table_selection_does_not_prefer_sales_named_tables(): - service = AskService(pipelines={}) - tables = [ - { - "name": "dbo_agent_events", - "columns": [ - {"name": "Region", "type": "varchar"}, - {"name": "Amount", "type": "decimal"}, - ], - }, - { - "name": "dbo_sales_archive", - "columns": [ - {"name": "Region", "type": "varchar"}, - {"name": "Amount", "type": "decimal"}, - ], - }, - ] - - selected = service._select_best_analytics_table( - tables, - required_dimensions=[("Region",)], - measure_candidates=("Amount",), - ) - - assert selected is not None - table, dimensions, measure, _ = selected - assert table["name"] == "dbo_agent_events" - assert dimensions == ["Region"] - assert measure == "Amount" - - def test_metadata_table_question_is_not_sql_or_chart_intent(): service = AskService(pipelines={}) From f4f84a68fdb97e405d758d358cdf1feb63e8e01c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:45:48 +0530 Subject: [PATCH 0352/1087] Revert "Prevent standalone asks from reusing thread SQL history" This reverts commit a7f3de2022847628509af94e7b2f290e64644a8d. --- wren-ai-service/src/web/v1/services/ask.py | 26 ++++---------- .../test_ask_heuristic_text_to_sql.py | 35 ++----------------- 2 files changed, 9 insertions(+), 52 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f2940b3ad2..12cf0396b0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -247,15 +247,6 @@ def _should_reuse_historical_question_sql( ) -> bool: return bool(histories) and self._needs_conversation_context(query) - def _sql_generation_histories_for_query( - self, - query: str, - histories: list[AskHistory] | None, - ) -> list[AskHistory]: - if not histories or not self._needs_conversation_context(query): - return [] - return histories - def _rewrite_query_for_text_to_sql(self, query: str) -> str: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -4248,10 +4239,7 @@ async def ask( if completed_retrieval_result: _retrieval_result = completed_retrieval_result - sql_generation_histories = self._sql_generation_histories_for_query( - sql_user_query, - histories, - ) + sql_generation_histories = histories if self._is_data_analysis_query( sql_user_query ) and not self._needs_conversation_context(sql_user_query): @@ -4276,7 +4264,7 @@ async def ask( intent_reasoning=intent_reasoning, retrieved_tables=table_names, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) if sql_generation_histories: @@ -4335,7 +4323,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) if not self._is_stopped(query_id, self._ask_results) and not api_results: @@ -4347,7 +4335,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) try: @@ -4479,7 +4467,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) if allow_sql_diagnosis: @@ -4554,7 +4542,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" @@ -4586,7 +4574,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index e79451aeb0..1ae7350aeb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -30,7 +30,6 @@ async def run(self, **kwargs): def test_independent_question_does_not_reuse_historical_sql(): service = AskService(pipelines={}) - histories = [AskHistory(question="previous", sql="SELECT 1")] assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", @@ -38,47 +37,17 @@ def test_independent_question_does_not_reuse_historical_sql(): ) assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", - histories, + [AskHistory(question="previous", sql="SELECT 1")], ) - assert service._sql_generation_histories_for_query( - "Show monthly order count by market.", - histories, - ) == [] def test_contextual_followup_can_reuse_historical_sql(): service = AskService(pipelines={}) - histories = [AskHistory(question="previous", sql="SELECT 1")] assert service._should_reuse_historical_question_sql( "Use the same table and show it by month.", - histories, + [AskHistory(question="previous", sql="SELECT 1")], ) - assert service._sql_generation_histories_for_query( - "Use the same table and show it by month.", - histories, - ) == histories - - -def test_independent_question_with_thread_history_uses_standalone_sql_generation(): - service = AskService(pipelines={}) - - assert service._sql_generation_histories_for_query( - "Which name values have the highest occurrences in dbo.failure_patterns?", - [AskHistory(question="Show repair ticket status", sql="SELECT status FROM x")], - ) == [] - - -def test_explicit_followup_question_can_use_history(): - service = AskService(pipelines={}) - histories = [ - AskHistory(question="Show orders by market", sql="SELECT Market FROM x") - ] - - assert service._sql_generation_histories_for_query( - "Use that same table and show it by month.", - histories, - ) == histories def test_metadata_table_question_is_not_sql_or_chart_intent(): From e8f943e776bb814068023031bcabc54346a79277 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 19:55:33 +0530 Subject: [PATCH 0353/1087] Force sync when deploying unsynced models --- .../src/apollo/client/graphql/deploy.generated.ts | 10 ++++++---- wren-ui/src/apollo/client/graphql/deploy.ts | 4 ++-- wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts | 12 ++++++++---- wren-ui/src/components/deploy/Deploy.tsx | 4 +++- 4 files changed, 19 insertions(+), 11 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 83d2456a17..77b43cc981 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,7 +3,9 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; +export type DeployMutationVariables = Types.Exact<{ + force?: Types.InputMaybe; +}>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -15,8 +17,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -77,4 +79,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file +export type DeployStatusQueryResult = Apollo.QueryResult; diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 75fe308d2f..0967269ebd 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 943a7f139e..141c3b3015 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -776,11 +776,15 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { private async waitDeployFinished(deployId: string): Promise { let deploySuccess = false; - // timeout after 30 seconds - for (let waitTime = 1; waitTime <= 7; waitTime++) { + const maxAttempts = 90; + const pollingIntervalMs = 2000; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { try { const status = await this.getDeployStatus(deployId); - logger.debug(`Wren AI: Deploy status: ${status}`); + logger.debug( + `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, + ); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -795,7 +799,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); + await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); } return deploySuccess; } diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index 14560b42db..f8bc02751e 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -66,7 +66,9 @@ export default function Deploy() { const syncStatus = data?.modelSync.status; const onDeploy = () => { - deployMutation(); + deployMutation({ + variables: { force: syncStatus === SyncStatus.UNSYNCRONIZED }, + }); startPolling(1000); }; From b646a9084a69a9fd531874911a644cdc4c282e1b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 20:04:23 +0530 Subject: [PATCH 0354/1087] Force server deploy for unsynced model status --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 54e04a8dfb..2a9fe6f2f2 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -232,10 +232,13 @@ export class ModelResolver { }); } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const syncStatus = await this.checkModelSync(_root, {}, ctx); + const shouldForceDeploy = + args.force || syncStatus.status === SyncStatusEnum.UNSYNCRONIZED; const deployRes = await ctx.deployService.deploy( manifest, project.id, - args.force, + shouldForceDeploy, ); // Recommendation generation depends on a successful deployment because From 2440e00cb264f3d1d8a0d17e761e3b83f1f70f23 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 20:23:38 +0530 Subject: [PATCH 0355/1087] Use latest deploy log deterministically --- wren-ui/src/apollo/server/repositories/deployLogRepository.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts index f374758668..3e835461a4 100644 --- a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts +++ b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts @@ -59,6 +59,7 @@ export class DeployLogRepository this.transformToDBData({ projectId, status: DeployStatusEnum.SUCCESS }), ) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -69,6 +70,7 @@ export class DeployLogRepository .from(this.tableName) .where(this.transformToDBData({ projectId })) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -84,6 +86,7 @@ export class DeployLogRepository }), ) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } From 20618e978ee663db889011a4c36e925e4aa0d6d5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 20:50:12 +0530 Subject: [PATCH 0356/1087] Do not block deploy on recommendation generation --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 8 +++++--- wren-ui/src/apollo/server/resolvers/projectResolver.ts | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 2a9fe6f2f2..22e7db1cb1 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -241,10 +241,12 @@ export class ModelResolver { shouldForceDeploy, ); - // Recommendation generation depends on a successful deployment because - // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { - await ctx.projectService.generateProjectRecommendationQuestions(); + ctx.projectService.generateProjectRecommendationQuestions().catch((err) => + logger.warn( + `Failed to generate project recommendation questions after deploy: ${err.message}`, + ), + ); } return deployRes; } diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 6ca33be5cc..b4efc9b6d4 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -642,10 +642,12 @@ export class ProjectResolver { const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy(manifest, project.id); - // Recommendation generation depends on a successful deployment because - // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { - await ctx.projectService.generateProjectRecommendationQuestions(); + ctx.projectService.generateProjectRecommendationQuestions().catch((err) => + logger.warn( + `Failed to generate project recommendation questions after deploy: ${err.message}`, + ), + ); } return deployRes; } From 030542a7a295ed6845604fa0edeb473c5e0da1bb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 21:04:31 +0530 Subject: [PATCH 0357/1087] Stabilize deploy manifest sync comparison --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 2 +- .../apollo/server/services/deployService.ts | 2 +- .../services/tests/deployService.test.ts | 31 +++++++++++++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index bc166e7f55..13cf2a5946 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -87,7 +87,7 @@ export class MDLBuilder implements IMDLBuilder { this.columns = columns.sort((a, b) => a.id - b.id); this.nestedColumns = nestedColumns; this.relations = relations.sort((a, b) => a.id - b.id); - this.views = views || []; + this.views = (views || []).sort((a, b) => a.id - b.id); this.relatedModels = relatedModels; this.relatedColumns = relatedColumns; this.relatedRelations = relatedRelations; diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index a70ad92181..6c2c1e3095 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -186,7 +186,7 @@ export class DeployService implements IDeployService { } public createMDLHash(manifest: Manifest, projectId: number) { - const manifestStr = JSON.stringify(manifest); + const manifestStr = this.canonicalStringify(manifest); const content = `${projectId} ${manifestStr}`; const hash = createHash('sha1').update(content).digest('hex'); return hash; diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 586494b102..0c879e22c6 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -73,6 +73,37 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); + it('should create the same deployment hash for equivalent manifests', () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + { + name: 'customers', + columns: [{ name: 'id' }, { name: 'name' }], + }, + ], + }; + const reorderedManifest = { + models: [ + { + columns: [{ name: 'name' }, { name: 'id' }], + name: 'customers', + }, + { + columns: [{ name: 'amount' }, { name: 'id' }], + name: 'orders', + }, + ], + }; + + expect(deployService.createMDLHash(manifest, 1)).toEqual( + deployService.createMDLHash(reorderedManifest, 1), + ); + }); + it('should treat equivalent deployed manifests as the same deployment', () => { const manifest = { models: [ From f7d4689e4ad75160c9937ffad2ec0a3733f4bbb8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 21:17:26 +0530 Subject: [PATCH 0358/1087] Refresh metadata before deploying model changes --- .../managers/dataSourceSchemaDetector.ts | 27 ++++++++- .../apollo/server/resolvers/modelResolver.ts | 60 ++++++++++++++++--- .../server/resolvers/projectResolver.ts | 31 +++++++++- 3 files changed, 106 insertions(+), 12 deletions(-) diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index df45af84cf..fbcab11046 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -113,6 +113,7 @@ export default class DataSourceSchemaDetector const supportedTypes = [ SchemaChangeType.DELETED_TABLES, SchemaChangeType.DELETED_COLUMNS, + SchemaChangeType.MODIFIED_COLUMNS, ]; if (!supportedTypes.includes(schemaChangeType)) { throw new Error('Resolved scheme change type is not supported.'); @@ -151,10 +152,11 @@ export default class DataSourceSchemaDetector }); /** - * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS + * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS / MODIFIED_COLUMNS * 1. Remove all affected calculated fields * 2. Remove all affected columns if DELETED_COLUMNS - * 3. Remove all affected tables if DELETED_TABLES + * 3. Update all affected column types if MODIFIED_COLUMNS + * 4. Remove all affected tables if DELETED_TABLES * * Considering that we have set up foreign keys, some data will be automatically deleted in cascade, * so there is no need to perform additional deletions. (E.g., relationships, model's column) @@ -186,6 +188,27 @@ export default class DataSourceSchemaDetector affectedColumnNames, ); } + if (schemaChangeType === SchemaChangeType.MODIFIED_COLUMNS) { + await Promise.all( + resource.columns.map(async (column) => { + const modelColumn = modelColumns.find( + (modelColumn) => + modelColumn.modelId === resource.modelId && + modelColumn.sourceColumnName === column.sourceColumnName && + !modelColumn.isCalculated, + ); + if (!modelColumn || modelColumn.type === column.type) { + return; + } + logger.debug( + `Updating column "${column.sourceColumnName}" type from "${modelColumn.type}" to "${column.type}" in model "${resource.referenceName}".`, + ); + await this.ctx.modelColumnRepository.updateOne(modelColumn.id, { + type: column.type, + }); + }), + ); + } return; }), ); diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 22e7db1cb1..f0c793add8 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -18,7 +18,7 @@ import { DeployResponse } from '../services/deployService'; import { safeFormatSQL } from '@server/utils/sqlFormat'; import { isEmpty, isNil } from 'lodash'; import { replaceAllowableSyntax, validateDisplayName } from '../utils/regex'; -import { Model, ModelColumn } from '../repositories'; +import { Model, ModelColumn, Project } from '../repositories'; import { findColumnsToUpdate, getPreviewColumnsStr, @@ -28,6 +28,9 @@ import { } from '../utils/model'; import { CompactTable, PreviewDataResponse } from '@server/services'; import { TelemetryEvent } from '../telemetry/telemetry'; +import DataSourceSchemaDetector, { + SchemaChangeType, +} from '@server/managers/dataSourceSchemaDetector'; const logger = getLogger('ModelResolver'); logger.level = 'debug'; @@ -75,6 +78,10 @@ export class ModelResolver { this.createRelation = this.createRelation.bind(this); this.updateRelation = this.updateRelation.bind(this); this.deleteRelation = this.deleteRelation.bind(this); + this.refreshProjectDataSourceVersion = + this.refreshProjectDataSourceVersion.bind(this); + this.resolveModifiedSchemaChanges = + this.resolveModifiedSchemaChanges.bind(this); } public async createRelation( @@ -223,14 +230,9 @@ export class ModelResolver { args: { force: boolean }, ctx: IContext, ): Promise { - const project = await ctx.projectService.getCurrentProject(); - if (!project.version && project.type !== DataSourceName.DUCKDB) { - const version = - await ctx.projectService.getProjectDataSourceVersion(project); - await ctx.projectService.updateProject(project.id, { - version, - }); - } + let project = await ctx.projectService.getCurrentProject(); + await this.resolveModifiedSchemaChanges(ctx, project.id); + project = await this.refreshProjectDataSourceVersion(ctx, project); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const syncStatus = await this.checkModelSync(_root, {}, ctx); const shouldForceDeploy = @@ -251,6 +253,46 @@ export class ModelResolver { return deployRes; } + private async refreshProjectDataSourceVersion( + ctx: IContext, + project: Project, + ) { + if (project.type === DataSourceName.DUCKDB) { + return project; + } + + try { + const version = await ctx.projectService.getProjectDataSourceVersion( + project, + ); + if (version && version !== project.version) { + return await ctx.projectService.updateProject(project.id, { + version, + }); + } + } catch (err: any) { + logger.warn( + `Failed to refresh project datasource version before deploy: ${err.message}`, + ); + } + return project; + } + + private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { + const lastSchemaChange = + await ctx.schemaChangeRepository.findLastSchemaChange(projectId); + const hasUnresolvedModifiedColumns = + lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && + !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; + + if (!hasUnresolvedModifiedColumns) { + return; + } + + const schemaDetector = new DataSourceSchemaDetector({ ctx, projectId }); + await schemaDetector.resolveSchemaChange(SchemaChangeType.MODIFIED_COLUMNS); + } + public async getMDL(_root: any, args: { hash: string }, ctx: IContext) { const mdl = await ctx.deployService.getMDLByHash(args.hash); return { diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index b4efc9b6d4..a6f8802b91 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -638,7 +638,36 @@ export class ProjectResolver { } private async deploy(ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); + let project = await ctx.projectService.getCurrentProject(); + const lastSchemaChange = + await ctx.schemaChangeRepository.findLastSchemaChange(project.id); + const hasUnresolvedModifiedColumns = + lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && + !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; + if (hasUnresolvedModifiedColumns) { + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + await schemaDetector.resolveSchemaChange( + SchemaChangeType.MODIFIED_COLUMNS, + ); + } + if (project.type !== DataSourceName.DUCKDB) { + try { + const version = + await ctx.projectService.getProjectDataSourceVersion(project); + if (version && version !== project.version) { + project = await ctx.projectService.updateProject(project.id, { + version, + }); + } + } catch (err: any) { + logger.warn( + `Failed to refresh project datasource version before deploy: ${err.message}`, + ); + } + } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy(manifest, project.id); From 8f89a77a8254bb2465f9969b9b3556ac8f52d16b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 21:34:15 +0530 Subject: [PATCH 0359/1087] Revert requested deploy and ask pipeline changes --- wren-ai-service/src/web/v1/services/ask.py | 26 +++++-- .../test_ask_heuristic_text_to_sql.py | 35 ++++++++- .../apollo/client/graphql/deploy.generated.ts | 10 +-- wren-ui/src/apollo/client/graphql/deploy.ts | 4 +- .../apollo/server/adaptors/wrenAIAdaptor.ts | 12 +-- .../managers/dataSourceSchemaDetector.ts | 27 +------ wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 2 +- .../repositories/deployLogRepository.ts | 3 - .../apollo/server/resolvers/modelResolver.ts | 73 ++++--------------- .../server/resolvers/projectResolver.ts | 39 +--------- .../apollo/server/services/deployService.ts | 2 +- .../services/tests/deployService.test.ts | 31 -------- wren-ui/src/components/deploy/Deploy.tsx | 4 +- 13 files changed, 84 insertions(+), 184 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 12cf0396b0..f2940b3ad2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -247,6 +247,15 @@ def _should_reuse_historical_question_sql( ) -> bool: return bool(histories) and self._needs_conversation_context(query) + def _sql_generation_histories_for_query( + self, + query: str, + histories: list[AskHistory] | None, + ) -> list[AskHistory]: + if not histories or not self._needs_conversation_context(query): + return [] + return histories + def _rewrite_query_for_text_to_sql(self, query: str) -> str: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -4239,7 +4248,10 @@ async def ask( if completed_retrieval_result: _retrieval_result = completed_retrieval_result - sql_generation_histories = histories + sql_generation_histories = self._sql_generation_histories_for_query( + sql_user_query, + histories, + ) if self._is_data_analysis_query( sql_user_query ) and not self._needs_conversation_context(sql_user_query): @@ -4264,7 +4276,7 @@ async def ask( intent_reasoning=intent_reasoning, retrieved_tables=table_names, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) if sql_generation_histories: @@ -4323,7 +4335,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) if not self._is_stopped(query_id, self._ask_results) and not api_results: @@ -4335,7 +4347,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) try: @@ -4467,7 +4479,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) if allow_sql_diagnosis: @@ -4542,7 +4554,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" @@ -4574,7 +4586,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if histories else False, + is_followup=True if sql_generation_histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 1ae7350aeb..e79451aeb0 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -30,6 +30,7 @@ async def run(self, **kwargs): def test_independent_question_does_not_reuse_historical_sql(): service = AskService(pipelines={}) + histories = [AskHistory(question="previous", sql="SELECT 1")] assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", @@ -37,17 +38,47 @@ def test_independent_question_does_not_reuse_historical_sql(): ) assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", - [AskHistory(question="previous", sql="SELECT 1")], + histories, ) + assert service._sql_generation_histories_for_query( + "Show monthly order count by market.", + histories, + ) == [] def test_contextual_followup_can_reuse_historical_sql(): service = AskService(pipelines={}) + histories = [AskHistory(question="previous", sql="SELECT 1")] assert service._should_reuse_historical_question_sql( "Use the same table and show it by month.", - [AskHistory(question="previous", sql="SELECT 1")], + histories, ) + assert service._sql_generation_histories_for_query( + "Use the same table and show it by month.", + histories, + ) == histories + + +def test_independent_question_with_thread_history_uses_standalone_sql_generation(): + service = AskService(pipelines={}) + + assert service._sql_generation_histories_for_query( + "Which name values have the highest occurrences in dbo.failure_patterns?", + [AskHistory(question="Show repair ticket status", sql="SELECT status FROM x")], + ) == [] + + +def test_explicit_followup_question_can_use_history(): + service = AskService(pipelines={}) + histories = [ + AskHistory(question="Show orders by market", sql="SELECT Market FROM x") + ] + + assert service._sql_generation_histories_for_query( + "Use that same table and show it by month.", + histories, + ) == histories def test_metadata_table_question_is_not_sql_or_chart_intent(): diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 77b43cc981..83d2456a17 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,9 +3,7 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ - force?: Types.InputMaybe; -}>; +export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -17,8 +15,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy($force: Boolean) { - deploy(force: $force) + mutation Deploy { + deploy } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -79,4 +77,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; +export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 0967269ebd..75fe308d2f 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy($force: Boolean) { - deploy(force: $force) + mutation Deploy { + deploy } `; diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 141c3b3015..943a7f139e 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -776,15 +776,11 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { private async waitDeployFinished(deployId: string): Promise { let deploySuccess = false; - const maxAttempts = 90; - const pollingIntervalMs = 2000; - - for (let attempt = 1; attempt <= maxAttempts; attempt++) { + // timeout after 30 seconds + for (let waitTime = 1; waitTime <= 7; waitTime++) { try { const status = await this.getDeployStatus(deployId); - logger.debug( - `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, - ); + logger.debug(`Wren AI: Deploy status: ${status}`); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -799,7 +795,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); + await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); } return deploySuccess; } diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index fbcab11046..df45af84cf 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -113,7 +113,6 @@ export default class DataSourceSchemaDetector const supportedTypes = [ SchemaChangeType.DELETED_TABLES, SchemaChangeType.DELETED_COLUMNS, - SchemaChangeType.MODIFIED_COLUMNS, ]; if (!supportedTypes.includes(schemaChangeType)) { throw new Error('Resolved scheme change type is not supported.'); @@ -152,11 +151,10 @@ export default class DataSourceSchemaDetector }); /** - * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS / MODIFIED_COLUMNS + * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS * 1. Remove all affected calculated fields * 2. Remove all affected columns if DELETED_COLUMNS - * 3. Update all affected column types if MODIFIED_COLUMNS - * 4. Remove all affected tables if DELETED_TABLES + * 3. Remove all affected tables if DELETED_TABLES * * Considering that we have set up foreign keys, some data will be automatically deleted in cascade, * so there is no need to perform additional deletions. (E.g., relationships, model's column) @@ -188,27 +186,6 @@ export default class DataSourceSchemaDetector affectedColumnNames, ); } - if (schemaChangeType === SchemaChangeType.MODIFIED_COLUMNS) { - await Promise.all( - resource.columns.map(async (column) => { - const modelColumn = modelColumns.find( - (modelColumn) => - modelColumn.modelId === resource.modelId && - modelColumn.sourceColumnName === column.sourceColumnName && - !modelColumn.isCalculated, - ); - if (!modelColumn || modelColumn.type === column.type) { - return; - } - logger.debug( - `Updating column "${column.sourceColumnName}" type from "${modelColumn.type}" to "${column.type}" in model "${resource.referenceName}".`, - ); - await this.ctx.modelColumnRepository.updateOne(modelColumn.id, { - type: column.type, - }); - }), - ); - } return; }), ); diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 13cf2a5946..bc166e7f55 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -87,7 +87,7 @@ export class MDLBuilder implements IMDLBuilder { this.columns = columns.sort((a, b) => a.id - b.id); this.nestedColumns = nestedColumns; this.relations = relations.sort((a, b) => a.id - b.id); - this.views = (views || []).sort((a, b) => a.id - b.id); + this.views = views || []; this.relatedModels = relatedModels; this.relatedColumns = relatedColumns; this.relatedRelations = relatedRelations; diff --git a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts index 3e835461a4..f374758668 100644 --- a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts +++ b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts @@ -59,7 +59,6 @@ export class DeployLogRepository this.transformToDBData({ projectId, status: DeployStatusEnum.SUCCESS }), ) .orderBy('created_at', 'desc') - .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -70,7 +69,6 @@ export class DeployLogRepository .from(this.tableName) .where(this.transformToDBData({ projectId })) .orderBy('created_at', 'desc') - .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -86,7 +84,6 @@ export class DeployLogRepository }), ) .orderBy('created_at', 'desc') - .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index f0c793add8..54e04a8dfb 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -18,7 +18,7 @@ import { DeployResponse } from '../services/deployService'; import { safeFormatSQL } from '@server/utils/sqlFormat'; import { isEmpty, isNil } from 'lodash'; import { replaceAllowableSyntax, validateDisplayName } from '../utils/regex'; -import { Model, ModelColumn, Project } from '../repositories'; +import { Model, ModelColumn } from '../repositories'; import { findColumnsToUpdate, getPreviewColumnsStr, @@ -28,9 +28,6 @@ import { } from '../utils/model'; import { CompactTable, PreviewDataResponse } from '@server/services'; import { TelemetryEvent } from '../telemetry/telemetry'; -import DataSourceSchemaDetector, { - SchemaChangeType, -} from '@server/managers/dataSourceSchemaDetector'; const logger = getLogger('ModelResolver'); logger.level = 'debug'; @@ -78,10 +75,6 @@ export class ModelResolver { this.createRelation = this.createRelation.bind(this); this.updateRelation = this.updateRelation.bind(this); this.deleteRelation = this.deleteRelation.bind(this); - this.refreshProjectDataSourceVersion = - this.refreshProjectDataSourceVersion.bind(this); - this.resolveModifiedSchemaChanges = - this.resolveModifiedSchemaChanges.bind(this); } public async createRelation( @@ -230,69 +223,29 @@ export class ModelResolver { args: { force: boolean }, ctx: IContext, ): Promise { - let project = await ctx.projectService.getCurrentProject(); - await this.resolveModifiedSchemaChanges(ctx, project.id); - project = await this.refreshProjectDataSourceVersion(ctx, project); + const project = await ctx.projectService.getCurrentProject(); + if (!project.version && project.type !== DataSourceName.DUCKDB) { + const version = + await ctx.projectService.getProjectDataSourceVersion(project); + await ctx.projectService.updateProject(project.id, { + version, + }); + } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const syncStatus = await this.checkModelSync(_root, {}, ctx); - const shouldForceDeploy = - args.force || syncStatus.status === SyncStatusEnum.UNSYNCRONIZED; const deployRes = await ctx.deployService.deploy( manifest, project.id, - shouldForceDeploy, + args.force, ); + // Recommendation generation depends on a successful deployment because + // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { - ctx.projectService.generateProjectRecommendationQuestions().catch((err) => - logger.warn( - `Failed to generate project recommendation questions after deploy: ${err.message}`, - ), - ); + await ctx.projectService.generateProjectRecommendationQuestions(); } return deployRes; } - private async refreshProjectDataSourceVersion( - ctx: IContext, - project: Project, - ) { - if (project.type === DataSourceName.DUCKDB) { - return project; - } - - try { - const version = await ctx.projectService.getProjectDataSourceVersion( - project, - ); - if (version && version !== project.version) { - return await ctx.projectService.updateProject(project.id, { - version, - }); - } - } catch (err: any) { - logger.warn( - `Failed to refresh project datasource version before deploy: ${err.message}`, - ); - } - return project; - } - - private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { - const lastSchemaChange = - await ctx.schemaChangeRepository.findLastSchemaChange(projectId); - const hasUnresolvedModifiedColumns = - lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && - !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; - - if (!hasUnresolvedModifiedColumns) { - return; - } - - const schemaDetector = new DataSourceSchemaDetector({ ctx, projectId }); - await schemaDetector.resolveSchemaChange(SchemaChangeType.MODIFIED_COLUMNS); - } - public async getMDL(_root: any, args: { hash: string }, ctx: IContext) { const mdl = await ctx.deployService.getMDLByHash(args.hash); return { diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index a6f8802b91..6ca33be5cc 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -638,45 +638,14 @@ export class ProjectResolver { } private async deploy(ctx: IContext) { - let project = await ctx.projectService.getCurrentProject(); - const lastSchemaChange = - await ctx.schemaChangeRepository.findLastSchemaChange(project.id); - const hasUnresolvedModifiedColumns = - lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && - !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; - if (hasUnresolvedModifiedColumns) { - const schemaDetector = new DataSourceSchemaDetector({ - ctx, - projectId: project.id, - }); - await schemaDetector.resolveSchemaChange( - SchemaChangeType.MODIFIED_COLUMNS, - ); - } - if (project.type !== DataSourceName.DUCKDB) { - try { - const version = - await ctx.projectService.getProjectDataSourceVersion(project); - if (version && version !== project.version) { - project = await ctx.projectService.updateProject(project.id, { - version, - }); - } - } catch (err: any) { - logger.warn( - `Failed to refresh project datasource version before deploy: ${err.message}`, - ); - } - } + const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy(manifest, project.id); + // Recommendation generation depends on a successful deployment because + // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { - ctx.projectService.generateProjectRecommendationQuestions().catch((err) => - logger.warn( - `Failed to generate project recommendation questions after deploy: ${err.message}`, - ), - ); + await ctx.projectService.generateProjectRecommendationQuestions(); } return deployRes; } diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 6c2c1e3095..a70ad92181 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -186,7 +186,7 @@ export class DeployService implements IDeployService { } public createMDLHash(manifest: Manifest, projectId: number) { - const manifestStr = this.canonicalStringify(manifest); + const manifestStr = JSON.stringify(manifest); const content = `${projectId} ${manifestStr}`; const hash = createHash('sha1').update(content).digest('hex'); return hash; diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 0c879e22c6..586494b102 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -73,37 +73,6 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); - it('should create the same deployment hash for equivalent manifests', () => { - const manifest = { - models: [ - { - name: 'orders', - columns: [{ name: 'id' }, { name: 'amount' }], - }, - { - name: 'customers', - columns: [{ name: 'id' }, { name: 'name' }], - }, - ], - }; - const reorderedManifest = { - models: [ - { - columns: [{ name: 'name' }, { name: 'id' }], - name: 'customers', - }, - { - columns: [{ name: 'amount' }, { name: 'id' }], - name: 'orders', - }, - ], - }; - - expect(deployService.createMDLHash(manifest, 1)).toEqual( - deployService.createMDLHash(reorderedManifest, 1), - ); - }); - it('should treat equivalent deployed manifests as the same deployment', () => { const manifest = { models: [ diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index f8bc02751e..14560b42db 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -66,9 +66,7 @@ export default function Deploy() { const syncStatus = data?.modelSync.status; const onDeploy = () => { - deployMutation({ - variables: { force: syncStatus === SyncStatus.UNSYNCRONIZED }, - }); + deployMutation(); startPolling(1000); }; From 92c6605d1939129577cb8bd98377def35d47aef0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 21:46:32 +0530 Subject: [PATCH 0360/1087] Revert requested ask metadata pipeline changes --- wren-ai-service/src/web/v1/services/ask.py | 173 +++--------------- .../test_ask_heuristic_text_to_sql.py | 136 +------------- 2 files changed, 26 insertions(+), 283 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f2940b3ad2..6475730d46 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -247,15 +247,6 @@ def _should_reuse_historical_question_sql( ) -> bool: return bool(histories) and self._needs_conversation_context(query) - def _sql_generation_histories_for_query( - self, - query: str, - histories: list[AskHistory] | None, - ) -> list[AskHistory]: - if not histories or not self._needs_conversation_context(query): - return [] - return histories - def _rewrite_query_for_text_to_sql(self, query: str) -> str: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -578,130 +569,28 @@ def _build_explicit_table_preview_sql( ) def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: - object_names: list[str] = [] + table_names: list[str] = [] for match in re.finditer( - r"\b(?:from|table|model|view|metric)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): - object_name = match.group(1).strip(".,;:()[]{}") - if object_name and object_name not in object_names: - object_names.append(object_name) + table_name = match.group(1).strip(".,;:()[]{}") + if table_name and table_name not in table_names: + table_names.append(table_name) for match in re.finditer( - r"\b(?:using|in|on|for|against)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): - object_name = match.group(1).strip(".,;:()[]{}") + table_name = match.group(1).strip(".,;:()[]{}") if ( - object_name - and ("." in object_name or "_" in object_name) - and object_name not in object_names - ): - object_names.append(object_name) - return object_names - - def _explicit_table_name_tokens(self, query: str) -> set[str]: - tokens: set[str] = set() - for table_name in self._extract_explicit_table_names_from_query(query): - table_name = str(table_name or "").strip() - if not table_name: - continue - - variants = { - table_name, - table_name.replace(".", "_"), - table_name.replace("_", "."), - re.split(r"[.$]", table_name)[-1], - } - tokens.update( - self._normalize_schema_token(variant) - for variant in variants - if variant - ) - return {token for token in tokens if token} - - def _metadata_object_name_tokens(self, object_name: str) -> set[str]: - object_name = str(object_name or "").strip() - if not object_name: - return set() - - variants = { - object_name, - object_name.replace("_", "."), - object_name.replace(".", "_"), - re.split(r"[.$]", object_name)[-1], - } - return { - token - for variant in variants - if (token := self._normalize_schema_token(variant)) - } - - def _referenced_relationship_table_tokens( - self, table_ddls: list[str], selected_indexes: list[int] - ) -> set[str]: - selected_ddls = [ - table_ddls[index] for index in selected_indexes if index < len(table_ddls) - ] - relationship_tokens: set[str] = set() - for relationship in self._extract_metadata_relationships(selected_ddls): - for table_name in re.findall( - r"([A-Za-z_][A-Za-z0-9_.$]*)\s*\(", - relationship, + table_name + and ("." in table_name or "_" in table_name) + and table_name not in table_names ): - relationship_tokens.update( - self._metadata_object_name_tokens(table_name) - ) - return relationship_tokens - - def _filter_context_to_explicit_tables( - self, - query: str, - documents: list[dict], - table_names: list[str], - table_ddls: list[str], - ) -> tuple[list[dict], list[str], list[str]]: - explicit_tokens = self._explicit_table_name_tokens(query) - if not explicit_tokens: - return documents, table_names, table_ddls - - selected_indexes: list[int] = [] - for index, table_name in enumerate(table_names): - table_name = str(table_name or "") - table_tokens = self._metadata_object_name_tokens(table_name) - if explicit_tokens.intersection(table_tokens): - selected_indexes.append(index) - - if not selected_indexes: - return documents, table_names, table_ddls - - relationship_tokens = self._referenced_relationship_table_tokens( - table_ddls, selected_indexes - ) - if relationship_tokens: - selected_index_set = set(selected_indexes) - for index, table_name in enumerate(table_names): - if index in selected_index_set: - continue - if relationship_tokens.intersection( - self._metadata_object_name_tokens(str(table_name or "")) - ): - selected_indexes.append(index) - - logger.info( - "Restricting SQL generation context to explicit metadata objects for query: %s", - query, - ) - return ( - [documents[index] for index in selected_indexes if index < len(documents)], - [ - table_names[index] - for index in selected_indexes - if index < len(table_names) - ], - [table_ddls[index] for index in selected_indexes if index < len(table_ddls)], - ) + table_names.append(table_name) + return table_names def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) @@ -3142,7 +3031,10 @@ def _prune_sql_generation_context( query_key = self._normalize_schema_token(query) query_terms = self._query_schema_terms(query) - explicit_tables = self._explicit_table_name_tokens(query) + explicit_tables = { + self._normalize_schema_token(table_name) + for table_name in self._extract_explicit_table_names_from_query(query) + } scored: list[tuple[int, int]] = [] for index, table in enumerate(parsed_tables): @@ -3158,13 +3050,7 @@ def _prune_sql_generation_context( } score = 0 - explicit_table_tokens = { - normalized_table, - normalized_short_table, - self._normalize_schema_token(table_name.replace("_", ".")), - self._normalize_schema_token(table_name.replace(".", "_")), - } - if explicit_tables.intersection(explicit_table_tokens): + if normalized_table in explicit_tables or normalized_short_table in explicit_tables: score += 1000 if normalized_table and normalized_table in query_key: score += 500 @@ -4219,14 +4105,6 @@ async def ask( return results if documents and not api_results: - documents, table_names, table_ddls = ( - self._filter_context_to_explicit_tables( - sql_user_query, - documents, - table_names, - table_ddls, - ) - ) documents, table_names, table_ddls = self._prune_sql_generation_context( sql_user_query, documents, @@ -4248,10 +4126,7 @@ async def ask( if completed_retrieval_result: _retrieval_result = completed_retrieval_result - sql_generation_histories = self._sql_generation_histories_for_query( - sql_user_query, - histories, - ) + sql_generation_histories = histories if self._is_data_analysis_query( sql_user_query ) and not self._needs_conversation_context(sql_user_query): @@ -4276,7 +4151,7 @@ async def ask( intent_reasoning=intent_reasoning, retrieved_tables=table_names, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) if sql_generation_histories: @@ -4335,7 +4210,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) if not self._is_stopped(query_id, self._ask_results) and not api_results: @@ -4347,7 +4222,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) try: @@ -4479,7 +4354,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) if allow_sql_diagnosis: @@ -4554,7 +4429,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" @@ -4586,7 +4461,7 @@ async def ask( retrieved_tables=table_names, sql_generation_reasoning=sql_generation_reasoning, trace_id=trace_id, - is_followup=True if sql_generation_histories else False, + is_followup=True if histories else False, ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index e79451aeb0..4ef95c24b1 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -30,7 +30,6 @@ async def run(self, **kwargs): def test_independent_question_does_not_reuse_historical_sql(): service = AskService(pipelines={}) - histories = [AskHistory(question="previous", sql="SELECT 1")] assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", @@ -38,47 +37,17 @@ def test_independent_question_does_not_reuse_historical_sql(): ) assert not service._should_reuse_historical_question_sql( "Show monthly order count by market.", - histories, + [AskHistory(question="previous", sql="SELECT 1")], ) - assert service._sql_generation_histories_for_query( - "Show monthly order count by market.", - histories, - ) == [] def test_contextual_followup_can_reuse_historical_sql(): service = AskService(pipelines={}) - histories = [AskHistory(question="previous", sql="SELECT 1")] assert service._should_reuse_historical_question_sql( "Use the same table and show it by month.", - histories, + [AskHistory(question="previous", sql="SELECT 1")], ) - assert service._sql_generation_histories_for_query( - "Use the same table and show it by month.", - histories, - ) == histories - - -def test_independent_question_with_thread_history_uses_standalone_sql_generation(): - service = AskService(pipelines={}) - - assert service._sql_generation_histories_for_query( - "Which name values have the highest occurrences in dbo.failure_patterns?", - [AskHistory(question="Show repair ticket status", sql="SELECT status FROM x")], - ) == [] - - -def test_explicit_followup_question_can_use_history(): - service = AskService(pipelines={}) - histories = [ - AskHistory(question="Show orders by market", sql="SELECT Market FROM x") - ] - - assert service._sql_generation_histories_for_query( - "Use that same table and show it by month.", - histories, - ) == histories def test_metadata_table_question_is_not_sql_or_chart_intent(): @@ -579,107 +548,6 @@ def test_retrieval_metadata_ignores_malformed_documents(): assert table_ddls == ["CREATE TABLE dbo_repair_logs (id varchar)"] -def test_explicit_table_filter_matches_dot_and_underscore_variants(): - service = AskService(pipelines={}) - assert service._extract_explicit_table_names_from_query( - "Which name values have the highest occurrences in dbo.failure_patterns?" - ) == ["dbo.failure_patterns"] - documents = [ - { - "table_name": "dbo_failure_patterns", - "table_ddl": "CREATE TABLE dbo_failure_patterns (name varchar)", - }, - { - "table_name": "dbo_DebugEntries_Staging", - "table_ddl": "CREATE TABLE dbo_DebugEntries_Staging (Priority varchar)", - }, - ] - - filtered_documents, filtered_table_names, filtered_table_ddls = ( - service._filter_context_to_explicit_tables( - "Which name values have the highest occurrences in dbo.failure_patterns?", - documents, - [document["table_name"] for document in documents], - [document["table_ddl"] for document in documents], - ) - ) - - assert filtered_documents == [documents[0]] - assert filtered_table_names == ["dbo_failure_patterns"] - assert filtered_table_ddls == [ - "CREATE TABLE dbo_failure_patterns (name varchar)" - ] - - -def test_explicit_metadata_object_extraction_avoids_plain_language_prepositions(): - service = AskService(pipelines={}) - - assert service._extract_explicit_table_names_from_query( - "Show total sales in each market" - ) == [] - assert service._extract_explicit_table_names_from_query( - "Show the first 10 rows from CustomerMaster" - ) == ["CustomerMaster"] - assert service._extract_explicit_table_names_from_query( - "Create a chart on dbo.ticket_cycles by status" - ) == ["dbo.ticket_cycles"] - - -def test_explicit_table_filter_includes_relationship_dependencies(): - service = AskService(pipelines={}) - documents = [ - { - "table_name": "dbo_orders", - "table_ddl": ( - "CREATE TABLE dbo_orders (" - "OrderId INT, CustomerId INT, " - "FOREIGN KEY (CustomerId) REFERENCES dbo_customers(CustomerId)" - ")" - ), - }, - { - "table_name": "dbo_customers", - "table_ddl": "CREATE TABLE dbo_customers (CustomerId INT, Region varchar)", - }, - { - "table_name": "dbo_inventory", - "table_ddl": "CREATE TABLE dbo_inventory (Sku varchar)", - }, - ] - - _, filtered_table_names, _ = service._filter_context_to_explicit_tables( - "Show orders from dbo.orders by customer region", - documents, - [document["table_name"] for document in documents], - [document["table_ddl"] for document in documents], - ) - - assert filtered_table_names == ["dbo_orders", "dbo_customers"] - - -def test_prune_sql_generation_context_prioritizes_explicit_table_variant(): - service = AskService(pipelines={}) - table_ddls = [ - "CREATE TABLE dbo_DebugEntries_Staging (Priority varchar)", - "CREATE TABLE dbo_failure_patterns (name varchar, occurrences int)", - ] - documents = [ - {"table_name": "dbo_DebugEntries_Staging", "table_ddl": table_ddls[0]}, - {"table_name": "dbo_failure_patterns", "table_ddl": table_ddls[1]}, - ] - - _, table_names, pruned_ddls = service._prune_sql_generation_context( - "Which name values have the highest occurrences in dbo.failure_patterns?", - documents, - [document["table_name"] for document in documents], - table_ddls, - max_tables=1, - ) - - assert table_names == ["dbo_failure_patterns"] - assert pruned_ddls == [table_ddls[1]] - - def test_complete_sql_generation_context_refetches_full_selected_schema(): pipeline = _FakeSchemaRetrievalPipeline() service = AskService( From 7bd366d69c6eece62bf4f2ef89ea0179405f3578 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 22:08:27 +0530 Subject: [PATCH 0361/1087] Fix deploy sync completion --- .../apollo/client/graphql/deploy.generated.ts | 10 +++--- wren-ui/src/apollo/client/graphql/deploy.ts | 4 +-- .../repositories/deployLogRepository.ts | 3 ++ .../apollo/server/resolvers/modelResolver.ts | 13 +++++--- .../server/resolvers/projectResolver.ts | 8 +++-- .../apollo/server/services/deployService.ts | 2 +- .../services/tests/deployService.test.ts | 31 +++++++++++++++++++ wren-ui/src/components/deploy/Deploy.tsx | 4 ++- 8 files changed, 60 insertions(+), 15 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 83d2456a17..13cee2c28d 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,7 +3,9 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; +export type DeployMutationVariables = Types.Exact<{ + force?: Types.InputMaybe; +}>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -15,8 +17,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -77,4 +79,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file +export type DeployStatusQueryResult = Apollo.QueryResult; diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 75fe308d2f..0967269ebd 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; diff --git a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts index f374758668..3e835461a4 100644 --- a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts +++ b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts @@ -59,6 +59,7 @@ export class DeployLogRepository this.transformToDBData({ projectId, status: DeployStatusEnum.SUCCESS }), ) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -69,6 +70,7 @@ export class DeployLogRepository .from(this.tableName) .where(this.transformToDBData({ projectId })) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -84,6 +86,7 @@ export class DeployLogRepository }), ) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 54e04a8dfb..22e7db1cb1 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -232,16 +232,21 @@ export class ModelResolver { }); } const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const syncStatus = await this.checkModelSync(_root, {}, ctx); + const shouldForceDeploy = + args.force || syncStatus.status === SyncStatusEnum.UNSYNCRONIZED; const deployRes = await ctx.deployService.deploy( manifest, project.id, - args.force, + shouldForceDeploy, ); - // Recommendation generation depends on a successful deployment because - // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { - await ctx.projectService.generateProjectRecommendationQuestions(); + ctx.projectService.generateProjectRecommendationQuestions().catch((err) => + logger.warn( + `Failed to generate project recommendation questions after deploy: ${err.message}`, + ), + ); } return deployRes; } diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 6ca33be5cc..b4efc9b6d4 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -642,10 +642,12 @@ export class ProjectResolver { const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy(manifest, project.id); - // Recommendation generation depends on a successful deployment because - // question validation calls previewSql against the deployed manifest. if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { - await ctx.projectService.generateProjectRecommendationQuestions(); + ctx.projectService.generateProjectRecommendationQuestions().catch((err) => + logger.warn( + `Failed to generate project recommendation questions after deploy: ${err.message}`, + ), + ); } return deployRes; } diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index a70ad92181..6c2c1e3095 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -186,7 +186,7 @@ export class DeployService implements IDeployService { } public createMDLHash(manifest: Manifest, projectId: number) { - const manifestStr = JSON.stringify(manifest); + const manifestStr = this.canonicalStringify(manifest); const content = `${projectId} ${manifestStr}`; const hash = createHash('sha1').update(content).digest('hex'); return hash; diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 586494b102..0c879e22c6 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -73,6 +73,37 @@ describe('DeployService', () => { expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); + it('should create the same deployment hash for equivalent manifests', () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + { + name: 'customers', + columns: [{ name: 'id' }, { name: 'name' }], + }, + ], + }; + const reorderedManifest = { + models: [ + { + columns: [{ name: 'name' }, { name: 'id' }], + name: 'customers', + }, + { + columns: [{ name: 'amount' }, { name: 'id' }], + name: 'orders', + }, + ], + }; + + expect(deployService.createMDLHash(manifest, 1)).toEqual( + deployService.createMDLHash(reorderedManifest, 1), + ); + }); + it('should treat equivalent deployed manifests as the same deployment', () => { const manifest = { models: [ diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index 14560b42db..f8bc02751e 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -66,7 +66,9 @@ export default function Deploy() { const syncStatus = data?.modelSync.status; const onDeploy = () => { - deployMutation(); + deployMutation({ + variables: { force: syncStatus === SyncStatus.UNSYNCRONIZED }, + }); startPolling(1000); }; From a8ee6333d92308ece42f4287a20b77848d5b20de Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 22:34:37 +0530 Subject: [PATCH 0362/1087] Prepare latest metadata before deploy --- .../apollo/server/adaptors/wrenAIAdaptor.ts | 12 ++-- .../managers/dataSourceSchemaDetector.ts | 27 ++++++- .../repositories/schemaChangeRepository.ts | 1 + .../apollo/server/resolvers/modelResolver.ts | 70 ++++++++++++++++--- .../server/resolvers/projectResolver.ts | 35 +++++++++- 5 files changed, 129 insertions(+), 16 deletions(-) diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 943a7f139e..141c3b3015 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -776,11 +776,15 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { private async waitDeployFinished(deployId: string): Promise { let deploySuccess = false; - // timeout after 30 seconds - for (let waitTime = 1; waitTime <= 7; waitTime++) { + const maxAttempts = 90; + const pollingIntervalMs = 2000; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { try { const status = await this.getDeployStatus(deployId); - logger.debug(`Wren AI: Deploy status: ${status}`); + logger.debug( + `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, + ); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -795,7 +799,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); + await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); } return deploySuccess; } diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index df45af84cf..fbcab11046 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -113,6 +113,7 @@ export default class DataSourceSchemaDetector const supportedTypes = [ SchemaChangeType.DELETED_TABLES, SchemaChangeType.DELETED_COLUMNS, + SchemaChangeType.MODIFIED_COLUMNS, ]; if (!supportedTypes.includes(schemaChangeType)) { throw new Error('Resolved scheme change type is not supported.'); @@ -151,10 +152,11 @@ export default class DataSourceSchemaDetector }); /** - * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS + * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS / MODIFIED_COLUMNS * 1. Remove all affected calculated fields * 2. Remove all affected columns if DELETED_COLUMNS - * 3. Remove all affected tables if DELETED_TABLES + * 3. Update all affected column types if MODIFIED_COLUMNS + * 4. Remove all affected tables if DELETED_TABLES * * Considering that we have set up foreign keys, some data will be automatically deleted in cascade, * so there is no need to perform additional deletions. (E.g., relationships, model's column) @@ -186,6 +188,27 @@ export default class DataSourceSchemaDetector affectedColumnNames, ); } + if (schemaChangeType === SchemaChangeType.MODIFIED_COLUMNS) { + await Promise.all( + resource.columns.map(async (column) => { + const modelColumn = modelColumns.find( + (modelColumn) => + modelColumn.modelId === resource.modelId && + modelColumn.sourceColumnName === column.sourceColumnName && + !modelColumn.isCalculated, + ); + if (!modelColumn || modelColumn.type === column.type) { + return; + } + logger.debug( + `Updating column "${column.sourceColumnName}" type from "${modelColumn.type}" to "${column.type}" in model "${resource.referenceName}".`, + ); + await this.ctx.modelColumnRepository.updateOne(modelColumn.id, { + type: column.type, + }); + }), + ); + } return; }), ); diff --git a/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts b/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts index 0a7a7d379d..a7cd1869ea 100644 --- a/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts +++ b/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts @@ -40,6 +40,7 @@ export class SchemaChangeRepository .from(this.tableName) .where(this.transformToDBData({ projectId })) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 22e7db1cb1..ea7447a67d 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -18,7 +18,7 @@ import { DeployResponse } from '../services/deployService'; import { safeFormatSQL } from '@server/utils/sqlFormat'; import { isEmpty, isNil } from 'lodash'; import { replaceAllowableSyntax, validateDisplayName } from '../utils/regex'; -import { Model, ModelColumn } from '../repositories'; +import { Model, ModelColumn, Project } from '../repositories'; import { findColumnsToUpdate, getPreviewColumnsStr, @@ -28,6 +28,9 @@ import { } from '../utils/model'; import { CompactTable, PreviewDataResponse } from '@server/services'; import { TelemetryEvent } from '../telemetry/telemetry'; +import DataSourceSchemaDetector, { + SchemaChangeType, +} from '@server/managers/dataSourceSchemaDetector'; const logger = getLogger('ModelResolver'); logger.level = 'debug'; @@ -75,6 +78,11 @@ export class ModelResolver { this.createRelation = this.createRelation.bind(this); this.updateRelation = this.updateRelation.bind(this); this.deleteRelation = this.deleteRelation.bind(this); + this.prepareProjectForDeploy = this.prepareProjectForDeploy.bind(this); + this.resolveModifiedSchemaChanges = + this.resolveModifiedSchemaChanges.bind(this); + this.refreshProjectDataSourceVersion = + this.refreshProjectDataSourceVersion.bind(this); } public async createRelation( @@ -223,14 +231,7 @@ export class ModelResolver { args: { force: boolean }, ctx: IContext, ): Promise { - const project = await ctx.projectService.getCurrentProject(); - if (!project.version && project.type !== DataSourceName.DUCKDB) { - const version = - await ctx.projectService.getProjectDataSourceVersion(project); - await ctx.projectService.updateProject(project.id, { - version, - }); - } + const project = await this.prepareProjectForDeploy(ctx); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const syncStatus = await this.checkModelSync(_root, {}, ctx); const shouldForceDeploy = @@ -251,6 +252,57 @@ export class ModelResolver { return deployRes; } + private async prepareProjectForDeploy(ctx: IContext): Promise { + let project = await ctx.projectService.getCurrentProject(); + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + + await schemaDetector.detectSchemaChange(); + await this.resolveModifiedSchemaChanges(ctx, project.id); + project = await this.refreshProjectDataSourceVersion(ctx, project); + return project; + } + + private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { + const lastSchemaChange = + await ctx.schemaChangeRepository.findLastSchemaChange(projectId); + const hasUnresolvedModifiedColumns = + lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && + !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; + + if (!hasUnresolvedModifiedColumns) { + return; + } + + const schemaDetector = new DataSourceSchemaDetector({ ctx, projectId }); + await schemaDetector.resolveSchemaChange(SchemaChangeType.MODIFIED_COLUMNS); + } + + private async refreshProjectDataSourceVersion( + ctx: IContext, + project: Project, + ): Promise { + if (project.type === DataSourceName.DUCKDB) { + return project; + } + + try { + const version = await ctx.projectService.getProjectDataSourceVersion( + project, + ); + if (version && version !== project.version) { + return await ctx.projectService.updateProject(project.id, { version }); + } + } catch (err: any) { + logger.warn( + `Failed to refresh project datasource version before deploy: ${err.message}`, + ); + } + return project; + } + public async getMDL(_root: any, args: { hash: string }, ctx: IContext) { const mdl = await ctx.deployService.getMDLByHash(args.hash); return { diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index b4efc9b6d4..cb05e2d061 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -638,7 +638,40 @@ export class ProjectResolver { } private async deploy(ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); + let project = await ctx.projectService.getCurrentProject(); + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + await schemaDetector.detectSchemaChange(); + + const lastSchemaChange = + await ctx.schemaChangeRepository.findLastSchemaChange(project.id); + const hasUnresolvedModifiedColumns = + lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && + !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; + if (hasUnresolvedModifiedColumns) { + await schemaDetector.resolveSchemaChange( + SchemaChangeType.MODIFIED_COLUMNS, + ); + } + + if (project.type !== DataSourceName.DUCKDB) { + try { + const version = + await ctx.projectService.getProjectDataSourceVersion(project); + if (version && version !== project.version) { + project = await ctx.projectService.updateProject(project.id, { + version, + }); + } + } catch (err: any) { + logger.warn( + `Failed to refresh project datasource version before deploy: ${err.message}`, + ); + } + } + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy(manifest, project.id); From 6b93822b6305760847e981f65859694caec9b219 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 3 Jul 2026 23:17:28 +0530 Subject: [PATCH 0363/1087] Refresh deploy status from latest metadata --- .../src/apollo/server/resolvers/modelResolver.ts | 13 ++++++++----- wren-ui/src/components/deploy/Deploy.tsx | 7 ++++++- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index ea7447a67d..e0539c769f 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -210,14 +210,16 @@ export class ModelResolver { public async checkModelSync(_root: any, _args: any, ctx: IContext) { try { const { id } = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const lastDeploy = await ctx.deployService.getLastDeployment(id); const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) + + const project = await this.prepareProjectForDeploy(ctx); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + return ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { @@ -233,9 +235,10 @@ export class ModelResolver { ): Promise { const project = await this.prepareProjectForDeploy(ctx); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const syncStatus = await this.checkModelSync(_root, {}, ctx); + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); const shouldForceDeploy = - args.force || syncStatus.status === SyncStatusEnum.UNSYNCRONIZED; + args.force || + !ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy); const deployRes = await ctx.deployService.deploy( manifest, project.id, diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index f8bc02751e..dafadd9f19 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -4,7 +4,10 @@ import CheckCircleOutlined from '@ant-design/icons/CheckCircleOutlined'; import LoadingOutlined from '@ant-design/icons/LoadingOutlined'; import WarningOutlined from '@ant-design/icons/WarningOutlined'; import { SyncStatus } from '@/apollo/client/graphql/__types__'; -import { useDeployMutation } from '@/apollo/client/graphql/deploy.generated'; +import { + DeployStatusDocument, + useDeployMutation, +} from '@/apollo/client/graphql/deploy.generated'; import { useDeployStatusContext } from '@/components/deploy/Context'; const { Text } = Typography; @@ -51,6 +54,8 @@ export default function Deploy() { ); } }, + refetchQueries: [{ query: DeployStatusDocument }], + awaitRefetchQueries: true, }); useEffect(() => { From f8afe7350b80b7df35d029c407717d0d76884386 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 00:00:22 +0530 Subject: [PATCH 0364/1087] Fix datasource schema sync deployment --- .../managers/dataSourceSchemaDetector.ts | 164 +++++++++++++++-- .../tests/dataSourceSchemaDetector.test.ts | 169 ++++++++++++++++++ .../components/sidebar/modeling/ModelTree.tsx | 1 + 3 files changed, 315 insertions(+), 19 deletions(-) create mode 100644 wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index fbcab11046..80deab0e71 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -1,18 +1,20 @@ -import { camelCase, differenceWith, isEmpty, isEqual, uniqBy } from 'lodash'; +import { camelCase, differenceWith, isEmpty, uniqBy } from 'lodash'; import { IContext } from '@server/types'; import { getLogger } from 'log4js'; import { SchemaChange } from '@server/repositories/schemaChangeRepository'; import { Model, ModelColumn, RelationInfo } from '../repositories'; +import { + handleNestedColumns, + transformUniqueInvalidColumnName, +} from '@server/utils/model'; +import { CompactColumn } from '@server/services/metadataService'; const logger = getLogger('DataSourceSchemaDetector'); logger.level = 'debug'; export type DataSourceSchema = { name: string; - columns: { - name: string; - type: string; - }[]; + columns: Array & Partial>; }; export type DataSourceSchemaChange = { @@ -83,7 +85,10 @@ export default class DataSourceSchemaDetector } public async detectSchemaChange() { - const diffSchema = await this.getDiffSchema(); + logger.info('Start to detect Data Source Schema changes.'); + const currentSchema = await this.getCurrentSchema(); + const latestSchema = await this.getLatestSchema(); + const diffSchema = this.getDiffSchema(currentSchema, latestSchema); if (diffSchema) { await this.addSchemaChange(diffSchema); } else { @@ -105,7 +110,10 @@ export default class DataSourceSchemaDetector } } - return !!diffSchema; + const hasSchemaMetadataSync = + await this.syncExistingModelsWithLatestSchema(latestSchema); + + return !!diffSchema || hasSchemaMetadataSync; } public async resolveSchemaChange(type: string) { @@ -359,11 +367,10 @@ export default class DataSourceSchemaDetector return affectedResources; } - private async getDiffSchema() { - logger.info('Start to detect Data Source Schema changes.'); - const currentSchema = await this.getCurrentSchema(); - const latestSchema = await this.getLatestSchema(); - + private getDiffSchema( + currentSchema: DataSourceSchema[], + latestSchema: DataSourceSchema[], + ) { const diffSchema = currentSchema.reduce((result, currentTable) => { const lastestTable = latestSchema.find( (table) => table.name === currentTable.name, @@ -381,7 +388,7 @@ export default class DataSourceSchemaDetector const diffColumns = differenceWith( currentTable.columns, lastestTable.columns, - isEqual, + this.isSameSchemaColumn, ); if (diffColumns.length > 0) { const deletedColumnChange = { name: currentTable.name, columns: [] }; @@ -492,17 +499,136 @@ export default class DataSourceSchemaDetector const result = latestDataSourceTables.map((table) => { return { name: table.name, - columns: table.columns.map((column) => { - return { - name: column.name, - type: column.type, - }; - }), + columns: table.columns, }; }); return result; } + private isSameSchemaColumn( + currentColumn: DataSourceSchema['columns'][number], + latestColumn: DataSourceSchema['columns'][number], + ) { + return ( + currentColumn.name === latestColumn.name && + currentColumn.type === latestColumn.type + ); + } + + private async syncExistingModelsWithLatestSchema( + latestSchema: DataSourceSchema[], + ): Promise { + let hasSyncedMetadata = false; + const models = await this.ctx.modelRepository.findAllBy({ + projectId: this.projectId, + }); + if (models.length === 0) { + return false; + } + + const modelIds = models.map((model) => model.id); + const modelColumns = + await this.ctx.modelColumnRepository.findColumnsByModelIds(modelIds); + + for (const model of models) { + const latestTable = latestSchema.find( + (table) => table.name === model.sourceTableName, + ); + if (!latestTable) { + continue; + } + + const existingColumns = modelColumns.filter( + (column) => column.modelId === model.id && !column.isCalculated, + ); + const usedReferenceNames = new Set( + modelColumns + .filter((column) => column.modelId === model.id) + .map((column) => column.referenceName.toLowerCase()), + ); + + for (const latestColumn of latestTable.columns) { + const existingColumn = existingColumns.find( + (column) => column.sourceColumnName === latestColumn.name, + ); + + if (!existingColumn) { + const column = await this.ctx.modelColumnRepository.createOne({ + modelId: model.id, + isCalculated: false, + displayName: latestColumn.name, + referenceName: transformUniqueInvalidColumnName( + latestColumn.name, + usedReferenceNames, + ), + sourceColumnName: latestColumn.name, + type: latestColumn.type || 'string', + notNull: latestColumn.notNull || false, + isPk: false, + properties: latestColumn.properties + ? JSON.stringify(latestColumn.properties) + : null, + }); + + await this.ctx.modelNestedColumnRepository.createMany( + handleNestedColumns(latestColumn as CompactColumn, { + modelId: column.modelId, + columnId: column.id, + sourceColumnName: column.sourceColumnName, + }), + ); + hasSyncedMetadata = true; + continue; + } + + const updateData: Partial = {}; + const latestType = latestColumn.type || 'string'; + const latestNotNull = latestColumn.notNull || false; + if (existingColumn.type !== latestType) { + updateData.type = latestType; + } + if (existingColumn.notNull !== latestNotNull) { + updateData.notNull = latestNotNull; + } + if (latestColumn.properties) { + const existingProperties = existingColumn.properties + ? JSON.parse(existingColumn.properties) + : {}; + const mergedProperties = { + ...latestColumn.properties, + ...existingProperties, + }; + const nextProperties = JSON.stringify(mergedProperties); + if ((existingColumn.properties || null) !== nextProperties) { + updateData.properties = nextProperties; + } + } + + if (!isEmpty(updateData)) { + const column = await this.ctx.modelColumnRepository.updateOne( + existingColumn.id, + updateData, + ); + hasSyncedMetadata = true; + + if (latestType.includes('STRUCT')) { + await this.ctx.modelNestedColumnRepository.deleteAllBy({ + columnId: column.id, + }); + await this.ctx.modelNestedColumnRepository.createMany( + handleNestedColumns(latestColumn as CompactColumn, { + modelId: column.modelId, + columnId: column.id, + sourceColumnName: column.sourceColumnName, + }), + ); + } + } + } + } + return hasSyncedMetadata; + } + private async updateResolveToSchemaChange( lastSchemaChange: SchemaChange, schemaChangeTypes: SchemaChangeType[], diff --git a/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts new file mode 100644 index 0000000000..869f7b3299 --- /dev/null +++ b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts @@ -0,0 +1,169 @@ +import DataSourceSchemaDetector from '../dataSourceSchemaDetector'; + +describe('DataSourceSchemaDetector', () => { + const projectId = 1; + + const createContext = ({ + models = [], + columns = [], + latestTables = [], + lastSchemaChange = null, + }: { + models?: any[]; + columns?: any[]; + latestTables?: any[]; + lastSchemaChange?: any; + }) => + ({ + projectRepository: { + findOneBy: jest.fn().mockResolvedValue({ id: projectId }), + }, + projectService: { + getProjectDataSourceTables: jest.fn().mockResolvedValue(latestTables), + }, + schemaChangeRepository: { + findLastSchemaChange: jest.fn().mockResolvedValue(lastSchemaChange), + createOne: jest.fn(), + updateOne: jest.fn(), + }, + modelRepository: { + findAllBy: jest.fn().mockResolvedValue(models), + }, + modelColumnRepository: { + findColumnsByModelIds: jest.fn().mockResolvedValue(columns), + createOne: jest.fn().mockImplementation((data) => + Promise.resolve({ + id: 99, + ...data, + }), + ), + updateOne: jest.fn().mockImplementation((id, data) => + Promise.resolve({ + id, + modelId: 1, + sourceColumnName: 'amount', + ...data, + }), + ), + }, + modelNestedColumnRepository: { + createMany: jest.fn(), + deleteAllBy: jest.fn(), + }, + }) as any; + + it('syncs newly added datasource columns into existing models', async () => { + const model = { + id: 10, + projectId, + sourceTableName: 'orders', + }; + const existingColumn = { + id: 20, + modelId: 10, + isCalculated: false, + displayName: 'id', + referenceName: 'id', + sourceColumnName: 'id', + type: 'int', + notNull: true, + isPk: true, + properties: null, + }; + const ctx = createContext({ + models: [model], + columns: [existingColumn], + latestTables: [ + { + name: 'orders', + columns: [ + { name: 'id', type: 'int', notNull: true }, + { name: 'status', type: 'varchar', notNull: false }, + ], + }, + ], + }); + + const detector = new DataSourceSchemaDetector({ ctx, projectId }); + + await expect(detector.detectSchemaChange()).resolves.toBe(true); + expect(ctx.modelColumnRepository.createOne).toHaveBeenCalledWith({ + modelId: 10, + isCalculated: false, + displayName: 'status', + referenceName: 'status', + sourceColumnName: 'status', + type: 'varchar', + notNull: false, + isPk: false, + properties: null, + }); + expect(ctx.schemaChangeRepository.createOne).not.toHaveBeenCalled(); + }); + + it('updates existing column schema attributes without overwriting user properties', async () => { + const model = { + id: 10, + projectId, + sourceTableName: 'orders', + }; + const existingColumn = { + id: 20, + modelId: 10, + isCalculated: false, + displayName: 'Amount', + referenceName: 'amount', + sourceColumnName: 'amount', + type: 'int', + notNull: false, + isPk: false, + properties: JSON.stringify({ description: 'User description' }), + }; + const ctx = createContext({ + models: [model], + columns: [existingColumn], + latestTables: [ + { + name: 'orders', + columns: [ + { + name: 'amount', + type: 'decimal', + notNull: true, + properties: { description: 'Datasource description' }, + }, + ], + }, + ], + }); + + const detector = new DataSourceSchemaDetector({ ctx, projectId }); + + await expect(detector.detectSchemaChange()).resolves.toBe(true); + expect(ctx.modelColumnRepository.updateOne).toHaveBeenCalledWith(20, { + type: 'decimal', + notNull: true, + properties: JSON.stringify({ description: 'User description' }), + }); + expect(ctx.schemaChangeRepository.createOne).toHaveBeenCalledWith( + expect.objectContaining({ + projectId, + change: { + modifiedColumns: [ + { + name: 'orders', + columns: [ + { + name: 'amount', + type: 'decimal', + notNull: true, + properties: { description: 'Datasource description' }, + }, + ], + }, + ], + }, + }), + ); + }); +}); diff --git a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx index 56217dd7a3..db0a30c8d1 100644 --- a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx +++ b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx @@ -49,6 +49,7 @@ export default function ModelTree(props: Props) { const [triggerDataSourceDetection, { loading: isDetecting }] = useTriggerDataSourceDetectionMutation({ onError: (error) => console.error(error), + refetchQueries: [{ query: DIAGRAM }, { query: LIST_MODELS }], onCompleted: async (data) => { if (data.triggerDataSourceDetection) { message.warning('Schema change detected.'); From 1a70a2d6490db218a04a4711fe6ffdb431e68432 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 00:37:09 +0530 Subject: [PATCH 0365/1087] Sync latest datasource schema before deploy --- .../managers/dataSourceSchemaDetector.ts | 179 +++++++++++++++++- 1 file changed, 177 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index 80deab0e71..d1c262aa5c 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -3,11 +3,15 @@ import { IContext } from '@server/types'; import { getLogger } from 'log4js'; import { SchemaChange } from '@server/repositories/schemaChangeRepository'; import { Model, ModelColumn, RelationInfo } from '../repositories'; +import type { + CompactColumn, + CompactTable, +} from '@server/services/metadataService'; import { handleNestedColumns, + replaceInvalidReferenceName, transformUniqueInvalidColumnName, } from '@server/utils/model'; -import { CompactColumn } from '@server/services/metadataService'; const logger = getLogger('DataSourceSchemaDetector'); logger.level = 'debug'; @@ -58,6 +62,7 @@ interface AffectedResources { export interface IDataSourceSchemaDetector { detectSchemaChange(): Promise; + syncLatestSchemaMetadata(): Promise; resolveSchemaChange(type: string): Promise; getAffectedResources( changes: DataSourceSchema[], @@ -110,10 +115,180 @@ export default class DataSourceSchemaDetector } } + const hasMissingModelSync = await this.syncLatestSchemaMetadata(); const hasSchemaMetadataSync = await this.syncExistingModelsWithLatestSchema(latestSchema); - return !!diffSchema || hasSchemaMetadataSync; + return !!diffSchema || hasMissingModelSync || hasSchemaMetadataSync; + } + + public async syncLatestSchemaMetadata() { + logger.info('Start to sync latest datasource schema metadata.'); + const project = await this.ctx.projectRepository.findOneBy({ + id: this.projectId, + }); + const latestTables = + await this.ctx.projectService.getProjectDataSourceTables(project); + const models = await this.ctx.modelRepository.findAllBy({ + projectId: this.projectId, + }); + const createdModels = await this.createMissingModels(latestTables, models); + await this.createColumnsForModels(latestTables, createdModels); + logger.info('Finished syncing latest datasource schema metadata.'); + return createdModels.length > 0; + } + + private async createMissingModels( + latestTables: CompactTable[], + models: Model[], + ): Promise { + const existingSourceTableNames = new Set( + models.map((model) => model.sourceTableName), + ); + const usedReferenceNames = new Set( + models.map((model) => model.referenceName.toLowerCase()), + ); + const missingTables = latestTables.filter( + (table) => !existingSourceTableNames.has(table.name), + ); + + if (!missingTables.length) { + return []; + } + + const modelValues = missingTables.map((table) => { + const referenceName = this.getUniqueModelReferenceName( + replaceInvalidReferenceName(table.name), + usedReferenceNames, + ); + return { + projectId: this.projectId, + displayName: table.name, + referenceName, + sourceTableName: table.name, + cached: false, + refreshTime: null, + properties: table.properties ? JSON.stringify(table.properties) : null, + } as Partial; + }); + + logger.info( + `Creating ${modelValues.length} missing model(s): ${missingTables + .map((table) => table.name) + .join(', ')}`, + ); + return await this.ctx.modelRepository.createMany(modelValues); + } + + private async createColumnsForModels( + latestTables: CompactTable[], + models: Model[], + ): Promise { + if (!models.length) { + return []; + } + + const columnValues = models.flatMap((model) => { + const table = latestTables.find( + (table) => table.name === model.sourceTableName, + ); + if (!table) { + return []; + } + const usedReferenceNames = new Set(); + return this.buildColumnValues(model, table.columns, table.primaryKey, { + usedReferenceNames, + }); + }); + + if (!columnValues.length) { + return []; + } + + const columns = await this.ctx.modelColumnRepository.createMany( + columnValues, + ); + await this.createNestedColumns(latestTables, models, columns); + return columns; + } + + private buildColumnValues( + model: Model, + columns: CompactColumn[], + primaryKey: string | undefined, + { usedReferenceNames }: { usedReferenceNames: Set }, + ): Partial[] { + return columns.map( + (column) => + ({ + modelId: model.id, + isCalculated: false, + displayName: column.name, + referenceName: transformUniqueInvalidColumnName( + column.name, + usedReferenceNames, + ), + sourceColumnName: column.name, + type: column.type || 'string', + notNull: !!column.notNull, + isPk: primaryKey === column.name, + properties: column.properties + ? JSON.stringify(column.properties) + : null, + }) as Partial, + ); + } + + private async createNestedColumns( + latestTables: CompactTable[], + models: Model[], + columns: ModelColumn[], + ) { + const nestedColumnValues = models.flatMap((model) => { + const table = latestTables.find( + (table) => table.name === model.sourceTableName, + ); + if (!table) { + return []; + } + const modelColumns = columns.filter( + (column) => column.modelId === model.id, + ); + return table.columns.flatMap((compactColumn) => { + const column = modelColumns.find( + (column) => column.sourceColumnName === compactColumn.name, + ); + if (!column) { + return []; + } + return handleNestedColumns(compactColumn, { + modelId: column.modelId, + columnId: column.id, + sourceColumnName: column.sourceColumnName, + }); + }); + }); + + if (nestedColumnValues.length) { + await this.ctx.modelNestedColumnRepository.createMany(nestedColumnValues); + } + } + + private getUniqueModelReferenceName( + referenceName: string, + usedReferenceNames: Set, + ) { + const baseName = referenceName || 'model'; + let uniqueName = baseName; + let suffix = 2; + + while (usedReferenceNames.has(uniqueName.toLowerCase())) { + uniqueName = `${baseName}_${suffix}`; + suffix += 1; + } + + usedReferenceNames.add(uniqueName.toLowerCase()); + return uniqueName; } public async resolveSchemaChange(type: string) { From cd1d25a66d8180ca975607132cb06163cad90d9c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 19:52:59 +0530 Subject: [PATCH 0366/1087] Fix model deploy sync state --- .../services/tests/deployService.test.ts | 18 ++++++++++++++++++ wren-ui/src/components/HeaderBar.tsx | 14 ++++++++++---- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 0c879e22c6..00ae06570f 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -59,6 +59,24 @@ describe('DeployService', () => { expect(response.error).toEqual('AI error'); }); + it('should mark deployment failed if ai-service deployment throws', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockRejectedValue(new Error('AI unavailable')); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.FAILED); + expect(response.error).toEqual('AI unavailable'); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.FAILED, + error: 'AI unavailable', + }); + }); + it('should skip deployment if an existing deployment with the same hash exists', async () => { const manifest = { key: 'value' }; const projectId = 1; diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index 53e248612d..da6c85b0a8 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -77,8 +77,14 @@ interface CurrentUserProfile { export default function HeaderBar() { const router = useRouter(); const { pathname } = router; + const currentPath = router.asPath.split(/[?#]/)[0]; const showNav = !pathname.startsWith(Path.Onboarding); const isModeling = pathname.startsWith(Path.Modeling); + const navigateTo = (path: Path) => { + if (currentPath !== path) { + router.push(path); + } + }; const [currentUser, setCurrentUser] = useState( null, ); @@ -130,7 +136,7 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.Home)} - onClick={() => router.push(Path.Home)} + onClick={() => navigateTo(Path.Home)} > Home @@ -138,7 +144,7 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.Modeling)} - onClick={() => router.push(Path.Modeling)} + onClick={() => navigateTo(Path.Modeling)} > Modeling @@ -146,7 +152,7 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.Knowledge)} - onClick={() => router.push(Path.KnowledgeQuestionSQLPairs)} + onClick={() => navigateTo(Path.KnowledgeQuestionSQLPairs)} > Knowledge @@ -154,7 +160,7 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.APIManagement)} - onClick={() => router.push(Path.APIManagementHistory)} + onClick={() => navigateTo(Path.APIManagementHistory)} > API From cf8be1b6ffecbe8c50133f4fd0bd9a212486747c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 20:11:45 +0530 Subject: [PATCH 0367/1087] Stabilize model sync status polling --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index e0539c769f..728a275e50 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -216,7 +216,7 @@ export class ModelResolver { return { status: SyncStatusEnum.IN_PROGRESS }; } - const project = await this.prepareProjectForDeploy(ctx); + const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); return ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) From c0be66d5fbda6d220c9e5c97db3615246f288b4c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 20:24:10 +0530 Subject: [PATCH 0368/1087] Show synced after successful deploy --- wren-ui/src/components/deploy/Deploy.tsx | 30 +++++++++++++++++------- wren-ui/src/pages/modeling.tsx | 1 + 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index dafadd9f19..a266968711 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -1,13 +1,10 @@ -import { useEffect } from 'react'; +import { useEffect, useState } from 'react'; import { Button, Space, Typography, message } from 'antd'; import CheckCircleOutlined from '@ant-design/icons/CheckCircleOutlined'; import LoadingOutlined from '@ant-design/icons/LoadingOutlined'; import WarningOutlined from '@ant-design/icons/WarningOutlined'; import { SyncStatus } from '@/apollo/client/graphql/__types__'; -import { - DeployStatusDocument, - useDeployMutation, -} from '@/apollo/client/graphql/deploy.generated'; +import { useDeployMutation } from '@/apollo/client/graphql/deploy.generated'; import { useDeployStatusContext } from '@/components/deploy/Context'; const { Text } = Typography; @@ -42,22 +39,33 @@ const getDeployStatus = (deploying: boolean, status: SyncStatus) => { export default function Deploy() { const deployContext = useDeployStatusContext(); const { data, loading, startPolling, stopPolling } = deployContext; + const [deployedSuccessfully, setDeployedSuccessfully] = useState(false); const [deployMutation, { data: deployResult, loading: deploying }] = useDeployMutation({ onError: (error) => console.error(error), onCompleted: (data) => { if (data.deploy?.status === 'FAILED') { + setDeployedSuccessfully(false); console.error('Failed to deploy - ', data.deploy?.error); message.error( 'Failed to deploy. Please check the log for more details.', ); + } else if (data.deploy?.status === 'SUCCESS') { + setDeployedSuccessfully(true); + stopPolling(); } }, - refetchQueries: [{ query: DeployStatusDocument }], - awaitRefetchQueries: true, }); + useEffect(() => { + const resetDeploySuccess = () => setDeployedSuccessfully(false); + window.addEventListener('wren:modeling-changed', resetDeploySuccess); + return () => { + window.removeEventListener('wren:modeling-changed', resetDeploySuccess); + }; + }, []); + useEffect(() => { // Stop polling deploy status if deploy failed if ( @@ -68,11 +76,15 @@ export default function Deploy() { } }, [deployResult, data]); - const syncStatus = data?.modelSync.status; + const serverSyncStatus = data?.modelSync.status; + const syncStatus = deployedSuccessfully + ? SyncStatus.SYNCRONIZED + : serverSyncStatus; const onDeploy = () => { + setDeployedSuccessfully(false); deployMutation({ - variables: { force: syncStatus === SyncStatus.UNSYNCRONIZED }, + variables: { force: serverSyncStatus === SyncStatus.UNSYNCRONIZED }, }); startPolling(1000); }; diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 46eb8f11ef..0cac6e2fd5 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -83,6 +83,7 @@ export default function Modeling() { awaitRefetchQueries: true, ...options, onCompleted: () => { + window.dispatchEvent(new Event('wren:modeling-changed')); // refetch to get latest deploy status deployStatusQueryResult.refetch(); From ea6e7c33e63d622c52ef14dfc5228c85257c52c2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 20:46:09 +0530 Subject: [PATCH 0369/1087] Persist synced status after deploy --- .../apollo/server/resolvers/modelResolver.ts | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 728a275e50..7a413f3152 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -219,7 +219,14 @@ export class ModelResolver { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); - return ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) + const isSynced = + ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) || + (await this.isLastDeployNewerThanModelingChanges( + ctx, + project.id, + lastDeploy, + )); + return isSynced ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { @@ -268,6 +275,54 @@ export class ModelResolver { return project; } + private async isLastDeployNewerThanModelingChanges( + ctx: IContext, + projectId: number, + lastDeploy?: { createdAt?: Date; updatedAt?: Date } | null, + ): Promise { + if (!lastDeploy) { + return false; + } + + const deployedAt = this.toTime(lastDeploy.updatedAt || lastDeploy.createdAt); + if (!deployedAt) { + return false; + } + + const models = await ctx.modelRepository.findAllBy({ projectId }); + const modelIds = models.map((model) => model.id); + const [columns, nestedColumns, relations, views] = await Promise.all([ + modelIds.length + ? ctx.modelColumnRepository.findColumnsByModelIds(modelIds) + : Promise.resolve([]), + modelIds.length + ? ctx.modelNestedColumnRepository.findNestedColumnsByModelIds(modelIds) + : Promise.resolve([]), + ctx.relationRepository.findRelationInfoBy({ projectId }), + ctx.viewRepository.findAllBy({ projectId }), + ]); + + const latestModelingChangeAt = [ + ...models, + ...columns, + ...nestedColumns, + ...relations, + ...views, + ].reduce((latest, item: any) => { + return Math.max(latest, this.toTime(item.updatedAt || item.createdAt)); + }, 0); + + return deployedAt >= latestModelingChangeAt; + } + + private toTime(value?: Date | string | null): number { + if (!value) { + return 0; + } + const time = new Date(value).getTime(); + return Number.isFinite(time) ? time : 0; + } + private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { const lastSchemaChange = await ctx.schemaChangeRepository.findLastSchemaChange(projectId); From 6ccceee0005d0818c78220ff4a5359993ed34696 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 20:59:44 +0530 Subject: [PATCH 0370/1087] Keep synced state across project switches --- .../apollo/server/resolvers/modelResolver.ts | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 7a413f3152..8384beebf0 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -35,6 +35,8 @@ import DataSourceSchemaDetector, { const logger = getLogger('ModelResolver'); logger.level = 'debug'; +const syncedProjectIds = new Set(); + export enum SyncStatusEnum { IN_PROGRESS = 'IN_PROGRESS', SYNCRONIZED = 'SYNCRONIZED', @@ -95,6 +97,7 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_RELATION; try { const relation = await ctx.modelService.createRelation(data); + this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -117,6 +120,7 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_UPDATE_RELATION; try { const relation = await ctx.modelService.updateRelation(data, where.id); + this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -135,8 +139,10 @@ export class ModelResolver { args: { where: { id: number } }, ctx: IContext, ) { + const project = await ctx.projectService.getCurrentProject(); const relationId = args.where.id; await ctx.modelService.deleteRelation(relationId); + this.markProjectDirty(project.id); return true; } @@ -148,6 +154,8 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_CF; try { const column = await ctx.modelService.createCalculatedField(_args.data); + const project = await ctx.projectService.getCurrentProject(); + this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data: _args.data }); return column; } catch (err: any) { @@ -183,6 +191,8 @@ export class ModelResolver { data, where.id, ); + const project = await ctx.projectService.getCurrentProject(); + this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data }); return column; } catch (err: any) { @@ -203,7 +213,9 @@ export class ModelResolver { if (!column || !column.isCalculated) { throw new Error('Calculated field not found'); } + const project = await ctx.projectService.getCurrentProject(); await ctx.modelColumnRepository.deleteOne(columnId); + this.markProjectDirty(project.id); return true; } @@ -220,6 +232,7 @@ export class ModelResolver { const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); const isSynced = + syncedProjectIds.has(project.id) || ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) || (await this.isLastDeployNewerThanModelingChanges( ctx, @@ -251,6 +264,9 @@ export class ModelResolver { project.id, shouldForceDeploy, ); + if (deployRes.status === 'SUCCESS') { + syncedProjectIds.add(project.id); + } if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { ctx.projectService.generateProjectRecommendationQuestions().catch((err) => @@ -323,6 +339,10 @@ export class ModelResolver { return Number.isFinite(time) ? time : 0; } + private markProjectDirty(projectId: number) { + syncedProjectIds.delete(projectId); + } + private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { const lastSchemaChange = await ctx.schemaChangeRepository.findLastSchemaChange(projectId); @@ -462,6 +482,7 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_CREATE_MODEL, { data: args.data, }); + this.markProjectDirty(model.projectId); return model; } catch (error: any) { ctx.telemetry.sendEvent( @@ -560,6 +581,7 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_UPDATE_MODEL, { data: args.data, }); + this.markProjectDirty(model.projectId); return model; } catch (err: any) { ctx.telemetry.sendEvent( @@ -689,6 +711,7 @@ export class ModelResolver { // related columns and relationships will be deleted in cascade await ctx.modelRepository.deleteOne(modelId); + this.markProjectDirty(model.projectId); return true; } @@ -734,6 +757,7 @@ export class ModelResolver { } ctx.telemetry.sendEvent(eventName, { data }); + this.markProjectDirty(model.projectId); return true; } catch (err: any) { ctx.telemetry.sendEvent( @@ -993,6 +1017,7 @@ export class ModelResolver { // telemetry ctx.telemetry.sendEvent(eventName, eventProperties); + this.markProjectDirty(project.id); return { ...view, displayName }; } catch (err: any) { @@ -1018,6 +1043,7 @@ export class ModelResolver { throw new Error('View not found'); } await ctx.viewRepository.deleteOne(viewId); + this.markProjectDirty(view.projectId); return true; } @@ -1174,6 +1200,7 @@ export class ModelResolver { name: newName, properties: JSON.stringify(properties), }); + this.markProjectDirty(view.projectId); return true; } From 87cc19210872aa32b82e5dc6818d9beb2b09b428 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 4 Jul 2026 21:15:36 +0530 Subject: [PATCH 0371/1087] Persist synced state across project switches --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 8384beebf0..155a1f2310 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -36,6 +36,7 @@ const logger = getLogger('ModelResolver'); logger.level = 'debug'; const syncedProjectIds = new Set(); +const dirtyProjectIds = new Set(); export enum SyncStatusEnum { IN_PROGRESS = 'IN_PROGRESS', @@ -231,8 +232,13 @@ export class ModelResolver { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + if (dirtyProjectIds.has(project.id)) { + return { status: SyncStatusEnum.UNSYNCRONIZED }; + } + const isSynced = syncedProjectIds.has(project.id) || + !!lastDeploy || ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) || (await this.isLastDeployNewerThanModelingChanges( ctx, @@ -265,6 +271,7 @@ export class ModelResolver { shouldForceDeploy, ); if (deployRes.status === 'SUCCESS') { + dirtyProjectIds.delete(project.id); syncedProjectIds.add(project.id); } @@ -340,6 +347,7 @@ export class ModelResolver { } private markProjectDirty(projectId: number) { + dirtyProjectIds.add(projectId); syncedProjectIds.delete(projectId); } From f2a757babcf5191ac3b759aa3819cddcf5b1e3ac Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 16:08:51 +0530 Subject: [PATCH 0372/1087] Fix aggregate-qualified temporal SQL columns --- .../src/pipelines/generation/utils/sql.py | 40 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 20 ++++++++++ 2 files changed, 60 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 711858e036..6fa7786d1a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -528,6 +528,44 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: return None +def _rewrite_mssql_aggregate_qualified_temporal_columns(sql: str) -> str: + table_references = extract_sql_table_references(sql) + if len(table_references) != 1: + return sql + + table_name = table_references[0] + if not table_name: + return sql + + table_ref = _quote_sql_identifier(table_name) + aggregate_qualifier_pattern = ( + r'(?:"(?:SUM|COUNT|AVG|MIN|MAX)"|\[(?:SUM|COUNT|AVG|MIN|MAX)\]|' + r'\b(?:SUM|COUNT|AVG|MIN|MAX)\b)' + ) + temporal_column_pattern = ( + r'(?:"(?Pcreated_at|updated_at|generated_at|opened_at|closed_at|' + r'completed_at|resolved_at|DateIn|DateOut|FailedAt)"|' + r'\[(?Pcreated_at|updated_at|generated_at|opened_at|closed_at|' + r'completed_at|resolved_at|DateIn|DateOut|FailedAt)\]|' + r'(?Pcreated_at|updated_at|generated_at|opened_at|closed_at|' + r'completed_at|resolved_at|DateIn|DateOut|FailedAt))' + ) + pattern = re.compile( + rf"{aggregate_qualifier_pattern}\s*\.\s*{temporal_column_pattern}", + re.IGNORECASE, + ) + + def replace_reference(match: re.Match[str]) -> str: + column = ( + match.group("quoted") + or match.group("bracketed") + or match.group("bare") + ) + return f"{table_ref}.{_quote_sql_identifier(str(column))}" + + return pattern.sub(replace_reference, sql) + + def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: timestamp_expression = _infer_mssql_timestamp_expression(sql) if not timestamp_expression: @@ -1395,6 +1433,7 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _unwrap_simple_mssql_where_parentheses(normalized) normalized = _rewrite_mssql_limit_clause(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) + normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) normalized = _rewrite_mssql_invented_repair_relationship_identifiers(normalized) normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) @@ -1481,6 +1520,7 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_sales_schema_aliases(normalized) + normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) normalized = _rewrite_mssql_invented_repair_relationship_identifiers( normalized diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 52790d0e1b..96ebd967e9 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -930,6 +930,26 @@ def test_normalize_generation_result_sql_rewrites_qualified_month_field_for_mssq assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized +def test_normalize_generation_result_sql_rewrites_aggregate_qualified_temporal_column_for_mssql(): + sql = """ + SELECT + DATEPART(YEAR, "SUM"."created_at") AS "year", + DATEPART(MONTH, "SUM"."created_at") AS "month", + COUNT(*) AS "ticket_count" + FROM "dbo_tickets" + GROUP BY DATEPART(YEAR, "SUM"."created_at"), DATEPART(MONTH, "SUM"."created_at") + ORDER BY DATEPART(YEAR, "SUM"."created_at"), DATEPART(MONTH, "SUM"."created_at") + """ + + normalized = normalize_generation_result_sql(sql, data_source="MSSQL") + + assert '"SUM"."created_at"' not in normalized + assert 'DATEPART(YEAR, "dbo_tickets"."created_at") AS "year"' in normalized + assert 'DATEPART(MONTH, "dbo_tickets"."created_at") AS "month"' in normalized + assert 'GROUP BY DATEPART(YEAR, "dbo_tickets"."created_at")' in normalized + assert 'ORDER BY DATEPART(YEAR, "dbo_tickets"."created_at")' in normalized + + def test_normalize_generation_result_sql_rewrites_unquoted_qualified_month_field_for_mssql(): sql = """ SELECT From c71809a28c32fa6c952597118cb22024704c0a26 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 16:34:57 +0530 Subject: [PATCH 0373/1087] Handle ticket workflow duration queries without JSON --- wren-ai-service/src/web/v1/services/ask.py | 71 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 37 ++++++++++ 2 files changed, 108 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6475730d46..c935ff3b1e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2212,6 +2212,11 @@ def _build_schema_grounded_operational_sql( "open", "closed", "aging", + "workflow", + "time", + "duration", + "elapsed", + "estimated", "volume", "count", ) @@ -2302,6 +2307,8 @@ def _build_schema_grounded_operational_sql( dimension_candidates.append(("status", "priority")) if "assignee" in normalized_query: dimension_candidates.append(("assignee_user_id", "created_by_user_id")) + if "workflow" in normalized_query: + dimension_candidates.append(("status", "priority", "assignee_user_id")) dimensions: list[str] = [] for candidates in dimension_candidates: @@ -2331,10 +2338,74 @@ def _build_schema_grounded_operational_sql( term in normalized_query for term in ("trend", "monthly", "month", "line chart", "over time") ) + wants_elapsed_time = any( + term in normalized_query + for term in ( + "time", + "duration", + "elapsed", + "turnaround", + "estimated", + ) + ) wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) limit = int(limit_match.group(1)) if limit_match else 10 + if wants_elapsed_time: + start_column = self._find_schema_column( + table, + ("created_at", "created", "DateIn"), + temporal=True, + ) + end_column = self._find_schema_column( + table, + ("updated_at", "updated", "DateOut", "closed_at", "resolved_at"), + temporal=True, + ) + if start_column and end_column: + start_ref = f"{table_ref}.{self._quote_sql_identifier(start_column)}" + end_ref = f"{table_ref}.{self._quote_sql_identifier(end_column)}" + duration_expr = f"DATEDIFF('second', {start_ref}, {end_ref})" + if not dimensions: + fallback_dimension = self._find_first_schema_column( + table, + ( + "status", + "priority", + "assignee_user_id", + "created_by_user_id", + "org_id", + ), + ) + if fallback_dimension: + dimensions.append(fallback_dimension) + if dimensions: + dimension = dimensions[0] + dimension_ref = ( + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + ) + dimension_alias = ( + "workflow" if "workflow" in normalized_query else dimension + ) + return ( + f"SELECT {dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_alias)}, " + f'SUM({duration_expr}) AS "total_time_seconds" ' + f"FROM {table_ref} " + f"WHERE {start_ref} IS NOT NULL " + f"AND {end_ref} IS NOT NULL " + f"AND {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f'ORDER BY "total_time_seconds" DESC' + ) + return ( + f'SELECT SUM({duration_expr}) AS "total_time_seconds" ' + f"FROM {table_ref} " + f"WHERE {start_ref} IS NOT NULL " + f"AND {end_ref} IS NOT NULL" + ) + if wants_trend and date_column: date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" select_parts = [ diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2f43719ce3..79ba71f639 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -702,6 +702,43 @@ def test_build_schema_grounded_sql_for_ticket_throughput_trend(): ) +def test_build_schema_grounded_sql_for_ticket_workflow_total_time_uses_first_class_columns(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "What is the estimated total time for each workflow?", + [ + """ + CREATE TABLE dbo_tickets ( + id VARCHAR, + org_id VARCHAR, + title VARCHAR, + description VARCHAR, + status VARCHAR, + priority VARCHAR, + assignee_user_id VARCHAR, + data VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tickets"."status" AS "workflow", ' + 'SUM(DATEDIFF(\'second\', "dbo_tickets"."created_at", ' + '"dbo_tickets"."updated_at")) AS "total_time_seconds" ' + 'FROM "dbo_tickets" ' + 'WHERE "dbo_tickets"."created_at" IS NOT NULL ' + 'AND "dbo_tickets"."updated_at" IS NOT NULL ' + 'AND "dbo_tickets"."status" IS NOT NULL ' + 'GROUP BY "dbo_tickets"."status" ' + 'ORDER BY "total_time_seconds" DESC' + ) + assert "data" not in sql + assert "JSON" not in sql + + def test_build_audit_log_activity_sql_uses_existing_condition_columns(): service = AskService.__new__(AskService) sql = service._build_audit_log_activity_sql( From 900a3f4fd2b7e2b8c0b634878d7c74a4d47182c1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 17:37:04 +0530 Subject: [PATCH 0374/1087] Guard historical SQL reuse by question match --- wren-ai-service/src/web/v1/services/ask.py | 83 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 14 ++++ 2 files changed, 97 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c935ff3b1e..44eda8b671 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -100,6 +100,39 @@ class AskResultResponse(_AskResultResponse): class AskService: + _HISTORICAL_QUESTION_STOP_WORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "can", + "chart", + "create", + "each", + "for", + "from", + "give", + "graph", + "how", + "in", + "is", + "me", + "of", + "on", + "please", + "show", + "the", + "to", + "total", + "what", + "which", + "with", + } + def __init__( self, pipelines: Dict[str, BasicPipeline], @@ -142,6 +175,44 @@ def _is_stopped(self, query_id: str, container: dict): return False + @classmethod + def _normalize_historical_question_text(cls, question: str | None) -> str: + return " ".join(re.findall(r"[a-z0-9]+", (question or "").lower())) + + @classmethod + def _historical_question_tokens(cls, question: str | None) -> set[str]: + normalized = cls._normalize_historical_question_text(question) + return { + token + for token in normalized.split() + if len(token) > 1 and token not in cls._HISTORICAL_QUESTION_STOP_WORDS + } + + @classmethod + def _is_reusable_historical_question( + cls, query: str | None, historical_question: str | None + ) -> bool: + normalized_query = cls._normalize_historical_question_text(query) + normalized_historical_question = cls._normalize_historical_question_text( + historical_question + ) + if not normalized_query or not normalized_historical_question: + return False + if normalized_query == normalized_historical_question: + return True + + query_tokens = cls._historical_question_tokens(normalized_query) + historical_tokens = cls._historical_question_tokens( + normalized_historical_question + ) + if len(query_tokens) < 2 or len(historical_tokens) < 2: + return False + + overlap = query_tokens & historical_tokens + coverage = len(overlap) / max(len(query_tokens), 1) + similarity = len(overlap) / len(query_tokens | historical_tokens) + return coverage >= 0.8 and similarity >= 0.72 + def _is_greeting_query(self, query: str) -> bool: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) greeting_patterns = { @@ -3711,6 +3782,18 @@ async def ask( valid_historical_results = [] for result in historical_question_result: + historical_question_text = result.get("question") + if not self._is_reusable_historical_question( + user_query, historical_question_text + ): + logger.info( + "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", + query_id, + user_query, + historical_question_text, + ) + continue + sql_statement = result.get("statement") if not self._is_valid_select_sql(sql_statement): logger.warning( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 79ba71f639..290996f708 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -797,3 +797,17 @@ def test_build_validated_ask_result_from_sql_uses_local_schema_validation(): assert result is not None assert result.sql == 'SELECT "dbo_tblSales"."SalesPerson" FROM "dbo_tblSales"' + + +def test_reusable_historical_question_allows_exact_recommended_question(): + assert AskService._is_reusable_historical_question( + "How many tickets are currently open?", + "How many tickets are currently open?", + ) + + +def test_reusable_historical_question_rejects_materially_different_agent_question(): + assert not AskService._is_reusable_historical_question( + "What is the estimated total time for each workflow?", + "How many tickets are currently open?", + ) From eb7d90708c2366e532b0354212e26224bdf1c626 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 17:45:50 +0530 Subject: [PATCH 0375/1087] Normalize sales period column references --- .../src/pipelines/generation/utils/sql.py | 3 +++ wren-ai-service/src/web/v1/services/ask.py | 9 +++++++++ .../pipelines/generation/test_sql_utils.py | 18 +++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 20 +++++++++++++++++++ 4 files changed, 50 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 6fa7786d1a..fb8593c5b1 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2622,6 +2622,9 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: ), "invoicequantity": ("Qty", "Quantity", "InvoiceQty", "InvoiceCount"), "otddate": ("InvDate", "OrdDate", "OrderDate", "InvoiceDate", "Date"), + "period": ("timeid", "TimeID", "TimeId", "YearInd", "Year", "Date"), + "periodid": ("timeid", "TimeID", "TimeId"), + "timeid": ("timeid", "TimeID", "TimeId"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 44eda8b671..cb1b90c978 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -8,6 +8,10 @@ from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + normalize_sql_column_references_to_schema, +) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -3262,6 +3266,11 @@ def _build_validated_ask_result_from_sql( sql: Optional[str], table_ddls: list[str], ) -> Optional[AskResult]: + if isinstance(sql, str): + sql = normalize_sql_column_references_to_schema( + sql, + construct_valid_table_columns(table_ddls), + ) ask_result = self._build_ask_result_from_sql(sql) if not ask_result: return None diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 96ebd967e9..56e8321248 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -306,6 +306,24 @@ def test_normalize_sql_column_references_to_schema_maps_sales_business_aliases() ) == [] +def test_normalize_sql_column_references_to_schema_maps_period_to_timeid(): + sql = 'SELECT "dbo_tblFactSales"."Period" FROM "dbo_tblFactSales"' + + normalized = normalize_sql_column_references_to_schema( + sql, + { + "dbo_tblFactSales": [ + "account", + "customerpo", + "timeid", + "amount", + ] + }, + ) + + assert normalized == 'SELECT "dbo_tblFactSales"."timeid" FROM "dbo_tblFactSales"' + + def test_normalize_sql_column_references_to_schema_maps_debug_business_aliases(): sql = ( 'SELECT COUNT("FixLogId") AS "FixLogCount" ' diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 290996f708..2c8f88d7ee 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -799,6 +799,26 @@ def test_build_validated_ask_result_from_sql_uses_local_schema_validation(): assert result.sql == 'SELECT "dbo_tblSales"."SalesPerson" FROM "dbo_tblSales"' +def test_build_validated_ask_result_from_sql_normalizes_column_case_to_schema(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + 'SELECT "dbo_tblFactSales"."TimeID" FROM "dbo_tblFactSales"', + [ + """ + CREATE TABLE dbo_tblFactSales ( + account VARCHAR, + timeid VARCHAR, + amount DOUBLE + ); + """ + ], + ) + + assert result is not None + assert result.sql == 'SELECT "dbo_tblFactSales"."timeid" FROM "dbo_tblFactSales"' + + def test_reusable_historical_question_allows_exact_recommended_question(): assert AskService._is_reusable_historical_question( "How many tickets are currently open?", From daafd30b21992615d7d5886f27b212fd63fe8567 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 18:03:10 +0530 Subject: [PATCH 0376/1087] Fallback project lookup for orphaned thread responses --- .../apollo/server/services/askingService.ts | 5 +++- .../services/tests/askingService.test.ts | 25 ++++++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 01eb9e68e5..36b564c3a1 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1355,7 +1355,10 @@ export class AskingService implements IAskingService { id: threadResponse.threadId, }); if (!thread) { - throw new Error(`Thread ${threadResponse.threadId} not found`); + logger.warn( + `Thread ${threadResponse.threadId} for response ${threadResponse.id} not found; falling back to current project`, + ); + return this.projectService.getCurrentProject(); } return this.projectService.getProjectById(thread.projectId); diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index e6fa374c14..cf1706db51 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -1,4 +1,4 @@ -import { constructCteSql } from '../askingService'; +import { AskingService, constructCteSql } from '../askingService'; describe('AskingService', () => { describe('utility: constructCteSql', () => { @@ -94,4 +94,27 @@ describe('AskingService', () => { ); }); }); + + describe('project lookup for thread response', () => { + test('falls back to current project when parent thread is missing', async () => { + const currentProject = { id: 2, type: 'mssql' }; + const service = Object.create(AskingService.prototype) as any; + service.threadRepository = { + findOneBy: jest.fn().mockResolvedValue(null), + }; + service.projectService = { + getCurrentProject: jest.fn().mockResolvedValue(currentProject), + getProjectById: jest.fn(), + }; + + const project = await service.getProjectForThreadResponse({ + id: 10, + threadId: 530, + }); + + expect(project).toBe(currentProject); + expect(service.projectService.getCurrentProject).toHaveBeenCalledTimes(1); + expect(service.projectService.getProjectById).not.toHaveBeenCalled(); + }); + }); }); From 516f262d9caacd631dbc319a5009be57838ea2b8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 18:12:38 +0530 Subject: [PATCH 0377/1087] Normalize sales customer column references --- .../src/pipelines/generation/utils/sql.py | 5 +++ .../pipelines/generation/test_sql_utils.py | 36 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 21 +++++++++++ 3 files changed, 62 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index fb8593c5b1..4426655c42 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2625,6 +2625,11 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "period": ("timeid", "TimeID", "TimeId", "YearInd", "Year", "Date"), "periodid": ("timeid", "TimeID", "TimeId"), "timeid": ("timeid", "TimeID", "TimeId"), + "customer": ("account", "Customer", "CustName", "CustNo", "customerpo"), + "customers": ("account", "Customer", "CustName", "CustNo", "customerpo"), + "customername": ("account", "Customer", "CustName", "CustNo", "customerpo"), + "customerid": ("account", "Customer", "CustNo", "customerpo"), + "customeraccount": ("account", "Customer", "CustName", "CustNo"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 56e8321248..a8b7ae9774 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -324,6 +324,42 @@ def test_normalize_sql_column_references_to_schema_maps_period_to_timeid(): assert normalized == 'SELECT "dbo_tblFactSales"."timeid" FROM "dbo_tblFactSales"' +def test_normalize_sql_column_references_to_schema_maps_customer_to_account(): + sql = 'SELECT "dbo_tblFactSales"."customer" FROM "dbo_tblFactSales"' + + normalized = normalize_sql_column_references_to_schema( + sql, + { + "dbo_tblFactSales": [ + "account", + "customerpo", + "timeid", + "amount", + ] + }, + ) + + assert normalized == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' + + +def test_normalize_sql_column_references_to_schema_maps_unqualified_customer_to_account(): + sql = 'SELECT "customer" FROM "dbo_tblFactSales"' + + normalized = normalize_sql_column_references_to_schema( + sql, + { + "dbo_tblFactSales": [ + "account", + "customerpo", + "timeid", + "amount", + ] + }, + ) + + assert normalized == 'SELECT "account" FROM "dbo_tblFactSales"' + + def test_normalize_sql_column_references_to_schema_maps_debug_business_aliases(): sql = ( 'SELECT COUNT("FixLogId") AS "FixLogCount" ' diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2c8f88d7ee..9c7aac10e3 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -819,6 +819,27 @@ def test_build_validated_ask_result_from_sql_normalizes_column_case_to_schema(): assert result.sql == 'SELECT "dbo_tblFactSales"."timeid" FROM "dbo_tblFactSales"' +def test_build_validated_ask_result_from_sql_normalizes_customer_to_account(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + 'SELECT "dbo_tblFactSales"."customer" FROM "dbo_tblFactSales"', + [ + """ + CREATE TABLE dbo_tblFactSales ( + account VARCHAR, + customerpo VARCHAR, + timeid VARCHAR, + amount DOUBLE + ); + """ + ], + ) + + assert result is not None + assert result.sql == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' + + def test_reusable_historical_question_allows_exact_recommended_question(): assert AskService._is_reusable_historical_question( "How many tickets are currently open?", From 80462c8a4a0ad565beb6d050f32beef8d520b68f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 18:38:59 +0530 Subject: [PATCH 0378/1087] Validate SQL against question intent --- wren-ai-service/src/web/v1/services/ask.py | 247 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 47 ++++ 2 files changed, 288 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cb1b90c978..b541eb853b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -505,6 +505,203 @@ def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: return tables + _INTENT_STOPWORDS = { + "a", + "an", + "and", + "are", + "as", + "based", + "be", + "by", + "can", + "chart", + "correct", + "data", + "different", + "do", + "does", + "each", + "for", + "from", + "give", + "how", + "in", + "is", + "it", + "list", + "many", + "me", + "of", + "on", + "or", + "per", + "question", + "rate", + "records", + "reduce", + "show", + "system", + "taken", + "the", + "there", + "to", + "total", + "type", + "types", + "what", + "which", + "with", + } + + def _intent_tokens(self, text: str) -> set[str]: + tokens: set[str] = set() + for raw_token in re.findall(r"[A-Za-z][A-Za-z0-9_]*", text or ""): + split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) + for token in re.findall(r"[A-Za-z0-9]+", split_token.lower()): + if len(token) <= 2 or token in self._INTENT_STOPWORDS: + continue + tokens.add(token) + if token.endswith("ies") and len(token) > 4: + tokens.add(token[:-3] + "y") + elif token.endswith("s") and len(token) > 3: + tokens.add(token[:-1]) + return tokens + + def _schema_name_tokens(self, name: str) -> set[str]: + spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(name or "")) + return { + token + for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) + if len(token) > 1 + } + + def _table_for_sql_reference( + self, table_reference: str, valid_tables: dict[str, dict[str, Any]] + ) -> dict[str, Any] | None: + table_key = str(table_reference or "").lower() + if table_key in valid_tables: + return valid_tables[table_key] + suffix_key = table_key.split(".")[-1] + for valid_table_name, table in valid_tables.items(): + if valid_table_name.split(".")[-1] == suffix_key: + return table + return None + + def _sql_matches_question_intent( + self, + sql: str, + query: str | None, + schema_tables: list[dict[str, Any]], + ) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + expects_dimension = bool( + re.search( + r"\b(?:by|per|each|which|different|type|types|category|" + r"categories|status|source)\b", + normalized_query, + ) + ) + if not expects_dimension: + return True + + question_tokens = self._intent_tokens(query or "") + if not question_tokens: + return True + + valid_tables = { + str(table.get("name") or "").lower(): table + for table in schema_tables + if table.get("name") + } + if not valid_tables: + return True + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, + ) + referenced_tables = [ + next(value for value in match.groupdict().values() if value) + for match in table_reference_pattern.finditer(sql) + ] + if not referenced_tables: + return True + + qualified_column_pattern = re.compile( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_$]*))", + flags=re.IGNORECASE, + ) + referenced_columns_by_table: dict[str, set[str]] = {} + for match in qualified_column_pattern.finditer(sql): + table_reference = ( + match.group("table_quoted") + or match.group("table_bracketed") + or match.group("table_bare") + or "" + ).lower() + column_reference = ( + match.group("column_quoted") + or match.group("column_bracketed") + or match.group("column_bare") + or "" + ) + referenced_columns_by_table.setdefault(table_reference, set()).add( + column_reference + ) + + for table_reference in referenced_tables: + table = self._table_for_sql_reference(table_reference, valid_tables) + if not table: + continue + + columns = [ + column for column in table.get("columns", []) if column.get("name") + ] + intent_matching_columns = [ + str(column.get("name")) + for column in columns + if self._schema_name_tokens(str(column.get("name"))) & question_tokens + ] + if not intent_matching_columns: + continue + + table_key = str(table_reference or "").lower() + referenced_columns = referenced_columns_by_table.get( + table_key + ) or referenced_columns_by_table.get( + table_key.split(".")[-1], + set(), + ) + referenced_column_tokens = ( + set().union( + *[ + self._schema_name_tokens(column_name) + for column_name in referenced_columns + ] + ) + if referenced_columns + else set() + ) + + if not referenced_column_tokens & question_tokens: + logger.warning( + "Ignoring SQL because selected columns do not match question intent. " + "query=%s table=%s matching_schema_columns=%s referenced_columns=%s sql=%s", + query, + table.get("name"), + intent_matching_columns, + sorted(referenced_columns), + sql, + ) + return False + + return True + def _is_numeric_schema_type(self, column_type: str) -> bool: return bool( re.search( @@ -2316,6 +2513,16 @@ def _build_schema_grounded_operational_sql( score += 5 if "repair" in normalized_query and "repair" in normalized_table: score += 5 + if "failure" in normalized_query and "failure" in normalized_table: + score += 8 + if any(term in normalized_query for term in ("error", "failure")) and any( + self._find_schema_column(table, candidates) + for candidates in ( + ("failure_code", "FailureSys", "failure", "failure_type"), + ("category", "name", "description"), + ) + ): + score += 6 if self._find_schema_column( table, ("created_at", "updated_at", "DateIn", "DateOut", "created", "date"), @@ -2350,6 +2557,19 @@ def _build_schema_grounded_operational_sql( ) dimension_candidates: list[tuple[str, ...]] = [] + if any(term in normalized_query for term in ("failure", "failures", "error")): + dimension_candidates.append( + ( + "failure_code", + "FailureSys", + "failure", + "failure_type", + "failure_category", + "category", + "name", + "description", + ) + ) if "manufacturing" in normalized_query or "unit" in normalized_query: dimension_candidates.append( ( @@ -3265,6 +3485,7 @@ def _build_validated_ask_result_from_sql( self, sql: Optional[str], table_ddls: list[str], + query: str | None = None, ) -> Optional[AskResult]: if isinstance(sql, str): sql = normalize_sql_column_references_to_schema( @@ -3356,6 +3577,13 @@ def _build_validated_ask_result_from_sql( ) return None + if not self._sql_matches_question_intent( + ask_result.sql, + query, + schema_tables, + ): + return None + return ask_result def _build_failed_text_to_sql_response( @@ -3692,6 +3920,7 @@ async def ask( if ask_result := self._build_validated_ask_result_from_sql( heuristic_sql, table_ddls, + user_query, ): api_results = [ask_result] if not self._is_stopped(query_id, self._ask_results): @@ -4213,6 +4442,7 @@ async def ask( ask_result = self._build_validated_ask_result_from_sql( heuristic_sql, table_ddls, + user_query, ) if not ask_result: invalid_sql = heuristic_sql @@ -4480,14 +4710,16 @@ async def ask( if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - if ask_result := self._build_ask_result_from_sql( - sql_valid_result.get("sql") + if ask_result := self._build_validated_ask_result_from_sql( + sql_valid_result.get("sql"), + table_ddls, + sql_user_query, ): api_results = [ask_result] else: invalid_sql = sql_valid_result.get("sql") error_message = ( - "SQL generation did not produce a valid SELECT statement." + "SQL generation did not produce SQL that matches the active datasource schema and question intent." ) elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" @@ -4563,14 +4795,16 @@ async def ask( if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - if ask_result := self._build_ask_result_from_sql( - valid_generation_result.get("sql") + if ask_result := self._build_validated_ask_result_from_sql( + valid_generation_result.get("sql"), + table_ddls, + sql_user_query, ): api_results = [ask_result] break invalid_sql = valid_generation_result.get("sql") error_message = ( - "SQL correction did not produce a valid SELECT statement." + "SQL correction did not produce SQL that matches the active datasource schema and question intent." ) failed_dry_run_result = sql_correction_results["post_process"][ @@ -4608,6 +4842,7 @@ async def ask( ask_result = self._build_validated_ask_result_from_sql( heuristic_sql, table_ddls, + user_query, ) if not ask_result: invalid_sql = heuristic_sql diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 9c7aac10e3..e4cf7932b1 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -840,6 +840,53 @@ def test_build_validated_ask_result_from_sql_normalizes_customer_to_account(): assert result.sql == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' +def test_build_schema_grounded_operational_sql_prefers_failure_column_for_error_rate(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_operational_sql( + "What is the error rate in the repair data collection system for different types of failures?", + [ + { + "name": "dbo_repair_logs", + "columns": [ + {"name": "status", "type": "varchar"}, + {"name": "failure_code", "type": "varchar"}, + {"name": "created_at", "type": "timestamp"}, + ], + } + ], + ) + + assert sql is not None + assert '"dbo_repair_logs"."failure_code" AS "failure_code"' in sql + assert '"dbo_repair_logs"."status" AS "status"' not in sql + + +def test_build_validated_ask_result_rejects_status_when_failure_field_matches_question(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_repair_logs"."status" AS "status", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_repair_logs" ' + 'GROUP BY "dbo_repair_logs"."status"' + ), + [ + """ + CREATE TABLE dbo_repair_logs ( + status VARCHAR, + failure_code VARCHAR, + created_at TIMESTAMP + ); + """ + ], + "What is the error rate for different types of failures?", + ) + + assert result is None + + def test_reusable_historical_question_allows_exact_recommended_question(): assert AskService._is_reusable_historical_question( "How many tickets are currently open?", From f5c29e8621a5be5af754de8a2059a670892f3fa1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 18:59:41 +0530 Subject: [PATCH 0379/1087] Normalize generated table references to active schema --- .../src/pipelines/generation/utils/sql.py | 115 ++++++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 6 + .../pipelines/generation/test_sql_utils.py | 18 +++ .../pytest/services/test_ask_sales_sql.py | 27 ++++ 4 files changed, 166 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4426655c42..2797f6be83 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1574,6 +1574,10 @@ async def run( cleaned_generation_result = normalize_generation_result_sql( cleaned_generation_result, data_source=data_source ) + cleaned_generation_result = normalize_sql_table_references_to_schema( + cleaned_generation_result, + valid_table_names or [], + ) cleaned_generation_result = normalize_sql_column_references_to_schema( cleaned_generation_result, valid_table_columns or {}, @@ -2661,6 +2665,12 @@ def _table_reference_suffixes(table_reference: str) -> list[str]: return [".".join(parts[index:]) for index in range(len(parts))] +def _quote_table_reference(table_reference: str) -> str: + return ".".join( + _quote_sql_identifier(part) for part in _split_table_reference(table_reference) + ) + + def extract_sql_table_references(sql: str) -> list[str]: references = [] for match in _SQL_TABLE_REFERENCE_PATTERN.finditer(sql): @@ -2709,6 +2719,111 @@ def find_invalid_table_references(sql: str, valid_table_names: list[str]) -> lis return sorted(set(invalid_references)) +def _find_schema_table_alias( + requested_table: str, valid_table_names: list[str] +) -> str | None: + requested_suffixes = _table_reference_suffixes(requested_table) + requested_candidates = [ + suffix for suffix in requested_suffixes if suffix and len(suffix) > 2 + ] + if not requested_candidates: + return None + + valid_candidates = [ + str(table_name) + for table_name in valid_table_names or [] + if table_name is not None and str(table_name).strip() + ] + valid_by_lower = {table.lower(): table for table in valid_candidates} + for candidate in requested_candidates: + exact = valid_by_lower.get(candidate.lower()) + if exact: + return exact + + scored: list[tuple[int, int, str]] = [] + for valid_table in valid_candidates: + valid_suffixes = _table_reference_suffixes(valid_table) + valid_keys = [valid_table, *valid_suffixes] + for requested_key in requested_candidates: + requested_compact = _compact_sql_identifier(requested_key) + if not requested_compact: + continue + for valid_key in valid_keys: + valid_compact = _compact_sql_identifier(valid_key) + if not valid_compact: + continue + score = 0 + if requested_compact == valid_compact: + score = 1000 + len(valid_compact) + elif valid_compact.endswith(requested_compact): + score = 800 + len(requested_compact) + elif requested_compact.endswith(valid_compact): + score = 700 + len(valid_compact) + elif ( + len(requested_compact) >= 6 + and valid_compact.startswith(requested_compact) + ): + score = 600 + len(requested_compact) + elif ( + len(valid_compact) >= 6 + and requested_compact.startswith(valid_compact) + ): + score = 500 + len(valid_compact) + if score: + scored.append((score, len(valid_compact), valid_table)) + + if not scored: + return None + + scored.sort(reverse=True) + best_score = scored[0][0] + best_tables = {table for score, _, table in scored if score == best_score} + if len(best_tables) != 1: + return None + return scored[0][2] + + +def normalize_sql_table_references_to_schema( + sql: str, valid_table_names: list[str] +) -> str: + if not sql or not valid_table_names: + return sql + + replacements: dict[str, str] = {} + for table_reference in extract_sql_table_references(sql): + canonical_table = _find_schema_table_alias(table_reference, valid_table_names) + if not canonical_table or canonical_table == table_reference: + continue + replacements[table_reference] = canonical_table + + if not replacements: + return sql + + normalized_sql = sql + for requested_table, canonical_table in sorted( + replacements.items(), key=lambda item: len(item[0]), reverse=True + ): + requested_parts = _split_table_reference(requested_table) + if not requested_parts: + continue + quoted_requested = r"\s*\.\s*".join( + re.escape(_quote_sql_identifier(part)) for part in requested_parts + ) + bare_requested = r"\s*\.\s*".join( + re.escape(part) for part in requested_parts + ) + table_pattern = re.compile( + rf"(? dict[str, str]: diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b541eb853b..d5a1fb411d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -10,7 +10,9 @@ from src.core.pipeline import BasicPipeline from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, + construct_valid_table_names, normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, ) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -3488,6 +3490,10 @@ def _build_validated_ask_result_from_sql( query: str | None = None, ) -> Optional[AskResult]: if isinstance(sql, str): + sql = normalize_sql_table_references_to_schema( + sql, + construct_valid_table_names(table_ddls), + ) sql = normalize_sql_column_references_to_schema( sql, construct_valid_table_columns(table_ddls), diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index a8b7ae9774..ad4ea43f62 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -10,6 +10,7 @@ normalize_data_source, normalize_generation_result_sql, normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, get_sql_generation_system_prompt, get_text_to_sql_rules, ) @@ -80,6 +81,23 @@ def test_column_validation_allows_valid_unqualified_projection_for_single_table( ) == [] +def test_normalize_sql_table_references_to_schema_maps_unique_prefix_table(): + sql = ( + 'SELECT "public"."dbo_failure"."created_at" ' + 'FROM "public"."dbo_failure"' + ) + + normalized = normalize_sql_table_references_to_schema( + sql, + ["dbo_failure_patterns"], + ) + + assert normalized == ( + 'SELECT "dbo_failure_patterns"."created_at" ' + 'FROM "dbo_failure_patterns"' + ) + + def test_construct_valid_table_columns_adds_qualified_suffix_tables(): documents = [ 'CREATE TABLE "wrenai"."public"."dbo_repair_logs" ("warning_signals" INTEGER);', diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index e4cf7932b1..1fc3c7bab5 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -840,6 +840,33 @@ def test_build_validated_ask_result_from_sql_normalizes_customer_to_account(): assert result.sql == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' +def test_build_validated_ask_result_from_sql_normalizes_active_table_reference(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "public"."dbo_failure"."created_at" ' + 'FROM "public"."dbo_failure"' + ), + [ + """ + CREATE TABLE dbo_failure_patterns ( + id VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP + ); + """ + ], + "Show monthly record count by created_at in dbo.failure_patterns", + ) + + assert result is not None + assert result.sql == ( + 'SELECT "dbo_failure_patterns"."created_at" ' + 'FROM "dbo_failure_patterns"' + ) + + def test_build_schema_grounded_operational_sql_prefers_failure_column_for_error_rate(): service = AskService.__new__(AskService) From 6caf4366b8b4965d16f8bd1fe017005f72b9f52a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 19:13:29 +0530 Subject: [PATCH 0380/1087] Normalize quoted generated table references --- .../src/pipelines/generation/utils/sql.py | 32 +++++++++++++++-- .../pipelines/generation/test_sql_utils.py | 34 +++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2797f6be83..4bec10530f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2653,9 +2653,32 @@ def _find_semantic_column_alias( def _split_table_reference(table_reference: str) -> list[str]: + stripped = table_reference.strip() + if not stripped: + return [] + + is_multipart_quoted_reference = bool( + re.search(r'"\s*\.\s*"|\]\s*\.\s*\[|`\s*\.\s*`', stripped) + ) + if ( + not is_multipart_quoted_reference + and ( + (stripped.startswith('"') and stripped.endswith('"')) + or (stripped.startswith("[") and stripped.endswith("]")) + or (stripped.startswith("`") and stripped.endswith("`")) + ) + ): + normalized_identifier = _normalize_sql_identifier(stripped) + if "." in normalized_identifier: + return [ + part + for part in re.split(r"\s*\.\s*", normalized_identifier) + if part.strip() + ] + return [ _normalize_sql_identifier(part) - for part in re.split(r"\s*\.\s*", table_reference.strip()) + for part in re.split(r"\s*\.\s*", stripped) if part.strip() ] @@ -2809,11 +2832,16 @@ def normalize_sql_table_references_to_schema( quoted_requested = r"\s*\.\s*".join( re.escape(_quote_sql_identifier(part)) for part in requested_parts ) + single_quoted_requested = re.escape(_quote_sql_identifier(requested_table)) + bracketed_requested = re.escape(f"[{requested_table}]") + backticked_requested = re.escape(f"`{requested_table}`") bare_requested = r"\s*\.\s*".join( re.escape(part) for part in requested_parts ) table_pattern = re.compile( - rf"(? Date: Mon, 6 Jul 2026 19:21:03 +0530 Subject: [PATCH 0381/1087] Normalize preview SQL table references --- .../apollo/server/services/queryService.ts | 133 +++++++++++++++++- .../services/tests/queryService.test.ts | 66 +++++++++ 2 files changed, 195 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index df46ea7c28..f7381c045c 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -309,12 +309,40 @@ const normalizeSqlIdentifier = (identifier: string) => { return trimmed; }; -const splitTableReference = (tableReference: string) => - tableReference - .trim() +const splitTableReference = (tableReference: string) => { + const trimmed = tableReference.trim(); + if (!trimmed) { + return []; + } + + const isMultipartQuotedReference = + /"\s*\.\s*"|\]\s*\.\s*\[|`\s*\.\s*`/.test(trimmed); + if ( + !isMultipartQuotedReference && + ((trimmed.startsWith('"') && trimmed.endsWith('"')) || + (trimmed.startsWith('`') && trimmed.endsWith('`')) || + (trimmed.startsWith('[') && trimmed.endsWith(']'))) + ) { + const normalized = normalizeSqlIdentifier(trimmed); + if (normalized.includes('.')) { + return normalized.split(/\s*\.\s*/).filter(Boolean); + } + } + + return trimmed .split(/\s*\.\s*/) .map(normalizeSqlIdentifier) .filter(Boolean); +}; + +const quoteSqlIdentifier = (identifier: string) => + `"${identifier.replace(/"/g, '""')}"`; + +const quoteTableReference = (tableReference: string) => + splitTableReference(tableReference).map(quoteSqlIdentifier).join('.'); + +const escapeRegExp = (value: string) => + value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); const compactSqlIdentifier = (identifier: string) => normalizeSqlIdentifier(identifier).replace(/[^A-Za-z0-9]/g, '').toLowerCase(); @@ -661,6 +689,98 @@ const findSqlReferenceValidationErrors = ( return [...new Set(errors)]; }; +const addModelReferenceAlias = ( + aliases: Map, + parts: Array, + modelName?: string, +) => { + if (!modelName) { + return; + } + + const normalizedParts = parts + .filter((part): part is string => Boolean(part)) + .map((part) => part.toLowerCase()); + if (!normalizedParts.length) { + return; + } + + for (let index = 0; index < normalizedParts.length; index += 1) { + aliases.set(normalizedParts.slice(index).join('.'), modelName); + } +}; + +const getManifestTableReferenceAliases = (manifest?: Manifest) => { + const aliases = new Map(); + for (const model of manifest?.models || []) { + if (!model.name) { + continue; + } + + aliases.set(model.name.toLowerCase(), model.name); + if (model.tableReference?.table) { + addModelReferenceAlias( + aliases, + [ + model.tableReference.catalog, + model.tableReference.schema, + model.tableReference.table, + ], + model.name, + ); + } + } + return aliases; +}; + +const tableReferencePatternFor = (tableReference: string) => { + const parts = splitTableReference(tableReference); + if (!parts.length) { + return undefined; + } + + const multipartQuoted = parts + .map(quoteSqlIdentifier) + .map(escapeRegExp) + .join(String.raw`\s*\.\s*`); + const bare = parts.map(escapeRegExp).join(String.raw`\s*\.\s*`); + const singleQuoted = escapeRegExp(quoteSqlIdentifier(tableReference)); + const bracketed = escapeRegExp(`[${tableReference}]`); + const backticked = escapeRegExp(`\`${tableReference}\``); + return new RegExp( + String.raw`(? { + const aliases = getManifestTableReferenceAliases(manifest); + if (!aliases.size) { + return sql; + } + + let normalizedSql = sql; + const replacements = new Map(); + for (const reference of extractSqlTableReferences(sql)) { + const canonicalName = aliases.get(reference.toLowerCase()); + if (canonicalName && canonicalName.toLowerCase() !== reference.toLowerCase()) { + replacements.set(reference, canonicalName); + } + } + + for (const [reference, canonicalName] of [...replacements.entries()].sort( + ([left], [right]) => right.length - left.length, + )) { + const pattern = tableReferencePatternFor(reference); + if (!pattern) { + continue; + } + normalizedSql = normalizedSql.replace(pattern, quoteTableReference(canonicalName)); + } + + return normalizedSql; +}; + const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { const validNames = getManifestQueryableNames(manifest); if (!validNames.size) { @@ -737,7 +857,12 @@ export class QueryService implements IQueryService { } = options; const mdl = normalizeDeployedManifestForDatasource(rawMdl, project); const { type: dataSource, connectionInfo } = project; - const normalizedPreview = normalizePreviewSqlForIbis(sql, dataSource, limit); + const manifestNormalizedSql = normalizeSqlReferencesToManifest(sql, mdl); + const normalizedPreview = normalizePreviewSqlForIbis( + manifestNormalizedSql, + dataSource, + limit, + ); validateSqlReferencesManifest(normalizedPreview.sql, mdl); if (this.useEngine(dataSource)) { if (dryRun) { diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index e51b2bd27d..83cd0124ef 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -528,6 +528,72 @@ describe('QueryService', () => { expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); }); + it('should rewrite physical table references to active model names before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT "wrenai.public.dbo_failure"."created_at" FROM "wrenai.public.dbo_failure"', + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_failure_patterns', + tableReference: { + catalog: 'wrenai', + schema: 'public', + table: 'dbo_failure', + }, + columns: [{ name: 'created_at' }], + }, + ], + }, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', + expect.any(Object), + ); + }); + + it('should rewrite multipart physical table references to active model names before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT "wrenai"."public"."dbo_failure"."created_at" FROM "wrenai"."public"."dbo_failure"', + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_failure_patterns', + tableReference: { + catalog: 'wrenai', + schema: 'public', + table: 'dbo_failure', + }, + columns: [{ name: 'created_at' }], + }, + ], + }, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', + expect.any(Object), + ); + }); + it('should allow source tables referenced by active manifest refSql before ibis dry run', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ correlationId: '123', From 357281a9998bcec31bab83669111d744e24a65d5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 19:29:43 +0530 Subject: [PATCH 0382/1087] Normalize base pattern table references --- .../apollo/server/services/queryService.ts | 4 +++ .../services/tests/queryService.test.ts | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index f7381c045c..c1ce622784 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -718,6 +718,10 @@ const getManifestTableReferenceAliases = (manifest?: Manifest) => { } aliases.set(model.name.toLowerCase(), model.name); + const basePatternMatch = model.name.match(/^(.+)_patterns$/i); + if (basePatternMatch?.[1]) { + aliases.set(basePatternMatch[1].toLowerCase(), model.name); + } if (model.tableReference?.table) { addModelReferenceAlias( aliases, diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 83cd0124ef..e048da393a 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -594,6 +594,34 @@ describe('QueryService', () => { ); }); + it('should rewrite generated base pattern table references to active pattern model names before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT "dbo_failure"."created_at" FROM "dbo_failure"', + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_failure_patterns', + columns: [{ name: 'created_at' }], + }, + ], + }, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', + expect.any(Object), + ); + }); + it('should allow source tables referenced by active manifest refSql before ibis dry run', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ correlationId: '123', From ae0d09e835e038f183a6a2d66defa962ed8e21af Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 19:35:12 +0530 Subject: [PATCH 0383/1087] Normalize dotted dbo table references --- .../apollo/server/services/queryService.ts | 8 +++ .../services/tests/queryService.test.ts | 56 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index c1ce622784..09c3734120 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -718,9 +718,17 @@ const getManifestTableReferenceAliases = (manifest?: Manifest) => { } aliases.set(model.name.toLowerCase(), model.name); + const dboModelMatch = model.name.match(/^dbo_(.+)$/i); + if (dboModelMatch?.[1]) { + aliases.set(`dbo.${dboModelMatch[1]}`.toLowerCase(), model.name); + } const basePatternMatch = model.name.match(/^(.+)_patterns$/i); if (basePatternMatch?.[1]) { aliases.set(basePatternMatch[1].toLowerCase(), model.name); + const dboBasePatternMatch = basePatternMatch[1].match(/^dbo_(.+)$/i); + if (dboBasePatternMatch?.[1]) { + aliases.set(`dbo.${dboBasePatternMatch[1]}`.toLowerCase(), model.name); + } } if (model.tableReference?.table) { addModelReferenceAlias( diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index e048da393a..8aa5e6cd58 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -622,6 +622,62 @@ describe('QueryService', () => { ); }); + it('should rewrite dotted dbo pattern model references before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT "dbo"."failure_patterns"."created_at" FROM "dbo"."failure_patterns"', + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_failure_patterns', + columns: [{ name: 'created_at' }], + }, + ], + }, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', + expect.any(Object), + ); + }); + + it('should rewrite dotted dbo base pattern references before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT "dbo"."failure"."created_at" FROM "dbo"."failure"', + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_failure_patterns', + columns: [{ name: 'created_at' }], + }, + ], + }, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', + expect.any(Object), + ); + }); + it('should allow source tables referenced by active manifest refSql before ibis dry run', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ correlationId: '123', From 96142f9f79db2b2199433850616a2f721a41bb5d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 19:46:50 +0530 Subject: [PATCH 0384/1087] Normalize generated pattern column references --- .../src/pipelines/generation/utils/sql.py | 5 ++ .../pipelines/generation/test_sql_utils.py | 18 ++++ .../apollo/server/services/queryService.ts | 83 ++++++++++++++++++- .../services/tests/queryService.test.ts | 28 +++++++ 4 files changed, 133 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4bec10530f..bc3a8b288a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2635,6 +2635,11 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "customerid": ("account", "Customer", "CustNo", "customerpo"), "customeraccount": ("account", "Customer", "CustName", "CustNo"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), + "pattern": ("name", "category", "description", "id"), + "patterns": ("name", "category", "description", "id"), + "failurename": ("name", "category", "description", "id"), + "failurepattern": ("name", "category", "description", "id"), + "failurepatterns": ("name", "category", "description", "id"), "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 421566f5eb..8bddfdb9ef 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -394,6 +394,24 @@ def test_normalize_sql_column_references_to_schema_maps_customer_to_account(): assert normalized == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' +def test_normalize_sql_column_references_to_schema_maps_patterns_to_name(): + sql = 'SELECT "dbo_failure_patterns"."patterns" FROM "dbo_failure_patterns"' + + normalized = normalize_sql_column_references_to_schema( + sql, + { + "dbo_failure_patterns": [ + "id", + "name", + "category", + "created_at", + ] + }, + ) + + assert normalized == 'SELECT "dbo_failure_patterns"."name" FROM "dbo_failure_patterns"' + + def test_normalize_sql_column_references_to_schema_maps_unqualified_customer_to_account(): sql = 'SELECT "customer" FROM "dbo_tblFactSales"' diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 09c3734120..64dfc0c034 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -790,7 +790,88 @@ const normalizeSqlReferencesToManifest = (sql: string, manifest?: Manifest) => { normalizedSql = normalizedSql.replace(pattern, quoteTableReference(canonicalName)); } - return normalizedSql; + return normalizeSqlColumnReferencesToManifest(normalizedSql, manifest); +}; + +const getSemanticColumnAlias = ( + column: string, + validColumns: Set, +): string | undefined => { + const normalizedColumn = column.replace(/[^A-Za-z0-9]/g, '').toLowerCase(); + const candidatesByAlias: Record = { + pattern: ['name', 'category', 'description', 'id'], + patterns: ['name', 'category', 'description', 'id'], + failurename: ['name', 'category', 'description', 'id'], + failurepattern: ['name', 'category', 'description', 'id'], + failurepatterns: ['name', 'category', 'description', 'id'], + }; + + const candidates = candidatesByAlias[normalizedColumn] || []; + for (const candidate of candidates) { + for (const validColumn of validColumns) { + const normalizedValidColumn = validColumn + .replace(/[^A-Za-z0-9]/g, '') + .toLowerCase(); + if (normalizedValidColumn === candidate) { + return validColumn; + } + } + } + + if (normalizedColumn.endsWith('s')) { + const singularColumn = normalizedColumn.slice(0, -1); + for (const validColumn of validColumns) { + const normalizedValidColumn = validColumn + .replace(/[^A-Za-z0-9]/g, '') + .toLowerCase(); + if (normalizedValidColumn === singularColumn) { + return validColumn; + } + } + } + + return undefined; +}; + +const normalizeSqlColumnReferencesToManifest = ( + sql: string, + manifest?: Manifest, +) => { + const columnsByName = getManifestColumnsByQueryableName(manifest); + if (!columnsByName.size) { + return sql; + } + + const tableAliases = extractSqlTableAliases(sql); + const qualifiedColumnPattern = new RegExp( + String.raw`(${SQL_IDENTIFIER_PATTERN})\s*\.\s*(${SQL_IDENTIFIER_PATTERN})`, + 'gi', + ); + + return sql.replace( + qualifiedColumnPattern, + (match, qualifierReference: string, columnReference: string) => { + const qualifier = normalizeSqlIdentifier(qualifierReference).toLowerCase(); + const column = normalizeSqlIdentifier(columnReference); + const tableName = tableAliases.get(qualifier) || qualifier; + const validColumns = + columnsByName.get(tableName) || + columnsByName.get( + splitTableReference(tableName).pop()?.toLowerCase() || '', + ); + + if (!validColumns || validColumns.has(column.toLowerCase())) { + return match; + } + + const alias = getSemanticColumnAlias(column, validColumns); + if (!alias) { + return match; + } + + return `${qualifierReference}.${quoteSqlIdentifier(alias)}`; + }, + ); }; const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 8aa5e6cd58..c2d9f15964 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -678,6 +678,34 @@ describe('QueryService', () => { ); }); + it('should rewrite generated pattern column aliases to active model columns before ibis dry run', async () => { + mockIbisAdaptor.dryRun.mockResolvedValue({ + correlationId: '123', + processTime: '1s', + }); + + await queryService.preview( + 'SELECT "dbo_failure_patterns"."patterns" FROM "dbo_failure_patterns"', + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_failure_patterns', + columns: [{ name: 'name' }, { name: 'created_at' }], + }, + ], + }, + dryRun: true, + }, + ); + + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + 'SELECT "dbo_failure_patterns"."name" FROM "dbo_failure_patterns"', + expect.any(Object), + ); + }); + it('should allow source tables referenced by active manifest refSql before ibis dry run', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ correlationId: '123', From 51a00c816f199fe2cd3d6c196604c240a58a7797 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 19:55:49 +0530 Subject: [PATCH 0385/1087] Revert "Normalize generated pattern column references" This reverts commit 96142f9f79db2b2199433850616a2f721a41bb5d. --- .../src/pipelines/generation/utils/sql.py | 5 -- .../pipelines/generation/test_sql_utils.py | 18 ---- .../apollo/server/services/queryService.ts | 83 +------------------ .../services/tests/queryService.test.ts | 28 ------- 4 files changed, 1 insertion(+), 133 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index bc3a8b288a..4bec10530f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2635,11 +2635,6 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "customerid": ("account", "Customer", "CustNo", "customerpo"), "customeraccount": ("account", "Customer", "CustName", "CustNo"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), - "pattern": ("name", "category", "description", "id"), - "patterns": ("name", "category", "description", "id"), - "failurename": ("name", "category", "description", "id"), - "failurepattern": ("name", "category", "description", "id"), - "failurepatterns": ("name", "category", "description", "id"), "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 8bddfdb9ef..421566f5eb 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -394,24 +394,6 @@ def test_normalize_sql_column_references_to_schema_maps_customer_to_account(): assert normalized == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' -def test_normalize_sql_column_references_to_schema_maps_patterns_to_name(): - sql = 'SELECT "dbo_failure_patterns"."patterns" FROM "dbo_failure_patterns"' - - normalized = normalize_sql_column_references_to_schema( - sql, - { - "dbo_failure_patterns": [ - "id", - "name", - "category", - "created_at", - ] - }, - ) - - assert normalized == 'SELECT "dbo_failure_patterns"."name" FROM "dbo_failure_patterns"' - - def test_normalize_sql_column_references_to_schema_maps_unqualified_customer_to_account(): sql = 'SELECT "customer" FROM "dbo_tblFactSales"' diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 64dfc0c034..09c3734120 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -790,88 +790,7 @@ const normalizeSqlReferencesToManifest = (sql: string, manifest?: Manifest) => { normalizedSql = normalizedSql.replace(pattern, quoteTableReference(canonicalName)); } - return normalizeSqlColumnReferencesToManifest(normalizedSql, manifest); -}; - -const getSemanticColumnAlias = ( - column: string, - validColumns: Set, -): string | undefined => { - const normalizedColumn = column.replace(/[^A-Za-z0-9]/g, '').toLowerCase(); - const candidatesByAlias: Record = { - pattern: ['name', 'category', 'description', 'id'], - patterns: ['name', 'category', 'description', 'id'], - failurename: ['name', 'category', 'description', 'id'], - failurepattern: ['name', 'category', 'description', 'id'], - failurepatterns: ['name', 'category', 'description', 'id'], - }; - - const candidates = candidatesByAlias[normalizedColumn] || []; - for (const candidate of candidates) { - for (const validColumn of validColumns) { - const normalizedValidColumn = validColumn - .replace(/[^A-Za-z0-9]/g, '') - .toLowerCase(); - if (normalizedValidColumn === candidate) { - return validColumn; - } - } - } - - if (normalizedColumn.endsWith('s')) { - const singularColumn = normalizedColumn.slice(0, -1); - for (const validColumn of validColumns) { - const normalizedValidColumn = validColumn - .replace(/[^A-Za-z0-9]/g, '') - .toLowerCase(); - if (normalizedValidColumn === singularColumn) { - return validColumn; - } - } - } - - return undefined; -}; - -const normalizeSqlColumnReferencesToManifest = ( - sql: string, - manifest?: Manifest, -) => { - const columnsByName = getManifestColumnsByQueryableName(manifest); - if (!columnsByName.size) { - return sql; - } - - const tableAliases = extractSqlTableAliases(sql); - const qualifiedColumnPattern = new RegExp( - String.raw`(${SQL_IDENTIFIER_PATTERN})\s*\.\s*(${SQL_IDENTIFIER_PATTERN})`, - 'gi', - ); - - return sql.replace( - qualifiedColumnPattern, - (match, qualifierReference: string, columnReference: string) => { - const qualifier = normalizeSqlIdentifier(qualifierReference).toLowerCase(); - const column = normalizeSqlIdentifier(columnReference); - const tableName = tableAliases.get(qualifier) || qualifier; - const validColumns = - columnsByName.get(tableName) || - columnsByName.get( - splitTableReference(tableName).pop()?.toLowerCase() || '', - ); - - if (!validColumns || validColumns.has(column.toLowerCase())) { - return match; - } - - const alias = getSemanticColumnAlias(column, validColumns); - if (!alias) { - return match; - } - - return `${qualifierReference}.${quoteSqlIdentifier(alias)}`; - }, - ); + return normalizedSql; }; const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index c2d9f15964..8aa5e6cd58 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -678,34 +678,6 @@ describe('QueryService', () => { ); }); - it('should rewrite generated pattern column aliases to active model columns before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - 'SELECT "dbo_failure_patterns"."patterns" FROM "dbo_failure_patterns"', - { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_failure_patterns', - columns: [{ name: 'name' }, { name: 'created_at' }], - }, - ], - }, - dryRun: true, - }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT "dbo_failure_patterns"."name" FROM "dbo_failure_patterns"', - expect.any(Object), - ); - }); - it('should allow source tables referenced by active manifest refSql before ibis dry run', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ correlationId: '123', From 6370cbe255bd9045688677d411162706509850c3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 20:06:58 +0530 Subject: [PATCH 0386/1087] Require exact question match for SQL reuse --- wren-ai-service/src/web/v1/services/ask.py | 15 +-------------- .../tests/pytest/services/test_ask_sales_sql.py | 7 +++++++ 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d5a1fb411d..2d6107cb03 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -204,20 +204,7 @@ def _is_reusable_historical_question( ) if not normalized_query or not normalized_historical_question: return False - if normalized_query == normalized_historical_question: - return True - - query_tokens = cls._historical_question_tokens(normalized_query) - historical_tokens = cls._historical_question_tokens( - normalized_historical_question - ) - if len(query_tokens) < 2 or len(historical_tokens) < 2: - return False - - overlap = query_tokens & historical_tokens - coverage = len(overlap) / max(len(query_tokens), 1) - similarity = len(overlap) / len(query_tokens | historical_tokens) - return coverage >= 0.8 and similarity >= 0.72 + return normalized_query == normalized_historical_question def _is_greeting_query(self, query: str) -> bool: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 1fc3c7bab5..db4100abc3 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -926,3 +926,10 @@ def test_reusable_historical_question_rejects_materially_different_agent_questio "What is the estimated total time for each workflow?", "How many tickets are currently open?", ) + + +def test_reusable_historical_question_rejects_similar_but_different_failure_question(): + assert not AskService._is_reusable_historical_question( + "Which name values have the highest occurrences in dbo.failure_patterns?", + "What is the distribution of name in dbo.failure_patterns?", + ) From 154c0421f99845d259290c97863a2e9efdbb34ba Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 20:28:48 +0530 Subject: [PATCH 0387/1087] Limit follow-up history to contextual questions --- wren-ai-service/src/web/v1/services/ask.py | 42 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 14 +++++++ 2 files changed, 56 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2d6107cb03..a034f0c84d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -206,6 +206,41 @@ def _is_reusable_historical_question( return False return normalized_query == normalized_historical_question + @classmethod + def _should_use_histories_for_query(cls, query: str | None) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + contextual_prefixes = ( + "also ", + "and ", + "but ", + "for those ", + "for that ", + "for the same ", + "from that ", + "how about ", + "in that ", + "now ", + "same ", + "show more", + "show the same", + "then ", + "use that ", + "what about ", + "what if ", + ) + if normalized.startswith(contextual_prefixes): + return True + + contextual_patterns = ( + r"\b(previous|last|above|earlier|same|those|that|these|them|it|its|there)\b", + r"\b(add|break down|compare|filter|group|instead|only|sort|split)\b.+\b(by|to|with)\b", + r"\b(by|for|with)\s+(month|quarter|year|status|type|category|customer|market|region|country|division)\b", + ) + return any(re.search(pattern, normalized) for pattern in contextual_patterns) + def _is_greeting_query(self, query: str) -> bool: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) greeting_patterns = { @@ -3643,6 +3678,13 @@ async def ask( histories = ask_request.histories[: self._max_histories][ ::-1 ] # reverse the order of histories + if histories and not self._should_use_histories_for_query(user_query): + logger.info( + "Ignoring thread histories for independent question. query_id=%s query=%s", + query_id, + user_query, + ) + histories = [] rephrased_question = None intent_reasoning = None sql_generation_reasoning = None diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index db4100abc3..e653e31105 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -933,3 +933,17 @@ def test_reusable_historical_question_rejects_similar_but_different_failure_ques "Which name values have the highest occurrences in dbo.failure_patterns?", "What is the distribution of name in dbo.failure_patterns?", ) + + +def test_should_not_use_histories_for_independent_same_thread_question(): + assert not AskService._should_use_histories_for_query( + "Which name values have the highest occurrences in dbo.failure_patterns?" + ) + assert not AskService._should_use_histories_for_query( + "Show monthly record count by created_at in dbo.failure_patterns" + ) + + +def test_should_use_histories_for_contextual_followup_question(): + assert AskService._should_use_histories_for_query("What about by month?") + assert AskService._should_use_histories_for_query("Show the same for last year") From 04c8a1daa5774729e8440dcb93de0596fb6ac20c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 21:45:42 +0530 Subject: [PATCH 0388/1087] Generate table question SQL from active schema --- wren-ai-service/src/web/v1/services/ask.py | 248 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 100 +++++++ 2 files changed, 346 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a034f0c84d..25cc1a4785 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -811,6 +811,203 @@ def _find_first_schema_column( def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' + def _normalize_schema_identifier_key(self, value: str) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + + def _table_matches_query(self, table_name: str, query: str) -> bool: + normalized_query = self._normalize_schema_identifier_key(query) + normalized_table = self._normalize_schema_identifier_key(table_name) + short_table = re.split(r"[.$_]", str(table_name or ""))[-1] + normalized_short_table = self._normalize_schema_identifier_key(short_table) + return bool( + normalized_table + and normalized_table in normalized_query + or normalized_short_table + and normalized_short_table in normalized_query + ) + + def _find_best_schema_table_for_query( + self, query: str, tables: list[dict[str, Any]] + ) -> dict[str, Any] | None: + if not tables: + return None + + scored_tables: list[tuple[int, dict[str, Any]]] = [] + query_tokens = self._intent_tokens(query) + for table in tables: + table_name = str(table.get("name") or "") + if not table_name: + continue + + score = 0 + if self._table_matches_query(table_name, query): + score += 100 + + table_tokens = self._schema_name_tokens(table_name) + score += 8 * len(table_tokens & query_tokens) + + column_token_matches = 0 + for column in table.get("columns", []): + column_token_matches += len( + self._schema_name_tokens(str(column.get("name") or "")) + & query_tokens + ) + score += column_token_matches + + if score > 0: + scored_tables.append((score, table)) + + if scored_tables: + return sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] + if len(tables) == 1: + return tables[0] + return None + + def _query_mentions_column(self, query: str, column_name: str) -> bool: + normalized_query = self._normalize_schema_identifier_key(query) + normalized_column = self._normalize_schema_identifier_key(column_name) + return bool(normalized_column and normalized_column in normalized_query) + + def _find_dimension_column_for_query( + self, query: str, table: dict[str, Any] + ) -> str | None: + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + if column_name and self._query_mentions_column(query, column_name): + return column_name + + text_columns = [ + str(column.get("name")) + for column in table.get("columns", []) + if column.get("name") + and self._is_text_schema_type(str(column.get("type") or "")) + ] + for candidate in ("name", "category", "type", "status", "code"): + column = self._find_schema_column(table, (candidate,)) + if column in text_columns: + return column + return text_columns[0] if text_columns else None + + def _find_temporal_column_for_query( + self, query: str, table: dict[str, Any] + ) -> str | None: + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if ( + column_name + and self._is_temporal_schema_type(column_type) + and self._query_mentions_column(query, column_name) + ): + return column_name + + return self._find_first_schema_column( + table, + ( + "created_at", + "createdat", + "created", + "date", + "time", + "timestamp", + "updated_at", + ), + ) + + def _build_schema_grounded_table_question_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + tables = self._parse_schema_tables(table_ddls) + table = self._find_best_schema_table_for_query(query, tables) + if not table: + return None + + table_name = str(table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + limit = self._extract_requested_top_n(query, default_value=10) + + wants_monthly_count = any( + term in normalized + for term in ("monthly", "by month", "per month", "month-wise") + ) and any(term in normalized for term in ("count", "records", "rows")) + if wants_monthly_count: + date_column = self._find_temporal_column_for_query(query, table) + if not date_column: + return None + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " + f"DATEPART(MONTH, {date_ref}) ASC" + ) + + wants_total_count = ( + re.search(r"\bhow many\b", normalized) + or "record count" in normalized + or "count of records" in normalized + or "number of records" in normalized + ) and not re.search(r"\b(?:by|per|each|distribution|highest|top)\b", normalized) + if wants_total_count: + return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' + + wants_distribution = any( + term in normalized + for term in ( + "distribution", + "highest occurrence", + "highest occurrences", + "most occurrence", + "most occurrences", + "occurrences", + "top", + "common", + ) + ) + if wants_distribution: + dimension_column = self._find_dimension_column_for_query(query, table) + if not dimension_column: + return None + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" + + occurrence_column = self._find_schema_column( + table, + ("occurrences", "occurrence", "count", "total_count", "record_count"), + numeric=True, + ) + if occurrence_column and self._query_mentions_column( + query, occurrence_column + ): + metric_ref = f"{table_ref}.{self._quote_sql_identifier(occurrence_column)}" + return ( + f"SELECT TOP {limit} {dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_column)}, " + f"{metric_ref} AS {self._quote_sql_identifier(occurrence_column)} " + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"ORDER BY {metric_ref} DESC" + ) + + return ( + f"SELECT TOP {limit} {dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_column)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f"ORDER BY COUNT(*) DESC" + ) + + return None + def _build_explicit_table_preview_sql( self, query: str, table_ddls: list[str] ) -> tuple[str, str] | None: @@ -3835,6 +4032,31 @@ async def ask( table_names, ) + if table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ): + ask_result = self._build_validated_ask_result_from_sql( + table_question_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table question matched deployed schema.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = table_question_sql + if explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ): @@ -4383,8 +4605,30 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls + if not api_results and ( + table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ) + ): + logger.info( + "Using schema-grounded table question SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + table_question_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = table_question_sql + error_message = "Schema-grounded table SQL was not valid for the active datasource schema." + + if not api_results and ( + explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ) ): explicit_sql, explicit_table_name = explicit_table_preview logger.info( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index e653e31105..9b37c21b19 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -947,3 +947,103 @@ def test_should_not_use_histories_for_independent_same_thread_question(): def test_should_use_histories_for_contextual_followup_question(): assert AskService._should_use_histories_for_query("What about by month?") assert AskService._should_use_histories_for_query("Show the same for last year") + + +def test_build_schema_grounded_table_question_sql_for_record_count(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "How many records are in dbo.failure_patterns?", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == 'SELECT COUNT(*) AS "RecordCount" FROM "dbo_failure_patterns"' + + +def test_build_schema_grounded_table_question_sql_for_name_distribution(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "What is the distribution of name in dbo.failure_patterns?", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_failure_patterns"."name" AS "name", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."name" IS NOT NULL ' + 'GROUP BY "dbo_failure_patterns"."name" ' + 'ORDER BY COUNT(*) DESC' + ) + + +def test_build_schema_grounded_table_question_sql_for_highest_occurrences(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Which name values have the highest occurrences in dbo.failure_patterns?", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR, + occurrences INTEGER, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_failure_patterns"."name" AS "name", ' + '"dbo_failure_patterns"."occurrences" AS "occurrences" ' + 'FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."name" IS NOT NULL ' + 'ORDER BY "dbo_failure_patterns"."occurrences" DESC' + ) + + +def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Show monthly record count by created_at in dbo.failure_patterns", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ASC' + ) From 4502830d3c78cbf1568f0a33b5a5b998640e0f26 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 22:05:45 +0530 Subject: [PATCH 0389/1087] Reject SQL missing requested question concepts --- wren-ai-service/src/web/v1/services/ask.py | 75 ++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 96 +++++++++++++++++++ 2 files changed, 168 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 25cc1a4785..cdc7c05236 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -612,6 +612,55 @@ def _table_for_sql_reference( return table return None + def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return [] + + concept_groups: list[set[str]] = [] + if "product line" in normalized or "productline" in normalized: + concept_groups.append({"product", "prod", "line", "productline"}) + if "pcb" in normalized: + concept_groups.append({"pcb", "board"}) + if "critical" in normalized: + concept_groups.append({"critical", "severity", "priority"}) + if "cost" in normalized: + concept_groups.append({"cost", "amount", "expense", "impact"}) + if "quarterly" in normalized or "quarter" in normalized: + concept_groups.append({"quarter", "quarterly"}) + if "recurring" in normalized or "recurrence" in normalized: + concept_groups.append({"recurring", "recurrence", "occurrence", "occurrences", "count"}) + if "issue" in normalized or "issues" in normalized: + concept_groups.append({"issue", "issues", "failure", "failures", "problem", "defect"}) + + return concept_groups + + def _sql_covers_required_question_concepts( + self, + sql: str, + query: str | None, + referenced_column_tokens: set[str], + referenced_table_tokens: set[str], + ) -> bool: + sql_text = (sql or "").lower() + available_tokens = referenced_column_tokens | referenced_table_tokens + for concept_group in self._required_sql_concept_groups(query): + if concept_group & available_tokens: + continue + if any(token in sql_text for token in concept_group): + continue + logger.warning( + "Ignoring SQL because it does not cover required question concept. " + "query=%s required=%s referenced_column_tokens=%s referenced_table_tokens=%s sql=%s", + query, + sorted(concept_group), + sorted(referenced_column_tokens), + sorted(referenced_table_tokens), + sql, + ) + return False + return True + def _sql_matches_question_intent( self, sql: str, @@ -626,11 +675,10 @@ def _sql_matches_question_intent( normalized_query, ) ) - if not expects_dimension: - return True question_tokens = self._intent_tokens(query or "") - if not question_tokens: + required_concept_groups = self._required_sql_concept_groups(query) + if not question_tokens and not required_concept_groups: return True valid_tables = { @@ -653,6 +701,9 @@ def _sql_matches_question_intent( if not referenced_tables: return True + referenced_table_tokens = set().union( + *[self._schema_name_tokens(table) for table in referenced_tables] + ) qualified_column_pattern = re.compile( r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" @@ -678,6 +729,24 @@ def _sql_matches_question_intent( column_reference ) + all_referenced_column_tokens = set().union( + *[ + self._schema_name_tokens(column_name) + for columns in referenced_columns_by_table.values() + for column_name in columns + ] + ) if referenced_columns_by_table else set() + if not self._sql_covers_required_question_concepts( + sql, + query, + all_referenced_column_tokens, + referenced_table_tokens, + ): + return False + + if not expects_dimension: + return True + for table_reference in referenced_tables: table = self._table_for_sql_reference(table_reference, valid_tables) if not table: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 9b37c21b19..e06c418d39 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1047,3 +1047,99 @@ def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count() 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."created_at") ASC, ' 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ASC' ) + + +def test_build_validated_ask_result_rejects_status_for_product_line_pcb_question(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_repair_logs"."status" AS "status" ' + 'FROM "dbo_repair_logs"' + ), + [ + """ + CREATE TABLE dbo_repair_logs ( + status VARCHAR, + product_line VARCHAR, + pcb_issue VARCHAR, + created_at TIMESTAMP + ); + """ + ], + "Create a visualization of recurring PCB issues by product line.", + ) + + assert result is None + + +def test_build_validated_ask_result_rejects_status_for_repair_cost_question(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_repair_logs"."status" AS "status" ' + 'FROM "dbo_repair_logs"' + ), + [ + """ + CREATE TABLE dbo_repair_logs ( + status VARCHAR, + repair_cost DOUBLE, + created_at TIMESTAMP + ); + """ + ], + "Generate a quarterly repair cost analysis chart.", + ) + + assert result is None + + +def test_build_validated_ask_result_rejects_status_for_critical_repairs_question(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_repair_logs"."status" AS "status" ' + 'FROM "dbo_repair_logs"' + ), + [ + """ + CREATE TABLE dbo_repair_logs ( + status VARCHAR, + severity VARCHAR, + repair_id INTEGER + ); + """ + ], + "Create a stacked bar chart comparing critical vs non-critical repairs.", + ) + + assert result is None + + +def test_build_validated_ask_result_accepts_product_line_pcb_question_sql(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_repair_logs"."product_line" AS "product_line", ' + '"dbo_repair_logs"."pcb_issue" AS "pcb_issue", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_repair_logs" ' + 'GROUP BY "dbo_repair_logs"."product_line", ' + '"dbo_repair_logs"."pcb_issue"' + ), + [ + """ + CREATE TABLE dbo_repair_logs ( + product_line VARCHAR, + pcb_issue VARCHAR + ); + """ + ], + "Create a visualization of recurring PCB issues by product line.", + ) + + assert result is not None From 9d0e36cd94bf17d5bf82f4ab9b99b1387d6f050a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 22:27:57 +0530 Subject: [PATCH 0390/1087] Reject generated SQL with unknown bare columns --- wren-ai-service/src/web/v1/services/ask.py | 199 ++++++++++++++++-- .../pytest/services/test_ask_sales_sql.py | 54 +++++ 2 files changed, 238 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 cdc7c05236..235948f075 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -661,6 +661,151 @@ def _sql_covers_required_question_concepts( return False return True + def _invalid_unqualified_sql_identifiers( + self, sql: str, schema_tables: list[dict[str, Any]] + ) -> list[str]: + valid_columns = { + str(column.get("name") or "").lower() + for table in schema_tables + for column in table.get("columns", []) + if column.get("name") + } + valid_tables = { + str(table.get("name") or "").lower() + for table in schema_tables + if table.get("name") + } + valid_table_suffixes = {table_name.split(".")[-1] for table_name in valid_tables} + allowed_words = { + "and", + "as", + "asc", + "avg", + "by", + "case", + "cast", + "count", + "datepart", + "day", + "desc", + "distinct", + "else", + "end", + "from", + "group", + "having", + "in", + "is", + "join", + "limit", + "month", + "not", + "null", + "on", + "or", + "order", + "over", + "partition", + "quarter", + "select", + "sum", + "then", + "top", + "when", + "where", + "with", + "year", + } + aliases = { + (match.group("quoted") or match.group("bare") or "").lower() + for match in re.finditer( + r'\bAS\s+(?:"(?P[^"]+)"|(?P[A-Za-z_][A-Za-z0-9_]*))', + sql or "", + flags=re.IGNORECASE, + ) + } + + sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") + identifier_pattern = re.compile( + r'"(?P[^"]+)"|(?P\b[A-Za-z_][A-Za-z0-9_]*\b)' + ) + invalid: list[str] = [] + for match in identifier_pattern.finditer(sql_without_strings): + identifier = match.group("quoted") or match.group("bare") or "" + identifier_key = identifier.lower() + if not identifier_key: + continue + + before = sql_without_strings[: match.start()].rstrip() + after = sql_without_strings[match.end() :].lstrip() + if before.endswith(".") or after.startswith("."): + continue + + previous_word_match = re.search(r"([A-Za-z_][A-Za-z0-9_]*)\s*$", before) + previous_word = ( + previous_word_match.group(1).lower() if previous_word_match else "" + ) + if previous_word == "as": + continue + + order_by_context = bool( + re.search(r"\bORDER\s+BY\b[^)]*$", before, flags=re.IGNORECASE) + ) + if order_by_context and identifier_key in aliases: + continue + + if ( + identifier_key in valid_columns + or identifier_key in valid_tables + or identifier_key in valid_table_suffixes + or identifier_key in allowed_words + ): + continue + + if identifier not in invalid: + invalid.append(identifier) + + return invalid + + def _unqualified_valid_sql_column_tokens( + self, sql: str, schema_tables: list[dict[str, Any]] + ) -> set[str]: + valid_columns = { + str(column.get("name") or "").lower(): str(column.get("name") or "") + for table in schema_tables + for column in table.get("columns", []) + if column.get("name") + } + if not valid_columns: + return set() + + sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") + identifier_pattern = re.compile( + r'"(?P[^"]+)"|(?P\b[A-Za-z_][A-Za-z0-9_]*\b)' + ) + tokens: set[str] = set() + for match in identifier_pattern.finditer(sql_without_strings): + identifier = match.group("quoted") or match.group("bare") or "" + identifier_key = identifier.lower() + if identifier_key not in valid_columns: + continue + + before = sql_without_strings[: match.start()].rstrip() + after = sql_without_strings[match.end() :].lstrip() + if before.endswith(".") or after.startswith("."): + continue + + previous_word_match = re.search(r"([A-Za-z_][A-Za-z0-9_]*)\s*$", before) + previous_word = ( + previous_word_match.group(1).lower() if previous_word_match else "" + ) + if previous_word == "as": + continue + + tokens.update(self._schema_name_tokens(valid_columns[identifier_key])) + + return tokens + def _sql_matches_question_intent( self, sql: str, @@ -736,6 +881,9 @@ def _sql_matches_question_intent( for column_name in columns ] ) if referenced_columns_by_table else set() + all_referenced_column_tokens.update( + self._unqualified_valid_sql_column_tokens(sql, schema_tables) + ) if not self._sql_covers_required_question_concepts( sql, query, @@ -778,7 +926,7 @@ def _sql_matches_question_intent( ] ) if referenced_columns - else set() + else all_referenced_column_tokens ) if not referenced_column_tokens & question_tokens: @@ -3871,6 +4019,19 @@ def _build_validated_ask_result_from_sql( ) return None + invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( + ask_result.sql, + schema_tables, + ) + if invalid_unqualified_identifiers: + logger.warning( + "Ignoring SQL because it references unqualified fields outside the active schema. " + "invalid_identifiers=%s sql=%s", + invalid_unqualified_identifiers, + ask_result.sql, + ) + return None + if not self._sql_matches_question_intent( ask_result.sql, query, @@ -4725,14 +4886,18 @@ async def ask( "Using schema-grounded audit log activity SQL for query_id %s", query_id, ) - api_results = [ - AskResult( - **{ - "sql": audit_log_activity_sql, - "type": "llm", - } + ask_result = self._build_validated_ask_result_from_sql( + audit_log_activity_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = audit_log_activity_sql + error_message = ( + "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." ) - ] if not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( @@ -4743,14 +4908,18 @@ async def ask( "Using schema-grounded CWSales SQL for query_id %s", query_id, ) - api_results = [ - AskResult( - **{ - "sql": deterministic_sales_sql, - "type": "llm", - } + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sales_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = deterministic_sales_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - ] if unqueryable_metric_message := self._get_unqueryable_metric_message( user_query, table_ddls diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index e06c418d39..82eccc89bc 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1143,3 +1143,57 @@ def test_build_validated_ask_result_accepts_product_line_pcb_question_sql(): ) assert result is not None + + +def test_build_validated_ask_result_rejects_unqualified_invalid_columns(): + service = AskService.__new__(AskService) + table_ddls = [ + """ + CREATE TABLE dbo_repair_logs ( + created_at TIMESTAMP, + updated_at TIMESTAMP, + status VARCHAR, + name VARCHAR, + occurrences INTEGER + ); + """ + ] + + invalid_sqls = [ + 'SELECT DATEPART(MONTH, execution_date) AS "month", COUNT(*) AS "RecordCount" FROM "dbo_repair_logs" GROUP BY DATEPART(MONTH, execution_date)', + 'SELECT created_by AS "created_by", COUNT(*) AS "RecordCount" FROM "dbo_repair_logs" GROUP BY created_by', + 'SELECT physical_name AS "physical_name", occurrences FROM "dbo_repair_logs" ORDER BY occurrences DESC', + ] + + for sql in invalid_sqls: + assert ( + service._build_validated_ask_result_from_sql( + sql, + table_ddls, + "Show top 10 failure pattern names by occurrences.", + ) + is None + ) + + +def test_build_validated_ask_result_accepts_unqualified_valid_columns(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT name AS "name", occurrences AS "occurrences" ' + 'FROM "dbo_repair_logs" ' + 'ORDER BY occurrences DESC' + ), + [ + """ + CREATE TABLE dbo_repair_logs ( + name VARCHAR, + occurrences INTEGER + ); + """ + ], + "Show top 10 failure pattern names by occurrences.", + ) + + assert result is not None From f273af540dc234f3e59a544212a0f64a03bd88ec Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 23:31:35 +0530 Subject: [PATCH 0391/1087] Handle plural table prompts and latest records --- wren-ai-service/src/web/v1/services/ask.py | 38 ++++++++++++-- .../pytest/services/test_ask_sales_sql.py | 49 +++++++++++++++++++ 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 235948f075..d47e12ebb3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -745,7 +745,10 @@ def _invalid_unqualified_sql_identifiers( previous_word = ( previous_word_match.group(1).lower() if previous_word_match else "" ) - if previous_word == "as": + if previous_word == "as" and self._is_alias_identifier_position( + sql_without_strings, + match.start(), + ): continue order_by_context = bool( @@ -799,7 +802,10 @@ def _unqualified_valid_sql_column_tokens( previous_word = ( previous_word_match.group(1).lower() if previous_word_match else "" ) - if previous_word == "as": + if previous_word == "as" and self._is_alias_identifier_position( + sql_without_strings, + match.start(), + ): continue tokens.update(self._schema_name_tokens(valid_columns[identifier_key])) @@ -1083,7 +1089,17 @@ def _find_best_schema_table_for_query( def _query_mentions_column(self, query: str, column_name: str) -> bool: normalized_query = self._normalize_schema_identifier_key(query) normalized_column = self._normalize_schema_identifier_key(column_name) - return bool(normalized_column and normalized_column in normalized_query) + if not normalized_column: + return False + if normalized_column in normalized_query: + return True + if normalized_column.endswith("y"): + return f"{normalized_column[:-1]}ies" in normalized_query + return f"{normalized_column}s" in normalized_query + + def _is_alias_identifier_position(self, sql: str, start: int) -> bool: + before = sql[:start].rstrip() + return bool(re.search(r"\bAS\s*$", before, flags=re.IGNORECASE)) def _find_dimension_column_for_query( self, query: str, table: dict[str, Any] @@ -1147,6 +1163,22 @@ def _build_schema_grounded_table_question_sql( table_ref = self._quote_sql_identifier(table_name) limit = self._extract_requested_top_n(query, default_value=10) + wants_latest_records = any( + term in normalized + for term in ("latest", "recent", "newest", "last records", "latest records") + ) + if wants_latest_records: + date_column = self._find_temporal_column_for_query(query, table) + if not date_column: + return None + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT TOP {limit} * " + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"ORDER BY {date_ref} DESC" + ) + wants_monthly_count = any( term in normalized for term in ("monthly", "by month", "per month", "month-wise") diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 82eccc89bc..ad8e8f9d52 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1020,6 +1020,54 @@ def test_build_schema_grounded_table_question_sql_for_highest_occurrences(): ) +def test_build_schema_grounded_table_question_sql_for_plural_names_by_occurrences(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Show top 10 names by occurrences in dbo.failure_patterns.", + [ + """ + CREATE TABLE dbo_failure_patterns ( + name VARCHAR, + occurrences INTEGER, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_failure_patterns"."name" AS "name", ' + '"dbo_failure_patterns"."occurrences" AS "occurrences" ' + 'FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."name" IS NOT NULL ' + 'ORDER BY "dbo_failure_patterns"."occurrences" DESC' + ) + + +def test_build_schema_grounded_table_question_sql_for_latest_records_by_created_at(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Show the latest records from dbo.failure_patterns by created_at.", + [ + """ + CREATE TABLE dbo_failure_patterns ( + name VARCHAR, + occurrences INTEGER, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 * FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' + 'ORDER BY "dbo_failure_patterns"."created_at" DESC' + ) + + def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count(): service = AskService.__new__(AskService) @@ -1163,6 +1211,7 @@ def test_build_validated_ask_result_rejects_unqualified_invalid_columns(): 'SELECT DATEPART(MONTH, execution_date) AS "month", COUNT(*) AS "RecordCount" FROM "dbo_repair_logs" GROUP BY DATEPART(MONTH, execution_date)', 'SELECT created_by AS "created_by", COUNT(*) AS "RecordCount" FROM "dbo_repair_logs" GROUP BY created_by', 'SELECT physical_name AS "physical_name", occurrences FROM "dbo_repair_logs" ORDER BY occurrences DESC', + 'SELECT name AS "created_by" FROM "dbo_repair_logs" ORDER BY created_by', ] for sql in invalid_sqls: From b8852f9c7569655315c9fada69cfa15f47507991 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 23:38:19 +0530 Subject: [PATCH 0392/1087] Ground throughput trends in active unit columns --- wren-ai-service/src/web/v1/services/ask.py | 88 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 53 +++++++++++ 2 files changed, 141 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d47e12ebb3..10207e31e8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2336,6 +2336,59 @@ def _build_manufacturing_throughput_sql( if not (wants_throughput and wants_unit_breakdown): return None + tables = self._parse_schema_tables(table_ddls) + table = self._find_best_schema_table_for_query(query, tables) + if table: + unit_column = self._find_schema_column( + table, + ( + "BusinessUnit", + "business_unit", + "manufacturing_unit", + "manufacturingunit", + "unit", + "unit_name", + "BU", + "division", + ), + ) + if unit_column: + table_name = str(table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + unit_ref = f"{table_ref}.{self._quote_sql_identifier(unit_column)}" + timestamp_column = self._find_temporal_column_for_query(query, table) + + if timestamp_column and any( + term in normalized for term in ("trend", "monthly", "over time") + ): + timestamp_ref = ( + f"{table_ref}.{self._quote_sql_identifier(timestamp_column)}" + ) + return ( + f"SELECT {unit_ref} AS " + f"{self._quote_sql_identifier(unit_column)}, " + f"DATEPART(YEAR, {timestamp_ref}) AS \"year\", " + f"DATEPART(MONTH, {timestamp_ref}) AS \"month\", " + f'COUNT(*) AS "throughput" ' + f"FROM {table_ref} " + f"WHERE {unit_ref} IS NOT NULL " + f"AND {timestamp_ref} IS NOT NULL " + f"GROUP BY {unit_ref}, DATEPART(YEAR, {timestamp_ref}), " + f"DATEPART(MONTH, {timestamp_ref}) " + f"ORDER BY {unit_ref} ASC, DATEPART(YEAR, {timestamp_ref}) ASC, " + f"DATEPART(MONTH, {timestamp_ref}) ASC" + ) + + return ( + f"SELECT {unit_ref} AS " + f"{self._quote_sql_identifier(unit_column)}, " + f'COUNT(*) AS "throughput" ' + f"FROM {table_ref} " + f"WHERE {unit_ref} IS NOT NULL " + f"GROUP BY {unit_ref} " + f"ORDER BY COUNT(*) DESC" + ) + has_debug_entries = self._schema_contains( table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names ) @@ -3230,6 +3283,41 @@ def _get_unqueryable_metric_message( if not normalized_query: return None + if "throughput" in normalized_query and any( + term in normalized_query for term in ("manufacturing", "unit", "units") + ): + unit_field_patterns = ( + r"\bbusiness[_ ]?unit\b", + r"\bmanufacturing[_ ]?unit\b", + r"\bunit[_ ]?name\b", + r"\bunit\b", + r"\bbu\b", + r"\bdivision\b", + ) + has_unit_field = any( + re.search(pattern, column_name) + for pattern in unit_field_patterns + for column_name in schema_column_names + ) + has_temporal_field = any( + self._is_temporal_schema_type(str(column.get("type") or "")) + for table in self._parse_schema_tables(table_ddls) + for column in table.get("columns", []) + ) + if not has_unit_field: + return ( + "The active datasource does not expose a manufacturing unit, " + "business unit, unit, BU, or division column. I cannot build " + "throughput trends across manufacturing units without a " + "queryable unit field." + ) + if "trend" in normalized_query and not has_temporal_field: + return ( + "The active datasource does not expose a queryable date or " + "timestamp column. I cannot build a throughput trend without " + "a first-class temporal field." + ) + repair_cost_terms = ( "repair cost", "repair_cost", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index ad8e8f9d52..1d3cd8a197 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -702,6 +702,59 @@ def test_build_schema_grounded_sql_for_ticket_throughput_trend(): ) +def test_build_manufacturing_throughput_sql_uses_active_unit_and_date_columns(): + service = AskService.__new__(AskService) + + sql = service._build_manufacturing_throughput_sql( + "Show throughput trends across different manufacturing units.", + [ + """ + CREATE TABLE dbo_production_events ( + id INTEGER, + manufacturing_unit VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_production_events"."manufacturing_unit" AS "manufacturing_unit", ' + 'DATEPART(YEAR, "dbo_production_events"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_production_events"."created_at") AS "month", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_production_events" ' + 'WHERE "dbo_production_events"."manufacturing_unit" IS NOT NULL ' + 'AND "dbo_production_events"."created_at" IS NOT NULL ' + 'GROUP BY "dbo_production_events"."manufacturing_unit", ' + 'DATEPART(YEAR, "dbo_production_events"."created_at"), ' + 'DATEPART(MONTH, "dbo_production_events"."created_at") ' + 'ORDER BY "dbo_production_events"."manufacturing_unit" ASC, ' + 'DATEPART(YEAR, "dbo_production_events"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_production_events"."created_at") ASC' + ) + + +def test_get_unqueryable_metric_message_for_throughput_without_unit_column(): + service = AskService.__new__(AskService) + + message = service._get_unqueryable_metric_message( + "Show throughput trends across different manufacturing units.", + [ + """ + CREATE TABLE dbo_failure_patterns ( + name VARCHAR, + occurrences INTEGER, + created_at TIMESTAMP + ); + """ + ], + ) + + assert message is not None + assert "unit" in message + + def test_build_schema_grounded_sql_for_ticket_workflow_total_time_uses_first_class_columns(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From af7ce907efe5c4e915753bc5a623f460a500386c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 6 Jul 2026 23:47:24 +0530 Subject: [PATCH 0393/1087] Reject SQL when schema validation has no tables --- wren-ai-service/src/web/v1/services/ask.py | 6 ++++++ .../tests/pytest/services/test_ask_sales_sql.py | 12 ++++++++++++ 2 files changed, 18 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 10207e31e8..b7e73f36af 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4059,6 +4059,12 @@ def _build_validated_ask_result_from_sql( return None schema_tables = self._parse_schema_tables(table_ddls) + if not schema_tables: + logger.warning( + "Ignoring SQL because no active schema tables were available for validation. sql=%s", + ask_result.sql, + ) + return None valid_tables = { str(table.get("name") or "").lower(): table for table in schema_tables diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 1d3cd8a197..b28094b45a 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1299,3 +1299,15 @@ def test_build_validated_ask_result_accepts_unqualified_valid_columns(): ) assert result is not None + + +def test_build_validated_ask_result_rejects_sql_when_schema_is_missing(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + 'SELECT execution_date FROM "dbo_repair_logs"', + [], + "Generate a trend chart for average turnaround time by month.", + ) + + assert result is None From 8926a3fffd453ccad241c195d164524cb7377a57 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 00:12:52 +0530 Subject: [PATCH 0394/1087] Remove historical context from ask flow --- wren-ai-service/src/web/v1/services/ask.py | 153 +++------------------ wren-ui/src/pages/api/v1/ask.ts | 9 +- wren-ui/src/pages/api/v1/generate_sql.ts | 8 +- wren-ui/src/pages/api/v1/stream/ask.ts | 8 +- 4 files changed, 23 insertions(+), 155 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b7e73f36af..68ba41d1dc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4228,16 +4228,11 @@ async def ask( return results logger.info(f"Ask pipeline started for query_id: {query_id}") - histories = ask_request.histories[: self._max_histories][ - ::-1 - ] # reverse the order of histories - if histories and not self._should_use_histories_for_query(user_query): - logger.info( - "Ignoring thread histories for independent question. query_id=%s query=%s", - query_id, - user_query, - ) - histories = [] + # Always answer from the active datasource schema and current question. + # Prior thread context, saved SQL examples, and historical questions are + # intentionally not reused here to prevent stale or cross-datasource + # leakage from influencing a fresh answer. + histories: list[AskHistory] = [] rephrased_question = None intent_reasoning = None sql_generation_reasoning = None @@ -4575,132 +4570,24 @@ async def ask( query_id, ) - historical_question_result = [] - should_skip_pre_sql_retrieval = self._is_data_analysis_query( - user_query - ) - if should_skip_pre_sql_retrieval: - rephrased_question = user_query - intent_reasoning = ( - "Detected a deployed-data analytics question; skipping " - "intent classification and using SQL generation." - ) - sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - logger.info( - "Skipping pre-SQL retrieval for analytics query_id %s: %s", - query_id, - user_query, + if not api_results: + original_user_query = user_query + # Only user instructions are kept. Prior SQL samples are not + # reused for fresh questions because they can bias the model + # toward stale or unrelated queries. + instructions_task = await self._run_with_timeout( + "Instruction retrieval", + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), ) - if ( - not api_results - and not should_skip_pre_sql_retrieval - and self._should_reuse_historical_question_sql( - user_query, histories + sql_samples = [] + instructions = instructions_task["formatted_output"].get( + "documents", [] ) - ): - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - - try: - historical_question = await self._run_with_timeout( - "Historical question retrieval", - self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ), - timeout_seconds=min(understanding_timeout_seconds, 10), - ) - - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - except TimeoutError as exc: - logger.warning( - "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) - - valid_historical_results = [] - for result in historical_question_result: - historical_question_text = result.get("question") - if not self._is_reusable_historical_question( - user_query, historical_question_text - ): - logger.info( - "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", - query_id, - user_query, - historical_question_text, - ) - continue - - sql_statement = result.get("statement") - if not self._is_valid_select_sql(sql_statement): - logger.warning( - "Ignoring historical question without valid SQL for query_id %s", - query_id, - ) - continue - valid_historical_results.append( - AskResult( - **{ - "sql": sql_statement.strip(), - "type": "view" if result.get("viewId") else "llm", - "viewId": result.get("viewId"), - } - ) - ) - - if valid_historical_results: - api_results = valid_historical_results - sql_generation_reasoning = "" - elif not api_results and not should_skip_pre_sql_retrieval: - original_user_query = user_query - # Run both pipeline operations concurrently - try: - sql_samples_task, instructions_task = await self._run_with_timeout( - "SQL pair and instruction retrieval", - asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", - ), - ), - timeout_seconds=understanding_timeout_seconds, - ) - - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] - ) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) - except TimeoutError as exc: - logger.warning( - "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) - sql_samples = [] - instructions = [] if self._allow_intent_classification: try: diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index d2e98d0eb2..cdcceb703b 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -10,7 +10,6 @@ import { MAX_WAIT_TIME, isAskResultFinished, validateSummaryResult, - transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; import { AskResult, @@ -27,7 +26,6 @@ const logger = getLogger('API_ASK'); logger.level = 'debug'; const { - apiHistoryRepository, projectService, deployService, wrenAIAdaptor, @@ -75,16 +73,11 @@ export default async function handler( // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); - // Get conversation history if threadId is provided - const histories = threadId - ? await apiHistoryRepository.findAllBy({ threadId }) - : undefined; - // Step 1: Generate SQL const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, - histories: transformHistoryInput(histories) as any, + histories: undefined, configurations: { language: language || WrenAILanguage[project.language] || WrenAILanguage.EN, diff --git a/wren-ui/src/pages/api/v1/generate_sql.ts b/wren-ui/src/pages/api/v1/generate_sql.ts index fa5b3859e1..55adb6edfe 100644 --- a/wren-ui/src/pages/api/v1/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/generate_sql.ts @@ -12,7 +12,6 @@ import { MAX_WAIT_TIME, isAskResultFinished, validateAskResult, - transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; import { DataSourceName } from '@server/types'; @@ -20,7 +19,6 @@ const logger = getLogger('API_GENERATE_SQL'); logger.level = 'debug'; const { - apiHistoryRepository, projectService, deployService, wrenAIAdaptor, @@ -72,14 +70,10 @@ export default async function handler( ); } - // ask AI service to generate SQL - const histories = threadId - ? await apiHistoryRepository.findAllBy({ threadId }) - : undefined; const task = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, - histories: transformHistoryInput(histories) as any, + histories: undefined, configurations: { language: language || WrenAILanguage[project.language] || WrenAILanguage.EN, diff --git a/wren-ui/src/pages/api/v1/stream/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts index 12e53eb4db..98159f9048 100644 --- a/wren-ui/src/pages/api/v1/stream/ask.ts +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -8,7 +8,6 @@ import { MAX_WAIT_TIME, isAskResultFinished, validateSummaryResult, - transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; import { AskResult, @@ -136,11 +135,6 @@ export default async function handler( // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); - // Get conversation history if threadId is provided - const histories = threadId - ? await apiHistoryRepository.findAllBy({ threadId }) - : undefined; - // Step 1: Generate SQL sendStateUpdate(res, StateType.SQL_GENERATION_START, { question, @@ -151,7 +145,7 @@ export default async function handler( const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, - histories: transformHistoryInput(histories) as any, + histories: undefined, configurations: { language: language || WrenAILanguage[project.language] || WrenAILanguage.EN, From 1fff30cfd5a747aa38e99cc655316cadfb788d29 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 00:32:26 +0530 Subject: [PATCH 0395/1087] Use only schema-backed temporal columns --- wren-ai-service/src/web/v1/services/ask.py | 83 +++++++++++++++---- .../pytest/services/test_ask_sales_sql.py | 67 +++++++++++++++ 2 files changed, 136 insertions(+), 14 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 68ba41d1dc..3f91f65b03 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1031,6 +1031,14 @@ def _find_first_schema_column( return column return None + def _find_any_temporal_schema_column(self, table: dict[str, Any]) -> str | None: + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if column_name and self._is_temporal_schema_type(column_type): + return column_name + return None + def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' @@ -1134,18 +1142,7 @@ def _find_temporal_column_for_query( ): return column_name - return self._find_first_schema_column( - table, - ( - "created_at", - "createdat", - "created", - "date", - "time", - "timestamp", - "updated_at", - ), - ) + return self._find_any_temporal_schema_column(table) def _build_schema_grounded_table_question_sql( self, query: str, table_ddls: list[str] @@ -3171,16 +3168,47 @@ def _build_schema_grounded_operational_sql( limit = int(limit_match.group(1)) if limit_match else 10 if wants_elapsed_time: + temporal_columns = [ + str(column.get("name") or "") + for column in table.get("columns", []) + if column.get("name") + and self._is_temporal_schema_type(str(column.get("type") or "")) + ] start_column = self._find_schema_column( table, - ("created_at", "created", "DateIn"), + ( + "created_at", + "created", + "DateIn", + "execution_date", + "opened_at", + "started_at", + "start_date", + "begin_date", + ), temporal=True, ) end_column = self._find_schema_column( table, - ("updated_at", "updated", "DateOut", "closed_at", "resolved_at"), + ( + "updated_at", + "updated", + "DateOut", + "closed_at", + "resolved_at", + "completed_at", + "finished_at", + "end_date", + ), temporal=True, ) + if not start_column and temporal_columns: + start_column = temporal_columns[0] + if not end_column: + for candidate in temporal_columns: + if candidate.lower() != str(start_column or "").lower(): + end_column = candidate + break if start_column and end_column: start_ref = f"{table_ref}.{self._quote_sql_identifier(start_column)}" end_ref = f"{table_ref}.{self._quote_sql_identifier(end_column)}" @@ -3318,6 +3346,33 @@ def _get_unqueryable_metric_message( "a first-class temporal field." ) + if any( + term in normalized_query + for term in ( + "monthly", + "trend", + "turnaround", + "time", + "duration", + "elapsed", + "latest", + "recent", + "newest", + "last records", + ) + ): + has_temporal_field = any( + self._is_temporal_schema_type(str(column.get("type") or "")) + for table in self._parse_schema_tables(table_ddls) + for column in table.get("columns", []) + ) + if not has_temporal_field: + return ( + "The active datasource does not expose a queryable date or " + "timestamp column. I cannot build a time-based analysis " + "without a first-class temporal field." + ) + repair_cost_terms = ( "repair cost", "repair_cost", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index b28094b45a..75c0978b4b 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1150,6 +1150,54 @@ def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count() ) +def test_build_schema_grounded_table_question_sql_uses_any_temporal_column(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Show monthly record count for dbo.failure_patterns.", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR, + execution_date TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."execution_date") AS "year", ' + 'DATEPART(MONTH, "dbo_failure_patterns"."execution_date") AS "month", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."execution_date" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."execution_date"), ' + 'DATEPART(MONTH, "dbo_failure_patterns"."execution_date") ' + 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."execution_date") ASC, ' + 'DATEPART(MONTH, "dbo_failure_patterns"."execution_date") ASC' + ) + + +def test_build_schema_grounded_table_question_sql_rejects_time_question_without_temporal_column(): + service = AskService.__new__(AskService) + + assert ( + service._build_schema_grounded_table_question_sql( + "Show monthly record count for dbo.failure_patterns.", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR + ); + """ + ], + ) + is None + ) + + def test_build_validated_ask_result_rejects_status_for_product_line_pcb_question(): service = AskService.__new__(AskService) @@ -1311,3 +1359,22 @@ def test_build_validated_ask_result_rejects_sql_when_schema_is_missing(): ) assert result is None + + +def test_get_unqueryable_metric_message_rejects_turnaround_trend_without_temporal_column(): + service = AskService.__new__(AskService) + + message = service._get_unqueryable_metric_message( + "Generate a trend chart for average turnaround time by month.", + [ + """ + CREATE TABLE dbo_repair_logs ( + status VARCHAR, + repair_cost DOUBLE + ); + """ + ], + ) + + assert message is not None + assert "temporal field" in message From f56f045a1cb4b578fdca2bc48ae33c6c72c4330b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 00:49:05 +0530 Subject: [PATCH 0396/1087] Map generic temporal fields to schema columns --- .../src/pipelines/generation/utils/sql.py | 58 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 20 +++++++ 2 files changed, 78 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4bec10530f..2748aef836 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2652,6 +2652,57 @@ def _find_semantic_column_alias( return None +def _find_temporal_schema_column( + canonical_columns: dict[str, str], +) -> str | None: + temporal_aliases = ( + "created_at", + "updated_at", + "generated_at", + "opened_at", + "closed_at", + "completed_at", + "resolved_at", + "DateIn", + "DateOut", + "FailedAt", + "ModifiedAt", + "CreatedAt", + "execution_date", + "event_date", + "date", + "time", + "timestamp", + ) + for alias in temporal_aliases: + canonical = canonical_columns.get(_compact_sql_identifier(alias)) + if canonical: + return canonical + return next(iter(canonical_columns.values()), None) + + +def _is_temporal_identifier(identifier: str) -> bool: + normalized = _compact_sql_identifier(identifier) + return normalized in { + "createdat", + "updatedat", + "generatedat", + "openedat", + "closedat", + "completedat", + "resolvedat", + "datein", + "dateout", + "failedat", + "modifiedat", + "executiondate", + "eventdate", + "date", + "time", + "timestamp", + } + + def _split_table_reference(table_reference: str) -> list[str]: stripped = table_reference.strip() if not stripped: @@ -2932,6 +2983,11 @@ def replace_column_reference(match: re.Match[str]) -> str: canonical_column = canonical_columns.get( compact_column ) or _find_semantic_column_alias(normalized_column, canonical_columns) + if ( + not canonical_column + and _is_temporal_identifier(normalized_column) + ): + canonical_column = _find_temporal_schema_column(canonical_columns) if not canonical_column or canonical_column == normalized_column: return match.group(0) @@ -2968,6 +3024,8 @@ def replace_unqualified_identifier(match: re.Match[str]) -> str: canonical_column = canonical_columns.get( compact_identifier ) or _find_semantic_column_alias(identifier, canonical_columns) + if not canonical_column and _is_temporal_identifier(identifier): + canonical_column = _find_temporal_schema_column(canonical_columns) if not canonical_column: return match.group(0) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 421566f5eb..b7db27806f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -455,6 +455,26 @@ def test_normalize_sql_column_references_to_schema_maps_last_update_date_alias() ) == [] +def test_normalize_sql_column_references_to_schema_maps_created_at_to_temporal_column(): + sql = ( + 'SELECT DATEPART(YEAR, "dbo_DebugEntries_Staging2"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."created_at") AS "month" ' + 'FROM "dbo_DebugEntries_Staging2"' + ) + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_DebugEntries_Staging2": ["DebugEntryId", "DateIn", "DateOut"]}, + ) + + assert '"dbo_DebugEntries_Staging2"."created_at"' not in normalized + assert '"dbo_DebugEntries_Staging2"."DateIn"' in normalized + assert find_invalid_column_references( + normalized, + {"dbo_DebugEntries_Staging2": ["DebugEntryId", "DateIn", "DateOut"]}, + ) == [] + + def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid(): sql = 'SELECT "dbo_qSales1"."UnitPrice" FROM "dbo_qSales1"' normalized = normalize_sql_column_references_to_schema( From 7ed2895354d694ea8dead9f95531bea41748ca69 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 01:13:06 +0530 Subject: [PATCH 0397/1087] Revert "Map generic temporal fields to schema columns" This reverts commit f56f045a1cb4b578fdca2bc48ae33c6c72c4330b. --- .../src/pipelines/generation/utils/sql.py | 58 ------------------- .../pipelines/generation/test_sql_utils.py | 20 ------- 2 files changed, 78 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2748aef836..4bec10530f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2652,57 +2652,6 @@ def _find_semantic_column_alias( return None -def _find_temporal_schema_column( - canonical_columns: dict[str, str], -) -> str | None: - temporal_aliases = ( - "created_at", - "updated_at", - "generated_at", - "opened_at", - "closed_at", - "completed_at", - "resolved_at", - "DateIn", - "DateOut", - "FailedAt", - "ModifiedAt", - "CreatedAt", - "execution_date", - "event_date", - "date", - "time", - "timestamp", - ) - for alias in temporal_aliases: - canonical = canonical_columns.get(_compact_sql_identifier(alias)) - if canonical: - return canonical - return next(iter(canonical_columns.values()), None) - - -def _is_temporal_identifier(identifier: str) -> bool: - normalized = _compact_sql_identifier(identifier) - return normalized in { - "createdat", - "updatedat", - "generatedat", - "openedat", - "closedat", - "completedat", - "resolvedat", - "datein", - "dateout", - "failedat", - "modifiedat", - "executiondate", - "eventdate", - "date", - "time", - "timestamp", - } - - def _split_table_reference(table_reference: str) -> list[str]: stripped = table_reference.strip() if not stripped: @@ -2983,11 +2932,6 @@ def replace_column_reference(match: re.Match[str]) -> str: canonical_column = canonical_columns.get( compact_column ) or _find_semantic_column_alias(normalized_column, canonical_columns) - if ( - not canonical_column - and _is_temporal_identifier(normalized_column) - ): - canonical_column = _find_temporal_schema_column(canonical_columns) if not canonical_column or canonical_column == normalized_column: return match.group(0) @@ -3024,8 +2968,6 @@ def replace_unqualified_identifier(match: re.Match[str]) -> str: canonical_column = canonical_columns.get( compact_identifier ) or _find_semantic_column_alias(identifier, canonical_columns) - if not canonical_column and _is_temporal_identifier(identifier): - canonical_column = _find_temporal_schema_column(canonical_columns) if not canonical_column: return match.group(0) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index b7db27806f..421566f5eb 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -455,26 +455,6 @@ def test_normalize_sql_column_references_to_schema_maps_last_update_date_alias() ) == [] -def test_normalize_sql_column_references_to_schema_maps_created_at_to_temporal_column(): - sql = ( - 'SELECT DATEPART(YEAR, "dbo_DebugEntries_Staging2"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."created_at") AS "month" ' - 'FROM "dbo_DebugEntries_Staging2"' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_DebugEntries_Staging2": ["DebugEntryId", "DateIn", "DateOut"]}, - ) - - assert '"dbo_DebugEntries_Staging2"."created_at"' not in normalized - assert '"dbo_DebugEntries_Staging2"."DateIn"' in normalized - assert find_invalid_column_references( - normalized, - {"dbo_DebugEntries_Staging2": ["DebugEntryId", "DateIn", "DateOut"]}, - ) == [] - - def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid(): sql = 'SELECT "dbo_qSales1"."UnitPrice" FROM "dbo_qSales1"' normalized = normalize_sql_column_references_to_schema( From e1fa3a0ae25b890fa1d20f00a424e5c79d12b12f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 01:13:06 +0530 Subject: [PATCH 0398/1087] Revert "Use only schema-backed temporal columns" This reverts commit 1fff30cfd5a747aa38e99cc655316cadfb788d29. --- wren-ai-service/src/web/v1/services/ask.py | 83 ++++--------------- .../pytest/services/test_ask_sales_sql.py | 67 --------------- 2 files changed, 14 insertions(+), 136 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3f91f65b03..68ba41d1dc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1031,14 +1031,6 @@ def _find_first_schema_column( return column return None - def _find_any_temporal_schema_column(self, table: dict[str, Any]) -> str | None: - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if column_name and self._is_temporal_schema_type(column_type): - return column_name - return None - def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' @@ -1142,7 +1134,18 @@ def _find_temporal_column_for_query( ): return column_name - return self._find_any_temporal_schema_column(table) + return self._find_first_schema_column( + table, + ( + "created_at", + "createdat", + "created", + "date", + "time", + "timestamp", + "updated_at", + ), + ) def _build_schema_grounded_table_question_sql( self, query: str, table_ddls: list[str] @@ -3168,47 +3171,16 @@ def _build_schema_grounded_operational_sql( limit = int(limit_match.group(1)) if limit_match else 10 if wants_elapsed_time: - temporal_columns = [ - str(column.get("name") or "") - for column in table.get("columns", []) - if column.get("name") - and self._is_temporal_schema_type(str(column.get("type") or "")) - ] start_column = self._find_schema_column( table, - ( - "created_at", - "created", - "DateIn", - "execution_date", - "opened_at", - "started_at", - "start_date", - "begin_date", - ), + ("created_at", "created", "DateIn"), temporal=True, ) end_column = self._find_schema_column( table, - ( - "updated_at", - "updated", - "DateOut", - "closed_at", - "resolved_at", - "completed_at", - "finished_at", - "end_date", - ), + ("updated_at", "updated", "DateOut", "closed_at", "resolved_at"), temporal=True, ) - if not start_column and temporal_columns: - start_column = temporal_columns[0] - if not end_column: - for candidate in temporal_columns: - if candidate.lower() != str(start_column or "").lower(): - end_column = candidate - break if start_column and end_column: start_ref = f"{table_ref}.{self._quote_sql_identifier(start_column)}" end_ref = f"{table_ref}.{self._quote_sql_identifier(end_column)}" @@ -3346,33 +3318,6 @@ def _get_unqueryable_metric_message( "a first-class temporal field." ) - if any( - term in normalized_query - for term in ( - "monthly", - "trend", - "turnaround", - "time", - "duration", - "elapsed", - "latest", - "recent", - "newest", - "last records", - ) - ): - has_temporal_field = any( - self._is_temporal_schema_type(str(column.get("type") or "")) - for table in self._parse_schema_tables(table_ddls) - for column in table.get("columns", []) - ) - if not has_temporal_field: - return ( - "The active datasource does not expose a queryable date or " - "timestamp column. I cannot build a time-based analysis " - "without a first-class temporal field." - ) - repair_cost_terms = ( "repair cost", "repair_cost", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 75c0978b4b..b28094b45a 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1150,54 +1150,6 @@ def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count() ) -def test_build_schema_grounded_table_question_sql_uses_any_temporal_column(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Show monthly record count for dbo.failure_patterns.", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR, - execution_date TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."execution_date") AS "year", ' - 'DATEPART(MONTH, "dbo_failure_patterns"."execution_date") AS "month", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."execution_date" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."execution_date"), ' - 'DATEPART(MONTH, "dbo_failure_patterns"."execution_date") ' - 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."execution_date") ASC, ' - 'DATEPART(MONTH, "dbo_failure_patterns"."execution_date") ASC' - ) - - -def test_build_schema_grounded_table_question_sql_rejects_time_question_without_temporal_column(): - service = AskService.__new__(AskService) - - assert ( - service._build_schema_grounded_table_question_sql( - "Show monthly record count for dbo.failure_patterns.", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR - ); - """ - ], - ) - is None - ) - - def test_build_validated_ask_result_rejects_status_for_product_line_pcb_question(): service = AskService.__new__(AskService) @@ -1359,22 +1311,3 @@ def test_build_validated_ask_result_rejects_sql_when_schema_is_missing(): ) assert result is None - - -def test_get_unqueryable_metric_message_rejects_turnaround_trend_without_temporal_column(): - service = AskService.__new__(AskService) - - message = service._get_unqueryable_metric_message( - "Generate a trend chart for average turnaround time by month.", - [ - """ - CREATE TABLE dbo_repair_logs ( - status VARCHAR, - repair_cost DOUBLE - ); - """ - ], - ) - - assert message is not None - assert "temporal field" in message From dc8f313b2ccf00f25e3bb3ff4181245fb0f41ca2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 01:13:06 +0530 Subject: [PATCH 0399/1087] Revert "Remove historical context from ask flow" This reverts commit 8926a3fffd453ccad241c195d164524cb7377a57. --- wren-ai-service/src/web/v1/services/ask.py | 153 ++++++++++++++++++--- wren-ui/src/pages/api/v1/ask.ts | 9 +- wren-ui/src/pages/api/v1/generate_sql.ts | 8 +- wren-ui/src/pages/api/v1/stream/ask.ts | 8 +- 4 files changed, 155 insertions(+), 23 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 68ba41d1dc..b7e73f36af 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4228,11 +4228,16 @@ async def ask( return results logger.info(f"Ask pipeline started for query_id: {query_id}") - # Always answer from the active datasource schema and current question. - # Prior thread context, saved SQL examples, and historical questions are - # intentionally not reused here to prevent stale or cross-datasource - # leakage from influencing a fresh answer. - histories: list[AskHistory] = [] + histories = ask_request.histories[: self._max_histories][ + ::-1 + ] # reverse the order of histories + if histories and not self._should_use_histories_for_query(user_query): + logger.info( + "Ignoring thread histories for independent question. query_id=%s query=%s", + query_id, + user_query, + ) + histories = [] rephrased_question = None intent_reasoning = None sql_generation_reasoning = None @@ -4570,24 +4575,132 @@ async def ask( query_id, ) - if not api_results: - original_user_query = user_query - # Only user instructions are kept. Prior SQL samples are not - # reused for fresh questions because they can bias the model - # toward stale or unrelated queries. - instructions_task = await self._run_with_timeout( - "Instruction retrieval", - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", - ), + historical_question_result = [] + should_skip_pre_sql_retrieval = self._is_data_analysis_query( + user_query + ) + if should_skip_pre_sql_retrieval: + rephrased_question = user_query + intent_reasoning = ( + "Detected a deployed-data analytics question; skipping " + "intent classification and using SQL generation." + ) + sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + logger.info( + "Skipping pre-SQL retrieval for analytics query_id %s: %s", + query_id, + user_query, ) - sql_samples = [] - instructions = instructions_task["formatted_output"].get( - "documents", [] + if ( + not api_results + and not should_skip_pre_sql_retrieval + and self._should_reuse_historical_question_sql( + user_query, histories ) + ): + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + + try: + historical_question = await self._run_with_timeout( + "Historical question retrieval", + self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ), + timeout_seconds=min(understanding_timeout_seconds, 10), + ) + + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] + except TimeoutError as exc: + logger.warning( + "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) + + valid_historical_results = [] + for result in historical_question_result: + historical_question_text = result.get("question") + if not self._is_reusable_historical_question( + user_query, historical_question_text + ): + logger.info( + "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", + query_id, + user_query, + historical_question_text, + ) + continue + + sql_statement = result.get("statement") + if not self._is_valid_select_sql(sql_statement): + logger.warning( + "Ignoring historical question without valid SQL for query_id %s", + query_id, + ) + continue + valid_historical_results.append( + AskResult( + **{ + "sql": sql_statement.strip(), + "type": "view" if result.get("viewId") else "llm", + "viewId": result.get("viewId"), + } + ) + ) + + if valid_historical_results: + api_results = valid_historical_results + sql_generation_reasoning = "" + elif not api_results and not should_skip_pre_sql_retrieval: + original_user_query = user_query + # Run both pipeline operations concurrently + try: + sql_samples_task, instructions_task = await self._run_with_timeout( + "SQL pair and instruction retrieval", + asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), + ), + timeout_seconds=understanding_timeout_seconds, + ) + + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] + ) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) + except TimeoutError as exc: + logger.warning( + "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) + sql_samples = [] + instructions = [] if self._allow_intent_classification: try: diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index cdcceb703b..d2e98d0eb2 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -10,6 +10,7 @@ import { MAX_WAIT_TIME, isAskResultFinished, validateSummaryResult, + transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; import { AskResult, @@ -26,6 +27,7 @@ const logger = getLogger('API_ASK'); logger.level = 'debug'; const { + apiHistoryRepository, projectService, deployService, wrenAIAdaptor, @@ -73,11 +75,16 @@ export default async function handler( // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); + // Get conversation history if threadId is provided + const histories = threadId + ? await apiHistoryRepository.findAllBy({ threadId }) + : undefined; + // Step 1: Generate SQL const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, - histories: undefined, + histories: transformHistoryInput(histories) as any, configurations: { language: language || WrenAILanguage[project.language] || WrenAILanguage.EN, diff --git a/wren-ui/src/pages/api/v1/generate_sql.ts b/wren-ui/src/pages/api/v1/generate_sql.ts index 55adb6edfe..fa5b3859e1 100644 --- a/wren-ui/src/pages/api/v1/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/generate_sql.ts @@ -12,6 +12,7 @@ import { MAX_WAIT_TIME, isAskResultFinished, validateAskResult, + transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; import { DataSourceName } from '@server/types'; @@ -19,6 +20,7 @@ const logger = getLogger('API_GENERATE_SQL'); logger.level = 'debug'; const { + apiHistoryRepository, projectService, deployService, wrenAIAdaptor, @@ -70,10 +72,14 @@ export default async function handler( ); } + // ask AI service to generate SQL + const histories = threadId + ? await apiHistoryRepository.findAllBy({ threadId }) + : undefined; const task = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, - histories: undefined, + histories: transformHistoryInput(histories) as any, configurations: { language: language || WrenAILanguage[project.language] || WrenAILanguage.EN, diff --git a/wren-ui/src/pages/api/v1/stream/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts index 98159f9048..12e53eb4db 100644 --- a/wren-ui/src/pages/api/v1/stream/ask.ts +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -8,6 +8,7 @@ import { MAX_WAIT_TIME, isAskResultFinished, validateSummaryResult, + transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; import { AskResult, @@ -135,6 +136,11 @@ export default async function handler( // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); + // Get conversation history if threadId is provided + const histories = threadId + ? await apiHistoryRepository.findAllBy({ threadId }) + : undefined; + // Step 1: Generate SQL sendStateUpdate(res, StateType.SQL_GENERATION_START, { question, @@ -145,7 +151,7 @@ export default async function handler( const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, - histories: undefined, + histories: transformHistoryInput(histories) as any, configurations: { language: language || WrenAILanguage[project.language] || WrenAILanguage.EN, From 495e03381c38acd6d04eabdb226c7900ad67f232 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 01:26:48 +0530 Subject: [PATCH 0400/1087] Revert "Reject SQL when schema validation has no tables" This reverts commit af7ce907efe5c4e915753bc5a623f460a500386c. --- wren-ai-service/src/web/v1/services/ask.py | 6 ------ .../tests/pytest/services/test_ask_sales_sql.py | 12 ------------ 2 files changed, 18 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b7e73f36af..10207e31e8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4059,12 +4059,6 @@ def _build_validated_ask_result_from_sql( return None schema_tables = self._parse_schema_tables(table_ddls) - if not schema_tables: - logger.warning( - "Ignoring SQL because no active schema tables were available for validation. sql=%s", - ask_result.sql, - ) - return None valid_tables = { str(table.get("name") or "").lower(): table for table in schema_tables diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index b28094b45a..1d3cd8a197 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1299,15 +1299,3 @@ def test_build_validated_ask_result_accepts_unqualified_valid_columns(): ) assert result is not None - - -def test_build_validated_ask_result_rejects_sql_when_schema_is_missing(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - 'SELECT execution_date FROM "dbo_repair_logs"', - [], - "Generate a trend chart for average turnaround time by month.", - ) - - assert result is None From 95ca25a820632f827489fea7a15a14a6cc01855f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 11:18:54 +0530 Subject: [PATCH 0401/1087] Handle order date distribution questions --- wren-ai-service/src/web/v1/services/ask.py | 80 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 43 ++++++++++ 2 files changed, 123 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 10207e31e8..e2e37ba418 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1516,6 +1516,36 @@ def _append_not_null_filters( return f"{where_clause.rstrip()} AND {' AND '.join(conditions)} " return f" WHERE {' AND '.join(conditions)} " + def _build_schema_literal_filter_conditions( + self, + query: str, + table: dict[str, Any], + table_ref: str, + ) -> list[str]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + conditions: list[str] = [] + + if "backlog" in normalized_query: + filter_column = self._find_schema_column( + table, + ( + "Category", + "OrderCategory", + "Order Category", + "Status", + "OrderStatus", + "Order Status", + "Stage", + "OrderStage", + "Order Stage", + ), + ) + if filter_column: + filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" + conditions.append(f"{filter_ref} = 'Backlog'") + + return conditions + def _select_best_analytics_table( self, tables: list[dict[str, Any]], @@ -1740,6 +1770,16 @@ def _build_schema_grounded_analytics_sql( or "over time" in normalized_query or "last 12 months" in normalized_query ) + wants_date_distribution = ( + any( + term in normalized_query + for term in ("distribution", "breakdown", "split") + ) + and any( + term in normalized_query + for term in ("date", "dates", "orddate", "order date", "order dates") + ) + ) wants_order_count_metric = ( any(term in normalized_query for term in ("order", "orders", "new order", "new orders")) and not any( @@ -1765,6 +1805,7 @@ def _build_schema_grounded_analytics_sql( ) wants_date = ( wants_trend + or wants_date_distribution or mentions_date_column or "this year" in normalized_query or bool(re.search(r"\b20\d{2}\b", normalized_query)) @@ -1816,6 +1857,45 @@ def _build_schema_grounded_analytics_sql( f"ORDER BY {metric_ref} DESC" ) + if wants_date_distribution and date_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + select_parts = [ + f"DATEPART(YEAR, {date_ref}) AS \"year\"", + f"DATEPART(MONTH, {date_ref}) AS \"month\"", + *[ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ], + 'COUNT(*) AS "OrderCount"', + ] + group_parts = [ + f"DATEPART(YEAR, {date_ref})", + f"DATEPART(MONTH, {date_ref})", + *dimension_refs, + ] + where_clause = self._append_not_null_filters( + self._build_date_filter(table_name, date_column, query), + [date_ref, *dimension_refs], + ) + extra_conditions = self._build_schema_literal_filter_conditions( + query, + table, + table_ref, + ) + if extra_conditions: + where_clause = ( + f"{where_clause.rstrip()} AND {' AND '.join(extra_conditions)} " + if where_clause.strip() + else f" WHERE {' AND '.join(extra_conditions)} " + ) + return ( + f"SELECT {', '.join(select_parts)} FROM {table_ref}" + f"{where_clause}" + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}), COUNT(*) DESC" + ) + wants_top_per_group = ( len(dimensions) >= 2 and any(term in normalized_query for term in ("highest", "top", "most")) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 1d3cd8a197..4d33d54d55 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -297,6 +297,49 @@ def test_build_schema_grounded_sales_sql_for_orders_by_dimensions(): ) +def test_build_schema_grounded_sales_sql_for_order_date_distribution_by_dimensions(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + 'In the "Backlog" category, what is the distribution of order dates ' + "(OrdDate) for each product type (ProdType) sold in each market " + "segment (Market), considering the salesperson responsible " + "(SalesPerson)?", + [ + """ + CREATE TABLE dbo_tblSales ( + Category VARCHAR, + Market VARCHAR, + ProdType VARCHAR, + SalesPerson VARCHAR, + SalesValue DOUBLE, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_tblSales"."OrdDate") AS "year", ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") AS "month", ' + '"dbo_tblSales"."SalesPerson" AS "SalesPerson", ' + '"dbo_tblSales"."Market" AS "Market", ' + '"dbo_tblSales"."ProdType" AS "ProdType", ' + 'COUNT(*) AS "OrderCount" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."OrdDate" IS NOT NULL ' + 'AND "dbo_tblSales"."SalesPerson" IS NOT NULL ' + 'AND "dbo_tblSales"."Market" IS NOT NULL ' + 'AND "dbo_tblSales"."ProdType" IS NOT NULL ' + 'AND "dbo_tblSales"."Category" = \'Backlog\' ' + 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), ' + '"dbo_tblSales"."SalesPerson", "dbo_tblSales"."Market", ' + '"dbo_tblSales"."ProdType" ' + 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' + 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), COUNT(*) DESC' + ) + + def test_build_schema_grounded_sales_sql_for_top_new_order_detail_rows(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 56fbae562c0e0d35847a77239154073c32cdfd2f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 11:33:33 +0530 Subject: [PATCH 0402/1087] Retry active schema for data questions --- wren-ai-service/src/web/v1/services/ask.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e2e37ba418..e6cbb0da59 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -5031,6 +5031,32 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + if not documents and self._is_data_analysis_query(user_query): + logger.info( + "Query-based schema retrieval returned no tables for data question; " + "retrying full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) From 094928c1eeaf4cc0913b7fd46f7fad95d259d0fd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 11:51:03 +0530 Subject: [PATCH 0403/1087] Use date buckets for monthly repair volume --- wren-ai-service/src/web/v1/services/ask.py | 88 +++++++++++++++---- .../pytest/services/test_ask_sales_sql.py | 66 ++++++++++++++ 2 files changed, 136 insertions(+), 18 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e6cbb0da59..c10e21c237 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1645,6 +1645,11 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql + if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( + query, table_ddls + ): + return monthly_repair_volume_sql + if categorical_count_sql := self._build_generic_categorical_count_sql( query, tables ): @@ -2792,26 +2797,73 @@ def _build_monthly_repair_volume_sql( if not wants_monthly_repairs: return None - has_repair_created_at = self._schema_has_table_column( - table_ddls, - "dbo_repair_logs", - "created_at", - table_names=table_names, - ) - if has_repair_created_at: - return ( - 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + tables = self._parse_schema_tables(table_ddls) + scored_tables: list[tuple[int, dict[str, Any], str]] = [] + for table in tables: + table_name = str(table.get("name") or "") + if table_names and table_name not in table_names: + continue + + date_column = self._find_schema_column( + table, + ( + "created_at", + "createdAt", + "created", + "DateIn", + "Date", + "repair_date", + "RepairDate", + "opened_at", + "started_at", + ), + temporal=True, ) + if not date_column: + date_column = self._find_any_temporal_schema_column(table) + if not date_column: + continue - return None + normalized_table = self._normalize_schema_token(table_name) + score = 0 + if "repair" in normalized_table: + score += 30 + if "debugentries" in normalized_table or "debugentry" in normalized_table: + score += 25 + if "log" in normalized_table: + score += 10 + if self._find_schema_column( + table, + ("DebugEntryId", "RepairId", "repair_id", "id"), + ): + score += 5 + scored_tables.append((score, table, date_column)) + + if not scored_tables: + return None + + _score, table, date_column = sorted( + scored_tables, + key=lambda item: item[0], + reverse=True, + )[0] + table_name = str(table.get("name") or "") + if not table_name: + return None + + table_ref = self._quote_sql_identifier(table_name) + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "repair_count" ' + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " + f"DATEPART(MONTH, {date_ref}) ASC" + ) def _is_direct_heuristic_sql_query(self, query: str) -> bool: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 4d33d54d55..7f77e50040 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -778,6 +778,72 @@ def test_build_manufacturing_throughput_sql_uses_active_unit_and_date_columns(): ) +def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Generate a line chart showing monthly repair volume for the last 12 months.", + [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' + ) + assert '"status"' not in sql + + +def test_build_monthly_repair_volume_sql_uses_debug_entry_date_column(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Generate a line chart showing monthly repair volume for the last 12 months.", + [ + """ + CREATE TABLE dbo_DebugEntries ( + DebugEntryId VARCHAR, + Status VARCHAR, + DateIn TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_DebugFixLogs ( + DebugEntryId VARCHAR, + FixId VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "year", ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_DebugEntries" ' + 'WHERE "dbo_DebugEntries"."DateIn" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn"), ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ' + 'ORDER BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC, ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' + ) + assert '"Status"' not in sql + + def test_get_unqueryable_metric_message_for_throughput_without_unit_column(): service = AskService.__new__(AskService) From 05fdb3549c3d2ae220594f75380fa47901dbf3b0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 12:35:26 +0530 Subject: [PATCH 0404/1087] Handle explicit synced table questions --- wren-ai-service/src/web/v1/services/ask.py | 47 ++++++++++++++----- .../pytest/services/test_ask_sales_sql.py | 46 ++++++++++++++++++ 2 files changed, 80 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c10e21c237..8910eb9d3b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1037,16 +1037,31 @@ def _quote_sql_identifier(self, identifier: str) -> str: def _normalize_schema_identifier_key(self, value: str) -> str: return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + def _schema_identifier_alias_keys(self, value: str) -> set[str]: + raw_value = str(value or "").strip() + base_key = self._normalize_schema_identifier_key(raw_value) + separator_normalized_key = self._normalize_schema_identifier_key( + re.sub(r"[.$]", "_", raw_value) + ) + compact_parts_key = "".join( + self._normalize_schema_identifier_key(part) + for part in re.split(r"[.$_]+", raw_value) + if part + ) + return { + key + for key in (base_key, separator_normalized_key, compact_parts_key) + if key + } + def _table_matches_query(self, table_name: str, query: str) -> bool: - normalized_query = self._normalize_schema_identifier_key(query) - normalized_table = self._normalize_schema_identifier_key(table_name) + query_keys = self._schema_identifier_alias_keys(query) short_table = re.split(r"[.$_]", str(table_name or ""))[-1] - normalized_short_table = self._normalize_schema_identifier_key(short_table) - return bool( - normalized_table - and normalized_table in normalized_query - or normalized_short_table - and normalized_short_table in normalized_query + table_keys = self._schema_identifier_alias_keys(table_name) + table_keys.update(self._schema_identifier_alias_keys(short_table)) + return any( + table_key and any(table_key in query_key for query_key in query_keys) + for table_key in table_keys ) def _find_best_schema_table_for_query( @@ -1109,17 +1124,17 @@ def _find_dimension_column_for_query( if column_name and self._query_mentions_column(query, column_name): return column_name - text_columns = [ + candidate_columns = [ str(column.get("name")) for column in table.get("columns", []) if column.get("name") - and self._is_text_schema_type(str(column.get("type") or "")) + and not self._is_temporal_schema_type(str(column.get("type") or "")) ] for candidate in ("name", "category", "type", "status", "code"): column = self._find_schema_column(table, (candidate,)) - if column in text_columns: + if column in candidate_columns: return column - return text_columns[0] if text_columns else None + return candidate_columns[0] if candidate_columns else None def _find_temporal_column_for_query( self, query: str, table: dict[str, Any] @@ -1280,6 +1295,7 @@ def _build_explicit_table_preview_sql( return None normalized_query_key = re.sub(r"[^a-z0-9]", "", normalized_query.lower()) + normalized_query_keys = self._schema_identifier_alias_keys(normalized_query) scored_tables: list[tuple[int, str]] = [] for table in tables: table_name = table.get("name") @@ -1289,7 +1305,12 @@ def _build_explicit_table_preview_sql( normalized_table = re.sub(r"[^a-z0-9]", "", table_name.lower()) if not normalized_table: continue - if normalized_table in normalized_query_key: + table_keys = self._schema_identifier_alias_keys(table_name) + if normalized_table in normalized_query_key or any( + table_key in query_key + for table_key in table_keys + for query_key in normalized_query_keys + ): scored_tables.append((100 + len(normalized_table), table_name)) continue diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 7f77e50040..3aa808cc7e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1130,6 +1130,25 @@ def test_build_schema_grounded_table_question_sql_for_record_count(): assert sql == 'SELECT COUNT(*) AS "RecordCount" FROM "dbo_failure_patterns"' +def test_build_schema_grounded_table_question_sql_matches_dot_table_to_underscore_table(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "How many records are in dbo.ytblTarrifsRec?", + [ + """ + CREATE TABLE dbo_ytblTarrifsRec ( + EntrySummaryNumber2 VARCHAR, + LiquidationStatus VARCHAR, + LiquidationDate TIMESTAMP + ); + """ + ], + ) + + assert sql == 'SELECT COUNT(*) AS "RecordCount" FROM "dbo_ytblTarrifsRec"' + + def test_build_schema_grounded_table_question_sql_for_name_distribution(): service = AskService.__new__(AskService) @@ -1156,6 +1175,33 @@ def test_build_schema_grounded_table_question_sql_for_name_distribution(): ) +def test_build_schema_grounded_table_question_sql_for_numeric_column_distribution(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "What is the distribution of EntrySummaryNumber2 in dbo.ytblTarrifsRec?", + [ + """ + CREATE TABLE dbo_ytblTarrifsRec ( + EntrySummaryNumber2 BIGINT, + LiquidationStatus VARCHAR, + LiquidationDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_ytblTarrifsRec"."EntrySummaryNumber2" AS ' + '"EntrySummaryNumber2", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ytblTarrifsRec" ' + 'WHERE "dbo_ytblTarrifsRec"."EntrySummaryNumber2" IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsRec"."EntrySummaryNumber2" ' + 'ORDER BY COUNT(*) DESC' + ) + + def test_build_schema_grounded_table_question_sql_for_highest_occurrences(): service = AskService.__new__(AskService) From b8339d6b6a8e345f0ae64fb0a57c42b074ed6e6c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 13:29:59 +0530 Subject: [PATCH 0405/1087] Support country revenue questions --- wren-ai-service/src/web/v1/services/ask.py | 2 ++ .../pytest/services/test_ask_sales_sql.py | 25 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8910eb9d3b..bbf68f5692 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1743,6 +1743,8 @@ def _build_schema_grounded_analytics_sql( dimension_candidates.append(("Market", "MarketType", "Region")) if "region" in normalized_query: dimension_candidates.append(("Region", "Market", "Area", "Territory")) + if "country" in normalized_query or "countries" in normalized_query: + dimension_candidates.append(("Country", "CountryName", "Nation", "Market")) if "division" in normalized_query: dimension_candidates.append(("Division",)) if ( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 3aa808cc7e..d404f75016 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -519,6 +519,31 @@ def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): ) +def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_country(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Which countries have the highest order revenue?", + [ + """ + CREATE TABLE dbo_tblSales ( + Country VARCHAR, + OrderValue DOUBLE, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblSales"."Country" AS "Country", ' + 'SUM("dbo_tblSales"."OrderValue") AS "TotalOrderValue" ' + 'FROM "dbo_tblSales" ' + 'WHERE "dbo_tblSales"."Country" IS NOT NULL ' + 'GROUP BY "dbo_tblSales"."Country" ' + 'ORDER BY SUM("dbo_tblSales"."OrderValue") DESC' + ) + + def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 8a56b0c8bf7565b679b9fbd61cec3460819bfa50 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 13:47:23 +0530 Subject: [PATCH 0406/1087] Handle order value ranking questions --- wren-ai-service/src/web/v1/services/ask.py | 22 +++++++- .../pytest/services/test_ask_sales_sql.py | 50 +++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bbf68f5692..1c070bb5a1 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1740,7 +1740,7 @@ def _build_schema_grounded_analytics_sql( if "business unit" in normalized_query or "bu" in normalized_query: dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) if "market" in normalized_query: - dimension_candidates.append(("Market", "MarketType", "Region")) + dimension_candidates.append(("Market", "MarketType", "MarketName", "Region", "Country")) if "region" in normalized_query: dimension_candidates.append(("Region", "Market", "Area", "Territory")) if "country" in normalized_query or "countries" in normalized_query: @@ -1777,8 +1777,10 @@ def _build_schema_grounded_analytics_sql( "SalesValue", "FXSalesValue", "Revenue", + "TotalRevenue", "Amount", "Value", + "TotalOrderValue", "Cost", "Qty", "Quantity", @@ -2062,6 +2064,22 @@ def _build_schema_grounded_analytics_sql( limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) limit = int(limit_match.group(1)) if limit_match else 10 top_clause = f"TOP {limit} " if wants_top else "" + sort_direction = ( + "ASC" + if any( + term in normalized_query + for term in ( + "losing", + "lowest", + "least", + "bottom", + "declining", + "underperforming", + "smallest", + ) + ) + else "DESC" + ) date_filter = ( self._build_date_filter(table_name, date_column, query) if date_column @@ -2079,7 +2097,7 @@ def _build_schema_grounded_analytics_sql( f"FROM {table_ref}" f"{self._append_not_null_filters(date_filter, dimension_refs)} " f"GROUP BY {', '.join(dimension_refs)} " - f"ORDER BY {metric_expr} DESC" + f"ORDER BY {metric_expr} {sort_direction}" ) def _build_contribution_sql( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index d404f75016..fc08ee084a 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -544,6 +544,56 @@ def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_country(): ) +def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_prefixed_country(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Which countries have the highest order revenue?", + [ + """ + CREATE TABLE dbo_xStageLoad8 ( + col_07_Country VARCHAR, + TotalOrderValue DOUBLE, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_xStageLoad8"."col_07_Country" AS "col_07_Country", ' + 'SUM("dbo_xStageLoad8"."TotalOrderValue") AS "TotalTotalOrderValue" ' + 'FROM "dbo_xStageLoad8" ' + 'WHERE "dbo_xStageLoad8"."col_07_Country" IS NOT NULL ' + 'GROUP BY "dbo_xStageLoad8"."col_07_Country" ' + 'ORDER BY SUM("dbo_xStageLoad8"."TotalOrderValue") DESC' + ) + + +def test_build_schema_grounded_sales_sql_for_losing_order_value_by_market(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Which markets are losing order value?", + [ + """ + CREATE TABLE dbo_tblStageNewOrders ( + Market VARCHAR, + TotalOrderValue DOUBLE, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblStageNewOrders"."Market" AS "Market", ' + 'SUM("dbo_tblStageNewOrders"."TotalOrderValue") AS "TotalTotalOrderValue" ' + 'FROM "dbo_tblStageNewOrders" ' + 'WHERE "dbo_tblStageNewOrders"."Market" IS NOT NULL ' + 'GROUP BY "dbo_tblStageNewOrders"."Market" ' + 'ORDER BY SUM("dbo_tblStageNewOrders"."TotalOrderValue") ASC' + ) + + def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 7c503ed4880ec8d53afd603960c876385032363b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 14:09:08 +0530 Subject: [PATCH 0407/1087] Count restored destination databases by active column --- wren-ai-service/src/web/v1/services/ask.py | 17 +++++++++++ .../pytest/services/test_ask_sales_sql.py | 30 +++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 1c070bb5a1..88c5341b03 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2186,7 +2186,10 @@ def _build_generic_categorical_count_sql( "frequency", "group by", "grouped by", + "most often", + "often", "pie chart", + "restored", "status", "type", "category", @@ -2247,6 +2250,20 @@ def _build_generic_categorical_count_sql( score += 80 if "type" in normalized_query and "type" in normalized_column: score += 70 + if ( + "destination" in normalized_query + and "destination" in normalized_column + and ( + "database" in normalized_query + or "databases" in normalized_query + ) + and ( + "name" in normalized_column + or "phys" in normalized_column + or "db" in normalized_column + ) + ): + score += 140 if any(pattern == normalized_column for pattern in low_value_column_patterns): score -= 100 elif any( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index fc08ee084a..e3abde7150 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -677,6 +677,36 @@ def test_build_schema_grounded_sql_counts_categorical_status_values(): ) +def test_build_schema_grounded_sql_counts_destination_databases_restored_most_often(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Which destination databases were restored most often?", + [ + """ + CREATE TABLE dbo_db_policies ( + id VARCHAR, + policy_name VARCHAR, + destination_phys_name VARCHAR, + restore_date TIMESTAMP, + restore_type VARCHAR, + status VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_db_policies"."destination_phys_name" AS ' + '"destination_phys_name", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_db_policies" ' + 'WHERE "dbo_db_policies"."destination_phys_name" IS NOT NULL ' + 'GROUP BY "dbo_db_policies"."destination_phys_name" ' + 'ORDER BY COUNT(*) DESC' + ) + assert "destination_database_name" not in sql + + def test_build_schema_grounded_sales_sql_for_yoy_waterfall_dimensions(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From a0b7c66e8a33991a955af7452990d3c2dbdcfc62 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 15:07:38 +0530 Subject: [PATCH 0408/1087] Feed active datasource metadata to LLM prompts --- wren-ai-service/src/globals.py | 2 + .../pipelines/generation/chart_generation.py | 20 ++- .../generation/followup_sql_generation.py | 24 ++-- .../followup_sql_generation_reasoning.py | 8 +- .../src/pipelines/generation/sql_answer.py | 14 ++ .../pipelines/generation/sql_correction.py | 17 ++- .../pipelines/generation/sql_generation.py | 20 +-- .../generation/sql_generation_reasoning.py | 5 +- .../src/pipelines/generation/utils/sql.py | 73 ++--------- .../retrieval/db_schema_retrieval.py | 124 +----------------- wren-ai-service/src/web/v1/services/ask.py | 91 +++++++++++-- wren-ai-service/src/web/v1/services/chart.py | 26 ++++ .../src/web/v1/services/sql_answer.py | 26 ++++ .../retrieval/test_db_schema_retrieval.py | 75 +++++++++-- 14 files changed, 292 insertions(+), 233 deletions(-) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index d5cf1bb6fb..3c8ab0ab30 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -189,6 +189,7 @@ def create_service_container( chart_service=services.ChartService( pipelines={ "sql_executor": _sql_executor_pipeline, + "db_schema_retrieval": _db_schema_retrieval_pipeline, "chart_generation": generation.ChartGeneration( **pipe_components["chart_generation"], ), @@ -206,6 +207,7 @@ def create_service_container( ), sql_answer_service=services.SqlAnswerService( pipelines={ + "db_schema_retrieval": _db_schema_retrieval_pipeline, "preprocess_sql_data": retrieval.PreprocessSqlData( **pipe_components["preprocess_sql_data"], ), diff --git a/wren-ai-service/src/pipelines/generation/chart_generation.py b/wren-ai-service/src/pipelines/generation/chart_generation.py index 1c22161977..1b0c159655 100644 --- a/wren-ai-service/src/pipelines/generation/chart_generation.py +++ b/wren-ai-service/src/pipelines/generation/chart_generation.py @@ -42,6 +42,16 @@ """ chart_generation_user_prompt_template = """ +{% if documents %} +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource. Use it to +understand the schema, tables, columns, metrics, views, and relationships behind +the SQL before choosing a chart. +{% for document in documents %} + {{ document }} +{% endfor %} +{% endif %} + ### INPUT ### Question: {{ query }} SQL: {{ sql }} @@ -70,6 +80,7 @@ def prompt( language: str, custom_instruction: str, prompt_builder: PromptBuilder, + documents: list[str] | None = None, ) -> dict: sample_data = preprocess_data.get("sample_data") sample_column_values = preprocess_data.get("sample_column_values") @@ -81,6 +92,7 @@ def prompt( sample_column_values=sample_column_values, language=language, custom_instruction=custom_instruction, + documents=documents or [], ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -140,7 +152,11 @@ def __init__( "post_processor": ChartGenerationPostProcessor(), } - with open("src/pipelines/generation/utils/vega-lite-schema-v5.json", "r") as f: + with open( + "src/pipelines/generation/utils/vega-lite-schema-v5.json", + "r", + encoding="utf-8", + ) as f: _vega_schema = orjson.loads(f.read()) self._configs = { @@ -160,6 +176,7 @@ async def run( language: str, remove_data_from_chart_schema: bool = True, custom_instruction: Optional[str] = None, + contexts: Optional[list[str]] = None, ) -> dict: logger.info("Chart Generation pipeline is running...") return await self._pipe.execute( @@ -171,6 +188,7 @@ async def run( "language": language, "remove_data_from_chart_schema": remove_data_from_chart_schema, "custom_instruction": custom_instruction or "", + "documents": contexts or [], **self._components, **self._configs, }, 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 a49281d04c..c9f8b23537 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import logging import sys from typing import TYPE_CHECKING, Any @@ -26,6 +28,7 @@ from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost + if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory else: @@ -42,7 +45,10 @@ ### TARGET DATA SOURCE ### {{ data_source }} -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +interpreting intent and generating SQL. {% for document in documents %} {{ document }} {% endfor %} @@ -93,15 +99,13 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -### BUSINESS ANALYTICS TERM MAPPING ### -If the user asks about PCB repair trends, repair volume, repair counts, debug hours, -turnaround time, resolved entries, failure category, sales performance, salesperson -ranking, top customers, customer growth, revenue, margin, orders, or invoices, map -those business terms to the closest explicit table and column names in DATABASE SCHEMA -and VALID TABLE NAMES. -Never reuse table or column names from SQL SAMPLES or chat history unless those exact -names also appear in DATABASE SCHEMA or VALID TABLE NAMES for the active datasource. -Do not SUM or AVG string columns. +### INTENT AND SCHEMA GROUNDING ### +Interpret the user's business terms by matching them to explicit tables, columns, +metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Never reuse table +or column names from SQL SAMPLES or chat history unless those exact names also +appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. +Only apply aggregate functions to columns whose active metadata type supports that +operation. ### REASONING PLAN ### {{ sql_generation_reasoning }} diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index ebd3923a5d..abbbb81d56 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import logging import sys @@ -17,6 +19,7 @@ ) from src.utils import trace_cost from src.web.v1.services import Configuration + if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory else: @@ -26,7 +29,10 @@ sql_generation_reasoning_user_prompt_template = """ -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +planning SQL. {% for document in documents %} {{ document }} {% endfor %} diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 948cff0291..7ed59d5b63 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -42,6 +42,16 @@ """ sql_to_answer_user_prompt_template = """ +{% if documents %} +### Active Datasource Metadata ### +This is the complete deployed metadata for the active datasource. Use it to +understand the schema, tables, columns, metrics, views, and relationships behind +the SQL before answering. +{% for document in documents %} + {{ document }} +{% endfor %} +{% endif %} + ### Inputs ### User's question: {{ query }} SQL: {{ sql }} @@ -67,6 +77,7 @@ def prompt( current_time: str, custom_instruction: str, prompt_builder: PromptBuilder, + documents: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -75,6 +86,7 @@ def prompt( language=language, current_time=current_time, custom_instruction=custom_instruction, + documents=documents or [], ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -159,6 +171,7 @@ async def run( current_time: str = Configuration().show_current_time(), query_id: Optional[str] = None, custom_instruction: Optional[str] = None, + contexts: Optional[list[str]] = None, ) -> dict: logger.info("Sql_Answer Generation pipeline is running...") return await self._pipe.execute( @@ -171,6 +184,7 @@ async def run( "current_time": current_time, "query_id": query_id, "custom_instruction": custom_instruction or "", + "documents": contexts or [], **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 6c4743f761..daae37e02e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -64,7 +64,10 @@ def get_sql_correction_system_prompt( {{ data_source }} {% if documents %} -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +correcting SQL. {% for document in documents %} {{ document }} {% endfor %} @@ -106,12 +109,12 @@ def get_sql_correction_system_prompt( Error Message: {{ invalid_generation_result.error }} ### CORRECTION GROUNDING ### -Use DATABASE SCHEMA and VALID TABLE NAMES as the source of truth. If the invalid SQL -uses a table such as bookexamples.sales, sales, orders, or customers that is not listed -above, replace it with an explicitly listed table only when the listed schema supports -the user's request. For sales performance, salesperson ranking, customer growth, revenue, -margin, orders, or invoices, use only the active datasource's exposed sales/customer -tables and numeric columns. Do not SUM or AVG string columns. +Use ACTIVE DATASOURCE METADATA and VALID TABLE NAMES as the source of truth. If the +invalid SQL references a table or column not listed above, replace it only when the +active datasource metadata clearly contains an equivalent object that supports the +user's request. Do not invent tables, columns, joins, metrics, or relationships. +Only apply aggregate functions to columns whose active metadata type supports that +operation. Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 61451fae15..59bea2279c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -33,7 +33,10 @@ ### TARGET DATA SOURCE ### {{ data_source }} -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +interpreting intent and generating SQL. {% for document in documents %} {{ document }} {% endfor %} @@ -84,15 +87,12 @@ ### QUESTION ### User's Question: {{ query }} -### BUSINESS ANALYTICS TERM MAPPING ### -If the user asks about PCB repair trends, repair volume, repair counts, debug hours, -turnaround time, resolved entries, failure category, sales performance, salesperson -ranking, top customers, customer growth, revenue, margin, orders, or invoices, map -those business terms to the closest explicit table and column names in DATABASE SCHEMA -and VALID TABLE NAMES. -Do not answer with general guidance when a SQL aggregation, comparison, trend, or chart is requested. -Never reuse table or column names from SQL SAMPLES unless those exact names also appear -in DATABASE SCHEMA or VALID TABLE NAMES for the active datasource. +### INTENT AND SCHEMA GROUNDING ### +Interpret the user's business terms by matching them to explicit tables, columns, +metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Do not answer with +general guidance when the question can be answered with SQL over the active metadata. +Never reuse table or column names from SQL SAMPLES unless those exact names also +appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. {% if sql_generation_reasoning %} ### REASONING PLAN ### 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 00b731cb2c..f91a4288e4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -22,7 +22,10 @@ sql_generation_reasoning_user_prompt_template = """ -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +planning SQL. {% for document in documents %} {{ document }} {% endfor %} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4bec10530f..0c1f7e58fa 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1419,13 +1419,7 @@ def replace_clause(match: re.Match[str]) -> str: def _references_known_hallucination_prone_schema(sql: str) -> bool: - return bool( - re.search( - r"\b(?:dbo_repair_logs|dbo_ticket_cycles|dbo_DebugEntries|dbo_reports|dbo_knowledge_articles|dbo_kb_articles)\b", - sql, - flags=re.IGNORECASE, - ) - ) + return False def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: @@ -1435,12 +1429,6 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) - normalized = _rewrite_mssql_invented_repair_relationship_identifiers(normalized) - normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) - normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) - normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) - normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) - normalized = _rewrite_mssql_invented_failure_category(normalized) normalized = _rewrite_mssql_invented_report_fields(normalized) normalized = _rewrite_mssql_invented_ticket_metrics(normalized) normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) @@ -1519,20 +1507,8 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) - normalized = _rewrite_mssql_sales_schema_aliases(normalized) normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) normalized = _rewrite_mssql_invented_date_identifiers(normalized) - normalized = _rewrite_mssql_invented_repair_relationship_identifiers( - normalized - ) - normalized = _rewrite_mssql_repair_log_turnaround_trend_shape(normalized) - normalized = _rewrite_mssql_ticket_cycle_turnaround_shape(normalized) - normalized = _rewrite_mssql_repair_log_throughput_shape(normalized) - normalized = _rewrite_mssql_invented_pcb_throughput_identifiers(normalized) - normalized = _rewrite_mssql_invented_failure_category(normalized) - normalized = _rewrite_mssql_invented_report_fields(normalized) - normalized = _rewrite_mssql_invented_ticket_metrics(normalized) - normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -1846,12 +1822,11 @@ async def _classify_generation_result( - DON'T USE "TO_CHAR" function in the generated SQL query. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- Never invent foreign key columns or relationship fields such as "FailurePatternID", "FailurePatternId", "TicketID", or "
ID" unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. +- Never invent foreign key columns or relationship fields from table names unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. - Never invent time bucket columns such as "MONTH", "YEAR", "DAY", "month", "year", or "date" unless that exact column appears in the DATABASE SCHEMA. For monthly, yearly, or daily trends, apply a supported date/time bucket function from SQL FUNCTIONS to a real timestamp column from the selected table. - Every generated SQL query must be grounded only in the connected datasource metadata, deployed semantic model definitions, relationships, and DATABASE SCHEMA shown in the prompt. Do not use table names, column names, join paths, JSON keys, or business dimensions that are not explicitly present in that context. -- For synced repair-log schemas, if "dbo_repair_logs" contains "created_at" and the user asks for monthly repair volume or repair trends, count repair rows and bucket "dbo_repair_logs"."created_at". Do not select, group by, or order by "dbo_repair_logs"."MONTH" or bare "MONTH" unless the schema explicitly contains that column. -- For repair counts grouped by failure category, prefer the richest explicit category field exposed by the connected datasource. If the schema includes "dbo_DebugEntries", "dbo_DebugFixLogs", and "dbo_DebugFixes", group by "dbo_DebugFixes"."Description" after joining "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" and "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id". Otherwise use "dbo_repair_logs"."failure_code" only when that column appears in the schema. Do not invent "failure_category" unless it appears in the DATABASE SCHEMA. -- For repair SLA compliance dashboard/chart requests, do not invent "DAY", "MONTH", "turnaround_time", "sla_due_at", or due-date fields. If the schema only exposes "dbo_repair_logs"."status" and no explicit SLA/duration/deadline column, return a status distribution using "dbo_repair_logs"."status" and COUNT(*) so the UI can render a grounded chart. +- For trend questions, choose an explicit timestamp/date column from the active schema and bucket it with supported SQL FUNCTIONS. Do not select, group by, or order by invented time bucket columns unless they explicitly appear in the schema. +- For grouped count questions, choose an explicit dimension column from the active schema. Do not invent category/status/type columns unless they explicitly appear in the schema. - For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. - For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. - Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. @@ -1865,30 +1840,10 @@ async def _classify_generation_result( - DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. - DO NOT use JSON extraction functions or operators such as JSON_VALUE, JSON_QUERY, JSON_EXTRACT, JSON_EXTRACT_SCALAR, JSON_EXTRACT_ARRAY, json_value, json_extract, ->, or ->>. The MSSQL Wren/Ibis runtime does not support them. - If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. -- If a requested metric such as debug hours, risk score, repair cost, or turnaround time is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. -- Never invent JSON-derived columns such as "repair_date", "repair_status", or "failure_code" unless they are explicitly listed as columns in the DATABASE SCHEMA. -- For repair trend or repair volume questions, prefer explicit timestamp columns such as "created_at", "updated_at", "opened_at", or "closed_at" only when those exact columns appear in the selected table schema. -- For repair SLA compliance charts on "dbo_repair_logs", use "dbo_repair_logs"."status" as the compliance/status dimension when no explicit SLA, due-date, duration, or turnaround column appears in the DATABASE SCHEMA. Never use invented "DAY", "MONTH", or "turnaround_time" fields for SLA compliance. -- For repair counts grouped by failure category, use explicit exposed fields and schema-backed joins only. Prefer "dbo_DebugFixes"."Description" joined through "dbo_DebugFixLogs" when "dbo_DebugEntries"."DebugEntryId", "dbo_DebugFixLogs"."DebugEntryId", "dbo_DebugFixLogs"."FixId", and "dbo_DebugFixes"."Id" are present. Otherwise use "dbo_repair_logs"."failure_code" when present. Do not invent "dbo_repair_logs"."FailurePatternID"; only join to "dbo_failure_patterns" when an explicit join key or relationship exists in the DATABASE SCHEMA. -- For PCB/debug-entry failure charts, do not join "dbo_DebugEntries"."DebugEntryId" to "dbo_failure_patterns"."id"; those fields have incompatible types. If both "dbo_DebugEntries"."FailureSys" and "dbo_failure_patterns"."id" exist, join "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". -- For PCB synced database questions: - - Use "dbo_DebugEntries" for debug/PCB event records when the schema contains it. - - Use columns such as "Material", "WorkOrder", "SerialNumber", "FailedAt", "DateIn", "DateOut", "Hours", "Priority", "Actions", "Notes", and "FailureSys" only when they appear in the schema. - - Use "dbo_failure_patterns" for failure names, categories, severity, trend, occurrence counts, daily pattern summaries, and cost impact when those columns appear in the schema. - - For throughput trends across manufacturing/business units, use "dbo_DebugEntries"."BusinessUnit" as the unit dimension and a real debug-entry timestamp such as "dbo_DebugEntries"."DateIn" or "dbo_DebugEntries"."FailedAt" for the trend bucket. Do not use "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH", or invented manufacturing/date fields. - - For top/common PCB failure questions, first prefer grouping by "dbo_DebugFixes"."Description" and counting rows through the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join when those tables and join columns are in the schema. Otherwise prefer grouping by "dbo_failure_patterns"."name" or "dbo_failure_patterns"."category" and counting "dbo_DebugEntries"."DebugEntryId" after joining "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id". - - If a useful aggregate already exists in "dbo_failure_patterns" such as "occurrences", it can be used directly for top failure pattern questions without joining event rows. - - For requests such as "show top 10 most common PCB failures", "bar chart of failures by category", or "count of repairs grouped by failure category", generate SQL first. Do not answer with general charting guidance. Return the categorical failure field plus a count metric that can drive a bar chart. - - For failure-category charts, prefer one of these patterns depending on schema availability: - 1. `GROUP BY "dbo_DebugFixes"."Description"` and `COUNT(*)` using the explicit "dbo_DebugEntries" -> "dbo_DebugFixLogs" -> "dbo_DebugFixes" join - 2. `GROUP BY "dbo_failure_patterns"."category"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` - 3. `GROUP BY "dbo_failure_patterns"."name"` and `COUNT("dbo_DebugEntries"."DebugEntryId")` - 4. `GROUP BY "dbo_repair_logs"."failure_code"` and `COUNT(*)` - - For chart-oriented questions, ensure the final SELECT contains only the chart-ready dimension and metric columns. Avoid prose-like outputs or helper columns. -- For knowledge article tables: - - Use "created_at" for year/month trend buckets. Do not select, group by, or order by invented "YEAR" or "MONTH" columns. - - In "dbo_knowledge_articles", use "helpful" and "views" for effectiveness-style questions, and use "author" for creator/author groupings. Do not invent "effectiveness_score" or "created_by". - - In "dbo_kb_articles", use "created_by_user_id" for creator groupings. Do not invent "created_by" or "author" unless those exact columns appear in the schema. +- If a requested metric is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. +- Never invent JSON-derived columns unless they are explicitly listed as columns in the DATABASE SCHEMA. +- For trend or volume questions, use explicit timestamp/date columns only when those exact columns appear in the selected table schema. +- For grouped count questions, use explicit exposed dimension fields and only join tables when an explicit join key or relationship exists in the DATABASE SCHEMA. - DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. - Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. - Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. @@ -2237,12 +2192,12 @@ def get_sql_generation_system_prompt( 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. 4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. -6. YOU MUST ONLY use table names and column names that are explicitly present in the DATABASE SCHEMA or VALID TABLE NAMES sections. -7. SQL SAMPLES are examples of style only. NEVER reuse a sample table or column name unless that exact table or column also appears in the active DATABASE SCHEMA or VALID TABLE NAMES sections. -8. NEVER invent generic table names such as repair_logs, repair_log, sales_data, sales, orders, customers, users, tickets, events, or transactions unless that exact table name is present in the DATABASE SCHEMA or VALID TABLE NAMES sections. -9. If the user asks about a business concept such as repairs, PCB, cost, turnaround time, volume, sales performance, salesperson ranking, customer growth, revenue, margin, orders, or invoices, map it to the closest explicit table and column names from the provided schema. Do not create a new table name from the business concept. -10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the DATABASE SCHEMA. Do not aggregate text/string columns as numeric values. -11. Do not prefix table names with catalog or schema names unless the DATABASE SCHEMA or VALID TABLE NAMES section shows the table name with that exact prefix. +6. YOU MUST ONLY use table names and column names that are explicitly present in the ACTIVE DATASOURCE METADATA, DATABASE SCHEMA, or VALID TABLE NAMES sections. +7. SQL SAMPLES are examples of style only. NEVER reuse a sample table or column name unless that exact table or column also appears in the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES sections. +8. NEVER invent generic table names from the user's business terms unless that exact table name is present in the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES sections. +9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. +10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. +11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 59eca37846..107c5c3479 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -126,72 +126,7 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline def expand_business_terms_for_retrieval(query: str) -> str: - normalized = (query or "").lower() - analytics_terms = { - "business unit", - "pcb", - "repair", - "debug", - "turnaround", - "failure", - "failure code", - "failure category", - "failure pattern", - "resolved", - "trend", - "volume", - "count", - "counts", - "average", - "avg", - "chart", - "month", - "monthly", - "sales", - "sale", - "revenue", - "customer", - "customers", - "salesperson", - "sales person", - "sales rep", - "performance", - "ranking", - "rank", - "top", - "bottom", - "growth", - "fastest growing", - "order", - "orders", - "invoice", - "invoices", - "margin", - "profit", - "quantity", - "qty", - "amount", - "value", - "year", - "yearly", - } - if not any(term in normalized for term in analytics_terms): - return query - - return "\n".join( - [ - query, - "Business analytics aliases:", - "throughput trend volume count counts average total ranking top bottom grouped distribution", - "business unit manufacturing unit department location site plant team region category status", - "repair trends repair volume repair counts debug entries debug fixes failure code", - "monthly trend quarter grouped by month bar chart line chart", - "top common failures most common failure categories", - "failure patterns category occurrences material workorder serial number", - "sales revenue amount sales value sales performance salesperson ranking", - "customer sales top customers customer growth orders invoices margin quantity", - ] - ) + return query def _is_project_wide_analysis_query(query: str) -> bool: @@ -317,52 +252,6 @@ async def dbschema_retrieval( dbschema_retriever: Any, tables: Optional[list[str]] = None, ) -> list[Document]: - table_names = [] - if tables: - table_names.extend(tables) - else: - retrieved_tables = table_retrieval.get("documents", []) - for table in retrieved_tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) - - table_name_conditions = [ - {"field": "name", "operator": "==", "value": table_name} - for table_name in table_names - ] - - if table_name_conditions: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } - - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) - - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - documents = results["documents"] - if project_id and _is_project_wide_analysis_query(query): - all_project_results = await dbschema_retriever.run( - query_embedding=[], - filters={ - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": project_id}, - ], - }, - ) - documents = _dedupe_documents( - documents + all_project_results.get("documents", []) - ) - return documents - filters = { "operator": "AND", "conditions": [ @@ -375,7 +264,7 @@ async def dbschema_retrieval( ) logger.info( - "No table-description matches found; falling back to all deployed schema for project_id %s", + "Loading complete deployed schema metadata for active project_id %s", project_id, ) results = await dbschema_retriever.run(query_embedding=[], filters=filters) @@ -460,15 +349,6 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) - if _token_count > context_window_size or enable_column_pruning: - return { - "db_schemas": [], - "tokens": _token_count, - "has_calculated_field": has_calculated_field, - "has_metric": has_metric, - "has_json_field": has_json_field, - } - return { "db_schemas": retrieval_results, "tokens": _token_count, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 88c5341b03..172404cd23 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1031,6 +1031,14 @@ def _find_first_schema_column( return column return None + def _find_any_temporal_schema_column(self, table: dict[str, Any]) -> str | None: + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if column_name and self._is_temporal_schema_type(column_type): + return column_name + return None + def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' @@ -1149,18 +1157,7 @@ def _find_temporal_column_for_query( ): return column_name - return self._find_first_schema_column( - table, - ( - "created_at", - "createdat", - "created", - "date", - "time", - "timestamp", - "updated_at", - ), - ) + return self._find_any_temporal_schema_column(table) def _build_schema_grounded_table_question_sql( self, query: str, table_ddls: list[str] @@ -1356,6 +1353,8 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: return table_names def _build_direct_orders_sales_sql(self, query: str) -> str | None: + return None + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None @@ -2924,6 +2923,8 @@ def _build_monthly_repair_volume_sql( ) def _is_direct_heuristic_sql_query(self, query: str) -> bool: + return False + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return False @@ -2967,6 +2968,8 @@ def _build_heuristic_text_to_sql_fallback( table_ddls: list[str], table_names: Optional[list[str]] = None, ) -> str | None: + return None + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None @@ -3361,16 +3364,47 @@ def _build_schema_grounded_operational_sql( limit = int(limit_match.group(1)) if limit_match else 10 if wants_elapsed_time: + temporal_columns = [ + str(column.get("name") or "") + for column in table.get("columns", []) + if column.get("name") + and self._is_temporal_schema_type(str(column.get("type") or "")) + ] start_column = self._find_schema_column( table, - ("created_at", "created", "DateIn"), + ( + "created_at", + "created", + "DateIn", + "execution_date", + "opened_at", + "started_at", + "start_date", + "begin_date", + ), temporal=True, ) end_column = self._find_schema_column( table, - ("updated_at", "updated", "DateOut", "closed_at", "resolved_at"), + ( + "updated_at", + "updated", + "DateOut", + "closed_at", + "resolved_at", + "completed_at", + "finished_at", + "end_date", + ), temporal=True, ) + if not start_column and temporal_columns: + start_column = temporal_columns[0] + if not end_column: + for candidate in temporal_columns: + if candidate.lower() != str(start_column or "").lower(): + end_column = candidate + break if start_column and end_column: start_ref = f"{table_ref}.{self._quote_sql_identifier(start_column)}" end_ref = f"{table_ref}.{self._quote_sql_identifier(end_column)}" @@ -3508,6 +3542,33 @@ def _get_unqueryable_metric_message( "a first-class temporal field." ) + if any( + term in normalized_query + for term in ( + "monthly", + "trend", + "turnaround", + "time", + "duration", + "elapsed", + "latest", + "recent", + "newest", + "last records", + ) + ): + has_temporal_field = any( + self._is_temporal_schema_type(str(column.get("type") or "")) + for table in self._parse_schema_tables(table_ddls) + for column in table.get("columns", []) + ) + if not has_temporal_field: + return ( + "The active datasource does not expose a queryable date or " + "timestamp column. I cannot build a time-based analysis " + "without a first-class temporal field." + ) + repair_cost_terms = ( "repair cost", "repair_cost", @@ -3578,6 +3639,8 @@ def _get_unqueryable_metric_message( def _build_schema_grounded_sales_sql( self, query: str, table_ddls: list[str] ) -> str | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) normalized_schema = "\n".join( ddl for ddl in table_ddls or [] if isinstance(ddl, str) diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index 3392cad921..70c0ea13c7 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -80,6 +80,26 @@ def _is_stopped(self, query_id: str): return False + async def _load_active_schema_contexts( + self, query: str, project_id: Optional[str] + ) -> list[str]: + retrieval_pipeline = self._pipelines.get("db_schema_retrieval") + if not retrieval_pipeline: + return [] + + retrieval_result = await retrieval_pipeline.run( + query=query, + project_id=project_id, + ) + documents = retrieval_result.get("construct_retrieval_results", {}).get( + "retrieval_results", [] + ) + return [ + document["table_ddl"] + for document in documents + if isinstance(document, dict) and document.get("table_ddl") + ] + @observe(name="Generate Chart") @trace_metadata async def chart( @@ -154,6 +174,11 @@ async def chart( results["chart_result"] = deterministic_chart_result return results + schema_contexts = await self._load_active_schema_contexts( + chart_request.query, + chart_request.project_id, + ) + chart_generation_result = await self._pipelines["chart_generation"].run( query=chart_request.query, sql=chart_request.sql, @@ -161,6 +186,7 @@ async def chart( language=chart_request.configurations.language, remove_data_from_chart_schema=chart_request.remove_data_from_chart_schema, custom_instruction=chart_request.custom_instruction, + contexts=schema_contexts, ) chart_result = chart_generation_result["post_process"]["results"] diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index c907387300..596c3a89cb 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -53,6 +53,26 @@ def __init__( maxsize=maxsize, ttl=ttl ) + async def _load_active_schema_contexts( + self, query: str, project_id: Optional[str] + ) -> list[str]: + retrieval_pipeline = self._pipelines.get("db_schema_retrieval") + if not retrieval_pipeline: + return [] + + retrieval_result = await retrieval_pipeline.run( + query=query, + project_id=project_id, + ) + documents = retrieval_result.get("construct_retrieval_results", {}).get( + "retrieval_results", [] + ) + return [ + document["table_ddl"] + for document in documents + if isinstance(document, dict) and document.get("table_ddl") + ] + @observe(name="SQL Answer") @trace_metadata async def sql_answer( @@ -93,6 +113,11 @@ async def sql_answer( trace_id=trace_id, ) + schema_contexts = await self._load_active_schema_contexts( + sql_answer_request.query, + sql_answer_request.project_id, + ) + asyncio.create_task( self._pipelines["sql_answer"].run( query=sql_answer_request.query, @@ -102,6 +127,7 @@ async def sql_answer( current_time=sql_answer_request.configurations.show_current_time(), query_id=query_id, custom_instruction=sql_answer_request.custom_instruction, + contexts=schema_contexts, ) ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 80f8244f59..5b6694bb45 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1,5 +1,9 @@ +import pytest +from haystack import Document + from src.pipelines.retrieval.db_schema_retrieval import ( _is_project_wide_analysis_query, + dbschema_retrieval, expand_business_terms_for_retrieval, ) @@ -14,17 +18,72 @@ def test_project_wide_analysis_query_ignores_empty_query(): assert not _is_project_wide_analysis_query("") -def test_expand_business_terms_for_retrieval_includes_sales_aliases(): - expanded = expand_business_terms_for_retrieval( - "Create a SalesPerson performance ranking chart" - ) +def test_expand_business_terms_for_retrieval_does_not_add_datasource_specific_aliases(): + query = "Create a SalesPerson performance ranking chart" - assert "salesperson ranking" in expanded - assert "customer growth" in expanded - assert "Create a SalesPerson performance ranking chart" in expanded + assert expand_business_terms_for_retrieval(query) == query -def test_expand_business_terms_for_retrieval_leaves_non_analytics_query_unchanged(): +def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): query = "Explain what this workspace does" assert expand_business_terms_for_retrieval(query) == query + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_loads_complete_active_project_schema(): + class Retriever: + def __init__(self): + self.filters = None + + async def run(self, query_embedding, filters): + self.filters = filters + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "orders", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "orders"}, + ), + Document( + content=str( + { + "type": "TABLE", + "name": "customers", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "customers"}, + ), + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + query="total orders", + table_retrieval={ + "documents": [ + Document( + content=str({"name": "orders"}), + meta={"type": "TABLE_DESCRIPTION", "name": "orders"}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + ) + + assert [document.meta["name"] for document in documents] == ["orders", "customers"] + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } From 32d6eb385181145201dcd9fcdf5fdf9851c3af59 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 15:50:34 +0530 Subject: [PATCH 0409/1087] Fix order datasource schema-grounded SQL --- wren-ai-service/src/web/v1/services/ask.py | 210 +++++++++++++++------ 1 file changed, 154 insertions(+), 56 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 172404cd23..f28d8c5f48 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1618,6 +1618,10 @@ def _select_best_analytics_table( continue if "sales" in table_name: score += 5 + if "order" in table_name: + score += 4 + if "invoice" in table_name or "inv" in table_name: + score += 3 if "stage" in table_name: score -= 8 @@ -1675,6 +1679,16 @@ def _build_schema_grounded_analytics_sql( ): return categorical_count_sql + wants_count_metric = any( + term in normalized_query for term in ("count", "counts", "volume", "how many") + ) and not any( + term in normalized_query + for term in ("revenue", "sales value", "amount", "value", "quantity", "qty") + ) + wants_average_metric = any( + term in normalized_query for term in ("average", "avg", "mean") + ) + wants_monthly_count = ( "monthly" in normalized_query and any(term in normalized_query for term in ("count", "volume")) @@ -1735,7 +1749,9 @@ def _build_schema_grounded_analytics_sql( dimension_candidates: list[tuple[str, ...]] = [] if "salesperson" in normalized_query or "sales person" in normalized_query: - dimension_candidates.append(("SalesPerson", "Sales Rep", "SalesRep")) + dimension_candidates.append( + ("SalesPerson", "Salesman", "Sales Rep", "SalesRep", "Rep", "Owner") + ) if "business unit" in normalized_query or "bu" in normalized_query: dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) if "market" in normalized_query: @@ -1755,49 +1771,84 @@ def _build_schema_grounded_analytics_sql( dimension_candidates.append(("ProdType", "ProductType", "Product Type")) elif "product" in normalized_query: dimension_candidates.append( - ("ProdName", "Product", "ProductName", "Item", "ProdCode") + ( + "ProdName", + "Product", + "ProductName", + "ProductDescription", + "Item", + "ItemName", + "ProdCode", + "ProductCode", + "PartNo", + "SKU", + ) ) if ( "customer" in normalized_query or "custname" in compact_query or "custno" in compact_query ): - dimension_candidates.append(("Customer", "CustName", "CustNo")) - - if not dimension_candidates: - return None + dimension_candidates.append( + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + "Client", + "ClientName", + ) + ) measure_candidates = ( + "Qty", + "Quantity", + "SalesQty", + "OrderQty", + "InvoiceQty", + ) if any(term in normalized_query for term in ("quantity", "qty")) else ( + "SalesValue", + "FXSalesValue", + "Revenue", + "NetSales", + "SalesAmount", "NewOrderValue", "NewOrdersValue", "InvoiceValue", "InvoiceAmount", + "InvoiceAmt", "OrderValue", - "SalesValue", - "FXSalesValue", - "Revenue", "TotalRevenue", "Amount", "Value", "TotalOrderValue", "Cost", - "Qty", - "Quantity", ) if "invoice" in normalized_query: measure_candidates = ( "InvoiceValue", "InvoiceAmount", + "InvoiceAmt", + "InvValue", + "InvAmount", "SalesValue", "FXSalesValue", "Value", "Amount", ) + if wants_count_metric: + measure_candidates = () wants_trend = ( "trend" in normalized_query or "line chart" in normalized_query or "over time" in normalized_query or "last 12 months" in normalized_query + or "by month" in normalized_query + or "monthly" in normalized_query ) wants_date_distribution = ( any( @@ -1817,6 +1868,13 @@ def _build_schema_grounded_analytics_sql( ) ) wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) + wants_top = wants_top or any( + term in normalized_query + for term in ("top ", "highest", "largest", "most ", "best ") + ) + wants_time_bucket = wants_trend or bool( + re.search(r"\bby\s+(?:month|year|quarter|date)\b", normalized_query) + ) wants_detail_rows = ( wants_top and ("new order" in normalized_query or "orders" in normalized_query) @@ -1840,12 +1898,54 @@ def _build_schema_grounded_analytics_sql( or bool(re.search(r"\b20\d{2}\b", normalized_query)) ) + if not dimension_candidates and wants_time_bucket: + selected = self._select_best_analytics_table( + tables, + [], + measure_candidates, + wants_date=True, + allow_count_metric=wants_count_metric, + ) + if not selected: + return None + + table, _dimensions, measure, date_column = selected + table_name = table.get("name") + if not (table_name and date_column): + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + if wants_count_metric or not measure: + metric_expr = "COUNT(*)" + metric_alias = "RecordCount" + elif wants_average_metric: + metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Average{measure}" + else: + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Total{measure}" + + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " + f"FROM {table_ref}" + f"{self._build_date_filter(table_name, date_column, query)} " + f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" + ) + + if not dimension_candidates: + return None + selected = self._select_best_analytics_table( tables, dimension_candidates, measure_candidates, wants_date=wants_date, - allow_count_metric=wants_order_count_metric, + allow_count_metric=wants_order_count_metric or wants_count_metric, ) if not selected: return None @@ -2014,7 +2114,7 @@ def _build_schema_grounded_analytics_sql( if wants_trend and date_column: date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - if wants_order_count_metric: + if wants_order_count_metric or wants_count_metric: order_column = self._find_schema_column( table, ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), @@ -2025,6 +2125,9 @@ def _build_schema_grounded_analytics_sql( else "COUNT(*)" ) metric_alias = "OrderCount" + elif wants_average_metric: + metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Average{measure}" else: metric_expr = ( f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" @@ -2054,12 +2157,23 @@ def _build_schema_grounded_analytics_sql( f"DATEPART(MONTH, {date_ref})" ) - metric_expr = ( - f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - if measure - else "COUNT(*)" - ) - metric_alias = f"Total{measure}" if measure else "OrderCount" + if wants_order_count_metric or wants_count_metric or not measure: + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), + ) + metric_expr = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" + if order_column and (wants_order_count_metric or wants_count_metric) + else "COUNT(*)" + ) + metric_alias = "OrderCount" if order_column else "RecordCount" + elif wants_average_metric: + metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Average{measure}" + else: + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Total{measure}" limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) limit = int(limit_match.group(1)) if limit_match else 10 top_clause = f"TOP {limit} " if wants_top else "" @@ -3639,8 +3753,6 @@ def _get_unqueryable_metric_message( def _build_schema_grounded_sales_sql( self, query: str, table_ddls: list[str] ) -> str | None: - return None - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) normalized_schema = "\n".join( ddl for ddl in table_ddls or [] if isinstance(ddl, str) @@ -3648,44 +3760,30 @@ def _build_schema_grounded_sales_sql( if not normalized_query or not normalized_schema: return None - asks_for_salesperson_performance = ( - any( - term in normalized_query - for term in ( - "salesperson performance", - "sales person performance", - "sales rep performance", - "salesperson ranking", - "sales person ranking", - "sales rep ranking", - ) - ) - or ( - "salesperson" in normalized_query - and any(term in normalized_query for term in ("performance", "ranking")) + if not any( + term in normalized_query + for term in ( + "amount", + "average order value", + "customer", + "invoice", + "order", + "orders", + "product", + "quantity", + "qty", + "revenue", + "sale", + "sales", + "salesperson", + "sales person", + "trend", + "value", ) - ) - if not asks_for_salesperson_performance: - return self._build_schema_grounded_analytics_sql(query, table_ddls) - - required_schema_terms = ( - "create table dbo_tblsales", - "salesperson", - "salesvalue", - ) - if not all(term in normalized_schema for term in required_schema_terms): + ): return None - limit = 20 if re.search(r"\btop\s+20\b", normalized_query) else 10 - return ( - f'SELECT TOP {limit} ' - f'"dbo_tblSales"."SalesPerson" AS "SalesPerson", ' - f'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - f'FROM "dbo_tblSales" ' - f'WHERE "dbo_tblSales"."SalesPerson" IS NOT NULL ' - f'GROUP BY "dbo_tblSales"."SalesPerson" ' - f'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' - ) + return self._build_schema_grounded_analytics_sql(query, table_ddls) async def _run_with_timeout( self, From 4462bcfdd0ce866a0458a7b4fc6f01fc844d7619 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 16:44:46 +0530 Subject: [PATCH 0410/1087] Fix PCB schema-grounded SQL fallback --- wren-ai-service/src/web/v1/services/ask.py | 90 ++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f28d8c5f48..f6c3fbe35e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -3300,6 +3300,7 @@ def _build_schema_grounded_operational_sql( "ticket", "repair", "failure", + "pcb", "component", "board", "throughput", @@ -3345,6 +3346,40 @@ def _build_schema_grounded_operational_sql( score += 5 if "failure" in normalized_query and "failure" in normalized_table: score += 8 + if any( + term in normalized_query + for term in ("business unit", "business units", "unit", "units") + ) and self._find_schema_column( + table, + ( + "BusinessUnit", + "Business_Unit", + "Business Unit", + "manufacturing_unit", + "ManufacturingUnit", + "unit", + "BU", + "Division", + ), + ): + score += 15 + if any( + term in normalized_query + for term in ("product line", "product family", "product", "products") + ) and self._find_schema_column( + table, + ( + "Product_Family", + "ProductFamily", + "Product Family", + "ProductLine", + "Product_Line", + "Product", + "ProdType", + "Material", + ), + ): + score += 15 if any(term in normalized_query for term in ("error", "failure")) and any( self._find_schema_column(table, candidates) for candidates in ( @@ -3403,15 +3438,37 @@ def _build_schema_grounded_operational_sql( if "manufacturing" in normalized_query or "unit" in normalized_query: dimension_candidates.append( ( + "BusinessUnit", + "Business_Unit", + "Business Unit", "manufacturing_unit", "manufacturing unit", + "ManufacturingUnit", "unit", + "BU", + "Division", "assignee_user_id", "created_by_user_id", "org_id", "status", ) ) + if any( + term in normalized_query + for term in ("product line", "product family", "product", "products") + ): + dimension_candidates.append( + ( + "Product_Family", + "ProductFamily", + "Product Family", + "ProductLine", + "Product_Line", + "Product", + "ProdType", + "Material", + ) + ) if "component" in normalized_query: dimension_candidates.append( ("component", "component_type", "board_type", "title", "status") @@ -5396,6 +5453,39 @@ async def ask( "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." ) + if not api_results and any( + term in user_query.lower() + for term in ( + "pcb", + "repair", + "failure", + "business unit", + "business units", + "product line", + "product family", + ) + ): + operational_sql = self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + if operational_sql: + logger.info( + "Using schema-grounded operational SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + operational_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = operational_sql + error_message = ( + "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." + ) + if not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls From 7cc09293842ec8494af46ea568c47c96b830482f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 17:14:43 +0530 Subject: [PATCH 0411/1087] Prefer repair schema SQL for PCB questions --- wren-ai-service/src/web/v1/services/ask.py | 35 ++++++++++++++-------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f6c3fbe35e..fde452cfdd 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1653,6 +1653,16 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + if repair_failure_count_sql := self._build_repair_failure_count_sql( + query, table_ddls + ): + return repair_failure_count_sql + + if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( + query, table_ddls + ): + return monthly_repair_volume_sql + if operational_sql := self._build_schema_grounded_operational_sql( query, tables ): @@ -1669,11 +1679,6 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( - query, table_ddls - ): - return monthly_repair_volume_sql - if categorical_count_sql := self._build_generic_categorical_count_sql( query, tables ): @@ -3520,14 +3525,18 @@ def _build_schema_grounded_operational_sql( term in normalized_query for term in ("trend", "monthly", "month", "line chart", "over time") ) - wants_elapsed_time = any( - term in normalized_query - for term in ( - "time", - "duration", - "elapsed", - "turnaround", - "estimated", + wants_elapsed_time = ( + not wants_trend + and any( + term in normalized_query + for term in ( + "duration", + "elapsed", + "turnaround", + "estimated", + "time spent", + "time taken", + ) ) ) wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) From 7baf2916063199fda6e80d35d77bb89e0f6e04a8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 17:55:41 +0530 Subject: [PATCH 0412/1087] Allow PCB repair schema intent validation --- wren-ai-service/src/web/v1/services/ask.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index fde452cfdd..9ffdaea7eb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -621,7 +621,17 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: - concept_groups.append({"pcb", "board"}) + concept_groups.append( + { + "pcb", + "board", + "repair", + "repairs", + "debug", + "failure", + "failures", + } + ) if "critical" in normalized: concept_groups.append({"critical", "severity", "priority"}) if "cost" in normalized: From 69bd4060601ce2176a2c1507a85d2df89523042f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 18:28:44 +0530 Subject: [PATCH 0413/1087] Revert "Allow PCB repair schema intent validation" This reverts commit 7baf2916063199fda6e80d35d77bb89e0f6e04a8. --- wren-ai-service/src/web/v1/services/ask.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9ffdaea7eb..fde452cfdd 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -621,17 +621,7 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: - concept_groups.append( - { - "pcb", - "board", - "repair", - "repairs", - "debug", - "failure", - "failures", - } - ) + concept_groups.append({"pcb", "board"}) if "critical" in normalized: concept_groups.append({"critical", "severity", "priority"}) if "cost" in normalized: From 0a1dd76b166566cde44f0ef60c9de8e4b618951f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 20:07:54 +0530 Subject: [PATCH 0414/1087] Reapply "Allow PCB repair schema intent validation" This reverts commit 69bd4060601ce2176a2c1507a85d2df89523042f. --- wren-ai-service/src/web/v1/services/ask.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index fde452cfdd..9ffdaea7eb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -621,7 +621,17 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: - concept_groups.append({"pcb", "board"}) + concept_groups.append( + { + "pcb", + "board", + "repair", + "repairs", + "debug", + "failure", + "failures", + } + ) if "critical" in normalized: concept_groups.append({"critical", "severity", "priority"}) if "cost" in normalized: From 401f323e723bbb3aa9707cacb5ad7b109a2f1dc5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 7 Jul 2026 23:50:12 +0530 Subject: [PATCH 0415/1087] Support unique customer ranking by market --- wren-ai-service/src/web/v1/services/ask.py | 75 ++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9ffdaea7eb..e56a26859b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1912,6 +1912,81 @@ def _build_schema_grounded_analytics_sql( or "this year" in normalized_query or bool(re.search(r"\b20\d{2}\b", normalized_query)) ) + wants_unique_customers_by_group = ( + any( + term in normalized_query + for term in ("unique customer", "unique customers") + ) + and "customer" in normalized_query + and "division" in normalized_query + and "market" in normalized_query + and any(term in normalized_query for term in ("highest", "top", "most")) + and any(term in normalized_query for term in ("each", "per ")) + ) + if wants_unique_customers_by_group: + selected = self._select_best_analytics_table( + tables, + [ + ("Market", "MarketType", "MarketName", "Region", "Country"), + ("Division",), + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + "Client", + "ClientName", + ), + ], + (), + wants_date=False, + allow_count_metric=True, + ) + if selected: + table, dimensions, _measure, _date_column = selected + table_name = table.get("name") + if table_name and len(dimensions) >= 3: + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + market, division, customer = dimensions[:3] + market_ref = f"{table_ref}.{self._quote_sql_identifier(market)}" + division_ref = ( + f"{table_ref}.{self._quote_sql_identifier(division)}" + ) + customer_ref = ( + f"{table_ref}.{self._quote_sql_identifier(customer)}" + ) + where_clause = self._append_not_null_filters( + "", + [market_ref, division_ref, customer_ref], + ) + return ( + "WITH grouped_results AS (" + f"SELECT {market_ref} AS {self._quote_sql_identifier(market)}, " + f"{division_ref} AS {self._quote_sql_identifier(division)}, " + f"COUNT(DISTINCT {customer_ref}) AS \"UniqueCustomerCount\" " + f"FROM {table_ref}" + f"{where_clause} " + f"GROUP BY {market_ref}, {division_ref}" + "), ranked_results AS (" + f"SELECT {self._quote_sql_identifier(market)}, " + f"{self._quote_sql_identifier(division)}, " + "\"UniqueCustomerCount\", " + f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(market)} " + "ORDER BY \"UniqueCustomerCount\" DESC) AS \"rank\" " + "FROM grouped_results" + ") " + f"SELECT {self._quote_sql_identifier(market)}, " + f"{self._quote_sql_identifier(division)}, " + "\"UniqueCustomerCount\" " + "FROM ranked_results " + "WHERE \"rank\" = 1 " + "ORDER BY \"UniqueCustomerCount\" DESC" + ) if not dimension_candidates and wants_time_bucket: selected = self._select_best_analytics_table( From fe41af45558b11d8391fdfbdca3d96511a4fc876 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 00:38:04 +0530 Subject: [PATCH 0416/1087] Use active metadata for generic SQL safety --- .../generation/question_recommendation.py | 11 +- wren-ai-service/src/web/v1/services/ask.py | 121 +++++++++++++----- .../v1/services/question_recommendation.py | 6 +- 3 files changed, 97 insertions(+), 41 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index 9b2556f69a..5a5b3e3e2d 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -160,7 +160,10 @@ {% endif %} {% if documents %} -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +Use only this latest deployed metadata from the active datasource when generating +recommended questions. Do not reuse tables, columns, or business terms from prior +questions unless they are answerable from this metadata. {% for document in documents %} {{ document }} {% endfor %} @@ -180,12 +183,6 @@ def prompt( max_categories: int, prompt_builder: PromptBuilder, ) -> dict: - """ - If previous_questions is provided, the MDL is omitted to allow the LLM to focus on - generating recommendations based on the question history. This helps provide more - contextually relevant questions that build on previous questions. - """ - _prompt = prompt_builder.run( documents=documents, previous_questions=previous_questions, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e56a26859b..098bd44f4f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -19,6 +19,10 @@ logger = logging.getLogger("wren-ai-service") +NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( + "No relevant data found in the active datasource for this question." +) + async def _return_value(value): return value @@ -4689,6 +4693,28 @@ def _build_failed_text_to_sql_response( is_followup=is_followup, ) + def _build_no_relevant_active_datasource_response( + self, + trace_id: Optional[str], + *, + rephrased_question: Optional[str] = None, + intent_reasoning: Optional[str] = None, + retrieved_tables: Optional[list[str]] = None, + sql_generation_reasoning: Optional[str] = None, + is_followup: bool = False, + ) -> AskResultResponse: + return self._build_failed_text_to_sql_response( + trace_id, + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=retrieved_tables, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=None, + is_followup=is_followup, + code="NO_RELEVANT_DATA", + ) + @observe(name="Ask Question") @trace_metadata async def ask( @@ -5547,6 +5573,32 @@ async def ask( "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." ) + if ( + not api_results + and self._is_data_analysis_query(user_query) + and ( + schema_grounded_sql := self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + ) + ): + logger.info( + "Using generic schema-grounded analytics SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + schema_grounded_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = schema_grounded_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." + ) + if not api_results and any( term in user_query.lower() for term in ( @@ -5647,18 +5699,18 @@ async def ask( error_message = "Heuristic SQL fallback was not valid for the active datasource schema." if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( - self._build_failed_text_to_sql_response( + self._build_no_relevant_active_datasource_response( trace_id, - error_message, rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, retrieved_tables=table_names, - invalid_sql=invalid_sql, is_followup=True if histories else False, ) ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = error_message + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) results["metadata"]["type"] = "TEXT_TO_SQL" return results api_results = [ask_result] @@ -5679,19 +5731,19 @@ async def ask( logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_DATA", - message="No relevant data", - ), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + is_followup=True if histories else False, + ) ) results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) results["metadata"]["type"] = "TEXT_TO_SQL" return results @@ -6065,23 +6117,28 @@ async def ask( logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_SQL", - message=error_message or "No relevant SQL", - ), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - trace_id=trace_id, - is_followup=True if histories else False, + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + is_followup=True if histories else False, + ) ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = error_message + if error_message or invalid_sql: + logger.info( + "Suppressed technical SQL failure for query_id %s. " + "error=%s invalid_sql=%s", + query_id, + error_message, + invalid_sql, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) results["metadata"]["type"] = "TEXT_TO_SQL" return results diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 4432c65ce6..d93e456d2a 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -396,10 +396,12 @@ async def recommend(self, input: Request, **kwargs) -> Event: trace_id = kwargs.get("trace_id") try: - mdl = orjson.loads(input.mdl) + orjson.loads(input.mdl) retrieval_result = await self._pipelines["db_schema_retrieval"].run( - tables=[model["name"] for model in mdl["models"]], + query="", + histories=[], project_id=input.project_id, + enable_column_pruning=False, ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) From 151351fe3f03929fc0b596cfe614fa75160bd6ee Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 00:59:53 +0530 Subject: [PATCH 0417/1087] Validate answer and chart SQL against active metadata --- wren-ai-service/src/web/v1/services/chart.py | 80 ++++++++++++++++--- .../src/web/v1/services/sql_answer.py | 66 +++++++++++++-- 2 files changed, 126 insertions(+), 20 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index 70c0ea13c7..efb3656512 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,12 +6,24 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + find_invalid_column_references, + find_invalid_table_references, + normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, +) from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest logger = logging.getLogger("wren-ai-service") +NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( + "No relevant data found in the active datasource for this question." +) + # POST /v1/charts class ChartRequest(BaseRequest): @@ -81,15 +93,17 @@ def _is_stopped(self, query_id: str): return False async def _load_active_schema_contexts( - self, query: str, project_id: Optional[str] + self, project_id: Optional[str] ) -> list[str]: retrieval_pipeline = self._pipelines.get("db_schema_retrieval") if not retrieval_pipeline: return [] retrieval_result = await retrieval_pipeline.run( - query=query, + query="", + histories=[], project_id=project_id, + enable_column_pruning=False, ) documents = retrieval_result.get("construct_retrieval_results", {}).get( "retrieval_results", [] @@ -100,6 +114,25 @@ async def _load_active_schema_contexts( if isinstance(document, dict) and document.get("table_ddl") ] + def _normalize_and_validate_sql( + self, sql: str, schema_contexts: list[str] + ) -> str | None: + valid_table_names = construct_valid_table_names(schema_contexts) + valid_table_columns = construct_valid_table_columns(schema_contexts) + normalized_sql = normalize_sql_table_references_to_schema( + sql, + valid_table_names, + ) + normalized_sql = normalize_sql_column_references_to_schema( + normalized_sql, + valid_table_columns, + ) + if find_invalid_table_references(normalized_sql, valid_table_names): + return None + if find_invalid_column_references(normalized_sql, valid_table_columns): + return None + return normalized_sql + @observe(name="Generate Chart") @trace_metadata async def chart( @@ -120,6 +153,27 @@ async def chart( try: query_id = chart_request.query_id execute_sql_error_message = None + schema_contexts = await self._load_active_schema_contexts( + chart_request.project_id, + ) + normalized_sql = self._normalize_and_validate_sql( + chart_request.sql, + schema_contexts, + ) + if not normalized_sql: + self._chart_results[query_id] = ChartResultResponse( + status="failed", + error=ChartError( + code="OTHERS", + message=NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, + ), + trace_id=trace_id, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) + return results if not chart_request.data: self._chart_results[query_id] = ChartResultResponse( @@ -129,7 +183,7 @@ async def chart( execute_sql_result = ( await self._pipelines["sql_executor"].run( - sql=chart_request.sql, + sql=normalized_sql, project_id=chart_request.project_id, ) )["execute_sql"] @@ -147,12 +201,19 @@ async def chart( status="failed", error=ChartError( code="OTHERS", - message=execute_sql_error_message, + message=NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, ), trace_id=trace_id, ) - results["metadata"]["error_type"] = "OTHERS" - results["metadata"]["error_message"] = execute_sql_error_message + logger.info( + "Suppressed chart SQL execution failure for query_id %s: %s", + query_id, + execute_sql_error_message, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) return results self._chart_results[query_id] = ChartResultResponse( @@ -174,14 +235,9 @@ async def chart( results["chart_result"] = deterministic_chart_result return results - schema_contexts = await self._load_active_schema_contexts( - chart_request.query, - chart_request.project_id, - ) - chart_generation_result = await self._pipelines["chart_generation"].run( query=chart_request.query, - sql=chart_request.sql, + sql=normalized_sql, data=sql_data, language=chart_request.configurations.language, remove_data_from_chart_schema=chart_request.remove_data_from_chart_schema, diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index 596c3a89cb..7442f2f884 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -7,11 +7,23 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + find_invalid_column_references, + find_invalid_table_references, + normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, +) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent logger = logging.getLogger("wren-ai-service") +NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( + "No relevant data found in the active datasource for this question." +) + # POST /v1/sql-answers class SqlAnswerRequest(BaseRequest): @@ -54,15 +66,17 @@ def __init__( ) async def _load_active_schema_contexts( - self, query: str, project_id: Optional[str] + self, project_id: Optional[str] ) -> list[str]: retrieval_pipeline = self._pipelines.get("db_schema_retrieval") if not retrieval_pipeline: return [] retrieval_result = await retrieval_pipeline.run( - query=query, + query="", + histories=[], project_id=project_id, + enable_column_pruning=False, ) documents = retrieval_result.get("construct_retrieval_results", {}).get( "retrieval_results", [] @@ -73,6 +87,25 @@ async def _load_active_schema_contexts( if isinstance(document, dict) and document.get("table_ddl") ] + def _normalize_and_validate_sql( + self, sql: str, schema_contexts: list[str] + ) -> str | None: + valid_table_names = construct_valid_table_names(schema_contexts) + valid_table_columns = construct_valid_table_columns(schema_contexts) + normalized_sql = normalize_sql_table_references_to_schema( + sql, + valid_table_names, + ) + normalized_sql = normalize_sql_column_references_to_schema( + normalized_sql, + valid_table_columns, + ) + if find_invalid_table_references(normalized_sql, valid_table_names): + return None + if find_invalid_column_references(normalized_sql, valid_table_columns): + return None + return normalized_sql + @observe(name="SQL Answer") @trace_metadata async def sql_answer( @@ -99,6 +132,28 @@ async def sql_answer( trace_id=trace_id, ) + schema_contexts = await self._load_active_schema_contexts( + sql_answer_request.project_id, + ) + normalized_sql = self._normalize_and_validate_sql( + sql_answer_request.sql, + schema_contexts, + ) + if not normalized_sql: + self._sql_answer_results[query_id] = SqlAnswerResultResponse( + status="failed", + error=SqlAnswerResultResponse.SqlAnswerError( + code="OTHERS", + message=NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, + ), + trace_id=trace_id, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) + return results + preprocessed_sql_data = self._pipelines["preprocess_sql_data"].run( sql_data=sql_answer_request.sql_data, )["preprocess"] @@ -113,15 +168,10 @@ async def sql_answer( trace_id=trace_id, ) - schema_contexts = await self._load_active_schema_contexts( - sql_answer_request.query, - sql_answer_request.project_id, - ) - asyncio.create_task( self._pipelines["sql_answer"].run( query=sql_answer_request.query, - sql=sql_answer_request.sql, + sql=normalized_sql, sql_data=preprocessed_sql_data.get("sql_data", {}), language=sql_answer_request.configurations.language, current_time=sql_answer_request.configurations.show_current_time(), From e5278fcb1fcc1f94abaaeee3ea17685ed75f9ceb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 14:07:44 +0530 Subject: [PATCH 0418/1087] Fix recommended question SQL context --- .../apollo/server/services/askingService.ts | 55 +++++++-- .../services/tests/askingService.test.ts | 114 ++++++++++++++++++ 2 files changed, 158 insertions(+), 11 deletions(-) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 36b564c3a1..f5b1115cae 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -799,7 +799,14 @@ export class AskingService implements IAskingService { */ public async createThread(input: AskingDetailTaskInput): Promise { // 1. create a thread and the first thread response - const { id } = await this.projectService.getCurrentProject(); + const project = await this.projectService.getCurrentProject(); + const { id } = project; + const trackedAskingResult = + await this.createTrackedAskingResultFromShortcutSql(input, { + projectId: id, + language: WrenAILanguage[project.language] || WrenAILanguage.EN, + }); + const thread = await this.threadRepository.createOne({ projectId: id, summary: input.question, @@ -808,16 +815,16 @@ export class AskingService implements IAskingService { const threadResponse = await this.threadResponseRepository.createOne({ threadId: thread.id, question: input.question, - sql: input.sql, - askingTaskId: input.trackedAskingResult?.taskId, answerDetail: input.answerDetail, + sql: trackedAskingResult ? undefined : input.sql, + askingTaskId: trackedAskingResult?.taskId, }); // if queryId is provided, update asking task - if (input.trackedAskingResult?.taskId) { + if (trackedAskingResult?.taskId) { await this.askingTaskTracker.bindThreadResponse( - input.trackedAskingResult.taskId, - input.trackedAskingResult.queryId, + trackedAskingResult.taskId, + trackedAskingResult.queryId, thread.id, threadResponse.id, ); @@ -858,19 +865,26 @@ export class AskingService implements IAskingService { ): Promise { const thread = await this.ensureThreadInCurrentProject(threadId); + const project = await this.projectService.getProjectById(thread.projectId); + const trackedAskingResult = + await this.createTrackedAskingResultFromShortcutSql(input, { + projectId: thread.projectId, + language: WrenAILanguage[project.language] || WrenAILanguage.EN, + }); + const threadResponse = await this.threadResponseRepository.createOne({ threadId: thread.id, question: input.question, - sql: input.sql, - askingTaskId: input.trackedAskingResult?.taskId, answerDetail: input.answerDetail, + sql: trackedAskingResult ? undefined : input.sql, + askingTaskId: trackedAskingResult?.taskId, }); // if queryId is provided, update asking task - if (input.trackedAskingResult?.taskId) { + if (trackedAskingResult?.taskId) { await this.askingTaskTracker.bindThreadResponse( - input.trackedAskingResult.taskId, - input.trackedAskingResult.queryId, + trackedAskingResult.taskId, + trackedAskingResult.queryId, thread.id, threadResponse.id, ); @@ -1364,6 +1378,25 @@ export class AskingService implements IAskingService { return this.projectService.getProjectById(thread.projectId); } + private async createTrackedAskingResultFromShortcutSql( + input: AskingDetailTaskInput, + payload: AskingPayload, + ): Promise { + if (input.trackedAskingResult || !input.sql || !input.question) { + return input.trackedAskingResult; + } + + const task = await this.createAskingTask( + { question: input.question }, + payload, + ); + const trackedAskingResult = await this.getAskingTask(task.id); + if (!trackedAskingResult?.taskId) { + throw new Error(`Asking task ${task.id} not found`); + } + return trackedAskingResult; + } + public async adjustThreadResponseWithSQL( threadResponseId: number, input: AdjustmentSqlInput, diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index cf1706db51..d33c9e166d 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -117,4 +117,118 @@ describe('AskingService', () => { expect(service.projectService.getProjectById).not.toHaveBeenCalled(); }); }); + + describe('recommendation question shortcut SQL', () => { + const trackedAskingResult = { + taskId: 42, + queryId: 'ask-query-id', + question: 'Show monthly record count', + status: 'UNDERSTANDING', + response: null, + error: null, + }; + + const createService = () => { + const service = Object.create(AskingService.prototype) as any; + service.projectService = { + getCurrentProject: jest.fn().mockResolvedValue({ + id: 1, + language: 'EN', + }), + getProjectById: jest.fn().mockResolvedValue({ + id: 1, + language: 'EN', + }), + }; + service.deployService = { + getLastDeployment: jest.fn().mockResolvedValue({ + hash: 'latest-deploy-hash', + }), + }; + service.threadRepository = { + createOne: jest.fn().mockResolvedValue({ id: 7, projectId: 1 }), + findOneBy: jest.fn().mockResolvedValue({ id: 7, projectId: 1 }), + }; + service.threadResponseRepository = { + createOne: jest.fn().mockResolvedValue({ id: 11, threadId: 7 }), + getResponsesWithThread: jest.fn().mockResolvedValue([]), + }; + service.askingTaskTracker = { + createAskingTask: jest.fn().mockResolvedValue({ + queryId: trackedAskingResult.queryId, + }), + getAskingResult: jest.fn().mockResolvedValue(trackedAskingResult), + bindThreadResponse: jest.fn().mockResolvedValue(undefined), + }; + return service; + }; + + test('creates a normal asking task for a new thread instead of storing shortcut SQL', async () => { + const service = createService(); + + await service.createThread({ + question: trackedAskingResult.question, + sql: 'SELECT stale_recommendation_sql', + }); + + expect(service.deployService.getLastDeployment).toHaveBeenCalledWith(1); + expect(service.askingTaskTracker.createAskingTask).toHaveBeenCalledWith({ + query: trackedAskingResult.question, + histories: null, + deployId: 'latest-deploy-hash', + projectId: '1', + configurations: { language: 'EN' }, + rerunFromCancelled: undefined, + previousTaskId: undefined, + threadResponseId: undefined, + }); + expect(service.threadResponseRepository.createOne).toHaveBeenCalledWith({ + threadId: 7, + question: trackedAskingResult.question, + sql: undefined, + askingTaskId: trackedAskingResult.taskId, + }); + expect(service.askingTaskTracker.bindThreadResponse).toHaveBeenCalledWith( + trackedAskingResult.taskId, + trackedAskingResult.queryId, + 7, + 11, + ); + }); + + test('creates a standalone asking task for current-thread recommendations', async () => { + const service = createService(); + service.threadResponseRepository.getResponsesWithThread.mockResolvedValue([ + { id: 1, question: 'Previous question', sql: 'SELECT 1' }, + ]); + + await service.createThreadResponse( + { + question: trackedAskingResult.question, + sql: 'SELECT stale_recommendation_sql', + }, + 7, + ); + + expect(service.projectService.getProjectById).toHaveBeenCalledWith(1); + expect(service.askingTaskTracker.createAskingTask).toHaveBeenCalledWith( + expect.objectContaining({ + query: trackedAskingResult.question, + histories: null, + deployId: 'latest-deploy-hash', + projectId: '1', + configurations: { language: 'EN' }, + }), + ); + expect( + service.threadResponseRepository.getResponsesWithThread, + ).not.toHaveBeenCalled(); + expect(service.threadResponseRepository.createOne).toHaveBeenCalledWith({ + threadId: 7, + question: trackedAskingResult.question, + sql: undefined, + askingTaskId: trackedAskingResult.taskId, + }); + }); + }); }); From ffbc00e4d0ed23dabc80bc17f1d28d94e8397bce Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 14:40:33 +0530 Subject: [PATCH 0419/1087] Fix explicit table recommendation asks --- wren-ai-service/src/web/v1/services/ask.py | 26 ++++++++++++++++--- .../pytest/services/test_ask_sales_sql.py | 8 ++++++ .../apollo/server/services/askingService.ts | 2 +- .../services/tests/askingService.test.ts | 26 +++++++++++++++++++ 4 files changed, 57 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 098bd44f4f..dbd172442c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1208,7 +1208,7 @@ def _build_schema_grounded_table_question_sql( wants_monthly_count = any( term in normalized for term in ("monthly", "by month", "per month", "month-wise") - ) and any(term in normalized for term in ("count", "records", "rows")) + ) and any(term in normalized for term in ("count", "record", "records", "rows")) if wants_monthly_count: date_column = self._find_temporal_column_for_query(query, table) if not date_column: @@ -1345,13 +1345,14 @@ def _build_explicit_table_preview_sql( def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names: list[str] = [] for match in re.finditer( - r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\b(?:from|in|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - if table_name and table_name not in table_names: - table_names.append(table_name) + for candidate in self._explicit_table_name_candidates(table_name): + if candidate and candidate not in table_names: + table_names.append(candidate) for match in re.finditer( r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", @@ -1366,6 +1367,23 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names.append(table_name) return table_names + def _explicit_table_name_candidates(self, table_name: str) -> list[str]: + table_name = str(table_name or "").strip(".,;:()[]{}") + if not table_name: + return [] + + candidates = [table_name] + dotted_parts = [part for part in re.split(r"[.$]", table_name) if part] + if len(dotted_parts) > 1: + candidates.append("_".join(dotted_parts)) + candidates.append(dotted_parts[-1]) + + unique_candidates: list[str] = [] + for candidate in candidates: + if candidate and candidate not in unique_candidates: + unique_candidates.append(candidate) + return unique_candidates + def _build_direct_orders_sales_sql(self, query: str) -> str | None: return None diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index e3abde7150..2893956565 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -135,6 +135,14 @@ def test_extract_explicit_table_names_from_query(): ) == ["tblNewOrders"] +def test_extract_explicit_table_names_from_in_clause_adds_deployed_table_candidate(): + service = AskService.__new__(AskService) + + assert service._extract_explicit_table_names_from_query( + "Show monthly record count by created_at in dbo.failure_patterns." + ) == ["dbo.failure_patterns", "dbo_failure_patterns", "failure_patterns"] + + def test_extract_explicit_table_names_from_using_clause(): service = AskService.__new__(AskService) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index f5b1115cae..fe7025d9c9 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1382,7 +1382,7 @@ export class AskingService implements IAskingService { input: AskingDetailTaskInput, payload: AskingPayload, ): Promise { - if (input.trackedAskingResult || !input.sql || !input.question) { + if (input.trackedAskingResult || !input.question) { return input.trackedAskingResult; } diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index d33c9e166d..55c01871c7 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -230,5 +230,31 @@ describe('AskingService', () => { askingTaskId: trackedAskingResult.taskId, }); }); + + test('also creates a normal asking task when direct question payload has no SQL', async () => { + const service = createService(); + + await service.createThreadResponse( + { + question: trackedAskingResult.question, + }, + 7, + ); + + expect(service.askingTaskTracker.createAskingTask).toHaveBeenCalledWith( + expect.objectContaining({ + query: trackedAskingResult.question, + histories: null, + deployId: 'latest-deploy-hash', + projectId: '1', + }), + ); + expect(service.threadResponseRepository.createOne).toHaveBeenCalledWith({ + threadId: 7, + question: trackedAskingResult.question, + sql: undefined, + askingTaskId: trackedAskingResult.taskId, + }); + }); }); }); From c83aa0a37b26e0e5ad3f726c2c5500dd648ac8d7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 14:56:38 +0530 Subject: [PATCH 0420/1087] Reject columns outside referenced table schema --- wren-ai-service/src/web/v1/services/ask.py | 40 ++++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 32 +++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index dbd172442c..bd73776519 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -678,12 +678,50 @@ def _sql_covers_required_question_concepts( def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: - valid_columns = { + all_valid_columns = { str(column.get("name") or "").lower() for table in schema_tables for column in table.get("columns", []) if column.get("name") } + columns_by_table: dict[str, set[str]] = {} + for table in schema_tables: + table_name = str(table.get("name") or "").lower() + if not table_name: + continue + table_columns = { + str(column.get("name") or "").lower() + for column in table.get("columns", []) + if column.get("name") + } + columns_by_table[table_name] = table_columns + columns_by_table[table_name.split(".")[-1]] = table_columns + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, + ) + referenced_table_keys = { + ( + next(value for value in match.groupdict().values() if value) or "" + ).lower() + for match in table_reference_pattern.finditer(sql or "") + } + referenced_column_sets = [ + columns + for table_key in referenced_table_keys + for columns in [ + columns_by_table.get(table_key) + or columns_by_table.get(table_key.split(".")[-1]) + ] + if columns is not None + ] + valid_columns = ( + referenced_column_sets[0] + if len(referenced_column_sets) == 1 + else all_valid_columns + ) valid_tables = { str(table.get("name") or "").lower() for table in schema_tables diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2893956565..569ab4f383 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1546,6 +1546,38 @@ def test_build_validated_ask_result_rejects_unqualified_invalid_columns(): ) +def test_build_validated_ask_result_rejects_column_from_different_active_table(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT policy_category_id AS "policy_category_id", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_knowledge_articles" ' + "GROUP BY policy_category_id" + ), + [ + """ + CREATE TABLE dbo_knowledge_articles ( + id VARCHAR, + policy_id VARCHAR, + category VARCHAR, + created_at TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_policies ( + id VARCHAR, + policy_category_id VARCHAR + ); + """, + ], + "Show knowledge article count by policy category.", + ) + + assert result is None + + def test_build_validated_ask_result_accepts_unqualified_valid_columns(): service = AskService.__new__(AskService) From 7ee464b9ab0e6632ceb157f321665d16836861c8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 15:10:41 +0530 Subject: [PATCH 0421/1087] Prefer unit tables for throughput trends --- wren-ai-service/src/web/v1/services/ask.py | 129 +++++++++++++++--- .../pytest/services/test_ask_sales_sql.py | 48 +++++++ 2 files changed, 160 insertions(+), 17 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bd73776519..196095eeb6 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2743,30 +2743,83 @@ def _build_manufacturing_throughput_sql( return None tables = self._parse_schema_tables(table_ddls) - table = self._find_best_schema_table_for_query(query, tables) + unit_candidates = ( + "BusinessUnit", + "Business_Unit", + "Business Unit", + "manufacturing_unit", + "manufacturing unit", + "manufacturingunit", + "ManufacturingUnit", + "unit", + "unit_name", + "BU", + "division", + ) + wants_temporal_trend = any( + term in normalized for term in ("trend", "trends", "monthly", "over time") + ) + scored_unit_tables: list[tuple[int, dict[str, Any], str, str | None]] = [] + for candidate_table in tables: + unit_column = self._find_schema_column(candidate_table, unit_candidates) + if not unit_column: + continue + + timestamp_column = self._find_temporal_column_for_query( + query, candidate_table + ) + if wants_temporal_trend and not timestamp_column: + continue + + table_name = str(candidate_table.get("name") or "") + normalized_table_name = table_name.lower() + score = 100 + if timestamp_column: + score += 40 + if any( + token in normalized_table_name + for token in ( + "debug", + "entry", + "entries", + "production", + "event", + "events", + "repair", + "manufacturing", + ) + ): + score += 20 + if self._table_matches_query(table_name, query): + score += 15 + scored_unit_tables.append( + (score, candidate_table, unit_column, timestamp_column) + ) + + table = None + preferred_unit_column = None + preferred_timestamp_column = None + if scored_unit_tables: + _, table, preferred_unit_column, preferred_timestamp_column = sorted( + scored_unit_tables, key=lambda item: item[0], reverse=True + )[0] + else: + table = self._find_best_schema_table_for_query(query, tables) + if table: - unit_column = self._find_schema_column( - table, - ( - "BusinessUnit", - "business_unit", - "manufacturing_unit", - "manufacturingunit", - "unit", - "unit_name", - "BU", - "division", - ), + unit_column = preferred_unit_column or self._find_schema_column( + table, unit_candidates ) if unit_column: table_name = str(table.get("name") or "") table_ref = self._quote_sql_identifier(table_name) unit_ref = f"{table_ref}.{self._quote_sql_identifier(unit_column)}" - timestamp_column = self._find_temporal_column_for_query(query, table) + timestamp_column = ( + preferred_timestamp_column + or self._find_temporal_column_for_query(query, table) + ) - if timestamp_column and any( - term in normalized for term in ("trend", "monthly", "over time") - ): + if timestamp_column and wants_temporal_trend: timestamp_ref = ( f"{table_ref}.{self._quote_sql_identifier(timestamp_column)}" ) @@ -4521,6 +4574,24 @@ def _prune_sql_generation_context( self._normalize_schema_token(table_name) for table_name in self._extract_explicit_table_names_from_query(query) } + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + wants_throughput_by_unit = "throughput" in normalized_query and any( + term in normalized_query + for term in ( + "manufacturing unit", + "manufacturing units", + "business unit", + "business units", + "different unit", + "different units", + "unit", + "units", + ) + ) + wants_temporal_trend = any( + term in normalized_query + for term in ("trend", "trends", "monthly", "over time", "by month") + ) scored: list[tuple[int, int]] = [] for index, table in enumerate(parsed_tables): @@ -4555,6 +4626,30 @@ def _prune_sql_generation_context( elif term in column_term or column_term in term: score += 25 + if wants_throughput_by_unit: + unit_column = self._find_schema_column( + table, + ( + "BusinessUnit", + "Business_Unit", + "Business Unit", + "manufacturing_unit", + "manufacturing unit", + "ManufacturingUnit", + "unit", + "unit_name", + "BU", + "division", + ), + ) + temporal_column = self._find_temporal_column_for_query(query, table) + if unit_column: + score += 500 + if unit_column and temporal_column: + score += 300 + if wants_temporal_trend and unit_column and temporal_column: + score += 300 + if score > 0: scored.append((score, index)) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 569ab4f383..6ebd1d3653 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -891,6 +891,54 @@ def test_build_manufacturing_throughput_sql_uses_active_unit_and_date_columns(): ) +def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Show throughput trends across different manufacturing units.", + [ + """ + CREATE TABLE dbo_ai_job_queue ( + database_name VARCHAR, + status VARCHAR, + payload VARCHAR, + error VARCHAR, + started_at TIMESTAMP, + completed_at TIMESTAMP, + created_at TIMESTAMP, + updated_at TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_DebugEntries ( + DebugEntryId VARCHAR, + BusinessUnit VARCHAR, + DateIn TIMESTAMP, + Status VARCHAR + ); + """, + ], + ) + + assert sql == ( + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "BusinessUnit", ' + 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "year", ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries" ' + 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' + 'AND "dbo_DebugEntries"."DateIn" IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' + 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn"), ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ' + 'ORDER BY "dbo_DebugEntries"."BusinessUnit" ASC, ' + 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC, ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' + ) + assert "dbo_ai_job_queue" not in sql + assert "database_name" not in sql + + def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): service = AskService.__new__(AskService) From 259fdc1cf986814fa02fda0e77c4bef8da04afba Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 15:32:35 +0530 Subject: [PATCH 0422/1087] Validate unqualified SQL columns against metadata --- .../src/pipelines/generation/utils/sql.py | 152 +++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 63 +++++++ .../apollo/server/services/queryService.ts | 159 ++++++++++++++++++ .../services/tests/queryService.test.ts | 110 ++++++++++++ 4 files changed, 484 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 0c1f7e58fa..eb5498ea44 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2975,6 +2975,148 @@ def replace_unqualified_identifier(match: re.Match[str]) -> str: ) +_SQL_IDENTIFIER_KEYWORDS = { + "all", + "and", + "as", + "asc", + "avg", + "between", + "by", + "case", + "cast", + "coalesce", + "count", + "current_date", + "date", + "dateadd", + "datediff", + "datepart", + "day", + "desc", + "distinct", + "else", + "end", + "extract", + "false", + "from", + "group", + "having", + "hour", + "in", + "is", + "join", + "last", + "like", + "limit", + "minute", + "month", + "not", + "null", + "nulls", + "on", + "or", + "order", + "over", + "partition", + "quarter", + "second", + "select", + "sum", + "then", + "top", + "true", + "when", + "week", + "where", + "with", + "year", +} + + +def _extract_projection_aliases(sql: str) -> set[str]: + aliases: set[str] = set() + for start, end in _find_select_list_spans(sql): + for item in _split_top_level_select_items(sql[start:end]): + alias_match = re.search( + rf"\s+(?:AS\s+)?(?P{_SQL_IDENTIFIER_PATTERN})\s*$", + item, + flags=re.IGNORECASE, + ) + if not alias_match: + continue + expression = item[: alias_match.start()].strip() + alias = _normalize_sql_identifier(alias_match.group("alias")) + if expression and alias: + aliases.add(alias.lower()) + return aliases + + +def _find_invalid_unqualified_column_references_for_single_table( + sql: str, + table_name: str, + valid_columns: set[str], + valid_compact_columns: set[str], + aliases: dict[str, str], +) -> list[str]: + table_aliases = { + alias.lower() + for alias, alias_table in aliases.items() + if str(alias_table).lower() == str(table_name).lower() + } + table_aliases.update( + suffix.lower() for suffix in _table_reference_suffixes(str(table_name)) + ) + projection_aliases = _extract_projection_aliases(sql) + sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") + invalid_references: list[str] = [] + + identifier_pattern = re.compile( + r'"(?P[^"]+)"|`(?P[^`]+)`|\[(?P[^\]]+)\]|(?P\b[A-Za-z_][A-Za-z0-9_$]*\b)' + ) + for match in identifier_pattern.finditer(sql_without_strings): + identifier = ( + match.group("quoted") + or match.group("backticked") + or match.group("bracketed") + or match.group("bare") + or "" + ) + normalized_identifier = identifier.lower() + if not normalized_identifier: + continue + + before = sql_without_strings[: match.start()].rstrip() + after = sql_without_strings[match.end() :].lstrip() + if before.endswith(".") or after.startswith("."): + continue + if after.startswith("("): + continue + + previous_word_match = re.search(r"([A-Za-z_][A-Za-z0-9_$]*)\s*$", before) + previous_word = ( + previous_word_match.group(1).lower() if previous_word_match else "" + ) + if previous_word == "as": + continue + + if ( + normalized_identifier in _SQL_IDENTIFIER_KEYWORDS + or normalized_identifier in table_aliases + or normalized_identifier in projection_aliases + ): + continue + if ( + normalized_identifier in valid_columns + or _compact_sql_identifier(identifier) in valid_compact_columns + ): + continue + if identifier not in invalid_references: + invalid_references.append(identifier) + + return invalid_references + + def find_invalid_column_references( sql: str, valid_table_columns: dict[str, list[str]] ) -> list[str]: @@ -3038,6 +3180,16 @@ def find_invalid_column_references( ): invalid_references.append(column) + invalid_references.extend( + _find_invalid_unqualified_column_references_for_single_table( + sql, + str(table_name), + valid_columns, + valid_compact_columns, + aliases, + ) + ) + return sorted(set(invalid_references)) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 421566f5eb..47cc307e55 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -469,6 +469,69 @@ def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid ) == ["dbo_qSales1.UnitPrice"] +def test_find_invalid_column_references_rejects_unqualified_function_arguments(): + sql = ( + 'SELECT user_id customer, COUNT(*) issue_count, ' + 'AVG(DATEDIFF(day, created_at, resolved_at)) avg_resolution_time ' + 'FROM "dbo_kb_article_feedback" ' + "WHERE details LIKE '%device%' AND helpful = 0 " + "GROUP BY user_id " + "ORDER BY issue_count DESC NULLS LAST" + ) + + assert find_invalid_column_references( + sql, + { + "dbo_kb_article_feedback": [ + "detail_id", + "history_id", + "target_query_expression", + "execution_date", + "result", + "result_detail", + "exception_message", + "exception", + "id", + "org_id", + "article_id", + "user_id", + "helpful", + "reasons", + "details", + "created_at", + ] + }, + ) == ["resolved_at"] + + +def test_find_invalid_column_references_rejects_unqualified_where_columns(): + sql = ( + 'SELECT COUNT(*) AS "count" ' + 'FROM "dbo_ai_workflows" ' + "WHERE role = 'admin' " + "GROUP BY debug_user" + ) + + assert find_invalid_column_references( + sql, + { + "dbo_ai_workflows": [ + "id", + "org_id", + "repair_id", + "workflow_name", + "priority", + "steps", + "estimated_total_min", + "source_inspection_id", + "created_by_user_id", + "created_at", + "updated_at", + ] + }, + ) == ["debug_user", "role"] + + def test_normalize_sql_column_references_to_schema_maps_kb_article_aliases(): sql = ( 'SELECT "dbo_kb_articles"."article_type", COUNT(*) AS "RecordCount" ' diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 09c3734120..ce64720615 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -600,6 +600,158 @@ const extractSimpleProjectionColumns = (sql: string) => { return columns; }; +const SQL_IDENTIFIER_KEYWORDS = new Set([ + 'all', + 'and', + 'as', + 'asc', + 'avg', + 'between', + 'by', + 'case', + 'cast', + 'coalesce', + 'count', + 'current_date', + 'date', + 'dateadd', + 'datediff', + 'datepart', + 'day', + 'desc', + 'distinct', + 'else', + 'end', + 'extract', + 'false', + 'from', + 'group', + 'having', + 'hour', + 'in', + 'is', + 'join', + 'last', + 'like', + 'limit', + 'minute', + 'month', + 'not', + 'null', + 'nulls', + 'on', + 'or', + 'order', + 'over', + 'partition', + 'quarter', + 'second', + 'select', + 'sum', + 'then', + 'top', + 'true', + 'when', + 'week', + 'where', + 'with', + 'year', +]); + +const extractProjectionAliases = (sql: string) => { + const aliases = new Set(); + const selectPattern = /\bSELECT\b(?.*?)(?=\bFROM\b)/gis; + let match: RegExpExecArray | null; + while ((match = selectPattern.exec(sql))) { + const body = match.groups?.body || ''; + splitTopLevelSqlList(body).forEach((item) => { + const aliasMatch = item.match( + new RegExp( + String.raw`\s+(?:AS\s+)?(?${SQL_IDENTIFIER_PATTERN})\s*$`, + 'i', + ), + ); + if (!aliasMatch?.groups?.alias || aliasMatch.index === undefined) { + return; + } + const expression = item.slice(0, aliasMatch.index).trim(); + if (expression) { + aliases.add( + normalizeSqlIdentifier(aliasMatch.groups.alias).toLowerCase(), + ); + } + }); + } + return aliases; +}; + +const findInvalidUnqualifiedColumnReferencesForSingleSchema = ( + sql: string, + schema: ManifestModelSchema, + aliases: Map, +) => { + const tableAliases = new Set(); + aliases.forEach((aliasSchema, alias) => { + if (aliasSchema === schema) { + tableAliases.add(alias.toLowerCase()); + } + }); + const projectionAliases = extractProjectionAliases(sql); + const sqlWithoutStrings = sql.replace(/'(?:''|[^'])*'/g, "''"); + const invalidReferences: string[] = []; + const identifierPattern = + /"(?[^"]+)"|`(?[^`]+)`|\[(?[^\]]+)\]|(?\b[A-Za-z_][A-Za-z0-9_$]*\b)/g; + + let match: RegExpExecArray | null; + while ((match = identifierPattern.exec(sqlWithoutStrings))) { + const identifier = + match.groups?.quoted || + match.groups?.backticked || + match.groups?.bracketed || + match.groups?.bare || + ''; + const normalizedIdentifier = identifier.toLowerCase(); + if (!normalizedIdentifier) { + continue; + } + + const before = sqlWithoutStrings.slice(0, match.index).trimEnd(); + const after = sqlWithoutStrings.slice(match.index + match[0].length).trimStart(); + if (before.endsWith('.') || after.startsWith('.')) { + continue; + } + if (after.startsWith('(')) { + continue; + } + + const previousWord = before + .match(/([A-Za-z_][A-Za-z0-9_$]*)\s*$/)?.[1] + ?.toLowerCase(); + if (previousWord === 'as') { + continue; + } + + if ( + SQL_IDENTIFIER_KEYWORDS.has(normalizedIdentifier) || + tableAliases.has(normalizedIdentifier) || + projectionAliases.has(normalizedIdentifier) + ) { + continue; + } + if ( + schema.columns.has(normalizedIdentifier) || + schema.columns.has(compactSqlIdentifier(identifier)) + ) { + continue; + } + if (!invalidReferences.includes(identifier)) { + invalidReferences.push(identifier); + } + } + + return invalidReferences; +}; + const findSqlReferenceValidationErrors = ( sql: string, manifest?: Manifest, @@ -684,6 +836,13 @@ const findSqlReferenceValidationErrors = ( errors.push(column); } }); + errors.push( + ...findInvalidUnqualifiedColumnReferencesForSingleSchema( + sql, + schema, + aliases, + ), + ); } return [...new Set(errors)]; diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 8aa5e6cd58..8b98c6f378 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -457,6 +457,116 @@ describe('QueryService', () => { expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); }); + it('should reject unqualified function arguments outside a single active manifest table', async () => { + await expect( + queryService.preview( + [ + 'SELECT user_id customer, COUNT(*) issue_count,', + 'AVG(DATEDIFF(day, created_at, resolved_at)) avg_resolution_time', + 'FROM "dbo_kb_article_feedback"', + "WHERE details LIKE '%device%' AND helpful = 0", + 'GROUP BY user_id', + 'ORDER BY issue_count DESC NULLS LAST', + ].join(' '), + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_kb_article_feedback', + tableReference: { table: 'dbo_kb_article_feedback' }, + columns: [ + { name: 'detail_id', type: 'string', isCalculated: false }, + { name: 'history_id', type: 'string', isCalculated: false }, + { + name: 'target_query_expression', + type: 'string', + isCalculated: false, + }, + { + name: 'execution_date', + type: 'timestamp', + isCalculated: false, + }, + { name: 'result', type: 'string', isCalculated: false }, + { name: 'result_detail', type: 'string', isCalculated: false }, + { + name: 'exception_message', + type: 'string', + isCalculated: false, + }, + { name: 'exception', type: 'string', isCalculated: false }, + { name: 'id', type: 'string', isCalculated: false }, + { name: 'org_id', type: 'string', isCalculated: false }, + { name: 'article_id', type: 'string', isCalculated: false }, + { name: 'user_id', type: 'string', isCalculated: false }, + { name: 'helpful', type: 'integer', isCalculated: false }, + { name: 'reasons', type: 'string', isCalculated: false }, + { name: 'details', type: 'string', isCalculated: false }, + { name: 'created_at', type: 'timestamp', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }, + ), + ).rejects.toThrow( + 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: resolved_at', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + + it('should reject unqualified where and group columns outside a single active manifest table', async () => { + await expect( + queryService.preview( + 'SELECT COUNT(*) AS "count" FROM "dbo_ai_workflows" WHERE role = \'admin\' GROUP BY debug_user', + { + project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'dbo_ai_workflows', + tableReference: { table: 'dbo_ai_workflows' }, + columns: [ + { name: 'id', type: 'string', isCalculated: false }, + { name: 'org_id', type: 'string', isCalculated: false }, + { name: 'repair_id', type: 'string', isCalculated: false }, + { name: 'workflow_name', type: 'string', isCalculated: false }, + { name: 'priority', type: 'string', isCalculated: false }, + { name: 'steps', type: 'string', isCalculated: false }, + { + name: 'estimated_total_min', + type: 'integer', + isCalculated: false, + }, + { + name: 'source_inspection_id', + type: 'string', + isCalculated: false, + }, + { + name: 'created_by_user_id', + type: 'string', + isCalculated: false, + }, + { name: 'created_at', type: 'timestamp', isCalculated: false }, + { name: 'updated_at', type: 'timestamp', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }, + ), + ).rejects.toThrow( + 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: debug_user, role', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + it('should reject numeric aggregates on non-numeric manifest columns before ibis planning', async () => { await expect( queryService.preview('SELECT AVG("orders"."quantity") FROM "orders"', { From 0053bcf38a101722f8d9137c46f3636fbc399621 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 15:52:25 +0530 Subject: [PATCH 0423/1087] Use throughput builder for manufacturing unit trends --- wren-ai-service/src/web/v1/services/ask.py | 29 ++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 34 +++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 196095eeb6..096b441aef 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1733,6 +1733,11 @@ def _build_schema_grounded_analytics_sql( ): return monthly_repair_volume_sql + if throughput_sql := self._build_manufacturing_throughput_sql( + query, table_ddls + ): + return throughput_sql + if operational_sql := self._build_schema_grounded_operational_sql( query, tables ): @@ -2755,6 +2760,16 @@ def _build_manufacturing_throughput_sql( "unit_name", "BU", "division", + "Debug_Shelf", + "DebugShelf", + "debug shelf", + "shelf", + "station", + "workstation", + "work_center", + "workcenter", + "line", + "cell", ) wants_temporal_trend = any( term in normalized for term in ("trend", "trends", "monthly", "over time") @@ -4042,6 +4057,10 @@ def _build_schema_grounded_sales_sql( "salesperson", "sales person", "trend", + "throughput", + "manufacturing", + "unit", + "units", "value", ) ): @@ -4640,6 +4659,16 @@ def _prune_sql_generation_context( "unit_name", "BU", "division", + "Debug_Shelf", + "DebugShelf", + "debug shelf", + "shelf", + "station", + "workstation", + "work_center", + "workcenter", + "line", + "cell", ), ) temporal_column = self._find_temporal_column_for_query(query, table) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 6ebd1d3653..699cca6a3e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -939,6 +939,40 @@ def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables assert "database_name" not in sql +def test_build_manufacturing_throughput_sql_uses_debug_shelf_unit_column(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Show throughput trends across different manufacturing units.", + [ + """ + CREATE TABLE dbo_DebugEntries_Staging2 ( + id INTEGER, + Debug_Shelf VARCHAR, + last_update_date TIMESTAMP, + status VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_DebugEntries_Staging2"."Debug_Shelf" AS "Debug_Shelf", ' + 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date") AS "year", ' + 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") AS "month", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries_Staging2" ' + 'WHERE "dbo_DebugEntries_Staging2"."Debug_Shelf" IS NOT NULL ' + 'AND "dbo_DebugEntries_Staging2"."last_update_date" IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries_Staging2"."Debug_Shelf", ' + 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date"), ' + 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") ' + 'ORDER BY "dbo_DebugEntries_Staging2"."Debug_Shelf" ASC, ' + 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date") ASC, ' + 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") ASC' + ) + + def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): service = AskService.__new__(AskService) From 7a0a9a5c1393ab4756ed4bf46b3162d88dbcfa3e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 16:29:46 +0530 Subject: [PATCH 0424/1087] Keep explicit table analytics out of preview fallback --- wren-ai-service/src/web/v1/services/ask.py | 23 ++++++++ .../pytest/services/test_ask_sales_sql.py | 55 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 096b441aef..f9247773f9 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1219,6 +1219,23 @@ def _build_schema_grounded_table_question_sql( return None tables = self._parse_schema_tables(table_ddls) + explicit_table_names = self._extract_explicit_table_names_from_query(query) + if explicit_table_names: + explicit_keys = { + key + for explicit_table_name in explicit_table_names + for key in self._schema_identifier_alias_keys(explicit_table_name) + } + explicit_tables = [ + table + for table in tables + if explicit_keys.intersection( + self._schema_identifier_alias_keys(str(table.get("name") or "")) + ) + ] + if explicit_tables: + tables = explicit_tables + table = self._find_best_schema_table_for_query(query, tables) if not table: return None @@ -1338,6 +1355,12 @@ def _build_explicit_table_preview_sql( r"\b(?:rows?|records?|data)\b", normalized_query, flags=re.IGNORECASE ): return None + if re.search( + r"\b(?:count|counts|monthly|month|trend|trends|by|per|each|distribution|group(?:ed)?|aggregate)\b", + normalized_query, + flags=re.IGNORECASE, + ): + return None tables = self._parse_schema_tables(table_ddls) if not tables: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 699cca6a3e..7930c553ae 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -127,6 +127,24 @@ def test_build_explicit_table_preview_sql_for_show_data_prompt(): assert result == ('SELECT TOP 10 * FROM "CustomerMaster"', "CustomerMaster") +def test_build_explicit_table_preview_sql_ignores_monthly_count_question(): + service = AskService.__new__(AskService) + + result = service._build_explicit_table_preview_sql( + "Show monthly record count by created_at in dbo.failure_patterns.", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + created_at TIMESTAMP + ); + """ + ], + ) + + assert result is None + + def test_extract_explicit_table_names_from_query(): service = AskService.__new__(AskService) @@ -1500,6 +1518,43 @@ def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count() ) +def test_build_schema_grounded_table_question_sql_prefers_explicit_table_alias(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Show monthly record count by created_at in dbo.failure_patterns.", + [ + """ + CREATE TABLE dbo_knowledge_articles ( + id INTEGER, + policy_category_id INTEGER, + created_at TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR, + created_at TIMESTAMP + ); + """, + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ASC' + ) + assert "dbo_knowledge_articles" not in sql + + def test_build_validated_ask_result_rejects_status_for_product_line_pcb_question(): service = AskService.__new__(AskService) From e8c75256b5bbd020733aeffe0ed51edfc80cf665 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 16:44:17 +0530 Subject: [PATCH 0425/1087] Route monthly record counts through schema SQL --- wren-ai-service/src/web/v1/services/ask.py | 8 +++++ .../pytest/services/test_ask_sales_sql.py | 30 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f9247773f9..d4fb128985 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4082,6 +4082,14 @@ def _build_schema_grounded_sales_sql( "trend", "throughput", "manufacturing", + "monthly", + "month", + "count", + "counts", + "record", + "records", + "row", + "rows", "unit", "units", "value", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 7930c553ae..cda21fd212 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1518,6 +1518,36 @@ def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count() ) +def test_build_schema_grounded_sales_sql_handles_how_monthly_record_count(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "how monthly record count by created_at in dbo.failure_patterns.", + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + name VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_failure_patterns" ' + 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' + 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ASC' + ) + assert '"dbo_failure"."patterns"' not in sql + + def test_build_schema_grounded_table_question_sql_prefers_explicit_table_alias(): service = AskService.__new__(AskService) From 90808d695c1a090091a35d39f8f5f9f324273847 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 17:01:13 +0530 Subject: [PATCH 0426/1087] Reject SQL for wrong explicit table --- wren-ai-service/src/web/v1/services/ask.py | 44 ++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 59 +++++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d4fb128985..700ef86c78 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1124,6 +1124,47 @@ def _table_matches_query(self, table_name: str, query: str) -> bool: for table_key in table_keys ) + def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: + keys: set[str] = set() + for table_name in self._extract_explicit_table_names_from_query(query or ""): + keys.update(self._schema_identifier_alias_keys(table_name)) + return keys + + def _sql_references_explicit_table( + self, + sql: str, + query: str | None, + ) -> bool: + explicit_table_keys = self._explicit_table_alias_keys_from_query(query) + if not explicit_table_keys: + return True + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, + ) + referenced_tables = [ + next(value for value in match.groupdict().values() if value) + for match in table_reference_pattern.finditer(sql) + ] + + for table_reference in referenced_tables: + reference_keys = self._schema_identifier_alias_keys(table_reference) + short_reference = re.split(r"[.$_]", str(table_reference or ""))[-1] + reference_keys.update(self._schema_identifier_alias_keys(short_reference)) + if explicit_table_keys.intersection(reference_keys): + return True + + logger.warning( + "Ignoring SQL because it does not reference the explicitly requested table. " + "query=%s referenced_tables=%s sql=%s", + query, + referenced_tables, + sql, + ) + return False + def _find_best_schema_table_for_query( self, query: str, tables: list[dict[str, Any]] ) -> dict[str, Any] | None: @@ -4869,6 +4910,9 @@ def _build_validated_ask_result_from_sql( ) return None + if not self._sql_references_explicit_table(ask_result.sql, query): + return None + if not self._sql_matches_question_intent( ask_result.sql, query, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index cda21fd212..ea40d53f44 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -172,6 +172,65 @@ def test_extract_explicit_table_names_from_using_clause(): ) == [] +def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): + service = AskService.__new__(AskService) + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date") AS "year", ' + 'DATEPART(MONTH, "dbo_knowledge_articles"."last_run_date") AS "month", ' + '"dbo_knowledge_articles"."category" AS "category", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_knowledge_articles" ' + 'GROUP BY DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date"), ' + 'DATEPART(MONTH, "dbo_knowledge_articles"."last_run_date"), ' + '"dbo_knowledge_articles"."category"' + ), + [ + """ + CREATE TABLE dbo_knowledge_articles ( + id INTEGER, + last_run_date TIMESTAMP, + category VARCHAR + ); + """, + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + created_at TIMESTAMP + ); + """, + ], + "show monthly record count by created_at in dbo.failure_patterns.", + ) + + assert result is None + + +def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): + service = AskService.__new__(AskService) + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_failure_patterns" ' + 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at")' + ), + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + created_at TIMESTAMP + ); + """ + ], + "show monthly record count by created_at in dbo.failure_patterns.", + ) + + assert result is not None + + def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) From df3115b09728e35ef5a139a9f809a3971492d04a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 17:22:34 +0530 Subject: [PATCH 0427/1087] Filter explicit table retrieval context --- wren-ai-service/src/web/v1/services/ask.py | 160 +++++++++++++++--- .../pytest/services/test_ask_sales_sql.py | 35 ++++ 2 files changed, 172 insertions(+), 23 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 700ef86c78..51a5ea877c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1165,6 +1165,65 @@ def _sql_references_explicit_table( ) return False + def _filter_retrieval_metadata_for_explicit_query( + self, + query: str, + documents: list[dict], + ) -> tuple[list[dict], list[str], list[str]]: + explicit_table_keys = self._explicit_table_alias_keys_from_query(query) + if not explicit_table_keys: + table_names = [ + table_name + for document in documents + if isinstance(table_name := document.get("table_name"), str) + and table_name.strip() + ] + table_ddls = [ + table_ddl + for document in documents + if isinstance(table_ddl := document.get("table_ddl"), str) + and table_ddl.strip() + ] + return documents, table_names, table_ddls + + matched_documents: list[dict] = [] + for document in documents: + candidate_names = [] + if isinstance(table_name := document.get("table_name"), str): + candidate_names.append(table_name) + if isinstance(table_ddl := document.get("table_ddl"), str): + candidate_names.extend( + str(table.get("name") or "") + for table in self._parse_schema_tables([table_ddl]) + if table.get("name") + ) + + candidate_keys: set[str] = set() + for candidate_name in candidate_names: + candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) + candidate_keys.update( + self._schema_identifier_alias_keys( + re.split(r"[.$_]", str(candidate_name or ""))[-1] + ) + ) + + if explicit_table_keys.intersection(candidate_keys): + matched_documents.append(document) + + table_names = [ + table_name + for document in matched_documents + if isinstance(table_name := document.get("table_name"), str) + and table_name.strip() + ] + table_ddls = [ + table_ddl + for document in matched_documents + if isinstance(table_ddl := document.get("table_ddl"), str) + and table_ddl.strip() + ] + return matched_documents, table_names, table_ddls + def _find_best_schema_table_for_query( self, query: str, tables: list[dict[str, Any]] ) -> dict[str, Any] | None: @@ -5156,6 +5215,41 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) + ) + if not documents: + logger.info( + "Explicit table retrieval did not return the requested table for query_id %s; " + "loading full active schema.", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval for explicit table", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + all_documents, _, _ = self._extract_retrieval_metadata( + retrieval_result + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + all_documents, + ) + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) @@ -5223,27 +5317,27 @@ async def ask( user_query, table_ddls ) ): - api_results = [ - AskResult( - **{ - "sql": deterministic_sql, - "type": "llm", - } - ) - ] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sql, + table_ddls, + user_query, ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = deterministic_sql if not documents: error_message = ( @@ -5704,10 +5798,17 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - if not documents: - explicit_table_names = self._extract_explicit_table_names_from_query( - user_query + explicit_table_names = self._extract_explicit_table_names_from_query( + user_query + ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) ) + if not documents: if explicit_table_names: logger.info( "Retrying schema retrieval for explicit tables query_id %s: %s", @@ -5734,6 +5835,12 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) + ) if not documents and self._is_data_analysis_query(user_query): logger.info( "Query-based schema retrieval returned no tables for data question; " @@ -5760,6 +5867,13 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) + ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index ea40d53f44..c6c2494303 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -231,6 +231,41 @@ def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): assert result is not None +def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): + service = AskService.__new__(AskService) + documents = [ + { + "table_name": "dbo_knowledge_articles", + "table_ddl": """ + CREATE TABLE dbo_knowledge_articles ( + id INTEGER, + last_run_date TIMESTAMP + ); + """, + }, + { + "table_name": "dbo_failure_patterns", + "table_ddl": """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + created_at TIMESTAMP + ); + """, + }, + ] + + filtered_documents, table_names, table_ddls = ( + service._filter_retrieval_metadata_for_explicit_query( + "show monthly record count by created_at in dbo.failure_patterns.", + documents, + ) + ) + + assert filtered_documents == [documents[1]] + assert table_names == ["dbo_failure_patterns"] + assert table_ddls == [documents[1]["table_ddl"]] + + def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) From dfddd42f72b2aa1f48033beacb14bf0d3629ebe4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 17:52:44 +0530 Subject: [PATCH 0428/1087] Avoid debug tables for manufacturing throughput --- wren-ai-service/src/web/v1/services/ask.py | 150 ++++++++++++++---- .../pytest/services/test_ask_sales_sql.py | 88 ++++++++-- 2 files changed, 188 insertions(+), 50 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 51a5ea877c..bb4439ecf6 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -904,6 +904,27 @@ def _sql_matches_question_intent( if not referenced_tables: return True + asks_manufacturing_throughput = "throughput" in normalized_query and any( + term in normalized_query + for term in ("manufacturing", "manufacturing unit", "manufacturing units") + ) + if ( + asks_manufacturing_throughput + and not self._query_allows_internal_debug_table(query) + and any( + self._is_internal_debug_table_name(table_reference) + for table_reference in referenced_tables + ) + ): + logger.warning( + "Ignoring SQL because manufacturing throughput query selected an internal debug/staging table. " + "query=%s referenced_tables=%s sql=%s", + query, + referenced_tables, + sql, + ) + return False + referenced_table_tokens = set().union( *[self._schema_name_tokens(table) for table in referenced_tables] ) @@ -1114,6 +1135,36 @@ def _schema_identifier_alias_keys(self, value: str) -> set[str]: if key } + def _query_allows_internal_debug_table(self, query: str | None) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + return any( + term in normalized + for term in ( + "debug", + "debugging", + "debug entries", + "debugentries", + "debug shelf", + "staging", + "internal", + ) + ) + + def _is_internal_debug_table_name(self, table_name: str) -> bool: + normalized = str(table_name or "").lower() + return any( + token in normalized + for token in ( + "debug", + "debugentries", + "staging", + "ai_job", + "sysdiagram", + "knex", + "migration", + ) + ) + def _table_matches_query(self, table_name: str, query: str) -> bool: query_keys = self._schema_identifier_alias_keys(query) short_table = re.split(r"[.$_]", str(table_name or ""))[-1] @@ -2864,6 +2915,8 @@ def _build_manufacturing_throughput_sql( "business units", "different unit", "different units", + "debug shelf", + "debug shelves", ) ) @@ -2871,7 +2924,7 @@ def _build_manufacturing_throughput_sql( return None tables = self._parse_schema_tables(table_ddls) - unit_candidates = ( + unit_candidates = [ "BusinessUnit", "Business_Unit", "Business Unit", @@ -2883,17 +2936,22 @@ def _build_manufacturing_throughput_sql( "unit_name", "BU", "division", - "Debug_Shelf", - "DebugShelf", - "debug shelf", - "shelf", "station", "workstation", "work_center", "workcenter", "line", "cell", - ) + ] + if self._query_allows_internal_debug_table(query): + unit_candidates.extend( + ( + "Debug_Shelf", + "DebugShelf", + "debug shelf", + "shelf", + ) + ) wants_temporal_trend = any( term in normalized for term in ("trend", "trends", "monthly", "over time") ) @@ -2911,28 +2969,38 @@ def _build_manufacturing_throughput_sql( table_name = str(candidate_table.get("name") or "") normalized_table_name = table_name.lower() + if self._is_internal_debug_table_name( + table_name + ) and not self._query_allows_internal_debug_table(query): + continue + score = 100 if timestamp_column: score += 40 if any( token in normalized_table_name for token in ( - "debug", - "entry", - "entries", "production", "event", "events", "repair", "manufacturing", + "factory", + "throughput", ) ): - score += 20 + score += 50 + if any( + token in normalized_table_name + for token in ("debug", "entry", "entries", "staging") + ): + score -= 80 if self._table_matches_query(table_name, query): score += 15 - scored_unit_tables.append( - (score, candidate_table, unit_column, timestamp_column) - ) + if score > 0: + scored_unit_tables.append( + (score, candidate_table, unit_column, timestamp_column) + ) table = None preferred_unit_column = None @@ -2992,6 +3060,8 @@ def _build_manufacturing_throughput_sql( has_business_unit = self._schema_contains( table_ddls, r"\bBusinessUnit\b", table_names=table_names ) + if not self._query_allows_internal_debug_table(query): + return None if not (has_debug_entries and has_business_unit): return None @@ -4777,30 +4847,36 @@ def _prune_sql_generation_context( score += 25 if wants_throughput_by_unit: + unit_column_candidates = [ + "BusinessUnit", + "Business_Unit", + "Business Unit", + "manufacturing_unit", + "manufacturing unit", + "ManufacturingUnit", + "unit", + "unit_name", + "BU", + "division", + "station", + "workstation", + "work_center", + "workcenter", + "line", + "cell", + ] + if self._query_allows_internal_debug_table(query): + unit_column_candidates.extend( + ( + "Debug_Shelf", + "DebugShelf", + "debug shelf", + "shelf", + ) + ) unit_column = self._find_schema_column( table, - ( - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "manufacturing unit", - "ManufacturingUnit", - "unit", - "unit_name", - "BU", - "division", - "Debug_Shelf", - "DebugShelf", - "debug shelf", - "shelf", - "station", - "workstation", - "work_center", - "workcenter", - "line", - "cell", - ), + unit_column_candidates, ) temporal_column = self._find_temporal_column_for_query(query, table) if unit_column: @@ -4809,6 +4885,10 @@ def _prune_sql_generation_context( score += 300 if wants_temporal_trend and unit_column and temporal_column: score += 300 + if self._is_internal_debug_table_name( + table_name + ) and not self._query_allows_internal_debug_table(query): + score -= 1200 if score > 0: scored.append((score, index)) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index c6c2494303..bd4d02238a 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1003,7 +1003,7 @@ def test_build_manufacturing_throughput_sql_uses_active_unit_and_date_columns(): ) -def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables(): +def test_build_manufacturing_throughput_sql_prefers_production_unit_table_over_admin_tables(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( @@ -1022,6 +1022,13 @@ def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables ); """, """ + CREATE TABLE dbo_production_events ( + id INTEGER, + manufacturing_unit VARCHAR, + created_at TIMESTAMP + ); + """, + """ CREATE TABLE dbo_DebugEntries ( DebugEntryId VARCHAR, BusinessUnit VARCHAR, @@ -1033,28 +1040,29 @@ def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables ) assert sql == ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "BusinessUnit", ' - 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "year", ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month", ' + 'SELECT "dbo_production_events"."manufacturing_unit" AS "manufacturing_unit", ' + 'DATEPART(YEAR, "dbo_production_events"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_production_events"."created_at") AS "month", ' 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' - 'AND "dbo_DebugEntries"."DateIn" IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' - 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn"), ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ' - 'ORDER BY "dbo_DebugEntries"."BusinessUnit" ASC, ' - 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC, ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' + 'FROM "dbo_production_events" ' + 'WHERE "dbo_production_events"."manufacturing_unit" IS NOT NULL ' + 'AND "dbo_production_events"."created_at" IS NOT NULL ' + 'GROUP BY "dbo_production_events"."manufacturing_unit", ' + 'DATEPART(YEAR, "dbo_production_events"."created_at"), ' + 'DATEPART(MONTH, "dbo_production_events"."created_at") ' + 'ORDER BY "dbo_production_events"."manufacturing_unit" ASC, ' + 'DATEPART(YEAR, "dbo_production_events"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_production_events"."created_at") ASC' ) assert "dbo_ai_job_queue" not in sql assert "database_name" not in sql + assert "dbo_DebugEntries" not in sql -def test_build_manufacturing_throughput_sql_uses_debug_shelf_unit_column(): +def test_build_manufacturing_throughput_sql_does_not_use_debug_shelf_for_manufacturing_units(): service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( + sql = service._build_manufacturing_throughput_sql( "Show throughput trends across different manufacturing units.", [ """ @@ -1068,6 +1076,26 @@ def test_build_manufacturing_throughput_sql_uses_debug_shelf_unit_column(): ], ) + assert sql is None + + +def test_build_manufacturing_throughput_sql_uses_debug_shelf_when_requested(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Show debug shelf throughput trends.", + [ + """ + CREATE TABLE dbo_DebugEntries_Staging2 ( + id INTEGER, + Debug_Shelf VARCHAR, + last_update_date TIMESTAMP, + status VARCHAR + ); + """ + ], + ) + assert sql == ( 'SELECT "dbo_DebugEntries_Staging2"."Debug_Shelf" AS "Debug_Shelf", ' 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date") AS "year", ' @@ -1085,6 +1113,36 @@ def test_build_manufacturing_throughput_sql_uses_debug_shelf_unit_column(): ) +def test_build_validated_ask_result_rejects_debug_staging_for_manufacturing_units(): + service = AskService.__new__(AskService) + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_DebugEntries_Staging2"."Debug_Shelf" AS "Debug_Shelf", ' + 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date") AS "year", ' + 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") AS "month", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries_Staging2" ' + 'WHERE "dbo_DebugEntries_Staging2"."Debug_Shelf" IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries_Staging2"."Debug_Shelf", ' + 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date"), ' + 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date")' + ), + [ + """ + CREATE TABLE dbo_DebugEntries_Staging2 ( + id INTEGER, + Debug_Shelf VARCHAR, + last_update_date TIMESTAMP, + status VARCHAR + ); + """ + ], + "Show throughput trends across different manufacturing units.", + ) + + assert result is None + + def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): service = AskService.__new__(AskService) From 5fcc4ea5b7a022c6beb49273ee05cceb6745b270 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:37 +0530 Subject: [PATCH 0429/1087] Revert "Avoid debug tables for manufacturing throughput" This reverts commit dfddd42f72b2aa1f48033beacb14bf0d3629ebe4. --- wren-ai-service/src/web/v1/services/ask.py | 150 ++++-------------- .../pytest/services/test_ask_sales_sql.py | 88 ++-------- 2 files changed, 50 insertions(+), 188 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bb4439ecf6..51a5ea877c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -904,27 +904,6 @@ def _sql_matches_question_intent( if not referenced_tables: return True - asks_manufacturing_throughput = "throughput" in normalized_query and any( - term in normalized_query - for term in ("manufacturing", "manufacturing unit", "manufacturing units") - ) - if ( - asks_manufacturing_throughput - and not self._query_allows_internal_debug_table(query) - and any( - self._is_internal_debug_table_name(table_reference) - for table_reference in referenced_tables - ) - ): - logger.warning( - "Ignoring SQL because manufacturing throughput query selected an internal debug/staging table. " - "query=%s referenced_tables=%s sql=%s", - query, - referenced_tables, - sql, - ) - return False - referenced_table_tokens = set().union( *[self._schema_name_tokens(table) for table in referenced_tables] ) @@ -1135,36 +1114,6 @@ def _schema_identifier_alias_keys(self, value: str) -> set[str]: if key } - def _query_allows_internal_debug_table(self, query: str | None) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - return any( - term in normalized - for term in ( - "debug", - "debugging", - "debug entries", - "debugentries", - "debug shelf", - "staging", - "internal", - ) - ) - - def _is_internal_debug_table_name(self, table_name: str) -> bool: - normalized = str(table_name or "").lower() - return any( - token in normalized - for token in ( - "debug", - "debugentries", - "staging", - "ai_job", - "sysdiagram", - "knex", - "migration", - ) - ) - def _table_matches_query(self, table_name: str, query: str) -> bool: query_keys = self._schema_identifier_alias_keys(query) short_table = re.split(r"[.$_]", str(table_name or ""))[-1] @@ -2915,8 +2864,6 @@ def _build_manufacturing_throughput_sql( "business units", "different unit", "different units", - "debug shelf", - "debug shelves", ) ) @@ -2924,7 +2871,7 @@ def _build_manufacturing_throughput_sql( return None tables = self._parse_schema_tables(table_ddls) - unit_candidates = [ + unit_candidates = ( "BusinessUnit", "Business_Unit", "Business Unit", @@ -2936,22 +2883,17 @@ def _build_manufacturing_throughput_sql( "unit_name", "BU", "division", + "Debug_Shelf", + "DebugShelf", + "debug shelf", + "shelf", "station", "workstation", "work_center", "workcenter", "line", "cell", - ] - if self._query_allows_internal_debug_table(query): - unit_candidates.extend( - ( - "Debug_Shelf", - "DebugShelf", - "debug shelf", - "shelf", - ) - ) + ) wants_temporal_trend = any( term in normalized for term in ("trend", "trends", "monthly", "over time") ) @@ -2969,38 +2911,28 @@ def _build_manufacturing_throughput_sql( table_name = str(candidate_table.get("name") or "") normalized_table_name = table_name.lower() - if self._is_internal_debug_table_name( - table_name - ) and not self._query_allows_internal_debug_table(query): - continue - score = 100 if timestamp_column: score += 40 if any( token in normalized_table_name for token in ( + "debug", + "entry", + "entries", "production", "event", "events", "repair", "manufacturing", - "factory", - "throughput", ) ): - score += 50 - if any( - token in normalized_table_name - for token in ("debug", "entry", "entries", "staging") - ): - score -= 80 + score += 20 if self._table_matches_query(table_name, query): score += 15 - if score > 0: - scored_unit_tables.append( - (score, candidate_table, unit_column, timestamp_column) - ) + scored_unit_tables.append( + (score, candidate_table, unit_column, timestamp_column) + ) table = None preferred_unit_column = None @@ -3060,8 +2992,6 @@ def _build_manufacturing_throughput_sql( has_business_unit = self._schema_contains( table_ddls, r"\bBusinessUnit\b", table_names=table_names ) - if not self._query_allows_internal_debug_table(query): - return None if not (has_debug_entries and has_business_unit): return None @@ -4847,36 +4777,30 @@ def _prune_sql_generation_context( score += 25 if wants_throughput_by_unit: - unit_column_candidates = [ - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "manufacturing unit", - "ManufacturingUnit", - "unit", - "unit_name", - "BU", - "division", - "station", - "workstation", - "work_center", - "workcenter", - "line", - "cell", - ] - if self._query_allows_internal_debug_table(query): - unit_column_candidates.extend( - ( - "Debug_Shelf", - "DebugShelf", - "debug shelf", - "shelf", - ) - ) unit_column = self._find_schema_column( table, - unit_column_candidates, + ( + "BusinessUnit", + "Business_Unit", + "Business Unit", + "manufacturing_unit", + "manufacturing unit", + "ManufacturingUnit", + "unit", + "unit_name", + "BU", + "division", + "Debug_Shelf", + "DebugShelf", + "debug shelf", + "shelf", + "station", + "workstation", + "work_center", + "workcenter", + "line", + "cell", + ), ) temporal_column = self._find_temporal_column_for_query(query, table) if unit_column: @@ -4885,10 +4809,6 @@ def _prune_sql_generation_context( score += 300 if wants_temporal_trend and unit_column and temporal_column: score += 300 - if self._is_internal_debug_table_name( - table_name - ) and not self._query_allows_internal_debug_table(query): - score -= 1200 if score > 0: scored.append((score, index)) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index bd4d02238a..c6c2494303 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1003,7 +1003,7 @@ def test_build_manufacturing_throughput_sql_uses_active_unit_and_date_columns(): ) -def test_build_manufacturing_throughput_sql_prefers_production_unit_table_over_admin_tables(): +def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( @@ -1022,13 +1022,6 @@ def test_build_manufacturing_throughput_sql_prefers_production_unit_table_over_a ); """, """ - CREATE TABLE dbo_production_events ( - id INTEGER, - manufacturing_unit VARCHAR, - created_at TIMESTAMP - ); - """, - """ CREATE TABLE dbo_DebugEntries ( DebugEntryId VARCHAR, BusinessUnit VARCHAR, @@ -1040,50 +1033,29 @@ def test_build_manufacturing_throughput_sql_prefers_production_unit_table_over_a ) assert sql == ( - 'SELECT "dbo_production_events"."manufacturing_unit" AS "manufacturing_unit", ' - 'DATEPART(YEAR, "dbo_production_events"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_production_events"."created_at") AS "month", ' + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "BusinessUnit", ' + 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "year", ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month", ' 'COUNT(*) AS "throughput" ' - 'FROM "dbo_production_events" ' - 'WHERE "dbo_production_events"."manufacturing_unit" IS NOT NULL ' - 'AND "dbo_production_events"."created_at" IS NOT NULL ' - 'GROUP BY "dbo_production_events"."manufacturing_unit", ' - 'DATEPART(YEAR, "dbo_production_events"."created_at"), ' - 'DATEPART(MONTH, "dbo_production_events"."created_at") ' - 'ORDER BY "dbo_production_events"."manufacturing_unit" ASC, ' - 'DATEPART(YEAR, "dbo_production_events"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_production_events"."created_at") ASC' + 'FROM "dbo_DebugEntries" ' + 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' + 'AND "dbo_DebugEntries"."DateIn" IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' + 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn"), ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ' + 'ORDER BY "dbo_DebugEntries"."BusinessUnit" ASC, ' + 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC, ' + 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' ) assert "dbo_ai_job_queue" not in sql assert "database_name" not in sql - assert "dbo_DebugEntries" not in sql - - -def test_build_manufacturing_throughput_sql_does_not_use_debug_shelf_for_manufacturing_units(): - service = AskService.__new__(AskService) - - sql = service._build_manufacturing_throughput_sql( - "Show throughput trends across different manufacturing units.", - [ - """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - id INTEGER, - Debug_Shelf VARCHAR, - last_update_date TIMESTAMP, - status VARCHAR - ); - """ - ], - ) - - assert sql is None -def test_build_manufacturing_throughput_sql_uses_debug_shelf_when_requested(): +def test_build_manufacturing_throughput_sql_uses_debug_shelf_unit_column(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( - "Show debug shelf throughput trends.", + "Show throughput trends across different manufacturing units.", [ """ CREATE TABLE dbo_DebugEntries_Staging2 ( @@ -1113,36 +1085,6 @@ def test_build_manufacturing_throughput_sql_uses_debug_shelf_when_requested(): ) -def test_build_validated_ask_result_rejects_debug_staging_for_manufacturing_units(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_DebugEntries_Staging2"."Debug_Shelf" AS "Debug_Shelf", ' - 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date") AS "year", ' - 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") AS "month", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries_Staging2" ' - 'WHERE "dbo_DebugEntries_Staging2"."Debug_Shelf" IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries_Staging2"."Debug_Shelf", ' - 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date"), ' - 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date")' - ), - [ - """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - id INTEGER, - Debug_Shelf VARCHAR, - last_update_date TIMESTAMP, - status VARCHAR - ); - """ - ], - "Show throughput trends across different manufacturing units.", - ) - - assert result is None - - def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): service = AskService.__new__(AskService) From 41ca796524d6e5a8624fdbe588f7725eb9aa987b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:37 +0530 Subject: [PATCH 0430/1087] Revert "Filter explicit table retrieval context" This reverts commit df3115b09728e35ef5a139a9f809a3971492d04a. --- wren-ai-service/src/web/v1/services/ask.py | 160 +++--------------- .../pytest/services/test_ask_sales_sql.py | 35 ---- 2 files changed, 23 insertions(+), 172 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 51a5ea877c..700ef86c78 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1165,65 +1165,6 @@ def _sql_references_explicit_table( ) return False - def _filter_retrieval_metadata_for_explicit_query( - self, - query: str, - documents: list[dict], - ) -> tuple[list[dict], list[str], list[str]]: - explicit_table_keys = self._explicit_table_alias_keys_from_query(query) - if not explicit_table_keys: - table_names = [ - table_name - for document in documents - if isinstance(table_name := document.get("table_name"), str) - and table_name.strip() - ] - table_ddls = [ - table_ddl - for document in documents - if isinstance(table_ddl := document.get("table_ddl"), str) - and table_ddl.strip() - ] - return documents, table_names, table_ddls - - matched_documents: list[dict] = [] - for document in documents: - candidate_names = [] - if isinstance(table_name := document.get("table_name"), str): - candidate_names.append(table_name) - if isinstance(table_ddl := document.get("table_ddl"), str): - candidate_names.extend( - str(table.get("name") or "") - for table in self._parse_schema_tables([table_ddl]) - if table.get("name") - ) - - candidate_keys: set[str] = set() - for candidate_name in candidate_names: - candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) - candidate_keys.update( - self._schema_identifier_alias_keys( - re.split(r"[.$_]", str(candidate_name or ""))[-1] - ) - ) - - if explicit_table_keys.intersection(candidate_keys): - matched_documents.append(document) - - table_names = [ - table_name - for document in matched_documents - if isinstance(table_name := document.get("table_name"), str) - and table_name.strip() - ] - table_ddls = [ - table_ddl - for document in matched_documents - if isinstance(table_ddl := document.get("table_ddl"), str) - and table_ddl.strip() - ] - return matched_documents, table_names, table_ddls - def _find_best_schema_table_for_query( self, query: str, tables: list[dict[str, Any]] ) -> dict[str, Any] | None: @@ -5215,41 +5156,6 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - ) - ) - if not documents: - logger.info( - "Explicit table retrieval did not return the requested table for query_id %s; " - "loading full active schema.", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval for explicit table", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - all_documents, _, _ = self._extract_retrieval_metadata( - retrieval_result - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - all_documents, - ) - ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) @@ -5317,27 +5223,27 @@ async def ask( user_query, table_ddls ) ): - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, + api_results = [ + AskResult( + **{ + "sql": deterministic_sql, + "type": "llm", + } ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = deterministic_sql + ] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if not documents: error_message = ( @@ -5798,17 +5704,10 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - explicit_table_names = self._extract_explicit_table_names_from_query( - user_query - ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - ) - ) if not documents: + explicit_table_names = self._extract_explicit_table_names_from_query( + user_query + ) if explicit_table_names: logger.info( "Retrying schema retrieval for explicit tables query_id %s: %s", @@ -5835,12 +5734,6 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - ) - ) if not documents and self._is_data_analysis_query(user_query): logger.info( "Query-based schema retrieval returned no tables for data question; " @@ -5867,13 +5760,6 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - ) - ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index c6c2494303..ea40d53f44 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -231,41 +231,6 @@ def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): assert result is not None -def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): - service = AskService.__new__(AskService) - documents = [ - { - "table_name": "dbo_knowledge_articles", - "table_ddl": """ - CREATE TABLE dbo_knowledge_articles ( - id INTEGER, - last_run_date TIMESTAMP - ); - """, - }, - { - "table_name": "dbo_failure_patterns", - "table_ddl": """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - created_at TIMESTAMP - ); - """, - }, - ] - - filtered_documents, table_names, table_ddls = ( - service._filter_retrieval_metadata_for_explicit_query( - "show monthly record count by created_at in dbo.failure_patterns.", - documents, - ) - ) - - assert filtered_documents == [documents[1]] - assert table_names == ["dbo_failure_patterns"] - assert table_ddls == [documents[1]["table_ddl"]] - - def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) From a0be718a8b83eb800d69193e34abfc00e3c01051 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:38 +0530 Subject: [PATCH 0431/1087] Revert "Reject SQL for wrong explicit table" This reverts commit 90808d695c1a090091a35d39f8f5f9f324273847. --- wren-ai-service/src/web/v1/services/ask.py | 44 -------------- .../pytest/services/test_ask_sales_sql.py | 59 ------------------- 2 files changed, 103 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 700ef86c78..d4fb128985 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1124,47 +1124,6 @@ def _table_matches_query(self, table_name: str, query: str) -> bool: for table_key in table_keys ) - def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: - keys: set[str] = set() - for table_name in self._extract_explicit_table_names_from_query(query or ""): - keys.update(self._schema_identifier_alias_keys(table_name)) - return keys - - def _sql_references_explicit_table( - self, - sql: str, - query: str | None, - ) -> bool: - explicit_table_keys = self._explicit_table_alias_keys_from_query(query) - if not explicit_table_keys: - return True - - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_tables = [ - next(value for value in match.groupdict().values() if value) - for match in table_reference_pattern.finditer(sql) - ] - - for table_reference in referenced_tables: - reference_keys = self._schema_identifier_alias_keys(table_reference) - short_reference = re.split(r"[.$_]", str(table_reference or ""))[-1] - reference_keys.update(self._schema_identifier_alias_keys(short_reference)) - if explicit_table_keys.intersection(reference_keys): - return True - - logger.warning( - "Ignoring SQL because it does not reference the explicitly requested table. " - "query=%s referenced_tables=%s sql=%s", - query, - referenced_tables, - sql, - ) - return False - def _find_best_schema_table_for_query( self, query: str, tables: list[dict[str, Any]] ) -> dict[str, Any] | None: @@ -4910,9 +4869,6 @@ def _build_validated_ask_result_from_sql( ) return None - if not self._sql_references_explicit_table(ask_result.sql, query): - return None - if not self._sql_matches_question_intent( ask_result.sql, query, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index ea40d53f44..cda21fd212 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -172,65 +172,6 @@ def test_extract_explicit_table_names_from_using_clause(): ) == [] -def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date") AS "year", ' - 'DATEPART(MONTH, "dbo_knowledge_articles"."last_run_date") AS "month", ' - '"dbo_knowledge_articles"."category" AS "category", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_knowledge_articles" ' - 'GROUP BY DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date"), ' - 'DATEPART(MONTH, "dbo_knowledge_articles"."last_run_date"), ' - '"dbo_knowledge_articles"."category"' - ), - [ - """ - CREATE TABLE dbo_knowledge_articles ( - id INTEGER, - last_run_date TIMESTAMP, - category VARCHAR - ); - """, - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - created_at TIMESTAMP - ); - """, - ], - "show monthly record count by created_at in dbo.failure_patterns.", - ) - - assert result is None - - -def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_failure_patterns" ' - 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at")' - ), - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - created_at TIMESTAMP - ); - """ - ], - "show monthly record count by created_at in dbo.failure_patterns.", - ) - - assert result is not None - - def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) From 15ae2ef5d8687019915ef22cccf8e144bd5f2f09 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:38 +0530 Subject: [PATCH 0432/1087] Revert "Route monthly record counts through schema SQL" This reverts commit e8c75256b5bbd020733aeffe0ed51edfc80cf665. --- wren-ai-service/src/web/v1/services/ask.py | 8 ----- .../pytest/services/test_ask_sales_sql.py | 30 ------------------- 2 files changed, 38 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d4fb128985..f9247773f9 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4082,14 +4082,6 @@ def _build_schema_grounded_sales_sql( "trend", "throughput", "manufacturing", - "monthly", - "month", - "count", - "counts", - "record", - "records", - "row", - "rows", "unit", "units", "value", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index cda21fd212..7930c553ae 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1518,36 +1518,6 @@ def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count() ) -def test_build_schema_grounded_sales_sql_handles_how_monthly_record_count(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "how monthly record count by created_at in dbo.failure_patterns.", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ASC' - ) - assert '"dbo_failure"."patterns"' not in sql - - def test_build_schema_grounded_table_question_sql_prefers_explicit_table_alias(): service = AskService.__new__(AskService) From 61c3f0855dbc227e7ffed477ab103c9a76f17856 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:38 +0530 Subject: [PATCH 0433/1087] Revert "Keep explicit table analytics out of preview fallback" This reverts commit 7a0a9a5c1393ab4756ed4bf46b3162d88dbcfa3e. --- wren-ai-service/src/web/v1/services/ask.py | 23 -------- .../pytest/services/test_ask_sales_sql.py | 55 ------------------- 2 files changed, 78 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f9247773f9..096b441aef 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1219,23 +1219,6 @@ def _build_schema_grounded_table_question_sql( return None tables = self._parse_schema_tables(table_ddls) - explicit_table_names = self._extract_explicit_table_names_from_query(query) - if explicit_table_names: - explicit_keys = { - key - for explicit_table_name in explicit_table_names - for key in self._schema_identifier_alias_keys(explicit_table_name) - } - explicit_tables = [ - table - for table in tables - if explicit_keys.intersection( - self._schema_identifier_alias_keys(str(table.get("name") or "")) - ) - ] - if explicit_tables: - tables = explicit_tables - table = self._find_best_schema_table_for_query(query, tables) if not table: return None @@ -1355,12 +1338,6 @@ def _build_explicit_table_preview_sql( r"\b(?:rows?|records?|data)\b", normalized_query, flags=re.IGNORECASE ): return None - if re.search( - r"\b(?:count|counts|monthly|month|trend|trends|by|per|each|distribution|group(?:ed)?|aggregate)\b", - normalized_query, - flags=re.IGNORECASE, - ): - return None tables = self._parse_schema_tables(table_ddls) if not tables: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 7930c553ae..699cca6a3e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -127,24 +127,6 @@ def test_build_explicit_table_preview_sql_for_show_data_prompt(): assert result == ('SELECT TOP 10 * FROM "CustomerMaster"', "CustomerMaster") -def test_build_explicit_table_preview_sql_ignores_monthly_count_question(): - service = AskService.__new__(AskService) - - result = service._build_explicit_table_preview_sql( - "Show monthly record count by created_at in dbo.failure_patterns.", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - created_at TIMESTAMP - ); - """ - ], - ) - - assert result is None - - def test_extract_explicit_table_names_from_query(): service = AskService.__new__(AskService) @@ -1518,43 +1500,6 @@ def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count() ) -def test_build_schema_grounded_table_question_sql_prefers_explicit_table_alias(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Show monthly record count by created_at in dbo.failure_patterns.", - [ - """ - CREATE TABLE dbo_knowledge_articles ( - id INTEGER, - policy_category_id INTEGER, - created_at TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR, - created_at TIMESTAMP - ); - """, - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ASC' - ) - assert "dbo_knowledge_articles" not in sql - - def test_build_validated_ask_result_rejects_status_for_product_line_pcb_question(): service = AskService.__new__(AskService) From 85fb51fa75164faa210c9647b97e45f2f8806685 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:38 +0530 Subject: [PATCH 0434/1087] Revert "Use throughput builder for manufacturing unit trends" This reverts commit 0053bcf38a101722f8d9137c46f3636fbc399621. --- wren-ai-service/src/web/v1/services/ask.py | 29 ---------------- .../pytest/services/test_ask_sales_sql.py | 34 ------------------- 2 files changed, 63 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 096b441aef..196095eeb6 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1733,11 +1733,6 @@ def _build_schema_grounded_analytics_sql( ): return monthly_repair_volume_sql - if throughput_sql := self._build_manufacturing_throughput_sql( - query, table_ddls - ): - return throughput_sql - if operational_sql := self._build_schema_grounded_operational_sql( query, tables ): @@ -2760,16 +2755,6 @@ def _build_manufacturing_throughput_sql( "unit_name", "BU", "division", - "Debug_Shelf", - "DebugShelf", - "debug shelf", - "shelf", - "station", - "workstation", - "work_center", - "workcenter", - "line", - "cell", ) wants_temporal_trend = any( term in normalized for term in ("trend", "trends", "monthly", "over time") @@ -4057,10 +4042,6 @@ def _build_schema_grounded_sales_sql( "salesperson", "sales person", "trend", - "throughput", - "manufacturing", - "unit", - "units", "value", ) ): @@ -4659,16 +4640,6 @@ def _prune_sql_generation_context( "unit_name", "BU", "division", - "Debug_Shelf", - "DebugShelf", - "debug shelf", - "shelf", - "station", - "workstation", - "work_center", - "workcenter", - "line", - "cell", ), ) temporal_column = self._find_temporal_column_for_query(query, table) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 699cca6a3e..6ebd1d3653 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -939,40 +939,6 @@ def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables assert "database_name" not in sql -def test_build_manufacturing_throughput_sql_uses_debug_shelf_unit_column(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Show throughput trends across different manufacturing units.", - [ - """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - id INTEGER, - Debug_Shelf VARCHAR, - last_update_date TIMESTAMP, - status VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_DebugEntries_Staging2"."Debug_Shelf" AS "Debug_Shelf", ' - 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date") AS "year", ' - 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") AS "month", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries_Staging2" ' - 'WHERE "dbo_DebugEntries_Staging2"."Debug_Shelf" IS NOT NULL ' - 'AND "dbo_DebugEntries_Staging2"."last_update_date" IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries_Staging2"."Debug_Shelf", ' - 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date"), ' - 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") ' - 'ORDER BY "dbo_DebugEntries_Staging2"."Debug_Shelf" ASC, ' - 'DATEPART(YEAR, "dbo_DebugEntries_Staging2"."last_update_date") ASC, ' - 'DATEPART(MONTH, "dbo_DebugEntries_Staging2"."last_update_date") ASC' - ) - - def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): service = AskService.__new__(AskService) From ecf1dd98bf6c4cdd582d23b6e2ea62df4001c9ee Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:38 +0530 Subject: [PATCH 0435/1087] Revert "Validate unqualified SQL columns against metadata" This reverts commit 259fdc1cf986814fa02fda0e77c4bef8da04afba. --- .../src/pipelines/generation/utils/sql.py | 152 ----------------- .../pipelines/generation/test_sql_utils.py | 63 ------- .../apollo/server/services/queryService.ts | 159 ------------------ .../services/tests/queryService.test.ts | 110 ------------ 4 files changed, 484 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index eb5498ea44..0c1f7e58fa 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2975,148 +2975,6 @@ def replace_unqualified_identifier(match: re.Match[str]) -> str: ) -_SQL_IDENTIFIER_KEYWORDS = { - "all", - "and", - "as", - "asc", - "avg", - "between", - "by", - "case", - "cast", - "coalesce", - "count", - "current_date", - "date", - "dateadd", - "datediff", - "datepart", - "day", - "desc", - "distinct", - "else", - "end", - "extract", - "false", - "from", - "group", - "having", - "hour", - "in", - "is", - "join", - "last", - "like", - "limit", - "minute", - "month", - "not", - "null", - "nulls", - "on", - "or", - "order", - "over", - "partition", - "quarter", - "second", - "select", - "sum", - "then", - "top", - "true", - "when", - "week", - "where", - "with", - "year", -} - - -def _extract_projection_aliases(sql: str) -> set[str]: - aliases: set[str] = set() - for start, end in _find_select_list_spans(sql): - for item in _split_top_level_select_items(sql[start:end]): - alias_match = re.search( - rf"\s+(?:AS\s+)?(?P{_SQL_IDENTIFIER_PATTERN})\s*$", - item, - flags=re.IGNORECASE, - ) - if not alias_match: - continue - expression = item[: alias_match.start()].strip() - alias = _normalize_sql_identifier(alias_match.group("alias")) - if expression and alias: - aliases.add(alias.lower()) - return aliases - - -def _find_invalid_unqualified_column_references_for_single_table( - sql: str, - table_name: str, - valid_columns: set[str], - valid_compact_columns: set[str], - aliases: dict[str, str], -) -> list[str]: - table_aliases = { - alias.lower() - for alias, alias_table in aliases.items() - if str(alias_table).lower() == str(table_name).lower() - } - table_aliases.update( - suffix.lower() for suffix in _table_reference_suffixes(str(table_name)) - ) - projection_aliases = _extract_projection_aliases(sql) - sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") - invalid_references: list[str] = [] - - identifier_pattern = re.compile( - r'"(?P[^"]+)"|`(?P[^`]+)`|\[(?P[^\]]+)\]|(?P\b[A-Za-z_][A-Za-z0-9_$]*\b)' - ) - for match in identifier_pattern.finditer(sql_without_strings): - identifier = ( - match.group("quoted") - or match.group("backticked") - or match.group("bracketed") - or match.group("bare") - or "" - ) - normalized_identifier = identifier.lower() - if not normalized_identifier: - continue - - before = sql_without_strings[: match.start()].rstrip() - after = sql_without_strings[match.end() :].lstrip() - if before.endswith(".") or after.startswith("."): - continue - if after.startswith("("): - continue - - previous_word_match = re.search(r"([A-Za-z_][A-Za-z0-9_$]*)\s*$", before) - previous_word = ( - previous_word_match.group(1).lower() if previous_word_match else "" - ) - if previous_word == "as": - continue - - if ( - normalized_identifier in _SQL_IDENTIFIER_KEYWORDS - or normalized_identifier in table_aliases - or normalized_identifier in projection_aliases - ): - continue - if ( - normalized_identifier in valid_columns - or _compact_sql_identifier(identifier) in valid_compact_columns - ): - continue - if identifier not in invalid_references: - invalid_references.append(identifier) - - return invalid_references - - def find_invalid_column_references( sql: str, valid_table_columns: dict[str, list[str]] ) -> list[str]: @@ -3180,16 +3038,6 @@ def find_invalid_column_references( ): invalid_references.append(column) - invalid_references.extend( - _find_invalid_unqualified_column_references_for_single_table( - sql, - str(table_name), - valid_columns, - valid_compact_columns, - aliases, - ) - ) - return sorted(set(invalid_references)) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 47cc307e55..421566f5eb 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -469,69 +469,6 @@ def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid ) == ["dbo_qSales1.UnitPrice"] -def test_find_invalid_column_references_rejects_unqualified_function_arguments(): - sql = ( - 'SELECT user_id customer, COUNT(*) issue_count, ' - 'AVG(DATEDIFF(day, created_at, resolved_at)) avg_resolution_time ' - 'FROM "dbo_kb_article_feedback" ' - "WHERE details LIKE '%device%' AND helpful = 0 " - "GROUP BY user_id " - "ORDER BY issue_count DESC NULLS LAST" - ) - - assert find_invalid_column_references( - sql, - { - "dbo_kb_article_feedback": [ - "detail_id", - "history_id", - "target_query_expression", - "execution_date", - "result", - "result_detail", - "exception_message", - "exception", - "id", - "org_id", - "article_id", - "user_id", - "helpful", - "reasons", - "details", - "created_at", - ] - }, - ) == ["resolved_at"] - - -def test_find_invalid_column_references_rejects_unqualified_where_columns(): - sql = ( - 'SELECT COUNT(*) AS "count" ' - 'FROM "dbo_ai_workflows" ' - "WHERE role = 'admin' " - "GROUP BY debug_user" - ) - - assert find_invalid_column_references( - sql, - { - "dbo_ai_workflows": [ - "id", - "org_id", - "repair_id", - "workflow_name", - "priority", - "steps", - "estimated_total_min", - "source_inspection_id", - "created_by_user_id", - "created_at", - "updated_at", - ] - }, - ) == ["debug_user", "role"] - - def test_normalize_sql_column_references_to_schema_maps_kb_article_aliases(): sql = ( 'SELECT "dbo_kb_articles"."article_type", COUNT(*) AS "RecordCount" ' diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index ce64720615..09c3734120 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -600,158 +600,6 @@ const extractSimpleProjectionColumns = (sql: string) => { return columns; }; -const SQL_IDENTIFIER_KEYWORDS = new Set([ - 'all', - 'and', - 'as', - 'asc', - 'avg', - 'between', - 'by', - 'case', - 'cast', - 'coalesce', - 'count', - 'current_date', - 'date', - 'dateadd', - 'datediff', - 'datepart', - 'day', - 'desc', - 'distinct', - 'else', - 'end', - 'extract', - 'false', - 'from', - 'group', - 'having', - 'hour', - 'in', - 'is', - 'join', - 'last', - 'like', - 'limit', - 'minute', - 'month', - 'not', - 'null', - 'nulls', - 'on', - 'or', - 'order', - 'over', - 'partition', - 'quarter', - 'second', - 'select', - 'sum', - 'then', - 'top', - 'true', - 'when', - 'week', - 'where', - 'with', - 'year', -]); - -const extractProjectionAliases = (sql: string) => { - const aliases = new Set(); - const selectPattern = /\bSELECT\b(?.*?)(?=\bFROM\b)/gis; - let match: RegExpExecArray | null; - while ((match = selectPattern.exec(sql))) { - const body = match.groups?.body || ''; - splitTopLevelSqlList(body).forEach((item) => { - const aliasMatch = item.match( - new RegExp( - String.raw`\s+(?:AS\s+)?(?${SQL_IDENTIFIER_PATTERN})\s*$`, - 'i', - ), - ); - if (!aliasMatch?.groups?.alias || aliasMatch.index === undefined) { - return; - } - const expression = item.slice(0, aliasMatch.index).trim(); - if (expression) { - aliases.add( - normalizeSqlIdentifier(aliasMatch.groups.alias).toLowerCase(), - ); - } - }); - } - return aliases; -}; - -const findInvalidUnqualifiedColumnReferencesForSingleSchema = ( - sql: string, - schema: ManifestModelSchema, - aliases: Map, -) => { - const tableAliases = new Set(); - aliases.forEach((aliasSchema, alias) => { - if (aliasSchema === schema) { - tableAliases.add(alias.toLowerCase()); - } - }); - const projectionAliases = extractProjectionAliases(sql); - const sqlWithoutStrings = sql.replace(/'(?:''|[^'])*'/g, "''"); - const invalidReferences: string[] = []; - const identifierPattern = - /"(?[^"]+)"|`(?[^`]+)`|\[(?[^\]]+)\]|(?\b[A-Za-z_][A-Za-z0-9_$]*\b)/g; - - let match: RegExpExecArray | null; - while ((match = identifierPattern.exec(sqlWithoutStrings))) { - const identifier = - match.groups?.quoted || - match.groups?.backticked || - match.groups?.bracketed || - match.groups?.bare || - ''; - const normalizedIdentifier = identifier.toLowerCase(); - if (!normalizedIdentifier) { - continue; - } - - const before = sqlWithoutStrings.slice(0, match.index).trimEnd(); - const after = sqlWithoutStrings.slice(match.index + match[0].length).trimStart(); - if (before.endsWith('.') || after.startsWith('.')) { - continue; - } - if (after.startsWith('(')) { - continue; - } - - const previousWord = before - .match(/([A-Za-z_][A-Za-z0-9_$]*)\s*$/)?.[1] - ?.toLowerCase(); - if (previousWord === 'as') { - continue; - } - - if ( - SQL_IDENTIFIER_KEYWORDS.has(normalizedIdentifier) || - tableAliases.has(normalizedIdentifier) || - projectionAliases.has(normalizedIdentifier) - ) { - continue; - } - if ( - schema.columns.has(normalizedIdentifier) || - schema.columns.has(compactSqlIdentifier(identifier)) - ) { - continue; - } - if (!invalidReferences.includes(identifier)) { - invalidReferences.push(identifier); - } - } - - return invalidReferences; -}; - const findSqlReferenceValidationErrors = ( sql: string, manifest?: Manifest, @@ -836,13 +684,6 @@ const findSqlReferenceValidationErrors = ( errors.push(column); } }); - errors.push( - ...findInvalidUnqualifiedColumnReferencesForSingleSchema( - sql, - schema, - aliases, - ), - ); } return [...new Set(errors)]; diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 8b98c6f378..8aa5e6cd58 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -457,116 +457,6 @@ describe('QueryService', () => { expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); }); - it('should reject unqualified function arguments outside a single active manifest table', async () => { - await expect( - queryService.preview( - [ - 'SELECT user_id customer, COUNT(*) issue_count,', - 'AVG(DATEDIFF(day, created_at, resolved_at)) avg_resolution_time', - 'FROM "dbo_kb_article_feedback"', - "WHERE details LIKE '%device%' AND helpful = 0", - 'GROUP BY user_id', - 'ORDER BY issue_count DESC NULLS LAST', - ].join(' '), - { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_kb_article_feedback', - tableReference: { table: 'dbo_kb_article_feedback' }, - columns: [ - { name: 'detail_id', type: 'string', isCalculated: false }, - { name: 'history_id', type: 'string', isCalculated: false }, - { - name: 'target_query_expression', - type: 'string', - isCalculated: false, - }, - { - name: 'execution_date', - type: 'timestamp', - isCalculated: false, - }, - { name: 'result', type: 'string', isCalculated: false }, - { name: 'result_detail', type: 'string', isCalculated: false }, - { - name: 'exception_message', - type: 'string', - isCalculated: false, - }, - { name: 'exception', type: 'string', isCalculated: false }, - { name: 'id', type: 'string', isCalculated: false }, - { name: 'org_id', type: 'string', isCalculated: false }, - { name: 'article_id', type: 'string', isCalculated: false }, - { name: 'user_id', type: 'string', isCalculated: false }, - { name: 'helpful', type: 'integer', isCalculated: false }, - { name: 'reasons', type: 'string', isCalculated: false }, - { name: 'details', type: 'string', isCalculated: false }, - { name: 'created_at', type: 'timestamp', isCalculated: false }, - ], - }, - ], - }, - dryRun: true, - }, - ), - ).rejects.toThrow( - 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: resolved_at', - ); - - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); - }); - - it('should reject unqualified where and group columns outside a single active manifest table', async () => { - await expect( - queryService.preview( - 'SELECT COUNT(*) AS "count" FROM "dbo_ai_workflows" WHERE role = \'admin\' GROUP BY debug_user', - { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_ai_workflows', - tableReference: { table: 'dbo_ai_workflows' }, - columns: [ - { name: 'id', type: 'string', isCalculated: false }, - { name: 'org_id', type: 'string', isCalculated: false }, - { name: 'repair_id', type: 'string', isCalculated: false }, - { name: 'workflow_name', type: 'string', isCalculated: false }, - { name: 'priority', type: 'string', isCalculated: false }, - { name: 'steps', type: 'string', isCalculated: false }, - { - name: 'estimated_total_min', - type: 'integer', - isCalculated: false, - }, - { - name: 'source_inspection_id', - type: 'string', - isCalculated: false, - }, - { - name: 'created_by_user_id', - type: 'string', - isCalculated: false, - }, - { name: 'created_at', type: 'timestamp', isCalculated: false }, - { name: 'updated_at', type: 'timestamp', isCalculated: false }, - ], - }, - ], - }, - dryRun: true, - }, - ), - ).rejects.toThrow( - 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: debug_user, role', - ); - - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); - }); - it('should reject numeric aggregates on non-numeric manifest columns before ibis planning', async () => { await expect( queryService.preview('SELECT AVG("orders"."quantity") FROM "orders"', { From c10fe83fa159743611786c4d77945516eeaf9fe8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:39 +0530 Subject: [PATCH 0436/1087] Revert "Prefer unit tables for throughput trends" This reverts commit 7ee464b9ab0e6632ceb157f321665d16836861c8. --- wren-ai-service/src/web/v1/services/ask.py | 129 +++--------------- .../pytest/services/test_ask_sales_sql.py | 48 ------- 2 files changed, 17 insertions(+), 160 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 196095eeb6..bd73776519 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2743,83 +2743,30 @@ def _build_manufacturing_throughput_sql( return None tables = self._parse_schema_tables(table_ddls) - unit_candidates = ( - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "manufacturing unit", - "manufacturingunit", - "ManufacturingUnit", - "unit", - "unit_name", - "BU", - "division", - ) - wants_temporal_trend = any( - term in normalized for term in ("trend", "trends", "monthly", "over time") - ) - scored_unit_tables: list[tuple[int, dict[str, Any], str, str | None]] = [] - for candidate_table in tables: - unit_column = self._find_schema_column(candidate_table, unit_candidates) - if not unit_column: - continue - - timestamp_column = self._find_temporal_column_for_query( - query, candidate_table - ) - if wants_temporal_trend and not timestamp_column: - continue - - table_name = str(candidate_table.get("name") or "") - normalized_table_name = table_name.lower() - score = 100 - if timestamp_column: - score += 40 - if any( - token in normalized_table_name - for token in ( - "debug", - "entry", - "entries", - "production", - "event", - "events", - "repair", - "manufacturing", - ) - ): - score += 20 - if self._table_matches_query(table_name, query): - score += 15 - scored_unit_tables.append( - (score, candidate_table, unit_column, timestamp_column) - ) - - table = None - preferred_unit_column = None - preferred_timestamp_column = None - if scored_unit_tables: - _, table, preferred_unit_column, preferred_timestamp_column = sorted( - scored_unit_tables, key=lambda item: item[0], reverse=True - )[0] - else: - table = self._find_best_schema_table_for_query(query, tables) - + table = self._find_best_schema_table_for_query(query, tables) if table: - unit_column = preferred_unit_column or self._find_schema_column( - table, unit_candidates + unit_column = self._find_schema_column( + table, + ( + "BusinessUnit", + "business_unit", + "manufacturing_unit", + "manufacturingunit", + "unit", + "unit_name", + "BU", + "division", + ), ) if unit_column: table_name = str(table.get("name") or "") table_ref = self._quote_sql_identifier(table_name) unit_ref = f"{table_ref}.{self._quote_sql_identifier(unit_column)}" - timestamp_column = ( - preferred_timestamp_column - or self._find_temporal_column_for_query(query, table) - ) + timestamp_column = self._find_temporal_column_for_query(query, table) - if timestamp_column and wants_temporal_trend: + if timestamp_column and any( + term in normalized for term in ("trend", "monthly", "over time") + ): timestamp_ref = ( f"{table_ref}.{self._quote_sql_identifier(timestamp_column)}" ) @@ -4574,24 +4521,6 @@ def _prune_sql_generation_context( self._normalize_schema_token(table_name) for table_name in self._extract_explicit_table_names_from_query(query) } - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - wants_throughput_by_unit = "throughput" in normalized_query and any( - term in normalized_query - for term in ( - "manufacturing unit", - "manufacturing units", - "business unit", - "business units", - "different unit", - "different units", - "unit", - "units", - ) - ) - wants_temporal_trend = any( - term in normalized_query - for term in ("trend", "trends", "monthly", "over time", "by month") - ) scored: list[tuple[int, int]] = [] for index, table in enumerate(parsed_tables): @@ -4626,30 +4555,6 @@ def _prune_sql_generation_context( elif term in column_term or column_term in term: score += 25 - if wants_throughput_by_unit: - unit_column = self._find_schema_column( - table, - ( - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "manufacturing unit", - "ManufacturingUnit", - "unit", - "unit_name", - "BU", - "division", - ), - ) - temporal_column = self._find_temporal_column_for_query(query, table) - if unit_column: - score += 500 - if unit_column and temporal_column: - score += 300 - if wants_temporal_trend and unit_column and temporal_column: - score += 300 - if score > 0: scored.append((score, index)) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 6ebd1d3653..569ab4f383 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -891,54 +891,6 @@ def test_build_manufacturing_throughput_sql_uses_active_unit_and_date_columns(): ) -def test_build_manufacturing_throughput_sql_prefers_unit_table_over_admin_tables(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Show throughput trends across different manufacturing units.", - [ - """ - CREATE TABLE dbo_ai_job_queue ( - database_name VARCHAR, - status VARCHAR, - payload VARCHAR, - error VARCHAR, - started_at TIMESTAMP, - completed_at TIMESTAMP, - created_at TIMESTAMP, - updated_at TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_DebugEntries ( - DebugEntryId VARCHAR, - BusinessUnit VARCHAR, - DateIn TIMESTAMP, - Status VARCHAR - ); - """, - ], - ) - - assert sql == ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "BusinessUnit", ' - 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "year", ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' - 'AND "dbo_DebugEntries"."DateIn" IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' - 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn"), ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ' - 'ORDER BY "dbo_DebugEntries"."BusinessUnit" ASC, ' - 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC, ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' - ) - assert "dbo_ai_job_queue" not in sql - assert "database_name" not in sql - - def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): service = AskService.__new__(AskService) From bd7234da38cbdddf32acf8eec0da07a23d78f04f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:39 +0530 Subject: [PATCH 0437/1087] Revert "Reject columns outside referenced table schema" This reverts commit c83aa0a37b26e0e5ad3f726c2c5500dd648ac8d7. --- wren-ai-service/src/web/v1/services/ask.py | 40 +------------------ .../pytest/services/test_ask_sales_sql.py | 32 --------------- 2 files changed, 1 insertion(+), 71 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bd73776519..dbd172442c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -678,50 +678,12 @@ def _sql_covers_required_question_concepts( def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: - all_valid_columns = { + valid_columns = { str(column.get("name") or "").lower() for table in schema_tables for column in table.get("columns", []) if column.get("name") } - columns_by_table: dict[str, set[str]] = {} - for table in schema_tables: - table_name = str(table.get("name") or "").lower() - if not table_name: - continue - table_columns = { - str(column.get("name") or "").lower() - for column in table.get("columns", []) - if column.get("name") - } - columns_by_table[table_name] = table_columns - columns_by_table[table_name.split(".")[-1]] = table_columns - - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_table_keys = { - ( - next(value for value in match.groupdict().values() if value) or "" - ).lower() - for match in table_reference_pattern.finditer(sql or "") - } - referenced_column_sets = [ - columns - for table_key in referenced_table_keys - for columns in [ - columns_by_table.get(table_key) - or columns_by_table.get(table_key.split(".")[-1]) - ] - if columns is not None - ] - valid_columns = ( - referenced_column_sets[0] - if len(referenced_column_sets) == 1 - else all_valid_columns - ) valid_tables = { str(table.get("name") or "").lower() for table in schema_tables diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 569ab4f383..2893956565 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1546,38 +1546,6 @@ def test_build_validated_ask_result_rejects_unqualified_invalid_columns(): ) -def test_build_validated_ask_result_rejects_column_from_different_active_table(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT policy_category_id AS "policy_category_id", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_knowledge_articles" ' - "GROUP BY policy_category_id" - ), - [ - """ - CREATE TABLE dbo_knowledge_articles ( - id VARCHAR, - policy_id VARCHAR, - category VARCHAR, - created_at TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_policies ( - id VARCHAR, - policy_category_id VARCHAR - ); - """, - ], - "Show knowledge article count by policy category.", - ) - - assert result is None - - def test_build_validated_ask_result_accepts_unqualified_valid_columns(): service = AskService.__new__(AskService) From 9215ca225b37e68ca961bf2e88367b91945b16cc Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 18:20:39 +0530 Subject: [PATCH 0438/1087] Revert "Fix explicit table recommendation asks" This reverts commit ffbc00e4d0ed23dabc80bc17f1d28d94e8397bce. --- wren-ai-service/src/web/v1/services/ask.py | 26 +++---------------- .../pytest/services/test_ask_sales_sql.py | 8 ------ .../apollo/server/services/askingService.ts | 2 +- .../services/tests/askingService.test.ts | 26 ------------------- 4 files changed, 5 insertions(+), 57 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index dbd172442c..098bd44f4f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1208,7 +1208,7 @@ def _build_schema_grounded_table_question_sql( wants_monthly_count = any( term in normalized for term in ("monthly", "by month", "per month", "month-wise") - ) and any(term in normalized for term in ("count", "record", "records", "rows")) + ) and any(term in normalized for term in ("count", "records", "rows")) if wants_monthly_count: date_column = self._find_temporal_column_for_query(query, table) if not date_column: @@ -1345,14 +1345,13 @@ def _build_explicit_table_preview_sql( def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names: list[str] = [] for match in re.finditer( - r"\b(?:from|in|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - for candidate in self._explicit_table_name_candidates(table_name): - if candidate and candidate not in table_names: - table_names.append(candidate) + if table_name and table_name not in table_names: + table_names.append(table_name) for match in re.finditer( r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", @@ -1367,23 +1366,6 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names.append(table_name) return table_names - def _explicit_table_name_candidates(self, table_name: str) -> list[str]: - table_name = str(table_name or "").strip(".,;:()[]{}") - if not table_name: - return [] - - candidates = [table_name] - dotted_parts = [part for part in re.split(r"[.$]", table_name) if part] - if len(dotted_parts) > 1: - candidates.append("_".join(dotted_parts)) - candidates.append(dotted_parts[-1]) - - unique_candidates: list[str] = [] - for candidate in candidates: - if candidate and candidate not in unique_candidates: - unique_candidates.append(candidate) - return unique_candidates - def _build_direct_orders_sales_sql(self, query: str) -> str | None: return None diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2893956565..e3abde7150 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -135,14 +135,6 @@ def test_extract_explicit_table_names_from_query(): ) == ["tblNewOrders"] -def test_extract_explicit_table_names_from_in_clause_adds_deployed_table_candidate(): - service = AskService.__new__(AskService) - - assert service._extract_explicit_table_names_from_query( - "Show monthly record count by created_at in dbo.failure_patterns." - ) == ["dbo.failure_patterns", "dbo_failure_patterns", "failure_patterns"] - - def test_extract_explicit_table_names_from_using_clause(): service = AskService.__new__(AskService) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index fe7025d9c9..f5b1115cae 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1382,7 +1382,7 @@ export class AskingService implements IAskingService { input: AskingDetailTaskInput, payload: AskingPayload, ): Promise { - if (input.trackedAskingResult || !input.question) { + if (input.trackedAskingResult || !input.sql || !input.question) { return input.trackedAskingResult; } diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index 55c01871c7..d33c9e166d 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -230,31 +230,5 @@ describe('AskingService', () => { askingTaskId: trackedAskingResult.taskId, }); }); - - test('also creates a normal asking task when direct question payload has no SQL', async () => { - const service = createService(); - - await service.createThreadResponse( - { - question: trackedAskingResult.question, - }, - 7, - ); - - expect(service.askingTaskTracker.createAskingTask).toHaveBeenCalledWith( - expect.objectContaining({ - query: trackedAskingResult.question, - histories: null, - deployId: 'latest-deploy-hash', - projectId: '1', - }), - ); - expect(service.threadResponseRepository.createOne).toHaveBeenCalledWith({ - threadId: 7, - question: trackedAskingResult.question, - sql: undefined, - askingTaskId: trackedAskingResult.taskId, - }); - }); }); }); From f05f172494e3e5f625550b59ddc6917219c6aa31 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 8 Jul 2026 23:36:42 +0530 Subject: [PATCH 0439/1087] Validate asks against active schema context --- wren-ai-service/src/web/v1/services/ask.py | 276 +++++++++++++----- .../test_ask_heuristic_text_to_sql.py | 4 +- .../pytest/services/test_ask_sales_sql.py | 93 ++++++ 3 files changed, 299 insertions(+), 74 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 098bd44f4f..2e95e4f91e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -348,7 +348,7 @@ def _should_reuse_historical_question_sql( query: str, histories: list[AskHistory] | None, ) -> bool: - return bool(histories) and self._needs_conversation_context(query) + return False def _rewrite_query_for_text_to_sql(self, query: str) -> str: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) @@ -1076,6 +1076,86 @@ def _schema_identifier_alias_keys(self, value: str) -> set[str]: if key } + def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: + keys: set[str] = set() + for table_name in self._extract_explicit_table_names_from_query(query or ""): + keys.update(self._schema_identifier_alias_keys(table_name)) + return keys + + def _filter_retrieval_metadata_for_explicit_query( + self, + query: str, + documents: list[dict], + ) -> tuple[list[dict], list[str], list[str]]: + explicit_table_keys = self._explicit_table_alias_keys_from_query(query) + if not explicit_table_keys: + table_names, table_ddls = self._metadata_from_documents(documents) + return documents, table_names, table_ddls + + matched_documents: list[dict] = [] + for document in documents: + candidate_names = [] + if isinstance(table_name := document.get("table_name"), str): + candidate_names.append(table_name) + if isinstance(table_ddl := document.get("table_ddl"), str): + candidate_names.extend( + str(table.get("name") or "") + for table in self._parse_schema_tables([table_ddl]) + if table.get("name") + ) + + candidate_keys: set[str] = set() + for candidate_name in candidate_names: + candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) + candidate_keys.update( + self._schema_identifier_alias_keys( + re.split(r"[.$_]", str(candidate_name or ""))[-1] + ) + ) + if explicit_table_keys.intersection(candidate_keys): + matched_documents.append(document) + + table_names, table_ddls = self._metadata_from_documents(matched_documents) + return matched_documents, table_names, table_ddls + + def _sql_references_explicit_table( + self, + sql: str, + query: str | None, + ) -> bool: + explicit_table_keys = self._explicit_table_alias_keys_from_query(query) + if not explicit_table_keys: + return True + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, + ) + referenced_tables = [ + next(value for value in match.groupdict().values() if value) + for match in table_reference_pattern.finditer(sql) + ] + + for table_reference in referenced_tables: + reference_keys = self._schema_identifier_alias_keys(table_reference) + reference_keys.update( + self._schema_identifier_alias_keys( + re.split(r"[.$_]", str(table_reference or ""))[-1] + ) + ) + if explicit_table_keys.intersection(reference_keys): + return True + + logger.warning( + "Ignoring SQL because it does not reference the explicitly requested table. " + "query=%s referenced_tables=%s sql=%s", + query, + referenced_tables, + sql, + ) + return False + def _table_matches_query(self, table_name: str, query: str) -> bool: query_keys = self._schema_identifier_alias_keys(query) short_table = re.split(r"[.$_]", str(table_name or ""))[-1] @@ -4000,6 +4080,11 @@ def _extract_retrieval_metadata( self, retrieval_result: dict ) -> tuple[list[dict], list[str], list[str]]: documents = self._extract_retrieval_documents(retrieval_result) + return documents, *self._metadata_from_documents(documents) + + def _metadata_from_documents( + self, documents: list[dict] + ) -> tuple[list[str], list[str]]: table_names = [ table_name for document in documents @@ -4012,7 +4097,7 @@ def _extract_retrieval_metadata( if isinstance(table_ddl := document.get("table_ddl"), str) and table_ddl.strip() ] - return documents, table_names, table_ddls + return table_names, table_ddls async def _complete_sql_generation_context( self, @@ -4658,6 +4743,9 @@ def _build_validated_ask_result_from_sql( ) return None + if not self._sql_references_explicit_table(ask_result.sql, query): + return None + if not self._sql_matches_question_intent( ask_result.sql, query, @@ -4889,7 +4977,7 @@ async def ask( query=user_query, tables=explicit_table_names, project_id=ask_request.project_id, - histories=histories, + histories=[], enable_column_pruning=False, ), timeout_seconds=min( @@ -4901,6 +4989,41 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) + ) + if not documents: + logger.info( + "Explicit table retrieval did not return requested active-schema table; " + "loading full active schema. query_id=%s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval for explicit table", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + all_documents, _, _ = self._extract_retrieval_metadata( + retrieval_result + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + all_documents, + ) + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) @@ -4941,54 +5064,54 @@ async def ask( explicit_sql, explicit_table_name = explicit_table_preview if explicit_table_name not in table_names: table_names.append(explicit_table_name) - api_results = [ - AskResult( - **{ - "sql": explicit_sql, - "type": "llm", - } - ) - ] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table preview request matched deployed schema.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table preview request matched deployed schema.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = explicit_sql if documents and ( deterministic_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) ): - api_results = [ - AskResult( - **{ - "sql": deterministic_sql, - "type": "llm", - } - ) - ] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sql, + table_ddls, + user_query, ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = deterministic_sql if not documents: error_message = ( @@ -5030,7 +5153,7 @@ async def ask( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( query=user_query, - histories=histories, + histories=[], project_id=ask_request.project_id, enable_column_pruning=False, ), @@ -5077,23 +5200,10 @@ async def ask( if explicit_group_count_sql := self._build_explicit_group_count_sql( user_query ): - table_column_reference = ( - self._extract_explicit_table_column_reference(user_query) - ) - table_names = ( - [table_column_reference[0]] if table_column_reference else [] - ) - api_results = [ - AskResult( - **{ - "sql": explicit_group_count_sql, - "type": "llm", - } - ) - ] + invalid_sql = explicit_group_count_sql rephrased_question = user_query logger.info( - "Using explicit table-column grouped count SQL for query_id %s", + "Deferring explicit grouped count SQL until active schema validation for query_id %s", query_id, ) @@ -5414,7 +5524,7 @@ async def ask( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( query=sql_user_query, - histories=histories, + histories=[], project_id=ask_request.project_id, enable_column_pruning=( enable_column_pruning @@ -5449,10 +5559,17 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - if not documents: - explicit_table_names = self._extract_explicit_table_names_from_query( - user_query + explicit_table_names = self._extract_explicit_table_names_from_query( + user_query + ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) ) + if not documents: if explicit_table_names: logger.info( "Retrying schema retrieval for explicit tables query_id %s: %s", @@ -5465,7 +5582,7 @@ async def ask( query=user_query, tables=explicit_table_names, project_id=ask_request.project_id, - histories=histories, + histories=[], enable_column_pruning=enable_column_pruning, ), timeout_seconds=min( @@ -5479,6 +5596,12 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) + ) if not documents and self._is_data_analysis_query(user_query): logger.info( "Query-based schema retrieval returned no tables for data question; " @@ -5505,6 +5628,13 @@ async def ask( documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + ) + ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) @@ -5542,14 +5672,16 @@ async def ask( ) if explicit_table_name not in table_names: table_names.append(explicit_table_name) - api_results = [ - AskResult( - **{ - "sql": explicit_sql, - "type": "llm", - } - ) - ] + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = explicit_sql + error_message = "Explicit table preview SQL was not valid for the active datasource schema." if not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index 4ef95c24b1..c9ed99629f 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -41,10 +41,10 @@ def test_independent_question_does_not_reuse_historical_sql(): ) -def test_contextual_followup_can_reuse_historical_sql(): +def test_contextual_followup_does_not_reuse_historical_sql(): service = AskService(pipelines={}) - assert service._should_reuse_historical_question_sql( + assert not service._should_reuse_historical_question_sql( "Use the same table and show it by month.", [AskHistory(question="previous", sql="SELECT 1")], ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index e3abde7150..73dceee9b0 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -146,6 +146,99 @@ def test_extract_explicit_table_names_from_using_clause(): ) == [] +def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): + service = AskService.__new__(AskService) + documents = [ + { + "table_name": "dbo_knowledge_articles", + "table_ddl": """ + CREATE TABLE dbo_knowledge_articles ( + id INTEGER, + last_run_date TIMESTAMP, + category VARCHAR + ); + """, + }, + { + "table_name": "dbo_failure_patterns", + "table_ddl": """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + created_at TIMESTAMP + ); + """, + }, + ] + + filtered_documents, table_names, table_ddls = ( + service._filter_retrieval_metadata_for_explicit_query( + "show monthly record count by created_at in dbo.failure_patterns.", + documents, + ) + ) + + assert filtered_documents == [documents[1]] + assert table_names == ["dbo_failure_patterns"] + assert table_ddls == [documents[1]["table_ddl"]] + + +def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): + service = AskService.__new__(AskService) + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date") AS "year", ' + '"dbo_knowledge_articles"."category" AS "category", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_knowledge_articles" ' + 'GROUP BY DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date"), ' + '"dbo_knowledge_articles"."category"' + ), + [ + """ + CREATE TABLE dbo_knowledge_articles ( + id INTEGER, + last_run_date TIMESTAMP, + category VARCHAR + ); + """, + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + created_at TIMESTAMP + ); + """, + ], + "show monthly record count by created_at in dbo.failure_patterns.", + ) + + assert result is None + + +def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): + service = AskService.__new__(AskService) + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_failure_patterns" ' + 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' + 'DATEPART(MONTH, "dbo_failure_patterns"."created_at")' + ), + [ + """ + CREATE TABLE dbo_failure_patterns ( + id INTEGER, + created_at TIMESTAMP + ); + """ + ], + "show monthly record count by created_at in dbo.failure_patterns.", + ) + + assert result is not None + + def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) From e651af7cbceef51d52259e0d4ca4b59b0a746aab Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 01:04:10 +0530 Subject: [PATCH 0440/1087] Set schema change timestamps on deploy sync --- .../repositories/schemaChangeRepository.ts | 41 +++++++++++++++++-- 1 file changed, 38 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts b/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts index a7cd1869ea..846e578c50 100644 --- a/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts +++ b/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -17,8 +21,8 @@ export interface SchemaChange { projectId: number; // Reference to project.id change: DataSourceSchemaChange; // Schema change resolve: DataSourceSchemaResolve; // Save resolve - createdAt: string; // Created at - updateAt: string; // Updated at + createdAt?: Date | string; // Created at + updatedAt?: Date | string; // Updated at } export interface ISchemaChangeRepository @@ -34,6 +38,28 @@ export class SchemaChangeRepository super({ knexPg, tableName: 'schema_change' }); } + public async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ) { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ) { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + public async findLastSchemaChange(projectId: number) { const res = await this.knex .select('*') @@ -76,4 +102,13 @@ export class SchemaChangeRepository }) as SchemaChange; return formattedData; }; + + private withTimestamps(data: Partial): Partial { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + } } From 676ad988bce404b2b958c358679ff337d0804df9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 01:18:11 +0530 Subject: [PATCH 0441/1087] Force deploy when project metadata is newer --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 155a1f2310..7860c17b23 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -262,8 +262,15 @@ export class ModelResolver { const project = await this.prepareProjectForDeploy(ctx); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + const hasModelingChangesAfterDeploy = + !(await this.isLastDeployNewerThanModelingChanges( + ctx, + project.id, + lastDeploy, + )); const shouldForceDeploy = args.force || + hasModelingChangesAfterDeploy || !ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy); const deployRes = await ctx.deployService.deploy( manifest, From 87bc302dcee7f667bd012697790168937e647805 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 02:08:39 +0530 Subject: [PATCH 0442/1087] Refresh deploy timestamp when hash is unchanged --- wren-ui/src/apollo/server/services/deployService.ts | 4 ++++ .../src/apollo/server/services/tests/deployService.test.ts | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 6c2c1e3095..746e5fa12e 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -109,6 +109,10 @@ export class DeployService implements IDeployService { await this.deployLogRepository.findLastProjectDeployLog(projectId); if (lastDeploy && lastDeploy.hash === hash) { logger.log(`Model has been deployed, hash: ${hash}`); + await this.deployLogRepository.updateOne(lastDeploy.id, { + status: DeployStatusEnum.SUCCESS, + error: null, + }); return { status: DeployStatusEnum.SUCCESS }; } } diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 00ae06570f..8788ef3f57 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -82,6 +82,7 @@ describe('DeployService', () => { const projectId = 1; mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ + id: 123, hash: deployService.createMDLHash(manifest, 1), }); @@ -89,6 +90,10 @@ describe('DeployService', () => { expect(response.status).toEqual(DeployStatusEnum.SUCCESS); expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.SUCCESS, + error: null, + }); }); it('should create the same deployment hash for equivalent manifests', () => { From 5d84909f5281b9e98dfd9e13e0f74fa1d953e1d4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 03:08:04 +0530 Subject: [PATCH 0443/1087] Require current deploy for synced model status --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 7860c17b23..f537ba8052 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -35,7 +35,6 @@ import DataSourceSchemaDetector, { const logger = getLogger('ModelResolver'); logger.level = 'debug'; -const syncedProjectIds = new Set(); const dirtyProjectIds = new Set(); export enum SyncStatusEnum { @@ -236,15 +235,14 @@ export class ModelResolver { return { status: SyncStatusEnum.UNSYNCRONIZED }; } - const isSynced = - syncedProjectIds.has(project.id) || - !!lastDeploy || - ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) || + const lastDeployIsCurrent = + ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) && (await this.isLastDeployNewerThanModelingChanges( ctx, project.id, lastDeploy, )); + const isSynced = lastDeployIsCurrent; return isSynced ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; @@ -279,7 +277,6 @@ export class ModelResolver { ); if (deployRes.status === 'SUCCESS') { dirtyProjectIds.delete(project.id); - syncedProjectIds.add(project.id); } if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { @@ -355,7 +352,6 @@ export class ModelResolver { private markProjectDirty(projectId: number) { dirtyProjectIds.add(projectId); - syncedProjectIds.delete(projectId); } private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { From fde07d590fb96bdee2d613df36e947038e2fac73 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 13:10:07 +0530 Subject: [PATCH 0444/1087] Preserve sync state on datasource version refresh --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 5 ++++- wren-ui/src/apollo/server/resolvers/projectResolver.ts | 1 + 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index f537ba8052..09c8a0b26b 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -382,7 +382,10 @@ export class ModelResolver { project, ); if (version && version !== project.version) { - return await ctx.projectService.updateProject(project.id, { version }); + return await ctx.projectService.updateProject(project.id, { + version, + updatedAt: project.updatedAt, + }); } } catch (err: any) { logger.warn( diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index cb05e2d061..5253d5d197 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -663,6 +663,7 @@ export class ProjectResolver { if (version && version !== project.version) { project = await ctx.projectService.updateProject(project.id, { version, + updatedAt: project.updatedAt, }); } } catch (err: any) { From e2d68213b047d6a2d110fa49f5c12a4cc419277f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 13:21:31 +0530 Subject: [PATCH 0445/1087] Derive sync status from persisted deploy state --- .../src/apollo/server/resolvers/modelResolver.ts | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 09c8a0b26b..9a5fff7237 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -35,8 +35,6 @@ import DataSourceSchemaDetector, { const logger = getLogger('ModelResolver'); logger.level = 'debug'; -const dirtyProjectIds = new Set(); - export enum SyncStatusEnum { IN_PROGRESS = 'IN_PROGRESS', SYNCRONIZED = 'SYNCRONIZED', @@ -231,10 +229,6 @@ export class ModelResolver { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); - if (dirtyProjectIds.has(project.id)) { - return { status: SyncStatusEnum.UNSYNCRONIZED }; - } - const lastDeployIsCurrent = ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) && (await this.isLastDeployNewerThanModelingChanges( @@ -275,10 +269,6 @@ export class ModelResolver { project.id, shouldForceDeploy, ); - if (deployRes.status === 'SUCCESS') { - dirtyProjectIds.delete(project.id); - } - if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { ctx.projectService.generateProjectRecommendationQuestions().catch((err) => logger.warn( @@ -350,8 +340,9 @@ export class ModelResolver { return Number.isFinite(time) ? time : 0; } - private markProjectDirty(projectId: number) { - dirtyProjectIds.add(projectId); + private markProjectDirty(_projectId: number) { + // Sync status is derived from persisted deploy/model state so project + // switching cannot leave stale in-memory dirty flags behind. } private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { From ae05fc3183b3c3b104ac975e4c231c013258b9aa Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 13:35:07 +0530 Subject: [PATCH 0446/1087] Keep datasource switch synced across nullability refresh --- .../apollo/server/resolvers/modelResolver.ts | 56 ++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 9a5fff7237..6b8f67ab03 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -229,13 +229,16 @@ export class ModelResolver { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); - const lastDeployIsCurrent = + const isExactCurrentDeploy = ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) && (await this.isLastDeployNewerThanModelingChanges( ctx, project.id, lastDeploy, )); + const lastDeployIsCurrent = + isExactCurrentDeploy || + this.isSameDeploymentIgnoringColumnNullability(manifest, lastDeploy); const isSynced = lastDeployIsCurrent; return isSynced ? { status: SyncStatusEnum.SYNCRONIZED } @@ -340,6 +343,57 @@ export class ModelResolver { return Number.isFinite(time) ? time : 0; } + private isSameDeploymentIgnoringColumnNullability( + manifest: any, + lastDeploy?: { manifest?: any } | null, + ): boolean { + if (!lastDeploy?.manifest) { + return false; + } + + return ( + this.stableStringify(this.omitColumnNullability(lastDeploy.manifest)) === + this.stableStringify(this.omitColumnNullability(manifest)) + ); + } + + private omitColumnNullability(value: any): any { + if (Array.isArray(value)) { + return value.map((item) => this.omitColumnNullability(item)); + } + if (!value || typeof value !== 'object') { + return value; + } + + const result: Record = {}; + for (const key of Object.keys(value)) { + if (key === 'notNull') { + continue; + } + result[key] = this.omitColumnNullability(value[key]); + } + return result; + } + + private stableStringify(value: any): string { + if (Array.isArray(value)) { + const serializedItems = value.map((item) => this.stableStringify(item)); + if (value.every((item) => item && typeof item === 'object')) { + serializedItems.sort(); + } + return `[${serializedItems.join(',')}]`; + } + + if (value && typeof value === 'object') { + return `{${Object.keys(value) + .sort() + .map((key) => `${JSON.stringify(key)}:${this.stableStringify(value[key])}`) + .join(',')}}`; + } + + return JSON.stringify(value); + } + private markProjectDirty(_projectId: number) { // Sync status is derived from persisted deploy/model state so project // switching cannot leave stale in-memory dirty flags behind. From b635f7a45ad7376ebd7eace288926353957e9934 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 13:51:57 +0530 Subject: [PATCH 0447/1087] Normalize schema change detection identifiers --- .../managers/dataSourceSchemaDetector.ts | 175 +++++++++++++++--- .../tests/dataSourceSchemaDetector.test.ts | 43 +++++ .../server/resolvers/projectResolver.ts | 41 +++- 3 files changed, 227 insertions(+), 32 deletions(-) diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index d1c262aa5c..11b80eaeef 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -138,18 +138,53 @@ export default class DataSourceSchemaDetector return createdModels.length > 0; } + public async filterCurrentSchemaChange( + schemaChange: DataSourceSchemaChange, + ): Promise { + const latestSchema = await this.getLatestSchema(); + const filteredChange: DataSourceSchemaChange = {}; + + const deletedTables = schemaChange[SchemaChangeType.DELETED_TABLES]?.filter( + (table) => + !latestSchema.some((latestTable) => + this.isSameIdentifier(latestTable.name, table.name), + ), + ); + if (deletedTables?.length) { + filteredChange[SchemaChangeType.DELETED_TABLES] = deletedTables; + } + + const deletedColumns = this.filterMissingColumns( + schemaChange[SchemaChangeType.DELETED_COLUMNS], + latestSchema, + ); + if (deletedColumns.length) { + filteredChange[SchemaChangeType.DELETED_COLUMNS] = deletedColumns; + } + + const modifiedColumns = this.filterModifiedColumns( + schemaChange[SchemaChangeType.MODIFIED_COLUMNS], + latestSchema, + ); + if (modifiedColumns.length) { + filteredChange[SchemaChangeType.MODIFIED_COLUMNS] = modifiedColumns; + } + + return filteredChange; + } + private async createMissingModels( latestTables: CompactTable[], models: Model[], ): Promise { - const existingSourceTableNames = new Set( - models.map((model) => model.sourceTableName), - ); const usedReferenceNames = new Set( models.map((model) => model.referenceName.toLowerCase()), ); const missingTables = latestTables.filter( - (table) => !existingSourceTableNames.has(table.name), + (table) => + !models.some((model) => + this.isSameIdentifier(model.sourceTableName, table.name), + ), ); if (!missingTables.length) { @@ -189,8 +224,8 @@ export default class DataSourceSchemaDetector } const columnValues = models.flatMap((model) => { - const table = latestTables.find( - (table) => table.name === model.sourceTableName, + const table = latestTables.find((table) => + this.isSameIdentifier(table.name, model.sourceTableName), ); if (!table) { return []; @@ -245,8 +280,8 @@ export default class DataSourceSchemaDetector columns: ModelColumn[], ) { const nestedColumnValues = models.flatMap((model) => { - const table = latestTables.find( - (table) => table.name === model.sourceTableName, + const table = latestTables.find((table) => + this.isSameIdentifier(table.name, model.sourceTableName), ); if (!table) { return []; @@ -255,8 +290,8 @@ export default class DataSourceSchemaDetector (column) => column.modelId === model.id, ); return table.columns.flatMap((compactColumn) => { - const column = modelColumns.find( - (column) => column.sourceColumnName === compactColumn.name, + const column = modelColumns.find((column) => + this.isSameIdentifier(column.sourceColumnName, compactColumn.name), ); if (!column) { return []; @@ -377,7 +412,10 @@ export default class DataSourceSchemaDetector const modelColumn = modelColumns.find( (modelColumn) => modelColumn.modelId === resource.modelId && - modelColumn.sourceColumnName === column.sourceColumnName && + this.isSameIdentifier( + modelColumn.sourceColumnName, + column.sourceColumnName, + ) && !modelColumn.isCalculated, ); if (!modelColumn || modelColumn.type === column.type) { @@ -440,13 +478,14 @@ export default class DataSourceSchemaDetector ) { const affectedModels = models.filter( (model) => - changes.findIndex((table) => table.name === model.sourceTableName) !== - -1, + changes.findIndex((table) => + this.isSameIdentifier(table.name, model.sourceTableName), + ) !== -1, ); const affectedResources = affectedModels.map((model) => { - const affectedColumns = changes.find( - (table) => table.name === model.sourceTableName, + const affectedColumns = changes.find((table) => + this.isSameIdentifier(table.name, model.sourceTableName), ).columns; const allCalculatedFields = modelColumns.filter( @@ -457,7 +496,7 @@ export default class DataSourceSchemaDetector (result, column) => { const affectedColumn = modelColumns.find( (modelColumn) => - modelColumn.sourceColumnName === column.name && + this.isSameIdentifier(modelColumn.sourceColumnName, column.name) && modelColumn.modelId === model.id, ); @@ -547,8 +586,8 @@ export default class DataSourceSchemaDetector latestSchema: DataSourceSchema[], ) { const diffSchema = currentSchema.reduce((result, currentTable) => { - const lastestTable = latestSchema.find( - (table) => table.name === currentTable.name, + const lastestTable = latestSchema.find((table) => + this.isSameIdentifier(table.name, currentTable.name), ); // If the table is not found in the latest schema, it means the table has been deleted. if (!lastestTable) { @@ -570,8 +609,8 @@ export default class DataSourceSchemaDetector const modifiedColumnChange = { name: currentTable.name, columns: [] }; for (const currentColumn of diffColumns) { - const latestColumn = lastestTable.columns.find( - (column) => column.name === currentColumn.name, + const latestColumn = lastestTable.columns.find((column) => + this.isSameIdentifier(column.name, currentColumn.name), ); // If the column is not found in the latest schema, it means the column has been deleted. if (!latestColumn) { @@ -685,7 +724,7 @@ export default class DataSourceSchemaDetector latestColumn: DataSourceSchema['columns'][number], ) { return ( - currentColumn.name === latestColumn.name && + this.isSameIdentifier(currentColumn.name, latestColumn.name) && currentColumn.type === latestColumn.type ); } @@ -706,8 +745,8 @@ export default class DataSourceSchemaDetector await this.ctx.modelColumnRepository.findColumnsByModelIds(modelIds); for (const model of models) { - const latestTable = latestSchema.find( - (table) => table.name === model.sourceTableName, + const latestTable = latestSchema.find((table) => + this.isSameIdentifier(table.name, model.sourceTableName), ); if (!latestTable) { continue; @@ -723,8 +762,8 @@ export default class DataSourceSchemaDetector ); for (const latestColumn of latestTable.columns) { - const existingColumn = existingColumns.find( - (column) => column.sourceColumnName === latestColumn.name, + const existingColumn = existingColumns.find((column) => + this.isSameIdentifier(column.sourceColumnName, latestColumn.name), ); if (!existingColumn) { @@ -819,4 +858,90 @@ export default class DataSourceSchemaDetector }); logger.info(`Schema change "${schemaChangeTypes}" resolved successfully.`); } + + private filterMissingColumns( + tables: DataSourceSchema[] | undefined, + latestSchema: DataSourceSchema[], + ): DataSourceSchema[] { + if (!tables?.length) { + return []; + } + + return tables + .map((table) => { + const latestTable = latestSchema.find((latestTable) => + this.isSameIdentifier(latestTable.name, table.name), + ); + if (!latestTable) { + return table; + } + + const columns = table.columns.filter( + (column) => + !latestTable.columns.some((latestColumn) => + this.isSameIdentifier(latestColumn.name, column.name), + ), + ); + return { ...table, columns }; + }) + .filter((table) => table.columns.length > 0); + } + + private filterModifiedColumns( + tables: DataSourceSchema[] | undefined, + latestSchema: DataSourceSchema[], + ): DataSourceSchema[] { + if (!tables?.length) { + return []; + } + + return tables + .map((table) => { + const latestTable = latestSchema.find((latestTable) => + this.isSameIdentifier(latestTable.name, table.name), + ); + if (!latestTable) { + return table; + } + + const columns = table.columns.filter((column) => { + const latestColumn = latestTable.columns.find((latestColumn) => + this.isSameIdentifier(latestColumn.name, column.name), + ); + return !!latestColumn && latestColumn.type !== column.type; + }); + return { ...table, columns }; + }) + .filter((table) => table.columns.length > 0); + } + + private isSameIdentifier(left?: string, right?: string): boolean { + const normalizedLeft = this.normalizeIdentifier(left); + const normalizedRight = this.normalizeIdentifier(right); + if (normalizedLeft === normalizedRight) { + return true; + } + + const leftIsQualified = normalizedLeft.includes('.'); + const rightIsQualified = normalizedRight.includes('.'); + if (leftIsQualified && rightIsQualified) { + return false; + } + + return ( + this.getUnqualifiedIdentifier(normalizedLeft) === + this.getUnqualifiedIdentifier(normalizedRight) + ); + } + + private normalizeIdentifier(value?: string): string { + return String(value || '') + .replace(/[\[\]"`]/g, '') + .trim() + .toLowerCase(); + } + + private getUnqualifiedIdentifier(value: string): string { + return value.split('.').pop() || value; + } } diff --git a/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts index 869f7b3299..df6ad864a6 100644 --- a/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts +++ b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts @@ -166,4 +166,47 @@ describe('DataSourceSchemaDetector', () => { }), ); }); + + it('does not report schema changes for qualified or case-only identifier differences', async () => { + const model = { + id: 10, + projectId, + sourceTableName: 'dbo.Repair_Logs', + }; + const existingColumn = { + id: 20, + modelId: 10, + isCalculated: false, + displayName: 'Created At', + referenceName: 'created_at', + sourceColumnName: 'created_at', + type: 'datetime', + notNull: false, + isPk: false, + properties: null, + }; + const ctx = createContext({ + models: [model], + columns: [existingColumn], + latestTables: [ + { + name: '[repair_logs]', + columns: [ + { + name: '[Created_At]', + type: 'datetime', + notNull: false, + }, + ], + }, + ], + }); + + const detector = new DataSourceSchemaDetector({ ctx, projectId }); + + await expect(detector.detectSchemaChange()).resolves.toBe(false); + expect(ctx.schemaChangeRepository.createOne).not.toHaveBeenCalled(); + expect(ctx.modelColumnRepository.createOne).not.toHaveBeenCalled(); + expect(ctx.modelColumnRepository.updateOne).not.toHaveBeenCalled(); + }); }); diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 5253d5d197..6df3ba00d4 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -29,7 +29,7 @@ import { getRelations, sampleDatasets, } from '@server/data'; -import { snakeCase } from 'lodash'; +import { isEmpty, snakeCase } from 'lodash'; import { CompactTable, ProjectData } from '../services'; import { DuckDBPrepareOptions } from '@server/adaptors/wrenEngineAdaptor'; import DataSourceSchemaDetector, { @@ -544,6 +544,38 @@ export class ProjectResolver { }; } + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + const currentChange = await schemaDetector.filterCurrentSchemaChange( + lastSchemaChange.change, + ); + const staleResolvedTypes = Object.values(SchemaChangeType).filter((type) => { + const isResolved = lastSchemaChange.resolve[type]; + return !isResolved && !!lastSchemaChange.change[type] && !currentChange[type]; + }); + if (staleResolvedTypes.length) { + await ctx.schemaChangeRepository.updateOne(lastSchemaChange.id, { + resolve: { + ...lastSchemaChange.resolve, + ...staleResolvedTypes.reduce( + (result, type) => ({ ...result, [type]: true }), + {}, + ), + }, + }); + } + + if (isEmpty(currentChange)) { + return { + deletedTables: null, + deletedColumns: null, + modifiedColumns: null, + lastSchemaChangeTime: lastSchemaChange.createdAt, + }; + } + const models = await ctx.modelRepository.findAllBy({ projectId: project.id, }); @@ -555,15 +587,10 @@ export class ProjectResolver { modelIds, }); - const schemaDetector = new DataSourceSchemaDetector({ - ctx, - projectId: project.id, - }); - const resolves = lastSchemaChange.resolve; const unresolvedChanges = Object.keys(resolves).reduce((result, key) => { const isResolved = resolves[key]; - const changes = lastSchemaChange.change[key]; + const changes = currentChange[key]; // return if resolved or no changes if (isResolved || !changes) return result; From 137722a299165051e8733fb6acd2c51133ce050c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 14:22:17 +0530 Subject: [PATCH 0448/1087] Fix schema resolve callback variables access --- wren-ui/src/components/sidebar/modeling/ModelTree.tsx | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx index db0a30c8d1..7435b376dd 100644 --- a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx +++ b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo, useState } from 'react'; +import { useEffect, useMemo, useRef, useState } from 'react'; import { message } from 'antd'; import { DataNode } from 'antd/es/tree'; import { DiagramModel } from '@/utils/data'; @@ -46,6 +46,7 @@ export default function ModelTree(props: Props) { const { onOpenModelDrawer, models } = props; const schemaChangeModal = useModalAction(); + const pendingResolveType = useRef(null); const [triggerDataSourceDetection, { loading: isDetecting }] = useTriggerDataSourceDetectionMutation({ onError: (error) => console.error(error), @@ -62,8 +63,8 @@ export default function ModelTree(props: Props) { const [resolveSchemaChange, { loading: isResolving }] = useResolveSchemaChangeMutation({ onError: (error) => console.error(error), - onCompleted: async (_, options) => { - const { type } = options.variables?.where; + onCompleted: async () => { + const type = pendingResolveType.current; if (type === SchemaChangeType.DELETED_TABLES) { message.success('Source table deleted resolved successfully.'); } else if (type === SchemaChangeType.DELETED_COLUMNS) { @@ -75,6 +76,7 @@ export default function ModelTree(props: Props) { if (!getHasSchemaChange(data.schemaChange)) { schemaChangeModal.closeModal(); } + pendingResolveType.current = null; }, refetchQueries: [{ query: DIAGRAM }, { query: LIST_MODELS }], }); @@ -90,6 +92,7 @@ export default function ModelTree(props: Props) { schemaChangeModal.openModal(); }; const onResolveSchemaChange = (type: SchemaChangeType) => { + pendingResolveType.current = type; resolveSchemaChange({ variables: { where: { type } } }); }; From 6164d16d2021e3b5cc5ca528314b6598e59c8f69 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 14:28:13 +0530 Subject: [PATCH 0449/1087] Revert "Fix schema resolve callback variables access" This reverts commit 137722a299165051e8733fb6acd2c51133ce050c. --- wren-ui/src/components/sidebar/modeling/ModelTree.tsx | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx index 7435b376dd..db0a30c8d1 100644 --- a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx +++ b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo, useRef, useState } from 'react'; +import { useEffect, useMemo, useState } from 'react'; import { message } from 'antd'; import { DataNode } from 'antd/es/tree'; import { DiagramModel } from '@/utils/data'; @@ -46,7 +46,6 @@ export default function ModelTree(props: Props) { const { onOpenModelDrawer, models } = props; const schemaChangeModal = useModalAction(); - const pendingResolveType = useRef(null); const [triggerDataSourceDetection, { loading: isDetecting }] = useTriggerDataSourceDetectionMutation({ onError: (error) => console.error(error), @@ -63,8 +62,8 @@ export default function ModelTree(props: Props) { const [resolveSchemaChange, { loading: isResolving }] = useResolveSchemaChangeMutation({ onError: (error) => console.error(error), - onCompleted: async () => { - const type = pendingResolveType.current; + onCompleted: async (_, options) => { + const { type } = options.variables?.where; if (type === SchemaChangeType.DELETED_TABLES) { message.success('Source table deleted resolved successfully.'); } else if (type === SchemaChangeType.DELETED_COLUMNS) { @@ -76,7 +75,6 @@ export default function ModelTree(props: Props) { if (!getHasSchemaChange(data.schemaChange)) { schemaChangeModal.closeModal(); } - pendingResolveType.current = null; }, refetchQueries: [{ query: DIAGRAM }, { query: LIST_MODELS }], }); @@ -92,7 +90,6 @@ export default function ModelTree(props: Props) { schemaChangeModal.openModal(); }; const onResolveSchemaChange = (type: SchemaChangeType) => { - pendingResolveType.current = type; resolveSchemaChange({ variables: { where: { type } } }); }; From c3ac8a719d41d358eb109f548ce3e7a59798053c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 14:28:13 +0530 Subject: [PATCH 0450/1087] Revert "Normalize schema change detection identifiers" This reverts commit b635f7a45ad7376ebd7eace288926353957e9934. --- .../managers/dataSourceSchemaDetector.ts | 175 +++--------------- .../tests/dataSourceSchemaDetector.test.ts | 43 ----- .../server/resolvers/projectResolver.ts | 41 +--- 3 files changed, 32 insertions(+), 227 deletions(-) diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index 11b80eaeef..d1c262aa5c 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -138,53 +138,18 @@ export default class DataSourceSchemaDetector return createdModels.length > 0; } - public async filterCurrentSchemaChange( - schemaChange: DataSourceSchemaChange, - ): Promise { - const latestSchema = await this.getLatestSchema(); - const filteredChange: DataSourceSchemaChange = {}; - - const deletedTables = schemaChange[SchemaChangeType.DELETED_TABLES]?.filter( - (table) => - !latestSchema.some((latestTable) => - this.isSameIdentifier(latestTable.name, table.name), - ), - ); - if (deletedTables?.length) { - filteredChange[SchemaChangeType.DELETED_TABLES] = deletedTables; - } - - const deletedColumns = this.filterMissingColumns( - schemaChange[SchemaChangeType.DELETED_COLUMNS], - latestSchema, - ); - if (deletedColumns.length) { - filteredChange[SchemaChangeType.DELETED_COLUMNS] = deletedColumns; - } - - const modifiedColumns = this.filterModifiedColumns( - schemaChange[SchemaChangeType.MODIFIED_COLUMNS], - latestSchema, - ); - if (modifiedColumns.length) { - filteredChange[SchemaChangeType.MODIFIED_COLUMNS] = modifiedColumns; - } - - return filteredChange; - } - private async createMissingModels( latestTables: CompactTable[], models: Model[], ): Promise { + const existingSourceTableNames = new Set( + models.map((model) => model.sourceTableName), + ); const usedReferenceNames = new Set( models.map((model) => model.referenceName.toLowerCase()), ); const missingTables = latestTables.filter( - (table) => - !models.some((model) => - this.isSameIdentifier(model.sourceTableName, table.name), - ), + (table) => !existingSourceTableNames.has(table.name), ); if (!missingTables.length) { @@ -224,8 +189,8 @@ export default class DataSourceSchemaDetector } const columnValues = models.flatMap((model) => { - const table = latestTables.find((table) => - this.isSameIdentifier(table.name, model.sourceTableName), + const table = latestTables.find( + (table) => table.name === model.sourceTableName, ); if (!table) { return []; @@ -280,8 +245,8 @@ export default class DataSourceSchemaDetector columns: ModelColumn[], ) { const nestedColumnValues = models.flatMap((model) => { - const table = latestTables.find((table) => - this.isSameIdentifier(table.name, model.sourceTableName), + const table = latestTables.find( + (table) => table.name === model.sourceTableName, ); if (!table) { return []; @@ -290,8 +255,8 @@ export default class DataSourceSchemaDetector (column) => column.modelId === model.id, ); return table.columns.flatMap((compactColumn) => { - const column = modelColumns.find((column) => - this.isSameIdentifier(column.sourceColumnName, compactColumn.name), + const column = modelColumns.find( + (column) => column.sourceColumnName === compactColumn.name, ); if (!column) { return []; @@ -412,10 +377,7 @@ export default class DataSourceSchemaDetector const modelColumn = modelColumns.find( (modelColumn) => modelColumn.modelId === resource.modelId && - this.isSameIdentifier( - modelColumn.sourceColumnName, - column.sourceColumnName, - ) && + modelColumn.sourceColumnName === column.sourceColumnName && !modelColumn.isCalculated, ); if (!modelColumn || modelColumn.type === column.type) { @@ -478,14 +440,13 @@ export default class DataSourceSchemaDetector ) { const affectedModels = models.filter( (model) => - changes.findIndex((table) => - this.isSameIdentifier(table.name, model.sourceTableName), - ) !== -1, + changes.findIndex((table) => table.name === model.sourceTableName) !== + -1, ); const affectedResources = affectedModels.map((model) => { - const affectedColumns = changes.find((table) => - this.isSameIdentifier(table.name, model.sourceTableName), + const affectedColumns = changes.find( + (table) => table.name === model.sourceTableName, ).columns; const allCalculatedFields = modelColumns.filter( @@ -496,7 +457,7 @@ export default class DataSourceSchemaDetector (result, column) => { const affectedColumn = modelColumns.find( (modelColumn) => - this.isSameIdentifier(modelColumn.sourceColumnName, column.name) && + modelColumn.sourceColumnName === column.name && modelColumn.modelId === model.id, ); @@ -586,8 +547,8 @@ export default class DataSourceSchemaDetector latestSchema: DataSourceSchema[], ) { const diffSchema = currentSchema.reduce((result, currentTable) => { - const lastestTable = latestSchema.find((table) => - this.isSameIdentifier(table.name, currentTable.name), + const lastestTable = latestSchema.find( + (table) => table.name === currentTable.name, ); // If the table is not found in the latest schema, it means the table has been deleted. if (!lastestTable) { @@ -609,8 +570,8 @@ export default class DataSourceSchemaDetector const modifiedColumnChange = { name: currentTable.name, columns: [] }; for (const currentColumn of diffColumns) { - const latestColumn = lastestTable.columns.find((column) => - this.isSameIdentifier(column.name, currentColumn.name), + const latestColumn = lastestTable.columns.find( + (column) => column.name === currentColumn.name, ); // If the column is not found in the latest schema, it means the column has been deleted. if (!latestColumn) { @@ -724,7 +685,7 @@ export default class DataSourceSchemaDetector latestColumn: DataSourceSchema['columns'][number], ) { return ( - this.isSameIdentifier(currentColumn.name, latestColumn.name) && + currentColumn.name === latestColumn.name && currentColumn.type === latestColumn.type ); } @@ -745,8 +706,8 @@ export default class DataSourceSchemaDetector await this.ctx.modelColumnRepository.findColumnsByModelIds(modelIds); for (const model of models) { - const latestTable = latestSchema.find((table) => - this.isSameIdentifier(table.name, model.sourceTableName), + const latestTable = latestSchema.find( + (table) => table.name === model.sourceTableName, ); if (!latestTable) { continue; @@ -762,8 +723,8 @@ export default class DataSourceSchemaDetector ); for (const latestColumn of latestTable.columns) { - const existingColumn = existingColumns.find((column) => - this.isSameIdentifier(column.sourceColumnName, latestColumn.name), + const existingColumn = existingColumns.find( + (column) => column.sourceColumnName === latestColumn.name, ); if (!existingColumn) { @@ -858,90 +819,4 @@ export default class DataSourceSchemaDetector }); logger.info(`Schema change "${schemaChangeTypes}" resolved successfully.`); } - - private filterMissingColumns( - tables: DataSourceSchema[] | undefined, - latestSchema: DataSourceSchema[], - ): DataSourceSchema[] { - if (!tables?.length) { - return []; - } - - return tables - .map((table) => { - const latestTable = latestSchema.find((latestTable) => - this.isSameIdentifier(latestTable.name, table.name), - ); - if (!latestTable) { - return table; - } - - const columns = table.columns.filter( - (column) => - !latestTable.columns.some((latestColumn) => - this.isSameIdentifier(latestColumn.name, column.name), - ), - ); - return { ...table, columns }; - }) - .filter((table) => table.columns.length > 0); - } - - private filterModifiedColumns( - tables: DataSourceSchema[] | undefined, - latestSchema: DataSourceSchema[], - ): DataSourceSchema[] { - if (!tables?.length) { - return []; - } - - return tables - .map((table) => { - const latestTable = latestSchema.find((latestTable) => - this.isSameIdentifier(latestTable.name, table.name), - ); - if (!latestTable) { - return table; - } - - const columns = table.columns.filter((column) => { - const latestColumn = latestTable.columns.find((latestColumn) => - this.isSameIdentifier(latestColumn.name, column.name), - ); - return !!latestColumn && latestColumn.type !== column.type; - }); - return { ...table, columns }; - }) - .filter((table) => table.columns.length > 0); - } - - private isSameIdentifier(left?: string, right?: string): boolean { - const normalizedLeft = this.normalizeIdentifier(left); - const normalizedRight = this.normalizeIdentifier(right); - if (normalizedLeft === normalizedRight) { - return true; - } - - const leftIsQualified = normalizedLeft.includes('.'); - const rightIsQualified = normalizedRight.includes('.'); - if (leftIsQualified && rightIsQualified) { - return false; - } - - return ( - this.getUnqualifiedIdentifier(normalizedLeft) === - this.getUnqualifiedIdentifier(normalizedRight) - ); - } - - private normalizeIdentifier(value?: string): string { - return String(value || '') - .replace(/[\[\]"`]/g, '') - .trim() - .toLowerCase(); - } - - private getUnqualifiedIdentifier(value: string): string { - return value.split('.').pop() || value; - } } diff --git a/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts index df6ad864a6..869f7b3299 100644 --- a/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts +++ b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts @@ -166,47 +166,4 @@ describe('DataSourceSchemaDetector', () => { }), ); }); - - it('does not report schema changes for qualified or case-only identifier differences', async () => { - const model = { - id: 10, - projectId, - sourceTableName: 'dbo.Repair_Logs', - }; - const existingColumn = { - id: 20, - modelId: 10, - isCalculated: false, - displayName: 'Created At', - referenceName: 'created_at', - sourceColumnName: 'created_at', - type: 'datetime', - notNull: false, - isPk: false, - properties: null, - }; - const ctx = createContext({ - models: [model], - columns: [existingColumn], - latestTables: [ - { - name: '[repair_logs]', - columns: [ - { - name: '[Created_At]', - type: 'datetime', - notNull: false, - }, - ], - }, - ], - }); - - const detector = new DataSourceSchemaDetector({ ctx, projectId }); - - await expect(detector.detectSchemaChange()).resolves.toBe(false); - expect(ctx.schemaChangeRepository.createOne).not.toHaveBeenCalled(); - expect(ctx.modelColumnRepository.createOne).not.toHaveBeenCalled(); - expect(ctx.modelColumnRepository.updateOne).not.toHaveBeenCalled(); - }); }); diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 6df3ba00d4..5253d5d197 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -29,7 +29,7 @@ import { getRelations, sampleDatasets, } from '@server/data'; -import { isEmpty, snakeCase } from 'lodash'; +import { snakeCase } from 'lodash'; import { CompactTable, ProjectData } from '../services'; import { DuckDBPrepareOptions } from '@server/adaptors/wrenEngineAdaptor'; import DataSourceSchemaDetector, { @@ -544,38 +544,6 @@ export class ProjectResolver { }; } - const schemaDetector = new DataSourceSchemaDetector({ - ctx, - projectId: project.id, - }); - const currentChange = await schemaDetector.filterCurrentSchemaChange( - lastSchemaChange.change, - ); - const staleResolvedTypes = Object.values(SchemaChangeType).filter((type) => { - const isResolved = lastSchemaChange.resolve[type]; - return !isResolved && !!lastSchemaChange.change[type] && !currentChange[type]; - }); - if (staleResolvedTypes.length) { - await ctx.schemaChangeRepository.updateOne(lastSchemaChange.id, { - resolve: { - ...lastSchemaChange.resolve, - ...staleResolvedTypes.reduce( - (result, type) => ({ ...result, [type]: true }), - {}, - ), - }, - }); - } - - if (isEmpty(currentChange)) { - return { - deletedTables: null, - deletedColumns: null, - modifiedColumns: null, - lastSchemaChangeTime: lastSchemaChange.createdAt, - }; - } - const models = await ctx.modelRepository.findAllBy({ projectId: project.id, }); @@ -587,10 +555,15 @@ export class ProjectResolver { modelIds, }); + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + const resolves = lastSchemaChange.resolve; const unresolvedChanges = Object.keys(resolves).reduce((result, key) => { const isResolved = resolves[key]; - const changes = currentChange[key]; + const changes = lastSchemaChange.change[key]; // return if resolved or no changes if (isResolved || !changes) return result; From 96153eb06166c5bef58cbf9089af33f45263c39b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 15:16:20 +0530 Subject: [PATCH 0451/1087] Keep model synced after navigation --- .../apollo/server/resolvers/modelResolver.ts | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 6b8f67ab03..6b7320f6ea 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -227,19 +227,12 @@ export class ModelResolver { } const project = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const lastDeploy = await ctx.deployService.getLastDeployment(project.id); - const isExactCurrentDeploy = - ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy) && - (await this.isLastDeployNewerThanModelingChanges( - ctx, - project.id, - lastDeploy, - )); - const lastDeployIsCurrent = - isExactCurrentDeploy || - this.isSameDeploymentIgnoringColumnNullability(manifest, lastDeploy); - const isSynced = lastDeployIsCurrent; + const isSynced = await this.isLastDeployNewerThanModelingChanges( + ctx, + project.id, + lastDeploy, + ); return isSynced ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; From 006934c9d0f5969d28bd40819872aa4fdadb22ef Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 15:22:39 +0530 Subject: [PATCH 0452/1087] Keep synced state after schema refresh --- .../apollo/server/resolvers/modelResolver.ts | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 6b7320f6ea..b7e91310d2 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -35,6 +35,8 @@ import DataSourceSchemaDetector, { const logger = getLogger('ModelResolver'); logger.level = 'debug'; +const dirtyProjectIds = new Set(); + export enum SyncStatusEnum { IN_PROGRESS = 'IN_PROGRESS', SYNCRONIZED = 'SYNCRONIZED', @@ -227,13 +229,12 @@ export class ModelResolver { } const project = await ctx.projectService.getCurrentProject(); + if (dirtyProjectIds.has(project.id)) { + return { status: SyncStatusEnum.UNSYNCRONIZED }; + } + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); - const isSynced = await this.isLastDeployNewerThanModelingChanges( - ctx, - project.id, - lastDeploy, - ); - return isSynced + return lastDeploy ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { @@ -265,6 +266,9 @@ export class ModelResolver { project.id, shouldForceDeploy, ); + if (deployRes.status === 'SUCCESS') { + dirtyProjectIds.delete(project.id); + } if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { ctx.projectService.generateProjectRecommendationQuestions().catch((err) => logger.warn( @@ -387,9 +391,8 @@ export class ModelResolver { return JSON.stringify(value); } - private markProjectDirty(_projectId: number) { - // Sync status is derived from persisted deploy/model state so project - // switching cannot leave stale in-memory dirty flags behind. + private markProjectDirty(projectId: number) { + dirtyProjectIds.add(projectId); } private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { From ba5b51b8cbc97da15ce46e5c6f90854b6b1ff370 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 15:40:03 +0530 Subject: [PATCH 0453/1087] Reject invalid unqualified SQL columns --- .../src/pipelines/generation/utils/sql.py | 137 ++++++++++++++ .../pipelines/generation/test_sql_utils.py | 14 ++ .../apollo/server/services/queryService.ts | 170 ++++++++++++++++++ .../services/tests/queryService.test.ts | 53 ++++++ 4 files changed, 374 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 0c1f7e58fa..8316d24e86 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2494,6 +2494,46 @@ def construct_valid_table_columns( "limit", "union", } +_SQL_NON_COLUMN_IDENTIFIERS = { + *_SQL_RESERVED_ALIASES, + "and", + "as", + "asc", + "between", + "by", + "case", + "cast", + "count", + "date", + "dateadd", + "datediff", + "datepart", + "day", + "desc", + "distinct", + "else", + "end", + "false", + "from", + "getdate", + "hour", + "in", + "is", + "like", + "max", + "min", + "month", + "not", + "null", + "or", + "select", + "sum", + "then", + "top", + "true", + "when", + "year", +} def _normalize_sql_identifier(identifier: str) -> str: @@ -2515,6 +2555,89 @@ def _compact_sql_identifier(identifier: str) -> str: return re.sub(r"[^a-z0-9]", "", str(identifier or "").lower()) +def _strip_sql_literals(sql: str) -> str: + without_strings = re.sub(r"'(?:''|[^'])*'", " ", sql) + return re.sub(r"\b\d+(?:\.\d+)?\b", " ", without_strings) + + +def _extract_clause_bodies( + sql: str, clause: str, terminators: list[str] +) -> list[str]: + terminator_pattern = "|".join(rf"\b{terminator}\b" for terminator in terminators) + pattern = re.compile( + rf"\b{clause}\b(?P.*?)(?={terminator_pattern}|$)", + flags=re.IGNORECASE | re.DOTALL, + ) + return [match.group("body") or "" for match in pattern.finditer(sql)] + + +def _find_unqualified_column_candidates(sql: str) -> set[str]: + bodies = [ + *_extract_clause_bodies( + sql, + r"WHERE", + [r"GROUP\s+BY", r"ORDER\s+BY", "HAVING", "LIMIT", "FETCH", "UNION"], + ), + *_extract_clause_bodies( + sql, + r"HAVING", + [r"GROUP\s+BY", r"ORDER\s+BY", "LIMIT", "FETCH", "UNION"], + ), + *_extract_clause_bodies( + sql, + r"ON", + [ + "WHERE", + r"GROUP\s+BY", + r"ORDER\s+BY", + "HAVING", + "JOIN", + "LIMIT", + "FETCH", + "UNION", + ], + ), + *_extract_clause_bodies( + sql, + r"GROUP\s+BY", + [r"ORDER\s+BY", "HAVING", "LIMIT", "FETCH", "UNION"], + ), + ] + candidates: set[str] = set() + identifier_pattern = re.compile(_SQL_IDENTIFIER_PATTERN, flags=re.IGNORECASE) + + for body in bodies: + searchable_body = _strip_sql_literals(body) + for match in identifier_pattern.finditer(searchable_body): + token = match.group(0) + before = searchable_body[: match.start()].rstrip() + after = searchable_body[match.end() :].lstrip() + identifier = _normalize_sql_identifier(token) + normalized = identifier.lower() + if ( + not identifier + or normalized in _SQL_NON_COLUMN_IDENTIFIERS + or before.endswith(".") + or after.startswith(".") + or after.startswith("(") + ): + continue + candidates.add(identifier) + + function_argument_pattern = re.compile( + rf"\b[A-Za-z_][A-Za-z0-9_$]*\s*\(\s*" + rf"(?:DISTINCT\s+)?(?P{_SQL_IDENTIFIER_PATTERN})" + rf"(?:\s*\.\s*(?P{_SQL_IDENTIFIER_PATTERN}))?", + flags=re.IGNORECASE, + ) + for match in function_argument_pattern.finditer(_strip_sql_literals(sql)): + column = _normalize_sql_identifier(match.group("column") or match.group("arg")) + if column and column != "*" and column.lower() not in _SQL_NON_COLUMN_IDENTIFIERS: + candidates.add(column) + + return candidates + + def _sql_identifier_alias_candidates(identifier: str) -> set[str]: normalized = _normalize_sql_identifier(identifier) candidates = {normalized} @@ -3037,6 +3160,20 @@ def find_invalid_column_references( and _compact_sql_identifier(column) not in valid_compact_columns ): invalid_references.append(column) + table_aliases = { + alias.lower() + for alias, alias_table in aliases.items() + if alias_table == table_name + } + for column in _find_unqualified_column_candidates(sql): + normalized_column = column.lower() + if ( + normalized_column in table_aliases + or normalized_column in valid_columns + or _compact_sql_identifier(column) in valid_compact_columns + ): + continue + invalid_references.append(column) return sorted(set(invalid_references)) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 421566f5eb..fa46c2ca7a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -74,6 +74,20 @@ def test_column_validation_rejects_invalid_unqualified_projection_alias(): ) == ["categories"] +def test_column_validation_rejects_invalid_unqualified_filter_column(): + assert find_invalid_column_references( + 'SELECT id FROM "policies" WHERE policy_category_id = 1', + {"policies": ["id", "policy_name"]}, + ) == ["policy_category_id"] + + +def test_column_validation_rejects_invalid_unqualified_function_argument(): + assert find_invalid_column_references( + 'SELECT COUNT(policy_category_id) FROM "policies"', + {"policies": ["id", "policy_name"]}, + ) == ["policy_category_id"] + + def test_column_validation_allows_valid_unqualified_projection_for_single_table(): assert find_invalid_column_references( 'SELECT status AS repair_status FROM "dbo_repair_logs"', diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 09c3734120..4ae56a81e8 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -532,6 +532,55 @@ const isNumericColumnType = (type?: string) => const isDefined = (value: T | undefined | null): value is T => value !== undefined && value !== null; +const SQL_NON_COLUMN_IDENTIFIERS = new Set([ + 'and', + 'as', + 'asc', + 'between', + 'by', + 'case', + 'cast', + 'count', + 'date', + 'dateadd', + 'datediff', + 'datepart', + 'day', + 'desc', + 'distinct', + 'else', + 'end', + 'false', + 'from', + 'getdate', + 'group', + 'having', + 'hour', + 'in', + 'is', + 'join', + 'left', + 'like', + 'limit', + 'max', + 'min', + 'month', + 'not', + 'null', + 'on', + 'or', + 'order', + 'right', + 'select', + 'sum', + 'then', + 'top', + 'true', + 'when', + 'where', + 'year', +]); + const splitTopLevelSqlList = (body: string) => { const items: string[] = []; let current = ''; @@ -600,6 +649,109 @@ const extractSimpleProjectionColumns = (sql: string) => { return columns; }; +const stripSqlLiterals = (sql: string) => + sql + .replace(/'(?:''|[^'])*'/g, ' ') + .replace(/\b\d+(?:\.\d+)?\b/g, ' '); + +const extractClauseBody = ( + sql: string, + clause: string, + terminators: string[], +) => { + const pattern = new RegExp( + String.raw`\b${clause}\b(?.*?)(?=${ + terminators.map((word) => String.raw`\b${word}\b`).join('|') + }|$)`, + 'gis', + ); + const bodies: string[] = []; + let match: RegExpExecArray | null; + while ((match = pattern.exec(sql))) { + bodies.push(match.groups?.body || ''); + } + return bodies; +}; + +const extractPotentialUnqualifiedColumnReferences = (sql: string) => { + const bodies = [ + ...extractClauseBody(sql, 'WHERE', [ + 'GROUP\\s+BY', + 'ORDER\\s+BY', + 'HAVING', + 'LIMIT', + 'FETCH', + 'UNION', + ]), + ...extractClauseBody(sql, 'HAVING', [ + 'GROUP\\s+BY', + 'ORDER\\s+BY', + 'LIMIT', + 'FETCH', + 'UNION', + ]), + ...extractClauseBody(sql, 'ON', [ + 'WHERE', + 'GROUP\\s+BY', + 'ORDER\\s+BY', + 'HAVING', + 'JOIN', + 'LIMIT', + 'FETCH', + 'UNION', + ]), + ...extractClauseBody(sql, 'GROUP\\s+BY', [ + 'ORDER\\s+BY', + 'HAVING', + 'LIMIT', + 'FETCH', + 'UNION', + ]), + ]; + const identifiers = new Set(); + const identifierPattern = new RegExp(SQL_IDENTIFIER_PATTERN, 'gi'); + + for (const body of bodies) { + const searchableBody = stripSqlLiterals(body); + let match: RegExpExecArray | null; + while ((match = identifierPattern.exec(searchableBody))) { + const token = match[0]; + const before = searchableBody.slice(0, match.index).trimEnd(); + const after = searchableBody.slice(match.index + token.length).trimStart(); + const identifier = normalizeSqlIdentifier(token); + const normalized = identifier.toLowerCase(); + if ( + !identifier || + SQL_NON_COLUMN_IDENTIFIERS.has(normalized) || + before.endsWith('.') || + after.startsWith('.') || + after.startsWith('(') + ) { + continue; + } + identifiers.add(identifier); + } + } + + const functionArgumentPattern = new RegExp( + String.raw`\b[A-Za-z_][A-Za-z0-9_$]*\s*\(\s*(?:DISTINCT\s+)?(${SQL_IDENTIFIER_PATTERN})(?:\s*\.\s*(${SQL_IDENTIFIER_PATTERN}))?`, + 'gi', + ); + let match: RegExpExecArray | null; + while ((match = functionArgumentPattern.exec(stripSqlLiterals(sql)))) { + const column = normalizeSqlIdentifier(match[2] || match[1]); + if ( + column && + column !== '*' && + !SQL_NON_COLUMN_IDENTIFIERS.has(column.toLowerCase()) + ) { + identifiers.add(column); + } + } + + return [...identifiers]; +}; + const findSqlReferenceValidationErrors = ( sql: string, manifest?: Manifest, @@ -684,6 +836,24 @@ const findSqlReferenceValidationErrors = ( errors.push(column); } }); + + const tableAliases = new Set( + [...aliases.entries()] + .filter(([, aliasSchema]) => aliasSchema === schema) + .map(([alias]) => alias.toLowerCase()), + ); + const validCompactColumns = new Set([...schema.columns.keys()]); + extractPotentialUnqualifiedColumnReferences(sql).forEach((column) => { + const normalizedColumn = column.toLowerCase(); + if ( + tableAliases.has(normalizedColumn) || + schema.columns.has(normalizedColumn) || + validCompactColumns.has(compactSqlIdentifier(column)) + ) { + return; + } + errors.push(column); + }); } return [...new Set(errors)]; diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 8aa5e6cd58..38c5d2fdef 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -457,6 +457,59 @@ describe('QueryService', () => { expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); }); + it('should reject unqualified filter columns outside a single active manifest table', async () => { + await expect( + queryService.preview( + 'SELECT id FROM "policies" WHERE policy_category_id = 1', + { + project: { type: DataSourceName.MSSQL, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'policies', + tableReference: { table: 'policies' }, + columns: [ + { name: 'id', type: 'integer', isCalculated: false }, + { name: 'policy_name', type: 'string', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }, + ), + ).rejects.toThrow( + 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: policy_category_id', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + + it('should reject unknown unqualified function argument columns before ibis planning', async () => { + await expect( + queryService.preview('SELECT COUNT(policy_category_id) FROM "policies"', { + project: { type: DataSourceName.MSSQL, connectionInfo: {} }, + manifest: { + models: [ + { + name: 'policies', + tableReference: { table: 'policies' }, + columns: [ + { name: 'id', type: 'integer', isCalculated: false }, + { name: 'policy_name', type: 'string', isCalculated: false }, + ], + }, + ], + }, + dryRun: true, + }), + ).rejects.toThrow( + 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: policy_category_id', + ); + + expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + }); + it('should reject numeric aggregates on non-numeric manifest columns before ibis planning', async () => { await expect( queryService.preview('SELECT AVG("orders"."quantity") FROM "orders"', { From 2049ab435e4a4e8653cda8585aeb23c2223e8c90 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 16:23:46 +0530 Subject: [PATCH 0454/1087] Remove stale datasource columns during schema sync --- .../managers/dataSourceSchemaDetector.ts | 25 ++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index d1c262aa5c..b8e41b2a3e 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -716,9 +716,32 @@ export default class DataSourceSchemaDetector const existingColumns = modelColumns.filter( (column) => column.modelId === model.id && !column.isCalculated, ); + const latestColumnNames = new Set( + latestTable.columns.map((column) => column.name), + ); + const staleColumnNames = existingColumns + .filter((column) => !latestColumnNames.has(column.sourceColumnName)) + .map((column) => column.sourceColumnName); + if (staleColumnNames.length) { + logger.info( + `Removing stale datasource column metadata "${staleColumnNames.join( + ', ', + )}" from model "${model.referenceName}".`, + ); + await this.ctx.modelColumnRepository.deleteAllBySourceColumnNames( + model.id, + staleColumnNames, + ); + hasSyncedMetadata = true; + } + const usedReferenceNames = new Set( modelColumns - .filter((column) => column.modelId === model.id) + .filter( + (column) => + column.modelId === model.id && + !staleColumnNames.includes(column.sourceColumnName), + ) .map((column) => column.referenceName.toLowerCase()), ); From 8aa400016bf5d2cb57850659d4b26efc7c939a55 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 17:09:20 +0530 Subject: [PATCH 0455/1087] Route explicit repair log count questions --- wren-ai-service/src/web/v1/services/ask.py | 32 +++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 40 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2e95e4f91e..9cbfa693be 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1314,6 +1314,34 @@ def _build_schema_grounded_table_question_sql( if wants_total_count: return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' + wants_grouped_count = ( + any( + term in normalized + for term in ( + "count", + "counts", + "record count", + "number of", + "how many", + ) + ) + and re.search(r"\b(?:by|per|each|grouped by|group by)\b", normalized) + ) + if wants_grouped_count: + dimension_column = self._find_dimension_column_for_query(query, table) + if not dimension_column: + return None + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" + return ( + f"SELECT {dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_column)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f"ORDER BY COUNT(*) DESC" + ) + wants_distribution = any( term in normalized for term in ( @@ -1444,6 +1472,10 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: and table_name not in table_names ): table_names.append(table_name) + if re.search(r"\brepair\s+logs?\b", query or "", flags=re.IGNORECASE): + for table_name in ("repair_logs", "dbo_repair_logs"): + if table_name not in table_names: + table_names.append(table_name) return table_names def _build_direct_orders_sales_sql(self, query: str) -> str | None: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 73dceee9b0..f1ac4575c8 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -146,6 +146,14 @@ def test_extract_explicit_table_names_from_using_clause(): ) == [] +def test_extract_explicit_table_names_from_repair_logs_phrase(): + service = AskService.__new__(AskService) + + assert service._extract_explicit_table_names_from_query( + "Count repair logs by failure_code in repair logs." + ) == ["repair_logs", "dbo_repair_logs"] + + def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): service = AskService.__new__(AskService) documents = [ @@ -1373,6 +1381,38 @@ def test_build_schema_grounded_table_question_sql_for_name_distribution(): ) +def test_build_schema_grounded_table_question_sql_for_repair_log_failure_code_count(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Count repair logs by failure_code in repair logs.", + [ + """ + CREATE TABLE dbo_repair_logs ( + org_id VARCHAR, + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP, + updated_at TIMESTAMP, + data JSON + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_repair_logs"."failure_code" AS "failure_code", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' + 'GROUP BY "dbo_repair_logs"."failure_code" ' + 'ORDER BY COUNT(*) DESC' + ) + + def test_build_schema_grounded_table_question_sql_for_numeric_column_distribution(): service = AskService.__new__(AskService) From 142d8521d9262579d36ad1b167085a050f211bba Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 17:29:13 +0530 Subject: [PATCH 0456/1087] Route PCB repair and ticket questions --- wren-ai-service/src/web/v1/services/ask.py | 215 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 157 +++++++++++++ 2 files changed, 371 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9cbfa693be..fa01b27c12 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1460,6 +1460,14 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_name = match.group(1).strip(".,;:()[]{}") if table_name and table_name not in table_names: table_names.append(table_name) + for match in re.finditer( + r"\bin\s+(?:the\s+)?([A-Za-z_][A-Za-z0-9_.$]*)\s+table\b", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if table_name and table_name not in table_names: + table_names.append(table_name) for match in re.finditer( r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", @@ -1472,10 +1480,18 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: and table_name not in table_names ): table_names.append(table_name) - if re.search(r"\brepair\s+logs?\b", query or "", flags=re.IGNORECASE): + if re.search( + r"\b(?:repair\s+logs?|repair\s+tickets?|board\s+models?)\b", + query or "", + flags=re.IGNORECASE, + ): for table_name in ("repair_logs", "dbo_repair_logs"): if table_name not in table_names: table_names.append(table_name) + if re.search(r"\bticket\s+labels?\b", query or "", flags=re.IGNORECASE): + for table_name in ("ticket_labels", "dbo_ticket_labels"): + if table_name not in table_names: + table_names.append(table_name) return table_names def _build_direct_orders_sales_sql(self, query: str) -> str | None: @@ -1779,6 +1795,9 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): + return pcb_direct_sql + if repair_failure_count_sql := self._build_repair_failure_count_sql( query, table_ddls ): @@ -3874,6 +3893,200 @@ def _build_schema_grounded_operational_sql( return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' + def _build_pcb_direct_question_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + repair_table = next( + ( + table + for table in tables + if str(table.get("name") or "").lower() == "dbo_repair_logs" + ), + None, + ) + ticket_label_table = next( + ( + table + for table in tables + if str(table.get("name") or "").lower() == "dbo_ticket_labels" + ), + None, + ) + limit = self._extract_requested_top_n(query, default_value=10) + + if ticket_label_table and "ticket" in normalized and "label" in normalized: + label_column = self._find_first_schema_column( + ticket_label_table, + ("name", "label", "title", "value", "id"), + ) + if label_column: + table_name = str(ticket_label_table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + label_ref = f"{table_ref}.{self._quote_sql_identifier(label_column)}" + return ( + f"SELECT TOP {limit} {label_ref} AS " + f"{self._quote_sql_identifier(label_column)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {label_ref} IS NOT NULL " + f"GROUP BY {label_ref} " + f"ORDER BY COUNT(*) DESC" + ) + + if not repair_table: + return None + + table_name = str(repair_table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + board_model_column = self._find_schema_column( + repair_table, ("board_model", "boardModel", "board model", "product") + ) + failure_code_column = self._find_schema_column( + repair_table, ("failure_code", "failureCode", "failure code", "failure") + ) + created_at_column = self._find_schema_column( + repair_table, + ("created_at", "createdAt", "created", "date_received", "dateReceived"), + temporal=True, + ) + priority_column = self._find_schema_column(repair_table, ("priority",)) + status_column = self._find_schema_column(repair_table, ("status",)) + id_column = self._find_schema_column(repair_table, ("id", "repair_id")) + + asks_board_model_distribution = ( + "board model" in normalized + and any(term in normalized for term in ("distribution", "over time", "trend")) + ) + if asks_board_model_distribution and board_model_column and created_at_column: + board_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" + date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" + return ( + f"SELECT {board_ref} AS " + f"{self._quote_sql_identifier(board_model_column)}, " + f"DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {board_ref} IS NOT NULL " + f"AND {date_ref} IS NOT NULL " + f"GROUP BY {board_ref}, DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " + f"DATEPART(MONTH, {date_ref}) ASC, {board_ref} ASC" + ) + + asks_recurring_failures_by_product = ( + "recurring" in normalized + and "failure" in normalized + and ("product" in normalized or "pcb" in normalized) + ) + if ( + asks_recurring_failures_by_product + and board_model_column + and failure_code_column + ): + product_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" + failure_ref = f"{table_ref}.{self._quote_sql_identifier(failure_code_column)}" + return ( + f"SELECT {product_ref} AS " + f"{self._quote_sql_identifier(board_model_column)}, " + f"{failure_ref} AS " + f"{self._quote_sql_identifier(failure_code_column)}, " + f'COUNT(*) AS "failure_count" ' + f"FROM {table_ref} " + f"WHERE {product_ref} IS NOT NULL " + f"AND {failure_ref} IS NOT NULL " + f"GROUP BY {product_ref}, {failure_ref} " + f'ORDER BY "failure_count" DESC' + ) + + asks_highest_priority_repairs = ( + "repair" in normalized + and "priority" in normalized + and any(term in normalized for term in ("highest", "top", "high priority")) + ) + if asks_highest_priority_repairs and priority_column: + priority_ref = f"{table_ref}.{self._quote_sql_identifier(priority_column)}" + select_refs = [] + for column in ( + id_column, + board_model_column, + failure_code_column, + status_column, + priority_column, + created_at_column, + ): + if column and column not in select_refs: + select_refs.append(column) + select_sql = ", ".join( + f"{table_ref}.{self._quote_sql_identifier(column)} AS " + f"{self._quote_sql_identifier(column)}" + for column in select_refs + ) + return ( + f"SELECT TOP {limit} {select_sql} " + f"FROM {table_ref} " + f"WHERE {priority_ref} IS NOT NULL " + f"ORDER BY CASE LOWER({priority_ref}) " + f"WHEN 'critical' THEN 1 " + f"WHEN 'high' THEN 2 " + f"WHEN 'medium' THEN 3 " + f"WHEN 'low' THEN 4 " + f"ELSE 5 END" + ) + + asks_repair_ticket_distribution = ( + "repair" in normalized + and "ticket" in normalized + and ( + "distribution" in normalized + or "again distribution" in normalized + or "aging distribution" in normalized + ) + ) + if asks_repair_ticket_distribution: + if "aging" in normalized and created_at_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" + age_bucket = ( + f"CASE " + f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 7 THEN '0-7 days' " + f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 30 THEN '8-30 days' " + f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 90 THEN '31-90 days' " + f"ELSE '90+ days' END" + ) + return ( + f'SELECT {age_bucket} AS "age_bucket", ' + f'COUNT(*) AS "ticket_count" ' + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"GROUP BY {age_bucket} " + f'ORDER BY "ticket_count" DESC' + ) + distribution_column = status_column or priority_column or failure_code_column + if distribution_column: + dimension_ref = ( + f"{table_ref}.{self._quote_sql_identifier(distribution_column)}" + ) + return ( + f"SELECT {dimension_ref} AS " + f"{self._quote_sql_identifier(distribution_column)}, " + f'COUNT(*) AS "ticket_count" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f'ORDER BY "ticket_count" DESC' + ) + + return None + def _get_unqueryable_metric_message( self, query: str, table_ddls: list[str] ) -> str | None: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index f1ac4575c8..66b0694202 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -154,6 +154,17 @@ def test_extract_explicit_table_names_from_repair_logs_phrase(): ) == ["repair_logs", "dbo_repair_logs"] +def test_extract_explicit_table_names_from_pcb_repair_phrases(): + service = AskService.__new__(AskService) + + assert service._extract_explicit_table_names_from_query( + "How many different board models are present in the dbo.repair_logs table?" + ) == ["dbo.repair_logs", "repair_logs", "dbo_repair_logs"] + assert service._extract_explicit_table_names_from_query( + "Display top 10 ticket labels." + ) == ["ticket_labels", "dbo_ticket_labels"] + + def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): service = AskService.__new__(AskService) documents = [ @@ -1413,6 +1424,152 @@ def test_build_schema_grounded_table_question_sql_for_repair_log_failure_code_co ) +def test_build_pcb_direct_question_sql_for_board_model_distribution_over_time(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "How many different board models are present in the dbo.repair_logs table, and what is their distribution over time?", + [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_repair_logs"."board_model" AS "board_model", ' + 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."board_model" IS NOT NULL ' + 'AND "dbo_repair_logs"."created_at" IS NOT NULL ' + 'GROUP BY "dbo_repair_logs"."board_model", ' + 'DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' + 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' + 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC, ' + '"dbo_repair_logs"."board_model" ASC' + ) + + +def test_build_pcb_direct_question_sql_for_recurring_pcb_failures_by_product(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Show recurring PCB failures by product.", + [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_repair_logs"."board_model" AS "board_model", ' + '"dbo_repair_logs"."failure_code" AS "failure_code", ' + 'COUNT(*) AS "failure_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."board_model" IS NOT NULL ' + 'AND "dbo_repair_logs"."failure_code" IS NOT NULL ' + 'GROUP BY "dbo_repair_logs"."board_model", ' + '"dbo_repair_logs"."failure_code" ' + 'ORDER BY "failure_count" DESC' + ) + + +def test_build_pcb_direct_question_sql_for_highest_priority_repairs(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Which repair logs have the highest priority?", + [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql is not None + assert sql.startswith('SELECT TOP 10 "dbo_repair_logs"."id" AS "id"') + assert 'FROM "dbo_repair_logs"' in sql + assert 'CASE LOWER("dbo_repair_logs"."priority")' in sql + assert "WHEN 'critical' THEN 1" in sql + assert "WHEN 'high' THEN 2" in sql + + +def test_build_pcb_direct_question_sql_for_repair_ticket_distribution(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Show repair ticket again distribution.", + [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_repair_logs"."status" AS "status", ' + 'COUNT(*) AS "ticket_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."status" IS NOT NULL ' + 'GROUP BY "dbo_repair_logs"."status" ' + 'ORDER BY "ticket_count" DESC' + ) + + +def test_build_pcb_direct_question_sql_for_top_ticket_labels(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Display top 10 ticket labels.", + [ + """ + CREATE TABLE dbo_ticket_labels ( + id VARCHAR, + name VARCHAR, + created_at TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_ticket_labels"."name" AS "name", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ticket_labels" ' + 'WHERE "dbo_ticket_labels"."name" IS NOT NULL ' + 'GROUP BY "dbo_ticket_labels"."name" ' + 'ORDER BY COUNT(*) DESC' + ) + + def test_build_schema_grounded_table_question_sql_for_numeric_column_distribution(): service = AskService.__new__(AskService) From 20f7312d3a197c0d9e8f48550b7c61a9058fc05e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 17:38:08 +0530 Subject: [PATCH 0457/1087] Retry full schema for data questions --- wren-ai-service/src/web/v1/services/ask.py | 89 +++++++++++++++++++++- 1 file changed, 87 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index fa01b27c12..2d391eb866 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -6031,8 +6031,93 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if unqueryable_metric_message := self._get_unqueryable_metric_message( - user_query, table_ddls + should_retry_full_schema = ( + not api_results + and self._is_data_analysis_query(user_query) + and "db_schema_retrieval" in self._pipelines + ) + if should_retry_full_schema: + logger.info( + "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retry", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 30, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + full_documents, full_table_names, full_table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + full_documents, full_table_names, full_table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + full_documents, + ) + ) + if full_documents: + documents, table_names, table_ddls = ( + full_documents, + full_table_names, + full_table_ddls, + ) + logger.info( + "Using full active deployed schema retry for query_id %s: %s", + query_id, + table_names, + ) + + full_schema_preview = self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + full_schema_sql_candidates = ( + self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ), + full_schema_preview[0] if full_schema_preview else None, + self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ), + self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ), + self._build_schema_grounded_sales_sql( + user_query, table_ddls + ), + ) + for full_schema_sql in full_schema_sql_candidates: + if not full_schema_sql: + continue + ask_result = self._build_validated_ask_result_from_sql( + full_schema_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + break + invalid_sql = full_schema_sql + error_message = ( + "Full-schema grounded SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and ( + unqueryable_metric_message := self._get_unqueryable_metric_message( + user_query, table_ddls + ) ): logger.info( "ask pipeline - NO_RELEVANT_SQL due to unqueryable metric: %s", From 92fa826722c80e1deaca0392c25832eebc51612c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 18:44:03 +0530 Subject: [PATCH 0458/1087] Honor explicit in-table questions --- wren-ai-service/src/web/v1/services/ask.py | 12 ++++++ .../pytest/services/test_ask_sales_sql.py | 42 +++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2d391eb866..4d95e6dae5 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1468,6 +1468,18 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_name = match.group(1).strip(".,;:()[]{}") if table_name and table_name not in table_names: table_names.append(table_name) + for match in re.finditer( + r"\bin\s+([A-Za-z_][A-Za-z0-9_.$]*)", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if ( + table_name + and ("." in table_name or "_" in table_name) + and table_name not in table_names + ): + table_names.append(table_name) for match in re.finditer( r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 66b0694202..42d46a3483 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -146,6 +146,17 @@ def test_extract_explicit_table_names_from_using_clause(): ) == [] +def test_extract_explicit_table_names_from_in_clause(): + service = AskService.__new__(AskService) + + assert service._extract_explicit_table_names_from_query( + "Which customers have the highest number of orders in dbo.tblNewOrders?" + ) == ["dbo.tblNewOrders"] + assert service._extract_explicit_table_names_from_query( + "Which customers have the highest number of orders in market?" + ) == [] + + def test_extract_explicit_table_names_from_repair_logs_phrase(): service = AskService.__new__(AskService) @@ -1293,6 +1304,37 @@ def test_build_validated_ask_result_rejects_status_when_failure_field_matches_qu assert result is None +def test_build_validated_ask_result_rejects_sql_for_wrong_explicit_table(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_qMarginSales"."Customer" AS "Customer", ' + 'COUNT(DISTINCT "dbo_qMarginSales"."OrdNo") AS "OrderCount" ' + 'FROM "dbo_qMarginSales" ' + 'WHERE "dbo_qMarginSales"."Customer" IS NOT NULL ' + 'GROUP BY "dbo_qMarginSales"."Customer"' + ), + [ + """ + CREATE TABLE dbo_tblNewOrders ( + Customer VARCHAR, + OrdNo VARCHAR + ); + """, + """ + CREATE TABLE dbo_qMarginSales ( + Customer VARCHAR, + OrdNo VARCHAR + ); + """, + ], + "Which customers have the highest number of orders in dbo.tblNewOrders?", + ) + + assert result is None + + def test_reusable_historical_question_allows_exact_recommended_question(): assert AskService._is_reusable_historical_question( "How many tickets are currently open?", From d681ff16832ad8577014051ec9ab615aee266df8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 18:57:37 +0530 Subject: [PATCH 0459/1087] Build ranked counts for explicit tables --- wren-ai-service/src/web/v1/services/ask.py | 26 +++++++++++++++--- .../pytest/services/test_ask_sales_sql.py | 27 +++++++++++++++++++ 2 files changed, 49 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 4d95e6dae5..a66812b5ae 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1327,19 +1327,37 @@ def _build_schema_grounded_table_question_sql( ) and re.search(r"\b(?:by|per|each|grouped by|group by)\b", normalized) ) - if wants_grouped_count: + wants_ranked_count = any( + term in normalized for term in ("highest", "top", "most", "largest") + ) and any( + term in normalized + for term in ("count", "counts", "number of", "orders", "records", "rows") + ) + if wants_grouped_count or wants_ranked_count: dimension_column = self._find_dimension_column_for_query(query, table) if not dimension_column: return None dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" + count_column = None + if any(term in normalized for term in ("order", "orders")): + count_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "OrderID", "id"), + ) + count_expression = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(count_column)})" + if count_column + else "COUNT(*)" + ) return ( - f"SELECT {dimension_ref} AS " + f"SELECT TOP {limit} {dimension_ref} AS " f"{self._quote_sql_identifier(dimension_column)}, " - f'COUNT(*) AS "RecordCount" ' + f'{count_expression} AS "RecordCount" ' f"FROM {table_ref} " f"WHERE {dimension_ref} IS NOT NULL " + f"AND LTRIM(RTRIM({dimension_ref})) <> '' " f"GROUP BY {dimension_ref} " - f"ORDER BY COUNT(*) DESC" + f"ORDER BY {count_expression} DESC" ) wants_distribution = any( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 42d46a3483..d479f8217f 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1466,6 +1466,33 @@ def test_build_schema_grounded_table_question_sql_for_repair_log_failure_code_co ) +def test_build_schema_grounded_table_question_sql_for_highest_orders_by_customer(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "Which customers have the highest number of orders in dbo.tblNewOrders?", + [ + """ + CREATE TABLE dbo_tblNewOrders ( + Customer VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "dbo_tblNewOrders"."Customer" AS "Customer", ' + 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "RecordCount" ' + 'FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."Customer" IS NOT NULL ' + 'AND LTRIM(RTRIM("dbo_tblNewOrders"."Customer")) <> \'\' ' + 'GROUP BY "dbo_tblNewOrders"."Customer" ' + 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' + ) + + def test_build_pcb_direct_question_sql_for_board_model_distribution_over_time(): service = AskService.__new__(AskService) From 5c7e0d5b78905956d5e45a35074b96f930d4232f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 19:09:31 +0530 Subject: [PATCH 0460/1087] Keep related tables in SQL context --- wren-ai-service/src/web/v1/services/ask.py | 102 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 42 ++++++++ 2 files changed, 139 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a66812b5ae..d6116e6050 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4865,12 +4865,22 @@ def _prune_sql_generation_context( if not scored: return documents, table_names, table_ddls[:max_tables] - selected_indexes = [ - index - for _, index in sorted(scored, key=lambda item: item[0], reverse=True)[ - :max_tables - ] + sorted_scored_indexes = [ + index for _, index in sorted(scored, key=lambda item: item[0], reverse=True) ] + core_limit = max(1, max_tables - 2) if max_tables > 2 else 1 + selected_indexes = sorted_scored_indexes[:core_limit] + selected_indexes = self._expand_pruned_context_with_related_tables( + selected_indexes, + parsed_tables, + table_ddls, + max_tables=max_tables, + ) + for index in sorted_scored_indexes: + if len(selected_indexes) >= max_tables: + break + if index not in selected_indexes: + selected_indexes.append(index) selected_indexes = sorted(selected_indexes) pruned_documents = [ documents[index] for index in selected_indexes if index < len(documents) @@ -4890,6 +4900,88 @@ def _prune_sql_generation_context( ) return pruned_documents, pruned_table_names, pruned_table_ddls + def _expand_pruned_context_with_related_tables( + self, + selected_indexes: list[int], + parsed_tables: list[dict[str, Any]], + table_ddls: list[str], + *, + max_tables: int, + ) -> list[int]: + if len(selected_indexes) >= max_tables: + return selected_indexes[:max_tables] + + selected: list[int] = list(dict.fromkeys(selected_indexes)) + selected_set = set(selected) + + def join_key_columns(table: dict[str, Any]) -> set[str]: + keys = set() + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + normalized = self._normalize_schema_token(column_name) + if not normalized: + continue + if ( + normalized == "id" + or normalized.endswith("id") + or normalized.endswith("no") + or normalized.endswith("number") + or normalized.endswith("code") + or normalized.endswith("key") + ): + keys.add(normalized) + return keys + + selected_table_names = { + self._normalize_schema_token( + str(parsed_tables[index].get("name") or "") + ) + for index in selected + if index < len(parsed_tables) + } + selected_join_keys: set[str] = set() + for index in selected: + if index < len(parsed_tables): + selected_join_keys.update(join_key_columns(parsed_tables[index])) + + candidates: list[tuple[int, int]] = [] + for index, table in enumerate(parsed_tables): + if index in selected_set: + continue + + table_name = str(table.get("name") or "") + normalized_table_name = self._normalize_schema_token(table_name) + ddl = table_ddls[index] if index < len(table_ddls) else "" + normalized_ddl = self._normalize_schema_token(ddl) + table_join_keys = join_key_columns(table) + + score = 0 + shared_keys = selected_join_keys & table_join_keys + if shared_keys: + score += 20 + 5 * len(shared_keys) + if normalized_table_name and any( + selected_table + and ( + selected_table in normalized_ddl + or normalized_table_name in selected_table + ) + for selected_table in selected_table_names + ): + score += 40 + if re.search(r"\b(?:foreign\s+key|references)\b", ddl, flags=re.IGNORECASE): + score += 15 + + if score > 0: + candidates.append((score, index)) + + for _, index in sorted(candidates, key=lambda item: item[0], reverse=True): + if len(selected) >= max_tables: + break + selected.append(index) + selected_set.add(index) + + return selected + def _is_valid_select_sql(self, sql: Optional[str]) -> bool: if not isinstance(sql, str): return False diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index d479f8217f..2b31929abb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -321,6 +321,48 @@ def test_prune_sql_generation_context_prefers_referenced_table_and_columns(): assert pruned_ddls == [table_ddls[2]] +def test_prune_sql_generation_context_keeps_related_join_table(): + service = AskService.__new__(AskService) + table_ddls = [ + """ + CREATE TABLE dbo_Customers ( + CustomerId VARCHAR, + CustomerName VARCHAR + ); + """, + """ + CREATE TABLE dbo_Products ( + ProductId VARCHAR, + ProductName VARCHAR + ); + """, + """ + CREATE TABLE dbo_Orders ( + OrderId VARCHAR, + CustomerId VARCHAR, + ProductId VARCHAR, + OrderDate TIMESTAMP, + FOREIGN KEY (CustomerId) REFERENCES dbo_Customers(CustomerId) + ); + """, + ] + documents = [ + {"table_name": "dbo_Customers", "table_ddl": table_ddls[0]}, + {"table_name": "dbo_Products", "table_ddl": table_ddls[1]}, + {"table_name": "dbo_Orders", "table_ddl": table_ddls[2]}, + ] + + _, table_names, _ = service._prune_sql_generation_context( + "Which customer names have the highest number of orders?", + documents, + [document["table_name"] for document in documents], + table_ddls, + max_tables=2, + ) + + assert table_names == ["dbo_Customers", "dbo_Orders"] + + def test_build_schema_grounded_sales_sql_for_top_markets(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 7e545887dbf923aa27b8ca64ae04e03e55f377fa Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 20:49:30 +0530 Subject: [PATCH 0461/1087] Score analytics tables by user intent --- wren-ai-service/src/web/v1/services/ask.py | 55 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 29 ++++++++++ 2 files changed, 84 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d6116e6050..9db38bd5a8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1745,7 +1745,9 @@ def _select_best_analytics_table( measure_candidates: tuple[str, ...], wants_date: bool = False, allow_count_metric: bool = False, + query: str = "", ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] ] = [] @@ -1796,6 +1798,54 @@ def _select_best_analytics_table( score += 3 if "stage" in table_name: score -= 8 + if any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ): + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber"), + ) + if "order" in table_name: + score += 30 + if "neworder" in self._normalize_schema_token(table_name): + score += 15 + if order_column: + score += 12 + if "margin" in table_name and "margin" not in normalized_query: + score -= 12 + if "customer" in normalized_query: + if "customer" in table_name or "account" in table_name: + score += 16 + if self._find_schema_column( + table, + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + ), + ): + score += 10 + if any( + term in normalized_query for term in ("product", "products", "item") + ): + if "product" in table_name or "item" in table_name: + score += 16 + if any( + term in normalized_query + for term in ("sales", "revenue", "value", "amount") + ): + if "sales" in table_name: + score += 12 + if "invoice" in normalized_query and ( + "invoice" in table_name or "inv" in table_name + ): + score += 20 scored.append((score, table, dimensions, measure, date_column)) @@ -2110,6 +2160,7 @@ def _build_schema_grounded_analytics_sql( (), wants_date=False, allow_count_metric=True, + query=query, ) if selected: table, dimensions, _measure, _date_column = selected @@ -2160,6 +2211,7 @@ def _build_schema_grounded_analytics_sql( measure_candidates, wants_date=True, allow_count_metric=wants_count_metric, + query=query, ) if not selected: return None @@ -2201,6 +2253,7 @@ def _build_schema_grounded_analytics_sql( measure_candidates, wants_date=wants_date, allow_count_metric=wants_order_count_metric or wants_count_metric, + query=query, ) if not selected: return None @@ -2507,6 +2560,7 @@ def _build_contribution_sql( "Amount", ), wants_date=False, + query=query, ) if not selected: return None @@ -2763,6 +2817,7 @@ def _build_yoy_sales_change_sql( "Amount", ), wants_date=False, + query=query, ) if not selected: return None diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2b31929abb..3c26e551c6 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1535,6 +1535,35 @@ def test_build_schema_grounded_table_question_sql_for_highest_orders_by_customer ) +def test_schema_grounded_analytics_prefers_order_table_for_order_count_question(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Which customers have the highest number of orders?", + [ + """ + CREATE TABLE dbo_qMarginSales ( + Customer VARCHAR, + OrdNo VARCHAR, + SalesValue DOUBLE + ); + """, + """ + CREATE TABLE dbo_tblNewOrders ( + Customer VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """, + ], + ) + + assert sql is not None + assert 'FROM "dbo_tblNewOrders"' in sql + assert 'FROM "dbo_qMarginSales"' not in sql + assert 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo")' in sql + + def test_build_pcb_direct_question_sql_for_board_model_distribution_over_time(): service = AskService.__new__(AskService) From 6df8e2acb41db0e408dc91ec1435c8dc2b6a11ae Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 21:17:36 +0530 Subject: [PATCH 0462/1087] Revert "Score analytics tables by user intent" This reverts commit 7e545887dbf923aa27b8ca64ae04e03e55f377fa. --- wren-ai-service/src/web/v1/services/ask.py | 55 ------------------- .../pytest/services/test_ask_sales_sql.py | 29 ---------- 2 files changed, 84 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9db38bd5a8..d6116e6050 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1745,9 +1745,7 @@ def _select_best_analytics_table( measure_candidates: tuple[str, ...], wants_date: bool = False, allow_count_metric: bool = False, - query: str = "", ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] ] = [] @@ -1798,54 +1796,6 @@ def _select_best_analytics_table( score += 3 if "stage" in table_name: score -= 8 - if any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ): - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber"), - ) - if "order" in table_name: - score += 30 - if "neworder" in self._normalize_schema_token(table_name): - score += 15 - if order_column: - score += 12 - if "margin" in table_name and "margin" not in normalized_query: - score -= 12 - if "customer" in normalized_query: - if "customer" in table_name or "account" in table_name: - score += 16 - if self._find_schema_column( - table, - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - ), - ): - score += 10 - if any( - term in normalized_query for term in ("product", "products", "item") - ): - if "product" in table_name or "item" in table_name: - score += 16 - if any( - term in normalized_query - for term in ("sales", "revenue", "value", "amount") - ): - if "sales" in table_name: - score += 12 - if "invoice" in normalized_query and ( - "invoice" in table_name or "inv" in table_name - ): - score += 20 scored.append((score, table, dimensions, measure, date_column)) @@ -2160,7 +2110,6 @@ def _build_schema_grounded_analytics_sql( (), wants_date=False, allow_count_metric=True, - query=query, ) if selected: table, dimensions, _measure, _date_column = selected @@ -2211,7 +2160,6 @@ def _build_schema_grounded_analytics_sql( measure_candidates, wants_date=True, allow_count_metric=wants_count_metric, - query=query, ) if not selected: return None @@ -2253,7 +2201,6 @@ def _build_schema_grounded_analytics_sql( measure_candidates, wants_date=wants_date, allow_count_metric=wants_order_count_metric or wants_count_metric, - query=query, ) if not selected: return None @@ -2560,7 +2507,6 @@ def _build_contribution_sql( "Amount", ), wants_date=False, - query=query, ) if not selected: return None @@ -2817,7 +2763,6 @@ def _build_yoy_sales_change_sql( "Amount", ), wants_date=False, - query=query, ) if not selected: return None diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 3c26e551c6..2b31929abb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1535,35 +1535,6 @@ def test_build_schema_grounded_table_question_sql_for_highest_orders_by_customer ) -def test_schema_grounded_analytics_prefers_order_table_for_order_count_question(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Which customers have the highest number of orders?", - [ - """ - CREATE TABLE dbo_qMarginSales ( - Customer VARCHAR, - OrdNo VARCHAR, - SalesValue DOUBLE - ); - """, - """ - CREATE TABLE dbo_tblNewOrders ( - Customer VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """, - ], - ) - - assert sql is not None - assert 'FROM "dbo_tblNewOrders"' in sql - assert 'FROM "dbo_qMarginSales"' not in sql - assert 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo")' in sql - - def test_build_pcb_direct_question_sql_for_board_model_distribution_over_time(): service = AskService.__new__(AskService) From cd9ac75efcafb80e83c9e22f331b7f0fe5c5d580 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 22:00:40 +0530 Subject: [PATCH 0463/1087] Require SQL to cover question concepts --- wren-ai-service/src/web/v1/services/ask.py | 81 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 71 ++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d6116e6050..742be55b02 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -622,6 +622,87 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: return [] concept_groups: list[set[str]] = [] + compact_query = re.sub(r"[^a-z0-9]", "", normalized) + if ( + "customer" in normalized + or "custname" in compact_query + or "custno" in compact_query + ): + concept_groups.append( + { + "account", + "client", + "cust", + "customer", + "customername", + "custname", + "custno", + } + ) + if "market segment" in normalized: + concept_groups.append( + {"market", "marketsegment", "segment", "markettype", "region"} + ) + elif "market" in normalized: + concept_groups.append({"market", "markettype", "marketname", "segment"}) + if "region" in normalized: + concept_groups.append( + {"region", "regional", "market", "area", "territory", "country"} + ) + if "salesperson" in normalized or "sales person" in normalized: + concept_groups.append( + {"salesperson", "salesman", "salesrep", "rep", "owner"} + ) + if "product type" in normalized or "producttype" in compact_query: + concept_groups.append({"prodtype", "producttype", "product", "type"}) + elif "product" in normalized: + concept_groups.append( + { + "item", + "part", + "prod", + "prodcode", + "prodname", + "product", + "productcode", + "productname", + "sku", + } + ) + if "quantity" in normalized or re.search(r"\bqty\b", normalized): + concept_groups.append( + {"qty", "quantity", "salesqty", "orderqty", "invoiceqty"} + ) + if any( + term in normalized + for term in ( + "sales", + "sale", + "revenue", + "amount", + "value", + "order value", + "sales value", + ) + ) and not any( + term in normalized for term in ("salesperson", "sales person", "sales rep") + ): + concept_groups.append( + { + "amount", + "fxsalesvalue", + "invoiceamount", + "invoicevalue", + "net", + "revenue", + "salesamount", + "salesvalue", + "total", + "value", + } + ) + if any(term in normalized for term in ("order priority", "priorities")): + concept_groups.append({"priority", "orderpriority"}) if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2b31929abb..c9324dd7c4 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,6 +269,77 @@ def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): assert result is not None +def test_sql_intent_rejects_partial_customer_market_quantity_answer(): + service = AskService.__new__(AskService) + + assert not service._sql_matches_question_intent( + ( + 'SELECT "dbo_tblStageNewOrders3"."Market" AS "Market", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_tblStageNewOrders3" ' + 'GROUP BY "dbo_tblStageNewOrders3"."Market"' + ), + "Show total order quantity by customer and market segment.", + [ + { + "name": "dbo_tblStageNewOrders3", + "columns": [ + {"name": "Market", "type": "VARCHAR"}, + {"name": "Customer", "type": "VARCHAR"}, + {"name": "OrderQty", "type": "FLOAT"}, + ], + } + ], + ) + + +def test_sql_intent_accepts_customer_market_quantity_answer(): + service = AskService.__new__(AskService) + + assert service._sql_matches_question_intent( + ( + 'SELECT "dbo_tblNewOrders"."Customer" AS "Customer", ' + '"dbo_tblNewOrders"."Market" AS "Market", ' + 'SUM("dbo_tblNewOrders"."OrderQty") AS "TotalOrderQty" ' + 'FROM "dbo_tblNewOrders" ' + 'GROUP BY "dbo_tblNewOrders"."Customer", "dbo_tblNewOrders"."Market"' + ), + "Show total order quantity by customer and market segment.", + [ + { + "name": "dbo_tblNewOrders", + "columns": [ + {"name": "Customer", "type": "VARCHAR"}, + {"name": "Market", "type": "VARCHAR"}, + {"name": "OrderQty", "type": "FLOAT"}, + ], + } + ], + ) + + +def test_sql_intent_rejects_product_sales_question_without_sales_or_region(): + service = AskService.__new__(AskService) + + assert not service._sql_matches_question_intent( + ( + 'SELECT "dbo_products"."products" AS "products" ' + 'FROM "dbo_products"' + ), + "Which products have the highest sales by region?", + [ + { + "name": "dbo_products", + "columns": [ + {"name": "products", "type": "VARCHAR"}, + {"name": "Region", "type": "VARCHAR"}, + {"name": "SalesValue", "type": "FLOAT"}, + ], + } + ], + ) + + def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) From 925824d86f55bede8332a84cf117bfc733ea0b2e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 22:19:19 +0530 Subject: [PATCH 0464/1087] Revert "Require SQL to cover question concepts" This reverts commit cd9ac75efcafb80e83c9e22f331b7f0fe5c5d580. --- wren-ai-service/src/web/v1/services/ask.py | 81 ------------------- .../pytest/services/test_ask_sales_sql.py | 71 ---------------- 2 files changed, 152 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 742be55b02..d6116e6050 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -622,87 +622,6 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: return [] concept_groups: list[set[str]] = [] - compact_query = re.sub(r"[^a-z0-9]", "", normalized) - if ( - "customer" in normalized - or "custname" in compact_query - or "custno" in compact_query - ): - concept_groups.append( - { - "account", - "client", - "cust", - "customer", - "customername", - "custname", - "custno", - } - ) - if "market segment" in normalized: - concept_groups.append( - {"market", "marketsegment", "segment", "markettype", "region"} - ) - elif "market" in normalized: - concept_groups.append({"market", "markettype", "marketname", "segment"}) - if "region" in normalized: - concept_groups.append( - {"region", "regional", "market", "area", "territory", "country"} - ) - if "salesperson" in normalized or "sales person" in normalized: - concept_groups.append( - {"salesperson", "salesman", "salesrep", "rep", "owner"} - ) - if "product type" in normalized or "producttype" in compact_query: - concept_groups.append({"prodtype", "producttype", "product", "type"}) - elif "product" in normalized: - concept_groups.append( - { - "item", - "part", - "prod", - "prodcode", - "prodname", - "product", - "productcode", - "productname", - "sku", - } - ) - if "quantity" in normalized or re.search(r"\bqty\b", normalized): - concept_groups.append( - {"qty", "quantity", "salesqty", "orderqty", "invoiceqty"} - ) - if any( - term in normalized - for term in ( - "sales", - "sale", - "revenue", - "amount", - "value", - "order value", - "sales value", - ) - ) and not any( - term in normalized for term in ("salesperson", "sales person", "sales rep") - ): - concept_groups.append( - { - "amount", - "fxsalesvalue", - "invoiceamount", - "invoicevalue", - "net", - "revenue", - "salesamount", - "salesvalue", - "total", - "value", - } - ) - if any(term in normalized for term in ("order priority", "priorities")): - concept_groups.append({"priority", "orderpriority"}) if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index c9324dd7c4..2b31929abb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,77 +269,6 @@ def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): assert result is not None -def test_sql_intent_rejects_partial_customer_market_quantity_answer(): - service = AskService.__new__(AskService) - - assert not service._sql_matches_question_intent( - ( - 'SELECT "dbo_tblStageNewOrders3"."Market" AS "Market", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_tblStageNewOrders3" ' - 'GROUP BY "dbo_tblStageNewOrders3"."Market"' - ), - "Show total order quantity by customer and market segment.", - [ - { - "name": "dbo_tblStageNewOrders3", - "columns": [ - {"name": "Market", "type": "VARCHAR"}, - {"name": "Customer", "type": "VARCHAR"}, - {"name": "OrderQty", "type": "FLOAT"}, - ], - } - ], - ) - - -def test_sql_intent_accepts_customer_market_quantity_answer(): - service = AskService.__new__(AskService) - - assert service._sql_matches_question_intent( - ( - 'SELECT "dbo_tblNewOrders"."Customer" AS "Customer", ' - '"dbo_tblNewOrders"."Market" AS "Market", ' - 'SUM("dbo_tblNewOrders"."OrderQty") AS "TotalOrderQty" ' - 'FROM "dbo_tblNewOrders" ' - 'GROUP BY "dbo_tblNewOrders"."Customer", "dbo_tblNewOrders"."Market"' - ), - "Show total order quantity by customer and market segment.", - [ - { - "name": "dbo_tblNewOrders", - "columns": [ - {"name": "Customer", "type": "VARCHAR"}, - {"name": "Market", "type": "VARCHAR"}, - {"name": "OrderQty", "type": "FLOAT"}, - ], - } - ], - ) - - -def test_sql_intent_rejects_product_sales_question_without_sales_or_region(): - service = AskService.__new__(AskService) - - assert not service._sql_matches_question_intent( - ( - 'SELECT "dbo_products"."products" AS "products" ' - 'FROM "dbo_products"' - ), - "Which products have the highest sales by region?", - [ - { - "name": "dbo_products", - "columns": [ - {"name": "products", "type": "VARCHAR"}, - {"name": "Region", "type": "VARCHAR"}, - {"name": "SalesValue", "type": "FLOAT"}, - ], - } - ], - ) - - def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) From 6cdf2a2966a3b685f28e3f9cfffdd68a6f523ca4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 22:35:11 +0530 Subject: [PATCH 0465/1087] Reapply "Score analytics tables by user intent" This reverts commit 6df8e2acb41db0e408dc91ec1435c8dc2b6a11ae. --- wren-ai-service/src/web/v1/services/ask.py | 55 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 29 ++++++++++ 2 files changed, 84 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d6116e6050..9db38bd5a8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1745,7 +1745,9 @@ def _select_best_analytics_table( measure_candidates: tuple[str, ...], wants_date: bool = False, allow_count_metric: bool = False, + query: str = "", ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] ] = [] @@ -1796,6 +1798,54 @@ def _select_best_analytics_table( score += 3 if "stage" in table_name: score -= 8 + if any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ): + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber"), + ) + if "order" in table_name: + score += 30 + if "neworder" in self._normalize_schema_token(table_name): + score += 15 + if order_column: + score += 12 + if "margin" in table_name and "margin" not in normalized_query: + score -= 12 + if "customer" in normalized_query: + if "customer" in table_name or "account" in table_name: + score += 16 + if self._find_schema_column( + table, + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + ), + ): + score += 10 + if any( + term in normalized_query for term in ("product", "products", "item") + ): + if "product" in table_name or "item" in table_name: + score += 16 + if any( + term in normalized_query + for term in ("sales", "revenue", "value", "amount") + ): + if "sales" in table_name: + score += 12 + if "invoice" in normalized_query and ( + "invoice" in table_name or "inv" in table_name + ): + score += 20 scored.append((score, table, dimensions, measure, date_column)) @@ -2110,6 +2160,7 @@ def _build_schema_grounded_analytics_sql( (), wants_date=False, allow_count_metric=True, + query=query, ) if selected: table, dimensions, _measure, _date_column = selected @@ -2160,6 +2211,7 @@ def _build_schema_grounded_analytics_sql( measure_candidates, wants_date=True, allow_count_metric=wants_count_metric, + query=query, ) if not selected: return None @@ -2201,6 +2253,7 @@ def _build_schema_grounded_analytics_sql( measure_candidates, wants_date=wants_date, allow_count_metric=wants_order_count_metric or wants_count_metric, + query=query, ) if not selected: return None @@ -2507,6 +2560,7 @@ def _build_contribution_sql( "Amount", ), wants_date=False, + query=query, ) if not selected: return None @@ -2763,6 +2817,7 @@ def _build_yoy_sales_change_sql( "Amount", ), wants_date=False, + query=query, ) if not selected: return None diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 2b31929abb..3c26e551c6 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -1535,6 +1535,35 @@ def test_build_schema_grounded_table_question_sql_for_highest_orders_by_customer ) +def test_schema_grounded_analytics_prefers_order_table_for_order_count_question(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Which customers have the highest number of orders?", + [ + """ + CREATE TABLE dbo_qMarginSales ( + Customer VARCHAR, + OrdNo VARCHAR, + SalesValue DOUBLE + ); + """, + """ + CREATE TABLE dbo_tblNewOrders ( + Customer VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """, + ], + ) + + assert sql is not None + assert 'FROM "dbo_tblNewOrders"' in sql + assert 'FROM "dbo_qMarginSales"' not in sql + assert 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo")' in sql + + def test_build_pcb_direct_question_sql_for_board_model_distribution_over_time(): service = AskService.__new__(AskService) From f92cabe7afec57d8bc0052eeea028eefb9184f5a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 9 Jul 2026 23:12:57 +0530 Subject: [PATCH 0466/1087] Validate SQL covers requested metadata concepts --- wren-ai-service/src/web/v1/services/ask.py | 84 ++++++++- .../pytest/services/test_ask_sales_sql.py | 159 ++++++++++++++++++ 2 files changed, 240 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9db38bd5a8..fe38f23044 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -622,6 +622,87 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: return [] concept_groups: list[set[str]] = [] + compact_query = re.sub(r"[^a-z0-9]", "", normalized) + if ( + "customer" in normalized + or "custname" in compact_query + or "custno" in compact_query + ): + concept_groups.append( + { + "account", + "client", + "cust", + "customer", + "customername", + "custname", + "custno", + } + ) + if "market segment" in normalized: + concept_groups.append( + {"market", "marketsegment", "markettype", "region", "segment"} + ) + elif "market" in normalized: + concept_groups.append({"market", "marketname", "markettype", "segment"}) + if "region" in normalized: + concept_groups.append( + {"area", "country", "market", "region", "regional", "territory"} + ) + if "salesperson" in normalized or "sales person" in normalized: + concept_groups.append( + {"owner", "rep", "salesman", "salesperson", "salesrep"} + ) + if "product type" in normalized or "producttype" in compact_query: + concept_groups.append({"prodtype", "product", "producttype", "type"}) + elif "product" in normalized: + concept_groups.append( + { + "item", + "part", + "prod", + "prodcode", + "prodname", + "product", + "productcode", + "productname", + "sku", + } + ) + if "quantity" in normalized or re.search(r"\bqty\b", normalized): + concept_groups.append( + {"invoiceqty", "orderqty", "qty", "quantity", "salesqty"} + ) + if any( + term in normalized + for term in ( + "amount", + "order value", + "revenue", + "sale", + "sales", + "sales value", + "value", + ) + ) and not any( + term in normalized for term in ("salesperson", "sales person", "sales rep") + ): + concept_groups.append( + { + "amount", + "fxsalesvalue", + "invoiceamount", + "invoicevalue", + "net", + "revenue", + "salesamount", + "salesvalue", + "total", + "value", + } + ) + if any(term in normalized for term in ("order priority", "priorities")): + concept_groups.append({"orderpriority", "priority"}) if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: @@ -656,13 +737,10 @@ def _sql_covers_required_question_concepts( referenced_column_tokens: set[str], referenced_table_tokens: set[str], ) -> bool: - sql_text = (sql or "").lower() available_tokens = referenced_column_tokens | referenced_table_tokens for concept_group in self._required_sql_concept_groups(query): if concept_group & available_tokens: continue - if any(token in sql_text for token in concept_group): - continue logger.warning( "Ignoring SQL because it does not cover required question concept. " "query=%s required=%s referenced_column_tokens=%s referenced_table_tokens=%s sql=%s", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 3c26e551c6..bf6d6373db 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,6 +269,115 @@ def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): assert result is not None +def test_sql_intent_rejects_partial_customer_market_quantity_answer(): + service = AskService.__new__(AskService) + + assert not service._sql_matches_question_intent( + ( + 'SELECT "dbo_tblStageNewOrders3"."Market" AS "Market", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_tblStageNewOrders3" ' + 'GROUP BY "dbo_tblStageNewOrders3"."Market"' + ), + "Show total order quantity by customer and market segment.", + [ + { + "name": "dbo_tblStageNewOrders3", + "columns": [ + {"name": "Market", "type": "VARCHAR"}, + {"name": "Customer", "type": "VARCHAR"}, + {"name": "OrderQty", "type": "FLOAT"}, + ], + } + ], + ) + + +def test_sql_intent_accepts_customer_market_quantity_answer(): + service = AskService.__new__(AskService) + + assert service._sql_matches_question_intent( + ( + 'SELECT "dbo_tblNewOrders"."Customer" AS "Customer", ' + '"dbo_tblNewOrders"."Market" AS "Market", ' + 'SUM("dbo_tblNewOrders"."OrderQty") AS "TotalOrderQty" ' + 'FROM "dbo_tblNewOrders" ' + 'GROUP BY "dbo_tblNewOrders"."Customer", "dbo_tblNewOrders"."Market"' + ), + "Show total order quantity by customer and market segment.", + [ + { + "name": "dbo_tblNewOrders", + "columns": [ + {"name": "Customer", "type": "VARCHAR"}, + {"name": "Market", "type": "VARCHAR"}, + {"name": "OrderQty", "type": "FLOAT"}, + ], + } + ], + ) + + +def test_sql_intent_rejects_product_sales_question_without_sales_or_region(): + service = AskService.__new__(AskService) + + assert not service._sql_matches_question_intent( + 'SELECT "dbo_products"."products" AS "products" FROM "dbo_products"', + "Which products have the highest sales by region?", + [ + { + "name": "dbo_products", + "columns": [ + {"name": "products", "type": "VARCHAR"}, + {"name": "Region", "type": "VARCHAR"}, + {"name": "SalesValue", "type": "FLOAT"}, + ], + } + ], + ) + + +def test_sql_intent_accepts_joined_product_sales_by_region_answer(): + service = AskService.__new__(AskService) + + assert service._sql_matches_question_intent( + ( + 'SELECT "dbo_products"."ProductName" AS "ProductName", ' + '"dbo_regions"."Region" AS "Region", ' + 'SUM("dbo_order_lines"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_order_lines" ' + 'JOIN "dbo_products" ON "dbo_order_lines"."ProductId" = "dbo_products"."ProductId" ' + 'JOIN "dbo_regions" ON "dbo_order_lines"."RegionId" = "dbo_regions"."RegionId" ' + 'GROUP BY "dbo_products"."ProductName", "dbo_regions"."Region"' + ), + "Which products have the highest sales by region?", + [ + { + "name": "dbo_order_lines", + "columns": [ + {"name": "ProductId", "type": "VARCHAR"}, + {"name": "RegionId", "type": "VARCHAR"}, + {"name": "SalesValue", "type": "FLOAT"}, + ], + }, + { + "name": "dbo_products", + "columns": [ + {"name": "ProductId", "type": "VARCHAR"}, + {"name": "ProductName", "type": "VARCHAR"}, + ], + }, + { + "name": "dbo_regions", + "columns": [ + {"name": "RegionId", "type": "VARCHAR"}, + {"name": "Region", "type": "VARCHAR"}, + ], + }, + ], + ) + + def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) @@ -363,6 +472,56 @@ def test_prune_sql_generation_context_keeps_related_join_table(): assert table_names == ["dbo_Customers", "dbo_Orders"] +def test_prune_sql_generation_context_keeps_multiple_join_tables(): + service = AskService.__new__(AskService) + table_ddls = [ + """ + CREATE TABLE dbo_Products ( + ProductId VARCHAR, + ProductName VARCHAR + ); + """, + """ + CREATE TABLE dbo_Regions ( + RegionId VARCHAR, + Region VARCHAR + ); + """, + """ + CREATE TABLE dbo_OrderLines ( + OrderLineId VARCHAR, + ProductId VARCHAR, + RegionId VARCHAR, + SalesValue FLOAT, + FOREIGN KEY (ProductId) REFERENCES dbo_Products(ProductId), + FOREIGN KEY (RegionId) REFERENCES dbo_Regions(RegionId) + ); + """, + """ + CREATE TABLE dbo_Unrelated ( + id VARCHAR, + notes VARCHAR + ); + """, + ] + documents = [ + {"table_name": "dbo_Products", "table_ddl": table_ddls[0]}, + {"table_name": "dbo_Regions", "table_ddl": table_ddls[1]}, + {"table_name": "dbo_OrderLines", "table_ddl": table_ddls[2]}, + {"table_name": "dbo_Unrelated", "table_ddl": table_ddls[3]}, + ] + + _, table_names, _ = service._prune_sql_generation_context( + "Which products have the highest sales by region?", + documents, + [document["table_name"] for document in documents], + table_ddls, + max_tables=3, + ) + + assert table_names == ["dbo_Products", "dbo_Regions", "dbo_OrderLines"] + + def test_build_schema_grounded_sales_sql_for_top_markets(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 3539d46e642e9aa49dce49a7e5c4931dc692e54a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 00:15:48 +0530 Subject: [PATCH 0467/1087] Revert "Validate SQL covers requested metadata concepts" This reverts commit f92cabe7afec57d8bc0052eeea028eefb9184f5a. --- wren-ai-service/src/web/v1/services/ask.py | 84 +-------- .../pytest/services/test_ask_sales_sql.py | 159 ------------------ 2 files changed, 3 insertions(+), 240 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index fe38f23044..9db38bd5a8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -622,87 +622,6 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: return [] concept_groups: list[set[str]] = [] - compact_query = re.sub(r"[^a-z0-9]", "", normalized) - if ( - "customer" in normalized - or "custname" in compact_query - or "custno" in compact_query - ): - concept_groups.append( - { - "account", - "client", - "cust", - "customer", - "customername", - "custname", - "custno", - } - ) - if "market segment" in normalized: - concept_groups.append( - {"market", "marketsegment", "markettype", "region", "segment"} - ) - elif "market" in normalized: - concept_groups.append({"market", "marketname", "markettype", "segment"}) - if "region" in normalized: - concept_groups.append( - {"area", "country", "market", "region", "regional", "territory"} - ) - if "salesperson" in normalized or "sales person" in normalized: - concept_groups.append( - {"owner", "rep", "salesman", "salesperson", "salesrep"} - ) - if "product type" in normalized or "producttype" in compact_query: - concept_groups.append({"prodtype", "product", "producttype", "type"}) - elif "product" in normalized: - concept_groups.append( - { - "item", - "part", - "prod", - "prodcode", - "prodname", - "product", - "productcode", - "productname", - "sku", - } - ) - if "quantity" in normalized or re.search(r"\bqty\b", normalized): - concept_groups.append( - {"invoiceqty", "orderqty", "qty", "quantity", "salesqty"} - ) - if any( - term in normalized - for term in ( - "amount", - "order value", - "revenue", - "sale", - "sales", - "sales value", - "value", - ) - ) and not any( - term in normalized for term in ("salesperson", "sales person", "sales rep") - ): - concept_groups.append( - { - "amount", - "fxsalesvalue", - "invoiceamount", - "invoicevalue", - "net", - "revenue", - "salesamount", - "salesvalue", - "total", - "value", - } - ) - if any(term in normalized for term in ("order priority", "priorities")): - concept_groups.append({"orderpriority", "priority"}) if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: @@ -737,10 +656,13 @@ def _sql_covers_required_question_concepts( referenced_column_tokens: set[str], referenced_table_tokens: set[str], ) -> bool: + sql_text = (sql or "").lower() available_tokens = referenced_column_tokens | referenced_table_tokens for concept_group in self._required_sql_concept_groups(query): if concept_group & available_tokens: continue + if any(token in sql_text for token in concept_group): + continue logger.warning( "Ignoring SQL because it does not cover required question concept. " "query=%s required=%s referenced_column_tokens=%s referenced_table_tokens=%s sql=%s", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index bf6d6373db..3c26e551c6 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,115 +269,6 @@ def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): assert result is not None -def test_sql_intent_rejects_partial_customer_market_quantity_answer(): - service = AskService.__new__(AskService) - - assert not service._sql_matches_question_intent( - ( - 'SELECT "dbo_tblStageNewOrders3"."Market" AS "Market", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_tblStageNewOrders3" ' - 'GROUP BY "dbo_tblStageNewOrders3"."Market"' - ), - "Show total order quantity by customer and market segment.", - [ - { - "name": "dbo_tblStageNewOrders3", - "columns": [ - {"name": "Market", "type": "VARCHAR"}, - {"name": "Customer", "type": "VARCHAR"}, - {"name": "OrderQty", "type": "FLOAT"}, - ], - } - ], - ) - - -def test_sql_intent_accepts_customer_market_quantity_answer(): - service = AskService.__new__(AskService) - - assert service._sql_matches_question_intent( - ( - 'SELECT "dbo_tblNewOrders"."Customer" AS "Customer", ' - '"dbo_tblNewOrders"."Market" AS "Market", ' - 'SUM("dbo_tblNewOrders"."OrderQty") AS "TotalOrderQty" ' - 'FROM "dbo_tblNewOrders" ' - 'GROUP BY "dbo_tblNewOrders"."Customer", "dbo_tblNewOrders"."Market"' - ), - "Show total order quantity by customer and market segment.", - [ - { - "name": "dbo_tblNewOrders", - "columns": [ - {"name": "Customer", "type": "VARCHAR"}, - {"name": "Market", "type": "VARCHAR"}, - {"name": "OrderQty", "type": "FLOAT"}, - ], - } - ], - ) - - -def test_sql_intent_rejects_product_sales_question_without_sales_or_region(): - service = AskService.__new__(AskService) - - assert not service._sql_matches_question_intent( - 'SELECT "dbo_products"."products" AS "products" FROM "dbo_products"', - "Which products have the highest sales by region?", - [ - { - "name": "dbo_products", - "columns": [ - {"name": "products", "type": "VARCHAR"}, - {"name": "Region", "type": "VARCHAR"}, - {"name": "SalesValue", "type": "FLOAT"}, - ], - } - ], - ) - - -def test_sql_intent_accepts_joined_product_sales_by_region_answer(): - service = AskService.__new__(AskService) - - assert service._sql_matches_question_intent( - ( - 'SELECT "dbo_products"."ProductName" AS "ProductName", ' - '"dbo_regions"."Region" AS "Region", ' - 'SUM("dbo_order_lines"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_order_lines" ' - 'JOIN "dbo_products" ON "dbo_order_lines"."ProductId" = "dbo_products"."ProductId" ' - 'JOIN "dbo_regions" ON "dbo_order_lines"."RegionId" = "dbo_regions"."RegionId" ' - 'GROUP BY "dbo_products"."ProductName", "dbo_regions"."Region"' - ), - "Which products have the highest sales by region?", - [ - { - "name": "dbo_order_lines", - "columns": [ - {"name": "ProductId", "type": "VARCHAR"}, - {"name": "RegionId", "type": "VARCHAR"}, - {"name": "SalesValue", "type": "FLOAT"}, - ], - }, - { - "name": "dbo_products", - "columns": [ - {"name": "ProductId", "type": "VARCHAR"}, - {"name": "ProductName", "type": "VARCHAR"}, - ], - }, - { - "name": "dbo_regions", - "columns": [ - {"name": "RegionId", "type": "VARCHAR"}, - {"name": "Region", "type": "VARCHAR"}, - ], - }, - ], - ) - - def test_needs_conversation_context_only_for_true_followups(): service = AskService.__new__(AskService) @@ -472,56 +363,6 @@ def test_prune_sql_generation_context_keeps_related_join_table(): assert table_names == ["dbo_Customers", "dbo_Orders"] -def test_prune_sql_generation_context_keeps_multiple_join_tables(): - service = AskService.__new__(AskService) - table_ddls = [ - """ - CREATE TABLE dbo_Products ( - ProductId VARCHAR, - ProductName VARCHAR - ); - """, - """ - CREATE TABLE dbo_Regions ( - RegionId VARCHAR, - Region VARCHAR - ); - """, - """ - CREATE TABLE dbo_OrderLines ( - OrderLineId VARCHAR, - ProductId VARCHAR, - RegionId VARCHAR, - SalesValue FLOAT, - FOREIGN KEY (ProductId) REFERENCES dbo_Products(ProductId), - FOREIGN KEY (RegionId) REFERENCES dbo_Regions(RegionId) - ); - """, - """ - CREATE TABLE dbo_Unrelated ( - id VARCHAR, - notes VARCHAR - ); - """, - ] - documents = [ - {"table_name": "dbo_Products", "table_ddl": table_ddls[0]}, - {"table_name": "dbo_Regions", "table_ddl": table_ddls[1]}, - {"table_name": "dbo_OrderLines", "table_ddl": table_ddls[2]}, - {"table_name": "dbo_Unrelated", "table_ddl": table_ddls[3]}, - ] - - _, table_names, _ = service._prune_sql_generation_context( - "Which products have the highest sales by region?", - documents, - [document["table_name"] for document in documents], - table_ddls, - max_tables=3, - ) - - assert table_names == ["dbo_Products", "dbo_Regions", "dbo_OrderLines"] - - def test_build_schema_grounded_sales_sql_for_top_markets(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( From 69735dc05270168cfcb526fafd6a1fe389694054 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 10 Jul 2026 01:12:46 +0530 Subject: [PATCH 0468/1087] Fix multi-table schema context for SQL generation --- .../pipelines/generation/sql_generation.py | 6 + .../src/pipelines/generation/utils/sql.py | 2 + wren-ai-service/src/web/v1/services/ask.py | 201 +++++++++++++++++- .../test_ask_heuristic_text_to_sql.py | 55 +++++ 4 files changed, 263 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 59bea2279c..a67ada372f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -93,6 +93,12 @@ general guidance when the question can be answered with SQL over the active metadata. Never reuse table or column names from SQL SAMPLES unless those exact names also appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. +Before writing SQL, identify every business entity and measure requested by the full +question. If those entities live in different tables, include every required table and +join only through explicit relationships or matching key columns shown in ACTIVE +DATASOURCE METADATA. Do not stop after the first matching table. If the active +metadata does not contain the table, column, or relationship needed to answer, do not +invent it. {% if sql_generation_reasoning %} ### REASONING PLAN ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 8316d24e86..d4b5093ced 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2198,6 +2198,8 @@ def get_sql_generation_system_prompt( 9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. 10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. 11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. +12. Analyze the complete user question before writing SQL. Identify all requested business entities and measures, include all required schema tables, and join only through explicit relationships or schema-backed key columns. Do not stop at the first matching table when the question requires multiple entities. +13. If the active metadata does not contain enough schema information to answer the question, return no invented SQL; schema validation will reject hallucinated tables, columns, and joins. {text_to_sql_rules} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9db38bd5a8..9a92bf0adc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4858,6 +4858,171 @@ def _query_schema_terms(self, query: str) -> set[str]: } return {term for term in terms if term} + def _table_concept_match_score( + self, + table: dict[str, Any], + concept_group: set[str], + ) -> int: + concept_keys = { + self._normalize_schema_token(concept) + for concept in concept_group + if concept + } + concept_keys = {key for key in concept_keys if key} + if not concept_keys: + return 0 + + table_name = str(table.get("name") or "") + table_key = self._normalize_schema_token(table_name) + table_tokens = { + self._normalize_schema_token(token) + for token in self._schema_name_tokens(table_name) + } + table_tokens = {token for token in table_tokens if token} + + score = 0 + for concept_key in concept_keys: + if concept_key == table_key: + score += 180 + elif concept_key in table_tokens: + score += 120 + elif concept_key and concept_key in table_key: + score += 80 + + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_key = self._normalize_schema_token(column_name) + column_tokens = { + self._normalize_schema_token(token) + for token in self._schema_name_tokens(column_name) + } + column_tokens = {token for token in column_tokens if token} + + for concept_key in concept_keys: + if concept_key == column_key: + score += 120 + elif concept_key in column_tokens: + score += 100 + elif concept_key and concept_key in column_key: + score += 45 + + return score + + def _required_concept_table_indexes( + self, + query: str, + parsed_tables: list[dict[str, Any]], + table_ddls: list[str] | None = None, + ) -> list[int]: + selected_indexes: list[int] = [] + selected_set: set[int] = set() + relationship_indexes = self._relationship_table_indexes( + table_ddls or [], + parsed_tables, + ) + + for concept_group in self._required_sql_concept_groups(query): + candidates = [] + for index, table in enumerate(parsed_tables): + score = self._table_concept_match_score(table, concept_group) + if score <= 0: + continue + + neighbor_indexes: set[int] = set() + for key in self._table_alias_keys(str(table.get("name") or "")): + neighbor_indexes.update(relationship_indexes.get(key, set())) + if neighbor_indexes & selected_set: + score += 175 + + candidates.append((score, index)) + candidates = [ + (score, index) for score, index in candidates if score > 0 + ] + if not candidates: + continue + + score, index = max(candidates, key=lambda item: item[0]) + if index not in selected_set: + selected_indexes.append(index) + selected_set.add(index) + logger.info( + "Preserving schema table for required query concept. " + "query=%s concept=%s table=%s score=%s", + query, + sorted(concept_group), + parsed_tables[index].get("name"), + score, + ) + + return selected_indexes + + def _table_alias_keys(self, table_name: str) -> set[str]: + raw_table_name = str(table_name or "") + keys = self._schema_identifier_alias_keys(raw_table_name) + short_name = re.split(r"[.$]", raw_table_name)[-1] + keys.update(self._schema_identifier_alias_keys(short_name)) + if "_" in short_name: + keys.update(self._schema_identifier_alias_keys(short_name.split("_", 1)[1])) + return {key for key in keys if key} + + def _relationship_table_indexes( + self, + table_ddls: list[str], + parsed_tables: list[dict[str, Any]], + ) -> dict[str, set[int]]: + relationship_indexes: dict[str, set[int]] = {} + table_key_to_index: dict[str, int] = {} + + for index, table in enumerate(parsed_tables): + for key in self._table_alias_keys(str(table.get("name") or "")): + table_key_to_index[key] = index + + for source_index, ddl in enumerate(table_ddls or []): + if not isinstance(ddl, str): + continue + if source_index >= len(parsed_tables): + continue + + source_name = str(parsed_tables[source_index].get("name") or "") + source_keys = self._table_alias_keys(source_name) + for relationship_match in re.finditer( + r"\bFOREIGN\s+KEY\s*\([^)]+\)\s+REFERENCES\s+" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" + r"\s*\([^)]+\)", + ddl, + flags=re.IGNORECASE, + ): + target_name = next( + ( + value + for value in relationship_match.groupdict().values() + if value + ), + "", + ) + target_index = None + for target_key in self._table_alias_keys(target_name): + if target_key in table_key_to_index: + target_index = table_key_to_index[target_key] + break + if target_index is None: + continue + + target_keys = self._table_alias_keys( + str(parsed_tables[target_index].get("name") or target_name) + ) + for source_key in source_keys: + relationship_indexes.setdefault(source_key, set()).add( + target_index + ) + for target_key in target_keys: + relationship_indexes.setdefault(target_key, set()).add( + source_index + ) + + return relationship_indexes + def _prune_sql_generation_context( self, query: str, @@ -4924,7 +5089,16 @@ def _prune_sql_generation_context( index for _, index in sorted(scored, key=lambda item: item[0], reverse=True) ] core_limit = max(1, max_tables - 2) if max_tables > 2 else 1 - selected_indexes = sorted_scored_indexes[:core_limit] + selected_indexes = self._required_concept_table_indexes( + query, + parsed_tables, + table_ddls, + ) + for index in sorted_scored_indexes: + if len(selected_indexes) >= core_limit: + break + if index not in selected_indexes: + selected_indexes.append(index) selected_indexes = self._expand_pruned_context_with_related_tables( selected_indexes, parsed_tables, @@ -4968,6 +5142,31 @@ def _expand_pruned_context_with_related_tables( selected: list[int] = list(dict.fromkeys(selected_indexes)) selected_set = set(selected) + relationship_indexes = self._relationship_table_indexes( + table_ddls, + parsed_tables, + ) + + def add_relationship_neighbors() -> None: + for selected_index in list(selected): + if selected_index >= len(parsed_tables): + continue + table_name = str(parsed_tables[selected_index].get("name") or "") + neighbor_indexes: set[int] = set() + for key in self._table_alias_keys(table_name): + neighbor_indexes.update(relationship_indexes.get(key, set())) + + for neighbor_index in sorted(neighbor_indexes): + if len(selected) >= max_tables: + return + if neighbor_index in selected_set: + continue + selected.append(neighbor_index) + selected_set.add(neighbor_index) + + add_relationship_neighbors() + if len(selected) >= max_tables: + return selected[:max_tables] def join_key_columns(table: dict[str, Any]) -> set[str]: keys = set() diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index c9ed99629f..eb4ab941b3 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -594,3 +594,58 @@ def test_complete_sql_generation_context_refetches_full_selected_schema(): "enable_column_pruning": False, } ] + + +def test_prune_sql_generation_context_preserves_all_required_business_entities(): + service = AskService(pipelines={}) + documents = [ + { + "table_name": "dbo_SalesQuantityByRegion", + "table_ddl": """ + CREATE TABLE dbo_SalesQuantityByRegion ( + Region VARCHAR, + SalesQuantity INT + ); + """, + }, + { + "table_name": "dbo_SalesQuantityByCustomer", + "table_ddl": """ + CREATE TABLE dbo_SalesQuantityByCustomer ( + CustomerName VARCHAR, + SalesQuantity INT + ); + """, + }, + { + "table_name": "dbo_InventoryItems", + "table_ddl": """ + CREATE TABLE dbo_InventoryItems ( + ItemID INT PRIMARY KEY, + SKU VARCHAR, + ItemDescription VARCHAR + ); + """, + }, + { + "table_name": "dbo_SalesOrderDetails", + "table_ddl": """ + CREATE TABLE dbo_SalesOrderDetails ( + SalesOrderDetailID INT PRIMARY KEY, + ItemID INT, + Quantity INT, + FOREIGN KEY (ItemID) REFERENCES dbo_InventoryItems(ItemID) + ); + """, + }, + ] + + _, table_names, _ = service._prune_sql_generation_context( + "Show me the products based on sales quantity.", + documents, + [document["table_name"] for document in documents], + [document["table_ddl"] for document in documents], + max_tables=2, + ) + + assert table_names == ["dbo_InventoryItems", "dbo_SalesOrderDetails"] From 0f8b9c774fd2beea1322cf49467c57a70bcae562 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 10 Jul 2026 01:29:55 +0530 Subject: [PATCH 0469/1087] Revert "Fix multi-table schema context for SQL generation" This reverts commit 69735dc05270168cfcb526fafd6a1fe389694054. --- .../pipelines/generation/sql_generation.py | 6 - .../src/pipelines/generation/utils/sql.py | 2 - wren-ai-service/src/web/v1/services/ask.py | 201 +----------------- .../test_ask_heuristic_text_to_sql.py | 55 ----- 4 files changed, 1 insertion(+), 263 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index a67ada372f..59bea2279c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -93,12 +93,6 @@ general guidance when the question can be answered with SQL over the active metadata. Never reuse table or column names from SQL SAMPLES unless those exact names also appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. -Before writing SQL, identify every business entity and measure requested by the full -question. If those entities live in different tables, include every required table and -join only through explicit relationships or matching key columns shown in ACTIVE -DATASOURCE METADATA. Do not stop after the first matching table. If the active -metadata does not contain the table, column, or relationship needed to answer, do not -invent it. {% if sql_generation_reasoning %} ### REASONING PLAN ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index d4b5093ced..8316d24e86 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2198,8 +2198,6 @@ def get_sql_generation_system_prompt( 9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. 10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. 11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. -12. Analyze the complete user question before writing SQL. Identify all requested business entities and measures, include all required schema tables, and join only through explicit relationships or schema-backed key columns. Do not stop at the first matching table when the question requires multiple entities. -13. If the active metadata does not contain enough schema information to answer the question, return no invented SQL; schema validation will reject hallucinated tables, columns, and joins. {text_to_sql_rules} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9a92bf0adc..9db38bd5a8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4858,171 +4858,6 @@ def _query_schema_terms(self, query: str) -> set[str]: } return {term for term in terms if term} - def _table_concept_match_score( - self, - table: dict[str, Any], - concept_group: set[str], - ) -> int: - concept_keys = { - self._normalize_schema_token(concept) - for concept in concept_group - if concept - } - concept_keys = {key for key in concept_keys if key} - if not concept_keys: - return 0 - - table_name = str(table.get("name") or "") - table_key = self._normalize_schema_token(table_name) - table_tokens = { - self._normalize_schema_token(token) - for token in self._schema_name_tokens(table_name) - } - table_tokens = {token for token in table_tokens if token} - - score = 0 - for concept_key in concept_keys: - if concept_key == table_key: - score += 180 - elif concept_key in table_tokens: - score += 120 - elif concept_key and concept_key in table_key: - score += 80 - - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_key = self._normalize_schema_token(column_name) - column_tokens = { - self._normalize_schema_token(token) - for token in self._schema_name_tokens(column_name) - } - column_tokens = {token for token in column_tokens if token} - - for concept_key in concept_keys: - if concept_key == column_key: - score += 120 - elif concept_key in column_tokens: - score += 100 - elif concept_key and concept_key in column_key: - score += 45 - - return score - - def _required_concept_table_indexes( - self, - query: str, - parsed_tables: list[dict[str, Any]], - table_ddls: list[str] | None = None, - ) -> list[int]: - selected_indexes: list[int] = [] - selected_set: set[int] = set() - relationship_indexes = self._relationship_table_indexes( - table_ddls or [], - parsed_tables, - ) - - for concept_group in self._required_sql_concept_groups(query): - candidates = [] - for index, table in enumerate(parsed_tables): - score = self._table_concept_match_score(table, concept_group) - if score <= 0: - continue - - neighbor_indexes: set[int] = set() - for key in self._table_alias_keys(str(table.get("name") or "")): - neighbor_indexes.update(relationship_indexes.get(key, set())) - if neighbor_indexes & selected_set: - score += 175 - - candidates.append((score, index)) - candidates = [ - (score, index) for score, index in candidates if score > 0 - ] - if not candidates: - continue - - score, index = max(candidates, key=lambda item: item[0]) - if index not in selected_set: - selected_indexes.append(index) - selected_set.add(index) - logger.info( - "Preserving schema table for required query concept. " - "query=%s concept=%s table=%s score=%s", - query, - sorted(concept_group), - parsed_tables[index].get("name"), - score, - ) - - return selected_indexes - - def _table_alias_keys(self, table_name: str) -> set[str]: - raw_table_name = str(table_name or "") - keys = self._schema_identifier_alias_keys(raw_table_name) - short_name = re.split(r"[.$]", raw_table_name)[-1] - keys.update(self._schema_identifier_alias_keys(short_name)) - if "_" in short_name: - keys.update(self._schema_identifier_alias_keys(short_name.split("_", 1)[1])) - return {key for key in keys if key} - - def _relationship_table_indexes( - self, - table_ddls: list[str], - parsed_tables: list[dict[str, Any]], - ) -> dict[str, set[int]]: - relationship_indexes: dict[str, set[int]] = {} - table_key_to_index: dict[str, int] = {} - - for index, table in enumerate(parsed_tables): - for key in self._table_alias_keys(str(table.get("name") or "")): - table_key_to_index[key] = index - - for source_index, ddl in enumerate(table_ddls or []): - if not isinstance(ddl, str): - continue - if source_index >= len(parsed_tables): - continue - - source_name = str(parsed_tables[source_index].get("name") or "") - source_keys = self._table_alias_keys(source_name) - for relationship_match in re.finditer( - r"\bFOREIGN\s+KEY\s*\([^)]+\)\s+REFERENCES\s+" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" - r"\s*\([^)]+\)", - ddl, - flags=re.IGNORECASE, - ): - target_name = next( - ( - value - for value in relationship_match.groupdict().values() - if value - ), - "", - ) - target_index = None - for target_key in self._table_alias_keys(target_name): - if target_key in table_key_to_index: - target_index = table_key_to_index[target_key] - break - if target_index is None: - continue - - target_keys = self._table_alias_keys( - str(parsed_tables[target_index].get("name") or target_name) - ) - for source_key in source_keys: - relationship_indexes.setdefault(source_key, set()).add( - target_index - ) - for target_key in target_keys: - relationship_indexes.setdefault(target_key, set()).add( - source_index - ) - - return relationship_indexes - def _prune_sql_generation_context( self, query: str, @@ -5089,16 +4924,7 @@ def _prune_sql_generation_context( index for _, index in sorted(scored, key=lambda item: item[0], reverse=True) ] core_limit = max(1, max_tables - 2) if max_tables > 2 else 1 - selected_indexes = self._required_concept_table_indexes( - query, - parsed_tables, - table_ddls, - ) - for index in sorted_scored_indexes: - if len(selected_indexes) >= core_limit: - break - if index not in selected_indexes: - selected_indexes.append(index) + selected_indexes = sorted_scored_indexes[:core_limit] selected_indexes = self._expand_pruned_context_with_related_tables( selected_indexes, parsed_tables, @@ -5142,31 +4968,6 @@ def _expand_pruned_context_with_related_tables( selected: list[int] = list(dict.fromkeys(selected_indexes)) selected_set = set(selected) - relationship_indexes = self._relationship_table_indexes( - table_ddls, - parsed_tables, - ) - - def add_relationship_neighbors() -> None: - for selected_index in list(selected): - if selected_index >= len(parsed_tables): - continue - table_name = str(parsed_tables[selected_index].get("name") or "") - neighbor_indexes: set[int] = set() - for key in self._table_alias_keys(table_name): - neighbor_indexes.update(relationship_indexes.get(key, set())) - - for neighbor_index in sorted(neighbor_indexes): - if len(selected) >= max_tables: - return - if neighbor_index in selected_set: - continue - selected.append(neighbor_index) - selected_set.add(neighbor_index) - - add_relationship_neighbors() - if len(selected) >= max_tables: - return selected[:max_tables] def join_key_columns(table: dict[str, Any]) -> set[str]: keys = set() diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py index eb4ab941b3..c9ed99629f 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py @@ -594,58 +594,3 @@ def test_complete_sql_generation_context_refetches_full_selected_schema(): "enable_column_pruning": False, } ] - - -def test_prune_sql_generation_context_preserves_all_required_business_entities(): - service = AskService(pipelines={}) - documents = [ - { - "table_name": "dbo_SalesQuantityByRegion", - "table_ddl": """ - CREATE TABLE dbo_SalesQuantityByRegion ( - Region VARCHAR, - SalesQuantity INT - ); - """, - }, - { - "table_name": "dbo_SalesQuantityByCustomer", - "table_ddl": """ - CREATE TABLE dbo_SalesQuantityByCustomer ( - CustomerName VARCHAR, - SalesQuantity INT - ); - """, - }, - { - "table_name": "dbo_InventoryItems", - "table_ddl": """ - CREATE TABLE dbo_InventoryItems ( - ItemID INT PRIMARY KEY, - SKU VARCHAR, - ItemDescription VARCHAR - ); - """, - }, - { - "table_name": "dbo_SalesOrderDetails", - "table_ddl": """ - CREATE TABLE dbo_SalesOrderDetails ( - SalesOrderDetailID INT PRIMARY KEY, - ItemID INT, - Quantity INT, - FOREIGN KEY (ItemID) REFERENCES dbo_InventoryItems(ItemID) - ); - """, - }, - ] - - _, table_names, _ = service._prune_sql_generation_context( - "Show me the products based on sales quantity.", - documents, - [document["table_name"] for document in documents], - [document["table_ddl"] for document in documents], - max_tables=2, - ) - - assert table_names == ["dbo_InventoryItems", "dbo_SalesOrderDetails"] From dd5e7149e624953f5f1138561d6db57c48984c29 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 12:43:31 +0530 Subject: [PATCH 0470/1087] Improve semantic SQL intent validation --- .../generation/followup_sql_generation.py | 20 + .../followup_sql_generation_reasoning.py | 13 + .../pipelines/generation/sql_correction.py | 18 + .../pipelines/generation/sql_generation.py | 20 + .../generation/sql_generation_reasoning.py | 13 + .../pipelines/generation/sql_regeneration.py | 25 + .../src/pipelines/generation/utils/sql.py | 569 ++++++++++++++++++ .../retrieval/db_schema_retrieval.py | 73 ++- wren-ai-service/src/web/v1/services/ask.py | 82 ++- .../src/web/v1/services/ask_feedback.py | 29 +- .../pipelines/generation/test_sql_utils.py | 176 ++++++ .../retrieval/test_db_schema_retrieval.py | 71 +++ 12 files changed, 1086 insertions(+), 23 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 c9f8b23537..a4616ff881 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -99,6 +99,14 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +This is the pre-generation semantic analysis of the follow-up request against the +active deployed schema. Use it as a contract for table, column, metric, dimension, +filter, time, relationship, aggregation, and ranking selection. +{{ schema_intent_analysis }} +{% endif %} + ### INTENT AND SCHEMA GROUNDING ### Interpret the user's business terms by matching them to explicit tables, columns, metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Never reuse table @@ -106,6 +114,10 @@ appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. Only apply aggregate functions to columns whose active metadata type supports that operation. +Before writing SQL, validate that the selected schema elements directly support every +key entity, metric, dimension, filter, time range, relationship, and aggregation in +the follow-up question. If the schema cannot support the requested information, do +not replace the request with a generic COUNT(*) or unrelated table query. ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -129,6 +141,7 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -154,6 +167,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -186,9 +200,11 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: list[str], data_source: str, + query: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -198,6 +214,8 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -249,6 +267,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -268,6 +287,7 @@ async def run( "has_metric": has_metric, "has_json_field": has_json_field, "sql_functions": sql_functions, + "schema_intent_analysis": schema_intent_analysis, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index abbbb81d56..9759617bb9 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -67,6 +67,15 @@ Language: {{ language }} Current Time: {{ current_time }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +Use this semantic analysis as the planning contract for entities, metrics, +dimensions, filters, joins, time constraints, aggregations, ranking, and analytical +intent. If it shows missing or ambiguous requirements, state that limitation in the +plan instead of planning unrelated SQL. +{{ schema_intent_analysis }} +{% endif %} + Let's think step by step. """ @@ -81,6 +90,7 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -92,6 +102,7 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -185,6 +196,7 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("Followup SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -197,6 +209,7 @@ async def run( "instructions": instructions or [], "configuration": configuration, "query_id": query_id, + "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index daae37e02e..be63983227 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -102,6 +102,12 @@ def get_sql_correction_system_prompt( {% if query %} User's Question: {{ query }} {% endif %} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +This is the semantic contract for the corrected SQL. Preserve this intent while +fixing syntax or planner errors. +{{ schema_intent_analysis }} +{% endif %} {% if invalid_generation_result.original_sql %} Original SQL: {{ invalid_generation_result.original_sql }} {% endif %} @@ -115,6 +121,10 @@ def get_sql_correction_system_prompt( user's request. Do not invent tables, columns, joins, metrics, or relationships. Only apply aggregate functions to columns whose active metadata type supports that operation. +Before returning corrected SQL, validate that it still directly supports every key +entity, metric, dimension, filter, time range, relationship, and aggregation in the +user's question. Do not replace an unsupported request with a generic COUNT(*) or +unrelated table query. Let's think step by step. """ @@ -130,6 +140,7 @@ def prompt( query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -137,6 +148,7 @@ def prompt( documents=documents, valid_table_names=construct_valid_table_names(documents), invalid_generation_result=invalid_generation_result, + schema_intent_analysis=schema_intent_analysis, instructions=construct_instructions( instructions=instructions, ), @@ -169,9 +181,11 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: List[Document], data_source: str, + query: str | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -181,6 +195,8 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -227,6 +243,7 @@ async def run( allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -238,6 +255,7 @@ async def run( "invalid_generation_result": invalid_generation_result, "documents": contexts, "query": query, + "schema_intent_analysis": schema_intent_analysis, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 59bea2279c..9280f3e8f1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -87,12 +87,24 @@ ### QUESTION ### User's Question: {{ query }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +This is the pre-generation semantic analysis of the user's request against the +active deployed schema. Use it as a contract for table, column, metric, dimension, +filter, time, relationship, aggregation, and ranking selection. +{{ schema_intent_analysis }} +{% endif %} + ### INTENT AND SCHEMA GROUNDING ### Interpret the user's business terms by matching them to explicit tables, columns, metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Do not answer with general guidance when the question can be answered with SQL over the active metadata. Never reuse table or column names from SQL SAMPLES unless those exact names also appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. +Before writing SQL, validate that the selected schema elements directly support every +key entity, metric, dimension, filter, time range, relationship, and aggregation in +the question. If the schema cannot support the requested information, do not replace +the request with a generic COUNT(*) or unrelated table query. {% if sql_generation_reasoning %} ### REASONING PLAN ### @@ -118,6 +130,7 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: schema_context = "\n".join(documents or []).lower() has_pcb_context = any( @@ -157,6 +170,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -185,10 +199,12 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: list[str], data_source: str, + query: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -199,6 +215,8 @@ async def post_process( allow_data_preview=allow_data_preview, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -250,6 +268,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -268,6 +287,7 @@ async def run( "has_metric": has_metric, "has_json_field": has_json_field, "sql_functions": sql_functions, + "schema_intent_analysis": schema_intent_analysis, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), 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 f91a4288e4..4db6239cd6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -52,6 +52,15 @@ Language: {{ language }} Current Time: {{ current_time }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +Use this semantic analysis as the planning contract for entities, metrics, +dimensions, filters, joins, time constraints, aggregations, ranking, and analytical +intent. If it shows missing or ambiguous requirements, state that limitation in the +plan instead of planning unrelated SQL. +{{ schema_intent_analysis }} +{% endif %} + Let's think step by step. """ @@ -65,6 +74,7 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -75,6 +85,7 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -163,6 +174,7 @@ async def run( instructions: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -174,6 +186,7 @@ async def run( "instructions": instructions or [], "configuration": configuration, "query_id": query_id, + "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 8b3eaafcb1..bef428a21c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -44,6 +44,10 @@ def get_sql_regeneration_system_prompt( please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. While generating the new SQL query, you should use the original SQL query as a reference. While generating the new SQL query, make sure to use the database schema to generate the SQL query. +Before returning SQL, validate that the selected schema elements directly support +the key entities, metrics, dimensions, filters, time ranges, relationships, and +aggregations from the user's question or reasoning. Do not replace an unsupported +request with a generic COUNT(*) or unrelated table query. {text_to_sql_rules} @@ -102,6 +106,15 @@ def get_sql_regeneration_system_prompt( {% endif %} ### QUESTION ### +{% if query %} +User's Question: {{ query }} +{% endif %} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +This is the semantic contract for regenerated SQL. Preserve this intent while +improving the original SQL. +{{ schema_intent_analysis }} +{% endif %} SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} @@ -117,6 +130,8 @@ def prompt( sql: str, prompt_builder: PromptBuilder, data_source: str, + query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -129,6 +144,8 @@ def prompt( sql=sql, data_source=data_source, documents=documents, + query=query, + schema_intent_analysis=schema_intent_analysis, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -176,7 +193,9 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: list[str], data_source: str, + query: str | None = None, project_id: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), @@ -184,6 +203,8 @@ async def post_process( data_source=data_source, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -228,6 +249,8 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -237,6 +260,8 @@ async def run( "documents": contexts, "sql_generation_reasoning": sql_generation_reasoning, "sql": sql, + "query": query, + "schema_intent_analysis": schema_intent_analysis, "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 8316d24e86..db9450226f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1543,6 +1543,8 @@ async def run( allow_data_preview: bool = False, valid_table_names: list[str] | None = None, valid_table_columns: dict[str, list[str]] | None = None, + query: str | None = None, + semantic_analysis: dict[str, Any] | None = None, ) -> dict: try: cleaned_generation_result = extract_sql_generation_result(replies[0]) @@ -1617,6 +1619,24 @@ async def run( }, } + intent_validation_error = validate_sql_intent_alignment( + query, + cleaned_generation_result, + valid_table_columns or {}, + semantic_analysis=semantic_analysis, + ) + if intent_validation_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_INTENT_VALIDATION", + "error": intent_validation_error, + "correlation_id": "", + }, + } + if normalize_data_source( data_source ) == "MSSQL" and contains_unsupported_mssql_json_access( @@ -2198,6 +2218,9 @@ def get_sql_generation_system_prompt( 9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. 10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. 11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. +12. Before generating SQL, validate that the selected schema elements directly support all key entities, metrics, dimensions, filters, time ranges, relationships, and aggregations mentioned or implied by the question. +13. Do not answer a specific business metric, trend, summary, comparison, dashboard, or analysis request with a generic record-count query unless the user explicitly asks only for record count. +14. If the required information cannot be derived from the available active schema, return the closest schema-grounded limitation instead of inventing unrelated SQL. {text_to_sql_rules} @@ -3185,6 +3208,552 @@ def format_valid_table_columns(valid_table_columns: dict[str, list[str]]) -> str ) +_PLAIN_COUNT_SQL_PATTERN = re.compile( + r"^\s*SELECT\s+COUNT\s*\(\s*\*\s*\)(?:\s+AS\s+" + r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*))?\s+' + rf"FROM\s+{_SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SQL_IDENTIFIER_PATTERN})*" + r"(?:\s+WHERE\s+.+?)?\s*(?:ORDER\s+BY\s+.+?)?(?:LIMIT\s+\d+\s*)?$", + flags=re.IGNORECASE | re.DOTALL, +) +_AGGREGATE_PATTERN = re.compile( + r"\b(?:SUM|AVG|MIN|MAX|COUNT)\s*\(", + flags=re.IGNORECASE, +) +_SPECIFIC_METRIC_TERM_GROUPS: dict[str, tuple[str, ...]] = { + "revenue": ("revenue", "sales", "sale", "amount", "value", "price", "total"), + "sales": ("sales", "sale", "revenue", "amount", "value", "price"), + "profit": ("profit", "margin", "income", "earnings"), + "cost": ("cost", "expense", "spend", "charge"), + "amount": ("amount", "value", "price", "total", "sum"), + "value": ("value", "amount", "price", "total"), + "quantity": ("quantity", "qty", "volume", "units", "count"), + "average": ("average", "avg", "mean"), + "avg": ("avg", "average", "mean"), + "rate": ("rate", "ratio", "percent", "percentage"), + "ratio": ("ratio", "rate", "percent", "percentage"), + "percentage": ("percentage", "percent", "pct", "rate"), + "percent": ("percent", "percentage", "pct", "rate"), + "duration": ("duration", "turnaround", "elapsed", "cycle", "leadtime", "time"), + "turnaround": ("turnaround", "duration", "elapsed", "cycle", "leadtime", "time"), +} +_ANALYSIS_TERMS = { + "analysis", + "analyze", + "dashboard", + "summary", + "summarize", + "trend", + "compare", + "comparison", + "breakdown", + "distribution", + "performance", + "ranking", + "top", + "bottom", + "highest", + "lowest", +} +_COUNT_VOLUME_TERMS = { + "count", + "counts", + "number", + "volume", + "records", + "record", + "rows", + "row", + "how many", +} +_TEMPORAL_TERMS = { + "date", + "timestamp", + "month", + "monthly", + "year", + "yearly", + "week", + "weekly", + "day", + "daily", + "quarter", + "quarterly", + "trend", + "over time", +} +_TEMPORAL_IDENTIFIER_TERMS = { + "date", + "time", + "timestamp", + "month", + "year", + "week", + "day", + "quarter", + "created", + "updated", + "modified", + "started", + "ended", + "closed", + "approved", +} +_DIMENSION_TERMS = { + "category", + "type", + "status", + "source", + "region", + "country", + "market", + "customer", + "product", + "salesperson", + "owner", + "assignee", + "division", + "department", + "location", +} + + +def _contains_phrase(text: str, terms: set[str] | tuple[str, ...]) -> bool: + normalized = f" {str(text or '').lower()} " + return any(f" {term.lower()} " in normalized for term in terms if " " in term) or any( + re.search(rf"\b{re.escape(term.lower())}\b", normalized) + for term in terms + if " " not in term + ) + + +def _query_requests_specific_metric(query: str) -> bool: + normalized = str(query or "").lower() + if not normalized: + return False + + return any( + re.search(rf"\b{re.escape(term)}\b", normalized) + for term in _SPECIFIC_METRIC_TERM_GROUPS + ) + + +def _query_requests_count_volume(query: str) -> bool: + return _contains_phrase(query, _COUNT_VOLUME_TERMS) + + +def _query_requests_time_analysis(query: str) -> bool: + normalized = str(query or "").lower() + if _contains_phrase(normalized, _TEMPORAL_TERMS): + return True + + return bool( + re.search( + r"\b(?:last|next|previous|prior|this)\s+" + r"(?:\d+\s+)?(?:day|week|month|quarter|year)s?\b", + normalized, + ) + or re.search(r"\b(?:between|since|before|after)\b", normalized) + ) + + +def _query_requests_time_bucket(query: str) -> bool: + normalized = str(query or "").lower() + return bool( + _contains_phrase( + normalized, + { + "daily", + "weekly", + "monthly", + "quarterly", + "yearly", + "trend", + "over time", + "time series", + }, + ) + or re.search( + r"\b(?:by|per|for each)\s+" + r"(?:day|week|month|quarter|year)s?\b", + normalized, + ) + ) + + +def _query_requests_grouped_analysis(query: str) -> bool: + normalized = str(query or "").lower() + return bool( + re.search(r"\b(?:by|per|across)\s+[A-Za-z_][A-Za-z0-9_ -]*", normalized) + or re.search(r"\bfor each\s+[A-Za-z_][A-Za-z0-9_ -]*", normalized) + or _contains_phrase(normalized, {"breakdown", "distribution", "grouped"}) + ) + + +def _sql_is_plain_count(sql: str) -> bool: + if re.search(r"\bGROUP\s+BY\b", sql, flags=re.IGNORECASE): + return False + return bool(_PLAIN_COUNT_SQL_PATTERN.match(sql or "")) + + +def _compacted_schema_identifiers( + valid_table_columns: dict[str, list[str]], +) -> set[str]: + identifiers: set[str] = set() + for table, columns in valid_table_columns.items(): + identifiers.add(_compact_sql_identifier(table)) + for column in columns or []: + identifiers.add(_compact_sql_identifier(column)) + return {identifier for identifier in identifiers if identifier} + + +def _compacted_sql_identifiers(sql: str) -> set[str]: + identifiers = { + _compact_sql_identifier(identifier) + for identifier in re.findall( + r'"([^"]+)"|`([^`]+)`|\[([^\]]+)\]|\b([A-Za-z_][A-Za-z0-9_$]*)\b', + sql or "", + ) + for identifier in identifier + if identifier + } + return {identifier for identifier in identifiers if identifier} + + +def _sql_references_term_group(sql: str, terms: tuple[str, ...]) -> bool: + sql_identifiers = _compacted_sql_identifiers(sql) + compact_terms = {_compact_sql_identifier(term) for term in terms} + return any( + term + and any(term in identifier or identifier in term for identifier in sql_identifiers) + for term in compact_terms + ) + + +def _schema_supports_term_group( + valid_table_columns: dict[str, list[str]], + terms: tuple[str, ...], +) -> bool: + schema_identifiers = _compacted_schema_identifiers(valid_table_columns) + compact_terms = {_compact_sql_identifier(term) for term in terms} + return any( + term + and any(term in identifier or identifier in term for identifier in schema_identifiers) + for term in compact_terms + ) + + +def _missing_metric_support( + query: str, + sql: str, + valid_table_columns: dict[str, list[str]], +) -> list[str]: + normalized_query = str(query or "").lower() + missing_terms: list[str] = [] + + for term, group in _SPECIFIC_METRIC_TERM_GROUPS.items(): + if not re.search(rf"\b{re.escape(term)}\b", normalized_query): + continue + if _sql_references_term_group(sql, group): + continue + if not _schema_supports_term_group(valid_table_columns, group): + missing_terms.append(term) + + return sorted(set(missing_terms)) + + +def _sql_has_temporal_reference( + sql: str, + valid_table_columns: dict[str, list[str]], +) -> bool: + if re.search( + r"\b(?:DATEPART|DATE_TRUNC|DATETRUNC|EXTRACT|TO_TIMESTAMP|CAST)\s*\(", + sql or "", + flags=re.IGNORECASE, + ): + return True + + return _sql_references_term_group( + sql, + tuple(_TEMPORAL_IDENTIFIER_TERMS), + ) or any( + _schema_supports_term_group({table: [column]}, tuple(_TEMPORAL_IDENTIFIER_TERMS)) + and _sql_references_term_group(sql, (column,)) + for table, columns in valid_table_columns.items() + for column in columns + ) + + +def _semantic_analysis_items( + semantic_analysis: dict[str, Any] | None, + key: str, +) -> list[str]: + if not isinstance(semantic_analysis, dict): + return [] + + value = semantic_analysis.get(key) + if isinstance(value, str): + return [value] if value.strip() else [] + if isinstance(value, list): + return [ + str(item).strip() + for item in value + if item is not None and str(item).strip() + ] + return [] + + +def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: + if not isinstance(semantic_analysis, dict) or not semantic_analysis: + return False + semantic_keys = { + "analytical_intent", + "entities", + "identifiers", + "metrics", + "dimensions", + "filters", + "aggregations", + "relationships", + "time_constraints", + "ranking", + "supported_schema_objects", + "missing_requirements", + "ambiguous_requirements", + "support_reasoning", + } + return any(semantic_analysis.get(key) for key in semantic_keys) + + +def get_schema_intent_analysis_error( + semantic_analysis: dict[str, Any] | None, +) -> str | None: + if not _has_semantic_analysis(semantic_analysis): + return None + + missing_requirements = _semantic_analysis_items( + semantic_analysis, "missing_requirements" + ) + if missing_requirements: + return ( + "The active datasource schema does not expose the information needed " + "to answer the request: " + f"{', '.join(missing_requirements)}. I cannot generate unrelated SQL." + ) + + ambiguous_requirements = _semantic_analysis_items( + semantic_analysis, "ambiguous_requirements" + ) + if ambiguous_requirements: + return ( + "The request has multiple equally plausible schema interpretations: " + f"{', '.join(ambiguous_requirements)}. Please clarify which one to use." + ) + + if semantic_analysis.get("is_fully_supported") is False: + support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() + if support_reasoning: + return ( + "The selected schema does not fully support the request: " + f"{support_reasoning}" + ) + return ( + "The selected schema does not fully support every required component " + "of the request. I cannot generate unrelated SQL." + ) + + return None + + +def _semantic_analysis_requests_record_count( + semantic_analysis: dict[str, Any], +) -> bool: + analytical_intent = str(semantic_analysis.get("analytical_intent") or "").lower() + if analytical_intent == "record_count": + return True + + semantic_text = " ".join( + item + for key in ("metrics", "aggregations") + for item in _semantic_analysis_items(semantic_analysis, key) + ) + return bool( + re.search( + r"\b(?:count|number of|volume|record count|row count|count records|count rows|number of records)\b", + semantic_text.lower(), + ) + ) + + +def _validate_sql_against_semantic_analysis( + semantic_analysis: dict[str, Any] | None, + sql: str, + valid_table_columns: dict[str, list[str]], +) -> str | None: + if not _has_semantic_analysis(semantic_analysis): + return None + + if analysis_error := get_schema_intent_analysis_error(semantic_analysis): + return analysis_error + + analytical_intent = str( + semantic_analysis.get("analytical_intent") or "" + ).strip().lower() + metrics = _semantic_analysis_items(semantic_analysis, "metrics") + dimensions = _semantic_analysis_items(semantic_analysis, "dimensions") + aggregations = _semantic_analysis_items(semantic_analysis, "aggregations") + time_constraints = _semantic_analysis_items( + semantic_analysis, "time_constraints" + ) + ranking = _semantic_analysis_items(semantic_analysis, "ranking") + requests_record_count = _semantic_analysis_requests_record_count( + semantic_analysis + ) + analytical_sql_intents = { + "summary", + "comparison", + "trend", + "dashboard", + "kpi", + "ranking", + } + + if _sql_is_plain_count(sql) and ( + (metrics and not requests_record_count) + or dimensions + or ranking + or (analytical_intent in analytical_sql_intents and not requests_record_count) + ): + return ( + "Generated SQL answers with a generic record count, but the semantic " + "analysis requires specific business metrics, dimensions, ranking, " + "or analytical calculations from the active schema." + ) + + if time_constraints and not _sql_has_temporal_reference(sql, valid_table_columns): + return ( + "Generated SQL does not use a temporal field or supported date/time " + "expression, but the semantic analysis identified time constraints " + "or trend requirements." + ) + + if (dimensions or analytical_intent == "trend") and _AGGREGATE_PATTERN.search( + sql or "" + ): + has_grouping = bool(re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE)) + if not has_grouping: + return ( + "Generated SQL aggregates results without grouping by the " + "dimensions or time grain identified in the semantic analysis." + ) + + if ranking and not re.search( + r"\b(?:ORDER\s+BY|LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", + sql or "", + flags=re.IGNORECASE, + ): + return ( + "Generated SQL does not include sorting or limiting logic required " + "by the ranking intent." + ) + + if aggregations and not _AGGREGATE_PATTERN.search(sql or ""): + return ( + "Generated SQL does not include the aggregation required by the " + "semantic analysis." + ) + + return None + + +def validate_sql_intent_alignment( + query: str | None, + sql: str, + valid_table_columns: dict[str, list[str]] | None = None, + semantic_analysis: dict[str, Any] | None = None, +) -> str | None: + valid_table_columns = valid_table_columns or {} + + semantic_validation_error = _validate_sql_against_semantic_analysis( + semantic_analysis, + sql, + valid_table_columns, + ) + if semantic_validation_error: + return semantic_validation_error + + if not query: + return None + + normalized_query = str(query or "").lower() + asks_specific_metric = _query_requests_specific_metric(normalized_query) + asks_count_volume = _query_requests_count_volume(normalized_query) + asks_time_analysis = _query_requests_time_analysis(normalized_query) + asks_time_bucket = _query_requests_time_bucket(normalized_query) + asks_grouped_analysis = _query_requests_grouped_analysis(normalized_query) + asks_analysis = _contains_phrase(normalized_query, _ANALYSIS_TERMS) + + if _sql_is_plain_count(sql) and ( + asks_specific_metric + or asks_time_bucket + or asks_grouped_analysis + or (asks_analysis and not asks_count_volume) + ): + return ( + "Generated SQL answers with a generic record count, but the question " + "asks for a specific metric, trend, grouping, comparison, dashboard, " + "or analysis. Select schema elements that directly support the " + "requested business intent, or report that the schema does not expose them." + ) + + if asks_time_analysis and not _sql_has_temporal_reference(sql, valid_table_columns): + return ( + "Generated SQL does not use a temporal field or supported date/time " + "expression, but the question asks for a time range or trend. The " + "active schema must expose a relevant date/time column to answer this." + ) + + if (asks_grouped_analysis or asks_time_bucket) and _AGGREGATE_PATTERN.search( + sql or "" + ): + has_grouping = bool(re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE)) + if not has_grouping: + return ( + "Generated SQL aggregates results without grouping by the requested " + "dimension. Use an explicit schema column for the requested grouping, " + "or report that the schema does not expose that dimension." + ) + + missing_metric_terms = _missing_metric_support( + normalized_query, + sql, + valid_table_columns, + ) + if missing_metric_terms: + missing = ", ".join(missing_metric_terms) + return ( + "The active schema does not expose columns or metrics that directly " + f"support the requested business term(s): {missing}. Do not generate " + "an unrelated query." + ) + + if asks_grouped_analysis and _contains_phrase(normalized_query, _DIMENSION_TERMS): + missing_dimensions = [ + term + for term in _DIMENSION_TERMS + if re.search(rf"\b{re.escape(term)}\b", normalized_query) + and not _sql_references_term_group(sql, (term,)) + and not _schema_supports_term_group(valid_table_columns, (term,)) + ] + if missing_dimensions: + return ( + "The active schema does not expose the requested grouping " + f"dimension(s): {', '.join(sorted(missing_dimensions))}. Do not " + "generate an unrelated query." + ) + + return None + + def construct_ask_history_messages( histories: list[Any] | list[dict], ) -> list[ChatMessage]: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 107c5c3479..102f5dcd92 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -10,7 +10,7 @@ from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import BaseModel, Field from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider @@ -31,23 +31,47 @@ table_columns_selection_system_prompt = """ ### TASK ### -You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. +You are a highly skilled data analyst. Your goal is to examine the provided active deployed database schema, interpret the posed question, and identify the specific tables, columns, metrics, views, and relationships required to construct an accurate SQL query. The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### -1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. -2. For each table, provide a clear and concise reasoning for why specific columns are selected. -3. List each reason as part of a step-by-step chain of thought, justifying the inclusion of each column. -4. If a "." is included in columns, put the name before the first dot into chosen columns. -5. The number of columns chosen must match the number of reasoning. -6. Final chosen columns must be only column names, don't prefix it with table names. -7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +1. First perform a semantic analysis of the user's request. Identify intended business entities, identifiers, descriptive attributes, metrics, dimensions, filters, aggregations, relationships, time constraints, ranking requirements, and analytical intent such as retrieval, detailed records, summary, comparison, trend analysis, dashboard, KPI, ranking, or record count. +2. Map each business term to explicit schema objects only when the active schema directly supports that term. Distinguish entities such as customer/order/invoice/product from identifiers such as order ID or invoice number, descriptive attributes, and measurable metrics such as amount, quantity, cost, profit, revenue, or duration. +3. Select tables and columns by semantic fit to the full request, not by isolated keyword overlap or commonly used default tables. +4. Include join keys and relationship columns needed to connect selected tables. Do not invent relationships or foreign keys. +5. If the schema does not support a requested entity, metric, dimension, filter, time range, aggregation, or ranking requirement, record it in `missing_requirements`. +6. If multiple schema interpretations are equally plausible and the question does not disambiguate them, record them in `ambiguous_requirements`. +7. Set `is_fully_supported` to false when any required request component is missing or ambiguous. +8. For each selected table, provide a concise reason for why the table is semantically relevant. +9. For each selected column, provide a concise reason for why the column is necessary. +10. If a "." is included in columns, put the name before the first dot into chosen columns. +11. The number of columns chosen must match the number of reasoning. +12. Final chosen columns must be only column names, don't prefix it with table names. +13. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +14. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: { + "semantic_analysis": { + "analytical_intent": "retrieval | detailed_records | summary | comparison | trend | dashboard | kpi | ranking | record_count | other", + "entities": ["business entities requested by the user"], + "identifiers": ["identifier fields requested by the user"], + "metrics": ["business metrics or measures requested by the user"], + "dimensions": ["grouping or descriptive dimensions requested by the user"], + "filters": ["filters or predicates requested by the user"], + "aggregations": ["aggregation or calculation requirements"], + "relationships": ["required joins or relationships"], + "time_constraints": ["time filters, grains, or trend requirements"], + "ranking": ["top/bottom/order/limit requirements"], + "supported_schema_objects": ["table.column or metric names that directly support the request"], + "missing_requirements": ["required request components not supported by the schema"], + "ambiguous_requirements": ["request components with multiple equally plausible schema mappings"], + "is_fully_supported": true, + "support_reasoning": "Concise explanation of whether the selected schema fully supports the request" + }, "results": [ { "table_selection_reason": "Reason for selecting tablename1", @@ -82,6 +106,7 @@ - Each table key must list only the columns relevant to answering the question. - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. +- Populate `semantic_analysis` before `results`; use it to verify the selected schema directly supports the request. - Be logical, concise, and ensure the output strictly follows the required JSON format. - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. @@ -355,6 +380,7 @@ def check_using_db_schemas_without_pruning( "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, + "semantic_analysis": {}, } @@ -405,9 +431,9 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] + retrieval_payload = orjson.loads(filter_columns_in_tables["replies"][0]) + columns_and_tables_needed = retrieval_payload.get("results", []) + semantic_analysis = retrieval_payload.get("semantic_analysis") or {} # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -467,6 +493,7 @@ def construct_retrieval_results( "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, + "semantic_analysis": semantic_analysis, } else: retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] @@ -478,6 +505,9 @@ def construct_retrieval_results( ], "has_metric": check_using_db_schemas_without_pruning["has_metric"], "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], + "semantic_analysis": check_using_db_schemas_without_pruning.get( + "semantic_analysis", {} + ), } @@ -493,7 +523,26 @@ class MatchingTable(BaseModel): table_selection_reason: str +class SemanticAnalysis(BaseModel): + analytical_intent: str = "" + entities: list[str] = Field(default_factory=list) + identifiers: list[str] = Field(default_factory=list) + metrics: list[str] = Field(default_factory=list) + dimensions: list[str] = Field(default_factory=list) + filters: list[str] = Field(default_factory=list) + aggregations: list[str] = Field(default_factory=list) + relationships: list[str] = Field(default_factory=list) + time_constraints: list[str] = Field(default_factory=list) + ranking: list[str] = Field(default_factory=list) + supported_schema_objects: list[str] = Field(default_factory=list) + missing_requirements: list[str] = Field(default_factory=list) + ambiguous_requirements: list[str] = Field(default_factory=list) + is_fully_supported: bool | None = None + support_reasoning: str = "" + + class RetrievalResults(BaseModel): + semantic_analysis: SemanticAnalysis | None = None results: list[MatchingTable] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9db38bd5a8..9731665919 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -11,6 +11,7 @@ from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, construct_valid_table_names, + get_schema_intent_analysis_error, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, ) @@ -1288,7 +1289,7 @@ def _build_schema_grounded_table_question_sql( wants_monthly_count = any( term in normalized for term in ("monthly", "by month", "per month", "month-wise") - ) and any(term in normalized for term in ("count", "records", "rows")) + ) and any(term in normalized for term in ("count", "record", "records", "rows")) if wants_monthly_count: date_column = self._find_temporal_column_for_query(query, table) if not date_column: @@ -1471,21 +1472,23 @@ def _build_explicit_table_preview_sql( def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names: list[str] = [] for match in re.finditer( - r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\b(?:from|in|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - if table_name and table_name not in table_names: - table_names.append(table_name) + for candidate in self._explicit_table_name_candidates(table_name): + if candidate and candidate not in table_names: + table_names.append(candidate) for match in re.finditer( r"\bin\s+(?:the\s+)?([A-Za-z_][A-Za-z0-9_.$]*)\s+table\b", query or "", flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - if table_name and table_name not in table_names: - table_names.append(table_name) + for candidate in self._explicit_table_name_candidates(table_name): + if candidate and candidate not in table_names: + table_names.append(candidate) for match in re.finditer( r"\bin\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", @@ -1495,9 +1498,10 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: if ( table_name and ("." in table_name or "_" in table_name) - and table_name not in table_names ): - table_names.append(table_name) + for candidate in self._explicit_table_name_candidates(table_name): + if candidate and candidate not in table_names: + table_names.append(candidate) for match in re.finditer( r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", @@ -1507,9 +1511,10 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: if ( table_name and ("." in table_name or "_" in table_name) - and table_name not in table_names ): - table_names.append(table_name) + for candidate in self._explicit_table_name_candidates(table_name): + if candidate and candidate not in table_names: + table_names.append(candidate) if re.search( r"\b(?:repair\s+logs?|repair\s+tickets?|board\s+models?)\b", query or "", @@ -1524,6 +1529,24 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names.append(table_name) return table_names + def _explicit_table_name_candidates(self, table_name: str) -> list[str]: + table_name = str(table_name or "").strip(".,;:()[]{}") + if not table_name: + return [] + + candidates = [table_name] + dotted_parts = [part for part in re.split(r"[.$]", table_name) if part] + if len(dotted_parts) > 1: + underscored = "_".join(dotted_parts) + candidates.append(underscored) + candidates.append(dotted_parts[-1]) + + normalized_candidates: list[str] = [] + for candidate in candidates: + if candidate and candidate not in normalized_candidates: + normalized_candidates.append(candidate) + return normalized_candidates + def _build_direct_orders_sales_sql(self, query: str) -> str | None: return None @@ -5298,6 +5321,7 @@ async def ask( allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) + schema_intent_analysis: dict[str, Any] = {} try: sql_user_query = user_query @@ -5978,6 +6002,9 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6015,6 +6042,9 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6319,6 +6349,32 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + if semantic_support_error := get_schema_intent_analysis_error( + schema_intent_analysis + ): + logger.info( + "ask pipeline - NO_RELEVANT_SQL due to schema intent analysis: %s", + user_query, + ) + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_SQL", + message=semantic_support_error, + ), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = semantic_support_error + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names @@ -6451,6 +6507,7 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, + schema_intent_analysis=schema_intent_analysis, ), ) ).get("post_process", {}) @@ -6473,6 +6530,7 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, + schema_intent_analysis=schema_intent_analysis, ), ) ).get("post_process", {}) @@ -6561,6 +6619,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, ), ) else: @@ -6580,6 +6639,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, ), ) except TimeoutError as generation_timeout: @@ -6617,6 +6677,7 @@ async def ask( if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", + "SCHEMA_INTENT_VALIDATION", }: invalid_sql = failed_dry_run_result.get("sql", invalid_sql) error_message = failed_dry_run_result.get( @@ -6678,6 +6739,7 @@ async def ask( sql_functions=sql_functions, sql_knowledge=sql_knowledge, query=sql_user_query, + schema_intent_analysis=schema_intent_analysis, ), ) diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 25044de18c..a9a0c1eabe 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -7,6 +7,7 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import get_schema_intent_analysis_error from src.utils import trace_metadata from src.web.v1.services import BaseRequest from src.web.v1.services.ask import AskError, AskResult @@ -106,6 +107,7 @@ async def ask_feedback( invalid_sql = None sql_knowledge = None allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval + schema_intent_analysis = {} try: if not self._is_stopped(query_id, self._ask_feedback_results): @@ -159,6 +161,9 @@ async def ask_feedback( ) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] sql_samples = sql_samples_task["formatted_output"].get("documents", []) @@ -167,6 +172,21 @@ async def ask_feedback( ) if not self._is_stopped(query_id, self._ask_feedback_results): + if semantic_support_error := get_schema_intent_analysis_error( + schema_intent_analysis + ): + self._ask_feedback_results[query_id] = AskFeedbackResultResponse( + status="failed", + error=AskError( + code="NO_RELEVANT_SQL", + message=semantic_support_error, + ), + trace_id=trace_id, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = semantic_support_error + return results + self._ask_feedback_results[query_id] = AskFeedbackResultResponse( status="generating", trace_id=trace_id, @@ -178,6 +198,7 @@ async def ask_feedback( contexts=table_ddls, sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, + query=ask_feedback_request.question, project_id=ask_feedback_request.project_id, sql_samples=sql_samples, instructions=instructions, @@ -186,6 +207,7 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -202,7 +224,10 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] != "TIME_OUT": + if failed_dry_run_result["type"] not in { + "TIME_OUT", + "SCHEMA_INTENT_VALIDATION", + }: original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] @@ -248,6 +273,8 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + query=ask_feedback_request.question, + schema_intent_analysis=schema_intent_analysis, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index fa46c2ca7a..e2fc3f7922 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -5,6 +5,7 @@ extract_sql_generation_result, find_invalid_column_references, find_invalid_table_references, + get_schema_intent_analysis_error, get_json_field_instructions, get_metric_instructions, normalize_data_source, @@ -13,6 +14,7 @@ normalize_sql_table_references_to_schema, get_sql_generation_system_prompt, get_text_to_sql_rules, + validate_sql_intent_alignment, ) @@ -1550,3 +1552,177 @@ def test_normalize_generation_result_sql_rewrites_kb_article_created_by_for_mssq assert "created_by," not in normalized assert "GROUP BY created_by" not in normalized assert '"created_by_user_id"' in normalized + + +def test_validate_sql_intent_alignment_rejects_generic_count_for_metric_request(): + error = validate_sql_intent_alignment( + "Show revenue trend by month", + 'SELECT COUNT(*) AS "RecordCount" FROM "orders"', + {"orders": ["created_at", "order_id"]}, + ) + + assert error is not None + assert "generic record count" in error + + +def test_validate_sql_intent_alignment_rejects_trend_without_temporal_field(): + error = validate_sql_intent_alignment( + "Show monthly order volume", + 'SELECT "orders"."status", COUNT(*) AS "RecordCount" ' + 'FROM "orders" GROUP BY "orders"."status"', + {"orders": ["status", "order_id"]}, + ) + + assert error is not None + assert "temporal field" in error + + +def test_validate_sql_intent_alignment_allows_schema_supported_metric_trend(): + error = validate_sql_intent_alignment( + "Show revenue trend by month", + 'SELECT DATEPART(YEAR, "orders"."created_at") AS "year", ' + 'DATEPART(MONTH, "orders"."created_at") AS "month", ' + 'SUM("orders"."revenue") AS "revenue" ' + 'FROM "orders" ' + 'GROUP BY DATEPART(YEAR, "orders"."created_at"), ' + 'DATEPART(MONTH, "orders"."created_at")', + {"orders": ["created_at", "revenue"]}, + ) + + assert error is None + + +def test_validate_sql_intent_alignment_allows_explicit_record_count(): + error = validate_sql_intent_alignment( + "How many records are in orders?", + 'SELECT COUNT(*) AS "RecordCount" FROM "orders"', + {"orders": ["id"]}, + ) + + assert error is None + + +def test_validate_sql_intent_alignment_allows_temporal_record_count_filter(): + error = validate_sql_intent_alignment( + "How many orders were created last month?", + 'SELECT COUNT(*) AS "RecordCount" FROM "orders" ' + 'WHERE "orders"."created_at" >= \'2026-06-01 00:00:00\' ' + 'AND "orders"."created_at" < \'2026-07-01 00:00:00\'', + {"orders": ["created_at", "id"]}, + ) + + assert error is None + + +def test_validate_sql_intent_alignment_rejects_time_bucket_without_grouping(): + error = validate_sql_intent_alignment( + "Show monthly order volume", + 'SELECT COUNT(*) AS "RecordCount" FROM "orders" ' + 'WHERE "orders"."created_at" IS NOT NULL', + {"orders": ["created_at", "id"]}, + ) + + assert error is not None + assert "generic record count" in error + + +def test_validate_sql_intent_alignment_allows_duration_metric_without_trend(): + error = validate_sql_intent_alignment( + "Show average turnaround time by status", + 'SELECT "repairs"."status", AVG("repairs"."turnaround_hours") ' + 'AS "avg_turnaround_hours" FROM "repairs" GROUP BY "repairs"."status"', + {"repairs": ["status", "turnaround_hours"]}, + ) + + assert error is None + + +def test_get_schema_intent_analysis_error_reports_missing_requirements(): + error = get_schema_intent_analysis_error( + { + "missing_requirements": ["net income metric"], + "support_reasoning": "No metric maps to net income.", + } + ) + + assert error is not None + assert "net income metric" in error + + +def test_get_schema_intent_analysis_error_requests_clarification_for_ambiguity(): + error = get_schema_intent_analysis_error( + { + "ambiguous_requirements": [ + "amount could map to gross_amount or net_amount" + ], + } + ) + + assert error is not None + assert "Please clarify" in error + + +def test_get_schema_intent_analysis_error_reports_unsupported_analysis(): + error = get_schema_intent_analysis_error( + { + "is_fully_supported": False, + "support_reasoning": "No relationship connects invoices to products.", + } + ) + + assert error is not None + assert "No relationship connects invoices to products" in error + + +def test_validate_sql_intent_alignment_uses_semantic_analysis_for_metric_count_mismatch(): + error = validate_sql_intent_alignment( + "Show invoice amount by customer", + 'SELECT COUNT(*) AS "RecordCount" FROM "invoices"', + {"invoices": ["customer_id", "invoice_amount"]}, + semantic_analysis={ + "analytical_intent": "summary", + "entities": ["invoice", "customer"], + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "is_fully_supported": True, + }, + ) + + assert error is not None + assert "generic record count" in error + + +def test_validate_sql_intent_alignment_uses_semantic_analysis_for_ranking(): + error = validate_sql_intent_alignment( + "Top customers by invoice amount", + 'SELECT "invoices"."customer_id", SUM("invoices"."invoice_amount") ' + 'AS "total_invoice_amount" FROM "invoices" GROUP BY "invoices"."customer_id"', + {"invoices": ["customer_id", "invoice_amount"]}, + semantic_analysis={ + "analytical_intent": "ranking", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "ranking": ["top customers by invoice amount"], + "is_fully_supported": True, + }, + ) + + assert error is not None + assert "ranking intent" in error + + +def test_validate_sql_intent_alignment_allows_semantic_count_metric(): + error = validate_sql_intent_alignment( + "Show total order count", + 'SELECT COUNT(*) AS "order_count" FROM "orders"', + {"orders": ["id"]}, + semantic_analysis={ + "analytical_intent": "summary", + "entities": ["order"], + "metrics": ["order count"], + "aggregations": ["count orders"], + "is_fully_supported": True, + }, + ) + + assert error is None diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 5b6694bb45..6adbad941d 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -3,6 +3,7 @@ from src.pipelines.retrieval.db_schema_retrieval import ( _is_project_wide_analysis_query, + construct_retrieval_results, dbschema_retrieval, expand_business_terms_for_retrieval, ) @@ -87,3 +88,73 @@ async def run(self, query_embedding, filters): {"field": "project_id", "operator": "==", "value": "project-1"}, ], } + + +def test_construct_retrieval_results_preserves_semantic_analysis(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={"db_schemas": []}, + filter_columns_in_tables={ + "replies": [ + """ + { + "semantic_analysis": { + "analytical_intent": "summary", + "entities": ["invoice"], + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "is_fully_supported": true + }, + "results": [ + { + "table_name": "invoices", + "table_selection_reason": "Contains invoice facts.", + "table_contents": { + "chain_of_thought_reasoning": [ + "Needed to group by customer.", + "Needed to sum invoice amount." + ], + "columns": ["customer_id", "invoice_amount"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "invoices", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "customer_id", + "data_type": "varchar", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "invoice_amount", + "data_type": "double", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "internal_note", + "data_type": "varchar", + "comment": "", + "is_primary_key": False, + }, + ], + } + ], + dbschema_retrieval=[], + ) + + assert result["semantic_analysis"]["metrics"] == ["invoice amount"] + assert result["retrieval_results"][0]["table_name"] == "invoices" + assert "invoice_amount" in result["retrieval_results"][0]["table_ddl"] + assert "internal_note" not in result["retrieval_results"][0]["table_ddl"] From e87181d3a659ee4bd043589af648a6fce4dd9533 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 14:42:12 +0530 Subject: [PATCH 0471/1087] Strengthen semantic SQL intent validation --- .../src/pipelines/generation/utils/sql.py | 311 ++++++++++++++++++ .../retrieval/db_schema_retrieval.py | 48 ++- .../pipelines/generation/test_sql_utils.py | 144 ++++++++ .../retrieval/test_db_schema_retrieval.py | 22 ++ 4 files changed, 520 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index db9450226f..a2d0b3f02d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3502,6 +3502,20 @@ def _semantic_analysis_items( return [] +def _semantic_analysis_dict_items( + semantic_analysis: dict[str, Any] | None, + key: str, +) -> list[dict[str, Any]]: + if not isinstance(semantic_analysis, dict): + return [] + + value = semantic_analysis.get(key) + if not isinstance(value, list): + return [] + + return [item for item in value if isinstance(item, dict)] + + def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: if not isinstance(semantic_analysis, dict) or not semantic_analysis: return False @@ -3517,6 +3531,8 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: "time_constraints", "ranking", "supported_schema_objects", + "concept_mappings", + "interpretations", "missing_requirements", "ambiguous_requirements", "support_reasoning", @@ -3524,6 +3540,42 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: return any(semantic_analysis.get(key) for key in semantic_keys) +def _schema_interpretation_clarification_error( + semantic_analysis: dict[str, Any], +) -> str | None: + interpretations = _semantic_analysis_dict_items( + semantic_analysis, "interpretations" + ) + if not interpretations: + return None + + selected_interpretations = [ + str(interpretation.get("description") or "").strip() + for interpretation in interpretations + if interpretation.get("is_selected") is True + and str(interpretation.get("description") or "").strip() + ] + if len(selected_interpretations) > 1: + return ( + "The request has multiple selected schema interpretations: " + f"{', '.join(selected_interpretations)}. Please clarify which one to use." + ) + + clarification_interpretations = [ + str(interpretation.get("description") or "").strip() + for interpretation in interpretations + if interpretation.get("needs_clarification") is True + and str(interpretation.get("description") or "").strip() + ] + if clarification_interpretations: + return ( + "The request needs clarification before SQL generation: " + f"{', '.join(clarification_interpretations)}." + ) + + return None + + def get_schema_intent_analysis_error( semantic_analysis: dict[str, Any] | None, ) -> str | None: @@ -3549,6 +3601,11 @@ def get_schema_intent_analysis_error( f"{', '.join(ambiguous_requirements)}. Please clarify which one to use." ) + if interpretation_error := _schema_interpretation_clarification_error( + semantic_analysis + ): + return interpretation_error + if semantic_analysis.get("is_fully_supported") is False: support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() if support_reasoning: @@ -3584,6 +3641,253 @@ def _semantic_analysis_requests_record_count( ) +def _semantic_concept_mappings( + semantic_analysis: dict[str, Any] | None, +) -> list[dict[str, Any]]: + return _semantic_analysis_dict_items(semantic_analysis, "concept_mappings") + + +def _schema_object_parts(schema_object: str) -> list[str]: + return [ + _normalize_sql_identifier(part.strip()) + for part in str(schema_object or "").split(".") + if part.strip() + ] + + +def _schema_object_table_matches( + referenced_table: str, + expected_table: str, +) -> bool: + referenced_suffixes = { + suffix.lower() for suffix in _table_reference_suffixes(referenced_table) + } + expected_suffixes = { + suffix.lower() for suffix in _table_reference_suffixes(expected_table) + } + return bool(referenced_suffixes & expected_suffixes) + + +def _sql_contains_identifier(sql: str, identifier: str) -> bool: + identifier = _normalize_sql_identifier(str(identifier or "").strip()) + if not identifier: + return False + + escaped = re.escape(identifier) + quoted_identifier_pattern = rf'(?:"{escaped}"|`{escaped}`|\[{escaped}\])' + bare_identifier_pattern = rf"(? bool: + table_name = str(table_name or "").strip() + if not table_name: + return False + + table_candidates = {table_name, *_table_reference_suffixes(table_name)} + return any(_sql_contains_identifier(sql, candidate) for candidate in table_candidates) + + +def _sql_references_schema_object( + sql: str, + schema_object: str, + valid_table_columns: dict[str, list[str]], +) -> bool: + parts = _schema_object_parts(schema_object) + if not parts: + return False + + if len(parts) == 1: + return _sql_contains_identifier(sql, parts[0]) + + expected_column = parts[-1] + expected_table = ".".join(parts[:-1]) + aliases = _extract_table_aliases(sql, valid_table_columns) + + for match in _SQL_QUALIFIED_COLUMN_PATTERN.finditer(sql or ""): + qualifier = _normalize_sql_identifier(match.group("qualifier")) + column = _normalize_sql_identifier(match.group("column")) + referenced_table = aliases.get(qualifier.lower(), qualifier) + if ( + column.lower() == expected_column.lower() + and _schema_object_table_matches(referenced_table, expected_table) + ): + return True + + if _sql_references_table(sql, expected_table) and _sql_contains_identifier( + sql, expected_column + ): + return True + + return False + + +def _sql_references_schema_object_table( + sql: str, + schema_object: str, +) -> bool: + parts = _schema_object_parts(schema_object) + if len(parts) < 2: + return _sql_references_table(sql, schema_object) + return _sql_references_table(sql, ".".join(parts[:-1])) + + +def _mapping_schema_objects(mapping: dict[str, Any]) -> list[str]: + value = mapping.get("schema_objects") + if isinstance(value, str): + return [value] if value.strip() else [] + if isinstance(value, list): + return [ + str(item).strip() + for item in value + if item is not None and str(item).strip() + ] + return [] + + +def _mapping_concept_type(mapping: dict[str, Any]) -> str: + return str(mapping.get("concept_type") or "").strip().lower() + + +def _mapping_request_concept(mapping: dict[str, Any]) -> str: + return str(mapping.get("request_concept") or "requested concept").strip() + + +def _requested_aggregate_functions(*texts: str) -> set[str]: + joined = " ".join(text for text in texts if text).lower() + aggregate_terms = { + "SUM": r"\b(?:sum|total)\b", + "AVG": r"\b(?:avg|average|mean)\b", + "MIN": r"\b(?:min|minimum|lowest|smallest)\b", + "MAX": r"\b(?:max|maximum|highest|largest)\b", + "COUNT": r"\b(?:count|number of|how many|volume)\b", + } + return { + function + for function, pattern in aggregate_terms.items() + if re.search(pattern, joined, flags=re.IGNORECASE) + } + + +def _sql_has_aggregate_function(sql: str, function_name: str) -> bool: + return bool( + re.search( + rf"\b{re.escape(function_name)}\s*\(", + sql or "", + flags=re.IGNORECASE, + ) + ) + + +def _validate_sql_against_concept_mappings( + semantic_analysis: dict[str, Any], + sql: str, + valid_table_columns: dict[str, list[str]], +) -> str | None: + mappings = _semantic_concept_mappings(semantic_analysis) + if not mappings: + return None + + requests_record_count = _semantic_analysis_requests_record_count( + semantic_analysis + ) + aggregation_text = " ".join( + _semantic_analysis_items(semantic_analysis, "aggregations") + ) + has_grouping = bool(re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE)) + has_ordering = bool( + re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE) + ) + has_limit = bool( + re.search( + r"\b(?:LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", + sql or "", + flags=re.IGNORECASE, + ) + ) + + for mapping in mappings: + if mapping.get("required_in_sql") is False: + continue + + concept_type = _mapping_concept_type(mapping) + request_concept = _mapping_request_concept(mapping) + schema_objects = _mapping_schema_objects(mapping) + if not schema_objects: + return ( + "The semantic analysis did not map the required " + f"{concept_type or 'concept'} '{request_concept}' to a schema " + "object. I cannot generate unrelated SQL." + ) + + references_concept = any( + _sql_references_schema_object(sql, schema_object, valid_table_columns) + for schema_object in schema_objects + ) + if ( + not references_concept + and concept_type == "metric" + and requests_record_count + and _sql_is_plain_count(sql) + ): + references_concept = any( + _sql_references_schema_object_table(sql, schema_object) + for schema_object in schema_objects + ) + + if not references_concept: + return ( + "Generated SQL does not reference schema objects mapped to the " + f"required {concept_type or 'concept'} '{request_concept}': " + f"{', '.join(schema_objects)}." + ) + + if concept_type == "metric": + if _sql_is_plain_count(sql) and not requests_record_count: + return ( + "Generated SQL answers with a generic record count, but the " + f"requested metric '{request_concept}' maps to " + f"{', '.join(schema_objects)} and must be retrieved or " + "calculated from that schema object." + ) + + requested_functions = _requested_aggregate_functions( + request_concept, + str(mapping.get("mapping_reason") or ""), + aggregation_text, + ) + for function_name in requested_functions: + if function_name == "COUNT" and requests_record_count: + continue + if not _sql_has_aggregate_function(sql, function_name): + return ( + "Generated SQL does not use the aggregation required " + f"for metric '{request_concept}': {function_name}." + ) + + if concept_type in {"dimension", "time"} and _AGGREGATE_PATTERN.search( + sql or "" + ) and not has_grouping: + return ( + "Generated SQL aggregates results without grouping by the " + f"required {concept_type} '{request_concept}'." + ) + + if concept_type == "ranking" and (not has_ordering or not has_limit): + return ( + "Generated SQL does not include sorting and limiting logic " + f"required by ranking concept '{request_concept}'." + ) + + return None + + def _validate_sql_against_semantic_analysis( semantic_analysis: dict[str, Any] | None, sql: str, @@ -3595,6 +3899,13 @@ def _validate_sql_against_semantic_analysis( if analysis_error := get_schema_intent_analysis_error(semantic_analysis): return analysis_error + if mapping_error := _validate_sql_against_concept_mappings( + semantic_analysis, + sql, + valid_table_columns, + ): + return mapping_error + analytical_intent = str( semantic_analysis.get("analytical_intent") or "" ).strip().lower() diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 102f5dcd92..8dbb2c1dfd 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -45,11 +45,13 @@ 7. Set `is_fully_supported` to false when any required request component is missing or ambiguous. 8. For each selected table, provide a concise reason for why the table is semantically relevant. 9. For each selected column, provide a concise reason for why the column is necessary. -10. If a "." is included in columns, put the name before the first dot into chosen columns. -11. The number of columns chosen must match the number of reasoning. -12. Final chosen columns must be only column names, don't prefix it with table names. -13. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -14. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. +11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance and mark the selected interpretation only when it is clearly the best supported one. +12. If a "." is included in columns, put the name before the first dot into chosen columns. +13. The number of columns chosen must match the number of reasoning. +14. Final chosen columns must be only column names, don't prefix it with table names. +15. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +16. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -67,6 +69,24 @@ "time_constraints": ["time filters, grains, or trend requirements"], "ranking": ["top/bottom/order/limit requirements"], "supported_schema_objects": ["table.column or metric names that directly support the request"], + "concept_mappings": [ + { + "request_concept": "business concept from the user request", + "concept_type": "entity | identifier | dimension | metric | filter | time | aggregation | ranking | relationship | comparison", + "schema_objects": ["table.column, table, view, metric, or relationship object that directly supports the concept"], + "required_in_sql": true, + "confidence": 0.0, + "mapping_reason": "Why these schema objects semantically support the concept" + } + ], + "interpretations": [ + { + "description": "Possible interpretation of the request", + "schema_objects": ["schema objects used by this interpretation"], + "confidence": 0.0, + "is_selected": true + } + ], "missing_requirements": ["required request components not supported by the schema"], "ambiguous_requirements": ["request components with multiple equally plausible schema mappings"], "is_fully_supported": true, @@ -523,6 +543,22 @@ class MatchingTable(BaseModel): table_selection_reason: str +class SemanticConceptMapping(BaseModel): + request_concept: str = "" + concept_type: str = "" + schema_objects: list[str] = Field(default_factory=list) + required_in_sql: bool = True + confidence: float | None = None + mapping_reason: str = "" + + +class SemanticInterpretation(BaseModel): + description: str = "" + schema_objects: list[str] = Field(default_factory=list) + confidence: float | None = None + is_selected: bool = False + + class SemanticAnalysis(BaseModel): analytical_intent: str = "" entities: list[str] = Field(default_factory=list) @@ -535,6 +571,8 @@ class SemanticAnalysis(BaseModel): time_constraints: list[str] = Field(default_factory=list) ranking: list[str] = Field(default_factory=list) supported_schema_objects: list[str] = Field(default_factory=list) + concept_mappings: list[SemanticConceptMapping] = Field(default_factory=list) + interpretations: list[SemanticInterpretation] = Field(default_factory=list) missing_requirements: list[str] = Field(default_factory=list) ambiguous_requirements: list[str] = Field(default_factory=list) is_fully_supported: bool | None = None diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index e2fc3f7922..bc68923326 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1726,3 +1726,147 @@ def test_validate_sql_intent_alignment_allows_semantic_count_metric(): ) assert error is None + + +def test_validate_sql_intent_alignment_rejects_unmapped_metric_substitution(): + error = validate_sql_intent_alignment( + "Show invoice amount by customer", + 'SELECT "invoices"."customer_id", COUNT(*) AS "invoice_amount" ' + 'FROM "invoices" GROUP BY "invoices"."customer_id"', + {"invoices": ["customer_id", "invoice_amount"]}, + semantic_analysis={ + "analytical_intent": "summary", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "aggregations": ["sum invoice amount"], + "concept_mappings": [ + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["invoices.invoice_amount"], + "required_in_sql": True, + "confidence": 0.95, + } + ], + "is_fully_supported": True, + }, + ) + + assert error is not None + assert "invoice amount" in error + assert "invoices.invoice_amount" in error + + +def test_validate_sql_intent_alignment_allows_mapped_metric_sql(): + error = validate_sql_intent_alignment( + "Show total invoice amount by customer", + 'SELECT "invoices"."customer_id", ' + 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' + 'FROM "invoices" GROUP BY "invoices"."customer_id"', + {"invoices": ["customer_id", "invoice_amount"]}, + semantic_analysis={ + "analytical_intent": "summary", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "aggregations": ["sum invoice amount"], + "concept_mappings": [ + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["invoices.invoice_amount"], + "required_in_sql": True, + "confidence": 0.95, + }, + { + "request_concept": "customer", + "concept_type": "dimension", + "schema_objects": ["invoices.customer_id"], + "required_in_sql": True, + "confidence": 0.9, + }, + ], + "is_fully_supported": True, + }, + ) + + assert error is None + + +def test_validate_sql_intent_alignment_rejects_missing_mapped_dimension(): + error = validate_sql_intent_alignment( + "Show total invoice amount by customer", + 'SELECT SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' + 'FROM "invoices"', + {"invoices": ["customer_id", "invoice_amount"]}, + semantic_analysis={ + "analytical_intent": "summary", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "aggregations": ["sum invoice amount"], + "concept_mappings": [ + { + "request_concept": "customer", + "concept_type": "dimension", + "schema_objects": ["invoices.customer_id"], + "required_in_sql": True, + } + ], + "is_fully_supported": True, + }, + ) + + assert error is not None + assert "customer" in error + + +def test_validate_sql_intent_alignment_rejects_mapped_ranking_without_limit(): + error = validate_sql_intent_alignment( + "Top customers by invoice amount", + 'SELECT "invoices"."customer_id", ' + 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' + 'FROM "invoices" GROUP BY "invoices"."customer_id" ' + 'ORDER BY "total_invoice_amount" DESC', + {"invoices": ["customer_id", "invoice_amount"]}, + semantic_analysis={ + "analytical_intent": "ranking", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "ranking": ["top customers by invoice amount"], + "aggregations": ["sum invoice amount"], + "concept_mappings": [ + { + "request_concept": "top customers", + "concept_type": "ranking", + "schema_objects": ["invoices.customer_id"], + "required_in_sql": True, + } + ], + "is_fully_supported": True, + }, + ) + + assert error is not None + assert "sorting and limiting" in error + + +def test_get_schema_intent_analysis_error_rejects_multiple_selected_interpretations(): + error = get_schema_intent_analysis_error( + { + "interpretations": [ + { + "description": "Use gross amount", + "confidence": 0.88, + "is_selected": True, + }, + { + "description": "Use net amount", + "confidence": 0.87, + "is_selected": True, + }, + ], + "is_fully_supported": True, + } + ) + + assert error is not None + assert "multiple selected schema interpretations" in error diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 6adbad941d..afb2d06c8b 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -102,6 +102,24 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): "entities": ["invoice"], "metrics": ["invoice amount"], "dimensions": ["customer"], + "concept_mappings": [ + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["invoices.invoice_amount"], + "required_in_sql": true, + "confidence": 0.95, + "mapping_reason": "invoice_amount stores invoice value" + } + ], + "interpretations": [ + { + "description": "Summarize invoice amount by customer", + "schema_objects": ["invoices.customer_id", "invoices.invoice_amount"], + "confidence": 0.9, + "is_selected": true + } + ], "is_fully_supported": true }, "results": [ @@ -155,6 +173,10 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): ) assert result["semantic_analysis"]["metrics"] == ["invoice amount"] + assert result["semantic_analysis"]["concept_mappings"][0]["schema_objects"] == [ + "invoices.invoice_amount" + ] + assert result["semantic_analysis"]["interpretations"][0]["is_selected"] is True assert result["retrieval_results"][0]["table_name"] == "invoices" assert "invoice_amount" in result["retrieval_results"][0]["table_ddl"] assert "internal_note" not in result["retrieval_results"][0]["table_ddl"] From 188b8a5628950957f9ee29079ebdc21e4de98be7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 15:07:08 +0530 Subject: [PATCH 0472/1087] Make semantic schema mapping drive SQL generation --- .../generation/followup_sql_generation.py | 14 +- .../followup_sql_generation_reasoning.py | 19 ++- .../pipelines/generation/sql_correction.py | 11 +- .../pipelines/generation/sql_generation.py | 14 +- .../generation/sql_generation_reasoning.py | 19 ++- .../pipelines/generation/sql_regeneration.py | 11 +- .../src/pipelines/generation/utils/sql.py | 116 ++++++++++++++ .../retrieval/db_schema_retrieval.py | 2 +- wren-ai-service/src/web/v1/services/ask.py | 144 ++++++++++++++++++ .../src/web/v1/services/ask_feedback.py | 96 +++++++++++- .../pipelines/generation/test_sql_utils.py | 64 ++++++++ 11 files changed, 470 insertions(+), 40 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 a4616ff881..f0f2412891 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -17,6 +17,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -99,12 +100,9 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -This is the pre-generation semantic analysis of the follow-up request against the -active deployed schema. Use it as a contract for table, column, metric, dimension, -filter, time, relationship, aggregation, and ranking selection. -{{ schema_intent_analysis }} +{% if semantic_schema_contract %} +### SEMANTIC SCHEMA CONTRACT ### +{{ semantic_schema_contract }} {% endif %} ### INTENT AND SCHEMA GROUNDING ### @@ -167,7 +165,9 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - schema_intent_analysis=schema_intent_analysis, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 9759617bb9..0374caf609 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -15,6 +15,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_semantic_schema_contract, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -67,13 +68,13 @@ Language: {{ language }} Current Time: {{ current_time }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -Use this semantic analysis as the planning contract for entities, metrics, -dimensions, filters, joins, time constraints, aggregations, ranking, and analytical -intent. If it shows missing or ambiguous requirements, state that limitation in the -plan instead of planning unrelated SQL. -{{ schema_intent_analysis }} +{% if semantic_schema_contract %} +### SEMANTIC SCHEMA CONTRACT ### +Use this contract for entities, metrics, dimensions, filters, joins, time +constraints, aggregations, ranking, and analytical intent. If it shows missing or +ambiguous requirements, state that limitation in the plan instead of planning +unrelated SQL. +{{ semantic_schema_contract }} {% endif %} Let's think step by step. @@ -102,7 +103,9 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), - schema_intent_analysis=schema_intent_analysis, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index be63983227..d1e2e93333 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_sql_generation_model_kwargs, @@ -102,11 +103,11 @@ def get_sql_correction_system_prompt( {% if query %} User's Question: {{ query }} {% endif %} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### +{% if semantic_schema_contract %} +### SEMANTIC SCHEMA CONTRACT ### This is the semantic contract for the corrected SQL. Preserve this intent while fixing syntax or planner errors. -{{ schema_intent_analysis }} +{{ semantic_schema_contract }} {% endif %} {% if invalid_generation_result.original_sql %} Original SQL: {{ invalid_generation_result.original_sql }} @@ -148,7 +149,9 @@ def prompt( documents=documents, valid_table_names=construct_valid_table_names(documents), invalid_generation_result=invalid_generation_result, - schema_intent_analysis=schema_intent_analysis, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 9280f3e8f1..9c568abec0 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -87,12 +88,9 @@ ### QUESTION ### User's Question: {{ query }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -This is the pre-generation semantic analysis of the user's request against the -active deployed schema. Use it as a contract for table, column, metric, dimension, -filter, time, relationship, aggregation, and ranking selection. -{{ schema_intent_analysis }} +{% if semantic_schema_contract %} +### SEMANTIC SCHEMA CONTRACT ### +{{ semantic_schema_contract }} {% endif %} ### INTENT AND SCHEMA GROUNDING ### @@ -170,7 +168,9 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - schema_intent_analysis=schema_intent_analysis, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} 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 4db6239cd6..7d14e29184 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -13,6 +13,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_semantic_schema_contract, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -52,13 +53,13 @@ Language: {{ language }} Current Time: {{ current_time }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -Use this semantic analysis as the planning contract for entities, metrics, -dimensions, filters, joins, time constraints, aggregations, ranking, and analytical -intent. If it shows missing or ambiguous requirements, state that limitation in the -plan instead of planning unrelated SQL. -{{ schema_intent_analysis }} +{% if semantic_schema_contract %} +### SEMANTIC SCHEMA CONTRACT ### +Use this contract for entities, metrics, dimensions, filters, joins, time +constraints, aggregations, ranking, and analytical intent. If it shows missing or +ambiguous requirements, state that limitation in the plan instead of planning +unrelated SQL. +{{ semantic_schema_contract }} {% endif %} Let's think step by step. @@ -85,7 +86,9 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), - schema_intent_analysis=schema_intent_analysis, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index bef428a21c..4f0668ff0f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -109,11 +110,11 @@ def get_sql_regeneration_system_prompt( {% if query %} User's Question: {{ query }} {% endif %} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### +{% if semantic_schema_contract %} +### SEMANTIC SCHEMA CONTRACT ### This is the semantic contract for regenerated SQL. Preserve this intent while improving the original SQL. -{{ schema_intent_analysis }} +{{ semantic_schema_contract }} {% endif %} SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} @@ -145,7 +146,9 @@ def prompt( data_source=data_source, documents=documents, query=query, - schema_intent_analysis=schema_intent_analysis, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a2d0b3f02d..97ae55f7c3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2221,6 +2221,7 @@ def get_sql_generation_system_prompt( 12. Before generating SQL, validate that the selected schema elements directly support all key entities, metrics, dimensions, filters, time ranges, relationships, and aggregations mentioned or implied by the question. 13. Do not answer a specific business metric, trend, summary, comparison, dashboard, or analysis request with a generic record-count query unless the user explicitly asks only for record count. 14. If the required information cannot be derived from the available active schema, return the closest schema-grounded limitation instead of inventing unrelated SQL. +15. If a SEMANTIC SCHEMA CONTRACT is provided, it is the primary source of truth for selecting tables, columns, metrics, joins, filters, grouping, sorting, and date logic. Generate SQL from the highest-confidence validated concept-to-schema mappings in that contract and do not independently infer substitute schema objects. {text_to_sql_rules} @@ -2299,6 +2300,121 @@ def construct_instructions( return _instructions +def _format_semantic_list(label: str, values: list[str]) -> list[str]: + if not values: + return [] + return [f"{label}: {', '.join(values)}"] + + +def construct_semantic_schema_contract( + semantic_analysis: dict[str, Any] | None, +) -> str: + if not _has_semantic_analysis(semantic_analysis): + return "" + + lines: list[str] = [ + "Use this semantic schema contract as the primary source of truth for SQL generation.", + "Generate SQL only from schema objects listed here or in the selected retrieval metadata.", + "Do not infer alternative tables, columns, metrics, joins, or identifiers independently.", + ] + + analytical_intent = str( + semantic_analysis.get("analytical_intent") or "" + ).strip() + if analytical_intent: + lines.append(f"Analytical intent: {analytical_intent}") + + for label, key in ( + ("Entities", "entities"), + ("Identifiers", "identifiers"), + ("Metrics", "metrics"), + ("Dimensions", "dimensions"), + ("Filters", "filters"), + ("Time constraints", "time_constraints"), + ("Aggregations", "aggregations"), + ("Ranking", "ranking"), + ("Relationships", "relationships"), + ("Supported schema objects", "supported_schema_objects"), + ): + lines.extend(_format_semantic_list(label, _semantic_analysis_items(semantic_analysis, key))) + + concept_mappings = _semantic_concept_mappings(semantic_analysis) + if concept_mappings: + lines.append("Required concept-to-schema mappings:") + for mapping in concept_mappings: + schema_objects = _mapping_schema_objects(mapping) + required = "required" if mapping.get("required_in_sql") is not False else "optional" + confidence = mapping.get("confidence") + confidence_text = ( + f", confidence={confidence}" + if confidence is not None and str(confidence).strip() + else "" + ) + mapping_reason = str(mapping.get("mapping_reason") or "").strip() + reason_text = f" Reason: {mapping_reason}" if mapping_reason else "" + lines.append( + "- " + f"{_mapping_request_concept(mapping)} " + f"({_mapping_concept_type(mapping) or 'concept'}, {required}" + f"{confidence_text}) -> {', '.join(schema_objects) or 'NO_MAPPING'}." + f"{reason_text}" + ) + + interpretations = _semantic_analysis_dict_items( + semantic_analysis, "interpretations" + ) + if interpretations: + lines.append("Ranked schema interpretations:") + for interpretation in interpretations: + description = str(interpretation.get("description") or "").strip() + if not description: + continue + selected = "selected" if interpretation.get("is_selected") is True else "candidate" + confidence = interpretation.get("confidence") + confidence_text = ( + f", confidence={confidence}" + if confidence is not None and str(confidence).strip() + else "" + ) + schema_objects = interpretation.get("schema_objects") + if isinstance(schema_objects, list): + schema_text = ", ".join( + str(item).strip() + for item in schema_objects + if item is not None and str(item).strip() + ) + else: + schema_text = "" + schema_suffix = f" Objects: {schema_text}." if schema_text else "" + lines.append( + f"- {description} ({selected}{confidence_text}).{schema_suffix}" + ) + + missing_requirements = _semantic_analysis_items( + semantic_analysis, "missing_requirements" + ) + if missing_requirements: + lines.append(f"Missing requirements: {', '.join(missing_requirements)}") + + ambiguous_requirements = _semantic_analysis_items( + semantic_analysis, "ambiguous_requirements" + ) + if ambiguous_requirements: + lines.append(f"Ambiguous requirements: {', '.join(ambiguous_requirements)}") + + support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() + if support_reasoning: + lines.append(f"Support reasoning: {support_reasoning}") + + lines.append( + "Validation requirement: every required mapping must be represented in the SQL. " + "Do not substitute identifiers for metrics, entities for identifiers, or COUNT(*) " + "for a requested business measure unless the semantic intent explicitly requests a record count." + ) + + return "\n".join(lines) + + def _parse_semantic_metadata_content(content: str) -> Any | None: content = content.strip() if not content: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 8dbb2c1dfd..0fc19d58b1 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -46,7 +46,7 @@ 8. For each selected table, provide a concise reason for why the table is semantically relevant. 9. For each selected column, provide a concise reason for why the column is necessary. 10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. -11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance and mark the selected interpretation only when it is clearly the best supported one. +11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. 12. If a "." is included in columns, put the name before the first dot into chosen columns. 13. The number of columns chosen must match the number of reasoning. 14. Final chosen columns must be only column names, don't prefix it with table names. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9731665919..5377479691 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -6673,7 +6673,151 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: + if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": + semantic_retry_error = failed_dry_run_result.get("error", "") + semantic_retry_query = ( + f"{sql_user_query}\n\n" + "The previous generated SQL failed semantic validation: " + f"{semantic_retry_error}\n" + "Re-run semantic schema retrieval and choose the next " + "highest-confidence concept-to-schema mapping that " + "directly supports the user's requested entities, " + "metrics, dimensions, filters, time constraints, " + "aggregations, relationships, and ranking." + ) + logger.info( + "Retrying semantic schema retrieval after intent validation failure for query_id %s", + query_id, + ) + try: + retry_retrieval_result = await self._run_with_timeout( + "Semantic schema retrieval retry", + self._pipelines["db_schema_retrieval"].run( + query=semantic_retry_query, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=enable_column_pruning, + ), + timeout_seconds=self._schema_retrieval_timeout_seconds, + ) + retry_construct_result = retry_retrieval_result.get( + "construct_retrieval_results", {} + ) + retry_schema_intent_analysis = retry_construct_result.get( + "semantic_analysis", {} + ) + retry_documents, retry_table_names, retry_table_ddls = ( + self._extract_retrieval_metadata( + retry_retrieval_result + ) + ) + retry_support_error = get_schema_intent_analysis_error( + retry_schema_intent_analysis + ) + if retry_documents and not retry_support_error: + _retrieval_result = retry_construct_result + schema_intent_analysis = retry_schema_intent_analysis + documents = retry_documents + table_names = retry_table_names + table_ddls = retry_table_ddls + has_calculated_field = _retrieval_result.get( + "has_calculated_field", False + ) + has_metric = _retrieval_result.get("has_metric", False) + has_json_field = _retrieval_result.get( + "has_json_field", False + ) + + if histories: + text_to_sql_generation_results = ( + await self._run_with_timeout( + "Follow-up SQL generation after semantic retry", + self._pipelines[ + "followup_sql_generation" + ].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), + ) + ) + else: + text_to_sql_generation_results = ( + await self._run_with_timeout( + "SQL generation after semantic retry", + self._pipelines["sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), + ) + ) + + if sql_valid_result := text_to_sql_generation_results[ + "post_process" + ]["valid_generation_result"]: + if ask_result := self._build_validated_ask_result_from_sql( + sql_valid_result.get("sql"), + table_ddls, + sql_user_query, + ): + api_results = [ask_result] + else: + invalid_sql = sql_valid_result.get("sql") + error_message = ( + "SQL generation after semantic retrieval retry did not produce SQL that matches the active datasource schema and question intent." + ) + else: + failed_dry_run_result = ( + text_to_sql_generation_results["post_process"][ + "invalid_generation_result" + ] + ) + invalid_sql = failed_dry_run_result.get( + "sql", invalid_sql + ) + error_message = failed_dry_run_result.get( + "error", error_message + ) + elif retry_support_error: + error_message = retry_support_error + else: + error_message = ( + "Semantic schema retrieval retry did not find a supported mapping for the request." + ) + except Exception as retry_error: + logger.warning( + "Semantic schema retrieval retry failed for query_id %s: %s", + query_id, + retry_error, + ) + while current_sql_correction_retries < max_sql_correction_retries: + if api_results: + break if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index a9a0c1eabe..cbd6f7d169 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -224,7 +224,101 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] not in { + if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": + semantic_retry_error = failed_dry_run_result.get("error", "") + semantic_retry_query = ( + f"{ask_feedback_request.question}\n\n" + "The previous regenerated SQL failed semantic validation: " + f"{semantic_retry_error}\n" + "Re-run semantic schema retrieval and choose the next " + "highest-confidence concept-to-schema mapping that " + "directly supports the user's requested entities, " + "metrics, dimensions, filters, time constraints, " + "aggregations, relationships, and ranking." + ) + logger.info( + "Retrying semantic schema retrieval after feedback intent validation failure for query_id %s", + query_id, + ) + retry_retrieval_result = await self._pipelines[ + "db_schema_retrieval" + ].run( + query=semantic_retry_query, + histories=[], + project_id=ask_feedback_request.project_id, + enable_column_pruning=enable_column_pruning, + ) + retry_construct_result = retry_retrieval_result.get( + "construct_retrieval_results", {} + ) + retry_schema_intent_analysis = retry_construct_result.get( + "semantic_analysis", {} + ) + retry_documents = retry_construct_result.get( + "retrieval_results", [] + ) + retry_table_ddls = [ + document.get("table_ddl") + for document in retry_documents + if isinstance(document, dict) + and document.get("table_ddl") + ] + retry_support_error = get_schema_intent_analysis_error( + retry_schema_intent_analysis + ) + if retry_table_ddls and not retry_support_error: + schema_intent_analysis = retry_schema_intent_analysis + documents = retry_documents + table_ddls = retry_table_ddls + has_calculated_field = retry_construct_result.get( + "has_calculated_field", False + ) + has_metric = retry_construct_result.get( + "has_metric", False + ) + has_json_field = retry_construct_result.get( + "has_json_field", False + ) + + text_to_sql_generation_results = await self._pipelines[ + "sql_regeneration" + ].run( + contexts=table_ddls, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + sql=ask_feedback_request.sql, + query=ask_feedback_request.question, + project_id=ask_feedback_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ) + + if sql_valid_result := text_to_sql_generation_results[ + "post_process" + ]["valid_generation_result"]: + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } + ) + ] + else: + failed_dry_run_result = text_to_sql_generation_results[ + "post_process" + ]["invalid_generation_result"] + elif retry_support_error: + error_message = retry_support_error + + if api_results: + pass + elif failed_dry_run_result["type"] not in { "TIME_OUT", "SCHEMA_INTENT_VALIDATION", }: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index bc68923326..4d8797ee50 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,5 +1,6 @@ from src.pipelines.generation.utils.sql import ( contains_unsupported_mssql_json_access, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, extract_sql_generation_result, @@ -1870,3 +1871,66 @@ def test_get_schema_intent_analysis_error_rejects_multiple_selected_interpretati assert error is not None assert "multiple selected schema interpretations" in error + + +def test_construct_semantic_schema_contract_prioritizes_concept_mappings(): + contract = construct_semantic_schema_contract( + { + "analytical_intent": "summary", + "entities": ["invoice"], + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "aggregations": ["sum invoice amount"], + "supported_schema_objects": [ + "invoices.customer_id", + "invoices.invoice_amount", + ], + "concept_mappings": [ + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["invoices.invoice_amount"], + "required_in_sql": True, + "confidence": 0.95, + "mapping_reason": "invoice_amount stores the requested measure", + } + ], + "interpretations": [ + { + "description": "Summarize invoice amount by customer", + "schema_objects": [ + "invoices.customer_id", + "invoices.invoice_amount", + ], + "confidence": 0.9, + "is_selected": True, + } + ], + "is_fully_supported": True, + } + ) + + assert "primary source of truth" in contract + assert "Required concept-to-schema mappings" in contract + assert "invoice amount (metric, required, confidence=0.95)" in contract + assert "invoices.invoice_amount" in contract + assert "Ranked schema interpretations" in contract + assert "selected" in contract + assert "Do not substitute identifiers for metrics" in contract + + +def test_construct_semantic_schema_contract_allows_legacy_analysis_without_mappings(): + contract = construct_semantic_schema_contract( + { + "analytical_intent": "trend", + "metrics": ["order volume"], + "time_constraints": ["monthly"], + "supported_schema_objects": ["orders.created_at", "orders.id"], + "is_fully_supported": True, + } + ) + + assert "Analytical intent: trend" in contract + assert "Metrics: order volume" in contract + assert "orders.created_at" in contract + assert "Required concept-to-schema mappings" not in contract From fea5587d422c945a986be633a38240186c1b4e9b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 15:33:49 +0530 Subject: [PATCH 0473/1087] Retry SQL generation with alternate semantic mappings --- .../retrieval/db_schema_retrieval.py | 115 ++++++- wren-ai-service/src/web/v1/services/ask.py | 314 ++++++++++++------ .../src/web/v1/services/ask_feedback.py | 175 +++++++--- .../retrieval/test_db_schema_retrieval.py | 58 ++++ 4 files changed, 508 insertions(+), 154 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 0fc19d58b1..f45e24a5b0 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -47,11 +47,14 @@ 9. For each selected column, provide a concise reason for why the column is necessary. 10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. 11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. -12. If a "." is included in columns, put the name before the first dot into chosen columns. -13. The number of columns chosen must match the number of reasoning. -14. Final chosen columns must be only column names, don't prefix it with table names. -15. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -16. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +12. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. +13. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. +14. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. +15. If a "." is included in columns, put the name before the first dot into chosen columns. +16. The number of columns chosen must match the number of reasoning. +17. Final chosen columns must be only column names, don't prefix it with table names. +18. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +19. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -69,6 +72,17 @@ "time_constraints": ["time filters, grains, or trend requirements"], "ranking": ["top/bottom/order/limit requirements"], "supported_schema_objects": ["table.column or metric names that directly support the request"], + "candidate_schema_scores": [ + { + "candidate_id": "candidate-1", + "schema_objects": ["table.column objects included in this candidate"], + "covered_concepts": ["request concepts this candidate supports"], + "missing_concepts": ["request concepts this candidate cannot support"], + "confidence": 0.0, + "is_complete": true, + "selection_reason": "Why this candidate is accepted or rejected" + } + ], "concept_mappings": [ { "request_concept": "business concept from the user request", @@ -145,6 +159,17 @@ ### INPUT ### {{ question }} + +{% if semantic_retry_context %} +### RETRY CONTEXT ### +Previous semantic SQL validation failed. Discard the previous contract and do not reuse rejected schema mappings unless no other complete candidate exists. +Validation failure: {{ semantic_retry_context.validation_error }} +Retry attempt: {{ semantic_retry_context.retry_attempt }} +Rejected schema objects: +{% for schema_object in semantic_retry_context.rejected_schema_objects %} +- {{ schema_object }} +{% endfor %} +{% endif %} """ @@ -394,6 +419,16 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) + if enable_column_pruning or _token_count > context_window_size: + return { + "db_schemas": [], + "tokens": _token_count, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + "semantic_analysis": {}, + } + return { "db_schemas": retrieval_results, "tokens": _token_count, @@ -411,6 +446,7 @@ def prompt( prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, histories: list[AskHistory], + semantic_retry_context: dict[str, Any] | None = None, ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -424,7 +460,11 @@ def prompt( query = "\n".join(previous_query_summaries) + "\n" + query - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) + _prompt = prompt_builder.run( + question=query, + db_schemas=db_schemas, + semantic_retry_context=semantic_retry_context or {}, + ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: return {} @@ -454,6 +494,7 @@ def construct_retrieval_results( retrieval_payload = orjson.loads(filter_columns_in_tables["replies"][0]) columns_and_tables_needed = retrieval_payload.get("results", []) semantic_analysis = retrieval_payload.get("semantic_analysis") or {} + _log_semantic_retrieval_decision(semantic_analysis) # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -531,6 +572,51 @@ def construct_retrieval_results( } +def _semantic_log_items(semantic_analysis: dict[str, Any], key: str) -> list[str]: + value = semantic_analysis.get(key) + if isinstance(value, str): + return [value] if value.strip() else [] + if isinstance(value, list): + return [ + str(item).strip() + for item in value + if item is not None and str(item).strip() + ] + return [] + + +def _log_semantic_retrieval_decision(semantic_analysis: dict[str, Any]) -> None: + if not isinstance(semantic_analysis, dict) or not semantic_analysis: + logger.info("semantic_retrieval_decision=no_semantic_analysis") + return + + concepts = { + "intent": semantic_analysis.get("analytical_intent"), + "entities": _semantic_log_items(semantic_analysis, "entities"), + "identifiers": _semantic_log_items(semantic_analysis, "identifiers"), + "metrics": _semantic_log_items(semantic_analysis, "metrics"), + "dimensions": _semantic_log_items(semantic_analysis, "dimensions"), + "filters": _semantic_log_items(semantic_analysis, "filters"), + "time_constraints": _semantic_log_items( + semantic_analysis, "time_constraints" + ), + "aggregations": _semantic_log_items(semantic_analysis, "aggregations"), + "ranking": _semantic_log_items(semantic_analysis, "ranking"), + } + candidate_scores = semantic_analysis.get("candidate_schema_scores") or [] + selected_contract = { + "supported_schema_objects": semantic_analysis.get( + "supported_schema_objects", [] + ), + "concept_mappings": semantic_analysis.get("concept_mappings", []), + "is_fully_supported": semantic_analysis.get("is_fully_supported"), + "support_reasoning": semantic_analysis.get("support_reasoning"), + } + logger.info("semantic_retrieval_concepts=%s", concepts) + logger.info("semantic_retrieval_candidate_scores=%s", candidate_scores) + logger.info("semantic_retrieval_selected_contract=%s", selected_contract) + + ## End of Pipeline class MatchingTableContents(BaseModel): chain_of_thought_reasoning: list[str] @@ -559,6 +645,16 @@ class SemanticInterpretation(BaseModel): is_selected: bool = False +class SemanticCandidateSchemaScore(BaseModel): + candidate_id: str = "" + schema_objects: list[str] = Field(default_factory=list) + covered_concepts: list[str] = Field(default_factory=list) + missing_concepts: list[str] = Field(default_factory=list) + confidence: float | None = None + is_complete: bool = False + selection_reason: str = "" + + class SemanticAnalysis(BaseModel): analytical_intent: str = "" entities: list[str] = Field(default_factory=list) @@ -571,6 +667,9 @@ class SemanticAnalysis(BaseModel): time_constraints: list[str] = Field(default_factory=list) ranking: list[str] = Field(default_factory=list) supported_schema_objects: list[str] = Field(default_factory=list) + candidate_schema_scores: list[SemanticCandidateSchemaScore] = Field( + default_factory=list + ) concept_mappings: list[SemanticConceptMapping] = Field(default_factory=list) interpretations: list[SemanticInterpretation] = Field(default_factory=list) missing_requirements: list[str] = Field(default_factory=list) @@ -649,8 +748,11 @@ async def run( project_id: Optional[str] = None, histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, + semantic_retry_context: Optional[dict[str, Any]] = None, ): logger.info("Ask Retrieval pipeline is running...") + if semantic_retry_context: + logger.info("semantic_retrieval_retry_context=%s", semantic_retry_context) return await self._pipe.execute( ["construct_retrieval_results"], inputs={ @@ -659,6 +761,7 @@ async def run( "project_id": project_id or "", "histories": histories or [], "enable_column_pruning": enable_column_pruning, + "semantic_retry_context": semantic_retry_context or {}, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 5377479691..16023abdea 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4452,6 +4452,66 @@ def _metadata_from_documents( ] return table_names, table_ddls + @staticmethod + def _semantic_schema_objects(semantic_analysis: dict[str, Any] | None) -> list[str]: + if not isinstance(semantic_analysis, dict): + return [] + + schema_objects: list[str] = [] + supported_schema_objects = semantic_analysis.get("supported_schema_objects") + if isinstance(supported_schema_objects, list): + schema_objects.extend( + str(schema_object).strip() + for schema_object in supported_schema_objects + if schema_object is not None and str(schema_object).strip() + ) + + concept_mappings = semantic_analysis.get("concept_mappings") + if isinstance(concept_mappings, list): + for mapping in concept_mappings: + if not isinstance(mapping, dict): + continue + mapping_schema_objects = mapping.get("schema_objects") + if isinstance(mapping_schema_objects, list): + schema_objects.extend( + str(schema_object).strip() + for schema_object in mapping_schema_objects + if schema_object is not None and str(schema_object).strip() + ) + + interpretations = semantic_analysis.get("interpretations") + if isinstance(interpretations, list): + for interpretation in interpretations: + if not isinstance(interpretation, dict): + continue + if interpretation.get("is_selected") is not True: + continue + interpretation_schema_objects = interpretation.get("schema_objects") + if isinstance(interpretation_schema_objects, list): + schema_objects.extend( + str(schema_object).strip() + for schema_object in interpretation_schema_objects + if schema_object is not None and str(schema_object).strip() + ) + + return sorted(set(schema_objects)) + + @staticmethod + def _semantic_retry_context( + semantic_analysis: dict[str, Any] | None, + validation_error: str, + retry_attempt: int, + rejected_schema_objects: set[str], + ) -> dict[str, Any]: + rejected_schema_objects.update( + AskService._semantic_schema_objects(semantic_analysis) + ) + return { + "validation_error": validation_error, + "retry_attempt": retry_attempt, + "rejected_schema_objects": sorted(rejected_schema_objects), + } + async def _complete_sql_generation_context( self, *, @@ -6673,21 +6733,36 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": + rejected_schema_objects: set[str] = set() + semantic_retry_attempt = 0 + while ( + failed_dry_run_result + and failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION" + and semantic_retry_attempt < 3 + and not api_results + ): + semantic_retry_attempt += 1 semantic_retry_error = failed_dry_run_result.get("error", "") + retry_context = self._semantic_retry_context( + schema_intent_analysis, + semantic_retry_error, + semantic_retry_attempt, + rejected_schema_objects, + ) semantic_retry_query = ( f"{sql_user_query}\n\n" "The previous generated SQL failed semantic validation: " f"{semantic_retry_error}\n" - "Re-run semantic schema retrieval and choose the next " - "highest-confidence concept-to-schema mapping that " - "directly supports the user's requested entities, " - "metrics, dimensions, filters, time constraints, " - "aggregations, relationships, and ranking." + "Perform a fresh semantic retrieval. Select the next " + "highest-confidence complete concept-to-schema mapping. " + "Do not reuse rejected schema objects from RETRY CONTEXT." ) logger.info( - "Retrying semantic schema retrieval after intent validation failure for query_id %s", + "semantic_retry_attempt=%s query_id=%s validation_failure=%s rejected_schema_objects=%s", + semantic_retry_attempt, query_id, + semantic_retry_error, + sorted(rejected_schema_objects), ) try: retry_retrieval_result = await self._run_with_timeout( @@ -6696,7 +6771,8 @@ async def ask( query=semantic_retry_query, histories=histories, project_id=ask_request.project_id, - enable_column_pruning=enable_column_pruning, + enable_column_pruning=True, + semantic_retry_context=retry_context, ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -6706,6 +6782,25 @@ async def ask( retry_schema_intent_analysis = retry_construct_result.get( "semantic_analysis", {} ) + retry_selected_objects = set( + self._semantic_schema_objects( + retry_schema_intent_analysis + ) + ) + if retry_selected_objects and retry_selected_objects.issubset( + rejected_schema_objects + ): + error_message = ( + "Semantic schema retrieval retry selected only previously rejected schema objects." + ) + logger.info( + "semantic_retry_rejected_repeated_contract query_id=%s attempt=%s schema_objects=%s", + query_id, + semantic_retry_attempt, + sorted(retry_selected_objects), + ) + break + retry_documents, retry_table_names, retry_table_ddls = ( self._extract_retrieval_metadata( retry_retrieval_result @@ -6714,110 +6809,141 @@ async def ask( retry_support_error = get_schema_intent_analysis_error( retry_schema_intent_analysis ) - if retry_documents and not retry_support_error: - _retrieval_result = retry_construct_result - schema_intent_analysis = retry_schema_intent_analysis - documents = retry_documents - table_names = retry_table_names - table_ddls = retry_table_ddls - has_calculated_field = _retrieval_result.get( - "has_calculated_field", False + if not retry_documents or retry_support_error: + error_message = retry_support_error or ( + "Semantic schema retrieval retry did not find a supported mapping for the request." ) - has_metric = _retrieval_result.get("has_metric", False) - has_json_field = _retrieval_result.get( - "has_json_field", False + logger.info( + "semantic_retry_candidate_rejected query_id=%s attempt=%s reason=%s selected_objects=%s", + query_id, + semantic_retry_attempt, + error_message, + sorted(retry_selected_objects), ) + rejected_schema_objects.update(retry_selected_objects) + break - if histories: - text_to_sql_generation_results = ( - await self._run_with_timeout( - "Follow-up SQL generation after semantic retry", - self._pipelines[ - "followup_sql_generation" - ].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), - ) - ) - else: - text_to_sql_generation_results = ( - await self._run_with_timeout( - "SQL generation after semantic retry", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), - ) - ) + _retrieval_result = retry_construct_result + schema_intent_analysis = retry_schema_intent_analysis + rejected_schema_objects.update(retry_selected_objects) + documents = retry_documents + table_names = retry_table_names + table_ddls = retry_table_ddls + has_calculated_field = _retrieval_result.get( + "has_calculated_field", False + ) + has_metric = _retrieval_result.get("has_metric", False) + has_json_field = _retrieval_result.get( + "has_json_field", False + ) + logger.info( + "semantic_retry_candidate_selected query_id=%s attempt=%s table_names=%s schema_objects=%s", + query_id, + semantic_retry_attempt, + table_names, + sorted(retry_selected_objects), + ) - if sql_valid_result := text_to_sql_generation_results[ - "post_process" - ]["valid_generation_result"]: - if ask_result := self._build_validated_ask_result_from_sql( - sql_valid_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - else: - invalid_sql = sql_valid_result.get("sql") - error_message = ( - "SQL generation after semantic retrieval retry did not produce SQL that matches the active datasource schema and question intent." - ) - else: - failed_dry_run_result = ( - text_to_sql_generation_results["post_process"][ - "invalid_generation_result" - ] + if sql_generation_histories: + text_to_sql_generation_results = ( + await self._run_with_timeout( + "Follow-up SQL generation after semantic retry", + self._pipelines[ + "followup_sql_generation" + ].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=sql_generation_histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), ) - invalid_sql = failed_dry_run_result.get( - "sql", invalid_sql + ) + else: + text_to_sql_generation_results = ( + await self._run_with_timeout( + "SQL generation after semantic retry", + self._pipelines["sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), ) - error_message = failed_dry_run_result.get( - "error", error_message + ) + + if sql_valid_result := text_to_sql_generation_results[ + "post_process" + ]["valid_generation_result"]: + if ask_result := self._build_validated_ask_result_from_sql( + sql_valid_result.get("sql"), + table_ddls, + sql_user_query, + ): + api_results = [ask_result] + logger.info( + "semantic_retry_candidate_accepted query_id=%s attempt=%s", + query_id, + semantic_retry_attempt, ) - elif retry_support_error: - error_message = retry_support_error - else: + break + invalid_sql = sql_valid_result.get("sql") error_message = ( - "Semantic schema retrieval retry did not find a supported mapping for the request." + "SQL generation after semantic retrieval retry did not produce SQL that matches the active datasource schema and question intent." ) + break + + failed_dry_run_result = ( + text_to_sql_generation_results["post_process"][ + "invalid_generation_result" + ] + ) + invalid_sql = failed_dry_run_result.get( + "sql", invalid_sql + ) + error_message = failed_dry_run_result.get( + "error", error_message + ) + logger.info( + "semantic_retry_candidate_failed_validation query_id=%s attempt=%s error=%s", + query_id, + semantic_retry_attempt, + error_message, + ) except Exception as retry_error: logger.warning( - "Semantic schema retrieval retry failed for query_id %s: %s", + "Semantic schema retrieval retry failed for query_id %s attempt %s: %s", query_id, + semantic_retry_attempt, retry_error, ) + break while current_sql_correction_retries < max_sql_correction_retries: if api_results: break + if not failed_dry_run_result: + break if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index cbd6f7d169..d26ce51269 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -10,7 +10,7 @@ from src.pipelines.generation.utils.sql import get_schema_intent_analysis_error from src.utils import trace_metadata from src.web.v1.services import BaseRequest -from src.web.v1.services.ask import AskError, AskResult +from src.web.v1.services.ask import AskError, AskResult, AskService logger = logging.getLogger("wren-ai-service") @@ -224,21 +224,36 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": + rejected_schema_objects: set[str] = set() + semantic_retry_attempt = 0 + while ( + failed_dry_run_result + and failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION" + and semantic_retry_attempt < 3 + and not api_results + ): + semantic_retry_attempt += 1 semantic_retry_error = failed_dry_run_result.get("error", "") + retry_context = AskService._semantic_retry_context( + schema_intent_analysis, + semantic_retry_error, + semantic_retry_attempt, + rejected_schema_objects, + ) semantic_retry_query = ( f"{ask_feedback_request.question}\n\n" "The previous regenerated SQL failed semantic validation: " f"{semantic_retry_error}\n" - "Re-run semantic schema retrieval and choose the next " - "highest-confidence concept-to-schema mapping that " - "directly supports the user's requested entities, " - "metrics, dimensions, filters, time constraints, " - "aggregations, relationships, and ranking." + "Perform a fresh semantic retrieval. Select the next " + "highest-confidence complete concept-to-schema mapping. " + "Do not reuse rejected schema objects from RETRY CONTEXT." ) logger.info( - "Retrying semantic schema retrieval after feedback intent validation failure for query_id %s", + "feedback_semantic_retry_attempt=%s query_id=%s validation_failure=%s rejected_schema_objects=%s", + semantic_retry_attempt, query_id, + semantic_retry_error, + sorted(rejected_schema_objects), ) retry_retrieval_result = await self._pipelines[ "db_schema_retrieval" @@ -246,7 +261,8 @@ async def ask_feedback( query=semantic_retry_query, histories=[], project_id=ask_feedback_request.project_id, - enable_column_pruning=enable_column_pruning, + enable_column_pruning=True, + semantic_retry_context=retry_context, ) retry_construct_result = retry_retrieval_result.get( "construct_retrieval_results", {} @@ -254,6 +270,25 @@ async def ask_feedback( retry_schema_intent_analysis = retry_construct_result.get( "semantic_analysis", {} ) + retry_selected_objects = set( + AskService._semantic_schema_objects( + retry_schema_intent_analysis + ) + ) + if retry_selected_objects and retry_selected_objects.issubset( + rejected_schema_objects + ): + error_message = ( + "Semantic schema retrieval retry selected only previously rejected schema objects." + ) + logger.info( + "feedback_semantic_retry_rejected_repeated_contract query_id=%s attempt=%s schema_objects=%s", + query_id, + semantic_retry_attempt, + sorted(retry_selected_objects), + ) + break + retry_documents = retry_construct_result.get( "retrieval_results", [] ) @@ -266,59 +301,91 @@ async def ask_feedback( retry_support_error = get_schema_intent_analysis_error( retry_schema_intent_analysis ) - if retry_table_ddls and not retry_support_error: - schema_intent_analysis = retry_schema_intent_analysis - documents = retry_documents - table_ddls = retry_table_ddls - has_calculated_field = retry_construct_result.get( - "has_calculated_field", False - ) - has_metric = retry_construct_result.get( - "has_metric", False + if not retry_table_ddls or retry_support_error: + error_message = retry_support_error or ( + "Semantic schema retrieval retry did not find a supported mapping for the request." ) - has_json_field = retry_construct_result.get( - "has_json_field", False + logger.info( + "feedback_semantic_retry_candidate_rejected query_id=%s attempt=%s reason=%s selected_objects=%s", + query_id, + semantic_retry_attempt, + error_message, + sorted(retry_selected_objects), ) + rejected_schema_objects.update(retry_selected_objects) + break + + schema_intent_analysis = retry_schema_intent_analysis + rejected_schema_objects.update(retry_selected_objects) + documents = retry_documents + table_ddls = retry_table_ddls + has_calculated_field = retry_construct_result.get( + "has_calculated_field", False + ) + has_metric = retry_construct_result.get("has_metric", False) + has_json_field = retry_construct_result.get( + "has_json_field", False + ) + logger.info( + "feedback_semantic_retry_candidate_selected query_id=%s attempt=%s schema_objects=%s", + query_id, + semantic_retry_attempt, + sorted(retry_selected_objects), + ) - text_to_sql_generation_results = await self._pipelines[ - "sql_regeneration" - ].run( - contexts=table_ddls, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, - sql=ask_feedback_request.sql, - query=ask_feedback_request.question, - project_id=ask_feedback_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, + text_to_sql_generation_results = await self._pipelines[ + "sql_regeneration" + ].run( + contexts=table_ddls, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + sql=ask_feedback_request.sql, + query=ask_feedback_request.question, + project_id=ask_feedback_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ) + + if sql_valid_result := text_to_sql_generation_results[ + "post_process" + ]["valid_generation_result"]: + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } + ) + ] + logger.info( + "feedback_semantic_retry_candidate_accepted query_id=%s attempt=%s", + query_id, + semantic_retry_attempt, ) + break - if sql_valid_result := text_to_sql_generation_results[ - "post_process" - ]["valid_generation_result"]: - api_results = [ - AskResult( - **{ - "sql": sql_valid_result.get("sql"), - "type": "llm", - } - ) - ] - else: - failed_dry_run_result = text_to_sql_generation_results[ - "post_process" - ]["invalid_generation_result"] - elif retry_support_error: - error_message = retry_support_error + failed_dry_run_result = text_to_sql_generation_results[ + "post_process" + ]["invalid_generation_result"] + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) + logger.info( + "feedback_semantic_retry_candidate_failed_validation query_id=%s attempt=%s error=%s", + query_id, + semantic_retry_attempt, + error_message, + ) if api_results: pass - elif failed_dry_run_result["type"] not in { + elif failed_dry_run_result and failed_dry_run_result["type"] not in { "TIME_OUT", "SCHEMA_INTENT_VALIDATION", }: diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index afb2d06c8b..ed67f4e9aa 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -6,6 +6,7 @@ construct_retrieval_results, dbschema_retrieval, expand_business_terms_for_retrieval, + prompt, ) @@ -102,6 +103,17 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): "entities": ["invoice"], "metrics": ["invoice amount"], "dimensions": ["customer"], + "candidate_schema_scores": [ + { + "candidate_id": "candidate-1", + "schema_objects": ["invoices.customer_id", "invoices.invoice_amount"], + "covered_concepts": ["invoice amount", "customer"], + "missing_concepts": [], + "confidence": 0.95, + "is_complete": true, + "selection_reason": "Complete invoice amount by customer mapping." + } + ], "concept_mappings": [ { "request_concept": "invoice amount", @@ -176,7 +188,53 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): assert result["semantic_analysis"]["concept_mappings"][0]["schema_objects"] == [ "invoices.invoice_amount" ] + assert result["semantic_analysis"]["candidate_schema_scores"][0]["is_complete"] assert result["semantic_analysis"]["interpretations"][0]["is_selected"] is True assert result["retrieval_results"][0]["table_name"] == "invoices" assert "invoice_amount" in result["retrieval_results"][0]["table_ddl"] assert "internal_note" not in result["retrieval_results"][0]["table_ddl"] + + +def test_prompt_includes_semantic_retry_context(): + class PromptBuilder: + def run(self, **kwargs): + retry_context = kwargs["semantic_retry_context"] + return { + "prompt": ( + f"retry={retry_context['retry_attempt']} " + f"error={retry_context['validation_error']} " + f"rejected={','.join(retry_context['rejected_schema_objects'])}" + ) + } + + result = prompt( + query="Top customers by invoice amount", + construct_db_schemas=[ + { + "type": "TABLE", + "name": "invoices", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "invoice_amount", + "data_type": "double", + "comment": "", + "is_primary_key": False, + } + ], + } + ], + prompt_builder=PromptBuilder(), + check_using_db_schemas_without_pruning={"db_schemas": []}, + histories=[], + semantic_retry_context={ + "validation_error": "Generic count did not answer invoice amount", + "retry_attempt": 2, + "rejected_schema_objects": ["dbo_ytblES002_1.Name_of_Reported_Received"], + }, + ) + + assert "retry=2" in result["prompt"] + assert "Generic count" in result["prompt"] + assert "dbo_ytblES002_1.Name_of_Reported_Received" in result["prompt"] From 11a631f13e7558f4bac839bea7daaad1e5201b7f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 16:09:00 +0530 Subject: [PATCH 0474/1087] Keep semantic SQL requests out of legacy fallbacks --- wren-ai-service/src/web/v1/services/ask.py | 163 ++++++++++++++++++--- 1 file changed, 143 insertions(+), 20 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 16023abdea..00c45f2fa7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4496,6 +4496,31 @@ def _semantic_schema_objects(semantic_analysis: dict[str, Any] | None) -> list[s return sorted(set(schema_objects)) + @staticmethod + def _has_semantic_contract(semantic_analysis: dict[str, Any] | None) -> bool: + if not isinstance(semantic_analysis, dict) or not semantic_analysis: + return False + semantic_keys = { + "analytical_intent", + "entities", + "identifiers", + "metrics", + "dimensions", + "filters", + "aggregations", + "relationships", + "time_constraints", + "ranking", + "supported_schema_objects", + "candidate_schema_scores", + "concept_mappings", + "interpretations", + "missing_requirements", + "ambiguous_requirements", + "support_reasoning", + } + return any(semantic_analysis.get(key) for key in semantic_keys) + @staticmethod def _semantic_retry_context( semantic_analysis: dict[str, Any] | None, @@ -5382,6 +5407,8 @@ async def ask( sql_knowledge = None understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) schema_intent_analysis: dict[str, Any] = {} + semantic_pipeline_active = False + semantic_retrieval_activated = False try: sql_user_query = user_query @@ -6026,16 +6053,25 @@ async def ask( ) try: + semantic_retrieval_activated = ( + enable_column_pruning + or self._is_data_analysis_query(user_query) + ) + logger.info( + "semantic_retrieval_activation query_id=%s active=%s reason=%s", + query_id, + semantic_retrieval_activated, + "analytics_or_column_pruning" + if semantic_retrieval_activated + else "standard_retrieval", + ) retrieval_result = await self._run_with_timeout( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( query=sql_user_query, histories=[], project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) - ), + enable_column_pruning=semantic_retrieval_activated, ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -6150,8 +6186,23 @@ async def ask( logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) + semantic_pipeline_active = self._has_semantic_contract( + schema_intent_analysis + ) or semantic_retrieval_activated + logger.info( + "sql_generation_pipeline_decision query_id=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s", + query_id, + semantic_pipeline_active, + self._has_semantic_contract(schema_intent_analysis), + self._semantic_schema_objects(schema_intent_analysis), + ) + if semantic_pipeline_active: + logger.info( + "legacy_sql_fallbacks_disabled query_id=%s reason=semantic_pipeline_active", + query_id, + ) - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6171,7 +6222,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6195,7 +6246,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6219,6 +6270,7 @@ async def ask( if ( not api_results + and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and ( schema_grounded_sql := self._build_schema_grounded_analytics_sql( @@ -6243,7 +6295,7 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if not api_results and any( + if not semantic_pipeline_active and not api_results and any( term in user_query.lower() for term in ( "pcb", @@ -6276,7 +6328,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6300,6 +6352,7 @@ async def ask( should_retry_full_schema = ( not api_results + and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and "db_schema_retrieval" in self._pipelines ) @@ -6435,7 +6488,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if not documents: + if not documents and not semantic_pipeline_active: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): @@ -6502,6 +6555,34 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + if semantic_pipeline_active and not documents: + semantic_failure_message = error_message or ( + "Semantic schema retrieval did not find a complete schema mapping for the request." + ) + logger.info( + "semantic_pipeline_no_supported_documents query_id=%s message=%s", + query_id, + semantic_failure_message, + ) + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_failed_text_to_sql_response( + trace_id, + semantic_failure_message, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + is_followup=True if histories else False, + code="NO_RELEVANT_SQL", + ) + ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = semantic_failure_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if documents and not api_results: documents, table_names, table_ddls = self._prune_sql_generation_context( sql_user_query, @@ -6529,13 +6610,19 @@ async def ask( sql_user_query ) and not self._needs_conversation_context(sql_user_query): sql_generation_histories = [] - allow_sql_generation_reasoning = False allow_sql_knowledge_retrieval = False max_sql_correction_retries = min(max_sql_correction_retries, 1) - logger.info( - "Using fast standalone SQL generation path for query_id %s", - query_id, - ) + if semantic_pipeline_active: + logger.info( + "fast_standalone_sql_generation_disabled query_id=%s reason=semantic_pipeline_active", + query_id, + ) + else: + allow_sql_generation_reasoning = False + logger.info( + "Using fast standalone SQL generation path for query_id %s", + query_id, + ) if ( not self._is_stopped(query_id, self._ask_results) @@ -6703,11 +6790,18 @@ async def ask( ), ) except TimeoutError as generation_timeout: - logger.warning( - "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", - query_id, - generation_timeout, - ) + if semantic_pipeline_active: + logger.warning( + "Semantic SQL generation timed out for query_id %s; legacy fallbacks remain disabled: %s", + query_id, + generation_timeout, + ) + else: + logger.warning( + "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + query_id, + generation_timeout, + ) text_to_sql_generation_results = { "post_process": { "valid_generation_result": None, @@ -7052,6 +7146,35 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: + if semantic_pipeline_active: + semantic_failure_message = error_message or ( + "No valid semantic schema contract could satisfy the request after semantic retrieval retries." + ) + logger.info( + "semantic_pipeline_exhausted query_id=%s message=%s rejected_sql=%s", + query_id, + semantic_failure_message, + invalid_sql, + ) + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_failed_text_to_sql_response( + trace_id, + semantic_failure_message, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + is_followup=True if histories else False, + code="NO_RELEVANT_SQL", + ) + ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = semantic_failure_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): From 97d09d73076b4e25169cb3bcd2241a8d80db1cb0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 16:55:21 +0530 Subject: [PATCH 0475/1087] Add generic semantic schema candidate ranking --- .../src/pipelines/generation/utils/sql.py | 92 +++++ .../retrieval/db_schema_retrieval.py | 389 ++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 43 ++ .../retrieval/test_db_schema_retrieval.py | 114 ++++- 4 files changed, 637 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 97ae55f7c3..2062de554d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3647,6 +3647,7 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: "time_constraints", "ranking", "supported_schema_objects", + "candidate_schema_scores", "concept_mappings", "interpretations", "missing_requirements", @@ -3692,6 +3693,87 @@ def _schema_interpretation_clarification_error( return None +def _semantic_candidate_scores( + semantic_analysis: dict[str, Any] | None, +) -> list[dict[str, Any]]: + return _semantic_analysis_dict_items(semantic_analysis, "candidate_schema_scores") + + +def _semantic_candidate_support_error( + semantic_analysis: dict[str, Any], +) -> str | None: + candidate_scores = _semantic_candidate_scores(semantic_analysis) + if not candidate_scores: + return None + + complete_candidates = [ + candidate + for candidate in candidate_scores + if candidate.get("is_complete") is True + ] + if complete_candidates: + return None + + incomplete_with_missing = [ + candidate + for candidate in candidate_scores + if candidate.get("missing_concepts") + ] + if not incomplete_with_missing: + return None + + missing_concepts = [] + for candidate in incomplete_with_missing[:3]: + candidate_id = str(candidate.get("candidate_id") or "candidate").strip() + missing = candidate.get("missing_concepts") + if isinstance(missing, list): + missing_text = ", ".join( + str(item).strip() + for item in missing + if item is not None and str(item).strip() + ) + else: + missing_text = str(missing or "").strip() + if missing_text: + missing_concepts.append(f"{candidate_id}: {missing_text}") + + if not missing_concepts: + return None + + return ( + "Semantic schema retrieval did not find a complete schema mapping for " + "the request. Missing concepts: " + f"{'; '.join(missing_concepts)}. I cannot generate unrelated SQL." + ) + + +def _required_concept_mapping_support_error( + semantic_analysis: dict[str, Any], +) -> str | None: + unsupported_required_concepts = [] + for mapping in _semantic_concept_mappings(semantic_analysis): + if mapping.get("required_in_sql") is False: + continue + if _mapping_schema_objects(mapping): + continue + + request_concept = _mapping_request_concept(mapping) + concept_type = _mapping_concept_type(mapping) + if request_concept: + unsupported_required_concepts.append( + f"{request_concept} ({concept_type or 'concept'})" + ) + + if not unsupported_required_concepts: + return None + + return ( + "The semantic contract did not map required request concepts to active " + "schema objects: " + f"{', '.join(unsupported_required_concepts)}. I cannot generate unrelated SQL." + ) + + def get_schema_intent_analysis_error( semantic_analysis: dict[str, Any] | None, ) -> str | None: @@ -3722,6 +3804,16 @@ def get_schema_intent_analysis_error( ): return interpretation_error + if candidate_support_error := _semantic_candidate_support_error( + semantic_analysis + ): + return candidate_support_error + + if required_mapping_error := _required_concept_mapping_support_error( + semantic_analysis + ): + return required_mapping_error + if semantic_analysis.get("is_fully_supported") is False: support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() if support_reasoning: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index f45e24a5b0..bfbbae7b50 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,5 +1,6 @@ import ast import logging +import re import sys from typing import TYPE_CHECKING, Any, Optional @@ -157,6 +158,29 @@ {{ db_schema }} {% endfor %} +{% if semantic_candidate_context %} +### PRE-RANKED SEMANTIC SCHEMA CANDIDATES ### +These candidates were scored generically from the active datasource schema metadata and the user's full request. +Use them as retrieval evidence, but still verify complete concept coverage before selecting a contract. +Prefer candidates that cover all requested entities, identifiers, metrics, dimensions, filters, time constraints, aggregations, and ranking requirements. +Do not select a high lexical match when it misses a required business concept. + +{% for candidate in semantic_candidate_context %} +- candidate_id: {{ candidate.candidate_id }} + table_name: {{ candidate.table_name }} + confidence: {{ candidate.confidence }} + coverage_score: {{ candidate.coverage_score }} + matched_query_terms: {{ candidate.matched_query_terms }} + missing_query_terms: {{ candidate.missing_query_terms }} + rejected_by_retry: {{ candidate.rejected_by_retry }} + selection_reason: {{ candidate.selection_reason }} + matched_columns: +{% for column in candidate.matched_columns %} + - {{ column.column_name }} (score={{ column.score }}, data_type={{ column.data_type }}, matched_terms={{ column.matched_terms }}) +{% endfor %} +{% endfor %} +{% endif %} + ### INPUT ### {{ question }} @@ -259,6 +283,361 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: return deduped +_SEMANTIC_TOKEN_STOPWORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "for", + "from", + "give", + "have", + "how", + "in", + "is", + "me", + "of", + "on", + "or", + "show", + "that", + "the", + "to", + "with", + "dbo", + "tbl", + "table", + "view", + "dim", + "fact", + "stage", + "stg", +} + +_NUMERIC_SCHEMA_TERMS = { + "amount", + "avg", + "average", + "balance", + "cost", + "count", + "gross", + "margin", + "measure", + "metric", + "net", + "price", + "profit", + "quantity", + "rate", + "revenue", + "sales", + "sum", + "total", + "value", +} + +_TEMPORAL_SCHEMA_TERMS = { + "date", + "day", + "month", + "monthly", + "quarter", + "time", + "week", + "year", +} + +_RANKING_SCHEMA_TERMS = { + "bottom", + "highest", + "least", + "lowest", + "most", + "rank", + "ranking", + "top", +} + + +def _semantic_tokens(value: Any) -> set[str]: + text = str(value or "") + if not text: + return set() + + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) + text = re.sub(r"[^A-Za-z0-9]+", " ", text) + tokens = { + token.lower() + for token in text.split() + if len(token) > 1 and token.lower() not in _SEMANTIC_TOKEN_STOPWORDS + } + for token in list(tokens): + if token.endswith("ies") and len(token) > 4: + tokens.add(f"{token[:-3]}y") + elif token.endswith("s") and len(token) > 3: + tokens.add(token[:-1]) + return tokens + + +def _schema_comment_text(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, list): + return " ".join(_schema_comment_text(item) for item in value) + if isinstance(value, dict): + return " ".join(_schema_comment_text(item) for item in value.values()) + return str(value) + + +def _column_tokens(column: dict[str, Any]) -> set[str]: + tokens = set() + for key in ("name", "display_name", "alias", "comment", "description", "data_type"): + tokens.update(_semantic_tokens(column.get(key))) + tokens.update(_semantic_tokens(_schema_comment_text(column.get("properties")))) + return tokens + + +def _table_tokens(table_schema: dict[str, Any]) -> set[str]: + tokens = set() + for key in ("name", "display_name", "alias", "comment", "description"): + tokens.update(_semantic_tokens(table_schema.get(key))) + for column in table_schema.get("columns", []) or []: + if isinstance(column, dict): + tokens.update(_column_tokens(column)) + return tokens + + +def _query_semantic_terms(query: str) -> dict[str, set[str]]: + tokens = _semantic_tokens(query) + return { + "all": tokens, + "metric": tokens & _NUMERIC_SCHEMA_TERMS, + "time": tokens & _TEMPORAL_SCHEMA_TERMS, + "ranking": tokens & _RANKING_SCHEMA_TERMS, + } + + +def _is_numeric_column(column: dict[str, Any]) -> bool: + data_type = str(column.get("data_type") or "").lower() + return bool( + re.search( + r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|number|double|float|real|money)\b", + data_type, + ) + ) + + +def _is_identifier_column(column: dict[str, Any]) -> bool: + tokens = _column_tokens(column) + return bool(tokens & {"code", "id", "identifier", "key", "no", "number"}) + + +def _is_temporal_column(column: dict[str, Any]) -> bool: + data_type = str(column.get("data_type") or "").lower() + return bool(re.search(r"\b(?:date|time|timestamp|datetime)\b", data_type)) + + +def _normalized_schema_object(value: Any) -> str: + return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) + + +def _rejected_schema_objects(semantic_retry_context: dict[str, Any] | None) -> set[str]: + if not isinstance(semantic_retry_context, dict): + return set() + + rejected = semantic_retry_context.get("rejected_schema_objects") + if not isinstance(rejected, list): + return set() + + return { + _normalized_schema_object(item) + for item in rejected + if item is not None and str(item).strip() + } + + +def _schema_object_was_rejected( + table_name: str, + column_name: str | None, + rejected_schema_objects: set[str], +) -> bool: + if not rejected_schema_objects: + return False + + table_key = _normalized_schema_object(table_name) + object_key = _normalized_schema_object( + f"{table_name}.{column_name}" if column_name else table_name + ) + return any( + rejected_key + and ( + rejected_key == table_key + or rejected_key == object_key + or rejected_key.endswith(object_key) + or object_key.endswith(rejected_key) + ) + for rejected_key in rejected_schema_objects + ) + + +def rank_semantic_schema_candidates( + query: str, + construct_db_schemas: list[dict], + semantic_retry_context: dict[str, Any] | None = None, + max_candidates: int = 15, + max_columns_per_candidate: int = 8, +) -> list[dict[str, Any]]: + query_terms = _query_semantic_terms(query) + all_query_terms = query_terms["all"] + if not all_query_terms: + return [] + + rejected_schema_objects = _rejected_schema_objects(semantic_retry_context) + candidates: list[dict[str, Any]] = [] + + for table_schema in construct_db_schemas: + if table_schema.get("type") != "TABLE": + continue + + table_name = str(table_schema.get("name") or "").strip() + if not table_name: + continue + + table_term_matches = _table_tokens(table_schema) & all_query_terms + matched_columns = [] + table_rejected = _schema_object_was_rejected( + table_name, None, rejected_schema_objects + ) + + for column in table_schema.get("columns", []) or []: + if not isinstance(column, dict): + continue + + column_name = str(column.get("name") or "").strip() + if not column_name: + continue + + tokens = _column_tokens(column) + matched_terms = sorted(tokens & all_query_terms) + score = float(len(matched_terms) * 3) + + if query_terms["metric"] and _is_numeric_column(column): + score += 0.3 if _is_identifier_column(column) else 1.5 + if tokens & query_terms["metric"]: + score += 2.0 + if query_terms["time"] and _is_temporal_column(column): + score += 1.5 + if tokens & query_terms["time"]: + score += 2.0 + if query_terms["ranking"] and matched_terms: + score += 0.5 + + rejected = _schema_object_was_rejected( + table_name, column_name, rejected_schema_objects + ) + if rejected: + score -= 5.0 + + if score > 0 or matched_terms: + matched_columns.append( + { + "column_name": column_name, + "score": round(max(score, 0.0), 3), + "matched_terms": matched_terms, + "data_type": str(column.get("data_type") or ""), + "rejected_by_retry": rejected, + } + ) + + matched_columns.sort( + key=lambda item: (item["score"], len(item["matched_terms"])), + reverse=True, + ) + matched_columns = matched_columns[:max_columns_per_candidate] + + covered_terms = set(table_term_matches) + for column in matched_columns: + covered_terms.update(column["matched_terms"]) + + if not covered_terms and not table_rejected: + continue + + coverage_score = len(covered_terms) / max(len(all_query_terms), 1) + raw_score = ( + len(table_term_matches) * 2.0 + + sum(column["score"] for column in matched_columns) + + coverage_score * 4.0 + ) + if table_rejected: + raw_score -= 6.0 + column_lookup = { + str(column.get("name") or ""): column + for column in table_schema.get("columns", []) or [] + if isinstance(column, dict) + } + has_metric_support = any( + query_terms["metric"] & set(column["matched_terms"]) + or ( + _is_numeric_column(column_lookup.get(column["column_name"], {})) + and not _is_identifier_column( + column_lookup.get(column["column_name"], {}) + ) + ) + for column in matched_columns + ) + if query_terms["metric"] and not has_metric_support: + raw_score -= 2.0 + + confidence = min(max(raw_score / 20.0, 0.0), 0.99) + selection_reason = ( + "Covers " + f"{len(covered_terms)} of {len(all_query_terms)} significant request terms" + ) + if table_rejected: + selection_reason += "; penalized because it was rejected by semantic validation" + if query_terms["metric"] and not any( + set(column["matched_terms"]) & query_terms["metric"] + for column in matched_columns + ): + selection_reason += "; metric term coverage is weak" + + candidates.append( + { + "candidate_id": f"candidate-{len(candidates) + 1}", + "table_name": table_name, + "confidence": round(confidence, 3), + "coverage_score": round(coverage_score, 3), + "matched_query_terms": sorted(covered_terms), + "missing_query_terms": sorted(all_query_terms - covered_terms), + "matched_columns": matched_columns, + "rejected_by_retry": table_rejected + or any(column["rejected_by_retry"] for column in matched_columns), + "selection_reason": selection_reason, + } + ) + + candidates.sort( + key=lambda item: ( + item["rejected_by_retry"] is False, + item["confidence"], + item["coverage_score"], + ), + reverse=True, + ) + + for index, candidate in enumerate(candidates[:max_candidates], start=1): + candidate["candidate_id"] = f"candidate-{index}" + + return candidates[:max_candidates] + + @observe(capture_input=False, capture_output=False) async def embedding( query: str, @@ -459,10 +838,20 @@ def prompt( ) query = "\n".join(previous_query_summaries) + "\n" + query + semantic_candidate_context = rank_semantic_schema_candidates( + query=query, + construct_db_schemas=construct_db_schemas, + semantic_retry_context=semantic_retry_context, + ) + logger.info( + "semantic_retrieval_pre_ranked_candidates=%s", + semantic_candidate_context, + ) _prompt = prompt_builder.run( question=query, db_schemas=db_schemas, + semantic_candidate_context=semantic_candidate_context, semantic_retry_context=semantic_retry_context or {}, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 4d8797ee50..d22e45f36c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1675,6 +1675,49 @@ def test_get_schema_intent_analysis_error_reports_unsupported_analysis(): assert "No relationship connects invoices to products" in error +def test_get_schema_intent_analysis_error_rejects_incomplete_semantic_candidates(): + error = get_schema_intent_analysis_error( + { + "candidate_schema_scores": [ + { + "candidate_id": "candidate-1", + "schema_objects": ["refunds.refund_amount"], + "covered_concepts": ["amount"], + "missing_concepts": ["customer", "invoice amount"], + "confidence": 0.62, + "is_complete": False, + "selection_reason": "Only amount matched.", + } + ], + "is_fully_supported": True, + } + ) + + assert error is not None + assert "complete schema mapping" in error + assert "invoice amount" in error + + +def test_get_schema_intent_analysis_error_rejects_required_unmapped_concepts(): + error = get_schema_intent_analysis_error( + { + "concept_mappings": [ + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": [], + "required_in_sql": True, + } + ], + "is_fully_supported": True, + } + ) + + assert error is not None + assert "invoice amount" in error + assert "did not map required request concepts" in error + + def test_validate_sql_intent_alignment_uses_semantic_analysis_for_metric_count_mismatch(): error = validate_sql_intent_alignment( "Show invoice amount by customer", diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index ed67f4e9aa..a70f26aecb 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -7,6 +7,7 @@ dbschema_retrieval, expand_business_terms_for_retrieval, prompt, + rank_semantic_schema_candidates, ) @@ -32,6 +33,114 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query +def test_rank_semantic_schema_candidates_prefers_complete_business_concept_coverage(): + candidates = rank_semantic_schema_candidates( + query="Show top 10 customers by invoice amount", + construct_db_schemas=[ + { + "type": "TABLE", + "name": "dbo_ytblES002_1", + "comment": "", + "columns": [ + { + "name": "Name_of_Reported_Received", + "data_type": "varchar", + "comment": "reported received name", + }, + { + "name": "Refund_Amount", + "data_type": "decimal", + "comment": "refund amount", + }, + ], + }, + { + "type": "TABLE", + "name": "dbo_tblFactSales", + "comment": "invoice sales facts by customer", + "columns": [ + { + "name": "invoice", + "data_type": "varchar", + "comment": "invoice identifier", + }, + { + "name": "customer_id", + "data_type": "varchar", + "comment": "customer identifier", + }, + { + "name": "Amount_Received", + "data_type": "decimal", + "comment": "invoice amount received", + }, + ], + }, + ], + ) + + assert candidates[0]["table_name"] == "dbo_tblFactSales" + assert "customer" in candidates[0]["matched_query_terms"] + assert "amount" in candidates[0]["matched_query_terms"] + assert candidates[0]["confidence"] > candidates[1]["confidence"] + + +def test_rank_semantic_schema_candidates_penalizes_retry_rejected_schema_objects(): + candidates = rank_semantic_schema_candidates( + query="Show top 10 customers by invoice amount", + construct_db_schemas=[ + { + "type": "TABLE", + "name": "dbo_tblFactSales", + "comment": "invoice sales facts by customer", + "columns": [ + { + "name": "customer_id", + "data_type": "varchar", + "comment": "customer identifier", + }, + { + "name": "Amount_Received", + "data_type": "decimal", + "comment": "invoice amount received", + }, + ], + }, + { + "type": "TABLE", + "name": "dbo_qSales", + "comment": "invoice analytics by account", + "columns": [ + { + "name": "Account", + "data_type": "varchar", + "comment": "customer account", + }, + { + "name": "InvoiceAmount", + "data_type": "decimal", + "comment": "invoice amount", + }, + ], + }, + ], + semantic_retry_context={ + "rejected_schema_objects": [ + "dbo_tblFactSales.customer_id", + "dbo_tblFactSales.Amount_Received", + ] + }, + ) + + assert candidates[0]["table_name"] == "dbo_qSales" + rejected_candidate = next( + candidate + for candidate in candidates + if candidate["table_name"] == "dbo_tblFactSales" + ) + assert rejected_candidate["rejected_by_retry"] is True + + @pytest.mark.asyncio async def test_dbschema_retrieval_loads_complete_active_project_schema(): class Retriever: @@ -199,11 +308,13 @@ def test_prompt_includes_semantic_retry_context(): class PromptBuilder: def run(self, **kwargs): retry_context = kwargs["semantic_retry_context"] + candidate_context = kwargs["semantic_candidate_context"] return { "prompt": ( f"retry={retry_context['retry_attempt']} " f"error={retry_context['validation_error']} " - f"rejected={','.join(retry_context['rejected_schema_objects'])}" + f"rejected={','.join(retry_context['rejected_schema_objects'])} " + f"candidates={candidate_context[0]['table_name']}" ) } @@ -238,3 +349,4 @@ def run(self, **kwargs): assert "retry=2" in result["prompt"] assert "Generic count" in result["prompt"] assert "dbo_ytblES002_1.Name_of_Reported_Received" in result["prompt"] + assert "candidates=invoices" in result["prompt"] From 3257287045198a4b9d73e610ae842b2a56a07b10 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 19:24:43 +0530 Subject: [PATCH 0476/1087] Fix semantic entity mapping and ranking SQL validation --- .../src/pipelines/generation/utils/sql.py | 180 +++++++++++++++++- .../retrieval/db_schema_retrieval.py | 51 ++++- .../pipelines/generation/test_sql_utils.py | 120 ++++++++++++ .../retrieval/test_db_schema_retrieval.py | 42 ++++ 4 files changed, 374 insertions(+), 19 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2062de554d..8a7ce77c9d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1080,7 +1080,7 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: if re.match(r"\s*SELECT\s+DISTINCT\b", without_limit, flags=re.IGNORECASE): return re.sub( r"\bSELECT\s+DISTINCT\b", - f"SELECT DISTINCT TOP {limit}", + f"SELECT DISTINCT TOP ({limit})", without_limit, count=1, flags=re.IGNORECASE, @@ -1089,7 +1089,7 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: if re.match(r"\s*SELECT\b", without_limit, flags=re.IGNORECASE): return re.sub( r"\bSELECT\b", - f"SELECT TOP {limit}", + f"SELECT TOP ({limit})", without_limit, count=1, flags=re.IGNORECASE, @@ -1458,7 +1458,7 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: return re.sub( r"\bSELECT\s+(DISTINCT\s+)?", - lambda match: f"{match.group(0)}TOP {limit} ", + lambda match: f"{match.group(0)}TOP ({limit}) ", without_limit, count=1, flags=re.IGNORECASE, @@ -1484,8 +1484,21 @@ def _normalize_identifier_quote_syntax(sql: str) -> str: return normalized +def _normalize_mssql_top_clause(sql: str) -> str: + return re.sub( + r"\bSELECT\s+(DISTINCT\s+)?TOP\s+(\d+)\b", + lambda match: ( + f"SELECT {match.group(1) or ''}TOP ({match.group(2)})" + ), + sql, + count=1, + flags=re.IGNORECASE, + ) + + def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = _normalize_identifier_quote_syntax(sql) + normalized = _normalize_mssql_top_clause(normalized) normalized_data_source = normalize_data_source(data_source) if normalized_data_source == "MSSQL": @@ -2411,6 +2424,15 @@ def construct_semantic_schema_contract( "Do not substitute identifiers for metrics, entities for identifiers, or COUNT(*) " "for a requested business measure unless the semantic intent explicitly requests a record count." ) + lines.append( + "Ranking requirement: for top/bottom/ranked questions, include ORDER BY on the mapped metric " + "and a row limit such as LIMIT N, FETCH FIRST N ROWS ONLY, or MSSQL/Wren-safe TOP (N)." + ) + lines.append( + "Schema safety requirement: use only tables and columns present in this contract or the " + "retrieved DATABASE SCHEMA. Do not add unrequested date filters or common timestamp columns " + "such as created_at unless they are explicitly listed and the user asked for a time constraint." + ) return "\n".join(lines) @@ -3632,6 +3654,80 @@ def _semantic_analysis_dict_items( return [item for item in value if isinstance(item, dict)] +_SEMANTIC_CONCEPT_STOPWORDS = { + "a", + "an", + "and", + "as", + "by", + "for", + "from", + "in", + "of", + "on", + "or", + "the", + "to", + "with", + "dbo", + "tbl", + "table", + "view", + "dim", + "fact", +} + +_SEMANTIC_CONCEPT_SYNONYMS = { + "acct": {"account", "customer"}, + "account": {"acct", "customer", "client"}, + "accounts": {"acct", "account", "customer", "client"}, + "amt": {"amount", "value", "total"}, + "amount": {"amt", "value", "total"}, + "bill": {"invoice"}, + "billing": {"invoice"}, + "client": {"account", "customer"}, + "clients": {"account", "customer"}, + "cust": {"customer", "client", "account"}, + "customer": {"cust", "client", "account"}, + "customers": {"cust", "client", "account", "customer"}, + "desc": {"description", "name"}, + "description": {"desc", "name"}, + "inv": {"invoice"}, + "invoice": {"inv", "bill", "billing"}, + "invoices": {"inv", "invoice", "bill", "billing"}, + "name": {"description", "label"}, + "no": {"number", "identifier", "id"}, + "num": {"number", "identifier", "id"}, + "number": {"no", "num", "identifier", "id"}, + "qty": {"quantity"}, + "quantity": {"qty"}, + "total": {"amount", "value", "sum"}, + "value": {"amount", "total"}, +} + + +def _semantic_concept_tokens(value: Any) -> set[str]: + text = str(value or "") + if not text: + return set() + + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) + text = re.sub(r"[^A-Za-z0-9]+", " ", text) + tokens = { + token.lower() + for token in text.split() + if len(token) > 1 and token.lower() not in _SEMANTIC_CONCEPT_STOPWORDS + } + for token in list(tokens): + if token.endswith("ies") and len(token) > 4: + tokens.add(f"{token[:-3]}y") + elif token.endswith("s") and len(token) > 3: + tokens.add(token[:-1]) + for token in list(tokens): + tokens.update(_SEMANTIC_CONCEPT_SYNONYMS.get(token, set())) + return tokens + + def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: if not isinstance(semantic_analysis, dict) or not semantic_analysis: return False @@ -3750,15 +3846,42 @@ def _semantic_candidate_support_error( def _required_concept_mapping_support_error( semantic_analysis: dict[str, Any], ) -> str | None: + schema_bound_concept_types = { + "dimension", + "entity", + "filter", + "identifier", + "metric", + "relationship", + "time", + } + mapped_schema_objects = [ + schema_object + for mapping in _semantic_concept_mappings(semantic_analysis) + for schema_object in _mapping_schema_objects(mapping) + ] + mapped_schema_objects.extend( + _semantic_analysis_items(semantic_analysis, "supported_schema_objects") + ) + unsupported_required_concepts = [] for mapping in _semantic_concept_mappings(semantic_analysis): if mapping.get("required_in_sql") is False: continue + concept_type = _mapping_concept_type(mapping) + if concept_type not in schema_bound_concept_types: + continue if _mapping_schema_objects(mapping): continue request_concept = _mapping_request_concept(mapping) - concept_type = _mapping_concept_type(mapping) + concept_tokens = _semantic_concept_tokens(request_concept) + if concept_type == "entity" and concept_tokens: + if any( + concept_tokens & _semantic_concept_tokens(schema_object) + for schema_object in mapped_schema_objects + ): + continue if request_concept: unsupported_required_concepts.append( f"{request_concept} ({concept_type or 'concept'})" @@ -4014,11 +4137,28 @@ def _validate_sql_against_concept_mappings( ) has_limit = bool( re.search( - r"\b(?:LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", + r"\b(?:LIMIT|TOP\s*(?:\(\s*)?\d+|FETCH\s+FIRST)\b", sql or "", flags=re.IGNORECASE, ) ) + schema_bound_concept_types = { + "dimension", + "entity", + "filter", + "identifier", + "metric", + "relationship", + "time", + } + mapped_schema_objects = [ + schema_object + for mapping in mappings + for schema_object in _mapping_schema_objects(mapping) + ] + mapped_schema_objects.extend( + _semantic_analysis_items(semantic_analysis, "supported_schema_objects") + ) for mapping in mappings: if mapping.get("required_in_sql") is False: @@ -4028,6 +4168,15 @@ def _validate_sql_against_concept_mappings( request_concept = _mapping_request_concept(mapping) schema_objects = _mapping_schema_objects(mapping) if not schema_objects: + if concept_type not in schema_bound_concept_types: + continue + concept_tokens = _semantic_concept_tokens(request_concept) + if concept_type == "entity" and concept_tokens: + if any( + concept_tokens & _semantic_concept_tokens(schema_object) + for schema_object in mapped_schema_objects + ): + continue return ( "The semantic analysis did not map the required " f"{concept_type or 'concept'} '{request_concept}' to a schema " @@ -4165,11 +4314,22 @@ def _validate_sql_against_semantic_analysis( "dimensions or time grain identified in the semantic analysis." ) - if ranking and not re.search( - r"\b(?:ORDER\s+BY|LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", - sql or "", - flags=re.IGNORECASE, - ): + if ranking: + has_ranking_order = bool( + re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE) + ) + has_ranking_limit = bool( + re.search( + r"\b(?:LIMIT|TOP\s*(?:\(\s*)?\d+|FETCH\s+FIRST)\b", + sql or "", + flags=re.IGNORECASE, + ) + ) + else: + has_ranking_order = True + has_ranking_limit = True + + if ranking and (not has_ranking_order or not has_ranking_limit): return ( "Generated SQL does not include sorting or limiting logic required " "by the ranking intent." diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index bfbbae7b50..60dea2a78d 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -47,15 +47,18 @@ 8. For each selected table, provide a concise reason for why the table is semantically relevant. 9. For each selected column, provide a concise reason for why the column is necessary. 10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. -11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. -12. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. -13. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. -14. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. -15. If a "." is included in columns, put the name before the first dot into chosen columns. -16. The number of columns chosen must match the number of reasoning. -17. Final chosen columns must be only column names, don't prefix it with table names. -18. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -19. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +11. Broad business entities can be satisfied by the selected table plus the best descriptive or identifier column for that entity. For example, an entity such as customer may map to a customer name, customer number, account, client, or similar descriptive/identifier column when that is the active schema's representation. Do not leave an entity unmapped when a selected dimension or identifier column represents it. +12. When the user asks for a top/bottom/ranking query, map the ranked dimension and ranked metric separately. The SQL generator must be able to ORDER BY the metric and limit rows. +13. Do not add filters or time constraints that are not requested or implied by the user. Only map date/time concepts when the user asks for a time period, trend, date filter, or date dimension. +14. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. +15. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. +16. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. +17. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. +18. If a "." is included in columns, put the name before the first dot into chosen columns. +19. The number of columns chosen must match the number of reasoning. +20. Final chosen columns must be only column names, don't prefix it with table names. +21. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +22. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -363,6 +366,34 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: "top", } +_GENERIC_SEMANTIC_SYNONYMS = { + "acct": {"account", "customer"}, + "account": {"acct", "customer", "client"}, + "accounts": {"acct", "account", "customer", "client"}, + "amt": {"amount", "value", "total"}, + "amount": {"amt", "value", "total"}, + "bill": {"invoice"}, + "billing": {"invoice"}, + "client": {"account", "customer"}, + "clients": {"account", "customer"}, + "cust": {"customer", "client", "account"}, + "customer": {"cust", "client", "account"}, + "customers": {"cust", "client", "account", "customer"}, + "desc": {"description", "name"}, + "description": {"desc", "name"}, + "inv": {"invoice"}, + "invoice": {"inv", "bill", "billing"}, + "invoices": {"inv", "invoice", "bill", "billing"}, + "name": {"description", "label"}, + "no": {"number", "identifier", "id"}, + "num": {"number", "identifier", "id"}, + "number": {"no", "num", "identifier", "id"}, + "qty": {"quantity"}, + "quantity": {"qty"}, + "total": {"amount", "value", "sum"}, + "value": {"amount", "total"}, +} + def _semantic_tokens(value: Any) -> set[str]: text = str(value or "") @@ -381,6 +412,8 @@ def _semantic_tokens(value: Any) -> set[str]: tokens.add(f"{token[:-3]}y") elif token.endswith("s") and len(token) > 3: tokens.add(token[:-1]) + for token in list(tokens): + tokens.update(_GENERIC_SEMANTIC_SYNONYMS.get(token, set())) return tokens diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index d22e45f36c..b54a0335c8 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1032,6 +1032,15 @@ def test_normalize_generation_result_sql_rewrites_bare_month_field_for_mssql(): assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized +def test_normalize_generation_result_sql_rewrites_top_limit_with_parentheses(): + normalized = normalize_generation_result_sql( + 'SELECT TOP 10 "sales"."CustName" FROM "sales"', + data_source="MSSQL", + ) + + assert normalized.startswith('SELECT TOP (10) "sales"."CustName"') + + def test_normalize_generation_result_sql_rewrites_qualified_month_field_for_mssql(): sql = """ SELECT @@ -1718,6 +1727,37 @@ def test_get_schema_intent_analysis_error_rejects_required_unmapped_concepts(): assert "did not map required request concepts" in error +def test_get_schema_intent_analysis_error_allows_entity_covered_by_dimension_mapping(): + error = get_schema_intent_analysis_error( + { + "concept_mappings": [ + { + "request_concept": "customer", + "concept_type": "entity", + "schema_objects": [], + "required_in_sql": True, + }, + { + "request_concept": "customer name", + "concept_type": "dimension", + "schema_objects": ["sales.CustName"], + "required_in_sql": True, + }, + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["sales.InvAmt"], + "required_in_sql": True, + }, + ], + "supported_schema_objects": ["sales.CustName", "sales.InvAmt"], + "is_fully_supported": True, + } + ) + + assert error is None + + def test_validate_sql_intent_alignment_uses_semantic_analysis_for_metric_count_mismatch(): error = validate_sql_intent_alignment( "Show invoice amount by customer", @@ -1893,6 +1933,86 @@ def test_validate_sql_intent_alignment_rejects_mapped_ranking_without_limit(): assert "sorting and limiting" in error +def test_validate_sql_intent_alignment_accepts_top_without_parentheses_for_ranking(): + error = validate_sql_intent_alignment( + "Top customers by invoice amount", + 'SELECT TOP 10 "sales"."CustName", SUM("sales"."InvAmt") AS "invoice_amount" ' + 'FROM "sales" GROUP BY "sales"."CustName" ORDER BY SUM("sales"."InvAmt") DESC', + {"sales": ["CustName", "InvAmt"]}, + semantic_analysis={ + "analytical_intent": "ranking", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "ranking": ["top 10 customers by invoice amount"], + "aggregations": ["sum invoice amount"], + "concept_mappings": [ + { + "request_concept": "customer", + "concept_type": "entity", + "schema_objects": [], + "required_in_sql": True, + }, + { + "request_concept": "customer name", + "concept_type": "dimension", + "schema_objects": ["sales.CustName"], + "required_in_sql": True, + }, + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["sales.InvAmt"], + "required_in_sql": True, + }, + { + "request_concept": "top 10", + "concept_type": "ranking", + "schema_objects": [], + "required_in_sql": True, + }, + ], + "supported_schema_objects": ["sales.CustName", "sales.InvAmt"], + "is_fully_supported": True, + }, + ) + + assert error is None + + +def test_validate_sql_intent_alignment_rejects_ranking_without_limit_even_with_order(): + error = validate_sql_intent_alignment( + "Top customers by invoice amount", + 'SELECT "sales"."CustName", SUM("sales"."InvAmt") AS "invoice_amount" ' + 'FROM "sales" GROUP BY "sales"."CustName" ORDER BY SUM("sales"."InvAmt") DESC', + {"sales": ["CustName", "InvAmt"]}, + semantic_analysis={ + "analytical_intent": "ranking", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "ranking": ["top 10 customers by invoice amount"], + "aggregations": ["sum invoice amount"], + "concept_mappings": [ + { + "request_concept": "customer name", + "concept_type": "dimension", + "schema_objects": ["sales.CustName"], + "required_in_sql": True, + }, + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["sales.InvAmt"], + "required_in_sql": True, + }, + ], + "is_fully_supported": True, + }, + ) + + assert error is not None + assert "sorting or limiting" in error + + def test_get_schema_intent_analysis_error_rejects_multiple_selected_interpretations(): error = get_schema_intent_analysis_error( { diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index a70f26aecb..03f7330392 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -85,6 +85,48 @@ def test_rank_semantic_schema_candidates_prefers_complete_business_concept_cover assert candidates[0]["confidence"] > candidates[1]["confidence"] +def test_rank_semantic_schema_candidates_matches_generic_abbreviations(): + candidates = rank_semantic_schema_candidates( + query="Show the top 10 customers by invoice amount", + construct_db_schemas=[ + { + "type": "TABLE", + "name": "sales_summary", + "comment": "", + "columns": [ + { + "name": "CustName", + "data_type": "varchar", + "comment": "customer display name", + }, + { + "name": "InvAmt", + "data_type": "decimal", + "comment": "invoice amount", + }, + ], + }, + { + "type": "TABLE", + "name": "refund_summary", + "comment": "", + "columns": [ + { + "name": "Refund_Amount", + "data_type": "decimal", + "comment": "refund amount", + } + ], + }, + ], + ) + + assert candidates[0]["table_name"] == "sales_summary" + assert {"customer", "invoice", "amount"} <= set( + candidates[0]["matched_query_terms"] + ) + + def test_rank_semantic_schema_candidates_penalizes_retry_rejected_schema_objects(): candidates = rank_semantic_schema_candidates( query="Show top 10 customers by invoice amount", From 3aac131f9526b08e49aff696b3bb3144d6068e67 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 20:11:07 +0530 Subject: [PATCH 0477/1087] Retry semantic SQL after schema validation failures --- .../src/pipelines/generation/utils/sql.py | 45 +++++++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 35 +++++++++++++-- .../pipelines/generation/test_sql_utils.py | 20 +++++++++ 3 files changed, 97 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 8a7ce77c9d..6c77f59599 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1496,6 +1496,28 @@ def _normalize_mssql_top_clause(sql: str) -> str: ) +def _extract_placeholder_schema_references(sql: str) -> list[str]: + placeholders = re.findall(r"<\s*([^<>]+?)\s*>", sql or "") + placeholder_refs = [ + str(placeholder).strip() + for placeholder in placeholders + if str(placeholder).strip() + ] + + placeholder_names = re.findall( + r"\b(?:dbo_)?(?:table|column|schema|database|field|metric|dimension|date|amount|customer)_?name\b", + sql or "", + flags=re.IGNORECASE, + ) + placeholder_refs.extend( + str(name).strip() + for name in placeholder_names + if str(name).strip() + ) + + return sorted(set(placeholder_refs)) + + def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = _normalize_identifier_quote_syntax(sql) normalized = _normalize_mssql_top_clause(normalized) @@ -1586,6 +1608,27 @@ async def run( }, } + placeholder_schema_references = _extract_placeholder_schema_references( + cleaned_generation_result + ) + if placeholder_schema_references: + invalid_placeholder_list = ", ".join(placeholder_schema_references) + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_VALIDATION", + "error": ( + "Generated SQL contains placeholder schema references " + f"that are not active datasource objects: {invalid_placeholder_list}. " + "Use only concrete table and column names from the active metadata." + ), + "invalid_schema_objects": placeholder_schema_references, + "correlation_id": "", + }, + } + invalid_table_references = find_invalid_table_references( cleaned_generation_result, valid_table_names or [], @@ -1605,6 +1648,7 @@ async def run( "Use only these valid table names exactly as shown: " f"{valid_table_list}" ), + "invalid_schema_objects": invalid_table_references, "correlation_id": "", }, } @@ -1628,6 +1672,7 @@ async def run( "Use only these valid table columns exactly as shown: " f"{valid_column_list}" ), + "invalid_schema_objects": invalid_column_references, "correlation_id": "", }, } diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 00c45f2fa7..1706632923 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -6829,14 +6829,30 @@ async def ask( ]["invalid_generation_result"]: rejected_schema_objects: set[str] = set() semantic_retry_attempt = 0 + semantic_retriable_validation_types = { + "SCHEMA_INTENT_VALIDATION", + "SCHEMA_VALIDATION", + } while ( failed_dry_run_result - and failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION" + and semantic_pipeline_active + and failed_dry_run_result["type"] + in semantic_retriable_validation_types and semantic_retry_attempt < 3 and not api_results ): semantic_retry_attempt += 1 semantic_retry_error = failed_dry_run_result.get("error", "") + invalid_schema_objects = failed_dry_run_result.get( + "invalid_schema_objects" + ) + if isinstance(invalid_schema_objects, list): + rejected_schema_objects.update( + str(schema_object).strip() + for schema_object in invalid_schema_objects + if schema_object is not None + and str(schema_object).strip() + ) retry_context = self._semantic_retry_context( schema_intent_analysis, semantic_retry_error, @@ -6893,7 +6909,8 @@ async def ask( semantic_retry_attempt, sorted(retry_selected_objects), ) - break + rejected_schema_objects.update(retry_selected_objects) + continue retry_documents, retry_table_names, retry_table_ddls = ( self._extract_retrieval_metadata( @@ -6915,7 +6932,7 @@ async def ask( sorted(retry_selected_objects), ) rejected_schema_objects.update(retry_selected_objects) - break + continue _retrieval_result = retry_construct_result schema_intent_analysis = retry_schema_intent_analysis @@ -7012,6 +7029,8 @@ async def ask( "invalid_generation_result" ] ) + if not failed_dry_run_result: + break invalid_sql = failed_dry_run_result.get( "sql", invalid_sql ) @@ -7038,6 +7057,16 @@ async def ask( break if not failed_dry_run_result: break + if ( + semantic_pipeline_active + and failed_dry_run_result["type"] + in semantic_retriable_validation_types + ): + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) + break if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index b54a0335c8..a44bb652b7 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,4 +1,7 @@ +import pytest + from src.pipelines.generation.utils.sql import ( + SQLGenPostProcessor, contains_unsupported_mssql_json_access, construct_semantic_schema_contract, construct_valid_table_columns, @@ -51,6 +54,23 @@ def test_schema_validation_allows_qualified_suffix_table_references(): ) == [] +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_placeholder_schema_references_before_dry_run(): + post_processor = SQLGenPostProcessor(engine=None) + + result = await post_processor.run( + replies=['SELECT * FROM '], + valid_table_names=["customers"], + valid_table_columns={"customers": ["CustName"]}, + query="Show all customers", + ) + + invalid_result = result["invalid_generation_result"] + assert invalid_result["type"] == "SCHEMA_VALIDATION" + assert invalid_result["invalid_schema_objects"] == ["table_name"] + assert "placeholder schema references" in invalid_result["error"] + + def test_column_validation_allows_qualified_suffix_table_references(): sql = ( 'SELECT "wrenai"."public"."dbo_repair_logs"."warning_signals" ' From fab63187a129d9b051e6256201b0e27fef9a8092 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 10 Jul 2026 20:58:58 +0530 Subject: [PATCH 0478/1087] Improve semantic NL-to-SQL planning and validation --- .../generation/followup_sql_generation.py | 20 +- .../followup_sql_generation_reasoning.py | 22 +- .../pipelines/generation/sql_correction.py | 10 +- .../pipelines/generation/sql_generation.py | 23 +- .../generation/sql_generation_reasoning.py | 26 +- .../pipelines/generation/sql_regeneration.py | 14 +- .../src/pipelines/generation/utils/sql.py | 630 +++++++++++++++++- .../retrieval/db_schema_retrieval.py | 54 +- .../pipelines/generation/test_sql_utils.py | 121 ++++ 9 files changed, 874 insertions(+), 46 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 f0f2412891..b5772cab37 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -2,6 +2,7 @@ import logging import sys +from datetime import datetime from typing import TYPE_CHECKING, Any from hamilton import base @@ -99,6 +100,7 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +Current Time: {{ current_time }} {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### @@ -112,10 +114,21 @@ appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. Only apply aggregate functions to columns whose active metadata type supports that operation. +Reuse the previous SQL and summary for unresolved references such as same, that, +those, previous, it, filter, sort, period, table, chart, or metric. Apply only the +requested follow-up change instead of regenerating the entire analysis from scratch. +Resolve business synonyms and different phrasings through active metadata, foreign +keys, and semantic relationships. For multi-table follow-ups, connect tables with +explicit INNER JOIN or LEFT JOIN conditions from trusted relationships only. +Resolve relative date phrases against Current Time. Infer aggregation, sorting, +limits, and chart/dashboard datasets from the follow-up wording. +Avoid SELECT *. Select only required columns, keep existing filters unless the user +changes them, push new filters before joins where possible, and limit rows for +retrieval/ranking requests. Before writing SQL, validate that the selected schema elements directly support every -key entity, metric, dimension, filter, time range, relationship, and aggregation in -the follow-up question. If the schema cannot support the requested information, do -not replace the request with a generic COUNT(*) or unrelated table query. +key entity, measure, dimension, filter, time range, join, sorting, chart, dashboard, +and aggregation in the follow-up question. If the schema cannot support the requested +information, do not replace the request with a generic COUNT(*) or unrelated table query. ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -143,6 +156,7 @@ def prompt( ) -> dict: _prompt = prompt_builder.run( query=query, + current_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), data_source=data_source, documents=documents, valid_table_names=construct_valid_table_names(documents), diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 0374caf609..a3999f3bc6 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -70,13 +70,27 @@ {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### -Use this contract for entities, metrics, dimensions, filters, joins, time -constraints, aggregations, ranking, and analytical intent. If it shows missing or -ambiguous requirements, state that limitation in the plan instead of planning -unrelated SQL. +Use this contract for entities, measures, dimensions, filters, joins, normalized +date ranges, aggregations, sorting, ranking, chart requirements, dashboard +requirements, and analytical intent. If it shows missing or ambiguous requirements, +state that limitation in the plan instead of planning unrelated SQL. {{ semantic_schema_contract }} {% endif %} +### FOLLOW-UP PLANNING REQUIREMENTS ### +Use query history to resolve references such as same, that, previous, it, those, +them, filter, period, metric, table, result, or chart. Reuse the previous schema +context and apply only the requested change unless the user asks for a new analysis. +Resolve synonyms from active metadata, metrics, views, foreign keys, and semantic +relationships. Identify required entities, measures, dimensions, filters, date +windows, sort direction, top/bottom limits, joins, and chart/dashboard dataset +shape. +For multi-table follow-ups, choose a trusted join path from foreign keys or semantic +relationships and include the join keys. If no path exists, say clarification or +schema support is required. +Resolve relative dates against Current Time. Plan only required columns and avoid +SELECT *. + Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index d1e2e93333..d0f20521cd 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -122,10 +122,14 @@ def get_sql_correction_system_prompt( user's request. Do not invent tables, columns, joins, metrics, or relationships. Only apply aggregate functions to columns whose active metadata type supports that operation. +Preserve the user's business intent, including entities, measures, dimensions, +filters, date ranges, joins, sorting, top/bottom limits, chart/dashboard dataset +shape, and aggregation. Fix obvious table/column/alias/GROUP BY/JOIN mistakes only +when the active metadata supports the correction. Before returning corrected SQL, validate that it still directly supports every key -entity, metric, dimension, filter, time range, relationship, and aggregation in the -user's question. Do not replace an unsupported request with a generic COUNT(*) or -unrelated table query. +entity, measure, dimension, filter, time range, relationship, sorting requirement, +chart/dashboard requirement, and aggregation in the user's question. Do not replace +an unsupported request with a generic COUNT(*) or unrelated table query. Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 9c568abec0..11a2646744 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -1,5 +1,6 @@ import logging import sys +from datetime import datetime from typing import Any from hamilton import base @@ -87,6 +88,7 @@ ### QUESTION ### User's Question: {{ query }} +Current Time: {{ current_time }} {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### @@ -99,10 +101,24 @@ general guidance when the question can be answered with SQL over the active metadata. Never reuse table or column names from SQL SAMPLES unless those exact names also appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. +Resolve business synonyms and different phrasings through schema names, column names, +descriptions, metadata, metrics, views, foreign keys, and semantic relationships. +For multi-table questions, choose only the tables needed and connect them with +explicit INNER JOIN or LEFT JOIN conditions from foreign keys or semantic +relationships. If no trustworthy relationship exists, do not invent a join. +Resolve relative date phrases against Current Time and express them as concrete +filter predicates on real temporal columns. +Infer aggregation from language: total/sum -> SUM when a measure is requested, +average -> AVG, number/count/how many -> COUNT, min/max/highest/lowest -> MIN/MAX +or ranking as appropriate. +For chart or dashboard requests, return the optimized aggregated dataset needed +for the chart/KPI/summary, with suitable grouping, sorting, and limits. +Avoid SELECT *. Select only required columns, push filters before joins where +possible, and limit rows for retrieval/ranking requests. Before writing SQL, validate that the selected schema elements directly support every -key entity, metric, dimension, filter, time range, relationship, and aggregation in -the question. If the schema cannot support the requested information, do not replace -the request with a generic COUNT(*) or unrelated table query. +key entity, measure, dimension, filter, time range, join, sorting, chart, dashboard, +and aggregation in the question. If the schema cannot support the requested +information, do not replace the request with a generic COUNT(*) or unrelated table query. {% if sql_generation_reasoning %} ### REASONING PLAN ### @@ -145,6 +161,7 @@ def prompt( ) _prompt = prompt_builder.run( query=query, + current_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), data_source=data_source, documents=documents, has_pcb_context=has_pcb_context, 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 7d14e29184..ef9486360a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -55,13 +55,31 @@ {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### -Use this contract for entities, metrics, dimensions, filters, joins, time -constraints, aggregations, ranking, and analytical intent. If it shows missing or -ambiguous requirements, state that limitation in the plan instead of planning -unrelated SQL. +Use this contract for entities, measures, dimensions, filters, joins, normalized +date ranges, aggregations, sorting, ranking, chart requirements, dashboard +requirements, and analytical intent. If it shows missing or ambiguous requirements, +state that limitation in the plan instead of planning unrelated SQL. {{ semantic_schema_contract }} {% endif %} +### PLANNING REQUIREMENTS ### +Build a schema-grounded SQL plan, not example-specific SQL. Resolve synonyms from +table names, column names, descriptions, metadata, metrics, views, foreign keys, +and semantic relationships. Identify the exact business entities, measures, +dimensions, filters, date windows, sort direction, top/bottom limits, joins, and +chart/dashboard requirements needed by the question. +For multi-table questions, choose a trusted join path from foreign keys or semantic +relationships and include the join keys. If no path exists, say clarification or +schema support is required. +Resolve relative dates such as today, yesterday, this/last week, this/last month, +this/last quarter, this/last year, last 30 days, last 90 days, and rolling 12 +months against Current Time. +Infer aggregation from language: total/sum -> SUM for measures, average -> AVG, +number/count/how many -> COUNT, min/max -> MIN/MAX, top/bottom/highest/lowest -> +ORDER BY with a limit. +Plan only required columns. Avoid SELECT *. Preserve conversational context for +follow-up wording and apply only the requested change. + Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 4f0668ff0f..87ae969af6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -46,9 +46,11 @@ def get_sql_regeneration_system_prompt( While generating the new SQL query, you should use the original SQL query as a reference. While generating the new SQL query, make sure to use the database schema to generate the SQL query. Before returning SQL, validate that the selected schema elements directly support -the key entities, metrics, dimensions, filters, time ranges, relationships, and -aggregations from the user's question or reasoning. Do not replace an unsupported -request with a generic COUNT(*) or unrelated table query. +the key entities, measures, dimensions, filters, time ranges, relationships, +sorting, chart/dashboard requirements, and aggregations from the user's question +or reasoning. Reuse the previous semantic context and original SQL where correct, +but fix schema, join, alias, GROUP BY, aggregation, and filter mistakes. Do not +replace an unsupported request with a generic COUNT(*) or unrelated table query. {text_to_sql_rules} @@ -119,6 +121,12 @@ def get_sql_regeneration_system_prompt( SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} +### REGENERATION REQUIREMENTS ### +Preserve the business question and semantic schema contract. Use active metadata +to correct only the parts that are wrong: tables, columns, aliases, joins, filters, +GROUP BY, HAVING, ORDER BY, aggregation, limits, date predicates, and chart/dashboard +dataset shape. Avoid SELECT * and do not add unrelated columns or tables. + Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 6c77f59599..af16f96e72 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2276,7 +2276,7 @@ def get_sql_generation_system_prompt( 9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. 10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. 11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. -12. Before generating SQL, validate that the selected schema elements directly support all key entities, metrics, dimensions, filters, time ranges, relationships, and aggregations mentioned or implied by the question. +12. Before generating SQL, validate that the selected schema elements directly support all key entities, measures, dimensions, filters, sorting, chart shape, time ranges, relationships, and aggregations mentioned or implied by the question. 13. Do not answer a specific business metric, trend, summary, comparison, dashboard, or analysis request with a generic record-count query unless the user explicitly asks only for record count. 14. If the required information cannot be derived from the available active schema, return the closest schema-grounded limitation instead of inventing unrelated SQL. 15. If a SEMANTIC SCHEMA CONTRACT is provided, it is the primary source of truth for selecting tables, columns, metrics, joins, filters, grouping, sorting, and date logic. Generate SQL from the highest-confidence validated concept-to-schema mappings in that contract and do not independently infer substitute schema objects. @@ -2389,9 +2389,14 @@ def construct_semantic_schema_contract( ("Dimensions", "dimensions"), ("Filters", "filters"), ("Time constraints", "time_constraints"), + ("Date ranges", "date_ranges"), ("Aggregations", "aggregations"), ("Ranking", "ranking"), + ("Sorting", "sorting"), + ("Chart requirements", "chart_requirements"), + ("Dashboard requirements", "dashboard_requirements"), ("Relationships", "relationships"), + ("Join paths", "join_paths"), ("Supported schema objects", "supported_schema_objects"), ): lines.extend(_format_semantic_list(label, _semantic_analysis_items(semantic_analysis, key))) @@ -3498,6 +3503,536 @@ def format_valid_table_columns(valid_table_columns: dict[str, list[str]]) -> str "department", "location", } +_BUSINESS_INTENT_STOP_WORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "chart", + "create", + "data", + "each", + "for", + "from", + "give", + "graph", + "how", + "in", + "is", + "me", + "of", + "on", + "or", + "please", + "show", + "table", + "the", + "to", + "with", +} +_SCHEMA_CONCEPT_GENERIC_TOKENS = { + "amount", + "at", + "code", + "count", + "date", + "day", + "description", + "dt", + "id", + "ids", + "key", + "month", + "name", + "no", + "number", + "num", + "pct", + "percent", + "percentage", + "price", + "quantity", + "qty", + "rate", + "time", + "timestamp", + "total", + "type", + "value", + "year", +} +_TEMPORAL_SCHEMA_TOKENS = { + "at", + "created", + "date", + "day", + "ended", + "modified", + "month", + "started", + "time", + "timestamp", + "updated", + "week", + "year", +} +_CHART_TERMS = { + "area chart", + "bar chart", + "chart", + "donut chart", + "graph", + "line chart", + "pie chart", + "plot", + "scatter plot", + "visualization", +} +_RANKING_TERMS = { + "bottom", + "highest", + "largest", + "lowest", + "most", + "rank", + "ranked", + "ranking", + "smallest", + "top", +} +_SORTING_TERMS = { + "ascending", + "descending", + "order by", + "ordered by", + "sort", + "sorted", +} + + +def _singularize_business_token(token: str) -> str: + token = token.lower() + if len(token) > 4 and token.endswith("ies"): + return f"{token[:-3]}y" + if len(token) > 3 and token.endswith("es"): + return token[:-2] + if len(token) > 3 and token.endswith("s"): + return token[:-1] + return token + + +def _business_intent_tokens(text: str | None) -> list[str]: + spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(text or "")) + return [ + _singularize_business_token(token) + for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) + if token and token not in _BUSINESS_INTENT_STOP_WORDS + ] + + +def _schema_column_concept_keys(column_name: str) -> set[tuple[str, ...]]: + tokens = tuple(_business_intent_tokens(column_name)) + if not tokens: + return set() + + keys: set[tuple[str, ...]] = {tokens} + token_set = set(tokens) + informative_tokens = tuple( + token for token in tokens if token not in _SCHEMA_CONCEPT_GENERIC_TOKENS + ) + metric_tokens = { + "amount", + "cost", + "count", + "price", + "profit", + "quantity", + "qty", + "rate", + "revenue", + "sales", + "total", + "value", + } + identifier_suffix_tokens = {"code", "id", "ids", "key", "no", "number"} + if informative_tokens and ( + len(informative_tokens) > 1 + or not token_set.intersection(metric_tokens) + and ( + not token_set.intersection(identifier_suffix_tokens) + or bool(set(informative_tokens).intersection(_DIMENSION_TERMS)) + ) + ): + keys.add(informative_tokens) + + if len(tokens) > 1: + without_suffix = tuple( + token + for token in tokens + if token not in {"id", "ids", "key", "name", "number", "no", "code"} + ) + if without_suffix and ( + len(without_suffix) > 1 + or bool(set(without_suffix).intersection(_DIMENSION_TERMS)) + ): + keys.add(without_suffix) + + for token in tokens: + if token in _SCHEMA_CONCEPT_GENERIC_TOKENS or token in _DIMENSION_TERMS: + keys.add((token,)) + + return {key for key in keys if key} + + +def _query_mentions_concept( + concept_tokens: tuple[str, ...], + query_tokens: set[str], + asks_time_analysis: bool, +) -> bool: + if not concept_tokens: + return False + + if set(concept_tokens).issubset(query_tokens): + return True + + if set(concept_tokens).intersection({"code", "id", "ids", "key", "no", "number"}): + return False + + meaningful_tokens = [ + token + for token in concept_tokens + if token not in _SCHEMA_CONCEPT_GENERIC_TOKENS + ] + if len(meaningful_tokens) != 1: + return False + + token = meaningful_tokens[0] + if token not in query_tokens: + return False + + if set(concept_tokens) & _TEMPORAL_SCHEMA_TOKENS and not asks_time_analysis: + return False + + return len(token) >= 4 or token in _DIMENSION_TERMS + + +def _infer_query_schema_requirements( + query: str, + valid_table_columns: dict[str, list[str]], +) -> dict[str, set[str]]: + query_tokens = set(_business_intent_tokens(query)) + if not query_tokens: + return {} + + asks_time_analysis = _query_requests_time_analysis(query) + requirements: dict[str, set[str]] = {} + + for table_name, columns in valid_table_columns.items(): + if not table_name: + continue + table_tokens = tuple( + token + for token in _business_intent_tokens(table_name) + if token not in {"dbo", "public", "schema", "table", "tbl"} + ) + table_informative_tokens = tuple( + token + for token in table_tokens + if token not in _SCHEMA_CONCEPT_GENERIC_TOKENS + ) + for concept_tokens in {table_tokens, table_informative_tokens}: + if _query_mentions_concept( + concept_tokens, + query_tokens, + asks_time_analysis, + ): + requirements.setdefault(" ".join(concept_tokens), set()).add( + str(table_name) + ) + + for column_name in columns or []: + if not column_name: + continue + for concept_tokens in _schema_column_concept_keys(str(column_name)): + if not _query_mentions_concept( + concept_tokens, + query_tokens, + asks_time_analysis, + ): + continue + concept = " ".join(concept_tokens) + requirements.setdefault(concept, set()).add( + f"{table_name}.{column_name}" + ) + + return { + concept: objects + for concept, objects in requirements.items() + if objects and not _is_ambiguous_schema_concept(concept, objects) + } + + +def _is_ambiguous_schema_concept(concept: str, schema_objects: set[str]) -> bool: + if len(schema_objects) <= 1: + return False + + concept_compact = _compact_sql_identifier(concept) + column_names = { + _schema_object_parts(schema_object)[-1] + for schema_object in schema_objects + if _schema_object_parts(schema_object) + } + compact_columns = { + _compact_sql_identifier(column_name) for column_name in column_names + } + + if concept_compact in compact_columns: + return False + + return len(compact_columns) > 1 + + +def _sql_references_any_schema_object( + sql: str, + schema_objects: set[str], + valid_table_columns: dict[str, list[str]], +) -> bool: + return any( + _sql_references_schema_object(sql, schema_object, valid_table_columns) + for schema_object in schema_objects + ) + + +def _sql_has_ordering(sql: str) -> bool: + return bool(re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE)) + + +def _sql_has_limit(sql: str) -> bool: + return bool( + re.search( + r"\b(?:LIMIT\s+\d+|TOP\s*\(?\s*\d+\s*\)?|FETCH\s+FIRST\s+\d+)\b", + sql or "", + flags=re.IGNORECASE, + ) + ) + + +def _sql_has_filter_predicate(sql: str) -> bool: + return bool( + re.search(r"\b(?:WHERE|HAVING)\b", sql or "", flags=re.IGNORECASE) + ) + + +def _query_requests_ranking(query: str) -> bool: + return _contains_phrase(query, _RANKING_TERMS) + + +def _query_requests_sorting(query: str) -> bool: + return _query_requests_ranking(query) or _contains_phrase(query, _SORTING_TERMS) + + +def _query_requests_limited_ranking(query: str) -> bool: + normalized = str(query or "").lower() + return bool( + re.search(r"\b(?:top|bottom)\s+\d+\b", normalized) + or re.search(r"\b(?:top|bottom|highest|lowest|largest|smallest)\b", normalized) + ) + + +def _query_requests_chart(query: str) -> bool: + return _contains_phrase(query, _CHART_TERMS) + + +def _query_requests_time_filter(query: str) -> bool: + normalized = str(query or "").lower() + return bool( + re.search( + r"\b(?:last|next|previous|prior|this)\s+" + r"(?:\d+\s+)?(?:day|week|month|quarter|year)s?\b", + normalized, + ) + or re.search(r"\b(?:between|since|before|after)\b", normalized) + ) + + +def _query_requests_literal_filter(query: str) -> bool: + normalized = str(query or "").lower() + if re.search(r"'[^']+'|\"[^\"]+\"", query or ""): + return True + return bool( + re.search( + r"\b(?:where|only|exclude|excluding|filtered by|filter by|with status|" + r"with category|for status|for category|for region|for country|" + r"for market|for customer|for product)\b", + normalized, + ) + ) + + +def _requested_query_aggregate_functions(query: str) -> set[str]: + normalized = str(query or "").lower() + functions: set[str] = set() + + if re.search(r"\b(?:avg|average|mean)\b", normalized): + functions.add("AVG") + if re.search(r"\b(?:min|minimum)\b", normalized): + functions.add("MIN") + if re.search(r"\b(?:max|maximum)\b", normalized): + functions.add("MAX") + if re.search(r"\b(?:count|number of|how many|record count|row count)\b", normalized): + functions.add("COUNT") + elif re.search(r"\bsum\b", normalized) or ( + re.search(r"\btotal\b", normalized) + and _query_requests_specific_metric(normalized) + ): + functions.add("SUM") + + return functions + + +def _simple_groupable_expression_key(expression: str) -> str | None: + expression = _strip_projection_alias( + re.sub(r"^\s*DISTINCT\s+", "", expression, flags=re.IGNORECASE) + ).strip() + expression = re.sub( + r"^\s*TOP\s*\(?\s*\d+\s*\)?\s+", + "", + expression, + flags=re.IGNORECASE, + ) + if not expression or expression == "*": + return None + if _AGGREGATE_PATTERN.search(expression): + return None + if re.search(r"\b(?:CASE|SELECT|OVER)\b", expression, flags=re.IGNORECASE): + return None + + return re.sub(r"\s+", "", expression).strip().lower() + + +def _validate_group_by_for_aggregates(sql: str) -> str | None: + if not _AGGREGATE_PATTERN.search(sql or ""): + return None + + select_spans = _find_select_list_spans(sql) + if not select_spans: + return None + + group_bodies = _extract_clause_bodies( + sql, + r"GROUP\s+BY", + ["HAVING", r"ORDER\s+BY", "LIMIT", "FETCH", "UNION"], + ) + if not group_bodies: + return None + + grouped_keys = { + key + for body in group_bodies + for item in _split_top_level_select_items(body) + if (key := _simple_groupable_expression_key(item)) + } + if not grouped_keys: + return None + + for start, end in select_spans: + for item in _split_top_level_select_items(sql[start:end]): + key = _simple_groupable_expression_key(item) + if not key or key in grouped_keys: + continue + return ( + "Generated SQL selects a non-aggregated expression that is not " + f"present in GROUP BY: {item.strip()}." + ) + + return None + + +def _validate_query_schema_requirements( + query: str, + sql: str, + valid_table_columns: dict[str, list[str]], +) -> str | None: + if group_by_error := _validate_group_by_for_aggregates(sql): + return group_by_error + + schema_requirements = _infer_query_schema_requirements( + query, + valid_table_columns, + ) + for concept, schema_objects in sorted(schema_requirements.items()): + if _sql_references_any_schema_object( + sql, + schema_objects, + valid_table_columns, + ): + continue + return ( + "Generated SQL does not reference schema objects that match the " + f"requested business concept '{concept}': " + f"{', '.join(sorted(schema_objects))}." + ) + + for function_name in _requested_query_aggregate_functions(query): + if function_name == "COUNT" and _sql_is_plain_count(sql): + continue + if not _sql_has_aggregate_function(sql, function_name): + return ( + "Generated SQL does not use the aggregation requested by the " + f"business question: {function_name}." + ) + + asks_ranking = _query_requests_ranking(query) + asks_sorting = _query_requests_sorting(query) + if asks_sorting and not _sql_has_ordering(sql): + return ( + "Generated SQL does not include ORDER BY required by the requested " + "sorting or ranking intent." + ) + + if asks_ranking and _query_requests_limited_ranking(query) and not _sql_has_limit(sql): + return ( + "Generated SQL does not include a limit/TOP clause required by the " + "requested top/bottom ranking intent." + ) + + if ( + (_query_requests_time_filter(query) or _query_requests_literal_filter(query)) + and not _sql_has_filter_predicate(sql) + ): + return ( + "Generated SQL does not include a WHERE or HAVING predicate required " + "by the requested filter or time range." + ) + + if _query_requests_chart(query): + if (asks_ranking or _query_requests_grouped_analysis(query)) and not re.search( + r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE + ): + return ( + "Generated SQL does not group the dataset required for the " + "requested chart dimension." + ) + if asks_ranking and not _sql_has_ordering(sql): + return ( + "Generated SQL does not sort the dataset required for the " + "requested chart ranking." + ) + + if re.search(r"\bJOIN\b", sql or "", flags=re.IGNORECASE): + if re.search(r"\bCROSS\s+JOIN\b", sql or "", flags=re.IGNORECASE): + return None + if not re.search(r"\b(?:ON|USING)\b", sql or "", flags=re.IGNORECASE): + return ( + "Generated SQL joins tables without an ON or USING condition. " + "Use explicit schema relationships for JOINs." + ) + + return None def _contains_phrase(text: str, terms: set[str] | tuple[str, ...]) -> bool: @@ -3786,7 +4321,12 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: "aggregations", "relationships", "time_constraints", + "date_ranges", "ranking", + "sorting", + "chart_requirements", + "dashboard_requirements", + "join_paths", "supported_schema_objects", "candidate_schema_scores", "concept_mappings", @@ -4070,6 +4610,40 @@ def _sql_references_table(sql: str, table_name: str) -> bool: return any(_sql_contains_identifier(sql, candidate) for candidate in table_candidates) +def _sql_references_unqualified_column_expression(sql: str, column_name: str) -> bool: + expected_column = _normalize_sql_identifier(str(column_name or "")) + if not expected_column: + return False + + expected_lower = expected_column.lower() + expected_compact = _compact_sql_identifier(expected_column) + candidate_columns = { + _normalize_sql_identifier(column) + for column in _find_unqualified_column_candidates(sql) + } + + for start, end in _find_select_list_spans(sql): + for item in _split_top_level_select_items(sql[start:end]): + expression = _strip_projection_alias( + re.sub(r"^\s*DISTINCT\s+", "", item, flags=re.IGNORECASE) + ) + searchable_expression = _strip_sql_literals(expression) + for match in re.finditer(_SQL_IDENTIFIER_PATTERN, searchable_expression): + before = searchable_expression[: match.start()].rstrip() + after = searchable_expression[match.end() :].lstrip() + if before.endswith(".") or after.startswith("."): + continue + identifier = _normalize_sql_identifier(match.group(0)) + if identifier and identifier.lower() not in _SQL_NON_COLUMN_IDENTIFIERS: + candidate_columns.add(identifier) + + return any( + candidate.lower() == expected_lower + or _compact_sql_identifier(candidate) == expected_compact + for candidate in candidate_columns + ) + + def _sql_references_schema_object( sql: str, schema_object: str, @@ -4099,7 +4673,7 @@ def _sql_references_schema_object( if _sql_references_table(sql, expected_table) and _sql_contains_identifier( sql, expected_column ): - return True + return _sql_references_unqualified_column_expression(sql, expected_column) return False @@ -4317,7 +4891,16 @@ def _validate_sql_against_semantic_analysis( time_constraints = _semantic_analysis_items( semantic_analysis, "time_constraints" ) + date_ranges = _semantic_analysis_items(semantic_analysis, "date_ranges") ranking = _semantic_analysis_items(semantic_analysis, "ranking") + sorting = _semantic_analysis_items(semantic_analysis, "sorting") + chart_requirements = _semantic_analysis_items( + semantic_analysis, "chart_requirements" + ) + dashboard_requirements = _semantic_analysis_items( + semantic_analysis, "dashboard_requirements" + ) + join_paths = _semantic_analysis_items(semantic_analysis, "join_paths") requests_record_count = _semantic_analysis_requests_record_count( semantic_analysis ) @@ -4342,7 +4925,9 @@ def _validate_sql_against_semantic_analysis( "or analytical calculations from the active schema." ) - if time_constraints and not _sql_has_temporal_reference(sql, valid_table_columns): + if (time_constraints or date_ranges) and not _sql_has_temporal_reference( + sql, valid_table_columns + ): return ( "Generated SQL does not use a temporal field or supported date/time " "expression, but the semantic analysis identified time constraints " @@ -4359,7 +4944,7 @@ def _validate_sql_against_semantic_analysis( "dimensions or time grain identified in the semantic analysis." ) - if ranking: + if ranking or sorting: has_ranking_order = bool( re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE) ) @@ -4374,12 +4959,39 @@ def _validate_sql_against_semantic_analysis( has_ranking_order = True has_ranking_limit = True - if ranking and (not has_ranking_order or not has_ranking_limit): + if (ranking or sorting) and not has_ranking_order: + return ( + "Generated SQL does not include sorting logic required " + "by the ranking intent or sorting intent." + ) + + if ranking and not has_ranking_limit: return ( "Generated SQL does not include sorting or limiting logic required " "by the ranking intent." ) + if join_paths and re.search(r"\bJOIN\b", sql or "", flags=re.IGNORECASE): + if not re.search(r"\b(?:ON|USING)\b", sql or "", flags=re.IGNORECASE): + return ( + "Generated SQL does not use explicit join conditions required " + "by the semantic join path." + ) + + if chart_requirements and _AGGREGATE_PATTERN.search(sql or "") and dimensions: + if not re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE): + return ( + "Generated SQL does not group the aggregated dataset required " + "by the chart requirements." + ) + + if dashboard_requirements and _sql_is_plain_count(sql) and not requests_record_count: + return ( + "Generated SQL answers a dashboard request with only a generic record " + "count. Use the KPI, chart, summary, and insight requirements from " + "the semantic analysis." + ) + if aggregations and not _AGGREGATE_PATTERN.search(sql or ""): return ( "Generated SQL does not include the aggregation required by the " @@ -4447,6 +5059,14 @@ def validate_sql_intent_alignment( "or report that the schema does not expose that dimension." ) + schema_requirement_error = _validate_query_schema_requirements( + normalized_query, + sql, + valid_table_columns, + ) + if schema_requirement_error: + return schema_requirement_error + missing_metric_terms = _missing_metric_support( normalized_query, sql, diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 60dea2a78d..e30119374b 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -37,28 +37,30 @@ The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### -1. First perform a semantic analysis of the user's request. Identify intended business entities, identifiers, descriptive attributes, metrics, dimensions, filters, aggregations, relationships, time constraints, ranking requirements, and analytical intent such as retrieval, detailed records, summary, comparison, trend analysis, dashboard, KPI, ranking, or record count. +1. First perform a semantic analysis of the user's request. Identify intended business entities, identifiers, descriptive attributes, metrics, dimensions, filters, aggregations, relationships, time constraints, ranking requirements, chart requirements, dashboard/KPI requirements, and analytical intent such as retrieval, detailed records, summary, comparison, trend analysis, dashboard, KPI, ranking, or record count. 2. Map each business term to explicit schema objects only when the active schema directly supports that term. Distinguish entities such as customer/order/invoice/product from identifiers such as order ID or invoice number, descriptive attributes, and measurable metrics such as amount, quantity, cost, profit, revenue, or duration. -3. Select tables and columns by semantic fit to the full request, not by isolated keyword overlap or commonly used default tables. -4. Include join keys and relationship columns needed to connect selected tables. Do not invent relationships or foreign keys. -5. If the schema does not support a requested entity, metric, dimension, filter, time range, aggregation, or ranking requirement, record it in `missing_requirements`. -6. If multiple schema interpretations are equally plausible and the question does not disambiguate them, record them in `ambiguous_requirements`. -7. Set `is_fully_supported` to false when any required request component is missing or ambiguous. -8. For each selected table, provide a concise reason for why the table is semantically relevant. -9. For each selected column, provide a concise reason for why the column is necessary. -10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. -11. Broad business entities can be satisfied by the selected table plus the best descriptive or identifier column for that entity. For example, an entity such as customer may map to a customer name, customer number, account, client, or similar descriptive/identifier column when that is the active schema's representation. Do not leave an entity unmapped when a selected dimension or identifier column represents it. -12. When the user asks for a top/bottom/ranking query, map the ranked dimension and ranked metric separately. The SQL generator must be able to ORDER BY the metric and limit rows. -13. Do not add filters or time constraints that are not requested or implied by the user. Only map date/time concepts when the user asks for a time period, trend, date filter, or date dimension. -14. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. -15. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. -16. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. -17. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. -18. If a "." is included in columns, put the name before the first dot into chosen columns. -19. The number of columns chosen must match the number of reasoning. -20. Final chosen columns must be only column names, don't prefix it with table names. -21. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -22. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +3. Understand synonyms and different phrasings by using table names, column names, descriptions, metadata, metrics, views, foreign keys, and semantic relationships. Do not rely on hardcoded examples or default tables. +4. Select tables and columns by semantic fit to the full request, not by isolated keyword overlap or commonly used default tables. +5. Include join keys and relationship columns needed to connect selected tables. Prefer explicit foreign keys and semantic relationships. For multi-table questions, identify a join path; if no trustworthy path exists, record the missing relationship instead of selecting unrelated tables. +6. Normalize relative date language such as today, yesterday, this/last week, this/last month, this/last quarter, this/last year, last 30 days, last 90 days, and rolling 12 months into date requirements using Current Time when available. +7. If the schema does not support a requested entity, metric, dimension, filter, time range, aggregation, chart, dashboard, or ranking requirement, record it in `missing_requirements`. +8. If multiple schema interpretations are equally plausible and the question does not disambiguate them, record them in `ambiguous_requirements`. +9. Set `is_fully_supported` to false when any required request component is missing or ambiguous. +10. For each selected table, provide a concise reason for why the table is semantically relevant. +11. For each selected column, provide a concise reason for why the column is necessary. +12. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. +13. Broad business entities can be satisfied by the selected table plus the best descriptive or identifier column for that entity. For example, an entity such as customer may map to a customer name, customer number, account, client, or similar descriptive/identifier column when that is the active schema's representation. Do not leave an entity unmapped when a selected dimension or identifier column represents it. +14. When the user asks for a top/bottom/ranking query, map the ranked dimension and ranked metric separately. The SQL generator must be able to ORDER BY the metric and limit rows. +15. Do not add filters or time constraints that are not requested or implied by the user. Only map date/time concepts when the user asks for a time period, trend, date filter, or date dimension. +16. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. +17. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/chart/dashboard/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. +18. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. +19. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. +20. If a "." is included in columns, put the name before the first dot into chosen columns. +21. The number of columns chosen must match the number of reasoning. +22. Final chosen columns must be only column names, don't prefix it with table names. +23. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +24. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -74,7 +76,12 @@ "aggregations": ["aggregation or calculation requirements"], "relationships": ["required joins or relationships"], "time_constraints": ["time filters, grains, or trend requirements"], + "date_ranges": ["normalized relative date requirements such as start/end dates or rolling windows"], "ranking": ["top/bottom/order/limit requirements"], + "sorting": ["sort fields and direction requirements"], + "chart_requirements": ["requested or inferred chart type, x/y encodings, series, and grain"], + "dashboard_requirements": ["requested KPI, chart, summary, and insight sections"], + "join_paths": ["foreign key or semantic relationship path needed to connect selected tables"], "supported_schema_objects": ["table.column or metric names that directly support the request"], "candidate_schema_scores": [ { @@ -1087,7 +1094,12 @@ class SemanticAnalysis(BaseModel): aggregations: list[str] = Field(default_factory=list) relationships: list[str] = Field(default_factory=list) time_constraints: list[str] = Field(default_factory=list) + date_ranges: list[str] = Field(default_factory=list) ranking: list[str] = Field(default_factory=list) + sorting: list[str] = Field(default_factory=list) + chart_requirements: list[str] = Field(default_factory=list) + dashboard_requirements: list[str] = Field(default_factory=list) + join_paths: list[str] = Field(default_factory=list) supported_schema_objects: list[str] = Field(default_factory=list) candidate_schema_scores: list[SemanticCandidateSchemaScore] = Field( default_factory=list diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index a44bb652b7..f0b4a522f2 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1832,6 +1832,105 @@ def test_validate_sql_intent_alignment_allows_semantic_count_metric(): assert error is None +def test_validate_sql_intent_alignment_rejects_missing_schema_backed_measure(): + error = validate_sql_intent_alignment( + "Show total invoice amount by customer", + 'SELECT "invoices"."customer_id", SUM("invoices"."tax_amount") ' + 'AS "total_invoice_amount" FROM "invoices" ' + 'GROUP BY "invoices"."customer_id"', + {"invoices": ["customer_id", "invoice_amount", "tax_amount"]}, + ) + + assert error is not None + assert "invoice amount" in error + assert "invoices.invoice_amount" in error + + +def test_validate_sql_intent_alignment_rejects_missing_requested_aggregation(): + error = validate_sql_intent_alignment( + "Show average order value by region", + 'SELECT "orders"."region", SUM("orders"."order_value") AS "order_value" ' + 'FROM "orders" GROUP BY "orders"."region"', + {"orders": ["region", "order_value"]}, + ) + + assert error is not None + assert "AVG" in error + + +def test_validate_sql_intent_alignment_rejects_ranking_without_order_or_limit(): + error = validate_sql_intent_alignment( + "Show top 10 customers by invoice amount", + 'SELECT "invoices"."customer_id", ' + 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' + 'FROM "invoices" GROUP BY "invoices"."customer_id"', + {"invoices": ["customer_id", "invoice_amount"]}, + ) + + assert error is not None + assert "ORDER BY" in error + + +def test_validate_sql_intent_alignment_rejects_time_filter_without_predicate(): + error = validate_sql_intent_alignment( + "Show monthly order count for the last 12 months", + 'SELECT DATEPART(YEAR, "orders"."created_at") AS "year", ' + 'DATEPART(MONTH, "orders"."created_at") AS "month", ' + 'COUNT(*) AS "order_count" FROM "orders" ' + 'GROUP BY DATEPART(YEAR, "orders"."created_at"), ' + 'DATEPART(MONTH, "orders"."created_at")', + {"orders": ["id", "created_at"]}, + ) + + assert error is not None + assert "WHERE or HAVING" in error + + +def test_validate_sql_intent_alignment_rejects_missing_group_by_column(): + error = validate_sql_intent_alignment( + "Show total invoice amount by customer and region", + 'SELECT "invoices"."customer_id", "invoices"."region", ' + 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' + 'FROM "invoices" GROUP BY "invoices"."customer_id"', + {"invoices": ["customer_id", "region", "invoice_amount"]}, + ) + + assert error is not None + assert "not present in GROUP BY" in error + + +def test_validate_sql_intent_alignment_allows_schema_backed_business_question(): + error = validate_sql_intent_alignment( + "Show top 10 customers by total invoice amount last month as a bar chart", + 'SELECT TOP 10 "invoices"."customer_id", ' + 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' + 'FROM "invoices" ' + 'WHERE "invoices"."created_at" >= \'2026-06-01 00:00:00\' ' + 'AND "invoices"."created_at" < \'2026-07-01 00:00:00\' ' + 'GROUP BY "invoices"."customer_id" ' + 'ORDER BY "total_invoice_amount" DESC', + {"invoices": ["customer_id", "invoice_amount", "created_at"]}, + ) + + assert error is None + + +def test_validate_sql_intent_alignment_rejects_join_without_condition(): + error = validate_sql_intent_alignment( + "Show order amount by customer region", + 'SELECT "customers"."region", SUM("orders"."order_amount") ' + 'AS "total_order_amount" FROM "orders" JOIN "customers" ' + 'GROUP BY "customers"."region"', + { + "orders": ["customer_id", "order_amount"], + "customers": ["customer_id", "region"], + }, + ) + + assert error is not None + assert "ON or USING" in error + + def test_validate_sql_intent_alignment_rejects_unmapped_metric_substitution(): error = validate_sql_intent_alignment( "Show invoice amount by customer", @@ -2102,6 +2201,28 @@ def test_construct_semantic_schema_contract_prioritizes_concept_mappings(): assert "Do not substitute identifiers for metrics" in contract +def test_construct_semantic_schema_contract_includes_generic_planning_requirements(): + contract = construct_semantic_schema_contract( + { + "analytical_intent": "dashboard", + "metrics": ["invoice amount"], + "dimensions": ["customer"], + "date_ranges": ["last month: 2026-06-01 to 2026-07-01"], + "sorting": ["invoice amount descending"], + "chart_requirements": ["bar chart by customer"], + "dashboard_requirements": ["KPI total invoice amount"], + "join_paths": ["invoices.customer_id -> customers.id"], + "is_fully_supported": True, + } + ) + + assert "Date ranges: last month" in contract + assert "Sorting: invoice amount descending" in contract + assert "Chart requirements: bar chart by customer" in contract + assert "Dashboard requirements: KPI total invoice amount" in contract + assert "Join paths: invoices.customer_id -> customers.id" in contract + + def test_construct_semantic_schema_contract_allows_legacy_analysis_without_mappings(): contract = construct_semantic_schema_contract( { From 762a71a109bd856da3eb15470d22d14bfdd867a7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:42:57 +0530 Subject: [PATCH 0479/1087] Revert "Retry semantic SQL after schema validation failures" This reverts commit 3aac131f9526b08e49aff696b3bb3144d6068e67. --- .../src/pipelines/generation/utils/sql.py | 45 ------------------- wren-ai-service/src/web/v1/services/ask.py | 35 ++------------- .../pipelines/generation/test_sql_utils.py | 20 --------- 3 files changed, 3 insertions(+), 97 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index af16f96e72..c6f00ff3ad 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1496,28 +1496,6 @@ def _normalize_mssql_top_clause(sql: str) -> str: ) -def _extract_placeholder_schema_references(sql: str) -> list[str]: - placeholders = re.findall(r"<\s*([^<>]+?)\s*>", sql or "") - placeholder_refs = [ - str(placeholder).strip() - for placeholder in placeholders - if str(placeholder).strip() - ] - - placeholder_names = re.findall( - r"\b(?:dbo_)?(?:table|column|schema|database|field|metric|dimension|date|amount|customer)_?name\b", - sql or "", - flags=re.IGNORECASE, - ) - placeholder_refs.extend( - str(name).strip() - for name in placeholder_names - if str(name).strip() - ) - - return sorted(set(placeholder_refs)) - - def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = _normalize_identifier_quote_syntax(sql) normalized = _normalize_mssql_top_clause(normalized) @@ -1608,27 +1586,6 @@ async def run( }, } - placeholder_schema_references = _extract_placeholder_schema_references( - cleaned_generation_result - ) - if placeholder_schema_references: - invalid_placeholder_list = ", ".join(placeholder_schema_references) - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_VALIDATION", - "error": ( - "Generated SQL contains placeholder schema references " - f"that are not active datasource objects: {invalid_placeholder_list}. " - "Use only concrete table and column names from the active metadata." - ), - "invalid_schema_objects": placeholder_schema_references, - "correlation_id": "", - }, - } - invalid_table_references = find_invalid_table_references( cleaned_generation_result, valid_table_names or [], @@ -1648,7 +1605,6 @@ async def run( "Use only these valid table names exactly as shown: " f"{valid_table_list}" ), - "invalid_schema_objects": invalid_table_references, "correlation_id": "", }, } @@ -1672,7 +1628,6 @@ async def run( "Use only these valid table columns exactly as shown: " f"{valid_column_list}" ), - "invalid_schema_objects": invalid_column_references, "correlation_id": "", }, } diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 1706632923..00c45f2fa7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -6829,30 +6829,14 @@ async def ask( ]["invalid_generation_result"]: rejected_schema_objects: set[str] = set() semantic_retry_attempt = 0 - semantic_retriable_validation_types = { - "SCHEMA_INTENT_VALIDATION", - "SCHEMA_VALIDATION", - } while ( failed_dry_run_result - and semantic_pipeline_active - and failed_dry_run_result["type"] - in semantic_retriable_validation_types + and failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION" and semantic_retry_attempt < 3 and not api_results ): semantic_retry_attempt += 1 semantic_retry_error = failed_dry_run_result.get("error", "") - invalid_schema_objects = failed_dry_run_result.get( - "invalid_schema_objects" - ) - if isinstance(invalid_schema_objects, list): - rejected_schema_objects.update( - str(schema_object).strip() - for schema_object in invalid_schema_objects - if schema_object is not None - and str(schema_object).strip() - ) retry_context = self._semantic_retry_context( schema_intent_analysis, semantic_retry_error, @@ -6909,8 +6893,7 @@ async def ask( semantic_retry_attempt, sorted(retry_selected_objects), ) - rejected_schema_objects.update(retry_selected_objects) - continue + break retry_documents, retry_table_names, retry_table_ddls = ( self._extract_retrieval_metadata( @@ -6932,7 +6915,7 @@ async def ask( sorted(retry_selected_objects), ) rejected_schema_objects.update(retry_selected_objects) - continue + break _retrieval_result = retry_construct_result schema_intent_analysis = retry_schema_intent_analysis @@ -7029,8 +7012,6 @@ async def ask( "invalid_generation_result" ] ) - if not failed_dry_run_result: - break invalid_sql = failed_dry_run_result.get( "sql", invalid_sql ) @@ -7057,16 +7038,6 @@ async def ask( break if not failed_dry_run_result: break - if ( - semantic_pipeline_active - and failed_dry_run_result["type"] - in semantic_retriable_validation_types - ): - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) - break if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index f0b4a522f2..df8dd53825 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,7 +1,4 @@ -import pytest - from src.pipelines.generation.utils.sql import ( - SQLGenPostProcessor, contains_unsupported_mssql_json_access, construct_semantic_schema_contract, construct_valid_table_columns, @@ -54,23 +51,6 @@ def test_schema_validation_allows_qualified_suffix_table_references(): ) == [] -@pytest.mark.asyncio -async def test_sql_post_processor_rejects_placeholder_schema_references_before_dry_run(): - post_processor = SQLGenPostProcessor(engine=None) - - result = await post_processor.run( - replies=['SELECT * FROM '], - valid_table_names=["customers"], - valid_table_columns={"customers": ["CustName"]}, - query="Show all customers", - ) - - invalid_result = result["invalid_generation_result"] - assert invalid_result["type"] == "SCHEMA_VALIDATION" - assert invalid_result["invalid_schema_objects"] == ["table_name"] - assert "placeholder schema references" in invalid_result["error"] - - def test_column_validation_allows_qualified_suffix_table_references(): sql = ( 'SELECT "wrenai"."public"."dbo_repair_logs"."warning_signals" ' From ac2eca68063b68d7077bbc24ebad48be35cc6639 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:46:28 +0530 Subject: [PATCH 0480/1087] Revert "Fix semantic entity mapping and ranking SQL validation" This reverts commit 3257287045198a4b9d73e610ae842b2a56a07b10. --- .../src/pipelines/generation/utils/sql.py | 185 ++---------------- .../retrieval/db_schema_retrieval.py | 51 +---- .../pipelines/generation/test_sql_utils.py | 120 ------------ .../retrieval/test_db_schema_retrieval.py | 42 ---- 4 files changed, 23 insertions(+), 375 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c6f00ff3ad..896db1848d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1080,7 +1080,7 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: if re.match(r"\s*SELECT\s+DISTINCT\b", without_limit, flags=re.IGNORECASE): return re.sub( r"\bSELECT\s+DISTINCT\b", - f"SELECT DISTINCT TOP ({limit})", + f"SELECT DISTINCT TOP {limit}", without_limit, count=1, flags=re.IGNORECASE, @@ -1089,7 +1089,7 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: if re.match(r"\s*SELECT\b", without_limit, flags=re.IGNORECASE): return re.sub( r"\bSELECT\b", - f"SELECT TOP ({limit})", + f"SELECT TOP {limit}", without_limit, count=1, flags=re.IGNORECASE, @@ -1458,7 +1458,7 @@ def _rewrite_mssql_limit_clause(sql: str) -> str: return re.sub( r"\bSELECT\s+(DISTINCT\s+)?", - lambda match: f"{match.group(0)}TOP ({limit}) ", + lambda match: f"{match.group(0)}TOP {limit} ", without_limit, count=1, flags=re.IGNORECASE, @@ -1484,21 +1484,8 @@ def _normalize_identifier_quote_syntax(sql: str) -> str: return normalized -def _normalize_mssql_top_clause(sql: str) -> str: - return re.sub( - r"\bSELECT\s+(DISTINCT\s+)?TOP\s+(\d+)\b", - lambda match: ( - f"SELECT {match.group(1) or ''}TOP ({match.group(2)})" - ), - sql, - count=1, - flags=re.IGNORECASE, - ) - - def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: normalized = _normalize_identifier_quote_syntax(sql) - normalized = _normalize_mssql_top_clause(normalized) normalized_data_source = normalize_data_source(data_source) if normalized_data_source == "MSSQL": @@ -2429,15 +2416,6 @@ def construct_semantic_schema_contract( "Do not substitute identifiers for metrics, entities for identifiers, or COUNT(*) " "for a requested business measure unless the semantic intent explicitly requests a record count." ) - lines.append( - "Ranking requirement: for top/bottom/ranked questions, include ORDER BY on the mapped metric " - "and a row limit such as LIMIT N, FETCH FIRST N ROWS ONLY, or MSSQL/Wren-safe TOP (N)." - ) - lines.append( - "Schema safety requirement: use only tables and columns present in this contract or the " - "retrieved DATABASE SCHEMA. Do not add unrequested date filters or common timestamp columns " - "such as created_at unless they are explicitly listed and the user asked for a time constraint." - ) return "\n".join(lines) @@ -4189,80 +4167,6 @@ def _semantic_analysis_dict_items( return [item for item in value if isinstance(item, dict)] -_SEMANTIC_CONCEPT_STOPWORDS = { - "a", - "an", - "and", - "as", - "by", - "for", - "from", - "in", - "of", - "on", - "or", - "the", - "to", - "with", - "dbo", - "tbl", - "table", - "view", - "dim", - "fact", -} - -_SEMANTIC_CONCEPT_SYNONYMS = { - "acct": {"account", "customer"}, - "account": {"acct", "customer", "client"}, - "accounts": {"acct", "account", "customer", "client"}, - "amt": {"amount", "value", "total"}, - "amount": {"amt", "value", "total"}, - "bill": {"invoice"}, - "billing": {"invoice"}, - "client": {"account", "customer"}, - "clients": {"account", "customer"}, - "cust": {"customer", "client", "account"}, - "customer": {"cust", "client", "account"}, - "customers": {"cust", "client", "account", "customer"}, - "desc": {"description", "name"}, - "description": {"desc", "name"}, - "inv": {"invoice"}, - "invoice": {"inv", "bill", "billing"}, - "invoices": {"inv", "invoice", "bill", "billing"}, - "name": {"description", "label"}, - "no": {"number", "identifier", "id"}, - "num": {"number", "identifier", "id"}, - "number": {"no", "num", "identifier", "id"}, - "qty": {"quantity"}, - "quantity": {"qty"}, - "total": {"amount", "value", "sum"}, - "value": {"amount", "total"}, -} - - -def _semantic_concept_tokens(value: Any) -> set[str]: - text = str(value or "") - if not text: - return set() - - text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) - text = re.sub(r"[^A-Za-z0-9]+", " ", text) - tokens = { - token.lower() - for token in text.split() - if len(token) > 1 and token.lower() not in _SEMANTIC_CONCEPT_STOPWORDS - } - for token in list(tokens): - if token.endswith("ies") and len(token) > 4: - tokens.add(f"{token[:-3]}y") - elif token.endswith("s") and len(token) > 3: - tokens.add(token[:-1]) - for token in list(tokens): - tokens.update(_SEMANTIC_CONCEPT_SYNONYMS.get(token, set())) - return tokens - - def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: if not isinstance(semantic_analysis, dict) or not semantic_analysis: return False @@ -4386,42 +4290,15 @@ def _semantic_candidate_support_error( def _required_concept_mapping_support_error( semantic_analysis: dict[str, Any], ) -> str | None: - schema_bound_concept_types = { - "dimension", - "entity", - "filter", - "identifier", - "metric", - "relationship", - "time", - } - mapped_schema_objects = [ - schema_object - for mapping in _semantic_concept_mappings(semantic_analysis) - for schema_object in _mapping_schema_objects(mapping) - ] - mapped_schema_objects.extend( - _semantic_analysis_items(semantic_analysis, "supported_schema_objects") - ) - unsupported_required_concepts = [] for mapping in _semantic_concept_mappings(semantic_analysis): if mapping.get("required_in_sql") is False: continue - concept_type = _mapping_concept_type(mapping) - if concept_type not in schema_bound_concept_types: - continue if _mapping_schema_objects(mapping): continue request_concept = _mapping_request_concept(mapping) - concept_tokens = _semantic_concept_tokens(request_concept) - if concept_type == "entity" and concept_tokens: - if any( - concept_tokens & _semantic_concept_tokens(schema_object) - for schema_object in mapped_schema_objects - ): - continue + concept_type = _mapping_concept_type(mapping) if request_concept: unsupported_required_concepts.append( f"{request_concept} ({concept_type or 'concept'})" @@ -4711,28 +4588,11 @@ def _validate_sql_against_concept_mappings( ) has_limit = bool( re.search( - r"\b(?:LIMIT|TOP\s*(?:\(\s*)?\d+|FETCH\s+FIRST)\b", + r"\b(?:LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", sql or "", flags=re.IGNORECASE, ) ) - schema_bound_concept_types = { - "dimension", - "entity", - "filter", - "identifier", - "metric", - "relationship", - "time", - } - mapped_schema_objects = [ - schema_object - for mapping in mappings - for schema_object in _mapping_schema_objects(mapping) - ] - mapped_schema_objects.extend( - _semantic_analysis_items(semantic_analysis, "supported_schema_objects") - ) for mapping in mappings: if mapping.get("required_in_sql") is False: @@ -4742,15 +4602,6 @@ def _validate_sql_against_concept_mappings( request_concept = _mapping_request_concept(mapping) schema_objects = _mapping_schema_objects(mapping) if not schema_objects: - if concept_type not in schema_bound_concept_types: - continue - concept_tokens = _semantic_concept_tokens(request_concept) - if concept_type == "entity" and concept_tokens: - if any( - concept_tokens & _semantic_concept_tokens(schema_object) - for schema_object in mapped_schema_objects - ): - continue return ( "The semantic analysis did not map the required " f"{concept_type or 'concept'} '{request_concept}' to a schema " @@ -4899,28 +4750,20 @@ def _validate_sql_against_semantic_analysis( "dimensions or time grain identified in the semantic analysis." ) - if ranking or sorting: - has_ranking_order = bool( - re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE) - ) - has_ranking_limit = bool( - re.search( - r"\b(?:LIMIT|TOP\s*(?:\(\s*)?\d+|FETCH\s+FIRST)\b", - sql or "", - flags=re.IGNORECASE, - ) - ) - else: - has_ranking_order = True - has_ranking_limit = True - - if (ranking or sorting) and not has_ranking_order: + has_sorting_order = bool( + re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE) + ) + if sorting and not has_sorting_order: return ( "Generated SQL does not include sorting logic required " "by the ranking intent or sorting intent." ) - if ranking and not has_ranking_limit: + if ranking and not re.search( + r"\b(?:ORDER\s+BY|LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", + sql or "", + flags=re.IGNORECASE, + ): return ( "Generated SQL does not include sorting or limiting logic required " "by the ranking intent." diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e30119374b..7bafb7612e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -49,18 +49,15 @@ 10. For each selected table, provide a concise reason for why the table is semantically relevant. 11. For each selected column, provide a concise reason for why the column is necessary. 12. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. -13. Broad business entities can be satisfied by the selected table plus the best descriptive or identifier column for that entity. For example, an entity such as customer may map to a customer name, customer number, account, client, or similar descriptive/identifier column when that is the active schema's representation. Do not leave an entity unmapped when a selected dimension or identifier column represents it. -14. When the user asks for a top/bottom/ranking query, map the ranked dimension and ranked metric separately. The SQL generator must be able to ORDER BY the metric and limit rows. -15. Do not add filters or time constraints that are not requested or implied by the user. Only map date/time concepts when the user asks for a time period, trend, date filter, or date dimension. -16. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. -17. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/chart/dashboard/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. -18. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. -19. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. -20. If a "." is included in columns, put the name before the first dot into chosen columns. -21. The number of columns chosen must match the number of reasoning. -22. Final chosen columns must be only column names, don't prefix it with table names. -23. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -24. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +13. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. +14. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/chart/dashboard/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. +15. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. +16. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. +17. If a "." is included in columns, put the name before the first dot into chosen columns. +18. The number of columns chosen must match the number of reasoning. +19. Final chosen columns must be only column names, don't prefix it with table names. +20. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +21. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -373,34 +370,6 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: "top", } -_GENERIC_SEMANTIC_SYNONYMS = { - "acct": {"account", "customer"}, - "account": {"acct", "customer", "client"}, - "accounts": {"acct", "account", "customer", "client"}, - "amt": {"amount", "value", "total"}, - "amount": {"amt", "value", "total"}, - "bill": {"invoice"}, - "billing": {"invoice"}, - "client": {"account", "customer"}, - "clients": {"account", "customer"}, - "cust": {"customer", "client", "account"}, - "customer": {"cust", "client", "account"}, - "customers": {"cust", "client", "account", "customer"}, - "desc": {"description", "name"}, - "description": {"desc", "name"}, - "inv": {"invoice"}, - "invoice": {"inv", "bill", "billing"}, - "invoices": {"inv", "invoice", "bill", "billing"}, - "name": {"description", "label"}, - "no": {"number", "identifier", "id"}, - "num": {"number", "identifier", "id"}, - "number": {"no", "num", "identifier", "id"}, - "qty": {"quantity"}, - "quantity": {"qty"}, - "total": {"amount", "value", "sum"}, - "value": {"amount", "total"}, -} - def _semantic_tokens(value: Any) -> set[str]: text = str(value or "") @@ -419,8 +388,6 @@ def _semantic_tokens(value: Any) -> set[str]: tokens.add(f"{token[:-3]}y") elif token.endswith("s") and len(token) > 3: tokens.add(token[:-1]) - for token in list(tokens): - tokens.update(_GENERIC_SEMANTIC_SYNONYMS.get(token, set())) return tokens diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index df8dd53825..3a8c8dd861 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1032,15 +1032,6 @@ def test_normalize_generation_result_sql_rewrites_bare_month_field_for_mssql(): assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized -def test_normalize_generation_result_sql_rewrites_top_limit_with_parentheses(): - normalized = normalize_generation_result_sql( - 'SELECT TOP 10 "sales"."CustName" FROM "sales"', - data_source="MSSQL", - ) - - assert normalized.startswith('SELECT TOP (10) "sales"."CustName"') - - def test_normalize_generation_result_sql_rewrites_qualified_month_field_for_mssql(): sql = """ SELECT @@ -1727,37 +1718,6 @@ def test_get_schema_intent_analysis_error_rejects_required_unmapped_concepts(): assert "did not map required request concepts" in error -def test_get_schema_intent_analysis_error_allows_entity_covered_by_dimension_mapping(): - error = get_schema_intent_analysis_error( - { - "concept_mappings": [ - { - "request_concept": "customer", - "concept_type": "entity", - "schema_objects": [], - "required_in_sql": True, - }, - { - "request_concept": "customer name", - "concept_type": "dimension", - "schema_objects": ["sales.CustName"], - "required_in_sql": True, - }, - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["sales.InvAmt"], - "required_in_sql": True, - }, - ], - "supported_schema_objects": ["sales.CustName", "sales.InvAmt"], - "is_fully_supported": True, - } - ) - - assert error is None - - def test_validate_sql_intent_alignment_uses_semantic_analysis_for_metric_count_mismatch(): error = validate_sql_intent_alignment( "Show invoice amount by customer", @@ -2032,86 +1992,6 @@ def test_validate_sql_intent_alignment_rejects_mapped_ranking_without_limit(): assert "sorting and limiting" in error -def test_validate_sql_intent_alignment_accepts_top_without_parentheses_for_ranking(): - error = validate_sql_intent_alignment( - "Top customers by invoice amount", - 'SELECT TOP 10 "sales"."CustName", SUM("sales"."InvAmt") AS "invoice_amount" ' - 'FROM "sales" GROUP BY "sales"."CustName" ORDER BY SUM("sales"."InvAmt") DESC', - {"sales": ["CustName", "InvAmt"]}, - semantic_analysis={ - "analytical_intent": "ranking", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "ranking": ["top 10 customers by invoice amount"], - "aggregations": ["sum invoice amount"], - "concept_mappings": [ - { - "request_concept": "customer", - "concept_type": "entity", - "schema_objects": [], - "required_in_sql": True, - }, - { - "request_concept": "customer name", - "concept_type": "dimension", - "schema_objects": ["sales.CustName"], - "required_in_sql": True, - }, - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["sales.InvAmt"], - "required_in_sql": True, - }, - { - "request_concept": "top 10", - "concept_type": "ranking", - "schema_objects": [], - "required_in_sql": True, - }, - ], - "supported_schema_objects": ["sales.CustName", "sales.InvAmt"], - "is_fully_supported": True, - }, - ) - - assert error is None - - -def test_validate_sql_intent_alignment_rejects_ranking_without_limit_even_with_order(): - error = validate_sql_intent_alignment( - "Top customers by invoice amount", - 'SELECT "sales"."CustName", SUM("sales"."InvAmt") AS "invoice_amount" ' - 'FROM "sales" GROUP BY "sales"."CustName" ORDER BY SUM("sales"."InvAmt") DESC', - {"sales": ["CustName", "InvAmt"]}, - semantic_analysis={ - "analytical_intent": "ranking", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "ranking": ["top 10 customers by invoice amount"], - "aggregations": ["sum invoice amount"], - "concept_mappings": [ - { - "request_concept": "customer name", - "concept_type": "dimension", - "schema_objects": ["sales.CustName"], - "required_in_sql": True, - }, - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["sales.InvAmt"], - "required_in_sql": True, - }, - ], - "is_fully_supported": True, - }, - ) - - assert error is not None - assert "sorting or limiting" in error - - def test_get_schema_intent_analysis_error_rejects_multiple_selected_interpretations(): error = get_schema_intent_analysis_error( { diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 03f7330392..a70f26aecb 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -85,48 +85,6 @@ def test_rank_semantic_schema_candidates_prefers_complete_business_concept_cover assert candidates[0]["confidence"] > candidates[1]["confidence"] -def test_rank_semantic_schema_candidates_matches_generic_abbreviations(): - candidates = rank_semantic_schema_candidates( - query="Show the top 10 customers by invoice amount", - construct_db_schemas=[ - { - "type": "TABLE", - "name": "sales_summary", - "comment": "", - "columns": [ - { - "name": "CustName", - "data_type": "varchar", - "comment": "customer display name", - }, - { - "name": "InvAmt", - "data_type": "decimal", - "comment": "invoice amount", - }, - ], - }, - { - "type": "TABLE", - "name": "refund_summary", - "comment": "", - "columns": [ - { - "name": "Refund_Amount", - "data_type": "decimal", - "comment": "refund amount", - } - ], - }, - ], - ) - - assert candidates[0]["table_name"] == "sales_summary" - assert {"customer", "invoice", "amount"} <= set( - candidates[0]["matched_query_terms"] - ) - - def test_rank_semantic_schema_candidates_penalizes_retry_rejected_schema_objects(): candidates = rank_semantic_schema_candidates( query="Show top 10 customers by invoice amount", From 9e38a3d15ec230614439f16a13c01084b4aa87bb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:57:22 +0530 Subject: [PATCH 0481/1087] Revert "Improve semantic NL-to-SQL planning and validation" This reverts commit fab63187a129d9b051e6256201b0e27fef9a8092. --- .../generation/followup_sql_generation.py | 20 +- .../followup_sql_generation_reasoning.py | 22 +- .../pipelines/generation/sql_correction.py | 10 +- .../pipelines/generation/sql_generation.py | 23 +- .../generation/sql_generation_reasoning.py | 26 +- .../pipelines/generation/sql_regeneration.py | 14 +- .../src/pipelines/generation/utils/sql.py | 629 +----------------- .../retrieval/db_schema_retrieval.py | 48 +- .../pipelines/generation/test_sql_utils.py | 121 ---- 9 files changed, 41 insertions(+), 872 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 b5772cab37..f0f2412891 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -2,7 +2,6 @@ import logging import sys -from datetime import datetime from typing import TYPE_CHECKING, Any from hamilton import base @@ -100,7 +99,6 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Current Time: {{ current_time }} {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### @@ -114,21 +112,10 @@ appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. Only apply aggregate functions to columns whose active metadata type supports that operation. -Reuse the previous SQL and summary for unresolved references such as same, that, -those, previous, it, filter, sort, period, table, chart, or metric. Apply only the -requested follow-up change instead of regenerating the entire analysis from scratch. -Resolve business synonyms and different phrasings through active metadata, foreign -keys, and semantic relationships. For multi-table follow-ups, connect tables with -explicit INNER JOIN or LEFT JOIN conditions from trusted relationships only. -Resolve relative date phrases against Current Time. Infer aggregation, sorting, -limits, and chart/dashboard datasets from the follow-up wording. -Avoid SELECT *. Select only required columns, keep existing filters unless the user -changes them, push new filters before joins where possible, and limit rows for -retrieval/ranking requests. Before writing SQL, validate that the selected schema elements directly support every -key entity, measure, dimension, filter, time range, join, sorting, chart, dashboard, -and aggregation in the follow-up question. If the schema cannot support the requested -information, do not replace the request with a generic COUNT(*) or unrelated table query. +key entity, metric, dimension, filter, time range, relationship, and aggregation in +the follow-up question. If the schema cannot support the requested information, do +not replace the request with a generic COUNT(*) or unrelated table query. ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -156,7 +143,6 @@ def prompt( ) -> dict: _prompt = prompt_builder.run( query=query, - current_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), data_source=data_source, documents=documents, valid_table_names=construct_valid_table_names(documents), diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index a3999f3bc6..0374caf609 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -70,27 +70,13 @@ {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### -Use this contract for entities, measures, dimensions, filters, joins, normalized -date ranges, aggregations, sorting, ranking, chart requirements, dashboard -requirements, and analytical intent. If it shows missing or ambiguous requirements, -state that limitation in the plan instead of planning unrelated SQL. +Use this contract for entities, metrics, dimensions, filters, joins, time +constraints, aggregations, ranking, and analytical intent. If it shows missing or +ambiguous requirements, state that limitation in the plan instead of planning +unrelated SQL. {{ semantic_schema_contract }} {% endif %} -### FOLLOW-UP PLANNING REQUIREMENTS ### -Use query history to resolve references such as same, that, previous, it, those, -them, filter, period, metric, table, result, or chart. Reuse the previous schema -context and apply only the requested change unless the user asks for a new analysis. -Resolve synonyms from active metadata, metrics, views, foreign keys, and semantic -relationships. Identify required entities, measures, dimensions, filters, date -windows, sort direction, top/bottom limits, joins, and chart/dashboard dataset -shape. -For multi-table follow-ups, choose a trusted join path from foreign keys or semantic -relationships and include the join keys. If no path exists, say clarification or -schema support is required. -Resolve relative dates against Current Time. Plan only required columns and avoid -SELECT *. - Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index d0f20521cd..d1e2e93333 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -122,14 +122,10 @@ def get_sql_correction_system_prompt( user's request. Do not invent tables, columns, joins, metrics, or relationships. Only apply aggregate functions to columns whose active metadata type supports that operation. -Preserve the user's business intent, including entities, measures, dimensions, -filters, date ranges, joins, sorting, top/bottom limits, chart/dashboard dataset -shape, and aggregation. Fix obvious table/column/alias/GROUP BY/JOIN mistakes only -when the active metadata supports the correction. Before returning corrected SQL, validate that it still directly supports every key -entity, measure, dimension, filter, time range, relationship, sorting requirement, -chart/dashboard requirement, and aggregation in the user's question. Do not replace -an unsupported request with a generic COUNT(*) or unrelated table query. +entity, metric, dimension, filter, time range, relationship, and aggregation in the +user's question. Do not replace an unsupported request with a generic COUNT(*) or +unrelated table query. Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 11a2646744..9c568abec0 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -1,6 +1,5 @@ import logging import sys -from datetime import datetime from typing import Any from hamilton import base @@ -88,7 +87,6 @@ ### QUESTION ### User's Question: {{ query }} -Current Time: {{ current_time }} {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### @@ -101,24 +99,10 @@ general guidance when the question can be answered with SQL over the active metadata. Never reuse table or column names from SQL SAMPLES unless those exact names also appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. -Resolve business synonyms and different phrasings through schema names, column names, -descriptions, metadata, metrics, views, foreign keys, and semantic relationships. -For multi-table questions, choose only the tables needed and connect them with -explicit INNER JOIN or LEFT JOIN conditions from foreign keys or semantic -relationships. If no trustworthy relationship exists, do not invent a join. -Resolve relative date phrases against Current Time and express them as concrete -filter predicates on real temporal columns. -Infer aggregation from language: total/sum -> SUM when a measure is requested, -average -> AVG, number/count/how many -> COUNT, min/max/highest/lowest -> MIN/MAX -or ranking as appropriate. -For chart or dashboard requests, return the optimized aggregated dataset needed -for the chart/KPI/summary, with suitable grouping, sorting, and limits. -Avoid SELECT *. Select only required columns, push filters before joins where -possible, and limit rows for retrieval/ranking requests. Before writing SQL, validate that the selected schema elements directly support every -key entity, measure, dimension, filter, time range, join, sorting, chart, dashboard, -and aggregation in the question. If the schema cannot support the requested -information, do not replace the request with a generic COUNT(*) or unrelated table query. +key entity, metric, dimension, filter, time range, relationship, and aggregation in +the question. If the schema cannot support the requested information, do not replace +the request with a generic COUNT(*) or unrelated table query. {% if sql_generation_reasoning %} ### REASONING PLAN ### @@ -161,7 +145,6 @@ def prompt( ) _prompt = prompt_builder.run( query=query, - current_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), data_source=data_source, documents=documents, has_pcb_context=has_pcb_context, 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 ef9486360a..7d14e29184 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -55,31 +55,13 @@ {% if semantic_schema_contract %} ### SEMANTIC SCHEMA CONTRACT ### -Use this contract for entities, measures, dimensions, filters, joins, normalized -date ranges, aggregations, sorting, ranking, chart requirements, dashboard -requirements, and analytical intent. If it shows missing or ambiguous requirements, -state that limitation in the plan instead of planning unrelated SQL. +Use this contract for entities, metrics, dimensions, filters, joins, time +constraints, aggregations, ranking, and analytical intent. If it shows missing or +ambiguous requirements, state that limitation in the plan instead of planning +unrelated SQL. {{ semantic_schema_contract }} {% endif %} -### PLANNING REQUIREMENTS ### -Build a schema-grounded SQL plan, not example-specific SQL. Resolve synonyms from -table names, column names, descriptions, metadata, metrics, views, foreign keys, -and semantic relationships. Identify the exact business entities, measures, -dimensions, filters, date windows, sort direction, top/bottom limits, joins, and -chart/dashboard requirements needed by the question. -For multi-table questions, choose a trusted join path from foreign keys or semantic -relationships and include the join keys. If no path exists, say clarification or -schema support is required. -Resolve relative dates such as today, yesterday, this/last week, this/last month, -this/last quarter, this/last year, last 30 days, last 90 days, and rolling 12 -months against Current Time. -Infer aggregation from language: total/sum -> SUM for measures, average -> AVG, -number/count/how many -> COUNT, min/max -> MIN/MAX, top/bottom/highest/lowest -> -ORDER BY with a limit. -Plan only required columns. Avoid SELECT *. Preserve conversational context for -follow-up wording and apply only the requested change. - Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 87ae969af6..4f0668ff0f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -46,11 +46,9 @@ def get_sql_regeneration_system_prompt( While generating the new SQL query, you should use the original SQL query as a reference. While generating the new SQL query, make sure to use the database schema to generate the SQL query. Before returning SQL, validate that the selected schema elements directly support -the key entities, measures, dimensions, filters, time ranges, relationships, -sorting, chart/dashboard requirements, and aggregations from the user's question -or reasoning. Reuse the previous semantic context and original SQL where correct, -but fix schema, join, alias, GROUP BY, aggregation, and filter mistakes. Do not -replace an unsupported request with a generic COUNT(*) or unrelated table query. +the key entities, metrics, dimensions, filters, time ranges, relationships, and +aggregations from the user's question or reasoning. Do not replace an unsupported +request with a generic COUNT(*) or unrelated table query. {text_to_sql_rules} @@ -121,12 +119,6 @@ def get_sql_regeneration_system_prompt( SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} -### REGENERATION REQUIREMENTS ### -Preserve the business question and semantic schema contract. Use active metadata -to correct only the parts that are wrong: tables, columns, aliases, joins, filters, -GROUP BY, HAVING, ORDER BY, aggregation, limits, date predicates, and chart/dashboard -dataset shape. Avoid SELECT * and do not add unrelated columns or tables. - Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 896db1848d..2062de554d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2218,7 +2218,7 @@ def get_sql_generation_system_prompt( 9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. 10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. 11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. -12. Before generating SQL, validate that the selected schema elements directly support all key entities, measures, dimensions, filters, sorting, chart shape, time ranges, relationships, and aggregations mentioned or implied by the question. +12. Before generating SQL, validate that the selected schema elements directly support all key entities, metrics, dimensions, filters, time ranges, relationships, and aggregations mentioned or implied by the question. 13. Do not answer a specific business metric, trend, summary, comparison, dashboard, or analysis request with a generic record-count query unless the user explicitly asks only for record count. 14. If the required information cannot be derived from the available active schema, return the closest schema-grounded limitation instead of inventing unrelated SQL. 15. If a SEMANTIC SCHEMA CONTRACT is provided, it is the primary source of truth for selecting tables, columns, metrics, joins, filters, grouping, sorting, and date logic. Generate SQL from the highest-confidence validated concept-to-schema mappings in that contract and do not independently infer substitute schema objects. @@ -2331,14 +2331,9 @@ def construct_semantic_schema_contract( ("Dimensions", "dimensions"), ("Filters", "filters"), ("Time constraints", "time_constraints"), - ("Date ranges", "date_ranges"), ("Aggregations", "aggregations"), ("Ranking", "ranking"), - ("Sorting", "sorting"), - ("Chart requirements", "chart_requirements"), - ("Dashboard requirements", "dashboard_requirements"), ("Relationships", "relationships"), - ("Join paths", "join_paths"), ("Supported schema objects", "supported_schema_objects"), ): lines.extend(_format_semantic_list(label, _semantic_analysis_items(semantic_analysis, key))) @@ -3436,536 +3431,6 @@ def format_valid_table_columns(valid_table_columns: dict[str, list[str]]) -> str "department", "location", } -_BUSINESS_INTENT_STOP_WORDS = { - "a", - "an", - "and", - "are", - "as", - "at", - "be", - "by", - "chart", - "create", - "data", - "each", - "for", - "from", - "give", - "graph", - "how", - "in", - "is", - "me", - "of", - "on", - "or", - "please", - "show", - "table", - "the", - "to", - "with", -} -_SCHEMA_CONCEPT_GENERIC_TOKENS = { - "amount", - "at", - "code", - "count", - "date", - "day", - "description", - "dt", - "id", - "ids", - "key", - "month", - "name", - "no", - "number", - "num", - "pct", - "percent", - "percentage", - "price", - "quantity", - "qty", - "rate", - "time", - "timestamp", - "total", - "type", - "value", - "year", -} -_TEMPORAL_SCHEMA_TOKENS = { - "at", - "created", - "date", - "day", - "ended", - "modified", - "month", - "started", - "time", - "timestamp", - "updated", - "week", - "year", -} -_CHART_TERMS = { - "area chart", - "bar chart", - "chart", - "donut chart", - "graph", - "line chart", - "pie chart", - "plot", - "scatter plot", - "visualization", -} -_RANKING_TERMS = { - "bottom", - "highest", - "largest", - "lowest", - "most", - "rank", - "ranked", - "ranking", - "smallest", - "top", -} -_SORTING_TERMS = { - "ascending", - "descending", - "order by", - "ordered by", - "sort", - "sorted", -} - - -def _singularize_business_token(token: str) -> str: - token = token.lower() - if len(token) > 4 and token.endswith("ies"): - return f"{token[:-3]}y" - if len(token) > 3 and token.endswith("es"): - return token[:-2] - if len(token) > 3 and token.endswith("s"): - return token[:-1] - return token - - -def _business_intent_tokens(text: str | None) -> list[str]: - spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(text or "")) - return [ - _singularize_business_token(token) - for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) - if token and token not in _BUSINESS_INTENT_STOP_WORDS - ] - - -def _schema_column_concept_keys(column_name: str) -> set[tuple[str, ...]]: - tokens = tuple(_business_intent_tokens(column_name)) - if not tokens: - return set() - - keys: set[tuple[str, ...]] = {tokens} - token_set = set(tokens) - informative_tokens = tuple( - token for token in tokens if token not in _SCHEMA_CONCEPT_GENERIC_TOKENS - ) - metric_tokens = { - "amount", - "cost", - "count", - "price", - "profit", - "quantity", - "qty", - "rate", - "revenue", - "sales", - "total", - "value", - } - identifier_suffix_tokens = {"code", "id", "ids", "key", "no", "number"} - if informative_tokens and ( - len(informative_tokens) > 1 - or not token_set.intersection(metric_tokens) - and ( - not token_set.intersection(identifier_suffix_tokens) - or bool(set(informative_tokens).intersection(_DIMENSION_TERMS)) - ) - ): - keys.add(informative_tokens) - - if len(tokens) > 1: - without_suffix = tuple( - token - for token in tokens - if token not in {"id", "ids", "key", "name", "number", "no", "code"} - ) - if without_suffix and ( - len(without_suffix) > 1 - or bool(set(without_suffix).intersection(_DIMENSION_TERMS)) - ): - keys.add(without_suffix) - - for token in tokens: - if token in _SCHEMA_CONCEPT_GENERIC_TOKENS or token in _DIMENSION_TERMS: - keys.add((token,)) - - return {key for key in keys if key} - - -def _query_mentions_concept( - concept_tokens: tuple[str, ...], - query_tokens: set[str], - asks_time_analysis: bool, -) -> bool: - if not concept_tokens: - return False - - if set(concept_tokens).issubset(query_tokens): - return True - - if set(concept_tokens).intersection({"code", "id", "ids", "key", "no", "number"}): - return False - - meaningful_tokens = [ - token - for token in concept_tokens - if token not in _SCHEMA_CONCEPT_GENERIC_TOKENS - ] - if len(meaningful_tokens) != 1: - return False - - token = meaningful_tokens[0] - if token not in query_tokens: - return False - - if set(concept_tokens) & _TEMPORAL_SCHEMA_TOKENS and not asks_time_analysis: - return False - - return len(token) >= 4 or token in _DIMENSION_TERMS - - -def _infer_query_schema_requirements( - query: str, - valid_table_columns: dict[str, list[str]], -) -> dict[str, set[str]]: - query_tokens = set(_business_intent_tokens(query)) - if not query_tokens: - return {} - - asks_time_analysis = _query_requests_time_analysis(query) - requirements: dict[str, set[str]] = {} - - for table_name, columns in valid_table_columns.items(): - if not table_name: - continue - table_tokens = tuple( - token - for token in _business_intent_tokens(table_name) - if token not in {"dbo", "public", "schema", "table", "tbl"} - ) - table_informative_tokens = tuple( - token - for token in table_tokens - if token not in _SCHEMA_CONCEPT_GENERIC_TOKENS - ) - for concept_tokens in {table_tokens, table_informative_tokens}: - if _query_mentions_concept( - concept_tokens, - query_tokens, - asks_time_analysis, - ): - requirements.setdefault(" ".join(concept_tokens), set()).add( - str(table_name) - ) - - for column_name in columns or []: - if not column_name: - continue - for concept_tokens in _schema_column_concept_keys(str(column_name)): - if not _query_mentions_concept( - concept_tokens, - query_tokens, - asks_time_analysis, - ): - continue - concept = " ".join(concept_tokens) - requirements.setdefault(concept, set()).add( - f"{table_name}.{column_name}" - ) - - return { - concept: objects - for concept, objects in requirements.items() - if objects and not _is_ambiguous_schema_concept(concept, objects) - } - - -def _is_ambiguous_schema_concept(concept: str, schema_objects: set[str]) -> bool: - if len(schema_objects) <= 1: - return False - - concept_compact = _compact_sql_identifier(concept) - column_names = { - _schema_object_parts(schema_object)[-1] - for schema_object in schema_objects - if _schema_object_parts(schema_object) - } - compact_columns = { - _compact_sql_identifier(column_name) for column_name in column_names - } - - if concept_compact in compact_columns: - return False - - return len(compact_columns) > 1 - - -def _sql_references_any_schema_object( - sql: str, - schema_objects: set[str], - valid_table_columns: dict[str, list[str]], -) -> bool: - return any( - _sql_references_schema_object(sql, schema_object, valid_table_columns) - for schema_object in schema_objects - ) - - -def _sql_has_ordering(sql: str) -> bool: - return bool(re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE)) - - -def _sql_has_limit(sql: str) -> bool: - return bool( - re.search( - r"\b(?:LIMIT\s+\d+|TOP\s*\(?\s*\d+\s*\)?|FETCH\s+FIRST\s+\d+)\b", - sql or "", - flags=re.IGNORECASE, - ) - ) - - -def _sql_has_filter_predicate(sql: str) -> bool: - return bool( - re.search(r"\b(?:WHERE|HAVING)\b", sql or "", flags=re.IGNORECASE) - ) - - -def _query_requests_ranking(query: str) -> bool: - return _contains_phrase(query, _RANKING_TERMS) - - -def _query_requests_sorting(query: str) -> bool: - return _query_requests_ranking(query) or _contains_phrase(query, _SORTING_TERMS) - - -def _query_requests_limited_ranking(query: str) -> bool: - normalized = str(query or "").lower() - return bool( - re.search(r"\b(?:top|bottom)\s+\d+\b", normalized) - or re.search(r"\b(?:top|bottom|highest|lowest|largest|smallest)\b", normalized) - ) - - -def _query_requests_chart(query: str) -> bool: - return _contains_phrase(query, _CHART_TERMS) - - -def _query_requests_time_filter(query: str) -> bool: - normalized = str(query or "").lower() - return bool( - re.search( - r"\b(?:last|next|previous|prior|this)\s+" - r"(?:\d+\s+)?(?:day|week|month|quarter|year)s?\b", - normalized, - ) - or re.search(r"\b(?:between|since|before|after)\b", normalized) - ) - - -def _query_requests_literal_filter(query: str) -> bool: - normalized = str(query or "").lower() - if re.search(r"'[^']+'|\"[^\"]+\"", query or ""): - return True - return bool( - re.search( - r"\b(?:where|only|exclude|excluding|filtered by|filter by|with status|" - r"with category|for status|for category|for region|for country|" - r"for market|for customer|for product)\b", - normalized, - ) - ) - - -def _requested_query_aggregate_functions(query: str) -> set[str]: - normalized = str(query or "").lower() - functions: set[str] = set() - - if re.search(r"\b(?:avg|average|mean)\b", normalized): - functions.add("AVG") - if re.search(r"\b(?:min|minimum)\b", normalized): - functions.add("MIN") - if re.search(r"\b(?:max|maximum)\b", normalized): - functions.add("MAX") - if re.search(r"\b(?:count|number of|how many|record count|row count)\b", normalized): - functions.add("COUNT") - elif re.search(r"\bsum\b", normalized) or ( - re.search(r"\btotal\b", normalized) - and _query_requests_specific_metric(normalized) - ): - functions.add("SUM") - - return functions - - -def _simple_groupable_expression_key(expression: str) -> str | None: - expression = _strip_projection_alias( - re.sub(r"^\s*DISTINCT\s+", "", expression, flags=re.IGNORECASE) - ).strip() - expression = re.sub( - r"^\s*TOP\s*\(?\s*\d+\s*\)?\s+", - "", - expression, - flags=re.IGNORECASE, - ) - if not expression or expression == "*": - return None - if _AGGREGATE_PATTERN.search(expression): - return None - if re.search(r"\b(?:CASE|SELECT|OVER)\b", expression, flags=re.IGNORECASE): - return None - - return re.sub(r"\s+", "", expression).strip().lower() - - -def _validate_group_by_for_aggregates(sql: str) -> str | None: - if not _AGGREGATE_PATTERN.search(sql or ""): - return None - - select_spans = _find_select_list_spans(sql) - if not select_spans: - return None - - group_bodies = _extract_clause_bodies( - sql, - r"GROUP\s+BY", - ["HAVING", r"ORDER\s+BY", "LIMIT", "FETCH", "UNION"], - ) - if not group_bodies: - return None - - grouped_keys = { - key - for body in group_bodies - for item in _split_top_level_select_items(body) - if (key := _simple_groupable_expression_key(item)) - } - if not grouped_keys: - return None - - for start, end in select_spans: - for item in _split_top_level_select_items(sql[start:end]): - key = _simple_groupable_expression_key(item) - if not key or key in grouped_keys: - continue - return ( - "Generated SQL selects a non-aggregated expression that is not " - f"present in GROUP BY: {item.strip()}." - ) - - return None - - -def _validate_query_schema_requirements( - query: str, - sql: str, - valid_table_columns: dict[str, list[str]], -) -> str | None: - if group_by_error := _validate_group_by_for_aggregates(sql): - return group_by_error - - schema_requirements = _infer_query_schema_requirements( - query, - valid_table_columns, - ) - for concept, schema_objects in sorted(schema_requirements.items()): - if _sql_references_any_schema_object( - sql, - schema_objects, - valid_table_columns, - ): - continue - return ( - "Generated SQL does not reference schema objects that match the " - f"requested business concept '{concept}': " - f"{', '.join(sorted(schema_objects))}." - ) - - for function_name in _requested_query_aggregate_functions(query): - if function_name == "COUNT" and _sql_is_plain_count(sql): - continue - if not _sql_has_aggregate_function(sql, function_name): - return ( - "Generated SQL does not use the aggregation requested by the " - f"business question: {function_name}." - ) - - asks_ranking = _query_requests_ranking(query) - asks_sorting = _query_requests_sorting(query) - if asks_sorting and not _sql_has_ordering(sql): - return ( - "Generated SQL does not include ORDER BY required by the requested " - "sorting or ranking intent." - ) - - if asks_ranking and _query_requests_limited_ranking(query) and not _sql_has_limit(sql): - return ( - "Generated SQL does not include a limit/TOP clause required by the " - "requested top/bottom ranking intent." - ) - - if ( - (_query_requests_time_filter(query) or _query_requests_literal_filter(query)) - and not _sql_has_filter_predicate(sql) - ): - return ( - "Generated SQL does not include a WHERE or HAVING predicate required " - "by the requested filter or time range." - ) - - if _query_requests_chart(query): - if (asks_ranking or _query_requests_grouped_analysis(query)) and not re.search( - r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE - ): - return ( - "Generated SQL does not group the dataset required for the " - "requested chart dimension." - ) - if asks_ranking and not _sql_has_ordering(sql): - return ( - "Generated SQL does not sort the dataset required for the " - "requested chart ranking." - ) - - if re.search(r"\bJOIN\b", sql or "", flags=re.IGNORECASE): - if re.search(r"\bCROSS\s+JOIN\b", sql or "", flags=re.IGNORECASE): - return None - if not re.search(r"\b(?:ON|USING)\b", sql or "", flags=re.IGNORECASE): - return ( - "Generated SQL joins tables without an ON or USING condition. " - "Use explicit schema relationships for JOINs." - ) - - return None def _contains_phrase(text: str, terms: set[str] | tuple[str, ...]) -> bool: @@ -4180,12 +3645,7 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: "aggregations", "relationships", "time_constraints", - "date_ranges", "ranking", - "sorting", - "chart_requirements", - "dashboard_requirements", - "join_paths", "supported_schema_objects", "candidate_schema_scores", "concept_mappings", @@ -4442,40 +3902,6 @@ def _sql_references_table(sql: str, table_name: str) -> bool: return any(_sql_contains_identifier(sql, candidate) for candidate in table_candidates) -def _sql_references_unqualified_column_expression(sql: str, column_name: str) -> bool: - expected_column = _normalize_sql_identifier(str(column_name or "")) - if not expected_column: - return False - - expected_lower = expected_column.lower() - expected_compact = _compact_sql_identifier(expected_column) - candidate_columns = { - _normalize_sql_identifier(column) - for column in _find_unqualified_column_candidates(sql) - } - - for start, end in _find_select_list_spans(sql): - for item in _split_top_level_select_items(sql[start:end]): - expression = _strip_projection_alias( - re.sub(r"^\s*DISTINCT\s+", "", item, flags=re.IGNORECASE) - ) - searchable_expression = _strip_sql_literals(expression) - for match in re.finditer(_SQL_IDENTIFIER_PATTERN, searchable_expression): - before = searchable_expression[: match.start()].rstrip() - after = searchable_expression[match.end() :].lstrip() - if before.endswith(".") or after.startswith("."): - continue - identifier = _normalize_sql_identifier(match.group(0)) - if identifier and identifier.lower() not in _SQL_NON_COLUMN_IDENTIFIERS: - candidate_columns.add(identifier) - - return any( - candidate.lower() == expected_lower - or _compact_sql_identifier(candidate) == expected_compact - for candidate in candidate_columns - ) - - def _sql_references_schema_object( sql: str, schema_object: str, @@ -4505,7 +3931,7 @@ def _sql_references_schema_object( if _sql_references_table(sql, expected_table) and _sql_contains_identifier( sql, expected_column ): - return _sql_references_unqualified_column_expression(sql, expected_column) + return True return False @@ -4697,16 +4123,7 @@ def _validate_sql_against_semantic_analysis( time_constraints = _semantic_analysis_items( semantic_analysis, "time_constraints" ) - date_ranges = _semantic_analysis_items(semantic_analysis, "date_ranges") ranking = _semantic_analysis_items(semantic_analysis, "ranking") - sorting = _semantic_analysis_items(semantic_analysis, "sorting") - chart_requirements = _semantic_analysis_items( - semantic_analysis, "chart_requirements" - ) - dashboard_requirements = _semantic_analysis_items( - semantic_analysis, "dashboard_requirements" - ) - join_paths = _semantic_analysis_items(semantic_analysis, "join_paths") requests_record_count = _semantic_analysis_requests_record_count( semantic_analysis ) @@ -4731,9 +4148,7 @@ def _validate_sql_against_semantic_analysis( "or analytical calculations from the active schema." ) - if (time_constraints or date_ranges) and not _sql_has_temporal_reference( - sql, valid_table_columns - ): + if time_constraints and not _sql_has_temporal_reference(sql, valid_table_columns): return ( "Generated SQL does not use a temporal field or supported date/time " "expression, but the semantic analysis identified time constraints " @@ -4750,15 +4165,6 @@ def _validate_sql_against_semantic_analysis( "dimensions or time grain identified in the semantic analysis." ) - has_sorting_order = bool( - re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE) - ) - if sorting and not has_sorting_order: - return ( - "Generated SQL does not include sorting logic required " - "by the ranking intent or sorting intent." - ) - if ranking and not re.search( r"\b(?:ORDER\s+BY|LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", sql or "", @@ -4769,27 +4175,6 @@ def _validate_sql_against_semantic_analysis( "by the ranking intent." ) - if join_paths and re.search(r"\bJOIN\b", sql or "", flags=re.IGNORECASE): - if not re.search(r"\b(?:ON|USING)\b", sql or "", flags=re.IGNORECASE): - return ( - "Generated SQL does not use explicit join conditions required " - "by the semantic join path." - ) - - if chart_requirements and _AGGREGATE_PATTERN.search(sql or "") and dimensions: - if not re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE): - return ( - "Generated SQL does not group the aggregated dataset required " - "by the chart requirements." - ) - - if dashboard_requirements and _sql_is_plain_count(sql) and not requests_record_count: - return ( - "Generated SQL answers a dashboard request with only a generic record " - "count. Use the KPI, chart, summary, and insight requirements from " - "the semantic analysis." - ) - if aggregations and not _AGGREGATE_PATTERN.search(sql or ""): return ( "Generated SQL does not include the aggregation required by the " @@ -4857,14 +4242,6 @@ def validate_sql_intent_alignment( "or report that the schema does not expose that dimension." ) - schema_requirement_error = _validate_query_schema_requirements( - normalized_query, - sql, - valid_table_columns, - ) - if schema_requirement_error: - return schema_requirement_error - missing_metric_terms = _missing_metric_support( normalized_query, sql, diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 7bafb7612e..bfbbae7b50 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -37,27 +37,25 @@ The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### -1. First perform a semantic analysis of the user's request. Identify intended business entities, identifiers, descriptive attributes, metrics, dimensions, filters, aggregations, relationships, time constraints, ranking requirements, chart requirements, dashboard/KPI requirements, and analytical intent such as retrieval, detailed records, summary, comparison, trend analysis, dashboard, KPI, ranking, or record count. +1. First perform a semantic analysis of the user's request. Identify intended business entities, identifiers, descriptive attributes, metrics, dimensions, filters, aggregations, relationships, time constraints, ranking requirements, and analytical intent such as retrieval, detailed records, summary, comparison, trend analysis, dashboard, KPI, ranking, or record count. 2. Map each business term to explicit schema objects only when the active schema directly supports that term. Distinguish entities such as customer/order/invoice/product from identifiers such as order ID or invoice number, descriptive attributes, and measurable metrics such as amount, quantity, cost, profit, revenue, or duration. -3. Understand synonyms and different phrasings by using table names, column names, descriptions, metadata, metrics, views, foreign keys, and semantic relationships. Do not rely on hardcoded examples or default tables. -4. Select tables and columns by semantic fit to the full request, not by isolated keyword overlap or commonly used default tables. -5. Include join keys and relationship columns needed to connect selected tables. Prefer explicit foreign keys and semantic relationships. For multi-table questions, identify a join path; if no trustworthy path exists, record the missing relationship instead of selecting unrelated tables. -6. Normalize relative date language such as today, yesterday, this/last week, this/last month, this/last quarter, this/last year, last 30 days, last 90 days, and rolling 12 months into date requirements using Current Time when available. -7. If the schema does not support a requested entity, metric, dimension, filter, time range, aggregation, chart, dashboard, or ranking requirement, record it in `missing_requirements`. -8. If multiple schema interpretations are equally plausible and the question does not disambiguate them, record them in `ambiguous_requirements`. -9. Set `is_fully_supported` to false when any required request component is missing or ambiguous. -10. For each selected table, provide a concise reason for why the table is semantically relevant. -11. For each selected column, provide a concise reason for why the column is necessary. -12. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. -13. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. -14. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/chart/dashboard/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. -15. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. -16. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. -17. If a "." is included in columns, put the name before the first dot into chosen columns. -18. The number of columns chosen must match the number of reasoning. -19. Final chosen columns must be only column names, don't prefix it with table names. -20. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -21. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +3. Select tables and columns by semantic fit to the full request, not by isolated keyword overlap or commonly used default tables. +4. Include join keys and relationship columns needed to connect selected tables. Do not invent relationships or foreign keys. +5. If the schema does not support a requested entity, metric, dimension, filter, time range, aggregation, or ranking requirement, record it in `missing_requirements`. +6. If multiple schema interpretations are equally plausible and the question does not disambiguate them, record them in `ambiguous_requirements`. +7. Set `is_fully_supported` to false when any required request component is missing or ambiguous. +8. For each selected table, provide a concise reason for why the table is semantically relevant. +9. For each selected column, provide a concise reason for why the column is necessary. +10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. +11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. +12. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. +13. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. +14. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. +15. If a "." is included in columns, put the name before the first dot into chosen columns. +16. The number of columns chosen must match the number of reasoning. +17. Final chosen columns must be only column names, don't prefix it with table names. +18. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +19. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -73,12 +71,7 @@ "aggregations": ["aggregation or calculation requirements"], "relationships": ["required joins or relationships"], "time_constraints": ["time filters, grains, or trend requirements"], - "date_ranges": ["normalized relative date requirements such as start/end dates or rolling windows"], "ranking": ["top/bottom/order/limit requirements"], - "sorting": ["sort fields and direction requirements"], - "chart_requirements": ["requested or inferred chart type, x/y encodings, series, and grain"], - "dashboard_requirements": ["requested KPI, chart, summary, and insight sections"], - "join_paths": ["foreign key or semantic relationship path needed to connect selected tables"], "supported_schema_objects": ["table.column or metric names that directly support the request"], "candidate_schema_scores": [ { @@ -1061,12 +1054,7 @@ class SemanticAnalysis(BaseModel): aggregations: list[str] = Field(default_factory=list) relationships: list[str] = Field(default_factory=list) time_constraints: list[str] = Field(default_factory=list) - date_ranges: list[str] = Field(default_factory=list) ranking: list[str] = Field(default_factory=list) - sorting: list[str] = Field(default_factory=list) - chart_requirements: list[str] = Field(default_factory=list) - dashboard_requirements: list[str] = Field(default_factory=list) - join_paths: list[str] = Field(default_factory=list) supported_schema_objects: list[str] = Field(default_factory=list) candidate_schema_scores: list[SemanticCandidateSchemaScore] = Field( default_factory=list diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 3a8c8dd861..d22e45f36c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1772,105 +1772,6 @@ def test_validate_sql_intent_alignment_allows_semantic_count_metric(): assert error is None -def test_validate_sql_intent_alignment_rejects_missing_schema_backed_measure(): - error = validate_sql_intent_alignment( - "Show total invoice amount by customer", - 'SELECT "invoices"."customer_id", SUM("invoices"."tax_amount") ' - 'AS "total_invoice_amount" FROM "invoices" ' - 'GROUP BY "invoices"."customer_id"', - {"invoices": ["customer_id", "invoice_amount", "tax_amount"]}, - ) - - assert error is not None - assert "invoice amount" in error - assert "invoices.invoice_amount" in error - - -def test_validate_sql_intent_alignment_rejects_missing_requested_aggregation(): - error = validate_sql_intent_alignment( - "Show average order value by region", - 'SELECT "orders"."region", SUM("orders"."order_value") AS "order_value" ' - 'FROM "orders" GROUP BY "orders"."region"', - {"orders": ["region", "order_value"]}, - ) - - assert error is not None - assert "AVG" in error - - -def test_validate_sql_intent_alignment_rejects_ranking_without_order_or_limit(): - error = validate_sql_intent_alignment( - "Show top 10 customers by invoice amount", - 'SELECT "invoices"."customer_id", ' - 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' - 'FROM "invoices" GROUP BY "invoices"."customer_id"', - {"invoices": ["customer_id", "invoice_amount"]}, - ) - - assert error is not None - assert "ORDER BY" in error - - -def test_validate_sql_intent_alignment_rejects_time_filter_without_predicate(): - error = validate_sql_intent_alignment( - "Show monthly order count for the last 12 months", - 'SELECT DATEPART(YEAR, "orders"."created_at") AS "year", ' - 'DATEPART(MONTH, "orders"."created_at") AS "month", ' - 'COUNT(*) AS "order_count" FROM "orders" ' - 'GROUP BY DATEPART(YEAR, "orders"."created_at"), ' - 'DATEPART(MONTH, "orders"."created_at")', - {"orders": ["id", "created_at"]}, - ) - - assert error is not None - assert "WHERE or HAVING" in error - - -def test_validate_sql_intent_alignment_rejects_missing_group_by_column(): - error = validate_sql_intent_alignment( - "Show total invoice amount by customer and region", - 'SELECT "invoices"."customer_id", "invoices"."region", ' - 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' - 'FROM "invoices" GROUP BY "invoices"."customer_id"', - {"invoices": ["customer_id", "region", "invoice_amount"]}, - ) - - assert error is not None - assert "not present in GROUP BY" in error - - -def test_validate_sql_intent_alignment_allows_schema_backed_business_question(): - error = validate_sql_intent_alignment( - "Show top 10 customers by total invoice amount last month as a bar chart", - 'SELECT TOP 10 "invoices"."customer_id", ' - 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' - 'FROM "invoices" ' - 'WHERE "invoices"."created_at" >= \'2026-06-01 00:00:00\' ' - 'AND "invoices"."created_at" < \'2026-07-01 00:00:00\' ' - 'GROUP BY "invoices"."customer_id" ' - 'ORDER BY "total_invoice_amount" DESC', - {"invoices": ["customer_id", "invoice_amount", "created_at"]}, - ) - - assert error is None - - -def test_validate_sql_intent_alignment_rejects_join_without_condition(): - error = validate_sql_intent_alignment( - "Show order amount by customer region", - 'SELECT "customers"."region", SUM("orders"."order_amount") ' - 'AS "total_order_amount" FROM "orders" JOIN "customers" ' - 'GROUP BY "customers"."region"', - { - "orders": ["customer_id", "order_amount"], - "customers": ["customer_id", "region"], - }, - ) - - assert error is not None - assert "ON or USING" in error - - def test_validate_sql_intent_alignment_rejects_unmapped_metric_substitution(): error = validate_sql_intent_alignment( "Show invoice amount by customer", @@ -2061,28 +1962,6 @@ def test_construct_semantic_schema_contract_prioritizes_concept_mappings(): assert "Do not substitute identifiers for metrics" in contract -def test_construct_semantic_schema_contract_includes_generic_planning_requirements(): - contract = construct_semantic_schema_contract( - { - "analytical_intent": "dashboard", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "date_ranges": ["last month: 2026-06-01 to 2026-07-01"], - "sorting": ["invoice amount descending"], - "chart_requirements": ["bar chart by customer"], - "dashboard_requirements": ["KPI total invoice amount"], - "join_paths": ["invoices.customer_id -> customers.id"], - "is_fully_supported": True, - } - ) - - assert "Date ranges: last month" in contract - assert "Sorting: invoice amount descending" in contract - assert "Chart requirements: bar chart by customer" in contract - assert "Dashboard requirements: KPI total invoice amount" in contract - assert "Join paths: invoices.customer_id -> customers.id" in contract - - def test_construct_semantic_schema_contract_allows_legacy_analysis_without_mappings(): contract = construct_semantic_schema_contract( { From 97af000b47c40ff61deec07eef2f19a40017ca2d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:57:39 +0530 Subject: [PATCH 0482/1087] Revert "Add generic semantic schema candidate ranking" This reverts commit 97d09d73076b4e25169cb3bcd2241a8d80db1cb0. --- .../src/pipelines/generation/utils/sql.py | 92 ----- .../retrieval/db_schema_retrieval.py | 389 ------------------ .../pipelines/generation/test_sql_utils.py | 43 -- .../retrieval/test_db_schema_retrieval.py | 114 +---- 4 files changed, 1 insertion(+), 637 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2062de554d..97ae55f7c3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3647,7 +3647,6 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: "time_constraints", "ranking", "supported_schema_objects", - "candidate_schema_scores", "concept_mappings", "interpretations", "missing_requirements", @@ -3693,87 +3692,6 @@ def _schema_interpretation_clarification_error( return None -def _semantic_candidate_scores( - semantic_analysis: dict[str, Any] | None, -) -> list[dict[str, Any]]: - return _semantic_analysis_dict_items(semantic_analysis, "candidate_schema_scores") - - -def _semantic_candidate_support_error( - semantic_analysis: dict[str, Any], -) -> str | None: - candidate_scores = _semantic_candidate_scores(semantic_analysis) - if not candidate_scores: - return None - - complete_candidates = [ - candidate - for candidate in candidate_scores - if candidate.get("is_complete") is True - ] - if complete_candidates: - return None - - incomplete_with_missing = [ - candidate - for candidate in candidate_scores - if candidate.get("missing_concepts") - ] - if not incomplete_with_missing: - return None - - missing_concepts = [] - for candidate in incomplete_with_missing[:3]: - candidate_id = str(candidate.get("candidate_id") or "candidate").strip() - missing = candidate.get("missing_concepts") - if isinstance(missing, list): - missing_text = ", ".join( - str(item).strip() - for item in missing - if item is not None and str(item).strip() - ) - else: - missing_text = str(missing or "").strip() - if missing_text: - missing_concepts.append(f"{candidate_id}: {missing_text}") - - if not missing_concepts: - return None - - return ( - "Semantic schema retrieval did not find a complete schema mapping for " - "the request. Missing concepts: " - f"{'; '.join(missing_concepts)}. I cannot generate unrelated SQL." - ) - - -def _required_concept_mapping_support_error( - semantic_analysis: dict[str, Any], -) -> str | None: - unsupported_required_concepts = [] - for mapping in _semantic_concept_mappings(semantic_analysis): - if mapping.get("required_in_sql") is False: - continue - if _mapping_schema_objects(mapping): - continue - - request_concept = _mapping_request_concept(mapping) - concept_type = _mapping_concept_type(mapping) - if request_concept: - unsupported_required_concepts.append( - f"{request_concept} ({concept_type or 'concept'})" - ) - - if not unsupported_required_concepts: - return None - - return ( - "The semantic contract did not map required request concepts to active " - "schema objects: " - f"{', '.join(unsupported_required_concepts)}. I cannot generate unrelated SQL." - ) - - def get_schema_intent_analysis_error( semantic_analysis: dict[str, Any] | None, ) -> str | None: @@ -3804,16 +3722,6 @@ def get_schema_intent_analysis_error( ): return interpretation_error - if candidate_support_error := _semantic_candidate_support_error( - semantic_analysis - ): - return candidate_support_error - - if required_mapping_error := _required_concept_mapping_support_error( - semantic_analysis - ): - return required_mapping_error - if semantic_analysis.get("is_fully_supported") is False: support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() if support_reasoning: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index bfbbae7b50..f45e24a5b0 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,6 +1,5 @@ import ast import logging -import re import sys from typing import TYPE_CHECKING, Any, Optional @@ -158,29 +157,6 @@ {{ db_schema }} {% endfor %} -{% if semantic_candidate_context %} -### PRE-RANKED SEMANTIC SCHEMA CANDIDATES ### -These candidates were scored generically from the active datasource schema metadata and the user's full request. -Use them as retrieval evidence, but still verify complete concept coverage before selecting a contract. -Prefer candidates that cover all requested entities, identifiers, metrics, dimensions, filters, time constraints, aggregations, and ranking requirements. -Do not select a high lexical match when it misses a required business concept. - -{% for candidate in semantic_candidate_context %} -- candidate_id: {{ candidate.candidate_id }} - table_name: {{ candidate.table_name }} - confidence: {{ candidate.confidence }} - coverage_score: {{ candidate.coverage_score }} - matched_query_terms: {{ candidate.matched_query_terms }} - missing_query_terms: {{ candidate.missing_query_terms }} - rejected_by_retry: {{ candidate.rejected_by_retry }} - selection_reason: {{ candidate.selection_reason }} - matched_columns: -{% for column in candidate.matched_columns %} - - {{ column.column_name }} (score={{ column.score }}, data_type={{ column.data_type }}, matched_terms={{ column.matched_terms }}) -{% endfor %} -{% endfor %} -{% endif %} - ### INPUT ### {{ question }} @@ -283,361 +259,6 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: return deduped -_SEMANTIC_TOKEN_STOPWORDS = { - "a", - "an", - "and", - "are", - "as", - "at", - "be", - "by", - "for", - "from", - "give", - "have", - "how", - "in", - "is", - "me", - "of", - "on", - "or", - "show", - "that", - "the", - "to", - "with", - "dbo", - "tbl", - "table", - "view", - "dim", - "fact", - "stage", - "stg", -} - -_NUMERIC_SCHEMA_TERMS = { - "amount", - "avg", - "average", - "balance", - "cost", - "count", - "gross", - "margin", - "measure", - "metric", - "net", - "price", - "profit", - "quantity", - "rate", - "revenue", - "sales", - "sum", - "total", - "value", -} - -_TEMPORAL_SCHEMA_TERMS = { - "date", - "day", - "month", - "monthly", - "quarter", - "time", - "week", - "year", -} - -_RANKING_SCHEMA_TERMS = { - "bottom", - "highest", - "least", - "lowest", - "most", - "rank", - "ranking", - "top", -} - - -def _semantic_tokens(value: Any) -> set[str]: - text = str(value or "") - if not text: - return set() - - text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) - text = re.sub(r"[^A-Za-z0-9]+", " ", text) - tokens = { - token.lower() - for token in text.split() - if len(token) > 1 and token.lower() not in _SEMANTIC_TOKEN_STOPWORDS - } - for token in list(tokens): - if token.endswith("ies") and len(token) > 4: - tokens.add(f"{token[:-3]}y") - elif token.endswith("s") and len(token) > 3: - tokens.add(token[:-1]) - return tokens - - -def _schema_comment_text(value: Any) -> str: - if value is None: - return "" - if isinstance(value, str): - return value - if isinstance(value, list): - return " ".join(_schema_comment_text(item) for item in value) - if isinstance(value, dict): - return " ".join(_schema_comment_text(item) for item in value.values()) - return str(value) - - -def _column_tokens(column: dict[str, Any]) -> set[str]: - tokens = set() - for key in ("name", "display_name", "alias", "comment", "description", "data_type"): - tokens.update(_semantic_tokens(column.get(key))) - tokens.update(_semantic_tokens(_schema_comment_text(column.get("properties")))) - return tokens - - -def _table_tokens(table_schema: dict[str, Any]) -> set[str]: - tokens = set() - for key in ("name", "display_name", "alias", "comment", "description"): - tokens.update(_semantic_tokens(table_schema.get(key))) - for column in table_schema.get("columns", []) or []: - if isinstance(column, dict): - tokens.update(_column_tokens(column)) - return tokens - - -def _query_semantic_terms(query: str) -> dict[str, set[str]]: - tokens = _semantic_tokens(query) - return { - "all": tokens, - "metric": tokens & _NUMERIC_SCHEMA_TERMS, - "time": tokens & _TEMPORAL_SCHEMA_TERMS, - "ranking": tokens & _RANKING_SCHEMA_TERMS, - } - - -def _is_numeric_column(column: dict[str, Any]) -> bool: - data_type = str(column.get("data_type") or "").lower() - return bool( - re.search( - r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|number|double|float|real|money)\b", - data_type, - ) - ) - - -def _is_identifier_column(column: dict[str, Any]) -> bool: - tokens = _column_tokens(column) - return bool(tokens & {"code", "id", "identifier", "key", "no", "number"}) - - -def _is_temporal_column(column: dict[str, Any]) -> bool: - data_type = str(column.get("data_type") or "").lower() - return bool(re.search(r"\b(?:date|time|timestamp|datetime)\b", data_type)) - - -def _normalized_schema_object(value: Any) -> str: - return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) - - -def _rejected_schema_objects(semantic_retry_context: dict[str, Any] | None) -> set[str]: - if not isinstance(semantic_retry_context, dict): - return set() - - rejected = semantic_retry_context.get("rejected_schema_objects") - if not isinstance(rejected, list): - return set() - - return { - _normalized_schema_object(item) - for item in rejected - if item is not None and str(item).strip() - } - - -def _schema_object_was_rejected( - table_name: str, - column_name: str | None, - rejected_schema_objects: set[str], -) -> bool: - if not rejected_schema_objects: - return False - - table_key = _normalized_schema_object(table_name) - object_key = _normalized_schema_object( - f"{table_name}.{column_name}" if column_name else table_name - ) - return any( - rejected_key - and ( - rejected_key == table_key - or rejected_key == object_key - or rejected_key.endswith(object_key) - or object_key.endswith(rejected_key) - ) - for rejected_key in rejected_schema_objects - ) - - -def rank_semantic_schema_candidates( - query: str, - construct_db_schemas: list[dict], - semantic_retry_context: dict[str, Any] | None = None, - max_candidates: int = 15, - max_columns_per_candidate: int = 8, -) -> list[dict[str, Any]]: - query_terms = _query_semantic_terms(query) - all_query_terms = query_terms["all"] - if not all_query_terms: - return [] - - rejected_schema_objects = _rejected_schema_objects(semantic_retry_context) - candidates: list[dict[str, Any]] = [] - - for table_schema in construct_db_schemas: - if table_schema.get("type") != "TABLE": - continue - - table_name = str(table_schema.get("name") or "").strip() - if not table_name: - continue - - table_term_matches = _table_tokens(table_schema) & all_query_terms - matched_columns = [] - table_rejected = _schema_object_was_rejected( - table_name, None, rejected_schema_objects - ) - - for column in table_schema.get("columns", []) or []: - if not isinstance(column, dict): - continue - - column_name = str(column.get("name") or "").strip() - if not column_name: - continue - - tokens = _column_tokens(column) - matched_terms = sorted(tokens & all_query_terms) - score = float(len(matched_terms) * 3) - - if query_terms["metric"] and _is_numeric_column(column): - score += 0.3 if _is_identifier_column(column) else 1.5 - if tokens & query_terms["metric"]: - score += 2.0 - if query_terms["time"] and _is_temporal_column(column): - score += 1.5 - if tokens & query_terms["time"]: - score += 2.0 - if query_terms["ranking"] and matched_terms: - score += 0.5 - - rejected = _schema_object_was_rejected( - table_name, column_name, rejected_schema_objects - ) - if rejected: - score -= 5.0 - - if score > 0 or matched_terms: - matched_columns.append( - { - "column_name": column_name, - "score": round(max(score, 0.0), 3), - "matched_terms": matched_terms, - "data_type": str(column.get("data_type") or ""), - "rejected_by_retry": rejected, - } - ) - - matched_columns.sort( - key=lambda item: (item["score"], len(item["matched_terms"])), - reverse=True, - ) - matched_columns = matched_columns[:max_columns_per_candidate] - - covered_terms = set(table_term_matches) - for column in matched_columns: - covered_terms.update(column["matched_terms"]) - - if not covered_terms and not table_rejected: - continue - - coverage_score = len(covered_terms) / max(len(all_query_terms), 1) - raw_score = ( - len(table_term_matches) * 2.0 - + sum(column["score"] for column in matched_columns) - + coverage_score * 4.0 - ) - if table_rejected: - raw_score -= 6.0 - column_lookup = { - str(column.get("name") or ""): column - for column in table_schema.get("columns", []) or [] - if isinstance(column, dict) - } - has_metric_support = any( - query_terms["metric"] & set(column["matched_terms"]) - or ( - _is_numeric_column(column_lookup.get(column["column_name"], {})) - and not _is_identifier_column( - column_lookup.get(column["column_name"], {}) - ) - ) - for column in matched_columns - ) - if query_terms["metric"] and not has_metric_support: - raw_score -= 2.0 - - confidence = min(max(raw_score / 20.0, 0.0), 0.99) - selection_reason = ( - "Covers " - f"{len(covered_terms)} of {len(all_query_terms)} significant request terms" - ) - if table_rejected: - selection_reason += "; penalized because it was rejected by semantic validation" - if query_terms["metric"] and not any( - set(column["matched_terms"]) & query_terms["metric"] - for column in matched_columns - ): - selection_reason += "; metric term coverage is weak" - - candidates.append( - { - "candidate_id": f"candidate-{len(candidates) + 1}", - "table_name": table_name, - "confidence": round(confidence, 3), - "coverage_score": round(coverage_score, 3), - "matched_query_terms": sorted(covered_terms), - "missing_query_terms": sorted(all_query_terms - covered_terms), - "matched_columns": matched_columns, - "rejected_by_retry": table_rejected - or any(column["rejected_by_retry"] for column in matched_columns), - "selection_reason": selection_reason, - } - ) - - candidates.sort( - key=lambda item: ( - item["rejected_by_retry"] is False, - item["confidence"], - item["coverage_score"], - ), - reverse=True, - ) - - for index, candidate in enumerate(candidates[:max_candidates], start=1): - candidate["candidate_id"] = f"candidate-{index}" - - return candidates[:max_candidates] - - @observe(capture_input=False, capture_output=False) async def embedding( query: str, @@ -838,20 +459,10 @@ def prompt( ) query = "\n".join(previous_query_summaries) + "\n" + query - semantic_candidate_context = rank_semantic_schema_candidates( - query=query, - construct_db_schemas=construct_db_schemas, - semantic_retry_context=semantic_retry_context, - ) - logger.info( - "semantic_retrieval_pre_ranked_candidates=%s", - semantic_candidate_context, - ) _prompt = prompt_builder.run( question=query, db_schemas=db_schemas, - semantic_candidate_context=semantic_candidate_context, semantic_retry_context=semantic_retry_context or {}, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index d22e45f36c..4d8797ee50 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1675,49 +1675,6 @@ def test_get_schema_intent_analysis_error_reports_unsupported_analysis(): assert "No relationship connects invoices to products" in error -def test_get_schema_intent_analysis_error_rejects_incomplete_semantic_candidates(): - error = get_schema_intent_analysis_error( - { - "candidate_schema_scores": [ - { - "candidate_id": "candidate-1", - "schema_objects": ["refunds.refund_amount"], - "covered_concepts": ["amount"], - "missing_concepts": ["customer", "invoice amount"], - "confidence": 0.62, - "is_complete": False, - "selection_reason": "Only amount matched.", - } - ], - "is_fully_supported": True, - } - ) - - assert error is not None - assert "complete schema mapping" in error - assert "invoice amount" in error - - -def test_get_schema_intent_analysis_error_rejects_required_unmapped_concepts(): - error = get_schema_intent_analysis_error( - { - "concept_mappings": [ - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": [], - "required_in_sql": True, - } - ], - "is_fully_supported": True, - } - ) - - assert error is not None - assert "invoice amount" in error - assert "did not map required request concepts" in error - - def test_validate_sql_intent_alignment_uses_semantic_analysis_for_metric_count_mismatch(): error = validate_sql_intent_alignment( "Show invoice amount by customer", diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index a70f26aecb..ed67f4e9aa 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -7,7 +7,6 @@ dbschema_retrieval, expand_business_terms_for_retrieval, prompt, - rank_semantic_schema_candidates, ) @@ -33,114 +32,6 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query -def test_rank_semantic_schema_candidates_prefers_complete_business_concept_coverage(): - candidates = rank_semantic_schema_candidates( - query="Show top 10 customers by invoice amount", - construct_db_schemas=[ - { - "type": "TABLE", - "name": "dbo_ytblES002_1", - "comment": "", - "columns": [ - { - "name": "Name_of_Reported_Received", - "data_type": "varchar", - "comment": "reported received name", - }, - { - "name": "Refund_Amount", - "data_type": "decimal", - "comment": "refund amount", - }, - ], - }, - { - "type": "TABLE", - "name": "dbo_tblFactSales", - "comment": "invoice sales facts by customer", - "columns": [ - { - "name": "invoice", - "data_type": "varchar", - "comment": "invoice identifier", - }, - { - "name": "customer_id", - "data_type": "varchar", - "comment": "customer identifier", - }, - { - "name": "Amount_Received", - "data_type": "decimal", - "comment": "invoice amount received", - }, - ], - }, - ], - ) - - assert candidates[0]["table_name"] == "dbo_tblFactSales" - assert "customer" in candidates[0]["matched_query_terms"] - assert "amount" in candidates[0]["matched_query_terms"] - assert candidates[0]["confidence"] > candidates[1]["confidence"] - - -def test_rank_semantic_schema_candidates_penalizes_retry_rejected_schema_objects(): - candidates = rank_semantic_schema_candidates( - query="Show top 10 customers by invoice amount", - construct_db_schemas=[ - { - "type": "TABLE", - "name": "dbo_tblFactSales", - "comment": "invoice sales facts by customer", - "columns": [ - { - "name": "customer_id", - "data_type": "varchar", - "comment": "customer identifier", - }, - { - "name": "Amount_Received", - "data_type": "decimal", - "comment": "invoice amount received", - }, - ], - }, - { - "type": "TABLE", - "name": "dbo_qSales", - "comment": "invoice analytics by account", - "columns": [ - { - "name": "Account", - "data_type": "varchar", - "comment": "customer account", - }, - { - "name": "InvoiceAmount", - "data_type": "decimal", - "comment": "invoice amount", - }, - ], - }, - ], - semantic_retry_context={ - "rejected_schema_objects": [ - "dbo_tblFactSales.customer_id", - "dbo_tblFactSales.Amount_Received", - ] - }, - ) - - assert candidates[0]["table_name"] == "dbo_qSales" - rejected_candidate = next( - candidate - for candidate in candidates - if candidate["table_name"] == "dbo_tblFactSales" - ) - assert rejected_candidate["rejected_by_retry"] is True - - @pytest.mark.asyncio async def test_dbschema_retrieval_loads_complete_active_project_schema(): class Retriever: @@ -308,13 +199,11 @@ def test_prompt_includes_semantic_retry_context(): class PromptBuilder: def run(self, **kwargs): retry_context = kwargs["semantic_retry_context"] - candidate_context = kwargs["semantic_candidate_context"] return { "prompt": ( f"retry={retry_context['retry_attempt']} " f"error={retry_context['validation_error']} " - f"rejected={','.join(retry_context['rejected_schema_objects'])} " - f"candidates={candidate_context[0]['table_name']}" + f"rejected={','.join(retry_context['rejected_schema_objects'])}" ) } @@ -349,4 +238,3 @@ def run(self, **kwargs): assert "retry=2" in result["prompt"] assert "Generic count" in result["prompt"] assert "dbo_ytblES002_1.Name_of_Reported_Received" in result["prompt"] - assert "candidates=invoices" in result["prompt"] From e6417feca39b879bec9676f46f890a5078648ab9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:57:52 +0530 Subject: [PATCH 0483/1087] Revert "Keep semantic SQL requests out of legacy fallbacks" This reverts commit 11a631f13e7558f4bac839bea7daaad1e5201b7f. --- wren-ai-service/src/web/v1/services/ask.py | 163 +++------------------ 1 file changed, 20 insertions(+), 143 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 00c45f2fa7..16023abdea 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4496,31 +4496,6 @@ def _semantic_schema_objects(semantic_analysis: dict[str, Any] | None) -> list[s return sorted(set(schema_objects)) - @staticmethod - def _has_semantic_contract(semantic_analysis: dict[str, Any] | None) -> bool: - if not isinstance(semantic_analysis, dict) or not semantic_analysis: - return False - semantic_keys = { - "analytical_intent", - "entities", - "identifiers", - "metrics", - "dimensions", - "filters", - "aggregations", - "relationships", - "time_constraints", - "ranking", - "supported_schema_objects", - "candidate_schema_scores", - "concept_mappings", - "interpretations", - "missing_requirements", - "ambiguous_requirements", - "support_reasoning", - } - return any(semantic_analysis.get(key) for key in semantic_keys) - @staticmethod def _semantic_retry_context( semantic_analysis: dict[str, Any] | None, @@ -5407,8 +5382,6 @@ async def ask( sql_knowledge = None understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) schema_intent_analysis: dict[str, Any] = {} - semantic_pipeline_active = False - semantic_retrieval_activated = False try: sql_user_query = user_query @@ -6053,25 +6026,16 @@ async def ask( ) try: - semantic_retrieval_activated = ( - enable_column_pruning - or self._is_data_analysis_query(user_query) - ) - logger.info( - "semantic_retrieval_activation query_id=%s active=%s reason=%s", - query_id, - semantic_retrieval_activated, - "analytics_or_column_pruning" - if semantic_retrieval_activated - else "standard_retrieval", - ) retrieval_result = await self._run_with_timeout( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( query=sql_user_query, histories=[], project_id=ask_request.project_id, - enable_column_pruning=semantic_retrieval_activated, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -6186,23 +6150,8 @@ async def ask( logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) - semantic_pipeline_active = self._has_semantic_contract( - schema_intent_analysis - ) or semantic_retrieval_activated - logger.info( - "sql_generation_pipeline_decision query_id=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s", - query_id, - semantic_pipeline_active, - self._has_semantic_contract(schema_intent_analysis), - self._semantic_schema_objects(schema_intent_analysis), - ) - if semantic_pipeline_active: - logger.info( - "legacy_sql_fallbacks_disabled query_id=%s reason=semantic_pipeline_active", - query_id, - ) - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6222,7 +6171,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6246,7 +6195,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6270,7 +6219,6 @@ async def ask( if ( not api_results - and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and ( schema_grounded_sql := self._build_schema_grounded_analytics_sql( @@ -6295,7 +6243,7 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if not semantic_pipeline_active and not api_results and any( + if not api_results and any( term in user_query.lower() for term in ( "pcb", @@ -6328,7 +6276,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6352,7 +6300,6 @@ async def ask( should_retry_full_schema = ( not api_results - and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and "db_schema_retrieval" in self._pipelines ) @@ -6488,7 +6435,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if not documents and not semantic_pipeline_active: + if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): @@ -6555,34 +6502,6 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if semantic_pipeline_active and not documents: - semantic_failure_message = error_message or ( - "Semantic schema retrieval did not find a complete schema mapping for the request." - ) - logger.info( - "semantic_pipeline_no_supported_documents query_id=%s message=%s", - query_id, - semantic_failure_message, - ) - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_failed_text_to_sql_response( - trace_id, - semantic_failure_message, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - is_followup=True if histories else False, - code="NO_RELEVANT_SQL", - ) - ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = semantic_failure_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - if documents and not api_results: documents, table_names, table_ddls = self._prune_sql_generation_context( sql_user_query, @@ -6610,19 +6529,13 @@ async def ask( sql_user_query ) and not self._needs_conversation_context(sql_user_query): sql_generation_histories = [] + allow_sql_generation_reasoning = False allow_sql_knowledge_retrieval = False max_sql_correction_retries = min(max_sql_correction_retries, 1) - if semantic_pipeline_active: - logger.info( - "fast_standalone_sql_generation_disabled query_id=%s reason=semantic_pipeline_active", - query_id, - ) - else: - allow_sql_generation_reasoning = False - logger.info( - "Using fast standalone SQL generation path for query_id %s", - query_id, - ) + logger.info( + "Using fast standalone SQL generation path for query_id %s", + query_id, + ) if ( not self._is_stopped(query_id, self._ask_results) @@ -6790,18 +6703,11 @@ async def ask( ), ) except TimeoutError as generation_timeout: - if semantic_pipeline_active: - logger.warning( - "Semantic SQL generation timed out for query_id %s; legacy fallbacks remain disabled: %s", - query_id, - generation_timeout, - ) - else: - logger.warning( - "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", - query_id, - generation_timeout, - ) + logger.warning( + "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + query_id, + generation_timeout, + ) text_to_sql_generation_results = { "post_process": { "valid_generation_result": None, @@ -7146,35 +7052,6 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if semantic_pipeline_active: - semantic_failure_message = error_message or ( - "No valid semantic schema contract could satisfy the request after semantic retrieval retries." - ) - logger.info( - "semantic_pipeline_exhausted query_id=%s message=%s rejected_sql=%s", - query_id, - semantic_failure_message, - invalid_sql, - ) - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_failed_text_to_sql_response( - trace_id, - semantic_failure_message, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - is_followup=True if histories else False, - code="NO_RELEVANT_SQL", - ) - ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = semantic_failure_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): From 04acaefd2aafd65ac4fdc9445e4a76c9c432e63e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:58:34 +0530 Subject: [PATCH 0484/1087] Revert "Retry SQL generation with alternate semantic mappings" This reverts commit fea5587d422c945a986be633a38240186c1b4e9b. --- .../retrieval/db_schema_retrieval.py | 115 +------ wren-ai-service/src/web/v1/services/ask.py | 314 ++++++------------ .../src/web/v1/services/ask_feedback.py | 175 +++------- .../retrieval/test_db_schema_retrieval.py | 58 ---- 4 files changed, 154 insertions(+), 508 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index f45e24a5b0..0fc19d58b1 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -47,14 +47,11 @@ 9. For each selected column, provide a concise reason for why the column is necessary. 10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. 11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. -12. Populate `candidate_schema_scores` with ranked candidates. Score each candidate by full concept coverage, semantic fit, relationship viability, metric validity, and whether it satisfies filters/time/ranking/aggregation requirements. Reject partial lexical matches even when a table or column name looks similar. -13. Select only the highest-confidence candidate whose mappings completely cover all required concepts. If no candidate fully covers the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting a partial mapping. -14. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Do not select those objects again unless every complete candidate is exhausted; explain any reuse in `support_reasoning`. -15. If a "." is included in columns, put the name before the first dot into chosen columns. -16. The number of columns chosen must match the number of reasoning. -17. Final chosen columns must be only column names, don't prefix it with table names. -18. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -19. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +12. If a "." is included in columns, put the name before the first dot into chosen columns. +13. The number of columns chosen must match the number of reasoning. +14. Final chosen columns must be only column names, don't prefix it with table names. +15. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +16. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -72,17 +69,6 @@ "time_constraints": ["time filters, grains, or trend requirements"], "ranking": ["top/bottom/order/limit requirements"], "supported_schema_objects": ["table.column or metric names that directly support the request"], - "candidate_schema_scores": [ - { - "candidate_id": "candidate-1", - "schema_objects": ["table.column objects included in this candidate"], - "covered_concepts": ["request concepts this candidate supports"], - "missing_concepts": ["request concepts this candidate cannot support"], - "confidence": 0.0, - "is_complete": true, - "selection_reason": "Why this candidate is accepted or rejected" - } - ], "concept_mappings": [ { "request_concept": "business concept from the user request", @@ -159,17 +145,6 @@ ### INPUT ### {{ question }} - -{% if semantic_retry_context %} -### RETRY CONTEXT ### -Previous semantic SQL validation failed. Discard the previous contract and do not reuse rejected schema mappings unless no other complete candidate exists. -Validation failure: {{ semantic_retry_context.validation_error }} -Retry attempt: {{ semantic_retry_context.retry_attempt }} -Rejected schema objects: -{% for schema_object in semantic_retry_context.rejected_schema_objects %} -- {{ schema_object }} -{% endfor %} -{% endif %} """ @@ -419,16 +394,6 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) - if enable_column_pruning or _token_count > context_window_size: - return { - "db_schemas": [], - "tokens": _token_count, - "has_calculated_field": has_calculated_field, - "has_metric": has_metric, - "has_json_field": has_json_field, - "semantic_analysis": {}, - } - return { "db_schemas": retrieval_results, "tokens": _token_count, @@ -446,7 +411,6 @@ def prompt( prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, histories: list[AskHistory], - semantic_retry_context: dict[str, Any] | None = None, ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -460,11 +424,7 @@ def prompt( query = "\n".join(previous_query_summaries) + "\n" + query - _prompt = prompt_builder.run( - question=query, - db_schemas=db_schemas, - semantic_retry_context=semantic_retry_context or {}, - ) + _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: return {} @@ -494,7 +454,6 @@ def construct_retrieval_results( retrieval_payload = orjson.loads(filter_columns_in_tables["replies"][0]) columns_and_tables_needed = retrieval_payload.get("results", []) semantic_analysis = retrieval_payload.get("semantic_analysis") or {} - _log_semantic_retrieval_decision(semantic_analysis) # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -572,51 +531,6 @@ def construct_retrieval_results( } -def _semantic_log_items(semantic_analysis: dict[str, Any], key: str) -> list[str]: - value = semantic_analysis.get(key) - if isinstance(value, str): - return [value] if value.strip() else [] - if isinstance(value, list): - return [ - str(item).strip() - for item in value - if item is not None and str(item).strip() - ] - return [] - - -def _log_semantic_retrieval_decision(semantic_analysis: dict[str, Any]) -> None: - if not isinstance(semantic_analysis, dict) or not semantic_analysis: - logger.info("semantic_retrieval_decision=no_semantic_analysis") - return - - concepts = { - "intent": semantic_analysis.get("analytical_intent"), - "entities": _semantic_log_items(semantic_analysis, "entities"), - "identifiers": _semantic_log_items(semantic_analysis, "identifiers"), - "metrics": _semantic_log_items(semantic_analysis, "metrics"), - "dimensions": _semantic_log_items(semantic_analysis, "dimensions"), - "filters": _semantic_log_items(semantic_analysis, "filters"), - "time_constraints": _semantic_log_items( - semantic_analysis, "time_constraints" - ), - "aggregations": _semantic_log_items(semantic_analysis, "aggregations"), - "ranking": _semantic_log_items(semantic_analysis, "ranking"), - } - candidate_scores = semantic_analysis.get("candidate_schema_scores") or [] - selected_contract = { - "supported_schema_objects": semantic_analysis.get( - "supported_schema_objects", [] - ), - "concept_mappings": semantic_analysis.get("concept_mappings", []), - "is_fully_supported": semantic_analysis.get("is_fully_supported"), - "support_reasoning": semantic_analysis.get("support_reasoning"), - } - logger.info("semantic_retrieval_concepts=%s", concepts) - logger.info("semantic_retrieval_candidate_scores=%s", candidate_scores) - logger.info("semantic_retrieval_selected_contract=%s", selected_contract) - - ## End of Pipeline class MatchingTableContents(BaseModel): chain_of_thought_reasoning: list[str] @@ -645,16 +559,6 @@ class SemanticInterpretation(BaseModel): is_selected: bool = False -class SemanticCandidateSchemaScore(BaseModel): - candidate_id: str = "" - schema_objects: list[str] = Field(default_factory=list) - covered_concepts: list[str] = Field(default_factory=list) - missing_concepts: list[str] = Field(default_factory=list) - confidence: float | None = None - is_complete: bool = False - selection_reason: str = "" - - class SemanticAnalysis(BaseModel): analytical_intent: str = "" entities: list[str] = Field(default_factory=list) @@ -667,9 +571,6 @@ class SemanticAnalysis(BaseModel): time_constraints: list[str] = Field(default_factory=list) ranking: list[str] = Field(default_factory=list) supported_schema_objects: list[str] = Field(default_factory=list) - candidate_schema_scores: list[SemanticCandidateSchemaScore] = Field( - default_factory=list - ) concept_mappings: list[SemanticConceptMapping] = Field(default_factory=list) interpretations: list[SemanticInterpretation] = Field(default_factory=list) missing_requirements: list[str] = Field(default_factory=list) @@ -748,11 +649,8 @@ async def run( project_id: Optional[str] = None, histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, - semantic_retry_context: Optional[dict[str, Any]] = None, ): logger.info("Ask Retrieval pipeline is running...") - if semantic_retry_context: - logger.info("semantic_retrieval_retry_context=%s", semantic_retry_context) return await self._pipe.execute( ["construct_retrieval_results"], inputs={ @@ -761,7 +659,6 @@ async def run( "project_id": project_id or "", "histories": histories or [], "enable_column_pruning": enable_column_pruning, - "semantic_retry_context": semantic_retry_context or {}, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 16023abdea..5377479691 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4452,66 +4452,6 @@ def _metadata_from_documents( ] return table_names, table_ddls - @staticmethod - def _semantic_schema_objects(semantic_analysis: dict[str, Any] | None) -> list[str]: - if not isinstance(semantic_analysis, dict): - return [] - - schema_objects: list[str] = [] - supported_schema_objects = semantic_analysis.get("supported_schema_objects") - if isinstance(supported_schema_objects, list): - schema_objects.extend( - str(schema_object).strip() - for schema_object in supported_schema_objects - if schema_object is not None and str(schema_object).strip() - ) - - concept_mappings = semantic_analysis.get("concept_mappings") - if isinstance(concept_mappings, list): - for mapping in concept_mappings: - if not isinstance(mapping, dict): - continue - mapping_schema_objects = mapping.get("schema_objects") - if isinstance(mapping_schema_objects, list): - schema_objects.extend( - str(schema_object).strip() - for schema_object in mapping_schema_objects - if schema_object is not None and str(schema_object).strip() - ) - - interpretations = semantic_analysis.get("interpretations") - if isinstance(interpretations, list): - for interpretation in interpretations: - if not isinstance(interpretation, dict): - continue - if interpretation.get("is_selected") is not True: - continue - interpretation_schema_objects = interpretation.get("schema_objects") - if isinstance(interpretation_schema_objects, list): - schema_objects.extend( - str(schema_object).strip() - for schema_object in interpretation_schema_objects - if schema_object is not None and str(schema_object).strip() - ) - - return sorted(set(schema_objects)) - - @staticmethod - def _semantic_retry_context( - semantic_analysis: dict[str, Any] | None, - validation_error: str, - retry_attempt: int, - rejected_schema_objects: set[str], - ) -> dict[str, Any]: - rejected_schema_objects.update( - AskService._semantic_schema_objects(semantic_analysis) - ) - return { - "validation_error": validation_error, - "retry_attempt": retry_attempt, - "rejected_schema_objects": sorted(rejected_schema_objects), - } - async def _complete_sql_generation_context( self, *, @@ -6733,36 +6673,21 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - rejected_schema_objects: set[str] = set() - semantic_retry_attempt = 0 - while ( - failed_dry_run_result - and failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION" - and semantic_retry_attempt < 3 - and not api_results - ): - semantic_retry_attempt += 1 + if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": semantic_retry_error = failed_dry_run_result.get("error", "") - retry_context = self._semantic_retry_context( - schema_intent_analysis, - semantic_retry_error, - semantic_retry_attempt, - rejected_schema_objects, - ) semantic_retry_query = ( f"{sql_user_query}\n\n" "The previous generated SQL failed semantic validation: " f"{semantic_retry_error}\n" - "Perform a fresh semantic retrieval. Select the next " - "highest-confidence complete concept-to-schema mapping. " - "Do not reuse rejected schema objects from RETRY CONTEXT." + "Re-run semantic schema retrieval and choose the next " + "highest-confidence concept-to-schema mapping that " + "directly supports the user's requested entities, " + "metrics, dimensions, filters, time constraints, " + "aggregations, relationships, and ranking." ) logger.info( - "semantic_retry_attempt=%s query_id=%s validation_failure=%s rejected_schema_objects=%s", - semantic_retry_attempt, + "Retrying semantic schema retrieval after intent validation failure for query_id %s", query_id, - semantic_retry_error, - sorted(rejected_schema_objects), ) try: retry_retrieval_result = await self._run_with_timeout( @@ -6771,8 +6696,7 @@ async def ask( query=semantic_retry_query, histories=histories, project_id=ask_request.project_id, - enable_column_pruning=True, - semantic_retry_context=retry_context, + enable_column_pruning=enable_column_pruning, ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -6782,25 +6706,6 @@ async def ask( retry_schema_intent_analysis = retry_construct_result.get( "semantic_analysis", {} ) - retry_selected_objects = set( - self._semantic_schema_objects( - retry_schema_intent_analysis - ) - ) - if retry_selected_objects and retry_selected_objects.issubset( - rejected_schema_objects - ): - error_message = ( - "Semantic schema retrieval retry selected only previously rejected schema objects." - ) - logger.info( - "semantic_retry_rejected_repeated_contract query_id=%s attempt=%s schema_objects=%s", - query_id, - semantic_retry_attempt, - sorted(retry_selected_objects), - ) - break - retry_documents, retry_table_names, retry_table_ddls = ( self._extract_retrieval_metadata( retry_retrieval_result @@ -6809,141 +6714,110 @@ async def ask( retry_support_error = get_schema_intent_analysis_error( retry_schema_intent_analysis ) - if not retry_documents or retry_support_error: - error_message = retry_support_error or ( - "Semantic schema retrieval retry did not find a supported mapping for the request." + if retry_documents and not retry_support_error: + _retrieval_result = retry_construct_result + schema_intent_analysis = retry_schema_intent_analysis + documents = retry_documents + table_names = retry_table_names + table_ddls = retry_table_ddls + has_calculated_field = _retrieval_result.get( + "has_calculated_field", False ) - logger.info( - "semantic_retry_candidate_rejected query_id=%s attempt=%s reason=%s selected_objects=%s", - query_id, - semantic_retry_attempt, - error_message, - sorted(retry_selected_objects), + has_metric = _retrieval_result.get("has_metric", False) + has_json_field = _retrieval_result.get( + "has_json_field", False ) - rejected_schema_objects.update(retry_selected_objects) - break - - _retrieval_result = retry_construct_result - schema_intent_analysis = retry_schema_intent_analysis - rejected_schema_objects.update(retry_selected_objects) - documents = retry_documents - table_names = retry_table_names - table_ddls = retry_table_ddls - has_calculated_field = _retrieval_result.get( - "has_calculated_field", False - ) - has_metric = _retrieval_result.get("has_metric", False) - has_json_field = _retrieval_result.get( - "has_json_field", False - ) - logger.info( - "semantic_retry_candidate_selected query_id=%s attempt=%s table_names=%s schema_objects=%s", - query_id, - semantic_retry_attempt, - table_names, - sorted(retry_selected_objects), - ) - if sql_generation_histories: - text_to_sql_generation_results = ( - await self._run_with_timeout( - "Follow-up SQL generation after semantic retry", - self._pipelines[ - "followup_sql_generation" - ].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=sql_generation_histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), + if histories: + text_to_sql_generation_results = ( + await self._run_with_timeout( + "Follow-up SQL generation after semantic retry", + self._pipelines[ + "followup_sql_generation" + ].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), + ) ) - ) - else: - text_to_sql_generation_results = ( - await self._run_with_timeout( - "SQL generation after semantic retry", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), + else: + text_to_sql_generation_results = ( + await self._run_with_timeout( + "SQL generation after semantic retry", + self._pipelines["sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), + ) ) - ) - if sql_valid_result := text_to_sql_generation_results[ - "post_process" - ]["valid_generation_result"]: - if ask_result := self._build_validated_ask_result_from_sql( - sql_valid_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - logger.info( - "semantic_retry_candidate_accepted query_id=%s attempt=%s", - query_id, - semantic_retry_attempt, + if sql_valid_result := text_to_sql_generation_results[ + "post_process" + ]["valid_generation_result"]: + if ask_result := self._build_validated_ask_result_from_sql( + sql_valid_result.get("sql"), + table_ddls, + sql_user_query, + ): + api_results = [ask_result] + else: + invalid_sql = sql_valid_result.get("sql") + error_message = ( + "SQL generation after semantic retrieval retry did not produce SQL that matches the active datasource schema and question intent." + ) + else: + failed_dry_run_result = ( + text_to_sql_generation_results["post_process"][ + "invalid_generation_result" + ] + ) + invalid_sql = failed_dry_run_result.get( + "sql", invalid_sql + ) + error_message = failed_dry_run_result.get( + "error", error_message ) - break - invalid_sql = sql_valid_result.get("sql") + elif retry_support_error: + error_message = retry_support_error + else: error_message = ( - "SQL generation after semantic retrieval retry did not produce SQL that matches the active datasource schema and question intent." + "Semantic schema retrieval retry did not find a supported mapping for the request." ) - break - - failed_dry_run_result = ( - text_to_sql_generation_results["post_process"][ - "invalid_generation_result" - ] - ) - invalid_sql = failed_dry_run_result.get( - "sql", invalid_sql - ) - error_message = failed_dry_run_result.get( - "error", error_message - ) - logger.info( - "semantic_retry_candidate_failed_validation query_id=%s attempt=%s error=%s", - query_id, - semantic_retry_attempt, - error_message, - ) except Exception as retry_error: logger.warning( - "Semantic schema retrieval retry failed for query_id %s attempt %s: %s", + "Semantic schema retrieval retry failed for query_id %s: %s", query_id, - semantic_retry_attempt, retry_error, ) - break while current_sql_correction_retries < max_sql_correction_retries: if api_results: break - if not failed_dry_run_result: - break if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index d26ce51269..cbd6f7d169 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -10,7 +10,7 @@ from src.pipelines.generation.utils.sql import get_schema_intent_analysis_error from src.utils import trace_metadata from src.web.v1.services import BaseRequest -from src.web.v1.services.ask import AskError, AskResult, AskService +from src.web.v1.services.ask import AskError, AskResult logger = logging.getLogger("wren-ai-service") @@ -224,36 +224,21 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - rejected_schema_objects: set[str] = set() - semantic_retry_attempt = 0 - while ( - failed_dry_run_result - and failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION" - and semantic_retry_attempt < 3 - and not api_results - ): - semantic_retry_attempt += 1 + if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": semantic_retry_error = failed_dry_run_result.get("error", "") - retry_context = AskService._semantic_retry_context( - schema_intent_analysis, - semantic_retry_error, - semantic_retry_attempt, - rejected_schema_objects, - ) semantic_retry_query = ( f"{ask_feedback_request.question}\n\n" "The previous regenerated SQL failed semantic validation: " f"{semantic_retry_error}\n" - "Perform a fresh semantic retrieval. Select the next " - "highest-confidence complete concept-to-schema mapping. " - "Do not reuse rejected schema objects from RETRY CONTEXT." + "Re-run semantic schema retrieval and choose the next " + "highest-confidence concept-to-schema mapping that " + "directly supports the user's requested entities, " + "metrics, dimensions, filters, time constraints, " + "aggregations, relationships, and ranking." ) logger.info( - "feedback_semantic_retry_attempt=%s query_id=%s validation_failure=%s rejected_schema_objects=%s", - semantic_retry_attempt, + "Retrying semantic schema retrieval after feedback intent validation failure for query_id %s", query_id, - semantic_retry_error, - sorted(rejected_schema_objects), ) retry_retrieval_result = await self._pipelines[ "db_schema_retrieval" @@ -261,8 +246,7 @@ async def ask_feedback( query=semantic_retry_query, histories=[], project_id=ask_feedback_request.project_id, - enable_column_pruning=True, - semantic_retry_context=retry_context, + enable_column_pruning=enable_column_pruning, ) retry_construct_result = retry_retrieval_result.get( "construct_retrieval_results", {} @@ -270,25 +254,6 @@ async def ask_feedback( retry_schema_intent_analysis = retry_construct_result.get( "semantic_analysis", {} ) - retry_selected_objects = set( - AskService._semantic_schema_objects( - retry_schema_intent_analysis - ) - ) - if retry_selected_objects and retry_selected_objects.issubset( - rejected_schema_objects - ): - error_message = ( - "Semantic schema retrieval retry selected only previously rejected schema objects." - ) - logger.info( - "feedback_semantic_retry_rejected_repeated_contract query_id=%s attempt=%s schema_objects=%s", - query_id, - semantic_retry_attempt, - sorted(retry_selected_objects), - ) - break - retry_documents = retry_construct_result.get( "retrieval_results", [] ) @@ -301,91 +266,59 @@ async def ask_feedback( retry_support_error = get_schema_intent_analysis_error( retry_schema_intent_analysis ) - if not retry_table_ddls or retry_support_error: - error_message = retry_support_error or ( - "Semantic schema retrieval retry did not find a supported mapping for the request." + if retry_table_ddls and not retry_support_error: + schema_intent_analysis = retry_schema_intent_analysis + documents = retry_documents + table_ddls = retry_table_ddls + has_calculated_field = retry_construct_result.get( + "has_calculated_field", False ) - logger.info( - "feedback_semantic_retry_candidate_rejected query_id=%s attempt=%s reason=%s selected_objects=%s", - query_id, - semantic_retry_attempt, - error_message, - sorted(retry_selected_objects), + has_metric = retry_construct_result.get( + "has_metric", False + ) + has_json_field = retry_construct_result.get( + "has_json_field", False ) - rejected_schema_objects.update(retry_selected_objects) - break - - schema_intent_analysis = retry_schema_intent_analysis - rejected_schema_objects.update(retry_selected_objects) - documents = retry_documents - table_ddls = retry_table_ddls - has_calculated_field = retry_construct_result.get( - "has_calculated_field", False - ) - has_metric = retry_construct_result.get("has_metric", False) - has_json_field = retry_construct_result.get( - "has_json_field", False - ) - logger.info( - "feedback_semantic_retry_candidate_selected query_id=%s attempt=%s schema_objects=%s", - query_id, - semantic_retry_attempt, - sorted(retry_selected_objects), - ) - - text_to_sql_generation_results = await self._pipelines[ - "sql_regeneration" - ].run( - contexts=table_ddls, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, - sql=ask_feedback_request.sql, - query=ask_feedback_request.question, - project_id=ask_feedback_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ) - if sql_valid_result := text_to_sql_generation_results[ - "post_process" - ]["valid_generation_result"]: - api_results = [ - AskResult( - **{ - "sql": sql_valid_result.get("sql"), - "type": "llm", - } - ) - ] - logger.info( - "feedback_semantic_retry_candidate_accepted query_id=%s attempt=%s", - query_id, - semantic_retry_attempt, + text_to_sql_generation_results = await self._pipelines[ + "sql_regeneration" + ].run( + contexts=table_ddls, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + sql=ask_feedback_request.sql, + query=ask_feedback_request.question, + project_id=ask_feedback_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, ) - break - failed_dry_run_result = text_to_sql_generation_results[ - "post_process" - ]["invalid_generation_result"] - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) - logger.info( - "feedback_semantic_retry_candidate_failed_validation query_id=%s attempt=%s error=%s", - query_id, - semantic_retry_attempt, - error_message, - ) + if sql_valid_result := text_to_sql_generation_results[ + "post_process" + ]["valid_generation_result"]: + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } + ) + ] + else: + failed_dry_run_result = text_to_sql_generation_results[ + "post_process" + ]["invalid_generation_result"] + elif retry_support_error: + error_message = retry_support_error if api_results: pass - elif failed_dry_run_result and failed_dry_run_result["type"] not in { + elif failed_dry_run_result["type"] not in { "TIME_OUT", "SCHEMA_INTENT_VALIDATION", }: diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index ed67f4e9aa..afb2d06c8b 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -6,7 +6,6 @@ construct_retrieval_results, dbschema_retrieval, expand_business_terms_for_retrieval, - prompt, ) @@ -103,17 +102,6 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): "entities": ["invoice"], "metrics": ["invoice amount"], "dimensions": ["customer"], - "candidate_schema_scores": [ - { - "candidate_id": "candidate-1", - "schema_objects": ["invoices.customer_id", "invoices.invoice_amount"], - "covered_concepts": ["invoice amount", "customer"], - "missing_concepts": [], - "confidence": 0.95, - "is_complete": true, - "selection_reason": "Complete invoice amount by customer mapping." - } - ], "concept_mappings": [ { "request_concept": "invoice amount", @@ -188,53 +176,7 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): assert result["semantic_analysis"]["concept_mappings"][0]["schema_objects"] == [ "invoices.invoice_amount" ] - assert result["semantic_analysis"]["candidate_schema_scores"][0]["is_complete"] assert result["semantic_analysis"]["interpretations"][0]["is_selected"] is True assert result["retrieval_results"][0]["table_name"] == "invoices" assert "invoice_amount" in result["retrieval_results"][0]["table_ddl"] assert "internal_note" not in result["retrieval_results"][0]["table_ddl"] - - -def test_prompt_includes_semantic_retry_context(): - class PromptBuilder: - def run(self, **kwargs): - retry_context = kwargs["semantic_retry_context"] - return { - "prompt": ( - f"retry={retry_context['retry_attempt']} " - f"error={retry_context['validation_error']} " - f"rejected={','.join(retry_context['rejected_schema_objects'])}" - ) - } - - result = prompt( - query="Top customers by invoice amount", - construct_db_schemas=[ - { - "type": "TABLE", - "name": "invoices", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "invoice_amount", - "data_type": "double", - "comment": "", - "is_primary_key": False, - } - ], - } - ], - prompt_builder=PromptBuilder(), - check_using_db_schemas_without_pruning={"db_schemas": []}, - histories=[], - semantic_retry_context={ - "validation_error": "Generic count did not answer invoice amount", - "retry_attempt": 2, - "rejected_schema_objects": ["dbo_ytblES002_1.Name_of_Reported_Received"], - }, - ) - - assert "retry=2" in result["prompt"] - assert "Generic count" in result["prompt"] - assert "dbo_ytblES002_1.Name_of_Reported_Received" in result["prompt"] From da40654c11b083ffd9b48cc4dc76774f56321f69 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:58:48 +0530 Subject: [PATCH 0485/1087] Revert "Make semantic schema mapping drive SQL generation" This reverts commit 188b8a5628950957f9ee29079ebdc21e4de98be7. --- .../generation/followup_sql_generation.py | 14 +- .../followup_sql_generation_reasoning.py | 19 +-- .../pipelines/generation/sql_correction.py | 11 +- .../pipelines/generation/sql_generation.py | 14 +- .../generation/sql_generation_reasoning.py | 19 +-- .../pipelines/generation/sql_regeneration.py | 11 +- .../src/pipelines/generation/utils/sql.py | 116 -------------- .../retrieval/db_schema_retrieval.py | 2 +- wren-ai-service/src/web/v1/services/ask.py | 144 ------------------ .../src/web/v1/services/ask_feedback.py | 96 +----------- .../pipelines/generation/test_sql_utils.py | 64 -------- 11 files changed, 40 insertions(+), 470 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 f0f2412891..a4616ff881 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -17,7 +17,6 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -100,9 +99,12 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -{% if semantic_schema_contract %} -### SEMANTIC SCHEMA CONTRACT ### -{{ semantic_schema_contract }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +This is the pre-generation semantic analysis of the follow-up request against the +active deployed schema. Use it as a contract for table, column, metric, dimension, +filter, time, relationship, aggregation, and ranking selection. +{{ schema_intent_analysis }} {% endif %} ### INTENT AND SCHEMA GROUNDING ### @@ -165,9 +167,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 0374caf609..9759617bb9 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -15,7 +15,6 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, - construct_semantic_schema_contract, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -68,13 +67,13 @@ Language: {{ language }} Current Time: {{ current_time }} -{% if semantic_schema_contract %} -### SEMANTIC SCHEMA CONTRACT ### -Use this contract for entities, metrics, dimensions, filters, joins, time -constraints, aggregations, ranking, and analytical intent. If it shows missing or -ambiguous requirements, state that limitation in the plan instead of planning -unrelated SQL. -{{ semantic_schema_contract }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +Use this semantic analysis as the planning contract for entities, metrics, +dimensions, filters, joins, time constraints, aggregations, ranking, and analytical +intent. If it shows missing or ambiguous requirements, state that limitation in the +plan instead of planning unrelated SQL. +{{ schema_intent_analysis }} {% endif %} Let's think step by step. @@ -103,9 +102,7 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index d1e2e93333..be63983227 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,7 +15,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_sql_generation_model_kwargs, @@ -103,11 +102,11 @@ def get_sql_correction_system_prompt( {% if query %} User's Question: {{ query }} {% endif %} -{% if semantic_schema_contract %} -### SEMANTIC SCHEMA CONTRACT ### +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### This is the semantic contract for the corrected SQL. Preserve this intent while fixing syntax or planner errors. -{{ semantic_schema_contract }} +{{ schema_intent_analysis }} {% endif %} {% if invalid_generation_result.original_sql %} Original SQL: {{ invalid_generation_result.original_sql }} @@ -149,9 +148,7 @@ def prompt( documents=documents, valid_table_names=construct_valid_table_names(documents), invalid_generation_result=invalid_generation_result, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), + schema_intent_analysis=schema_intent_analysis, instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 9c568abec0..9280f3e8f1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,7 +14,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -88,9 +87,12 @@ ### QUESTION ### User's Question: {{ query }} -{% if semantic_schema_contract %} -### SEMANTIC SCHEMA CONTRACT ### -{{ semantic_schema_contract }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +This is the pre-generation semantic analysis of the user's request against the +active deployed schema. Use it as a contract for table, column, metric, dimension, +filter, time, relationship, aggregation, and ranking selection. +{{ schema_intent_analysis }} {% endif %} ### INTENT AND SCHEMA GROUNDING ### @@ -168,9 +170,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} 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 7d14e29184..4db6239cd6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -13,7 +13,6 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, - construct_semantic_schema_contract, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -53,13 +52,13 @@ Language: {{ language }} Current Time: {{ current_time }} -{% if semantic_schema_contract %} -### SEMANTIC SCHEMA CONTRACT ### -Use this contract for entities, metrics, dimensions, filters, joins, time -constraints, aggregations, ranking, and analytical intent. If it shows missing or -ambiguous requirements, state that limitation in the plan instead of planning -unrelated SQL. -{{ semantic_schema_contract }} +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### +Use this semantic analysis as the planning contract for entities, metrics, +dimensions, filters, joins, time constraints, aggregations, ranking, and analytical +intent. If it shows missing or ambiguous requirements, state that limitation in the +plan instead of planning unrelated SQL. +{{ schema_intent_analysis }} {% endif %} Let's think step by step. @@ -86,9 +85,7 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), + schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 4f0668ff0f..bef428a21c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,7 +14,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -110,11 +109,11 @@ def get_sql_regeneration_system_prompt( {% if query %} User's Question: {{ query }} {% endif %} -{% if semantic_schema_contract %} -### SEMANTIC SCHEMA CONTRACT ### +{% if schema_intent_analysis %} +### SCHEMA INTENT ANALYSIS ### This is the semantic contract for regenerated SQL. Preserve this intent while improving the original SQL. -{{ semantic_schema_contract }} +{{ schema_intent_analysis }} {% endif %} SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} @@ -146,9 +145,7 @@ def prompt( data_source=data_source, documents=documents, query=query, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), + schema_intent_analysis=schema_intent_analysis, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 97ae55f7c3..a2d0b3f02d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2221,7 +2221,6 @@ def get_sql_generation_system_prompt( 12. Before generating SQL, validate that the selected schema elements directly support all key entities, metrics, dimensions, filters, time ranges, relationships, and aggregations mentioned or implied by the question. 13. Do not answer a specific business metric, trend, summary, comparison, dashboard, or analysis request with a generic record-count query unless the user explicitly asks only for record count. 14. If the required information cannot be derived from the available active schema, return the closest schema-grounded limitation instead of inventing unrelated SQL. -15. If a SEMANTIC SCHEMA CONTRACT is provided, it is the primary source of truth for selecting tables, columns, metrics, joins, filters, grouping, sorting, and date logic. Generate SQL from the highest-confidence validated concept-to-schema mappings in that contract and do not independently infer substitute schema objects. {text_to_sql_rules} @@ -2300,121 +2299,6 @@ def construct_instructions( return _instructions -def _format_semantic_list(label: str, values: list[str]) -> list[str]: - if not values: - return [] - return [f"{label}: {', '.join(values)}"] - - -def construct_semantic_schema_contract( - semantic_analysis: dict[str, Any] | None, -) -> str: - if not _has_semantic_analysis(semantic_analysis): - return "" - - lines: list[str] = [ - "Use this semantic schema contract as the primary source of truth for SQL generation.", - "Generate SQL only from schema objects listed here or in the selected retrieval metadata.", - "Do not infer alternative tables, columns, metrics, joins, or identifiers independently.", - ] - - analytical_intent = str( - semantic_analysis.get("analytical_intent") or "" - ).strip() - if analytical_intent: - lines.append(f"Analytical intent: {analytical_intent}") - - for label, key in ( - ("Entities", "entities"), - ("Identifiers", "identifiers"), - ("Metrics", "metrics"), - ("Dimensions", "dimensions"), - ("Filters", "filters"), - ("Time constraints", "time_constraints"), - ("Aggregations", "aggregations"), - ("Ranking", "ranking"), - ("Relationships", "relationships"), - ("Supported schema objects", "supported_schema_objects"), - ): - lines.extend(_format_semantic_list(label, _semantic_analysis_items(semantic_analysis, key))) - - concept_mappings = _semantic_concept_mappings(semantic_analysis) - if concept_mappings: - lines.append("Required concept-to-schema mappings:") - for mapping in concept_mappings: - schema_objects = _mapping_schema_objects(mapping) - required = "required" if mapping.get("required_in_sql") is not False else "optional" - confidence = mapping.get("confidence") - confidence_text = ( - f", confidence={confidence}" - if confidence is not None and str(confidence).strip() - else "" - ) - mapping_reason = str(mapping.get("mapping_reason") or "").strip() - reason_text = f" Reason: {mapping_reason}" if mapping_reason else "" - lines.append( - "- " - f"{_mapping_request_concept(mapping)} " - f"({_mapping_concept_type(mapping) or 'concept'}, {required}" - f"{confidence_text}) -> {', '.join(schema_objects) or 'NO_MAPPING'}." - f"{reason_text}" - ) - - interpretations = _semantic_analysis_dict_items( - semantic_analysis, "interpretations" - ) - if interpretations: - lines.append("Ranked schema interpretations:") - for interpretation in interpretations: - description = str(interpretation.get("description") or "").strip() - if not description: - continue - selected = "selected" if interpretation.get("is_selected") is True else "candidate" - confidence = interpretation.get("confidence") - confidence_text = ( - f", confidence={confidence}" - if confidence is not None and str(confidence).strip() - else "" - ) - schema_objects = interpretation.get("schema_objects") - if isinstance(schema_objects, list): - schema_text = ", ".join( - str(item).strip() - for item in schema_objects - if item is not None and str(item).strip() - ) - else: - schema_text = "" - schema_suffix = f" Objects: {schema_text}." if schema_text else "" - lines.append( - f"- {description} ({selected}{confidence_text}).{schema_suffix}" - ) - - missing_requirements = _semantic_analysis_items( - semantic_analysis, "missing_requirements" - ) - if missing_requirements: - lines.append(f"Missing requirements: {', '.join(missing_requirements)}") - - ambiguous_requirements = _semantic_analysis_items( - semantic_analysis, "ambiguous_requirements" - ) - if ambiguous_requirements: - lines.append(f"Ambiguous requirements: {', '.join(ambiguous_requirements)}") - - support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() - if support_reasoning: - lines.append(f"Support reasoning: {support_reasoning}") - - lines.append( - "Validation requirement: every required mapping must be represented in the SQL. " - "Do not substitute identifiers for metrics, entities for identifiers, or COUNT(*) " - "for a requested business measure unless the semantic intent explicitly requests a record count." - ) - - return "\n".join(lines) - - def _parse_semantic_metadata_content(content: str) -> Any | None: content = content.strip() if not content: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 0fc19d58b1..8dbb2c1dfd 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -46,7 +46,7 @@ 8. For each selected table, provide a concise reason for why the table is semantically relevant. 9. For each selected column, provide a concise reason for why the column is necessary. 10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. -11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance, confidence, and schema support; mark the selected interpretation only when it is clearly the best supported one. Keep non-selected high-confidence interpretations so the SQL pipeline can retry the next-best mapping if validation fails. +11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance and mark the selected interpretation only when it is clearly the best supported one. 12. If a "." is included in columns, put the name before the first dot into chosen columns. 13. The number of columns chosen must match the number of reasoning. 14. Final chosen columns must be only column names, don't prefix it with table names. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 5377479691..9731665919 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -6673,151 +6673,7 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": - semantic_retry_error = failed_dry_run_result.get("error", "") - semantic_retry_query = ( - f"{sql_user_query}\n\n" - "The previous generated SQL failed semantic validation: " - f"{semantic_retry_error}\n" - "Re-run semantic schema retrieval and choose the next " - "highest-confidence concept-to-schema mapping that " - "directly supports the user's requested entities, " - "metrics, dimensions, filters, time constraints, " - "aggregations, relationships, and ranking." - ) - logger.info( - "Retrying semantic schema retrieval after intent validation failure for query_id %s", - query_id, - ) - try: - retry_retrieval_result = await self._run_with_timeout( - "Semantic schema retrieval retry", - self._pipelines["db_schema_retrieval"].run( - query=semantic_retry_query, - histories=histories, - project_id=ask_request.project_id, - enable_column_pruning=enable_column_pruning, - ), - timeout_seconds=self._schema_retrieval_timeout_seconds, - ) - retry_construct_result = retry_retrieval_result.get( - "construct_retrieval_results", {} - ) - retry_schema_intent_analysis = retry_construct_result.get( - "semantic_analysis", {} - ) - retry_documents, retry_table_names, retry_table_ddls = ( - self._extract_retrieval_metadata( - retry_retrieval_result - ) - ) - retry_support_error = get_schema_intent_analysis_error( - retry_schema_intent_analysis - ) - if retry_documents and not retry_support_error: - _retrieval_result = retry_construct_result - schema_intent_analysis = retry_schema_intent_analysis - documents = retry_documents - table_names = retry_table_names - table_ddls = retry_table_ddls - has_calculated_field = _retrieval_result.get( - "has_calculated_field", False - ) - has_metric = _retrieval_result.get("has_metric", False) - has_json_field = _retrieval_result.get( - "has_json_field", False - ) - - if histories: - text_to_sql_generation_results = ( - await self._run_with_timeout( - "Follow-up SQL generation after semantic retry", - self._pipelines[ - "followup_sql_generation" - ].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), - ) - ) - else: - text_to_sql_generation_results = ( - await self._run_with_timeout( - "SQL generation after semantic retry", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), - ) - ) - - if sql_valid_result := text_to_sql_generation_results[ - "post_process" - ]["valid_generation_result"]: - if ask_result := self._build_validated_ask_result_from_sql( - sql_valid_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - else: - invalid_sql = sql_valid_result.get("sql") - error_message = ( - "SQL generation after semantic retrieval retry did not produce SQL that matches the active datasource schema and question intent." - ) - else: - failed_dry_run_result = ( - text_to_sql_generation_results["post_process"][ - "invalid_generation_result" - ] - ) - invalid_sql = failed_dry_run_result.get( - "sql", invalid_sql - ) - error_message = failed_dry_run_result.get( - "error", error_message - ) - elif retry_support_error: - error_message = retry_support_error - else: - error_message = ( - "Semantic schema retrieval retry did not find a supported mapping for the request." - ) - except Exception as retry_error: - logger.warning( - "Semantic schema retrieval retry failed for query_id %s: %s", - query_id, - retry_error, - ) - while current_sql_correction_retries < max_sql_correction_retries: - if api_results: - break if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index cbd6f7d169..a9a0c1eabe 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -224,101 +224,7 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] == "SCHEMA_INTENT_VALIDATION": - semantic_retry_error = failed_dry_run_result.get("error", "") - semantic_retry_query = ( - f"{ask_feedback_request.question}\n\n" - "The previous regenerated SQL failed semantic validation: " - f"{semantic_retry_error}\n" - "Re-run semantic schema retrieval and choose the next " - "highest-confidence concept-to-schema mapping that " - "directly supports the user's requested entities, " - "metrics, dimensions, filters, time constraints, " - "aggregations, relationships, and ranking." - ) - logger.info( - "Retrying semantic schema retrieval after feedback intent validation failure for query_id %s", - query_id, - ) - retry_retrieval_result = await self._pipelines[ - "db_schema_retrieval" - ].run( - query=semantic_retry_query, - histories=[], - project_id=ask_feedback_request.project_id, - enable_column_pruning=enable_column_pruning, - ) - retry_construct_result = retry_retrieval_result.get( - "construct_retrieval_results", {} - ) - retry_schema_intent_analysis = retry_construct_result.get( - "semantic_analysis", {} - ) - retry_documents = retry_construct_result.get( - "retrieval_results", [] - ) - retry_table_ddls = [ - document.get("table_ddl") - for document in retry_documents - if isinstance(document, dict) - and document.get("table_ddl") - ] - retry_support_error = get_schema_intent_analysis_error( - retry_schema_intent_analysis - ) - if retry_table_ddls and not retry_support_error: - schema_intent_analysis = retry_schema_intent_analysis - documents = retry_documents - table_ddls = retry_table_ddls - has_calculated_field = retry_construct_result.get( - "has_calculated_field", False - ) - has_metric = retry_construct_result.get( - "has_metric", False - ) - has_json_field = retry_construct_result.get( - "has_json_field", False - ) - - text_to_sql_generation_results = await self._pipelines[ - "sql_regeneration" - ].run( - contexts=table_ddls, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, - sql=ask_feedback_request.sql, - query=ask_feedback_request.question, - project_id=ask_feedback_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ) - - if sql_valid_result := text_to_sql_generation_results[ - "post_process" - ]["valid_generation_result"]: - api_results = [ - AskResult( - **{ - "sql": sql_valid_result.get("sql"), - "type": "llm", - } - ) - ] - else: - failed_dry_run_result = text_to_sql_generation_results[ - "post_process" - ]["invalid_generation_result"] - elif retry_support_error: - error_message = retry_support_error - - if api_results: - pass - elif failed_dry_run_result["type"] not in { + if failed_dry_run_result["type"] not in { "TIME_OUT", "SCHEMA_INTENT_VALIDATION", }: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 4d8797ee50..bc68923326 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,6 +1,5 @@ from src.pipelines.generation.utils.sql import ( contains_unsupported_mssql_json_access, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, extract_sql_generation_result, @@ -1871,66 +1870,3 @@ def test_get_schema_intent_analysis_error_rejects_multiple_selected_interpretati assert error is not None assert "multiple selected schema interpretations" in error - - -def test_construct_semantic_schema_contract_prioritizes_concept_mappings(): - contract = construct_semantic_schema_contract( - { - "analytical_intent": "summary", - "entities": ["invoice"], - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "aggregations": ["sum invoice amount"], - "supported_schema_objects": [ - "invoices.customer_id", - "invoices.invoice_amount", - ], - "concept_mappings": [ - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["invoices.invoice_amount"], - "required_in_sql": True, - "confidence": 0.95, - "mapping_reason": "invoice_amount stores the requested measure", - } - ], - "interpretations": [ - { - "description": "Summarize invoice amount by customer", - "schema_objects": [ - "invoices.customer_id", - "invoices.invoice_amount", - ], - "confidence": 0.9, - "is_selected": True, - } - ], - "is_fully_supported": True, - } - ) - - assert "primary source of truth" in contract - assert "Required concept-to-schema mappings" in contract - assert "invoice amount (metric, required, confidence=0.95)" in contract - assert "invoices.invoice_amount" in contract - assert "Ranked schema interpretations" in contract - assert "selected" in contract - assert "Do not substitute identifiers for metrics" in contract - - -def test_construct_semantic_schema_contract_allows_legacy_analysis_without_mappings(): - contract = construct_semantic_schema_contract( - { - "analytical_intent": "trend", - "metrics": ["order volume"], - "time_constraints": ["monthly"], - "supported_schema_objects": ["orders.created_at", "orders.id"], - "is_fully_supported": True, - } - ) - - assert "Analytical intent: trend" in contract - assert "Metrics: order volume" in contract - assert "orders.created_at" in contract - assert "Required concept-to-schema mappings" not in contract From cbefde96028a0f38645fc14ff50df256911c8c0d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:59:01 +0530 Subject: [PATCH 0486/1087] Revert "Strengthen semantic SQL intent validation" This reverts commit e87181d3a659ee4bd043589af648a6fce4dd9533. --- .../src/pipelines/generation/utils/sql.py | 311 ------------------ .../retrieval/db_schema_retrieval.py | 48 +-- .../pipelines/generation/test_sql_utils.py | 144 -------- .../retrieval/test_db_schema_retrieval.py | 22 -- 4 files changed, 5 insertions(+), 520 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a2d0b3f02d..db9450226f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3502,20 +3502,6 @@ def _semantic_analysis_items( return [] -def _semantic_analysis_dict_items( - semantic_analysis: dict[str, Any] | None, - key: str, -) -> list[dict[str, Any]]: - if not isinstance(semantic_analysis, dict): - return [] - - value = semantic_analysis.get(key) - if not isinstance(value, list): - return [] - - return [item for item in value if isinstance(item, dict)] - - def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: if not isinstance(semantic_analysis, dict) or not semantic_analysis: return False @@ -3531,8 +3517,6 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: "time_constraints", "ranking", "supported_schema_objects", - "concept_mappings", - "interpretations", "missing_requirements", "ambiguous_requirements", "support_reasoning", @@ -3540,42 +3524,6 @@ def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: return any(semantic_analysis.get(key) for key in semantic_keys) -def _schema_interpretation_clarification_error( - semantic_analysis: dict[str, Any], -) -> str | None: - interpretations = _semantic_analysis_dict_items( - semantic_analysis, "interpretations" - ) - if not interpretations: - return None - - selected_interpretations = [ - str(interpretation.get("description") or "").strip() - for interpretation in interpretations - if interpretation.get("is_selected") is True - and str(interpretation.get("description") or "").strip() - ] - if len(selected_interpretations) > 1: - return ( - "The request has multiple selected schema interpretations: " - f"{', '.join(selected_interpretations)}. Please clarify which one to use." - ) - - clarification_interpretations = [ - str(interpretation.get("description") or "").strip() - for interpretation in interpretations - if interpretation.get("needs_clarification") is True - and str(interpretation.get("description") or "").strip() - ] - if clarification_interpretations: - return ( - "The request needs clarification before SQL generation: " - f"{', '.join(clarification_interpretations)}." - ) - - return None - - def get_schema_intent_analysis_error( semantic_analysis: dict[str, Any] | None, ) -> str | None: @@ -3601,11 +3549,6 @@ def get_schema_intent_analysis_error( f"{', '.join(ambiguous_requirements)}. Please clarify which one to use." ) - if interpretation_error := _schema_interpretation_clarification_error( - semantic_analysis - ): - return interpretation_error - if semantic_analysis.get("is_fully_supported") is False: support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() if support_reasoning: @@ -3641,253 +3584,6 @@ def _semantic_analysis_requests_record_count( ) -def _semantic_concept_mappings( - semantic_analysis: dict[str, Any] | None, -) -> list[dict[str, Any]]: - return _semantic_analysis_dict_items(semantic_analysis, "concept_mappings") - - -def _schema_object_parts(schema_object: str) -> list[str]: - return [ - _normalize_sql_identifier(part.strip()) - for part in str(schema_object or "").split(".") - if part.strip() - ] - - -def _schema_object_table_matches( - referenced_table: str, - expected_table: str, -) -> bool: - referenced_suffixes = { - suffix.lower() for suffix in _table_reference_suffixes(referenced_table) - } - expected_suffixes = { - suffix.lower() for suffix in _table_reference_suffixes(expected_table) - } - return bool(referenced_suffixes & expected_suffixes) - - -def _sql_contains_identifier(sql: str, identifier: str) -> bool: - identifier = _normalize_sql_identifier(str(identifier or "").strip()) - if not identifier: - return False - - escaped = re.escape(identifier) - quoted_identifier_pattern = rf'(?:"{escaped}"|`{escaped}`|\[{escaped}\])' - bare_identifier_pattern = rf"(? bool: - table_name = str(table_name or "").strip() - if not table_name: - return False - - table_candidates = {table_name, *_table_reference_suffixes(table_name)} - return any(_sql_contains_identifier(sql, candidate) for candidate in table_candidates) - - -def _sql_references_schema_object( - sql: str, - schema_object: str, - valid_table_columns: dict[str, list[str]], -) -> bool: - parts = _schema_object_parts(schema_object) - if not parts: - return False - - if len(parts) == 1: - return _sql_contains_identifier(sql, parts[0]) - - expected_column = parts[-1] - expected_table = ".".join(parts[:-1]) - aliases = _extract_table_aliases(sql, valid_table_columns) - - for match in _SQL_QUALIFIED_COLUMN_PATTERN.finditer(sql or ""): - qualifier = _normalize_sql_identifier(match.group("qualifier")) - column = _normalize_sql_identifier(match.group("column")) - referenced_table = aliases.get(qualifier.lower(), qualifier) - if ( - column.lower() == expected_column.lower() - and _schema_object_table_matches(referenced_table, expected_table) - ): - return True - - if _sql_references_table(sql, expected_table) and _sql_contains_identifier( - sql, expected_column - ): - return True - - return False - - -def _sql_references_schema_object_table( - sql: str, - schema_object: str, -) -> bool: - parts = _schema_object_parts(schema_object) - if len(parts) < 2: - return _sql_references_table(sql, schema_object) - return _sql_references_table(sql, ".".join(parts[:-1])) - - -def _mapping_schema_objects(mapping: dict[str, Any]) -> list[str]: - value = mapping.get("schema_objects") - if isinstance(value, str): - return [value] if value.strip() else [] - if isinstance(value, list): - return [ - str(item).strip() - for item in value - if item is not None and str(item).strip() - ] - return [] - - -def _mapping_concept_type(mapping: dict[str, Any]) -> str: - return str(mapping.get("concept_type") or "").strip().lower() - - -def _mapping_request_concept(mapping: dict[str, Any]) -> str: - return str(mapping.get("request_concept") or "requested concept").strip() - - -def _requested_aggregate_functions(*texts: str) -> set[str]: - joined = " ".join(text for text in texts if text).lower() - aggregate_terms = { - "SUM": r"\b(?:sum|total)\b", - "AVG": r"\b(?:avg|average|mean)\b", - "MIN": r"\b(?:min|minimum|lowest|smallest)\b", - "MAX": r"\b(?:max|maximum|highest|largest)\b", - "COUNT": r"\b(?:count|number of|how many|volume)\b", - } - return { - function - for function, pattern in aggregate_terms.items() - if re.search(pattern, joined, flags=re.IGNORECASE) - } - - -def _sql_has_aggregate_function(sql: str, function_name: str) -> bool: - return bool( - re.search( - rf"\b{re.escape(function_name)}\s*\(", - sql or "", - flags=re.IGNORECASE, - ) - ) - - -def _validate_sql_against_concept_mappings( - semantic_analysis: dict[str, Any], - sql: str, - valid_table_columns: dict[str, list[str]], -) -> str | None: - mappings = _semantic_concept_mappings(semantic_analysis) - if not mappings: - return None - - requests_record_count = _semantic_analysis_requests_record_count( - semantic_analysis - ) - aggregation_text = " ".join( - _semantic_analysis_items(semantic_analysis, "aggregations") - ) - has_grouping = bool(re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE)) - has_ordering = bool( - re.search(r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE) - ) - has_limit = bool( - re.search( - r"\b(?:LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", - sql or "", - flags=re.IGNORECASE, - ) - ) - - for mapping in mappings: - if mapping.get("required_in_sql") is False: - continue - - concept_type = _mapping_concept_type(mapping) - request_concept = _mapping_request_concept(mapping) - schema_objects = _mapping_schema_objects(mapping) - if not schema_objects: - return ( - "The semantic analysis did not map the required " - f"{concept_type or 'concept'} '{request_concept}' to a schema " - "object. I cannot generate unrelated SQL." - ) - - references_concept = any( - _sql_references_schema_object(sql, schema_object, valid_table_columns) - for schema_object in schema_objects - ) - if ( - not references_concept - and concept_type == "metric" - and requests_record_count - and _sql_is_plain_count(sql) - ): - references_concept = any( - _sql_references_schema_object_table(sql, schema_object) - for schema_object in schema_objects - ) - - if not references_concept: - return ( - "Generated SQL does not reference schema objects mapped to the " - f"required {concept_type or 'concept'} '{request_concept}': " - f"{', '.join(schema_objects)}." - ) - - if concept_type == "metric": - if _sql_is_plain_count(sql) and not requests_record_count: - return ( - "Generated SQL answers with a generic record count, but the " - f"requested metric '{request_concept}' maps to " - f"{', '.join(schema_objects)} and must be retrieved or " - "calculated from that schema object." - ) - - requested_functions = _requested_aggregate_functions( - request_concept, - str(mapping.get("mapping_reason") or ""), - aggregation_text, - ) - for function_name in requested_functions: - if function_name == "COUNT" and requests_record_count: - continue - if not _sql_has_aggregate_function(sql, function_name): - return ( - "Generated SQL does not use the aggregation required " - f"for metric '{request_concept}': {function_name}." - ) - - if concept_type in {"dimension", "time"} and _AGGREGATE_PATTERN.search( - sql or "" - ) and not has_grouping: - return ( - "Generated SQL aggregates results without grouping by the " - f"required {concept_type} '{request_concept}'." - ) - - if concept_type == "ranking" and (not has_ordering or not has_limit): - return ( - "Generated SQL does not include sorting and limiting logic " - f"required by ranking concept '{request_concept}'." - ) - - return None - - def _validate_sql_against_semantic_analysis( semantic_analysis: dict[str, Any] | None, sql: str, @@ -3899,13 +3595,6 @@ def _validate_sql_against_semantic_analysis( if analysis_error := get_schema_intent_analysis_error(semantic_analysis): return analysis_error - if mapping_error := _validate_sql_against_concept_mappings( - semantic_analysis, - sql, - valid_table_columns, - ): - return mapping_error - analytical_intent = str( semantic_analysis.get("analytical_intent") or "" ).strip().lower() diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 8dbb2c1dfd..102f5dcd92 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -45,13 +45,11 @@ 7. Set `is_fully_supported` to false when any required request component is missing or ambiguous. 8. For each selected table, provide a concise reason for why the table is semantically relevant. 9. For each selected column, provide a concise reason for why the column is necessary. -10. Populate `concept_mappings` for every important concept in the request. Each mapping must classify the concept, list only directly supporting schema objects, state whether it must appear in SQL, and include a confidence score between 0 and 1. -11. Populate `interpretations` when the request has more than one plausible schema interpretation. Rank interpretations by semantic relevance and mark the selected interpretation only when it is clearly the best supported one. -12. If a "." is included in columns, put the name before the first dot into chosen columns. -13. The number of columns chosen must match the number of reasoning. -14. Final chosen columns must be only column names, don't prefix it with table names. -15. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -16. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +10. If a "." is included in columns, put the name before the first dot into chosen columns. +11. The number of columns chosen must match the number of reasoning. +12. Final chosen columns must be only column names, don't prefix it with table names. +13. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +14. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -69,24 +67,6 @@ "time_constraints": ["time filters, grains, or trend requirements"], "ranking": ["top/bottom/order/limit requirements"], "supported_schema_objects": ["table.column or metric names that directly support the request"], - "concept_mappings": [ - { - "request_concept": "business concept from the user request", - "concept_type": "entity | identifier | dimension | metric | filter | time | aggregation | ranking | relationship | comparison", - "schema_objects": ["table.column, table, view, metric, or relationship object that directly supports the concept"], - "required_in_sql": true, - "confidence": 0.0, - "mapping_reason": "Why these schema objects semantically support the concept" - } - ], - "interpretations": [ - { - "description": "Possible interpretation of the request", - "schema_objects": ["schema objects used by this interpretation"], - "confidence": 0.0, - "is_selected": true - } - ], "missing_requirements": ["required request components not supported by the schema"], "ambiguous_requirements": ["request components with multiple equally plausible schema mappings"], "is_fully_supported": true, @@ -543,22 +523,6 @@ class MatchingTable(BaseModel): table_selection_reason: str -class SemanticConceptMapping(BaseModel): - request_concept: str = "" - concept_type: str = "" - schema_objects: list[str] = Field(default_factory=list) - required_in_sql: bool = True - confidence: float | None = None - mapping_reason: str = "" - - -class SemanticInterpretation(BaseModel): - description: str = "" - schema_objects: list[str] = Field(default_factory=list) - confidence: float | None = None - is_selected: bool = False - - class SemanticAnalysis(BaseModel): analytical_intent: str = "" entities: list[str] = Field(default_factory=list) @@ -571,8 +535,6 @@ class SemanticAnalysis(BaseModel): time_constraints: list[str] = Field(default_factory=list) ranking: list[str] = Field(default_factory=list) supported_schema_objects: list[str] = Field(default_factory=list) - concept_mappings: list[SemanticConceptMapping] = Field(default_factory=list) - interpretations: list[SemanticInterpretation] = Field(default_factory=list) missing_requirements: list[str] = Field(default_factory=list) ambiguous_requirements: list[str] = Field(default_factory=list) is_fully_supported: bool | None = None diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index bc68923326..e2fc3f7922 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1726,147 +1726,3 @@ def test_validate_sql_intent_alignment_allows_semantic_count_metric(): ) assert error is None - - -def test_validate_sql_intent_alignment_rejects_unmapped_metric_substitution(): - error = validate_sql_intent_alignment( - "Show invoice amount by customer", - 'SELECT "invoices"."customer_id", COUNT(*) AS "invoice_amount" ' - 'FROM "invoices" GROUP BY "invoices"."customer_id"', - {"invoices": ["customer_id", "invoice_amount"]}, - semantic_analysis={ - "analytical_intent": "summary", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "aggregations": ["sum invoice amount"], - "concept_mappings": [ - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["invoices.invoice_amount"], - "required_in_sql": True, - "confidence": 0.95, - } - ], - "is_fully_supported": True, - }, - ) - - assert error is not None - assert "invoice amount" in error - assert "invoices.invoice_amount" in error - - -def test_validate_sql_intent_alignment_allows_mapped_metric_sql(): - error = validate_sql_intent_alignment( - "Show total invoice amount by customer", - 'SELECT "invoices"."customer_id", ' - 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' - 'FROM "invoices" GROUP BY "invoices"."customer_id"', - {"invoices": ["customer_id", "invoice_amount"]}, - semantic_analysis={ - "analytical_intent": "summary", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "aggregations": ["sum invoice amount"], - "concept_mappings": [ - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["invoices.invoice_amount"], - "required_in_sql": True, - "confidence": 0.95, - }, - { - "request_concept": "customer", - "concept_type": "dimension", - "schema_objects": ["invoices.customer_id"], - "required_in_sql": True, - "confidence": 0.9, - }, - ], - "is_fully_supported": True, - }, - ) - - assert error is None - - -def test_validate_sql_intent_alignment_rejects_missing_mapped_dimension(): - error = validate_sql_intent_alignment( - "Show total invoice amount by customer", - 'SELECT SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' - 'FROM "invoices"', - {"invoices": ["customer_id", "invoice_amount"]}, - semantic_analysis={ - "analytical_intent": "summary", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "aggregations": ["sum invoice amount"], - "concept_mappings": [ - { - "request_concept": "customer", - "concept_type": "dimension", - "schema_objects": ["invoices.customer_id"], - "required_in_sql": True, - } - ], - "is_fully_supported": True, - }, - ) - - assert error is not None - assert "customer" in error - - -def test_validate_sql_intent_alignment_rejects_mapped_ranking_without_limit(): - error = validate_sql_intent_alignment( - "Top customers by invoice amount", - 'SELECT "invoices"."customer_id", ' - 'SUM("invoices"."invoice_amount") AS "total_invoice_amount" ' - 'FROM "invoices" GROUP BY "invoices"."customer_id" ' - 'ORDER BY "total_invoice_amount" DESC', - {"invoices": ["customer_id", "invoice_amount"]}, - semantic_analysis={ - "analytical_intent": "ranking", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "ranking": ["top customers by invoice amount"], - "aggregations": ["sum invoice amount"], - "concept_mappings": [ - { - "request_concept": "top customers", - "concept_type": "ranking", - "schema_objects": ["invoices.customer_id"], - "required_in_sql": True, - } - ], - "is_fully_supported": True, - }, - ) - - assert error is not None - assert "sorting and limiting" in error - - -def test_get_schema_intent_analysis_error_rejects_multiple_selected_interpretations(): - error = get_schema_intent_analysis_error( - { - "interpretations": [ - { - "description": "Use gross amount", - "confidence": 0.88, - "is_selected": True, - }, - { - "description": "Use net amount", - "confidence": 0.87, - "is_selected": True, - }, - ], - "is_fully_supported": True, - } - ) - - assert error is not None - assert "multiple selected schema interpretations" in error diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index afb2d06c8b..6adbad941d 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -102,24 +102,6 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): "entities": ["invoice"], "metrics": ["invoice amount"], "dimensions": ["customer"], - "concept_mappings": [ - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["invoices.invoice_amount"], - "required_in_sql": true, - "confidence": 0.95, - "mapping_reason": "invoice_amount stores invoice value" - } - ], - "interpretations": [ - { - "description": "Summarize invoice amount by customer", - "schema_objects": ["invoices.customer_id", "invoices.invoice_amount"], - "confidence": 0.9, - "is_selected": true - } - ], "is_fully_supported": true }, "results": [ @@ -173,10 +155,6 @@ def test_construct_retrieval_results_preserves_semantic_analysis(): ) assert result["semantic_analysis"]["metrics"] == ["invoice amount"] - assert result["semantic_analysis"]["concept_mappings"][0]["schema_objects"] == [ - "invoices.invoice_amount" - ] - assert result["semantic_analysis"]["interpretations"][0]["is_selected"] is True assert result["retrieval_results"][0]["table_name"] == "invoices" assert "invoice_amount" in result["retrieval_results"][0]["table_ddl"] assert "internal_note" not in result["retrieval_results"][0]["table_ddl"] From 3ec70f623b435492da73735d178484551f3ff742 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 10 Jul 2026 22:59:16 +0530 Subject: [PATCH 0487/1087] Revert "Improve semantic SQL intent validation" This reverts commit dd5e7149e624953f5f1138561d6db57c48984c29. --- .../generation/followup_sql_generation.py | 20 - .../followup_sql_generation_reasoning.py | 13 - .../pipelines/generation/sql_correction.py | 18 - .../pipelines/generation/sql_generation.py | 20 - .../generation/sql_generation_reasoning.py | 13 - .../pipelines/generation/sql_regeneration.py | 25 - .../src/pipelines/generation/utils/sql.py | 569 ------------------ .../retrieval/db_schema_retrieval.py | 73 +-- wren-ai-service/src/web/v1/services/ask.py | 82 +-- .../src/web/v1/services/ask_feedback.py | 29 +- .../pipelines/generation/test_sql_utils.py | 176 ------ .../retrieval/test_db_schema_retrieval.py | 71 --- 12 files changed, 23 insertions(+), 1086 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 a4616ff881..c9f8b23537 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -99,14 +99,6 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -This is the pre-generation semantic analysis of the follow-up request against the -active deployed schema. Use it as a contract for table, column, metric, dimension, -filter, time, relationship, aggregation, and ranking selection. -{{ schema_intent_analysis }} -{% endif %} - ### INTENT AND SCHEMA GROUNDING ### Interpret the user's business terms by matching them to explicit tables, columns, metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Never reuse table @@ -114,10 +106,6 @@ appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. Only apply aggregate functions to columns whose active metadata type supports that operation. -Before writing SQL, validate that the selected schema elements directly support every -key entity, metric, dimension, filter, time range, relationship, and aggregation in -the follow-up question. If the schema cannot support the requested information, do -not replace the request with a generic COUNT(*) or unrelated table query. ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -141,7 +129,6 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -167,7 +154,6 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -200,11 +186,9 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: list[str], data_source: str, - query: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -214,8 +198,6 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -267,7 +249,6 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -287,7 +268,6 @@ async def run( "has_metric": has_metric, "has_json_field": has_json_field, "sql_functions": sql_functions, - "schema_intent_analysis": schema_intent_analysis, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 9759617bb9..abbbb81d56 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -67,15 +67,6 @@ Language: {{ language }} Current Time: {{ current_time }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -Use this semantic analysis as the planning contract for entities, metrics, -dimensions, filters, joins, time constraints, aggregations, ranking, and analytical -intent. If it shows missing or ambiguous requirements, state that limitation in the -plan instead of planning unrelated SQL. -{{ schema_intent_analysis }} -{% endif %} - Let's think step by step. """ @@ -90,7 +81,6 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -102,7 +92,6 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), - schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -196,7 +185,6 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("Followup SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -209,7 +197,6 @@ async def run( "instructions": instructions or [], "configuration": configuration, "query_id": query_id, - "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index be63983227..daae37e02e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -102,12 +102,6 @@ def get_sql_correction_system_prompt( {% if query %} User's Question: {{ query }} {% endif %} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -This is the semantic contract for the corrected SQL. Preserve this intent while -fixing syntax or planner errors. -{{ schema_intent_analysis }} -{% endif %} {% if invalid_generation_result.original_sql %} Original SQL: {{ invalid_generation_result.original_sql }} {% endif %} @@ -121,10 +115,6 @@ def get_sql_correction_system_prompt( user's request. Do not invent tables, columns, joins, metrics, or relationships. Only apply aggregate functions to columns whose active metadata type supports that operation. -Before returning corrected SQL, validate that it still directly supports every key -entity, metric, dimension, filter, time range, relationship, and aggregation in the -user's question. Do not replace an unsupported request with a generic COUNT(*) or -unrelated table query. Let's think step by step. """ @@ -140,7 +130,6 @@ def prompt( query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -148,7 +137,6 @@ def prompt( documents=documents, valid_table_names=construct_valid_table_names(documents), invalid_generation_result=invalid_generation_result, - schema_intent_analysis=schema_intent_analysis, instructions=construct_instructions( instructions=instructions, ), @@ -181,11 +169,9 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: List[Document], data_source: str, - query: str | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -195,8 +181,6 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -243,7 +227,6 @@ async def run( allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -255,7 +238,6 @@ async def run( "invalid_generation_result": invalid_generation_result, "documents": contexts, "query": query, - "schema_intent_analysis": schema_intent_analysis, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 9280f3e8f1..59bea2279c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -87,24 +87,12 @@ ### QUESTION ### User's Question: {{ query }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -This is the pre-generation semantic analysis of the user's request against the -active deployed schema. Use it as a contract for table, column, metric, dimension, -filter, time, relationship, aggregation, and ranking selection. -{{ schema_intent_analysis }} -{% endif %} - ### INTENT AND SCHEMA GROUNDING ### Interpret the user's business terms by matching them to explicit tables, columns, metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Do not answer with general guidance when the question can be answered with SQL over the active metadata. Never reuse table or column names from SQL SAMPLES unless those exact names also appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. -Before writing SQL, validate that the selected schema elements directly support every -key entity, metric, dimension, filter, time range, relationship, and aggregation in -the question. If the schema cannot support the requested information, do not replace -the request with a generic COUNT(*) or unrelated table query. {% if sql_generation_reasoning %} ### REASONING PLAN ### @@ -130,7 +118,6 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: schema_context = "\n".join(documents or []).lower() has_pcb_context = any( @@ -170,7 +157,6 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -199,12 +185,10 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: list[str], data_source: str, - query: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -215,8 +199,6 @@ async def post_process( allow_data_preview=allow_data_preview, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -268,7 +250,6 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -287,7 +268,6 @@ async def run( "has_metric": has_metric, "has_json_field": has_json_field, "sql_functions": sql_functions, - "schema_intent_analysis": schema_intent_analysis, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), 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 4db6239cd6..f91a4288e4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -52,15 +52,6 @@ Language: {{ language }} Current Time: {{ current_time }} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -Use this semantic analysis as the planning contract for entities, metrics, -dimensions, filters, joins, time constraints, aggregations, ranking, and analytical -intent. If it shows missing or ambiguous requirements, state that limitation in the -plan instead of planning unrelated SQL. -{{ schema_intent_analysis }} -{% endif %} - Let's think step by step. """ @@ -74,7 +65,6 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -85,7 +75,6 @@ def prompt( ), language=configuration.language, current_time=configuration.show_current_time(), - schema_intent_analysis=schema_intent_analysis, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -174,7 +163,6 @@ async def run( instructions: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -186,7 +174,6 @@ async def run( "instructions": instructions or [], "configuration": configuration, "query_id": query_id, - "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index bef428a21c..8b3eaafcb1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -44,10 +44,6 @@ def get_sql_regeneration_system_prompt( please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. While generating the new SQL query, you should use the original SQL query as a reference. While generating the new SQL query, make sure to use the database schema to generate the SQL query. -Before returning SQL, validate that the selected schema elements directly support -the key entities, metrics, dimensions, filters, time ranges, relationships, and -aggregations from the user's question or reasoning. Do not replace an unsupported -request with a generic COUNT(*) or unrelated table query. {text_to_sql_rules} @@ -106,15 +102,6 @@ def get_sql_regeneration_system_prompt( {% endif %} ### QUESTION ### -{% if query %} -User's Question: {{ query }} -{% endif %} -{% if schema_intent_analysis %} -### SCHEMA INTENT ANALYSIS ### -This is the semantic contract for regenerated SQL. Preserve this intent while -improving the original SQL. -{{ schema_intent_analysis }} -{% endif %} SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} @@ -130,8 +117,6 @@ def prompt( sql: str, prompt_builder: PromptBuilder, data_source: str, - query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -144,8 +129,6 @@ def prompt( sql=sql, data_source=data_source, documents=documents, - query=query, - schema_intent_analysis=schema_intent_analysis, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -193,9 +176,7 @@ async def post_process( post_processor: SQLGenPostProcessor, documents: list[str], data_source: str, - query: str | None = None, project_id: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), @@ -203,8 +184,6 @@ async def post_process( data_source=data_source, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -249,8 +228,6 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -260,8 +237,6 @@ async def run( "documents": contexts, "sql_generation_reasoning": sql_generation_reasoning, "sql": sql, - "query": query, - "schema_intent_analysis": schema_intent_analysis, "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index db9450226f..8316d24e86 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1543,8 +1543,6 @@ async def run( allow_data_preview: bool = False, valid_table_names: list[str] | None = None, valid_table_columns: dict[str, list[str]] | None = None, - query: str | None = None, - semantic_analysis: dict[str, Any] | None = None, ) -> dict: try: cleaned_generation_result = extract_sql_generation_result(replies[0]) @@ -1619,24 +1617,6 @@ async def run( }, } - intent_validation_error = validate_sql_intent_alignment( - query, - cleaned_generation_result, - valid_table_columns or {}, - semantic_analysis=semantic_analysis, - ) - if intent_validation_error: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_INTENT_VALIDATION", - "error": intent_validation_error, - "correlation_id": "", - }, - } - if normalize_data_source( data_source ) == "MSSQL" and contains_unsupported_mssql_json_access( @@ -2218,9 +2198,6 @@ def get_sql_generation_system_prompt( 9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. 10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. 11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. -12. Before generating SQL, validate that the selected schema elements directly support all key entities, metrics, dimensions, filters, time ranges, relationships, and aggregations mentioned or implied by the question. -13. Do not answer a specific business metric, trend, summary, comparison, dashboard, or analysis request with a generic record-count query unless the user explicitly asks only for record count. -14. If the required information cannot be derived from the available active schema, return the closest schema-grounded limitation instead of inventing unrelated SQL. {text_to_sql_rules} @@ -3208,552 +3185,6 @@ def format_valid_table_columns(valid_table_columns: dict[str, list[str]]) -> str ) -_PLAIN_COUNT_SQL_PATTERN = re.compile( - r"^\s*SELECT\s+COUNT\s*\(\s*\*\s*\)(?:\s+AS\s+" - r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*))?\s+' - rf"FROM\s+{_SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SQL_IDENTIFIER_PATTERN})*" - r"(?:\s+WHERE\s+.+?)?\s*(?:ORDER\s+BY\s+.+?)?(?:LIMIT\s+\d+\s*)?$", - flags=re.IGNORECASE | re.DOTALL, -) -_AGGREGATE_PATTERN = re.compile( - r"\b(?:SUM|AVG|MIN|MAX|COUNT)\s*\(", - flags=re.IGNORECASE, -) -_SPECIFIC_METRIC_TERM_GROUPS: dict[str, tuple[str, ...]] = { - "revenue": ("revenue", "sales", "sale", "amount", "value", "price", "total"), - "sales": ("sales", "sale", "revenue", "amount", "value", "price"), - "profit": ("profit", "margin", "income", "earnings"), - "cost": ("cost", "expense", "spend", "charge"), - "amount": ("amount", "value", "price", "total", "sum"), - "value": ("value", "amount", "price", "total"), - "quantity": ("quantity", "qty", "volume", "units", "count"), - "average": ("average", "avg", "mean"), - "avg": ("avg", "average", "mean"), - "rate": ("rate", "ratio", "percent", "percentage"), - "ratio": ("ratio", "rate", "percent", "percentage"), - "percentage": ("percentage", "percent", "pct", "rate"), - "percent": ("percent", "percentage", "pct", "rate"), - "duration": ("duration", "turnaround", "elapsed", "cycle", "leadtime", "time"), - "turnaround": ("turnaround", "duration", "elapsed", "cycle", "leadtime", "time"), -} -_ANALYSIS_TERMS = { - "analysis", - "analyze", - "dashboard", - "summary", - "summarize", - "trend", - "compare", - "comparison", - "breakdown", - "distribution", - "performance", - "ranking", - "top", - "bottom", - "highest", - "lowest", -} -_COUNT_VOLUME_TERMS = { - "count", - "counts", - "number", - "volume", - "records", - "record", - "rows", - "row", - "how many", -} -_TEMPORAL_TERMS = { - "date", - "timestamp", - "month", - "monthly", - "year", - "yearly", - "week", - "weekly", - "day", - "daily", - "quarter", - "quarterly", - "trend", - "over time", -} -_TEMPORAL_IDENTIFIER_TERMS = { - "date", - "time", - "timestamp", - "month", - "year", - "week", - "day", - "quarter", - "created", - "updated", - "modified", - "started", - "ended", - "closed", - "approved", -} -_DIMENSION_TERMS = { - "category", - "type", - "status", - "source", - "region", - "country", - "market", - "customer", - "product", - "salesperson", - "owner", - "assignee", - "division", - "department", - "location", -} - - -def _contains_phrase(text: str, terms: set[str] | tuple[str, ...]) -> bool: - normalized = f" {str(text or '').lower()} " - return any(f" {term.lower()} " in normalized for term in terms if " " in term) or any( - re.search(rf"\b{re.escape(term.lower())}\b", normalized) - for term in terms - if " " not in term - ) - - -def _query_requests_specific_metric(query: str) -> bool: - normalized = str(query or "").lower() - if not normalized: - return False - - return any( - re.search(rf"\b{re.escape(term)}\b", normalized) - for term in _SPECIFIC_METRIC_TERM_GROUPS - ) - - -def _query_requests_count_volume(query: str) -> bool: - return _contains_phrase(query, _COUNT_VOLUME_TERMS) - - -def _query_requests_time_analysis(query: str) -> bool: - normalized = str(query or "").lower() - if _contains_phrase(normalized, _TEMPORAL_TERMS): - return True - - return bool( - re.search( - r"\b(?:last|next|previous|prior|this)\s+" - r"(?:\d+\s+)?(?:day|week|month|quarter|year)s?\b", - normalized, - ) - or re.search(r"\b(?:between|since|before|after)\b", normalized) - ) - - -def _query_requests_time_bucket(query: str) -> bool: - normalized = str(query or "").lower() - return bool( - _contains_phrase( - normalized, - { - "daily", - "weekly", - "monthly", - "quarterly", - "yearly", - "trend", - "over time", - "time series", - }, - ) - or re.search( - r"\b(?:by|per|for each)\s+" - r"(?:day|week|month|quarter|year)s?\b", - normalized, - ) - ) - - -def _query_requests_grouped_analysis(query: str) -> bool: - normalized = str(query or "").lower() - return bool( - re.search(r"\b(?:by|per|across)\s+[A-Za-z_][A-Za-z0-9_ -]*", normalized) - or re.search(r"\bfor each\s+[A-Za-z_][A-Za-z0-9_ -]*", normalized) - or _contains_phrase(normalized, {"breakdown", "distribution", "grouped"}) - ) - - -def _sql_is_plain_count(sql: str) -> bool: - if re.search(r"\bGROUP\s+BY\b", sql, flags=re.IGNORECASE): - return False - return bool(_PLAIN_COUNT_SQL_PATTERN.match(sql or "")) - - -def _compacted_schema_identifiers( - valid_table_columns: dict[str, list[str]], -) -> set[str]: - identifiers: set[str] = set() - for table, columns in valid_table_columns.items(): - identifiers.add(_compact_sql_identifier(table)) - for column in columns or []: - identifiers.add(_compact_sql_identifier(column)) - return {identifier for identifier in identifiers if identifier} - - -def _compacted_sql_identifiers(sql: str) -> set[str]: - identifiers = { - _compact_sql_identifier(identifier) - for identifier in re.findall( - r'"([^"]+)"|`([^`]+)`|\[([^\]]+)\]|\b([A-Za-z_][A-Za-z0-9_$]*)\b', - sql or "", - ) - for identifier in identifier - if identifier - } - return {identifier for identifier in identifiers if identifier} - - -def _sql_references_term_group(sql: str, terms: tuple[str, ...]) -> bool: - sql_identifiers = _compacted_sql_identifiers(sql) - compact_terms = {_compact_sql_identifier(term) for term in terms} - return any( - term - and any(term in identifier or identifier in term for identifier in sql_identifiers) - for term in compact_terms - ) - - -def _schema_supports_term_group( - valid_table_columns: dict[str, list[str]], - terms: tuple[str, ...], -) -> bool: - schema_identifiers = _compacted_schema_identifiers(valid_table_columns) - compact_terms = {_compact_sql_identifier(term) for term in terms} - return any( - term - and any(term in identifier or identifier in term for identifier in schema_identifiers) - for term in compact_terms - ) - - -def _missing_metric_support( - query: str, - sql: str, - valid_table_columns: dict[str, list[str]], -) -> list[str]: - normalized_query = str(query or "").lower() - missing_terms: list[str] = [] - - for term, group in _SPECIFIC_METRIC_TERM_GROUPS.items(): - if not re.search(rf"\b{re.escape(term)}\b", normalized_query): - continue - if _sql_references_term_group(sql, group): - continue - if not _schema_supports_term_group(valid_table_columns, group): - missing_terms.append(term) - - return sorted(set(missing_terms)) - - -def _sql_has_temporal_reference( - sql: str, - valid_table_columns: dict[str, list[str]], -) -> bool: - if re.search( - r"\b(?:DATEPART|DATE_TRUNC|DATETRUNC|EXTRACT|TO_TIMESTAMP|CAST)\s*\(", - sql or "", - flags=re.IGNORECASE, - ): - return True - - return _sql_references_term_group( - sql, - tuple(_TEMPORAL_IDENTIFIER_TERMS), - ) or any( - _schema_supports_term_group({table: [column]}, tuple(_TEMPORAL_IDENTIFIER_TERMS)) - and _sql_references_term_group(sql, (column,)) - for table, columns in valid_table_columns.items() - for column in columns - ) - - -def _semantic_analysis_items( - semantic_analysis: dict[str, Any] | None, - key: str, -) -> list[str]: - if not isinstance(semantic_analysis, dict): - return [] - - value = semantic_analysis.get(key) - if isinstance(value, str): - return [value] if value.strip() else [] - if isinstance(value, list): - return [ - str(item).strip() - for item in value - if item is not None and str(item).strip() - ] - return [] - - -def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: - if not isinstance(semantic_analysis, dict) or not semantic_analysis: - return False - semantic_keys = { - "analytical_intent", - "entities", - "identifiers", - "metrics", - "dimensions", - "filters", - "aggregations", - "relationships", - "time_constraints", - "ranking", - "supported_schema_objects", - "missing_requirements", - "ambiguous_requirements", - "support_reasoning", - } - return any(semantic_analysis.get(key) for key in semantic_keys) - - -def get_schema_intent_analysis_error( - semantic_analysis: dict[str, Any] | None, -) -> str | None: - if not _has_semantic_analysis(semantic_analysis): - return None - - missing_requirements = _semantic_analysis_items( - semantic_analysis, "missing_requirements" - ) - if missing_requirements: - return ( - "The active datasource schema does not expose the information needed " - "to answer the request: " - f"{', '.join(missing_requirements)}. I cannot generate unrelated SQL." - ) - - ambiguous_requirements = _semantic_analysis_items( - semantic_analysis, "ambiguous_requirements" - ) - if ambiguous_requirements: - return ( - "The request has multiple equally plausible schema interpretations: " - f"{', '.join(ambiguous_requirements)}. Please clarify which one to use." - ) - - if semantic_analysis.get("is_fully_supported") is False: - support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() - if support_reasoning: - return ( - "The selected schema does not fully support the request: " - f"{support_reasoning}" - ) - return ( - "The selected schema does not fully support every required component " - "of the request. I cannot generate unrelated SQL." - ) - - return None - - -def _semantic_analysis_requests_record_count( - semantic_analysis: dict[str, Any], -) -> bool: - analytical_intent = str(semantic_analysis.get("analytical_intent") or "").lower() - if analytical_intent == "record_count": - return True - - semantic_text = " ".join( - item - for key in ("metrics", "aggregations") - for item in _semantic_analysis_items(semantic_analysis, key) - ) - return bool( - re.search( - r"\b(?:count|number of|volume|record count|row count|count records|count rows|number of records)\b", - semantic_text.lower(), - ) - ) - - -def _validate_sql_against_semantic_analysis( - semantic_analysis: dict[str, Any] | None, - sql: str, - valid_table_columns: dict[str, list[str]], -) -> str | None: - if not _has_semantic_analysis(semantic_analysis): - return None - - if analysis_error := get_schema_intent_analysis_error(semantic_analysis): - return analysis_error - - analytical_intent = str( - semantic_analysis.get("analytical_intent") or "" - ).strip().lower() - metrics = _semantic_analysis_items(semantic_analysis, "metrics") - dimensions = _semantic_analysis_items(semantic_analysis, "dimensions") - aggregations = _semantic_analysis_items(semantic_analysis, "aggregations") - time_constraints = _semantic_analysis_items( - semantic_analysis, "time_constraints" - ) - ranking = _semantic_analysis_items(semantic_analysis, "ranking") - requests_record_count = _semantic_analysis_requests_record_count( - semantic_analysis - ) - analytical_sql_intents = { - "summary", - "comparison", - "trend", - "dashboard", - "kpi", - "ranking", - } - - if _sql_is_plain_count(sql) and ( - (metrics and not requests_record_count) - or dimensions - or ranking - or (analytical_intent in analytical_sql_intents and not requests_record_count) - ): - return ( - "Generated SQL answers with a generic record count, but the semantic " - "analysis requires specific business metrics, dimensions, ranking, " - "or analytical calculations from the active schema." - ) - - if time_constraints and not _sql_has_temporal_reference(sql, valid_table_columns): - return ( - "Generated SQL does not use a temporal field or supported date/time " - "expression, but the semantic analysis identified time constraints " - "or trend requirements." - ) - - if (dimensions or analytical_intent == "trend") and _AGGREGATE_PATTERN.search( - sql or "" - ): - has_grouping = bool(re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE)) - if not has_grouping: - return ( - "Generated SQL aggregates results without grouping by the " - "dimensions or time grain identified in the semantic analysis." - ) - - if ranking and not re.search( - r"\b(?:ORDER\s+BY|LIMIT|TOP\s*\(|FETCH\s+FIRST)\b", - sql or "", - flags=re.IGNORECASE, - ): - return ( - "Generated SQL does not include sorting or limiting logic required " - "by the ranking intent." - ) - - if aggregations and not _AGGREGATE_PATTERN.search(sql or ""): - return ( - "Generated SQL does not include the aggregation required by the " - "semantic analysis." - ) - - return None - - -def validate_sql_intent_alignment( - query: str | None, - sql: str, - valid_table_columns: dict[str, list[str]] | None = None, - semantic_analysis: dict[str, Any] | None = None, -) -> str | None: - valid_table_columns = valid_table_columns or {} - - semantic_validation_error = _validate_sql_against_semantic_analysis( - semantic_analysis, - sql, - valid_table_columns, - ) - if semantic_validation_error: - return semantic_validation_error - - if not query: - return None - - normalized_query = str(query or "").lower() - asks_specific_metric = _query_requests_specific_metric(normalized_query) - asks_count_volume = _query_requests_count_volume(normalized_query) - asks_time_analysis = _query_requests_time_analysis(normalized_query) - asks_time_bucket = _query_requests_time_bucket(normalized_query) - asks_grouped_analysis = _query_requests_grouped_analysis(normalized_query) - asks_analysis = _contains_phrase(normalized_query, _ANALYSIS_TERMS) - - if _sql_is_plain_count(sql) and ( - asks_specific_metric - or asks_time_bucket - or asks_grouped_analysis - or (asks_analysis and not asks_count_volume) - ): - return ( - "Generated SQL answers with a generic record count, but the question " - "asks for a specific metric, trend, grouping, comparison, dashboard, " - "or analysis. Select schema elements that directly support the " - "requested business intent, or report that the schema does not expose them." - ) - - if asks_time_analysis and not _sql_has_temporal_reference(sql, valid_table_columns): - return ( - "Generated SQL does not use a temporal field or supported date/time " - "expression, but the question asks for a time range or trend. The " - "active schema must expose a relevant date/time column to answer this." - ) - - if (asks_grouped_analysis or asks_time_bucket) and _AGGREGATE_PATTERN.search( - sql or "" - ): - has_grouping = bool(re.search(r"\bGROUP\s+BY\b", sql or "", flags=re.IGNORECASE)) - if not has_grouping: - return ( - "Generated SQL aggregates results without grouping by the requested " - "dimension. Use an explicit schema column for the requested grouping, " - "or report that the schema does not expose that dimension." - ) - - missing_metric_terms = _missing_metric_support( - normalized_query, - sql, - valid_table_columns, - ) - if missing_metric_terms: - missing = ", ".join(missing_metric_terms) - return ( - "The active schema does not expose columns or metrics that directly " - f"support the requested business term(s): {missing}. Do not generate " - "an unrelated query." - ) - - if asks_grouped_analysis and _contains_phrase(normalized_query, _DIMENSION_TERMS): - missing_dimensions = [ - term - for term in _DIMENSION_TERMS - if re.search(rf"\b{re.escape(term)}\b", normalized_query) - and not _sql_references_term_group(sql, (term,)) - and not _schema_supports_term_group(valid_table_columns, (term,)) - ] - if missing_dimensions: - return ( - "The active schema does not expose the requested grouping " - f"dimension(s): {', '.join(sorted(missing_dimensions))}. Do not " - "generate an unrelated query." - ) - - return None - - def construct_ask_history_messages( histories: list[Any] | list[dict], ) -> list[ChatMessage]: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 102f5dcd92..107c5c3479 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -10,7 +10,7 @@ from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel, Field +from pydantic import BaseModel from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider @@ -31,47 +31,23 @@ table_columns_selection_system_prompt = """ ### TASK ### -You are a highly skilled data analyst. Your goal is to examine the provided active deployed database schema, interpret the posed question, and identify the specific tables, columns, metrics, views, and relationships required to construct an accurate SQL query. +You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### -1. First perform a semantic analysis of the user's request. Identify intended business entities, identifiers, descriptive attributes, metrics, dimensions, filters, aggregations, relationships, time constraints, ranking requirements, and analytical intent such as retrieval, detailed records, summary, comparison, trend analysis, dashboard, KPI, ranking, or record count. -2. Map each business term to explicit schema objects only when the active schema directly supports that term. Distinguish entities such as customer/order/invoice/product from identifiers such as order ID or invoice number, descriptive attributes, and measurable metrics such as amount, quantity, cost, profit, revenue, or duration. -3. Select tables and columns by semantic fit to the full request, not by isolated keyword overlap or commonly used default tables. -4. Include join keys and relationship columns needed to connect selected tables. Do not invent relationships or foreign keys. -5. If the schema does not support a requested entity, metric, dimension, filter, time range, aggregation, or ranking requirement, record it in `missing_requirements`. -6. If multiple schema interpretations are equally plausible and the question does not disambiguate them, record them in `ambiguous_requirements`. -7. Set `is_fully_supported` to false when any required request component is missing or ambiguous. -8. For each selected table, provide a concise reason for why the table is semantically relevant. -9. For each selected column, provide a concise reason for why the column is necessary. -10. If a "." is included in columns, put the name before the first dot into chosen columns. -11. The number of columns chosen must match the number of reasoning. -12. Final chosen columns must be only column names, don't prefix it with table names. -13. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -14. If the schema cannot answer the question, return the closest directly relevant schema objects only if they explain the limitation. Do not select unrelated fallback tables. +1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. +2. For each table, provide a clear and concise reasoning for why specific columns are selected. +3. List each reason as part of a step-by-step chain of thought, justifying the inclusion of each column. +4. If a "." is included in columns, put the name before the first dot into chosen columns. +5. The number of columns chosen must match the number of reasoning. +6. Final chosen columns must be only column names, don't prefix it with table names. +7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: { - "semantic_analysis": { - "analytical_intent": "retrieval | detailed_records | summary | comparison | trend | dashboard | kpi | ranking | record_count | other", - "entities": ["business entities requested by the user"], - "identifiers": ["identifier fields requested by the user"], - "metrics": ["business metrics or measures requested by the user"], - "dimensions": ["grouping or descriptive dimensions requested by the user"], - "filters": ["filters or predicates requested by the user"], - "aggregations": ["aggregation or calculation requirements"], - "relationships": ["required joins or relationships"], - "time_constraints": ["time filters, grains, or trend requirements"], - "ranking": ["top/bottom/order/limit requirements"], - "supported_schema_objects": ["table.column or metric names that directly support the request"], - "missing_requirements": ["required request components not supported by the schema"], - "ambiguous_requirements": ["request components with multiple equally plausible schema mappings"], - "is_fully_supported": true, - "support_reasoning": "Concise explanation of whether the selected schema fully supports the request" - }, "results": [ { "table_selection_reason": "Reason for selecting tablename1", @@ -106,7 +82,6 @@ - Each table key must list only the columns relevant to answering the question. - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. -- Populate `semantic_analysis` before `results`; use it to verify the selected schema directly supports the request. - Be logical, concise, and ensure the output strictly follows the required JSON format. - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. @@ -380,7 +355,6 @@ def check_using_db_schemas_without_pruning( "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, - "semantic_analysis": {}, } @@ -431,9 +405,9 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - retrieval_payload = orjson.loads(filter_columns_in_tables["replies"][0]) - columns_and_tables_needed = retrieval_payload.get("results", []) - semantic_analysis = retrieval_payload.get("semantic_analysis") or {} + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + )["results"] # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -493,7 +467,6 @@ def construct_retrieval_results( "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, - "semantic_analysis": semantic_analysis, } else: retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] @@ -505,9 +478,6 @@ def construct_retrieval_results( ], "has_metric": check_using_db_schemas_without_pruning["has_metric"], "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], - "semantic_analysis": check_using_db_schemas_without_pruning.get( - "semantic_analysis", {} - ), } @@ -523,26 +493,7 @@ class MatchingTable(BaseModel): table_selection_reason: str -class SemanticAnalysis(BaseModel): - analytical_intent: str = "" - entities: list[str] = Field(default_factory=list) - identifiers: list[str] = Field(default_factory=list) - metrics: list[str] = Field(default_factory=list) - dimensions: list[str] = Field(default_factory=list) - filters: list[str] = Field(default_factory=list) - aggregations: list[str] = Field(default_factory=list) - relationships: list[str] = Field(default_factory=list) - time_constraints: list[str] = Field(default_factory=list) - ranking: list[str] = Field(default_factory=list) - supported_schema_objects: list[str] = Field(default_factory=list) - missing_requirements: list[str] = Field(default_factory=list) - ambiguous_requirements: list[str] = Field(default_factory=list) - is_fully_supported: bool | None = None - support_reasoning: str = "" - - class RetrievalResults(BaseModel): - semantic_analysis: SemanticAnalysis | None = None results: list[MatchingTable] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9731665919..9db38bd5a8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -11,7 +11,6 @@ from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, construct_valid_table_names, - get_schema_intent_analysis_error, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, ) @@ -1289,7 +1288,7 @@ def _build_schema_grounded_table_question_sql( wants_monthly_count = any( term in normalized for term in ("monthly", "by month", "per month", "month-wise") - ) and any(term in normalized for term in ("count", "record", "records", "rows")) + ) and any(term in normalized for term in ("count", "records", "rows")) if wants_monthly_count: date_column = self._find_temporal_column_for_query(query, table) if not date_column: @@ -1472,23 +1471,21 @@ def _build_explicit_table_preview_sql( def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names: list[str] = [] for match in re.finditer( - r"\b(?:from|in|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - for candidate in self._explicit_table_name_candidates(table_name): - if candidate and candidate not in table_names: - table_names.append(candidate) + if table_name and table_name not in table_names: + table_names.append(table_name) for match in re.finditer( r"\bin\s+(?:the\s+)?([A-Za-z_][A-Za-z0-9_.$]*)\s+table\b", query or "", flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - for candidate in self._explicit_table_name_candidates(table_name): - if candidate and candidate not in table_names: - table_names.append(candidate) + if table_name and table_name not in table_names: + table_names.append(table_name) for match in re.finditer( r"\bin\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", @@ -1498,10 +1495,9 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: if ( table_name and ("." in table_name or "_" in table_name) + and table_name not in table_names ): - for candidate in self._explicit_table_name_candidates(table_name): - if candidate and candidate not in table_names: - table_names.append(candidate) + table_names.append(table_name) for match in re.finditer( r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", query or "", @@ -1511,10 +1507,9 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: if ( table_name and ("." in table_name or "_" in table_name) + and table_name not in table_names ): - for candidate in self._explicit_table_name_candidates(table_name): - if candidate and candidate not in table_names: - table_names.append(candidate) + table_names.append(table_name) if re.search( r"\b(?:repair\s+logs?|repair\s+tickets?|board\s+models?)\b", query or "", @@ -1529,24 +1524,6 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: table_names.append(table_name) return table_names - def _explicit_table_name_candidates(self, table_name: str) -> list[str]: - table_name = str(table_name or "").strip(".,;:()[]{}") - if not table_name: - return [] - - candidates = [table_name] - dotted_parts = [part for part in re.split(r"[.$]", table_name) if part] - if len(dotted_parts) > 1: - underscored = "_".join(dotted_parts) - candidates.append(underscored) - candidates.append(dotted_parts[-1]) - - normalized_candidates: list[str] = [] - for candidate in candidates: - if candidate and candidate not in normalized_candidates: - normalized_candidates.append(candidate) - return normalized_candidates - def _build_direct_orders_sales_sql(self, query: str) -> str | None: return None @@ -5321,7 +5298,6 @@ async def ask( allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) - schema_intent_analysis: dict[str, Any] = {} try: sql_user_query = user_query @@ -6002,9 +5978,6 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6042,9 +6015,6 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6349,32 +6319,6 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if semantic_support_error := get_schema_intent_analysis_error( - schema_intent_analysis - ): - logger.info( - "ask pipeline - NO_RELEVANT_SQL due to schema intent analysis: %s", - user_query, - ) - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_SQL", - message=semantic_support_error, - ), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = semantic_support_error - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names @@ -6507,7 +6451,6 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, - schema_intent_analysis=schema_intent_analysis, ), ) ).get("post_process", {}) @@ -6530,7 +6473,6 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, - schema_intent_analysis=schema_intent_analysis, ), ) ).get("post_process", {}) @@ -6619,7 +6561,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, ), ) else: @@ -6639,7 +6580,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, ), ) except TimeoutError as generation_timeout: @@ -6677,7 +6617,6 @@ async def ask( if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", - "SCHEMA_INTENT_VALIDATION", }: invalid_sql = failed_dry_run_result.get("sql", invalid_sql) error_message = failed_dry_run_result.get( @@ -6739,7 +6678,6 @@ async def ask( sql_functions=sql_functions, sql_knowledge=sql_knowledge, query=sql_user_query, - schema_intent_analysis=schema_intent_analysis, ), ) diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index a9a0c1eabe..25044de18c 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -7,7 +7,6 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import get_schema_intent_analysis_error from src.utils import trace_metadata from src.web.v1.services import BaseRequest from src.web.v1.services.ask import AskError, AskResult @@ -107,7 +106,6 @@ async def ask_feedback( invalid_sql = None sql_knowledge = None allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - schema_intent_analysis = {} try: if not self._is_stopped(query_id, self._ask_feedback_results): @@ -161,9 +159,6 @@ async def ask_feedback( ) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] sql_samples = sql_samples_task["formatted_output"].get("documents", []) @@ -172,21 +167,6 @@ async def ask_feedback( ) if not self._is_stopped(query_id, self._ask_feedback_results): - if semantic_support_error := get_schema_intent_analysis_error( - schema_intent_analysis - ): - self._ask_feedback_results[query_id] = AskFeedbackResultResponse( - status="failed", - error=AskError( - code="NO_RELEVANT_SQL", - message=semantic_support_error, - ), - trace_id=trace_id, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = semantic_support_error - return results - self._ask_feedback_results[query_id] = AskFeedbackResultResponse( status="generating", trace_id=trace_id, @@ -198,7 +178,6 @@ async def ask_feedback( contexts=table_ddls, sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, - query=ask_feedback_request.question, project_id=ask_feedback_request.project_id, sql_samples=sql_samples, instructions=instructions, @@ -207,7 +186,6 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -224,10 +202,7 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] not in { - "TIME_OUT", - "SCHEMA_INTENT_VALIDATION", - }: + if failed_dry_run_result["type"] != "TIME_OUT": original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] @@ -273,8 +248,6 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - query=ask_feedback_request.question, - schema_intent_analysis=schema_intent_analysis, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index e2fc3f7922..fa46c2ca7a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -5,7 +5,6 @@ extract_sql_generation_result, find_invalid_column_references, find_invalid_table_references, - get_schema_intent_analysis_error, get_json_field_instructions, get_metric_instructions, normalize_data_source, @@ -14,7 +13,6 @@ normalize_sql_table_references_to_schema, get_sql_generation_system_prompt, get_text_to_sql_rules, - validate_sql_intent_alignment, ) @@ -1552,177 +1550,3 @@ def test_normalize_generation_result_sql_rewrites_kb_article_created_by_for_mssq assert "created_by," not in normalized assert "GROUP BY created_by" not in normalized assert '"created_by_user_id"' in normalized - - -def test_validate_sql_intent_alignment_rejects_generic_count_for_metric_request(): - error = validate_sql_intent_alignment( - "Show revenue trend by month", - 'SELECT COUNT(*) AS "RecordCount" FROM "orders"', - {"orders": ["created_at", "order_id"]}, - ) - - assert error is not None - assert "generic record count" in error - - -def test_validate_sql_intent_alignment_rejects_trend_without_temporal_field(): - error = validate_sql_intent_alignment( - "Show monthly order volume", - 'SELECT "orders"."status", COUNT(*) AS "RecordCount" ' - 'FROM "orders" GROUP BY "orders"."status"', - {"orders": ["status", "order_id"]}, - ) - - assert error is not None - assert "temporal field" in error - - -def test_validate_sql_intent_alignment_allows_schema_supported_metric_trend(): - error = validate_sql_intent_alignment( - "Show revenue trend by month", - 'SELECT DATEPART(YEAR, "orders"."created_at") AS "year", ' - 'DATEPART(MONTH, "orders"."created_at") AS "month", ' - 'SUM("orders"."revenue") AS "revenue" ' - 'FROM "orders" ' - 'GROUP BY DATEPART(YEAR, "orders"."created_at"), ' - 'DATEPART(MONTH, "orders"."created_at")', - {"orders": ["created_at", "revenue"]}, - ) - - assert error is None - - -def test_validate_sql_intent_alignment_allows_explicit_record_count(): - error = validate_sql_intent_alignment( - "How many records are in orders?", - 'SELECT COUNT(*) AS "RecordCount" FROM "orders"', - {"orders": ["id"]}, - ) - - assert error is None - - -def test_validate_sql_intent_alignment_allows_temporal_record_count_filter(): - error = validate_sql_intent_alignment( - "How many orders were created last month?", - 'SELECT COUNT(*) AS "RecordCount" FROM "orders" ' - 'WHERE "orders"."created_at" >= \'2026-06-01 00:00:00\' ' - 'AND "orders"."created_at" < \'2026-07-01 00:00:00\'', - {"orders": ["created_at", "id"]}, - ) - - assert error is None - - -def test_validate_sql_intent_alignment_rejects_time_bucket_without_grouping(): - error = validate_sql_intent_alignment( - "Show monthly order volume", - 'SELECT COUNT(*) AS "RecordCount" FROM "orders" ' - 'WHERE "orders"."created_at" IS NOT NULL', - {"orders": ["created_at", "id"]}, - ) - - assert error is not None - assert "generic record count" in error - - -def test_validate_sql_intent_alignment_allows_duration_metric_without_trend(): - error = validate_sql_intent_alignment( - "Show average turnaround time by status", - 'SELECT "repairs"."status", AVG("repairs"."turnaround_hours") ' - 'AS "avg_turnaround_hours" FROM "repairs" GROUP BY "repairs"."status"', - {"repairs": ["status", "turnaround_hours"]}, - ) - - assert error is None - - -def test_get_schema_intent_analysis_error_reports_missing_requirements(): - error = get_schema_intent_analysis_error( - { - "missing_requirements": ["net income metric"], - "support_reasoning": "No metric maps to net income.", - } - ) - - assert error is not None - assert "net income metric" in error - - -def test_get_schema_intent_analysis_error_requests_clarification_for_ambiguity(): - error = get_schema_intent_analysis_error( - { - "ambiguous_requirements": [ - "amount could map to gross_amount or net_amount" - ], - } - ) - - assert error is not None - assert "Please clarify" in error - - -def test_get_schema_intent_analysis_error_reports_unsupported_analysis(): - error = get_schema_intent_analysis_error( - { - "is_fully_supported": False, - "support_reasoning": "No relationship connects invoices to products.", - } - ) - - assert error is not None - assert "No relationship connects invoices to products" in error - - -def test_validate_sql_intent_alignment_uses_semantic_analysis_for_metric_count_mismatch(): - error = validate_sql_intent_alignment( - "Show invoice amount by customer", - 'SELECT COUNT(*) AS "RecordCount" FROM "invoices"', - {"invoices": ["customer_id", "invoice_amount"]}, - semantic_analysis={ - "analytical_intent": "summary", - "entities": ["invoice", "customer"], - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "is_fully_supported": True, - }, - ) - - assert error is not None - assert "generic record count" in error - - -def test_validate_sql_intent_alignment_uses_semantic_analysis_for_ranking(): - error = validate_sql_intent_alignment( - "Top customers by invoice amount", - 'SELECT "invoices"."customer_id", SUM("invoices"."invoice_amount") ' - 'AS "total_invoice_amount" FROM "invoices" GROUP BY "invoices"."customer_id"', - {"invoices": ["customer_id", "invoice_amount"]}, - semantic_analysis={ - "analytical_intent": "ranking", - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "ranking": ["top customers by invoice amount"], - "is_fully_supported": True, - }, - ) - - assert error is not None - assert "ranking intent" in error - - -def test_validate_sql_intent_alignment_allows_semantic_count_metric(): - error = validate_sql_intent_alignment( - "Show total order count", - 'SELECT COUNT(*) AS "order_count" FROM "orders"', - {"orders": ["id"]}, - semantic_analysis={ - "analytical_intent": "summary", - "entities": ["order"], - "metrics": ["order count"], - "aggregations": ["count orders"], - "is_fully_supported": True, - }, - ) - - assert error is None diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 6adbad941d..5b6694bb45 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -3,7 +3,6 @@ from src.pipelines.retrieval.db_schema_retrieval import ( _is_project_wide_analysis_query, - construct_retrieval_results, dbschema_retrieval, expand_business_terms_for_retrieval, ) @@ -88,73 +87,3 @@ async def run(self, query_embedding, filters): {"field": "project_id", "operator": "==", "value": "project-1"}, ], } - - -def test_construct_retrieval_results_preserves_semantic_analysis(): - result = construct_retrieval_results( - check_using_db_schemas_without_pruning={"db_schemas": []}, - filter_columns_in_tables={ - "replies": [ - """ - { - "semantic_analysis": { - "analytical_intent": "summary", - "entities": ["invoice"], - "metrics": ["invoice amount"], - "dimensions": ["customer"], - "is_fully_supported": true - }, - "results": [ - { - "table_name": "invoices", - "table_selection_reason": "Contains invoice facts.", - "table_contents": { - "chain_of_thought_reasoning": [ - "Needed to group by customer.", - "Needed to sum invoice amount." - ], - "columns": ["customer_id", "invoice_amount"] - } - } - ] - } - """ - ] - }, - construct_db_schemas=[ - { - "type": "TABLE", - "name": "invoices", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "customer_id", - "data_type": "varchar", - "comment": "", - "is_primary_key": False, - }, - { - "type": "COLUMN", - "name": "invoice_amount", - "data_type": "double", - "comment": "", - "is_primary_key": False, - }, - { - "type": "COLUMN", - "name": "internal_note", - "data_type": "varchar", - "comment": "", - "is_primary_key": False, - }, - ], - } - ], - dbschema_retrieval=[], - ) - - assert result["semantic_analysis"]["metrics"] == ["invoice amount"] - assert result["retrieval_results"][0]["table_name"] == "invoices" - assert "invoice_amount" in result["retrieval_results"][0]["table_ddl"] - assert "internal_note" not in result["retrieval_results"][0]["table_ddl"] From 91dae570f8f52678d7cb8f268d5f80e462696d73 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 11 Jul 2026 00:10:47 +0530 Subject: [PATCH 0488/1087] Add deterministic semantic SQL compiler --- .../generation/followup_sql_generation.py | 33 +- .../src/pipelines/generation/semantic_sql.py | 1213 +++++++++++++++++ .../pipelines/generation/sql_generation.py | 34 +- .../pipelines/generation/test_semantic_sql.py | 156 +++ 4 files changed, 1434 insertions(+), 2 deletions(-) create mode 100644 wren-ai-service/src/pipelines/generation/semantic_sql.py create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py 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 c9f8b23537..e6e6b0c4c3 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -13,6 +13,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata +from src.pipelines.generation.semantic_sql import compile_semantic_sql from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_ask_history_messages, @@ -253,6 +254,36 @@ async def run( logger.info("Follow-Up SQL Generation pipeline is running...") metadata = await retrieve_metadata(project_id or "", self._retriever) + data_source = metadata.get("data_source", "local_file") + + deterministic_result = compile_semantic_sql( + query=query, + documents=contexts, + semantic_analysis=schema_intent_analysis, + data_source=data_source, + ) + if deterministic_result: + logger.info( + "Deterministic semantic SQL compiler produced follow-up SQL for query: %s", + query, + ) + post_process_result = await self._components["post_processor"].run( + [deterministic_result.sql], + project_id=project_id, + use_dry_plan=use_dry_plan, + data_source=data_source, + allow_dry_plan_fallback=allow_dry_plan_fallback, + valid_table_names=construct_valid_table_names(contexts), + valid_table_columns=construct_valid_table_columns(contexts), + query=query, + semantic_analysis=schema_intent_analysis, + ) + if post_process_result.get("valid_generation_result"): + return {"post_process": post_process_result} + logger.info( + "Deterministic follow-up SQL rejected by post processor; falling back to LLM. error=%s", + post_process_result.get("invalid_generation_result", {}).get("error"), + ) return await self._pipe.execute( ["post_process"], @@ -270,7 +301,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "sql_knowledge": sql_knowledge, **self._components, }, diff --git a/wren-ai-service/src/pipelines/generation/semantic_sql.py b/wren-ai-service/src/pipelines/generation/semantic_sql.py new file mode 100644 index 0000000000..0677de24f5 --- /dev/null +++ b/wren-ai-service/src/pipelines/generation/semantic_sql.py @@ -0,0 +1,1213 @@ +from __future__ import annotations + +import logging +import json +import re +from ast import literal_eval +from collections import deque +from dataclasses import dataclass, field +from datetime import date, datetime, timedelta +from typing import Any, Literal + +import sqlparse + +logger = logging.getLogger("wren-ai-service") + + +Aggregate = Literal["COUNT", "SUM", "AVG", "MIN", "MAX"] +JoinType = Literal["INNER JOIN", "LEFT JOIN"] + + +@dataclass(frozen=True) +class Intent: + question_type: str + chart_requested: bool = False + chart_type: str = "auto" + ranking: bool = False + top_n: int | None = None + bottom_n: int | None = None + distinct: bool = False + aggregation: Aggregate | None = None + needs_sql: bool = True + + +@dataclass(frozen=True) +class ColumnRef: + table: str + column: str + data_type: str = "" + description: str = "" + + @property + def sql(self) -> str: + return f'{quote_identifier(self.table)}.{quote_identifier(self.column)}' + + @property + def object_name(self) -> str: + return f"{self.table}.{self.column}" + + +@dataclass(frozen=True) +class Relationship: + left_table: str + left_column: str + right_table: str + right_column: str + join_type: JoinType = "INNER JOIN" + cardinality: str = "" + source: str = "schema" + + +@dataclass(frozen=True) +class MetricDefinition: + name: str + column: ColumnRef + aggregation: Aggregate + description: str = "" + synonyms: tuple[str, ...] = () + allowed_dimensions: tuple[str, ...] = () + formula: str | None = None + grain: str | None = None + join_requirements: tuple[str, ...] = () + + +@dataclass(frozen=True) +class FilterDefinition: + column: ColumnRef + operator: str + value: Any = None + end_value: Any = None + + +@dataclass(frozen=True) +class SortDefinition: + expression: str + direction: Literal["ASC", "DESC"] = "DESC" + + +@dataclass +class SemanticPlan: + intent: Intent + entities: list[str] = field(default_factory=list) + metrics: list[MetricDefinition] = field(default_factory=list) + aggregation: Aggregate | None = None + filters: list[FilterDefinition] = field(default_factory=list) + group_by: list[ColumnRef] = field(default_factory=list) + sort: list[SortDefinition] = field(default_factory=list) + limit: int | None = None + base_table: str | None = None + joins: list[Relationship] = field(default_factory=list) + chart_type: str = "" + warnings: list[str] = field(default_factory=list) + + @property + def is_complete(self) -> bool: + if not self.base_table: + return False + if self.intent.aggregation and not self.metrics and self.intent.aggregation != "COUNT": + return False + required_tables = { + ref.table + for ref in [ + *[metric.column for metric in self.metrics], + *self.group_by, + *[filter_.column for filter_ in self.filters], + ] + } + connected_tables = {self.base_table} + for join in self.joins: + connected_tables.add(join.left_table) + connected_tables.add(join.right_table) + return required_tables.issubset(connected_tables) + + +@dataclass +class SQLValidationResult: + valid: bool + errors: list[str] = field(default_factory=list) + + +@dataclass +class CompileResult: + sql: str + plan: SemanticPlan + validation: SQLValidationResult + + +@dataclass +class SchemaCatalog: + tables: dict[str, list[ColumnRef]] = field(default_factory=dict) + relationships: list[Relationship] = field(default_factory=list) + + def columns(self) -> list[ColumnRef]: + return [column for columns in self.tables.values() for column in columns] + + def table_for_column(self, column: ColumnRef) -> str: + return column.table + + def get_column(self, table: str, column: str) -> ColumnRef | None: + for candidate in self.tables.get(table, []): + if candidate.column.lower() == column.lower(): + return candidate + return None + + +_STOPWORDS = { + "a", + "an", + "and", + "as", + "at", + "by", + "chart", + "for", + "from", + "give", + "graph", + "in", + "last", + "me", + "of", + "on", + "per", + "show", + "the", + "this", + "to", + "with", +} +_GENERIC_COLUMN_TOKENS = { + "amount", + "at", + "code", + "date", + "day", + "id", + "key", + "month", + "name", + "no", + "number", + "time", + "total", + "type", + "value", + "year", +} +_SYNONYMS = { + "acct": {"account", "customer", "client"}, + "account": {"acct", "customer", "client"}, + "amount": {"amt", "value", "total"}, + "avg": {"average", "mean"}, + "bill": {"invoice"}, + "billing": {"invoice"}, + "client": {"customer", "account"}, + "cost": {"expense", "spend"}, + "cust": {"customer", "client", "account"}, + "customer": {"cust", "client", "account"}, + "gmv": {"revenue", "sales", "amount"}, + "invoice": {"inv", "bill", "billing"}, + "profit": {"margin", "income", "earnings"}, + "qty": {"quantity", "units"}, + "quantity": {"qty", "units"}, + "revenue": {"sales", "amount", "value", "gmv"}, + "sale": {"sales", "revenue", "amount"}, + "sales": {"sale", "revenue", "amount", "gmv"}, + "total": {"sum", "amount", "value"}, + "value": {"amount", "total"}, +} +_NUMERIC_TYPES = { + "bigint", + "decimal", + "double", + "float", + "int", + "integer", + "numeric", + "real", + "smallint", +} +_TEMPORAL_TYPES = {"date", "datetime", "timestamp", "time"} +_TEXT_TYPES = {"char", "string", "text", "varchar"} + + +def quote_identifier(identifier: str) -> str: + return f'"{str(identifier).replace(chr(34), chr(34) + chr(34))}"' + + +def tokenize(value: Any) -> set[str]: + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value or "")) + raw_tokens = { + token.lower() + for token in re.findall(r"[A-Za-z0-9]+", text) + if len(token) > 1 and token.lower() not in _STOPWORDS + } + tokens = set(raw_tokens) + for token in list(raw_tokens): + if len(token) > 4 and token.endswith("ies"): + tokens.add(f"{token[:-3]}y") + elif len(token) > 3 and token.endswith("s"): + tokens.add(token[:-1]) + for token in list(tokens): + tokens.update(_SYNONYMS.get(token, set())) + return tokens + + +def compact(value: Any) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + + +def is_numeric_type(data_type: str) -> bool: + normalized = data_type.lower() + return any(type_name in normalized for type_name in _NUMERIC_TYPES) + + +def is_temporal_type(data_type: str) -> bool: + normalized = data_type.lower() + return any(type_name in normalized for type_name in _TEMPORAL_TYPES) + + +def is_text_type(data_type: str) -> bool: + normalized = data_type.lower() + return any(type_name in normalized for type_name in _TEXT_TYPES) + + +class IntentDetector: + def detect(self, query: str) -> Intent: + normalized = query.lower() + top_match = re.search(r"\btop\s+(\d+)\b", normalized) + bottom_match = re.search(r"\bbottom\s+(\d+)\b", normalized) + aggregation = self._detect_aggregation(normalized) + chart_type = self._detect_chart_type(normalized) + question_type = "retrieval" + if aggregation: + question_type = "aggregation" + if re.search(r"\btrend|over time|monthly|weekly|daily|yearly\b", normalized): + question_type = "trend" + if top_match or bottom_match or re.search( + r"\b(highest|lowest|largest|smallest|rank|ranking)\b", normalized + ): + question_type = "ranking" + if "dashboard" in normalized or "kpi" in normalized: + question_type = "dashboard" + + return Intent( + question_type=question_type, + chart_requested=bool(chart_type), + chart_type=chart_type or "auto", + ranking=question_type == "ranking", + top_n=( + int(top_match.group(1)) + if top_match + else 10 + if re.search(r"\btop\b", normalized) + else None + ), + bottom_n=int(bottom_match.group(1)) if bottom_match else None, + distinct=bool(re.search(r"\bdistinct|unique\b", normalized)), + aggregation=aggregation, + ) + + def _detect_aggregation(self, normalized_query: str) -> Aggregate | None: + if re.search(r"\b(avg|average|mean)\b", normalized_query): + return "AVG" + if re.search(r"\b(count|number of|how many)\b", normalized_query): + return "COUNT" + if re.search(r"\b(min|minimum)\b", normalized_query): + return "MIN" + if re.search(r"\b(max|maximum)\b", normalized_query): + return "MAX" + if re.search(r"\b(sum|total)\b", normalized_query): + return "SUM" + return None + + def _detect_chart_type(self, normalized_query: str) -> str: + checks = ( + ("line", ("line chart", "line graph", "trend")), + ("pie", ("pie chart", "donut chart", "part to whole")), + ("scatter", ("scatter", "correlation")), + ("bar", ("bar chart", "bar graph", "column chart")), + ("card", ("kpi", "single kpi")), + ) + for chart_type, terms in checks: + if any(term in normalized_query for term in terms): + return chart_type + if re.search(r"\b(chart|graph|plot|visuali[sz]e)\b", normalized_query): + return "auto" + return "" + + +class SchemaParser: + def parse(self, documents: list[str]) -> SchemaCatalog: + catalog = SchemaCatalog() + for document in documents or []: + text = str(document) + self._parse_create_table(text, catalog) + self._parse_metadata_document(text, catalog) + return catalog + + def _parse_create_table(self, ddl: str, catalog: SchemaCatalog) -> None: + table_match = re.search( + r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|`(?P[^`]+)`|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", + ddl, + flags=re.IGNORECASE, + ) + if not table_match: + return + table_name = next(value for value in table_match.groupdict().values() if value) + body = self._table_body(ddl, table_match.end()) + columns: list[ColumnRef] = [] + for definition in self._split_definitions(body): + stripped = definition.strip().rstrip(",") + if not stripped: + continue + fk = re.search( + r"FOREIGN\s+KEY\s*\((?P[^\)]+)\)\s+REFERENCES\s+" + r'(?:"(?P[^"]+)"|`(?P[^`]+)`|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\((?P[^\)]+)\)", + stripped, + flags=re.IGNORECASE, + ) + if fk: + right_table = next( + value + for value in ( + fk.group("rq"), + fk.group("rb"), + fk.group("rs"), + fk.group("rbare"), + ) + if value + ) + catalog.relationships.append( + Relationship( + left_table=table_name, + left_column=self._clean_identifier(fk.group("left")), + right_table=right_table, + right_column=self._clean_identifier(fk.group("right")), + ) + ) + continue + if re.match( + r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY|INDEX)\b", + stripped, + flags=re.IGNORECASE, + ): + continue + column_match = re.match( + r'(?:"(?P[^"]+)"|`(?P[^`]+)`|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_$]*))" + r"\s+(?P[A-Za-z_][A-Za-z0-9_]*(?:\([^\)]*\))?)", + stripped, + ) + if column_match: + column_name = next( + value + for key, value in column_match.groupdict().items() + if key != "type" and value + ) + columns.append( + ColumnRef( + table=table_name, + column=column_name, + data_type=column_match.group("type") or "", + ) + ) + if columns: + catalog.tables[table_name] = columns + + def _parse_metadata_document(self, text: str, catalog: SchemaCatalog) -> None: + metadata = self._literal_document(text) + if not isinstance(metadata, dict): + return + for model in [ + *(metadata.get("models") or []), + *(metadata.get("views") or []), + ]: + if not isinstance(model, dict) or not model.get("name"): + continue + table_name = str(model["name"]) + columns: list[ColumnRef] = [] + for column in [ + *(model.get("columns") or []), + *(model.get("calculatedFields") or []), + ]: + if not isinstance(column, dict) or not column.get("name"): + continue + columns.append( + ColumnRef( + table=table_name, + column=str(column["name"]), + data_type=str( + column.get("data_type") + or column.get("type") + or column.get("dataType") + or "" + ), + description=str(column.get("comment") or column.get("description") or ""), + ) + ) + if columns: + catalog.tables[table_name] = columns + reference_name = model.get("referenceName") + if reference_name and columns: + catalog.tables[str(reference_name)] = [ + ColumnRef( + table=str(reference_name), + column=column.column, + data_type=column.data_type, + description=column.description, + ) + for column in columns + ] + for relationship in metadata.get("relationships") or []: + parsed = self._relationship_from_metadata(relationship) + if parsed: + catalog.relationships.append(parsed) + + def _literal_document(self, text: str) -> Any: + stripped = text.strip() + if not stripped.startswith("{"): + return None + try: + return json.loads(stripped) + except json.JSONDecodeError: + pass + try: + return literal_eval(stripped) + except (SyntaxError, ValueError): + return None + + def _relationship_from_metadata(self, relationship: Any) -> Relationship | None: + if not isinstance(relationship, dict): + return None + condition = str(relationship.get("condition") or "") + match = re.search( + r'(?P[A-Za-z_][A-Za-z0-9_.$]*)\.(?P[A-Za-z_][A-Za-z0-9_$]*)\s*=\s*' + r'(?P[A-Za-z_][A-Za-z0-9_.$]*)\.(?P[A-Za-z_][A-Za-z0-9_$]*)', + condition, + ) + if not match: + return None + return Relationship( + left_table=match.group("left_table"), + left_column=match.group("left_column"), + right_table=match.group("right_table"), + right_column=match.group("right_column"), + cardinality=str(relationship.get("joinType") or relationship.get("type") or ""), + source="semantic_metadata", + ) + + def _table_body(self, ddl: str, start: int) -> str: + depth = 1 + cursor = start + while cursor < len(ddl) and depth > 0: + if ddl[cursor] == "(": + depth += 1 + elif ddl[cursor] == ")": + depth -= 1 + cursor += 1 + return ddl[start : cursor - 1] + + def _split_definitions(self, body: str) -> list[str]: + definitions: list[str] = [] + depth = 0 + start = 0 + for index, char in enumerate(body): + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + elif char == "," and depth == 0: + definitions.append(body[start:index]) + start = index + 1 + definitions.append(body[start:]) + return definitions + + def _clean_identifier(self, value: str) -> str: + return str(value or "").strip().strip('"`[]') + + +class MetricRegistry: + def __init__(self, catalog: SchemaCatalog, semantic_analysis: dict[str, Any] | None): + self.catalog = catalog + self.semantic_analysis = semantic_analysis or {} + self.metrics = self._build_metrics() + + def resolve_metric(self, query: str, intent: Intent) -> MetricDefinition | None: + mappings = self.semantic_analysis.get("concept_mappings") or [] + for mapping in mappings: + if not isinstance(mapping, dict): + continue + if str(mapping.get("concept_type", "")).lower() != "metric": + continue + for schema_object in mapping.get("schema_objects") or []: + column = self._column_from_schema_object(str(schema_object)) + if column: + return MetricDefinition( + name=str(mapping.get("request_concept") or column.column), + column=column, + aggregation=intent.aggregation or self._default_aggregation(column), + synonyms=tuple(tokenize(mapping.get("request_concept"))), + ) + + query_tokens = tokenize(query) + scored = [ + (self._score_metric(metric, query_tokens), metric) + for metric in self.metrics + ] + scored = [(score, metric) for score, metric in scored if score > 0] + if not scored: + return None + scored.sort(key=lambda item: item[0], reverse=True) + metric = scored[0][1] + if intent.aggregation: + return MetricDefinition( + name=metric.name, + column=metric.column, + aggregation=intent.aggregation, + description=metric.description, + synonyms=metric.synonyms, + allowed_dimensions=metric.allowed_dimensions, + formula=metric.formula, + grain=metric.grain, + join_requirements=metric.join_requirements, + ) + return metric + + def _build_metrics(self) -> list[MetricDefinition]: + metrics: list[MetricDefinition] = [] + for column in self.catalog.columns(): + if not is_numeric_type(column.data_type): + continue + tokens = tokenize(f"{column.table} {column.column}") + if not tokens.intersection( + { + "amount", + "balance", + "cost", + "gmv", + "margin", + "price", + "profit", + "quantity", + "rate", + "revenue", + "sale", + "sales", + "score", + "total", + "value", + } + ): + continue + metrics.append( + MetricDefinition( + name=humanize(column.column), + column=column, + aggregation=self._default_aggregation(column), + synonyms=tuple(tokens), + ) + ) + return metrics + + def _column_from_schema_object(self, schema_object: str) -> ColumnRef | None: + parts = [part.strip().strip('"`[]') for part in schema_object.split(".")] + if len(parts) < 2: + return None + table = ".".join(parts[:-1]) + column = parts[-1] + direct = self.catalog.get_column(table, column) + if direct: + return direct + for table_name in self.catalog.tables: + if table_name.lower().endswith(table.lower()): + candidate = self.catalog.get_column(table_name, column) + if candidate: + return candidate + return None + + def _default_aggregation(self, column: ColumnRef) -> Aggregate: + tokens = tokenize(column.column) + if tokens.intersection({"avg", "average", "mean", "rate", "ratio", "percent"}): + return "AVG" + return "SUM" + + def _score_metric(self, metric: MetricDefinition, query_tokens: set[str]) -> int: + metric_tokens = set(metric.synonyms) | tokenize(metric.name) + score = len(metric_tokens.intersection(query_tokens)) * 10 + compact_query = compact(" ".join(query_tokens)) + compact_metric = compact(metric.name) + if compact_metric and compact_metric in compact_query: + score += 40 + return score + + +class EntityResolver: + def __init__(self, catalog: SchemaCatalog, semantic_analysis: dict[str, Any] | None): + self.catalog = catalog + self.semantic_analysis = semantic_analysis or {} + self._metric_registry = MetricRegistry(catalog, {}) + + def resolve_dimensions(self, query: str, metric: MetricDefinition | None) -> list[ColumnRef]: + mapped = self._dimensions_from_semantic_analysis() + if mapped: + return mapped + requested_terms = self._requested_grouping_terms(query) + if not requested_terms and metric: + return [] + if not requested_terms: + requested_terms = tokenize(query) + scored: list[tuple[int, ColumnRef]] = [] + for column in self.catalog.columns(): + if metric and column == metric.column: + continue + if is_numeric_type(column.data_type): + if self._looks_identifier(column) and not requested_terms.intersection( + {"id", "key", "number", "no"} + ): + continue + if not self._looks_identifier(column): + continue + column_tokens = tokenize(f"{column.table} {column.column}") + score = len(column_tokens.intersection(requested_terms)) * 10 + if is_text_type(column.data_type): + score += 5 + if compact(column.column) in compact(" ".join(requested_terms)): + score += 20 + if requested_terms and score: + scored.append((score, column)) + scored.sort(key=lambda item: item[0], reverse=True) + if not scored: + return [] + best_score = scored[0][0] + return [column for score, column in scored[:3] if score == best_score] + + def resolve_temporal_column(self, query: str, preferred_table: str | None) -> ColumnRef | None: + if not self._requests_date_filter(query): + return None + candidates = [ + column + for column in self.catalog.columns() + if is_temporal_type(column.data_type) + or tokenize(column.column).intersection( + {"date", "time", "created", "updated", "month", "year"} + ) + ] + if preferred_table: + candidates.sort(key=lambda column: column.table != preferred_table) + return candidates[0] if candidates else None + + def _dimensions_from_semantic_analysis(self) -> list[ColumnRef]: + dimensions: list[ColumnRef] = [] + mappings = self.semantic_analysis.get("concept_mappings") or [] + for mapping in mappings: + if not isinstance(mapping, dict): + continue + if str(mapping.get("concept_type", "")).lower() not in { + "dimension", + "entity", + "identifier", + }: + continue + for schema_object in mapping.get("schema_objects") or []: + column = self._metric_registry._column_from_schema_object(str(schema_object)) + if column and column not in dimensions: + dimensions.append(column) + return dimensions + + def _requested_grouping_terms(self, query: str) -> set[str]: + normalized = query.lower() + terms: set[str] = set() + for match in re.finditer( + r"\b(?:by|per|for each|group(?:ed)? by)\s+([A-Za-z0-9_ ]+)", + normalized, + ): + phrase = re.split( + r"\b(?:and|with|where|order|sort|top|bottom|last|this|limit)\b", + match.group(1), + maxsplit=1, + )[0] + terms.update(tokenize(phrase)) + ranking_entity = re.search( + r"\b(?:top|bottom)\s+(?:\d+\s+)?([A-Za-z0-9_ ]+?)\s+by\b", + normalized, + ) + if ranking_entity: + terms.update(tokenize(ranking_entity.group(1))) + return terms + + def _requests_date_filter(self, query: str) -> bool: + return bool( + re.search( + r"\b(today|yesterday|this|last|rolling|past|previous)\s+" + r"(?:\d+\s+)?(?:day|week|month|quarter|year)s?\b", + query.lower(), + ) + ) + + def _looks_identifier(self, column: ColumnRef) -> bool: + return bool(tokenize(column.column).intersection({"id", "key", "number", "no"})) + + +class RelationshipGraph: + def __init__(self, catalog: SchemaCatalog): + self.catalog = catalog + + def join_path(self, required_tables: set[str], base_table: str) -> list[Relationship] | None: + joins: list[Relationship] = [] + connected = {base_table} + for table in sorted(required_tables - connected): + path = self._shortest_path(connected, table) + if not path: + return None + joins.extend(path) + for relationship in path: + connected.add(relationship.left_table) + connected.add(relationship.right_table) + return self._dedupe(joins) + + def _shortest_path(self, sources: set[str], target: str) -> list[Relationship] | None: + queue: deque[tuple[str, list[Relationship]]] = deque( + (source, []) for source in sources + ) + seen = set(sources) + while queue: + table, path = queue.popleft() + if table == target: + return path + for relationship in self._neighbors(table): + next_table = ( + relationship.right_table + if relationship.left_table == table + else relationship.left_table + ) + if next_table in seen: + continue + seen.add(next_table) + queue.append((next_table, [*path, relationship])) + return None + + def _neighbors(self, table: str) -> list[Relationship]: + return [ + relationship + for relationship in self.catalog.relationships + if relationship.left_table == table or relationship.right_table == table + ] + + def _dedupe(self, relationships: list[Relationship]) -> list[Relationship]: + deduped: list[Relationship] = [] + seen: set[tuple[str, str, str, str]] = set() + for relationship in relationships: + key = ( + relationship.left_table, + relationship.left_column, + relationship.right_table, + relationship.right_column, + ) + if key in seen: + continue + seen.add(key) + deduped.append(relationship) + return deduped + + +class SemanticPlanner: + def __init__(self, catalog: SchemaCatalog, semantic_analysis: dict[str, Any] | None): + self.catalog = catalog + self.semantic_analysis = semantic_analysis or {} + self.metric_registry = MetricRegistry(catalog, semantic_analysis) + self.entity_resolver = EntityResolver(catalog, semantic_analysis) + self.relationship_graph = RelationshipGraph(catalog) + + def build_plan(self, query: str, now: datetime | None = None) -> SemanticPlan | None: + intent = IntentDetector().detect(query) + metric = self.metric_registry.resolve_metric(query, intent) + aggregation = intent.aggregation or (metric.aggregation if metric else None) + dimensions = self.entity_resolver.resolve_dimensions(query, metric) + if aggregation == "COUNT" and not self._has_explicit_grouping(query): + dimensions = [] + temporal_column = self.entity_resolver.resolve_temporal_column( + query, + metric.column.table if metric else (dimensions[0].table if dimensions else None), + ) + filters = [] + if temporal_column: + date_filter = DateResolver(now or datetime.now()).resolve(query, temporal_column) + if date_filter: + filters.append(date_filter) + + required_refs = [ + *([metric.column] if metric else []), + *dimensions, + *[filter_.column for filter_ in filters], + ] + base_table = self._choose_base_table(required_refs) + if not base_table and aggregation == "COUNT": + base_table = self._resolve_table_from_query(query) + if not base_table: + return None + required_tables = {ref.table for ref in required_refs} + joins = self.relationship_graph.join_path(required_tables, base_table) + if joins is None: + return None + + metric_expression = self._metric_expression(metric, aggregation) + sort = [] + if intent.ranking and metric_expression: + sort.append(SortDefinition(metric_expression, "ASC" if intent.bottom_n else "DESC")) + limit = intent.top_n or intent.bottom_n + chart_type = ChartRuleEngine().select_chart(intent, dimensions, [metric] if metric else []) + + plan = SemanticPlan( + intent=intent, + entities=[humanize(dimension.column) for dimension in dimensions], + metrics=[metric] if metric else [], + aggregation=aggregation, + filters=filters, + group_by=dimensions, + sort=sort, + limit=limit, + base_table=base_table, + joins=joins, + chart_type=chart_type, + ) + return plan if plan.is_complete else None + + def _choose_base_table(self, refs: list[ColumnRef]) -> str | None: + if not refs: + return None + table_counts: dict[str, int] = {} + for ref in refs: + table_counts[ref.table] = table_counts.get(ref.table, 0) + 1 + return sorted(table_counts.items(), key=lambda item: item[1], reverse=True)[0][0] + + def _resolve_table_from_query(self, query: str) -> str | None: + query_tokens = tokenize(query) + scored: list[tuple[int, str]] = [] + for table in self.catalog.tables: + table_tokens = tokenize(table) + score = len(table_tokens.intersection(query_tokens)) * 10 + if compact(table) in compact(query): + score += 30 + if score: + scored.append((score, table)) + if not scored and len(self.catalog.tables) == 1: + return next(iter(self.catalog.tables)) + scored.sort(key=lambda item: item[0], reverse=True) + return scored[0][1] if scored else None + + def _has_explicit_grouping(self, query: str) -> bool: + return bool( + re.search( + r"\b(?:by|per|for each|group(?:ed)? by|top|bottom|rank|ranking)\b", + query.lower(), + ) + ) + + def _metric_expression( + self, + metric: MetricDefinition | None, + aggregation: Aggregate | None, + ) -> str: + if aggregation == "COUNT" and not metric: + return "COUNT(*)" + if not metric or not aggregation: + return "" + return f"{aggregation}({metric.column.sql})" + + +class DateResolver: + def __init__(self, now: datetime): + self.today = now.date() + + def resolve(self, query: str, column: ColumnRef) -> FilterDefinition | None: + normalized = query.lower() + start: date | None = None + end: date | None = None + if "today" in normalized: + start = self.today + end = self.today + timedelta(days=1) + elif "yesterday" in normalized: + start = self.today - timedelta(days=1) + end = self.today + elif "last month" in normalized: + current_month = self.today.replace(day=1) + end = current_month + start = add_months(current_month, -1) + elif "this month" in normalized: + start = self.today.replace(day=1) + end = add_months(start, 1) + elif "last year" in normalized: + start = date(self.today.year - 1, 1, 1) + end = date(self.today.year, 1, 1) + elif "this year" in normalized: + start = date(self.today.year, 1, 1) + end = date(self.today.year + 1, 1, 1) + elif "last week" in normalized: + this_week = self.today - timedelta(days=self.today.weekday()) + start = this_week - timedelta(days=7) + end = this_week + elif "this week" in normalized: + start = self.today - timedelta(days=self.today.weekday()) + end = start + timedelta(days=7) + elif "last quarter" in normalized or "this quarter" in normalized: + quarter_month = ((self.today.month - 1) // 3) * 3 + 1 + this_quarter = date(self.today.year, quarter_month, 1) + if "last quarter" in normalized: + end = this_quarter + start = add_months(this_quarter, -3) + else: + start = this_quarter + end = add_months(this_quarter, 3) + else: + rolling_match = re.search( + r"\b(?:last|past|rolling)\s+(\d+)\s+(day|week|month|year)s?\b", + normalized, + ) + if rolling_match: + amount = int(rolling_match.group(1)) + unit = rolling_match.group(2) + end = self.today + timedelta(days=1) + if unit == "day": + start = self.today - timedelta(days=amount) + elif unit == "week": + start = self.today - timedelta(days=amount * 7) + elif unit == "month": + start = add_months(self.today, -amount) + elif unit == "year": + start = add_months(self.today, -amount * 12) + if not start or not end: + return None + return FilterDefinition( + column=column, + operator="BETWEEN_CLOSED_OPEN", + value=start.isoformat(), + end_value=end.isoformat(), + ) + + +class SQLCompiler: + def compile(self, plan: SemanticPlan, data_source: str = "") -> str: + if not plan.base_table: + raise ValueError("Semantic plan has no base table") + select_items = self._select_items(plan) + from_clause = f"FROM {quote_identifier(plan.base_table)}" + join_clause = self._join_clause(plan) + where_clause = self._where_clause(plan) + group_clause = self._group_clause(plan) + order_clause = self._order_clause(plan) + limit_clause = self._limit_clause(plan, data_source) + top_clause = "" + if normalize_data_source(data_source) == "MSSQL" and plan.limit: + top_clause = f" TOP {plan.limit}" + limit_clause = "" + return " ".join( + part + for part in ( + f"SELECT{top_clause} {', '.join(select_items)}", + from_clause, + join_clause, + where_clause, + group_clause, + order_clause, + limit_clause, + ) + if part + ) + + def _select_items(self, plan: SemanticPlan) -> list[str]: + items = [ + f"{dimension.sql} AS {quote_identifier(safe_alias(dimension.column))}" + for dimension in plan.group_by + ] + if plan.metrics: + for metric in plan.metrics: + aggregation = plan.aggregation or metric.aggregation + alias = safe_alias(f"{aggregation.lower()}_{metric.column.column}") + items.append(f"{aggregation}({metric.column.sql}) AS {quote_identifier(alias)}") + elif plan.aggregation == "COUNT": + items.append('COUNT(*) AS "count"') + if plan.intent.distinct and not plan.aggregation and not plan.metrics: + return [ + f"DISTINCT {dimension.sql} AS {quote_identifier(safe_alias(dimension.column))}" + for dimension in plan.group_by + ] + return items + + def _join_clause(self, plan: SemanticPlan) -> str: + clauses = [] + joined = {plan.base_table} + for relationship in plan.joins: + if relationship.left_table in joined: + join_table = relationship.right_table + on_left = f"{quote_identifier(relationship.left_table)}.{quote_identifier(relationship.left_column)}" + on_right = f"{quote_identifier(relationship.right_table)}.{quote_identifier(relationship.right_column)}" + else: + join_table = relationship.left_table + on_left = f"{quote_identifier(relationship.left_table)}.{quote_identifier(relationship.left_column)}" + on_right = f"{quote_identifier(relationship.right_table)}.{quote_identifier(relationship.right_column)}" + clauses.append( + f"{relationship.join_type} {quote_identifier(join_table)} ON {on_left} = {on_right}" + ) + joined.add(join_table) + return " ".join(clauses) + + def _where_clause(self, plan: SemanticPlan) -> str: + conditions = [] + for filter_ in plan.filters: + if filter_.operator == "BETWEEN_CLOSED_OPEN": + conditions.append( + f"{filter_.column.sql} >= '{filter_.value}' AND {filter_.column.sql} < '{filter_.end_value}'" + ) + return f"WHERE {' AND '.join(conditions)}" if conditions else "" + + def _group_clause(self, plan: SemanticPlan) -> str: + if not plan.group_by or not plan.aggregation: + return "" + return "GROUP BY " + ", ".join(dimension.sql for dimension in plan.group_by) + + def _order_clause(self, plan: SemanticPlan) -> str: + if not plan.sort: + return "" + return "ORDER BY " + ", ".join( + f"{sort.expression} {sort.direction}" for sort in plan.sort + ) + + def _limit_clause(self, plan: SemanticPlan, data_source: str) -> str: + if not plan.limit or normalize_data_source(data_source) == "MSSQL": + return "" + return f"LIMIT {plan.limit}" + + +class SQLAstValidator: + def validate(self, sql: str, plan: SemanticPlan) -> SQLValidationResult: + errors: list[str] = [] + parsed = sqlparse.parse(sql) + if len(parsed) != 1: + errors.append("SQL must contain exactly one statement") + if not sqlparse.tokens.DML: + errors.append("SQL parser unavailable") + if not re.match(r"^\s*SELECT\b", sql, flags=re.IGNORECASE): + errors.append("SQL must be a SELECT statement") + if "*" in sql and plan.metrics: + errors.append("Metric queries must not use SELECT *") + if plan.joins and not re.search(r"\b(?:ON|USING)\b", sql, flags=re.IGNORECASE): + errors.append("Joins must include ON or USING") + if plan.aggregation and plan.group_by: + group_text = self._clause(sql, "GROUP BY", ["HAVING", "ORDER BY", "LIMIT", "FETCH"]) + for dimension in plan.group_by: + if quote_identifier(dimension.column) not in group_text: + errors.append(f"Missing GROUP BY column: {dimension.object_name}") + if plan.limit and not re.search( + r"\b(?:LIMIT\s+\d+|TOP\s+\d+|FETCH\s+FIRST\s+\d+)\b", + sql, + flags=re.IGNORECASE, + ): + errors.append("Missing limit/TOP for ranked plan") + if plan.filters and not re.search(r"\bWHERE\b", sql, flags=re.IGNORECASE): + errors.append("Missing WHERE for filtered plan") + return SQLValidationResult(valid=not errors, errors=errors) + + def _clause(self, sql: str, clause: str, terminators: list[str]) -> str: + terminator_pattern = "|".join(rf"\b{terminator}\b" for terminator in terminators) + match = re.search( + rf"\b{clause}\b(?P.*?)(?={terminator_pattern}|$)", + sql, + flags=re.IGNORECASE | re.DOTALL, + ) + return match.group("body") if match else "" + + +class ExecutionValidator: + def validate_result_shape( + self, + plan: SemanticPlan, + rows: list[dict[str, Any]], + ) -> SQLValidationResult: + if not rows: + return SQLValidationResult(valid=True) + expected_columns = {safe_alias(dimension.column) for dimension in plan.group_by} + for metric in plan.metrics: + aggregation = plan.aggregation or metric.aggregation + expected_columns.add(safe_alias(f"{aggregation.lower()}_{metric.column.column}")) + actual_columns = set(rows[0].keys()) + missing = expected_columns - actual_columns + errors = [f"Missing result column: {column}" for column in sorted(missing)] + return SQLValidationResult(valid=not errors, errors=errors) + + +class ChartRuleEngine: + def select_chart( + self, + intent: Intent, + dimensions: list[ColumnRef], + metrics: list[MetricDefinition], + ) -> str: + if intent.chart_type and intent.chart_type != "auto": + return intent.chart_type + if intent.ranking: + return "bar" + if len(metrics) >= 2: + return "scatter" + if len(metrics) == 1 and not dimensions: + return "card" + if any(is_temporal_type(dimension.data_type) for dimension in dimensions): + return "line" + if "percent" in " ".join(tokenize(" ".join(d.column for d in dimensions))): + return "pie" + if dimensions and metrics: + return "bar" + return "" + + +def add_months(value: date, months: int) -> date: + month_index = value.month - 1 + months + year = value.year + month_index // 12 + month = month_index % 12 + 1 + day = min(value.day, month_days(year, month)) + return date(year, month, day) + + +def month_days(year: int, month: int) -> int: + if month == 2: + return 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28 + return [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31][month - 1] + + +def humanize(value: str) -> str: + return re.sub(r"[_\s]+", " ", str(value or "")).strip().title() + + +def safe_alias(value: str) -> str: + alias = re.sub(r"[^A-Za-z0-9_]+", "_", str(value or "").strip()).strip("_") + return alias.lower() or "value" + + +def normalize_data_source(data_source: str | None) -> str: + normalized = (data_source or "").strip().upper().replace("-", "_").replace(" ", "_") + if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: + return "MSSQL" + return normalized + + +def compile_semantic_sql( + query: str, + documents: list[str], + semantic_analysis: dict[str, Any] | None = None, + data_source: str = "", + now: datetime | None = None, +) -> CompileResult | None: + catalog = SchemaParser().parse(documents) + if not catalog.tables: + return None + plan = SemanticPlanner(catalog, semantic_analysis).build_plan(query, now=now) + if not plan: + return None + sql = SQLCompiler().compile(plan, data_source=data_source) + validation = SQLAstValidator().validate(sql, plan) + if not validation.valid: + logger.info("deterministic_semantic_sql_validation_failed errors=%s", validation.errors) + return None + return CompileResult(sql=sql, plan=plan, validation=validation) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 59bea2279c..227ffdf709 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -11,6 +11,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata +from src.pipelines.generation.semantic_sql import compile_semantic_sql from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, @@ -254,6 +255,37 @@ async def run( logger.info("SQL Generation pipeline is running...") metadata = await retrieve_metadata(project_id or "", self._retriever) + data_source = metadata.get("data_source", "local_file") + + deterministic_result = compile_semantic_sql( + query=query, + documents=contexts, + semantic_analysis=schema_intent_analysis, + data_source=data_source, + ) + if deterministic_result: + logger.info( + "Deterministic semantic SQL compiler produced SQL for query: %s", + query, + ) + post_process_result = await self._components["post_processor"].run( + [deterministic_result.sql], + project_id=project_id, + use_dry_plan=use_dry_plan, + data_source=data_source, + allow_dry_plan_fallback=allow_dry_plan_fallback, + allow_data_preview=allow_data_preview, + valid_table_names=construct_valid_table_names(contexts), + valid_table_columns=construct_valid_table_columns(contexts), + query=query, + semantic_analysis=schema_intent_analysis, + ) + if post_process_result.get("valid_generation_result"): + return {"post_process": post_process_result} + logger.info( + "Deterministic semantic SQL rejected by post processor; falling back to LLM. error=%s", + post_process_result.get("invalid_generation_result", {}).get("error"), + ) return await self._pipe.execute( ["post_process"], @@ -270,7 +302,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, **self._components, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py b/wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py new file mode 100644 index 0000000000..ecbe66191a --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py @@ -0,0 +1,156 @@ +from datetime import datetime + +from src.pipelines.generation.semantic_sql import ( + IntentDetector, + SchemaParser, + compile_semantic_sql, +) + + +def test_intent_detector_returns_structured_business_intent(): + intent = IntentDetector().detect( + "Show top 10 customers by total invoice amount last month as a bar chart" + ) + + assert intent.question_type == "ranking" + assert intent.chart_requested is True + assert intent.chart_type == "bar" + assert intent.ranking is True + assert intent.top_n == 10 + assert intent.aggregation == "SUM" + + +def test_compile_semantic_sql_generates_ranked_aggregate_with_date_filter(): + documents = [ + """ + CREATE TABLE invoices ( + id INTEGER, + customer_id INTEGER, + customer_name VARCHAR, + invoice_amount DECIMAL(10,2), + invoice_date DATE + ); + """ + ] + + result = compile_semantic_sql( + "Show top 10 customers by total invoice amount last month as a bar chart", + documents, + now=datetime(2026, 7, 10), + ) + + assert result is not None + assert result.plan.intent.question_type == "ranking" + assert result.plan.chart_type == "bar" + assert result.plan.metrics[0].column.object_name == "invoices.invoice_amount" + assert [dimension.object_name for dimension in result.plan.group_by] == [ + "invoices.customer_name" + ] + assert 'SUM("invoices"."invoice_amount")' in result.sql + assert 'GROUP BY "invoices"."customer_name"' in result.sql + assert "\"invoices\".\"invoice_date\" >= '2026-06-01'" in result.sql + assert "\"invoices\".\"invoice_date\" < '2026-07-01'" in result.sql + assert "ORDER BY SUM(" in result.sql + assert "LIMIT 10" in result.sql + + +def test_compile_semantic_sql_resolves_join_path_from_foreign_keys(): + documents = [ + """ + CREATE TABLE orders ( + id INTEGER, + customer_id INTEGER, + order_amount DECIMAL(10,2), + FOREIGN KEY (customer_id) REFERENCES customers(id) + ); + """, + """ + CREATE TABLE customers ( + id INTEGER, + customer_name VARCHAR + ); + """, + ] + + result = compile_semantic_sql( + "Show total order amount by customer name", + documents, + ) + + assert result is not None + assert [join.left_table for join in result.plan.joins] == ["orders"] + assert [join.right_table for join in result.plan.joins] == ["customers"] + assert ( + 'INNER JOIN "customers" ON "orders"."customer_id" = "customers"."id"' + in result.sql + ) + assert 'GROUP BY "customers"."customer_name"' in result.sql + + +def test_compile_semantic_sql_uses_semantic_metadata_documents(): + documents = [ + """ + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "id", "type": "INTEGER"}, + {"name": "customer_id", "type": "INTEGER"}, + {"name": "order_value", "type": "DECIMAL"}, + {"name": "created_at", "type": "DATE"} + ] + }, + { + "name": "customers", + "columns": [ + {"name": "id", "type": "INTEGER"}, + {"name": "customer_name", "type": "VARCHAR"} + ] + } + ], + "relationships": [ + { + "condition": "orders.customer_id = customers.id", + "joinType": "MANY_TO_ONE", + "models": ["orders", "customers"] + } + ] + } + """ + ] + + result = compile_semantic_sql( + "Show total order value by customer name this month", + documents, + now=datetime(2026, 7, 10), + ) + + assert result is not None + assert 'SUM("orders"."order_value")' in result.sql + assert 'INNER JOIN "customers" ON "orders"."customer_id" = "customers"."id"' in result.sql + assert "\"orders\".\"created_at\" >= '2026-07-01'" in result.sql + assert "\"orders\".\"created_at\" < '2026-08-01'" in result.sql + + +def test_compile_semantic_sql_counts_entities_without_unrequested_grouping(): + result = compile_semantic_sql( + "How many customers are there?", + ["CREATE TABLE customers (id INTEGER, customer_name VARCHAR);"], + ) + + assert result is not None + assert result.sql == 'SELECT COUNT(*) AS "count" FROM "customers"' + + +def test_schema_parser_handles_single_line_create_table_definitions(): + catalog = SchemaParser().parse( + ['CREATE TABLE customers (id INTEGER, customer_name VARCHAR, created_at DATE);'] + ) + + assert list(catalog.tables) == ["customers"] + assert [column.column for column in catalog.tables["customers"]] == [ + "id", + "customer_name", + "created_at", + ] From 85e948c3fdf24dc8d4974b562dacba90dc4be712 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 11 Jul 2026 00:17:56 +0530 Subject: [PATCH 0489/1087] Refactor SQL generation around semantic schema contracts --- .../generation/followup_sql_generation.py | 15 + .../pipelines/generation/sql_correction.py | 15 + .../pipelines/generation/sql_generation.py | 15 + .../pipelines/generation/sql_regeneration.py | 17 + .../src/pipelines/generation/utils/sql.py | 440 +++++++++++++++ .../retrieval/db_schema_retrieval.py | 520 +++++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 505 +++++++++++++++-- .../src/web/v1/services/ask_feedback.py | 10 + .../pipelines/generation/test_sql_utils.py | 68 +++ .../retrieval/test_db_schema_retrieval.py | 66 +++ 10 files changed, 1610 insertions(+), 61 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 e6e6b0c4c3..ed5f419c5b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -18,6 +18,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -97,6 +98,10 @@ {% endfor %} {% endif %} +{% if semantic_schema_contract %} +{{ semantic_schema_contract }} +{% endif %} + ### QUESTION ### User's Follow-up Question: {{ query }} @@ -130,6 +135,7 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -155,6 +161,9 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -190,6 +199,8 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -199,6 +210,8 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -250,6 +263,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -303,6 +317,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": data_source, "sql_knowledge": sql_knowledge, + "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index daae37e02e..fff595b461 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_sql_generation_model_kwargs, @@ -98,6 +99,10 @@ def get_sql_correction_system_prompt( {% endfor %} {% endif %} +{% if semantic_schema_contract %} +{{ semantic_schema_contract }} +{% endif %} + ### QUESTION ### {% if query %} User's Question: {{ query }} @@ -130,6 +135,7 @@ def prompt( query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -141,6 +147,9 @@ def prompt( instructions=instructions, ), sql_functions=sql_functions, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -172,6 +181,8 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -181,6 +192,8 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -227,6 +240,7 @@ async def run( allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -238,6 +252,7 @@ async def run( "invalid_generation_result": invalid_generation_result, "documents": contexts, "query": query, + "schema_intent_analysis": schema_intent_analysis, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 227ffdf709..56d2d0d399 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -85,6 +86,10 @@ {% endfor %} {% endif %} +{% if semantic_schema_contract %} +{{ semantic_schema_contract }} +{% endif %} + ### QUESTION ### User's Question: {{ query }} @@ -119,6 +124,7 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: schema_context = "\n".join(documents or []).lower() has_pcb_context = any( @@ -158,6 +164,9 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -190,6 +199,8 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, + query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -200,6 +211,8 @@ async def post_process( allow_data_preview=allow_data_preview, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -251,6 +264,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -305,6 +319,7 @@ async def run( "data_source": data_source, "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, + "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 8b3eaafcb1..5c2089d065 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -101,6 +102,10 @@ def get_sql_regeneration_system_prompt( {% endfor %} {% endif %} +{% if semantic_schema_contract %} +{{ semantic_schema_contract }} +{% endif %} + ### QUESTION ### SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} @@ -124,6 +129,7 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( sql=sql, @@ -148,6 +154,9 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + semantic_schema_contract=construct_semantic_schema_contract( + schema_intent_analysis + ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -177,6 +186,8 @@ async def post_process( documents: list[str], data_source: str, project_id: str | None = None, + query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), @@ -184,6 +195,8 @@ async def post_process( data_source=data_source, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), + query=query, + semantic_analysis=schema_intent_analysis, ) @@ -228,6 +241,8 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + query: str | None = None, + schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -240,6 +255,8 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, + "query": query, + "schema_intent_analysis": schema_intent_analysis, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 8316d24e86..02839d802b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1543,6 +1543,8 @@ async def run( allow_data_preview: bool = False, valid_table_names: list[str] | None = None, valid_table_columns: dict[str, list[str]] | None = None, + query: str | None = None, + semantic_analysis: dict[str, Any] | None = None, ) -> dict: try: cleaned_generation_result = extract_sql_generation_result(replies[0]) @@ -1571,6 +1573,26 @@ async def run( }, } + placeholder_schema_references = _extract_placeholder_schema_references( + cleaned_generation_result + ) + if placeholder_schema_references: + invalid_placeholder_list = ", ".join(placeholder_schema_references) + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_VALIDATION", + "error": ( + "Generated SQL contains placeholder schema references " + f"that are not active datasource objects: {invalid_placeholder_list}." + ), + "invalid_schema_objects": placeholder_schema_references, + "correlation_id": "", + }, + } + invalid_table_references = find_invalid_table_references( cleaned_generation_result, valid_table_names or [], @@ -1590,6 +1612,7 @@ async def run( "Use only these valid table names exactly as shown: " f"{valid_table_list}" ), + "invalid_schema_objects": invalid_table_references, "correlation_id": "", }, } @@ -1613,6 +1636,25 @@ async def run( "Use only these valid table columns exactly as shown: " f"{valid_column_list}" ), + "invalid_schema_objects": invalid_column_references, + "correlation_id": "", + }, + } + + intent_validation_error = validate_sql_intent_alignment( + query, + cleaned_generation_result, + valid_table_columns or {}, + semantic_analysis=semantic_analysis, + ) + if intent_validation_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_INTENT_VALIDATION", + "error": intent_validation_error, "correlation_id": "", }, } @@ -2276,6 +2318,404 @@ def construct_instructions( return _instructions +def _semantic_analysis_items( + semantic_analysis: dict[str, Any] | None, + key: str, +) -> list[str]: + if not isinstance(semantic_analysis, dict): + return [] + value = semantic_analysis.get(key) + if isinstance(value, str): + return [value] if value.strip() else [] + if isinstance(value, list): + return [ + str(item).strip() + for item in value + if item is not None and str(item).strip() + ] + return [] + + +def _semantic_analysis_dict_items( + semantic_analysis: dict[str, Any] | None, + key: str, +) -> list[dict[str, Any]]: + if not isinstance(semantic_analysis, dict): + return [] + value = semantic_analysis.get(key) + if not isinstance(value, list): + return [] + return [item for item in value if isinstance(item, dict)] + + +def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: + if not isinstance(semantic_analysis, dict) or not semantic_analysis: + return False + semantic_keys = { + "analytical_intent", + "entities", + "identifiers", + "metrics", + "dimensions", + "filters", + "aggregations", + "relationships", + "time_constraints", + "ranking", + "sorting", + "requested_output", + "supported_schema_objects", + "candidate_schema_scores", + "concept_mappings", + "interpretations", + "missing_requirements", + "ambiguous_requirements", + "support_reasoning", + } + return any(semantic_analysis.get(key) for key in semantic_keys) + + +def _mapping_schema_objects(mapping: dict[str, Any]) -> list[str]: + value = mapping.get("schema_objects") + if isinstance(value, str): + return [value] if value.strip() else [] + if isinstance(value, list): + return [ + str(item).strip() + for item in value + if item is not None and str(item).strip() + ] + return [] + + +def _mapping_concept_type(mapping: dict[str, Any]) -> str: + return str(mapping.get("concept_type") or "").strip().lower() + + +def _mapping_request_concept(mapping: dict[str, Any]) -> str: + return str(mapping.get("request_concept") or "requested concept").strip() + + +def construct_semantic_schema_contract( + semantic_analysis: dict[str, Any] | None, +) -> str: + if not _has_semantic_analysis(semantic_analysis): + return "" + + lines = [ + "### SEMANTIC SCHEMA CONTRACT ###", + "Use this contract as the primary source of truth for SQL generation.", + "Generate SQL only from schema objects listed here and in the retrieved DATABASE SCHEMA.", + "Do not infer alternative tables, columns, joins, metrics, filters, or time fields independently.", + ] + analytical_intent = str( + semantic_analysis.get("analytical_intent") or "" + ).strip() + if analytical_intent: + lines.append(f"Analytical intent: {analytical_intent}") + + for label, key in ( + ("Entities", "entities"), + ("Identifiers", "identifiers"), + ("Metrics", "metrics"), + ("Dimensions", "dimensions"), + ("Filters", "filters"), + ("Time constraints", "time_constraints"), + ("Aggregations", "aggregations"), + ("Ranking", "ranking"), + ("Sorting", "sorting"), + ("Requested output", "requested_output"), + ("Relationships", "relationships"), + ("Supported schema objects", "supported_schema_objects"), + ): + values = _semantic_analysis_items(semantic_analysis, key) + if values: + lines.append(f"{label}: {', '.join(values)}") + + mappings = _semantic_analysis_dict_items(semantic_analysis, "concept_mappings") + if mappings: + lines.append("Required concept-to-schema mappings:") + for mapping in mappings: + schema_objects = _mapping_schema_objects(mapping) + required = "required" if mapping.get("required_in_sql") is not False else "optional" + confidence = mapping.get("confidence") + confidence_text = ( + f", confidence={confidence}" + if confidence is not None and str(confidence).strip() + else "" + ) + lines.append( + "- " + f"{_mapping_request_concept(mapping)} " + f"({_mapping_concept_type(mapping) or 'concept'}, {required}{confidence_text}) " + f"-> {', '.join(schema_objects) or 'NO_MAPPING'}." + ) + + interpretations = _semantic_analysis_dict_items( + semantic_analysis, "interpretations" + ) + if interpretations: + lines.append("Ranked schema interpretations:") + for interpretation in interpretations: + description = str(interpretation.get("description") or "").strip() + if not description: + continue + selected = "selected" if interpretation.get("is_selected") is True else "candidate" + schema_objects = interpretation.get("schema_objects") + schema_text = ( + ", ".join(str(item).strip() for item in schema_objects if item) + if isinstance(schema_objects, list) + else "" + ) + lines.append(f"- {description} ({selected}). Objects: {schema_text}") + + missing_requirements = _semantic_analysis_items( + semantic_analysis, "missing_requirements" + ) + if missing_requirements: + lines.append(f"Missing requirements: {', '.join(missing_requirements)}") + ambiguous_requirements = _semantic_analysis_items( + semantic_analysis, "ambiguous_requirements" + ) + if ambiguous_requirements: + lines.append(f"Ambiguous requirements: {', '.join(ambiguous_requirements)}") + support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() + if support_reasoning: + lines.append(f"Support reasoning: {support_reasoning}") + + lines.append( + "Validation requirement: every required mapped schema object must appear in SQL when it represents an entity, identifier, dimension, metric, filter, time field, relationship, or requested output. Ranking/sorting/aggregation concepts must appear as SQL shape." + ) + return "\n".join(lines) + + +def get_schema_intent_analysis_error( + semantic_analysis: dict[str, Any] | None, +) -> str | None: + if not _has_semantic_analysis(semantic_analysis): + return None + missing = _semantic_analysis_items(semantic_analysis, "missing_requirements") + if missing: + return ( + "The active datasource schema does not expose the information needed " + f"to answer the request: {', '.join(missing)}. I cannot generate unrelated SQL." + ) + ambiguous = _semantic_analysis_items(semantic_analysis, "ambiguous_requirements") + if ambiguous: + return ( + "The request has multiple equally plausible schema interpretations: " + f"{', '.join(ambiguous)}. Please clarify which one to use." + ) + candidates = _semantic_analysis_dict_items( + semantic_analysis, "candidate_schema_scores" + ) + if candidates and not any(candidate.get("is_complete") is True for candidate in candidates): + missing_concepts = [] + for candidate in candidates[:3]: + missing_value = candidate.get("missing_concepts") + if isinstance(missing_value, list) and missing_value: + missing_concepts.append( + f"{candidate.get('candidate_id') or 'candidate'}: " + + ", ".join(str(item) for item in missing_value if item) + ) + if missing_concepts: + return ( + "Semantic schema retrieval did not find a complete mapping for " + f"the request. Missing concepts: {'; '.join(missing_concepts)}." + ) + if semantic_analysis.get("is_fully_supported") is False: + reason = str(semantic_analysis.get("support_reasoning") or "").strip() + return reason or ( + "The selected schema does not fully support every required component of the request." + ) + schema_bound_types = { + "dimension", + "entity", + "filter", + "identifier", + "metric", + "relationship", + "time", + "output", + } + unsupported = [] + for mapping in _semantic_analysis_dict_items(semantic_analysis, "concept_mappings"): + if mapping.get("required_in_sql") is False: + continue + if _mapping_concept_type(mapping) not in schema_bound_types: + continue + if not _mapping_schema_objects(mapping): + unsupported.append(_mapping_request_concept(mapping)) + if unsupported: + return ( + "The semantic contract did not map required request concepts to active " + f"schema objects: {', '.join(unsupported)}." + ) + return None + + +def _extract_placeholder_schema_references(sql: str) -> list[str]: + placeholders = re.findall(r"<\s*([^<>]+?)\s*>", sql or "") + placeholder_names = re.findall( + r"\b(?:table|column|schema|database|field|metric|dimension|date|amount|entity)_?name\b", + sql or "", + flags=re.IGNORECASE, + ) + return sorted( + { + str(item).strip() + for item in [*placeholders, *placeholder_names] + if item is not None and str(item).strip() + } + ) + + +def _sql_references_table(sql: str, table_name: str) -> bool: + normalized_table = _compact_sql_identifier(table_name) + if not normalized_table: + return False + table_suffixes = { + _compact_sql_identifier(suffix) + for suffix in _table_reference_suffixes(table_name) + } + table_suffixes.add(normalized_table) + for referenced_table in extract_sql_table_references(sql or ""): + normalized_reference = _compact_sql_identifier(referenced_table) + if not normalized_reference: + continue + if normalized_reference in table_suffixes: + return True + if any( + normalized_reference.endswith(suffix) or suffix.endswith(normalized_reference) + for suffix in table_suffixes + if suffix + ): + return True + return False + + +def _sql_references_schema_object( + sql: str, + schema_object: str, + valid_table_columns: dict[str, list[str]], +) -> bool: + parts = [ + _normalize_sql_identifier(part.strip()) + for part in str(schema_object or "").split(".") + if part.strip() + ] + if not parts: + return False + if len(parts) == 1: + identifier = parts[0] + return bool( + re.search( + rf'(? bool: + return bool( + re.search(r"\bCOUNT\s*\(\s*\*\s*\)", sql or "", flags=re.IGNORECASE) + ) and not re.search( + r"\b(?:SUM|AVG|MIN|MAX)\s*\(", sql or "", flags=re.IGNORECASE + ) + + +def _semantic_requests_record_count(semantic_analysis: dict[str, Any]) -> bool: + intent = str(semantic_analysis.get("analytical_intent") or "").lower() + if intent == "record_count": + return True + text = " ".join( + item + for key in ("metrics", "aggregations", "requested_output") + for item in _semantic_analysis_items(semantic_analysis, key) + ).lower() + return bool(re.search(r"\b(?:count|number of|row count|record count)\b", text)) + + +def validate_sql_intent_alignment( + query: str | None, + sql: str, + valid_table_columns: dict[str, list[str]] | None = None, + semantic_analysis: dict[str, Any] | None = None, +) -> str | None: + valid_table_columns = valid_table_columns or {} + if analysis_error := get_schema_intent_analysis_error(semantic_analysis): + return analysis_error + if not _has_semantic_analysis(semantic_analysis): + return None + + requests_count = _semantic_requests_record_count(semantic_analysis) + mappings = _semantic_analysis_dict_items(semantic_analysis, "concept_mappings") + schema_bound_types = { + "dimension", + "entity", + "filter", + "identifier", + "metric", + "relationship", + "time", + "output", + } + for mapping in mappings: + if mapping.get("required_in_sql") is False: + continue + concept_type = _mapping_concept_type(mapping) + if concept_type not in schema_bound_types: + continue + schema_objects = _mapping_schema_objects(mapping) + if not schema_objects: + return ( + "The semantic analysis did not map the required " + f"{concept_type or 'concept'} '{_mapping_request_concept(mapping)}' " + "to a schema object." + ) + if not any( + _sql_references_schema_object(sql, schema_object, valid_table_columns) + for schema_object in schema_objects + ): + return ( + "Generated SQL does not reference schema objects mapped to the " + f"required {concept_type or 'concept'} " + f"'{_mapping_request_concept(mapping)}': {', '.join(schema_objects)}." + ) + if concept_type == "metric" and _sql_is_plain_count(sql) and not requests_count: + return ( + "Generated SQL answers with a generic record count, but the " + f"requested metric '{_mapping_request_concept(mapping)}' must be " + "retrieved or calculated from the mapped schema object." + ) + + if _semantic_analysis_items(semantic_analysis, "ranking") and not re.search( + r"\bORDER\s+BY\b.*\b(?:LIMIT|FETCH\s+FIRST|TOP\s*(?:\(\s*)?\d+)\b|\b(?:LIMIT|FETCH\s+FIRST|TOP\s*(?:\(\s*)?\d+)\b.*\bORDER\s+BY\b", + sql or "", + flags=re.IGNORECASE | re.DOTALL, + ): + return "Generated SQL does not include sorting and limiting logic required by the ranking intent." + if _semantic_analysis_items(semantic_analysis, "sorting") and not re.search( + r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE + ): + return "Generated SQL does not include sorting logic required by the semantic intent." + if _semantic_analysis_items(semantic_analysis, "aggregations") and not re.search( + r"\b(?:COUNT|SUM|AVG|MIN|MAX)\s*\(", sql or "", flags=re.IGNORECASE + ): + return "Generated SQL does not include the aggregation required by the semantic intent." + return None + + def _parse_semantic_metadata_content(content: str) -> Any | None: content = content.strip() if not content: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 107c5c3479..dd5c81a9b3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,5 +1,6 @@ import ast import logging +import re import sys from typing import TYPE_CHECKING, Any, Optional @@ -10,7 +11,7 @@ from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import BaseModel, Field from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider @@ -31,23 +32,78 @@ table_columns_selection_system_prompt = """ ### TASK ### -You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. +You are a highly skilled data analyst. Your goal is to examine the active deployed database schema, semantically interpret the user's question, and identify the exact schema objects required to construct an accurate SQL query. The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### -1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. -2. For each table, provide a clear and concise reasoning for why specific columns are selected. -3. List each reason as part of a step-by-step chain of thought, justifying the inclusion of each column. -4. If a "." is included in columns, put the name before the first dot into chosen columns. -5. The number of columns chosen must match the number of reasoning. -6. Final chosen columns must be only column names, don't prefix it with table names. -7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +1. First identify the user's semantic intent: requested entities, identifiers, metrics, dimensions, filters, aggregations, ranking, sorting, time constraints, joins/relationships, and requested output shape. +2. Map each required business concept only to tables, columns, metrics, views, or relationships that are explicitly present in the active schema metadata. +3. Use semantic similarity from table names, column names, comments/descriptions, data types, primary/foreign keys, relationships, and metric/view definitions. Do not use datasource-specific rules, hardcoded mappings, or default tables. +4. Rank candidate schema mappings by complete concept coverage, semantic fit, relationship viability, metric validity, data type compatibility, and support for filters/time/ranking/aggregation requirements. +5. Select only the highest-confidence complete candidate. If no candidate fully supports the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting unrelated fallback schema. +6. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Prefer the next-best complete candidate that excludes those objects. +7. Populate `concept_mappings` for every required concept. Each mapping must classify the concept, list directly supporting schema objects, mark whether it must appear in SQL, and include a confidence score. +8. Populate `candidate_schema_scores` with accepted and rejected candidates. Include covered and missing concepts and a concise selection reason. +9. Include join keys and relationship columns when multiple tables are needed. If no trustworthy join path exists, mark the relationship missing instead of inventing one. +10. Do not add filters or time constraints not requested or implied by the user. +11. If a "." is included in columns, put the name before the first dot into chosen columns. +12. The number of columns chosen must match the number of reasoning. +13. Final chosen columns must be only column names, don't prefix it with table names. +14. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: { + "semantic_analysis": { + "analytical_intent": "retrieval | detailed_records | summary | comparison | trend | dashboard | kpi | ranking | record_count | other", + "entities": ["requested business entities"], + "identifiers": ["requested identifiers"], + "metrics": ["requested measures or calculated metrics"], + "dimensions": ["requested grouping/descriptive dimensions"], + "filters": ["requested filters"], + "aggregations": ["requested aggregations/calculations"], + "relationships": ["required joins or relationship paths"], + "time_constraints": ["requested date filters, grains, or trend requirements"], + "ranking": ["top/bottom/rank/limit requirements"], + "sorting": ["requested ordering requirements"], + "requested_output": ["requested columns, charts, summaries, records, or KPIs"], + "supported_schema_objects": ["table.column, table, view, metric, or relationship objects selected"], + "candidate_schema_scores": [ + { + "candidate_id": "candidate-1", + "schema_objects": ["schema objects included in this candidate"], + "covered_concepts": ["request concepts this candidate supports"], + "missing_concepts": ["request concepts this candidate cannot support"], + "confidence": 0.0, + "is_complete": true, + "selection_reason": "Why this candidate is accepted or rejected" + } + ], + "concept_mappings": [ + { + "request_concept": "business concept from the user request", + "concept_type": "entity | identifier | dimension | metric | filter | time | aggregation | ranking | sorting | relationship | output", + "schema_objects": ["table.column, table, metric, view, or relationship object that directly supports the concept"], + "required_in_sql": true, + "confidence": 0.0, + "mapping_reason": "Why these schema objects directly support the concept" + } + ], + "interpretations": [ + { + "description": "Possible schema interpretation", + "schema_objects": ["schema objects used by this interpretation"], + "confidence": 0.0, + "is_selected": true + } + ], + "missing_requirements": ["required concepts not supported by active schema"], + "ambiguous_requirements": ["concepts with multiple equally plausible mappings"], + "is_fully_supported": true, + "support_reasoning": "Concise explanation of schema support" + }, "results": [ { "table_selection_reason": "Reason for selecting tablename1", @@ -82,6 +138,7 @@ - Each table key must list only the columns relevant to answering the question. - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. +- Populate `semantic_analysis` before `results` and keep both consistent. - Be logical, concise, and ensure the output strictly follows the required JSON format. - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. @@ -98,8 +155,38 @@ {{ db_schema }} {% endfor %} +{% if semantic_candidate_context %} +### GENERIC SEMANTIC CANDIDATE RANKING ### +The following candidates were scored from active schema metadata only. Use this as evidence, then validate complete concept coverage before selecting a semantic contract. +{% for candidate in semantic_candidate_context %} +- candidate_id: {{ candidate.candidate_id }} + table_name: {{ candidate.table_name }} + confidence: {{ candidate.confidence }} + coverage_score: {{ candidate.coverage_score }} + matched_query_terms: {{ candidate.matched_query_terms }} + missing_query_terms: {{ candidate.missing_query_terms }} + rejected_by_retry: {{ candidate.rejected_by_retry }} + selection_reason: {{ candidate.selection_reason }} + matched_columns: +{% for column in candidate.matched_columns %} + - {{ column.column_name }} (score={{ column.score }}, data_type={{ column.data_type }}, matched_terms={{ column.matched_terms }}) +{% endfor %} +{% endfor %} +{% endif %} + ### INPUT ### {{ question }} + +{% if semantic_retry_context %} +### RETRY CONTEXT ### +Previous semantic SQL validation failed. Discard the previous contract and retrieve the next-best complete schema mapping. +Validation failure: {{ semantic_retry_context.validation_error }} +Retry attempt: {{ semantic_retry_context.retry_attempt }} +Rejected schema objects: +{% for schema_object in semantic_retry_context.rejected_schema_objects %} +- {{ schema_object }} +{% endfor %} +{% endif %} """ @@ -189,6 +276,275 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: return deduped +_SEMANTIC_TOKEN_STOPWORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "for", + "from", + "give", + "have", + "how", + "in", + "is", + "me", + "of", + "on", + "or", + "show", + "that", + "the", + "to", + "with", + "table", + "view", +} + +_NUMERIC_INTENT_TERMS = { + "amount", + "avg", + "average", + "balance", + "cost", + "count", + "measure", + "metric", + "price", + "quantity", + "rate", + "sum", + "total", + "value", +} + +_TEMPORAL_INTENT_TERMS = { + "date", + "day", + "month", + "quarter", + "time", + "trend", + "week", + "year", +} + +_RANKING_INTENT_TERMS = { + "bottom", + "highest", + "least", + "lowest", + "most", + "rank", + "ranking", + "top", +} + + +def _semantic_tokens(value: Any) -> set[str]: + text = str(value or "") + if not text: + return set() + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) + text = re.sub(r"[^A-Za-z0-9]+", " ", text) + tokens = { + token.lower() + for token in text.split() + if len(token) > 1 and token.lower() not in _SEMANTIC_TOKEN_STOPWORDS + } + for token in list(tokens): + if token.endswith("ies") and len(token) > 4: + tokens.add(f"{token[:-3]}y") + elif token.endswith("s") and len(token) > 3: + tokens.add(token[:-1]) + return tokens + + +def _column_tokens(column: dict[str, Any]) -> set[str]: + tokens = set() + for key in ("name", "display_name", "alias", "comment", "description", "data_type"): + tokens.update(_semantic_tokens(column.get(key))) + return tokens + + +def _table_tokens(table_schema: dict[str, Any]) -> set[str]: + tokens = set() + for key in ("name", "display_name", "alias", "comment", "description"): + tokens.update(_semantic_tokens(table_schema.get(key))) + for column in table_schema.get("columns", []) or []: + if isinstance(column, dict): + tokens.update(_column_tokens(column)) + return tokens + + +def _is_numeric_column(column: dict[str, Any]) -> bool: + return bool( + re.search( + r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|number|double|float|real|money)\b", + str(column.get("data_type") or "").lower(), + ) + ) + + +def _is_temporal_column(column: dict[str, Any]) -> bool: + return bool( + re.search( + r"\b(?:date|time|timestamp|datetime)\b", + str(column.get("data_type") or "").lower(), + ) + ) + + +def _normalized_schema_object(value: Any) -> str: + return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) + + +def _rejected_schema_objects(semantic_retry_context: dict[str, Any] | None) -> set[str]: + if not isinstance(semantic_retry_context, dict): + return set() + rejected = semantic_retry_context.get("rejected_schema_objects") + if not isinstance(rejected, list): + return set() + return { + _normalized_schema_object(item) + for item in rejected + if item is not None and str(item).strip() + } + + +def _schema_object_was_rejected( + table_name: str, + column_name: str | None, + rejected_schema_objects: set[str], +) -> bool: + if not rejected_schema_objects: + return False + object_key = _normalized_schema_object( + f"{table_name}.{column_name}" if column_name else table_name + ) + table_key = _normalized_schema_object(table_name) + return any( + rejected_key + and ( + rejected_key == object_key + or rejected_key == table_key + or rejected_key.endswith(object_key) + or object_key.endswith(rejected_key) + ) + for rejected_key in rejected_schema_objects + ) + + +def rank_semantic_schema_candidates( + query: str, + construct_db_schemas: list[dict], + semantic_retry_context: dict[str, Any] | None = None, + max_candidates: int = 15, + max_columns_per_candidate: int = 8, +) -> list[dict[str, Any]]: + query_terms = _semantic_tokens(query) + if not query_terms: + return [] + numeric_terms = query_terms & _NUMERIC_INTENT_TERMS + temporal_terms = query_terms & _TEMPORAL_INTENT_TERMS + ranking_terms = query_terms & _RANKING_INTENT_TERMS + rejected_objects = _rejected_schema_objects(semantic_retry_context) + candidates: list[dict[str, Any]] = [] + + for table_schema in construct_db_schemas: + if table_schema.get("type") != "TABLE": + continue + table_name = str(table_schema.get("name") or "").strip() + if not table_name: + continue + table_matches = _table_tokens(table_schema) & query_terms + table_rejected = _schema_object_was_rejected( + table_name, None, rejected_objects + ) + matched_columns = [] + for column in table_schema.get("columns", []) or []: + if not isinstance(column, dict): + continue + column_name = str(column.get("name") or "").strip() + if not column_name: + continue + tokens = _column_tokens(column) + matched_terms = sorted(tokens & query_terms) + score = len(matched_terms) * 3.0 + if numeric_terms and _is_numeric_column(column): + score += 1.0 + if temporal_terms and _is_temporal_column(column): + score += 1.0 + if ranking_terms and matched_terms: + score += 0.5 + rejected = _schema_object_was_rejected( + table_name, column_name, rejected_objects + ) + if rejected: + score -= 5.0 + if score > 0 or matched_terms: + matched_columns.append( + { + "column_name": column_name, + "score": round(max(score, 0.0), 3), + "matched_terms": matched_terms, + "data_type": str(column.get("data_type") or ""), + "rejected_by_retry": rejected, + } + ) + + matched_columns.sort( + key=lambda item: (item["score"], len(item["matched_terms"])), + reverse=True, + ) + matched_columns = matched_columns[:max_columns_per_candidate] + covered_terms = set(table_matches) + for column in matched_columns: + covered_terms.update(column["matched_terms"]) + if not covered_terms and not table_rejected: + continue + coverage_score = len(covered_terms) / max(len(query_terms), 1) + raw_score = ( + len(table_matches) * 2 + + sum(column["score"] for column in matched_columns) + + coverage_score * 4 + ) + if table_rejected: + raw_score -= 6 + candidates.append( + { + "candidate_id": f"candidate-{len(candidates) + 1}", + "table_name": table_name, + "confidence": round(min(max(raw_score / 20, 0), 0.99), 3), + "coverage_score": round(coverage_score, 3), + "matched_query_terms": sorted(covered_terms), + "missing_query_terms": sorted(query_terms - covered_terms), + "matched_columns": matched_columns, + "rejected_by_retry": table_rejected + or any(column["rejected_by_retry"] for column in matched_columns), + "selection_reason": ( + f"Covers {len(covered_terms)} of {len(query_terms)} significant request terms" + ), + } + ) + + candidates.sort( + key=lambda item: ( + item["rejected_by_retry"] is False, + item["confidence"], + item["coverage_score"], + ), + reverse=True, + ) + for index, candidate in enumerate(candidates[:max_candidates], start=1): + candidate["candidate_id"] = f"candidate-{index}" + return candidates[:max_candidates] + + @observe(capture_input=False, capture_output=False) async def embedding( query: str, @@ -349,12 +705,22 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) + if enable_column_pruning or _token_count > context_window_size: + return { + "db_schemas": [], + "tokens": _token_count, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + "semantic_analysis": {}, + } return { "db_schemas": retrieval_results, "tokens": _token_count, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, + "semantic_analysis": {}, } @@ -365,6 +731,7 @@ def prompt( prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, histories: list[AskHistory], + semantic_retry_context: dict[str, Any] | None = None, ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -377,8 +744,22 @@ def prompt( ) query = "\n".join(previous_query_summaries) + "\n" + query + semantic_candidate_context = rank_semantic_schema_candidates( + query=query, + construct_db_schemas=construct_db_schemas, + semantic_retry_context=semantic_retry_context, + ) + logger.info( + "semantic_retrieval_pre_ranked_candidates=%s", + semantic_candidate_context, + ) - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) + _prompt = prompt_builder.run( + question=query, + db_schemas=db_schemas, + semantic_candidate_context=semantic_candidate_context, + semantic_retry_context=semantic_retry_context or {}, + ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: return {} @@ -405,9 +786,10 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] + retrieval_payload = orjson.loads(filter_columns_in_tables["replies"][0]) + columns_and_tables_needed = retrieval_payload.get("results", []) + semantic_analysis = retrieval_payload.get("semantic_analysis") or {} + _log_semantic_retrieval_decision(semantic_analysis) # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -467,6 +849,7 @@ def construct_retrieval_results( "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, + "semantic_analysis": semantic_analysis, } else: retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] @@ -478,9 +861,64 @@ def construct_retrieval_results( ], "has_metric": check_using_db_schemas_without_pruning["has_metric"], "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], + "semantic_analysis": check_using_db_schemas_without_pruning.get( + "semantic_analysis", {} + ), } +def _semantic_log_items(semantic_analysis: dict[str, Any], key: str) -> list[str]: + value = semantic_analysis.get(key) + if isinstance(value, str): + return [value] if value.strip() else [] + if isinstance(value, list): + return [ + str(item).strip() + for item in value + if item is not None and str(item).strip() + ] + return [] + + +def _log_semantic_retrieval_decision(semantic_analysis: dict[str, Any]) -> None: + if not isinstance(semantic_analysis, dict) or not semantic_analysis: + logger.info("semantic_retrieval_decision=no_semantic_analysis") + return + + logger.info( + "semantic_retrieval_concepts=%s", + { + "intent": semantic_analysis.get("analytical_intent"), + "entities": _semantic_log_items(semantic_analysis, "entities"), + "identifiers": _semantic_log_items(semantic_analysis, "identifiers"), + "metrics": _semantic_log_items(semantic_analysis, "metrics"), + "dimensions": _semantic_log_items(semantic_analysis, "dimensions"), + "filters": _semantic_log_items(semantic_analysis, "filters"), + "time_constraints": _semantic_log_items( + semantic_analysis, "time_constraints" + ), + "aggregations": _semantic_log_items(semantic_analysis, "aggregations"), + "ranking": _semantic_log_items(semantic_analysis, "ranking"), + "sorting": _semantic_log_items(semantic_analysis, "sorting"), + }, + ) + logger.info( + "semantic_retrieval_candidate_scores=%s", + semantic_analysis.get("candidate_schema_scores") or [], + ) + logger.info( + "semantic_retrieval_selected_contract=%s", + { + "supported_schema_objects": semantic_analysis.get( + "supported_schema_objects", [] + ), + "concept_mappings": semantic_analysis.get("concept_mappings", []), + "is_fully_supported": semantic_analysis.get("is_fully_supported"), + "support_reasoning": semantic_analysis.get("support_reasoning"), + }, + ) + + ## End of Pipeline class MatchingTableContents(BaseModel): chain_of_thought_reasoning: list[str] @@ -493,7 +931,59 @@ class MatchingTable(BaseModel): table_selection_reason: str +class SemanticConceptMapping(BaseModel): + request_concept: str = "" + concept_type: str = "" + schema_objects: list[str] = Field(default_factory=list) + required_in_sql: bool = True + confidence: float | None = None + mapping_reason: str = "" + + +class SemanticInterpretation(BaseModel): + description: str = "" + schema_objects: list[str] = Field(default_factory=list) + confidence: float | None = None + is_selected: bool = False + + +class SemanticCandidateSchemaScore(BaseModel): + candidate_id: str = "" + schema_objects: list[str] = Field(default_factory=list) + covered_concepts: list[str] = Field(default_factory=list) + missing_concepts: list[str] = Field(default_factory=list) + confidence: float | None = None + is_complete: bool = False + selection_reason: str = "" + + +class SemanticAnalysis(BaseModel): + analytical_intent: str = "" + entities: list[str] = Field(default_factory=list) + identifiers: list[str] = Field(default_factory=list) + metrics: list[str] = Field(default_factory=list) + dimensions: list[str] = Field(default_factory=list) + filters: list[str] = Field(default_factory=list) + aggregations: list[str] = Field(default_factory=list) + relationships: list[str] = Field(default_factory=list) + time_constraints: list[str] = Field(default_factory=list) + ranking: list[str] = Field(default_factory=list) + sorting: list[str] = Field(default_factory=list) + requested_output: list[str] = Field(default_factory=list) + supported_schema_objects: list[str] = Field(default_factory=list) + candidate_schema_scores: list[SemanticCandidateSchemaScore] = Field( + default_factory=list + ) + concept_mappings: list[SemanticConceptMapping] = Field(default_factory=list) + interpretations: list[SemanticInterpretation] = Field(default_factory=list) + missing_requirements: list[str] = Field(default_factory=list) + ambiguous_requirements: list[str] = Field(default_factory=list) + is_fully_supported: bool | None = None + support_reasoning: str = "" + + class RetrievalResults(BaseModel): + semantic_analysis: SemanticAnalysis | None = None results: list[MatchingTable] @@ -562,8 +1052,11 @@ async def run( project_id: Optional[str] = None, histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, + semantic_retry_context: Optional[dict[str, Any]] = None, ): logger.info("Ask Retrieval pipeline is running...") + if semantic_retry_context: + logger.info("semantic_retrieval_retry_context=%s", semantic_retry_context) return await self._pipe.execute( ["construct_retrieval_results"], inputs={ @@ -572,6 +1065,7 @@ async def run( "project_id": project_id or "", "histories": histories or [], "enable_column_pruning": enable_column_pruning, + "semantic_retry_context": semantic_retry_context or {}, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9db38bd5a8..e0f62e945a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -11,6 +11,7 @@ from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, construct_valid_table_names, + get_schema_intent_analysis_error, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, ) @@ -4429,6 +4430,94 @@ def _metadata_from_documents( ] return table_names, table_ddls + def _semantic_analysis_has_contract( + self, semantic_analysis: dict[str, Any] | None + ) -> bool: + if not isinstance(semantic_analysis, dict) or not semantic_analysis: + return False + if semantic_analysis.get("supported_schema_objects"): + return True + mappings = semantic_analysis.get("concept_mappings") + if isinstance(mappings, list) and any( + isinstance(mapping, dict) and mapping.get("schema_objects") + for mapping in mappings + ): + return True + candidates = semantic_analysis.get("candidate_schema_scores") + if isinstance(candidates, list) and any( + isinstance(candidate, dict) + and candidate.get("is_complete") is True + and candidate.get("schema_objects") + for candidate in candidates + ): + return True + return False + + def _semantic_schema_objects( + self, semantic_analysis: dict[str, Any] | None + ) -> list[str]: + if not isinstance(semantic_analysis, dict): + return [] + schema_objects: list[str] = [] + supported = semantic_analysis.get("supported_schema_objects") + if isinstance(supported, list): + schema_objects.extend(str(item).strip() for item in supported if item) + mappings = semantic_analysis.get("concept_mappings") + if isinstance(mappings, list): + for mapping in mappings: + if not isinstance(mapping, dict): + continue + mapped_objects = mapping.get("schema_objects") + if isinstance(mapped_objects, str) and mapped_objects.strip(): + schema_objects.append(mapped_objects.strip()) + elif isinstance(mapped_objects, list): + schema_objects.extend( + str(item).strip() for item in mapped_objects if item + ) + return list(dict.fromkeys(item for item in schema_objects if item)) + + def _semantic_retry_context( + self, + *, + semantic_analysis: dict[str, Any] | None, + invalid_generation_result: dict[str, Any] | None, + retry_attempt: int, + ) -> dict[str, Any]: + invalid_generation_result = invalid_generation_result or {} + rejected_schema_objects = self._semantic_schema_objects(semantic_analysis) + invalid_schema_objects = invalid_generation_result.get("invalid_schema_objects") + if isinstance(invalid_schema_objects, list): + rejected_schema_objects.extend( + str(item).strip() for item in invalid_schema_objects if item + ) + elif isinstance(invalid_schema_objects, str) and invalid_schema_objects.strip(): + rejected_schema_objects.append(invalid_schema_objects.strip()) + + selected_candidates = [] + if isinstance(semantic_analysis, dict): + candidates = semantic_analysis.get("candidate_schema_scores") + if isinstance(candidates, list): + selected_candidates = [ + candidate + for candidate in candidates + if isinstance(candidate, dict) + and ( + candidate.get("is_complete") is True + or candidate.get("schema_objects") + ) + ][:3] + + return { + "retry_attempt": retry_attempt, + "failure_type": invalid_generation_result.get("type"), + "failure_reason": invalid_generation_result.get("error"), + "rejected_sql": invalid_generation_result.get("sql"), + "rejected_schema_objects": list( + dict.fromkeys(item for item in rejected_schema_objects if item) + ), + "rejected_candidates": selected_candidates, + } + async def _complete_sql_generation_context( self, *, @@ -5280,6 +5369,9 @@ async def ask( table_names = [] table_ddls = [] _retrieval_result = {} + schema_intent_analysis: dict[str, Any] = {} + semantic_pipeline_active = False + semantic_contract_available = False error_message = None invalid_sql = None allow_sql_generation_reasoning = ( @@ -5400,7 +5492,7 @@ async def ask( tables=explicit_table_names, project_id=ask_request.project_id, histories=[], - enable_column_pruning=False, + enable_column_pruning=True, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -5426,10 +5518,10 @@ async def ask( retrieval_result = await self._run_with_timeout( "Full active schema retrieval for explicit table", self._pipelines["db_schema_retrieval"].run( - query="", + query=user_query, project_id=ask_request.project_id, histories=[], - enable_column_pruning=False, + enable_column_pruning=True, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -5449,14 +5541,30 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) + semantic_pipeline_active = True + semantic_contract_available = self._semantic_analysis_has_contract( + schema_intent_analysis + ) + logger.info( + "sql_generation_pipeline_decision query_id=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s", + query_id, + semantic_pipeline_active, + semantic_contract_available, + self._semantic_schema_objects(schema_intent_analysis), + ) logger.info( "Retrieved explicit tables for query_id %s: %s", query_id, table_names, ) - if table_question_sql := self._build_schema_grounded_table_question_sql( + if not semantic_pipeline_active and ( + table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls + ) ): ask_result = self._build_validated_ask_result_from_sql( table_question_sql, @@ -5480,8 +5588,10 @@ async def ask( return results invalid_sql = table_question_sql - if explicit_table_preview := self._build_explicit_table_preview_sql( + if not semantic_pipeline_active and ( + explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls + ) ): explicit_sql, explicit_table_name = explicit_table_preview if explicit_table_name not in table_names: @@ -5508,7 +5618,7 @@ async def ask( return results invalid_sql = explicit_sql - if documents and ( + if not semantic_pipeline_active and documents and ( deterministic_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -5948,10 +6058,7 @@ async def ask( query=sql_user_query, histories=[], project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) - ), + enable_column_pruning=True, ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -5963,12 +6070,12 @@ async def ask( error, ) retrieval_result = await self._run_with_timeout( - "Deployed schema fallback retrieval", + "Semantic schema retry after retrieval timeout", self._pipelines["db_schema_retrieval"].run( - query="", + query=sql_user_query, histories=[], project_id=ask_request.project_id, - enable_column_pruning=False, + enable_column_pruning=True, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -5978,6 +6085,20 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) + semantic_pipeline_active = True + semantic_contract_available = self._semantic_analysis_has_contract( + schema_intent_analysis + ) + logger.info( + "sql_generation_pipeline_decision query_id=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s", + query_id, + semantic_pipeline_active, + semantic_contract_available, + self._semantic_schema_objects(schema_intent_analysis), + ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6005,7 +6126,7 @@ async def ask( tables=explicit_table_names, project_id=ask_request.project_id, histories=[], - enable_column_pruning=enable_column_pruning, + enable_column_pruning=True, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -6015,6 +6136,13 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) + semantic_pipeline_active = True + semantic_contract_available = self._semantic_analysis_has_contract( + schema_intent_analysis + ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6024,7 +6152,11 @@ async def ask( documents, ) ) - if not documents and self._is_data_analysis_query(user_query): + if ( + not semantic_pipeline_active + and not documents + and self._is_data_analysis_query(user_query) + ): logger.info( "Query-based schema retrieval returned no tables for data question; " "retrying full active deployed schema for query_id %s", @@ -6047,6 +6179,13 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) + semantic_pipeline_active = True + semantic_contract_available = self._semantic_analysis_has_contract( + schema_intent_analysis + ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6061,7 +6200,7 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6081,7 +6220,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6105,7 +6244,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6129,6 +6268,7 @@ async def ask( if ( not api_results + and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and ( schema_grounded_sql := self._build_schema_grounded_analytics_sql( @@ -6153,16 +6293,20 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if not api_results and any( - term in user_query.lower() - for term in ( - "pcb", - "repair", - "failure", - "business unit", - "business units", - "product line", - "product family", + if ( + not semantic_pipeline_active + and not api_results + and any( + term in user_query.lower() + for term in ( + "pcb", + "repair", + "failure", + "business unit", + "business units", + "product line", + "product family", + ) ) ): operational_sql = self._build_schema_grounded_analytics_sql( @@ -6186,7 +6330,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if not api_results and ( + if not semantic_pipeline_active and not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6210,6 +6354,7 @@ async def ask( should_retry_full_schema = ( not api_results + and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and "db_schema_retrieval" in self._pipelines ) @@ -6319,7 +6464,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if not documents: + if not semantic_pipeline_active and not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): @@ -6406,12 +6551,23 @@ async def ask( table_ddls=table_ddls, ) if completed_retrieval_result: - _retrieval_result = completed_retrieval_result + completed_semantic_analysis = completed_retrieval_result.get( + "semantic_analysis", {} + ) + if self._semantic_analysis_has_contract( + completed_semantic_analysis + ): + _retrieval_result = completed_retrieval_result + schema_intent_analysis = completed_semantic_analysis + semantic_pipeline_active = True + semantic_contract_available = True sql_generation_histories = histories if self._is_data_analysis_query( sql_user_query - ) and not self._needs_conversation_context(sql_user_query): + ) and not self._needs_conversation_context( + sql_user_query + ) and not semantic_pipeline_active: sql_generation_histories = [] allow_sql_generation_reasoning = False allow_sql_knowledge_retrieval = False @@ -6420,10 +6576,31 @@ async def ask( "Using fast standalone SQL generation path for query_id %s", query_id, ) + elif semantic_pipeline_active: + logger.info( + "fast_standalone_sql_generation_disabled query_id=%s reason=semantic_pipeline_active", + query_id, + ) + + if ( + semantic_pipeline_active + and not semantic_contract_available + and not api_results + ): + error_message = ( + "Semantic schema retrieval did not produce a complete active-schema " + "contract for this request. I cannot generate unrelated SQL." + ) + logger.info( + "semantic_pipeline_no_contract query_id=%s reason=%s", + query_id, + error_message, + ) if ( not self._is_stopped(query_id, self._ask_results) and not api_results + and (not semantic_pipeline_active or semantic_contract_available) and allow_sql_generation_reasoning ): self._ask_results[query_id] = AskResultResponse( @@ -6495,7 +6672,11 @@ async def ask( is_followup=True if histories else False, ) - if not self._is_stopped(query_id, self._ask_results) and not api_results: + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and (not semantic_pipeline_active or semantic_contract_available) + ): self._ask_results[query_id] = AskResultResponse( status="generating", type="TEXT_TO_SQL", @@ -6561,6 +6742,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, ), ) else: @@ -6580,12 +6762,14 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, ), ) except TimeoutError as generation_timeout: logger.warning( - "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + "SQL generation timed out for query_id %s; semantic_pipeline_active=%s error=%s", query_id, + semantic_pipeline_active, generation_timeout, ) text_to_sql_generation_results = { @@ -6613,7 +6797,210 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - while current_sql_correction_retries < max_sql_correction_retries: + semantic_validation_failure_types = { + "SCHEMA_INTENT_VALIDATION", + "SCHEMA_VALIDATION", + "DRY_RUN", + } + semantic_retry_attempt = 0 + max_semantic_retries = 3 + while ( + semantic_pipeline_active + and failed_dry_run_result + and failed_dry_run_result.get("type") + in semantic_validation_failure_types + and semantic_retry_attempt < max_semantic_retries + and not api_results + ): + semantic_retry_attempt += 1 + retry_context = self._semantic_retry_context( + semantic_analysis=schema_intent_analysis, + invalid_generation_result=failed_dry_run_result, + retry_attempt=semantic_retry_attempt, + ) + logger.info( + "semantic_sql_retry_start query_id=%s retry_attempt=%s failure_type=%s rejected_schema_objects=%s failure_reason=%s", + query_id, + semantic_retry_attempt, + failed_dry_run_result.get("type"), + retry_context.get("rejected_schema_objects"), + failed_dry_run_result.get("error"), + ) + + retrieval_result = await self._run_with_timeout( + "Semantic schema retry retrieval", + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=True, + semantic_retry_context=retry_context, + ), + timeout_seconds=self._schema_retrieval_timeout_seconds, + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) + semantic_pipeline_active = True + semantic_contract_available = self._semantic_analysis_has_contract( + schema_intent_analysis + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + has_calculated_field = _retrieval_result.get( + "has_calculated_field", False + ) + has_metric = _retrieval_result.get("has_metric", False) + has_json_field = _retrieval_result.get("has_json_field", False) + semantic_support_error = get_schema_intent_analysis_error( + schema_intent_analysis + ) + logger.info( + "semantic_sql_retry_contract query_id=%s retry_attempt=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s support_error=%s", + query_id, + semantic_retry_attempt, + semantic_pipeline_active, + semantic_contract_available, + self._semantic_schema_objects(schema_intent_analysis), + semantic_support_error, + ) + if ( + not semantic_contract_available + or not documents + or semantic_support_error + ): + invalid_sql = failed_dry_run_result.get( + "sql", invalid_sql + ) + error_message = semantic_support_error or ( + "Semantic schema retrieval did not produce a complete active-schema contract." + ) + logger.info( + "semantic_sql_retry_rejected query_id=%s retry_attempt=%s reason=%s", + query_id, + semantic_retry_attempt, + error_message, + ) + continue + + if sql_generation_histories: + retry_generation_results = await self._run_with_timeout( + "Follow-up semantic SQL retry generation", + self._pipelines["followup_sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=sql_generation_histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), + ) + else: + retry_generation_results = await self._run_with_timeout( + "Semantic SQL retry generation", + self._pipelines["sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_intent_analysis=schema_intent_analysis, + ), + ) + + retry_post_process = retry_generation_results.get( + "post_process", {} + ) + if retry_valid_result := retry_post_process.get( + "valid_generation_result" + ): + if ask_result := self._build_validated_ask_result_from_sql( + retry_valid_result.get("sql"), + table_ddls, + sql_user_query, + ): + logger.info( + "semantic_sql_retry_accepted query_id=%s retry_attempt=%s selected_schema_objects=%s", + query_id, + semantic_retry_attempt, + self._semantic_schema_objects( + schema_intent_analysis + ), + ) + api_results = [ask_result] + failed_dry_run_result = {} + break + invalid_sql = retry_valid_result.get("sql") + error_message = ( + "Semantic SQL retry generated SQL that failed active-schema validation." + ) + failed_dry_run_result = { + "type": "SCHEMA_INTENT_VALIDATION", + "sql": invalid_sql, + "original_sql": invalid_sql, + "error": error_message, + } + else: + failed_dry_run_result = retry_post_process.get( + "invalid_generation_result", {} + ) + invalid_sql = failed_dry_run_result.get( + "sql", invalid_sql + ) + error_message = failed_dry_run_result.get( + "error", error_message + ) + logger.info( + "semantic_sql_retry_failed query_id=%s retry_attempt=%s failure_type=%s error=%s", + query_id, + semantic_retry_attempt, + failed_dry_run_result.get("type"), + error_message, + ) + + if semantic_pipeline_active and not api_results: + current_sql_correction_retries = max_sql_correction_retries + if failed_dry_run_result: + invalid_sql = failed_dry_run_result.get( + "sql", invalid_sql + ) + error_message = failed_dry_run_result.get( + "error", error_message + ) + logger.info( + "semantic_pipeline_exhausted query_id=%s retry_attempts=%s error=%s invalid_sql=%s", + query_id, + semantic_retry_attempt, + error_message, + invalid_sql, + ) + + while ( + failed_dry_run_result + and not api_results + and current_sql_correction_retries < max_sql_correction_retries + ): if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", @@ -6678,6 +7065,7 @@ async def ask( sql_functions=sql_functions, sql_knowledge=sql_knowledge, query=sql_user_query, + schema_intent_analysis=schema_intent_analysis, ), ) @@ -6720,8 +7108,10 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names + if not semantic_pipeline_active and ( + heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ) ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", @@ -6756,16 +7146,31 @@ async def ask( logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - is_followup=True if histories else False, + if semantic_pipeline_active and error_message: + self._ask_results[query_id] = ( + self._build_failed_text_to_sql_response( + trace_id, + error_message, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + is_followup=True if histories else False, + code="NO_RELEVANT_SQL", + ) + ) + else: + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + is_followup=True if histories else False, + ) ) - ) if error_message or invalid_sql: logger.info( "Suppressed technical SQL failure for query_id %s. " @@ -6774,9 +7179,13 @@ async def ask( error_message, invalid_sql, ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_type"] = ( + "NO_RELEVANT_SQL" if semantic_pipeline_active else "NO_RELEVANT_DATA" + ) results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + error_message + if semantic_pipeline_active and error_message + else NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE ) results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 25044de18c..da7cce0d18 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -120,8 +120,11 @@ async def ask_feedback( instructions_task, ) = await asyncio.gather( self._pipelines["db_schema_retrieval"].run( + query=ask_feedback_request.question, tables=ask_feedback_request.tables, project_id=ask_feedback_request.project_id, + histories=[], + enable_column_pruning=True, ), self._pipelines["sql_pairs_retrieval"].run( query=ask_feedback_request.question, @@ -159,6 +162,9 @@ async def ask_feedback( ) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) + schema_intent_analysis = _retrieval_result.get( + "semantic_analysis", {} + ) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] sql_samples = sql_samples_task["formatted_output"].get("documents", []) @@ -186,6 +192,8 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + query=ask_feedback_request.question, + schema_intent_analysis=schema_intent_analysis, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -248,6 +256,8 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + query=ask_feedback_request.question, + schema_intent_analysis=schema_intent_analysis, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index fa46c2ca7a..050509e49f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -7,12 +7,14 @@ find_invalid_table_references, get_json_field_instructions, get_metric_instructions, + get_schema_intent_analysis_error, normalize_data_source, normalize_generation_result_sql, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, get_sql_generation_system_prompt, get_text_to_sql_rules, + validate_sql_intent_alignment, ) @@ -260,6 +262,72 @@ def test_schema_validation_ignores_null_table_metadata(): ) == [] +def test_schema_intent_validation_rejects_generic_count_for_requested_metric(): + semantic_analysis = { + "analytical_intent": "ranking", + "metrics": ["invoice amount"], + "concept_mappings": [ + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["dbo_tblFactSales.invoice_amount"], + "required_in_sql": True, + } + ], + } + + error = validate_sql_intent_alignment( + "Show top customers by invoice amount", + 'SELECT COUNT(*) AS "RecordCount" FROM "dbo_tblFactSales"', + {"dbo_tblFactSales": ["invoice_amount"]}, + semantic_analysis=semantic_analysis, + ) + + assert "generic record count" in error + + +def test_schema_intent_validation_requires_ranking_shape(): + semantic_analysis = { + "analytical_intent": "ranking", + "ranking": ["top 10"], + "concept_mappings": [ + { + "request_concept": "invoice amount", + "concept_type": "metric", + "schema_objects": ["dbo_tblFactSales.invoice_amount"], + "required_in_sql": True, + } + ], + } + + error = validate_sql_intent_alignment( + "Show top 10 customers by invoice amount", + 'SELECT SUM("dbo_tblFactSales"."invoice_amount") AS "invoice_amount" ' + 'FROM "dbo_tblFactSales"', + {"dbo_tblFactSales": ["invoice_amount"]}, + semantic_analysis=semantic_analysis, + ) + + assert "sorting and limiting" in error + + +def test_schema_intent_analysis_reports_missing_required_mapping(): + semantic_analysis = { + "concept_mappings": [ + { + "request_concept": "customer", + "concept_type": "entity", + "schema_objects": [], + "required_in_sql": True, + } + ] + } + + assert "did not map required request concepts" in get_schema_intent_analysis_error( + semantic_analysis + ) + + def test_schema_validation_ignores_null_column_metadata(): assert find_invalid_column_references( 'SELECT "dbo_tblSales"."Market" FROM "dbo_tblSales"', diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 5b6694bb45..b8ac2ff769 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -5,6 +5,7 @@ _is_project_wide_analysis_query, dbschema_retrieval, expand_business_terms_for_retrieval, + rank_semantic_schema_candidates, ) @@ -30,6 +31,71 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query +def test_rank_semantic_schema_candidates_prefers_complete_business_concept_coverage(): + schemas = [ + { + "type": "TABLE", + "name": "dbo_invoices", + "columns": [ + {"name": "invoice_amount", "data_type": "decimal"}, + {"name": "customer_name", "data_type": "varchar"}, + ], + }, + { + "type": "TABLE", + "name": "dbo_invoice_ids", + "columns": [ + {"name": "invoice_id", "data_type": "varchar"}, + ], + }, + ] + + candidates = rank_semantic_schema_candidates( + "Show top customers by invoice amount", + schemas, + ) + + assert candidates[0]["table_name"] == "dbo_invoices" + assert "customer" in candidates[0]["matched_query_terms"] + assert "invoice" in candidates[0]["matched_query_terms"] + assert "amount" in candidates[0]["matched_query_terms"] + + +def test_rank_semantic_schema_candidates_uses_retry_rejections_as_negative_feedback(): + schemas = [ + { + "type": "TABLE", + "name": "dbo_invoices", + "columns": [ + {"name": "invoice_amount", "data_type": "decimal"}, + {"name": "customer_name", "data_type": "varchar"}, + ], + }, + { + "type": "TABLE", + "name": "dbo_invoice_summary", + "columns": [ + {"name": "total_invoice_amount", "data_type": "decimal"}, + {"name": "customer_name", "data_type": "varchar"}, + ], + }, + ] + + candidates = rank_semantic_schema_candidates( + "Show top customers by invoice amount", + schemas, + semantic_retry_context={ + "rejected_schema_objects": [ + "dbo_invoices", + "dbo_invoices.invoice_amount", + ] + }, + ) + + assert candidates[0]["table_name"] == "dbo_invoice_summary" + assert candidates[0]["rejected_by_retry"] is False + + @pytest.mark.asyncio async def test_dbschema_retrieval_loads_complete_active_project_schema(): class Retriever: From d86f10d76273d0385b74b28bd55035184fbc5789 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 11 Jul 2026 00:27:29 +0530 Subject: [PATCH 0490/1087] Revert "Add deterministic semantic SQL compiler" This reverts commit 91dae570f8f52678d7cb8f268d5f80e462696d73. --- .../generation/followup_sql_generation.py | 33 +- .../src/pipelines/generation/semantic_sql.py | 1213 ----------------- .../pipelines/generation/sql_generation.py | 34 +- .../pipelines/generation/test_semantic_sql.py | 156 --- 4 files changed, 2 insertions(+), 1434 deletions(-) delete mode 100644 wren-ai-service/src/pipelines/generation/semantic_sql.py delete mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py 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 ed5f419c5b..eb39e7bd97 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -13,7 +13,6 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata -from src.pipelines.generation.semantic_sql import compile_semantic_sql from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_ask_history_messages, @@ -268,36 +267,6 @@ async def run( logger.info("Follow-Up SQL Generation pipeline is running...") metadata = await retrieve_metadata(project_id or "", self._retriever) - data_source = metadata.get("data_source", "local_file") - - deterministic_result = compile_semantic_sql( - query=query, - documents=contexts, - semantic_analysis=schema_intent_analysis, - data_source=data_source, - ) - if deterministic_result: - logger.info( - "Deterministic semantic SQL compiler produced follow-up SQL for query: %s", - query, - ) - post_process_result = await self._components["post_processor"].run( - [deterministic_result.sql], - project_id=project_id, - use_dry_plan=use_dry_plan, - data_source=data_source, - allow_dry_plan_fallback=allow_dry_plan_fallback, - valid_table_names=construct_valid_table_names(contexts), - valid_table_columns=construct_valid_table_columns(contexts), - query=query, - semantic_analysis=schema_intent_analysis, - ) - if post_process_result.get("valid_generation_result"): - return {"post_process": post_process_result} - logger.info( - "Deterministic follow-up SQL rejected by post processor; falling back to LLM. error=%s", - post_process_result.get("invalid_generation_result", {}).get("error"), - ) return await self._pipe.execute( ["post_process"], @@ -315,7 +284,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": data_source, + "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, "schema_intent_analysis": schema_intent_analysis, **self._components, diff --git a/wren-ai-service/src/pipelines/generation/semantic_sql.py b/wren-ai-service/src/pipelines/generation/semantic_sql.py deleted file mode 100644 index 0677de24f5..0000000000 --- a/wren-ai-service/src/pipelines/generation/semantic_sql.py +++ /dev/null @@ -1,1213 +0,0 @@ -from __future__ import annotations - -import logging -import json -import re -from ast import literal_eval -from collections import deque -from dataclasses import dataclass, field -from datetime import date, datetime, timedelta -from typing import Any, Literal - -import sqlparse - -logger = logging.getLogger("wren-ai-service") - - -Aggregate = Literal["COUNT", "SUM", "AVG", "MIN", "MAX"] -JoinType = Literal["INNER JOIN", "LEFT JOIN"] - - -@dataclass(frozen=True) -class Intent: - question_type: str - chart_requested: bool = False - chart_type: str = "auto" - ranking: bool = False - top_n: int | None = None - bottom_n: int | None = None - distinct: bool = False - aggregation: Aggregate | None = None - needs_sql: bool = True - - -@dataclass(frozen=True) -class ColumnRef: - table: str - column: str - data_type: str = "" - description: str = "" - - @property - def sql(self) -> str: - return f'{quote_identifier(self.table)}.{quote_identifier(self.column)}' - - @property - def object_name(self) -> str: - return f"{self.table}.{self.column}" - - -@dataclass(frozen=True) -class Relationship: - left_table: str - left_column: str - right_table: str - right_column: str - join_type: JoinType = "INNER JOIN" - cardinality: str = "" - source: str = "schema" - - -@dataclass(frozen=True) -class MetricDefinition: - name: str - column: ColumnRef - aggregation: Aggregate - description: str = "" - synonyms: tuple[str, ...] = () - allowed_dimensions: tuple[str, ...] = () - formula: str | None = None - grain: str | None = None - join_requirements: tuple[str, ...] = () - - -@dataclass(frozen=True) -class FilterDefinition: - column: ColumnRef - operator: str - value: Any = None - end_value: Any = None - - -@dataclass(frozen=True) -class SortDefinition: - expression: str - direction: Literal["ASC", "DESC"] = "DESC" - - -@dataclass -class SemanticPlan: - intent: Intent - entities: list[str] = field(default_factory=list) - metrics: list[MetricDefinition] = field(default_factory=list) - aggregation: Aggregate | None = None - filters: list[FilterDefinition] = field(default_factory=list) - group_by: list[ColumnRef] = field(default_factory=list) - sort: list[SortDefinition] = field(default_factory=list) - limit: int | None = None - base_table: str | None = None - joins: list[Relationship] = field(default_factory=list) - chart_type: str = "" - warnings: list[str] = field(default_factory=list) - - @property - def is_complete(self) -> bool: - if not self.base_table: - return False - if self.intent.aggregation and not self.metrics and self.intent.aggregation != "COUNT": - return False - required_tables = { - ref.table - for ref in [ - *[metric.column for metric in self.metrics], - *self.group_by, - *[filter_.column for filter_ in self.filters], - ] - } - connected_tables = {self.base_table} - for join in self.joins: - connected_tables.add(join.left_table) - connected_tables.add(join.right_table) - return required_tables.issubset(connected_tables) - - -@dataclass -class SQLValidationResult: - valid: bool - errors: list[str] = field(default_factory=list) - - -@dataclass -class CompileResult: - sql: str - plan: SemanticPlan - validation: SQLValidationResult - - -@dataclass -class SchemaCatalog: - tables: dict[str, list[ColumnRef]] = field(default_factory=dict) - relationships: list[Relationship] = field(default_factory=list) - - def columns(self) -> list[ColumnRef]: - return [column for columns in self.tables.values() for column in columns] - - def table_for_column(self, column: ColumnRef) -> str: - return column.table - - def get_column(self, table: str, column: str) -> ColumnRef | None: - for candidate in self.tables.get(table, []): - if candidate.column.lower() == column.lower(): - return candidate - return None - - -_STOPWORDS = { - "a", - "an", - "and", - "as", - "at", - "by", - "chart", - "for", - "from", - "give", - "graph", - "in", - "last", - "me", - "of", - "on", - "per", - "show", - "the", - "this", - "to", - "with", -} -_GENERIC_COLUMN_TOKENS = { - "amount", - "at", - "code", - "date", - "day", - "id", - "key", - "month", - "name", - "no", - "number", - "time", - "total", - "type", - "value", - "year", -} -_SYNONYMS = { - "acct": {"account", "customer", "client"}, - "account": {"acct", "customer", "client"}, - "amount": {"amt", "value", "total"}, - "avg": {"average", "mean"}, - "bill": {"invoice"}, - "billing": {"invoice"}, - "client": {"customer", "account"}, - "cost": {"expense", "spend"}, - "cust": {"customer", "client", "account"}, - "customer": {"cust", "client", "account"}, - "gmv": {"revenue", "sales", "amount"}, - "invoice": {"inv", "bill", "billing"}, - "profit": {"margin", "income", "earnings"}, - "qty": {"quantity", "units"}, - "quantity": {"qty", "units"}, - "revenue": {"sales", "amount", "value", "gmv"}, - "sale": {"sales", "revenue", "amount"}, - "sales": {"sale", "revenue", "amount", "gmv"}, - "total": {"sum", "amount", "value"}, - "value": {"amount", "total"}, -} -_NUMERIC_TYPES = { - "bigint", - "decimal", - "double", - "float", - "int", - "integer", - "numeric", - "real", - "smallint", -} -_TEMPORAL_TYPES = {"date", "datetime", "timestamp", "time"} -_TEXT_TYPES = {"char", "string", "text", "varchar"} - - -def quote_identifier(identifier: str) -> str: - return f'"{str(identifier).replace(chr(34), chr(34) + chr(34))}"' - - -def tokenize(value: Any) -> set[str]: - text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value or "")) - raw_tokens = { - token.lower() - for token in re.findall(r"[A-Za-z0-9]+", text) - if len(token) > 1 and token.lower() not in _STOPWORDS - } - tokens = set(raw_tokens) - for token in list(raw_tokens): - if len(token) > 4 and token.endswith("ies"): - tokens.add(f"{token[:-3]}y") - elif len(token) > 3 and token.endswith("s"): - tokens.add(token[:-1]) - for token in list(tokens): - tokens.update(_SYNONYMS.get(token, set())) - return tokens - - -def compact(value: Any) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - -def is_numeric_type(data_type: str) -> bool: - normalized = data_type.lower() - return any(type_name in normalized for type_name in _NUMERIC_TYPES) - - -def is_temporal_type(data_type: str) -> bool: - normalized = data_type.lower() - return any(type_name in normalized for type_name in _TEMPORAL_TYPES) - - -def is_text_type(data_type: str) -> bool: - normalized = data_type.lower() - return any(type_name in normalized for type_name in _TEXT_TYPES) - - -class IntentDetector: - def detect(self, query: str) -> Intent: - normalized = query.lower() - top_match = re.search(r"\btop\s+(\d+)\b", normalized) - bottom_match = re.search(r"\bbottom\s+(\d+)\b", normalized) - aggregation = self._detect_aggregation(normalized) - chart_type = self._detect_chart_type(normalized) - question_type = "retrieval" - if aggregation: - question_type = "aggregation" - if re.search(r"\btrend|over time|monthly|weekly|daily|yearly\b", normalized): - question_type = "trend" - if top_match or bottom_match or re.search( - r"\b(highest|lowest|largest|smallest|rank|ranking)\b", normalized - ): - question_type = "ranking" - if "dashboard" in normalized or "kpi" in normalized: - question_type = "dashboard" - - return Intent( - question_type=question_type, - chart_requested=bool(chart_type), - chart_type=chart_type or "auto", - ranking=question_type == "ranking", - top_n=( - int(top_match.group(1)) - if top_match - else 10 - if re.search(r"\btop\b", normalized) - else None - ), - bottom_n=int(bottom_match.group(1)) if bottom_match else None, - distinct=bool(re.search(r"\bdistinct|unique\b", normalized)), - aggregation=aggregation, - ) - - def _detect_aggregation(self, normalized_query: str) -> Aggregate | None: - if re.search(r"\b(avg|average|mean)\b", normalized_query): - return "AVG" - if re.search(r"\b(count|number of|how many)\b", normalized_query): - return "COUNT" - if re.search(r"\b(min|minimum)\b", normalized_query): - return "MIN" - if re.search(r"\b(max|maximum)\b", normalized_query): - return "MAX" - if re.search(r"\b(sum|total)\b", normalized_query): - return "SUM" - return None - - def _detect_chart_type(self, normalized_query: str) -> str: - checks = ( - ("line", ("line chart", "line graph", "trend")), - ("pie", ("pie chart", "donut chart", "part to whole")), - ("scatter", ("scatter", "correlation")), - ("bar", ("bar chart", "bar graph", "column chart")), - ("card", ("kpi", "single kpi")), - ) - for chart_type, terms in checks: - if any(term in normalized_query for term in terms): - return chart_type - if re.search(r"\b(chart|graph|plot|visuali[sz]e)\b", normalized_query): - return "auto" - return "" - - -class SchemaParser: - def parse(self, documents: list[str]) -> SchemaCatalog: - catalog = SchemaCatalog() - for document in documents or []: - text = str(document) - self._parse_create_table(text, catalog) - self._parse_metadata_document(text, catalog) - return catalog - - def _parse_create_table(self, ddl: str, catalog: SchemaCatalog) -> None: - table_match = re.search( - r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|`(?P[^`]+)`|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", - ddl, - flags=re.IGNORECASE, - ) - if not table_match: - return - table_name = next(value for value in table_match.groupdict().values() if value) - body = self._table_body(ddl, table_match.end()) - columns: list[ColumnRef] = [] - for definition in self._split_definitions(body): - stripped = definition.strip().rstrip(",") - if not stripped: - continue - fk = re.search( - r"FOREIGN\s+KEY\s*\((?P[^\)]+)\)\s+REFERENCES\s+" - r'(?:"(?P[^"]+)"|`(?P[^`]+)`|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\((?P[^\)]+)\)", - stripped, - flags=re.IGNORECASE, - ) - if fk: - right_table = next( - value - for value in ( - fk.group("rq"), - fk.group("rb"), - fk.group("rs"), - fk.group("rbare"), - ) - if value - ) - catalog.relationships.append( - Relationship( - left_table=table_name, - left_column=self._clean_identifier(fk.group("left")), - right_table=right_table, - right_column=self._clean_identifier(fk.group("right")), - ) - ) - continue - if re.match( - r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY|INDEX)\b", - stripped, - flags=re.IGNORECASE, - ): - continue - column_match = re.match( - r'(?:"(?P[^"]+)"|`(?P[^`]+)`|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_$]*))" - r"\s+(?P[A-Za-z_][A-Za-z0-9_]*(?:\([^\)]*\))?)", - stripped, - ) - if column_match: - column_name = next( - value - for key, value in column_match.groupdict().items() - if key != "type" and value - ) - columns.append( - ColumnRef( - table=table_name, - column=column_name, - data_type=column_match.group("type") or "", - ) - ) - if columns: - catalog.tables[table_name] = columns - - def _parse_metadata_document(self, text: str, catalog: SchemaCatalog) -> None: - metadata = self._literal_document(text) - if not isinstance(metadata, dict): - return - for model in [ - *(metadata.get("models") or []), - *(metadata.get("views") or []), - ]: - if not isinstance(model, dict) or not model.get("name"): - continue - table_name = str(model["name"]) - columns: list[ColumnRef] = [] - for column in [ - *(model.get("columns") or []), - *(model.get("calculatedFields") or []), - ]: - if not isinstance(column, dict) or not column.get("name"): - continue - columns.append( - ColumnRef( - table=table_name, - column=str(column["name"]), - data_type=str( - column.get("data_type") - or column.get("type") - or column.get("dataType") - or "" - ), - description=str(column.get("comment") or column.get("description") or ""), - ) - ) - if columns: - catalog.tables[table_name] = columns - reference_name = model.get("referenceName") - if reference_name and columns: - catalog.tables[str(reference_name)] = [ - ColumnRef( - table=str(reference_name), - column=column.column, - data_type=column.data_type, - description=column.description, - ) - for column in columns - ] - for relationship in metadata.get("relationships") or []: - parsed = self._relationship_from_metadata(relationship) - if parsed: - catalog.relationships.append(parsed) - - def _literal_document(self, text: str) -> Any: - stripped = text.strip() - if not stripped.startswith("{"): - return None - try: - return json.loads(stripped) - except json.JSONDecodeError: - pass - try: - return literal_eval(stripped) - except (SyntaxError, ValueError): - return None - - def _relationship_from_metadata(self, relationship: Any) -> Relationship | None: - if not isinstance(relationship, dict): - return None - condition = str(relationship.get("condition") or "") - match = re.search( - r'(?P[A-Za-z_][A-Za-z0-9_.$]*)\.(?P[A-Za-z_][A-Za-z0-9_$]*)\s*=\s*' - r'(?P[A-Za-z_][A-Za-z0-9_.$]*)\.(?P[A-Za-z_][A-Za-z0-9_$]*)', - condition, - ) - if not match: - return None - return Relationship( - left_table=match.group("left_table"), - left_column=match.group("left_column"), - right_table=match.group("right_table"), - right_column=match.group("right_column"), - cardinality=str(relationship.get("joinType") or relationship.get("type") or ""), - source="semantic_metadata", - ) - - def _table_body(self, ddl: str, start: int) -> str: - depth = 1 - cursor = start - while cursor < len(ddl) and depth > 0: - if ddl[cursor] == "(": - depth += 1 - elif ddl[cursor] == ")": - depth -= 1 - cursor += 1 - return ddl[start : cursor - 1] - - def _split_definitions(self, body: str) -> list[str]: - definitions: list[str] = [] - depth = 0 - start = 0 - for index, char in enumerate(body): - if char == "(": - depth += 1 - elif char == ")": - depth -= 1 - elif char == "," and depth == 0: - definitions.append(body[start:index]) - start = index + 1 - definitions.append(body[start:]) - return definitions - - def _clean_identifier(self, value: str) -> str: - return str(value or "").strip().strip('"`[]') - - -class MetricRegistry: - def __init__(self, catalog: SchemaCatalog, semantic_analysis: dict[str, Any] | None): - self.catalog = catalog - self.semantic_analysis = semantic_analysis or {} - self.metrics = self._build_metrics() - - def resolve_metric(self, query: str, intent: Intent) -> MetricDefinition | None: - mappings = self.semantic_analysis.get("concept_mappings") or [] - for mapping in mappings: - if not isinstance(mapping, dict): - continue - if str(mapping.get("concept_type", "")).lower() != "metric": - continue - for schema_object in mapping.get("schema_objects") or []: - column = self._column_from_schema_object(str(schema_object)) - if column: - return MetricDefinition( - name=str(mapping.get("request_concept") or column.column), - column=column, - aggregation=intent.aggregation or self._default_aggregation(column), - synonyms=tuple(tokenize(mapping.get("request_concept"))), - ) - - query_tokens = tokenize(query) - scored = [ - (self._score_metric(metric, query_tokens), metric) - for metric in self.metrics - ] - scored = [(score, metric) for score, metric in scored if score > 0] - if not scored: - return None - scored.sort(key=lambda item: item[0], reverse=True) - metric = scored[0][1] - if intent.aggregation: - return MetricDefinition( - name=metric.name, - column=metric.column, - aggregation=intent.aggregation, - description=metric.description, - synonyms=metric.synonyms, - allowed_dimensions=metric.allowed_dimensions, - formula=metric.formula, - grain=metric.grain, - join_requirements=metric.join_requirements, - ) - return metric - - def _build_metrics(self) -> list[MetricDefinition]: - metrics: list[MetricDefinition] = [] - for column in self.catalog.columns(): - if not is_numeric_type(column.data_type): - continue - tokens = tokenize(f"{column.table} {column.column}") - if not tokens.intersection( - { - "amount", - "balance", - "cost", - "gmv", - "margin", - "price", - "profit", - "quantity", - "rate", - "revenue", - "sale", - "sales", - "score", - "total", - "value", - } - ): - continue - metrics.append( - MetricDefinition( - name=humanize(column.column), - column=column, - aggregation=self._default_aggregation(column), - synonyms=tuple(tokens), - ) - ) - return metrics - - def _column_from_schema_object(self, schema_object: str) -> ColumnRef | None: - parts = [part.strip().strip('"`[]') for part in schema_object.split(".")] - if len(parts) < 2: - return None - table = ".".join(parts[:-1]) - column = parts[-1] - direct = self.catalog.get_column(table, column) - if direct: - return direct - for table_name in self.catalog.tables: - if table_name.lower().endswith(table.lower()): - candidate = self.catalog.get_column(table_name, column) - if candidate: - return candidate - return None - - def _default_aggregation(self, column: ColumnRef) -> Aggregate: - tokens = tokenize(column.column) - if tokens.intersection({"avg", "average", "mean", "rate", "ratio", "percent"}): - return "AVG" - return "SUM" - - def _score_metric(self, metric: MetricDefinition, query_tokens: set[str]) -> int: - metric_tokens = set(metric.synonyms) | tokenize(metric.name) - score = len(metric_tokens.intersection(query_tokens)) * 10 - compact_query = compact(" ".join(query_tokens)) - compact_metric = compact(metric.name) - if compact_metric and compact_metric in compact_query: - score += 40 - return score - - -class EntityResolver: - def __init__(self, catalog: SchemaCatalog, semantic_analysis: dict[str, Any] | None): - self.catalog = catalog - self.semantic_analysis = semantic_analysis or {} - self._metric_registry = MetricRegistry(catalog, {}) - - def resolve_dimensions(self, query: str, metric: MetricDefinition | None) -> list[ColumnRef]: - mapped = self._dimensions_from_semantic_analysis() - if mapped: - return mapped - requested_terms = self._requested_grouping_terms(query) - if not requested_terms and metric: - return [] - if not requested_terms: - requested_terms = tokenize(query) - scored: list[tuple[int, ColumnRef]] = [] - for column in self.catalog.columns(): - if metric and column == metric.column: - continue - if is_numeric_type(column.data_type): - if self._looks_identifier(column) and not requested_terms.intersection( - {"id", "key", "number", "no"} - ): - continue - if not self._looks_identifier(column): - continue - column_tokens = tokenize(f"{column.table} {column.column}") - score = len(column_tokens.intersection(requested_terms)) * 10 - if is_text_type(column.data_type): - score += 5 - if compact(column.column) in compact(" ".join(requested_terms)): - score += 20 - if requested_terms and score: - scored.append((score, column)) - scored.sort(key=lambda item: item[0], reverse=True) - if not scored: - return [] - best_score = scored[0][0] - return [column for score, column in scored[:3] if score == best_score] - - def resolve_temporal_column(self, query: str, preferred_table: str | None) -> ColumnRef | None: - if not self._requests_date_filter(query): - return None - candidates = [ - column - for column in self.catalog.columns() - if is_temporal_type(column.data_type) - or tokenize(column.column).intersection( - {"date", "time", "created", "updated", "month", "year"} - ) - ] - if preferred_table: - candidates.sort(key=lambda column: column.table != preferred_table) - return candidates[0] if candidates else None - - def _dimensions_from_semantic_analysis(self) -> list[ColumnRef]: - dimensions: list[ColumnRef] = [] - mappings = self.semantic_analysis.get("concept_mappings") or [] - for mapping in mappings: - if not isinstance(mapping, dict): - continue - if str(mapping.get("concept_type", "")).lower() not in { - "dimension", - "entity", - "identifier", - }: - continue - for schema_object in mapping.get("schema_objects") or []: - column = self._metric_registry._column_from_schema_object(str(schema_object)) - if column and column not in dimensions: - dimensions.append(column) - return dimensions - - def _requested_grouping_terms(self, query: str) -> set[str]: - normalized = query.lower() - terms: set[str] = set() - for match in re.finditer( - r"\b(?:by|per|for each|group(?:ed)? by)\s+([A-Za-z0-9_ ]+)", - normalized, - ): - phrase = re.split( - r"\b(?:and|with|where|order|sort|top|bottom|last|this|limit)\b", - match.group(1), - maxsplit=1, - )[0] - terms.update(tokenize(phrase)) - ranking_entity = re.search( - r"\b(?:top|bottom)\s+(?:\d+\s+)?([A-Za-z0-9_ ]+?)\s+by\b", - normalized, - ) - if ranking_entity: - terms.update(tokenize(ranking_entity.group(1))) - return terms - - def _requests_date_filter(self, query: str) -> bool: - return bool( - re.search( - r"\b(today|yesterday|this|last|rolling|past|previous)\s+" - r"(?:\d+\s+)?(?:day|week|month|quarter|year)s?\b", - query.lower(), - ) - ) - - def _looks_identifier(self, column: ColumnRef) -> bool: - return bool(tokenize(column.column).intersection({"id", "key", "number", "no"})) - - -class RelationshipGraph: - def __init__(self, catalog: SchemaCatalog): - self.catalog = catalog - - def join_path(self, required_tables: set[str], base_table: str) -> list[Relationship] | None: - joins: list[Relationship] = [] - connected = {base_table} - for table in sorted(required_tables - connected): - path = self._shortest_path(connected, table) - if not path: - return None - joins.extend(path) - for relationship in path: - connected.add(relationship.left_table) - connected.add(relationship.right_table) - return self._dedupe(joins) - - def _shortest_path(self, sources: set[str], target: str) -> list[Relationship] | None: - queue: deque[tuple[str, list[Relationship]]] = deque( - (source, []) for source in sources - ) - seen = set(sources) - while queue: - table, path = queue.popleft() - if table == target: - return path - for relationship in self._neighbors(table): - next_table = ( - relationship.right_table - if relationship.left_table == table - else relationship.left_table - ) - if next_table in seen: - continue - seen.add(next_table) - queue.append((next_table, [*path, relationship])) - return None - - def _neighbors(self, table: str) -> list[Relationship]: - return [ - relationship - for relationship in self.catalog.relationships - if relationship.left_table == table or relationship.right_table == table - ] - - def _dedupe(self, relationships: list[Relationship]) -> list[Relationship]: - deduped: list[Relationship] = [] - seen: set[tuple[str, str, str, str]] = set() - for relationship in relationships: - key = ( - relationship.left_table, - relationship.left_column, - relationship.right_table, - relationship.right_column, - ) - if key in seen: - continue - seen.add(key) - deduped.append(relationship) - return deduped - - -class SemanticPlanner: - def __init__(self, catalog: SchemaCatalog, semantic_analysis: dict[str, Any] | None): - self.catalog = catalog - self.semantic_analysis = semantic_analysis or {} - self.metric_registry = MetricRegistry(catalog, semantic_analysis) - self.entity_resolver = EntityResolver(catalog, semantic_analysis) - self.relationship_graph = RelationshipGraph(catalog) - - def build_plan(self, query: str, now: datetime | None = None) -> SemanticPlan | None: - intent = IntentDetector().detect(query) - metric = self.metric_registry.resolve_metric(query, intent) - aggregation = intent.aggregation or (metric.aggregation if metric else None) - dimensions = self.entity_resolver.resolve_dimensions(query, metric) - if aggregation == "COUNT" and not self._has_explicit_grouping(query): - dimensions = [] - temporal_column = self.entity_resolver.resolve_temporal_column( - query, - metric.column.table if metric else (dimensions[0].table if dimensions else None), - ) - filters = [] - if temporal_column: - date_filter = DateResolver(now or datetime.now()).resolve(query, temporal_column) - if date_filter: - filters.append(date_filter) - - required_refs = [ - *([metric.column] if metric else []), - *dimensions, - *[filter_.column for filter_ in filters], - ] - base_table = self._choose_base_table(required_refs) - if not base_table and aggregation == "COUNT": - base_table = self._resolve_table_from_query(query) - if not base_table: - return None - required_tables = {ref.table for ref in required_refs} - joins = self.relationship_graph.join_path(required_tables, base_table) - if joins is None: - return None - - metric_expression = self._metric_expression(metric, aggregation) - sort = [] - if intent.ranking and metric_expression: - sort.append(SortDefinition(metric_expression, "ASC" if intent.bottom_n else "DESC")) - limit = intent.top_n or intent.bottom_n - chart_type = ChartRuleEngine().select_chart(intent, dimensions, [metric] if metric else []) - - plan = SemanticPlan( - intent=intent, - entities=[humanize(dimension.column) for dimension in dimensions], - metrics=[metric] if metric else [], - aggregation=aggregation, - filters=filters, - group_by=dimensions, - sort=sort, - limit=limit, - base_table=base_table, - joins=joins, - chart_type=chart_type, - ) - return plan if plan.is_complete else None - - def _choose_base_table(self, refs: list[ColumnRef]) -> str | None: - if not refs: - return None - table_counts: dict[str, int] = {} - for ref in refs: - table_counts[ref.table] = table_counts.get(ref.table, 0) + 1 - return sorted(table_counts.items(), key=lambda item: item[1], reverse=True)[0][0] - - def _resolve_table_from_query(self, query: str) -> str | None: - query_tokens = tokenize(query) - scored: list[tuple[int, str]] = [] - for table in self.catalog.tables: - table_tokens = tokenize(table) - score = len(table_tokens.intersection(query_tokens)) * 10 - if compact(table) in compact(query): - score += 30 - if score: - scored.append((score, table)) - if not scored and len(self.catalog.tables) == 1: - return next(iter(self.catalog.tables)) - scored.sort(key=lambda item: item[0], reverse=True) - return scored[0][1] if scored else None - - def _has_explicit_grouping(self, query: str) -> bool: - return bool( - re.search( - r"\b(?:by|per|for each|group(?:ed)? by|top|bottom|rank|ranking)\b", - query.lower(), - ) - ) - - def _metric_expression( - self, - metric: MetricDefinition | None, - aggregation: Aggregate | None, - ) -> str: - if aggregation == "COUNT" and not metric: - return "COUNT(*)" - if not metric or not aggregation: - return "" - return f"{aggregation}({metric.column.sql})" - - -class DateResolver: - def __init__(self, now: datetime): - self.today = now.date() - - def resolve(self, query: str, column: ColumnRef) -> FilterDefinition | None: - normalized = query.lower() - start: date | None = None - end: date | None = None - if "today" in normalized: - start = self.today - end = self.today + timedelta(days=1) - elif "yesterday" in normalized: - start = self.today - timedelta(days=1) - end = self.today - elif "last month" in normalized: - current_month = self.today.replace(day=1) - end = current_month - start = add_months(current_month, -1) - elif "this month" in normalized: - start = self.today.replace(day=1) - end = add_months(start, 1) - elif "last year" in normalized: - start = date(self.today.year - 1, 1, 1) - end = date(self.today.year, 1, 1) - elif "this year" in normalized: - start = date(self.today.year, 1, 1) - end = date(self.today.year + 1, 1, 1) - elif "last week" in normalized: - this_week = self.today - timedelta(days=self.today.weekday()) - start = this_week - timedelta(days=7) - end = this_week - elif "this week" in normalized: - start = self.today - timedelta(days=self.today.weekday()) - end = start + timedelta(days=7) - elif "last quarter" in normalized or "this quarter" in normalized: - quarter_month = ((self.today.month - 1) // 3) * 3 + 1 - this_quarter = date(self.today.year, quarter_month, 1) - if "last quarter" in normalized: - end = this_quarter - start = add_months(this_quarter, -3) - else: - start = this_quarter - end = add_months(this_quarter, 3) - else: - rolling_match = re.search( - r"\b(?:last|past|rolling)\s+(\d+)\s+(day|week|month|year)s?\b", - normalized, - ) - if rolling_match: - amount = int(rolling_match.group(1)) - unit = rolling_match.group(2) - end = self.today + timedelta(days=1) - if unit == "day": - start = self.today - timedelta(days=amount) - elif unit == "week": - start = self.today - timedelta(days=amount * 7) - elif unit == "month": - start = add_months(self.today, -amount) - elif unit == "year": - start = add_months(self.today, -amount * 12) - if not start or not end: - return None - return FilterDefinition( - column=column, - operator="BETWEEN_CLOSED_OPEN", - value=start.isoformat(), - end_value=end.isoformat(), - ) - - -class SQLCompiler: - def compile(self, plan: SemanticPlan, data_source: str = "") -> str: - if not plan.base_table: - raise ValueError("Semantic plan has no base table") - select_items = self._select_items(plan) - from_clause = f"FROM {quote_identifier(plan.base_table)}" - join_clause = self._join_clause(plan) - where_clause = self._where_clause(plan) - group_clause = self._group_clause(plan) - order_clause = self._order_clause(plan) - limit_clause = self._limit_clause(plan, data_source) - top_clause = "" - if normalize_data_source(data_source) == "MSSQL" and plan.limit: - top_clause = f" TOP {plan.limit}" - limit_clause = "" - return " ".join( - part - for part in ( - f"SELECT{top_clause} {', '.join(select_items)}", - from_clause, - join_clause, - where_clause, - group_clause, - order_clause, - limit_clause, - ) - if part - ) - - def _select_items(self, plan: SemanticPlan) -> list[str]: - items = [ - f"{dimension.sql} AS {quote_identifier(safe_alias(dimension.column))}" - for dimension in plan.group_by - ] - if plan.metrics: - for metric in plan.metrics: - aggregation = plan.aggregation or metric.aggregation - alias = safe_alias(f"{aggregation.lower()}_{metric.column.column}") - items.append(f"{aggregation}({metric.column.sql}) AS {quote_identifier(alias)}") - elif plan.aggregation == "COUNT": - items.append('COUNT(*) AS "count"') - if plan.intent.distinct and not plan.aggregation and not plan.metrics: - return [ - f"DISTINCT {dimension.sql} AS {quote_identifier(safe_alias(dimension.column))}" - for dimension in plan.group_by - ] - return items - - def _join_clause(self, plan: SemanticPlan) -> str: - clauses = [] - joined = {plan.base_table} - for relationship in plan.joins: - if relationship.left_table in joined: - join_table = relationship.right_table - on_left = f"{quote_identifier(relationship.left_table)}.{quote_identifier(relationship.left_column)}" - on_right = f"{quote_identifier(relationship.right_table)}.{quote_identifier(relationship.right_column)}" - else: - join_table = relationship.left_table - on_left = f"{quote_identifier(relationship.left_table)}.{quote_identifier(relationship.left_column)}" - on_right = f"{quote_identifier(relationship.right_table)}.{quote_identifier(relationship.right_column)}" - clauses.append( - f"{relationship.join_type} {quote_identifier(join_table)} ON {on_left} = {on_right}" - ) - joined.add(join_table) - return " ".join(clauses) - - def _where_clause(self, plan: SemanticPlan) -> str: - conditions = [] - for filter_ in plan.filters: - if filter_.operator == "BETWEEN_CLOSED_OPEN": - conditions.append( - f"{filter_.column.sql} >= '{filter_.value}' AND {filter_.column.sql} < '{filter_.end_value}'" - ) - return f"WHERE {' AND '.join(conditions)}" if conditions else "" - - def _group_clause(self, plan: SemanticPlan) -> str: - if not plan.group_by or not plan.aggregation: - return "" - return "GROUP BY " + ", ".join(dimension.sql for dimension in plan.group_by) - - def _order_clause(self, plan: SemanticPlan) -> str: - if not plan.sort: - return "" - return "ORDER BY " + ", ".join( - f"{sort.expression} {sort.direction}" for sort in plan.sort - ) - - def _limit_clause(self, plan: SemanticPlan, data_source: str) -> str: - if not plan.limit or normalize_data_source(data_source) == "MSSQL": - return "" - return f"LIMIT {plan.limit}" - - -class SQLAstValidator: - def validate(self, sql: str, plan: SemanticPlan) -> SQLValidationResult: - errors: list[str] = [] - parsed = sqlparse.parse(sql) - if len(parsed) != 1: - errors.append("SQL must contain exactly one statement") - if not sqlparse.tokens.DML: - errors.append("SQL parser unavailable") - if not re.match(r"^\s*SELECT\b", sql, flags=re.IGNORECASE): - errors.append("SQL must be a SELECT statement") - if "*" in sql and plan.metrics: - errors.append("Metric queries must not use SELECT *") - if plan.joins and not re.search(r"\b(?:ON|USING)\b", sql, flags=re.IGNORECASE): - errors.append("Joins must include ON or USING") - if plan.aggregation and plan.group_by: - group_text = self._clause(sql, "GROUP BY", ["HAVING", "ORDER BY", "LIMIT", "FETCH"]) - for dimension in plan.group_by: - if quote_identifier(dimension.column) not in group_text: - errors.append(f"Missing GROUP BY column: {dimension.object_name}") - if plan.limit and not re.search( - r"\b(?:LIMIT\s+\d+|TOP\s+\d+|FETCH\s+FIRST\s+\d+)\b", - sql, - flags=re.IGNORECASE, - ): - errors.append("Missing limit/TOP for ranked plan") - if plan.filters and not re.search(r"\bWHERE\b", sql, flags=re.IGNORECASE): - errors.append("Missing WHERE for filtered plan") - return SQLValidationResult(valid=not errors, errors=errors) - - def _clause(self, sql: str, clause: str, terminators: list[str]) -> str: - terminator_pattern = "|".join(rf"\b{terminator}\b" for terminator in terminators) - match = re.search( - rf"\b{clause}\b(?P.*?)(?={terminator_pattern}|$)", - sql, - flags=re.IGNORECASE | re.DOTALL, - ) - return match.group("body") if match else "" - - -class ExecutionValidator: - def validate_result_shape( - self, - plan: SemanticPlan, - rows: list[dict[str, Any]], - ) -> SQLValidationResult: - if not rows: - return SQLValidationResult(valid=True) - expected_columns = {safe_alias(dimension.column) for dimension in plan.group_by} - for metric in plan.metrics: - aggregation = plan.aggregation or metric.aggregation - expected_columns.add(safe_alias(f"{aggregation.lower()}_{metric.column.column}")) - actual_columns = set(rows[0].keys()) - missing = expected_columns - actual_columns - errors = [f"Missing result column: {column}" for column in sorted(missing)] - return SQLValidationResult(valid=not errors, errors=errors) - - -class ChartRuleEngine: - def select_chart( - self, - intent: Intent, - dimensions: list[ColumnRef], - metrics: list[MetricDefinition], - ) -> str: - if intent.chart_type and intent.chart_type != "auto": - return intent.chart_type - if intent.ranking: - return "bar" - if len(metrics) >= 2: - return "scatter" - if len(metrics) == 1 and not dimensions: - return "card" - if any(is_temporal_type(dimension.data_type) for dimension in dimensions): - return "line" - if "percent" in " ".join(tokenize(" ".join(d.column for d in dimensions))): - return "pie" - if dimensions and metrics: - return "bar" - return "" - - -def add_months(value: date, months: int) -> date: - month_index = value.month - 1 + months - year = value.year + month_index // 12 - month = month_index % 12 + 1 - day = min(value.day, month_days(year, month)) - return date(year, month, day) - - -def month_days(year: int, month: int) -> int: - if month == 2: - return 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28 - return [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31][month - 1] - - -def humanize(value: str) -> str: - return re.sub(r"[_\s]+", " ", str(value or "")).strip().title() - - -def safe_alias(value: str) -> str: - alias = re.sub(r"[^A-Za-z0-9_]+", "_", str(value or "").strip()).strip("_") - return alias.lower() or "value" - - -def normalize_data_source(data_source: str | None) -> str: - normalized = (data_source or "").strip().upper().replace("-", "_").replace(" ", "_") - if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: - return "MSSQL" - return normalized - - -def compile_semantic_sql( - query: str, - documents: list[str], - semantic_analysis: dict[str, Any] | None = None, - data_source: str = "", - now: datetime | None = None, -) -> CompileResult | None: - catalog = SchemaParser().parse(documents) - if not catalog.tables: - return None - plan = SemanticPlanner(catalog, semantic_analysis).build_plan(query, now=now) - if not plan: - return None - sql = SQLCompiler().compile(plan, data_source=data_source) - validation = SQLAstValidator().validate(sql, plan) - if not validation.valid: - logger.info("deterministic_semantic_sql_validation_failed errors=%s", validation.errors) - return None - return CompileResult(sql=sql, plan=plan, validation=validation) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 56d2d0d399..af801b6594 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -11,7 +11,6 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata -from src.pipelines.generation.semantic_sql import compile_semantic_sql from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, @@ -269,37 +268,6 @@ async def run( logger.info("SQL Generation pipeline is running...") metadata = await retrieve_metadata(project_id or "", self._retriever) - data_source = metadata.get("data_source", "local_file") - - deterministic_result = compile_semantic_sql( - query=query, - documents=contexts, - semantic_analysis=schema_intent_analysis, - data_source=data_source, - ) - if deterministic_result: - logger.info( - "Deterministic semantic SQL compiler produced SQL for query: %s", - query, - ) - post_process_result = await self._components["post_processor"].run( - [deterministic_result.sql], - project_id=project_id, - use_dry_plan=use_dry_plan, - data_source=data_source, - allow_dry_plan_fallback=allow_dry_plan_fallback, - allow_data_preview=allow_data_preview, - valid_table_names=construct_valid_table_names(contexts), - valid_table_columns=construct_valid_table_columns(contexts), - query=query, - semantic_analysis=schema_intent_analysis, - ) - if post_process_result.get("valid_generation_result"): - return {"post_process": post_process_result} - logger.info( - "Deterministic semantic SQL rejected by post processor; falling back to LLM. error=%s", - post_process_result.get("invalid_generation_result", {}).get("error"), - ) return await self._pipe.execute( ["post_process"], @@ -316,7 +284,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": data_source, + "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, "schema_intent_analysis": schema_intent_analysis, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py b/wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py deleted file mode 100644 index ecbe66191a..0000000000 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_semantic_sql.py +++ /dev/null @@ -1,156 +0,0 @@ -from datetime import datetime - -from src.pipelines.generation.semantic_sql import ( - IntentDetector, - SchemaParser, - compile_semantic_sql, -) - - -def test_intent_detector_returns_structured_business_intent(): - intent = IntentDetector().detect( - "Show top 10 customers by total invoice amount last month as a bar chart" - ) - - assert intent.question_type == "ranking" - assert intent.chart_requested is True - assert intent.chart_type == "bar" - assert intent.ranking is True - assert intent.top_n == 10 - assert intent.aggregation == "SUM" - - -def test_compile_semantic_sql_generates_ranked_aggregate_with_date_filter(): - documents = [ - """ - CREATE TABLE invoices ( - id INTEGER, - customer_id INTEGER, - customer_name VARCHAR, - invoice_amount DECIMAL(10,2), - invoice_date DATE - ); - """ - ] - - result = compile_semantic_sql( - "Show top 10 customers by total invoice amount last month as a bar chart", - documents, - now=datetime(2026, 7, 10), - ) - - assert result is not None - assert result.plan.intent.question_type == "ranking" - assert result.plan.chart_type == "bar" - assert result.plan.metrics[0].column.object_name == "invoices.invoice_amount" - assert [dimension.object_name for dimension in result.plan.group_by] == [ - "invoices.customer_name" - ] - assert 'SUM("invoices"."invoice_amount")' in result.sql - assert 'GROUP BY "invoices"."customer_name"' in result.sql - assert "\"invoices\".\"invoice_date\" >= '2026-06-01'" in result.sql - assert "\"invoices\".\"invoice_date\" < '2026-07-01'" in result.sql - assert "ORDER BY SUM(" in result.sql - assert "LIMIT 10" in result.sql - - -def test_compile_semantic_sql_resolves_join_path_from_foreign_keys(): - documents = [ - """ - CREATE TABLE orders ( - id INTEGER, - customer_id INTEGER, - order_amount DECIMAL(10,2), - FOREIGN KEY (customer_id) REFERENCES customers(id) - ); - """, - """ - CREATE TABLE customers ( - id INTEGER, - customer_name VARCHAR - ); - """, - ] - - result = compile_semantic_sql( - "Show total order amount by customer name", - documents, - ) - - assert result is not None - assert [join.left_table for join in result.plan.joins] == ["orders"] - assert [join.right_table for join in result.plan.joins] == ["customers"] - assert ( - 'INNER JOIN "customers" ON "orders"."customer_id" = "customers"."id"' - in result.sql - ) - assert 'GROUP BY "customers"."customer_name"' in result.sql - - -def test_compile_semantic_sql_uses_semantic_metadata_documents(): - documents = [ - """ - { - "models": [ - { - "name": "orders", - "columns": [ - {"name": "id", "type": "INTEGER"}, - {"name": "customer_id", "type": "INTEGER"}, - {"name": "order_value", "type": "DECIMAL"}, - {"name": "created_at", "type": "DATE"} - ] - }, - { - "name": "customers", - "columns": [ - {"name": "id", "type": "INTEGER"}, - {"name": "customer_name", "type": "VARCHAR"} - ] - } - ], - "relationships": [ - { - "condition": "orders.customer_id = customers.id", - "joinType": "MANY_TO_ONE", - "models": ["orders", "customers"] - } - ] - } - """ - ] - - result = compile_semantic_sql( - "Show total order value by customer name this month", - documents, - now=datetime(2026, 7, 10), - ) - - assert result is not None - assert 'SUM("orders"."order_value")' in result.sql - assert 'INNER JOIN "customers" ON "orders"."customer_id" = "customers"."id"' in result.sql - assert "\"orders\".\"created_at\" >= '2026-07-01'" in result.sql - assert "\"orders\".\"created_at\" < '2026-08-01'" in result.sql - - -def test_compile_semantic_sql_counts_entities_without_unrequested_grouping(): - result = compile_semantic_sql( - "How many customers are there?", - ["CREATE TABLE customers (id INTEGER, customer_name VARCHAR);"], - ) - - assert result is not None - assert result.sql == 'SELECT COUNT(*) AS "count" FROM "customers"' - - -def test_schema_parser_handles_single_line_create_table_definitions(): - catalog = SchemaParser().parse( - ['CREATE TABLE customers (id INTEGER, customer_name VARCHAR, created_at DATE);'] - ) - - assert list(catalog.tables) == ["customers"] - assert [column.column for column in catalog.tables["customers"]] == [ - "id", - "customer_name", - "created_at", - ] From d387c1e1bc11cebcea636b17fa7e918748acce65 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 11 Jul 2026 01:14:42 +0530 Subject: [PATCH 0491/1087] Revert "Refactor SQL generation around semantic schema contracts" This reverts commit 85e948c3fdf24dc8d4974b562dacba90dc4be712. --- .../generation/followup_sql_generation.py | 15 - .../pipelines/generation/sql_correction.py | 15 - .../pipelines/generation/sql_generation.py | 15 - .../pipelines/generation/sql_regeneration.py | 17 - .../src/pipelines/generation/utils/sql.py | 440 --------------- .../retrieval/db_schema_retrieval.py | 520 +----------------- wren-ai-service/src/web/v1/services/ask.py | 505 ++--------------- .../src/web/v1/services/ask_feedback.py | 10 - .../pipelines/generation/test_sql_utils.py | 68 --- .../retrieval/test_db_schema_retrieval.py | 66 --- 10 files changed, 61 insertions(+), 1610 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 eb39e7bd97..c9f8b23537 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -17,7 +17,6 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -97,10 +96,6 @@ {% endfor %} {% endif %} -{% if semantic_schema_contract %} -{{ semantic_schema_contract }} -{% endif %} - ### QUESTION ### User's Follow-up Question: {{ query }} @@ -134,7 +129,6 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -160,9 +154,6 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -198,8 +189,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -209,8 +198,6 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -262,7 +249,6 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -286,7 +272,6 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, - "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index fff595b461..daae37e02e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,7 +15,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_sql_generation_model_kwargs, @@ -99,10 +98,6 @@ def get_sql_correction_system_prompt( {% endfor %} {% endif %} -{% if semantic_schema_contract %} -{{ semantic_schema_contract }} -{% endif %} - ### QUESTION ### {% if query %} User's Question: {{ query }} @@ -135,7 +130,6 @@ def prompt( query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -147,9 +141,6 @@ def prompt( instructions=instructions, ), sql_functions=sql_functions, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -181,8 +172,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -192,8 +181,6 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -240,7 +227,6 @@ async def run( allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -252,7 +238,6 @@ async def run( "invalid_generation_result": invalid_generation_result, "documents": contexts, "query": query, - "schema_intent_analysis": schema_intent_analysis, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index af801b6594..59bea2279c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,7 +14,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -85,10 +84,6 @@ {% endfor %} {% endif %} -{% if semantic_schema_contract %} -{{ semantic_schema_contract }} -{% endif %} - ### QUESTION ### User's Question: {{ query }} @@ -123,7 +118,6 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: schema_context = "\n".join(documents or []).lower() has_pcb_context = any( @@ -163,9 +157,6 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -198,8 +189,6 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, - query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -210,8 +199,6 @@ async def post_process( allow_data_preview=allow_data_preview, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -263,7 +250,6 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -287,7 +273,6 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, - "schema_intent_analysis": schema_intent_analysis, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 5c2089d065..8b3eaafcb1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,7 +14,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_semantic_schema_contract, construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, @@ -102,10 +101,6 @@ def get_sql_regeneration_system_prompt( {% endfor %} {% endif %} -{% if semantic_schema_contract %} -{{ semantic_schema_contract }} -{% endif %} - ### QUESTION ### SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} @@ -129,7 +124,6 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: _prompt = prompt_builder.run( sql=sql, @@ -154,9 +148,6 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, - semantic_schema_contract=construct_semantic_schema_contract( - schema_intent_analysis - ), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -186,8 +177,6 @@ async def post_process( documents: list[str], data_source: str, project_id: str | None = None, - query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), @@ -195,8 +184,6 @@ async def post_process( data_source=data_source, valid_table_names=construct_valid_table_names(documents), valid_table_columns=construct_valid_table_columns(documents), - query=query, - semantic_analysis=schema_intent_analysis, ) @@ -241,8 +228,6 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - query: str | None = None, - schema_intent_analysis: dict[str, Any] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -255,8 +240,6 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, - "query": query, - "schema_intent_analysis": schema_intent_analysis, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 02839d802b..8316d24e86 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1543,8 +1543,6 @@ async def run( allow_data_preview: bool = False, valid_table_names: list[str] | None = None, valid_table_columns: dict[str, list[str]] | None = None, - query: str | None = None, - semantic_analysis: dict[str, Any] | None = None, ) -> dict: try: cleaned_generation_result = extract_sql_generation_result(replies[0]) @@ -1573,26 +1571,6 @@ async def run( }, } - placeholder_schema_references = _extract_placeholder_schema_references( - cleaned_generation_result - ) - if placeholder_schema_references: - invalid_placeholder_list = ", ".join(placeholder_schema_references) - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_VALIDATION", - "error": ( - "Generated SQL contains placeholder schema references " - f"that are not active datasource objects: {invalid_placeholder_list}." - ), - "invalid_schema_objects": placeholder_schema_references, - "correlation_id": "", - }, - } - invalid_table_references = find_invalid_table_references( cleaned_generation_result, valid_table_names or [], @@ -1612,7 +1590,6 @@ async def run( "Use only these valid table names exactly as shown: " f"{valid_table_list}" ), - "invalid_schema_objects": invalid_table_references, "correlation_id": "", }, } @@ -1636,25 +1613,6 @@ async def run( "Use only these valid table columns exactly as shown: " f"{valid_column_list}" ), - "invalid_schema_objects": invalid_column_references, - "correlation_id": "", - }, - } - - intent_validation_error = validate_sql_intent_alignment( - query, - cleaned_generation_result, - valid_table_columns or {}, - semantic_analysis=semantic_analysis, - ) - if intent_validation_error: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_INTENT_VALIDATION", - "error": intent_validation_error, "correlation_id": "", }, } @@ -2318,404 +2276,6 @@ def construct_instructions( return _instructions -def _semantic_analysis_items( - semantic_analysis: dict[str, Any] | None, - key: str, -) -> list[str]: - if not isinstance(semantic_analysis, dict): - return [] - value = semantic_analysis.get(key) - if isinstance(value, str): - return [value] if value.strip() else [] - if isinstance(value, list): - return [ - str(item).strip() - for item in value - if item is not None and str(item).strip() - ] - return [] - - -def _semantic_analysis_dict_items( - semantic_analysis: dict[str, Any] | None, - key: str, -) -> list[dict[str, Any]]: - if not isinstance(semantic_analysis, dict): - return [] - value = semantic_analysis.get(key) - if not isinstance(value, list): - return [] - return [item for item in value if isinstance(item, dict)] - - -def _has_semantic_analysis(semantic_analysis: dict[str, Any] | None) -> bool: - if not isinstance(semantic_analysis, dict) or not semantic_analysis: - return False - semantic_keys = { - "analytical_intent", - "entities", - "identifiers", - "metrics", - "dimensions", - "filters", - "aggregations", - "relationships", - "time_constraints", - "ranking", - "sorting", - "requested_output", - "supported_schema_objects", - "candidate_schema_scores", - "concept_mappings", - "interpretations", - "missing_requirements", - "ambiguous_requirements", - "support_reasoning", - } - return any(semantic_analysis.get(key) for key in semantic_keys) - - -def _mapping_schema_objects(mapping: dict[str, Any]) -> list[str]: - value = mapping.get("schema_objects") - if isinstance(value, str): - return [value] if value.strip() else [] - if isinstance(value, list): - return [ - str(item).strip() - for item in value - if item is not None and str(item).strip() - ] - return [] - - -def _mapping_concept_type(mapping: dict[str, Any]) -> str: - return str(mapping.get("concept_type") or "").strip().lower() - - -def _mapping_request_concept(mapping: dict[str, Any]) -> str: - return str(mapping.get("request_concept") or "requested concept").strip() - - -def construct_semantic_schema_contract( - semantic_analysis: dict[str, Any] | None, -) -> str: - if not _has_semantic_analysis(semantic_analysis): - return "" - - lines = [ - "### SEMANTIC SCHEMA CONTRACT ###", - "Use this contract as the primary source of truth for SQL generation.", - "Generate SQL only from schema objects listed here and in the retrieved DATABASE SCHEMA.", - "Do not infer alternative tables, columns, joins, metrics, filters, or time fields independently.", - ] - analytical_intent = str( - semantic_analysis.get("analytical_intent") or "" - ).strip() - if analytical_intent: - lines.append(f"Analytical intent: {analytical_intent}") - - for label, key in ( - ("Entities", "entities"), - ("Identifiers", "identifiers"), - ("Metrics", "metrics"), - ("Dimensions", "dimensions"), - ("Filters", "filters"), - ("Time constraints", "time_constraints"), - ("Aggregations", "aggregations"), - ("Ranking", "ranking"), - ("Sorting", "sorting"), - ("Requested output", "requested_output"), - ("Relationships", "relationships"), - ("Supported schema objects", "supported_schema_objects"), - ): - values = _semantic_analysis_items(semantic_analysis, key) - if values: - lines.append(f"{label}: {', '.join(values)}") - - mappings = _semantic_analysis_dict_items(semantic_analysis, "concept_mappings") - if mappings: - lines.append("Required concept-to-schema mappings:") - for mapping in mappings: - schema_objects = _mapping_schema_objects(mapping) - required = "required" if mapping.get("required_in_sql") is not False else "optional" - confidence = mapping.get("confidence") - confidence_text = ( - f", confidence={confidence}" - if confidence is not None and str(confidence).strip() - else "" - ) - lines.append( - "- " - f"{_mapping_request_concept(mapping)} " - f"({_mapping_concept_type(mapping) or 'concept'}, {required}{confidence_text}) " - f"-> {', '.join(schema_objects) or 'NO_MAPPING'}." - ) - - interpretations = _semantic_analysis_dict_items( - semantic_analysis, "interpretations" - ) - if interpretations: - lines.append("Ranked schema interpretations:") - for interpretation in interpretations: - description = str(interpretation.get("description") or "").strip() - if not description: - continue - selected = "selected" if interpretation.get("is_selected") is True else "candidate" - schema_objects = interpretation.get("schema_objects") - schema_text = ( - ", ".join(str(item).strip() for item in schema_objects if item) - if isinstance(schema_objects, list) - else "" - ) - lines.append(f"- {description} ({selected}). Objects: {schema_text}") - - missing_requirements = _semantic_analysis_items( - semantic_analysis, "missing_requirements" - ) - if missing_requirements: - lines.append(f"Missing requirements: {', '.join(missing_requirements)}") - ambiguous_requirements = _semantic_analysis_items( - semantic_analysis, "ambiguous_requirements" - ) - if ambiguous_requirements: - lines.append(f"Ambiguous requirements: {', '.join(ambiguous_requirements)}") - support_reasoning = str(semantic_analysis.get("support_reasoning") or "").strip() - if support_reasoning: - lines.append(f"Support reasoning: {support_reasoning}") - - lines.append( - "Validation requirement: every required mapped schema object must appear in SQL when it represents an entity, identifier, dimension, metric, filter, time field, relationship, or requested output. Ranking/sorting/aggregation concepts must appear as SQL shape." - ) - return "\n".join(lines) - - -def get_schema_intent_analysis_error( - semantic_analysis: dict[str, Any] | None, -) -> str | None: - if not _has_semantic_analysis(semantic_analysis): - return None - missing = _semantic_analysis_items(semantic_analysis, "missing_requirements") - if missing: - return ( - "The active datasource schema does not expose the information needed " - f"to answer the request: {', '.join(missing)}. I cannot generate unrelated SQL." - ) - ambiguous = _semantic_analysis_items(semantic_analysis, "ambiguous_requirements") - if ambiguous: - return ( - "The request has multiple equally plausible schema interpretations: " - f"{', '.join(ambiguous)}. Please clarify which one to use." - ) - candidates = _semantic_analysis_dict_items( - semantic_analysis, "candidate_schema_scores" - ) - if candidates and not any(candidate.get("is_complete") is True for candidate in candidates): - missing_concepts = [] - for candidate in candidates[:3]: - missing_value = candidate.get("missing_concepts") - if isinstance(missing_value, list) and missing_value: - missing_concepts.append( - f"{candidate.get('candidate_id') or 'candidate'}: " - + ", ".join(str(item) for item in missing_value if item) - ) - if missing_concepts: - return ( - "Semantic schema retrieval did not find a complete mapping for " - f"the request. Missing concepts: {'; '.join(missing_concepts)}." - ) - if semantic_analysis.get("is_fully_supported") is False: - reason = str(semantic_analysis.get("support_reasoning") or "").strip() - return reason or ( - "The selected schema does not fully support every required component of the request." - ) - schema_bound_types = { - "dimension", - "entity", - "filter", - "identifier", - "metric", - "relationship", - "time", - "output", - } - unsupported = [] - for mapping in _semantic_analysis_dict_items(semantic_analysis, "concept_mappings"): - if mapping.get("required_in_sql") is False: - continue - if _mapping_concept_type(mapping) not in schema_bound_types: - continue - if not _mapping_schema_objects(mapping): - unsupported.append(_mapping_request_concept(mapping)) - if unsupported: - return ( - "The semantic contract did not map required request concepts to active " - f"schema objects: {', '.join(unsupported)}." - ) - return None - - -def _extract_placeholder_schema_references(sql: str) -> list[str]: - placeholders = re.findall(r"<\s*([^<>]+?)\s*>", sql or "") - placeholder_names = re.findall( - r"\b(?:table|column|schema|database|field|metric|dimension|date|amount|entity)_?name\b", - sql or "", - flags=re.IGNORECASE, - ) - return sorted( - { - str(item).strip() - for item in [*placeholders, *placeholder_names] - if item is not None and str(item).strip() - } - ) - - -def _sql_references_table(sql: str, table_name: str) -> bool: - normalized_table = _compact_sql_identifier(table_name) - if not normalized_table: - return False - table_suffixes = { - _compact_sql_identifier(suffix) - for suffix in _table_reference_suffixes(table_name) - } - table_suffixes.add(normalized_table) - for referenced_table in extract_sql_table_references(sql or ""): - normalized_reference = _compact_sql_identifier(referenced_table) - if not normalized_reference: - continue - if normalized_reference in table_suffixes: - return True - if any( - normalized_reference.endswith(suffix) or suffix.endswith(normalized_reference) - for suffix in table_suffixes - if suffix - ): - return True - return False - - -def _sql_references_schema_object( - sql: str, - schema_object: str, - valid_table_columns: dict[str, list[str]], -) -> bool: - parts = [ - _normalize_sql_identifier(part.strip()) - for part in str(schema_object or "").split(".") - if part.strip() - ] - if not parts: - return False - if len(parts) == 1: - identifier = parts[0] - return bool( - re.search( - rf'(? bool: - return bool( - re.search(r"\bCOUNT\s*\(\s*\*\s*\)", sql or "", flags=re.IGNORECASE) - ) and not re.search( - r"\b(?:SUM|AVG|MIN|MAX)\s*\(", sql or "", flags=re.IGNORECASE - ) - - -def _semantic_requests_record_count(semantic_analysis: dict[str, Any]) -> bool: - intent = str(semantic_analysis.get("analytical_intent") or "").lower() - if intent == "record_count": - return True - text = " ".join( - item - for key in ("metrics", "aggregations", "requested_output") - for item in _semantic_analysis_items(semantic_analysis, key) - ).lower() - return bool(re.search(r"\b(?:count|number of|row count|record count)\b", text)) - - -def validate_sql_intent_alignment( - query: str | None, - sql: str, - valid_table_columns: dict[str, list[str]] | None = None, - semantic_analysis: dict[str, Any] | None = None, -) -> str | None: - valid_table_columns = valid_table_columns or {} - if analysis_error := get_schema_intent_analysis_error(semantic_analysis): - return analysis_error - if not _has_semantic_analysis(semantic_analysis): - return None - - requests_count = _semantic_requests_record_count(semantic_analysis) - mappings = _semantic_analysis_dict_items(semantic_analysis, "concept_mappings") - schema_bound_types = { - "dimension", - "entity", - "filter", - "identifier", - "metric", - "relationship", - "time", - "output", - } - for mapping in mappings: - if mapping.get("required_in_sql") is False: - continue - concept_type = _mapping_concept_type(mapping) - if concept_type not in schema_bound_types: - continue - schema_objects = _mapping_schema_objects(mapping) - if not schema_objects: - return ( - "The semantic analysis did not map the required " - f"{concept_type or 'concept'} '{_mapping_request_concept(mapping)}' " - "to a schema object." - ) - if not any( - _sql_references_schema_object(sql, schema_object, valid_table_columns) - for schema_object in schema_objects - ): - return ( - "Generated SQL does not reference schema objects mapped to the " - f"required {concept_type or 'concept'} " - f"'{_mapping_request_concept(mapping)}': {', '.join(schema_objects)}." - ) - if concept_type == "metric" and _sql_is_plain_count(sql) and not requests_count: - return ( - "Generated SQL answers with a generic record count, but the " - f"requested metric '{_mapping_request_concept(mapping)}' must be " - "retrieved or calculated from the mapped schema object." - ) - - if _semantic_analysis_items(semantic_analysis, "ranking") and not re.search( - r"\bORDER\s+BY\b.*\b(?:LIMIT|FETCH\s+FIRST|TOP\s*(?:\(\s*)?\d+)\b|\b(?:LIMIT|FETCH\s+FIRST|TOP\s*(?:\(\s*)?\d+)\b.*\bORDER\s+BY\b", - sql or "", - flags=re.IGNORECASE | re.DOTALL, - ): - return "Generated SQL does not include sorting and limiting logic required by the ranking intent." - if _semantic_analysis_items(semantic_analysis, "sorting") and not re.search( - r"\bORDER\s+BY\b", sql or "", flags=re.IGNORECASE - ): - return "Generated SQL does not include sorting logic required by the semantic intent." - if _semantic_analysis_items(semantic_analysis, "aggregations") and not re.search( - r"\b(?:COUNT|SUM|AVG|MIN|MAX)\s*\(", sql or "", flags=re.IGNORECASE - ): - return "Generated SQL does not include the aggregation required by the semantic intent." - return None - - def _parse_semantic_metadata_content(content: str) -> Any | None: content = content.strip() if not content: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index dd5c81a9b3..107c5c3479 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,6 +1,5 @@ import ast import logging -import re import sys from typing import TYPE_CHECKING, Any, Optional @@ -11,7 +10,7 @@ from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel, Field +from pydantic import BaseModel from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider @@ -32,78 +31,23 @@ table_columns_selection_system_prompt = """ ### TASK ### -You are a highly skilled data analyst. Your goal is to examine the active deployed database schema, semantically interpret the user's question, and identify the exact schema objects required to construct an accurate SQL query. +You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### -1. First identify the user's semantic intent: requested entities, identifiers, metrics, dimensions, filters, aggregations, ranking, sorting, time constraints, joins/relationships, and requested output shape. -2. Map each required business concept only to tables, columns, metrics, views, or relationships that are explicitly present in the active schema metadata. -3. Use semantic similarity from table names, column names, comments/descriptions, data types, primary/foreign keys, relationships, and metric/view definitions. Do not use datasource-specific rules, hardcoded mappings, or default tables. -4. Rank candidate schema mappings by complete concept coverage, semantic fit, relationship viability, metric validity, data type compatibility, and support for filters/time/ranking/aggregation requirements. -5. Select only the highest-confidence complete candidate. If no candidate fully supports the request, set `is_fully_supported` to false and list missing or ambiguous requirements instead of selecting unrelated fallback schema. -6. If RETRY CONTEXT is provided, treat rejected schema objects as failed mappings. Prefer the next-best complete candidate that excludes those objects. -7. Populate `concept_mappings` for every required concept. Each mapping must classify the concept, list directly supporting schema objects, mark whether it must appear in SQL, and include a confidence score. -8. Populate `candidate_schema_scores` with accepted and rejected candidates. Include covered and missing concepts and a concise selection reason. -9. Include join keys and relationship columns when multiple tables are needed. If no trustworthy join path exists, mark the relationship missing instead of inventing one. -10. Do not add filters or time constraints not requested or implied by the user. -11. If a "." is included in columns, put the name before the first dot into chosen columns. -12. The number of columns chosen must match the number of reasoning. -13. Final chosen columns must be only column names, don't prefix it with table names. -14. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. +2. For each table, provide a clear and concise reasoning for why specific columns are selected. +3. List each reason as part of a step-by-step chain of thought, justifying the inclusion of each column. +4. If a "." is included in columns, put the name before the first dot into chosen columns. +5. The number of columns chosen must match the number of reasoning. +6. Final chosen columns must be only column names, don't prefix it with table names. +7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: { - "semantic_analysis": { - "analytical_intent": "retrieval | detailed_records | summary | comparison | trend | dashboard | kpi | ranking | record_count | other", - "entities": ["requested business entities"], - "identifiers": ["requested identifiers"], - "metrics": ["requested measures or calculated metrics"], - "dimensions": ["requested grouping/descriptive dimensions"], - "filters": ["requested filters"], - "aggregations": ["requested aggregations/calculations"], - "relationships": ["required joins or relationship paths"], - "time_constraints": ["requested date filters, grains, or trend requirements"], - "ranking": ["top/bottom/rank/limit requirements"], - "sorting": ["requested ordering requirements"], - "requested_output": ["requested columns, charts, summaries, records, or KPIs"], - "supported_schema_objects": ["table.column, table, view, metric, or relationship objects selected"], - "candidate_schema_scores": [ - { - "candidate_id": "candidate-1", - "schema_objects": ["schema objects included in this candidate"], - "covered_concepts": ["request concepts this candidate supports"], - "missing_concepts": ["request concepts this candidate cannot support"], - "confidence": 0.0, - "is_complete": true, - "selection_reason": "Why this candidate is accepted or rejected" - } - ], - "concept_mappings": [ - { - "request_concept": "business concept from the user request", - "concept_type": "entity | identifier | dimension | metric | filter | time | aggregation | ranking | sorting | relationship | output", - "schema_objects": ["table.column, table, metric, view, or relationship object that directly supports the concept"], - "required_in_sql": true, - "confidence": 0.0, - "mapping_reason": "Why these schema objects directly support the concept" - } - ], - "interpretations": [ - { - "description": "Possible schema interpretation", - "schema_objects": ["schema objects used by this interpretation"], - "confidence": 0.0, - "is_selected": true - } - ], - "missing_requirements": ["required concepts not supported by active schema"], - "ambiguous_requirements": ["concepts with multiple equally plausible mappings"], - "is_fully_supported": true, - "support_reasoning": "Concise explanation of schema support" - }, "results": [ { "table_selection_reason": "Reason for selecting tablename1", @@ -138,7 +82,6 @@ - Each table key must list only the columns relevant to answering the question. - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. -- Populate `semantic_analysis` before `results` and keep both consistent. - Be logical, concise, and ensure the output strictly follows the required JSON format. - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. @@ -155,38 +98,8 @@ {{ db_schema }} {% endfor %} -{% if semantic_candidate_context %} -### GENERIC SEMANTIC CANDIDATE RANKING ### -The following candidates were scored from active schema metadata only. Use this as evidence, then validate complete concept coverage before selecting a semantic contract. -{% for candidate in semantic_candidate_context %} -- candidate_id: {{ candidate.candidate_id }} - table_name: {{ candidate.table_name }} - confidence: {{ candidate.confidence }} - coverage_score: {{ candidate.coverage_score }} - matched_query_terms: {{ candidate.matched_query_terms }} - missing_query_terms: {{ candidate.missing_query_terms }} - rejected_by_retry: {{ candidate.rejected_by_retry }} - selection_reason: {{ candidate.selection_reason }} - matched_columns: -{% for column in candidate.matched_columns %} - - {{ column.column_name }} (score={{ column.score }}, data_type={{ column.data_type }}, matched_terms={{ column.matched_terms }}) -{% endfor %} -{% endfor %} -{% endif %} - ### INPUT ### {{ question }} - -{% if semantic_retry_context %} -### RETRY CONTEXT ### -Previous semantic SQL validation failed. Discard the previous contract and retrieve the next-best complete schema mapping. -Validation failure: {{ semantic_retry_context.validation_error }} -Retry attempt: {{ semantic_retry_context.retry_attempt }} -Rejected schema objects: -{% for schema_object in semantic_retry_context.rejected_schema_objects %} -- {{ schema_object }} -{% endfor %} -{% endif %} """ @@ -276,275 +189,6 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: return deduped -_SEMANTIC_TOKEN_STOPWORDS = { - "a", - "an", - "and", - "are", - "as", - "at", - "be", - "by", - "for", - "from", - "give", - "have", - "how", - "in", - "is", - "me", - "of", - "on", - "or", - "show", - "that", - "the", - "to", - "with", - "table", - "view", -} - -_NUMERIC_INTENT_TERMS = { - "amount", - "avg", - "average", - "balance", - "cost", - "count", - "measure", - "metric", - "price", - "quantity", - "rate", - "sum", - "total", - "value", -} - -_TEMPORAL_INTENT_TERMS = { - "date", - "day", - "month", - "quarter", - "time", - "trend", - "week", - "year", -} - -_RANKING_INTENT_TERMS = { - "bottom", - "highest", - "least", - "lowest", - "most", - "rank", - "ranking", - "top", -} - - -def _semantic_tokens(value: Any) -> set[str]: - text = str(value or "") - if not text: - return set() - text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) - text = re.sub(r"[^A-Za-z0-9]+", " ", text) - tokens = { - token.lower() - for token in text.split() - if len(token) > 1 and token.lower() not in _SEMANTIC_TOKEN_STOPWORDS - } - for token in list(tokens): - if token.endswith("ies") and len(token) > 4: - tokens.add(f"{token[:-3]}y") - elif token.endswith("s") and len(token) > 3: - tokens.add(token[:-1]) - return tokens - - -def _column_tokens(column: dict[str, Any]) -> set[str]: - tokens = set() - for key in ("name", "display_name", "alias", "comment", "description", "data_type"): - tokens.update(_semantic_tokens(column.get(key))) - return tokens - - -def _table_tokens(table_schema: dict[str, Any]) -> set[str]: - tokens = set() - for key in ("name", "display_name", "alias", "comment", "description"): - tokens.update(_semantic_tokens(table_schema.get(key))) - for column in table_schema.get("columns", []) or []: - if isinstance(column, dict): - tokens.update(_column_tokens(column)) - return tokens - - -def _is_numeric_column(column: dict[str, Any]) -> bool: - return bool( - re.search( - r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|number|double|float|real|money)\b", - str(column.get("data_type") or "").lower(), - ) - ) - - -def _is_temporal_column(column: dict[str, Any]) -> bool: - return bool( - re.search( - r"\b(?:date|time|timestamp|datetime)\b", - str(column.get("data_type") or "").lower(), - ) - ) - - -def _normalized_schema_object(value: Any) -> str: - return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) - - -def _rejected_schema_objects(semantic_retry_context: dict[str, Any] | None) -> set[str]: - if not isinstance(semantic_retry_context, dict): - return set() - rejected = semantic_retry_context.get("rejected_schema_objects") - if not isinstance(rejected, list): - return set() - return { - _normalized_schema_object(item) - for item in rejected - if item is not None and str(item).strip() - } - - -def _schema_object_was_rejected( - table_name: str, - column_name: str | None, - rejected_schema_objects: set[str], -) -> bool: - if not rejected_schema_objects: - return False - object_key = _normalized_schema_object( - f"{table_name}.{column_name}" if column_name else table_name - ) - table_key = _normalized_schema_object(table_name) - return any( - rejected_key - and ( - rejected_key == object_key - or rejected_key == table_key - or rejected_key.endswith(object_key) - or object_key.endswith(rejected_key) - ) - for rejected_key in rejected_schema_objects - ) - - -def rank_semantic_schema_candidates( - query: str, - construct_db_schemas: list[dict], - semantic_retry_context: dict[str, Any] | None = None, - max_candidates: int = 15, - max_columns_per_candidate: int = 8, -) -> list[dict[str, Any]]: - query_terms = _semantic_tokens(query) - if not query_terms: - return [] - numeric_terms = query_terms & _NUMERIC_INTENT_TERMS - temporal_terms = query_terms & _TEMPORAL_INTENT_TERMS - ranking_terms = query_terms & _RANKING_INTENT_TERMS - rejected_objects = _rejected_schema_objects(semantic_retry_context) - candidates: list[dict[str, Any]] = [] - - for table_schema in construct_db_schemas: - if table_schema.get("type") != "TABLE": - continue - table_name = str(table_schema.get("name") or "").strip() - if not table_name: - continue - table_matches = _table_tokens(table_schema) & query_terms - table_rejected = _schema_object_was_rejected( - table_name, None, rejected_objects - ) - matched_columns = [] - for column in table_schema.get("columns", []) or []: - if not isinstance(column, dict): - continue - column_name = str(column.get("name") or "").strip() - if not column_name: - continue - tokens = _column_tokens(column) - matched_terms = sorted(tokens & query_terms) - score = len(matched_terms) * 3.0 - if numeric_terms and _is_numeric_column(column): - score += 1.0 - if temporal_terms and _is_temporal_column(column): - score += 1.0 - if ranking_terms and matched_terms: - score += 0.5 - rejected = _schema_object_was_rejected( - table_name, column_name, rejected_objects - ) - if rejected: - score -= 5.0 - if score > 0 or matched_terms: - matched_columns.append( - { - "column_name": column_name, - "score": round(max(score, 0.0), 3), - "matched_terms": matched_terms, - "data_type": str(column.get("data_type") or ""), - "rejected_by_retry": rejected, - } - ) - - matched_columns.sort( - key=lambda item: (item["score"], len(item["matched_terms"])), - reverse=True, - ) - matched_columns = matched_columns[:max_columns_per_candidate] - covered_terms = set(table_matches) - for column in matched_columns: - covered_terms.update(column["matched_terms"]) - if not covered_terms and not table_rejected: - continue - coverage_score = len(covered_terms) / max(len(query_terms), 1) - raw_score = ( - len(table_matches) * 2 - + sum(column["score"] for column in matched_columns) - + coverage_score * 4 - ) - if table_rejected: - raw_score -= 6 - candidates.append( - { - "candidate_id": f"candidate-{len(candidates) + 1}", - "table_name": table_name, - "confidence": round(min(max(raw_score / 20, 0), 0.99), 3), - "coverage_score": round(coverage_score, 3), - "matched_query_terms": sorted(covered_terms), - "missing_query_terms": sorted(query_terms - covered_terms), - "matched_columns": matched_columns, - "rejected_by_retry": table_rejected - or any(column["rejected_by_retry"] for column in matched_columns), - "selection_reason": ( - f"Covers {len(covered_terms)} of {len(query_terms)} significant request terms" - ), - } - ) - - candidates.sort( - key=lambda item: ( - item["rejected_by_retry"] is False, - item["confidence"], - item["coverage_score"], - ), - reverse=True, - ) - for index, candidate in enumerate(candidates[:max_candidates], start=1): - candidate["candidate_id"] = f"candidate-{index}" - return candidates[:max_candidates] - - @observe(capture_input=False, capture_output=False) async def embedding( query: str, @@ -705,22 +349,12 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) - if enable_column_pruning or _token_count > context_window_size: - return { - "db_schemas": [], - "tokens": _token_count, - "has_calculated_field": has_calculated_field, - "has_metric": has_metric, - "has_json_field": has_json_field, - "semantic_analysis": {}, - } return { "db_schemas": retrieval_results, "tokens": _token_count, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, - "semantic_analysis": {}, } @@ -731,7 +365,6 @@ def prompt( prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, histories: list[AskHistory], - semantic_retry_context: dict[str, Any] | None = None, ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -744,22 +377,8 @@ def prompt( ) query = "\n".join(previous_query_summaries) + "\n" + query - semantic_candidate_context = rank_semantic_schema_candidates( - query=query, - construct_db_schemas=construct_db_schemas, - semantic_retry_context=semantic_retry_context, - ) - logger.info( - "semantic_retrieval_pre_ranked_candidates=%s", - semantic_candidate_context, - ) - _prompt = prompt_builder.run( - question=query, - db_schemas=db_schemas, - semantic_candidate_context=semantic_candidate_context, - semantic_retry_context=semantic_retry_context or {}, - ) + _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: return {} @@ -786,10 +405,9 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - retrieval_payload = orjson.loads(filter_columns_in_tables["replies"][0]) - columns_and_tables_needed = retrieval_payload.get("results", []) - semantic_analysis = retrieval_payload.get("semantic_analysis") or {} - _log_semantic_retrieval_decision(semantic_analysis) + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + )["results"] # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -849,7 +467,6 @@ def construct_retrieval_results( "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, - "semantic_analysis": semantic_analysis, } else: retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] @@ -861,64 +478,9 @@ def construct_retrieval_results( ], "has_metric": check_using_db_schemas_without_pruning["has_metric"], "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], - "semantic_analysis": check_using_db_schemas_without_pruning.get( - "semantic_analysis", {} - ), } -def _semantic_log_items(semantic_analysis: dict[str, Any], key: str) -> list[str]: - value = semantic_analysis.get(key) - if isinstance(value, str): - return [value] if value.strip() else [] - if isinstance(value, list): - return [ - str(item).strip() - for item in value - if item is not None and str(item).strip() - ] - return [] - - -def _log_semantic_retrieval_decision(semantic_analysis: dict[str, Any]) -> None: - if not isinstance(semantic_analysis, dict) or not semantic_analysis: - logger.info("semantic_retrieval_decision=no_semantic_analysis") - return - - logger.info( - "semantic_retrieval_concepts=%s", - { - "intent": semantic_analysis.get("analytical_intent"), - "entities": _semantic_log_items(semantic_analysis, "entities"), - "identifiers": _semantic_log_items(semantic_analysis, "identifiers"), - "metrics": _semantic_log_items(semantic_analysis, "metrics"), - "dimensions": _semantic_log_items(semantic_analysis, "dimensions"), - "filters": _semantic_log_items(semantic_analysis, "filters"), - "time_constraints": _semantic_log_items( - semantic_analysis, "time_constraints" - ), - "aggregations": _semantic_log_items(semantic_analysis, "aggregations"), - "ranking": _semantic_log_items(semantic_analysis, "ranking"), - "sorting": _semantic_log_items(semantic_analysis, "sorting"), - }, - ) - logger.info( - "semantic_retrieval_candidate_scores=%s", - semantic_analysis.get("candidate_schema_scores") or [], - ) - logger.info( - "semantic_retrieval_selected_contract=%s", - { - "supported_schema_objects": semantic_analysis.get( - "supported_schema_objects", [] - ), - "concept_mappings": semantic_analysis.get("concept_mappings", []), - "is_fully_supported": semantic_analysis.get("is_fully_supported"), - "support_reasoning": semantic_analysis.get("support_reasoning"), - }, - ) - - ## End of Pipeline class MatchingTableContents(BaseModel): chain_of_thought_reasoning: list[str] @@ -931,59 +493,7 @@ class MatchingTable(BaseModel): table_selection_reason: str -class SemanticConceptMapping(BaseModel): - request_concept: str = "" - concept_type: str = "" - schema_objects: list[str] = Field(default_factory=list) - required_in_sql: bool = True - confidence: float | None = None - mapping_reason: str = "" - - -class SemanticInterpretation(BaseModel): - description: str = "" - schema_objects: list[str] = Field(default_factory=list) - confidence: float | None = None - is_selected: bool = False - - -class SemanticCandidateSchemaScore(BaseModel): - candidate_id: str = "" - schema_objects: list[str] = Field(default_factory=list) - covered_concepts: list[str] = Field(default_factory=list) - missing_concepts: list[str] = Field(default_factory=list) - confidence: float | None = None - is_complete: bool = False - selection_reason: str = "" - - -class SemanticAnalysis(BaseModel): - analytical_intent: str = "" - entities: list[str] = Field(default_factory=list) - identifiers: list[str] = Field(default_factory=list) - metrics: list[str] = Field(default_factory=list) - dimensions: list[str] = Field(default_factory=list) - filters: list[str] = Field(default_factory=list) - aggregations: list[str] = Field(default_factory=list) - relationships: list[str] = Field(default_factory=list) - time_constraints: list[str] = Field(default_factory=list) - ranking: list[str] = Field(default_factory=list) - sorting: list[str] = Field(default_factory=list) - requested_output: list[str] = Field(default_factory=list) - supported_schema_objects: list[str] = Field(default_factory=list) - candidate_schema_scores: list[SemanticCandidateSchemaScore] = Field( - default_factory=list - ) - concept_mappings: list[SemanticConceptMapping] = Field(default_factory=list) - interpretations: list[SemanticInterpretation] = Field(default_factory=list) - missing_requirements: list[str] = Field(default_factory=list) - ambiguous_requirements: list[str] = Field(default_factory=list) - is_fully_supported: bool | None = None - support_reasoning: str = "" - - class RetrievalResults(BaseModel): - semantic_analysis: SemanticAnalysis | None = None results: list[MatchingTable] @@ -1052,11 +562,8 @@ async def run( project_id: Optional[str] = None, histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, - semantic_retry_context: Optional[dict[str, Any]] = None, ): logger.info("Ask Retrieval pipeline is running...") - if semantic_retry_context: - logger.info("semantic_retrieval_retry_context=%s", semantic_retry_context) return await self._pipe.execute( ["construct_retrieval_results"], inputs={ @@ -1065,7 +572,6 @@ async def run( "project_id": project_id or "", "histories": histories or [], "enable_column_pruning": enable_column_pruning, - "semantic_retry_context": semantic_retry_context or {}, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e0f62e945a..9db38bd5a8 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -11,7 +11,6 @@ from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, construct_valid_table_names, - get_schema_intent_analysis_error, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, ) @@ -4430,94 +4429,6 @@ def _metadata_from_documents( ] return table_names, table_ddls - def _semantic_analysis_has_contract( - self, semantic_analysis: dict[str, Any] | None - ) -> bool: - if not isinstance(semantic_analysis, dict) or not semantic_analysis: - return False - if semantic_analysis.get("supported_schema_objects"): - return True - mappings = semantic_analysis.get("concept_mappings") - if isinstance(mappings, list) and any( - isinstance(mapping, dict) and mapping.get("schema_objects") - for mapping in mappings - ): - return True - candidates = semantic_analysis.get("candidate_schema_scores") - if isinstance(candidates, list) and any( - isinstance(candidate, dict) - and candidate.get("is_complete") is True - and candidate.get("schema_objects") - for candidate in candidates - ): - return True - return False - - def _semantic_schema_objects( - self, semantic_analysis: dict[str, Any] | None - ) -> list[str]: - if not isinstance(semantic_analysis, dict): - return [] - schema_objects: list[str] = [] - supported = semantic_analysis.get("supported_schema_objects") - if isinstance(supported, list): - schema_objects.extend(str(item).strip() for item in supported if item) - mappings = semantic_analysis.get("concept_mappings") - if isinstance(mappings, list): - for mapping in mappings: - if not isinstance(mapping, dict): - continue - mapped_objects = mapping.get("schema_objects") - if isinstance(mapped_objects, str) and mapped_objects.strip(): - schema_objects.append(mapped_objects.strip()) - elif isinstance(mapped_objects, list): - schema_objects.extend( - str(item).strip() for item in mapped_objects if item - ) - return list(dict.fromkeys(item for item in schema_objects if item)) - - def _semantic_retry_context( - self, - *, - semantic_analysis: dict[str, Any] | None, - invalid_generation_result: dict[str, Any] | None, - retry_attempt: int, - ) -> dict[str, Any]: - invalid_generation_result = invalid_generation_result or {} - rejected_schema_objects = self._semantic_schema_objects(semantic_analysis) - invalid_schema_objects = invalid_generation_result.get("invalid_schema_objects") - if isinstance(invalid_schema_objects, list): - rejected_schema_objects.extend( - str(item).strip() for item in invalid_schema_objects if item - ) - elif isinstance(invalid_schema_objects, str) and invalid_schema_objects.strip(): - rejected_schema_objects.append(invalid_schema_objects.strip()) - - selected_candidates = [] - if isinstance(semantic_analysis, dict): - candidates = semantic_analysis.get("candidate_schema_scores") - if isinstance(candidates, list): - selected_candidates = [ - candidate - for candidate in candidates - if isinstance(candidate, dict) - and ( - candidate.get("is_complete") is True - or candidate.get("schema_objects") - ) - ][:3] - - return { - "retry_attempt": retry_attempt, - "failure_type": invalid_generation_result.get("type"), - "failure_reason": invalid_generation_result.get("error"), - "rejected_sql": invalid_generation_result.get("sql"), - "rejected_schema_objects": list( - dict.fromkeys(item for item in rejected_schema_objects if item) - ), - "rejected_candidates": selected_candidates, - } - async def _complete_sql_generation_context( self, *, @@ -5369,9 +5280,6 @@ async def ask( table_names = [] table_ddls = [] _retrieval_result = {} - schema_intent_analysis: dict[str, Any] = {} - semantic_pipeline_active = False - semantic_contract_available = False error_message = None invalid_sql = None allow_sql_generation_reasoning = ( @@ -5492,7 +5400,7 @@ async def ask( tables=explicit_table_names, project_id=ask_request.project_id, histories=[], - enable_column_pruning=True, + enable_column_pruning=False, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -5518,10 +5426,10 @@ async def ask( retrieval_result = await self._run_with_timeout( "Full active schema retrieval for explicit table", self._pipelines["db_schema_retrieval"].run( - query=user_query, + query="", project_id=ask_request.project_id, histories=[], - enable_column_pruning=True, + enable_column_pruning=False, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -5541,30 +5449,14 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) - semantic_pipeline_active = True - semantic_contract_available = self._semantic_analysis_has_contract( - schema_intent_analysis - ) - logger.info( - "sql_generation_pipeline_decision query_id=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s", - query_id, - semantic_pipeline_active, - semantic_contract_available, - self._semantic_schema_objects(schema_intent_analysis), - ) logger.info( "Retrieved explicit tables for query_id %s: %s", query_id, table_names, ) - if not semantic_pipeline_active and ( - table_question_sql := self._build_schema_grounded_table_question_sql( + if table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls - ) ): ask_result = self._build_validated_ask_result_from_sql( table_question_sql, @@ -5588,10 +5480,8 @@ async def ask( return results invalid_sql = table_question_sql - if not semantic_pipeline_active and ( - explicit_table_preview := self._build_explicit_table_preview_sql( + if explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls - ) ): explicit_sql, explicit_table_name = explicit_table_preview if explicit_table_name not in table_names: @@ -5618,7 +5508,7 @@ async def ask( return results invalid_sql = explicit_sql - if not semantic_pipeline_active and documents and ( + if documents and ( deterministic_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6058,7 +5948,10 @@ async def ask( query=sql_user_query, histories=[], project_id=ask_request.project_id, - enable_column_pruning=True, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -6070,12 +5963,12 @@ async def ask( error, ) retrieval_result = await self._run_with_timeout( - "Semantic schema retry after retrieval timeout", + "Deployed schema fallback retrieval", self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, + query="", histories=[], project_id=ask_request.project_id, - enable_column_pruning=True, + enable_column_pruning=False, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -6085,20 +5978,6 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) - semantic_pipeline_active = True - semantic_contract_available = self._semantic_analysis_has_contract( - schema_intent_analysis - ) - logger.info( - "sql_generation_pipeline_decision query_id=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s", - query_id, - semantic_pipeline_active, - semantic_contract_available, - self._semantic_schema_objects(schema_intent_analysis), - ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6126,7 +6005,7 @@ async def ask( tables=explicit_table_names, project_id=ask_request.project_id, histories=[], - enable_column_pruning=True, + enable_column_pruning=enable_column_pruning, ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, @@ -6136,13 +6015,6 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) - semantic_pipeline_active = True - semantic_contract_available = self._semantic_analysis_has_contract( - schema_intent_analysis - ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6152,11 +6024,7 @@ async def ask( documents, ) ) - if ( - not semantic_pipeline_active - and not documents - and self._is_data_analysis_query(user_query) - ): + if not documents and self._is_data_analysis_query(user_query): logger.info( "Query-based schema retrieval returned no tables for data question; " "retrying full active deployed schema for query_id %s", @@ -6179,13 +6047,6 @@ async def ask( _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) - semantic_pipeline_active = True - semantic_contract_available = self._semantic_analysis_has_contract( - schema_intent_analysis - ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) @@ -6200,7 +6061,7 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6220,7 +6081,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6244,7 +6105,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6268,7 +6129,6 @@ async def ask( if ( not api_results - and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and ( schema_grounded_sql := self._build_schema_grounded_analytics_sql( @@ -6293,20 +6153,16 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if ( - not semantic_pipeline_active - and not api_results - and any( - term in user_query.lower() - for term in ( - "pcb", - "repair", - "failure", - "business unit", - "business units", - "product line", - "product family", - ) + if not api_results and any( + term in user_query.lower() + for term in ( + "pcb", + "repair", + "failure", + "business unit", + "business units", + "product line", + "product family", ) ): operational_sql = self._build_schema_grounded_analytics_sql( @@ -6330,7 +6186,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if not semantic_pipeline_active and not api_results and ( + if not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6354,7 +6210,6 @@ async def ask( should_retry_full_schema = ( not api_results - and not semantic_pipeline_active and self._is_data_analysis_query(user_query) and "db_schema_retrieval" in self._pipelines ) @@ -6464,7 +6319,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if not semantic_pipeline_active and not documents: + if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): @@ -6551,23 +6406,12 @@ async def ask( table_ddls=table_ddls, ) if completed_retrieval_result: - completed_semantic_analysis = completed_retrieval_result.get( - "semantic_analysis", {} - ) - if self._semantic_analysis_has_contract( - completed_semantic_analysis - ): - _retrieval_result = completed_retrieval_result - schema_intent_analysis = completed_semantic_analysis - semantic_pipeline_active = True - semantic_contract_available = True + _retrieval_result = completed_retrieval_result sql_generation_histories = histories if self._is_data_analysis_query( sql_user_query - ) and not self._needs_conversation_context( - sql_user_query - ) and not semantic_pipeline_active: + ) and not self._needs_conversation_context(sql_user_query): sql_generation_histories = [] allow_sql_generation_reasoning = False allow_sql_knowledge_retrieval = False @@ -6576,31 +6420,10 @@ async def ask( "Using fast standalone SQL generation path for query_id %s", query_id, ) - elif semantic_pipeline_active: - logger.info( - "fast_standalone_sql_generation_disabled query_id=%s reason=semantic_pipeline_active", - query_id, - ) - - if ( - semantic_pipeline_active - and not semantic_contract_available - and not api_results - ): - error_message = ( - "Semantic schema retrieval did not produce a complete active-schema " - "contract for this request. I cannot generate unrelated SQL." - ) - logger.info( - "semantic_pipeline_no_contract query_id=%s reason=%s", - query_id, - error_message, - ) if ( not self._is_stopped(query_id, self._ask_results) and not api_results - and (not semantic_pipeline_active or semantic_contract_available) and allow_sql_generation_reasoning ): self._ask_results[query_id] = AskResultResponse( @@ -6672,11 +6495,7 @@ async def ask( is_followup=True if histories else False, ) - if ( - not self._is_stopped(query_id, self._ask_results) - and not api_results - and (not semantic_pipeline_active or semantic_contract_available) - ): + if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="generating", type="TEXT_TO_SQL", @@ -6742,7 +6561,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, ), ) else: @@ -6762,14 +6580,12 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, ), ) except TimeoutError as generation_timeout: logger.warning( - "SQL generation timed out for query_id %s; semantic_pipeline_active=%s error=%s", + "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", query_id, - semantic_pipeline_active, generation_timeout, ) text_to_sql_generation_results = { @@ -6797,210 +6613,7 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - semantic_validation_failure_types = { - "SCHEMA_INTENT_VALIDATION", - "SCHEMA_VALIDATION", - "DRY_RUN", - } - semantic_retry_attempt = 0 - max_semantic_retries = 3 - while ( - semantic_pipeline_active - and failed_dry_run_result - and failed_dry_run_result.get("type") - in semantic_validation_failure_types - and semantic_retry_attempt < max_semantic_retries - and not api_results - ): - semantic_retry_attempt += 1 - retry_context = self._semantic_retry_context( - semantic_analysis=schema_intent_analysis, - invalid_generation_result=failed_dry_run_result, - retry_attempt=semantic_retry_attempt, - ) - logger.info( - "semantic_sql_retry_start query_id=%s retry_attempt=%s failure_type=%s rejected_schema_objects=%s failure_reason=%s", - query_id, - semantic_retry_attempt, - failed_dry_run_result.get("type"), - retry_context.get("rejected_schema_objects"), - failed_dry_run_result.get("error"), - ) - - retrieval_result = await self._run_with_timeout( - "Semantic schema retry retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=True, - semantic_retry_context=retry_context, - ), - timeout_seconds=self._schema_retrieval_timeout_seconds, - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) - semantic_pipeline_active = True - semantic_contract_available = self._semantic_analysis_has_contract( - schema_intent_analysis - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - has_calculated_field = _retrieval_result.get( - "has_calculated_field", False - ) - has_metric = _retrieval_result.get("has_metric", False) - has_json_field = _retrieval_result.get("has_json_field", False) - semantic_support_error = get_schema_intent_analysis_error( - schema_intent_analysis - ) - logger.info( - "semantic_sql_retry_contract query_id=%s retry_attempt=%s semantic_pipeline_active=%s semantic_contract_available=%s selected_schema_objects=%s support_error=%s", - query_id, - semantic_retry_attempt, - semantic_pipeline_active, - semantic_contract_available, - self._semantic_schema_objects(schema_intent_analysis), - semantic_support_error, - ) - if ( - not semantic_contract_available - or not documents - or semantic_support_error - ): - invalid_sql = failed_dry_run_result.get( - "sql", invalid_sql - ) - error_message = semantic_support_error or ( - "Semantic schema retrieval did not produce a complete active-schema contract." - ) - logger.info( - "semantic_sql_retry_rejected query_id=%s retry_attempt=%s reason=%s", - query_id, - semantic_retry_attempt, - error_message, - ) - continue - - if sql_generation_histories: - retry_generation_results = await self._run_with_timeout( - "Follow-up semantic SQL retry generation", - self._pipelines["followup_sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=sql_generation_histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), - ) - else: - retry_generation_results = await self._run_with_timeout( - "Semantic SQL retry generation", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_intent_analysis=schema_intent_analysis, - ), - ) - - retry_post_process = retry_generation_results.get( - "post_process", {} - ) - if retry_valid_result := retry_post_process.get( - "valid_generation_result" - ): - if ask_result := self._build_validated_ask_result_from_sql( - retry_valid_result.get("sql"), - table_ddls, - sql_user_query, - ): - logger.info( - "semantic_sql_retry_accepted query_id=%s retry_attempt=%s selected_schema_objects=%s", - query_id, - semantic_retry_attempt, - self._semantic_schema_objects( - schema_intent_analysis - ), - ) - api_results = [ask_result] - failed_dry_run_result = {} - break - invalid_sql = retry_valid_result.get("sql") - error_message = ( - "Semantic SQL retry generated SQL that failed active-schema validation." - ) - failed_dry_run_result = { - "type": "SCHEMA_INTENT_VALIDATION", - "sql": invalid_sql, - "original_sql": invalid_sql, - "error": error_message, - } - else: - failed_dry_run_result = retry_post_process.get( - "invalid_generation_result", {} - ) - invalid_sql = failed_dry_run_result.get( - "sql", invalid_sql - ) - error_message = failed_dry_run_result.get( - "error", error_message - ) - logger.info( - "semantic_sql_retry_failed query_id=%s retry_attempt=%s failure_type=%s error=%s", - query_id, - semantic_retry_attempt, - failed_dry_run_result.get("type"), - error_message, - ) - - if semantic_pipeline_active and not api_results: - current_sql_correction_retries = max_sql_correction_retries - if failed_dry_run_result: - invalid_sql = failed_dry_run_result.get( - "sql", invalid_sql - ) - error_message = failed_dry_run_result.get( - "error", error_message - ) - logger.info( - "semantic_pipeline_exhausted query_id=%s retry_attempts=%s error=%s invalid_sql=%s", - query_id, - semantic_retry_attempt, - error_message, - invalid_sql, - ) - - while ( - failed_dry_run_result - and not api_results - and current_sql_correction_retries < max_sql_correction_retries - ): + while current_sql_correction_retries < max_sql_correction_retries: if failed_dry_run_result["type"] in { "TIME_OUT", "UNSUPPORTED_SQL", @@ -7065,7 +6678,6 @@ async def ask( sql_functions=sql_functions, sql_knowledge=sql_knowledge, query=sql_user_query, - schema_intent_analysis=schema_intent_analysis, ), ) @@ -7108,10 +6720,8 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if not semantic_pipeline_active and ( - heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ) + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", @@ -7146,31 +6756,16 @@ async def ask( logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - if semantic_pipeline_active and error_message: - self._ask_results[query_id] = ( - self._build_failed_text_to_sql_response( - trace_id, - error_message, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - is_followup=True if histories else False, - code="NO_RELEVANT_SQL", - ) - ) - else: - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - is_followup=True if histories else False, - ) + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + is_followup=True if histories else False, ) + ) if error_message or invalid_sql: logger.info( "Suppressed technical SQL failure for query_id %s. " @@ -7179,13 +6774,9 @@ async def ask( error_message, invalid_sql, ) - results["metadata"]["error_type"] = ( - "NO_RELEVANT_SQL" if semantic_pipeline_active else "NO_RELEVANT_DATA" - ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" results["metadata"]["error_message"] = ( - error_message - if semantic_pipeline_active and error_message - else NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE ) results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index da7cce0d18..25044de18c 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -120,11 +120,8 @@ async def ask_feedback( instructions_task, ) = await asyncio.gather( self._pipelines["db_schema_retrieval"].run( - query=ask_feedback_request.question, tables=ask_feedback_request.tables, project_id=ask_feedback_request.project_id, - histories=[], - enable_column_pruning=True, ), self._pipelines["sql_pairs_retrieval"].run( query=ask_feedback_request.question, @@ -162,9 +159,6 @@ async def ask_feedback( ) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - schema_intent_analysis = _retrieval_result.get( - "semantic_analysis", {} - ) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] sql_samples = sql_samples_task["formatted_output"].get("documents", []) @@ -192,8 +186,6 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - query=ask_feedback_request.question, - schema_intent_analysis=schema_intent_analysis, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -256,8 +248,6 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - query=ask_feedback_request.question, - schema_intent_analysis=schema_intent_analysis, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 050509e49f..fa46c2ca7a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -7,14 +7,12 @@ find_invalid_table_references, get_json_field_instructions, get_metric_instructions, - get_schema_intent_analysis_error, normalize_data_source, normalize_generation_result_sql, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, get_sql_generation_system_prompt, get_text_to_sql_rules, - validate_sql_intent_alignment, ) @@ -262,72 +260,6 @@ def test_schema_validation_ignores_null_table_metadata(): ) == [] -def test_schema_intent_validation_rejects_generic_count_for_requested_metric(): - semantic_analysis = { - "analytical_intent": "ranking", - "metrics": ["invoice amount"], - "concept_mappings": [ - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["dbo_tblFactSales.invoice_amount"], - "required_in_sql": True, - } - ], - } - - error = validate_sql_intent_alignment( - "Show top customers by invoice amount", - 'SELECT COUNT(*) AS "RecordCount" FROM "dbo_tblFactSales"', - {"dbo_tblFactSales": ["invoice_amount"]}, - semantic_analysis=semantic_analysis, - ) - - assert "generic record count" in error - - -def test_schema_intent_validation_requires_ranking_shape(): - semantic_analysis = { - "analytical_intent": "ranking", - "ranking": ["top 10"], - "concept_mappings": [ - { - "request_concept": "invoice amount", - "concept_type": "metric", - "schema_objects": ["dbo_tblFactSales.invoice_amount"], - "required_in_sql": True, - } - ], - } - - error = validate_sql_intent_alignment( - "Show top 10 customers by invoice amount", - 'SELECT SUM("dbo_tblFactSales"."invoice_amount") AS "invoice_amount" ' - 'FROM "dbo_tblFactSales"', - {"dbo_tblFactSales": ["invoice_amount"]}, - semantic_analysis=semantic_analysis, - ) - - assert "sorting and limiting" in error - - -def test_schema_intent_analysis_reports_missing_required_mapping(): - semantic_analysis = { - "concept_mappings": [ - { - "request_concept": "customer", - "concept_type": "entity", - "schema_objects": [], - "required_in_sql": True, - } - ] - } - - assert "did not map required request concepts" in get_schema_intent_analysis_error( - semantic_analysis - ) - - def test_schema_validation_ignores_null_column_metadata(): assert find_invalid_column_references( 'SELECT "dbo_tblSales"."Market" FROM "dbo_tblSales"', diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index b8ac2ff769..5b6694bb45 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -5,7 +5,6 @@ _is_project_wide_analysis_query, dbschema_retrieval, expand_business_terms_for_retrieval, - rank_semantic_schema_candidates, ) @@ -31,71 +30,6 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query -def test_rank_semantic_schema_candidates_prefers_complete_business_concept_coverage(): - schemas = [ - { - "type": "TABLE", - "name": "dbo_invoices", - "columns": [ - {"name": "invoice_amount", "data_type": "decimal"}, - {"name": "customer_name", "data_type": "varchar"}, - ], - }, - { - "type": "TABLE", - "name": "dbo_invoice_ids", - "columns": [ - {"name": "invoice_id", "data_type": "varchar"}, - ], - }, - ] - - candidates = rank_semantic_schema_candidates( - "Show top customers by invoice amount", - schemas, - ) - - assert candidates[0]["table_name"] == "dbo_invoices" - assert "customer" in candidates[0]["matched_query_terms"] - assert "invoice" in candidates[0]["matched_query_terms"] - assert "amount" in candidates[0]["matched_query_terms"] - - -def test_rank_semantic_schema_candidates_uses_retry_rejections_as_negative_feedback(): - schemas = [ - { - "type": "TABLE", - "name": "dbo_invoices", - "columns": [ - {"name": "invoice_amount", "data_type": "decimal"}, - {"name": "customer_name", "data_type": "varchar"}, - ], - }, - { - "type": "TABLE", - "name": "dbo_invoice_summary", - "columns": [ - {"name": "total_invoice_amount", "data_type": "decimal"}, - {"name": "customer_name", "data_type": "varchar"}, - ], - }, - ] - - candidates = rank_semantic_schema_candidates( - "Show top customers by invoice amount", - schemas, - semantic_retry_context={ - "rejected_schema_objects": [ - "dbo_invoices", - "dbo_invoices.invoice_amount", - ] - }, - ) - - assert candidates[0]["table_name"] == "dbo_invoice_summary" - assert candidates[0]["rejected_by_retry"] is False - - @pytest.mark.asyncio async def test_dbschema_retrieval_loads_complete_active_project_schema(): class Retriever: From 05272b25852093bc9a6dd1a6c9366828857cafb7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 12 Jul 2026 14:05:54 +0530 Subject: [PATCH 0492/1087] Fix scoped schema retrieval for ask flow --- .../retrieval/db_schema_retrieval.py | 49 ++- wren-ai-service/src/web/v1/services/ask.py | 284 ++++++++++++++---- .../retrieval/test_db_schema_retrieval.py | 37 ++- .../textBasedAnswerBackgroundTracker.ts | 2 +- 4 files changed, 307 insertions(+), 65 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 107c5c3479..ad5e3d4c21 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -189,6 +189,32 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: return deduped +def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: + normalized: list[str] = [] + for table_name in table_names or []: + if not isinstance(table_name, str): + continue + table_name = table_name.strip() + if table_name and table_name not in normalized: + normalized.append(table_name) + return normalized + + +def _extract_table_names_from_table_retrieval( + table_retrieval: dict, explicit_tables: Optional[list[str]] = None +) -> list[str]: + table_names = _normalize_table_names(explicit_tables) + for document in table_retrieval.get("documents") or []: + if not isinstance(document, Document): + continue + table_name = document.meta.get("name") + if isinstance(table_name, str): + table_name = table_name.strip() + if table_name and table_name not in table_names: + table_names.append(table_name) + return table_names + + @observe(capture_input=False, capture_output=False) async def embedding( query: str, @@ -252,6 +278,10 @@ async def dbschema_retrieval( dbschema_retriever: Any, tables: Optional[list[str]] = None, ) -> list[Document]: + selected_table_names = _extract_table_names_from_table_retrieval( + table_retrieval, tables + ) + filters = { "operator": "AND", "conditions": [ @@ -263,10 +293,21 @@ async def dbschema_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - logger.info( - "Loading complete deployed schema metadata for active project_id %s", - project_id, - ) + if selected_table_names: + filters["conditions"].append( + {"field": "name", "operator": "in", "value": selected_table_names} + ) + logger.info( + "Loading selected deployed schema metadata for active project_id %s tables=%s", + project_id, + selected_table_names, + ) + else: + logger.info( + "Loading complete deployed schema metadata for active project_id %s", + project_id, + ) + results = await dbschema_retriever.run(query_embedding=[], filters=filters) return results.get("documents", []) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9db38bd5a8..3328e5ae1a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -45,6 +45,10 @@ class AskRequest(BaseRequest): use_dry_plan: bool = False allow_dry_plan_fallback: bool = True custom_instruction: Optional[str] = None + explicit_tables: Optional[list[str]] = Field( + default=None, + validation_alias=AliasChoices("explicit_tables", "explicitTables"), + ) class AskResponse(BaseModel): @@ -784,6 +788,75 @@ def _invalid_unqualified_sql_identifiers( return invalid + def _invalid_sql_output_aliases( + self, sql: str, schema_tables: list[dict[str, Any]] + ) -> list[str]: + valid_columns = { + str(column.get("name") or "").lower() + for table in schema_tables + for column in table.get("columns", []) + if column.get("name") + } + allowed_alias_terms = { + "average", + "avg", + "category", + "count", + "current", + "customer", + "date", + "failure", + "market", + "month", + "name", + "order", + "previous", + "product", + "rank", + "record", + "repair", + "revenue", + "sales", + "status", + "sum", + "time", + "total", + "value", + "workflow", + "year", + } + select_match = re.search( + r"\bSELECT\b(?P.*?)\bFROM\b", + normalized_sql, + flags=re.IGNORECASE | re.DOTALL, + ) + if not select_match: + return True + select_clause = select_match.group("select").strip() + if not re.match(r"\b(?:TOP\s+\d+\s+)?DISTINCT\b", select_clause, flags=re.IGNORECASE): + logger.warning( + "Ignoring SQL because a unique/no-duplicate request lacks DISTINCT or GROUP BY. " + "query=%s sql=%s", + query, + sql, + ) + return False + + distinct_clause = re.sub( + r"^(?:TOP\s+\d+\s+)?DISTINCT\s+", + "", + select_clause, + flags=re.IGNORECASE, + ) + projected_items = [ + item.strip() + for item in re.split(r",(?![^()]*\))", distinct_clause) + if item.strip() + ] + non_aggregate_items = [ + item + for item in projected_items + if not re.search(r"\b(?:count|sum|avg|min|max)\s*\(", item, flags=re.IGNORECASE) + ] + if len(non_aggregate_items) > 1: + logger.warning( + "Ignoring SQL because DISTINCT covers multiple non-aggregate columns and can still duplicate the requested entity. " + "query=%s sql=%s", + query, + sql, + ) + return False + return True + def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1011,6 +1149,10 @@ def _sql_matches_question_intent( referenced_table_tokens, ): return False + if not self._sql_uses_required_measure_aggregation(sql, query): + return False + if not self._sql_satisfies_unique_entity_request(sql, query): + return False if not expects_dimension: return True diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 449b9d3b38..9aa943b4a6 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -110,6 +110,46 @@ def test_select_relevant_table_documents_limits_weak_extra_candidates(): assert "staging_audit" not in [document.meta["name"] for document in selected] +def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): + transaction_source = Document( + content="Invoice transaction fact rows with customer id and invoice amount.", + meta={"type": "TABLE_DESCRIPTION", "name": "invoice_fact"}, + score=0.95, + ) + reference_source = Document( + content="Customer master reference directory with customer names and accounts.", + meta={"type": "TABLE_DESCRIPTION", "name": "customer_master"}, + score=0.7, + ) + + documents = _rerank_table_documents( + "List customer names without duplicates.", + [transaction_source, reference_source], + ) + + assert documents[0].meta["name"] == "customer_master" + + +def test_rerank_table_documents_prefers_transaction_source_for_metric_question(): + reference_source = Document( + content="Product catalog reference table with names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "product_master"}, + score=0.95, + ) + transaction_source = Document( + content="Sales transaction fact table with product, amount, and revenue.", + meta={"type": "TABLE_DESCRIPTION", "name": "sales_fact"}, + score=0.7, + ) + + documents = _rerank_table_documents( + "Show total sales amount by product.", + [reference_source, transaction_source], + ) + + assert documents[0].meta["name"] == "sales_fact" + + @pytest.mark.asyncio async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 96bbc48534..f135b2b5f6 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -57,6 +57,88 @@ def test_direct_heuristic_gate_is_disabled_for_generic_schema_selection(): ) +def test_rewrite_query_for_text_to_sql_guides_total_amount_to_sum(): + service = AskService.__new__(AskService) + + rewritten = service._rewrite_query_for_text_to_sql( + "Show total invoice amount by currency." + ) + + assert "aggregate an exposed numeric measure with SUM" in rewritten + assert "use COUNT only for record-count questions" in rewritten + + +def test_validated_sql_rejects_count_for_total_amount_question(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_Invoices"."Currency" AS "Currency", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_Invoices" ' + 'GROUP BY "dbo_Invoices"."Currency"' + ), + [ + """ + CREATE TABLE dbo_Invoices ( + Currency VARCHAR, + InvoiceAmount DOUBLE + ); + """ + ], + "Show total invoice amount by currency.", + ) + + assert result is None + + +def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT DISTINCT "dbo_Customers"."CustomerName" AS "CustomerName", ' + '"dbo_Customers"."CustomerId" AS "CustomerId" ' + 'FROM "dbo_Customers"' + ), + [ + """ + CREATE TABLE dbo_Customers ( + CustomerName VARCHAR, + CustomerId VARCHAR + ); + """ + ], + "List customer names with no duplicates.", + ) + + assert result is None + + +def test_validated_sql_rejects_grouped_extra_columns_for_no_duplicates(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_Customers"."CustomerName" AS "CustomerName", ' + '"dbo_Customers"."CustomerId" AS "CustomerId" ' + 'FROM "dbo_Customers" ' + 'GROUP BY "dbo_Customers"."CustomerName", "dbo_Customers"."CustomerId"' + ), + [ + """ + CREATE TABLE dbo_Customers ( + CustomerName VARCHAR, + CustomerId VARCHAR + ); + """ + ], + "List customer names with no duplicates.", + ) + + assert result is None + + def test_build_direct_orders_sales_sql_for_top_new_orders_q1(): service = AskService.__new__(AskService) sql = service._build_direct_orders_sales_sql( From 4fb3c4e822e52fd695edfc396cca829872fab230 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 12 Jul 2026 23:01:59 +0530 Subject: [PATCH 0501/1087] Prevent full schema fallback for data retrieval timeouts --- wren-ai-service/src/web/v1/services/ask.py | 59 ++++++++++++------- .../pytest/services/test_ask_sales_sql.py | 8 +++ 2 files changed, 47 insertions(+), 20 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9c8a5c6b38..6b8a6b67e0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4656,6 +4656,11 @@ async def _run_with_timeout( except TimeoutError as exc: raise TimeoutError(f"{label} timed out after {timeout} seconds") from exc + def _should_retry_selected_schema_after_retrieval_timeout( + self, retrieval_table_names: Optional[list[str]] + ) -> bool: + return bool(retrieval_table_names) + def _build_greeting_response(self, query: str) -> str: return ( f"Hi. I can help with questions about your active datasource and Wren AI.\n\n" @@ -6271,26 +6276,40 @@ async def ask( timeout_seconds=self._schema_retrieval_timeout_seconds, ) except TimeoutError as error: - logger.warning( - "Schema retrieval timed out; falling back to deployed schemas. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - error, - ) - retrieval_result = await self._run_with_timeout( - "Deployed schema fallback retrieval", - self._pipelines["db_schema_retrieval"].run( - query="" if not retrieval_table_names else sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 30, - ), - ) + if not self._should_retry_selected_schema_after_retrieval_timeout( + retrieval_table_names + ): + logger.warning( + "Schema retrieval timed out for data query; not loading full project schema. " + "query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + error, + ) + retrieval_result = {"construct_retrieval_results": {}} + else: + logger.warning( + "Schema retrieval timed out; retrying only explicit selected schemas. " + "query_id=%s project_id=%s tables=%s error=%s", + query_id, + ask_request.project_id, + retrieval_table_names, + error, + ) + retrieval_result = await self._run_with_timeout( + "Selected schema fallback retrieval", + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 30, + ), + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index f135b2b5f6..6ca19324be 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -57,6 +57,14 @@ def test_direct_heuristic_gate_is_disabled_for_generic_schema_selection(): ) +def test_data_query_timeout_retry_does_not_allow_full_project_schema(): + service = AskService.__new__(AskService) + + assert not service._should_retry_selected_schema_after_retrieval_timeout(None) + assert not service._should_retry_selected_schema_after_retrieval_timeout([]) + assert service._should_retry_selected_schema_after_retrieval_timeout(["orders"]) + + def test_rewrite_query_for_text_to_sql_guides_total_amount_to_sum(): service = AskService.__new__(AskService) From d64be7e7504862de864b545a37d76c74d890a10b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 12 Jul 2026 23:43:35 +0530 Subject: [PATCH 0502/1087] Tighten scoped retrieval and metric intent handling --- .../retrieval/db_schema_retrieval.py | 16 +++- wren-ai-service/src/web/v1/services/ask.py | 78 ++++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 52 +++++++++++++ .../pytest/services/test_ask_sales_sql.py | 61 +++++++++++++++ 4 files changed, 203 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ae4ae10289..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -146,16 +146,23 @@ def expand_business_terms_for_retrieval(query: str) -> str: "markets", "order", "orders", + "product", + "products", + "category", + "categories", + "quantity", + "qty", "region", "regions", "sales", "salesperson", "sales person", + "sold", "value", ) ): expansions.append( - "transaction purchase billing account geography area representative amount value total metric money exchange currency" + "transaction purchase billing account geography area representative product item category sku quantity units sold amount value total metric money exchange currency" ) if any( @@ -546,6 +553,7 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query + query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -573,10 +581,14 @@ async def table_retrieval( ) if embedding: - return await table_retriever.run( + results = await table_retriever.run( query_embedding=embedding.get("embedding"), filters=base_filters, ) + results["documents"] = _select_relevant_table_documents( + query, results.get("documents") or [] + ) + return results if tables: logger.info("Loading explicit table descriptions: %s", tables) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6b8a6b67e0..13548f9620 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -23,6 +23,7 @@ NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( "No relevant data found in the active datasource for this question." ) +MAX_FORCED_EXPLICIT_TABLES = 5 async def _return_value(value): @@ -2185,11 +2186,16 @@ def _build_schema_grounded_analytics_sql( "market", "markets", "performance", + "product", + "products", + "quantity", + "qty", "revenue", "sale", "sales", "salesperson", "sales person", + "sold", ) ) if not is_sales_or_order_query: @@ -2220,7 +2226,20 @@ def _build_schema_grounded_analytics_sql( for term in ("count", "counts", "volume", "how many", "distribution") ) and not any( term in normalized_query - for term in ("revenue", "sales value", "amount", "value", "quantity", "qty") + for term in ( + "amount", + "cost", + "expense", + "quantity", + "qty", + "revenue", + "sale", + "sales", + "sold", + "sum", + "total", + "value", + ) ) wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") @@ -2313,6 +2332,32 @@ def _build_schema_grounded_analytics_sql( if "division" in normalized_query: dimension_candidates.append(("Division",)) if ( + ( + "category" in normalized_query + or "categories" in normalized_query + or "prodcategory" in compact_query + or "productcategory" in compact_query + ) + and "product" in normalized_query + and "product type" not in normalized_query + and "prodtype" not in compact_query + and "producttype" not in compact_query + ): + dimension_candidates.append( + ( + "ProductCategory", + "Product Category", + "ProdCategory", + "Category", + "ProductType", + "Product Type", + "ProdType", + "ProdName", + "Product", + "ProductName", + ) + ) + elif ( "product type" in normalized_query or "prodtype" in normalized_query or "producttype" in compact_query @@ -2357,15 +2402,25 @@ def _build_schema_grounded_analytics_sql( measure_candidates = ( "Qty", "Quantity", + "QtySold", + "SoldQty", + "QuantitySold", + "UnitsSold", + "ItemQty", "SalesQty", "OrderQty", + "OrderQuantity", "InvoiceQty", + "InvoiceQuantity", ) if any(term in normalized_query for term in ("quantity", "qty")) else ( + "Sales", "SalesValue", "FXSalesValue", "Revenue", "NetSales", + "TotalSales", "SalesAmount", + "SaleAmount", "NewOrderValue", "NewOrdersValue", "InvoiceValue", @@ -4661,6 +4716,21 @@ def _should_retry_selected_schema_after_retrieval_timeout( ) -> bool: return bool(retrieval_table_names) + def _forced_explicit_table_names( + self, table_names: list[str], *, source: str = "request" + ) -> list[str]: + if not table_names: + return [] + if len(table_names) <= MAX_FORCED_EXPLICIT_TABLES: + return table_names + + logger.info( + "Treating broad %s explicit_tables list as retrieval candidates, not a forced schema scope: %s", + source, + table_names, + ) + return [] + def _build_greeting_response(self, query: str) -> str: return ( f"Hi. I can help with questions about your active datasource and Wren AI.\n\n" @@ -5614,8 +5684,12 @@ async def ask( query_explicit_table_names = self._normalize_explicit_table_names( self._extract_explicit_table_names_from_query(user_query) ) + forced_request_explicit_table_names = self._forced_explicit_table_names( + request_explicit_table_names, + source="request", + ) explicit_table_names = ( - request_explicit_table_names or query_explicit_table_names + forced_request_explicit_table_names or query_explicit_table_names ) retrieval_table_names = explicit_table_names or None diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 9aa943b4a6..3ee705b621 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -110,6 +110,58 @@ def test_select_relevant_table_documents_limits_weak_extra_candidates(): assert "staging_audit" not in [document.meta["name"] for document in selected] +@pytest.mark.asyncio +async def test_table_retrieval_caps_embedding_results_before_schema_loading(): + documents = [ + Document( + content="Raw staging audit rows with load metadata.", + meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, + score=0.99, + ), + Document( + content="Invoice sales transactions with product categories and sales value.", + meta={"type": "TABLE_DESCRIPTION", "name": "sales_invoices"}, + score=0.8, + ), + Document( + content="Product catalog with product names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "products"}, + score=0.7, + ), + Document( + content="Customer account master data.", + meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, + score=0.6, + ), + Document( + content="Sales regions and market hierarchy.", + meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, + score=0.5, + ), + Document( + content="Exchange rate lookup by currency.", + meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, + score=0.4, + ), + ] + + class Retriever: + async def run(self, query_embedding, filters): + return {"documents": documents} + + result = await table_retrieval( + query="What is the distribution of sales across product categories?", + embedding={"embedding": [0.1, 0.2]}, + project_id="project-1", + tables=[], + table_retriever=Retriever(), + ) + + selected_names = [document.meta["name"] for document in result["documents"]] + assert 1 <= len(selected_names) <= 5 + assert "staging_audit" not in selected_names + + def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): transaction_source = Document( content="Invoice transaction fact rows with customer id and invoice amount.", diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 6ca19324be..cb02fdf47f 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -30,6 +30,67 @@ def test_build_schema_grounded_sales_sql_for_salesperson_performance(): assert "CustID" not in sql +def test_build_schema_grounded_sales_sql_for_sales_by_product_category(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "What is the distribution of sales across product categories?", + [ + """ + CREATE TABLE dbo_qSalesMargin ( + ProductCategory VARCHAR, + ProdName VARCHAR, + SalesValue DOUBLE, + OrdNo VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_qSalesMargin"."ProductCategory" AS "ProductCategory", ' + 'SUM("dbo_qSalesMargin"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_qSalesMargin" ' + 'WHERE "dbo_qSalesMargin"."ProductCategory" IS NOT NULL ' + 'GROUP BY "dbo_qSalesMargin"."ProductCategory" ' + 'ORDER BY SUM("dbo_qSalesMargin"."SalesValue") DESC' + ) + assert "COUNT" not in sql + + +def test_build_schema_grounded_sales_sql_for_total_quantity_sold_by_product(): + service = AskService.__new__(AskService) + sql = service._build_schema_grounded_sales_sql( + "Show total quantity sold by product.", + [ + """ + CREATE TABLE dbo_tblOrderLines ( + ProductName VARCHAR, + Quantity DOUBLE, + OrderNo VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblOrderLines"."ProductName" AS "ProductName", ' + 'SUM("dbo_tblOrderLines"."Quantity") AS "TotalQuantity" ' + 'FROM "dbo_tblOrderLines" ' + 'WHERE "dbo_tblOrderLines"."ProductName" IS NOT NULL ' + 'GROUP BY "dbo_tblOrderLines"."ProductName" ' + 'ORDER BY SUM("dbo_tblOrderLines"."Quantity") DESC' + ) + assert "COUNT" not in sql + + +def test_broad_request_explicit_tables_are_not_forced_schema_scope(): + service = AskService.__new__(AskService) + table_names = [f"table_{index}" for index in range(6)] + + assert service._forced_explicit_table_names(table_names) == [] + assert service._forced_explicit_table_names(table_names[:5]) == table_names[:5] + + def test_build_direct_orders_sales_sql_for_salesperson_order_count(): service = AskService.__new__(AskService) sql = service._build_direct_orders_sales_sql( From 0d91bf28c6fd8be6825f7e34c4bbf6ad23befb81 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 00:28:36 +0530 Subject: [PATCH 0503/1087] Reject detail rows for grouped count questions --- wren-ai-service/src/web/v1/services/ask.py | 68 +++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 50 ++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 13548f9620..f16f5405ff 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -768,6 +768,72 @@ def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> return False return True + def _sql_satisfies_count_ranking_request( + self, sql: str, query: str | None + ) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return True + + asks_for_count_metric = any( + term in normalized_query + for term in ( + "count", + "counts", + "how many", + "number of", + "record count", + ) + ) or ( + any(term in normalized_query for term in ("top", "most", "highest")) + and any( + term in normalized_query + for term in ("order", "orders", "record", "records", "row", "rows") + ) + ) + if not asks_for_count_metric: + return True + + asks_for_grouped_entity = any( + term in normalized_query + for term in ( + "category", + "customer", + "customers", + "currency", + "market", + "product", + "products", + "region", + "sales person", + "salesperson", + "source", + "status", + "type", + ) + ) + if not asks_for_grouped_entity: + return True + + normalized_sql = re.sub(r"\s+", " ", sql or "").lower() + if not re.search(r"\bcount\s*\(", normalized_sql): + logger.warning( + "Ignoring SQL because a count/ranking question was answered with detail rows. " + "query=%s sql=%s", + query, + sql, + ) + return False + if not re.search(r"\bgroup\s+by\b", normalized_sql): + logger.warning( + "Ignoring SQL because a grouped count/ranking question has no GROUP BY. " + "query=%s sql=%s", + query, + sql, + ) + return False + return True + def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> bool: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not re.search( @@ -1152,6 +1218,8 @@ def _sql_matches_question_intent( return False if not self._sql_uses_required_measure_aggregation(sql, query): return False + if not self._sql_satisfies_count_ranking_request(sql, query): + return False if not self._sql_satisfies_unique_entity_request(sql, query): return False diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index cb02fdf47f..9bd13569ac 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -161,6 +161,56 @@ def test_validated_sql_rejects_count_for_total_amount_question(): assert result is None +def test_validated_sql_rejects_detail_rows_for_customer_order_count_question(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT TOP 5 "dbo_tblNewOrders"."CustName" AS "CustName", ' + '"dbo_tblNewOrders"."OrdNo" AS "OrdNo" ' + 'FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."CustName" IS NOT NULL' + ), + [ + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + "From dbo_tblNewOrders, show the top 5 customers by order count using CustName.", + ) + + assert result is None + + +def test_schema_grounded_table_question_groups_top_customers_by_order_count(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_table_question_sql( + "From dbo_tblNewOrders, show the top 5 customers by order count using CustName.", + [ + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 5 "dbo_tblNewOrders"."CustName" AS "CustName", ' + 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "RecordCount" ' + 'FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."CustName" IS NOT NULL ' + 'AND LTRIM(RTRIM("dbo_tblNewOrders"."CustName")) <> \'\' ' + 'GROUP BY "dbo_tblNewOrders"."CustName" ' + 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' + ) + + def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) From 1c7e3d0f88f2d38db6cfff976e9d3118c21f0db6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 00:36:54 +0530 Subject: [PATCH 0504/1087] Restrict full schema retrieval to metadata questions --- wren-ai-service/src/web/v1/services/ask.py | 35 ++++--------------- .../pytest/services/test_ask_sales_sql.py | 24 +++++++++++++ 2 files changed, 31 insertions(+), 28 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f16f5405ff..4598adc0e9 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -4986,6 +4986,9 @@ def _get_metadata_question_kind(self, query: str) -> str | None: return None + def _should_load_full_schema_for_question(self, query: str | None) -> bool: + return bool(self._get_metadata_question_kind(query or "")) + def _find_metadata_table_matches( self, query: str, tables: list[dict[str, Any]] ) -> list[dict[str, Any]]: @@ -5877,36 +5880,12 @@ async def ask( explicit_table_names, ) ) - if not documents and not request_explicit_table_names: + if not documents: logger.info( "Explicit table retrieval did not return requested active-schema table; " - "loading full active schema. query_id=%s", + "not loading full active schema for data question. query_id=%s", query_id, ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval for explicit table", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - all_documents, _, _ = self._extract_retrieval_metadata( - retrieval_result - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - all_documents, - explicit_table_names, - ) - ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) @@ -6502,7 +6481,7 @@ async def ask( ) if ( not documents - and self._get_metadata_question_kind(user_query) + and self._should_load_full_schema_for_question(user_query) and not request_explicit_table_names ): logger.info( @@ -6691,7 +6670,7 @@ async def ask( should_retry_full_schema = ( not api_results - and self._get_metadata_question_kind(user_query) + and self._should_load_full_schema_for_question(user_query) and "db_schema_retrieval" in self._pipelines and not request_explicit_table_names and not table_names diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 9bd13569ac..fdd9d93f76 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -91,6 +91,30 @@ def test_broad_request_explicit_tables_are_not_forced_schema_scope(): assert service._forced_explicit_table_names(table_names[:5]) == table_names[:5] +def test_full_schema_loading_gate_allows_only_metadata_questions(): + service = AskService.__new__(AskService) + + assert service._should_load_full_schema_for_question( + "List all tables in this datasource." + ) + assert service._should_load_full_schema_for_question( + "How many deployed models are available?" + ) + assert service._should_load_full_schema_for_question( + "Show the schema metadata." + ) + + assert not service._should_load_full_schema_for_question( + "Show top 5 customers by order count." + ) + assert not service._should_load_full_schema_for_question( + "From dbo_tblNewOrders, show the top 5 customers by order count using CustName." + ) + assert not service._should_load_full_schema_for_question( + "What is the distribution of sales across product categories?" + ) + + def test_build_direct_orders_sales_sql_for_salesperson_order_count(): service = AskService.__new__(AskService) sql = service._build_direct_orders_sales_sql( From 97075b46d03a67b507bd8800624e13a2a25790fe Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 02:05:07 +0530 Subject: [PATCH 0505/1087] Require retrieved schema for SQL fallbacks --- wren-ai-service/src/web/v1/services/ask.py | 70 ++++++------------- .../src/web/v1/services/sql_answer.py | 7 +- .../pytest/services/test_ask_sales_sql.py | 34 +++++++++ .../tests/pytest/services/test_sql_answer.py | 61 ++++++++++++++++ 4 files changed, 119 insertions(+), 53 deletions(-) create mode 100644 wren-ai-service/tests/pytest/services/test_sql_answer.py diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 4598adc0e9..0223f584e2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -3986,6 +3986,18 @@ def _build_heuristic_text_to_sql_fallback( return None + def _can_use_schema_grounded_sql_fallback( + self, + documents: list[dict], + table_ddls: list[str], + query: str | None, + ) -> bool: + return bool( + documents + and table_ddls + and not self._should_load_full_schema_for_question(query) + ) + def _is_schema_grounded_query( self, query: str, db_schemas: Optional[list[str]] = None ) -> bool: @@ -6783,54 +6795,6 @@ async def ask( return results if not documents: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - is_followup=True if histories else False, - ) - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( @@ -7183,8 +7147,14 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names + if self._can_use_schema_grounded_sql_fallback( + documents, + table_ddls, + user_query, + ) and ( + heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ) ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index 7442f2f884..42eefc478a 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -66,14 +66,14 @@ def __init__( ) async def _load_active_schema_contexts( - self, project_id: Optional[str] + self, project_id: Optional[str], query: str ) -> list[str]: retrieval_pipeline = self._pipelines.get("db_schema_retrieval") - if not retrieval_pipeline: + if not retrieval_pipeline or not (query or "").strip(): return [] retrieval_result = await retrieval_pipeline.run( - query="", + query=query, histories=[], project_id=project_id, enable_column_pruning=False, @@ -134,6 +134,7 @@ async def sql_answer( schema_contexts = await self._load_active_schema_contexts( sql_answer_request.project_id, + sql_answer_request.query, ) normalized_sql = self._normalize_and_validate_sql( sql_answer_request.sql, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index fdd9d93f76..a64038fbb9 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -115,6 +115,40 @@ def test_full_schema_loading_gate_allows_only_metadata_questions(): ) +def test_schema_grounded_sql_fallback_requires_retrieved_documents(): + service = AskService.__new__(AskService) + + assert not service._can_use_schema_grounded_sql_fallback( + [], + [ + """ + CREATE TABLE dbo_orders ( + CustomerName VARCHAR, + OrderId VARCHAR + ); + """ + ], + "Show top customers by order count.", + ) + assert not service._can_use_schema_grounded_sql_fallback( + [{"table_name": "dbo_orders"}], + [], + "Show top customers by order count.", + ) + assert service._can_use_schema_grounded_sql_fallback( + [{"table_name": "dbo_orders"}], + [ + """ + CREATE TABLE dbo_orders ( + CustomerName VARCHAR, + OrderId VARCHAR + ); + """ + ], + "Show top customers by order count.", + ) + + def test_build_direct_orders_sales_sql_for_salesperson_order_count(): service = AskService.__new__(AskService) sql = service._build_direct_orders_sales_sql( diff --git a/wren-ai-service/tests/pytest/services/test_sql_answer.py b/wren-ai-service/tests/pytest/services/test_sql_answer.py new file mode 100644 index 0000000000..0d331122d0 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_sql_answer.py @@ -0,0 +1,61 @@ +import pytest + +from src.web.v1.services.sql_answer import SqlAnswerService + + +class _FakeSchemaRetrievalPipeline: + def __init__(self): + self.calls = [] + + async def run(self, **kwargs): + self.calls.append(kwargs) + return { + "construct_retrieval_results": { + "retrieval_results": [ + { + "table_name": "dbo_orders", + "table_ddl": ( + "CREATE TABLE dbo_orders (" + "CustomerName VARCHAR, OrderId VARCHAR)" + ), + } + ] + } + } + + +@pytest.mark.asyncio +async def test_sql_answer_loads_schema_context_with_user_query(): + retrieval = _FakeSchemaRetrievalPipeline() + service = SqlAnswerService({"db_schema_retrieval": retrieval}) + + contexts = await service._load_active_schema_contexts( + project_id="project-1", + query="Show top customers by order count.", + ) + + assert contexts == [ + "CREATE TABLE dbo_orders (CustomerName VARCHAR, OrderId VARCHAR)" + ] + assert retrieval.calls == [ + { + "query": "Show top customers by order count.", + "histories": [], + "project_id": "project-1", + "enable_column_pruning": False, + } + ] + + +@pytest.mark.asyncio +async def test_sql_answer_does_not_load_full_schema_without_query(): + retrieval = _FakeSchemaRetrievalPipeline() + service = SqlAnswerService({"db_schema_retrieval": retrieval}) + + contexts = await service._load_active_schema_contexts( + project_id="project-1", + query="", + ) + + assert contexts == [] + assert retrieval.calls == [] From dc336f23c01823f4d4a4038a98e65786f0a89b3f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 11:57:34 +0530 Subject: [PATCH 0506/1087] Match explicit tables across schema name forms --- wren-ai-service/src/web/v1/services/ask.py | 11 ++++ .../pytest/services/test_ask_sales_sql.py | 51 +++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0223f584e2..389d0e37fb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1862,7 +1862,18 @@ def _explicit_table_name_candidates(self, table_name: str) -> list[str]: separator_normalized = re.sub(r"[.$]", "_", table_name) if separator_normalized not in candidates: candidates.append(separator_normalized) + if "_" in table_name: + dotted_schema_name = re.sub( + r"^([A-Za-z_][A-Za-z0-9]*)_", + r"\1.", + table_name, + count=1, + ) + if dotted_schema_name not in candidates: + candidates.append(dotted_schema_name) short_name = re.split(r"[.$]", table_name)[-1] + if short_name == table_name and "_" in table_name: + short_name = table_name.split("_", 1)[-1] if short_name and short_name not in candidates: candidates.append(short_name) return candidates diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index a64038fbb9..01bafa19d8 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -446,6 +446,21 @@ def test_extract_explicit_table_names_from_pcb_repair_phrases(): ) == ["ticket_labels", "dbo_ticket_labels"] +def test_explicit_table_name_candidates_include_dotted_and_short_forms(): + service = AskService.__new__(AskService) + + assert service._explicit_table_name_candidates("dbo_tblNewOrders") == [ + "dbo_tblNewOrders", + "dbo.tblNewOrders", + "tblNewOrders", + ] + assert service._explicit_table_name_candidates("dbo.tblNewOrders") == [ + "dbo.tblNewOrders", + "dbo_tblNewOrders", + "tblNewOrders", + ] + + def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): service = AskService.__new__(AskService) documents = [ @@ -482,6 +497,42 @@ def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): assert table_ddls == [documents[1]["table_ddl"]] +def test_filter_retrieval_metadata_for_explicit_query_matches_dotted_table_name(): + service = AskService.__new__(AskService) + documents = [ + { + "table_name": "dbo.tblNewOrders", + "table_ddl": """ + CREATE TABLE "dbo.tblNewOrders" ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + }, + { + "table_name": "dbo_other", + "table_ddl": """ + CREATE TABLE dbo_other ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + }, + ] + + filtered_documents, table_names, table_ddls = ( + service._filter_retrieval_metadata_for_explicit_query( + "Show the top 5 CustName values from dbo_tblNewOrders by number of orders.", + documents, + ["dbo_tblNewOrders"], + ) + ) + + assert filtered_documents == [documents[0]] + assert table_names == ["dbo.tblNewOrders"] + assert table_ddls == [documents[0]["table_ddl"]] + + def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): service = AskService.__new__(AskService) result = service._build_validated_ask_result_from_sql( From 8cb4ee4e228f7093aa7d7ded20551f81ab847d98 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 12:46:01 +0530 Subject: [PATCH 0507/1087] Tighten retrieval ranking and intent validation --- .../retrieval/db_schema_retrieval.py | 33 ++- wren-ai-service/src/web/v1/services/ask.py | 234 +++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 20 ++ .../pytest/services/test_ask_sales_sql.py | 109 ++++++++ 4 files changed, 374 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..6b64c9b452 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -213,11 +213,15 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms = { - _normalize_retrieval_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") - if len(token) > 2 and token.lower() not in stop_words - } + terms: set[str] = set() + for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): + split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) + for token in re.findall(r"[A-Za-z0-9]+", split_token): + if len(token) <= 2 or token.lower() in stop_words: + continue + normalized_token = _normalize_retrieval_token(token) + if normalized_token: + terms.add(normalized_token) return {term for term in terms if term} @@ -243,7 +247,11 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - non_production_terms = ( + weak_non_production_terms = ( + "stage", + "staging", + ) + strong_non_production_terms = ( "archive", "backup", "copy", @@ -251,17 +259,18 @@ def _source_shape_score(query: str, document: Document) -> int: "development", "duplicate", "sample", - "stage", - "staging", "temp", "test", "tmp", ) - if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( - normalized_query, - non_production_terms, + if source_terms & set(strong_non_production_terms) and not _query_mentions_any( + normalized_query, strong_non_production_terms + ): + score -= 240 + if source_terms & set(weak_non_production_terms) and not _query_mentions_any( + normalized_query, weak_non_production_terms ): - score -= 60 + score -= 40 aggregation_terms = ( "amount", diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 389d0e37fb..ce42c81888 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -689,6 +689,8 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: concept_groups.append({"market", "region", "country", "territory"}) if "region" in normalized or "regions" in normalized: concept_groups.append({"region", "market", "area", "territory", "country"}) + if "country" in normalized or "countries" in normalized: + concept_groups.append({"country", "countries", "nation", "destination"}) if "quarterly" in normalized or "quarter" in normalized: concept_groups.append({"quarter", "quarterly"}) if "recurring" in normalized or "recurrence" in normalized: @@ -909,6 +911,123 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True + def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) + if not normalized_query: + return None + + match = re.search( + r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" + r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", + normalized_query, + flags=re.IGNORECASE, + ) + if not match: + return None + + phrase = match.group("phrase").strip(" .,;:()[]{}'\"") + phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) + if not phrase or len(phrase) < 3: + return None + if re.search( + r"\b(?:table|model|schema|column|columns|market|region|country|division|" + r"date|month|year|quarter|top|count|number|amount|value)\b", + phrase, + flags=re.IGNORECASE, + ): + return None + return phrase + + def _preferred_entity_lookup_columns( + self, query: str | None, table: dict[str, Any] + ) -> set[str]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + candidate_groups: list[tuple[str, ...]] = [] + if "account" in normalized_query: + candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) + if "company" in normalized_query: + candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) + candidate_groups.extend( + [ + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + ), + ("Client", "ClientName"), + ("Account", "AccountName"), + ("Company", "CompanyName"), + ("Name",), + ] + ) + + columns: set[str] = set() + for candidates in candidate_groups: + column = self._find_schema_column(table, candidates) + if column: + columns.add(column) + return columns + + def _sql_satisfies_entity_lookup_request( + self, + sql: str, + query: str | None, + referenced_tables: list[str], + referenced_columns_by_table: dict[str, set[str]], + valid_tables: dict[str, dict[str, Any]], + ) -> bool: + if not self._extract_entity_lookup_phrase(query): + return True + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if any( + term in normalized_query + for term in (" by ", " per ", " each ", "distribution", "top", "count") + ): + return True + + for table_reference in referenced_tables: + table = self._table_for_sql_reference(table_reference, valid_tables) + if not table: + continue + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + + table_key = str(table_reference or "").lower() + referenced_columns = referenced_columns_by_table.get( + table_key + ) or referenced_columns_by_table.get( + table_key.split(".")[-1], + set(), + ) + referenced_column_keys = { + self._normalize_schema_identifier_key(column) + for column in referenced_columns + } + preferred_column_keys = { + self._normalize_schema_identifier_key(column) + for column in preferred_columns + } + if referenced_column_keys & preferred_column_keys: + return True + + logger.warning( + "Ignoring SQL because entity lookup did not use available customer/name columns. " + "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", + query, + table.get("name"), + sorted(preferred_columns), + sorted(referenced_columns), + sql, + ) + return False + + return True + def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1222,6 +1341,14 @@ def _sql_matches_question_intent( return False if not self._sql_satisfies_unique_entity_request(sql, query): return False + if not self._sql_satisfies_entity_lookup_request( + sql, + query, + referenced_tables, + referenced_columns_by_table, + valid_tables, + ): + return False if not expects_dimension: return True @@ -2226,6 +2353,88 @@ def _select_best_analytics_table( date_column, ) + def _build_entity_lookup_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + lookup_phrase = self._extract_entity_lookup_phrase(query) + if not lookup_phrase: + return None + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + asks_for_orders = any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ) + + scored: list[tuple[int, dict[str, Any], str]] = [] + for table in tables: + table_name = str(table.get("name") or "") + if not table_name: + continue + + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + + preferred_column = sorted( + preferred_columns, + key=lambda column: ( + 0 + if self._normalize_schema_identifier_key(column) + in {"custname", "customername", "customer"} + else 1, + column.lower(), + ), + )[0] + + score = 20 + normalized_table = self._normalize_schema_token(table_name) + if asks_for_orders: + if "order" in normalized_table: + score += 40 + if "neworder" in normalized_table: + score += 20 + if self._find_schema_column( + table, ("OrdNo", "OrderNo", "OrderId", "NewOrderId") + ): + score += 25 + if "test" in normalized_table or "tmp" in normalized_table: + score -= 80 + if "dev" in normalized_table or "backup" in normalized_table: + score -= 60 + if "stage" in normalized_table: + score -= 10 + scored.append((score, table, preferred_column)) + + if not scored: + return None + + _score, table, filter_column = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + table_name = str(table.get("name") or "") + if not table_name: + return None + + table_ref = self._quote_sql_identifier(table_name) + filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" + escaped_phrase = lookup_phrase.replace("'", "''") + date_column = self._find_schema_column( + table, + ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), + temporal=True, + ) + order_clause = ( + f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" + if date_column + else "" + ) + return ( + f"SELECT TOP 500 * FROM {table_ref} " + f"WHERE {filter_ref} = '{escaped_phrase}'" + f"{order_clause}" + ) + def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2239,6 +2448,9 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): + return entity_lookup_sql + if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2294,16 +2506,7 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - if not is_sales_or_order_query: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not any( + asks_for_measure_value = any( term in normalized_query for term in ( "amount", @@ -2320,6 +2523,17 @@ def _build_schema_grounded_analytics_sql( "value", ) ) + + if not is_sales_or_order_query and not asks_for_measure_value: + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + + wants_count_metric = any( + term in normalized_query + for term in ("count", "counts", "volume", "how many", "distribution") + ) and not asks_for_measure_value wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 3ee705b621..bbb0c79aef 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -67,6 +67,26 @@ def test_rerank_table_documents_prefers_question_relevant_table_text(): assert documents[0].meta["name"] == "business_transactions" +def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): + test_load = Document( + content="Raw test load rows for order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ) + order_market_table = Document( + content="New order transaction records with market and customer fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.45, + ) + + documents = _rerank_table_documents( + "Show order distribution across markets.", + [test_load, order_market_table], + ) + + assert documents[0].meta["name"] == "dbo_xStageNewOrders" + + def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 01bafa19d8..5ba8ca74b8 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,6 +269,115 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) +def test_validated_sql_rejects_country_question_without_country_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' + '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' + 'ORDER BY COUNT(*) DESC' + ), + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + "Show the total commodity line value by country.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Show the total commodity line value by country.", + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'AS "Country_of_Ultimate_Destination_Code", ' + 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' + 'AS "TotalCommodity_Line_Value" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' + ) + + +def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT * FROM "dbo_tnoStageNewOrders" ' + 'WHERE "dbo_tnoStageNewOrders"."Division" = ' + "'Daimler Trucks North America'" + ), + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + "List orders for Daimler Trucks North America.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "List orders for Daimler Trucks North America.", + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """, + ], + ) + + assert sql == ( + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."CustName" = ' + "'Daimler Trucks North America' " + 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' + ) + + def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) From 05275616842654e07bcccb002fe6226603420d7d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 13:09:52 +0530 Subject: [PATCH 0508/1087] Reject unrequested test schema sources --- .../retrieval/db_schema_retrieval.py | 27 ++++++++++ wren-ai-service/src/web/v1/services/ask.py | 50 +++++++++++++++++++ .../retrieval/test_db_schema_retrieval.py | 22 ++++++++ .../pytest/services/test_ask_sales_sql.py | 26 ++++++++++ 4 files changed, 125 insertions(+) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6b64c9b452..5e24b26813 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -334,6 +334,26 @@ def _source_shape_score(query: str, document: Document) -> int: return score +def _is_unrequested_strong_non_production_source(query: str, document: Document) -> bool: + strong_non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "temp", + "test", + "tmp", + ) + source_terms = _retrieval_terms(_source_text(document)) + return bool( + source_terms & set(strong_non_production_terms) + and not _query_mentions_any(query or "", strong_non_production_terms) + ) + + def _document_relevance_score(document: Document, query_terms: set[str]) -> int: if not query_terms: return 0 @@ -438,6 +458,13 @@ def _select_relevant_table_documents( return documents[:max_tables] candidate_pool = [item for item in reranked if item[3] > 0] or reranked + production_pool = [ + item + for item in candidate_pool + if not _is_unrequested_strong_non_production_source(query, item[2]) + ] + if production_pool: + candidate_pool = production_pool selected = [ document for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ce42c81888..00f9315fcc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1028,6 +1028,50 @@ def _sql_satisfies_entity_lookup_request( return True + def _is_unrequested_non_production_table_reference( + self, table_name: str, query: str | None + ) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + normalized_table = self._normalize_schema_token(table_name) + strong_non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "temp", + "test", + "tmp", + ) + if not any(term in normalized_table for term in strong_non_production_terms): + return False + return not any( + re.search(rf"\b{re.escape(term)}\b", normalized_query) + for term in strong_non_production_terms + ) + + def _sql_avoids_unrequested_non_production_tables( + self, sql: str, query: str | None, referenced_tables: list[str] + ) -> bool: + invalid_tables = [ + table + for table in referenced_tables + if self._is_unrequested_non_production_table_reference(table, query) + ] + if not invalid_tables: + return True + + logger.warning( + "Ignoring SQL because it references unrequested non-production tables. " + "query=%s invalid_tables=%s sql=%s", + query, + invalid_tables, + sql, + ) + return False + def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1349,6 +1393,12 @@ def _sql_matches_question_intent( valid_tables, ): return False + if not self._sql_avoids_unrequested_non_production_tables( + sql, + query, + referenced_tables, + ): + return False if not expects_dimension: return True diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index bbb0c79aef..7c7d46ce89 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -130,6 +130,28 @@ def test_select_relevant_table_documents_limits_weak_extra_candidates(): assert "staging_audit" not in [document.meta["name"] for document in selected] +def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): + documents = [ + Document( + content="Raw test load rows with order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ), + Document( + content="New order transaction records with market and customer details.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.4, + ), + ] + + selected = _select_relevant_table_documents( + "Show order distribution across markets.", + documents, + ) + + assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] + + @pytest.mark.asyncio async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 5ba8ca74b8..c7ab4c9272 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -295,6 +295,32 @@ def test_validated_sql_rejects_country_question_without_country_column(): assert result is None +def test_validated_sql_rejects_unrequested_test_table_for_market_distribution(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT TOP 10 "dbo_xStageLoad8_Test"."Market" AS "Market", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_xStageLoad8_Test" ' + 'WHERE "dbo_xStageLoad8_Test"."Market" IS NOT NULL ' + 'GROUP BY "dbo_xStageLoad8_Test"."Market" ' + 'ORDER BY COUNT(*) DESC' + ), + [ + """ + CREATE TABLE dbo_xStageLoad8_Test ( + Market VARCHAR, + OrdNo VARCHAR + ); + """ + ], + "Show order distribution across markets.", + ) + + assert result is None + + def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): service = AskService.__new__(AskService) From 20e44de636b04caa4e30eda1e4059d78528dbc0f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 13:31:49 +0530 Subject: [PATCH 0509/1087] Revert "Reject unrequested test schema sources" This reverts commit 05275616842654e07bcccb002fe6226603420d7d. --- .../retrieval/db_schema_retrieval.py | 27 ---------- wren-ai-service/src/web/v1/services/ask.py | 50 ------------------- .../retrieval/test_db_schema_retrieval.py | 22 -------- .../pytest/services/test_ask_sales_sql.py | 26 ---------- 4 files changed, 125 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 5e24b26813..6b64c9b452 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -334,26 +334,6 @@ def _source_shape_score(query: str, document: Document) -> int: return score -def _is_unrequested_strong_non_production_source(query: str, document: Document) -> bool: - strong_non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "temp", - "test", - "tmp", - ) - source_terms = _retrieval_terms(_source_text(document)) - return bool( - source_terms & set(strong_non_production_terms) - and not _query_mentions_any(query or "", strong_non_production_terms) - ) - - def _document_relevance_score(document: Document, query_terms: set[str]) -> int: if not query_terms: return 0 @@ -458,13 +438,6 @@ def _select_relevant_table_documents( return documents[:max_tables] candidate_pool = [item for item in reranked if item[3] > 0] or reranked - production_pool = [ - item - for item in candidate_pool - if not _is_unrequested_strong_non_production_source(query, item[2]) - ] - if production_pool: - candidate_pool = production_pool selected = [ document for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 00f9315fcc..ce42c81888 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1028,50 +1028,6 @@ def _sql_satisfies_entity_lookup_request( return True - def _is_unrequested_non_production_table_reference( - self, table_name: str, query: str | None - ) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_table = self._normalize_schema_token(table_name) - strong_non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "temp", - "test", - "tmp", - ) - if not any(term in normalized_table for term in strong_non_production_terms): - return False - return not any( - re.search(rf"\b{re.escape(term)}\b", normalized_query) - for term in strong_non_production_terms - ) - - def _sql_avoids_unrequested_non_production_tables( - self, sql: str, query: str | None, referenced_tables: list[str] - ) -> bool: - invalid_tables = [ - table - for table in referenced_tables - if self._is_unrequested_non_production_table_reference(table, query) - ] - if not invalid_tables: - return True - - logger.warning( - "Ignoring SQL because it references unrequested non-production tables. " - "query=%s invalid_tables=%s sql=%s", - query, - invalid_tables, - sql, - ) - return False - def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1393,12 +1349,6 @@ def _sql_matches_question_intent( valid_tables, ): return False - if not self._sql_avoids_unrequested_non_production_tables( - sql, - query, - referenced_tables, - ): - return False if not expects_dimension: return True diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7c7d46ce89..bbb0c79aef 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -130,28 +130,6 @@ def test_select_relevant_table_documents_limits_weak_extra_candidates(): assert "staging_audit" not in [document.meta["name"] for document in selected] -def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): - documents = [ - Document( - content="Raw test load rows with order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ), - Document( - content="New order transaction records with market and customer details.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.4, - ), - ] - - selected = _select_relevant_table_documents( - "Show order distribution across markets.", - documents, - ) - - assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] - - @pytest.mark.asyncio async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index c7ab4c9272..5ba8ca74b8 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -295,32 +295,6 @@ def test_validated_sql_rejects_country_question_without_country_column(): assert result is None -def test_validated_sql_rejects_unrequested_test_table_for_market_distribution(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT TOP 10 "dbo_xStageLoad8_Test"."Market" AS "Market", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_xStageLoad8_Test" ' - 'WHERE "dbo_xStageLoad8_Test"."Market" IS NOT NULL ' - 'GROUP BY "dbo_xStageLoad8_Test"."Market" ' - 'ORDER BY COUNT(*) DESC' - ), - [ - """ - CREATE TABLE dbo_xStageLoad8_Test ( - Market VARCHAR, - OrdNo VARCHAR - ); - """ - ], - "Show order distribution across markets.", - ) - - assert result is None - - def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): service = AskService.__new__(AskService) From 1ba5c0b39c74f055c12e71c143ef975000a8a177 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 14:32:39 +0530 Subject: [PATCH 0510/1087] Add semantic concept coverage for retrieval --- .../retrieval/db_schema_retrieval.py | 92 ++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 106 +++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 20 ++++ .../pytest/services/test_ask_sales_sql.py | 57 ++++++++++ 4 files changed, 270 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6b64c9b452..06f4b47504 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -162,7 +162,16 @@ def expand_business_terms_for_retrieval(query: str) -> str: ) ): expansions.append( - "transaction purchase billing account geography area representative product item category sku quantity units sold amount value total metric money exchange currency" + "transaction purchase billing account geography customer client company name area representative product item category sku quantity units sold amount value total metric money exchange currency" + ) + + if re.search( + r"\b(?:show|list|find|get|display)\b.*\b(?:orders?|records?|rows?|transactions?)\b\s+" + r"(?:for|where|with)\s+\S+", + normalized, + ): + expansions.append( + "customer client account company name entity lookup identifier transaction order record" ) if any( @@ -230,6 +239,68 @@ def _query_mentions_any(query: str, terms: tuple[str, ...]) -> bool: return any(re.search(rf"\b{re.escape(term)}\b", normalized) for term in terms) +def _retrieval_concept_groups(query: str) -> list[set[str]]: + query_terms = _retrieval_terms(query) + if not query_terms: + return [] + + concept_specs: list[tuple[set[str], set[str]]] = [ + ( + {"order", "orders", "neworder", "purchase", "transaction"}, + {"order", "orders", "ord", "ordno", "orderid", "orderdate", "neworder", "purchase", "transaction"}, + ), + ( + {"customer", "customers", "client", "account", "company"}, + {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"}, + ), + ( + {"product", "products", "item", "sku", "category", "categories"}, + {"product", "products", "prod", "item", "sku", "category", "categories", "type", "name"}, + ), + ( + {"market", "markets", "region", "regions"}, + {"market", "markets", "region", "regions", "area", "territory", "country"}, + ), + ( + {"country", "countries", "destination"}, + {"country", "countries", "destination", "nation", "market", "region"}, + ), + ( + {"invoice", "invoices", "billing"}, + {"invoice", "invoices", "inv", "billing", "bill", "amount", "value"}, + ), + ( + {"amount", "value", "total", "sum", "revenue", "sales", "cost"}, + {"amount", "value", "total", "sum", "revenue", "sales", "sale", "cost", "price", "money"}, + ), + ( + {"quantity", "qty", "sold", "units"}, + {"quantity", "qty", "sold", "unit", "units", "volume"}, + ), + ( + {"date", "month", "monthly", "year", "quarter", "trend", "period"}, + {"date", "month", "year", "quarter", "time", "timestamp", "orddate", "invdate", "period"}, + ), + ] + + groups: list[set[str]] = [] + for triggers, aliases in concept_specs: + if query_terms & triggers: + groups.append(aliases) + + if re.search( + r"\b(?:show|list|find|get|display)\b.*\b(?:orders?|records?|rows?|transactions?)\b\s+" + r"(?:for|where|with)\s+\S+", + query or "", + flags=re.IGNORECASE, + ): + groups.append( + {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"} + ) + + return groups + + def _source_text(document: Document) -> str: return " ".join( str(part or "") @@ -247,6 +318,15 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 + concept_groups = _retrieval_concept_groups(query) + if concept_groups: + covered_groups = sum(1 for group in concept_groups if group & source_terms) + score += 18 * covered_groups + if covered_groups == len(concept_groups): + score += 30 + elif covered_groups == 0: + score -= 25 + weak_non_production_terms = ( "stage", "staging", @@ -315,6 +395,11 @@ def _source_shape_score(query: str, document: Document) -> int: r"employees?|entities|items?|names?|products?|suppliers?|users?|vendors?)\b", normalized_query, ) + entity_lookup_pattern = re.search( + r"\b(?:list|show|display|get|find)\b.*\b(?:orders?|records?|rows?|transactions?)\b\s+" + r"(?:for|where|with)\s+\S+", + normalized_query, + ) asks_for_aggregation = _query_mentions_any(normalized_query, aggregation_terms) or bool( re.search(r"\b(?:by|per|each|top|bottom|rank|ranking)\b", normalized_query) ) @@ -325,6 +410,11 @@ def _source_shape_score(query: str, document: Document) -> int: score += 35 if source_terms & set(transaction_source_terms): score -= 12 + elif entity_lookup_pattern: + if source_terms & {"customer", "cust", "custname", "client", "account", "company", "name"}: + score += 45 + if source_terms & set(transaction_source_terms): + score += 20 elif asks_for_aggregation: if source_terms & set(transaction_source_terms): score += 25 diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ce42c81888..2678400b73 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -641,11 +641,93 @@ def _intent_tokens(self, text: str) -> set[str]: def _schema_name_tokens(self, name: str) -> set[str]: spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(name or "")) - return { + tokens = { token for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) if len(token) > 1 } + normalized = self._normalize_schema_identifier_key(name) + if normalized: + tokens.add(normalized) + if ( + "date" in tokens + or "time" in tokens + or normalized + in { + "createdat", + "updatedat", + "createdon", + "updatedon", + "timestamp", + } + ): + tokens.update({"date", "time", "timestamp"}) + return tokens + + def _schema_terms_for_table(self, table: dict[str, Any]) -> set[str]: + terms = self._schema_name_tokens(str(table.get("name") or "")) + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + if column_name: + terms.update(self._schema_name_tokens(column_name)) + return terms + + def _semantic_concept_groups_for_query(self, query: str | None) -> list[set[str]]: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + query_tokens = self._intent_tokens(query or "") + if not normalized and not query_tokens: + return [] + + concept_specs: list[tuple[set[str], set[str]]] = [ + ( + {"order", "orders", "neworder", "purchase", "transaction"}, + {"order", "orders", "ord", "ordno", "orderid", "orderdate", "neworder", "purchase", "transaction"}, + ), + ( + {"customer", "customers", "client", "account", "company"}, + {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"}, + ), + ( + {"product", "products", "item", "sku", "category", "categories"}, + {"product", "products", "prod", "prodname", "item", "sku", "category", "categories", "type", "name"}, + ), + ( + {"market", "markets", "region", "regions"}, + {"market", "markets", "region", "regions", "area", "territory", "country"}, + ), + ( + {"country", "countries", "destination"}, + {"country", "countries", "destination", "nation", "market", "region"}, + ), + ( + {"invoice", "invoices", "billing"}, + {"invoice", "invoices", "inv", "billing", "bill", "amount", "value"}, + ), + ( + {"amount", "value", "total", "sum", "revenue", "sales", "sale", "cost"}, + {"amount", "value", "total", "sum", "revenue", "sales", "sale", "cost", "price", "money"}, + ), + ( + {"quantity", "qty", "sold", "units"}, + {"quantity", "qty", "sold", "unit", "units", "volume"}, + ), + ( + {"date", "month", "monthly", "year", "quarter", "trend", "period"}, + {"date", "month", "year", "quarter", "time", "timestamp", "orddate", "invdate", "period"}, + ), + ] + + groups: list[set[str]] = [] + for triggers, aliases in concept_specs: + if query_tokens & triggers: + groups.append(aliases) + + if self._extract_entity_lookup_phrase(query): + groups.append( + {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"} + ) + + return groups def _table_for_sql_reference( self, table_reference: str, valid_tables: dict[str, dict[str, Any]] @@ -664,7 +746,7 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: if not normalized: return [] - concept_groups: list[set[str]] = [] + concept_groups: list[set[str]] = self._semantic_concept_groups_for_query(query) if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: @@ -707,12 +789,17 @@ def _sql_covers_required_question_concepts( referenced_column_tokens: set[str], referenced_table_tokens: set[str], ) -> bool: - sql_text = (sql or "").lower() + sql_text = re.sub( + r"\b(?:order|group)\s+by\b|\bselect\b|\bfrom\b|\bwhere\b", + " ", + (sql or "").lower(), + ) + sql_text_tokens = self._intent_tokens(sql_text) available_tokens = referenced_column_tokens | referenced_table_tokens for concept_group in self._required_sql_concept_groups(query): if concept_group & available_tokens: continue - if any(token in sql_text for token in concept_group): + if concept_group & sql_text_tokens: continue logger.warning( "Ignoring SQL because it does not cover required question concept. " @@ -2237,6 +2324,7 @@ def _select_best_analytics_table( query: str = "", ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + concept_groups = self._semantic_concept_groups_for_query(query) scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] ] = [] @@ -2279,6 +2367,16 @@ def _select_best_analytics_table( table_name = str(table.get("name") or "").lower() if not table_name: continue + schema_terms = self._schema_terms_for_table(table) + if concept_groups: + covered_groups = sum( + 1 for concept_group in concept_groups if concept_group & schema_terms + ) + score += 14 * covered_groups + if covered_groups == len(concept_groups): + score += 25 + elif covered_groups == 0: + score -= 20 if "sales" in table_name: score += 5 if "tblsales" in self._normalize_schema_token(table_name): diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index bbb0c79aef..708e5526d8 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -87,6 +87,26 @@ def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_under assert documents[0].meta["name"] == "dbo_xStageNewOrders" +def test_rerank_table_documents_prefers_customer_capable_source_for_entity_lookup(): + generic_order_table = Document( + content="New order rows by division and market.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tnoStageNewOrders"}, + score=0.95, + ) + customer_order_table = Document( + content="New order transactions with customer name, order number, market, and customer purchase order.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, + score=0.45, + ) + + documents = _rerank_table_documents( + "Show me orders for Lockheed Martin.", + [generic_order_table, customer_order_table], + ) + + assert documents[0].meta["name"] == "dbo_tblNewOrders" + + def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 5ba8ca74b8..eaba1a3644 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -295,6 +295,32 @@ def test_validated_sql_rejects_country_question_without_country_column(): assert result is None +def test_validated_sql_rejects_order_distribution_without_order_concept(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT TOP 10 "dbo_xStageLoad8_Test"."Market" AS "Market", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_xStageLoad8_Test" ' + 'WHERE "dbo_xStageLoad8_Test"."Market" IS NOT NULL ' + 'GROUP BY "dbo_xStageLoad8_Test"."Market" ' + 'ORDER BY COUNT(*) DESC' + ), + [ + """ + CREATE TABLE dbo_xStageLoad8_Test ( + Market VARCHAR, + LoadId VARCHAR + ); + """ + ], + "Show order distribution across markets.", + ) + + assert result is None + + def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): service = AskService.__new__(AskService) @@ -323,6 +349,37 @@ def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): ) +def test_schema_grounded_analytics_prefers_full_concept_coverage_for_order_market_distribution(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Show order distribution across markets.", + [ + """ + CREATE TABLE dbo_xStageLoad8_Test ( + Market VARCHAR, + LoadId VARCHAR + ); + """, + """ + CREATE TABLE dbo_tblNewOrders ( + Market VARCHAR, + OrdNo VARCHAR + ); + """, + ], + ) + + assert sql == ( + 'SELECT "dbo_tblNewOrders"."Market" AS "Market", ' + 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "OrderCount" ' + 'FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."Market" IS NOT NULL ' + 'GROUP BY "dbo_tblNewOrders"."Market" ' + 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' + ) + + def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): service = AskService.__new__(AskService) From 164cd5d12fa6336876a268569c56c02cbfe4aa87 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 14:53:30 +0530 Subject: [PATCH 0511/1087] Prune SQL context to semantic contract --- wren-ai-service/src/web/v1/services/ask.py | 203 ++++++++++++++++++ .../pytest/services/test_ask_sales_sql.py | 108 ++++++++++ 2 files changed, 311 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2678400b73..cf645ab653 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -672,6 +672,100 @@ def _schema_terms_for_table(self, table: dict[str, Any]) -> set[str]: terms.update(self._schema_name_tokens(column_name)) return terms + def _schema_source_shape_score( + self, query: str | None, table: dict[str, Any] + ) -> int: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + table_terms = self._schema_terms_for_table(table) + + score = 0 + weak_non_production_terms = {"stage", "staging"} + strong_non_production_terms = { + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "temp", + "test", + "tmp", + } + if table_terms & strong_non_production_terms and not any( + re.search(rf"\b{re.escape(term)}\b", normalized_query) + for term in strong_non_production_terms + ): + score -= 80 + if table_terms & weak_non_production_terms and not any( + re.search(rf"\b{re.escape(term)}\b", normalized_query) + for term in weak_non_production_terms + ): + score -= 20 + + transaction_terms = { + "activity", + "detail", + "event", + "fact", + "history", + "invoice", + "line", + "order", + "orders", + "sale", + "sales", + "transaction", + } + reference_terms = { + "account", + "catalog", + "dimension", + "directory", + "entity", + "lookup", + "master", + "profile", + "reference", + } + if any( + term in normalized_query + for term in ( + "amount", + "average", + "count", + "distribution", + "rank", + "ranking", + "sum", + "top", + "total", + "trend", + "value", + ) + ): + if table_terms & transaction_terms: + score += 25 + if table_terms & reference_terms: + score += 5 + + if self._extract_entity_lookup_phrase(query): + if table_terms & transaction_terms: + score += 20 + if table_terms & { + "account", + "buyer", + "client", + "company", + "cust", + "customer", + "custname", + "name", + }: + score += 35 + + return score + def _semantic_concept_groups_for_query(self, query: str | None) -> list[set[str]]: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) query_tokens = self._intent_tokens(query or "") @@ -729,6 +823,93 @@ def _semantic_concept_groups_for_query(self, query: str | None) -> list[set[str] return groups + def _semantic_contract_score_table( + self, query: str | None, table: dict[str, Any] + ) -> tuple[int, set[int]]: + concept_groups = self._semantic_concept_groups_for_query(query) + table_terms = self._schema_terms_for_table(table) + covered_groups = { + index + for index, concept_group in enumerate(concept_groups) + if concept_group & table_terms + } + score = 100 * len(covered_groups) + if concept_groups and len(covered_groups) == len(concept_groups): + score += 80 + elif concept_groups and not covered_groups: + score -= 50 + score += self._schema_source_shape_score(query, table) + return score, covered_groups + + def _scope_retrieval_to_semantic_contract( + self, + query: str, + documents: list[dict], + table_names: list[str], + table_ddls: list[str], + *, + max_tables: int = 4, + ) -> tuple[list[dict], list[str], list[str]]: + concept_groups = self._semantic_concept_groups_for_query(query) + if not concept_groups or len(table_ddls) <= 1: + return documents, table_names, table_ddls + + parsed_tables = self._parse_schema_tables(table_ddls) + if not parsed_tables: + return documents, table_names, table_ddls + + scored: list[tuple[int, int, set[int]]] = [] + for index, table in enumerate(parsed_tables): + score, covered_groups = self._semantic_contract_score_table(query, table) + if covered_groups: + scored.append((score, index, covered_groups)) + + if not scored: + logger.warning( + "No retrieved schema tables cover the semantic contract; keeping original scoped retrieval. query=%s tables=%s", + query, + table_names, + ) + return documents, table_names, table_ddls + + scored = sorted(scored, key=lambda item: (item[0], len(item[2])), reverse=True) + best_score, best_index, best_coverage = scored[0] + all_group_indexes = set(range(len(concept_groups))) + + selected_indexes: list[int] = [] + selected_coverage: set[int] = set() + if best_coverage == all_group_indexes: + selected_indexes.append(best_index) + selected_coverage.update(best_coverage) + else: + for _score, index, coverage in scored: + new_coverage = coverage - selected_coverage + if not new_coverage and selected_indexes: + continue + selected_indexes.append(index) + selected_coverage.update(coverage) + if selected_coverage == all_group_indexes or len(selected_indexes) >= max_tables: + break + + if not selected_indexes: + selected_indexes = [best_index] + + selected_indexes = sorted(dict.fromkeys(selected_indexes)) + if len(selected_indexes) == len(table_ddls): + return documents, table_names, table_ddls + + logger.info( + "Scoped retrieved schema to semantic contract. query=%s before=%s after=%s", + query, + table_names, + [table_names[index] for index in selected_indexes if index < len(table_names)], + ) + return ( + [documents[index] for index in selected_indexes if index < len(documents)], + [table_names[index] for index in selected_indexes if index < len(table_names)], + [table_ddls[index] for index in selected_indexes if index < len(table_ddls)], + ) + def _table_for_sql_reference( self, table_reference: str, valid_tables: dict[str, dict[str, Any]] ) -> dict[str, Any] | None: @@ -2377,6 +2558,7 @@ def _select_best_analytics_table( score += 25 elif covered_groups == 0: score -= 20 + score += self._schema_source_shape_score(query, table) if "sales" in table_name: score += 5 if "tblsales" in self._normalize_schema_token(table_name): @@ -2866,6 +3048,9 @@ def _build_schema_grounded_analytics_sql( "volume", "how many", "number of", + "top", + "highest", + "most", "monthly", "over time", "last 12 months", @@ -6852,6 +7037,15 @@ async def ask( explicit_table_names, ) ) + if documents and not explicit_table_names: + documents, table_names, table_ddls = ( + self._scope_retrieval_to_semantic_contract( + sql_user_query, + documents, + table_names, + table_ddls, + ) + ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) @@ -7143,6 +7337,15 @@ async def ask( table_names, table_ddls, ) + if not explicit_table_names: + documents, table_names, table_ddls = ( + self._scope_retrieval_to_semantic_contract( + sql_user_query, + documents, + table_names, + table_ddls, + ) + ) ( documents, table_names, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index eaba1a3644..66a296da7f 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -380,6 +380,114 @@ def test_schema_grounded_analytics_prefers_full_concept_coverage_for_order_marke ) +def test_scope_retrieval_to_semantic_contract_prefers_production_full_coverage_table(): + service = AskService.__new__(AskService) + documents = [ + { + "table_name": "dbo_qSalesCubeDev", + "table_ddl": """ + CREATE TABLE dbo_qSalesCubeDev ( + Market VARCHAR, + OrdNo VARCHAR + ); + """, + }, + { + "table_name": "dbo_tblNewOrders", + "table_ddl": """ + CREATE TABLE dbo_tblNewOrders ( + Market VARCHAR, + OrdNo VARCHAR + ); + """, + }, + ] + + scoped_documents, scoped_table_names, scoped_table_ddls = ( + service._scope_retrieval_to_semantic_contract( + "Show order distribution across markets.", + documents, + ["dbo_qSalesCubeDev", "dbo_tblNewOrders"], + [document["table_ddl"] for document in documents], + ) + ) + + assert scoped_documents == [documents[1]] + assert scoped_table_names == ["dbo_tblNewOrders"] + assert scoped_table_ddls == [documents[1]["table_ddl"]] + + +def test_scope_retrieval_to_semantic_contract_keeps_multiple_tables_when_needed(): + service = AskService.__new__(AskService) + documents = [ + { + "table_name": "dbo_OrderFacts", + "table_ddl": """ + CREATE TABLE dbo_OrderFacts ( + Market VARCHAR, + OrdNo VARCHAR + ); + """, + }, + { + "table_name": "dbo_Customers", + "table_ddl": """ + CREATE TABLE dbo_Customers ( + CustNo VARCHAR, + CustName VARCHAR + ); + """, + }, + { + "table_name": "dbo_LoadAudit", + "table_ddl": """ + CREATE TABLE dbo_LoadAudit ( + LoadId VARCHAR, + Status VARCHAR + ); + """, + }, + ] + + scoped_documents, scoped_table_names, _scoped_table_ddls = ( + service._scope_retrieval_to_semantic_contract( + "Show order distribution by market and customer.", + documents, + ["dbo_OrderFacts", "dbo_Customers", "dbo_LoadAudit"], + [document["table_ddl"] for document in documents], + ) + ) + + assert scoped_documents == documents[:2] + assert scoped_table_names == ["dbo_OrderFacts", "dbo_Customers"] + + +def test_schema_grounded_analytics_counts_top_new_orders_by_business_unit(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Which business unit has the top 20 new orders this period?", + [ + """ + CREATE TABLE dbo_tblNewOrders ( + BU VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 20 "dbo_tblNewOrders"."BU" AS "BU", ' + 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "OrderCount" ' + 'FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."BU" IS NOT NULL ' + 'GROUP BY "dbo_tblNewOrders"."BU" ' + 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' + ) + + def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): service = AskService.__new__(AskService) From c0cb0667dfd397a78c857364ae419709ddaff2a0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 15:10:00 +0530 Subject: [PATCH 0512/1087] Tighten semantic table retrieval grounding --- .../retrieval/db_schema_retrieval.py | 104 ++++++++++++++---- .../retrieval/test_db_schema_retrieval.py | 79 +++++++++++++ 2 files changed, 162 insertions(+), 21 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 06f4b47504..22d412dd40 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -30,6 +30,23 @@ logger = logging.getLogger("wren-ai-service") MAX_RELEVANT_TABLE_CANDIDATES = 5 +MIN_TABLE_DESCRIPTION_CANDIDATE_WINDOW = 50 +WEAK_NON_PRODUCTION_TERMS = ( + "stage", + "staging", +) +STRONG_NON_PRODUCTION_TERMS = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "temp", + "test", + "tmp", +) table_columns_selection_system_prompt = """ @@ -327,28 +344,12 @@ def _source_shape_score(query: str, document: Document) -> int: elif covered_groups == 0: score -= 25 - weak_non_production_terms = ( - "stage", - "staging", - ) - strong_non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "temp", - "test", - "tmp", - ) - if source_terms & set(strong_non_production_terms) and not _query_mentions_any( - normalized_query, strong_non_production_terms + if source_terms & set(STRONG_NON_PRODUCTION_TERMS) and not _query_mentions_any( + normalized_query, STRONG_NON_PRODUCTION_TERMS ): score -= 240 - if source_terms & set(weak_non_production_terms) and not _query_mentions_any( - normalized_query, weak_non_production_terms + if source_terms & set(WEAK_NON_PRODUCTION_TERMS) and not _query_mentions_any( + normalized_query, WEAK_NON_PRODUCTION_TERMS ): score -= 40 @@ -463,6 +464,27 @@ def _semantic_score(document: Document) -> float: return 0.0 +def _retrieval_concept_coverage(query: str, document: Document) -> set[int]: + source_terms = _retrieval_terms(_source_text(document)) + return { + index + for index, concept_group in enumerate(_retrieval_concept_groups(query)) + if concept_group & source_terms + } + + +def _is_unrequested_strong_non_production_source( + query: str, document: Document +) -> bool: + source_terms = _retrieval_terms(_source_text(document)) + return bool(source_terms & set(STRONG_NON_PRODUCTION_TERMS)) and not ( + _query_mentions_any( + query or "", + STRONG_NON_PRODUCTION_TERMS, + ) + ) + + def _score_table_documents( query: str, documents: list[Document] ) -> list[tuple[float, int, Document, int, float]]: @@ -528,6 +550,42 @@ def _select_relevant_table_documents( return documents[:max_tables] candidate_pool = [item for item in reranked if item[3] > 0] or reranked + concept_groups = _retrieval_concept_groups(query) + if concept_groups: + coverage_by_index = { + index: _retrieval_concept_coverage(query, item[2]) + for index, item in enumerate(candidate_pool) + } + grounded_production_coverage: set[int] = set() + for index, item in enumerate(candidate_pool): + coverage = coverage_by_index[index] + if coverage and not _is_unrequested_strong_non_production_source( + query, item[2] + ): + grounded_production_coverage.update(coverage) + + if grounded_production_coverage: + filtered_pool = [] + excluded_names = [] + for index, item in enumerate(candidate_pool): + document = item[2] + coverage = coverage_by_index[index] + if ( + _is_unrequested_strong_non_production_source(query, document) + and coverage <= grounded_production_coverage + ): + excluded_names.append(document.meta.get("name")) + continue + filtered_pool.append(item) + if filtered_pool: + candidate_pool = filtered_pool + if excluded_names: + logger.info( + "Excluded unrequested non-production table candidates for query=%s names=%s", + query, + excluded_names, + ) + selected = [ document for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] @@ -1009,11 +1067,15 @@ def __init__( table_column_retrieval_size: int = 100, **kwargs, ): + table_description_candidate_window = max( + table_retrieval_size, + MIN_TABLE_DESCRIPTION_CANDIDATE_WINDOW, + ) self._components = { "embedder": embedder_provider.get_text_embedder(), "table_retriever": document_store_provider.get_retriever( document_store_provider.get_store(dataset_name="table_descriptions"), - top_k=table_retrieval_size, + top_k=table_description_candidate_window, ), "dbschema_retriever": document_store_provider.get_retriever( document_store_provider.get_store(), diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 708e5526d8..7da427d415 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,6 +2,7 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( + DbSchemaRetrieval, _is_project_wide_analysis_query, _rerank_table_documents, _select_relevant_table_documents, @@ -87,6 +88,46 @@ def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_under assert documents[0].meta["name"] == "dbo_xStageNewOrders" +def test_select_relevant_table_documents_excludes_unrequested_dev_equivalent(): + dev_cube = Document( + content="Development sales cube with order number and market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_qSalesCubeDev"}, + score=500, + ) + production_orders = Document( + content="New order transaction records with order number and market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, + score=0.2, + ) + + documents = _select_relevant_table_documents( + "Show order distribution across markets.", + [dev_cube, production_orders], + ) + + assert [document.meta["name"] for document in documents] == ["dbo_tblNewOrders"] + + +def test_select_relevant_table_documents_keeps_explicit_dev_request(): + dev_cube = Document( + content="Development sales cube with order number and market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_qSalesCubeDev"}, + score=500, + ) + production_orders = Document( + content="New order transaction records with order number and market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, + score=0.2, + ) + + documents = _select_relevant_table_documents( + "Show order distribution across markets in dev.", + [dev_cube, production_orders], + ) + + assert documents[0].meta["name"] == "dbo_qSalesCubeDev" + + def test_rerank_table_documents_prefers_customer_capable_source_for_entity_lookup(): generic_order_table = Document( content="New order rows by division and market.", @@ -272,6 +313,44 @@ async def run(self, query_embedding, filters): } +def test_db_schema_retrieval_fetches_wider_table_description_window(): + class LLMProvider: + def get_generator(self, **kwargs): + return object() + + def get_model(self): + return "gpt-4o-mini" + + def get_context_window_size(self): + return 1000 + + class EmbedderProvider: + def get_text_embedder(self): + return object() + + class DocumentStoreProvider: + def __init__(self): + self.retriever_top_k = [] + + def get_store(self, dataset_name=None): + return dataset_name or "default" + + def get_retriever(self, store, top_k): + self.retriever_top_k.append((store, top_k)) + return object() + + document_store_provider = DocumentStoreProvider() + + DbSchemaRetrieval( + llm_provider=LLMProvider(), + embedder_provider=EmbedderProvider(), + document_store_provider=document_store_provider, + table_retrieval_size=10, + ) + + assert document_store_provider.retriever_top_k[0] == ("table_descriptions", 50) + + @pytest.mark.asyncio async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: From aa158099961baaa519dc41988866bf1109777bf6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 15:24:00 +0530 Subject: [PATCH 0513/1087] Ground retrieval on complete business concepts --- .../retrieval/db_schema_retrieval.py | 23 +++++- wren-ai-service/src/web/v1/services/ask.py | 82 ++++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 22 ++++- .../pytest/services/test_ask_sales_sql.py | 44 ++++++++++ 4 files changed, 165 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 22d412dd40..84a6e70e93 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -30,7 +30,7 @@ logger = logging.getLogger("wren-ai-service") MAX_RELEVANT_TABLE_CANDIDATES = 5 -MIN_TABLE_DESCRIPTION_CANDIDATE_WINDOW = 50 +MIN_TABLE_DESCRIPTION_CANDIDATE_WINDOW = 100 WEAK_NON_PRODUCTION_TERMS = ( "stage", "staging", @@ -243,7 +243,9 @@ def _retrieval_terms(value: str) -> set[str]: for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) for token in re.findall(r"[A-Za-z0-9]+", split_token): - if len(token) <= 2 or token.lower() in stop_words: + if len(token) <= 2 and token.lower() not in {"bu"}: + continue + if token.lower() in stop_words: continue normalized_token = _normalize_retrieval_token(token) if normalized_token: @@ -278,6 +280,10 @@ def _retrieval_concept_groups(query: str) -> list[set[str]]: {"market", "markets", "region", "regions"}, {"market", "markets", "region", "regions", "area", "territory", "country"}, ), + ( + {"business", "unit", "units", "division"}, + {"business", "businessunit", "unit", "units", "bu", "division"}, + ), ( {"country", "countries", "destination"}, {"country", "countries", "destination", "nation", "market", "region"}, @@ -552,10 +558,23 @@ def _select_relevant_table_documents( candidate_pool = [item for item in reranked if item[3] > 0] or reranked concept_groups = _retrieval_concept_groups(query) if concept_groups: + all_concepts = set(range(len(concept_groups))) coverage_by_index = { index: _retrieval_concept_coverage(query, item[2]) for index, item in enumerate(candidate_pool) } + full_coverage_pool = [ + item + for index, item in enumerate(candidate_pool) + if coverage_by_index[index] == all_concepts + ] + if full_coverage_pool: + candidate_pool = full_coverage_pool + coverage_by_index = { + index: _retrieval_concept_coverage(query, item[2]) + for index, item in enumerate(candidate_pool) + } + grounded_production_coverage: set[int] = set() for index, item in enumerate(candidate_pool): coverage = coverage_by_index[index] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cf645ab653..7dc89d4824 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -789,6 +789,10 @@ def _semantic_concept_groups_for_query(self, query: str | None) -> list[set[str] {"market", "markets", "region", "regions"}, {"market", "markets", "region", "regions", "area", "territory", "country"}, ), + ( + {"business", "unit", "units", "division"}, + {"business", "businessunit", "unit", "units", "bu", "division"}, + ), ( {"country", "countries", "destination"}, {"country", "countries", "destination", "nation", "market", "region"}, @@ -950,6 +954,12 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: concept_groups.append({"currency", "curr", "money", "fx", "exchange"}) if "market" in normalized or "markets" in normalized: concept_groups.append({"market", "region", "country", "territory"}) + if ( + "business unit" in normalized + or "business units" in normalized + or re.search(r"\bbu\b", normalized) + ): + concept_groups.append({"business", "businessunit", "unit", "units", "bu", "division"}) if "region" in normalized or "regions" in normalized: concept_groups.append({"region", "market", "area", "territory", "country"}) if "country" in normalized or "countries" in normalized: @@ -1038,6 +1048,65 @@ def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> return False return True + def _is_grouped_metric_or_ranking_query(self, query: str | None) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return False + + has_metric_word = any( + term in normalized_query + for term in ( + "amount", + "average", + "avg", + "count", + "counts", + "distribution", + "how many", + "number of", + "record count", + "revenue", + "sum", + "total", + "value", + ) + ) + has_ranked_order_metric = any( + term in normalized_query + for term in ("top", "most", "highest", "largest", "rank", "ranking") + ) and any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ) + has_dimension_word = any( + term in normalized_query + for term in ( + "business unit", + "business units", + "category", + "categories", + "customer", + "customers", + "custname", + "currency", + "currencies", + "division", + "market", + "markets", + "product", + "products", + "region", + "regions", + "sales person", + "salesperson", + "source", + "status", + "type", + ) + ) or bool(re.search(r"\bbu\b", normalized_query)) + + return (has_metric_word or has_ranked_order_metric) and has_dimension_word + def _sql_satisfies_count_ranking_request( self, sql: str, query: str | None ) -> bool: @@ -1061,16 +1130,21 @@ def _sql_satisfies_count_ranking_request( for term in ("order", "orders", "record", "records", "row", "rows") ) ) - if not asks_for_count_metric: + asks_for_grouped_metric = self._is_grouped_metric_or_ranking_query(query) + if not asks_for_count_metric and not asks_for_grouped_metric: return True - asks_for_grouped_entity = any( + asks_for_grouped_entity = asks_for_grouped_metric or any( term in normalized_query for term in ( + "business unit", + "business units", "category", "customer", "customers", + "custname", "currency", + "division", "market", "product", "products", @@ -1081,7 +1155,7 @@ def _sql_satisfies_count_ranking_request( "status", "type", ) - ) + ) or bool(re.search(r"\bbu\b", normalized_query)) if not asks_for_grouped_entity: return True @@ -2139,6 +2213,8 @@ def _build_explicit_table_preview_sql( normalized_query = re.sub(r"\s+", " ", (query or "").strip()) if not normalized_query: return None + if self._is_grouped_metric_or_ranking_query(query): + return None if not re.search( r"\b(?:first|top|sample|preview|show|list)\b", diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7da427d415..3a45d6dab2 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -128,6 +128,26 @@ def test_select_relevant_table_documents_keeps_explicit_dev_request(): assert documents[0].meta["name"] == "dbo_qSalesCubeDev" +def test_select_relevant_table_documents_prefers_full_business_unit_order_coverage(): + weak_sales_margin = Document( + content="Sales margin facts with sales value, margin, order date, and product.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_qSalesMargin"}, + score=500, + ) + new_orders = Document( + content="New order records with BU, business unit, order number, order date, and customer.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, + score=0.2, + ) + + documents = _select_relevant_table_documents( + "Which business unit has the top 20 new orders this period?", + [weak_sales_margin, new_orders], + ) + + assert [document.meta["name"] for document in documents] == ["dbo_tblNewOrders"] + + def test_rerank_table_documents_prefers_customer_capable_source_for_entity_lookup(): generic_order_table = Document( content="New order rows by division and market.", @@ -348,7 +368,7 @@ def get_retriever(self, store, top_k): table_retrieval_size=10, ) - assert document_store_provider.retriever_top_k[0] == ("table_descriptions", 50) + assert document_store_provider.retriever_top_k[0] == ("table_descriptions", 100) @pytest.mark.asyncio diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 66a296da7f..0ee413f27c 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,6 +269,24 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) +def test_explicit_table_preview_does_not_handle_grouped_count_request(): + service = AskService.__new__(AskService) + + preview_sql = service._build_explicit_table_preview_sql( + "From dbo_tblNewOrders, show the top 5 customers by order count using CustName.", + [ + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + ) + + assert preview_sql is None + + def test_validated_sql_rejects_country_question_without_country_column(): service = AskService.__new__(AskService) @@ -488,6 +506,32 @@ def test_schema_grounded_analytics_counts_top_new_orders_by_business_unit(): ) +def test_validated_sql_rejects_business_unit_question_without_business_unit_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT TOP 20 "dbo_qSalesMargin"."ProductCategory" AS "ProductCategory", ' + 'COUNT(DISTINCT "dbo_qSalesMargin"."OrdNo") AS "OrderCount" ' + 'FROM "dbo_qSalesMargin" ' + 'GROUP BY "dbo_qSalesMargin"."ProductCategory" ' + 'ORDER BY COUNT(DISTINCT "dbo_qSalesMargin"."OrdNo") DESC' + ), + [ + """ + CREATE TABLE dbo_qSalesMargin ( + ProductCategory VARCHAR, + OrdNo VARCHAR, + OrderDate TIMESTAMP + ); + """ + ], + "Which business unit has the top 20 new orders this period?", + ) + + assert result is None + + def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): service = AskService.__new__(AskService) From ec5600239c72054ec6312497f062be879c4205bd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 15:32:48 +0530 Subject: [PATCH 0514/1087] Require table descriptions for schema retrieval --- .../retrieval/db_schema_retrieval.py | 33 +++++++--- .../retrieval/test_db_schema_retrieval.py | 60 ++++++++++++++----- 2 files changed, 70 insertions(+), 23 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 84a6e70e93..739c22e69e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -690,10 +690,30 @@ def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: return normalized +def _table_description_name_candidates(table_names: Optional[list[str]]) -> list[str]: + candidates: list[str] = [] + for table_name in _normalize_table_names(table_names): + raw_candidates = [table_name] + separator_normalized = re.sub(r"[.$]", "_", table_name) + raw_candidates.append(separator_normalized) + if "_" in table_name: + raw_candidates.append( + re.sub(r"^([A-Za-z_][A-Za-z0-9]*)_", r"\1.", table_name, count=1) + ) + if "." in table_name or "$" in table_name: + raw_candidates.append(re.split(r"[.$]", table_name)[-1]) + + for candidate in raw_candidates: + candidate = candidate.strip() + if candidate and candidate not in candidates: + candidates.append(candidate) + return candidates + + def _extract_table_names_from_table_retrieval( - table_retrieval: dict, explicit_tables: Optional[list[str]] = None + table_retrieval: dict, ) -> list[str]: - table_names = _normalize_table_names(explicit_tables) + table_names: list[str] = [] for document in table_retrieval.get("documents") or []: if not isinstance(document, Document): continue @@ -767,12 +787,13 @@ async def table_retrieval( return results if tables: - logger.info("Loading explicit table descriptions: %s", tables) + table_candidates = _table_description_name_candidates(tables) + logger.info("Loading explicit table descriptions: %s", table_candidates) explicit_filters = { **base_filters, "conditions": [ *base_filters["conditions"], - {"field": "name", "operator": "in", "value": tables}, + {"field": "name", "operator": "in", "value": table_candidates}, ], } return await table_retriever.run(query_embedding=[], filters=explicit_filters) @@ -788,9 +809,7 @@ async def dbschema_retrieval( dbschema_retriever: Any, tables: Optional[list[str]] = None, ) -> list[Document]: - selected_table_names = _extract_table_names_from_table_retrieval( - table_retrieval, tables - ) + selected_table_names = _extract_table_names_from_table_retrieval(table_retrieval) filters = { "operator": "AND", diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 3a45d6dab2..102f183a39 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -333,6 +333,44 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_table_retrieval_expands_explicit_table_name_forms_for_descriptions(): + class Retriever: + def __init__(self): + self.filters = None + + async def run(self, query_embedding, filters): + self.filters = filters + return {"documents": []} + + retriever = Retriever() + + await table_retrieval( + query="show failed repairs", + embedding={}, + project_id="project-1", + tables=["dbo.failure_patterns"], + table_retriever=retriever, + ) + + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + { + "field": "name", + "operator": "in", + "value": [ + "dbo.failure_patterns", + "dbo_failure_patterns", + "failure_patterns", + ], + }, + ], + } + + def test_db_schema_retrieval_fetches_wider_table_description_window(): class LLMProvider: def get_generator(self, **kwargs): @@ -489,18 +527,18 @@ def encode(self, value): @pytest.mark.asyncio -async def test_dbschema_retrieval_uses_explicit_tables_as_scope(): +async def test_dbschema_retrieval_does_not_use_explicit_tables_without_description(): class Retriever: def __init__(self): - self.filters = None + self.called = False async def run(self, query_embedding, filters): - self.filters = filters + self.called = True return {"documents": []} retriever = Retriever() - await dbschema_retrieval( + documents = await dbschema_retrieval( query="show failed repairs", table_retrieval={"documents": []}, project_id="project-1", @@ -508,15 +546,5 @@ async def run(self, query_embedding, filters): tables=["dbo.failure_patterns", "dbo_failure_patterns"], ) - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - { - "field": "name", - "operator": "in", - "value": ["dbo.failure_patterns", "dbo_failure_patterns"], - }, - ], - } + assert documents == [] + assert not retriever.called From 965b0c2ccd9cedfedecc69a7541d8b357b6a531d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 15:46:28 +0530 Subject: [PATCH 0515/1087] Handle generic market and entity distribution questions --- .../retrieval/db_schema_retrieval.py | 6 +- wren-ai-service/src/web/v1/services/ask.py | 99 ++++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 9 ++ .../pytest/services/test_ask_sales_sql.py | 53 ++++++++++ 4 files changed, 161 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 739c22e69e..0a805aee49 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -161,6 +161,8 @@ def expand_business_terms_for_retrieval(query: str) -> str: "invoices", "market", "markets", + "domestic", + "international", "order", "orders", "product", @@ -179,7 +181,7 @@ def expand_business_terms_for_retrieval(query: str) -> str: ) ): expansions.append( - "transaction purchase billing account geography customer client company name area representative product item category sku quantity units sold amount value total metric money exchange currency" + "transaction purchase billing account geography customer client company name area representative market domestic international product item category sku quantity units sold amount value total metric money exchange currency" ) if re.search( @@ -278,7 +280,7 @@ def _retrieval_concept_groups(query: str) -> list[set[str]]: ), ( {"market", "markets", "region", "regions"}, - {"market", "markets", "region", "regions", "area", "territory", "country"}, + {"market", "markets", "region", "regions", "area", "territory", "country", "domestic", "international"}, ), ( {"business", "unit", "units", "division"}, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7dc89d4824..3de52cabbf 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -787,7 +787,7 @@ def _semantic_concept_groups_for_query(self, query: str | None) -> list[set[str] ), ( {"market", "markets", "region", "regions"}, - {"market", "markets", "region", "regions", "area", "territory", "country"}, + {"market", "markets", "region", "regions", "area", "territory", "country", "domestic", "international"}, ), ( {"business", "unit", "units", "division"}, @@ -953,7 +953,7 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: if "currency" in normalized or "currencies" in normalized: concept_groups.append({"currency", "curr", "money", "fx", "exchange"}) if "market" in normalized or "markets" in normalized: - concept_groups.append({"market", "region", "country", "territory"}) + concept_groups.append({"market", "region", "country", "territory", "domestic", "international"}) if ( "business unit" in normalized or "business units" in normalized @@ -2791,6 +2791,94 @@ def _build_entity_lookup_sql( f"{order_clause}" ) + def _build_entity_distribution_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query or "distribution" not in normalized_query: + return None + + entity_candidates: tuple[str, ...] | None = None + entity_alias = "EntityCount" + if "customer" in normalized_query or "customers" in normalized_query: + entity_candidates = ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + "Client", + "ClientName", + ) + entity_alias = "CustomerCount" + elif "product" in normalized_query or "products" in normalized_query: + entity_candidates = ( + "Product", + "ProductName", + "ProdName", + "ProdCode", + "ProductCode", + "Item", + "ItemName", + "SKU", + ) + entity_alias = "ProductCount" + + dimension_candidates: tuple[str, ...] | None = None + if "market" in normalized_query or "markets" in normalized_query: + dimension_candidates = ("Market", "MarketType", "MarketName", "EndMarket", "Region", "Country") + elif "region" in normalized_query or "regions" in normalized_query: + dimension_candidates = ("Region", "Market", "Area", "Territory", "Country") + elif "country" in normalized_query or "countries" in normalized_query: + dimension_candidates = ("Country", "CountryName", "Nation", "Destination") + elif "division" in normalized_query: + dimension_candidates = ("Division",) + elif "business unit" in normalized_query or re.search(r"\bbu\b", normalized_query): + dimension_candidates = ("BusinessUnit", "Business Unit", "BU", "Division") + + if not entity_candidates or not dimension_candidates: + return None + + scored: list[tuple[int, dict[str, Any], str, str]] = [] + for table in tables: + dimension = self._find_schema_column(table, dimension_candidates) + entity = self._find_schema_column(table, entity_candidates) + if not dimension or not entity: + continue + + table_name = str(table.get("name") or "") + score = 20 + self._schema_source_shape_score(query, table) + if "order" in self._normalize_schema_token(table_name): + score += 15 + if "sales" in self._normalize_schema_token(table_name): + score += 10 + scored.append((score, table, dimension, entity)) + + if not scored: + return None + + _score, table, dimension, entity = sorted( + scored, + key=lambda item: item[0], + reverse=True, + )[0] + table_name = str(table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + entity_ref = f"{table_ref}.{self._quote_sql_identifier(entity)}" + where_clause = self._append_not_null_filters("", [dimension_ref, entity_ref]) + return ( + f"SELECT {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " + f"COUNT(DISTINCT {entity_ref}) AS {self._quote_sql_identifier(entity_alias)} " + f"FROM {table_ref}" + f"{where_clause} " + f"GROUP BY {dimension_ref} " + f"ORDER BY COUNT(DISTINCT {entity_ref}) DESC" + ) + def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2807,6 +2895,9 @@ def _build_schema_grounded_analytics_sql( if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): return entity_lookup_sql + if entity_distribution_sql := self._build_entity_distribution_sql(query, tables): + return entity_distribution_sql + if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2960,7 +3051,7 @@ def _build_schema_grounded_analytics_sql( if "business unit" in normalized_query or re.search(r"\bbu\b", normalized_query): dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) if "market" in normalized_query: - dimension_candidates.append(("Market", "MarketType", "MarketName", "Region", "Country")) + dimension_candidates.append(("Market", "MarketType", "MarketName", "EndMarket", "Region", "Country")) if "region" in normalized_query: dimension_candidates.append(("Region", "Market", "Area", "Territory")) if "currency" in normalized_query or "currencies" in normalized_query: @@ -3395,7 +3486,7 @@ def _build_schema_grounded_analytics_sql( rank_dimension = None if "market" in normalized_query: partition_dimension = self._find_schema_column( - table, ("Market", "MarketType", "Region") + table, ("Market", "MarketType", "MarketName", "EndMarket", "Region") ) if "region" in normalized_query and not partition_dimension: partition_dimension = self._find_schema_column( diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 102f183a39..a81a1202ed 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -42,6 +42,15 @@ def test_expand_business_terms_for_retrieval_adds_generic_currency_market_terms( assert "money exchange currency" in expanded_query +def test_expand_business_terms_for_retrieval_adds_domestic_international_market_terms(): + query = "Compare sales between domestic and international markets" + + expanded_query = expand_business_terms_for_retrieval(query) + + assert query in expanded_query + assert "domestic international" in expanded_query + + def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): query = "Explain what this workspace does" diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 0ee413f27c..f1fa8e0450 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -398,6 +398,59 @@ def test_schema_grounded_analytics_prefers_full_concept_coverage_for_order_marke ) +def test_schema_grounded_analytics_compares_sales_by_domestic_international_market(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Compare sales between domestic and international markets.", + [ + """ + CREATE TABLE dbo_qSalesMargin ( + EndMarket VARCHAR, + SalesValue DOUBLE, + OrdNo VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_qSalesMargin"."EndMarket" AS "EndMarket", ' + 'SUM("dbo_qSalesMargin"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_qSalesMargin" ' + 'WHERE "dbo_qSalesMargin"."EndMarket" IS NOT NULL ' + 'GROUP BY "dbo_qSalesMargin"."EndMarket" ' + 'ORDER BY SUM("dbo_qSalesMargin"."SalesValue") DESC' + ) + + +def test_schema_grounded_analytics_counts_customers_distribution_by_market(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Show customers distribution by market.", + [ + """ + CREATE TABLE dbo_tblNewOrders ( + Market VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_tblNewOrders"."Market" AS "Market", ' + 'COUNT(DISTINCT "dbo_tblNewOrders"."CustName") AS "CustomerCount" ' + 'FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."Market" IS NOT NULL ' + 'AND "dbo_tblNewOrders"."CustName" IS NOT NULL ' + 'GROUP BY "dbo_tblNewOrders"."Market" ' + 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."CustName") DESC' + ) + + def test_scope_retrieval_to_semantic_contract_prefers_production_full_coverage_table(): service = AskService.__new__(AskService) documents = [ From 1fbfec186436ae8c7329833c600ad501f8eb50eb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 16:09:49 +0530 Subject: [PATCH 0516/1087] Validate semantic comparison concepts before SQL execution --- wren-ai-service/src/web/v1/services/ask.py | 221 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 119 +++++++++- 2 files changed, 324 insertions(+), 16 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3de52cabbf..41186d82d0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1004,6 +1004,167 @@ def _sql_covers_required_question_concepts( return False return True + def _comparison_value_tokens_from_phrase(self, phrase: str) -> set[str]: + phrase = re.sub(r"\s+", " ", (phrase or "").strip().lower()) + if not phrase: + return set() + + dimension_words = { + "area", + "areas", + "business", + "categories", + "category", + "channels", + "channel", + "classes", + "class", + "countries", + "country", + "customer", + "customers", + "division", + "divisions", + "market", + "markets", + "product", + "products", + "region", + "regions", + "segment", + "segments", + "status", + "territories", + "territory", + "type", + "types", + "unit", + "units", + } + ignored_value_words = dimension_words | { + "amount", + "average", + "compare", + "comparison", + "count", + "counts", + "current", + "distribution", + "each", + "how", + "many", + "metric", + "metrics", + "number", + "order", + "orders", + "period", + "records", + "revenue", + "row", + "rows", + "sale", + "sales", + "sum", + "this", + "top", + "total", + "value", + "values", + } + + tokens = self._intent_tokens(phrase) - ignored_value_words + if not tokens: + return set() + + expanded = set(tokens) + if {"international", "intl"} & tokens: + expanded.update({"foreign", "global", "international", "intl", "overseas"}) + if {"domestic", "dom"} & tokens: + expanded.update({"domestic", "dom", "home", "local", "national"}) + return expanded + + def _requested_comparison_value_groups(self, query: str | None) -> list[set[str]]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return [] + + groups: list[set[str]] = [] + + for match in re.finditer( + r"\bbetween\s+(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)\s+" + r"(?:and|vs|versus)\s+" + r"(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)" + r"(?=$|[.?!,;]|\s+(?:by|for|from|in|over|during)\b)", + normalized_query, + ): + for phrase in (match.group("left"), match.group("right")): + tokens = self._comparison_value_tokens_from_phrase(phrase) + if tokens: + groups.append(tokens) + + if " between " not in normalized_query: + for match in re.finditer( + r"\bcompare\s+(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)\s+" + r"(?:and|with|to|vs|versus)\s+" + r"(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)" + r"(?=$|[.?!,;]|\s+(?:by|for|from|in|over|during)\b)", + normalized_query, + ): + for phrase in (match.group("left"), match.group("right")): + tokens = self._comparison_value_tokens_from_phrase(phrase) + if tokens: + groups.append(tokens) + + for match in re.finditer( + r"\b(?P[a-z0-9][a-z0-9 _/\-]{0,50}?)\s+" + r"(?:vs|versus)\s+" + r"(?P[a-z0-9][a-z0-9 _/\-]{0,50}?)" + r"(?=$|[.?!,;]|\s+(?:by|for|from|in|over|during)\b)", + normalized_query, + ): + for phrase in (match.group("left"), match.group("right")): + tokens = self._comparison_value_tokens_from_phrase(phrase) + if tokens: + groups.append(tokens) + + deduped: list[set[str]] = [] + seen: set[frozenset[str]] = set() + for group in groups: + key = frozenset(group) + if key in seen: + continue + seen.add(key) + deduped.append(group) + return deduped + + def _sql_satisfies_requested_comparison_values( + self, + sql: str, + query: str | None, + referenced_column_tokens: set[str], + ) -> bool: + comparison_groups = self._requested_comparison_value_groups(query) + if not comparison_groups: + return True + + sql_tokens = self._intent_tokens(sql or "") + grounded_tokens = sql_tokens | referenced_column_tokens + for comparison_group in comparison_groups: + if comparison_group & grounded_tokens: + continue + logger.warning( + "Ignoring SQL because a requested comparison value is not grounded " + "in the selected schema or generated SQL. query=%s required=%s " + "referenced_column_tokens=%s sql=%s", + query, + sorted(comparison_group), + sorted(referenced_column_tokens), + sql, + ) + return False + return True + def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> bool: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized_query: @@ -1288,7 +1449,20 @@ def _preferred_entity_lookup_columns( if "account" in normalized_query: candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) if "company" in normalized_query: - candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) + candidate_groups.append( + ( + "Company", + "CompanyName", + "CustName", + "CustomerName", + "ConsolidatedCustomer", + "SoldToCustomer", + "BillToCustomer", + "ShipToCustomer", + "EndCustomer", + "UltimateCustomer", + ) + ) candidate_groups.extend( [ ( @@ -1298,8 +1472,14 @@ def _preferred_entity_lookup_columns( "CustNo", "CustomerNo", "CustomerCode", + "ConsolidatedCustomer", + "SoldToCustomer", + "BillToCustomer", + "ShipToCustomer", + "EndCustomer", + "UltimateCustomer", ), - ("Client", "ClientName"), + ("Client", "ClientName", "Buyer", "BuyerName"), ("Account", "AccountName"), ("Company", "CompanyName"), ("Name",), @@ -1407,6 +1587,7 @@ def _invalid_unqualified_sql_identifiers( "is", "join", "limit", + "like", "month", "not", "null", @@ -1677,6 +1858,12 @@ def _sql_matches_question_intent( referenced_table_tokens, ): return False + if not self._sql_satisfies_requested_comparison_values( + sql, + query, + all_referenced_column_tokens, + ): + return False if not self._sql_uses_required_measure_aggregation(sql, query): return False if not self._sql_satisfies_count_ranking_request(sql, query): @@ -2737,7 +2924,17 @@ def _build_entity_lookup_sql( key=lambda column: ( 0 if self._normalize_schema_identifier_key(column) - in {"custname", "customername", "customer"} + in { + "billtocustomer", + "consolidatedcustomer", + "custname", + "customer", + "customername", + "endcustomer", + "shiptocustomer", + "soldtocustomer", + "ultimatecustomer", + } else 1, column.lower(), ), @@ -2787,7 +2984,7 @@ def _build_entity_lookup_sql( ) return ( f"SELECT TOP 500 * FROM {table_ref} " - f"WHERE {filter_ref} = '{escaped_phrase}'" + f"WHERE {filter_ref} LIKE '%{escaped_phrase}%'" f"{order_clause}" ) @@ -2808,10 +3005,18 @@ def _build_entity_distribution_sql( "CustNo", "CustomerNo", "CustomerCode", + "ConsolidatedCustomer", + "SoldToCustomer", + "BillToCustomer", + "ShipToCustomer", + "EndCustomer", + "UltimateCustomer", "Account", "AccountName", "Client", "ClientName", + "Buyer", + "BuyerName", ) entity_alias = "CustomerCount" elif "product" in normalized_query or "products" in normalized_query: @@ -3132,10 +3337,18 @@ def _build_schema_grounded_analytics_sql( "CustNo", "CustomerNo", "CustomerCode", + "ConsolidatedCustomer", + "SoldToCustomer", + "BillToCustomer", + "ShipToCustomer", + "EndCustomer", + "UltimateCustomer", "Account", "AccountName", "Client", "ClientName", + "Buyer", + "BuyerName", ) ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index f1fa8e0450..79b8fd0bfa 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -398,20 +398,21 @@ def test_schema_grounded_analytics_prefers_full_concept_coverage_for_order_marke ) -def test_schema_grounded_analytics_compares_sales_by_domestic_international_market(): +def test_sql_validation_rejects_unresolved_domestic_international_market_comparison(): service = AskService.__new__(AskService) + table_ddls = [ + """ + CREATE TABLE dbo_qSalesMargin ( + EndMarket VARCHAR, + SalesValue DOUBLE, + OrdNo VARCHAR + ); + """ + ] sql = service._build_schema_grounded_analytics_sql( "Compare sales between domestic and international markets.", - [ - """ - CREATE TABLE dbo_qSalesMargin ( - EndMarket VARCHAR, - SalesValue DOUBLE, - OrdNo VARCHAR - ); - """ - ], + table_ddls, ) assert sql == ( @@ -422,6 +423,47 @@ def test_schema_grounded_analytics_compares_sales_by_domestic_international_mark 'GROUP BY "dbo_qSalesMargin"."EndMarket" ' 'ORDER BY SUM("dbo_qSalesMargin"."SalesValue") DESC' ) + assert ( + service._build_validated_ask_result_from_sql( + sql, + table_ddls, + "Compare sales between domestic and international markets.", + ) + is None + ) + + +def test_sql_validation_accepts_grounded_domestic_international_market_mapping(): + service = AskService.__new__(AskService) + + table_ddls = [ + """ + CREATE TABLE dbo_qSalesMargin ( + EndMarket VARCHAR, + SalesValue DOUBLE + ); + """ + ] + sql = ( + 'SELECT CASE WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Domestic%\' ' + "THEN 'Domestic' " + 'WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Intl%\' ' + "THEN 'International' END AS \"MarketScope\", " + 'SUM("dbo_qSalesMargin"."SalesValue") AS "TotalSalesValue" ' + 'FROM "dbo_qSalesMargin" ' + 'WHERE "dbo_qSalesMargin"."EndMarket" LIKE \'%Domestic%\' ' + 'OR "dbo_qSalesMargin"."EndMarket" LIKE \'%Intl%\' ' + 'GROUP BY CASE WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Domestic%\' ' + "THEN 'Domestic' " + 'WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Intl%\' ' + "THEN 'International' END" + ) + + assert service._build_validated_ask_result_from_sql( + sql, + table_ddls, + "Compare sales between domestic and international markets.", + ) def test_schema_grounded_analytics_counts_customers_distribution_by_market(): @@ -451,6 +493,59 @@ def test_schema_grounded_analytics_counts_customers_distribution_by_market(): ) +def test_schema_grounded_analytics_counts_consolidated_customers_by_end_market(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Show customers distribution by market.", + [ + """ + CREATE TABLE dbo_qMarginSales ( + EndMarket VARCHAR, + ConsolidatedCustomer VARCHAR, + FXSales DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_qMarginSales"."EndMarket" AS "EndMarket", ' + 'COUNT(DISTINCT "dbo_qMarginSales"."ConsolidatedCustomer") ' + 'AS "CustomerCount" ' + 'FROM "dbo_qMarginSales" ' + 'WHERE "dbo_qMarginSales"."EndMarket" IS NOT NULL ' + 'AND "dbo_qMarginSales"."ConsolidatedCustomer" IS NOT NULL ' + 'GROUP BY "dbo_qMarginSales"."EndMarket" ' + 'ORDER BY COUNT(DISTINCT "dbo_qMarginSales"."ConsolidatedCustomer") DESC' + ) + + +def test_entity_lookup_uses_customer_like_match_for_company_phrase(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_analytics_sql( + "Show me orders for Lockheed Martin.", + [ + """ + CREATE TABLE dbo_tblNewOrders ( + ConsolidatedCustomer VARCHAR, + OrdNo VARCHAR, + OrdDate DATETIME, + Market VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."ConsolidatedCustomer" LIKE ' + "'%Lockheed Martin%' " + 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' + ) + + def test_scope_retrieval_to_semantic_contract_prefers_production_full_coverage_table(): service = AskService.__new__(AskService) documents = [ @@ -634,8 +729,8 @@ def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): assert sql == ( 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" = ' - "'Daimler Trucks North America' " + 'WHERE "dbo_tblNewOrders"."CustName" LIKE ' + "'%Daimler Trucks North America%' " 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' ) From 4f68dd8fc7b6fd2e7bcccf777ab59a91eb75a09a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 16:24:12 +0530 Subject: [PATCH 0517/1087] Enforce scoped SQL contract for data questions --- wren-ai-service/src/web/v1/services/ask.py | 345 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 143 +++++++- 2 files changed, 485 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 41186d82d0..760ca8d47b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2977,17 +2977,101 @@ def _build_entity_lookup_sql( ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), temporal=True, ) + projected_columns = self._entity_lookup_projection_columns( + query, + table, + filter_column=filter_column, + date_column=date_column, + ) + select_sql = ", ".join( + f"{table_ref}.{self._quote_sql_identifier(column)} AS {self._quote_sql_identifier(column)}" + for column in projected_columns + ) order_clause = ( f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" if date_column else "" ) return ( - f"SELECT TOP 500 * FROM {table_ref} " + f"SELECT TOP 500 {select_sql} FROM {table_ref} " f"WHERE {filter_ref} LIKE '%{escaped_phrase}%'" f"{order_clause}" ) + def _entity_lookup_projection_columns( + self, + query: str, + table: dict[str, Any], + *, + filter_column: str, + date_column: str | None, + ) -> list[str]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + selected: list[str] = [] + selected_keys: set[str] = set() + + def add_column(column: str | None) -> None: + if not column: + return + key = self._normalize_schema_identifier_key(column) + if not key or key in selected_keys: + return + selected.append(column) + selected_keys.add(key) + + def add_first(candidates: tuple[str, ...]) -> None: + add_column(self._find_schema_column(table, candidates)) + + asks_for_orders = any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ) + if asks_for_orders: + add_first(("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber")) + add_column(filter_column) + add_column(date_column) + add_first( + ( + "CustNo", + "CustomerNo", + "CustomerCode", + "AccountNo", + "AccountId", + "ClientNo", + ) + ) + add_first(("CustPO", "CustomerPO", "PONo", "PurchaseOrder")) + add_first(("Market", "MarketType", "MarketName", "EndMarket", "Region", "Country")) + add_first(("Division", "BusinessUnit", "Business Unit", "BU")) + add_first( + ( + "Product", + "ProductName", + "ProdName", + "ProdCode", + "ProductCode", + "Item", + "ItemName", + "ItemProductCode", + "SKU", + ) + ) + add_first( + ( + "SalesValue", + "OrderValue", + "NewOrderValue", + "Amount", + "Revenue", + "Cost", + "Value", + ) + ) + + if not selected: + add_column(filter_column) + return selected[:10] + def _build_entity_distribution_sql( self, query: str, tables: list[dict[str, Any]] ) -> str | None: @@ -6384,6 +6468,242 @@ def _is_valid_select_sql(self, sql: Optional[str]) -> bool: return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) + def _query_explicitly_requests_all_fields(self, query: str | None) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + patterns = ( + r"\bselect\s+\*", + r"\ball\s+(?:columns|fields)\b", + r"\bevery\s+(?:column|field)\b", + r"\b(?:whole|entire|complete)\s+(?:row|rows|record|records|table)\b", + r"\braw\s+(?:rows|records|data)\b", + r"\bpreview\s+(?:rows|records|data)\b", + r"\bsample\s+(?:rows|records|data)\b", + r"\bfirst\s+\d*\s*(?:rows|records)\b", + r"\btop\s+\d+\s+(?:rows|records)\b", + ) + return any(re.search(pattern, normalized) for pattern in patterns) + + def _sql_uses_select_star(self, sql: str) -> bool: + sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") + return bool( + re.search( + r"\bSELECT\s+(?:TOP\s+\d+\s+)?\*", + sql_without_strings, + flags=re.IGNORECASE, + ) + or re.search( + r"(?:^|,)\s*(?:\"[^\"]+\"|\[[^\]]+\]|`[^`]+`|" + r"[A-Za-z_][A-Za-z0-9_.$]*)\s*\.\s*\*", + sql_without_strings, + flags=re.IGNORECASE, + ) + ) + + def _sql_contains_placeholder_values(self, sql: str) -> bool: + normalized_sql = re.sub(r"\s+", " ", sql or "").lower() + placeholder_patterns = ( + r"\bdesired(?:_[a-z0-9]+)*\b", + r"\bspecified(?:_[a-z0-9]+)*\b", + r"\bplaceholder\b", + r"\byour_[a-z0-9_]+\b", + r"\bexample_[a-z0-9_]+\b", + r"<\s*[a-z0-9_ -]+\s*>", + r"\{\s*[a-z0-9_ -]+\s*\}", + r"\byyyy[-_/]?mm[-_/]?dd\b", + ) + return any( + re.search(pattern, normalized_sql, flags=re.IGNORECASE) + for pattern in placeholder_patterns + ) + + def _extract_schema_relationship_edges( + self, table_ddls: list[str] + ) -> set[tuple[tuple[str, str], tuple[str, str]]]: + edges: set[tuple[tuple[str, str], tuple[str, str]]] = set() + + def clean_identifier(value: str) -> str: + return str(value or "").strip().strip('"[]` ') + + def split_columns(value: str) -> list[str]: + return [ + clean_identifier(part) + for part in (value or "").split(",") + if clean_identifier(part) + ] + + for ddl in table_ddls or []: + if not isinstance(ddl, str): + continue + table_match = re.search( + r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", + ddl, + flags=re.IGNORECASE, + ) + if not table_match: + continue + source_table = clean_identifier( + next( + (value for value in table_match.groupdict().values() if value), + "", + ) + ) + if not source_table: + continue + + for relationship_match in re.finditer( + r"FOREIGN\s+KEY\s*\((?P[^)]+)\)\s+REFERENCES\s+" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" + r"\s*\((?P[^)]+)\)", + ddl, + flags=re.IGNORECASE, + ): + target_table = clean_identifier( + next( + ( + value + for key, value in relationship_match.groupdict().items() + if key + in { + "quoted", + "bracketed", + "backticked", + "bare", + } + and value + ), + "", + ) + ) + if not target_table: + continue + + source_columns = split_columns( + relationship_match.group("source_columns") + ) + target_columns = split_columns( + relationship_match.group("target_columns") + ) + for source_column, target_column in zip(source_columns, target_columns): + source_edge = ( + self._normalize_schema_token(source_table), + self._normalize_schema_token(source_column), + ) + target_edge = ( + self._normalize_schema_token(target_table), + self._normalize_schema_token(target_column), + ) + edges.add((source_edge, target_edge)) + edges.add((target_edge, source_edge)) + + return edges + + def _sql_joins_match_selected_relationships( + self, + sql: str, + table_ddls: list[str], + valid_tables: dict[str, dict[str, Any]], + ) -> bool: + sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") + if not re.search(r"\bJOIN\b", sql_without_strings, flags=re.IGNORECASE): + return True + + equality_pattern = re.compile( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' + r"(?P[A-Za-z_][A-Za-z0-9_$]*))\s*=\s*" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' + r"(?P[A-Za-z_][A-Za-z0-9_$]*))", + flags=re.IGNORECASE, + ) + equality_matches = list(equality_pattern.finditer(sql_without_strings)) + if not equality_matches: + logger.warning( + "Ignoring SQL because JOIN has no qualified equality condition. sql=%s", + sql, + ) + return False + + relationship_edges = self._extract_schema_relationship_edges(table_ddls) + if not relationship_edges: + return True + + invalid_edges: list[str] = [] + checked_edges = 0 + for match in equality_matches: + left_table_reference = ( + match.group("ltq") + or match.group("ltb") + or match.group("ltk") + or match.group("lt") + or "" + ) + left_column = ( + match.group("lcq") + or match.group("lcb") + or match.group("lck") + or match.group("lc") + or "" + ) + right_table_reference = ( + match.group("rtq") + or match.group("rtb") + or match.group("rtk") + or match.group("rt") + or "" + ) + right_column = ( + match.group("rcq") + or match.group("rcb") + or match.group("rck") + or match.group("rc") + or "" + ) + left_table = self._table_for_sql_reference( + left_table_reference, valid_tables + ) + right_table = self._table_for_sql_reference( + right_table_reference, valid_tables + ) + if not left_table or not right_table: + continue + + left_key = ( + self._normalize_schema_token(str(left_table.get("name") or "")), + self._normalize_schema_token(left_column), + ) + right_key = ( + self._normalize_schema_token(str(right_table.get("name") or "")), + self._normalize_schema_token(right_column), + ) + if left_key[0] == right_key[0]: + continue + + checked_edges += 1 + if (left_key, right_key) not in relationship_edges: + invalid_edges.append( + f"{left_table_reference}.{left_column} = " + f"{right_table_reference}.{right_column}" + ) + + if checked_edges and invalid_edges: + logger.warning( + "Ignoring SQL because JOIN condition is not grounded in selected schema relationships. " + "invalid_edges=%s sql=%s", + invalid_edges, + sql, + ) + return False + return True + def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: if not self._is_valid_select_sql(sql): return None @@ -6408,6 +6728,22 @@ def _build_validated_ask_result_from_sql( ask_result = self._build_ask_result_from_sql(sql) if not ask_result: return None + if self._sql_contains_placeholder_values(ask_result.sql): + logger.warning( + "Ignoring SQL because it contains placeholder filter values. sql=%s", + ask_result.sql, + ) + return None + if self._sql_uses_select_star( + ask_result.sql + ) and not self._query_explicitly_requests_all_fields(query): + logger.warning( + "Ignoring SQL because it uses SELECT * without an explicit all-fields request. " + "query=%s sql=%s", + query, + ask_result.sql, + ) + return None schema_tables = self._parse_schema_tables(table_ddls) valid_tables = { @@ -6490,6 +6826,13 @@ def _build_validated_ask_result_from_sql( ) return None + if not self._sql_joins_match_selected_relationships( + ask_result.sql, + table_ddls, + valid_tables, + ): + return None + invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( ask_result.sql, schema_tables, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 79b8fd0bfa..c56949f749 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -539,7 +539,11 @@ def test_entity_lookup_uses_customer_like_match_for_company_phrase(): ) assert sql == ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'SELECT TOP 500 "dbo_tblNewOrders"."OrdNo" AS "OrdNo", ' + '"dbo_tblNewOrders"."ConsolidatedCustomer" AS "ConsolidatedCustomer", ' + '"dbo_tblNewOrders"."OrdDate" AS "OrdDate", ' + '"dbo_tblNewOrders"."Market" AS "Market" ' + 'FROM "dbo_tblNewOrders" ' 'WHERE "dbo_tblNewOrders"."ConsolidatedCustomer" LIKE ' "'%Lockheed Martin%' " 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' @@ -728,13 +732,148 @@ def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): ) assert sql == ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'SELECT TOP 500 "dbo_tblNewOrders"."OrdNo" AS "OrdNo", ' + '"dbo_tblNewOrders"."CustName" AS "CustName", ' + '"dbo_tblNewOrders"."OrdDate" AS "OrdDate" ' + 'FROM "dbo_tblNewOrders" ' 'WHERE "dbo_tblNewOrders"."CustName" LIKE ' "'%Daimler Trucks North America%' " 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' ) +def test_validated_sql_rejects_select_star_without_explicit_all_fields_request(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."CustName" LIKE \'%Lockheed Martin%\'' + ), + [ + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """ + ], + "Show me orders for Lockheed Martin.", + ) + + assert result is None + + +def test_validated_sql_allows_select_star_for_explicit_all_fields_request(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + 'SELECT TOP 10 * FROM "dbo_tblNewOrders"', + [ + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """ + ], + "Show all columns from dbo_tblNewOrders.", + ) + + assert result is not None + + +def test_validated_sql_rejects_placeholder_filter_values(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_tblNewOrders"."OrdNo" AS "OrdNo" ' + 'FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."OrdDate" = \'desired_date\'' + ), + [ + """ + CREATE TABLE dbo_tblNewOrders ( + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """ + ], + "Show orders for the requested date.", + ) + + assert result is None + + +def test_validated_sql_rejects_join_not_grounded_in_selected_relationships(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_Orders"."OrdNo" AS "OrdNo", ' + '"dbo_Customers"."CustomerName" AS "CustomerName" ' + 'FROM "dbo_Orders" ' + 'JOIN "dbo_Customers" ' + 'ON "dbo_Orders"."ProductId" = "dbo_Customers"."CustomerId"' + ), + [ + """ + CREATE TABLE dbo_Orders ( + OrdNo VARCHAR, + CustomerId VARCHAR, + ProductId VARCHAR, + FOREIGN KEY (CustomerId) REFERENCES dbo_Customers(CustomerId) + ); + """, + """ + CREATE TABLE dbo_Customers ( + CustomerId VARCHAR, + CustomerName VARCHAR + ); + """, + ], + "Show customer names for orders.", + ) + + assert result is None + + +def test_validated_sql_accepts_join_grounded_in_selected_relationships(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_Orders"."OrdNo" AS "OrdNo", ' + '"dbo_Customers"."CustomerName" AS "CustomerName" ' + 'FROM "dbo_Orders" ' + 'JOIN "dbo_Customers" ' + 'ON "dbo_Orders"."CustomerId" = "dbo_Customers"."CustomerId"' + ), + [ + """ + CREATE TABLE dbo_Orders ( + OrdNo VARCHAR, + CustomerId VARCHAR, + ProductId VARCHAR, + FOREIGN KEY (CustomerId) REFERENCES dbo_Customers(CustomerId) + ); + """, + """ + CREATE TABLE dbo_Customers ( + CustomerId VARCHAR, + CustomerName VARCHAR + ); + """, + ], + "Show customer names for orders.", + ) + + assert result is not None + + def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) From 5748daefe895f54119b8de61b0ab4bc79e2bbb7b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 17:51:02 +0530 Subject: [PATCH 0518/1087] Revert "Enforce scoped SQL contract for data questions" This reverts commit 4f68dd8fc7b6fd2e7bcccf777ab59a91eb75a09a. --- wren-ai-service/src/web/v1/services/ask.py | 345 +----------------- .../pytest/services/test_ask_sales_sql.py | 143 +------- 2 files changed, 3 insertions(+), 485 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 760ca8d47b..41186d82d0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2977,101 +2977,17 @@ def _build_entity_lookup_sql( ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), temporal=True, ) - projected_columns = self._entity_lookup_projection_columns( - query, - table, - filter_column=filter_column, - date_column=date_column, - ) - select_sql = ", ".join( - f"{table_ref}.{self._quote_sql_identifier(column)} AS {self._quote_sql_identifier(column)}" - for column in projected_columns - ) order_clause = ( f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" if date_column else "" ) return ( - f"SELECT TOP 500 {select_sql} FROM {table_ref} " + f"SELECT TOP 500 * FROM {table_ref} " f"WHERE {filter_ref} LIKE '%{escaped_phrase}%'" f"{order_clause}" ) - def _entity_lookup_projection_columns( - self, - query: str, - table: dict[str, Any], - *, - filter_column: str, - date_column: str | None, - ) -> list[str]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - selected: list[str] = [] - selected_keys: set[str] = set() - - def add_column(column: str | None) -> None: - if not column: - return - key = self._normalize_schema_identifier_key(column) - if not key or key in selected_keys: - return - selected.append(column) - selected_keys.add(key) - - def add_first(candidates: tuple[str, ...]) -> None: - add_column(self._find_schema_column(table, candidates)) - - asks_for_orders = any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ) - if asks_for_orders: - add_first(("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber")) - add_column(filter_column) - add_column(date_column) - add_first( - ( - "CustNo", - "CustomerNo", - "CustomerCode", - "AccountNo", - "AccountId", - "ClientNo", - ) - ) - add_first(("CustPO", "CustomerPO", "PONo", "PurchaseOrder")) - add_first(("Market", "MarketType", "MarketName", "EndMarket", "Region", "Country")) - add_first(("Division", "BusinessUnit", "Business Unit", "BU")) - add_first( - ( - "Product", - "ProductName", - "ProdName", - "ProdCode", - "ProductCode", - "Item", - "ItemName", - "ItemProductCode", - "SKU", - ) - ) - add_first( - ( - "SalesValue", - "OrderValue", - "NewOrderValue", - "Amount", - "Revenue", - "Cost", - "Value", - ) - ) - - if not selected: - add_column(filter_column) - return selected[:10] - def _build_entity_distribution_sql( self, query: str, tables: list[dict[str, Any]] ) -> str | None: @@ -6468,242 +6384,6 @@ def _is_valid_select_sql(self, sql: Optional[str]) -> bool: return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) - def _query_explicitly_requests_all_fields(self, query: str | None) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - patterns = ( - r"\bselect\s+\*", - r"\ball\s+(?:columns|fields)\b", - r"\bevery\s+(?:column|field)\b", - r"\b(?:whole|entire|complete)\s+(?:row|rows|record|records|table)\b", - r"\braw\s+(?:rows|records|data)\b", - r"\bpreview\s+(?:rows|records|data)\b", - r"\bsample\s+(?:rows|records|data)\b", - r"\bfirst\s+\d*\s*(?:rows|records)\b", - r"\btop\s+\d+\s+(?:rows|records)\b", - ) - return any(re.search(pattern, normalized) for pattern in patterns) - - def _sql_uses_select_star(self, sql: str) -> bool: - sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") - return bool( - re.search( - r"\bSELECT\s+(?:TOP\s+\d+\s+)?\*", - sql_without_strings, - flags=re.IGNORECASE, - ) - or re.search( - r"(?:^|,)\s*(?:\"[^\"]+\"|\[[^\]]+\]|`[^`]+`|" - r"[A-Za-z_][A-Za-z0-9_.$]*)\s*\.\s*\*", - sql_without_strings, - flags=re.IGNORECASE, - ) - ) - - def _sql_contains_placeholder_values(self, sql: str) -> bool: - normalized_sql = re.sub(r"\s+", " ", sql or "").lower() - placeholder_patterns = ( - r"\bdesired(?:_[a-z0-9]+)*\b", - r"\bspecified(?:_[a-z0-9]+)*\b", - r"\bplaceholder\b", - r"\byour_[a-z0-9_]+\b", - r"\bexample_[a-z0-9_]+\b", - r"<\s*[a-z0-9_ -]+\s*>", - r"\{\s*[a-z0-9_ -]+\s*\}", - r"\byyyy[-_/]?mm[-_/]?dd\b", - ) - return any( - re.search(pattern, normalized_sql, flags=re.IGNORECASE) - for pattern in placeholder_patterns - ) - - def _extract_schema_relationship_edges( - self, table_ddls: list[str] - ) -> set[tuple[tuple[str, str], tuple[str, str]]]: - edges: set[tuple[tuple[str, str], tuple[str, str]]] = set() - - def clean_identifier(value: str) -> str: - return str(value or "").strip().strip('"[]` ') - - def split_columns(value: str) -> list[str]: - return [ - clean_identifier(part) - for part in (value or "").split(",") - if clean_identifier(part) - ] - - for ddl in table_ddls or []: - if not isinstance(ddl, str): - continue - table_match = re.search( - r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", - ddl, - flags=re.IGNORECASE, - ) - if not table_match: - continue - source_table = clean_identifier( - next( - (value for value in table_match.groupdict().values() if value), - "", - ) - ) - if not source_table: - continue - - for relationship_match in re.finditer( - r"FOREIGN\s+KEY\s*\((?P[^)]+)\)\s+REFERENCES\s+" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" - r"\s*\((?P[^)]+)\)", - ddl, - flags=re.IGNORECASE, - ): - target_table = clean_identifier( - next( - ( - value - for key, value in relationship_match.groupdict().items() - if key - in { - "quoted", - "bracketed", - "backticked", - "bare", - } - and value - ), - "", - ) - ) - if not target_table: - continue - - source_columns = split_columns( - relationship_match.group("source_columns") - ) - target_columns = split_columns( - relationship_match.group("target_columns") - ) - for source_column, target_column in zip(source_columns, target_columns): - source_edge = ( - self._normalize_schema_token(source_table), - self._normalize_schema_token(source_column), - ) - target_edge = ( - self._normalize_schema_token(target_table), - self._normalize_schema_token(target_column), - ) - edges.add((source_edge, target_edge)) - edges.add((target_edge, source_edge)) - - return edges - - def _sql_joins_match_selected_relationships( - self, - sql: str, - table_ddls: list[str], - valid_tables: dict[str, dict[str, Any]], - ) -> bool: - sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") - if not re.search(r"\bJOIN\b", sql_without_strings, flags=re.IGNORECASE): - return True - - equality_pattern = re.compile( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' - r"(?P[A-Za-z_][A-Za-z0-9_$]*))\s*=\s*" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|`(?P[^`]+)`|' - r"(?P[A-Za-z_][A-Za-z0-9_$]*))", - flags=re.IGNORECASE, - ) - equality_matches = list(equality_pattern.finditer(sql_without_strings)) - if not equality_matches: - logger.warning( - "Ignoring SQL because JOIN has no qualified equality condition. sql=%s", - sql, - ) - return False - - relationship_edges = self._extract_schema_relationship_edges(table_ddls) - if not relationship_edges: - return True - - invalid_edges: list[str] = [] - checked_edges = 0 - for match in equality_matches: - left_table_reference = ( - match.group("ltq") - or match.group("ltb") - or match.group("ltk") - or match.group("lt") - or "" - ) - left_column = ( - match.group("lcq") - or match.group("lcb") - or match.group("lck") - or match.group("lc") - or "" - ) - right_table_reference = ( - match.group("rtq") - or match.group("rtb") - or match.group("rtk") - or match.group("rt") - or "" - ) - right_column = ( - match.group("rcq") - or match.group("rcb") - or match.group("rck") - or match.group("rc") - or "" - ) - left_table = self._table_for_sql_reference( - left_table_reference, valid_tables - ) - right_table = self._table_for_sql_reference( - right_table_reference, valid_tables - ) - if not left_table or not right_table: - continue - - left_key = ( - self._normalize_schema_token(str(left_table.get("name") or "")), - self._normalize_schema_token(left_column), - ) - right_key = ( - self._normalize_schema_token(str(right_table.get("name") or "")), - self._normalize_schema_token(right_column), - ) - if left_key[0] == right_key[0]: - continue - - checked_edges += 1 - if (left_key, right_key) not in relationship_edges: - invalid_edges.append( - f"{left_table_reference}.{left_column} = " - f"{right_table_reference}.{right_column}" - ) - - if checked_edges and invalid_edges: - logger.warning( - "Ignoring SQL because JOIN condition is not grounded in selected schema relationships. " - "invalid_edges=%s sql=%s", - invalid_edges, - sql, - ) - return False - return True - def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: if not self._is_valid_select_sql(sql): return None @@ -6728,22 +6408,6 @@ def _build_validated_ask_result_from_sql( ask_result = self._build_ask_result_from_sql(sql) if not ask_result: return None - if self._sql_contains_placeholder_values(ask_result.sql): - logger.warning( - "Ignoring SQL because it contains placeholder filter values. sql=%s", - ask_result.sql, - ) - return None - if self._sql_uses_select_star( - ask_result.sql - ) and not self._query_explicitly_requests_all_fields(query): - logger.warning( - "Ignoring SQL because it uses SELECT * without an explicit all-fields request. " - "query=%s sql=%s", - query, - ask_result.sql, - ) - return None schema_tables = self._parse_schema_tables(table_ddls) valid_tables = { @@ -6826,13 +6490,6 @@ def _build_validated_ask_result_from_sql( ) return None - if not self._sql_joins_match_selected_relationships( - ask_result.sql, - table_ddls, - valid_tables, - ): - return None - invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( ask_result.sql, schema_tables, diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index c56949f749..79b8fd0bfa 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -539,11 +539,7 @@ def test_entity_lookup_uses_customer_like_match_for_company_phrase(): ) assert sql == ( - 'SELECT TOP 500 "dbo_tblNewOrders"."OrdNo" AS "OrdNo", ' - '"dbo_tblNewOrders"."ConsolidatedCustomer" AS "ConsolidatedCustomer", ' - '"dbo_tblNewOrders"."OrdDate" AS "OrdDate", ' - '"dbo_tblNewOrders"."Market" AS "Market" ' - 'FROM "dbo_tblNewOrders" ' + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' 'WHERE "dbo_tblNewOrders"."ConsolidatedCustomer" LIKE ' "'%Lockheed Martin%' " 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' @@ -732,148 +728,13 @@ def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): ) assert sql == ( - 'SELECT TOP 500 "dbo_tblNewOrders"."OrdNo" AS "OrdNo", ' - '"dbo_tblNewOrders"."CustName" AS "CustName", ' - '"dbo_tblNewOrders"."OrdDate" AS "OrdDate" ' - 'FROM "dbo_tblNewOrders" ' + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' 'WHERE "dbo_tblNewOrders"."CustName" LIKE ' "'%Daimler Trucks North America%' " 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' ) -def test_validated_sql_rejects_select_star_without_explicit_all_fields_request(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" LIKE \'%Lockheed Martin%\'' - ), - [ - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """ - ], - "Show me orders for Lockheed Martin.", - ) - - assert result is None - - -def test_validated_sql_allows_select_star_for_explicit_all_fields_request(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - 'SELECT TOP 10 * FROM "dbo_tblNewOrders"', - [ - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """ - ], - "Show all columns from dbo_tblNewOrders.", - ) - - assert result is not None - - -def test_validated_sql_rejects_placeholder_filter_values(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_tblNewOrders"."OrdNo" AS "OrdNo" ' - 'FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."OrdDate" = \'desired_date\'' - ), - [ - """ - CREATE TABLE dbo_tblNewOrders ( - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """ - ], - "Show orders for the requested date.", - ) - - assert result is None - - -def test_validated_sql_rejects_join_not_grounded_in_selected_relationships(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_Orders"."OrdNo" AS "OrdNo", ' - '"dbo_Customers"."CustomerName" AS "CustomerName" ' - 'FROM "dbo_Orders" ' - 'JOIN "dbo_Customers" ' - 'ON "dbo_Orders"."ProductId" = "dbo_Customers"."CustomerId"' - ), - [ - """ - CREATE TABLE dbo_Orders ( - OrdNo VARCHAR, - CustomerId VARCHAR, - ProductId VARCHAR, - FOREIGN KEY (CustomerId) REFERENCES dbo_Customers(CustomerId) - ); - """, - """ - CREATE TABLE dbo_Customers ( - CustomerId VARCHAR, - CustomerName VARCHAR - ); - """, - ], - "Show customer names for orders.", - ) - - assert result is None - - -def test_validated_sql_accepts_join_grounded_in_selected_relationships(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_Orders"."OrdNo" AS "OrdNo", ' - '"dbo_Customers"."CustomerName" AS "CustomerName" ' - 'FROM "dbo_Orders" ' - 'JOIN "dbo_Customers" ' - 'ON "dbo_Orders"."CustomerId" = "dbo_Customers"."CustomerId"' - ), - [ - """ - CREATE TABLE dbo_Orders ( - OrdNo VARCHAR, - CustomerId VARCHAR, - ProductId VARCHAR, - FOREIGN KEY (CustomerId) REFERENCES dbo_Customers(CustomerId) - ); - """, - """ - CREATE TABLE dbo_Customers ( - CustomerId VARCHAR, - CustomerName VARCHAR - ); - """, - ], - "Show customer names for orders.", - ) - - assert result is not None - - def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) From bd312762689ed7a65aa534dfc516eb5ee84c36eb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 17:51:02 +0530 Subject: [PATCH 0519/1087] Revert "Validate semantic comparison concepts before SQL execution" This reverts commit 1fbfec186436ae8c7329833c600ad501f8eb50eb. --- wren-ai-service/src/web/v1/services/ask.py | 221 +----------------- .../pytest/services/test_ask_sales_sql.py | 119 +--------- 2 files changed, 16 insertions(+), 324 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 41186d82d0..3de52cabbf 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1004,167 +1004,6 @@ def _sql_covers_required_question_concepts( return False return True - def _comparison_value_tokens_from_phrase(self, phrase: str) -> set[str]: - phrase = re.sub(r"\s+", " ", (phrase or "").strip().lower()) - if not phrase: - return set() - - dimension_words = { - "area", - "areas", - "business", - "categories", - "category", - "channels", - "channel", - "classes", - "class", - "countries", - "country", - "customer", - "customers", - "division", - "divisions", - "market", - "markets", - "product", - "products", - "region", - "regions", - "segment", - "segments", - "status", - "territories", - "territory", - "type", - "types", - "unit", - "units", - } - ignored_value_words = dimension_words | { - "amount", - "average", - "compare", - "comparison", - "count", - "counts", - "current", - "distribution", - "each", - "how", - "many", - "metric", - "metrics", - "number", - "order", - "orders", - "period", - "records", - "revenue", - "row", - "rows", - "sale", - "sales", - "sum", - "this", - "top", - "total", - "value", - "values", - } - - tokens = self._intent_tokens(phrase) - ignored_value_words - if not tokens: - return set() - - expanded = set(tokens) - if {"international", "intl"} & tokens: - expanded.update({"foreign", "global", "international", "intl", "overseas"}) - if {"domestic", "dom"} & tokens: - expanded.update({"domestic", "dom", "home", "local", "national"}) - return expanded - - def _requested_comparison_value_groups(self, query: str | None) -> list[set[str]]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return [] - - groups: list[set[str]] = [] - - for match in re.finditer( - r"\bbetween\s+(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)\s+" - r"(?:and|vs|versus)\s+" - r"(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)" - r"(?=$|[.?!,;]|\s+(?:by|for|from|in|over|during)\b)", - normalized_query, - ): - for phrase in (match.group("left"), match.group("right")): - tokens = self._comparison_value_tokens_from_phrase(phrase) - if tokens: - groups.append(tokens) - - if " between " not in normalized_query: - for match in re.finditer( - r"\bcompare\s+(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)\s+" - r"(?:and|with|to|vs|versus)\s+" - r"(?P[a-z0-9][a-z0-9 _/\-]{0,80}?)" - r"(?=$|[.?!,;]|\s+(?:by|for|from|in|over|during)\b)", - normalized_query, - ): - for phrase in (match.group("left"), match.group("right")): - tokens = self._comparison_value_tokens_from_phrase(phrase) - if tokens: - groups.append(tokens) - - for match in re.finditer( - r"\b(?P[a-z0-9][a-z0-9 _/\-]{0,50}?)\s+" - r"(?:vs|versus)\s+" - r"(?P[a-z0-9][a-z0-9 _/\-]{0,50}?)" - r"(?=$|[.?!,;]|\s+(?:by|for|from|in|over|during)\b)", - normalized_query, - ): - for phrase in (match.group("left"), match.group("right")): - tokens = self._comparison_value_tokens_from_phrase(phrase) - if tokens: - groups.append(tokens) - - deduped: list[set[str]] = [] - seen: set[frozenset[str]] = set() - for group in groups: - key = frozenset(group) - if key in seen: - continue - seen.add(key) - deduped.append(group) - return deduped - - def _sql_satisfies_requested_comparison_values( - self, - sql: str, - query: str | None, - referenced_column_tokens: set[str], - ) -> bool: - comparison_groups = self._requested_comparison_value_groups(query) - if not comparison_groups: - return True - - sql_tokens = self._intent_tokens(sql or "") - grounded_tokens = sql_tokens | referenced_column_tokens - for comparison_group in comparison_groups: - if comparison_group & grounded_tokens: - continue - logger.warning( - "Ignoring SQL because a requested comparison value is not grounded " - "in the selected schema or generated SQL. query=%s required=%s " - "referenced_column_tokens=%s sql=%s", - query, - sorted(comparison_group), - sorted(referenced_column_tokens), - sql, - ) - return False - return True - def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> bool: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized_query: @@ -1449,20 +1288,7 @@ def _preferred_entity_lookup_columns( if "account" in normalized_query: candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) if "company" in normalized_query: - candidate_groups.append( - ( - "Company", - "CompanyName", - "CustName", - "CustomerName", - "ConsolidatedCustomer", - "SoldToCustomer", - "BillToCustomer", - "ShipToCustomer", - "EndCustomer", - "UltimateCustomer", - ) - ) + candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) candidate_groups.extend( [ ( @@ -1472,14 +1298,8 @@ def _preferred_entity_lookup_columns( "CustNo", "CustomerNo", "CustomerCode", - "ConsolidatedCustomer", - "SoldToCustomer", - "BillToCustomer", - "ShipToCustomer", - "EndCustomer", - "UltimateCustomer", ), - ("Client", "ClientName", "Buyer", "BuyerName"), + ("Client", "ClientName"), ("Account", "AccountName"), ("Company", "CompanyName"), ("Name",), @@ -1587,7 +1407,6 @@ def _invalid_unqualified_sql_identifiers( "is", "join", "limit", - "like", "month", "not", "null", @@ -1858,12 +1677,6 @@ def _sql_matches_question_intent( referenced_table_tokens, ): return False - if not self._sql_satisfies_requested_comparison_values( - sql, - query, - all_referenced_column_tokens, - ): - return False if not self._sql_uses_required_measure_aggregation(sql, query): return False if not self._sql_satisfies_count_ranking_request(sql, query): @@ -2924,17 +2737,7 @@ def _build_entity_lookup_sql( key=lambda column: ( 0 if self._normalize_schema_identifier_key(column) - in { - "billtocustomer", - "consolidatedcustomer", - "custname", - "customer", - "customername", - "endcustomer", - "shiptocustomer", - "soldtocustomer", - "ultimatecustomer", - } + in {"custname", "customername", "customer"} else 1, column.lower(), ), @@ -2984,7 +2787,7 @@ def _build_entity_lookup_sql( ) return ( f"SELECT TOP 500 * FROM {table_ref} " - f"WHERE {filter_ref} LIKE '%{escaped_phrase}%'" + f"WHERE {filter_ref} = '{escaped_phrase}'" f"{order_clause}" ) @@ -3005,18 +2808,10 @@ def _build_entity_distribution_sql( "CustNo", "CustomerNo", "CustomerCode", - "ConsolidatedCustomer", - "SoldToCustomer", - "BillToCustomer", - "ShipToCustomer", - "EndCustomer", - "UltimateCustomer", "Account", "AccountName", "Client", "ClientName", - "Buyer", - "BuyerName", ) entity_alias = "CustomerCount" elif "product" in normalized_query or "products" in normalized_query: @@ -3337,18 +3132,10 @@ def _build_schema_grounded_analytics_sql( "CustNo", "CustomerNo", "CustomerCode", - "ConsolidatedCustomer", - "SoldToCustomer", - "BillToCustomer", - "ShipToCustomer", - "EndCustomer", - "UltimateCustomer", "Account", "AccountName", "Client", "ClientName", - "Buyer", - "BuyerName", ) ) diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 79b8fd0bfa..f1fa8e0450 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -398,21 +398,20 @@ def test_schema_grounded_analytics_prefers_full_concept_coverage_for_order_marke ) -def test_sql_validation_rejects_unresolved_domestic_international_market_comparison(): +def test_schema_grounded_analytics_compares_sales_by_domestic_international_market(): service = AskService.__new__(AskService) - table_ddls = [ - """ - CREATE TABLE dbo_qSalesMargin ( - EndMarket VARCHAR, - SalesValue DOUBLE, - OrdNo VARCHAR - ); - """ - ] sql = service._build_schema_grounded_analytics_sql( "Compare sales between domestic and international markets.", - table_ddls, + [ + """ + CREATE TABLE dbo_qSalesMargin ( + EndMarket VARCHAR, + SalesValue DOUBLE, + OrdNo VARCHAR + ); + """ + ], ) assert sql == ( @@ -423,47 +422,6 @@ def test_sql_validation_rejects_unresolved_domestic_international_market_compari 'GROUP BY "dbo_qSalesMargin"."EndMarket" ' 'ORDER BY SUM("dbo_qSalesMargin"."SalesValue") DESC' ) - assert ( - service._build_validated_ask_result_from_sql( - sql, - table_ddls, - "Compare sales between domestic and international markets.", - ) - is None - ) - - -def test_sql_validation_accepts_grounded_domestic_international_market_mapping(): - service = AskService.__new__(AskService) - - table_ddls = [ - """ - CREATE TABLE dbo_qSalesMargin ( - EndMarket VARCHAR, - SalesValue DOUBLE - ); - """ - ] - sql = ( - 'SELECT CASE WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Domestic%\' ' - "THEN 'Domestic' " - 'WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Intl%\' ' - "THEN 'International' END AS \"MarketScope\", " - 'SUM("dbo_qSalesMargin"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_qSalesMargin" ' - 'WHERE "dbo_qSalesMargin"."EndMarket" LIKE \'%Domestic%\' ' - 'OR "dbo_qSalesMargin"."EndMarket" LIKE \'%Intl%\' ' - 'GROUP BY CASE WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Domestic%\' ' - "THEN 'Domestic' " - 'WHEN "dbo_qSalesMargin"."EndMarket" LIKE \'%Intl%\' ' - "THEN 'International' END" - ) - - assert service._build_validated_ask_result_from_sql( - sql, - table_ddls, - "Compare sales between domestic and international markets.", - ) def test_schema_grounded_analytics_counts_customers_distribution_by_market(): @@ -493,59 +451,6 @@ def test_schema_grounded_analytics_counts_customers_distribution_by_market(): ) -def test_schema_grounded_analytics_counts_consolidated_customers_by_end_market(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Show customers distribution by market.", - [ - """ - CREATE TABLE dbo_qMarginSales ( - EndMarket VARCHAR, - ConsolidatedCustomer VARCHAR, - FXSales DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_qMarginSales"."EndMarket" AS "EndMarket", ' - 'COUNT(DISTINCT "dbo_qMarginSales"."ConsolidatedCustomer") ' - 'AS "CustomerCount" ' - 'FROM "dbo_qMarginSales" ' - 'WHERE "dbo_qMarginSales"."EndMarket" IS NOT NULL ' - 'AND "dbo_qMarginSales"."ConsolidatedCustomer" IS NOT NULL ' - 'GROUP BY "dbo_qMarginSales"."EndMarket" ' - 'ORDER BY COUNT(DISTINCT "dbo_qMarginSales"."ConsolidatedCustomer") DESC' - ) - - -def test_entity_lookup_uses_customer_like_match_for_company_phrase(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Show me orders for Lockheed Martin.", - [ - """ - CREATE TABLE dbo_tblNewOrders ( - ConsolidatedCustomer VARCHAR, - OrdNo VARCHAR, - OrdDate DATETIME, - Market VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."ConsolidatedCustomer" LIKE ' - "'%Lockheed Martin%' " - 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' - ) - - def test_scope_retrieval_to_semantic_contract_prefers_production_full_coverage_table(): service = AskService.__new__(AskService) documents = [ @@ -729,8 +634,8 @@ def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): assert sql == ( 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" LIKE ' - "'%Daimler Trucks North America%' " + 'WHERE "dbo_tblNewOrders"."CustName" = ' + "'Daimler Trucks North America' " 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' ) From f114144d3cb1b59ac9b9633b30c2aef9ad3fb275 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 18:42:21 +0530 Subject: [PATCH 0520/1087] Route asks through scoped schema SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3de52cabbf..e90266b1bb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2048,6 +2048,10 @@ def _find_temporal_column_for_query( def _build_schema_grounded_table_question_sql( self, query: str, table_ddls: list[str] ) -> str | None: + # Heuristic SQL skips semantic schema selection and can select a merely + # keyword-matching table. Let the scoped SQL-generation pipeline handle it. + return None + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None @@ -2210,6 +2214,10 @@ def _build_schema_grounded_table_question_sql( def _build_explicit_table_preview_sql( self, query: str, table_ddls: list[str] ) -> tuple[str, str] | None: + # A preview cannot safely infer which fields the user needs. The scoped + # generator must choose the required columns instead of emitting SELECT *. + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) if not normalized_query: return None From 878a2b245b36b37d385d14f4351dd61ead6134e5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 19:00:27 +0530 Subject: [PATCH 0521/1087] Revert "Route asks through scoped schema SQL generation" This reverts commit f114144d3cb1b59ac9b9633b30c2aef9ad3fb275. --- wren-ai-service/src/web/v1/services/ask.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e90266b1bb..3de52cabbf 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2048,10 +2048,6 @@ def _find_temporal_column_for_query( def _build_schema_grounded_table_question_sql( self, query: str, table_ddls: list[str] ) -> str | None: - # Heuristic SQL skips semantic schema selection and can select a merely - # keyword-matching table. Let the scoped SQL-generation pipeline handle it. - return None - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None @@ -2214,10 +2210,6 @@ def _build_schema_grounded_table_question_sql( def _build_explicit_table_preview_sql( self, query: str, table_ddls: list[str] ) -> tuple[str, str] | None: - # A preview cannot safely infer which fields the user needs. The scoped - # generator must choose the required columns instead of emitting SELECT *. - return None - normalized_query = re.sub(r"\s+", " ", (query or "").strip()) if not normalized_query: return None From 5fec368acdec9f78901e89fc90e016e638fe4cc9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 13 Jul 2026 19:28:23 +0530 Subject: [PATCH 0522/1087] Revert last scoped retrieval changes --- .../retrieval/db_schema_retrieval.py | 6 +- wren-ai-service/src/web/v1/services/ask.py | 99 +------------------ .../retrieval/test_db_schema_retrieval.py | 9 -- .../pytest/services/test_ask_sales_sql.py | 53 ---------- 4 files changed, 6 insertions(+), 161 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 0a805aee49..739c22e69e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -161,8 +161,6 @@ def expand_business_terms_for_retrieval(query: str) -> str: "invoices", "market", "markets", - "domestic", - "international", "order", "orders", "product", @@ -181,7 +179,7 @@ def expand_business_terms_for_retrieval(query: str) -> str: ) ): expansions.append( - "transaction purchase billing account geography customer client company name area representative market domestic international product item category sku quantity units sold amount value total metric money exchange currency" + "transaction purchase billing account geography customer client company name area representative product item category sku quantity units sold amount value total metric money exchange currency" ) if re.search( @@ -280,7 +278,7 @@ def _retrieval_concept_groups(query: str) -> list[set[str]]: ), ( {"market", "markets", "region", "regions"}, - {"market", "markets", "region", "regions", "area", "territory", "country", "domestic", "international"}, + {"market", "markets", "region", "regions", "area", "territory", "country"}, ), ( {"business", "unit", "units", "division"}, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3de52cabbf..7dc89d4824 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -787,7 +787,7 @@ def _semantic_concept_groups_for_query(self, query: str | None) -> list[set[str] ), ( {"market", "markets", "region", "regions"}, - {"market", "markets", "region", "regions", "area", "territory", "country", "domestic", "international"}, + {"market", "markets", "region", "regions", "area", "territory", "country"}, ), ( {"business", "unit", "units", "division"}, @@ -953,7 +953,7 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: if "currency" in normalized or "currencies" in normalized: concept_groups.append({"currency", "curr", "money", "fx", "exchange"}) if "market" in normalized or "markets" in normalized: - concept_groups.append({"market", "region", "country", "territory", "domestic", "international"}) + concept_groups.append({"market", "region", "country", "territory"}) if ( "business unit" in normalized or "business units" in normalized @@ -2791,94 +2791,6 @@ def _build_entity_lookup_sql( f"{order_clause}" ) - def _build_entity_distribution_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query or "distribution" not in normalized_query: - return None - - entity_candidates: tuple[str, ...] | None = None - entity_alias = "EntityCount" - if "customer" in normalized_query or "customers" in normalized_query: - entity_candidates = ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - "Client", - "ClientName", - ) - entity_alias = "CustomerCount" - elif "product" in normalized_query or "products" in normalized_query: - entity_candidates = ( - "Product", - "ProductName", - "ProdName", - "ProdCode", - "ProductCode", - "Item", - "ItemName", - "SKU", - ) - entity_alias = "ProductCount" - - dimension_candidates: tuple[str, ...] | None = None - if "market" in normalized_query or "markets" in normalized_query: - dimension_candidates = ("Market", "MarketType", "MarketName", "EndMarket", "Region", "Country") - elif "region" in normalized_query or "regions" in normalized_query: - dimension_candidates = ("Region", "Market", "Area", "Territory", "Country") - elif "country" in normalized_query or "countries" in normalized_query: - dimension_candidates = ("Country", "CountryName", "Nation", "Destination") - elif "division" in normalized_query: - dimension_candidates = ("Division",) - elif "business unit" in normalized_query or re.search(r"\bbu\b", normalized_query): - dimension_candidates = ("BusinessUnit", "Business Unit", "BU", "Division") - - if not entity_candidates or not dimension_candidates: - return None - - scored: list[tuple[int, dict[str, Any], str, str]] = [] - for table in tables: - dimension = self._find_schema_column(table, dimension_candidates) - entity = self._find_schema_column(table, entity_candidates) - if not dimension or not entity: - continue - - table_name = str(table.get("name") or "") - score = 20 + self._schema_source_shape_score(query, table) - if "order" in self._normalize_schema_token(table_name): - score += 15 - if "sales" in self._normalize_schema_token(table_name): - score += 10 - scored.append((score, table, dimension, entity)) - - if not scored: - return None - - _score, table, dimension, entity = sorted( - scored, - key=lambda item: item[0], - reverse=True, - )[0] - table_name = str(table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - entity_ref = f"{table_ref}.{self._quote_sql_identifier(entity)}" - where_clause = self._append_not_null_filters("", [dimension_ref, entity_ref]) - return ( - f"SELECT {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " - f"COUNT(DISTINCT {entity_ref}) AS {self._quote_sql_identifier(entity_alias)} " - f"FROM {table_ref}" - f"{where_clause} " - f"GROUP BY {dimension_ref} " - f"ORDER BY COUNT(DISTINCT {entity_ref}) DESC" - ) - def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2895,9 +2807,6 @@ def _build_schema_grounded_analytics_sql( if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): return entity_lookup_sql - if entity_distribution_sql := self._build_entity_distribution_sql(query, tables): - return entity_distribution_sql - if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -3051,7 +2960,7 @@ def _build_schema_grounded_analytics_sql( if "business unit" in normalized_query or re.search(r"\bbu\b", normalized_query): dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) if "market" in normalized_query: - dimension_candidates.append(("Market", "MarketType", "MarketName", "EndMarket", "Region", "Country")) + dimension_candidates.append(("Market", "MarketType", "MarketName", "Region", "Country")) if "region" in normalized_query: dimension_candidates.append(("Region", "Market", "Area", "Territory")) if "currency" in normalized_query or "currencies" in normalized_query: @@ -3486,7 +3395,7 @@ def _build_schema_grounded_analytics_sql( rank_dimension = None if "market" in normalized_query: partition_dimension = self._find_schema_column( - table, ("Market", "MarketType", "MarketName", "EndMarket", "Region") + table, ("Market", "MarketType", "Region") ) if "region" in normalized_query and not partition_dimension: partition_dimension = self._find_schema_column( diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index a81a1202ed..102f183a39 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -42,15 +42,6 @@ def test_expand_business_terms_for_retrieval_adds_generic_currency_market_terms( assert "money exchange currency" in expanded_query -def test_expand_business_terms_for_retrieval_adds_domestic_international_market_terms(): - query = "Compare sales between domestic and international markets" - - expanded_query = expand_business_terms_for_retrieval(query) - - assert query in expanded_query - assert "domestic international" in expanded_query - - def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): query = "Explain what this workspace does" diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index f1fa8e0450..0ee413f27c 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -398,59 +398,6 @@ def test_schema_grounded_analytics_prefers_full_concept_coverage_for_order_marke ) -def test_schema_grounded_analytics_compares_sales_by_domestic_international_market(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Compare sales between domestic and international markets.", - [ - """ - CREATE TABLE dbo_qSalesMargin ( - EndMarket VARCHAR, - SalesValue DOUBLE, - OrdNo VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_qSalesMargin"."EndMarket" AS "EndMarket", ' - 'SUM("dbo_qSalesMargin"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_qSalesMargin" ' - 'WHERE "dbo_qSalesMargin"."EndMarket" IS NOT NULL ' - 'GROUP BY "dbo_qSalesMargin"."EndMarket" ' - 'ORDER BY SUM("dbo_qSalesMargin"."SalesValue") DESC' - ) - - -def test_schema_grounded_analytics_counts_customers_distribution_by_market(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Show customers distribution by market.", - [ - """ - CREATE TABLE dbo_tblNewOrders ( - Market VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblNewOrders"."Market" AS "Market", ' - 'COUNT(DISTINCT "dbo_tblNewOrders"."CustName") AS "CustomerCount" ' - 'FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."Market" IS NOT NULL ' - 'AND "dbo_tblNewOrders"."CustName" IS NOT NULL ' - 'GROUP BY "dbo_tblNewOrders"."Market" ' - 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."CustName") DESC' - ) - - def test_scope_retrieval_to_semantic_contract_prefers_production_full_coverage_table(): service = AskService.__new__(AskService) documents = [ From af78319d832225ee2779c1c090fb8d9f0aa48156 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 01:00:33 +0530 Subject: [PATCH 0523/1087] Fail closed when semantic schema retrieval fails --- wren-ai-service/src/web/v1/services/ask.py | 219 --------------------- 1 file changed, 219 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7dc89d4824..9f02e2b377 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -6600,61 +6600,6 @@ async def ask( ) sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - logger.info( - "Retrieved tables for direct heuristic query_id %s: %s", - query_id, - table_names, - ) - - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using direct heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - if ask_result := self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ): - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - if explicit_group_count_sql := self._build_explicit_group_count_sql( user_query ): @@ -7075,44 +7020,6 @@ async def ask( explicit_table_names, ) ) - if ( - not documents - and self._should_load_full_schema_for_question(user_query) - and not request_explicit_table_names - ): - logger.info( - "Query-based schema retrieval returned no tables for data question; " - "retrying full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) if documents and not explicit_table_names: documents, table_names, table_ddls = ( self._scope_retrieval_to_semantic_contract( @@ -7273,92 +7180,6 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - should_retry_full_schema = ( - not api_results - and self._should_load_full_schema_for_question(user_query) - and "db_schema_retrieval" in self._pipelines - and not request_explicit_table_names - and not table_names - ) - if should_retry_full_schema: - logger.info( - "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retry", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 30, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - full_documents, full_table_names, full_table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - full_documents, full_table_names, full_table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - full_documents, - explicit_table_names, - ) - ) - if full_documents: - documents, table_names, table_ddls = ( - full_documents, - full_table_names, - full_table_ddls, - ) - logger.info( - "Using full active deployed schema retry for query_id %s: %s", - query_id, - table_names, - ) - - full_schema_preview = self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - full_schema_sql_candidates = ( - self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ), - full_schema_preview[0] if full_schema_preview else None, - self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ), - self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ), - self._build_schema_grounded_sales_sql( - user_query, table_ddls - ), - ) - for full_schema_sql in full_schema_sql_candidates: - if not full_schema_sql: - continue - ask_result = self._build_validated_ask_result_from_sql( - full_schema_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - break - invalid_sql = full_schema_sql - error_message = ( - "Full-schema grounded SQL was not valid for the active datasource schema and question intent." - ) - if not api_results and ( unqueryable_metric_message := self._get_unqueryable_metric_message( user_query, table_ddls @@ -7749,46 +7570,6 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if self._can_use_schema_grounded_sql_fallback( - documents, - table_ddls, - user_query, - ) and ( - heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ) - ): - logger.info( - "Using heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - else: - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( From 39a8371c0b80472665aa85bbe4995a060180372d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 01:17:21 +0530 Subject: [PATCH 0524/1087] Reject identifier columns for customer name queries --- wren-ai-service/src/web/v1/services/ask.py | 76 +++++++++++++++++++++- 1 file changed, 75 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9f02e2b377..51371d9d8d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1370,6 +1370,64 @@ def _sql_satisfies_entity_lookup_request( return True + def _sql_satisfies_named_entity_request( + self, + sql: str, + query: str | None, + referenced_tables: list[str], + valid_tables: dict[str, dict[str, Any]], + ) -> bool: + """Require a name field when the question explicitly asks for names. + + An account number, customer code, or generic ``Customer`` field may be a + useful identifier, but it is not evidence that the datasource exposes a + customer name. Returning such an identifier for a name request produces + plausible but incorrect answers. + """ + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + requests_customer_name = bool( + re.search( + r"\b(?:customer|customers|client|clients|account|accounts|company|companies)\s+names?\b" + r"|\bnames?\s+(?:of|for)\s+(?:customers|clients|accounts|companies)\b", + normalized_query, + ) + ) + if not requests_customer_name: + return True + + name_columns = ( + "CustomerName", + "CustName", + "ClientName", + "AccountName", + "CompanyName", + "Name", + ) + available_name_columns: set[str] = set() + for table_reference in referenced_tables: + table = self._table_for_sql_reference(table_reference, valid_tables) + if not table: + continue + if column := self._find_schema_column(table, name_columns): + available_name_columns.add(column) + + for column in available_name_columns: + escaped_column = re.escape(column) + if re.search( + rf'(?i)(?:"{escaped_column}"|\[{escaped_column}\]|\b{escaped_column}\b)', + sql, + ): + return True + + logger.warning( + "Ignoring SQL because a customer-name request was mapped to an identifier or " + "a datasource without a customer-name column. query=%s available_name_columns=%s sql=%s", + query, + sorted(available_name_columns), + sql, + ) + return False + def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1691,6 +1749,13 @@ def _sql_matches_question_intent( valid_tables, ): return False + if not self._sql_satisfies_named_entity_request( + sql, + query, + referenced_tables, + valid_tables, + ): + return False if not expects_dimension: return True @@ -2142,9 +2207,18 @@ def _build_schema_grounded_table_question_sql( else "COUNT(*)" ) top_clause = f"TOP {limit} " if wants_ranked_count else "" + dimension_type = next( + ( + str(column.get("type") or "") + for column in table.get("columns", []) + if self._normalize_schema_identifier_key(column.get("name")) + == self._normalize_schema_identifier_key(dimension_column) + ), + "", + ) nonblank_filter = ( f"AND LTRIM(RTRIM({dimension_ref})) <> '' " - if wants_ranked_count + if self._is_text_schema_type(dimension_type) else "" ) return ( From 620fd5082dd9c39434855aaef1129644c5409c15 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 01:38:10 +0530 Subject: [PATCH 0525/1087] Prevent unsafe customer and country column mappings --- .../src/pipelines/generation/utils/sql.py | 39 ++++++++++++++++--- wren-ai-service/src/web/v1/services/ask.py | 4 +- .../pipelines/generation/test_sql_utils.py | 26 +++++++++++++ 3 files changed, 63 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b60124120d..5e650d165e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2721,11 +2721,40 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "period": ("timeid", "TimeID", "TimeId", "YearInd", "Year", "Date"), "periodid": ("timeid", "TimeID", "TimeId"), "timeid": ("timeid", "TimeID", "TimeId"), - "customer": ("account", "Customer", "CustName", "CustNo", "customerpo"), - "customers": ("account", "Customer", "CustName", "CustNo", "customerpo"), - "customername": ("account", "Customer", "CustName", "CustNo", "customerpo"), - "customerid": ("account", "Customer", "CustNo", "customerpo"), - "customeraccount": ("account", "Customer", "CustName", "CustNo"), + "customer": ( + "Customer", + "CustomerName", + "CustName", + "Client", + "ClientName", + "Company", + "CompanyName", + ), + "customers": ( + "Customer", + "CustomerName", + "CustName", + "Client", + "ClientName", + "Company", + "CompanyName", + ), + "customername": ( + "CustomerName", + "CustName", + "ClientName", + "AccountName", + "CompanyName", + ), + "customerid": ( + "CustomerId", + "CustomerID", + "CustNo", + "CustomerNo", + "CustomerCode", + "customerpo", + ), + "customeraccount": ("account", "AccountName", "AcctNo", "CustNo"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 51371d9d8d..5a9c8a2438 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -3051,7 +3051,9 @@ def _build_schema_grounded_analytics_sql( ) ) if "country" in normalized_query or "countries" in normalized_query: - dimension_candidates.append(("Country", "CountryName", "Nation", "Market")) + dimension_candidates.append( + ("Country", "CountryName", "Nation", "Destination") + ) if "division" in normalized_query: dimension_candidates.append(("Division",)) if ( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6dc0b45205..8514b45dfc 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -397,6 +397,32 @@ def test_normalize_sql_column_references_to_schema_maps_sales_business_aliases() ) == [] +def test_customer_name_is_not_normalized_to_an_account_identifier(): + sql = 'SELECT "dbo_qSalesMargin"."CustomerName" FROM "dbo_qSalesMargin"' + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_qSalesMargin": ["Account"]}, + ) + + assert '"dbo_qSalesMargin"."Account"' not in normalized + assert find_invalid_column_references( + normalized, + {"dbo_qSalesMargin": ["Account"]}, + ) == ["dbo_qSalesMargin.CustomerName"] + + +def test_customer_name_is_normalized_to_a_real_customer_name_column(): + sql = 'SELECT "dbo_qSalesMargin"."CustomerName" FROM "dbo_qSalesMargin"' + + normalized = normalize_sql_column_references_to_schema( + sql, + {"dbo_qSalesMargin": ["CustName"]}, + ) + + assert normalized == 'SELECT "dbo_qSalesMargin"."CustName" FROM "dbo_qSalesMargin"' + + def test_normalize_sql_column_references_to_schema_maps_period_to_timeid(): sql = 'SELECT "dbo_tblFactSales"."Period" FROM "dbo_tblFactSales"' From 8a81b789a7d6fece96d2d8f6c4d3fad8b09a3f01 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 01:40:40 +0530 Subject: [PATCH 0526/1087] Use retrieved schema for all generic SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 149 +-------------------- 1 file changed, 1 insertion(+), 148 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 5a9c8a2438..15f7b9434b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -7109,153 +7109,6 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if not api_results and ( - table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded table question SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - table_question_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = table_question_sql - error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - - if not api_results and ( - explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - ): - explicit_sql, explicit_table_name = explicit_table_preview - logger.info( - "Using explicit table preview SQL for query_id %s and table %s", - query_id, - explicit_table_name, - ) - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = explicit_sql - error_message = "Explicit table preview SQL was not valid for the active datasource schema." - - if not api_results and ( - audit_log_activity_sql := self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ) - ): - logger.info( - "Using schema-grounded audit log activity SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - audit_log_activity_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = audit_log_activity_sql - error_message = ( - "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." - ) - - if ( - not api_results - and self._is_data_analysis_query(user_query) - and ( - schema_grounded_sql := self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - ) - ): - logger.info( - "Using generic schema-grounded analytics SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - schema_grounded_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = schema_grounded_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and any( - term in user_query.lower() - for term in ( - "pcb", - "repair", - "failure", - "business unit", - "business units", - "product line", - "product family", - ) - ): - operational_sql = self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - if operational_sql: - logger.info( - "Using schema-grounded operational SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - operational_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = operational_sql - error_message = ( - "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and ( - deterministic_sales_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded CWSales SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sales_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = deterministic_sales_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - if not api_results and ( unqueryable_metric_message := self._get_unqueryable_metric_message( user_query, table_ddls @@ -7510,7 +7363,7 @@ async def ask( ) except TimeoutError as generation_timeout: logger.warning( - "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + "SQL generation timed out for query_id %s; returning a generation failure: %s", query_id, generation_timeout, ) From ea91149061328ffc5c23814f419ae3f270f9e32f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 02:15:56 +0530 Subject: [PATCH 0527/1087] Simplify generic schema retrieval and preserve context --- .../pipelines/indexing/table_description.py | 18 ++- .../retrieval/db_schema_retrieval.py | 133 +++--------------- wren-ai-service/src/web/v1/services/ask.py | 35 ----- 3 files changed, 33 insertions(+), 153 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 8e1b875b49..20a9f5f689 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -70,10 +70,26 @@ def _additional_meta() -> Dict[str, Any]: def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[str]: def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: + def _column_description(column: Dict[str, Any]) -> str: + name = str(column.get("name") or "") + properties = column.get("properties") or {} + description = str(properties.get("description") or "").strip() + data_type = str(column.get("type") or "").strip() + parts = [name] + if data_type: + parts.append(f"type: {data_type}") + if description: + parts.append(description) + return " — ".join(parts) + return { "mdl_type": mdl_type, "name": payload.get("name"), - "columns": [column["name"] for column in payload.get("columns", [])], + "columns": [ + _column_description(column) + for column in payload.get("columns", []) + if column.get("name") + ], "properties": payload.get("properties", {}), } diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 739c22e69e..0771c10f7c 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -29,7 +29,7 @@ logger = logging.getLogger("wren-ai-service") -MAX_RELEVANT_TABLE_CANDIDATES = 5 +MAX_RELEVANT_TABLE_CANDIDATES = 10 MIN_TABLE_DESCRIPTION_CANDIDATE_WINDOW = 100 WEAK_NON_PRODUCTION_TERMS = ( "stage", @@ -146,68 +146,13 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline def expand_business_terms_for_retrieval(query: str) -> str: - normalized = (query or "").lower() - expansions: list[str] = [] - - if any( - term in normalized - for term in ( - "amount", - "currency", - "currencies", - "customer", - "customers", - "invoice", - "invoices", - "market", - "markets", - "order", - "orders", - "product", - "products", - "category", - "categories", - "quantity", - "qty", - "region", - "regions", - "sales", - "salesperson", - "sales person", - "sold", - "value", - ) - ): - expansions.append( - "transaction purchase billing account geography customer client company name area representative product item category sku quantity units sold amount value total metric money exchange currency" - ) - - if re.search( - r"\b(?:show|list|find|get|display)\b.*\b(?:orders?|records?|rows?|transactions?)\b\s+" - r"(?:for|where|with)\s+\S+", - normalized, - ): - expansions.append( - "customer client account company name entity lookup identifier transaction order record" - ) - - if any( - term in normalized - for term in ("defect", "failure", "issue", "repair", "resolved", "status") - ): - expansions.append( - "issue defect category status resolved created updated date timestamp event" - ) + """Keep retrieval grounded in the user's wording and indexed metadata. - if any(term in normalized for term in ("throughput", "production", "manufacturing")): - expansions.append( - "rate volume output capacity process unit group completed timestamp date" - ) - - if not expansions: - return query - - return f"{query}\n" + "\n".join(expansions) + Broad, global synonym expansion makes unrelated business questions share + the same embedding and lexical terms. Synonyms belong in deployed model + and column metadata, where their meaning is datasource-specific. + """ + return query def _normalize_retrieval_token(value: str) -> str: @@ -497,7 +442,7 @@ def _score_table_documents( if not documents: return [] - query_terms = _retrieval_terms(expand_business_terms_for_retrieval(query)) + query_terms = _retrieval_terms(query) if not query_terms: return [ (_semantic_score(document), -index, document, 0, _semantic_score(document)) @@ -508,8 +453,10 @@ def _score_table_documents( for index, document in enumerate(documents): lexical_score = _document_relevance_score(document, query_terms) semantic_score = _semantic_score(document) - source_shape_score = _source_shape_score(query, document) - combined_score = semantic_score + lexical_score + source_shape_score + # Qdrant similarity is the primary ranking signal. Lexical overlap is + # intentionally a small tie-breaker so physical names cannot overpower + # semantically relevant descriptions. + combined_score = semantic_score + min(lexical_score / 1000, 0.05) scored_documents.append( (combined_score, -index, document, lexical_score, semantic_score) ) @@ -555,59 +502,9 @@ def _select_relevant_table_documents( if not reranked: return documents[:max_tables] - candidate_pool = [item for item in reranked if item[3] > 0] or reranked - concept_groups = _retrieval_concept_groups(query) - if concept_groups: - all_concepts = set(range(len(concept_groups))) - coverage_by_index = { - index: _retrieval_concept_coverage(query, item[2]) - for index, item in enumerate(candidate_pool) - } - full_coverage_pool = [ - item - for index, item in enumerate(candidate_pool) - if coverage_by_index[index] == all_concepts - ] - if full_coverage_pool: - candidate_pool = full_coverage_pool - coverage_by_index = { - index: _retrieval_concept_coverage(query, item[2]) - for index, item in enumerate(candidate_pool) - } - - grounded_production_coverage: set[int] = set() - for index, item in enumerate(candidate_pool): - coverage = coverage_by_index[index] - if coverage and not _is_unrequested_strong_non_production_source( - query, item[2] - ): - grounded_production_coverage.update(coverage) - - if grounded_production_coverage: - filtered_pool = [] - excluded_names = [] - for index, item in enumerate(candidate_pool): - document = item[2] - coverage = coverage_by_index[index] - if ( - _is_unrequested_strong_non_production_source(query, document) - and coverage <= grounded_production_coverage - ): - excluded_names.append(document.meta.get("name")) - continue - filtered_pool.append(item) - if filtered_pool: - candidate_pool = filtered_pool - if excluded_names: - logger.info( - "Excluded unrequested non-production table candidates for query=%s names=%s", - query, - excluded_names, - ) - selected = [ document - for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] + for _score, _index, document, _lexical, _semantic in reranked[:max_tables] ] if len(selected) < len(documents): logger.info( @@ -763,6 +660,7 @@ async def table_retrieval( project_id: str, tables: list[str], table_retriever: Any, + max_tables: int = MAX_RELEVANT_TABLE_CANDIDATES, ) -> dict: base_filters = { "operator": "AND", @@ -782,7 +680,7 @@ async def table_retrieval( filters=base_filters, ) results["documents"] = _select_relevant_table_documents( - query, results.get("documents") or [] + query, results.get("documents") or [], max_tables=max_tables ) return results @@ -1115,6 +1013,7 @@ def __init__( document_store_provider.get_store(dataset_name="table_descriptions"), top_k=table_description_candidate_window, ), + "max_tables": max(table_retrieval_size, MAX_RELEVANT_TABLE_CANDIDATES), "dbschema_retriever": document_store_provider.get_retriever( document_store_provider.get_store(), top_k=table_column_retrieval_size, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 15f7b9434b..7c32453992 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -7096,15 +7096,6 @@ async def ask( explicit_table_names, ) ) - if documents and not explicit_table_names: - documents, table_names, table_ddls = ( - self._scope_retrieval_to_semantic_contract( - sql_user_query, - documents, - table_names, - table_ddls, - ) - ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) @@ -7157,21 +7148,6 @@ async def ask( return results if documents and not api_results: - documents, table_names, table_ddls = self._prune_sql_generation_context( - sql_user_query, - documents, - table_names, - table_ddls, - ) - if not explicit_table_names: - documents, table_names, table_ddls = ( - self._scope_retrieval_to_semantic_contract( - sql_user_query, - documents, - table_names, - table_ddls, - ) - ) ( documents, table_names, @@ -7188,17 +7164,6 @@ async def ask( _retrieval_result = completed_retrieval_result sql_generation_histories = histories - if self._is_data_analysis_query( - sql_user_query - ) and not self._needs_conversation_context(sql_user_query): - sql_generation_histories = [] - allow_sql_generation_reasoning = False - allow_sql_knowledge_retrieval = False - max_sql_correction_retries = min(max_sql_correction_retries, 1) - logger.info( - "Using fast standalone SQL generation path for query_id %s", - query_id, - ) if ( not self._is_stopped(query_id, self._ask_results) From 7d26c12e4b918b37b2f6a6e463b9986e0f972ab0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 03:05:09 +0530 Subject: [PATCH 0528/1087] Require filtered DDL before SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7c32453992..f69715fd62 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -5505,8 +5505,19 @@ def _extract_retrieval_documents(self, retrieval_result: dict) -> list[dict]: if not isinstance(document, dict): logger.warning("Ignoring malformed retrieval document: %s", document) continue - if not document.get("table_name") and not document.get("table_ddl"): - logger.warning("Ignoring retrieval document without table metadata") + table_name = document.get("table_name") + table_ddl = document.get("table_ddl") + if not isinstance(table_name, str) or not table_name.strip(): + logger.warning("Ignoring retrieval document without a table name") + continue + if not isinstance(table_ddl, str) or not table_ddl.strip(): + # A table name alone is not safe SQL-generation context. Treat this + # exactly like an unsuccessful schema retrieval rather than allowing + # the generator to infer columns from the rest of the project. + logger.warning( + "Ignoring retrieval document without filtered DDL for table %s", + table_name, + ) continue valid_documents.append(document) From 2ce7dcd9c19a3469a1b4103c1530b37f10f5be5c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 03:29:22 +0530 Subject: [PATCH 0529/1087] Bound generic SQL retrieval context --- .../retrieval/db_schema_retrieval.py | 33 ++++++++++++++++--- wren-ai-service/src/web/v1/services/ask.py | 8 ++--- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 0771c10f7c..72b6d218de 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -751,7 +751,14 @@ async def dbschema_retrieval( def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: db_schemas = {} for document in dbschema_retrieval: - content = ast.literal_eval(document.content) + try: + content = ast.literal_eval(document.content) + except (SyntaxError, ValueError): + logger.warning("Ignoring malformed schema index document: %s", document.meta) + continue + if not isinstance(content, dict) or not isinstance(content.get("type"), str): + logger.warning("Ignoring schema index document without a type: %s", document.meta) + continue if content["type"] == "TABLE": if document.meta["name"] not in db_schemas: db_schemas[document.meta["name"]] = content @@ -803,7 +810,14 @@ def check_using_db_schemas_without_pruning( has_json_field = True for document in dbschema_retrieval: - content = ast.literal_eval(document.content) + try: + content = ast.literal_eval(document.content) + except (SyntaxError, ValueError): + logger.warning("Ignoring malformed schema index document: %s", document.meta) + continue + if not isinstance(content, dict) or not isinstance(content.get("type"), str): + logger.warning("Ignoring schema index document without a type: %s", document.meta) + continue if content["type"] == "METRIC": retrieval_results.append( @@ -929,7 +943,14 @@ def construct_retrieval_results( for document in dbschema_retrieval: if document.meta["name"] in columns_and_tables_needed: - content = ast.literal_eval(document.content) + try: + content = ast.literal_eval(document.content) + except (SyntaxError, ValueError): + logger.warning("Ignoring malformed schema index document: %s", document.meta) + continue + if not isinstance(content, dict) or not isinstance(content.get("type"), str): + logger.warning("Ignoring schema index document without a type: %s", document.meta) + continue if content["type"] == "METRIC": retrieval_results.append( @@ -1013,7 +1034,11 @@ def __init__( document_store_provider.get_store(dataset_name="table_descriptions"), top_k=table_description_candidate_window, ), - "max_tables": max(table_retrieval_size, MAX_RELEVANT_TABLE_CANDIDATES), + # `top_k` above is intentionally a broad candidate window for reranking. + # The final schema context must, however, honour the configured limit. + # Otherwise every query can send ten complete table schemas to the LLM + # even when an installation configured a smaller retrieval size. + "max_tables": max(1, table_retrieval_size), "dbschema_retriever": document_store_provider.get_retriever( document_store_provider.get_store(), top_k=table_column_retrieval_size, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f69715fd62..44f2b0edf2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -7017,10 +7017,10 @@ async def ask( tables=retrieval_table_names, histories=[], project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) - ), + # Keep the SQL prompt bounded for every data question. + # This is the legacy behavior: select the relevant columns + # before SQL generation rather than sending full schemas. + enable_column_pruning=enable_column_pruning, ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) From a283e884988e6df8661bac46542b6b0ecfaa8469 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 03:40:21 +0530 Subject: [PATCH 0530/1087] Keep filtered retrieval context for SQL generation --- wren-ai-service/src/web/v1/services/ask.py | 20 +++++--------------- 1 file changed, 5 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 44f2b0edf2..4e1bf7e05b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -7158,21 +7158,11 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if documents and not api_results: - ( - documents, - table_names, - table_ddls, - completed_retrieval_result, - ) = await self._complete_sql_generation_context( - query=sql_user_query, - project_id=ask_request.project_id, - documents=documents, - table_names=table_names, - table_ddls=table_ddls, - ) - if completed_retrieval_result: - _retrieval_result = completed_retrieval_result + # The retrieval pipeline has already selected tables and, when enabled, + # pruned their columns. Do not reload those names as explicit tables: + # that turns the selected context back into full schemas and can send + # many unrelated tables to SQL generation. Legacy/v1 passes the + # filtered DDL produced above directly to reasoning and generation. sql_generation_histories = histories From 1d38eb2cb7e59ee771003e3d2b9b0e023fd930f7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 15:04:17 +0530 Subject: [PATCH 0531/1087] Restore Wren AI service flow from 97075b46 --- .../src/pipelines/generation/utils/sql.py | 39 +- .../pipelines/indexing/table_description.py | 18 +- .../retrieval/db_schema_retrieval.py | 309 ++--- wren-ai-service/src/web/v1/services/ask.py | 1136 +++++++---------- .../pipelines/generation/test_sql_utils.py | 26 - .../retrieval/test_db_schema_retrieval.py | 199 +-- .../pytest/services/test_ask_sales_sql.py | 369 ------ 7 files changed, 544 insertions(+), 1552 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5e650d165e..b60124120d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2721,40 +2721,11 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: "period": ("timeid", "TimeID", "TimeId", "YearInd", "Year", "Date"), "periodid": ("timeid", "TimeID", "TimeId"), "timeid": ("timeid", "TimeID", "TimeId"), - "customer": ( - "Customer", - "CustomerName", - "CustName", - "Client", - "ClientName", - "Company", - "CompanyName", - ), - "customers": ( - "Customer", - "CustomerName", - "CustName", - "Client", - "ClientName", - "Company", - "CompanyName", - ), - "customername": ( - "CustomerName", - "CustName", - "ClientName", - "AccountName", - "CompanyName", - ), - "customerid": ( - "CustomerId", - "CustomerID", - "CustNo", - "CustomerNo", - "CustomerCode", - "customerpo", - ), - "customeraccount": ("account", "AccountName", "AcctNo", "CustNo"), + "customer": ("account", "Customer", "CustName", "CustNo", "customerpo"), + "customers": ("account", "Customer", "CustName", "CustNo", "customerpo"), + "customername": ("account", "Customer", "CustName", "CustNo", "customerpo"), + "customerid": ("account", "Customer", "CustNo", "customerpo"), + "customeraccount": ("account", "Customer", "CustName", "CustNo"), "customerregion": ("Country", "Market", "Region", "CustomerRegion"), "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), } diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 20a9f5f689..8e1b875b49 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -70,26 +70,10 @@ def _additional_meta() -> Dict[str, Any]: def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[str]: def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: - def _column_description(column: Dict[str, Any]) -> str: - name = str(column.get("name") or "") - properties = column.get("properties") or {} - description = str(properties.get("description") or "").strip() - data_type = str(column.get("type") or "").strip() - parts = [name] - if data_type: - parts.append(f"type: {data_type}") - if description: - parts.append(description) - return " — ".join(parts) - return { "mdl_type": mdl_type, "name": payload.get("name"), - "columns": [ - _column_description(column) - for column in payload.get("columns", []) - if column.get("name") - ], + "columns": [column["name"] for column in payload.get("columns", [])], "properties": payload.get("properties", {}), } diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 72b6d218de..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -29,24 +29,7 @@ logger = logging.getLogger("wren-ai-service") -MAX_RELEVANT_TABLE_CANDIDATES = 10 -MIN_TABLE_DESCRIPTION_CANDIDATE_WINDOW = 100 -WEAK_NON_PRODUCTION_TERMS = ( - "stage", - "staging", -) -STRONG_NON_PRODUCTION_TERMS = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "temp", - "test", - "tmp", -) +MAX_RELEVANT_TABLE_CANDIDATES = 5 table_columns_selection_system_prompt = """ @@ -146,13 +129,59 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline def expand_business_terms_for_retrieval(query: str) -> str: - """Keep retrieval grounded in the user's wording and indexed metadata. + normalized = (query or "").lower() + expansions: list[str] = [] + + if any( + term in normalized + for term in ( + "amount", + "currency", + "currencies", + "customer", + "customers", + "invoice", + "invoices", + "market", + "markets", + "order", + "orders", + "product", + "products", + "category", + "categories", + "quantity", + "qty", + "region", + "regions", + "sales", + "salesperson", + "sales person", + "sold", + "value", + ) + ): + expansions.append( + "transaction purchase billing account geography area representative product item category sku quantity units sold amount value total metric money exchange currency" + ) - Broad, global synonym expansion makes unrelated business questions share - the same embedding and lexical terms. Synonyms belong in deployed model - and column metadata, where their meaning is datasource-specific. - """ - return query + if any( + term in normalized + for term in ("defect", "failure", "issue", "repair", "resolved", "status") + ): + expansions.append( + "issue defect category status resolved created updated date timestamp event" + ) + + if any(term in normalized for term in ("throughput", "production", "manufacturing")): + expansions.append( + "rate volume output capacity process unit group completed timestamp date" + ) + + if not expansions: + return query + + return f"{query}\n" + "\n".join(expansions) def _normalize_retrieval_token(value: str) -> str: @@ -184,17 +213,11 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms: set[str] = set() - for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): - split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) - for token in re.findall(r"[A-Za-z0-9]+", split_token): - if len(token) <= 2 and token.lower() not in {"bu"}: - continue - if token.lower() in stop_words: - continue - normalized_token = _normalize_retrieval_token(token) - if normalized_token: - terms.add(normalized_token) + terms = { + _normalize_retrieval_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") + if len(token) > 2 and token.lower() not in stop_words + } return {term for term in terms if term} @@ -203,72 +226,6 @@ def _query_mentions_any(query: str, terms: tuple[str, ...]) -> bool: return any(re.search(rf"\b{re.escape(term)}\b", normalized) for term in terms) -def _retrieval_concept_groups(query: str) -> list[set[str]]: - query_terms = _retrieval_terms(query) - if not query_terms: - return [] - - concept_specs: list[tuple[set[str], set[str]]] = [ - ( - {"order", "orders", "neworder", "purchase", "transaction"}, - {"order", "orders", "ord", "ordno", "orderid", "orderdate", "neworder", "purchase", "transaction"}, - ), - ( - {"customer", "customers", "client", "account", "company"}, - {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"}, - ), - ( - {"product", "products", "item", "sku", "category", "categories"}, - {"product", "products", "prod", "item", "sku", "category", "categories", "type", "name"}, - ), - ( - {"market", "markets", "region", "regions"}, - {"market", "markets", "region", "regions", "area", "territory", "country"}, - ), - ( - {"business", "unit", "units", "division"}, - {"business", "businessunit", "unit", "units", "bu", "division"}, - ), - ( - {"country", "countries", "destination"}, - {"country", "countries", "destination", "nation", "market", "region"}, - ), - ( - {"invoice", "invoices", "billing"}, - {"invoice", "invoices", "inv", "billing", "bill", "amount", "value"}, - ), - ( - {"amount", "value", "total", "sum", "revenue", "sales", "cost"}, - {"amount", "value", "total", "sum", "revenue", "sales", "sale", "cost", "price", "money"}, - ), - ( - {"quantity", "qty", "sold", "units"}, - {"quantity", "qty", "sold", "unit", "units", "volume"}, - ), - ( - {"date", "month", "monthly", "year", "quarter", "trend", "period"}, - {"date", "month", "year", "quarter", "time", "timestamp", "orddate", "invdate", "period"}, - ), - ] - - groups: list[set[str]] = [] - for triggers, aliases in concept_specs: - if query_terms & triggers: - groups.append(aliases) - - if re.search( - r"\b(?:show|list|find|get|display)\b.*\b(?:orders?|records?|rows?|transactions?)\b\s+" - r"(?:for|where|with)\s+\S+", - query or "", - flags=re.IGNORECASE, - ): - groups.append( - {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"} - ) - - return groups - - def _source_text(document: Document) -> str: return " ".join( str(part or "") @@ -286,23 +243,25 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - concept_groups = _retrieval_concept_groups(query) - if concept_groups: - covered_groups = sum(1 for group in concept_groups if group & source_terms) - score += 18 * covered_groups - if covered_groups == len(concept_groups): - score += 30 - elif covered_groups == 0: - score -= 25 - - if source_terms & set(STRONG_NON_PRODUCTION_TERMS) and not _query_mentions_any( - normalized_query, STRONG_NON_PRODUCTION_TERMS - ): - score -= 240 - if source_terms & set(WEAK_NON_PRODUCTION_TERMS) and not _query_mentions_any( - normalized_query, WEAK_NON_PRODUCTION_TERMS + non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "stage", + "staging", + "temp", + "test", + "tmp", + ) + if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( + normalized_query, + non_production_terms, ): - score -= 40 + score -= 60 aggregation_terms = ( "amount", @@ -347,11 +306,6 @@ def _source_shape_score(query: str, document: Document) -> int: r"employees?|entities|items?|names?|products?|suppliers?|users?|vendors?)\b", normalized_query, ) - entity_lookup_pattern = re.search( - r"\b(?:list|show|display|get|find)\b.*\b(?:orders?|records?|rows?|transactions?)\b\s+" - r"(?:for|where|with)\s+\S+", - normalized_query, - ) asks_for_aggregation = _query_mentions_any(normalized_query, aggregation_terms) or bool( re.search(r"\b(?:by|per|each|top|bottom|rank|ranking)\b", normalized_query) ) @@ -362,11 +316,6 @@ def _source_shape_score(query: str, document: Document) -> int: score += 35 if source_terms & set(transaction_source_terms): score -= 12 - elif entity_lookup_pattern: - if source_terms & {"customer", "cust", "custname", "client", "account", "company", "name"}: - score += 45 - if source_terms & set(transaction_source_terms): - score += 20 elif asks_for_aggregation: if source_terms & set(transaction_source_terms): score += 25 @@ -415,34 +364,13 @@ def _semantic_score(document: Document) -> float: return 0.0 -def _retrieval_concept_coverage(query: str, document: Document) -> set[int]: - source_terms = _retrieval_terms(_source_text(document)) - return { - index - for index, concept_group in enumerate(_retrieval_concept_groups(query)) - if concept_group & source_terms - } - - -def _is_unrequested_strong_non_production_source( - query: str, document: Document -) -> bool: - source_terms = _retrieval_terms(_source_text(document)) - return bool(source_terms & set(STRONG_NON_PRODUCTION_TERMS)) and not ( - _query_mentions_any( - query or "", - STRONG_NON_PRODUCTION_TERMS, - ) - ) - - def _score_table_documents( query: str, documents: list[Document] ) -> list[tuple[float, int, Document, int, float]]: if not documents: return [] - query_terms = _retrieval_terms(query) + query_terms = _retrieval_terms(expand_business_terms_for_retrieval(query)) if not query_terms: return [ (_semantic_score(document), -index, document, 0, _semantic_score(document)) @@ -453,10 +381,8 @@ def _score_table_documents( for index, document in enumerate(documents): lexical_score = _document_relevance_score(document, query_terms) semantic_score = _semantic_score(document) - # Qdrant similarity is the primary ranking signal. Lexical overlap is - # intentionally a small tie-breaker so physical names cannot overpower - # semantically relevant descriptions. - combined_score = semantic_score + min(lexical_score / 1000, 0.05) + source_shape_score = _source_shape_score(query, document) + combined_score = semantic_score + lexical_score + source_shape_score scored_documents.append( (combined_score, -index, document, lexical_score, semantic_score) ) @@ -502,9 +428,10 @@ def _select_relevant_table_documents( if not reranked: return documents[:max_tables] + candidate_pool = [item for item in reranked if item[3] > 0] or reranked selected = [ document - for _score, _index, document, _lexical, _semantic in reranked[:max_tables] + for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] ] if len(selected) < len(documents): logger.info( @@ -587,30 +514,10 @@ def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: return normalized -def _table_description_name_candidates(table_names: Optional[list[str]]) -> list[str]: - candidates: list[str] = [] - for table_name in _normalize_table_names(table_names): - raw_candidates = [table_name] - separator_normalized = re.sub(r"[.$]", "_", table_name) - raw_candidates.append(separator_normalized) - if "_" in table_name: - raw_candidates.append( - re.sub(r"^([A-Za-z_][A-Za-z0-9]*)_", r"\1.", table_name, count=1) - ) - if "." in table_name or "$" in table_name: - raw_candidates.append(re.split(r"[.$]", table_name)[-1]) - - for candidate in raw_candidates: - candidate = candidate.strip() - if candidate and candidate not in candidates: - candidates.append(candidate) - return candidates - - def _extract_table_names_from_table_retrieval( - table_retrieval: dict, + table_retrieval: dict, explicit_tables: Optional[list[str]] = None ) -> list[str]: - table_names: list[str] = [] + table_names = _normalize_table_names(explicit_tables) for document in table_retrieval.get("documents") or []: if not isinstance(document, Document): continue @@ -660,7 +567,6 @@ async def table_retrieval( project_id: str, tables: list[str], table_retriever: Any, - max_tables: int = MAX_RELEVANT_TABLE_CANDIDATES, ) -> dict: base_filters = { "operator": "AND", @@ -680,18 +586,17 @@ async def table_retrieval( filters=base_filters, ) results["documents"] = _select_relevant_table_documents( - query, results.get("documents") or [], max_tables=max_tables + query, results.get("documents") or [] ) return results if tables: - table_candidates = _table_description_name_candidates(tables) - logger.info("Loading explicit table descriptions: %s", table_candidates) + logger.info("Loading explicit table descriptions: %s", tables) explicit_filters = { **base_filters, "conditions": [ *base_filters["conditions"], - {"field": "name", "operator": "in", "value": table_candidates}, + {"field": "name", "operator": "in", "value": tables}, ], } return await table_retriever.run(query_embedding=[], filters=explicit_filters) @@ -707,7 +612,9 @@ async def dbschema_retrieval( dbschema_retriever: Any, tables: Optional[list[str]] = None, ) -> list[Document]: - selected_table_names = _extract_table_names_from_table_retrieval(table_retrieval) + selected_table_names = _extract_table_names_from_table_retrieval( + table_retrieval, tables + ) filters = { "operator": "AND", @@ -751,14 +658,7 @@ async def dbschema_retrieval( def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: db_schemas = {} for document in dbschema_retrieval: - try: - content = ast.literal_eval(document.content) - except (SyntaxError, ValueError): - logger.warning("Ignoring malformed schema index document: %s", document.meta) - continue - if not isinstance(content, dict) or not isinstance(content.get("type"), str): - logger.warning("Ignoring schema index document without a type: %s", document.meta) - continue + content = ast.literal_eval(document.content) if content["type"] == "TABLE": if document.meta["name"] not in db_schemas: db_schemas[document.meta["name"]] = content @@ -810,14 +710,7 @@ def check_using_db_schemas_without_pruning( has_json_field = True for document in dbschema_retrieval: - try: - content = ast.literal_eval(document.content) - except (SyntaxError, ValueError): - logger.warning("Ignoring malformed schema index document: %s", document.meta) - continue - if not isinstance(content, dict) or not isinstance(content.get("type"), str): - logger.warning("Ignoring schema index document without a type: %s", document.meta) - continue + content = ast.literal_eval(document.content) if content["type"] == "METRIC": retrieval_results.append( @@ -943,14 +836,7 @@ def construct_retrieval_results( for document in dbschema_retrieval: if document.meta["name"] in columns_and_tables_needed: - try: - content = ast.literal_eval(document.content) - except (SyntaxError, ValueError): - logger.warning("Ignoring malformed schema index document: %s", document.meta) - continue - if not isinstance(content, dict) or not isinstance(content.get("type"), str): - logger.warning("Ignoring schema index document without a type: %s", document.meta) - continue + content = ast.literal_eval(document.content) if content["type"] == "METRIC": retrieval_results.append( @@ -1024,21 +910,12 @@ def __init__( table_column_retrieval_size: int = 100, **kwargs, ): - table_description_candidate_window = max( - table_retrieval_size, - MIN_TABLE_DESCRIPTION_CANDIDATE_WINDOW, - ) self._components = { "embedder": embedder_provider.get_text_embedder(), "table_retriever": document_store_provider.get_retriever( document_store_provider.get_store(dataset_name="table_descriptions"), - top_k=table_description_candidate_window, + top_k=table_retrieval_size, ), - # `top_k` above is intentionally a broad candidate window for reranking. - # The final schema context must, however, honour the configured limit. - # Otherwise every query can send ten complete table schemas to the LLM - # even when an installation configured a smaller retrieval size. - "max_tables": max(1, table_retrieval_size), "dbschema_retriever": document_store_provider.get_retriever( document_store_provider.get_store(), top_k=table_column_retrieval_size, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 4e1bf7e05b..0223f584e2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -641,278 +641,11 @@ def _intent_tokens(self, text: str) -> set[str]: def _schema_name_tokens(self, name: str) -> set[str]: spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(name or "")) - tokens = { + return { token for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) if len(token) > 1 } - normalized = self._normalize_schema_identifier_key(name) - if normalized: - tokens.add(normalized) - if ( - "date" in tokens - or "time" in tokens - or normalized - in { - "createdat", - "updatedat", - "createdon", - "updatedon", - "timestamp", - } - ): - tokens.update({"date", "time", "timestamp"}) - return tokens - - def _schema_terms_for_table(self, table: dict[str, Any]) -> set[str]: - terms = self._schema_name_tokens(str(table.get("name") or "")) - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - if column_name: - terms.update(self._schema_name_tokens(column_name)) - return terms - - def _schema_source_shape_score( - self, query: str | None, table: dict[str, Any] - ) -> int: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - table_terms = self._schema_terms_for_table(table) - - score = 0 - weak_non_production_terms = {"stage", "staging"} - strong_non_production_terms = { - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "temp", - "test", - "tmp", - } - if table_terms & strong_non_production_terms and not any( - re.search(rf"\b{re.escape(term)}\b", normalized_query) - for term in strong_non_production_terms - ): - score -= 80 - if table_terms & weak_non_production_terms and not any( - re.search(rf"\b{re.escape(term)}\b", normalized_query) - for term in weak_non_production_terms - ): - score -= 20 - - transaction_terms = { - "activity", - "detail", - "event", - "fact", - "history", - "invoice", - "line", - "order", - "orders", - "sale", - "sales", - "transaction", - } - reference_terms = { - "account", - "catalog", - "dimension", - "directory", - "entity", - "lookup", - "master", - "profile", - "reference", - } - if any( - term in normalized_query - for term in ( - "amount", - "average", - "count", - "distribution", - "rank", - "ranking", - "sum", - "top", - "total", - "trend", - "value", - ) - ): - if table_terms & transaction_terms: - score += 25 - if table_terms & reference_terms: - score += 5 - - if self._extract_entity_lookup_phrase(query): - if table_terms & transaction_terms: - score += 20 - if table_terms & { - "account", - "buyer", - "client", - "company", - "cust", - "customer", - "custname", - "name", - }: - score += 35 - - return score - - def _semantic_concept_groups_for_query(self, query: str | None) -> list[set[str]]: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - query_tokens = self._intent_tokens(query or "") - if not normalized and not query_tokens: - return [] - - concept_specs: list[tuple[set[str], set[str]]] = [ - ( - {"order", "orders", "neworder", "purchase", "transaction"}, - {"order", "orders", "ord", "ordno", "orderid", "orderdate", "neworder", "purchase", "transaction"}, - ), - ( - {"customer", "customers", "client", "account", "company"}, - {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"}, - ), - ( - {"product", "products", "item", "sku", "category", "categories"}, - {"product", "products", "prod", "prodname", "item", "sku", "category", "categories", "type", "name"}, - ), - ( - {"market", "markets", "region", "regions"}, - {"market", "markets", "region", "regions", "area", "territory", "country"}, - ), - ( - {"business", "unit", "units", "division"}, - {"business", "businessunit", "unit", "units", "bu", "division"}, - ), - ( - {"country", "countries", "destination"}, - {"country", "countries", "destination", "nation", "market", "region"}, - ), - ( - {"invoice", "invoices", "billing"}, - {"invoice", "invoices", "inv", "billing", "bill", "amount", "value"}, - ), - ( - {"amount", "value", "total", "sum", "revenue", "sales", "sale", "cost"}, - {"amount", "value", "total", "sum", "revenue", "sales", "sale", "cost", "price", "money"}, - ), - ( - {"quantity", "qty", "sold", "units"}, - {"quantity", "qty", "sold", "unit", "units", "volume"}, - ), - ( - {"date", "month", "monthly", "year", "quarter", "trend", "period"}, - {"date", "month", "year", "quarter", "time", "timestamp", "orddate", "invdate", "period"}, - ), - ] - - groups: list[set[str]] = [] - for triggers, aliases in concept_specs: - if query_tokens & triggers: - groups.append(aliases) - - if self._extract_entity_lookup_phrase(query): - groups.append( - {"customer", "customers", "cust", "custname", "client", "account", "company", "buyer", "name"} - ) - - return groups - - def _semantic_contract_score_table( - self, query: str | None, table: dict[str, Any] - ) -> tuple[int, set[int]]: - concept_groups = self._semantic_concept_groups_for_query(query) - table_terms = self._schema_terms_for_table(table) - covered_groups = { - index - for index, concept_group in enumerate(concept_groups) - if concept_group & table_terms - } - score = 100 * len(covered_groups) - if concept_groups and len(covered_groups) == len(concept_groups): - score += 80 - elif concept_groups and not covered_groups: - score -= 50 - score += self._schema_source_shape_score(query, table) - return score, covered_groups - - def _scope_retrieval_to_semantic_contract( - self, - query: str, - documents: list[dict], - table_names: list[str], - table_ddls: list[str], - *, - max_tables: int = 4, - ) -> tuple[list[dict], list[str], list[str]]: - concept_groups = self._semantic_concept_groups_for_query(query) - if not concept_groups or len(table_ddls) <= 1: - return documents, table_names, table_ddls - - parsed_tables = self._parse_schema_tables(table_ddls) - if not parsed_tables: - return documents, table_names, table_ddls - - scored: list[tuple[int, int, set[int]]] = [] - for index, table in enumerate(parsed_tables): - score, covered_groups = self._semantic_contract_score_table(query, table) - if covered_groups: - scored.append((score, index, covered_groups)) - - if not scored: - logger.warning( - "No retrieved schema tables cover the semantic contract; keeping original scoped retrieval. query=%s tables=%s", - query, - table_names, - ) - return documents, table_names, table_ddls - - scored = sorted(scored, key=lambda item: (item[0], len(item[2])), reverse=True) - best_score, best_index, best_coverage = scored[0] - all_group_indexes = set(range(len(concept_groups))) - - selected_indexes: list[int] = [] - selected_coverage: set[int] = set() - if best_coverage == all_group_indexes: - selected_indexes.append(best_index) - selected_coverage.update(best_coverage) - else: - for _score, index, coverage in scored: - new_coverage = coverage - selected_coverage - if not new_coverage and selected_indexes: - continue - selected_indexes.append(index) - selected_coverage.update(coverage) - if selected_coverage == all_group_indexes or len(selected_indexes) >= max_tables: - break - - if not selected_indexes: - selected_indexes = [best_index] - - selected_indexes = sorted(dict.fromkeys(selected_indexes)) - if len(selected_indexes) == len(table_ddls): - return documents, table_names, table_ddls - - logger.info( - "Scoped retrieved schema to semantic contract. query=%s before=%s after=%s", - query, - table_names, - [table_names[index] for index in selected_indexes if index < len(table_names)], - ) - return ( - [documents[index] for index in selected_indexes if index < len(documents)], - [table_names[index] for index in selected_indexes if index < len(table_names)], - [table_ddls[index] for index in selected_indexes if index < len(table_ddls)], - ) def _table_for_sql_reference( self, table_reference: str, valid_tables: dict[str, dict[str, Any]] @@ -931,7 +664,7 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: if not normalized: return [] - concept_groups: list[set[str]] = self._semantic_concept_groups_for_query(query) + concept_groups: list[set[str]] = [] if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: @@ -954,16 +687,8 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: concept_groups.append({"currency", "curr", "money", "fx", "exchange"}) if "market" in normalized or "markets" in normalized: concept_groups.append({"market", "region", "country", "territory"}) - if ( - "business unit" in normalized - or "business units" in normalized - or re.search(r"\bbu\b", normalized) - ): - concept_groups.append({"business", "businessunit", "unit", "units", "bu", "division"}) if "region" in normalized or "regions" in normalized: concept_groups.append({"region", "market", "area", "territory", "country"}) - if "country" in normalized or "countries" in normalized: - concept_groups.append({"country", "countries", "nation", "destination"}) if "quarterly" in normalized or "quarter" in normalized: concept_groups.append({"quarter", "quarterly"}) if "recurring" in normalized or "recurrence" in normalized: @@ -980,17 +705,12 @@ def _sql_covers_required_question_concepts( referenced_column_tokens: set[str], referenced_table_tokens: set[str], ) -> bool: - sql_text = re.sub( - r"\b(?:order|group)\s+by\b|\bselect\b|\bfrom\b|\bwhere\b", - " ", - (sql or "").lower(), - ) - sql_text_tokens = self._intent_tokens(sql_text) + sql_text = (sql or "").lower() available_tokens = referenced_column_tokens | referenced_table_tokens for concept_group in self._required_sql_concept_groups(query): if concept_group & available_tokens: continue - if concept_group & sql_text_tokens: + if any(token in sql_text for token in concept_group): continue logger.warning( "Ignoring SQL because it does not cover required question concept. " @@ -1048,65 +768,6 @@ def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> return False return True - def _is_grouped_metric_or_ranking_query(self, query: str | None) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return False - - has_metric_word = any( - term in normalized_query - for term in ( - "amount", - "average", - "avg", - "count", - "counts", - "distribution", - "how many", - "number of", - "record count", - "revenue", - "sum", - "total", - "value", - ) - ) - has_ranked_order_metric = any( - term in normalized_query - for term in ("top", "most", "highest", "largest", "rank", "ranking") - ) and any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ) - has_dimension_word = any( - term in normalized_query - for term in ( - "business unit", - "business units", - "category", - "categories", - "customer", - "customers", - "custname", - "currency", - "currencies", - "division", - "market", - "markets", - "product", - "products", - "region", - "regions", - "sales person", - "salesperson", - "source", - "status", - "type", - ) - ) or bool(re.search(r"\bbu\b", normalized_query)) - - return (has_metric_word or has_ranked_order_metric) and has_dimension_word - def _sql_satisfies_count_ranking_request( self, sql: str, query: str | None ) -> bool: @@ -1130,21 +791,16 @@ def _sql_satisfies_count_ranking_request( for term in ("order", "orders", "record", "records", "row", "rows") ) ) - asks_for_grouped_metric = self._is_grouped_metric_or_ranking_query(query) - if not asks_for_count_metric and not asks_for_grouped_metric: + if not asks_for_count_metric: return True - asks_for_grouped_entity = asks_for_grouped_metric or any( + asks_for_grouped_entity = any( term in normalized_query for term in ( - "business unit", - "business units", "category", "customer", "customers", - "custname", "currency", - "division", "market", "product", "products", @@ -1155,7 +811,7 @@ def _sql_satisfies_count_ranking_request( "status", "type", ) - ) or bool(re.search(r"\bbu\b", normalized_query)) + ) if not asks_for_grouped_entity: return True @@ -1253,181 +909,6 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True - def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip()) - if not normalized_query: - return None - - match = re.search( - r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" - r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", - normalized_query, - flags=re.IGNORECASE, - ) - if not match: - return None - - phrase = match.group("phrase").strip(" .,;:()[]{}'\"") - phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) - if not phrase or len(phrase) < 3: - return None - if re.search( - r"\b(?:table|model|schema|column|columns|market|region|country|division|" - r"date|month|year|quarter|top|count|number|amount|value)\b", - phrase, - flags=re.IGNORECASE, - ): - return None - return phrase - - def _preferred_entity_lookup_columns( - self, query: str | None, table: dict[str, Any] - ) -> set[str]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - candidate_groups: list[tuple[str, ...]] = [] - if "account" in normalized_query: - candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) - if "company" in normalized_query: - candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) - candidate_groups.extend( - [ - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - ), - ("Client", "ClientName"), - ("Account", "AccountName"), - ("Company", "CompanyName"), - ("Name",), - ] - ) - - columns: set[str] = set() - for candidates in candidate_groups: - column = self._find_schema_column(table, candidates) - if column: - columns.add(column) - return columns - - def _sql_satisfies_entity_lookup_request( - self, - sql: str, - query: str | None, - referenced_tables: list[str], - referenced_columns_by_table: dict[str, set[str]], - valid_tables: dict[str, dict[str, Any]], - ) -> bool: - if not self._extract_entity_lookup_phrase(query): - return True - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if any( - term in normalized_query - for term in (" by ", " per ", " each ", "distribution", "top", "count") - ): - return True - - for table_reference in referenced_tables: - table = self._table_for_sql_reference(table_reference, valid_tables) - if not table: - continue - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - - table_key = str(table_reference or "").lower() - referenced_columns = referenced_columns_by_table.get( - table_key - ) or referenced_columns_by_table.get( - table_key.split(".")[-1], - set(), - ) - referenced_column_keys = { - self._normalize_schema_identifier_key(column) - for column in referenced_columns - } - preferred_column_keys = { - self._normalize_schema_identifier_key(column) - for column in preferred_columns - } - if referenced_column_keys & preferred_column_keys: - return True - - logger.warning( - "Ignoring SQL because entity lookup did not use available customer/name columns. " - "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", - query, - table.get("name"), - sorted(preferred_columns), - sorted(referenced_columns), - sql, - ) - return False - - return True - - def _sql_satisfies_named_entity_request( - self, - sql: str, - query: str | None, - referenced_tables: list[str], - valid_tables: dict[str, dict[str, Any]], - ) -> bool: - """Require a name field when the question explicitly asks for names. - - An account number, customer code, or generic ``Customer`` field may be a - useful identifier, but it is not evidence that the datasource exposes a - customer name. Returning such an identifier for a name request produces - plausible but incorrect answers. - """ - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - requests_customer_name = bool( - re.search( - r"\b(?:customer|customers|client|clients|account|accounts|company|companies)\s+names?\b" - r"|\bnames?\s+(?:of|for)\s+(?:customers|clients|accounts|companies)\b", - normalized_query, - ) - ) - if not requests_customer_name: - return True - - name_columns = ( - "CustomerName", - "CustName", - "ClientName", - "AccountName", - "CompanyName", - "Name", - ) - available_name_columns: set[str] = set() - for table_reference in referenced_tables: - table = self._table_for_sql_reference(table_reference, valid_tables) - if not table: - continue - if column := self._find_schema_column(table, name_columns): - available_name_columns.add(column) - - for column in available_name_columns: - escaped_column = re.escape(column) - if re.search( - rf'(?i)(?:"{escaped_column}"|\[{escaped_column}\]|\b{escaped_column}\b)', - sql, - ): - return True - - logger.warning( - "Ignoring SQL because a customer-name request was mapped to an identifier or " - "a datasource without a customer-name column. query=%s available_name_columns=%s sql=%s", - query, - sorted(available_name_columns), - sql, - ) - return False - def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1741,21 +1222,6 @@ def _sql_matches_question_intent( return False if not self._sql_satisfies_unique_entity_request(sql, query): return False - if not self._sql_satisfies_entity_lookup_request( - sql, - query, - referenced_tables, - referenced_columns_by_table, - valid_tables, - ): - return False - if not self._sql_satisfies_named_entity_request( - sql, - query, - referenced_tables, - valid_tables, - ): - return False if not expects_dimension: return True @@ -2207,18 +1673,9 @@ def _build_schema_grounded_table_question_sql( else "COUNT(*)" ) top_clause = f"TOP {limit} " if wants_ranked_count else "" - dimension_type = next( - ( - str(column.get("type") or "") - for column in table.get("columns", []) - if self._normalize_schema_identifier_key(column.get("name")) - == self._normalize_schema_identifier_key(dimension_column) - ), - "", - ) nonblank_filter = ( f"AND LTRIM(RTRIM({dimension_ref})) <> '' " - if self._is_text_schema_type(dimension_type) + if wants_ranked_count else "" ) return ( @@ -2287,8 +1744,6 @@ def _build_explicit_table_preview_sql( normalized_query = re.sub(r"\s+", " ", (query or "").strip()) if not normalized_query: return None - if self._is_grouped_metric_or_ranking_query(query): - return None if not re.search( r"\b(?:first|top|sample|preview|show|list)\b", @@ -2407,18 +1862,7 @@ def _explicit_table_name_candidates(self, table_name: str) -> list[str]: separator_normalized = re.sub(r"[.$]", "_", table_name) if separator_normalized not in candidates: candidates.append(separator_normalized) - if "_" in table_name: - dotted_schema_name = re.sub( - r"^([A-Za-z_][A-Za-z0-9]*)_", - r"\1.", - table_name, - count=1, - ) - if dotted_schema_name not in candidates: - candidates.append(dotted_schema_name) short_name = re.split(r"[.$]", table_name)[-1] - if short_name == table_name and "_" in table_name: - short_name = table_name.split("_", 1)[-1] if short_name and short_name not in candidates: candidates.append(short_name) return candidates @@ -2655,7 +2099,6 @@ def _select_best_analytics_table( query: str = "", ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - concept_groups = self._semantic_concept_groups_for_query(query) scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] ] = [] @@ -2698,17 +2141,6 @@ def _select_best_analytics_table( table_name = str(table.get("name") or "").lower() if not table_name: continue - schema_terms = self._schema_terms_for_table(table) - if concept_groups: - covered_groups = sum( - 1 for concept_group in concept_groups if concept_group & schema_terms - ) - score += 14 * covered_groups - if covered_groups == len(concept_groups): - score += 25 - elif covered_groups == 0: - score -= 20 - score += self._schema_source_shape_score(query, table) if "sales" in table_name: score += 5 if "tblsales" in self._normalize_schema_token(table_name): @@ -2783,88 +2215,6 @@ def _select_best_analytics_table( date_column, ) - def _build_entity_lookup_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - lookup_phrase = self._extract_entity_lookup_phrase(query) - if not lookup_phrase: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - asks_for_orders = any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ) - - scored: list[tuple[int, dict[str, Any], str]] = [] - for table in tables: - table_name = str(table.get("name") or "") - if not table_name: - continue - - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - - preferred_column = sorted( - preferred_columns, - key=lambda column: ( - 0 - if self._normalize_schema_identifier_key(column) - in {"custname", "customername", "customer"} - else 1, - column.lower(), - ), - )[0] - - score = 20 - normalized_table = self._normalize_schema_token(table_name) - if asks_for_orders: - if "order" in normalized_table: - score += 40 - if "neworder" in normalized_table: - score += 20 - if self._find_schema_column( - table, ("OrdNo", "OrderNo", "OrderId", "NewOrderId") - ): - score += 25 - if "test" in normalized_table or "tmp" in normalized_table: - score -= 80 - if "dev" in normalized_table or "backup" in normalized_table: - score -= 60 - if "stage" in normalized_table: - score -= 10 - scored.append((score, table, preferred_column)) - - if not scored: - return None - - _score, table, filter_column = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - table_name = str(table.get("name") or "") - if not table_name: - return None - - table_ref = self._quote_sql_identifier(table_name) - filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" - escaped_phrase = lookup_phrase.replace("'", "''") - date_column = self._find_schema_column( - table, - ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), - temporal=True, - ) - order_clause = ( - f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" - if date_column - else "" - ) - return ( - f"SELECT TOP 500 * FROM {table_ref} " - f"WHERE {filter_ref} = '{escaped_phrase}'" - f"{order_clause}" - ) - def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2878,9 +2228,6 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): - return entity_lookup_sql - if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2936,7 +2283,16 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - asks_for_measure_value = any( + if not is_sales_or_order_query: + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + + wants_count_metric = any( + term in normalized_query + for term in ("count", "counts", "volume", "how many", "distribution") + ) and not any( term in normalized_query for term in ( "amount", @@ -2948,22 +2304,11 @@ def _build_schema_grounded_analytics_sql( "sale", "sales", "sold", - "sum", - "total", - "value", - ) - ) - - if not is_sales_or_order_query and not asks_for_measure_value: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not asks_for_measure_value + "sum", + "total", + "value", + ) + ) wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") ) @@ -3051,9 +2396,7 @@ def _build_schema_grounded_analytics_sql( ) ) if "country" in normalized_query or "countries" in normalized_query: - dimension_candidates.append( - ("Country", "CountryName", "Nation", "Destination") - ) + dimension_candidates.append(("Country", "CountryName", "Nation", "Market")) if "division" in normalized_query: dimension_candidates.append(("Division",)) if ( @@ -3200,9 +2543,6 @@ def _build_schema_grounded_analytics_sql( "volume", "how many", "number of", - "top", - "highest", - "most", "monthly", "over time", "last 12 months", @@ -5505,19 +4845,8 @@ def _extract_retrieval_documents(self, retrieval_result: dict) -> list[dict]: if not isinstance(document, dict): logger.warning("Ignoring malformed retrieval document: %s", document) continue - table_name = document.get("table_name") - table_ddl = document.get("table_ddl") - if not isinstance(table_name, str) or not table_name.strip(): - logger.warning("Ignoring retrieval document without a table name") - continue - if not isinstance(table_ddl, str) or not table_ddl.strip(): - # A table name alone is not safe SQL-generation context. Treat this - # exactly like an unsuccessful schema retrieval rather than allowing - # the generator to infer columns from the rest of the project. - logger.warning( - "Ignoring retrieval document without filtered DDL for table %s", - table_name, - ) + if not document.get("table_name") and not document.get("table_ddl"): + logger.warning("Ignoring retrieval document without table metadata") continue valid_documents.append(document) @@ -6687,6 +6016,61 @@ async def ask( ) sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + logger.info( + "Retrieved tables for direct heuristic query_id %s: %s", + query_id, + table_names, + ) + + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using direct heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + if ask_result := self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ): + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + if explicit_group_count_sql := self._build_explicit_group_count_sql( user_query ): @@ -7017,10 +6401,10 @@ async def ask( tables=retrieval_table_names, histories=[], project_id=ask_request.project_id, - # Keep the SQL prompt bounded for every data question. - # This is the legacy behavior: select the relevant columns - # before SQL generation rather than sending full schemas. - enable_column_pruning=enable_column_pruning, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -7107,10 +6491,281 @@ async def ask( explicit_table_names, ) ) + if ( + not documents + and self._should_load_full_schema_for_question(user_query) + and not request_explicit_table_names + ): + logger.info( + "Query-based schema retrieval returned no tables for data question; " + "retrying full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) + if not api_results and ( + table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ) + ): + logger.info( + "Using schema-grounded table question SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + table_question_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = table_question_sql + error_message = "Schema-grounded table SQL was not valid for the active datasource schema." + + if not api_results and ( + explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + ): + explicit_sql, explicit_table_name = explicit_table_preview + logger.info( + "Using explicit table preview SQL for query_id %s and table %s", + query_id, + explicit_table_name, + ) + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = explicit_sql + error_message = "Explicit table preview SQL was not valid for the active datasource schema." + + if not api_results and ( + audit_log_activity_sql := self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ) + ): + logger.info( + "Using schema-grounded audit log activity SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + audit_log_activity_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = audit_log_activity_sql + error_message = ( + "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." + ) + + if ( + not api_results + and self._is_data_analysis_query(user_query) + and ( + schema_grounded_sql := self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + ) + ): + logger.info( + "Using generic schema-grounded analytics SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + schema_grounded_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = schema_grounded_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and any( + term in user_query.lower() + for term in ( + "pcb", + "repair", + "failure", + "business unit", + "business units", + "product line", + "product family", + ) + ): + operational_sql = self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + if operational_sql: + logger.info( + "Using schema-grounded operational SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + operational_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = operational_sql + error_message = ( + "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and ( + deterministic_sales_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ) + ): + logger.info( + "Using schema-grounded CWSales SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sales_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = deterministic_sales_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." + ) + + should_retry_full_schema = ( + not api_results + and self._should_load_full_schema_for_question(user_query) + and "db_schema_retrieval" in self._pipelines + and not request_explicit_table_names + and not table_names + ) + if should_retry_full_schema: + logger.info( + "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retry", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 30, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + full_documents, full_table_names, full_table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + full_documents, full_table_names, full_table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + full_documents, + explicit_table_names, + ) + ) + if full_documents: + documents, table_names, table_ddls = ( + full_documents, + full_table_names, + full_table_ddls, + ) + logger.info( + "Using full active deployed schema retry for query_id %s: %s", + query_id, + table_names, + ) + + full_schema_preview = self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + full_schema_sql_candidates = ( + self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ), + full_schema_preview[0] if full_schema_preview else None, + self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ), + self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ), + self._build_schema_grounded_sales_sql( + user_query, table_ddls + ), + ) + for full_schema_sql in full_schema_sql_candidates: + if not full_schema_sql: + continue + ask_result = self._build_validated_ask_result_from_sql( + full_schema_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + break + invalid_sql = full_schema_sql + error_message = ( + "Full-schema grounded SQL was not valid for the active datasource schema and question intent." + ) + if not api_results and ( unqueryable_metric_message := self._get_unqueryable_metric_message( user_query, table_ddls @@ -7158,13 +6813,40 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - # The retrieval pipeline has already selected tables and, when enabled, - # pruned their columns. Do not reload those names as explicit tables: - # that turns the selected context back into full schemas and can send - # many unrelated tables to SQL generation. Legacy/v1 passes the - # filtered DDL produced above directly to reasoning and generation. + if documents and not api_results: + documents, table_names, table_ddls = self._prune_sql_generation_context( + sql_user_query, + documents, + table_names, + table_ddls, + ) + ( + documents, + table_names, + table_ddls, + completed_retrieval_result, + ) = await self._complete_sql_generation_context( + query=sql_user_query, + project_id=ask_request.project_id, + documents=documents, + table_names=table_names, + table_ddls=table_ddls, + ) + if completed_retrieval_result: + _retrieval_result = completed_retrieval_result sql_generation_histories = histories + if self._is_data_analysis_query( + sql_user_query + ) and not self._needs_conversation_context(sql_user_query): + sql_generation_histories = [] + allow_sql_generation_reasoning = False + allow_sql_knowledge_retrieval = False + max_sql_correction_retries = min(max_sql_correction_retries, 1) + logger.info( + "Using fast standalone SQL generation path for query_id %s", + query_id, + ) if ( not self._is_stopped(query_id, self._ask_results) @@ -7329,7 +7011,7 @@ async def ask( ) except TimeoutError as generation_timeout: logger.warning( - "SQL generation timed out for query_id %s; returning a generation failure: %s", + "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", query_id, generation_timeout, ) @@ -7465,6 +7147,46 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: + if self._can_use_schema_grounded_sql_fallback( + documents, + table_ddls, + user_query, + ) and ( + heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ) + ): + logger.info( + "Using heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + else: + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 8514b45dfc..6dc0b45205 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -397,32 +397,6 @@ def test_normalize_sql_column_references_to_schema_maps_sales_business_aliases() ) == [] -def test_customer_name_is_not_normalized_to_an_account_identifier(): - sql = 'SELECT "dbo_qSalesMargin"."CustomerName" FROM "dbo_qSalesMargin"' - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_qSalesMargin": ["Account"]}, - ) - - assert '"dbo_qSalesMargin"."Account"' not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_qSalesMargin": ["Account"]}, - ) == ["dbo_qSalesMargin.CustomerName"] - - -def test_customer_name_is_normalized_to_a_real_customer_name_column(): - sql = 'SELECT "dbo_qSalesMargin"."CustomerName" FROM "dbo_qSalesMargin"' - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_qSalesMargin": ["CustName"]}, - ) - - assert normalized == 'SELECT "dbo_qSalesMargin"."CustName" FROM "dbo_qSalesMargin"' - - def test_normalize_sql_column_references_to_schema_maps_period_to_timeid(): sql = 'SELECT "dbo_tblFactSales"."Period" FROM "dbo_tblFactSales"' diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 102f183a39..3ee705b621 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,7 +2,6 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( - DbSchemaRetrieval, _is_project_wide_analysis_query, _rerank_table_documents, _select_relevant_table_documents, @@ -68,106 +67,6 @@ def test_rerank_table_documents_prefers_question_relevant_table_text(): assert documents[0].meta["name"] == "business_transactions" -def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): - test_load = Document( - content="Raw test load rows for order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ) - order_market_table = Document( - content="New order transaction records with market and customer fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.45, - ) - - documents = _rerank_table_documents( - "Show order distribution across markets.", - [test_load, order_market_table], - ) - - assert documents[0].meta["name"] == "dbo_xStageNewOrders" - - -def test_select_relevant_table_documents_excludes_unrequested_dev_equivalent(): - dev_cube = Document( - content="Development sales cube with order number and market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_qSalesCubeDev"}, - score=500, - ) - production_orders = Document( - content="New order transaction records with order number and market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, - score=0.2, - ) - - documents = _select_relevant_table_documents( - "Show order distribution across markets.", - [dev_cube, production_orders], - ) - - assert [document.meta["name"] for document in documents] == ["dbo_tblNewOrders"] - - -def test_select_relevant_table_documents_keeps_explicit_dev_request(): - dev_cube = Document( - content="Development sales cube with order number and market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_qSalesCubeDev"}, - score=500, - ) - production_orders = Document( - content="New order transaction records with order number and market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, - score=0.2, - ) - - documents = _select_relevant_table_documents( - "Show order distribution across markets in dev.", - [dev_cube, production_orders], - ) - - assert documents[0].meta["name"] == "dbo_qSalesCubeDev" - - -def test_select_relevant_table_documents_prefers_full_business_unit_order_coverage(): - weak_sales_margin = Document( - content="Sales margin facts with sales value, margin, order date, and product.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_qSalesMargin"}, - score=500, - ) - new_orders = Document( - content="New order records with BU, business unit, order number, order date, and customer.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, - score=0.2, - ) - - documents = _select_relevant_table_documents( - "Which business unit has the top 20 new orders this period?", - [weak_sales_margin, new_orders], - ) - - assert [document.meta["name"] for document in documents] == ["dbo_tblNewOrders"] - - -def test_rerank_table_documents_prefers_customer_capable_source_for_entity_lookup(): - generic_order_table = Document( - content="New order rows by division and market.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tnoStageNewOrders"}, - score=0.95, - ) - customer_order_table = Document( - content="New order transactions with customer name, order number, market, and customer purchase order.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_tblNewOrders"}, - score=0.45, - ) - - documents = _rerank_table_documents( - "Show me orders for Lockheed Martin.", - [generic_order_table, customer_order_table], - ) - - assert documents[0].meta["name"] == "dbo_tblNewOrders" - - def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( @@ -333,82 +232,6 @@ async def run(self, query_embedding, filters): } -@pytest.mark.asyncio -async def test_table_retrieval_expands_explicit_table_name_forms_for_descriptions(): - class Retriever: - def __init__(self): - self.filters = None - - async def run(self, query_embedding, filters): - self.filters = filters - return {"documents": []} - - retriever = Retriever() - - await table_retrieval( - query="show failed repairs", - embedding={}, - project_id="project-1", - tables=["dbo.failure_patterns"], - table_retriever=retriever, - ) - - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - { - "field": "name", - "operator": "in", - "value": [ - "dbo.failure_patterns", - "dbo_failure_patterns", - "failure_patterns", - ], - }, - ], - } - - -def test_db_schema_retrieval_fetches_wider_table_description_window(): - class LLMProvider: - def get_generator(self, **kwargs): - return object() - - def get_model(self): - return "gpt-4o-mini" - - def get_context_window_size(self): - return 1000 - - class EmbedderProvider: - def get_text_embedder(self): - return object() - - class DocumentStoreProvider: - def __init__(self): - self.retriever_top_k = [] - - def get_store(self, dataset_name=None): - return dataset_name or "default" - - def get_retriever(self, store, top_k): - self.retriever_top_k.append((store, top_k)) - return object() - - document_store_provider = DocumentStoreProvider() - - DbSchemaRetrieval( - llm_provider=LLMProvider(), - embedder_provider=EmbedderProvider(), - document_store_provider=document_store_provider, - table_retrieval_size=10, - ) - - assert document_store_provider.retriever_top_k[0] == ("table_descriptions", 100) - - @pytest.mark.asyncio async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: @@ -527,18 +350,18 @@ def encode(self, value): @pytest.mark.asyncio -async def test_dbschema_retrieval_does_not_use_explicit_tables_without_description(): +async def test_dbschema_retrieval_uses_explicit_tables_as_scope(): class Retriever: def __init__(self): - self.called = False + self.filters = None async def run(self, query_embedding, filters): - self.called = True + self.filters = filters return {"documents": []} retriever = Retriever() - documents = await dbschema_retrieval( + await dbschema_retrieval( query="show failed repairs", table_retrieval={"documents": []}, project_id="project-1", @@ -546,5 +369,15 @@ async def run(self, query_embedding, filters): tables=["dbo.failure_patterns", "dbo_failure_patterns"], ) - assert documents == [] - assert not retriever.called + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + { + "field": "name", + "operator": "in", + "value": ["dbo.failure_patterns", "dbo_failure_patterns"], + }, + ], + } diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 0ee413f27c..a64038fbb9 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,324 +269,6 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) -def test_explicit_table_preview_does_not_handle_grouped_count_request(): - service = AskService.__new__(AskService) - - preview_sql = service._build_explicit_table_preview_sql( - "From dbo_tblNewOrders, show the top 5 customers by order count using CustName.", - [ - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - ) - - assert preview_sql is None - - -def test_validated_sql_rejects_country_question_without_country_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' - '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' - 'ORDER BY COUNT(*) DESC' - ), - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - "Show the total commodity line value by country.", - ) - - assert result is None - - -def test_validated_sql_rejects_order_distribution_without_order_concept(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT TOP 10 "dbo_xStageLoad8_Test"."Market" AS "Market", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_xStageLoad8_Test" ' - 'WHERE "dbo_xStageLoad8_Test"."Market" IS NOT NULL ' - 'GROUP BY "dbo_xStageLoad8_Test"."Market" ' - 'ORDER BY COUNT(*) DESC' - ), - [ - """ - CREATE TABLE dbo_xStageLoad8_Test ( - Market VARCHAR, - LoadId VARCHAR - ); - """ - ], - "Show order distribution across markets.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Show the total commodity line value by country.", - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'AS "Country_of_Ultimate_Destination_Code", ' - 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' - 'AS "TotalCommodity_Line_Value" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' - ) - - -def test_schema_grounded_analytics_prefers_full_concept_coverage_for_order_market_distribution(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Show order distribution across markets.", - [ - """ - CREATE TABLE dbo_xStageLoad8_Test ( - Market VARCHAR, - LoadId VARCHAR - ); - """, - """ - CREATE TABLE dbo_tblNewOrders ( - Market VARCHAR, - OrdNo VARCHAR - ); - """, - ], - ) - - assert sql == ( - 'SELECT "dbo_tblNewOrders"."Market" AS "Market", ' - 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "OrderCount" ' - 'FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."Market" IS NOT NULL ' - 'GROUP BY "dbo_tblNewOrders"."Market" ' - 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' - ) - - -def test_scope_retrieval_to_semantic_contract_prefers_production_full_coverage_table(): - service = AskService.__new__(AskService) - documents = [ - { - "table_name": "dbo_qSalesCubeDev", - "table_ddl": """ - CREATE TABLE dbo_qSalesCubeDev ( - Market VARCHAR, - OrdNo VARCHAR - ); - """, - }, - { - "table_name": "dbo_tblNewOrders", - "table_ddl": """ - CREATE TABLE dbo_tblNewOrders ( - Market VARCHAR, - OrdNo VARCHAR - ); - """, - }, - ] - - scoped_documents, scoped_table_names, scoped_table_ddls = ( - service._scope_retrieval_to_semantic_contract( - "Show order distribution across markets.", - documents, - ["dbo_qSalesCubeDev", "dbo_tblNewOrders"], - [document["table_ddl"] for document in documents], - ) - ) - - assert scoped_documents == [documents[1]] - assert scoped_table_names == ["dbo_tblNewOrders"] - assert scoped_table_ddls == [documents[1]["table_ddl"]] - - -def test_scope_retrieval_to_semantic_contract_keeps_multiple_tables_when_needed(): - service = AskService.__new__(AskService) - documents = [ - { - "table_name": "dbo_OrderFacts", - "table_ddl": """ - CREATE TABLE dbo_OrderFacts ( - Market VARCHAR, - OrdNo VARCHAR - ); - """, - }, - { - "table_name": "dbo_Customers", - "table_ddl": """ - CREATE TABLE dbo_Customers ( - CustNo VARCHAR, - CustName VARCHAR - ); - """, - }, - { - "table_name": "dbo_LoadAudit", - "table_ddl": """ - CREATE TABLE dbo_LoadAudit ( - LoadId VARCHAR, - Status VARCHAR - ); - """, - }, - ] - - scoped_documents, scoped_table_names, _scoped_table_ddls = ( - service._scope_retrieval_to_semantic_contract( - "Show order distribution by market and customer.", - documents, - ["dbo_OrderFacts", "dbo_Customers", "dbo_LoadAudit"], - [document["table_ddl"] for document in documents], - ) - ) - - assert scoped_documents == documents[:2] - assert scoped_table_names == ["dbo_OrderFacts", "dbo_Customers"] - - -def test_schema_grounded_analytics_counts_top_new_orders_by_business_unit(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Which business unit has the top 20 new orders this period?", - [ - """ - CREATE TABLE dbo_tblNewOrders ( - BU VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 20 "dbo_tblNewOrders"."BU" AS "BU", ' - 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "OrderCount" ' - 'FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."BU" IS NOT NULL ' - 'GROUP BY "dbo_tblNewOrders"."BU" ' - 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' - ) - - -def test_validated_sql_rejects_business_unit_question_without_business_unit_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT TOP 20 "dbo_qSalesMargin"."ProductCategory" AS "ProductCategory", ' - 'COUNT(DISTINCT "dbo_qSalesMargin"."OrdNo") AS "OrderCount" ' - 'FROM "dbo_qSalesMargin" ' - 'GROUP BY "dbo_qSalesMargin"."ProductCategory" ' - 'ORDER BY COUNT(DISTINCT "dbo_qSalesMargin"."OrdNo") DESC' - ), - [ - """ - CREATE TABLE dbo_qSalesMargin ( - ProductCategory VARCHAR, - OrdNo VARCHAR, - OrderDate TIMESTAMP - ); - """ - ], - "Which business unit has the top 20 new orders this period?", - ) - - assert result is None - - -def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT * FROM "dbo_tnoStageNewOrders" ' - 'WHERE "dbo_tnoStageNewOrders"."Division" = ' - "'Daimler Trucks North America'" - ), - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - "List orders for Daimler Trucks North America.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "List orders for Daimler Trucks North America.", - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """, - ], - ) - - assert sql == ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" = ' - "'Daimler Trucks North America' " - 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' - ) - - def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) @@ -764,21 +446,6 @@ def test_extract_explicit_table_names_from_pcb_repair_phrases(): ) == ["ticket_labels", "dbo_ticket_labels"] -def test_explicit_table_name_candidates_include_dotted_and_short_forms(): - service = AskService.__new__(AskService) - - assert service._explicit_table_name_candidates("dbo_tblNewOrders") == [ - "dbo_tblNewOrders", - "dbo.tblNewOrders", - "tblNewOrders", - ] - assert service._explicit_table_name_candidates("dbo.tblNewOrders") == [ - "dbo.tblNewOrders", - "dbo_tblNewOrders", - "tblNewOrders", - ] - - def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): service = AskService.__new__(AskService) documents = [ @@ -815,42 +482,6 @@ def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): assert table_ddls == [documents[1]["table_ddl"]] -def test_filter_retrieval_metadata_for_explicit_query_matches_dotted_table_name(): - service = AskService.__new__(AskService) - documents = [ - { - "table_name": "dbo.tblNewOrders", - "table_ddl": """ - CREATE TABLE "dbo.tblNewOrders" ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - }, - { - "table_name": "dbo_other", - "table_ddl": """ - CREATE TABLE dbo_other ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - }, - ] - - filtered_documents, table_names, table_ddls = ( - service._filter_retrieval_metadata_for_explicit_query( - "Show the top 5 CustName values from dbo_tblNewOrders by number of orders.", - documents, - ["dbo_tblNewOrders"], - ) - ) - - assert filtered_documents == [documents[0]] - assert table_names == ["dbo.tblNewOrders"] - assert table_ddls == [documents[0]["table_ddl"]] - - def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): service = AskService.__new__(AskService) result = service._build_validated_ask_result_from_sql( From f2e92ed958f18bdd53daf894947a45b1266d9f4b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 16:51:47 +0530 Subject: [PATCH 0532/1087] Restore working code from 20e44de --- .../retrieval/db_schema_retrieval.py | 33 ++- wren-ai-service/src/web/v1/services/ask.py | 245 +++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 20 ++ .../pytest/services/test_ask_sales_sql.py | 160 ++++++++++++ 4 files changed, 436 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..6b64c9b452 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -213,11 +213,15 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms = { - _normalize_retrieval_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") - if len(token) > 2 and token.lower() not in stop_words - } + terms: set[str] = set() + for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): + split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) + for token in re.findall(r"[A-Za-z0-9]+", split_token): + if len(token) <= 2 or token.lower() in stop_words: + continue + normalized_token = _normalize_retrieval_token(token) + if normalized_token: + terms.add(normalized_token) return {term for term in terms if term} @@ -243,7 +247,11 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - non_production_terms = ( + weak_non_production_terms = ( + "stage", + "staging", + ) + strong_non_production_terms = ( "archive", "backup", "copy", @@ -251,17 +259,18 @@ def _source_shape_score(query: str, document: Document) -> int: "development", "duplicate", "sample", - "stage", - "staging", "temp", "test", "tmp", ) - if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( - normalized_query, - non_production_terms, + if source_terms & set(strong_non_production_terms) and not _query_mentions_any( + normalized_query, strong_non_production_terms + ): + score -= 240 + if source_terms & set(weak_non_production_terms) and not _query_mentions_any( + normalized_query, weak_non_production_terms ): - score -= 60 + score -= 40 aggregation_terms = ( "amount", diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0223f584e2..ce42c81888 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -689,6 +689,8 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: concept_groups.append({"market", "region", "country", "territory"}) if "region" in normalized or "regions" in normalized: concept_groups.append({"region", "market", "area", "territory", "country"}) + if "country" in normalized or "countries" in normalized: + concept_groups.append({"country", "countries", "nation", "destination"}) if "quarterly" in normalized or "quarter" in normalized: concept_groups.append({"quarter", "quarterly"}) if "recurring" in normalized or "recurrence" in normalized: @@ -909,6 +911,123 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True + def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) + if not normalized_query: + return None + + match = re.search( + r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" + r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", + normalized_query, + flags=re.IGNORECASE, + ) + if not match: + return None + + phrase = match.group("phrase").strip(" .,;:()[]{}'\"") + phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) + if not phrase or len(phrase) < 3: + return None + if re.search( + r"\b(?:table|model|schema|column|columns|market|region|country|division|" + r"date|month|year|quarter|top|count|number|amount|value)\b", + phrase, + flags=re.IGNORECASE, + ): + return None + return phrase + + def _preferred_entity_lookup_columns( + self, query: str | None, table: dict[str, Any] + ) -> set[str]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + candidate_groups: list[tuple[str, ...]] = [] + if "account" in normalized_query: + candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) + if "company" in normalized_query: + candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) + candidate_groups.extend( + [ + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + ), + ("Client", "ClientName"), + ("Account", "AccountName"), + ("Company", "CompanyName"), + ("Name",), + ] + ) + + columns: set[str] = set() + for candidates in candidate_groups: + column = self._find_schema_column(table, candidates) + if column: + columns.add(column) + return columns + + def _sql_satisfies_entity_lookup_request( + self, + sql: str, + query: str | None, + referenced_tables: list[str], + referenced_columns_by_table: dict[str, set[str]], + valid_tables: dict[str, dict[str, Any]], + ) -> bool: + if not self._extract_entity_lookup_phrase(query): + return True + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if any( + term in normalized_query + for term in (" by ", " per ", " each ", "distribution", "top", "count") + ): + return True + + for table_reference in referenced_tables: + table = self._table_for_sql_reference(table_reference, valid_tables) + if not table: + continue + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + + table_key = str(table_reference or "").lower() + referenced_columns = referenced_columns_by_table.get( + table_key + ) or referenced_columns_by_table.get( + table_key.split(".")[-1], + set(), + ) + referenced_column_keys = { + self._normalize_schema_identifier_key(column) + for column in referenced_columns + } + preferred_column_keys = { + self._normalize_schema_identifier_key(column) + for column in preferred_columns + } + if referenced_column_keys & preferred_column_keys: + return True + + logger.warning( + "Ignoring SQL because entity lookup did not use available customer/name columns. " + "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", + query, + table.get("name"), + sorted(preferred_columns), + sorted(referenced_columns), + sql, + ) + return False + + return True + def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1222,6 +1341,14 @@ def _sql_matches_question_intent( return False if not self._sql_satisfies_unique_entity_request(sql, query): return False + if not self._sql_satisfies_entity_lookup_request( + sql, + query, + referenced_tables, + referenced_columns_by_table, + valid_tables, + ): + return False if not expects_dimension: return True @@ -1862,7 +1989,18 @@ def _explicit_table_name_candidates(self, table_name: str) -> list[str]: separator_normalized = re.sub(r"[.$]", "_", table_name) if separator_normalized not in candidates: candidates.append(separator_normalized) + if "_" in table_name: + dotted_schema_name = re.sub( + r"^([A-Za-z_][A-Za-z0-9]*)_", + r"\1.", + table_name, + count=1, + ) + if dotted_schema_name not in candidates: + candidates.append(dotted_schema_name) short_name = re.split(r"[.$]", table_name)[-1] + if short_name == table_name and "_" in table_name: + short_name = table_name.split("_", 1)[-1] if short_name and short_name not in candidates: candidates.append(short_name) return candidates @@ -2215,6 +2353,88 @@ def _select_best_analytics_table( date_column, ) + def _build_entity_lookup_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + lookup_phrase = self._extract_entity_lookup_phrase(query) + if not lookup_phrase: + return None + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + asks_for_orders = any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ) + + scored: list[tuple[int, dict[str, Any], str]] = [] + for table in tables: + table_name = str(table.get("name") or "") + if not table_name: + continue + + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + + preferred_column = sorted( + preferred_columns, + key=lambda column: ( + 0 + if self._normalize_schema_identifier_key(column) + in {"custname", "customername", "customer"} + else 1, + column.lower(), + ), + )[0] + + score = 20 + normalized_table = self._normalize_schema_token(table_name) + if asks_for_orders: + if "order" in normalized_table: + score += 40 + if "neworder" in normalized_table: + score += 20 + if self._find_schema_column( + table, ("OrdNo", "OrderNo", "OrderId", "NewOrderId") + ): + score += 25 + if "test" in normalized_table or "tmp" in normalized_table: + score -= 80 + if "dev" in normalized_table or "backup" in normalized_table: + score -= 60 + if "stage" in normalized_table: + score -= 10 + scored.append((score, table, preferred_column)) + + if not scored: + return None + + _score, table, filter_column = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + table_name = str(table.get("name") or "") + if not table_name: + return None + + table_ref = self._quote_sql_identifier(table_name) + filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" + escaped_phrase = lookup_phrase.replace("'", "''") + date_column = self._find_schema_column( + table, + ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), + temporal=True, + ) + order_clause = ( + f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" + if date_column + else "" + ) + return ( + f"SELECT TOP 500 * FROM {table_ref} " + f"WHERE {filter_ref} = '{escaped_phrase}'" + f"{order_clause}" + ) + def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2228,6 +2448,9 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): + return entity_lookup_sql + if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2283,16 +2506,7 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - if not is_sales_or_order_query: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not any( + asks_for_measure_value = any( term in normalized_query for term in ( "amount", @@ -2309,6 +2523,17 @@ def _build_schema_grounded_analytics_sql( "value", ) ) + + if not is_sales_or_order_query and not asks_for_measure_value: + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + + wants_count_metric = any( + term in normalized_query + for term in ("count", "counts", "volume", "how many", "distribution") + ) and not asks_for_measure_value wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 3ee705b621..bbb0c79aef 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -67,6 +67,26 @@ def test_rerank_table_documents_prefers_question_relevant_table_text(): assert documents[0].meta["name"] == "business_transactions" +def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): + test_load = Document( + content="Raw test load rows for order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ) + order_market_table = Document( + content="New order transaction records with market and customer fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.45, + ) + + documents = _rerank_table_documents( + "Show order distribution across markets.", + [test_load, order_market_table], + ) + + assert documents[0].meta["name"] == "dbo_xStageNewOrders" + + def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index a64038fbb9..5ba8ca74b8 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,6 +269,115 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) +def test_validated_sql_rejects_country_question_without_country_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' + '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' + 'ORDER BY COUNT(*) DESC' + ), + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + "Show the total commodity line value by country.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Show the total commodity line value by country.", + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'AS "Country_of_Ultimate_Destination_Code", ' + 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' + 'AS "TotalCommodity_Line_Value" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' + ) + + +def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT * FROM "dbo_tnoStageNewOrders" ' + 'WHERE "dbo_tnoStageNewOrders"."Division" = ' + "'Daimler Trucks North America'" + ), + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + "List orders for Daimler Trucks North America.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "List orders for Daimler Trucks North America.", + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """, + ], + ) + + assert sql == ( + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."CustName" = ' + "'Daimler Trucks North America' " + 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' + ) + + def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) @@ -446,6 +555,21 @@ def test_extract_explicit_table_names_from_pcb_repair_phrases(): ) == ["ticket_labels", "dbo_ticket_labels"] +def test_explicit_table_name_candidates_include_dotted_and_short_forms(): + service = AskService.__new__(AskService) + + assert service._explicit_table_name_candidates("dbo_tblNewOrders") == [ + "dbo_tblNewOrders", + "dbo.tblNewOrders", + "tblNewOrders", + ] + assert service._explicit_table_name_candidates("dbo.tblNewOrders") == [ + "dbo.tblNewOrders", + "dbo_tblNewOrders", + "tblNewOrders", + ] + + def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): service = AskService.__new__(AskService) documents = [ @@ -482,6 +606,42 @@ def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): assert table_ddls == [documents[1]["table_ddl"]] +def test_filter_retrieval_metadata_for_explicit_query_matches_dotted_table_name(): + service = AskService.__new__(AskService) + documents = [ + { + "table_name": "dbo.tblNewOrders", + "table_ddl": """ + CREATE TABLE "dbo.tblNewOrders" ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + }, + { + "table_name": "dbo_other", + "table_ddl": """ + CREATE TABLE dbo_other ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + }, + ] + + filtered_documents, table_names, table_ddls = ( + service._filter_retrieval_metadata_for_explicit_query( + "Show the top 5 CustName values from dbo_tblNewOrders by number of orders.", + documents, + ["dbo_tblNewOrders"], + ) + ) + + assert filtered_documents == [documents[0]] + assert table_names == ["dbo.tblNewOrders"] + assert table_ddls == [documents[0]["table_ddl"]] + + def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): service = AskService.__new__(AskService) result = service._build_validated_ask_result_from_sql( From c59b2ba468496fc5bc15114731a9b3185922fba3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 17:15:52 +0530 Subject: [PATCH 0533/1087] Fallback to deployed schema retrieval --- .../retrieval/db_schema_retrieval.py | 71 ++++++++--- .../retrieval/test_db_schema_retrieval.py | 110 +++++++++++++++--- 2 files changed, 148 insertions(+), 33 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6b64c9b452..a0ac305d50 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -600,15 +600,20 @@ async def table_retrieval( return results if tables: - logger.info("Loading explicit table descriptions: %s", tables) - explicit_filters = { - **base_filters, - "conditions": [ - *base_filters["conditions"], - {"field": "name", "operator": "in", "value": tables}, - ], + normalized_tables = _normalize_table_names(tables) + logger.info( + "Using explicit table names without table-description lookup: %s", + normalized_tables, + ) + return { + "documents": [ + Document( + content=str({"name": table_name}), + meta={"type": "TABLE_DESCRIPTION", "name": table_name}, + ) + for table_name in normalized_tables + ] } - return await table_retriever.run(query_embedding=[], filters=explicit_filters) return {"documents": []} @@ -620,6 +625,7 @@ async def dbschema_retrieval( project_id: str, dbschema_retriever: Any, tables: Optional[list[str]] = None, + embedding: Optional[dict] = None, ) -> list[Document]: selected_table_names = _extract_table_names_from_table_retrieval( table_retrieval, tables @@ -651,13 +657,50 @@ async def dbschema_retrieval( project_id, ) else: - logger.info( - "No relevant table-description candidates found for active project_id %s; " - "skipping full schema loading for query=%s", - project_id, - query, + query_embedding = ( + embedding.get("embedding") if isinstance(embedding, dict) else None ) - return [] + if query_embedding: + logger.info( + "No table-description candidates found for active project_id %s; " + "falling back to deployed schema vector retrieval for query=%s", + project_id, + query, + ) + candidate_results = await dbschema_retriever.run( + query_embedding=query_embedding, + filters=filters, + ) + selected_table_names = _extract_table_names_from_table_retrieval( + candidate_results + )[:MAX_RELEVANT_TABLE_CANDIDATES] + if selected_table_names: + filters["conditions"].append( + {"field": "name", "operator": "in", "value": selected_table_names} + ) + logger.info( + "Loading deployed schema metadata from fallback candidates for " + "active project_id %s tables=%s", + project_id, + selected_table_names, + ) + else: + documents = candidate_results.get("documents") or [] + logger.info( + "No deployed schema fallback candidates found for active " + "project_id %s query=%s", + project_id, + query, + ) + return documents + else: + logger.info( + "No relevant table-description candidates found for active project_id %s; " + "skipping full schema loading for query=%s", + project_id, + query, + ) + return [] results = await dbschema_retriever.run(query_embedding=[], filters=filters) return results.get("documents", []) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index bbb0c79aef..8bbade7251 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -223,33 +223,23 @@ def test_rerank_table_documents_prefers_transaction_source_for_metric_question() @pytest.mark.asyncio -async def test_table_retrieval_fetches_explicit_table_descriptions(): +async def test_table_retrieval_uses_explicit_table_names_without_vector_lookup(): class Retriever: - def __init__(self): - self.filters = None - async def run(self, query_embedding, filters): - self.filters = filters - return {"documents": []} + raise AssertionError("explicit table names should not use vector lookup") - retriever = Retriever() - - await table_retrieval( + result = await table_retrieval( query="show rows", embedding={}, project_id="project-1", - tables=["orders"], - table_retriever=retriever, + tables=["orders", "orders", "customers"], + table_retriever=Retriever(), ) - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - {"field": "name", "operator": "in", "value": ["orders"]}, - ], - } + assert [document.meta["name"] for document in result["documents"]] == [ + "orders", + "customers", + ] @pytest.mark.asyncio @@ -335,6 +325,88 @@ async def run(self, query_embedding, filters): assert not retriever.called +@pytest.mark.asyncio +async def test_dbschema_retrieval_falls_back_to_deployed_schema_vector_search(): + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append((query_embedding, filters)) + if query_embedding: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "sales_orders", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "sales_orders"}, + score=0.9, + ) + ] + } + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "sales_orders", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "sales_orders"}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "name": "sales_orders", + "columns": [{"name": "amount"}], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "sales_orders"}, + ), + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + query="compare sales between countries", + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=retriever, + embedding={"embedding": [0.1, 0.2]}, + ) + + assert [document.meta["name"] for document in documents] == [ + "sales_orders", + "sales_orders", + ] + assert retriever.calls[0][0] == [0.1, 0.2] + assert retriever.calls[0][1] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert retriever.calls[1][0] == [] + assert retriever.calls[1][1] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "name", "operator": "in", "value": ["sales_orders"]}, + ], + } + + def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): class Encoding: def encode(self, value): From 8fdcce06994b4c5cb9bb70c35b7acf857fa18d3d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 18:16:53 +0530 Subject: [PATCH 0534/1087] Use explicit query table names for schema retrieval --- .../retrieval/db_schema_retrieval.py | 159 ++++++++++++++++-- .../retrieval/test_db_schema_retrieval.py | 136 +++++++++++++++ 2 files changed, 277 insertions(+), 18 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index a0ac305d50..471081179a 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -523,6 +523,84 @@ def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: return normalized +def _add_table_name_variants(table_names: list[str], candidate: str) -> None: + cleaned_candidate = str(candidate or "").strip().strip(".,;:()") + cleaned_candidate = ( + cleaned_candidate.replace("[", "") + .replace("]", "") + .replace('"', "") + .replace("`", "") + ) + if not cleaned_candidate: + return + + variants = [cleaned_candidate] + if "." in cleaned_candidate: + parts = [part for part in cleaned_candidate.split(".") if part] + if len(parts) >= 2: + qualified_name = ".".join(parts[-2:]) + underscored_name = "_".join(parts[-2:]) + bare_name = parts[-1] + variants.extend([qualified_name, underscored_name, bare_name]) + elif "_" in cleaned_candidate: + parts = [part for part in cleaned_candidate.split("_") if part] + if len(parts) >= 2: + variants.append(".".join(parts[-2:])) + variants.append(parts[-1]) + + for variant in variants: + if variant and variant not in table_names: + table_names.append(variant) + + +def _looks_like_explicit_table_reference(candidate: str) -> bool: + cleaned_candidate = str(candidate or "").strip() + if not cleaned_candidate: + return False + + if any(character in cleaned_candidate for character in (".", "_", "[", "]", "`")): + return True + + if re.match(r"(?i)^(dbo|tbl|xStage|ytbl)[A-Za-z0-9_]*$", cleaned_candidate): + return True + + return any(character.isupper() for character in cleaned_candidate[1:]) + + +def _extract_explicit_table_names_from_query(query: str) -> list[str]: + table_names: list[str] = [] + normalized_query = query or "" + + for match in re.finditer( + r"(? list[str]: @@ -577,6 +655,24 @@ async def table_retrieval( tables: list[str], table_retriever: Any, ) -> dict: + explicit_tables = _normalize_table_names( + [*(tables or []), *_extract_explicit_table_names_from_query(query)] + ) + if explicit_tables: + logger.info( + "Using explicit table names without table-description lookup: %s", + explicit_tables, + ) + return { + "documents": [ + Document( + content=str({"name": table_name}), + meta={"type": "TABLE_DESCRIPTION", "name": table_name}, + ) + for table_name in explicit_tables + ] + } + base_filters = { "operator": "AND", "conditions": [ @@ -599,22 +695,6 @@ async def table_retrieval( ) return results - if tables: - normalized_tables = _normalize_table_names(tables) - logger.info( - "Using explicit table names without table-description lookup: %s", - normalized_tables, - ) - return { - "documents": [ - Document( - content=str({"name": table_name}), - meta={"type": "TABLE_DESCRIPTION", "name": table_name}, - ) - for table_name in normalized_tables - ] - } - return {"documents": []} @@ -628,7 +708,8 @@ async def dbschema_retrieval( embedding: Optional[dict] = None, ) -> list[Document]: selected_table_names = _extract_table_names_from_table_retrieval( - table_retrieval, tables + table_retrieval, + [*(tables or []), *_extract_explicit_table_names_from_query(query)], ) filters = { @@ -703,7 +784,49 @@ async def dbschema_retrieval( return [] results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results.get("documents", []) + documents = results.get("documents", []) + if documents or not selected_table_names: + return documents + + query_embedding = embedding.get("embedding") if isinstance(embedding, dict) else None + if not query_embedding: + return documents + + fallback_filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + if project_id: + fallback_filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + logger.info( + "Selected schema names returned no deployed metadata for active project_id %s; " + "falling back to deployed schema vector retrieval for query=%s tables=%s", + project_id, + query, + selected_table_names, + ) + candidate_results = await dbschema_retriever.run( + query_embedding=query_embedding, + filters=fallback_filters, + ) + fallback_table_names = _extract_table_names_from_table_retrieval( + candidate_results + )[:MAX_RELEVANT_TABLE_CANDIDATES] + if not fallback_table_names: + return candidate_results.get("documents", []) + + fallback_filters["conditions"].append( + {"field": "name", "operator": "in", "value": fallback_table_names} + ) + fallback_results = await dbschema_retriever.run( + query_embedding=[], + filters=fallback_filters, + ) + return fallback_results.get("documents", []) @observe() diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 8bbade7251..bb0d25f00b 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,6 +2,7 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( + _extract_explicit_table_names_from_query, _is_project_wide_analysis_query, _rerank_table_documents, _select_relevant_table_documents, @@ -47,6 +48,20 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query +def test_extract_explicit_table_names_from_query_handles_schema_dot_reference(): + table_names = _extract_explicit_table_names_from_query( + "How many records are in dbo.xStageLoad8?" + ) + + assert table_names[:3] == ["dbo.xStageLoad8", "dbo_xStageLoad8", "xStageLoad8"] + + +def test_extract_explicit_table_names_from_query_ignores_plain_business_terms(): + assert _extract_explicit_table_names_from_query( + "Compare sales in countries by month" + ) == [] + + def test_rerank_table_documents_prefers_question_relevant_table_text(): generic_stage = Document( content="Generic imported staging records with product labels.", @@ -242,6 +257,27 @@ async def run(self, query_embedding, filters): ] +@pytest.mark.asyncio +async def test_table_retrieval_uses_query_table_names_before_embedding_lookup(): + class Retriever: + async def run(self, query_embedding, filters): + raise AssertionError("explicit query table name should not use vector lookup") + + result = await table_retrieval( + query="How many records are in dbo.xStageLoad8?", + embedding={"embedding": [0.1, 0.2]}, + project_id="project-1", + tables=[], + table_retriever=Retriever(), + ) + + assert [document.meta["name"] for document in result["documents"]] == [ + "dbo.xStageLoad8", + "dbo_xStageLoad8", + "xStageLoad8", + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: @@ -325,6 +361,39 @@ async def run(self, query_embedding, filters): assert not retriever.called +@pytest.mark.asyncio +async def test_dbschema_retrieval_uses_query_table_names_as_scope(): + class Retriever: + def __init__(self): + self.filters = None + + async def run(self, query_embedding, filters): + self.filters = filters + return {"documents": []} + + retriever = Retriever() + + await dbschema_retrieval( + query="How many records are in dbo.xStageLoad8?", + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=retriever, + ) + + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + { + "field": "name", + "operator": "in", + "value": ["dbo.xStageLoad8", "dbo_xStageLoad8", "xStageLoad8"], + }, + ], + } + + @pytest.mark.asyncio async def test_dbschema_retrieval_falls_back_to_deployed_schema_vector_search(): class Retriever: @@ -407,6 +476,73 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_dbschema_retrieval_falls_back_when_selected_schema_lookup_is_empty(): + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append((query_embedding, filters)) + if query_embedding: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "dbo_xStageLoad8", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "dbo_xStageLoad8"}, + ) + ] + } + if len(self.calls) == 1: + return {"documents": []} + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "dbo_xStageLoad8", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "dbo_xStageLoad8"}, + ) + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + query="How many records are in dbo.xStageLoad8?", + table_retrieval={ + "documents": [ + Document( + content=str({"name": "dbo.xStageLoad8"}), + meta={"type": "TABLE_DESCRIPTION", "name": "dbo.xStageLoad8"}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + embedding={"embedding": [0.1, 0.2]}, + ) + + assert [document.meta["name"] for document in documents] == ["dbo_xStageLoad8"] + assert retriever.calls[0][0] == [] + assert retriever.calls[1][0] == [0.1, 0.2] + assert retriever.calls[2][1]["conditions"][-1] == { + "field": "name", + "operator": "in", + "value": ["dbo_xStageLoad8"], + } + + def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): class Encoding: def encode(self, value): From 441d34c5647da70341adc486148b44b33e2f20b3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 18:41:01 +0530 Subject: [PATCH 0535/1087] Avoid explicit schema timeout failures --- wren-ai-service/src/web/v1/services/ask.py | 74 ++++++++++++++- .../v1/services/question_recommendation.py | 91 +++++++++++++++++-- 2 files changed, 153 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ce42c81888..af557237e0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2015,6 +2015,49 @@ def _normalize_explicit_table_names( normalized.append(candidate) return normalized + def _select_direct_explicit_sql_table_name(self, table_names: list[str]) -> str | None: + normalized_names = self._normalize_explicit_table_names(table_names) + if not normalized_names: + return None + + original_name = normalized_names[0] + if "." in original_name: + for table_name in normalized_names: + if "." not in table_name and "_" in table_name: + return table_name + return original_name + + def _build_direct_explicit_table_count_sql( + self, query: str, table_names: list[str] + ) -> tuple[str, str] | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized or not table_names: + return None + + wants_total_count = ( + re.search(r"\bhow many\b", normalized) + or "record count" in normalized + or "count of records" in normalized + or "number of records" in normalized + or "number of rows" in normalized + or "count rows" in normalized + ) + if not wants_total_count or re.search( + r"\b(?:by|per|each|distribution|highest|top|monthly|daily|weekly)\b", + normalized, + ): + return None + + table_name = self._select_direct_explicit_sql_table_name(table_names) + if not table_name: + return None + + return ( + f'SELECT COUNT(*) AS "RecordCount" FROM ' + f"{self._quote_sql_identifier(table_name)}", + table_name, + ) + def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -6084,6 +6127,30 @@ async def ask( return results if explicit_table_names: + if direct_count_sql := self._build_direct_explicit_table_count_sql( + user_query, + explicit_table_names, + ): + explicit_sql, explicit_table_name = direct_count_sql + ask_result = self._build_ask_result_from_sql(explicit_sql) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning=( + "Explicit table count request generated SQL directly." + ), + retrieved_tables=[explicit_table_name], + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -6104,7 +6171,7 @@ async def ask( timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 20, + 60, ), ) documents, table_names, table_ddls = ( @@ -6700,7 +6767,8 @@ async def ask( ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, - 20, + self._pipeline_timeout_seconds, + 60, ), ) _retrieval_result = retrieval_result.get( @@ -6737,7 +6805,7 @@ async def ask( timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 20, + 45, ), ) _retrieval_result = retrieval_result.get( diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index d93e456d2a..b5d66044fc 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -31,6 +31,8 @@ DEFAULT_VALIDATION_SQL_FUNCTION_CHARS = 2500 STRICT_VALIDATION_SQL_FUNCTION_ITEMS = 0 STRICT_VALIDATION_SQL_FUNCTION_CHARS = 0 +DEFAULT_RECOMMENDATION_RETRIEVAL_TIMEOUT_SECONDS = 20 +DEFAULT_RECOMMENDATION_VALIDATION_TIMEOUT_SECONDS = 20 DEFAULT_QUESTION_CATEGORIES = [ "Descriptive Questions", "Segmentation Questions", @@ -173,6 +175,42 @@ def _get_underfilled_categories( if len(response_questions.get(category, [])) < max_questions ] + def _add_recommended_questions( + self, + event_id: str, + candidates: list[dict], + *, + max_questions: int, + max_categories: int, + ) -> None: + current = self._cache[event_id] + questions = current.response.setdefault("questions", {}) + + for candidate in candidates: + category = candidate.get("category") + question = candidate.get("question") + if not isinstance(category, str) or not isinstance(question, str): + continue + if category not in questions and len(questions) >= max_categories: + continue + + category_questions = questions.setdefault(category, []) + existing = next( + ( + item + for item in category_questions + if item.get("question") == question + ), + None, + ) + if existing is not None: + existing.update(candidate) + continue + + if len(category_questions) >= max_questions: + continue + category_questions.append(candidate) + def _handle_exception( self, event_id: str, @@ -350,6 +388,11 @@ async def _instructions_retrieval() -> list[dict]: currnet_category = questions.setdefault(candidate["category"], []) + for existing_question in currnet_category: + if existing_question.get("question") == candidate["question"]: + existing_question.update({**candidate, "sql": valid_sql}) + return post_process + if len(currnet_category) >= max_questions: # Skip to update the questions for the category if it is already full return post_process @@ -373,6 +416,12 @@ class Request(BaseRequest): async def _recommend(self, request: dict): resp = await self._pipelines["question_recommendation"].run(**request) questions = resp.get("normalized", {}).get("questions", []) + self._add_recommended_questions( + request["event_id"], + questions, + max_questions=request["max_questions"], + max_categories=request["max_categories"], + ) validation_tasks = [ self._validate_question( question, @@ -385,7 +434,19 @@ async def _recommend(self, request: dict): for question in questions ] - await asyncio.gather(*validation_tasks, return_exceptions=True) + if not validation_tasks: + return + + try: + await asyncio.wait_for( + asyncio.gather(*validation_tasks, return_exceptions=True), + timeout=DEFAULT_RECOMMENDATION_VALIDATION_TIMEOUT_SECONDS, + ) + except TimeoutError: + logger.warning( + "Request %s: Question recommendation validation timed out; returning generated recommendations without full SQL validation", + request["event_id"], + ) @observe(name="Generate Question Recommendation") @trace_metadata @@ -397,14 +458,26 @@ async def recommend(self, input: Request, **kwargs) -> Event: try: orjson.loads(input.mdl) - retrieval_result = await self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=input.project_id, - enable_column_pruning=False, - ) - _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) - documents = _retrieval_result.get("retrieval_results", []) + try: + retrieval_result = await asyncio.wait_for( + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=input.project_id, + enable_column_pruning=False, + ), + timeout=DEFAULT_RECOMMENDATION_RETRIEVAL_TIMEOUT_SECONDS, + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents = _retrieval_result.get("retrieval_results", []) + except TimeoutError: + logger.warning( + "Request %s: Question recommendation schema retrieval timed out; continuing with generated recommendations from limited context", + input.event_id, + ) + documents = [] table_ddls = self._limit_text_items( [document.get("table_ddl") for document in documents], max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, From 0ace2e54e5319e6d9afa4463ac3c65a35215bc7a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 18:49:09 +0530 Subject: [PATCH 0536/1087] Generate explicit table SQL without schema wait --- wren-ai-service/src/web/v1/services/ask.py | 102 +++++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index af557237e0..255d06ef96 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -2058,6 +2058,84 @@ def _build_direct_explicit_table_count_sql( table_name, ) + def _extract_direct_explicit_projection_columns(self, query: str) -> list[str]: + columns: list[str] = [] + match = re.search( + r"\b(?:including|include|with)\s+(.+?)(?:\.|$)", + query or "", + flags=re.IGNORECASE, + ) + if not match: + return columns + + raw_columns = re.sub(r"\band\b", ",", match.group(1), flags=re.IGNORECASE) + stop_words = { + "columns", + "fields", + "including", + "include", + "with", + "and", + } + for raw_column in re.split(r",", raw_columns): + column = raw_column.strip().strip("`\"[] ") + if not column: + continue + if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", column): + continue + if column.lower() in stop_words: + continue + if column not in columns: + columns.append(column) + return columns + + def _build_direct_explicit_table_projection_sql( + self, query: str, table_names: list[str] + ) -> tuple[str, str] | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized or not table_names: + return None + + table_name = self._select_direct_explicit_sql_table_name(table_names) + if not table_name: + return None + + columns = self._extract_direct_explicit_projection_columns(query) + asks_for_preview = bool( + re.search(r"\b(?:preview|sample|first)\b", normalized) + or re.search(r"\b(?:rows?|records?|data)\b", normalized) + ) + if not columns and not asks_for_preview: + return None + + limit = self._extract_requested_top_n(query, default_value=500) + table_ref = self._quote_sql_identifier(table_name) + select_clause = ( + ", ".join(self._quote_sql_identifier(column) for column in columns) + if columns + else "*" + ) + + where_parts: list[str] = [] + for column in columns: + if re.search( + rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+empty", + normalized, + ): + column_ref = self._quote_sql_identifier(column) + where_parts.append(f"{column_ref} IS NOT NULL AND {column_ref} <> ''") + elif re.search( + rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+null", + normalized, + ): + where_parts.append(f"{self._quote_sql_identifier(column)} IS NOT NULL") + + where_clause = f" WHERE {' AND '.join(where_parts)}" if where_parts else "" + return ( + f"SELECT TOP {limit} {select_clause} FROM {table_ref}{where_clause}", + table_name, + ) + def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -6151,6 +6229,30 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + if direct_projection_sql := self._build_direct_explicit_table_projection_sql( + user_query, + explicit_table_names, + ): + explicit_sql, explicit_table_name = direct_projection_sql + ask_result = self._build_ask_result_from_sql(explicit_sql) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning=( + "Explicit table request generated SQL directly." + ), + retrieved_tables=[explicit_table_name], + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", From 7a63f1373cf252feeff9fe817354c59db5f74666 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 14 Jul 2026 18:50:38 +0530 Subject: [PATCH 0537/1087] Restore service code to 97075b46 working state --- .../retrieval/db_schema_retrieval.py | 237 ++-------- wren-ai-service/src/web/v1/services/ask.py | 421 +----------------- .../v1/services/question_recommendation.py | 91 +--- .../retrieval/test_db_schema_retrieval.py | 266 +---------- .../pytest/services/test_ask_sales_sql.py | 160 ------- 5 files changed, 72 insertions(+), 1103 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 471081179a..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -213,15 +213,11 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms: set[str] = set() - for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): - split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) - for token in re.findall(r"[A-Za-z0-9]+", split_token): - if len(token) <= 2 or token.lower() in stop_words: - continue - normalized_token = _normalize_retrieval_token(token) - if normalized_token: - terms.add(normalized_token) + terms = { + _normalize_retrieval_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") + if len(token) > 2 and token.lower() not in stop_words + } return {term for term in terms if term} @@ -247,11 +243,7 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - weak_non_production_terms = ( - "stage", - "staging", - ) - strong_non_production_terms = ( + non_production_terms = ( "archive", "backup", "copy", @@ -259,18 +251,17 @@ def _source_shape_score(query: str, document: Document) -> int: "development", "duplicate", "sample", + "stage", + "staging", "temp", "test", "tmp", ) - if source_terms & set(strong_non_production_terms) and not _query_mentions_any( - normalized_query, strong_non_production_terms - ): - score -= 240 - if source_terms & set(weak_non_production_terms) and not _query_mentions_any( - normalized_query, weak_non_production_terms + if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( + normalized_query, + non_production_terms, ): - score -= 40 + score -= 60 aggregation_terms = ( "amount", @@ -523,84 +514,6 @@ def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: return normalized -def _add_table_name_variants(table_names: list[str], candidate: str) -> None: - cleaned_candidate = str(candidate or "").strip().strip(".,;:()") - cleaned_candidate = ( - cleaned_candidate.replace("[", "") - .replace("]", "") - .replace('"', "") - .replace("`", "") - ) - if not cleaned_candidate: - return - - variants = [cleaned_candidate] - if "." in cleaned_candidate: - parts = [part for part in cleaned_candidate.split(".") if part] - if len(parts) >= 2: - qualified_name = ".".join(parts[-2:]) - underscored_name = "_".join(parts[-2:]) - bare_name = parts[-1] - variants.extend([qualified_name, underscored_name, bare_name]) - elif "_" in cleaned_candidate: - parts = [part for part in cleaned_candidate.split("_") if part] - if len(parts) >= 2: - variants.append(".".join(parts[-2:])) - variants.append(parts[-1]) - - for variant in variants: - if variant and variant not in table_names: - table_names.append(variant) - - -def _looks_like_explicit_table_reference(candidate: str) -> bool: - cleaned_candidate = str(candidate or "").strip() - if not cleaned_candidate: - return False - - if any(character in cleaned_candidate for character in (".", "_", "[", "]", "`")): - return True - - if re.match(r"(?i)^(dbo|tbl|xStage|ytbl)[A-Za-z0-9_]*$", cleaned_candidate): - return True - - return any(character.isupper() for character in cleaned_candidate[1:]) - - -def _extract_explicit_table_names_from_query(query: str) -> list[str]: - table_names: list[str] = [] - normalized_query = query or "" - - for match in re.finditer( - r"(? list[str]: @@ -655,24 +568,6 @@ async def table_retrieval( tables: list[str], table_retriever: Any, ) -> dict: - explicit_tables = _normalize_table_names( - [*(tables or []), *_extract_explicit_table_names_from_query(query)] - ) - if explicit_tables: - logger.info( - "Using explicit table names without table-description lookup: %s", - explicit_tables, - ) - return { - "documents": [ - Document( - content=str({"name": table_name}), - meta={"type": "TABLE_DESCRIPTION", "name": table_name}, - ) - for table_name in explicit_tables - ] - } - base_filters = { "operator": "AND", "conditions": [ @@ -695,6 +590,17 @@ async def table_retrieval( ) return results + if tables: + logger.info("Loading explicit table descriptions: %s", tables) + explicit_filters = { + **base_filters, + "conditions": [ + *base_filters["conditions"], + {"field": "name", "operator": "in", "value": tables}, + ], + } + return await table_retriever.run(query_embedding=[], filters=explicit_filters) + return {"documents": []} @@ -705,11 +611,9 @@ async def dbschema_retrieval( project_id: str, dbschema_retriever: Any, tables: Optional[list[str]] = None, - embedding: Optional[dict] = None, ) -> list[Document]: selected_table_names = _extract_table_names_from_table_retrieval( - table_retrieval, - [*(tables or []), *_extract_explicit_table_names_from_query(query)], + table_retrieval, tables ) filters = { @@ -738,95 +642,16 @@ async def dbschema_retrieval( project_id, ) else: - query_embedding = ( - embedding.get("embedding") if isinstance(embedding, dict) else None + logger.info( + "No relevant table-description candidates found for active project_id %s; " + "skipping full schema loading for query=%s", + project_id, + query, ) - if query_embedding: - logger.info( - "No table-description candidates found for active project_id %s; " - "falling back to deployed schema vector retrieval for query=%s", - project_id, - query, - ) - candidate_results = await dbschema_retriever.run( - query_embedding=query_embedding, - filters=filters, - ) - selected_table_names = _extract_table_names_from_table_retrieval( - candidate_results - )[:MAX_RELEVANT_TABLE_CANDIDATES] - if selected_table_names: - filters["conditions"].append( - {"field": "name", "operator": "in", "value": selected_table_names} - ) - logger.info( - "Loading deployed schema metadata from fallback candidates for " - "active project_id %s tables=%s", - project_id, - selected_table_names, - ) - else: - documents = candidate_results.get("documents") or [] - logger.info( - "No deployed schema fallback candidates found for active " - "project_id %s query=%s", - project_id, - query, - ) - return documents - else: - logger.info( - "No relevant table-description candidates found for active project_id %s; " - "skipping full schema loading for query=%s", - project_id, - query, - ) - return [] + return [] results = await dbschema_retriever.run(query_embedding=[], filters=filters) - documents = results.get("documents", []) - if documents or not selected_table_names: - return documents - - query_embedding = embedding.get("embedding") if isinstance(embedding, dict) else None - if not query_embedding: - return documents - - fallback_filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - } - if project_id: - fallback_filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) - logger.info( - "Selected schema names returned no deployed metadata for active project_id %s; " - "falling back to deployed schema vector retrieval for query=%s tables=%s", - project_id, - query, - selected_table_names, - ) - candidate_results = await dbschema_retriever.run( - query_embedding=query_embedding, - filters=fallback_filters, - ) - fallback_table_names = _extract_table_names_from_table_retrieval( - candidate_results - )[:MAX_RELEVANT_TABLE_CANDIDATES] - if not fallback_table_names: - return candidate_results.get("documents", []) - - fallback_filters["conditions"].append( - {"field": "name", "operator": "in", "value": fallback_table_names} - ) - fallback_results = await dbschema_retriever.run( - query_embedding=[], - filters=fallback_filters, - ) - return fallback_results.get("documents", []) + return results.get("documents", []) @observe() diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 255d06ef96..0223f584e2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -689,8 +689,6 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: concept_groups.append({"market", "region", "country", "territory"}) if "region" in normalized or "regions" in normalized: concept_groups.append({"region", "market", "area", "territory", "country"}) - if "country" in normalized or "countries" in normalized: - concept_groups.append({"country", "countries", "nation", "destination"}) if "quarterly" in normalized or "quarter" in normalized: concept_groups.append({"quarter", "quarterly"}) if "recurring" in normalized or "recurrence" in normalized: @@ -911,123 +909,6 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True - def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip()) - if not normalized_query: - return None - - match = re.search( - r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" - r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", - normalized_query, - flags=re.IGNORECASE, - ) - if not match: - return None - - phrase = match.group("phrase").strip(" .,;:()[]{}'\"") - phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) - if not phrase or len(phrase) < 3: - return None - if re.search( - r"\b(?:table|model|schema|column|columns|market|region|country|division|" - r"date|month|year|quarter|top|count|number|amount|value)\b", - phrase, - flags=re.IGNORECASE, - ): - return None - return phrase - - def _preferred_entity_lookup_columns( - self, query: str | None, table: dict[str, Any] - ) -> set[str]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - candidate_groups: list[tuple[str, ...]] = [] - if "account" in normalized_query: - candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) - if "company" in normalized_query: - candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) - candidate_groups.extend( - [ - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - ), - ("Client", "ClientName"), - ("Account", "AccountName"), - ("Company", "CompanyName"), - ("Name",), - ] - ) - - columns: set[str] = set() - for candidates in candidate_groups: - column = self._find_schema_column(table, candidates) - if column: - columns.add(column) - return columns - - def _sql_satisfies_entity_lookup_request( - self, - sql: str, - query: str | None, - referenced_tables: list[str], - referenced_columns_by_table: dict[str, set[str]], - valid_tables: dict[str, dict[str, Any]], - ) -> bool: - if not self._extract_entity_lookup_phrase(query): - return True - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if any( - term in normalized_query - for term in (" by ", " per ", " each ", "distribution", "top", "count") - ): - return True - - for table_reference in referenced_tables: - table = self._table_for_sql_reference(table_reference, valid_tables) - if not table: - continue - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - - table_key = str(table_reference or "").lower() - referenced_columns = referenced_columns_by_table.get( - table_key - ) or referenced_columns_by_table.get( - table_key.split(".")[-1], - set(), - ) - referenced_column_keys = { - self._normalize_schema_identifier_key(column) - for column in referenced_columns - } - preferred_column_keys = { - self._normalize_schema_identifier_key(column) - for column in preferred_columns - } - if referenced_column_keys & preferred_column_keys: - return True - - logger.warning( - "Ignoring SQL because entity lookup did not use available customer/name columns. " - "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", - query, - table.get("name"), - sorted(preferred_columns), - sorted(referenced_columns), - sql, - ) - return False - - return True - def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1341,14 +1222,6 @@ def _sql_matches_question_intent( return False if not self._sql_satisfies_unique_entity_request(sql, query): return False - if not self._sql_satisfies_entity_lookup_request( - sql, - query, - referenced_tables, - referenced_columns_by_table, - valid_tables, - ): - return False if not expects_dimension: return True @@ -1989,18 +1862,7 @@ def _explicit_table_name_candidates(self, table_name: str) -> list[str]: separator_normalized = re.sub(r"[.$]", "_", table_name) if separator_normalized not in candidates: candidates.append(separator_normalized) - if "_" in table_name: - dotted_schema_name = re.sub( - r"^([A-Za-z_][A-Za-z0-9]*)_", - r"\1.", - table_name, - count=1, - ) - if dotted_schema_name not in candidates: - candidates.append(dotted_schema_name) short_name = re.split(r"[.$]", table_name)[-1] - if short_name == table_name and "_" in table_name: - short_name = table_name.split("_", 1)[-1] if short_name and short_name not in candidates: candidates.append(short_name) return candidates @@ -2015,127 +1877,6 @@ def _normalize_explicit_table_names( normalized.append(candidate) return normalized - def _select_direct_explicit_sql_table_name(self, table_names: list[str]) -> str | None: - normalized_names = self._normalize_explicit_table_names(table_names) - if not normalized_names: - return None - - original_name = normalized_names[0] - if "." in original_name: - for table_name in normalized_names: - if "." not in table_name and "_" in table_name: - return table_name - return original_name - - def _build_direct_explicit_table_count_sql( - self, query: str, table_names: list[str] - ) -> tuple[str, str] | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized or not table_names: - return None - - wants_total_count = ( - re.search(r"\bhow many\b", normalized) - or "record count" in normalized - or "count of records" in normalized - or "number of records" in normalized - or "number of rows" in normalized - or "count rows" in normalized - ) - if not wants_total_count or re.search( - r"\b(?:by|per|each|distribution|highest|top|monthly|daily|weekly)\b", - normalized, - ): - return None - - table_name = self._select_direct_explicit_sql_table_name(table_names) - if not table_name: - return None - - return ( - f'SELECT COUNT(*) AS "RecordCount" FROM ' - f"{self._quote_sql_identifier(table_name)}", - table_name, - ) - - def _extract_direct_explicit_projection_columns(self, query: str) -> list[str]: - columns: list[str] = [] - match = re.search( - r"\b(?:including|include|with)\s+(.+?)(?:\.|$)", - query or "", - flags=re.IGNORECASE, - ) - if not match: - return columns - - raw_columns = re.sub(r"\band\b", ",", match.group(1), flags=re.IGNORECASE) - stop_words = { - "columns", - "fields", - "including", - "include", - "with", - "and", - } - for raw_column in re.split(r",", raw_columns): - column = raw_column.strip().strip("`\"[] ") - if not column: - continue - if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", column): - continue - if column.lower() in stop_words: - continue - if column not in columns: - columns.append(column) - return columns - - def _build_direct_explicit_table_projection_sql( - self, query: str, table_names: list[str] - ) -> tuple[str, str] | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized or not table_names: - return None - - table_name = self._select_direct_explicit_sql_table_name(table_names) - if not table_name: - return None - - columns = self._extract_direct_explicit_projection_columns(query) - asks_for_preview = bool( - re.search(r"\b(?:preview|sample|first)\b", normalized) - or re.search(r"\b(?:rows?|records?|data)\b", normalized) - ) - if not columns and not asks_for_preview: - return None - - limit = self._extract_requested_top_n(query, default_value=500) - table_ref = self._quote_sql_identifier(table_name) - select_clause = ( - ", ".join(self._quote_sql_identifier(column) for column in columns) - if columns - else "*" - ) - - where_parts: list[str] = [] - for column in columns: - if re.search( - rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+empty", - normalized, - ): - column_ref = self._quote_sql_identifier(column) - where_parts.append(f"{column_ref} IS NOT NULL AND {column_ref} <> ''") - elif re.search( - rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+null", - normalized, - ): - where_parts.append(f"{self._quote_sql_identifier(column)} IS NOT NULL") - - where_clause = f" WHERE {' AND '.join(where_parts)}" if where_parts else "" - return ( - f"SELECT TOP {limit} {select_clause} FROM {table_ref}{where_clause}", - table_name, - ) - def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -2474,88 +2215,6 @@ def _select_best_analytics_table( date_column, ) - def _build_entity_lookup_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - lookup_phrase = self._extract_entity_lookup_phrase(query) - if not lookup_phrase: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - asks_for_orders = any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ) - - scored: list[tuple[int, dict[str, Any], str]] = [] - for table in tables: - table_name = str(table.get("name") or "") - if not table_name: - continue - - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - - preferred_column = sorted( - preferred_columns, - key=lambda column: ( - 0 - if self._normalize_schema_identifier_key(column) - in {"custname", "customername", "customer"} - else 1, - column.lower(), - ), - )[0] - - score = 20 - normalized_table = self._normalize_schema_token(table_name) - if asks_for_orders: - if "order" in normalized_table: - score += 40 - if "neworder" in normalized_table: - score += 20 - if self._find_schema_column( - table, ("OrdNo", "OrderNo", "OrderId", "NewOrderId") - ): - score += 25 - if "test" in normalized_table or "tmp" in normalized_table: - score -= 80 - if "dev" in normalized_table or "backup" in normalized_table: - score -= 60 - if "stage" in normalized_table: - score -= 10 - scored.append((score, table, preferred_column)) - - if not scored: - return None - - _score, table, filter_column = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - table_name = str(table.get("name") or "") - if not table_name: - return None - - table_ref = self._quote_sql_identifier(table_name) - filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" - escaped_phrase = lookup_phrase.replace("'", "''") - date_column = self._find_schema_column( - table, - ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), - temporal=True, - ) - order_clause = ( - f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" - if date_column - else "" - ) - return ( - f"SELECT TOP 500 * FROM {table_ref} " - f"WHERE {filter_ref} = '{escaped_phrase}'" - f"{order_clause}" - ) - def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2569,9 +2228,6 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): - return entity_lookup_sql - if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2627,7 +2283,16 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - asks_for_measure_value = any( + if not is_sales_or_order_query: + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + + wants_count_metric = any( + term in normalized_query + for term in ("count", "counts", "volume", "how many", "distribution") + ) and not any( term in normalized_query for term in ( "amount", @@ -2644,17 +2309,6 @@ def _build_schema_grounded_analytics_sql( "value", ) ) - - if not is_sales_or_order_query and not asks_for_measure_value: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not asks_for_measure_value wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") ) @@ -6205,54 +5859,6 @@ async def ask( return results if explicit_table_names: - if direct_count_sql := self._build_direct_explicit_table_count_sql( - user_query, - explicit_table_names, - ): - explicit_sql, explicit_table_name = direct_count_sql - ask_result = self._build_ask_result_from_sql(explicit_sql) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning=( - "Explicit table count request generated SQL directly." - ), - retrieved_tables=[explicit_table_name], - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - if direct_projection_sql := self._build_direct_explicit_table_projection_sql( - user_query, - explicit_table_names, - ): - explicit_sql, explicit_table_name = direct_projection_sql - ask_result = self._build_ask_result_from_sql(explicit_sql) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning=( - "Explicit table request generated SQL directly." - ), - retrieved_tables=[explicit_table_name], - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -6273,7 +5879,7 @@ async def ask( timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 60, + 20, ), ) documents, table_names, table_ddls = ( @@ -6869,8 +6475,7 @@ async def ask( ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 60, + 20, ), ) _retrieval_result = retrieval_result.get( @@ -6907,7 +6512,7 @@ async def ask( timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 45, + 20, ), ) _retrieval_result = retrieval_result.get( diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index b5d66044fc..d93e456d2a 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -31,8 +31,6 @@ DEFAULT_VALIDATION_SQL_FUNCTION_CHARS = 2500 STRICT_VALIDATION_SQL_FUNCTION_ITEMS = 0 STRICT_VALIDATION_SQL_FUNCTION_CHARS = 0 -DEFAULT_RECOMMENDATION_RETRIEVAL_TIMEOUT_SECONDS = 20 -DEFAULT_RECOMMENDATION_VALIDATION_TIMEOUT_SECONDS = 20 DEFAULT_QUESTION_CATEGORIES = [ "Descriptive Questions", "Segmentation Questions", @@ -175,42 +173,6 @@ def _get_underfilled_categories( if len(response_questions.get(category, [])) < max_questions ] - def _add_recommended_questions( - self, - event_id: str, - candidates: list[dict], - *, - max_questions: int, - max_categories: int, - ) -> None: - current = self._cache[event_id] - questions = current.response.setdefault("questions", {}) - - for candidate in candidates: - category = candidate.get("category") - question = candidate.get("question") - if not isinstance(category, str) or not isinstance(question, str): - continue - if category not in questions and len(questions) >= max_categories: - continue - - category_questions = questions.setdefault(category, []) - existing = next( - ( - item - for item in category_questions - if item.get("question") == question - ), - None, - ) - if existing is not None: - existing.update(candidate) - continue - - if len(category_questions) >= max_questions: - continue - category_questions.append(candidate) - def _handle_exception( self, event_id: str, @@ -388,11 +350,6 @@ async def _instructions_retrieval() -> list[dict]: currnet_category = questions.setdefault(candidate["category"], []) - for existing_question in currnet_category: - if existing_question.get("question") == candidate["question"]: - existing_question.update({**candidate, "sql": valid_sql}) - return post_process - if len(currnet_category) >= max_questions: # Skip to update the questions for the category if it is already full return post_process @@ -416,12 +373,6 @@ class Request(BaseRequest): async def _recommend(self, request: dict): resp = await self._pipelines["question_recommendation"].run(**request) questions = resp.get("normalized", {}).get("questions", []) - self._add_recommended_questions( - request["event_id"], - questions, - max_questions=request["max_questions"], - max_categories=request["max_categories"], - ) validation_tasks = [ self._validate_question( question, @@ -434,19 +385,7 @@ async def _recommend(self, request: dict): for question in questions ] - if not validation_tasks: - return - - try: - await asyncio.wait_for( - asyncio.gather(*validation_tasks, return_exceptions=True), - timeout=DEFAULT_RECOMMENDATION_VALIDATION_TIMEOUT_SECONDS, - ) - except TimeoutError: - logger.warning( - "Request %s: Question recommendation validation timed out; returning generated recommendations without full SQL validation", - request["event_id"], - ) + await asyncio.gather(*validation_tasks, return_exceptions=True) @observe(name="Generate Question Recommendation") @trace_metadata @@ -458,26 +397,14 @@ async def recommend(self, input: Request, **kwargs) -> Event: try: orjson.loads(input.mdl) - try: - retrieval_result = await asyncio.wait_for( - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=input.project_id, - enable_column_pruning=False, - ), - timeout=DEFAULT_RECOMMENDATION_RETRIEVAL_TIMEOUT_SECONDS, - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents = _retrieval_result.get("retrieval_results", []) - except TimeoutError: - logger.warning( - "Request %s: Question recommendation schema retrieval timed out; continuing with generated recommendations from limited context", - input.event_id, - ) - documents = [] + retrieval_result = await self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=input.project_id, + enable_column_pruning=False, + ) + _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) + documents = _retrieval_result.get("retrieval_results", []) table_ddls = self._limit_text_items( [document.get("table_ddl") for document in documents], max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index bb0d25f00b..3ee705b621 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,7 +2,6 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( - _extract_explicit_table_names_from_query, _is_project_wide_analysis_query, _rerank_table_documents, _select_relevant_table_documents, @@ -48,20 +47,6 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query -def test_extract_explicit_table_names_from_query_handles_schema_dot_reference(): - table_names = _extract_explicit_table_names_from_query( - "How many records are in dbo.xStageLoad8?" - ) - - assert table_names[:3] == ["dbo.xStageLoad8", "dbo_xStageLoad8", "xStageLoad8"] - - -def test_extract_explicit_table_names_from_query_ignores_plain_business_terms(): - assert _extract_explicit_table_names_from_query( - "Compare sales in countries by month" - ) == [] - - def test_rerank_table_documents_prefers_question_relevant_table_text(): generic_stage = Document( content="Generic imported staging records with product labels.", @@ -82,26 +67,6 @@ def test_rerank_table_documents_prefers_question_relevant_table_text(): assert documents[0].meta["name"] == "business_transactions" -def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): - test_load = Document( - content="Raw test load rows for order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ) - order_market_table = Document( - content="New order transaction records with market and customer fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.45, - ) - - documents = _rerank_table_documents( - "Show order distribution across markets.", - [test_load, order_market_table], - ) - - assert documents[0].meta["name"] == "dbo_xStageNewOrders" - - def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( @@ -238,44 +203,33 @@ def test_rerank_table_documents_prefers_transaction_source_for_metric_question() @pytest.mark.asyncio -async def test_table_retrieval_uses_explicit_table_names_without_vector_lookup(): +async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: + def __init__(self): + self.filters = None + async def run(self, query_embedding, filters): - raise AssertionError("explicit table names should not use vector lookup") + self.filters = filters + return {"documents": []} - result = await table_retrieval( + retriever = Retriever() + + await table_retrieval( query="show rows", embedding={}, project_id="project-1", - tables=["orders", "orders", "customers"], - table_retriever=Retriever(), - ) - - assert [document.meta["name"] for document in result["documents"]] == [ - "orders", - "customers", - ] - - -@pytest.mark.asyncio -async def test_table_retrieval_uses_query_table_names_before_embedding_lookup(): - class Retriever: - async def run(self, query_embedding, filters): - raise AssertionError("explicit query table name should not use vector lookup") - - result = await table_retrieval( - query="How many records are in dbo.xStageLoad8?", - embedding={"embedding": [0.1, 0.2]}, - project_id="project-1", - tables=[], - table_retriever=Retriever(), + tables=["orders"], + table_retriever=retriever, ) - assert [document.meta["name"] for document in result["documents"]] == [ - "dbo.xStageLoad8", - "dbo_xStageLoad8", - "xStageLoad8", - ] + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "name", "operator": "in", "value": ["orders"]}, + ], + } @pytest.mark.asyncio @@ -361,188 +315,6 @@ async def run(self, query_embedding, filters): assert not retriever.called -@pytest.mark.asyncio -async def test_dbschema_retrieval_uses_query_table_names_as_scope(): - class Retriever: - def __init__(self): - self.filters = None - - async def run(self, query_embedding, filters): - self.filters = filters - return {"documents": []} - - retriever = Retriever() - - await dbschema_retrieval( - query="How many records are in dbo.xStageLoad8?", - table_retrieval={"documents": []}, - project_id="project-1", - dbschema_retriever=retriever, - ) - - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - { - "field": "name", - "operator": "in", - "value": ["dbo.xStageLoad8", "dbo_xStageLoad8", "xStageLoad8"], - }, - ], - } - - -@pytest.mark.asyncio -async def test_dbschema_retrieval_falls_back_to_deployed_schema_vector_search(): - class Retriever: - def __init__(self): - self.calls = [] - - async def run(self, query_embedding, filters): - self.calls.append((query_embedding, filters)) - if query_embedding: - return { - "documents": [ - Document( - content=str( - { - "type": "TABLE", - "name": "sales_orders", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "sales_orders"}, - score=0.9, - ) - ] - } - return { - "documents": [ - Document( - content=str( - { - "type": "TABLE", - "name": "sales_orders", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "sales_orders"}, - ), - Document( - content=str( - { - "type": "TABLE_COLUMNS", - "name": "sales_orders", - "columns": [{"name": "amount"}], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "sales_orders"}, - ), - ] - } - - retriever = Retriever() - - documents = await dbschema_retrieval( - query="compare sales between countries", - table_retrieval={"documents": []}, - project_id="project-1", - dbschema_retriever=retriever, - embedding={"embedding": [0.1, 0.2]}, - ) - - assert [document.meta["name"] for document in documents] == [ - "sales_orders", - "sales_orders", - ] - assert retriever.calls[0][0] == [0.1, 0.2] - assert retriever.calls[0][1] == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - ], - } - assert retriever.calls[1][0] == [] - assert retriever.calls[1][1] == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - {"field": "name", "operator": "in", "value": ["sales_orders"]}, - ], - } - - -@pytest.mark.asyncio -async def test_dbschema_retrieval_falls_back_when_selected_schema_lookup_is_empty(): - class Retriever: - def __init__(self): - self.calls = [] - - async def run(self, query_embedding, filters): - self.calls.append((query_embedding, filters)) - if query_embedding: - return { - "documents": [ - Document( - content=str( - { - "type": "TABLE", - "name": "dbo_xStageLoad8", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "dbo_xStageLoad8"}, - ) - ] - } - if len(self.calls) == 1: - return {"documents": []} - return { - "documents": [ - Document( - content=str( - { - "type": "TABLE", - "name": "dbo_xStageLoad8", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "dbo_xStageLoad8"}, - ) - ] - } - - retriever = Retriever() - - documents = await dbschema_retrieval( - query="How many records are in dbo.xStageLoad8?", - table_retrieval={ - "documents": [ - Document( - content=str({"name": "dbo.xStageLoad8"}), - meta={"type": "TABLE_DESCRIPTION", "name": "dbo.xStageLoad8"}, - ) - ] - }, - project_id="project-1", - dbschema_retriever=retriever, - embedding={"embedding": [0.1, 0.2]}, - ) - - assert [document.meta["name"] for document in documents] == ["dbo_xStageLoad8"] - assert retriever.calls[0][0] == [] - assert retriever.calls[1][0] == [0.1, 0.2] - assert retriever.calls[2][1]["conditions"][-1] == { - "field": "name", - "operator": "in", - "value": ["dbo_xStageLoad8"], - } - - def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): class Encoding: def encode(self, value): diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 5ba8ca74b8..a64038fbb9 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,115 +269,6 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) -def test_validated_sql_rejects_country_question_without_country_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' - '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' - 'ORDER BY COUNT(*) DESC' - ), - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - "Show the total commodity line value by country.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Show the total commodity line value by country.", - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'AS "Country_of_Ultimate_Destination_Code", ' - 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' - 'AS "TotalCommodity_Line_Value" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' - ) - - -def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT * FROM "dbo_tnoStageNewOrders" ' - 'WHERE "dbo_tnoStageNewOrders"."Division" = ' - "'Daimler Trucks North America'" - ), - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - "List orders for Daimler Trucks North America.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "List orders for Daimler Trucks North America.", - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """, - ], - ) - - assert sql == ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" = ' - "'Daimler Trucks North America' " - 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' - ) - - def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) @@ -555,21 +446,6 @@ def test_extract_explicit_table_names_from_pcb_repair_phrases(): ) == ["ticket_labels", "dbo_ticket_labels"] -def test_explicit_table_name_candidates_include_dotted_and_short_forms(): - service = AskService.__new__(AskService) - - assert service._explicit_table_name_candidates("dbo_tblNewOrders") == [ - "dbo_tblNewOrders", - "dbo.tblNewOrders", - "tblNewOrders", - ] - assert service._explicit_table_name_candidates("dbo.tblNewOrders") == [ - "dbo.tblNewOrders", - "dbo_tblNewOrders", - "tblNewOrders", - ] - - def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): service = AskService.__new__(AskService) documents = [ @@ -606,42 +482,6 @@ def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): assert table_ddls == [documents[1]["table_ddl"]] -def test_filter_retrieval_metadata_for_explicit_query_matches_dotted_table_name(): - service = AskService.__new__(AskService) - documents = [ - { - "table_name": "dbo.tblNewOrders", - "table_ddl": """ - CREATE TABLE "dbo.tblNewOrders" ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - }, - { - "table_name": "dbo_other", - "table_ddl": """ - CREATE TABLE dbo_other ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - }, - ] - - filtered_documents, table_names, table_ddls = ( - service._filter_retrieval_metadata_for_explicit_query( - "Show the top 5 CustName values from dbo_tblNewOrders by number of orders.", - documents, - ["dbo_tblNewOrders"], - ) - ) - - assert filtered_documents == [documents[0]] - assert table_names == ["dbo.tblNewOrders"] - assert table_ddls == [documents[0]["table_ddl"]] - - def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): service = AskService.__new__(AskService) result = service._build_validated_ask_result_from_sql( From ab9bb6861020317b78d3c70290f9f7d04efe57d3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 19:19:48 +0530 Subject: [PATCH 0538/1087] Validate explicit table SQL against metadata --- wren-ai-service/src/web/v1/services/ask.py | 40 ++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0223f584e2..b4ab872a91 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -5907,6 +5907,46 @@ async def ask( table_names, ) + for direct_sql_candidate in ( + self._build_direct_explicit_table_count_sql( + user_query, + table_names or explicit_table_names, + ), + self._build_direct_explicit_table_projection_sql( + user_query, + table_names or explicit_table_names, + ), + ): + if not direct_sql_candidate: + continue + explicit_sql, explicit_table_name = direct_sql_candidate + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, + ) + if ask_result: + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning=( + "Explicit table request matched deployed schema " + "and generated SQL directly." + ), + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = explicit_sql + if table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ): From 82be43ffa0dd4d1f7b0c9871ec81161b8833d97c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 19:54:43 +0530 Subject: [PATCH 0539/1087] Restore explicit table SQL helpers --- wren-ai-service/src/web/v1/services/ask.py | 123 +++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b4ab872a91..3aea12309f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1877,6 +1877,129 @@ def _normalize_explicit_table_names( normalized.append(candidate) return normalized + def _select_direct_explicit_sql_table_name( + self, table_names: list[str] + ) -> str | None: + normalized_names = self._normalize_explicit_table_names(table_names) + if not normalized_names: + return None + + original_name = normalized_names[0] + if "." in original_name: + for table_name in normalized_names: + if "." not in table_name and "_" in table_name: + return table_name + return original_name + + def _build_direct_explicit_table_count_sql( + self, query: str, table_names: list[str] + ) -> tuple[str, str] | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized or not table_names: + return None + + wants_total_count = ( + re.search(r"\bhow many\b", normalized) + or "record count" in normalized + or "count of records" in normalized + or "number of records" in normalized + or "number of rows" in normalized + or "count rows" in normalized + ) + if not wants_total_count or re.search( + r"\b(?:by|per|each|distribution|highest|top|monthly|daily|weekly)\b", + normalized, + ): + return None + + table_name = self._select_direct_explicit_sql_table_name(table_names) + if not table_name: + return None + + return ( + f'SELECT COUNT(*) AS "RecordCount" FROM ' + f"{self._quote_sql_identifier(table_name)}", + table_name, + ) + + def _extract_direct_explicit_projection_columns(self, query: str) -> list[str]: + columns: list[str] = [] + match = re.search( + r"\b(?:including|include|with)\s+(.+?)(?:\.|$)", + query or "", + flags=re.IGNORECASE, + ) + if not match: + return columns + + raw_columns = re.sub(r"\band\b", ",", match.group(1), flags=re.IGNORECASE) + stop_words = { + "columns", + "fields", + "including", + "include", + "with", + "and", + } + for raw_column in re.split(r",", raw_columns): + column = raw_column.strip().strip("`\"[] ") + if not column: + continue + if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", column): + continue + if column.lower() in stop_words: + continue + if column not in columns: + columns.append(column) + return columns + + def _build_direct_explicit_table_projection_sql( + self, query: str, table_names: list[str] + ) -> tuple[str, str] | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized or not table_names: + return None + + table_name = self._select_direct_explicit_sql_table_name(table_names) + if not table_name: + return None + + columns = self._extract_direct_explicit_projection_columns(query) + asks_for_preview = bool( + re.search(r"\b(?:preview|sample|first)\b", normalized) + or re.search(r"\b(?:rows?|records?|data)\b", normalized) + ) + if not columns and not asks_for_preview: + return None + + limit = self._extract_requested_top_n(query, default_value=500) + table_ref = self._quote_sql_identifier(table_name) + select_clause = ( + ", ".join(self._quote_sql_identifier(column) for column in columns) + if columns + else "*" + ) + + where_parts: list[str] = [] + for column in columns: + if re.search( + rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+empty", + normalized, + ): + column_ref = self._quote_sql_identifier(column) + where_parts.append(f"{column_ref} IS NOT NULL AND {column_ref} <> ''") + elif re.search( + rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+null", + normalized, + ): + where_parts.append(f"{self._quote_sql_identifier(column)} IS NOT NULL") + + where_clause = f" WHERE {' AND '.join(where_parts)}" if where_parts else "" + return ( + f"SELECT TOP {limit} {select_clause} FROM {table_ref}{where_clause}", + table_name, + ) + def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: From afff6c87d96ab9f8601c8cfe6942081a71b20e6b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 20:01:55 +0530 Subject: [PATCH 0540/1087] Revert last two explicit table changes --- wren-ai-service/src/web/v1/services/ask.py | 163 --------------------- 1 file changed, 163 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3aea12309f..0223f584e2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1877,129 +1877,6 @@ def _normalize_explicit_table_names( normalized.append(candidate) return normalized - def _select_direct_explicit_sql_table_name( - self, table_names: list[str] - ) -> str | None: - normalized_names = self._normalize_explicit_table_names(table_names) - if not normalized_names: - return None - - original_name = normalized_names[0] - if "." in original_name: - for table_name in normalized_names: - if "." not in table_name and "_" in table_name: - return table_name - return original_name - - def _build_direct_explicit_table_count_sql( - self, query: str, table_names: list[str] - ) -> tuple[str, str] | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized or not table_names: - return None - - wants_total_count = ( - re.search(r"\bhow many\b", normalized) - or "record count" in normalized - or "count of records" in normalized - or "number of records" in normalized - or "number of rows" in normalized - or "count rows" in normalized - ) - if not wants_total_count or re.search( - r"\b(?:by|per|each|distribution|highest|top|monthly|daily|weekly)\b", - normalized, - ): - return None - - table_name = self._select_direct_explicit_sql_table_name(table_names) - if not table_name: - return None - - return ( - f'SELECT COUNT(*) AS "RecordCount" FROM ' - f"{self._quote_sql_identifier(table_name)}", - table_name, - ) - - def _extract_direct_explicit_projection_columns(self, query: str) -> list[str]: - columns: list[str] = [] - match = re.search( - r"\b(?:including|include|with)\s+(.+?)(?:\.|$)", - query or "", - flags=re.IGNORECASE, - ) - if not match: - return columns - - raw_columns = re.sub(r"\band\b", ",", match.group(1), flags=re.IGNORECASE) - stop_words = { - "columns", - "fields", - "including", - "include", - "with", - "and", - } - for raw_column in re.split(r",", raw_columns): - column = raw_column.strip().strip("`\"[] ") - if not column: - continue - if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", column): - continue - if column.lower() in stop_words: - continue - if column not in columns: - columns.append(column) - return columns - - def _build_direct_explicit_table_projection_sql( - self, query: str, table_names: list[str] - ) -> tuple[str, str] | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized or not table_names: - return None - - table_name = self._select_direct_explicit_sql_table_name(table_names) - if not table_name: - return None - - columns = self._extract_direct_explicit_projection_columns(query) - asks_for_preview = bool( - re.search(r"\b(?:preview|sample|first)\b", normalized) - or re.search(r"\b(?:rows?|records?|data)\b", normalized) - ) - if not columns and not asks_for_preview: - return None - - limit = self._extract_requested_top_n(query, default_value=500) - table_ref = self._quote_sql_identifier(table_name) - select_clause = ( - ", ".join(self._quote_sql_identifier(column) for column in columns) - if columns - else "*" - ) - - where_parts: list[str] = [] - for column in columns: - if re.search( - rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+empty", - normalized, - ): - column_ref = self._quote_sql_identifier(column) - where_parts.append(f"{column_ref} IS NOT NULL AND {column_ref} <> ''") - elif re.search( - rf"\b{re.escape(column.lower())}\b\s+is\s+not\s+null", - normalized, - ): - where_parts.append(f"{self._quote_sql_identifier(column)} IS NOT NULL") - - where_clause = f" WHERE {' AND '.join(where_parts)}" if where_parts else "" - return ( - f"SELECT TOP {limit} {select_clause} FROM {table_ref}{where_clause}", - table_name, - ) - def _build_direct_orders_sales_sql(self, query: str) -> str | None: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: @@ -6030,46 +5907,6 @@ async def ask( table_names, ) - for direct_sql_candidate in ( - self._build_direct_explicit_table_count_sql( - user_query, - table_names or explicit_table_names, - ), - self._build_direct_explicit_table_projection_sql( - user_query, - table_names or explicit_table_names, - ), - ): - if not direct_sql_candidate: - continue - explicit_sql, explicit_table_name = direct_sql_candidate - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning=( - "Explicit table request matched deployed schema " - "and generated SQL directly." - ), - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = explicit_sql - if table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ): From ef584f1fb22af4ccd6144df9d81c6f6044ecab41 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 14 Jul 2026 21:14:14 +0530 Subject: [PATCH 0541/1087] Create WrenAI Fine-Tuning & Integration Playbook Added a comprehensive playbook for fine-tuning and integrating WrenAI within the Orders/Sales domain, detailing steps for data hygiene, semantics, instructions, and question-SQL pairs. --- WrenAI-FineTuning-Playbook.md | 494 ++++++++++++++++++++++++++++++++++ 1 file changed, 494 insertions(+) create mode 100644 WrenAI-FineTuning-Playbook.md diff --git a/WrenAI-FineTuning-Playbook.md b/WrenAI-FineTuning-Playbook.md new file mode 100644 index 0000000000..254c043189 --- /dev/null +++ b/WrenAI-FineTuning-Playbook.md @@ -0,0 +1,494 @@ +WrenAI Fine-Tuning & Integration Playbook — NetRay / Orders +Instance: http://10.104.74.13:3000 (self-hosted OSS) Scope: Orders/Sales domain + a separate Tariffs/Customs domain +Important framing: WrenAI is not model fine-tuning. There is no training run. You are doing context engineering — pushing your business knowledge into three layers Wren injects into the LLM prompt at query time (via vector retrieval). Do these in order; each layer compounds on the one before it. + +How to use this doc: Anything in is a placeholder your team must replace with real schema/business facts. The defaults are SyteLine-flavoured starting points — keep, edit, or delete them. Sections 4 and 5 are meant to be pasted directly into Wren's Knowledge tab. + + +0. Order of operations (the whole playbook in 8 steps) +Prune staging/junk tables out of the model. +Gold schema — point Wren only at clean views (best) or curated base tables. +Generate semantics (Modeling AI Assistant) → auto-fill descriptions. +Hand-correct the cryptic/ambiguous columns only. +Relationships — generate + verify join keys. +Instructions — paste Section 4 into Knowledge, edit placeholders. +Question-SQL pairs — paste Section 5, verify each runs, save. +Deploy, smoke-test, then wire the API into SyteRay (Section 6). + +Everything before "Deploy" is invisible to users until you hit Deploy. Deploy re-embeds the context into the vector store — it is not optional and it is the step people forget. + + +1. The three context layers (what actually moves accuracy) +Layer +Where in UI +What it holds +Fixes +Semantics (MDL) +Modeling tab +Table + column descriptions, types, relationships, calculated fields +"AI doesn't know what OrdNo / col_01_Division means" +Instructions +Knowledge tab +Reusable rules: terminology, filters, formatting, join rules +"AI counts cancelled orders in revenue", "money not rounded", inconsistent metric logic +Question-SQL pairs +Knowledge tab +Gold examples pinning a question → exact SQL +Complex/error-prone recurring questions (backlog cost, YoY growth) + + +Rule of thumb: facts about columns → Semantics. Rules about logic → Instructions. Whole gold answers → Q-SQL pairs. + + +2. Phase 0 — Data hygiene (do this first, it's 80% of the win) +Your model list is mostly staging junk: xStage, xStageLoad, xStageLoad2/3/4/5, xStageLoad8, xStageLoad8_Test, xStageNewOrders. Every one is a table the AI can wrongly select — that's exactly why it was recommending questions about dbo.xStageLoad8. + +2.1 Prune. In Modeling, remove every staging/load/test table from the model. Keep only business-meaningful tables. + +2.2 Build a gold layer (strongly recommended). In SQL Server, create a dedicated schema of clean, analytics-friendly views and point Wren only at those: + +CREATE SCHEMA gold; + +GO + +-- One row per order line, business-named columns, junk excluded + +CREATE VIEW gold.v_orders AS + +SELECT + + OrdNo AS order_number, + + CustNo AS customer_number, + + CustName AS customer_name, + + CustPO AS customer_po, + + Market AS market, + + AS division, + + AS segment, + + AS salesperson, + + AS product, + + AS order_status, + + AS order_date, + + AS invoice_date, + + AS qty_ordered, + + AS line_amount, -- pre-summed numeric measure + + AS is_backlog -- boolean, see §4 + +FROM dbo. + +WHERE ; + +Why gold views beat raw tables: + +The LLM reads clean names (order_status, not col_07_stat) → fewer wrong guesses. +You bake business logic (status filters, boolean flags, fiscal columns) into SQL once, so the AI doesn't reinvent it every query. +You control exactly what's exposed — no PII, no staging tables. + +2.3 Read-only DB user. Wren should connect via a wren_ro login with SELECT-only on gold (and nothing on staging). Never a write-capable account. + +2.4 Split domains into separate projects. dbo.ytblTarrifsFullA (Importer_Number, HTS, Entry_Date, Legal_Entity) is customs/tariff data — a different subject from sales. Mixing it into the Orders graph makes the AI blend customs columns into sales answers. Create a separate Wren project for Tariffs. Each project stays coherent and "opinionated," which is what reduces hallucination. + + +3. Phase 1 — Semantics (Modeling) +3.1 Auto-generate. Modeling page → Modeling AI Assistant (top-right) → Generate semantics. This fills the model + column Description fields across all tables from the schema. Then Generate relationships. + +3.2 Hand-correct only what the AI can't infer. The assistant handles obvious columns. You manually fix: + +Cryptic codes: col_01_Division, FY___Would_invoice_date, status/segment codes. +Ambiguous pairs: if two columns could both be "revenue" or "date," describe each precisely and say when each is used. +Measures vs dimensions vs IDs — phrasing matters (see table below). + +3.3 Description conventions (this is your data dictionary; phrasing drives behaviour): + +Column role +Write the description as… +Example +Measure (sum/avg) +"Total/Amount of … used for …" +line_amount → "Extended line amount in USD; sum for sales revenue." +Dimension (group/filter) +"Category of …" / "… segment, not geography" +market → "Business unit / market segment. NOT a country or region." +Date +"Date used for … filtering" +invoice_date → "Date the order line was invoiced; default date for revenue-by-period." +ID / key +"Unique identifier for …" +customer_number → "SyteLine customer code; use COUNT(DISTINCT) for customer counts." + + +Table-level description example for gold.v_orders: + +"One row per sales order line from SyteLine. Grain = order line. Use for sales, revenue, backlog, orders-by-market/division/salesperson analysis. Excludes voided and test orders." + +3.4 Relationships. Verify the auto-detected joins and add any missed keys explicitly (e.g. v_orders.customer_number → v_customers.customer_number, many-to-one). Explicit relationships = deterministic joins; without them the AI guesses. + +3.5 Deploy. + + +4. Phase 2 — INSTRUCTIONS (paste into Knowledge → Instructions) +Two types: Global (always applied) and Question-Matching (applied only when the question matches a topic/keyword). Add each block below as a separate instruction of the stated type. Edit every before deploying — a wrong rule is worse than no rule. +4A. Global instructions (always on) +[GLOBAL] Currency & rounding + +All monetary values are in USD. Always ROUND(value, 2) for revenue, averages, + +and percentages. Format large money values with thousands separators in summaries. + +[GLOBAL] Valid orders only + +Exclude orders where order_status IN () + +from any sales, revenue, backlog, or count calculation, unless the user explicitly + +asks about cancelled/quoted orders. + +[GLOBAL] Default date field + +For any sales/revenue question by time period, use invoice_date as the default + +date field. For "new orders" or "orders received," use order_date instead. + +[GLOBAL] Default time range + +If the user gives no date range, default to the last 90 days. Always state the + +range you assumed in the answer summary. + +[GLOBAL] Counting entities + +Headcount-style counts must use COUNT(DISTINCT ...), not COUNT(*). + +Customers = COUNT(DISTINCT customer_number). Orders = COUNT(DISTINCT order_number). + +[GLOBAL] Safe joins + +Use LEFT JOIN when joining optional/reference tables (products, salesperson, + +customer master) so order rows are never dropped when a lookup is missing. + +[GLOBAL] Grain awareness + +gold.v_orders is at ORDER-LINE grain. When counting or summing at the order level, + +aggregate to order_number first to avoid double counting. + +[GLOBAL] Fiscal calendar + +Our fiscal year runs . "FY", "quarter", "YTD", and + +"MTD" all refer to the fiscal calendar, not the calendar year. . +4B. Terminology instructions (Global — your business dictionary as rules) +[GLOBAL] Term: "Backlog" + +"Backlog" = open order lines not yet shipped/invoiced, i.e. + +is_backlog = 1 (or order_status = AND invoice_date IS NULL). + +"Backlog cost" / "total cost of backlog" = SUM(line_amount) over backlog lines. + +[GLOBAL] Term: "Market" vs "Division" vs "Segment" + +- market = business unit / market segment (e.g. "Honeywell BTP"). NOT geography. + +- division = . + +- segment = . + +When a user says "market," never map it to a country/region column. + +[GLOBAL] Term: "New orders" + +"New orders" = orders where order_date falls in the requested period, regardless + +of invoice status. Distinct from "sales/revenue" which uses invoice_date. + +[GLOBAL] Term: "Growing / declining market" + +Growth = period-over-period change in SUM(line_amount) by market. Default + +comparison is . Always show both periods and the % change. +4C. Question-Matching instructions (topic-scoped) +[MATCH: "year over year", "YoY", "vs last year", "growth"] + +Compute YoY as: current-period SUM(line_amount) vs the same period one fiscal + +year earlier, grouped by the requested dimension. Return both values and + +ROUND(((current-prior)/NULLIF(prior,0))*100, 2) AS pct_change. + +[MATCH: "salesperson", "sales rep", "who sold"] + +Attribute revenue via . Exclude house/unassigned + +accounts () unless explicitly asked. + +[MATCH: "underperforming", "underperform", "lagging"] + +"Underperforming" business units = those below for the period. State the benchmark used. + +[MATCH: chart / trend / over time] + +For time-series, order the x-axis chronologically and use a line chart. For + +"by market/division/product" rankings, use a horizontal bar chart sorted desc. + +Scoping discipline: Keep Globals few and universally true. Anything that only applies to one kind of question belongs in a Question-Matching instruction, or the AI over-applies it. + + +5. Phase 3 — QUESTION-SQL PAIRS (paste into Knowledge → Question-SQL Pairs) +These are drawn from the real questions already in your thread panel. Each SQL below is a template — replace columns to match your gold views, run it in Wren once, confirm the result, then Save. A pair with wrong SQL trains the AI wrongly, so verify before saving. + +Q1 — "Show total sales by market" + +SELECT market, + + ROUND(SUM(line_amount), 2) AS total_sales + +FROM gold.v_orders + +WHERE order_status NOT IN () + +GROUP BY market + +ORDER BY total_sales DESC; + +Q2 — "Which division has the highest sales?" + +SELECT TOP 1 division, + + ROUND(SUM(line_amount), 2) AS total_sales + +FROM gold.v_orders + +WHERE order_status NOT IN () + +GROUP BY division + +ORDER BY total_sales DESC; + +Q3 — "What is the total cost of all backlog?" + +SELECT ROUND(SUM(line_amount), 2) AS backlog_value + +FROM gold.v_orders + +WHERE is_backlog = 1; + +Q4 — "Break down month-to-date new orders" + +SELECT market, + + COUNT(DISTINCT order_number) AS new_orders, + + ROUND(SUM(line_amount), 2) AS order_value + +FROM gold.v_orders + +WHERE order_date >= DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1) + + AND order_date < DATEADD(DAY, 1, CAST(GETDATE() AS date)) + +GROUP BY market + +ORDER BY order_value DESC; + +Q5 — "Which markets are growing (YoY)?" + +WITH cur AS ( + + SELECT market, SUM(line_amount) AS amt + + FROM gold.v_orders + + WHERE invoice_date >= + + AND invoice_date < + + GROUP BY market), + +prior AS ( + + SELECT market, SUM(line_amount) AS amt + + FROM gold.v_orders + + WHERE invoice_date >= + + AND invoice_date < + + GROUP BY market) + +SELECT c.market, + + ROUND(c.amt,2) AS current_sales, + + ROUND(p.amt,2) AS prior_sales, + + ROUND(((c.amt - p.amt)/NULLIF(p.amt,0))*100, 2) AS pct_change + +FROM cur c LEFT JOIN prior p ON c.market = p.market + +ORDER BY pct_change DESC; + +Q6 — "Show new orders by CustName" + +SELECT customer_name, + + COUNT(DISTINCT order_number) AS orders, + + ROUND(SUM(line_amount), 2) AS order_value + +FROM gold.v_orders + +WHERE order_date >= + +GROUP BY customer_name + +ORDER BY order_value DESC; + +Q7 — "Which salesperson generated the most revenue?" + +SELECT TOP 10 salesperson, + + ROUND(SUM(line_amount), 2) AS revenue + +FROM gold.v_orders + +WHERE order_status NOT IN () + + AND salesperson <> + +GROUP BY salesperson + +ORDER BY revenue DESC; + +Add 3–5 more from your panel ("Which products contributed most," "Which business units are underperforming," "Which customers increased sales") the same way once the gold views are final. + + +6. Phase 4 — Integration (wiring Wren into SyteRay / your stack) +6.1 Get an API key +Wren UI → API tab → generate a key. All REST calls use header Authorization: Bearer . Base URL (self-hosted): http://10.104.74.21:3000/api/v1 + +Tier note: the REST Embedded AI API (generate_sql, generate_chart, streaming) is a governed/Agentic feature. If your OSS build's API tab exposes keys and these endpoints, use them (below). If not, the always-available OSS path is the GraphQL createAskingTask mutation (6.4). Check your API tab first. +6.2 Core REST endpoints +Generate SQL from a question: + +curl -X POST 'http://10.104.74.21:3000/api/v1/generate_sql' \ + + -H 'Authorization: Bearer ' \ + + -H 'Content-Type: application/json' \ + + -d '{ "projectId": , "question": "Show total sales by market" }' + +# → { "sql": "SELECT ...", "threadId": "..." } + +Other endpoints under the same base: + +POST /generate_chart — returns a Vega-Lite chart spec from a result set. +Streaming (SSE) — real-time token/step feedback for a chat UX. +Metadata introspection — list deployed models, columns, relationships, views (useful to render a schema picker in SyteRay). +Knowledge — read/manage instructions & Q-SQL pairs programmatically (supported tiers). + +Note: there are no webhooks — the client long-polls or uses SSE for async results. +6.3 Recommended SyteRay integration pattern +User asks a question in your UI. +SyteRay → generate_sql → gets governed SQL (which already respects your instructions + semantics). +Execute against the gold schema with the read-only user (either let Wren run it, or run it yourself for tighter control). +Optionally generate_chart for the visual. +Log threadId for audit; feed thumbs-up answers back as new Q-SQL pairs (closes the learning loop). + +This fits SyteRay cleanly: Wren becomes the text-to-SQL + governance layer; your policy engine / RBAC / audit trail wrap around it. Keep Wren's DB user scoped to gold so no agent can touch raw ERP tables. +6.4 OSS fallback (GraphQL asking task) +If REST embedded endpoints aren't in your build: submit questions via the createAskingTask GraphQL mutation and poll the task/thread for the answer. (Inspect the browser Network tab on the Home "Ask" flow to see the exact mutation shape your version uses.) +6.5 MCP (optional, for agent access) +Wren exposes a Model Context Protocol interface so agents (e.g. Claude, ChatGPT) query through your semantic layer instead of raw tables. On self-hosted, this runs via the Wren engine's MCP server. Useful if you want SyteRay's own agents to consult Wren as a governed data tool. +6.6 LLM configuration (self-hosted, LLM-agnostic) +Wren is LLM-agnostic. To keep ERP data on-prem, point it at a local model via LiteLLM/Ollama in ~/.wrenai/config.yaml: + +type: llm + +provider: litellm_llm + +models: + + - api_base: http://host.docker.internal:11434/v1 + + model: ollama_chat/ + + timeout: 600 + + kwargs: + + n: 1 + + temperature: 0 + +temperature: 0 for deterministic SQL. For accuracy, prefer a strong model (GPT-4o / o-series or a 70B-class local model); small models produce shakier SQL on messy ERP schemas. + + +7. Phase 5 — Maintenance loop (keep it accurate) +Schema change detection. Wren flags when tables/columns are added, renamed, removed, or retyped. Review after every ERP/gold-view change — renamed columns silently break Q-SQL pairs. +Smoke test. Keep a fixed list of your top ~10 questions. Re-run after any Deploy. If one regresses, fix the layer responsible (semantics vs instruction vs pair) — don't patch prompts ad hoc. +Feedback loop. Each week: take real user questions → if the SQL was right, Save as a Q-SQL pair; if it was almost right, add/refine an Instruction; if it picked the wrong table/column, fix the Semantic description. This is the "training set without training." +Version control. In newer Wren, this context lives in Git-friendly instructions.md + queries.yml. Even on your build, keep this document in Git as the source of truth and re-apply on rebuilds. + + +8. Rollout to the team +Owners + +__ owns Semantics (Modeling + gold views). +__ owns Instructions + Q-SQL pairs (Knowledge). +Weekly 30-min review: smoke test + feedback-loop triage. + +Guidance for people asking questions + +Use business terms from the dictionary ("backlog," "market," "new orders") — they're now defined for the AI. +Always read the generated SQL before trusting a number. Wrong-but-confident is the failure mode. +If an answer is wrong, don't just rephrase — report it so an owner fixes the underlying layer. +Sales questions → Orders project. Customs/HTS questions → Tariffs project. Don't cross them. + + +Appendix A — Data dictionary CSV template +Fill one row per exposed column; use it to drive/QA the Modeling descriptions. + +model,column,display_name,description,role,notes + +v_orders,order_number,Order Number,"SyteLine sales order number",id,"COUNT(DISTINCT) for order counts" + +v_orders,customer_number,Customer Number,"SyteLine customer code",id,"FK to v_customers" + +v_orders,market,Market,"Business unit / market segment; NOT geography",dimension,"" + +v_orders,line_amount,Line Amount,"Extended line amount USD; sum for revenue",measure,"" + +v_orders,invoice_date,Invoice Date,"Date invoiced; default date for revenue-by-period",date,"" + +v_orders,order_date,Order Date,"Date order received; use for new orders",date,"" + +v_orders,is_backlog,Is Backlog,"1 = open unshipped line",dimension,"boolean flag" +Appendix B — Instruction scoping cheat-sheet +Universally true, every query → Global. +True only for one topic/keyword → Question-Matching. +A whole correct answer to a recurring question → Question-SQL pair. +A fact about what a column is → Semantic description, not an instruction. + + From a3126e1bc8cd84362ffc7e7d6876c013e1b70b2a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 14 Jul 2026 16:20:11 +0530 Subject: [PATCH 0542/1087] Revert changes after 20e44de --- .../retrieval/db_schema_retrieval.py | 33 +- wren-ai-service/src/web/v1/services/ask.py | 374 +++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 20 + .../pytest/services/test_ask_sales_sql.py | 131 +++++- 4 files changed, 507 insertions(+), 51 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..6b64c9b452 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -213,11 +213,15 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms = { - _normalize_retrieval_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") - if len(token) > 2 and token.lower() not in stop_words - } + terms: set[str] = set() + for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): + split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) + for token in re.findall(r"[A-Za-z0-9]+", split_token): + if len(token) <= 2 or token.lower() in stop_words: + continue + normalized_token = _normalize_retrieval_token(token) + if normalized_token: + terms.add(normalized_token) return {term for term in terms if term} @@ -243,7 +247,11 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - non_production_terms = ( + weak_non_production_terms = ( + "stage", + "staging", + ) + strong_non_production_terms = ( "archive", "backup", "copy", @@ -251,17 +259,18 @@ def _source_shape_score(query: str, document: Document) -> int: "development", "duplicate", "sample", - "stage", - "staging", "temp", "test", "tmp", ) - if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( - normalized_query, - non_production_terms, + if source_terms & set(strong_non_production_terms) and not _query_mentions_any( + normalized_query, strong_non_production_terms + ): + score -= 240 + if source_terms & set(weak_non_production_terms) and not _query_mentions_any( + normalized_query, weak_non_production_terms ): - score -= 60 + score -= 40 aggregation_terms = ( "amount", diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0223f584e2..ffc988b4cc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -741,15 +741,14 @@ def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> ) ) asks_for_count = any( - term in normalized_query - for term in ( - "count", - "counts", - "how many", - "number of", - "record count", - "records", - "rows", + re.search(pattern, normalized_query) + for pattern in ( + r"\bcounts?\b", + r"\bhow many\b", + r"\bnumber of\b", + r"\brecord count\b", + r"\brecords?\b", + r"\brows?\b", ) ) if not asks_for_measure_sum or asks_for_count: @@ -909,6 +908,123 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True + def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) + if not normalized_query: + return None + + match = re.search( + r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" + r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", + normalized_query, + flags=re.IGNORECASE, + ) + if not match: + return None + + phrase = match.group("phrase").strip(" .,;:()[]{}'\"") + phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) + if not phrase or len(phrase) < 3: + return None + if re.search( + r"\b(?:table|model|schema|column|columns|market|region|country|division|" + r"date|month|year|quarter|top|count|number|amount|value)\b", + phrase, + flags=re.IGNORECASE, + ): + return None + return phrase + + def _preferred_entity_lookup_columns( + self, query: str | None, table: dict[str, Any] + ) -> set[str]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + candidate_groups: list[tuple[str, ...]] = [] + if "account" in normalized_query: + candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) + if "company" in normalized_query: + candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) + candidate_groups.extend( + [ + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + ), + ("Client", "ClientName"), + ("Account", "AccountName"), + ("Company", "CompanyName"), + ("Name",), + ] + ) + + columns: set[str] = set() + for candidates in candidate_groups: + column = self._find_schema_column(table, candidates) + if column: + columns.add(column) + return columns + + def _sql_satisfies_entity_lookup_request( + self, + sql: str, + query: str | None, + referenced_tables: list[str], + referenced_columns_by_table: dict[str, set[str]], + valid_tables: dict[str, dict[str, Any]], + ) -> bool: + if not self._extract_entity_lookup_phrase(query): + return True + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if any( + term in normalized_query + for term in (" by ", " per ", " each ", "distribution", "top", "count") + ): + return True + + for table_reference in referenced_tables: + table = self._table_for_sql_reference(table_reference, valid_tables) + if not table: + continue + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + + table_key = str(table_reference or "").lower() + referenced_columns = referenced_columns_by_table.get( + table_key + ) or referenced_columns_by_table.get( + table_key.split(".")[-1], + set(), + ) + referenced_column_keys = { + self._normalize_schema_identifier_key(column) + for column in referenced_columns + } + preferred_column_keys = { + self._normalize_schema_identifier_key(column) + for column in preferred_columns + } + if referenced_column_keys & preferred_column_keys: + return True + + logger.warning( + "Ignoring SQL because entity lookup did not use available customer/name columns. " + "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", + query, + table.get("name"), + sorted(preferred_columns), + sorted(referenced_columns), + sql, + ) + return False + + return True + def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1222,6 +1338,14 @@ def _sql_matches_question_intent( return False if not self._sql_satisfies_unique_entity_request(sql, query): return False + if not self._sql_satisfies_entity_lookup_request( + sql, + query, + referenced_tables, + referenced_columns_by_table, + valid_tables, + ): + return False if not expects_dimension: return True @@ -1372,6 +1496,145 @@ def _quote_sql_identifier(self, identifier: str) -> str: def _normalize_schema_identifier_key(self, value: str) -> str: return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + def _schema_identifier_words(self, value: str | None) -> list[str]: + raw_value = str(value or "") + normalized_parts = re.sub( + r"([a-z0-9])([A-Z])", + r"\1 \2", + raw_value.replace(".", " ").replace("$", " "), + ) + return [ + token.lower() + for token in re.split(r"[^A-Za-z0-9]+", normalized_parts) + if token + ] + + def _is_non_business_table_name(self, table_name: str | None) -> bool: + words = self._schema_identifier_words(table_name) + if words and words[0] in {"dbo", "public"}: + words = words[1:] + if not words: + return False + + normalized = "".join(words) + non_business_words = { + "archive", + "backup", + "bak", + "copy", + "dev", + "etl", + "import", + "landing", + "load", + "raw", + "scratch", + "snapshot", + "stage", + "staging", + "temp", + "test", + "tmp", + "work", + "wrk", + } + if any(word in non_business_words for word in words): + return True + + non_business_prefixes = ( + "tmp", + "temp", + "stage", + "staging", + "stg", + "load", + "raw", + "wrk", + "work", + "test", + ) + if any(normalized.startswith(prefix) for prefix in non_business_prefixes): + return True + + return any( + normalized.endswith(suffix) + for suffix in ("backup", "bak", "copy", "test", "tmp", "temp") + ) + + def _is_hygiene_excluded_table_name( + self, + table_name: str | None, + query: str | None = None, + ) -> bool: + return self._is_non_business_table_name(table_name) + + def _filter_schema_tables_for_query( + self, + query: str | None, + tables: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + return [ + table + for table in tables + if not self._is_hygiene_excluded_table_name( + str(table.get("name") or ""), + query, + ) + ] + + def _filter_sql_generation_context_by_hygiene( + self, + query: str | None, + documents: list[dict], + table_names: list[str], + table_ddls: list[str], + ) -> tuple[list[dict], list[str], list[str]]: + kept_documents: list[dict] = [] + kept_table_names: list[str] = [] + kept_table_ddls: list[str] = [] + + for index, document in enumerate(documents): + document_table_names: list[str] = [] + if isinstance(table_name := document.get("table_name"), str): + document_table_names.append(table_name) + if isinstance(table_ddl := document.get("table_ddl"), str): + document_table_names.extend( + str(table.get("name") or "") + for table in self._parse_schema_tables([table_ddl]) + if table.get("name") + ) + + if any( + self._is_hygiene_excluded_table_name(candidate, query) + for candidate in document_table_names + ): + continue + + kept_documents.append(document) + if index < len(table_names): + kept_table_names.append(table_names[index]) + if index < len(table_ddls): + kept_table_ddls.append(table_ddls[index]) + + if documents: + return kept_documents, kept_table_names, kept_table_ddls + + for table_name, table_ddl in zip(table_names, table_ddls): + parsed_names = [ + str(table.get("name") or "") + for table in self._parse_schema_tables([table_ddl]) + if table.get("name") + ] or [table_name] + if any( + self._is_hygiene_excluded_table_name(candidate, query) + for candidate in parsed_names + ): + continue + kept_table_names.append(table_name) + kept_table_ddls.append(table_ddl) + + return kept_documents, kept_table_names, kept_table_ddls + def _schema_identifier_alias_keys(self, value: str) -> set[str]: raw_value = str(value or "").strip() base_key = self._normalize_schema_identifier_key(raw_value) @@ -1492,6 +1755,7 @@ def _table_matches_query(self, table_name: str, query: str) -> bool: def _find_best_schema_table_for_query( self, query: str, tables: list[dict[str, Any]] ) -> dict[str, Any] | None: + tables = self._filter_schema_tables_for_query(query, tables) if not tables: return None @@ -1757,6 +2021,7 @@ def _build_explicit_table_preview_sql( return None tables = self._parse_schema_tables(table_ddls) + tables = self._filter_schema_tables_for_query(query, tables) if not tables: return None @@ -2089,6 +2354,53 @@ def _build_schema_literal_filter_conditions( return conditions + def _build_entity_lookup_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + phrase = self._extract_entity_lookup_phrase(query) + if not phrase: + return None + + scored: list[tuple[int, dict[str, Any], str, str | None]] = [] + for table in self._filter_schema_tables_for_query(query, tables): + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + lookup_column = sorted(preferred_columns)[0] + date_column = self._find_temporal_column_for_query(query, table) + table_name = str(table.get("name") or "") + if not table_name: + continue + + score = 10 + normalized_table = self._normalize_schema_token(table_name) + if any(term in normalized_table for term in ("order", "record", "event")): + score += 10 + if date_column: + score += 5 + scored.append((score, table, lookup_column, date_column)) + + if not scored: + return None + + _, table, lookup_column, date_column = sorted( + scored, + key=lambda item: item[0], + reverse=True, + )[0] + table_name = str(table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + lookup_ref = f"{table_ref}.{self._quote_sql_identifier(lookup_column)}" + escaped_phrase = phrase.replace("'", "''") + sql = ( + f"SELECT TOP 500 * FROM {table_ref} " + f"WHERE {lookup_ref} = '{escaped_phrase}'" + ) + if date_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + sql += f" ORDER BY {date_ref} DESC" + return sql + def _select_best_analytics_table( self, tables: list[dict[str, Any]], @@ -2099,6 +2411,7 @@ def _select_best_analytics_table( query: str = "", ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + tables = self._filter_schema_tables_for_query(query, tables) scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] ] = [] @@ -2149,8 +2462,8 @@ def _select_best_analytics_table( score += 4 if "invoice" in table_name or "inv" in table_name: score += 3 - if "stage" in table_name: - score -= 8 + if self._is_non_business_table_name(table_name): + score -= 25 if any( term in normalized_query for term in ("order", "orders", "new order", "new orders") @@ -2223,11 +2536,15 @@ def _build_schema_grounded_analytics_sql( return None tables = self._parse_schema_tables(table_ddls) + tables = self._filter_schema_tables_for_query(query, tables) if not tables: return None compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): + return entity_lookup_sql + if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2283,16 +2600,7 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - if not is_sales_or_order_query: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not any( + wants_measure_query = any( term in normalized_query for term in ( "amount", @@ -2309,6 +2617,21 @@ def _build_schema_grounded_analytics_sql( "value", ) ) + if not is_sales_or_order_query and not wants_measure_query: + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + + wants_count_metric = any( + re.search(pattern, normalized_query) + for pattern in ( + r"\bcounts?\b", + r"\bvolume\b", + r"\bhow many\b", + r"\bdistribution\b", + ) + ) and not wants_measure_query wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") ) @@ -5316,6 +5639,15 @@ def _prune_sql_generation_context( *, max_tables: int = 8, ) -> tuple[list[dict], list[str], list[str]]: + documents, table_names, table_ddls = ( + self._filter_sql_generation_context_by_hygiene( + query, + documents, + table_names, + table_ddls, + ) + ) + if len(table_ddls) <= max_tables: return documents, table_names, table_ddls diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 3ee705b621..bbb0c79aef 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -67,6 +67,26 @@ def test_rerank_table_documents_prefers_question_relevant_table_text(): assert documents[0].meta["name"] == "business_transactions" +def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): + test_load = Document( + content="Raw test load rows for order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ) + order_market_table = Document( + content="New order transaction records with market and customer fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.45, + ) + + documents = _rerank_table_documents( + "Show order distribution across markets.", + [test_load, order_market_table], + ) + + assert documents[0].meta["name"] == "dbo_xStageNewOrders" + + def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index a64038fbb9..3f493af91c 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,6 +269,115 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) +def test_validated_sql_rejects_country_question_without_country_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' + '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' + 'ORDER BY COUNT(*) DESC' + ), + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + "Show the total commodity line value by country.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Show the total commodity line value by country.", + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'AS "Country_of_Ultimate_Destination_Code", ' + 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' + 'AS "TotalCommodity_Line_Value" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' + ) + + +def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT * FROM "dbo_tnoStageNewOrders" ' + 'WHERE "dbo_tnoStageNewOrders"."Division" = ' + "'Daimler Trucks North America'" + ), + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + "List orders for Daimler Trucks North America.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "List orders for Daimler Trucks North America.", + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """, + ], + ) + + assert sql == ( + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."CustName" = ' + "'Daimler Trucks North America' " + 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' + ) + + def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) @@ -1058,7 +1167,7 @@ def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_country(): ) -def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_prefixed_country(): +def test_build_schema_grounded_sales_sql_ignores_load_table_for_revenue_question(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( "Which countries have the highest order revenue?", @@ -1073,17 +1182,10 @@ def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_prefixed_c ], ) - assert sql == ( - 'SELECT "dbo_xStageLoad8"."col_07_Country" AS "col_07_Country", ' - 'SUM("dbo_xStageLoad8"."TotalOrderValue") AS "TotalTotalOrderValue" ' - 'FROM "dbo_xStageLoad8" ' - 'WHERE "dbo_xStageLoad8"."col_07_Country" IS NOT NULL ' - 'GROUP BY "dbo_xStageLoad8"."col_07_Country" ' - 'ORDER BY SUM("dbo_xStageLoad8"."TotalOrderValue") DESC' - ) + assert sql is None -def test_build_schema_grounded_sales_sql_for_losing_order_value_by_market(): +def test_build_schema_grounded_sales_sql_ignores_staging_table_for_market_question(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( "Which markets are losing order value?", @@ -1098,14 +1200,7 @@ def test_build_schema_grounded_sales_sql_for_losing_order_value_by_market(): ], ) - assert sql == ( - 'SELECT "dbo_tblStageNewOrders"."Market" AS "Market", ' - 'SUM("dbo_tblStageNewOrders"."TotalOrderValue") AS "TotalTotalOrderValue" ' - 'FROM "dbo_tblStageNewOrders" ' - 'WHERE "dbo_tblStageNewOrders"."Market" IS NOT NULL ' - 'GROUP BY "dbo_tblStageNewOrders"."Market" ' - 'ORDER BY SUM("dbo_tblStageNewOrders"."TotalOrderValue") ASC' - ) + assert sql is None def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): From 75d13e7385e4c3dc9c1e5f7134b77bf4b3f8333a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 14 Jul 2026 18:42:37 +0530 Subject: [PATCH 0543/1087] Restore SQL retrieval code to working commit --- .../retrieval/db_schema_retrieval.py | 33 ++--- wren-ai-service/src/web/v1/services/ask.py | 128 ------------------ .../retrieval/test_db_schema_retrieval.py | 20 --- .../pytest/services/test_ask_sales_sql.py | 109 --------------- 4 files changed, 12 insertions(+), 278 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6b64c9b452..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -213,15 +213,11 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms: set[str] = set() - for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): - split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) - for token in re.findall(r"[A-Za-z0-9]+", split_token): - if len(token) <= 2 or token.lower() in stop_words: - continue - normalized_token = _normalize_retrieval_token(token) - if normalized_token: - terms.add(normalized_token) + terms = { + _normalize_retrieval_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") + if len(token) > 2 and token.lower() not in stop_words + } return {term for term in terms if term} @@ -247,11 +243,7 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - weak_non_production_terms = ( - "stage", - "staging", - ) - strong_non_production_terms = ( + non_production_terms = ( "archive", "backup", "copy", @@ -259,18 +251,17 @@ def _source_shape_score(query: str, document: Document) -> int: "development", "duplicate", "sample", + "stage", + "staging", "temp", "test", "tmp", ) - if source_terms & set(strong_non_production_terms) and not _query_mentions_any( - normalized_query, strong_non_production_terms - ): - score -= 240 - if source_terms & set(weak_non_production_terms) and not _query_mentions_any( - normalized_query, weak_non_production_terms + if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( + normalized_query, + non_production_terms, ): - score -= 40 + score -= 60 aggregation_terms = ( "amount", diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ffc988b4cc..b40306c215 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -908,123 +908,6 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True - def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip()) - if not normalized_query: - return None - - match = re.search( - r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" - r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", - normalized_query, - flags=re.IGNORECASE, - ) - if not match: - return None - - phrase = match.group("phrase").strip(" .,;:()[]{}'\"") - phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) - if not phrase or len(phrase) < 3: - return None - if re.search( - r"\b(?:table|model|schema|column|columns|market|region|country|division|" - r"date|month|year|quarter|top|count|number|amount|value)\b", - phrase, - flags=re.IGNORECASE, - ): - return None - return phrase - - def _preferred_entity_lookup_columns( - self, query: str | None, table: dict[str, Any] - ) -> set[str]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - candidate_groups: list[tuple[str, ...]] = [] - if "account" in normalized_query: - candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) - if "company" in normalized_query: - candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) - candidate_groups.extend( - [ - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - ), - ("Client", "ClientName"), - ("Account", "AccountName"), - ("Company", "CompanyName"), - ("Name",), - ] - ) - - columns: set[str] = set() - for candidates in candidate_groups: - column = self._find_schema_column(table, candidates) - if column: - columns.add(column) - return columns - - def _sql_satisfies_entity_lookup_request( - self, - sql: str, - query: str | None, - referenced_tables: list[str], - referenced_columns_by_table: dict[str, set[str]], - valid_tables: dict[str, dict[str, Any]], - ) -> bool: - if not self._extract_entity_lookup_phrase(query): - return True - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if any( - term in normalized_query - for term in (" by ", " per ", " each ", "distribution", "top", "count") - ): - return True - - for table_reference in referenced_tables: - table = self._table_for_sql_reference(table_reference, valid_tables) - if not table: - continue - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - - table_key = str(table_reference or "").lower() - referenced_columns = referenced_columns_by_table.get( - table_key - ) or referenced_columns_by_table.get( - table_key.split(".")[-1], - set(), - ) - referenced_column_keys = { - self._normalize_schema_identifier_key(column) - for column in referenced_columns - } - preferred_column_keys = { - self._normalize_schema_identifier_key(column) - for column in preferred_columns - } - if referenced_column_keys & preferred_column_keys: - return True - - logger.warning( - "Ignoring SQL because entity lookup did not use available customer/name columns. " - "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", - query, - table.get("name"), - sorted(preferred_columns), - sorted(referenced_columns), - sql, - ) - return False - - return True - def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1338,14 +1221,6 @@ def _sql_matches_question_intent( return False if not self._sql_satisfies_unique_entity_request(sql, query): return False - if not self._sql_satisfies_entity_lookup_request( - sql, - query, - referenced_tables, - referenced_columns_by_table, - valid_tables, - ): - return False if not expects_dimension: return True @@ -2542,9 +2417,6 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): - return entity_lookup_sql - if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index bbb0c79aef..3ee705b621 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -67,26 +67,6 @@ def test_rerank_table_documents_prefers_question_relevant_table_text(): assert documents[0].meta["name"] == "business_transactions" -def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): - test_load = Document( - content="Raw test load rows for order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ) - order_market_table = Document( - content="New order transaction records with market and customer fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.45, - ) - - documents = _rerank_table_documents( - "Show order distribution across markets.", - [test_load, order_market_table], - ) - - assert documents[0].meta["name"] == "dbo_xStageNewOrders" - - def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 3f493af91c..e367bdedfd 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,115 +269,6 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) -def test_validated_sql_rejects_country_question_without_country_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' - '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' - 'ORDER BY COUNT(*) DESC' - ), - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - "Show the total commodity line value by country.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Show the total commodity line value by country.", - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'AS "Country_of_Ultimate_Destination_Code", ' - 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' - 'AS "TotalCommodity_Line_Value" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' - ) - - -def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT * FROM "dbo_tnoStageNewOrders" ' - 'WHERE "dbo_tnoStageNewOrders"."Division" = ' - "'Daimler Trucks North America'" - ), - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - "List orders for Daimler Trucks North America.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "List orders for Daimler Trucks North America.", - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """, - ], - ) - - assert sql == ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" = ' - "'Daimler Trucks North America' " - 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' - ) - - def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) From a9ac692e39cf1b5160b88833e29207118a65259e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 02:02:33 +0530 Subject: [PATCH 0544/1087] Restore Wren AI files to 052756 working state --- WrenAI-FineTuning-Playbook.md | 494 ----------------- .../retrieval/db_schema_retrieval.py | 60 ++- wren-ai-service/src/web/v1/services/ask.py | 509 ++++++++++-------- .../retrieval/test_db_schema_retrieval.py | 42 ++ .../pytest/services/test_ask_sales_sql.py | 208 ++++++- 5 files changed, 584 insertions(+), 729 deletions(-) delete mode 100644 WrenAI-FineTuning-Playbook.md diff --git a/WrenAI-FineTuning-Playbook.md b/WrenAI-FineTuning-Playbook.md deleted file mode 100644 index 254c043189..0000000000 --- a/WrenAI-FineTuning-Playbook.md +++ /dev/null @@ -1,494 +0,0 @@ -WrenAI Fine-Tuning & Integration Playbook — NetRay / Orders -Instance: http://10.104.74.13:3000 (self-hosted OSS) Scope: Orders/Sales domain + a separate Tariffs/Customs domain -Important framing: WrenAI is not model fine-tuning. There is no training run. You are doing context engineering — pushing your business knowledge into three layers Wren injects into the LLM prompt at query time (via vector retrieval). Do these in order; each layer compounds on the one before it. - -How to use this doc: Anything in is a placeholder your team must replace with real schema/business facts. The defaults are SyteLine-flavoured starting points — keep, edit, or delete them. Sections 4 and 5 are meant to be pasted directly into Wren's Knowledge tab. - - -0. Order of operations (the whole playbook in 8 steps) -Prune staging/junk tables out of the model. -Gold schema — point Wren only at clean views (best) or curated base tables. -Generate semantics (Modeling AI Assistant) → auto-fill descriptions. -Hand-correct the cryptic/ambiguous columns only. -Relationships — generate + verify join keys. -Instructions — paste Section 4 into Knowledge, edit placeholders. -Question-SQL pairs — paste Section 5, verify each runs, save. -Deploy, smoke-test, then wire the API into SyteRay (Section 6). - -Everything before "Deploy" is invisible to users until you hit Deploy. Deploy re-embeds the context into the vector store — it is not optional and it is the step people forget. - - -1. The three context layers (what actually moves accuracy) -Layer -Where in UI -What it holds -Fixes -Semantics (MDL) -Modeling tab -Table + column descriptions, types, relationships, calculated fields -"AI doesn't know what OrdNo / col_01_Division means" -Instructions -Knowledge tab -Reusable rules: terminology, filters, formatting, join rules -"AI counts cancelled orders in revenue", "money not rounded", inconsistent metric logic -Question-SQL pairs -Knowledge tab -Gold examples pinning a question → exact SQL -Complex/error-prone recurring questions (backlog cost, YoY growth) - - -Rule of thumb: facts about columns → Semantics. Rules about logic → Instructions. Whole gold answers → Q-SQL pairs. - - -2. Phase 0 — Data hygiene (do this first, it's 80% of the win) -Your model list is mostly staging junk: xStage, xStageLoad, xStageLoad2/3/4/5, xStageLoad8, xStageLoad8_Test, xStageNewOrders. Every one is a table the AI can wrongly select — that's exactly why it was recommending questions about dbo.xStageLoad8. - -2.1 Prune. In Modeling, remove every staging/load/test table from the model. Keep only business-meaningful tables. - -2.2 Build a gold layer (strongly recommended). In SQL Server, create a dedicated schema of clean, analytics-friendly views and point Wren only at those: - -CREATE SCHEMA gold; - -GO - --- One row per order line, business-named columns, junk excluded - -CREATE VIEW gold.v_orders AS - -SELECT - - OrdNo AS order_number, - - CustNo AS customer_number, - - CustName AS customer_name, - - CustPO AS customer_po, - - Market AS market, - - AS division, - - AS segment, - - AS salesperson, - - AS product, - - AS order_status, - - AS order_date, - - AS invoice_date, - - AS qty_ordered, - - AS line_amount, -- pre-summed numeric measure - - AS is_backlog -- boolean, see §4 - -FROM dbo. - -WHERE ; - -Why gold views beat raw tables: - -The LLM reads clean names (order_status, not col_07_stat) → fewer wrong guesses. -You bake business logic (status filters, boolean flags, fiscal columns) into SQL once, so the AI doesn't reinvent it every query. -You control exactly what's exposed — no PII, no staging tables. - -2.3 Read-only DB user. Wren should connect via a wren_ro login with SELECT-only on gold (and nothing on staging). Never a write-capable account. - -2.4 Split domains into separate projects. dbo.ytblTarrifsFullA (Importer_Number, HTS, Entry_Date, Legal_Entity) is customs/tariff data — a different subject from sales. Mixing it into the Orders graph makes the AI blend customs columns into sales answers. Create a separate Wren project for Tariffs. Each project stays coherent and "opinionated," which is what reduces hallucination. - - -3. Phase 1 — Semantics (Modeling) -3.1 Auto-generate. Modeling page → Modeling AI Assistant (top-right) → Generate semantics. This fills the model + column Description fields across all tables from the schema. Then Generate relationships. - -3.2 Hand-correct only what the AI can't infer. The assistant handles obvious columns. You manually fix: - -Cryptic codes: col_01_Division, FY___Would_invoice_date, status/segment codes. -Ambiguous pairs: if two columns could both be "revenue" or "date," describe each precisely and say when each is used. -Measures vs dimensions vs IDs — phrasing matters (see table below). - -3.3 Description conventions (this is your data dictionary; phrasing drives behaviour): - -Column role -Write the description as… -Example -Measure (sum/avg) -"Total/Amount of … used for …" -line_amount → "Extended line amount in USD; sum for sales revenue." -Dimension (group/filter) -"Category of …" / "… segment, not geography" -market → "Business unit / market segment. NOT a country or region." -Date -"Date used for … filtering" -invoice_date → "Date the order line was invoiced; default date for revenue-by-period." -ID / key -"Unique identifier for …" -customer_number → "SyteLine customer code; use COUNT(DISTINCT) for customer counts." - - -Table-level description example for gold.v_orders: - -"One row per sales order line from SyteLine. Grain = order line. Use for sales, revenue, backlog, orders-by-market/division/salesperson analysis. Excludes voided and test orders." - -3.4 Relationships. Verify the auto-detected joins and add any missed keys explicitly (e.g. v_orders.customer_number → v_customers.customer_number, many-to-one). Explicit relationships = deterministic joins; without them the AI guesses. - -3.5 Deploy. - - -4. Phase 2 — INSTRUCTIONS (paste into Knowledge → Instructions) -Two types: Global (always applied) and Question-Matching (applied only when the question matches a topic/keyword). Add each block below as a separate instruction of the stated type. Edit every before deploying — a wrong rule is worse than no rule. -4A. Global instructions (always on) -[GLOBAL] Currency & rounding - -All monetary values are in USD. Always ROUND(value, 2) for revenue, averages, - -and percentages. Format large money values with thousands separators in summaries. - -[GLOBAL] Valid orders only - -Exclude orders where order_status IN () - -from any sales, revenue, backlog, or count calculation, unless the user explicitly - -asks about cancelled/quoted orders. - -[GLOBAL] Default date field - -For any sales/revenue question by time period, use invoice_date as the default - -date field. For "new orders" or "orders received," use order_date instead. - -[GLOBAL] Default time range - -If the user gives no date range, default to the last 90 days. Always state the - -range you assumed in the answer summary. - -[GLOBAL] Counting entities - -Headcount-style counts must use COUNT(DISTINCT ...), not COUNT(*). - -Customers = COUNT(DISTINCT customer_number). Orders = COUNT(DISTINCT order_number). - -[GLOBAL] Safe joins - -Use LEFT JOIN when joining optional/reference tables (products, salesperson, - -customer master) so order rows are never dropped when a lookup is missing. - -[GLOBAL] Grain awareness - -gold.v_orders is at ORDER-LINE grain. When counting or summing at the order level, - -aggregate to order_number first to avoid double counting. - -[GLOBAL] Fiscal calendar - -Our fiscal year runs . "FY", "quarter", "YTD", and - -"MTD" all refer to the fiscal calendar, not the calendar year. . -4B. Terminology instructions (Global — your business dictionary as rules) -[GLOBAL] Term: "Backlog" - -"Backlog" = open order lines not yet shipped/invoiced, i.e. - -is_backlog = 1 (or order_status = AND invoice_date IS NULL). - -"Backlog cost" / "total cost of backlog" = SUM(line_amount) over backlog lines. - -[GLOBAL] Term: "Market" vs "Division" vs "Segment" - -- market = business unit / market segment (e.g. "Honeywell BTP"). NOT geography. - -- division = . - -- segment = . - -When a user says "market," never map it to a country/region column. - -[GLOBAL] Term: "New orders" - -"New orders" = orders where order_date falls in the requested period, regardless - -of invoice status. Distinct from "sales/revenue" which uses invoice_date. - -[GLOBAL] Term: "Growing / declining market" - -Growth = period-over-period change in SUM(line_amount) by market. Default - -comparison is . Always show both periods and the % change. -4C. Question-Matching instructions (topic-scoped) -[MATCH: "year over year", "YoY", "vs last year", "growth"] - -Compute YoY as: current-period SUM(line_amount) vs the same period one fiscal - -year earlier, grouped by the requested dimension. Return both values and - -ROUND(((current-prior)/NULLIF(prior,0))*100, 2) AS pct_change. - -[MATCH: "salesperson", "sales rep", "who sold"] - -Attribute revenue via . Exclude house/unassigned - -accounts () unless explicitly asked. - -[MATCH: "underperforming", "underperform", "lagging"] - -"Underperforming" business units = those below for the period. State the benchmark used. - -[MATCH: chart / trend / over time] - -For time-series, order the x-axis chronologically and use a line chart. For - -"by market/division/product" rankings, use a horizontal bar chart sorted desc. - -Scoping discipline: Keep Globals few and universally true. Anything that only applies to one kind of question belongs in a Question-Matching instruction, or the AI over-applies it. - - -5. Phase 3 — QUESTION-SQL PAIRS (paste into Knowledge → Question-SQL Pairs) -These are drawn from the real questions already in your thread panel. Each SQL below is a template — replace columns to match your gold views, run it in Wren once, confirm the result, then Save. A pair with wrong SQL trains the AI wrongly, so verify before saving. - -Q1 — "Show total sales by market" - -SELECT market, - - ROUND(SUM(line_amount), 2) AS total_sales - -FROM gold.v_orders - -WHERE order_status NOT IN () - -GROUP BY market - -ORDER BY total_sales DESC; - -Q2 — "Which division has the highest sales?" - -SELECT TOP 1 division, - - ROUND(SUM(line_amount), 2) AS total_sales - -FROM gold.v_orders - -WHERE order_status NOT IN () - -GROUP BY division - -ORDER BY total_sales DESC; - -Q3 — "What is the total cost of all backlog?" - -SELECT ROUND(SUM(line_amount), 2) AS backlog_value - -FROM gold.v_orders - -WHERE is_backlog = 1; - -Q4 — "Break down month-to-date new orders" - -SELECT market, - - COUNT(DISTINCT order_number) AS new_orders, - - ROUND(SUM(line_amount), 2) AS order_value - -FROM gold.v_orders - -WHERE order_date >= DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1) - - AND order_date < DATEADD(DAY, 1, CAST(GETDATE() AS date)) - -GROUP BY market - -ORDER BY order_value DESC; - -Q5 — "Which markets are growing (YoY)?" - -WITH cur AS ( - - SELECT market, SUM(line_amount) AS amt - - FROM gold.v_orders - - WHERE invoice_date >= - - AND invoice_date < - - GROUP BY market), - -prior AS ( - - SELECT market, SUM(line_amount) AS amt - - FROM gold.v_orders - - WHERE invoice_date >= - - AND invoice_date < - - GROUP BY market) - -SELECT c.market, - - ROUND(c.amt,2) AS current_sales, - - ROUND(p.amt,2) AS prior_sales, - - ROUND(((c.amt - p.amt)/NULLIF(p.amt,0))*100, 2) AS pct_change - -FROM cur c LEFT JOIN prior p ON c.market = p.market - -ORDER BY pct_change DESC; - -Q6 — "Show new orders by CustName" - -SELECT customer_name, - - COUNT(DISTINCT order_number) AS orders, - - ROUND(SUM(line_amount), 2) AS order_value - -FROM gold.v_orders - -WHERE order_date >= - -GROUP BY customer_name - -ORDER BY order_value DESC; - -Q7 — "Which salesperson generated the most revenue?" - -SELECT TOP 10 salesperson, - - ROUND(SUM(line_amount), 2) AS revenue - -FROM gold.v_orders - -WHERE order_status NOT IN () - - AND salesperson <> - -GROUP BY salesperson - -ORDER BY revenue DESC; - -Add 3–5 more from your panel ("Which products contributed most," "Which business units are underperforming," "Which customers increased sales") the same way once the gold views are final. - - -6. Phase 4 — Integration (wiring Wren into SyteRay / your stack) -6.1 Get an API key -Wren UI → API tab → generate a key. All REST calls use header Authorization: Bearer . Base URL (self-hosted): http://10.104.74.21:3000/api/v1 - -Tier note: the REST Embedded AI API (generate_sql, generate_chart, streaming) is a governed/Agentic feature. If your OSS build's API tab exposes keys and these endpoints, use them (below). If not, the always-available OSS path is the GraphQL createAskingTask mutation (6.4). Check your API tab first. -6.2 Core REST endpoints -Generate SQL from a question: - -curl -X POST 'http://10.104.74.21:3000/api/v1/generate_sql' \ - - -H 'Authorization: Bearer ' \ - - -H 'Content-Type: application/json' \ - - -d '{ "projectId": , "question": "Show total sales by market" }' - -# → { "sql": "SELECT ...", "threadId": "..." } - -Other endpoints under the same base: - -POST /generate_chart — returns a Vega-Lite chart spec from a result set. -Streaming (SSE) — real-time token/step feedback for a chat UX. -Metadata introspection — list deployed models, columns, relationships, views (useful to render a schema picker in SyteRay). -Knowledge — read/manage instructions & Q-SQL pairs programmatically (supported tiers). - -Note: there are no webhooks — the client long-polls or uses SSE for async results. -6.3 Recommended SyteRay integration pattern -User asks a question in your UI. -SyteRay → generate_sql → gets governed SQL (which already respects your instructions + semantics). -Execute against the gold schema with the read-only user (either let Wren run it, or run it yourself for tighter control). -Optionally generate_chart for the visual. -Log threadId for audit; feed thumbs-up answers back as new Q-SQL pairs (closes the learning loop). - -This fits SyteRay cleanly: Wren becomes the text-to-SQL + governance layer; your policy engine / RBAC / audit trail wrap around it. Keep Wren's DB user scoped to gold so no agent can touch raw ERP tables. -6.4 OSS fallback (GraphQL asking task) -If REST embedded endpoints aren't in your build: submit questions via the createAskingTask GraphQL mutation and poll the task/thread for the answer. (Inspect the browser Network tab on the Home "Ask" flow to see the exact mutation shape your version uses.) -6.5 MCP (optional, for agent access) -Wren exposes a Model Context Protocol interface so agents (e.g. Claude, ChatGPT) query through your semantic layer instead of raw tables. On self-hosted, this runs via the Wren engine's MCP server. Useful if you want SyteRay's own agents to consult Wren as a governed data tool. -6.6 LLM configuration (self-hosted, LLM-agnostic) -Wren is LLM-agnostic. To keep ERP data on-prem, point it at a local model via LiteLLM/Ollama in ~/.wrenai/config.yaml: - -type: llm - -provider: litellm_llm - -models: - - - api_base: http://host.docker.internal:11434/v1 - - model: ollama_chat/ - - timeout: 600 - - kwargs: - - n: 1 - - temperature: 0 - -temperature: 0 for deterministic SQL. For accuracy, prefer a strong model (GPT-4o / o-series or a 70B-class local model); small models produce shakier SQL on messy ERP schemas. - - -7. Phase 5 — Maintenance loop (keep it accurate) -Schema change detection. Wren flags when tables/columns are added, renamed, removed, or retyped. Review after every ERP/gold-view change — renamed columns silently break Q-SQL pairs. -Smoke test. Keep a fixed list of your top ~10 questions. Re-run after any Deploy. If one regresses, fix the layer responsible (semantics vs instruction vs pair) — don't patch prompts ad hoc. -Feedback loop. Each week: take real user questions → if the SQL was right, Save as a Q-SQL pair; if it was almost right, add/refine an Instruction; if it picked the wrong table/column, fix the Semantic description. This is the "training set without training." -Version control. In newer Wren, this context lives in Git-friendly instructions.md + queries.yml. Even on your build, keep this document in Git as the source of truth and re-apply on rebuilds. - - -8. Rollout to the team -Owners - -__ owns Semantics (Modeling + gold views). -__ owns Instructions + Q-SQL pairs (Knowledge). -Weekly 30-min review: smoke test + feedback-loop triage. - -Guidance for people asking questions - -Use business terms from the dictionary ("backlog," "market," "new orders") — they're now defined for the AI. -Always read the generated SQL before trusting a number. Wrong-but-confident is the failure mode. -If an answer is wrong, don't just rephrase — report it so an owner fixes the underlying layer. -Sales questions → Orders project. Customs/HTS questions → Tariffs project. Don't cross them. - - -Appendix A — Data dictionary CSV template -Fill one row per exposed column; use it to drive/QA the Modeling descriptions. - -model,column,display_name,description,role,notes - -v_orders,order_number,Order Number,"SyteLine sales order number",id,"COUNT(DISTINCT) for order counts" - -v_orders,customer_number,Customer Number,"SyteLine customer code",id,"FK to v_customers" - -v_orders,market,Market,"Business unit / market segment; NOT geography",dimension,"" - -v_orders,line_amount,Line Amount,"Extended line amount USD; sum for revenue",measure,"" - -v_orders,invoice_date,Invoice Date,"Date invoiced; default date for revenue-by-period",date,"" - -v_orders,order_date,Order Date,"Date order received; use for new orders",date,"" - -v_orders,is_backlog,Is Backlog,"1 = open unshipped line",dimension,"boolean flag" -Appendix B — Instruction scoping cheat-sheet -Universally true, every query → Global. -True only for one topic/keyword → Question-Matching. -A whole correct answer to a recurring question → Question-SQL pair. -A fact about what a column is → Semantic description, not an instruction. - - diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..5e24b26813 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -213,11 +213,15 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms = { - _normalize_retrieval_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") - if len(token) > 2 and token.lower() not in stop_words - } + terms: set[str] = set() + for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): + split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) + for token in re.findall(r"[A-Za-z0-9]+", split_token): + if len(token) <= 2 or token.lower() in stop_words: + continue + normalized_token = _normalize_retrieval_token(token) + if normalized_token: + terms.add(normalized_token) return {term for term in terms if term} @@ -243,7 +247,11 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - non_production_terms = ( + weak_non_production_terms = ( + "stage", + "staging", + ) + strong_non_production_terms = ( "archive", "backup", "copy", @@ -251,17 +259,18 @@ def _source_shape_score(query: str, document: Document) -> int: "development", "duplicate", "sample", - "stage", - "staging", "temp", "test", "tmp", ) - if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( - normalized_query, - non_production_terms, + if source_terms & set(strong_non_production_terms) and not _query_mentions_any( + normalized_query, strong_non_production_terms + ): + score -= 240 + if source_terms & set(weak_non_production_terms) and not _query_mentions_any( + normalized_query, weak_non_production_terms ): - score -= 60 + score -= 40 aggregation_terms = ( "amount", @@ -325,6 +334,26 @@ def _source_shape_score(query: str, document: Document) -> int: return score +def _is_unrequested_strong_non_production_source(query: str, document: Document) -> bool: + strong_non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "temp", + "test", + "tmp", + ) + source_terms = _retrieval_terms(_source_text(document)) + return bool( + source_terms & set(strong_non_production_terms) + and not _query_mentions_any(query or "", strong_non_production_terms) + ) + + def _document_relevance_score(document: Document, query_terms: set[str]) -> int: if not query_terms: return 0 @@ -429,6 +458,13 @@ def _select_relevant_table_documents( return documents[:max_tables] candidate_pool = [item for item in reranked if item[3] > 0] or reranked + production_pool = [ + item + for item in candidate_pool + if not _is_unrequested_strong_non_production_source(query, item[2]) + ] + if production_pool: + candidate_pool = production_pool selected = [ document for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b40306c215..00f9315fcc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -689,6 +689,8 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: concept_groups.append({"market", "region", "country", "territory"}) if "region" in normalized or "regions" in normalized: concept_groups.append({"region", "market", "area", "territory", "country"}) + if "country" in normalized or "countries" in normalized: + concept_groups.append({"country", "countries", "nation", "destination"}) if "quarterly" in normalized or "quarter" in normalized: concept_groups.append({"quarter", "quarterly"}) if "recurring" in normalized or "recurrence" in normalized: @@ -741,14 +743,15 @@ def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> ) ) asks_for_count = any( - re.search(pattern, normalized_query) - for pattern in ( - r"\bcounts?\b", - r"\bhow many\b", - r"\bnumber of\b", - r"\brecord count\b", - r"\brecords?\b", - r"\brows?\b", + term in normalized_query + for term in ( + "count", + "counts", + "how many", + "number of", + "record count", + "records", + "rows", ) ) if not asks_for_measure_sum or asks_for_count: @@ -908,6 +911,167 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True + def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) + if not normalized_query: + return None + + match = re.search( + r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" + r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", + normalized_query, + flags=re.IGNORECASE, + ) + if not match: + return None + + phrase = match.group("phrase").strip(" .,;:()[]{}'\"") + phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) + if not phrase or len(phrase) < 3: + return None + if re.search( + r"\b(?:table|model|schema|column|columns|market|region|country|division|" + r"date|month|year|quarter|top|count|number|amount|value)\b", + phrase, + flags=re.IGNORECASE, + ): + return None + return phrase + + def _preferred_entity_lookup_columns( + self, query: str | None, table: dict[str, Any] + ) -> set[str]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + candidate_groups: list[tuple[str, ...]] = [] + if "account" in normalized_query: + candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) + if "company" in normalized_query: + candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) + candidate_groups.extend( + [ + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + ), + ("Client", "ClientName"), + ("Account", "AccountName"), + ("Company", "CompanyName"), + ("Name",), + ] + ) + + columns: set[str] = set() + for candidates in candidate_groups: + column = self._find_schema_column(table, candidates) + if column: + columns.add(column) + return columns + + def _sql_satisfies_entity_lookup_request( + self, + sql: str, + query: str | None, + referenced_tables: list[str], + referenced_columns_by_table: dict[str, set[str]], + valid_tables: dict[str, dict[str, Any]], + ) -> bool: + if not self._extract_entity_lookup_phrase(query): + return True + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if any( + term in normalized_query + for term in (" by ", " per ", " each ", "distribution", "top", "count") + ): + return True + + for table_reference in referenced_tables: + table = self._table_for_sql_reference(table_reference, valid_tables) + if not table: + continue + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + + table_key = str(table_reference or "").lower() + referenced_columns = referenced_columns_by_table.get( + table_key + ) or referenced_columns_by_table.get( + table_key.split(".")[-1], + set(), + ) + referenced_column_keys = { + self._normalize_schema_identifier_key(column) + for column in referenced_columns + } + preferred_column_keys = { + self._normalize_schema_identifier_key(column) + for column in preferred_columns + } + if referenced_column_keys & preferred_column_keys: + return True + + logger.warning( + "Ignoring SQL because entity lookup did not use available customer/name columns. " + "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", + query, + table.get("name"), + sorted(preferred_columns), + sorted(referenced_columns), + sql, + ) + return False + + return True + + def _is_unrequested_non_production_table_reference( + self, table_name: str, query: str | None + ) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + normalized_table = self._normalize_schema_token(table_name) + strong_non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "temp", + "test", + "tmp", + ) + if not any(term in normalized_table for term in strong_non_production_terms): + return False + return not any( + re.search(rf"\b{re.escape(term)}\b", normalized_query) + for term in strong_non_production_terms + ) + + def _sql_avoids_unrequested_non_production_tables( + self, sql: str, query: str | None, referenced_tables: list[str] + ) -> bool: + invalid_tables = [ + table + for table in referenced_tables + if self._is_unrequested_non_production_table_reference(table, query) + ] + if not invalid_tables: + return True + + logger.warning( + "Ignoring SQL because it references unrequested non-production tables. " + "query=%s invalid_tables=%s sql=%s", + query, + invalid_tables, + sql, + ) + return False + def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1221,6 +1385,20 @@ def _sql_matches_question_intent( return False if not self._sql_satisfies_unique_entity_request(sql, query): return False + if not self._sql_satisfies_entity_lookup_request( + sql, + query, + referenced_tables, + referenced_columns_by_table, + valid_tables, + ): + return False + if not self._sql_avoids_unrequested_non_production_tables( + sql, + query, + referenced_tables, + ): + return False if not expects_dimension: return True @@ -1371,145 +1549,6 @@ def _quote_sql_identifier(self, identifier: str) -> str: def _normalize_schema_identifier_key(self, value: str) -> str: return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - def _schema_identifier_words(self, value: str | None) -> list[str]: - raw_value = str(value or "") - normalized_parts = re.sub( - r"([a-z0-9])([A-Z])", - r"\1 \2", - raw_value.replace(".", " ").replace("$", " "), - ) - return [ - token.lower() - for token in re.split(r"[^A-Za-z0-9]+", normalized_parts) - if token - ] - - def _is_non_business_table_name(self, table_name: str | None) -> bool: - words = self._schema_identifier_words(table_name) - if words and words[0] in {"dbo", "public"}: - words = words[1:] - if not words: - return False - - normalized = "".join(words) - non_business_words = { - "archive", - "backup", - "bak", - "copy", - "dev", - "etl", - "import", - "landing", - "load", - "raw", - "scratch", - "snapshot", - "stage", - "staging", - "temp", - "test", - "tmp", - "work", - "wrk", - } - if any(word in non_business_words for word in words): - return True - - non_business_prefixes = ( - "tmp", - "temp", - "stage", - "staging", - "stg", - "load", - "raw", - "wrk", - "work", - "test", - ) - if any(normalized.startswith(prefix) for prefix in non_business_prefixes): - return True - - return any( - normalized.endswith(suffix) - for suffix in ("backup", "bak", "copy", "test", "tmp", "temp") - ) - - def _is_hygiene_excluded_table_name( - self, - table_name: str | None, - query: str | None = None, - ) -> bool: - return self._is_non_business_table_name(table_name) - - def _filter_schema_tables_for_query( - self, - query: str | None, - tables: list[dict[str, Any]], - ) -> list[dict[str, Any]]: - return [ - table - for table in tables - if not self._is_hygiene_excluded_table_name( - str(table.get("name") or ""), - query, - ) - ] - - def _filter_sql_generation_context_by_hygiene( - self, - query: str | None, - documents: list[dict], - table_names: list[str], - table_ddls: list[str], - ) -> tuple[list[dict], list[str], list[str]]: - kept_documents: list[dict] = [] - kept_table_names: list[str] = [] - kept_table_ddls: list[str] = [] - - for index, document in enumerate(documents): - document_table_names: list[str] = [] - if isinstance(table_name := document.get("table_name"), str): - document_table_names.append(table_name) - if isinstance(table_ddl := document.get("table_ddl"), str): - document_table_names.extend( - str(table.get("name") or "") - for table in self._parse_schema_tables([table_ddl]) - if table.get("name") - ) - - if any( - self._is_hygiene_excluded_table_name(candidate, query) - for candidate in document_table_names - ): - continue - - kept_documents.append(document) - if index < len(table_names): - kept_table_names.append(table_names[index]) - if index < len(table_ddls): - kept_table_ddls.append(table_ddls[index]) - - if documents: - return kept_documents, kept_table_names, kept_table_ddls - - for table_name, table_ddl in zip(table_names, table_ddls): - parsed_names = [ - str(table.get("name") or "") - for table in self._parse_schema_tables([table_ddl]) - if table.get("name") - ] or [table_name] - if any( - self._is_hygiene_excluded_table_name(candidate, query) - for candidate in parsed_names - ): - continue - kept_table_names.append(table_name) - kept_table_ddls.append(table_ddl) - - return kept_documents, kept_table_names, kept_table_ddls - def _schema_identifier_alias_keys(self, value: str) -> set[str]: raw_value = str(value or "").strip() base_key = self._normalize_schema_identifier_key(raw_value) @@ -1630,7 +1669,6 @@ def _table_matches_query(self, table_name: str, query: str) -> bool: def _find_best_schema_table_for_query( self, query: str, tables: list[dict[str, Any]] ) -> dict[str, Any] | None: - tables = self._filter_schema_tables_for_query(query, tables) if not tables: return None @@ -1896,7 +1934,6 @@ def _build_explicit_table_preview_sql( return None tables = self._parse_schema_tables(table_ddls) - tables = self._filter_schema_tables_for_query(query, tables) if not tables: return None @@ -2002,7 +2039,18 @@ def _explicit_table_name_candidates(self, table_name: str) -> list[str]: separator_normalized = re.sub(r"[.$]", "_", table_name) if separator_normalized not in candidates: candidates.append(separator_normalized) + if "_" in table_name: + dotted_schema_name = re.sub( + r"^([A-Za-z_][A-Za-z0-9]*)_", + r"\1.", + table_name, + count=1, + ) + if dotted_schema_name not in candidates: + candidates.append(dotted_schema_name) short_name = re.split(r"[.$]", table_name)[-1] + if short_name == table_name and "_" in table_name: + short_name = table_name.split("_", 1)[-1] if short_name and short_name not in candidates: candidates.append(short_name) return candidates @@ -2229,53 +2277,6 @@ def _build_schema_literal_filter_conditions( return conditions - def _build_entity_lookup_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - phrase = self._extract_entity_lookup_phrase(query) - if not phrase: - return None - - scored: list[tuple[int, dict[str, Any], str, str | None]] = [] - for table in self._filter_schema_tables_for_query(query, tables): - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - lookup_column = sorted(preferred_columns)[0] - date_column = self._find_temporal_column_for_query(query, table) - table_name = str(table.get("name") or "") - if not table_name: - continue - - score = 10 - normalized_table = self._normalize_schema_token(table_name) - if any(term in normalized_table for term in ("order", "record", "event")): - score += 10 - if date_column: - score += 5 - scored.append((score, table, lookup_column, date_column)) - - if not scored: - return None - - _, table, lookup_column, date_column = sorted( - scored, - key=lambda item: item[0], - reverse=True, - )[0] - table_name = str(table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - lookup_ref = f"{table_ref}.{self._quote_sql_identifier(lookup_column)}" - escaped_phrase = phrase.replace("'", "''") - sql = ( - f"SELECT TOP 500 * FROM {table_ref} " - f"WHERE {lookup_ref} = '{escaped_phrase}'" - ) - if date_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - sql += f" ORDER BY {date_ref} DESC" - return sql - def _select_best_analytics_table( self, tables: list[dict[str, Any]], @@ -2286,7 +2287,6 @@ def _select_best_analytics_table( query: str = "", ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - tables = self._filter_schema_tables_for_query(query, tables) scored: list[ tuple[int, dict[str, Any], list[str], str | None, str | None] ] = [] @@ -2337,8 +2337,8 @@ def _select_best_analytics_table( score += 4 if "invoice" in table_name or "inv" in table_name: score += 3 - if self._is_non_business_table_name(table_name): - score -= 25 + if "stage" in table_name: + score -= 8 if any( term in normalized_query for term in ("order", "orders", "new order", "new orders") @@ -2403,6 +2403,88 @@ def _select_best_analytics_table( date_column, ) + def _build_entity_lookup_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + lookup_phrase = self._extract_entity_lookup_phrase(query) + if not lookup_phrase: + return None + + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + asks_for_orders = any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ) + + scored: list[tuple[int, dict[str, Any], str]] = [] + for table in tables: + table_name = str(table.get("name") or "") + if not table_name: + continue + + preferred_columns = self._preferred_entity_lookup_columns(query, table) + if not preferred_columns: + continue + + preferred_column = sorted( + preferred_columns, + key=lambda column: ( + 0 + if self._normalize_schema_identifier_key(column) + in {"custname", "customername", "customer"} + else 1, + column.lower(), + ), + )[0] + + score = 20 + normalized_table = self._normalize_schema_token(table_name) + if asks_for_orders: + if "order" in normalized_table: + score += 40 + if "neworder" in normalized_table: + score += 20 + if self._find_schema_column( + table, ("OrdNo", "OrderNo", "OrderId", "NewOrderId") + ): + score += 25 + if "test" in normalized_table or "tmp" in normalized_table: + score -= 80 + if "dev" in normalized_table or "backup" in normalized_table: + score -= 60 + if "stage" in normalized_table: + score -= 10 + scored.append((score, table, preferred_column)) + + if not scored: + return None + + _score, table, filter_column = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + table_name = str(table.get("name") or "") + if not table_name: + return None + + table_ref = self._quote_sql_identifier(table_name) + filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" + escaped_phrase = lookup_phrase.replace("'", "''") + date_column = self._find_schema_column( + table, + ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), + temporal=True, + ) + order_clause = ( + f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" + if date_column + else "" + ) + return ( + f"SELECT TOP 500 * FROM {table_ref} " + f"WHERE {filter_ref} = '{escaped_phrase}'" + f"{order_clause}" + ) + def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2411,12 +2493,14 @@ def _build_schema_grounded_analytics_sql( return None tables = self._parse_schema_tables(table_ddls) - tables = self._filter_schema_tables_for_query(query, tables) if not tables: return None compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): + return entity_lookup_sql + if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2472,7 +2556,7 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - wants_measure_query = any( + asks_for_measure_value = any( term in normalized_query for term in ( "amount", @@ -2489,21 +2573,17 @@ def _build_schema_grounded_analytics_sql( "value", ) ) - if not is_sales_or_order_query and not wants_measure_query: + + if not is_sales_or_order_query and not asks_for_measure_value: if categorical_count_sql := self._build_generic_categorical_count_sql( query, tables ): return categorical_count_sql wants_count_metric = any( - re.search(pattern, normalized_query) - for pattern in ( - r"\bcounts?\b", - r"\bvolume\b", - r"\bhow many\b", - r"\bdistribution\b", - ) - ) and not wants_measure_query + term in normalized_query + for term in ("count", "counts", "volume", "how many", "distribution") + ) and not asks_for_measure_value wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") ) @@ -5511,15 +5591,6 @@ def _prune_sql_generation_context( *, max_tables: int = 8, ) -> tuple[list[dict], list[str], list[str]]: - documents, table_names, table_ddls = ( - self._filter_sql_generation_context_by_hygiene( - query, - documents, - table_names, - table_ddls, - ) - ) - if len(table_ddls) <= max_tables: return documents, table_names, table_ddls diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 3ee705b621..7c7d46ce89 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -67,6 +67,26 @@ def test_rerank_table_documents_prefers_question_relevant_table_text(): assert documents[0].meta["name"] == "business_transactions" +def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): + test_load = Document( + content="Raw test load rows for order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ) + order_market_table = Document( + content="New order transaction records with market and customer fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.45, + ) + + documents = _rerank_table_documents( + "Show order distribution across markets.", + [test_load, order_market_table], + ) + + assert documents[0].meta["name"] == "dbo_xStageNewOrders" + + def test_select_relevant_table_documents_limits_weak_extra_candidates(): documents = [ Document( @@ -110,6 +130,28 @@ def test_select_relevant_table_documents_limits_weak_extra_candidates(): assert "staging_audit" not in [document.meta["name"] for document in selected] +def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): + documents = [ + Document( + content="Raw test load rows with order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ), + Document( + content="New order transaction records with market and customer details.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.4, + ), + ] + + selected = _select_relevant_table_documents( + "Show order distribution across markets.", + documents, + ) + + assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] + + @pytest.mark.asyncio async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index e367bdedfd..c7ab4c9272 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -269,6 +269,141 @@ def test_schema_grounded_table_question_groups_top_customers_by_order_count(): ) +def test_validated_sql_rejects_country_question_without_country_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' + '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' + 'ORDER BY COUNT(*) DESC' + ), + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + "Show the total commodity line value by country.", + ) + + assert result is None + + +def test_validated_sql_rejects_unrequested_test_table_for_market_distribution(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT TOP 10 "dbo_xStageLoad8_Test"."Market" AS "Market", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_xStageLoad8_Test" ' + 'WHERE "dbo_xStageLoad8_Test"."Market" IS NOT NULL ' + 'GROUP BY "dbo_xStageLoad8_Test"."Market" ' + 'ORDER BY COUNT(*) DESC' + ), + [ + """ + CREATE TABLE dbo_xStageLoad8_Test ( + Market VARCHAR, + OrdNo VARCHAR + ); + """ + ], + "Show order distribution across markets.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "Show the total commodity line value by country.", + [ + """ + CREATE TABLE dbo_ytblTarrifsExportsA ( + Country_of_Ultimate_Destination_Code VARCHAR, + Commodity_Line_Value DOUBLE + ); + """ + ], + ) + + assert sql == ( + 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'AS "Country_of_Ultimate_Destination_Code", ' + 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' + 'AS "TotalCommodity_Line_Value" ' + 'FROM "dbo_ytblTarrifsExportsA" ' + 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'IS NOT NULL ' + 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' + 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' + ) + + +def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT * FROM "dbo_tnoStageNewOrders" ' + 'WHERE "dbo_tnoStageNewOrders"."Division" = ' + "'Daimler Trucks North America'" + ), + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """ + ], + "List orders for Daimler Trucks North America.", + ) + + assert result is None + + +def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): + service = AskService.__new__(AskService) + + sql = service._build_schema_grounded_sales_sql( + "List orders for Daimler Trucks North America.", + [ + """ + CREATE TABLE dbo_tnoStageNewOrders ( + Division VARCHAR, + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + """ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + OrdNo VARCHAR, + OrdDate TIMESTAMP + ); + """, + ], + ) + + assert sql == ( + 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' + 'WHERE "dbo_tblNewOrders"."CustName" = ' + "'Daimler Trucks North America' " + 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' + ) + + def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): service = AskService.__new__(AskService) @@ -446,6 +581,21 @@ def test_extract_explicit_table_names_from_pcb_repair_phrases(): ) == ["ticket_labels", "dbo_ticket_labels"] +def test_explicit_table_name_candidates_include_dotted_and_short_forms(): + service = AskService.__new__(AskService) + + assert service._explicit_table_name_candidates("dbo_tblNewOrders") == [ + "dbo_tblNewOrders", + "dbo.tblNewOrders", + "tblNewOrders", + ] + assert service._explicit_table_name_candidates("dbo.tblNewOrders") == [ + "dbo.tblNewOrders", + "dbo_tblNewOrders", + "tblNewOrders", + ] + + def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): service = AskService.__new__(AskService) documents = [ @@ -482,6 +632,42 @@ def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): assert table_ddls == [documents[1]["table_ddl"]] +def test_filter_retrieval_metadata_for_explicit_query_matches_dotted_table_name(): + service = AskService.__new__(AskService) + documents = [ + { + "table_name": "dbo.tblNewOrders", + "table_ddl": """ + CREATE TABLE "dbo.tblNewOrders" ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + }, + { + "table_name": "dbo_other", + "table_ddl": """ + CREATE TABLE dbo_other ( + CustName VARCHAR, + OrdNo VARCHAR + ); + """, + }, + ] + + filtered_documents, table_names, table_ddls = ( + service._filter_retrieval_metadata_for_explicit_query( + "Show the top 5 CustName values from dbo_tblNewOrders by number of orders.", + documents, + ["dbo_tblNewOrders"], + ) + ) + + assert filtered_documents == [documents[0]] + assert table_names == ["dbo.tblNewOrders"] + assert table_ddls == [documents[0]["table_ddl"]] + + def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): service = AskService.__new__(AskService) result = service._build_validated_ask_result_from_sql( @@ -1058,7 +1244,7 @@ def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_country(): ) -def test_build_schema_grounded_sales_sql_ignores_load_table_for_revenue_question(): +def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_prefixed_country(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( "Which countries have the highest order revenue?", @@ -1073,10 +1259,17 @@ def test_build_schema_grounded_sales_sql_ignores_load_table_for_revenue_question ], ) - assert sql is None + assert sql == ( + 'SELECT "dbo_xStageLoad8"."col_07_Country" AS "col_07_Country", ' + 'SUM("dbo_xStageLoad8"."TotalOrderValue") AS "TotalTotalOrderValue" ' + 'FROM "dbo_xStageLoad8" ' + 'WHERE "dbo_xStageLoad8"."col_07_Country" IS NOT NULL ' + 'GROUP BY "dbo_xStageLoad8"."col_07_Country" ' + 'ORDER BY SUM("dbo_xStageLoad8"."TotalOrderValue") DESC' + ) -def test_build_schema_grounded_sales_sql_ignores_staging_table_for_market_question(): +def test_build_schema_grounded_sales_sql_for_losing_order_value_by_market(): service = AskService.__new__(AskService) sql = service._build_schema_grounded_sales_sql( "Which markets are losing order value?", @@ -1091,7 +1284,14 @@ def test_build_schema_grounded_sales_sql_ignores_staging_table_for_market_questi ], ) - assert sql is None + assert sql == ( + 'SELECT "dbo_tblStageNewOrders"."Market" AS "Market", ' + 'SUM("dbo_tblStageNewOrders"."TotalOrderValue") AS "TotalTotalOrderValue" ' + 'FROM "dbo_tblStageNewOrders" ' + 'WHERE "dbo_tblStageNewOrders"."Market" IS NOT NULL ' + 'GROUP BY "dbo_tblStageNewOrders"."Market" ' + 'ORDER BY SUM("dbo_tblStageNewOrders"."TotalOrderValue") ASC' + ) def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): From 59f9710f77366e814af2cfe373c425df4b5f5d83 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 03:47:09 +0530 Subject: [PATCH 0545/1087] Fix model metadata update display name fallbacks --- .../server/resolvers/diagramResolver.ts | 12 ++++--- .../modeling/metadata/EditModelMetadata.tsx | 33 ++++++++++++++----- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index dab3331b83..b457514555 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -243,21 +243,23 @@ export class DiagramResolver { id: uuidv4(), relationId: relation.id, nodeType: NodeType.RELATION, - displayName, + displayName: displayName || referenceName, referenceName, type: relation.joinType as RelationType, fromModelId: relation.fromModelId, fromModelName: relation.fromModelName, - fromModelDisplayName: relation.fromModelDisplayName, + fromModelDisplayName: + relation.fromModelDisplayName || relation.fromModelName, fromColumnId: relation.fromColumnId, fromColumnName: relation.fromColumnName, - fromColumnDisplayName: relation.fromColumnDisplayName, + fromColumnDisplayName: + relation.fromColumnDisplayName || relation.fromColumnName, toModelId: relation.toModelId, toModelName: relation.toModelName, - toModelDisplayName: relation.toModelDisplayName, + toModelDisplayName: relation.toModelDisplayName || relation.toModelName, toColumnId: relation.toColumnId, toColumnName: relation.toColumnName, - toColumnDisplayName: relation.toColumnDisplayName, + toColumnDisplayName: relation.toColumnDisplayName || relation.toColumnName, description: properties?.description, }; } diff --git a/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx b/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx index c40410c30b..5a14c57071 100644 --- a/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx +++ b/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx @@ -36,6 +36,12 @@ const CalculatedFieldEditableTable = makeEditableBaseTable(CalculatedFieldTable); const RelationshipEditableTable = makeEditableBaseTable(RelationTable); +const getMetadataId = (item: any) => + item?.relationId || item?.columnId || item?.nestedColumnId; + +const getFieldDisplayName = (item: any) => + item?.displayName ?? item?.referenceName ?? item?.sourceColumnName ?? ''; + export default function EditModelMetadata(props: Props) { const { formNamespace, @@ -61,17 +67,26 @@ export default function EditModelMetadata(props: Props) { }); }; - const handleMetadataChange = (fieldsName: string) => (value: any[]) => { + const handleMetadataChange = (fieldsName: string) => (value: any[] = []) => { // bind changeable metadata values onChange({ - [fieldsName]: value.map((item) => ({ - id: item.relationId || item.columnId || item.nestedColumnId, - description: item.description, - // Only models & fields, nested fields have alias - ...([FIELDS_NAME.FIELDS, FIELDS_NAME.NESTED_FIELDS].includes(fieldsName) - ? { displayName: item.displayName } - : {}), - })), + [fieldsName]: value + .map((item) => { + const id = getMetadataId(item); + if (!id) return null; + + return { + id, + description: item?.description, + // Only models & fields, nested fields have alias + ...([FIELDS_NAME.FIELDS, FIELDS_NAME.NESTED_FIELDS].includes( + fieldsName, + ) + ? { displayName: getFieldDisplayName(item) } + : {}), + }; + }) + .filter(Boolean), }); }; From 7555239f737fa0ac3807d4c91eca8152ea9d1d64 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 04:19:41 +0530 Subject: [PATCH 0546/1087] Fix metadata update and MSSQL timestamps --- .../server/repositories/sqlPairRepository.ts | 37 ++++++++++++++++++- .../server/repositories/viewRepository.ts | 34 +++++++++++++++++ .../apollo/server/resolvers/modelResolver.ts | 21 +++++++++-- 3 files changed, 87 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts b/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts index 9f02f48232..e4b7f3c945 100644 --- a/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts +++ b/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; export interface SqlPair { id: number; // ID @@ -19,4 +23,35 @@ export class SqlPairRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'sql_pair' }); } + + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date().toISOString(), + }, + queryOptions, + ); + } + + private withTimestamps = (data: Partial): Partial => { + const now = new Date().toISOString(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; } diff --git a/wren-ui/src/apollo/server/repositories/viewRepository.ts b/wren-ui/src/apollo/server/repositories/viewRepository.ts index ff125c76ab..c056ebd0e0 100644 --- a/wren-ui/src/apollo/server/repositories/viewRepository.ts +++ b/wren-ui/src/apollo/server/repositories/viewRepository.ts @@ -2,6 +2,7 @@ import { Knex } from 'knex'; import { BaseRepository, IBasicRepository, + IQueryOptions, coerceBoolean, } from './baseRepository'; @@ -13,6 +14,8 @@ export interface View { cached: boolean; // View is cached or not refreshTime?: string; // Contain a number followed by a time unit (ns, us, ms, s, m, h, d). For example, "2h" properties?: string; // View properties, a json string, the description and displayName should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IViewRepository extends IBasicRepository {} @@ -25,6 +28,28 @@ export class ViewRepository super({ knexPg, tableName: 'view' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + protected override transformFromDBData = (data: any): View => { const view = this.defaultTransformFromDBData(data) as View; return { @@ -32,4 +57,13 @@ export class ViewRepository cached: coerceBoolean(view.cached), }; }; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index b7e91310d2..da7337453c 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -37,6 +37,9 @@ logger.level = 'debug'; const dirtyProjectIds = new Set(); +const isSameId = (left: string | number, right: string | number) => + String(left) === String(right); + export enum SyncStatusEnum { IN_PROGRESS = 'IN_PROGRESS', SYNCRONIZED = 'SYNCRONIZED', @@ -861,7 +864,10 @@ export class ModelResolver { const relationships = await ctx.relationRepository.findRelationsByIds(relationshipIds); for (const rel of relationships) { - const requestedMetadata = data.relationships.find((r) => r.id === rel.id); + const requestedMetadata = data.relationships.find((r) => + isSameId(r.id, rel.id), + ); + if (!requestedMetadata) continue; const relationMetadata: any = {}; @@ -888,8 +894,9 @@ export class ModelResolver { await ctx.modelColumnRepository.findColumnsByIds(calculatedFieldIds); for (const col of modelColumns) { const requestedMetadata = data.calculatedFields.find( - (c) => c.id === col.id, + (c) => isSameId(c.id, col.id), ); + if (!requestedMetadata) continue; const columnMetadata: any = {}; // check if description is empty @@ -917,7 +924,10 @@ export class ModelResolver { const modelColumns = await ctx.modelColumnRepository.findColumnsByIds(columnIds); for (const col of modelColumns) { - const requestedMetadata = data.columns.find((c) => c.id === col.id); + const requestedMetadata = data.columns.find((c) => + isSameId(c.id, col.id), + ); + if (!requestedMetadata) continue; // update metadata const columnMetadata: any = {}; @@ -952,7 +962,10 @@ export class ModelResolver { nestedColumnIds, ); for (const col of modelNestedColumns) { - const requestedMetadata = data.nestedColumns.find((c) => c.id === col.id); + const requestedMetadata = data.nestedColumns.find((c) => + isSameId(c.id, col.id), + ); + if (!requestedMetadata) continue; const nestedColumnMetadata: any = {}; From f2f442e814cbf1fa6e1855b63ebd55fbf818a171 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 04:52:06 +0530 Subject: [PATCH 0547/1087] Handle null metadata during indexing and diagram load --- wren-ai-service/src/pipelines/indexing/table_description.py | 5 ++++- wren-ui/src/apollo/server/resolvers/diagramResolver.ts | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 8e1b875b49..3c9f3b7877 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -24,7 +24,10 @@ @component class TableDescriptionChunker: - def _truncate_description(self, description: str) -> str: + def _truncate_description(self, description: str | None) -> str: + if description is None: + return "" + if len(description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: return description diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index b457514555..1223720f63 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -287,10 +287,13 @@ export class DiagramResolver { }; } - private parseProperties(properties?: string | null): Record { + private parseProperties(properties?: string | Record | null): Record { if (!properties) { return {}; } + if (typeof properties === 'object') { + return properties; + } try { const parsed = JSON.parse(properties); return parsed && typeof parsed === 'object' ? parsed : {}; From ca701bdd67ae11d63574b9967c98952cdfe78118 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 05:11:02 +0530 Subject: [PATCH 0548/1087] Revert table description null handling change --- wren-ai-service/src/pipelines/indexing/table_description.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 3c9f3b7877..8e1b875b49 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -24,10 +24,7 @@ @component class TableDescriptionChunker: - def _truncate_description(self, description: str | None) -> str: - if description is None: - return "" - + def _truncate_description(self, description: str) -> str: if len(description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: return description From c92e7e686d64070c4b978dd0e0d5ad40b13eb9c5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 05:11:42 +0530 Subject: [PATCH 0549/1087] Handle schema retrieval timeouts without failing ask flow --- wren-ai-service/src/web/v1/services/ask.py | 176 +++++++++++------- .../v1/services/question_recommendation.py | 73 +++++++- 2 files changed, 178 insertions(+), 71 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 00f9315fcc..37187263f9 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -5066,6 +5066,54 @@ async def _run_with_timeout( except TimeoutError as exc: raise TimeoutError(f"{label} timed out after {timeout} seconds") from exc + def _empty_schema_retrieval_result(self) -> dict[str, Any]: + return { + "construct_retrieval_results": { + "retrieval_results": [], + "has_calculated_field": False, + "has_metric": False, + "has_json_field": False, + "semantic_analysis": {}, + } + } + + async def _run_schema_retrieval( + self, + label: str, + *, + query: str, + project_id: Optional[str], + histories: Optional[list[AskHistory]] = None, + tables: Optional[list[str]] = None, + enable_column_pruning: bool = False, + timeout_seconds: Optional[int] = None, + query_id: Optional[str] = None, + ) -> dict[str, Any]: + try: + return await self._run_with_timeout( + label, + self._pipelines["db_schema_retrieval"].run( + query=query, + tables=tables, + project_id=project_id, + histories=histories or [], + enable_column_pruning=enable_column_pruning, + ), + timeout_seconds=timeout_seconds + or self._schema_retrieval_timeout_seconds, + ) + except TimeoutError as exc: + logger.warning( + "%s timed out; continuing without failing ask request. " + "query_id=%s project_id=%s tables=%s error=%s", + label, + query_id, + project_id, + tables, + exc, + ) + return self._empty_schema_retrieval_result() + def _should_retry_selected_schema_after_retrieval_timeout( self, retrieval_table_names: Optional[list[str]] ) -> bool: @@ -6092,19 +6140,18 @@ async def ask( is_followup=True if histories else False, general_type="DATA_ASSISTANCE", ) - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Metadata schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 20, + 60, ), + query_id=query_id, ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) @@ -6142,20 +6189,19 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 20, + 60, ), + query_id=query_id, ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) @@ -6298,14 +6344,13 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), + query=user_query, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + query_id=query_id, ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) @@ -6669,19 +6714,18 @@ async def ask( ) try: - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) - ), + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) ), timeout_seconds=self._schema_retrieval_timeout_seconds, + query_id=query_id, ) except TimeoutError as error: if not self._should_retry_selected_schema_after_retrieval_timeout( @@ -6704,19 +6748,18 @@ async def ask( retrieval_table_names, error, ) - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Selected schema fallback retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, timeout_seconds=min( self._schema_retrieval_timeout_seconds, - 30, + 60, ), + query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -6739,19 +6782,18 @@ async def ask( query_id, explicit_table_names, ) - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=enable_column_pruning, - ), + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, timeout_seconds=min( self._schema_retrieval_timeout_seconds, - 20, + 60, ), + query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -6776,19 +6818,18 @@ async def ask( "retrying full active deployed schema for query_id %s", query_id, ) - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Full active schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 20, + 60, ), + query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -6967,19 +7008,18 @@ async def ask( "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", query_id, ) - retrieval_result = await self._run_with_timeout( + retrieval_result = await self._run_schema_retrieval( "Full active schema retry", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 30, + 60, ), + query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index d93e456d2a..1f41eb1f3f 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -67,6 +67,65 @@ def __init__( self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval + async def _run_schema_retrieval( + self, + *, + query: str, + project_id: Optional[str], + histories: Optional[list] = None, + tables: Optional[list[str]] = None, + enable_column_pruning: bool = False, + timeout_seconds: int = 60, + ) -> dict: + try: + return await asyncio.wait_for( + self._pipelines["db_schema_retrieval"].run( + query=query, + tables=tables, + project_id=project_id, + histories=histories or [], + enable_column_pruning=enable_column_pruning, + ), + timeout=timeout_seconds, + ) + except TimeoutError as exc: + logger.warning( + "Question recommendation schema retrieval timed out; continuing with fallback context. " + "project_id=%s tables=%s error=%s", + project_id, + tables, + exc, + ) + return {"construct_retrieval_results": {"retrieval_results": []}} + + def _build_mdl_contexts(self, mdl: dict) -> list[str]: + contexts: list[str] = [] + for model in mdl.get("models", []): + name = model.get("name") + if not name: + continue + + columns = model.get("columns") or [] + column_lines = [] + for column in columns: + column_name = column.get("name") + if not column_name: + continue + column_type = ( + column.get("type") + or column.get("dataType") + or column.get("data_type") + or "TEXT" + ) + column_lines.append(f" {column_name} {column_type}") + + if column_lines: + contexts.append( + f"CREATE TABLE {name} (\n" + ",\n".join(column_lines) + "\n);" + ) + + return contexts + def _truncate_text(self, value: str, max_chars: int) -> str: if max_chars <= 0: return "" @@ -201,9 +260,10 @@ async def _validate_question( allow_data_preview: bool = True, ): async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: - retrieval_result = await self._pipelines["db_schema_retrieval"].run( + retrieval_result = await self._run_schema_retrieval( query=candidate["question"], project_id=project_id, + timeout_seconds=60, ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) @@ -396,12 +456,13 @@ async def recommend(self, input: Request, **kwargs) -> Event: trace_id = kwargs.get("trace_id") try: - orjson.loads(input.mdl) - retrieval_result = await self._pipelines["db_schema_retrieval"].run( + mdl = orjson.loads(input.mdl) + retrieval_result = await self._run_schema_retrieval( query="", histories=[], project_id=input.project_id, enable_column_pruning=False, + timeout_seconds=60, ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) @@ -410,6 +471,12 @@ async def recommend(self, input: Request, **kwargs) -> Event: max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, ) + if not table_ddls: + table_ddls = self._limit_text_items( + self._build_mdl_contexts(mdl), + max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, + max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, + ) request = { "contexts": table_ddls, From 6a60b5c64824ab50294d3ccd00a6f74d38bb11e4 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 05:22:33 +0530 Subject: [PATCH 0550/1087] Revert diagram properties parsing change --- wren-ui/src/apollo/server/resolvers/diagramResolver.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index 1223720f63..b457514555 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -287,13 +287,10 @@ export class DiagramResolver { }; } - private parseProperties(properties?: string | Record | null): Record { + private parseProperties(properties?: string | null): Record { if (!properties) { return {}; } - if (typeof properties === 'object') { - return properties; - } try { const parsed = JSON.parse(properties); return parsed && typeof parsed === 'object' ? parsed : {}; From d826214acd03f313162f9d6e533c0923fba3d06d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 06:08:29 +0530 Subject: [PATCH 0551/1087] Add Modeling AI Assistant workflows --- .../apollo/server/adaptors/wrenAIAdaptor.ts | 87 +++++ wren-ui/src/apollo/server/resolvers.ts | 4 + .../apollo/server/resolvers/modelResolver.ts | 53 +++ wren-ui/src/apollo/server/schema.ts | 9 + wren-ui/src/pages/modeling.tsx | 304 +++++++++++++++++- 5 files changed, 455 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 141c3b3015..5acc9af772 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -83,6 +83,19 @@ export interface IWrenAIAdaptor { queryId: string, ): Promise; + generateSemanticsDescription(input: { + manifest: any; + selectedModels: string[]; + userPrompt: string; + projectId: number; + }): Promise; + getSemanticsDescriptionResult(queryId: string): Promise; + generateRelationshipRecommendations(input: { + manifest: any; + projectId: number; + }): Promise; + getRelationshipRecommendationResult(queryId: string): Promise; + /** * Get text-based answer from SQL */ @@ -414,6 +427,80 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } } + public async generateSemanticsDescription(input: { + manifest: any; + selectedModels: string[]; + userPrompt: string; + projectId: number; + }): Promise { + try { + const res = await axios.post( + `${this.wrenAIBaseEndpoint}/v1/semantics-descriptions`, + { + mdl: JSON.stringify(input.manifest), + selected_models: input.selectedModels, + user_prompt: input.userPrompt, + project_id: String(input.projectId), + }, + ); + return { queryId: res.data.id }; + } catch (err: any) { + logger.debug( + `Got error when generating semantics descriptions: ${getAIServiceError(err)}`, + ); + throw err; + } + } + + public async getSemanticsDescriptionResult(queryId: string): Promise { + try { + const res = await axios.get( + `${this.wrenAIBaseEndpoint}/v1/semantics-descriptions/${queryId}`, + ); + return res.data; + } catch (err: any) { + logger.debug( + `Got error when getting semantics descriptions: ${getAIServiceError(err)}`, + ); + throw err; + } + } + + public async generateRelationshipRecommendations(input: { + manifest: any; + projectId: number; + }): Promise { + try { + const res = await axios.post( + `${this.wrenAIBaseEndpoint}/v1/relationship-recommendations`, + { + mdl: JSON.stringify(input.manifest), + project_id: String(input.projectId), + }, + ); + return { queryId: res.data.id }; + } catch (err: any) { + logger.debug( + `Got error when generating relationship recommendations: ${getAIServiceError(err)}`, + ); + throw err; + } + } + + public async getRelationshipRecommendationResult(queryId: string): Promise { + try { + const res = await axios.get( + `${this.wrenAIBaseEndpoint}/v1/relationship-recommendations/${queryId}`, + ); + return res.data; + } catch (err: any) { + logger.debug( + `Got error when getting relationship recommendations: ${getAIServiceError(err)}`, + ); + throw err; + } + } + public async createTextBasedAnswer( input: TextBasedAnswerInput, ): Promise { diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 8381952531..168bda10e1 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -57,6 +57,8 @@ const resolvers = { // Settings settings: projectResolver.getSettings, getMDL: modelResolver.getMDL, + modelingSemanticsResult: modelResolver.getModelingSemanticsResult, + modelingRelationshipsResult: modelResolver.getModelingRelationshipsResult, // Learning learningRecord: learningResolver.getLearningRecord, @@ -149,6 +151,8 @@ const resolvers = { previewViewData: modelResolver.previewViewData, validateView: modelResolver.validateView, updateViewMetadata: modelResolver.updateViewMetadata, + generateModelingSemantics: modelResolver.generateModelingSemantics, + generateModelingRelationships: modelResolver.generateModelingRelationships, // Settings resetCurrentProject: projectResolver.resetCurrentProject, diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index da7337453c..5fdfb334e2 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -57,6 +57,13 @@ export class ModelResolver { this.updateModelMetadata = this.updateModelMetadata.bind(this); this.deploy = this.deploy.bind(this); this.getMDL = this.getMDL.bind(this); + this.generateModelingSemantics = this.generateModelingSemantics.bind(this); + this.getModelingSemanticsResult = + this.getModelingSemanticsResult.bind(this); + this.generateModelingRelationships = + this.generateModelingRelationships.bind(this); + this.getModelingRelationshipsResult = + this.getModelingRelationshipsResult.bind(this); this.checkModelSync = this.checkModelSync.bind(this); // view @@ -447,6 +454,52 @@ export class ModelResolver { }; } + public async generateModelingSemantics( + _root: any, + args: { data: { selectedModels: string[]; userPrompt: string } }, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + return await ctx.wrenAIAdaptor.generateSemanticsDescription({ + manifest, + selectedModels: args.data.selectedModels, + userPrompt: args.data.userPrompt, + projectId: project.id, + }); + } + + public async getModelingSemanticsResult( + _root: any, + args: { queryId: string }, + ctx: IContext, + ) { + return await ctx.wrenAIAdaptor.getSemanticsDescriptionResult(args.queryId); + } + + public async generateModelingRelationships( + _root: any, + _args: any, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + return await ctx.wrenAIAdaptor.generateRelationshipRecommendations({ + manifest, + projectId: project.id, + }); + } + + public async getModelingRelationshipsResult( + _root: any, + args: { queryId: string }, + ctx: IContext, + ) { + return await ctx.wrenAIAdaptor.getRelationshipRecommendationResult( + args.queryId, + ); + } + public async listModels(_root: any, _args: any, ctx: IContext) { const { id: projectId } = await ctx.projectService.getCurrentProject(); const models = await ctx.modelRepository.findAllBy({ projectId }); diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 3caeed2d25..2700705cbd 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -368,6 +368,11 @@ export const typeDefs = gql` columns: [UpdateViewColumnMetadataInput!] } + input GenerateModelingSemanticsInput { + selectedModels: [String!]! + userPrompt: String! + } + type NestedFieldInfo { id: Int! displayName: String! @@ -1269,6 +1274,8 @@ export const typeDefs = gql` # System getMDL(hash: String!): GetMDLResult! + modelingSemanticsResult(queryId: String!): JSON! + modelingRelationshipsResult(queryId: String!): JSON! # Learning learningRecord: LearningRecord! @@ -1324,6 +1331,8 @@ export const typeDefs = gql` where: ViewWhereUniqueInput! data: UpdateViewMetadataInput! ): Boolean! + generateModelingSemantics(data: GenerateModelingSemanticsInput!): JSON! + generateModelingRelationships: JSON! # Relation createRelation(data: RelationInput!): JSON! diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 0cac6e2fd5..6f4f702798 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -1,8 +1,10 @@ import dynamic from 'next/dynamic'; import { useRouter } from 'next/router'; import { useSearchParams } from 'next/navigation'; -import { forwardRef, useEffect, useMemo, useRef } from 'react'; -import { message } from 'antd'; +import { forwardRef, useEffect, useMemo, useRef, useState } from 'react'; +import { Button, Checkbox, Dropdown, Input, Modal, Table, message } from 'antd'; +import { RobotOutlined } from '@ant-design/icons'; +import { gql, useApolloClient, useMutation } from '@apollo/client'; import styled from 'styled-components'; import { MORE_ACTION, NODE_TYPE } from '@/utils/enum'; import { editCalculatedField } from '@/utils/modelingHelper'; @@ -46,6 +48,30 @@ import { } from '@/apollo/client/graphql/relationship.generated'; import * as events from '@/utils/events'; +const GENERATE_MODELING_SEMANTICS = gql` + mutation GenerateModelingSemantics($data: GenerateModelingSemanticsInput!) { + generateModelingSemantics(data: $data) + } +`; + +const MODELING_SEMANTICS_RESULT = gql` + query ModelingSemanticsResult($queryId: String!) { + modelingSemanticsResult(queryId: $queryId) + } +`; + +const GENERATE_MODELING_RELATIONSHIPS = gql` + mutation GenerateModelingRelationships { + generateModelingRelationships + } +`; + +const MODELING_RELATIONSHIPS_RESULT = gql` + query ModelingRelationshipsResult($queryId: String!) { + modelingRelationshipsResult(queryId: $queryId) + } +`; + const Diagram = dynamic(() => import('@/components/diagram'), { ssr: false }); // https://github.com/vercel/next.js/issues/4957#issuecomment-413841689 const ForwardDiagram = forwardRef(function ForwardDiagram(props: any, ref) { @@ -57,10 +83,26 @@ const DiagramWrapper = styled.div` height: 100%; `; +const AssistantAction = styled.div` + position: absolute; + top: 16px; + right: 16px; + z-index: 10; +`; + export default function Modeling() { const router = useRouter(); const searchParams = useSearchParams(); + const apolloClient = useApolloClient(); const diagramRef = useRef(null); + const [assistantMode, setAssistantMode] = useState< + 'semantics' | 'relationships' | null + >(null); + const [assistantLoading, setAssistantLoading] = useState(false); + const [selectedModels, setSelectedModels] = useState([]); + const [semanticPrompt, setSemanticPrompt] = useState(''); + const [semanticResult, setSemanticResult] = useState([]); + const [relationshipResult, setRelationshipResult] = useState([]); const { data } = useDiagramQuery({ fetchPolicy: 'cache-and-network', @@ -200,6 +242,10 @@ export default function Modeling() { }, }), ); + const [generateModelingSemantics] = useMutation(GENERATE_MODELING_SEMANTICS); + const [generateModelingRelationships] = useMutation( + GENERATE_MODELING_RELATIONSHIPS, + ); const diagramData = useMemo(() => { if (!data) return null; @@ -378,6 +424,158 @@ export default function Modeling() { const modelLoading = modelCreating || modelUpdating; const relationshipLoading = relationshipUpdating || relationshipCreating; + const waitForAssistantResult = async ( + queryId: string, + query: any, + fieldName: string, + ) => { + for (let attempt = 0; attempt < 90; attempt += 1) { + const res = await apolloClient.query({ + query, + variables: { queryId }, + fetchPolicy: 'network-only', + }); + const payload = res.data?.[fieldName]; + if (payload?.status === 'finished') return payload.response || []; + if (payload?.status === 'failed') { + throw new Error(payload.error?.message || 'AI assistant failed.'); + } + await new Promise((resolve) => setTimeout(resolve, 2000)); + } + throw new Error('AI assistant timed out.'); + }; + + const openAssistant = (mode: 'semantics' | 'relationships') => { + setAssistantMode(mode); + setSelectedModels(diagramData?.models?.map((model) => model.referenceName) || []); + setSemanticResult([]); + setRelationshipResult([]); + }; + + const runAssistant = async () => { + try { + setAssistantLoading(true); + if (assistantMode === 'semantics') { + const res = await generateModelingSemantics({ + variables: { + data: { + selectedModels, + userPrompt: semanticPrompt || 'Describe this dataset for analytics.', + }, + }, + }); + const queryId = res.data?.generateModelingSemantics?.queryId; + const result = await waitForAssistantResult( + queryId, + MODELING_SEMANTICS_RESULT, + 'modelingSemanticsResult', + ); + setSemanticResult(result); + } + if (assistantMode === 'relationships') { + const res = await generateModelingRelationships(); + const queryId = res.data?.generateModelingRelationships?.queryId; + const result = await waitForAssistantResult( + queryId, + MODELING_RELATIONSHIPS_RESULT, + 'modelingRelationshipsResult', + ); + setRelationshipResult(result?.relationships || []); + } + } catch (error: any) { + message.error(error.message || 'Failed to run Modeling AI Assistant.'); + } finally { + setAssistantLoading(false); + } + }; + + const saveAssistantResult = async () => { + try { + if (!diagramData) return; + setAssistantLoading(true); + if (assistantMode === 'semantics') { + for (const model of semanticResult) { + const diagramModel = diagramData.models.find( + (item) => item.referenceName === model.name, + ); + if (!diagramModel) continue; + await updateModelMetadata({ + variables: { + where: { id: diagramModel.modelId }, + data: { + description: model.description, + columns: (model.columns || []) + .map((column) => { + const field = diagramModel.fields.find( + (item) => item.referenceName === column.name, + ); + return field + ? { + id: field.columnId, + displayName: field.displayName, + description: column.description, + } + : null; + }) + .filter(Boolean), + }, + }, + }); + } + } + if (assistantMode === 'relationships') { + for (const relationship of relationshipResult) { + const fromModel = diagramData.models.find( + (model) => model.referenceName === relationship.fromModel, + ); + const toModel = diagramData.models.find( + (model) => model.referenceName === relationship.toModel, + ); + const fromField = fromModel?.fields.find( + (field) => field.referenceName === relationship.fromColumn, + ); + const toField = toModel?.fields.find( + (field) => field.referenceName === relationship.toColumn, + ); + if (!fromModel || !toModel || !fromField || !toField) continue; + const alreadyExists = diagramData.models.some((model) => + model.relationFields.some((field) => { + const forward = + field.fromModelName === relationship.fromModel && + field.fromColumnName === relationship.fromColumn && + field.toModelName === relationship.toModel && + field.toColumnName === relationship.toColumn; + const reverse = + field.fromModelName === relationship.toModel && + field.fromColumnName === relationship.toColumn && + field.toModelName === relationship.fromModel && + field.toColumnName === relationship.fromColumn; + return forward || reverse; + }), + ); + if (alreadyExists) continue; + await createRelationshipMutation({ + variables: { + data: { + fromModelId: fromModel.modelId, + fromColumnId: fromField.columnId, + toModelId: toModel.modelId, + toColumnId: toField.columnId, + type: relationship.type, + }, + }, + }); + } + } + setAssistantMode(null); + message.success('Saved Modeling AI Assistant suggestions.'); + } catch (error: any) { + message.error(error.message || 'Failed to save assistant suggestions.'); + } finally { + setAssistantLoading(false); + } + }; + return ( + + + openAssistant(key as 'semantics' | 'relationships'), + }} + > + + + + setAssistantMode(null)} + > + {assistantMode === 'semantics' && ( + <> + ({ + label: model.displayName || model.referenceName, + value: model.referenceName, + }))} + onChange={(values) => setSelectedModels(values as string[])} + /> + setSemanticPrompt(event.target.value)} + /> + {!!semanticResult.length && ( +
( +
+ ), + }} + /> + )} + + )} + {assistantMode === 'relationships' && ( +
+ `${record.fromModel}.${record.fromColumn}-${record.toModel}.${record.toColumn}` + } + pagination={false} + dataSource={relationshipResult} + columns={[ + { + title: 'From', + render: (_value, record) => + `${record.fromModel}.${record.fromColumn}`, + }, + { + title: 'To', + render: (_value, record) => + `${record.toModel}.${record.toColumn}`, + }, + { title: 'Type', dataIndex: 'type', width: 150 }, + { title: 'Description', dataIndex: 'reason' }, + ]} + /> + )} + ); From 60631e098495ef180a4b58d9048438bab40ff732 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 06:16:20 +0530 Subject: [PATCH 0552/1087] Handle parsed diagram properties --- wren-ui/src/apollo/server/resolvers/diagramResolver.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index b457514555..f5f7ff1fb7 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -287,10 +287,15 @@ export class DiagramResolver { }; } - private parseProperties(properties?: string | null): Record { + private parseProperties( + properties?: string | Record | null, + ): Record { if (!properties) { return {}; } + if (typeof properties === 'object') { + return properties; + } try { const parsed = JSON.parse(properties); return parsed && typeof parsed === 'object' ? parsed : {}; From 6dc5967758ce6473011630050833301463c85524 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 10:20:25 +0530 Subject: [PATCH 0553/1087] Guard diagram manifest models --- wren-ui/src/apollo/server/resolvers/diagramResolver.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index f5f7ff1fb7..1dc4a7506f 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -89,7 +89,7 @@ export class DiagramResolver { const allColumns = modelColumns.filter( (column) => column.modelId === model.id, ); - const modelMDL = manifest.models.find( + const modelMDL = manifest.models?.find( (modelMDL) => modelMDL.name === model.referenceName, ); allColumns.forEach((column) => { From b0c42608b6d13df599a90a8a682770526bdcb1f1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 10:28:21 +0530 Subject: [PATCH 0554/1087] Improve semantic table retrieval coverage --- .../retrieval/db_schema_retrieval.py | 189 +++++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 39 ++++ 2 files changed, 223 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 5e24b26813..70fcdf3b92 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -225,6 +225,124 @@ def _retrieval_terms(value: str) -> set[str]: return {term for term in terms if term} +_BUSINESS_CONCEPT_TERMS: dict[str, set[str]] = { + "customer": { + "account", + "accounts", + "client", + "clients", + "cust", + "customer", + "customers", + }, + "order": { + "booking", + "bookings", + "ord", + "order", + "orders", + "purchase", + "transaction", + "transactions", + }, + "invoice": {"bill", "billing", "invoice", "invoices", "inv"}, + "sales": { + "amount", + "fxsales", + "margin", + "revenue", + "sale", + "sales", + "total", + "value", + }, + "product": { + "item", + "items", + "part", + "parts", + "product", + "products", + "sku", + }, + "quantity": {"qty", "quantity", "quantities", "unit", "units", "volume"}, + "market": { + "area", + "country", + "countries", + "domestic", + "international", + "intl", + "market", + "markets", + "mkt", + "region", + "territory", + }, + "time": { + "date", + "day", + "month", + "monthly", + "quarter", + "quarterly", + "time", + "week", + "year", + }, + "failure": { + "defect", + "failure", + "failures", + "issue", + "issues", + "pattern", + "patterns", + "problem", + "repair", + "status", + }, +} + + +def _business_concepts(value: str) -> set[str]: + terms = _retrieval_terms(value) + concepts: set[str] = set() + for concept, synonyms in _BUSINESS_CONCEPT_TERMS.items(): + if terms & {_normalize_retrieval_token(term) for term in synonyms}: + concepts.add(concept) + return concepts + + +def _concept_coverage_score( + query: str, document: Document +) -> tuple[int, set[str], set[str]]: + query_concepts = _business_concepts(query) + if not query_concepts: + return 0, set(), set() + + document_concepts = _business_concepts(_source_text(document)) + covered = query_concepts & document_concepts + missing = query_concepts - covered + score = 70 * len(covered) - 55 * len(missing) + + # Ranking/count questions are usually fact-style requests. Penalize reference + # tables that only contain a display name/id when the requested event concept + # is absent; this prevents unrelated tables from winning on generic columns. + normalized_query = (query or "").lower() + asks_for_ranked_count = bool( + re.search( + r"\b(?:top|highest|most|rank|ranking|number of|count)\b", + normalized_query, + ) + ) + if asks_for_ranked_count and {"order", "invoice", "sales"} & query_concepts: + if not ({"order", "invoice", "sales"} & document_concepts): + score -= 120 + + return score, covered, missing + + def _query_mentions_any(query: str, terms: tuple[str, ...]) -> bool: normalized = (query or "").lower() return any(re.search(rf"\b{re.escape(term)}\b", normalized) for term in terms) @@ -411,7 +529,20 @@ def _score_table_documents( lexical_score = _document_relevance_score(document, query_terms) semantic_score = _semantic_score(document) source_shape_score = _source_shape_score(query, document) - combined_score = semantic_score + lexical_score + source_shape_score + concept_score, covered_concepts, missing_concepts = _concept_coverage_score( + query, document + ) + combined_score = ( + semantic_score + lexical_score + source_shape_score + concept_score + ) + if covered_concepts or missing_concepts: + logger.debug( + "Table candidate concept coverage name=%s covered=%s missing=%s concept_score=%s", + document.meta.get("name"), + sorted(covered_concepts), + sorted(missing_concepts), + concept_score, + ) scored_documents.append( (combined_score, -index, document, lexical_score, semantic_score) ) @@ -457,7 +588,22 @@ def _select_relevant_table_documents( if not reranked: return documents[:max_tables] - candidate_pool = [item for item in reranked if item[3] > 0] or reranked + query_concepts = _business_concepts(query) + concept_pool = [ + item + for item in reranked + if not query_concepts + or (_business_concepts(_source_text(item[2])) & query_concepts) + ] + if query_concepts and not concept_pool: + logger.info( + "No table candidates covered requested business concepts; query=%s concepts=%s", + query, + sorted(query_concepts), + ) + return [] + + candidate_pool = concept_pool or [item for item in reranked if item[3] > 0] or reranked production_pool = [ item for item in candidate_pool @@ -550,10 +696,42 @@ def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: return normalized +def _table_name_lookup_variants(table_name: str) -> list[str]: + raw_name = str(table_name or "").strip().strip('"`[]') + if not raw_name: + return [] + + variants = [ + raw_name, + raw_name.replace(".", "_"), + raw_name.replace("$", "_"), + ] + parts = [part for part in re.split(r"[.$_\[\]`\"]+", raw_name) if part] + if len(parts) > 1: + variants.append(parts[-1]) + variants.append("_".join(parts[-2:])) + + deduped: list[str] = [] + for variant in variants: + variant = variant.strip() + if variant and variant not in deduped: + deduped.append(variant) + return deduped + + +def _expand_table_name_lookup_variants(table_names: Optional[list[str]]) -> list[str]: + expanded: list[str] = [] + for table_name in _normalize_table_names(table_names): + for variant in _table_name_lookup_variants(table_name): + if variant not in expanded: + expanded.append(variant) + return expanded + + def _extract_table_names_from_table_retrieval( table_retrieval: dict, explicit_tables: Optional[list[str]] = None ) -> list[str]: - table_names = _normalize_table_names(explicit_tables) + table_names = _expand_table_name_lookup_variants(explicit_tables) for document in table_retrieval.get("documents") or []: if not isinstance(document, Document): continue @@ -627,12 +805,13 @@ async def table_retrieval( return results if tables: - logger.info("Loading explicit table descriptions: %s", tables) + explicit_table_names = _expand_table_name_lookup_variants(tables) + logger.info("Loading explicit table descriptions: %s", explicit_table_names) explicit_filters = { **base_filters, "conditions": [ *base_filters["conditions"], - {"field": "name", "operator": "in", "value": tables}, + {"field": "name", "operator": "in", "value": explicit_table_names}, ], } return await table_retriever.run(query_embedding=[], filters=explicit_filters) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 37187263f9..3c5edbed46 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -665,6 +665,45 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: return [] concept_groups: list[set[str]] = [] + if re.search(r"\bcustomers?\b|\bclients?\b|\baccounts?\b|\bcust\b", normalized): + concept_groups.append( + {"account", "client", "cust", "customer", "customers", "name"} + ) + if re.search(r"\borders?\b|\border\s+count\b|\bbookings?\b", normalized): + concept_groups.append( + { + "booking", + "bookings", + "ord", + "order", + "orders", + "purchase", + "transaction", + } + ) + if re.search(r"\binvoices?\b|\bbilling\b|\bbills?\b", normalized): + concept_groups.append({"bill", "billing", "invoice", "invoices", "inv"}) + if re.search(r"\bsales?\b|\brevenue\b|\bamount\b|\bvalue\b", normalized): + concept_groups.append( + { + "amount", + "fxsales", + "margin", + "revenue", + "sale", + "sales", + "total", + "value", + } + ) + if re.search(r"\bproducts?\b|\bitems?\b|\bparts?\b|\bsku\b", normalized): + concept_groups.append( + {"item", "items", "part", "parts", "product", "products", "sku"} + ) + if re.search(r"\bquantity\b|\bquantities\b|\bqty\b|\bunits?\b", normalized): + concept_groups.append( + {"qty", "quantity", "quantities", "unit", "units", "volume"} + ) if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: From 0932befaaf3b6c5e4f041adb754c25cdc5982513 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 10:29:16 +0530 Subject: [PATCH 0555/1087] Bind diagram view transformer --- wren-ui/src/apollo/server/resolvers/diagramResolver.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index 1dc4a7506f..bc3bafd05f 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -132,7 +132,7 @@ export class DiagramResolver { return transformedModel; }); - const diagramViews = views.map(this.transformView); + const diagramViews = views.map((view) => this.transformView(view)); return { models: diagramModels, views: diagramViews }; } From 494fce0c05ce8eaff095580749d138a4de3bcb04 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 12:36:45 +0530 Subject: [PATCH 0556/1087] Fix Modeling AI Assistant compatibility regressions --- wren-ui/src/pages/modeling.tsx | 44 ++++++++++++++++++++++++++-------- 1 file changed, 34 insertions(+), 10 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 6f4f702798..cadbea22fe 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -2,7 +2,16 @@ import dynamic from 'next/dynamic'; import { useRouter } from 'next/router'; import { useSearchParams } from 'next/navigation'; import { forwardRef, useEffect, useMemo, useRef, useState } from 'react'; -import { Button, Checkbox, Dropdown, Input, Modal, Table, message } from 'antd'; +import { + Button, + Checkbox, + Dropdown, + Input, + Menu, + Modal, + Table, + message, +} from 'antd'; import { RobotOutlined } from '@ant-design/icons'; import { gql, useApolloClient, useMutation } from '@apollo/client'; import styled from 'styled-components'; @@ -429,6 +438,10 @@ export default function Modeling() { query: any, fieldName: string, ) => { + if (!queryId) { + throw new Error('AI assistant did not return a task id.'); + } + for (let attempt = 0; attempt < 90; attempt += 1) { const res = await apolloClient.query({ query, @@ -456,6 +469,9 @@ export default function Modeling() { try { setAssistantLoading(true); if (assistantMode === 'semantics') { + if (!selectedModels.length) { + throw new Error('Select at least one model.'); + } const res = await generateModelingSemantics({ variables: { data: { @@ -473,6 +489,9 @@ export default function Modeling() { setSemanticResult(result); } if (assistantMode === 'relationships') { + if (!diagramData?.models || diagramData.models.length < 2) { + throw new Error('At least two models are required.'); + } const res = await generateModelingRelationships(); const queryId = res.data?.generateModelingRelationships?.queryId; const result = await waitForAssistantResult( @@ -539,7 +558,8 @@ export default function Modeling() { ); if (!fromModel || !toModel || !fromField || !toField) continue; const alreadyExists = diagramData.models.some((model) => - model.relationFields.some((field) => { + (model.relationFields || []).some((field) => { + if (!field) return false; const forward = field.fromModelName === relationship.fromModel && field.fromColumnName === relationship.fromColumn && @@ -590,14 +610,18 @@ export default function Modeling() { - openAssistant(key as 'semantics' | 'relationships'), - }} + overlay={ + + openAssistant(key as 'semantics' | 'relationships') + } + > + Generate semantics + + Generate relationships + + + } > From 5fa32c4fa3d6edbe5eb2b42cec4719bb7d9ae751 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 12:53:52 +0530 Subject: [PATCH 0557/1087] Restore working ask query behavior --- .../retrieval/db_schema_retrieval.py | 249 +------ wren-ai-service/src/web/v1/services/ask.py | 683 ++++-------------- .../v1/services/question_recommendation.py | 73 +- .../src/web/v1/services/sql_answer.py | 7 +- 4 files changed, 179 insertions(+), 833 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 70fcdf3b92..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -213,136 +213,14 @@ def _retrieval_terms(value: str) -> set[str]: "which", "with", } - terms: set[str] = set() - for raw_token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or ""): - split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) - for token in re.findall(r"[A-Za-z0-9]+", split_token): - if len(token) <= 2 or token.lower() in stop_words: - continue - normalized_token = _normalize_retrieval_token(token) - if normalized_token: - terms.add(normalized_token) + terms = { + _normalize_retrieval_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") + if len(token) > 2 and token.lower() not in stop_words + } return {term for term in terms if term} -_BUSINESS_CONCEPT_TERMS: dict[str, set[str]] = { - "customer": { - "account", - "accounts", - "client", - "clients", - "cust", - "customer", - "customers", - }, - "order": { - "booking", - "bookings", - "ord", - "order", - "orders", - "purchase", - "transaction", - "transactions", - }, - "invoice": {"bill", "billing", "invoice", "invoices", "inv"}, - "sales": { - "amount", - "fxsales", - "margin", - "revenue", - "sale", - "sales", - "total", - "value", - }, - "product": { - "item", - "items", - "part", - "parts", - "product", - "products", - "sku", - }, - "quantity": {"qty", "quantity", "quantities", "unit", "units", "volume"}, - "market": { - "area", - "country", - "countries", - "domestic", - "international", - "intl", - "market", - "markets", - "mkt", - "region", - "territory", - }, - "time": { - "date", - "day", - "month", - "monthly", - "quarter", - "quarterly", - "time", - "week", - "year", - }, - "failure": { - "defect", - "failure", - "failures", - "issue", - "issues", - "pattern", - "patterns", - "problem", - "repair", - "status", - }, -} - - -def _business_concepts(value: str) -> set[str]: - terms = _retrieval_terms(value) - concepts: set[str] = set() - for concept, synonyms in _BUSINESS_CONCEPT_TERMS.items(): - if terms & {_normalize_retrieval_token(term) for term in synonyms}: - concepts.add(concept) - return concepts - - -def _concept_coverage_score( - query: str, document: Document -) -> tuple[int, set[str], set[str]]: - query_concepts = _business_concepts(query) - if not query_concepts: - return 0, set(), set() - - document_concepts = _business_concepts(_source_text(document)) - covered = query_concepts & document_concepts - missing = query_concepts - covered - score = 70 * len(covered) - 55 * len(missing) - - # Ranking/count questions are usually fact-style requests. Penalize reference - # tables that only contain a display name/id when the requested event concept - # is absent; this prevents unrelated tables from winning on generic columns. - normalized_query = (query or "").lower() - asks_for_ranked_count = bool( - re.search( - r"\b(?:top|highest|most|rank|ranking|number of|count)\b", - normalized_query, - ) - ) - if asks_for_ranked_count and {"order", "invoice", "sales"} & query_concepts: - if not ({"order", "invoice", "sales"} & document_concepts): - score -= 120 - - return score, covered, missing - - def _query_mentions_any(query: str, terms: tuple[str, ...]) -> bool: normalized = (query or "").lower() return any(re.search(rf"\b{re.escape(term)}\b", normalized) for term in terms) @@ -365,11 +243,7 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - weak_non_production_terms = ( - "stage", - "staging", - ) - strong_non_production_terms = ( + non_production_terms = ( "archive", "backup", "copy", @@ -377,18 +251,17 @@ def _source_shape_score(query: str, document: Document) -> int: "development", "duplicate", "sample", + "stage", + "staging", "temp", "test", "tmp", ) - if source_terms & set(strong_non_production_terms) and not _query_mentions_any( - normalized_query, strong_non_production_terms - ): - score -= 240 - if source_terms & set(weak_non_production_terms) and not _query_mentions_any( - normalized_query, weak_non_production_terms + if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( + normalized_query, + non_production_terms, ): - score -= 40 + score -= 60 aggregation_terms = ( "amount", @@ -452,26 +325,6 @@ def _source_shape_score(query: str, document: Document) -> int: return score -def _is_unrequested_strong_non_production_source(query: str, document: Document) -> bool: - strong_non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "temp", - "test", - "tmp", - ) - source_terms = _retrieval_terms(_source_text(document)) - return bool( - source_terms & set(strong_non_production_terms) - and not _query_mentions_any(query or "", strong_non_production_terms) - ) - - def _document_relevance_score(document: Document, query_terms: set[str]) -> int: if not query_terms: return 0 @@ -529,20 +382,7 @@ def _score_table_documents( lexical_score = _document_relevance_score(document, query_terms) semantic_score = _semantic_score(document) source_shape_score = _source_shape_score(query, document) - concept_score, covered_concepts, missing_concepts = _concept_coverage_score( - query, document - ) - combined_score = ( - semantic_score + lexical_score + source_shape_score + concept_score - ) - if covered_concepts or missing_concepts: - logger.debug( - "Table candidate concept coverage name=%s covered=%s missing=%s concept_score=%s", - document.meta.get("name"), - sorted(covered_concepts), - sorted(missing_concepts), - concept_score, - ) + combined_score = semantic_score + lexical_score + source_shape_score scored_documents.append( (combined_score, -index, document, lexical_score, semantic_score) ) @@ -588,29 +428,7 @@ def _select_relevant_table_documents( if not reranked: return documents[:max_tables] - query_concepts = _business_concepts(query) - concept_pool = [ - item - for item in reranked - if not query_concepts - or (_business_concepts(_source_text(item[2])) & query_concepts) - ] - if query_concepts and not concept_pool: - logger.info( - "No table candidates covered requested business concepts; query=%s concepts=%s", - query, - sorted(query_concepts), - ) - return [] - - candidate_pool = concept_pool or [item for item in reranked if item[3] > 0] or reranked - production_pool = [ - item - for item in candidate_pool - if not _is_unrequested_strong_non_production_source(query, item[2]) - ] - if production_pool: - candidate_pool = production_pool + candidate_pool = [item for item in reranked if item[3] > 0] or reranked selected = [ document for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] @@ -696,42 +514,10 @@ def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: return normalized -def _table_name_lookup_variants(table_name: str) -> list[str]: - raw_name = str(table_name or "").strip().strip('"`[]') - if not raw_name: - return [] - - variants = [ - raw_name, - raw_name.replace(".", "_"), - raw_name.replace("$", "_"), - ] - parts = [part for part in re.split(r"[.$_\[\]`\"]+", raw_name) if part] - if len(parts) > 1: - variants.append(parts[-1]) - variants.append("_".join(parts[-2:])) - - deduped: list[str] = [] - for variant in variants: - variant = variant.strip() - if variant and variant not in deduped: - deduped.append(variant) - return deduped - - -def _expand_table_name_lookup_variants(table_names: Optional[list[str]]) -> list[str]: - expanded: list[str] = [] - for table_name in _normalize_table_names(table_names): - for variant in _table_name_lookup_variants(table_name): - if variant not in expanded: - expanded.append(variant) - return expanded - - def _extract_table_names_from_table_retrieval( table_retrieval: dict, explicit_tables: Optional[list[str]] = None ) -> list[str]: - table_names = _expand_table_name_lookup_variants(explicit_tables) + table_names = _normalize_table_names(explicit_tables) for document in table_retrieval.get("documents") or []: if not isinstance(document, Document): continue @@ -805,13 +591,12 @@ async def table_retrieval( return results if tables: - explicit_table_names = _expand_table_name_lookup_variants(tables) - logger.info("Loading explicit table descriptions: %s", explicit_table_names) + logger.info("Loading explicit table descriptions: %s", tables) explicit_filters = { **base_filters, "conditions": [ *base_filters["conditions"], - {"field": "name", "operator": "in", "value": explicit_table_names}, + {"field": "name", "operator": "in", "value": tables}, ], } return await table_retriever.run(query_embedding=[], filters=explicit_filters) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3c5edbed46..13548f9620 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -665,45 +665,6 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: return [] concept_groups: list[set[str]] = [] - if re.search(r"\bcustomers?\b|\bclients?\b|\baccounts?\b|\bcust\b", normalized): - concept_groups.append( - {"account", "client", "cust", "customer", "customers", "name"} - ) - if re.search(r"\borders?\b|\border\s+count\b|\bbookings?\b", normalized): - concept_groups.append( - { - "booking", - "bookings", - "ord", - "order", - "orders", - "purchase", - "transaction", - } - ) - if re.search(r"\binvoices?\b|\bbilling\b|\bbills?\b", normalized): - concept_groups.append({"bill", "billing", "invoice", "invoices", "inv"}) - if re.search(r"\bsales?\b|\brevenue\b|\bamount\b|\bvalue\b", normalized): - concept_groups.append( - { - "amount", - "fxsales", - "margin", - "revenue", - "sale", - "sales", - "total", - "value", - } - ) - if re.search(r"\bproducts?\b|\bitems?\b|\bparts?\b|\bsku\b", normalized): - concept_groups.append( - {"item", "items", "part", "parts", "product", "products", "sku"} - ) - if re.search(r"\bquantity\b|\bquantities\b|\bqty\b|\bunits?\b", normalized): - concept_groups.append( - {"qty", "quantity", "quantities", "unit", "units", "volume"} - ) if "product line" in normalized or "productline" in normalized: concept_groups.append({"product", "prod", "line", "productline"}) if "pcb" in normalized: @@ -728,8 +689,6 @@ def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: concept_groups.append({"market", "region", "country", "territory"}) if "region" in normalized or "regions" in normalized: concept_groups.append({"region", "market", "area", "territory", "country"}) - if "country" in normalized or "countries" in normalized: - concept_groups.append({"country", "countries", "nation", "destination"}) if "quarterly" in normalized or "quarter" in normalized: concept_groups.append({"quarter", "quarterly"}) if "recurring" in normalized or "recurrence" in normalized: @@ -809,72 +768,6 @@ def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> return False return True - def _sql_satisfies_count_ranking_request( - self, sql: str, query: str | None - ) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return True - - asks_for_count_metric = any( - term in normalized_query - for term in ( - "count", - "counts", - "how many", - "number of", - "record count", - ) - ) or ( - any(term in normalized_query for term in ("top", "most", "highest")) - and any( - term in normalized_query - for term in ("order", "orders", "record", "records", "row", "rows") - ) - ) - if not asks_for_count_metric: - return True - - asks_for_grouped_entity = any( - term in normalized_query - for term in ( - "category", - "customer", - "customers", - "currency", - "market", - "product", - "products", - "region", - "sales person", - "salesperson", - "source", - "status", - "type", - ) - ) - if not asks_for_grouped_entity: - return True - - normalized_sql = re.sub(r"\s+", " ", sql or "").lower() - if not re.search(r"\bcount\s*\(", normalized_sql): - logger.warning( - "Ignoring SQL because a count/ranking question was answered with detail rows. " - "query=%s sql=%s", - query, - sql, - ) - return False - if not re.search(r"\bgroup\s+by\b", normalized_sql): - logger.warning( - "Ignoring SQL because a grouped count/ranking question has no GROUP BY. " - "query=%s sql=%s", - query, - sql, - ) - return False - return True - def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> bool: normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not re.search( @@ -950,167 +843,6 @@ def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> b return False return True - def _extract_entity_lookup_phrase(self, query: str | None) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip()) - if not normalized_query: - return None - - match = re.search( - r"\b(?:show|list|find|get|display)\b.*?\b(?:orders?|records?|rows?)\b\s+" - r"(?:for|where|with)\s+(?P.+?)(?:[?.!]|$)", - normalized_query, - flags=re.IGNORECASE, - ) - if not match: - return None - - phrase = match.group("phrase").strip(" .,;:()[]{}'\"") - phrase = re.sub(r"^(?:customer|client|account|company|name)\s+", "", phrase, flags=re.IGNORECASE) - if not phrase or len(phrase) < 3: - return None - if re.search( - r"\b(?:table|model|schema|column|columns|market|region|country|division|" - r"date|month|year|quarter|top|count|number|amount|value)\b", - phrase, - flags=re.IGNORECASE, - ): - return None - return phrase - - def _preferred_entity_lookup_columns( - self, query: str | None, table: dict[str, Any] - ) -> set[str]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - candidate_groups: list[tuple[str, ...]] = [] - if "account" in normalized_query: - candidate_groups.append(("Account", "AccountName", "AcctName", "AcctNo")) - if "company" in normalized_query: - candidate_groups.append(("Company", "CompanyName", "CustName", "CustomerName")) - candidate_groups.extend( - [ - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - ), - ("Client", "ClientName"), - ("Account", "AccountName"), - ("Company", "CompanyName"), - ("Name",), - ] - ) - - columns: set[str] = set() - for candidates in candidate_groups: - column = self._find_schema_column(table, candidates) - if column: - columns.add(column) - return columns - - def _sql_satisfies_entity_lookup_request( - self, - sql: str, - query: str | None, - referenced_tables: list[str], - referenced_columns_by_table: dict[str, set[str]], - valid_tables: dict[str, dict[str, Any]], - ) -> bool: - if not self._extract_entity_lookup_phrase(query): - return True - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if any( - term in normalized_query - for term in (" by ", " per ", " each ", "distribution", "top", "count") - ): - return True - - for table_reference in referenced_tables: - table = self._table_for_sql_reference(table_reference, valid_tables) - if not table: - continue - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - - table_key = str(table_reference or "").lower() - referenced_columns = referenced_columns_by_table.get( - table_key - ) or referenced_columns_by_table.get( - table_key.split(".")[-1], - set(), - ) - referenced_column_keys = { - self._normalize_schema_identifier_key(column) - for column in referenced_columns - } - preferred_column_keys = { - self._normalize_schema_identifier_key(column) - for column in preferred_columns - } - if referenced_column_keys & preferred_column_keys: - return True - - logger.warning( - "Ignoring SQL because entity lookup did not use available customer/name columns. " - "query=%s table=%s preferred_columns=%s referenced_columns=%s sql=%s", - query, - table.get("name"), - sorted(preferred_columns), - sorted(referenced_columns), - sql, - ) - return False - - return True - - def _is_unrequested_non_production_table_reference( - self, table_name: str, query: str | None - ) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_table = self._normalize_schema_token(table_name) - strong_non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "temp", - "test", - "tmp", - ) - if not any(term in normalized_table for term in strong_non_production_terms): - return False - return not any( - re.search(rf"\b{re.escape(term)}\b", normalized_query) - for term in strong_non_production_terms - ) - - def _sql_avoids_unrequested_non_production_tables( - self, sql: str, query: str | None, referenced_tables: list[str] - ) -> bool: - invalid_tables = [ - table - for table in referenced_tables - if self._is_unrequested_non_production_table_reference(table, query) - ] - if not invalid_tables: - return True - - logger.warning( - "Ignoring SQL because it references unrequested non-production tables. " - "query=%s invalid_tables=%s sql=%s", - query, - invalid_tables, - sql, - ) - return False - def _invalid_unqualified_sql_identifiers( self, sql: str, schema_tables: list[dict[str, Any]] ) -> list[str]: @@ -1420,24 +1152,8 @@ def _sql_matches_question_intent( return False if not self._sql_uses_required_measure_aggregation(sql, query): return False - if not self._sql_satisfies_count_ranking_request(sql, query): - return False if not self._sql_satisfies_unique_entity_request(sql, query): return False - if not self._sql_satisfies_entity_lookup_request( - sql, - query, - referenced_tables, - referenced_columns_by_table, - valid_tables, - ): - return False - if not self._sql_avoids_unrequested_non_production_tables( - sql, - query, - referenced_tables, - ): - return False if not expects_dimension: return True @@ -2078,18 +1794,7 @@ def _explicit_table_name_candidates(self, table_name: str) -> list[str]: separator_normalized = re.sub(r"[.$]", "_", table_name) if separator_normalized not in candidates: candidates.append(separator_normalized) - if "_" in table_name: - dotted_schema_name = re.sub( - r"^([A-Za-z_][A-Za-z0-9]*)_", - r"\1.", - table_name, - count=1, - ) - if dotted_schema_name not in candidates: - candidates.append(dotted_schema_name) short_name = re.split(r"[.$]", table_name)[-1] - if short_name == table_name and "_" in table_name: - short_name = table_name.split("_", 1)[-1] if short_name and short_name not in candidates: candidates.append(short_name) return candidates @@ -2442,88 +2147,6 @@ def _select_best_analytics_table( date_column, ) - def _build_entity_lookup_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - lookup_phrase = self._extract_entity_lookup_phrase(query) - if not lookup_phrase: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - asks_for_orders = any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ) - - scored: list[tuple[int, dict[str, Any], str]] = [] - for table in tables: - table_name = str(table.get("name") or "") - if not table_name: - continue - - preferred_columns = self._preferred_entity_lookup_columns(query, table) - if not preferred_columns: - continue - - preferred_column = sorted( - preferred_columns, - key=lambda column: ( - 0 - if self._normalize_schema_identifier_key(column) - in {"custname", "customername", "customer"} - else 1, - column.lower(), - ), - )[0] - - score = 20 - normalized_table = self._normalize_schema_token(table_name) - if asks_for_orders: - if "order" in normalized_table: - score += 40 - if "neworder" in normalized_table: - score += 20 - if self._find_schema_column( - table, ("OrdNo", "OrderNo", "OrderId", "NewOrderId") - ): - score += 25 - if "test" in normalized_table or "tmp" in normalized_table: - score -= 80 - if "dev" in normalized_table or "backup" in normalized_table: - score -= 60 - if "stage" in normalized_table: - score -= 10 - scored.append((score, table, preferred_column)) - - if not scored: - return None - - _score, table, filter_column = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - table_name = str(table.get("name") or "") - if not table_name: - return None - - table_ref = self._quote_sql_identifier(table_name) - filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" - escaped_phrase = lookup_phrase.replace("'", "''") - date_column = self._find_schema_column( - table, - ("OrdDate", "OrderDate", "NewOrderDate", "InvDate", "InvoiceDate", "Date"), - temporal=True, - ) - order_clause = ( - f" ORDER BY {table_ref}.{self._quote_sql_identifier(date_column)} DESC" - if date_column - else "" - ) - return ( - f"SELECT TOP 500 * FROM {table_ref} " - f"WHERE {filter_ref} = '{escaped_phrase}'" - f"{order_clause}" - ) - def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: @@ -2537,9 +2160,6 @@ def _build_schema_grounded_analytics_sql( compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - if entity_lookup_sql := self._build_entity_lookup_sql(query, tables): - return entity_lookup_sql - if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): return pcb_direct_sql @@ -2595,7 +2215,16 @@ def _build_schema_grounded_analytics_sql( if contribution_sql := self._build_contribution_sql(query, tables): return contribution_sql - asks_for_measure_value = any( + if not is_sales_or_order_query: + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + + wants_count_metric = any( + term in normalized_query + for term in ("count", "counts", "volume", "how many", "distribution") + ) and not any( term in normalized_query for term in ( "amount", @@ -2612,17 +2241,6 @@ def _build_schema_grounded_analytics_sql( "value", ) ) - - if not is_sales_or_order_query and not asks_for_measure_value: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not asks_for_measure_value wants_average_metric = any( term in normalized_query for term in ("average", "avg", "mean") ) @@ -4300,18 +3918,6 @@ def _build_heuristic_text_to_sql_fallback( return None - def _can_use_schema_grounded_sql_fallback( - self, - documents: list[dict], - table_ddls: list[str], - query: str | None, - ) -> bool: - return bool( - documents - and table_ddls - and not self._should_load_full_schema_for_question(query) - ) - def _is_schema_grounded_query( self, query: str, db_schemas: Optional[list[str]] = None ) -> bool: @@ -5105,54 +4711,6 @@ async def _run_with_timeout( except TimeoutError as exc: raise TimeoutError(f"{label} timed out after {timeout} seconds") from exc - def _empty_schema_retrieval_result(self) -> dict[str, Any]: - return { - "construct_retrieval_results": { - "retrieval_results": [], - "has_calculated_field": False, - "has_metric": False, - "has_json_field": False, - "semantic_analysis": {}, - } - } - - async def _run_schema_retrieval( - self, - label: str, - *, - query: str, - project_id: Optional[str], - histories: Optional[list[AskHistory]] = None, - tables: Optional[list[str]] = None, - enable_column_pruning: bool = False, - timeout_seconds: Optional[int] = None, - query_id: Optional[str] = None, - ) -> dict[str, Any]: - try: - return await self._run_with_timeout( - label, - self._pipelines["db_schema_retrieval"].run( - query=query, - tables=tables, - project_id=project_id, - histories=histories or [], - enable_column_pruning=enable_column_pruning, - ), - timeout_seconds=timeout_seconds - or self._schema_retrieval_timeout_seconds, - ) - except TimeoutError as exc: - logger.warning( - "%s timed out; continuing without failing ask request. " - "query_id=%s project_id=%s tables=%s error=%s", - label, - query_id, - project_id, - tables, - exc, - ) - return self._empty_schema_retrieval_result() - def _should_retry_selected_schema_after_retrieval_timeout( self, retrieval_table_names: Optional[list[str]] ) -> bool: @@ -5360,9 +4918,6 @@ def _get_metadata_question_kind(self, query: str) -> str | None: return None - def _should_load_full_schema_for_question(self, query: str | None) -> bool: - return bool(self._get_metadata_question_kind(query or "")) - def _find_metadata_table_matches( self, query: str, tables: list[dict[str, Any]] ) -> list[dict[str, Any]]: @@ -6179,18 +5734,19 @@ async def ask( is_followup=True if histories else False, general_type="DATA_ASSISTANCE", ) - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Metadata schema retrieval", - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 60, + 20, ), - query_id=query_id, ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) @@ -6228,19 +5784,20 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Explicit table schema retrieval", - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, + self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 60, + 20, ), - query_id=query_id, ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) @@ -6252,12 +5809,36 @@ async def ask( explicit_table_names, ) ) - if not documents: + if not documents and not request_explicit_table_names: logger.info( "Explicit table retrieval did not return requested active-schema table; " - "not loading full active schema for data question. query_id=%s", + "loading full active schema. query_id=%s", query_id, ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval for explicit table", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + all_documents, _, _ = self._extract_retrieval_metadata( + retrieval_result + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + all_documents, + explicit_table_names, + ) + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) @@ -6383,13 +5964,14 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Schema retrieval", - query=user_query, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - query_id=query_id, + self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) @@ -6753,18 +6335,19 @@ async def ask( ) try: - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Schema retrieval", - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), ), timeout_seconds=self._schema_retrieval_timeout_seconds, - query_id=query_id, ) except TimeoutError as error: if not self._should_retry_selected_schema_after_retrieval_timeout( @@ -6787,18 +6370,19 @@ async def ask( retrieval_table_names, error, ) - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Selected schema fallback retrieval", - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, - 60, + 30, ), - query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -6821,18 +6405,19 @@ async def ask( query_id, explicit_table_names, ) - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Explicit table schema retrieval", - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, + self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=enable_column_pruning, + ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, - 60, + 20, ), - query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -6849,7 +6434,7 @@ async def ask( ) if ( not documents - and self._should_load_full_schema_for_question(user_query) + and self._get_metadata_question_kind(user_query) and not request_explicit_table_names ): logger.info( @@ -6857,18 +6442,19 @@ async def ask( "retrying full active deployed schema for query_id %s", query_id, ) - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Full active schema retrieval", - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 60, + 20, ), - query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -7037,7 +6623,7 @@ async def ask( should_retry_full_schema = ( not api_results - and self._should_load_full_schema_for_question(user_query) + and self._get_metadata_question_kind(user_query) and "db_schema_retrieval" in self._pipelines and not request_explicit_table_names and not table_names @@ -7047,18 +6633,19 @@ async def ask( "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", query_id, ) - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._run_with_timeout( "Full active schema retry", - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), timeout_seconds=min( self._schema_retrieval_timeout_seconds, self._pipeline_timeout_seconds, - 60, + 30, ), - query_id=query_id, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -7149,6 +6736,54 @@ async def ask( return results if not documents: + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", + query_id, + user_query, + ) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + is_followup=True if histories else False, + ) + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( @@ -7501,14 +7136,8 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if self._can_use_schema_grounded_sql_fallback( - documents, - table_ddls, - user_query, - ) and ( - heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ) + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 1f41eb1f3f..d93e456d2a 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -67,65 +67,6 @@ def __init__( self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval - async def _run_schema_retrieval( - self, - *, - query: str, - project_id: Optional[str], - histories: Optional[list] = None, - tables: Optional[list[str]] = None, - enable_column_pruning: bool = False, - timeout_seconds: int = 60, - ) -> dict: - try: - return await asyncio.wait_for( - self._pipelines["db_schema_retrieval"].run( - query=query, - tables=tables, - project_id=project_id, - histories=histories or [], - enable_column_pruning=enable_column_pruning, - ), - timeout=timeout_seconds, - ) - except TimeoutError as exc: - logger.warning( - "Question recommendation schema retrieval timed out; continuing with fallback context. " - "project_id=%s tables=%s error=%s", - project_id, - tables, - exc, - ) - return {"construct_retrieval_results": {"retrieval_results": []}} - - def _build_mdl_contexts(self, mdl: dict) -> list[str]: - contexts: list[str] = [] - for model in mdl.get("models", []): - name = model.get("name") - if not name: - continue - - columns = model.get("columns") or [] - column_lines = [] - for column in columns: - column_name = column.get("name") - if not column_name: - continue - column_type = ( - column.get("type") - or column.get("dataType") - or column.get("data_type") - or "TEXT" - ) - column_lines.append(f" {column_name} {column_type}") - - if column_lines: - contexts.append( - f"CREATE TABLE {name} (\n" + ",\n".join(column_lines) + "\n);" - ) - - return contexts - def _truncate_text(self, value: str, max_chars: int) -> str: if max_chars <= 0: return "" @@ -260,10 +201,9 @@ async def _validate_question( allow_data_preview: bool = True, ): async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: - retrieval_result = await self._run_schema_retrieval( + retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, - timeout_seconds=60, ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) @@ -456,13 +396,12 @@ async def recommend(self, input: Request, **kwargs) -> Event: trace_id = kwargs.get("trace_id") try: - mdl = orjson.loads(input.mdl) - retrieval_result = await self._run_schema_retrieval( + orjson.loads(input.mdl) + retrieval_result = await self._pipelines["db_schema_retrieval"].run( query="", histories=[], project_id=input.project_id, enable_column_pruning=False, - timeout_seconds=60, ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) @@ -471,12 +410,6 @@ async def recommend(self, input: Request, **kwargs) -> Event: max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, ) - if not table_ddls: - table_ddls = self._limit_text_items( - self._build_mdl_contexts(mdl), - max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, - max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, - ) request = { "contexts": table_ddls, diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index 42eefc478a..7442f2f884 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -66,14 +66,14 @@ def __init__( ) async def _load_active_schema_contexts( - self, project_id: Optional[str], query: str + self, project_id: Optional[str] ) -> list[str]: retrieval_pipeline = self._pipelines.get("db_schema_retrieval") - if not retrieval_pipeline or not (query or "").strip(): + if not retrieval_pipeline: return [] retrieval_result = await retrieval_pipeline.run( - query=query, + query="", histories=[], project_id=project_id, enable_column_pruning=False, @@ -134,7 +134,6 @@ async def sql_answer( schema_contexts = await self._load_active_schema_contexts( sql_answer_request.project_id, - sql_answer_request.query, ) normalized_sql = self._normalize_and_validate_sql( sql_answer_request.sql, From 9bf9f62b21fdb0f35555c361395f1f21eacd59d5 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 13:50:25 +0530 Subject: [PATCH 0558/1087] Fix metadata deploy indexing null handling --- .../src/pipelines/indexing/db_schema.py | 31 ++++++++--- .../pipelines/indexing/table_description.py | 38 ++++++++++---- .../src/pipelines/indexing/utils/helper.py | 17 ++++--- .../web/v1/services/semantics_preparation.py | 42 ++++++++++++++- .../pipelines/indexing/test_db_schema.py | 51 +++++++++++++++++++ .../indexing/test_table_description.py | 25 +++++++++ 6 files changed, 178 insertions(+), 26 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 197d8def8c..8bdf0f17f1 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -32,11 +32,15 @@ @component class DDLChunker: - def _truncate_text(self, text: str, max_length: int) -> str: - if len(text) <= max_length: - return text + def _normalize_text(self, value: Any) -> str: + return "" if value is None else str(value) - return text[:max_length].rstrip() + "..." + def _truncate_text(self, text: Any, max_length: int) -> str: + normalized_text = self._normalize_text(text) + if len(normalized_text) <= max_length: + return normalized_text + + return normalized_text[:max_length].rstrip() + "..." def _serialize_table_columns_payload(self, columns: List[dict]) -> str: return str({"type": "TABLE_COLUMNS", "columns": columns}) @@ -174,9 +178,13 @@ async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: for column in model.get("columns", []) if column.get("isHidden") is not True ] + properties = model.get("properties") + if not isinstance(properties, dict): + properties = {} + return { "name": model.get("name", ""), - "properties": model.get("properties", {}), + "properties": properties, "columns": columns, "primaryKey": model.get("primaryKey", ""), } @@ -212,9 +220,14 @@ def _convert_models_and_relationships( ) -> List[Dict[str, str]]: def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) + if not isinstance(properties, dict): + properties = {} + display_name = properties.get("displayName", "") model_properties = { - "alias": clean_display_name(properties.get("displayName", "")), + "alias": clean_display_name( + "" if display_name is None else str(display_name) + ), "description": self._truncate_text( properties.get("description", ""), MAX_DB_SCHEMA_COMMENT_LENGTH, @@ -322,10 +335,14 @@ def _column_batch( def _convert_views(self, views: List[Dict[str, Any]]) -> List[Dict[str, str]]: def _payload(view: Dict[str, Any]) -> dict: + properties = view.get("properties") + if not isinstance(properties, dict): + properties = {} + return { "type": "VIEW", "comment": self._truncate_text( - f"/* {view['properties']} */\n" if "properties" in view else "", + f"/* {properties} */\n" if properties else "", MAX_DB_SCHEMA_COMMENT_LENGTH, ), "name": view["name"], diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 8e1b875b49..cb39b097a3 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -24,18 +24,26 @@ @component class TableDescriptionChunker: - def _truncate_description(self, description: str) -> str: - if len(description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: - return description + def _normalize_text(self, value: Any) -> str: + return "" if value is None else str(value) - return description[:MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH].rstrip() + "..." + def _truncate_description(self, description: Any) -> str: + normalized_description = self._normalize_text(description) + if len(normalized_description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: + return normalized_description - def _format_columns(self, columns: List[str]) -> str: - if len(columns) <= MAX_TABLE_DESCRIPTION_COLUMNS: - return ", ".join(columns) + return ( + normalized_description[:MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH].rstrip() + + "..." + ) + + def _format_columns(self, columns: List[Any]) -> str: + normalized_columns = [self._normalize_text(column) for column in columns] + if len(normalized_columns) <= MAX_TABLE_DESCRIPTION_COLUMNS: + return ", ".join(normalized_columns) - remaining_columns = len(columns) - MAX_TABLE_DESCRIPTION_COLUMNS - truncated_columns = columns[:MAX_TABLE_DESCRIPTION_COLUMNS] + [ + remaining_columns = len(normalized_columns) - MAX_TABLE_DESCRIPTION_COLUMNS + truncated_columns = normalized_columns[:MAX_TABLE_DESCRIPTION_COLUMNS] + [ f"... (+{remaining_columns} more columns)" ] return ", ".join(truncated_columns) @@ -70,11 +78,19 @@ def _additional_meta() -> Dict[str, Any]: def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[str]: def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: + properties = payload.get("properties") + if not isinstance(properties, dict): + properties = {} + return { "mdl_type": mdl_type, "name": payload.get("name"), - "columns": [column["name"] for column in payload.get("columns", [])], - "properties": payload.get("properties", {}), + "columns": [ + column.get("name", "") + for column in payload.get("columns", []) + if isinstance(column, dict) + ], + "properties": properties, } resources = ( diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 3829324a0a..8e0e2a6b25 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -29,10 +29,15 @@ def __call__(self, column: Dict[str, Any], **kwargs) -> Any: def _properties_comment(column: Dict[str, Any], **_) -> str: - props = column["properties"] + props = column.get("properties") + if not isinstance(props, dict): + props = {} + + display_name = props.get("displayName", "") + description = props.get("description", "") column_properties = { - "alias": clean_display_name(props.get("displayName", "")), - "description": props.get("description", ""), + "alias": clean_display_name("" if display_name is None else str(display_name)), + "description": "" if description is None else str(description), } # Add any nested columns if they exist @@ -56,8 +61,8 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_PREPROCESSORS = { "properties": Helper( - condition=lambda column, **_: "properties" in column, - helper=lambda column, **_: column.get("properties"), + condition=lambda column, **_: isinstance(column.get("properties"), dict), + helper=lambda column, **_: column.get("properties", {}), ), "relationship": Helper( condition=lambda column, **_: "relationship" in column, @@ -75,7 +80,7 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_COMMENT_HELPERS = { "properties": Helper( - condition=lambda column, **_: "properties" in column, + condition=lambda column, **_: isinstance(column.get("properties"), dict), helper=_properties_comment, ), "isCalculated": Helper( 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 35d931c4b6..23dd0994d1 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -62,8 +62,45 @@ def _parse_mdl(self, mdl: str) -> dict[str, Any]: parsed.setdefault("views", []) parsed.setdefault("metrics", []) parsed.setdefault("relationships", []) + self._normalize_mdl_metadata(parsed) return parsed + def _normalize_text(self, value: Any) -> str: + return "" if value is None else str(value) + + def _normalize_properties(self, payload: dict[str, Any]) -> None: + properties = payload.get("properties") + if not isinstance(properties, dict): + properties = {} + payload["properties"] = properties + + for key in ("description", "displayName"): + if key in properties: + properties[key] = self._normalize_text(properties[key]) + + def _normalize_resource_metadata(self, payload: dict[str, Any]) -> None: + self._normalize_properties(payload) + columns = payload.get("columns", []) + if not isinstance(columns, list): + payload["columns"] = [] + return + + for column in columns: + if not isinstance(column, dict): + continue + self._normalize_properties(column) + + def _normalize_mdl_metadata(self, mdl: dict[str, Any]) -> None: + for collection in ("models", "views", "metrics"): + resources = mdl.get(collection, []) + if not isinstance(resources, list): + mdl[collection] = [] + continue + + for resource in resources: + if isinstance(resource, dict): + self._normalize_resource_metadata(resource) + def _validate_mdl_integrity(self, mdl: dict[str, Any]) -> None: model_names = set() for model in mdl["models"]: @@ -185,10 +222,11 @@ async def prepare_semantics( try: mdl = self._parse_mdl(prepare_semantics_request.mdl) self._validate_mdl_integrity(mdl) - logger.info(f"MDL: {prepare_semantics_request.mdl}") + normalized_mdl = orjson.dumps(mdl).decode("utf-8") + logger.info(f"MDL: {normalized_mdl}") input = { - "mdl_str": prepare_semantics_request.mdl, + "mdl_str": normalized_mdl, "project_id": prepare_semantics_request.project_id, } diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py index 5bc8c1303f..3a1d3b05e1 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py @@ -199,6 +199,57 @@ async def test_column_with_properties(): ) +@pytest.mark.asyncio +async def test_null_metadata_properties_are_indexed_as_empty_text(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "user", + "properties": {"description": None, "displayName": None}, + "columns": [ + { + "name": "id", + "type": "INTEGER", + "properties": { + "displayName": None, + "description": None, + }, + } + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=1) + + assert len(actual["documents"]) == 2 + assert actual["documents"][0].content == str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "comment": '-- {"alias":"","description":""}\n ', + "name": "id", + "data_type": "INTEGER", + "is_primary_key": False, + } + ], + } + ) + assert actual["documents"][1].content == str( + { + "type": "TABLE", + "comment": "\n/* {'alias': '', 'description': ''} */\n", + "name": "user", + } + ) + + @pytest.mark.asyncio async def test_column_with_nested_columns(): chunker = DDLChunker() diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index 1585db75d5..244649e6b1 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -126,6 +126,31 @@ def test_table_description_missing_description(): assert document.content == str({"name": "user", "description": "", "columns": ""}) +def test_table_description_null_description(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + { + "name": "user", + "properties": {"description": None, "displayName": None}, + "columns": [{"name": "id"}, {"name": None}], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = chunker.run(mdl) + assert len(actual["documents"]) == 1 + + document: Document = actual["documents"][0] + assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "user"} + assert document.content == str( + {"name": "user", "description": "", "columns": "id, "} + ) + + def test_table_description_truncates_long_column_lists(): chunker = TableDescriptionChunker() mdl = { From 74879e42e4edd6b31921ec6f9f400f7d8def818e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 14:49:40 +0530 Subject: [PATCH 0559/1087] Fix Modeling AI Assistant result handling --- wren-ui/src/pages/modeling.tsx | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index cadbea22fe..3d28a94283 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -4,11 +4,11 @@ import { useSearchParams } from 'next/navigation'; import { forwardRef, useEffect, useMemo, useRef, useState } from 'react'; import { Button, - Checkbox, Dropdown, Input, Menu, Modal, + Select, Table, message, } from 'antd'; @@ -458,6 +458,23 @@ export default function Modeling() { throw new Error('AI assistant timed out.'); }; + const normalizeSemanticResult = (result: any): any[] => { + if (Array.isArray(result)) return result; + if (Array.isArray(result?.models)) return result.models; + if (Array.isArray(result?.semantics)) return result.semantics; + if (Array.isArray(result?.descriptions)) return result.descriptions; + return []; + }; + + const normalizeRelationshipResult = (result: any): any[] => { + if (Array.isArray(result)) return result; + if (Array.isArray(result?.relationships)) return result.relationships; + if (Array.isArray(result?.response?.relationships)) { + return result.response.relationships; + } + return []; + }; + const openAssistant = (mode: 'semantics' | 'relationships') => { setAssistantMode(mode); setSelectedModels(diagramData?.models?.map((model) => model.referenceName) || []); @@ -486,7 +503,7 @@ export default function Modeling() { MODELING_SEMANTICS_RESULT, 'modelingSemanticsResult', ); - setSemanticResult(result); + setSemanticResult(normalizeSemanticResult(result)); } if (assistantMode === 'relationships') { if (!diagramData?.models || diagramData.models.length < 2) { @@ -499,7 +516,7 @@ export default function Modeling() { MODELING_RELATIONSHIPS_RESULT, 'modelingRelationshipsResult', ); - setRelationshipResult(result?.relationships || []); + setRelationshipResult(normalizeRelationshipResult(result)); } } catch (error: any) { message.error(error.message || 'Failed to run Modeling AI Assistant.'); @@ -739,15 +756,16 @@ export default function Modeling() { > {assistantMode === 'semantics' && ( <> - ({ label: model.displayName || model.referenceName, value: model.referenceName, }))} - onChange={(values) => setSelectedModels(values as string[])} + onChange={(values) => setSelectedModels(values)} /> Date: Wed, 15 Jul 2026 16:07:14 +0530 Subject: [PATCH 0560/1087] Implement generate semantics assistant flow --- wren-ui/src/pages/modeling.tsx | 360 ++++++++++++++++++++++++++++++++- 1 file changed, 354 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 3d28a94283..499d056c56 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -99,6 +99,60 @@ const AssistantAction = styled.div` z-index: 10; `; +const AssistantPage = styled.div` + min-height: calc(100vh - 48px); + background: #f5f5f5; + padding: 72px 24px; +`; + +const AssistantCard = styled.div` + max-width: 1060px; + margin: 0 auto; + padding: 34px 72px; + background: #fff; + border: 1px solid #e5e7eb; + border-radius: 2px; +`; + +const AssistantBack = styled.button` + display: block; + max-width: 1060px; + margin: 0 auto 8px; + padding: 0; + border: 0; + background: transparent; + color: #5f6368; + cursor: pointer; +`; + +const AssistantTitle = styled.h1` + margin: 0 0 14px; + color: #3c4043; + font-size: 40px; + line-height: 1.2; +`; + +const AssistantDescription = styled.p` + margin: 0 0 26px; + color: #5f6368; + line-height: 1.6; +`; + +const AssistantFooter = styled.div` + display: flex; + justify-content: space-between; + margin-top: 28px; +`; + +const SemanticReviewCard = styled.div` + padding: 18px 16px 28px; + border-bottom: 1px solid #e5e7eb; + + &:last-child { + border-bottom: 0; + } +`; + export default function Modeling() { const router = useRouter(); const searchParams = useSearchParams(); @@ -110,6 +164,10 @@ export default function Modeling() { const [assistantLoading, setAssistantLoading] = useState(false); const [selectedModels, setSelectedModels] = useState([]); const [semanticPrompt, setSemanticPrompt] = useState(''); + const [semanticStep, setSemanticStep] = useState<'pick' | 'generate' | 'review'>( + 'pick', + ); + const [semanticSearch, setSemanticSearch] = useState(''); const [semanticResult, setSemanticResult] = useState([]); const [relationshipResult, setRelationshipResult] = useState([]); @@ -458,11 +516,38 @@ export default function Modeling() { throw new Error('AI assistant timed out.'); }; + const normalizeSemanticModel = (name: string, value: any): any => ({ + name: value?.name || name, + description: value?.description || value?.properties?.description || '', + columns: (value?.columns || []).map((column) => ({ + name: column?.name, + type: column?.type, + description: column?.description || column?.properties?.description || '', + })), + }); + const normalizeSemanticResult = (result: any): any[] => { - if (Array.isArray(result)) return result; - if (Array.isArray(result?.models)) return result.models; - if (Array.isArray(result?.semantics)) return result.semantics; - if (Array.isArray(result?.descriptions)) return result.descriptions; + if (Array.isArray(result)) { + return result.map((model) => normalizeSemanticModel(model?.name, model)); + } + if (Array.isArray(result?.models)) { + return result.models.map((model) => normalizeSemanticModel(model?.name, model)); + } + if (Array.isArray(result?.semantics)) { + return result.semantics.map((model) => + normalizeSemanticModel(model?.name, model), + ); + } + if (Array.isArray(result?.descriptions)) { + return result.descriptions.map((model) => + normalizeSemanticModel(model?.name, model), + ); + } + if (result && typeof result === 'object') { + return Object.entries(result).map(([name, value]) => + normalizeSemanticModel(name, value), + ); + } return []; }; @@ -477,6 +562,8 @@ export default function Modeling() { const openAssistant = (mode: 'semantics' | 'relationships') => { setAssistantMode(mode); + setSemanticStep('pick'); + setSemanticSearch(''); setSelectedModels(diagramData?.models?.map((model) => model.referenceName) || []); setSemanticResult([]); setRelationshipResult([]); @@ -503,7 +590,9 @@ export default function Modeling() { MODELING_SEMANTICS_RESULT, 'modelingSemanticsResult', ); - setSemanticResult(normalizeSemanticResult(result)); + const normalizedResult = normalizeSemanticResult(result); + setSemanticResult(normalizedResult); + setSemanticStep('review'); } if (assistantMode === 'relationships') { if (!diagramData?.models || diagramData.models.length < 2) { @@ -525,6 +614,42 @@ export default function Modeling() { } }; + const updateSemanticModelDescription = ( + modelName: string, + description: string, + ) => { + setSemanticResult((models) => + models.map((model) => + model.name === modelName ? { ...model, description } : model, + ), + ); + }; + + const updateSemanticColumnDescription = ( + modelName: string, + columnName: string, + description: string, + ) => { + setSemanticResult((models) => + models.map((model) => + model.name === modelName + ? { + ...model, + columns: (model.columns || []).map((column) => + column.name === columnName ? { ...column, description } : column, + ), + } + : model, + ), + ); + }; + + const closeAssistant = () => { + setAssistantMode(null); + setSemanticStep('pick'); + setSemanticSearch(''); + }; + const saveAssistantResult = async () => { try { if (!diagramData) return; @@ -604,7 +729,7 @@ export default function Modeling() { }); } } - setAssistantMode(null); + closeAssistant(); message.success('Saved Modeling AI Assistant suggestions.'); } catch (error: any) { message.error(error.message || 'Failed to save assistant suggestions.'); @@ -613,6 +738,229 @@ export default function Modeling() { } }; + const semanticModelOptions = (diagramData?.models || []).filter((model) => { + const keyword = semanticSearch.trim().toLowerCase(); + if (!keyword) return true; + return [model.displayName, model.referenceName] + .filter(Boolean) + .some((value) => String(value).toLowerCase().includes(keyword)); + }); + + if (assistantMode === 'semantics') { + return ( + + + + ← Back to modeling + + + {semanticStep === 'pick' && ( + <> + Pick models + + + Good semantics improve how AI understands and queries your + data. + {' '} + Select models to generate semantics with AI. Modeling AI + Assistant will help you create semantics that improve how AI + understands and queries your data. + +
+ {selectedModels.length}/{diagramData?.models?.length || 0}{' '} + model(s) +
+ setSemanticSearch(event.target.value)} + /> +
setSelectedModels(keys as string[]), + }} + columns={[ + { + title: 'Model name', + render: (_value, model) => + model.displayName || model.referenceName, + }, + ]} + /> + + + + + + )} + + {semanticStep === 'generate' && ( + <> + Generate semantics +

User Prompt

+ + Help AI better understand your data by providing a brief + description of your dataset's purpose. Modeling AI + Assistant will use this context to generate more relevant + semantics. + +
+ setSemanticPrompt(event.target.value)} + placeholder="Describe what this dataset represents and how it is used." + /> + +
+
+ Example prompt +
+ This dataset tracks operational records, users, events, and + business entities. It is used to answer analytical questions + about activity, performance, ownership, and trends. +
+
+ + + + + + )} + + {semanticStep === 'review' && ( + <> + Generate semantics +
+ setSemanticPrompt(event.target.value)} + /> + +
+
+
+ Generated semantics +
+ Review the semantics generated by AI. +
+
+ {semanticResult.map((model) => ( + +
+ {model.name} + + {(model.columns || []).length} column(s) + +
+
Description
+ + updateSemanticModelDescription( + model.name, + event.target.value, + ) + } + /> +
( + + updateSemanticColumnDescription( + model.name, + column.name, + event.target.value, + ) + } + /> + ), + }, + ]} + /> + + ))} + + + + + + + )} + + + + ); + } + return ( Date: Wed, 15 Jul 2026 17:10:37 +0530 Subject: [PATCH 0561/1087] Fix generate semantics completion handling --- wren-ai-service/src/globals.py | 1 + .../generation/semantics_description.py | 46 ++++++++++++++----- .../web/v1/routers/semantics_description.py | 15 ++++-- .../web/v1/services/semantics_description.py | 46 ++++++++++++++----- wren-ui/src/pages/modeling.tsx | 3 ++ 5 files changed, 83 insertions(+), 28 deletions(-) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 3c8ab0ab30..f2a0db24a6 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -90,6 +90,7 @@ def create_service_container( **pipe_components["semantics_description"], ) }, + generation_timeout_seconds=settings.pipeline_timeout_seconds, **query_cache, ), semantics_preparation_service=services.SemanticsPreparationService( diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index acc5fc8594..d197bf5017 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -96,29 +96,40 @@ def picked_models(mdl: dict, selected_models: list[str]) -> list[dict]: def relation_filter(column: dict) -> bool: return "relationship" not in column + def _properties(payload: dict) -> dict: + properties = payload.get("properties") + return properties if isinstance(properties, dict) else {} + + def _text(value) -> str: + return "" if value is None else str(value) + def column_formatter(columns: list[dict]) -> list[dict]: return [ { - "name": column["name"], - "type": column["type"], + "name": column.get("name", ""), + "type": column.get("type", ""), "properties": { - "description": column["properties"].get("description", ""), + "description": _text( + _properties(column).get("description", "") + ), "alias": clean_display_name( - column["properties"].get("displayName", "") + _text(_properties(column).get("displayName", "")) ), }, } - for column in columns + for column in columns or [] if relation_filter(column) ] def extract(model: dict) -> dict: return { - "name": model["name"], - "columns": column_formatter(model["columns"]), + "name": model.get("name", ""), + "columns": column_formatter(model.get("columns", [])), "properties": { - "description": model["properties"].get("description", ""), - "alias": clean_display_name(model["properties"].get("displayName", "")), + "description": _text(_properties(model).get("description", "")), + "alias": clean_display_name( + _text(_properties(model).get("displayName", "")) + ), }, } @@ -163,10 +174,18 @@ def wrapper(text: str) -> str: logger.error(f"Error decoding JSON: {e}") return {"models": []} # Return an empty list if JSON decoding fails - reply = generate.get("replies")[0] # Expecting only one reply + replies = generate.get("replies") or [] + if not replies: + return {} + + reply = replies[0] # Expecting only one reply normalized = wrapper(reply) - return {model["name"]: model for model in normalized["models"]} + return { + model["name"]: model + for model in normalized.get("models", []) + if isinstance(model, dict) and model.get("name") + } @observe(capture_input=False) @@ -179,7 +198,10 @@ def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: models = {model["name"]: model for model in picked_models} return { - name: {**data, "columns": _filter(data["columns"], models[name]["columns"])} + name: { + **data, + "columns": _filter(data.get("columns", []), models[name]["columns"]), + } for name, data in normalize.items() if name in models } diff --git a/wren-ai-service/src/web/v1/routers/semantics_description.py b/wren-ai-service/src/web/v1/routers/semantics_description.py index 16727bdd6c..3e36d299bf 100644 --- a/wren-ai-service/src/web/v1/routers/semantics_description.py +++ b/wren-ai-service/src/web/v1/routers/semantics_description.py @@ -74,19 +74,26 @@ def _formatter(response: Optional[dict]) -> Optional[list[dict]]: if response is None: return None + def _properties(payload: dict) -> dict: + properties = payload.get("properties") + return properties if isinstance(properties, dict) else {} + return [ { "name": model_name, "columns": [ { - "name": column["name"], - "description": column["properties"].get("description", ""), + "name": column.get("name", ""), + "type": column.get("type", ""), + "description": _properties(column).get("description", ""), } - for column in model_data["columns"] + for column in model_data.get("columns", []) + if isinstance(column, dict) ], - "description": model_data["properties"].get("description", ""), + "description": _properties(model_data).get("description", ""), } for model_name, model_data in response.items() + if isinstance(model_data, dict) ] return GetResponse( diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 85bf14b188..596d70a905 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,6 +1,6 @@ import asyncio import logging -from typing import Dict, Literal, Optional +from typing import Any, Dict, Literal, Optional import orjson from cachetools import TTLCache @@ -32,9 +32,11 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, + generation_timeout_seconds: int = 90, ): self._pipelines = pipelines self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) + self._generation_timeout_seconds = generation_timeout_seconds def _handle_exception( self, @@ -67,15 +69,25 @@ def _chunking( "language": request.configurations.language, } - chunks = [ - { - **model, - "columns": model["columns"][i : i + chunk_size], - } - for model in mdl_dict["models"] - if model["name"] in request.selected_models - for i in range(0, len(model["columns"]), chunk_size) - ] + chunks: list[dict[str, Any]] = [] + selected_models = set(request.selected_models) + for model in mdl_dict.get("models", []): + model_name = model.get("name") + if model_name not in selected_models: + continue + + columns = model.get("columns") or [] + if not columns: + chunks.append({**model, "columns": []}) + continue + + chunks.extend( + { + **model, + "columns": columns[i : i + chunk_size], + } + for i in range(0, len(columns), chunk_size) + ) return [ { @@ -87,8 +99,13 @@ def _chunking( ] async def _generate_task(self, request_id: str, chunk: dict): - resp = await self._pipelines["semantics_description"].run(**chunk) + resp = await asyncio.wait_for( + self._pipelines["semantics_description"].run(**chunk), + timeout=self._generation_timeout_seconds, + ) output = resp.get("output") + if not isinstance(output, dict): + raise ValueError("Semantics description pipeline returned no output") current = self[request_id] current.response = current.response or {} @@ -98,7 +115,8 @@ async def _generate_task(self, request_id: str, chunk: dict): current.response[key] = output[key] continue - current.response[key]["columns"].extend(output[key]["columns"]) + current.response[key].setdefault("columns", []) + current.response[key]["columns"].extend(output[key].get("columns", [])) @observe(name="Generate Semantics Description") @trace_metadata @@ -110,6 +128,10 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: mdl_dict = orjson.loads(request.mdl) chunks = self._chunking(mdl_dict, request) + if not chunks: + raise ValueError( + "No selected models matched the current semantic model metadata" + ) tasks = [self._generate_task(request.id, chunk) for chunk in chunks] await asyncio.gather(*tasks) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 499d056c56..d0ab06584e 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -591,6 +591,9 @@ export default function Modeling() { 'modelingSemanticsResult', ); const normalizedResult = normalizeSemanticResult(result); + if (!normalizedResult.length) { + throw new Error('AI assistant returned no semantic descriptions.'); + } setSemanticResult(normalizedResult); setSemanticStep('review'); } From cb9ec7fd9c33eb82331e21e6fb113b95c38a1f1a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 17:47:41 +0530 Subject: [PATCH 0562/1087] Batch generate semantics LLM calls --- .../web/v1/services/semantics_description.py | 41 +++++++++----- .../services/test_semantics_description.py | 56 ++++++++----------- 2 files changed, 50 insertions(+), 47 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 596d70a905..a74560e1a6 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -71,6 +71,25 @@ def _chunking( chunks: list[dict[str, Any]] = [] selected_models = set(request.selected_models) + current_models: list[dict[str, Any]] = [] + current_column_count = 0 + + def _flush_current_models(): + nonlocal current_models, current_column_count + if not current_models: + return + chunks.append({"models": current_models}) + current_models = [] + current_column_count = 0 + + def _append_model(model: dict[str, Any]): + nonlocal current_models, current_column_count + column_count = len(model.get("columns") or []) + if current_models and current_column_count + column_count > chunk_size: + _flush_current_models() + current_models.append(model) + current_column_count += column_count + for model in mdl_dict.get("models", []): model_name = model.get("name") if model_name not in selected_models: @@ -78,22 +97,19 @@ def _chunking( columns = model.get("columns") or [] if not columns: - chunks.append({**model, "columns": []}) + _append_model({**model, "columns": []}) continue - chunks.extend( - { - **model, - "columns": columns[i : i + chunk_size], - } - for i in range(0, len(columns), chunk_size) - ) + for i in range(0, len(columns), chunk_size): + _append_model({**model, "columns": columns[i : i + chunk_size]}) + + _flush_current_models() return [ { **template, - "mdl": {"models": [chunk]}, - "selected_models": [chunk["name"]], + "mdl": chunk, + "selected_models": [model["name"] for model in chunk["models"]], } for chunk in chunks ] @@ -132,9 +148,8 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: raise ValueError( "No selected models matched the current semantic model metadata" ) - tasks = [self._generate_task(request.id, chunk) for chunk in chunks] - - await asyncio.gather(*tasks) + for chunk in chunks: + await self._generate_task(request.id, chunk) self[request.id].status = "finished" self[request.id].trace_id = trace_id diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 4b85ed914f..e74ee26eae 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -139,12 +139,13 @@ async def test_batch_processing_with_multiple_models( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model3", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Mock pipeline responses for each chunk - service._pipelines["semantics_description"].run.side_effect = [ - {"output": {"model1": {"description": "Description 1"}}}, - {"output": {"model2": {"description": "Description 2"}}}, - {"output": {"model3": {"description": "Description 3"}}}, - ] + service._pipelines["semantics_description"].run.return_value = { + "output": { + "model1": {"description": "Description 1"}, + "model2": {"description": "Description 2"}, + "model3": {"description": "Description 3"}, + } + } await service.generate(request) response = service[request.id] @@ -158,10 +159,10 @@ async def test_batch_processing_with_multiple_models( } chunks = service._chunking(orjson.loads(request.mdl), request) - assert len(chunks) == 3 # Default chunk_size=1 + assert len(chunks) == 1 assert all("user_prompt" in chunk for chunk in chunks) assert all("mdl" in chunk for chunk in chunks) - assert [len(chunk["selected_models"]) for chunk in chunks] == [1, 1, 1] + assert [len(chunk["selected_models"]) for chunk in chunks] == [3] def test_batch_processing_with_custom_chunk_size( @@ -178,12 +179,10 @@ def test_batch_processing_with_custom_chunk_size( # Test chunking with custom chunk size chunks = service._chunking(orjson.loads(request.mdl), request, chunk_size=2) - assert len(chunks) == 4 - assert [len(chunk["selected_models"]) for chunk in chunks] == [1, 1, 1, 1] - assert chunks[0]["selected_models"] == ["model1"] - assert chunks[1]["selected_models"] == ["model2"] - assert chunks[2]["selected_models"] == ["model3"] - assert chunks[3]["selected_models"] == ["model4"] + assert len(chunks) == 2 + assert [len(chunk["selected_models"]) for chunk in chunks] == [2, 2] + assert chunks[0]["selected_models"] == ["model1", "model2"] + assert chunks[1]["selected_models"] == ["model3", "model4"] @pytest.mark.asyncio @@ -198,11 +197,9 @@ async def test_batch_processing_partial_failure( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Mock first chunk succeeds, second chunk fails - service._pipelines["semantics_description"].run.side_effect = [ - {"output": {"model1": {"description": "Description 1"}}}, - Exception("Failed processing model2"), - ] + service._pipelines["semantics_description"].run.side_effect = Exception( + "Failed processing selected models" + ) await service.generate(request) response = service[request.id] @@ -210,7 +207,7 @@ async def test_batch_processing_partial_failure( assert response.id == "test_id" assert response.status == "failed" assert response.error.code == "OTHERS" - assert "Failed processing model2" in response.error.message + assert "Failed processing selected models" in response.error.message @pytest.mark.asyncio @@ -227,22 +224,13 @@ async def test_concurrent_updates_no_race_condition( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model3", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model4", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model5", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Mock pipeline responses with delays to simulate concurrent execution - async def delayed_response(model_num, delay=0.1): - await asyncio.sleep(delay) # Add delay to increase chance of race condition - return { - "output": {f"model{model_num}": {"description": f"Description {model_num}"}} + service._pipelines["semantics_description"].run.return_value = { + "output": { + f"model{i}": {"description": f"Description {i}"} + for i in range(1, 6) } + } - service._pipelines["semantics_description"].run.side_effect = [ - await delayed_response(1), - await delayed_response(2), - await delayed_response(3), - await delayed_response(4), - await delayed_response(5), - ] - - # Generate response which will process chunks concurrently await service.generate(request) response = service[request.id] From ddfe0eb8ccc3fa0d24573cda33f50b6272ff46c2 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 18:13:28 +0530 Subject: [PATCH 0563/1087] Return semantics when LLM generation stalls --- .../web/v1/services/semantics_description.py | 72 +++++++++++++++++-- .../services/test_semantics_description.py | 28 ++++++++ 2 files changed, 94 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index a74560e1a6..00238a1ad5 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -36,7 +36,7 @@ def __init__( ): self._pipelines = pipelines self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) - self._generation_timeout_seconds = generation_timeout_seconds + self._generation_timeout_seconds = min(generation_timeout_seconds, 30) def _handle_exception( self, @@ -61,6 +61,57 @@ class GenerateRequest(BaseRequest): user_prompt: str mdl: str + def _properties(self, payload: dict[str, Any]) -> dict[str, Any]: + properties = payload.get("properties") + return properties if isinstance(properties, dict) else {} + + def _text(self, value: Any) -> str: + return "" if value is None else str(value) + + def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: + output: dict[str, Any] = {} + for model in chunk.get("mdl", {}).get("models", []): + model_name = self._text(model.get("name", "")) + if not model_name: + continue + + model_properties = self._properties(model) + model_description = self._text(model_properties.get("description", "")) + if not model_description: + model_description = ( + f"Represents {model_name.replace('_', ' ')} records in this dataset." + ) + + columns = [] + for column in model.get("columns", []) or []: + if column.get("relationship"): + continue + column_name = self._text(column.get("name", "")) + if not column_name: + continue + column_properties = self._properties(column) + column_description = self._text( + column_properties.get("description", "") + ) + if not column_description: + column_description = ( + f"{column_name.replace('_', ' ')} field from {model_name}." + ) + columns.append( + { + "name": column_name, + "type": self._text(column.get("type", "")), + "properties": {"description": column_description}, + } + ) + + output[model_name] = { + "name": model_name, + "columns": columns, + "properties": {"description": model_description}, + } + return output + def _chunking( self, mdl_dict: dict, request: GenerateRequest, chunk_size: int = 50 ) -> list[dict]: @@ -115,11 +166,20 @@ def _append_model(model: dict[str, Any]): ] async def _generate_task(self, request_id: str, chunk: dict): - resp = await asyncio.wait_for( - self._pipelines["semantics_description"].run(**chunk), - timeout=self._generation_timeout_seconds, - ) - output = resp.get("output") + try: + resp = await asyncio.wait_for( + self._pipelines["semantics_description"].run(**chunk), + timeout=self._generation_timeout_seconds, + ) + output = resp.get("output") + except TimeoutError: + logger.warning( + "Semantics description LLM call timed out after %s seconds; " + "returning metadata-based fallback descriptions.", + self._generation_timeout_seconds, + ) + output = self._fallback_output(chunk) + if not isinstance(output, dict): raise ValueError("Semantics description pipeline returned no output") diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index e74ee26eae..2b35d2982f 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -100,6 +100,34 @@ async def test_generate_semantics_description_with_exception( ) +@pytest.mark.asyncio +async def test_generate_semantics_description_with_llm_timeout_returns_fallback(): + mock_pipeline = AsyncMock() + + async def never_returns(**_): + await asyncio.sleep(1) + + mock_pipeline.run.side_effect = never_returns + service = SemanticsDescription( + pipelines={"semantics_description": mock_pipeline}, + generation_timeout_seconds=0.01, + ) + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["model1"], + mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', + ) + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.response["model1"]["properties"]["description"] + assert response.response["model1"]["columns"][0]["properties"]["description"] + + def test_get_semantics_description_result( service: SemanticsDescription, ): From 4f3c48130ddadfb338ad70728a0874e3266a43c1 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 18:38:21 +0530 Subject: [PATCH 0564/1087] Honor configured semantics LLM timeout --- .../src/web/v1/services/semantics_description.py | 14 +++++++++++++- .../pytest/services/test_semantics_description.py | 9 +++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 00238a1ad5..b0fbed77e3 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -36,7 +36,7 @@ def __init__( ): self._pipelines = pipelines self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) - self._generation_timeout_seconds = min(generation_timeout_seconds, 30) + self._generation_timeout_seconds = generation_timeout_seconds def _handle_exception( self, @@ -167,11 +167,23 @@ def _append_model(model: dict[str, Any]): async def _generate_task(self, request_id: str, chunk: dict): try: + logger.info( + "Calling configured LLM for semantics descriptions. " + "models=%s timeout_seconds=%s", + chunk.get("selected_models", []), + self._generation_timeout_seconds, + ) resp = await asyncio.wait_for( self._pipelines["semantics_description"].run(**chunk), timeout=self._generation_timeout_seconds, ) output = resp.get("output") + if not output: + logger.warning( + "Configured LLM returned empty semantics output; " + "returning metadata-based fallback descriptions." + ) + output = self._fallback_output(chunk) except TimeoutError: logger.warning( "Semantics description LLM call timed out after %s seconds; " diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 2b35d2982f..73edd26870 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -155,6 +155,15 @@ def test_get_non_existent_semantics_description_result( assert "not found" in result.error.message +def test_semantics_description_uses_configured_timeout(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + generation_timeout_seconds=123, + ) + + assert service._generation_timeout_seconds == 123 + + @pytest.mark.asyncio async def test_batch_processing_with_multiple_models( service: SemanticsDescription, From ccc6a2695964227e918ec6b95f52093b78e1b75e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 18:42:08 +0530 Subject: [PATCH 0565/1087] Implement relationship assistant review flow --- wren-ui/src/pages/modeling.tsx | 527 ++++++++++++++++++++++++++------- 1 file changed, 418 insertions(+), 109 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index d0ab06584e..f84b8b54e3 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -1,18 +1,25 @@ import dynamic from 'next/dynamic'; import { useRouter } from 'next/router'; import { useSearchParams } from 'next/navigation'; -import { forwardRef, useEffect, useMemo, useRef, useState } from 'react'; +import React, { forwardRef, useEffect, useMemo, useRef, useState } from 'react'; import { Button, Dropdown, Input, Menu, - Modal, Select, + Space, + Spin, Table, message, } from 'antd'; -import { RobotOutlined } from '@ant-design/icons'; +import { + DeleteOutlined, + EditOutlined, + RobotOutlined, + SaveOutlined, + TableOutlined, +} from '@ant-design/icons'; import { gql, useApolloClient, useMutation } from '@apollo/client'; import styled from 'styled-components'; import { MORE_ACTION, NODE_TYPE } from '@/utils/enum'; @@ -153,6 +160,70 @@ const SemanticReviewCard = styled.div` } `; +const RelationshipGroup = styled.div` + margin-top: 24px; + border: 1px solid #e5e7eb; + border-radius: 4px; + overflow: hidden; +`; + +const RelationshipGroupTitle = styled.div` + display: flex; + align-items: center; + gap: 8px; + padding: 14px 16px; + border-bottom: 1px solid #e5e7eb; +`; + +const AssistantCenter = styled.div` + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + min-height: 170px; + color: #0f3bff; +`; + +type AssistantRelationship = { + clientId: string; + fromModel: string; + fromColumn: string; + toModel: string; + toColumn: string; + type: string; + reason: string; +}; + +const RELATIONSHIP_TYPES = [ + { label: 'Many-to-one', value: 'MANY_TO_ONE' }, + { label: 'One-to-many', value: 'ONE_TO_MANY' }, + { label: 'One-to-one', value: 'ONE_TO_ONE' }, +]; + +const relationshipTypeLabel = (type: string) => + RELATIONSHIP_TYPES.find((item) => item.value === type)?.label || type; + +const normalizeRelationshipType = (type: string) => { + const normalized = String(type || '') + .trim() + .replace(/[\s-]+/g, '_') + .toUpperCase(); + if (normalized === 'MANY_TO_ONE') return 'MANY_TO_ONE'; + if (normalized === 'ONE_TO_MANY') return 'ONE_TO_MANY'; + if (normalized === 'ONE_TO_ONE') return 'ONE_TO_ONE'; + return type; +}; + +const parseQualifiedField = (value = '') => { + const [model, ...columnParts] = String(value).split('.'); + return { + model: model || '', + column: columnParts.join('.') || '', + }; +}; + +const renderIcon = (IconComponent) => React.createElement(IconComponent as any); + export default function Modeling() { const router = useRouter(); const searchParams = useSearchParams(); @@ -164,12 +235,21 @@ export default function Modeling() { const [assistantLoading, setAssistantLoading] = useState(false); const [selectedModels, setSelectedModels] = useState([]); const [semanticPrompt, setSemanticPrompt] = useState(''); - const [semanticStep, setSemanticStep] = useState<'pick' | 'generate' | 'review'>( - 'pick', - ); + const [semanticStep, setSemanticStep] = useState< + 'pick' | 'generate' | 'review' + >('pick'); const [semanticSearch, setSemanticSearch] = useState(''); const [semanticResult, setSemanticResult] = useState([]); - const [relationshipResult, setRelationshipResult] = useState([]); + const [relationshipResult, setRelationshipResult] = useState< + AssistantRelationship[] + >([]); + const [originalRelationshipResult, setOriginalRelationshipResult] = useState< + AssistantRelationship[] + >([]); + const [relationshipAutoStarted, setRelationshipAutoStarted] = useState(false); + const [editingRelationshipKey, setEditingRelationshipKey] = useState< + string | null + >(null); const { data } = useDiagramQuery({ fetchPolicy: 'cache-and-network', @@ -531,7 +611,9 @@ export default function Modeling() { return result.map((model) => normalizeSemanticModel(model?.name, model)); } if (Array.isArray(result?.models)) { - return result.models.map((model) => normalizeSemanticModel(model?.name, model)); + return result.models.map((model) => + normalizeSemanticModel(model?.name, model), + ); } if (Array.isArray(result?.semantics)) { return result.semantics.map((model) => @@ -551,22 +633,62 @@ export default function Modeling() { return []; }; - const normalizeRelationshipResult = (result: any): any[] => { - if (Array.isArray(result)) return result; - if (Array.isArray(result?.relationships)) return result.relationships; - if (Array.isArray(result?.response?.relationships)) { - return result.response.relationships; - } - return []; + const normalizeRelationshipResult = ( + result: any, + ): AssistantRelationship[] => { + const relationships = Array.isArray(result) + ? result + : Array.isArray(result?.relationships) + ? result.relationships + : Array.isArray(result?.response?.relationships) + ? result.response.relationships + : []; + + return relationships.map((relationship, index) => { + const from = parseQualifiedField( + relationship.from || relationship.fromField || '', + ); + const to = parseQualifiedField( + relationship.to || relationship.toField || '', + ); + const fromModel = relationship.fromModel || from.model; + const fromColumn = relationship.fromColumn || from.column; + const toModel = relationship.toModel || to.model; + const toColumn = relationship.toColumn || to.column; + + return { + clientId: + relationship.clientId || + [ + fromModel, + fromColumn, + toModel, + toColumn, + relationship.type, + index, + ].join(':'), + fromModel, + fromColumn, + toModel, + toColumn, + type: normalizeRelationshipType(relationship.type), + reason: relationship.reason || relationship.description || '', + }; + }); }; const openAssistant = (mode: 'semantics' | 'relationships') => { setAssistantMode(mode); setSemanticStep('pick'); setSemanticSearch(''); - setSelectedModels(diagramData?.models?.map((model) => model.referenceName) || []); + setSelectedModels( + diagramData?.models?.map((model) => model.referenceName) || [], + ); setSemanticResult([]); setRelationshipResult([]); + setOriginalRelationshipResult([]); + setRelationshipAutoStarted(false); + setEditingRelationshipKey(null); }; const runAssistant = async () => { @@ -580,7 +702,8 @@ export default function Modeling() { variables: { data: { selectedModels, - userPrompt: semanticPrompt || 'Describe this dataset for analytics.', + userPrompt: + semanticPrompt || 'Describe this dataset for analytics.', }, }, }); @@ -608,7 +731,9 @@ export default function Modeling() { MODELING_RELATIONSHIPS_RESULT, 'modelingRelationshipsResult', ); - setRelationshipResult(normalizeRelationshipResult(result)); + const normalizedResult = normalizeRelationshipResult(result); + setRelationshipResult(normalizedResult); + setOriginalRelationshipResult(normalizedResult); } } catch (error: any) { message.error(error.message || 'Failed to run Modeling AI Assistant.'); @@ -639,7 +764,9 @@ export default function Modeling() { ? { ...model, columns: (model.columns || []).map((column) => - column.name === columnName ? { ...column, description } : column, + column.name === columnName + ? { ...column, description } + : column, ), } : model, @@ -647,10 +774,50 @@ export default function Modeling() { ); }; + const updateRelationship = ( + clientId: string, + changes: Partial, + ) => { + setRelationshipResult((relationships) => + relationships.map((relationship) => + relationship.clientId === clientId + ? { ...relationship, ...changes } + : relationship, + ), + ); + }; + + const updateRelationshipField = ( + clientId: string, + side: 'from' | 'to', + value: string, + ) => { + const field = parseQualifiedField(value); + updateRelationship( + clientId, + side === 'from' + ? { fromModel: field.model, fromColumn: field.column } + : { toModel: field.model, toColumn: field.column }, + ); + }; + + const deleteSuggestedRelationship = (clientId: string) => { + setRelationshipResult((relationships) => + relationships.filter( + (relationship) => relationship.clientId !== clientId, + ), + ); + if (editingRelationshipKey === clientId) { + setEditingRelationshipKey(null); + } + }; + const closeAssistant = () => { setAssistantMode(null); setSemanticStep('pick'); setSemanticSearch(''); + setRelationshipAutoStarted(false); + setEditingRelationshipKey(null); }; const saveAssistantResult = async () => { @@ -749,6 +916,34 @@ export default function Modeling() { .some((value) => String(value).toLowerCase().includes(keyword)); }); + const relationshipFieldOptions = (diagramData?.models || []).flatMap( + (model) => + (model.fields || []).map((field) => ({ + label: `${model.referenceName}.${field.referenceName}`, + value: `${model.referenceName}.${field.referenceName}`, + })), + ); + + const relationshipGroups = relationshipResult.reduce< + Record + >((groups, relationship) => { + const key = relationship.fromModel || 'Unknown model'; + groups[key] = [...(groups[key] || []), relationship]; + return groups; + }, {}); + + const isRelationshipGenerating = + assistantMode === 'relationships' && + (!relationshipAutoStarted || assistantLoading) && + !relationshipResult.length; + + useEffect(() => { + if (assistantMode !== 'relationships') return; + if (relationshipAutoStarted) return; + setRelationshipAutoStarted(true); + runAssistant(); + }, [assistantMode, relationshipAutoStarted]); + if (assistantMode === 'semantics') { return ( @@ -946,7 +1141,9 @@ export default function Modeling() { ))} - +
+ editingRelationshipKey === record.clientId ? ( + + updateRelationshipField( + record.clientId, + 'to', + value, + ) + } + /> + ) : ( + `${record.toModel}.${record.toColumn}` + ), + }, + { + title: 'Type', + width: 170, + render: (_value, record) => + editingRelationshipKey === record.clientId ? ( + ({ - label: model.displayName || model.referenceName, - value: model.referenceName, - }))} - onChange={(values) => setSelectedModels(values)} - /> - setSemanticPrompt(event.target.value)} - /> - {!!semanticResult.length && ( -
( -
- ), - }} - /> - )} - - )} - {assistantMode === 'relationships' && ( -
- `${record.fromModel}.${record.fromColumn}-${record.toModel}.${record.toColumn}` - } - pagination={false} - dataSource={relationshipResult} - columns={[ - { - title: 'From', - render: (_value, record) => - `${record.fromModel}.${record.fromColumn}`, - }, - { - title: 'To', - render: (_value, record) => - `${record.toModel}.${record.toColumn}`, - }, - { title: 'Type', dataIndex: 'type', width: 150 }, - { title: 'Description', dataIndex: 'reason' }, - ]} - /> - )} - ); From 5b8fd57cf6a8554870fc10cb069c2369888d7538 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 18:51:19 +0530 Subject: [PATCH 0566/1087] Stop stale modeling assistant polling --- wren-ui/src/pages/modeling.tsx | 38 ++++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index f84b8b54e3..a099e43a07 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -223,12 +223,14 @@ const parseQualifiedField = (value = '') => { }; const renderIcon = (IconComponent) => React.createElement(IconComponent as any); +const ASSISTANT_CANCELLED = 'ASSISTANT_CANCELLED'; export default function Modeling() { const router = useRouter(); const searchParams = useSearchParams(); const apolloClient = useApolloClient(); const diagramRef = useRef(null); + const assistantRunIdRef = useRef(0); const [assistantMode, setAssistantMode] = useState< 'semantics' | 'relationships' | null >(null); @@ -575,20 +577,28 @@ export default function Modeling() { queryId: string, query: any, fieldName: string, + runId: number, ) => { if (!queryId) { throw new Error('AI assistant did not return a task id.'); } for (let attempt = 0; attempt < 90; attempt += 1) { + if (assistantRunIdRef.current !== runId) { + throw new Error(ASSISTANT_CANCELLED); + } const res = await apolloClient.query({ query, variables: { queryId }, fetchPolicy: 'network-only', }); + if (assistantRunIdRef.current !== runId) { + throw new Error(ASSISTANT_CANCELLED); + } const payload = res.data?.[fieldName]; - if (payload?.status === 'finished') return payload.response || []; - if (payload?.status === 'failed') { + const status = String(payload?.status || '').toLowerCase(); + if (status === 'finished') return payload.response || []; + if (status === 'failed') { throw new Error(payload.error?.message || 'AI assistant failed.'); } await new Promise((resolve) => setTimeout(resolve, 2000)); @@ -678,6 +688,7 @@ export default function Modeling() { }; const openAssistant = (mode: 'semantics' | 'relationships') => { + assistantRunIdRef.current += 1; setAssistantMode(mode); setSemanticStep('pick'); setSemanticSearch(''); @@ -692,6 +703,8 @@ export default function Modeling() { }; const runAssistant = async () => { + const runId = assistantRunIdRef.current + 1; + assistantRunIdRef.current = runId; try { setAssistantLoading(true); if (assistantMode === 'semantics') { @@ -712,7 +725,9 @@ export default function Modeling() { queryId, MODELING_SEMANTICS_RESULT, 'modelingSemanticsResult', + runId, ); + if (assistantRunIdRef.current !== runId) return; const normalizedResult = normalizeSemanticResult(result); if (!normalizedResult.length) { throw new Error('AI assistant returned no semantic descriptions.'); @@ -730,15 +745,21 @@ export default function Modeling() { queryId, MODELING_RELATIONSHIPS_RESULT, 'modelingRelationshipsResult', + runId, ); + if (assistantRunIdRef.current !== runId) return; const normalizedResult = normalizeRelationshipResult(result); setRelationshipResult(normalizedResult); setOriginalRelationshipResult(normalizedResult); } } catch (error: any) { - message.error(error.message || 'Failed to run Modeling AI Assistant.'); + if (error.message !== ASSISTANT_CANCELLED) { + message.error(error.message || 'Failed to run Modeling AI Assistant.'); + } } finally { - setAssistantLoading(false); + if (assistantRunIdRef.current === runId) { + setAssistantLoading(false); + } } }; @@ -813,7 +834,9 @@ export default function Modeling() { }; const closeAssistant = () => { + assistantRunIdRef.current += 1; setAssistantMode(null); + setAssistantLoading(false); setSemanticStep('pick'); setSemanticSearch(''); setRelationshipAutoStarted(false); @@ -944,6 +967,13 @@ export default function Modeling() { runAssistant(); }, [assistantMode, relationshipAutoStarted]); + useEffect( + () => () => { + assistantRunIdRef.current += 1; + }, + [], + ); + if (assistantMode === 'semantics') { return ( From 11f08cfdeb76030a28d0b36a73702aa46cb2a3e5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 18:59:08 +0530 Subject: [PATCH 0567/1087] Bound relationship recommendation generation --- .../services/relationship_recommendation.py | 125 +++++++++++++++++- .../test_relationship_recommendation.py | 76 +++++++++++ 2 files changed, 198 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py index 25c0bc5820..6aa6e9a1a0 100644 --- a/wren-ai-service/src/web/v1/services/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/services/relationship_recommendation.py @@ -1,5 +1,7 @@ +import asyncio import logging -from typing import Dict, Literal, Optional +import re +from typing import Any, Dict, Literal, Optional import orjson from cachetools import TTLCache @@ -35,11 +37,103 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, + generation_timeout_seconds: int = 45, ): self._pipelines = pipelines self._cache: Dict[str, RelationshipRecommendation.Resource] = TTLCache( maxsize=maxsize, ttl=ttl ) + self._generation_timeout_seconds = generation_timeout_seconds + + def _normalize_identifier(self, value: Any) -> str: + text = "" if value is None else str(value) + text = re.sub(r"[^a-zA-Z0-9]", "", text).lower() + return text[:-1] if text.endswith("s") else text + + def _fallback_relationships(self, mdl: dict) -> dict: + models = mdl.get("models", []) or [] + existing = { + ( + relationship.get("models", [None, None])[0], + relationship.get("condition", ""), + relationship.get("joinType", ""), + ) + for relationship in mdl.get("relationships", []) or [] + } + candidates = [] + + model_lookup = { + self._normalize_identifier(model.get("name")): model for model in models + } + + for from_model in models: + from_model_name = from_model.get("name") + if not from_model_name: + continue + + for column in from_model.get("columns", []) or []: + if column.get("relationship"): + continue + + from_column = column.get("name") + if not from_column: + continue + + normalized_column = self._normalize_identifier(from_column) + if not normalized_column.endswith("id") or normalized_column == "id": + continue + + target_key = normalized_column[:-2] + to_model = model_lookup.get(target_key) + if not to_model or to_model.get("name") == from_model_name: + continue + + to_model_name = to_model.get("name") + to_columns = to_model.get("columns", []) or [] + primary_key = to_model.get("primaryKey") + to_column = next( + ( + item.get("name") + for item in to_columns + if item.get("name") == primary_key + ), + None, + ) + to_column = to_column or next( + ( + item.get("name") + for item in to_columns + if self._normalize_identifier(item.get("name")) == "id" + ), + None, + ) + if not to_column: + continue + + signature = ( + from_model_name, + f"{from_model_name}.{from_column} = {to_model_name}.{to_column}", + "MANY_TO_ONE", + ) + if signature in existing: + continue + + candidates.append( + { + "name": f"{from_model_name}_{to_model_name}", + "fromModel": from_model_name, + "fromColumn": from_column, + "type": "MANY_TO_ONE", + "toModel": to_model_name, + "toColumn": to_column, + "reason": ( + f"{from_model_name}.{from_column} appears to reference " + f"{to_model_name}.{to_column}." + ), + } + ) + + return {"relationships": candidates} def _handle_exception( self, @@ -72,12 +166,37 @@ async def recommend(self, request: Input, **kwargs) -> Resource: "language": request.configurations.language, } - resp = await self._pipelines["relationship_recommendation"].run(**input) + try: + logger.info( + "Calling configured LLM for relationship recommendations. " + "timeout_seconds=%s", + self._generation_timeout_seconds, + ) + resp = await asyncio.wait_for( + self._pipelines["relationship_recommendation"].run(**input), + timeout=self._generation_timeout_seconds, + ) + response = resp.get("validated") + if not response or ( + "relationships" in response and not response.get("relationships") + ): + logger.warning( + "Configured LLM returned empty relationship recommendations; " + "returning metadata-based fallback relationships." + ) + response = self._fallback_relationships(mdl_dict) + except TimeoutError: + logger.warning( + "Relationship recommendation LLM call timed out after %s seconds; " + "returning metadata-based fallback relationships.", + self._generation_timeout_seconds, + ) + response = self._fallback_relationships(mdl_dict) self._cache[request.id] = self.Resource( id=request.id, status="finished", - response=resp.get("validated"), + response=response, trace_id=trace_id, request_from=request.request_from, ) diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py index 65672afde6..b902660c7a 100644 --- a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py +++ b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock import pytest @@ -16,6 +17,26 @@ def relationship_recommendation_service(mock_pipeline): return RelationshipRecommendation(pipelines) +@pytest.fixture +def mdl_with_project_relationship_candidate(): + return """ + { + "models": [ + { + "name": "project", + "primaryKey": "id", + "columns": [{"name": "id"}, {"name": "name"}] + }, + { + "name": "view", + "columns": [{"name": "id"}, {"name": "project_id"}] + } + ], + "relationships": [] + } + """ + + @pytest.mark.asyncio async def test_recommend_success(relationship_recommendation_service, mock_pipeline): request = RelationshipRecommendation.Input(id="test_id", mdl='{"key": "value"}') @@ -87,6 +108,61 @@ def test_getitem_not_found(relationship_recommendation_service): assert "not found" in response.error.message +@pytest.mark.asyncio +async def test_recommend_timeout_returns_fallback_relationships( + mock_pipeline, mdl_with_project_relationship_candidate +): + service = RelationshipRecommendation( + {"relationship_recommendation": mock_pipeline}, + generation_timeout_seconds=0.01, + ) + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_project_relationship_candidate + ) + + async def never_finishes(**_kwargs): + await asyncio.sleep(1) + + mock_pipeline.run.side_effect = never_finishes + + await service.recommend(request) + response = service[request.id] + + assert response.status == "finished" + assert response.response == { + "relationships": [ + { + "name": "view_project", + "fromModel": "view", + "fromColumn": "project_id", + "type": "MANY_TO_ONE", + "toModel": "project", + "toColumn": "id", + "reason": "view.project_id appears to reference project.id.", + } + ] + } + + +@pytest.mark.asyncio +async def test_recommend_empty_llm_result_returns_fallback_relationships( + relationship_recommendation_service, + mock_pipeline, + mdl_with_project_relationship_candidate, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_project_relationship_candidate + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response["relationships"][0]["fromModel"] == "view" + assert response.response["relationships"][0]["toModel"] == "project" + + def test_setitem(relationship_recommendation_service): id = "test_id" value = RelationshipRecommendation.Resource(id="test_id", status="finished") From db090bbbfb8571834b4d7a8f3304accc344d1048 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 19:11:53 +0530 Subject: [PATCH 0568/1087] Broaden relationship fallback matching --- .../services/relationship_recommendation.py | 36 ++++++++++++-- .../test_relationship_recommendation.py | 49 +++++++++++++++++++ 2 files changed, 82 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py index 6aa6e9a1a0..5034770cca 100644 --- a/wren-ai-service/src/web/v1/services/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/services/relationship_recommendation.py @@ -50,6 +50,35 @@ def _normalize_identifier(self, value: Any) -> str: text = re.sub(r"[^a-zA-Z0-9]", "", text).lower() return text[:-1] if text.endswith("s") else text + def _identifier_tokens(self, value: Any) -> list[str]: + text = "" if value is None else str(value) + text = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", text) + return [ + self._normalize_identifier(token) + for token in re.split(r"[^a-zA-Z0-9]+", text) + if token + ] + + def _model_aliases(self, model: dict) -> set[str]: + aliases: set[str] = set() + raw_values = [ + model.get("name"), + model.get("properties", {}).get("displayName"), + model.get("tableReference", {}).get("table"), + ] + + for value in raw_values: + normalized = self._normalize_identifier(value) + if normalized: + aliases.add(normalized) + + tokens = self._identifier_tokens(value) + if tokens: + aliases.add(tokens[-1]) + aliases.add("".join(tokens)) + + return aliases + def _fallback_relationships(self, mdl: dict) -> dict: models = mdl.get("models", []) or [] existing = { @@ -62,9 +91,10 @@ def _fallback_relationships(self, mdl: dict) -> dict: } candidates = [] - model_lookup = { - self._normalize_identifier(model.get("name")): model for model in models - } + model_lookup = {} + for model in models: + for alias in self._model_aliases(model): + model_lookup.setdefault(alias, model) for from_model in models: from_model_name = from_model.get("name") diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py index b902660c7a..8ee52b7cc5 100644 --- a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py +++ b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py @@ -37,6 +37,27 @@ def mdl_with_project_relationship_candidate(): """ +@pytest.fixture +def mdl_with_prefixed_project_model(): + return """ + { + "models": [ + { + "name": "dbo_project", + "primaryKey": "id", + "tableReference": {"schema": "dbo", "table": "project"}, + "columns": [{"name": "id"}, {"name": "name"}] + }, + { + "name": "dbo_view", + "columns": [{"name": "id"}, {"name": "project_id"}] + } + ], + "relationships": [] + } + """ + + @pytest.mark.asyncio async def test_recommend_success(relationship_recommendation_service, mock_pipeline): request = RelationshipRecommendation.Input(id="test_id", mdl='{"key": "value"}') @@ -163,6 +184,34 @@ async def test_recommend_empty_llm_result_returns_fallback_relationships( assert response.response["relationships"][0]["toModel"] == "project" +@pytest.mark.asyncio +async def test_recommend_fallback_matches_prefixed_model_name( + relationship_recommendation_service, + mock_pipeline, + mdl_with_prefixed_project_model, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_prefixed_project_model + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response["relationships"] == [ + { + "name": "dbo_view_dbo_project", + "fromModel": "dbo_view", + "fromColumn": "project_id", + "type": "MANY_TO_ONE", + "toModel": "dbo_project", + "toColumn": "id", + "reason": "dbo_view.project_id appears to reference dbo_project.id.", + } + ] + + def test_setitem(relationship_recommendation_service): id = "test_id" value = RelationshipRecommendation.Resource(id="test_id", status="finished") From 24a2713f5e450bcd43c29056a0b814ebd6761f53 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 19:17:40 +0530 Subject: [PATCH 0569/1087] Refresh diagram before saving assistant relationships --- wren-ui/src/pages/modeling.tsx | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index a099e43a07..54c1b645c3 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -878,11 +878,18 @@ export default function Modeling() { } } if (assistantMode === 'relationships') { + const latestDiagramResult = await apolloClient.query({ + query: DIAGRAM, + fetchPolicy: 'network-only', + }); + const latestDiagramData = + latestDiagramResult.data?.diagram || diagramData; + for (const relationship of relationshipResult) { - const fromModel = diagramData.models.find( + const fromModel = latestDiagramData.models.find( (model) => model.referenceName === relationship.fromModel, ); - const toModel = diagramData.models.find( + const toModel = latestDiagramData.models.find( (model) => model.referenceName === relationship.toModel, ); const fromField = fromModel?.fields.find( @@ -892,7 +899,7 @@ export default function Modeling() { (field) => field.referenceName === relationship.toColumn, ); if (!fromModel || !toModel || !fromField || !toField) continue; - const alreadyExists = diagramData.models.some((model) => + const alreadyExists = latestDiagramData.models.some((model) => (model.relationFields || []).some((field) => { if (!field) return false; const forward = From bf69c884f98e1e5d8a5ef95f84e52830cf73bf8f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 19:12:08 +0530 Subject: [PATCH 0570/1087] Bound semantics job status lifetime --- .../src/web/v1/services/semantics_description.py | 16 ++++++++++++++-- .../services/test_semantics_description.py | 11 +++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index b0fbed77e3..c772bfda2c 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -13,6 +13,9 @@ logger = logging.getLogger("wren-ai-service") +MAX_UI_WAIT_SECONDS = 180 +SEMANTICS_STATUS_TTL_BUFFER_SECONDS = 300 + class SemanticsDescription: class Resource(BaseModel, MetadataTraceable): @@ -35,8 +38,17 @@ def __init__( generation_timeout_seconds: int = 90, ): self._pipelines = pipelines - self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) - self._generation_timeout_seconds = generation_timeout_seconds + self._generation_timeout_seconds = min( + generation_timeout_seconds, + MAX_UI_WAIT_SECONDS - 30, + ) + self._cache: Dict[str, self.Resource] = TTLCache( + maxsize=maxsize, + ttl=max( + ttl, + self._generation_timeout_seconds + SEMANTICS_STATUS_TTL_BUFFER_SECONDS, + ), + ) def _handle_exception( self, diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 73edd26870..fcbcbf1eb0 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -164,6 +164,17 @@ def test_semantics_description_uses_configured_timeout(): assert service._generation_timeout_seconds == 123 +def test_semantics_description_caps_timeout_inside_ui_polling_window(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + ttl=120, + generation_timeout_seconds=600, + ) + + assert service._generation_timeout_seconds == 150 + assert service._cache.ttl >= 450 + + @pytest.mark.asyncio async def test_batch_processing_with_multiple_models( service: SemanticsDescription, From 0ffa5900f343ff4dad7955fa8660f6984eb4e7f7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 19:24:13 +0530 Subject: [PATCH 0571/1087] Handle partial assistant relationship saves --- wren-ui/src/pages/modeling.tsx | 63 +++++++++++++++++++++++++++------- 1 file changed, 50 insertions(+), 13 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 54c1b645c3..feb1cf09d8 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -365,6 +365,12 @@ export default function Modeling() { }, }), ); + const [createAssistantRelationshipMutation] = useCreateRelationshipMutation( + getBaseOptions({ + onError: null, + onCompleted: () => {}, + }), + ); const [deleteRelationshipMutation] = useDeleteRelationshipMutation( getBaseOptions({ @@ -878,6 +884,8 @@ export default function Modeling() { } } if (assistantMode === 'relationships') { + let createdCount = 0; + let skippedCount = 0; const latestDiagramResult = await apolloClient.query({ query: DIAGRAM, fetchPolicy: 'network-only', @@ -898,7 +906,10 @@ export default function Modeling() { const toField = toModel?.fields.find( (field) => field.referenceName === relationship.toColumn, ); - if (!fromModel || !toModel || !fromField || !toField) continue; + if (!fromModel || !toModel || !fromField || !toField) { + skippedCount += 1; + continue; + } const alreadyExists = latestDiagramData.models.some((model) => (model.relationFields || []).some((field) => { if (!field) return false; @@ -915,22 +926,48 @@ export default function Modeling() { return forward || reverse; }), ); - if (alreadyExists) continue; - await createRelationshipMutation({ - variables: { - data: { - fromModelId: fromModel.modelId, - fromColumnId: fromField.columnId, - toModelId: toModel.modelId, - toColumnId: toField.columnId, - type: relationship.type, + if (alreadyExists) { + skippedCount += 1; + continue; + } + try { + await createAssistantRelationshipMutation({ + variables: { + data: { + fromModelId: fromModel.modelId, + fromColumnId: fromField.columnId, + toModelId: toModel.modelId, + toColumnId: toField.columnId, + type: relationship.type, + }, }, - }, - }); + }); + createdCount += 1; + } catch (_error) { + skippedCount += 1; + } + } + + if (!createdCount) { + throw new Error( + skippedCount + ? 'No valid new relationships could be saved for the current models.' + : 'No relationship suggestions to save.', + ); + } + + if (skippedCount) { + message.warning( + `${createdCount} relationship(s) saved. ${skippedCount} invalid or duplicate suggestion(s) skipped.`, + ); + } else { + message.success('Saved Modeling AI Assistant suggestions.'); } } closeAssistant(); - message.success('Saved Modeling AI Assistant suggestions.'); + if (assistantMode !== 'relationships') { + message.success('Saved Modeling AI Assistant suggestions.'); + } } catch (error: any) { message.error(error.message || 'Failed to save assistant suggestions.'); } finally { From dd82fef1c3a7ed445a6f5add215a6c4e03176bd3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 19:34:03 +0530 Subject: [PATCH 0572/1087] Save assistant relationships by reference name --- wren-ui/src/apollo/server/resolvers.ts | 1 + .../apollo/server/resolvers/modelResolver.ts | 65 +++++++++++++ wren-ui/src/apollo/server/schema.ts | 9 ++ .../src/apollo/server/types/relationship.ts | 8 ++ wren-ui/src/pages/modeling.tsx | 92 +++++-------------- 5 files changed, 106 insertions(+), 69 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 168bda10e1..0227636968 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -153,6 +153,7 @@ const resolvers = { updateViewMetadata: modelResolver.updateViewMetadata, generateModelingSemantics: modelResolver.generateModelingSemantics, generateModelingRelationships: modelResolver.generateModelingRelationships, + saveModelingRelationships: modelResolver.saveModelingRelationships, // Settings resetCurrentProject: projectResolver.resetCurrentProject, diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 5fdfb334e2..a0a43f42ff 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -10,6 +10,7 @@ import { import { DataSourceName, IContext, + ModelingRelationshipData, RelationData, UpdateRelationData, } from '../types'; @@ -64,6 +65,7 @@ export class ModelResolver { this.generateModelingRelationships.bind(this); this.getModelingRelationshipsResult = this.getModelingRelationshipsResult.bind(this); + this.saveModelingRelationships = this.saveModelingRelationships.bind(this); this.checkModelSync = this.checkModelSync.bind(this); // view @@ -500,6 +502,69 @@ export class ModelResolver { ); } + public async saveModelingRelationships( + _root: any, + args: { data: ModelingRelationshipData[] }, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const models = await ctx.modelRepository.findAllBy({ + projectId: project.id, + }); + const modelIds = models.map((model) => model.id); + const columns = + await ctx.modelColumnRepository.findColumnsByModelIds(modelIds); + let createdCount = 0; + let skippedCount = 0; + + for (const relationship of args.data || []) { + const fromModel = models.find( + (model) => model.referenceName === relationship.fromModel, + ); + const toModel = models.find( + (model) => model.referenceName === relationship.toModel, + ); + const fromColumn = fromModel + ? columns.find( + (column) => + column.modelId === fromModel.id && + column.referenceName === relationship.fromColumn, + ) + : null; + const toColumn = toModel + ? columns.find( + (column) => + column.modelId === toModel.id && + column.referenceName === relationship.toColumn, + ) + : null; + + if (!fromModel || !toModel || !fromColumn || !toColumn) { + skippedCount += 1; + continue; + } + + try { + const savedRelation = await ctx.modelService.createRelation({ + fromModelId: fromModel.id, + fromColumnId: fromColumn.id, + toModelId: toModel.id, + toColumnId: toColumn.id, + type: relationship.type, + }); + this.markProjectDirty(savedRelation.projectId); + createdCount += 1; + } catch (err: any) { + logger.warn( + `Skip Modeling AI Assistant relationship ${relationship.fromModel}.${relationship.fromColumn} -> ${relationship.toModel}.${relationship.toColumn}: ${err.message}`, + ); + skippedCount += 1; + } + } + + return { createdCount, skippedCount }; + } + public async listModels(_root: any, _args: any, ctx: IContext) { const { id: projectId } = await ctx.projectService.getCurrentProject(); const models = await ctx.modelRepository.findAllBy({ projectId }); diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 2700705cbd..1c3dad1cc7 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -260,6 +260,14 @@ export const typeDefs = gql` relations: [RelationInput]! } + input ModelingRelationshipInput { + fromModel: String! + fromColumn: String! + toModel: String! + toColumn: String! + type: RelationType! + } + input SaveTablesInput { tables: [String!]! } @@ -1333,6 +1341,7 @@ export const typeDefs = gql` ): Boolean! generateModelingSemantics(data: GenerateModelingSemanticsInput!): JSON! generateModelingRelationships: JSON! + saveModelingRelationships(data: [ModelingRelationshipInput!]!): JSON! # Relation createRelation(data: RelationInput!): JSON! diff --git a/wren-ui/src/apollo/server/types/relationship.ts b/wren-ui/src/apollo/server/types/relationship.ts index c57103e443..de88970899 100644 --- a/wren-ui/src/apollo/server/types/relationship.ts +++ b/wren-ui/src/apollo/server/types/relationship.ts @@ -11,6 +11,14 @@ export interface UpdateRelationData { type: RelationType; } +export interface ModelingRelationshipData { + fromModel: string; + fromColumn: string; + toModel: string; + toColumn: string; + type: RelationType; +} + export interface AnalysisRelationInfo { name: string; fromModelId: number; diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index feb1cf09d8..97dcebd732 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -88,6 +88,12 @@ const MODELING_RELATIONSHIPS_RESULT = gql` } `; +const SAVE_MODELING_RELATIONSHIPS = gql` + mutation SaveModelingRelationships($data: [ModelingRelationshipInput!]!) { + saveModelingRelationships(data: $data) + } +`; + const Diagram = dynamic(() => import('@/components/diagram'), { ssr: false }); // https://github.com/vercel/next.js/issues/4957#issuecomment-413841689 const ForwardDiagram = forwardRef(function ForwardDiagram(props: any, ref) { @@ -365,13 +371,6 @@ export default function Modeling() { }, }), ); - const [createAssistantRelationshipMutation] = useCreateRelationshipMutation( - getBaseOptions({ - onError: null, - onCompleted: () => {}, - }), - ); - const [deleteRelationshipMutation] = useDeleteRelationshipMutation( getBaseOptions({ onCompleted: () => { @@ -401,6 +400,7 @@ export default function Modeling() { const [generateModelingRelationships] = useMutation( GENERATE_MODELING_RELATIONSHIPS, ); + const [saveModelingRelationships] = useMutation(SAVE_MODELING_RELATIONSHIPS); const diagramData = useMemo(() => { if (!data) return null; @@ -884,69 +884,23 @@ export default function Modeling() { } } if (assistantMode === 'relationships') { - let createdCount = 0; - let skippedCount = 0; - const latestDiagramResult = await apolloClient.query({ - query: DIAGRAM, - fetchPolicy: 'network-only', + const res = await saveModelingRelationships({ + variables: { + data: relationshipResult.map((relationship) => ({ + fromModel: relationship.fromModel, + fromColumn: relationship.fromColumn, + toModel: relationship.toModel, + toColumn: relationship.toColumn, + type: relationship.type, + })), + }, + refetchQueries, + awaitRefetchQueries: true, }); - const latestDiagramData = - latestDiagramResult.data?.diagram || diagramData; - - for (const relationship of relationshipResult) { - const fromModel = latestDiagramData.models.find( - (model) => model.referenceName === relationship.fromModel, - ); - const toModel = latestDiagramData.models.find( - (model) => model.referenceName === relationship.toModel, - ); - const fromField = fromModel?.fields.find( - (field) => field.referenceName === relationship.fromColumn, - ); - const toField = toModel?.fields.find( - (field) => field.referenceName === relationship.toColumn, - ); - if (!fromModel || !toModel || !fromField || !toField) { - skippedCount += 1; - continue; - } - const alreadyExists = latestDiagramData.models.some((model) => - (model.relationFields || []).some((field) => { - if (!field) return false; - const forward = - field.fromModelName === relationship.fromModel && - field.fromColumnName === relationship.fromColumn && - field.toModelName === relationship.toModel && - field.toColumnName === relationship.toColumn; - const reverse = - field.fromModelName === relationship.toModel && - field.fromColumnName === relationship.toColumn && - field.toModelName === relationship.fromModel && - field.toColumnName === relationship.fromColumn; - return forward || reverse; - }), - ); - if (alreadyExists) { - skippedCount += 1; - continue; - } - try { - await createAssistantRelationshipMutation({ - variables: { - data: { - fromModelId: fromModel.modelId, - fromColumnId: fromField.columnId, - toModelId: toModel.modelId, - toColumnId: toField.columnId, - type: relationship.type, - }, - }, - }); - createdCount += 1; - } catch (_error) { - skippedCount += 1; - } - } + const createdCount = + res.data?.saveModelingRelationships?.createdCount || 0; + const skippedCount = + res.data?.saveModelingRelationships?.skippedCount || 0; if (!createdCount) { throw new Error( From 783d11da143e054b02d686c504570e15d0b9bbeb Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 20:28:09 +0530 Subject: [PATCH 0573/1087] Generate selected semantics concurrently --- .../web/v1/services/semantics_description.py | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index c772bfda2c..53ec69da2f 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -15,6 +15,8 @@ MAX_UI_WAIT_SECONDS = 180 SEMANTICS_STATUS_TTL_BUFFER_SECONDS = 300 +SEMANTICS_MODEL_CHUNK_SIZE = 200 +SEMANTICS_MAX_CONCURRENT_LLM_CALLS = 3 class SemanticsDescription: @@ -125,7 +127,10 @@ def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: return output def _chunking( - self, mdl_dict: dict, request: GenerateRequest, chunk_size: int = 50 + self, + mdl_dict: dict, + request: GenerateRequest, + chunk_size: int = SEMANTICS_MODEL_CHUNK_SIZE, ) -> list[dict]: template = { "user_prompt": request.user_prompt, @@ -218,6 +223,15 @@ async def _generate_task(self, request_id: str, chunk: dict): current.response[key].setdefault("columns", []) current.response[key]["columns"].extend(output[key].get("columns", [])) + async def _generate_task_with_semaphore( + self, + semaphore: asyncio.Semaphore, + request_id: str, + chunk: dict, + ): + async with semaphore: + await self._generate_task(request_id, chunk) + @observe(name="Generate Semantics Description") @trace_metadata async def generate(self, request: GenerateRequest, **kwargs) -> Resource: @@ -232,8 +246,17 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: raise ValueError( "No selected models matched the current semantic model metadata" ) - for chunk in chunks: - await self._generate_task(request.id, chunk) + semaphore = asyncio.Semaphore(SEMANTICS_MAX_CONCURRENT_LLM_CALLS) + await asyncio.gather( + *[ + self._generate_task_with_semaphore( + semaphore, + request.id, + chunk, + ) + for chunk in chunks + ] + ) self[request.id].status = "finished" self[request.id].trace_id = trace_id From fa8c55edff0b7b1fa53ec9b90b6b52593d54f33a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 20:57:38 +0530 Subject: [PATCH 0574/1087] Support typed relationship fallback --- .../services/relationship_recommendation.py | 20 +++++++- .../test_relationship_recommendation.py | 50 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py index 5034770cca..4e52b5583c 100644 --- a/wren-ai-service/src/web/v1/services/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/services/relationship_recommendation.py @@ -79,6 +79,19 @@ def _model_aliases(self, model: dict) -> set[str]: return aliases + def _column_is_primary_key(self, model: dict, column_name: str) -> bool: + primary_key = model.get("primaryKey") + if primary_key and column_name == primary_key: + return True + + normalized_column = self._normalize_identifier(column_name) + return normalized_column == "id" + + def _fallback_relationship_type(self, from_model: dict, from_column: str) -> str: + if self._column_is_primary_key(from_model, from_column): + return "ONE_TO_ONE" + return "MANY_TO_ONE" + def _fallback_relationships(self, mdl: dict) -> dict: models = mdl.get("models", []) or [] existing = { @@ -140,10 +153,13 @@ def _fallback_relationships(self, mdl: dict) -> dict: if not to_column: continue + relationship_type = self._fallback_relationship_type( + from_model, from_column + ) signature = ( from_model_name, f"{from_model_name}.{from_column} = {to_model_name}.{to_column}", - "MANY_TO_ONE", + relationship_type, ) if signature in existing: continue @@ -153,7 +169,7 @@ def _fallback_relationships(self, mdl: dict) -> dict: "name": f"{from_model_name}_{to_model_name}", "fromModel": from_model_name, "fromColumn": from_column, - "type": "MANY_TO_ONE", + "type": relationship_type, "toModel": to_model_name, "toColumn": to_column, "reason": ( diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py index 8ee52b7cc5..de480e8ad6 100644 --- a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py +++ b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py @@ -58,6 +58,27 @@ def mdl_with_prefixed_project_model(): """ +@pytest.fixture +def mdl_with_one_to_one_profile_candidate(): + return """ + { + "models": [ + { + "name": "user", + "primaryKey": "id", + "columns": [{"name": "id"}, {"name": "email"}] + }, + { + "name": "profile", + "primaryKey": "user_id", + "columns": [{"name": "user_id"}, {"name": "display_name"}] + } + ], + "relationships": [] + } + """ + + @pytest.mark.asyncio async def test_recommend_success(relationship_recommendation_service, mock_pipeline): request = RelationshipRecommendation.Input(id="test_id", mdl='{"key": "value"}') @@ -182,6 +203,7 @@ async def test_recommend_empty_llm_result_returns_fallback_relationships( assert response.status == "finished" assert response.response["relationships"][0]["fromModel"] == "view" assert response.response["relationships"][0]["toModel"] == "project" + assert response.response["relationships"][0]["type"] == "MANY_TO_ONE" @pytest.mark.asyncio @@ -212,6 +234,34 @@ async def test_recommend_fallback_matches_prefixed_model_name( ] +@pytest.mark.asyncio +async def test_recommend_fallback_identifies_one_to_one_relationships( + relationship_recommendation_service, + mock_pipeline, + mdl_with_one_to_one_profile_candidate, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_one_to_one_profile_candidate + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response["relationships"] == [ + { + "name": "profile_user", + "fromModel": "profile", + "fromColumn": "user_id", + "type": "ONE_TO_ONE", + "toModel": "user", + "toColumn": "id", + "reason": "profile.user_id appears to reference user.id.", + } + ] + + def test_setitem(relationship_recommendation_service): id = "test_id" value = RelationshipRecommendation.Resource(id="test_id", status="finished") From 0c1bb024fbc5e109792d6f7617aaeb27d02bd96d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 21:20:41 +0530 Subject: [PATCH 0575/1087] Expand datasource relationship fallback --- .../services/relationship_recommendation.py | 262 +++++++++++++----- .../test_relationship_recommendation.py | 75 +++++ 2 files changed, 262 insertions(+), 75 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py index 4e52b5583c..673d7a621c 100644 --- a/wren-ai-service/src/web/v1/services/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/services/relationship_recommendation.py @@ -79,105 +79,217 @@ def _model_aliases(self, model: dict) -> set[str]: return aliases - def _column_is_primary_key(self, model: dict, column_name: str) -> bool: + def _model_columns(self, model: dict) -> list[dict]: + return [ + column + for column in model.get("columns", []) or [] + if column.get("name") and not column.get("relationship") + ] + + def _primary_key(self, model: dict) -> Optional[str]: primary_key = model.get("primaryKey") + columns = self._model_columns(model) + if primary_key and any(column.get("name") == primary_key for column in columns): + return primary_key + + model_aliases = self._model_aliases(model) + for column in columns: + normalized_column = self._normalize_identifier(column.get("name")) + if normalized_column == "id" or normalized_column in { + f"{alias}id" for alias in model_aliases + }: + return column.get("name") + + return None + + def _column_is_primary_key(self, model: dict, column_name: str) -> bool: + primary_key = self._primary_key(model) if primary_key and column_name == primary_key: return True - normalized_column = self._normalize_identifier(column_name) - return normalized_column == "id" + return False - def _fallback_relationship_type(self, from_model: dict, from_column: str) -> str: - if self._column_is_primary_key(from_model, from_column): + def _fallback_relationship_type( + self, from_model: dict, from_column: str, to_model: dict, to_column: str + ) -> str: + from_is_pk = self._column_is_primary_key(from_model, from_column) + to_is_pk = self._column_is_primary_key(to_model, to_column) + if from_is_pk and to_is_pk: return "ONE_TO_ONE" + if from_is_pk and not to_is_pk: + return "ONE_TO_MANY" return "MANY_TO_ONE" + def _relationship_signature( + self, + from_model_name: str, + from_column: str, + to_model_name: str, + to_column: str, + ) -> tuple[str, str, str, str]: + return (from_model_name, from_column, to_model_name, to_column) + + def _relationship_pair_signature( + self, + from_model_name: str, + from_column: str, + to_model_name: str, + to_column: str, + ) -> tuple[tuple[str, str], tuple[str, str]]: + left = (from_model_name, from_column) + right = (to_model_name, to_column) + return tuple(sorted([left, right])) + + def _existing_relationship_signatures( + self, mdl: dict + ) -> tuple[ + set[tuple[str, str, str, str]], set[tuple[tuple[str, str], tuple[str, str]]] + ]: + direct_signatures = set() + pair_signatures = set() + + for relationship in mdl.get("relationships", []) or []: + models = relationship.get("models", []) or [] + condition = relationship.get("condition", "") + if len(models) < 2 or not condition: + continue + + match = re.match( + r"\s*([^.=\s]+)\.([^.=\s]+)\s*=\s*([^.=\s]+)\.([^.=\s]+)\s*", + condition, + ) + if not match: + continue + + left_model, left_column, right_model, right_column = match.groups() + direct_signatures.add( + self._relationship_signature( + left_model, left_column, right_model, right_column + ) + ) + direct_signatures.add( + self._relationship_signature( + right_model, right_column, left_model, left_column + ) + ) + pair_signatures.add( + self._relationship_pair_signature( + left_model, left_column, right_model, right_column + ) + ) + + return direct_signatures, pair_signatures + def _fallback_relationships(self, mdl: dict) -> dict: models = mdl.get("models", []) or [] - existing = { - ( - relationship.get("models", [None, None])[0], - relationship.get("condition", ""), - relationship.get("joinType", ""), - ) - for relationship in mdl.get("relationships", []) or [] - } + existing, existing_pairs = self._existing_relationship_signatures(mdl) + seen = set(existing) + seen_pairs = set(existing_pairs) candidates = [] model_lookup = {} + primary_keys = {} for model in models: + primary_keys[model.get("name")] = self._primary_key(model) for alias in self._model_aliases(model): model_lookup.setdefault(alias, model) + def add_candidate( + from_model: dict, + from_column: str, + to_model: dict, + to_column: str, + ): + from_model_name = from_model.get("name") + to_model_name = to_model.get("name") + if not from_model_name or not to_model_name: + return + if from_model_name == to_model_name: + return + + signature = self._relationship_signature( + from_model_name, from_column, to_model_name, to_column + ) + pair_signature = self._relationship_pair_signature( + from_model_name, from_column, to_model_name, to_column + ) + if signature in seen or pair_signature in seen_pairs: + return + + relationship_type = self._fallback_relationship_type( + from_model, from_column, to_model, to_column + ) + reason = ( + f"{from_model_name}.{from_column} appears to reference " + f"{to_model_name}.{to_column}." + ) + if relationship_type == "ONE_TO_MANY": + reason = ( + f"{from_model_name}.{from_column} appears to be referenced by " + f"{to_model_name}.{to_column}." + ) + + seen.add(signature) + seen_pairs.add(pair_signature) + candidates.append( + { + "name": f"{from_model_name}_{to_model_name}", + "fromModel": from_model_name, + "fromColumn": from_column, + "type": relationship_type, + "toModel": to_model_name, + "toColumn": to_column, + "reason": reason, + } + ) + for from_model in models: from_model_name = from_model.get("name") if not from_model_name: continue - for column in from_model.get("columns", []) or []: - if column.get("relationship"): - continue - + for column in self._model_columns(from_model): from_column = column.get("name") - if not from_column: - continue - normalized_column = self._normalize_identifier(from_column) - if not normalized_column.endswith("id") or normalized_column == "id": - continue - - target_key = normalized_column[:-2] - to_model = model_lookup.get(target_key) - if not to_model or to_model.get("name") == from_model_name: - continue - - to_model_name = to_model.get("name") - to_columns = to_model.get("columns", []) or [] - primary_key = to_model.get("primaryKey") - to_column = next( - ( - item.get("name") - for item in to_columns - if item.get("name") == primary_key - ), - None, - ) - to_column = to_column or next( - ( - item.get("name") - for item in to_columns - if self._normalize_identifier(item.get("name")) == "id" - ), - None, - ) - if not to_column: - continue - - relationship_type = self._fallback_relationship_type( - from_model, from_column - ) - signature = ( - from_model_name, - f"{from_model_name}.{from_column} = {to_model_name}.{to_column}", - relationship_type, - ) - if signature in existing: - continue - - candidates.append( - { - "name": f"{from_model_name}_{to_model_name}", - "fromModel": from_model_name, - "fromColumn": from_column, - "type": relationship_type, - "toModel": to_model_name, - "toColumn": to_column, - "reason": ( - f"{from_model_name}.{from_column} appears to reference " - f"{to_model_name}.{to_column}." - ), - } - ) + target_keys = set() + if normalized_column.endswith("id") and normalized_column != "id": + target_keys.add(normalized_column[:-2]) + + for to_model in models: + to_model_name = to_model.get("name") + to_primary_key = primary_keys.get(to_model_name) + if ( + to_model_name == from_model_name + or not to_primary_key + or not self._column_is_primary_key(to_model, to_primary_key) + ): + continue + + to_primary_key_normalized = self._normalize_identifier( + to_primary_key + ) + if ( + normalized_column != "id" + and normalized_column == to_primary_key_normalized + ): + add_candidate( + to_model, + to_primary_key, + from_model, + from_column, + ) + + for target_key in target_keys: + to_model = model_lookup.get(target_key) + if not to_model or to_model.get("name") == from_model_name: + continue + + to_model_name = to_model.get("name") + to_column = primary_keys.get(to_model_name) + if not to_column: + continue + + add_candidate(from_model, from_column, to_model, to_column) return {"relationships": candidates} diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py index de480e8ad6..faea1d2451 100644 --- a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py +++ b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py @@ -79,6 +79,35 @@ def mdl_with_one_to_one_profile_candidate(): """ +@pytest.fixture +def mdl_with_shared_key_candidates(): + return """ + { + "models": [ + { + "name": "employees", + "primaryKey": "emp_no", + "columns": [{"name": "emp_no"}, {"name": "first_name"}] + }, + { + "name": "titles", + "columns": [{"name": "emp_no"}, {"name": "title"}] + }, + { + "name": "departments", + "primaryKey": "dept_no", + "columns": [{"name": "dept_no"}, {"name": "dept_name"}] + }, + { + "name": "dept_emp", + "columns": [{"name": "emp_no"}, {"name": "dept_no"}] + } + ], + "relationships": [] + } + """ + + @pytest.mark.asyncio async def test_recommend_success(relationship_recommendation_service, mock_pipeline): request = RelationshipRecommendation.Input(id="test_id", mdl='{"key": "value"}') @@ -262,6 +291,52 @@ async def test_recommend_fallback_identifies_one_to_one_relationships( ] +@pytest.mark.asyncio +async def test_recommend_fallback_scans_all_models_and_identifies_one_to_many( + relationship_recommendation_service, + mock_pipeline, + mdl_with_shared_key_candidates, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_shared_key_candidates + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response["relationships"] == [ + { + "name": "employees_titles", + "fromModel": "employees", + "fromColumn": "emp_no", + "type": "ONE_TO_MANY", + "toModel": "titles", + "toColumn": "emp_no", + "reason": "employees.emp_no appears to be referenced by titles.emp_no.", + }, + { + "name": "employees_dept_emp", + "fromModel": "employees", + "fromColumn": "emp_no", + "type": "ONE_TO_MANY", + "toModel": "dept_emp", + "toColumn": "emp_no", + "reason": "employees.emp_no appears to be referenced by dept_emp.emp_no.", + }, + { + "name": "departments_dept_emp", + "fromModel": "departments", + "fromColumn": "dept_no", + "type": "ONE_TO_MANY", + "toModel": "dept_emp", + "toColumn": "dept_no", + "reason": "departments.dept_no appears to be referenced by dept_emp.dept_no.", + }, + ] + + def test_setitem(relationship_recommendation_service): id = "test_id" value = RelationshipRecommendation.Resource(id="test_id", status="finished") From cd3ccaa0af6c2fb146c932f1f21239914b203168 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 21:22:13 +0530 Subject: [PATCH 0576/1087] Improve large semantic generation batching --- .../generation/semantics_description.py | 12 +- .../web/v1/services/semantics_description.py | 65 ++++++++++- .../services/test_semantics_description.py | 110 +++++++++++++++++- 3 files changed, 172 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index d197bf5017..ad72c659c5 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -36,11 +36,13 @@ ``` Your task is to update this JSON structure by adding a `description` field inside both the `properties` attribute of each `column` and the `model` itself. -Each `description` should be derived from a user-provided input that explains the purpose or context of the `model` and its respective columns. +Each `description` should be derived from the user-provided dataset context, the model name, column names, data types, aliases, and any existing descriptions. Follow these steps: -1. **For the `model`**: Prompt the user to provide a brief description of the model's overall purpose or its context. Insert this description in the `properties` field of the `model`. -2. **For each `column`**: Ask the user to describe each column's role or significance. Each column's description should be added under its respective `properties` field in the format: `'description': 'user-provided text'`. +1. **For the `model`**: Write a clear natural language business description of the model's purpose and what real-world records it represents. Insert this description in the `properties` field of the `model`. +2. **For each `column`**: Write a clear natural language business description of the column's meaning, not just its technical name. Each column's description should be added under its respective `properties` field in the format: `'description': 'business description'`. 3. Ensure that the output is a well-formatted JSON structure, preserving the input's original format and adding the appropriate `description` fields. +4. Avoid repeating technical table or column names as the whole description. Prefer business meaning such as identifiers, dates, amounts, statuses, dimensions, ownership, and operational usage. +5. Keep descriptions concise, factual, and useful for text-to-SQL retrieval. ### Output Format: @@ -77,7 +79,7 @@ } ``` -Make sure that the descriptions are concise, informative, and contextually appropriate based on the input provided by the user. +Make sure that the descriptions are concise, informative, business-friendly, and contextually appropriate based on the input provided by the user. """ user_prompt_template = """ @@ -86,7 +88,7 @@ Picked models: {{ picked_models }} Localization Language: {{ language }} -Please provide a brief description for the model and each column based on the user's prompt. +Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt. Do not omit selected models or columns. """ diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 53ec69da2f..eb644ca1c6 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -15,8 +15,8 @@ MAX_UI_WAIT_SECONDS = 180 SEMANTICS_STATUS_TTL_BUFFER_SECONDS = 300 -SEMANTICS_MODEL_CHUNK_SIZE = 200 -SEMANTICS_MAX_CONCURRENT_LLM_CALLS = 3 +SEMANTICS_MODEL_CHUNK_SIZE = 1000 +SEMANTICS_MAX_CONCURRENT_LLM_CALLS = 6 class SemanticsDescription: @@ -82,6 +82,13 @@ def _properties(self, payload: dict[str, Any]) -> dict[str, Any]: def _text(self, value: Any) -> str: return "" if value is None else str(value) + def _humanize_name(self, name: str) -> str: + return " ".join( + token + for token in name.replace(".", " ").replace("_", " ").split() + if token.lower() not in {"dbo", "public"} + ) or name + def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: output: dict[str, Any] = {} for model in chunk.get("mdl", {}).get("models", []): @@ -89,11 +96,12 @@ def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: if not model_name: continue + model_label = self._humanize_name(model_name) model_properties = self._properties(model) model_description = self._text(model_properties.get("description", "")) if not model_description: model_description = ( - f"Represents {model_name.replace('_', ' ')} records in this dataset." + f"Contains business records for {model_label}, used for reporting, analysis, and operational questions." ) columns = [] @@ -108,8 +116,9 @@ def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: column_properties.get("description", "") ) if not column_description: + column_label = self._humanize_name(column_name) column_description = ( - f"{column_name.replace('_', ' ')} field from {model_name}." + f"Stores the {column_label} value used to describe or analyze {model_label} records." ) columns.append( { @@ -126,6 +135,52 @@ def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: } return output + def _complete_output_with_fallback( + self, + output: dict[str, Any], + chunk: dict[str, Any], + ) -> dict[str, Any]: + fallback = self._fallback_output(chunk) + completed = dict(output) + + for model_name, fallback_model in fallback.items(): + model_output = completed.get(model_name) + if not isinstance(model_output, dict): + completed[model_name] = fallback_model + continue + + properties = model_output.get("properties") + if not isinstance(properties, dict): + properties = {} + model_output["properties"] = properties + if not properties.get("description"): + properties["description"] = self._text( + model_output.get("description") + ) or fallback_model["properties"]["description"] + + output_columns = { + column.get("name"): column + for column in model_output.get("columns", []) + if isinstance(column, dict) and column.get("name") + } + for fallback_column in fallback_model.get("columns", []): + column_name = fallback_column.get("name") + output_column = output_columns.get(column_name) + if not output_column: + model_output.setdefault("columns", []).append(fallback_column) + continue + + column_properties = output_column.get("properties") + if not isinstance(column_properties, dict): + column_properties = {} + output_column["properties"] = column_properties + if not column_properties.get("description"): + column_properties["description"] = self._text( + output_column.get("description") + ) or fallback_column["properties"]["description"] + + return completed + def _chunking( self, mdl_dict: dict, @@ -201,6 +256,8 @@ async def _generate_task(self, request_id: str, chunk: dict): "returning metadata-based fallback descriptions." ) output = self._fallback_output(chunk) + else: + output = self._complete_output_with_fallback(output, chunk) except TimeoutError: logger.warning( "Semantics description LLM call timed out after %s seconds; " diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index fcbcbf1eb0..4b54a712fd 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -42,7 +42,15 @@ async def test_generate_semantics_description( assert response.status == "finished" assert response.response == { "model1": { - "columns": [], + "columns": [ + { + "name": "column1", + "type": "varchar", + "properties": { + "description": "Stores the column1 value used to describe or analyze model1 records." + }, + } + ], "properties": {"description": "Test description"}, } } @@ -200,11 +208,12 @@ async def test_batch_processing_with_multiple_models( assert response.id == "test_id" assert response.status == "finished" - assert response.response == { - "model1": {"description": "Description 1"}, - "model2": {"description": "Description 2"}, - "model3": {"description": "Description 3"}, - } + assert response.response["model1"]["properties"]["description"] == "Description 1" + assert response.response["model2"]["properties"]["description"] == "Description 2" + assert response.response["model3"]["properties"]["description"] == "Description 3" + assert len(response.response["model1"]["columns"]) == 1 + assert len(response.response["model2"]["columns"]) == 1 + assert len(response.response["model3"]["columns"]) == 1 chunks = service._chunking(orjson.loads(request.mdl), request) assert len(chunks) == 1 @@ -233,6 +242,95 @@ def test_batch_processing_with_custom_chunk_size( assert chunks[1]["selected_models"] == ["model3", "model4"] +def test_default_batch_allows_large_column_groups( + service: SemanticsDescription, +): + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the models", + selected_models=["model1", "model2"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "model1", + "columns": [ + {"name": f"column_{index}", "type": "varchar"} + for index in range(500) + ], + }, + { + "name": "model2", + "columns": [ + {"name": f"field_{index}", "type": "varchar"} + for index in range(400) + ], + }, + ] + } + ).decode(), + ) + + chunks = service._chunking(orjson.loads(request.mdl), request) + + assert len(chunks) == 1 + assert chunks[0]["selected_models"] == ["model1", "model2"] + + +@pytest.mark.asyncio +async def test_partial_llm_output_is_completed_for_all_selected_columns( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the models", + selected_models=["orders", "customers"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "order_id", "type": "varchar"}, + {"name": "order_date", "type": "date"}, + ], + }, + { + "name": "customers", + "columns": [ + {"name": "customer_id", "type": "varchar"}, + ], + }, + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "orders": { + "name": "orders", + "columns": [ + { + "name": "order_id", + "properties": {"description": "Unique order identifier."}, + } + ], + "properties": {"description": "Customer purchase transactions."}, + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert set(response.response.keys()) == {"orders", "customers"} + assert len(response.response["orders"]["columns"]) == 2 + assert len(response.response["customers"]["columns"]) == 1 + assert response.response["customers"]["properties"]["description"] + + @pytest.mark.asyncio async def test_batch_processing_partial_failure( service: SemanticsDescription, From 118b812e098ae7e50fa7069ee65e311e66a4aab9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 21:37:30 +0530 Subject: [PATCH 0577/1087] Improve relationship recommendation descriptions --- .../generation/relationship_recommendation.py | 18 +-- .../services/relationship_recommendation.py | 150 +++++++++++++++++- .../test_relationship_recommendation.py | 65 +++++++- 3 files changed, 207 insertions(+), 26 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/relationship_recommendation.py b/wren-ai-service/src/pipelines/generation/relationship_recommendation.py index e0d0bed675..c4d73b4d6e 100644 --- a/wren-ai-service/src/pipelines/generation/relationship_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/relationship_recommendation.py @@ -27,13 +27,16 @@ - **type**: The type of relationship, which can be "MANY_TO_ONE", "ONE_TO_MANY" or "ONE_TO_ONE" only. - **toModel**: The name of the target model. - **toColumn**: The column in the target model that forms the relationship. -- **reason**: The reason for recommending this relationship. +- **reason**: A clear natural language business description of what the relationship means and why it is useful. Important guidelines: 1. Do not recommend relationships within the same model (fromModel and toModel must be different). 2. Only suggest relationships if there is a clear and beneficial reason to do so. 3. If there are no good relationships to recommend or if there are fewer than two models, return an empty list of relationships. 4. Use "MANY_TO_ONE" and "ONE_TO_MANY" instead of "MANY_TO_MANY" relationships. +5. Write the reason for business users. Do not merely repeat raw table names, model names, or column names. +6. Use available model names, display names, column names, descriptions, primary keys, and table context to explain how the entities are related. +7. Prefer descriptions like "Each order belongs to one customer, so revenue can be analyzed by customer." over descriptions like "orders.customer_id references customers.id." Output all relationships in the following JSON structure: @@ -46,7 +49,7 @@ "type": "", "toModel": "", "toColumn": "", - "reason": "" + "reason": "" } ... ] @@ -74,26 +77,17 @@ ## Start of Pipeline @observe(capture_input=False) def cleaned_models(mdl: dict) -> dict: - def remove_display_name(d: dict) -> dict: - if "properties" in d and isinstance(d["properties"], dict): - d["properties"] = d["properties"].copy() - d["properties"].pop("displayName", None) - return d - def column_filter(columns: list[dict]) -> list[dict]: filtered_columns = [] for column in columns: if "relationship" not in column: # Create a copy of the column to avoid modifying the original filtered_column = column.copy() - filtered_column = remove_display_name(filtered_column) filtered_columns.append(filtered_column) return filtered_columns return [ - remove_display_name( - {**model, "columns": column_filter(model.get("columns", []))} - ) + {**model, "columns": column_filter(model.get("columns", []))} for model in mdl.get("models", []) ] diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py index 673d7a621c..91dd309f31 100644 --- a/wren-ai-service/src/web/v1/services/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/services/relationship_recommendation.py @@ -86,6 +86,144 @@ def _model_columns(self, model: dict) -> list[dict]: if column.get("name") and not column.get("relationship") ] + def _humanize_identifier(self, value: Any) -> str: + text = "" if value is None else str(value) + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) + text = re.sub(r"[_\-.]+", " ", text) + text = re.sub(r"\b(id|pk|fk)\b", "", text, flags=re.IGNORECASE) + text = re.sub(r"\s+", " ", text).strip() + replacements = { + "dept": "department", + "emp": "employee", + "org": "organization", + "cust": "customer", + "prod": "product", + "dim": "dimension", + "fact": "fact", + } + parts = [replacements.get(part.lower(), part) for part in text.split()] + if len(parts) > 1 and parts[0].lower() in { + "dbo", + "public", + "stage", + "staging", + "tbl", + }: + parts = parts[1:] + if len(parts) > 1 and parts[0].lower() == "q": + parts = parts[1:] + text = " ".join(parts) + return text.lower() if text else "record" + + def _singularize_label(self, label: str) -> str: + if label.endswith("ies") and len(label) > 3: + return f"{label[:-3]}y" + if label.endswith(("ses", "xes", "zes", "ches", "shes")): + return label[:-2] + if ( + label.endswith("s") + and not label.endswith(("ss", "us", "is", "sales", "series")) + ): + return label[:-1] + return label + + def _model_label(self, model: dict) -> str: + properties = model.get("properties") or {} + return self._singularize_label( + self._humanize_identifier( + properties.get("displayName") + or model.get("tableReference", {}).get("table") + or model.get("name") + ) + ) + + def _pluralize_label(self, label: str) -> str: + if label.endswith("y") and label[-2:] not in {"ay", "ey", "iy", "oy", "uy"}: + return f"{label[:-1]}ies" + if label.endswith(("s", "x", "z", "ch", "sh")): + return f"{label}es" + return f"{label}s" + + def _relationship_description( + self, + from_model: dict, + to_model: dict, + relationship_type: str, + ) -> str: + from_label = self._model_label(from_model) + to_label = self._model_label(to_model) + from_plural = self._pluralize_label(from_label) + to_plural = self._pluralize_label(to_label) + + if relationship_type == "ONE_TO_ONE": + return ( + f"Each {from_label} is linked to one matching {to_label}, " + "connecting details that describe the same business record." + ) + if relationship_type == "ONE_TO_MANY": + return ( + f"Each {from_label} can be associated with multiple {to_plural}, " + f"supporting analysis of {to_plural} by {from_label}." + ) + return ( + f"Each {from_label} belongs to one {to_label}, " + f"so {from_plural} can be grouped and analyzed by {to_label}." + ) + + def _description_is_meaningful(self, value: Any) -> bool: + if not isinstance(value, str): + return False + + text = value.strip() + if len(text) < 24: + return False + + technical_patterns = [ + r"\bappears to reference\b", + r"\breferences\b", + r"\bforeign key\b", + r"\bprimary key\b", + r"\w+\.\w+", + ] + return not any( + re.search(pattern, text, flags=re.IGNORECASE) + for pattern in technical_patterns + ) + + def _ensure_relationship_descriptions(self, response: dict, mdl: dict) -> dict: + relationships = response.get("relationships") + if not isinstance(relationships, list): + return response + + models_by_name = { + model.get("name"): model + for model in mdl.get("models", []) or [] + if model.get("name") + } + normalized_relationships = [] + for relationship in relationships: + if not isinstance(relationship, dict): + continue + + from_model = models_by_name.get(relationship.get("fromModel")) + to_model = models_by_name.get(relationship.get("toModel")) + if not from_model or not to_model: + normalized_relationships.append(relationship) + continue + + reason = relationship.get("reason") + if not self._description_is_meaningful(reason): + relationship = { + **relationship, + "reason": self._relationship_description( + from_model, to_model, relationship.get("type", "MANY_TO_ONE") + ), + } + + normalized_relationships.append(relationship) + + return {**response, "relationships": normalized_relationships} + def _primary_key(self, model: dict) -> Optional[str]: primary_key = model.get("primaryKey") columns = self._model_columns(model) @@ -219,15 +357,9 @@ def add_candidate( relationship_type = self._fallback_relationship_type( from_model, from_column, to_model, to_column ) - reason = ( - f"{from_model_name}.{from_column} appears to reference " - f"{to_model_name}.{to_column}." + reason = self._relationship_description( + from_model, to_model, relationship_type ) - if relationship_type == "ONE_TO_MANY": - reason = ( - f"{from_model_name}.{from_column} appears to be referenced by " - f"{to_model_name}.{to_column}." - ) seen.add(signature) seen_pairs.add(pair_signature) @@ -351,6 +483,8 @@ async def recommend(self, request: Input, **kwargs) -> Resource: ) response = self._fallback_relationships(mdl_dict) + response = self._ensure_relationship_descriptions(response, mdl_dict) + self._cache[request.id] = self.Resource( id=request.id, status="finished", diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py index faea1d2451..4d733ff1bd 100644 --- a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py +++ b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py @@ -122,6 +122,41 @@ async def test_recommend_success(relationship_recommendation_service, mock_pipel mock_pipeline.run.assert_called_once_with(mdl={"key": "value"}, language="English") +@pytest.mark.asyncio +async def test_recommend_replaces_technical_llm_relationship_reason( + relationship_recommendation_service, + mock_pipeline, + mdl_with_project_relationship_candidate, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_project_relationship_candidate + ) + mock_pipeline.run.return_value = { + "validated": { + "relationships": [ + { + "name": "view_project", + "fromModel": "view", + "fromColumn": "project_id", + "type": "MANY_TO_ONE", + "toModel": "project", + "toColumn": "id", + "reason": "view.project_id references project.id.", + } + ] + } + } + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response["relationships"][0]["reason"] == ( + "Each view belongs to one project, so views can be grouped and analyzed " + "by project." + ) + + @pytest.mark.asyncio async def test_recommend_invalid_mdl(relationship_recommendation_service): request = RelationshipRecommendation.Input(id="test_id", mdl="invalid_json") @@ -209,7 +244,10 @@ async def never_finishes(**_kwargs): "type": "MANY_TO_ONE", "toModel": "project", "toColumn": "id", - "reason": "view.project_id appears to reference project.id.", + "reason": ( + "Each view belongs to one project, so views can be grouped " + "and analyzed by project." + ), } ] } @@ -258,7 +296,10 @@ async def test_recommend_fallback_matches_prefixed_model_name( "type": "MANY_TO_ONE", "toModel": "dbo_project", "toColumn": "id", - "reason": "dbo_view.project_id appears to reference dbo_project.id.", + "reason": ( + "Each view belongs to one project, so views can be grouped " + "and analyzed by project." + ), } ] @@ -286,7 +327,10 @@ async def test_recommend_fallback_identifies_one_to_one_relationships( "type": "ONE_TO_ONE", "toModel": "user", "toColumn": "id", - "reason": "profile.user_id appears to reference user.id.", + "reason": ( + "Each profile is linked to one matching user, connecting details " + "that describe the same business record." + ), } ] @@ -314,7 +358,10 @@ async def test_recommend_fallback_scans_all_models_and_identifies_one_to_many( "type": "ONE_TO_MANY", "toModel": "titles", "toColumn": "emp_no", - "reason": "employees.emp_no appears to be referenced by titles.emp_no.", + "reason": ( + "Each employee can be associated with multiple titles, supporting " + "analysis of titles by employee." + ), }, { "name": "employees_dept_emp", @@ -323,7 +370,10 @@ async def test_recommend_fallback_scans_all_models_and_identifies_one_to_many( "type": "ONE_TO_MANY", "toModel": "dept_emp", "toColumn": "emp_no", - "reason": "employees.emp_no appears to be referenced by dept_emp.emp_no.", + "reason": ( + "Each employee can be associated with multiple department employees, " + "supporting analysis of department employees by employee." + ), }, { "name": "departments_dept_emp", @@ -332,7 +382,10 @@ async def test_recommend_fallback_scans_all_models_and_identifies_one_to_many( "type": "ONE_TO_MANY", "toModel": "dept_emp", "toColumn": "dept_no", - "reason": "departments.dept_no appears to be referenced by dept_emp.dept_no.", + "reason": ( + "Each department can be associated with multiple department employees, " + "supporting analysis of department employees by department." + ), }, ] From a153cc5888b8ab1a643abfc02ca99f9f731c0689 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 22:44:35 +0530 Subject: [PATCH 0578/1087] Add metadata hygiene for AI grounding --- .../generation/followup_sql_generation.py | 4 + .../followup_sql_generation_reasoning.py | 2 + .../generation/question_recommendation.py | 3 + .../pipelines/generation/sql_generation.py | 4 + .../generation/sql_generation_reasoning.py | 2 + .../src/pipelines/metadata_hygiene.py | 130 ++++++++++++++++++ .../retrieval/db_schema_retrieval.py | 45 +++--- .../v1/services/question_recommendation.py | 8 +- .../retrieval/test_db_schema_retrieval.py | 61 ++++++++ .../pytest/pipelines/test_metadata_hygiene.py | 72 ++++++++++ 10 files changed, 305 insertions(+), 26 deletions(-) create mode 100644 wren-ai-service/src/pipelines/metadata_hygiene.py create mode 100644 wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py 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 c9f8b23537..f378f235c2 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -25,6 +25,7 @@ get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) +from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost @@ -130,6 +131,7 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: + documents = filter_business_schema_contexts(query, documents or []) _prompt = prompt_builder.run( query=query, data_source=data_source, @@ -185,11 +187,13 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, documents: list[str], + query: str, data_source: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: + documents = filter_business_schema_contexts(query, documents or []) return await post_processor.run( generate_sql_in_followup.get("replies"), project_id=project_id, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index abbbb81d56..2140c6fc09 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -17,6 +17,7 @@ construct_instructions, sql_generation_reasoning_system_prompt, ) +from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_cost from src.web.v1.services import Configuration @@ -82,6 +83,7 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: + documents = filter_business_schema_contexts(query, documents or []) _prompt = prompt_builder.run( query=query, documents=documents, diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index 5a5b3e3e2d..a03d86a669 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -12,6 +12,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines +from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_cost logger = logging.getLogger("wren-ai-service") @@ -183,6 +184,8 @@ def prompt( max_categories: int, prompt_builder: PromptBuilder, ) -> dict: + query_context = "\n".join(previous_questions or []) + documents = filter_business_schema_contexts(query_context, documents or []) _prompt = prompt_builder.run( documents=documents, previous_questions=previous_questions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 59bea2279c..cd51ac7993 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -22,6 +22,7 @@ get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) +from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost @@ -119,6 +120,7 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: + documents = filter_business_schema_contexts(query, documents or []) schema_context = "\n".join(documents or []).lower() has_pcb_context = any( term in schema_context @@ -184,12 +186,14 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, documents: list[str], + query: str, data_source: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, ) -> dict: + documents = filter_business_schema_contexts(query, documents or []) return await post_processor.run( generate_sql.get("replies"), project_id=project_id, 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 f91a4288e4..db0b953039 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -15,6 +15,7 @@ construct_instructions, sql_generation_reasoning_system_prompt, ) +from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_cost from src.web.v1.services import Configuration @@ -66,6 +67,7 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: + documents = filter_business_schema_contexts(query, documents or []) _prompt = prompt_builder.run( query=query, documents=documents, diff --git a/wren-ai-service/src/pipelines/metadata_hygiene.py b/wren-ai-service/src/pipelines/metadata_hygiene.py new file mode 100644 index 0000000000..25a26c0fd6 --- /dev/null +++ b/wren-ai-service/src/pipelines/metadata_hygiene.py @@ -0,0 +1,130 @@ +import re +from typing import Any + +from haystack import Document + + +NOISY_METADATA_TERMS = ( + "archive", + "audit", + "backup", + "cache", + "copy", + "debug", + "dev", + "duplicate", + "etl", + "import", + "load", + "log", + "migration", + "raw", + "sample", + "scratch", + "sync", + "sys", + "technical", + "temp", + "test", + "tmp", +) + + +EXPLICIT_NOISY_METADATA_TERMS = NOISY_METADATA_TERMS + ( + "temporary", + "duplicates", + "logs", + "stage", + "staging", + "staged", + "stages", + "tests", +) + + +def normalize_metadata_token(value: Any) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + + +def metadata_terms(value: Any) -> set[str]: + text = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", str(value or "")) + terms = { + normalize_metadata_token(token) + for token in re.split(r"[^A-Za-z0-9]+", text) + if token + } + return {term for term in terms if term} + + +def query_requests_noisy_metadata(query: str) -> bool: + query_terms = metadata_terms(query) + return bool(query_terms.intersection(EXPLICIT_NOISY_METADATA_TERMS)) + + +def is_noisy_metadata_name(value: Any) -> bool: + terms = metadata_terms(value) + if not terms: + return False + + compact_value = normalize_metadata_token(value) + if any( + compact_value.startswith(prefix) + for prefix in ("tmp", "temp", "test", "stg", "staging") + ): + return True + + return bool(terms.intersection(NOISY_METADATA_TERMS)) + + +def is_noisy_metadata_text(value: Any) -> bool: + terms = metadata_terms(value) + if terms.intersection(NOISY_METADATA_TERMS): + return True + + normalized = str(value or "").lower() + return any( + phrase in normalized + for phrase in ( + "raw load", + "load metadata", + "staging rows", + "temporary table", + "technical table", + ) + ) + + +def is_noisy_document(document: Document) -> bool: + name = document.meta.get("name", "") + description = document.meta.get("description", "") + return is_noisy_metadata_name(name) or is_noisy_metadata_text(description) + + +def filter_business_documents(query: str, documents: list[Document]) -> list[Document]: + if not documents or query_requests_noisy_metadata(query): + return documents + + filtered = [document for document in documents if not is_noisy_document(document)] + return filtered or documents + + +def _extract_context_name(context: str) -> str: + match = re.search( + r"\bCREATE\s+(?:TABLE|VIEW)\s+([^\s(]+)", + context or "", + flags=re.IGNORECASE, + ) + return match.group(1).strip("[]`\"") if match else "" + + +def is_noisy_schema_context(context: str) -> bool: + name = _extract_context_name(context) + return is_noisy_metadata_name(name) or is_noisy_metadata_text(context[:500]) + + +def filter_business_schema_contexts(query: str, contexts: list[str]) -> list[str]: + if not contexts or query_requests_noisy_metadata(query): + return contexts + + filtered = [context for context in contexts if not is_noisy_schema_context(context)] + return filtered or contexts diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..9b163d6712 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -21,6 +21,12 @@ get_engine_supported_data_type, normalize_data_type, ) +from src.pipelines.metadata_hygiene import ( + NOISY_METADATA_TERMS, + filter_business_documents, + metadata_terms, + query_requests_noisy_metadata, +) from src.utils import trace_cost if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory @@ -184,10 +190,6 @@ def expand_business_terms_for_retrieval(query: str) -> str: return f"{query}\n" + "\n".join(expansions) -def _normalize_retrieval_token(value: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - def _retrieval_terms(value: str) -> set[str]: stop_words = { "about", @@ -214,9 +216,9 @@ def _retrieval_terms(value: str) -> set[str]: "with", } terms = { - _normalize_retrieval_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") - if len(token) > 2 and token.lower() not in stop_words + term + for term in metadata_terms(value) + if len(term) > 2 and term.lower() not in stop_words } return {term for term in terms if term} @@ -243,23 +245,8 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "stage", - "staging", - "temp", - "test", - "tmp", - ) - if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( - normalized_query, - non_production_terms, + if source_terms.intersection(NOISY_METADATA_TERMS) and not query_requests_noisy_metadata( + normalized_query ): score -= 60 @@ -394,6 +381,7 @@ def _rerank_table_documents(query: str, documents: list[Document]) -> list[Docum if not documents: return documents + documents = filter_business_documents(query, documents) reranked = _score_table_documents(query, documents) if not reranked: return documents @@ -424,6 +412,7 @@ def _select_relevant_table_documents( if not documents or max_tables <= 0: return [] + documents = filter_business_documents(query, documents) reranked = _score_table_documents(query, documents) if not reranked: return documents[:max_tables] @@ -585,6 +574,9 @@ async def table_retrieval( query_embedding=embedding.get("embedding"), filters=base_filters, ) + results["documents"] = filter_business_documents( + query, results.get("documents") or [] + ) results["documents"] = _select_relevant_table_documents( query, results.get("documents") or [] ) @@ -651,7 +643,10 @@ async def dbschema_retrieval( return [] results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results.get("documents", []) + documents = results.get("documents", []) + if not tables: + documents = filter_business_documents(query, documents) + return documents @observe() diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index d93e456d2a..839fee9bbb 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -8,6 +8,7 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_metadata from src.web.v1.services import BaseRequest, MetadataTraceable @@ -405,8 +406,13 @@ async def recommend(self, input: Request, **kwargs) -> Event: ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) + raw_table_ddls = [document.get("table_ddl") for document in documents] + table_ddls_for_recommendation = filter_business_schema_contexts( + "\n".join(input.previous_questions or []), + raw_table_ddls, + ) table_ddls = self._limit_text_items( - [document.get("table_ddl") for document in documents], + table_ddls_for_recommendation, max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7c7d46ce89..05b12ca408 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -152,6 +152,28 @@ def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] +def test_select_relevant_table_documents_keeps_requested_test_candidate(): + documents = [ + Document( + content="Raw test load rows with order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ), + Document( + content="New order transaction records with market and customer details.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.4, + ), + ] + + selected = _select_relevant_table_documents( + "Show test load order distribution across markets.", + documents, + ) + + assert "dbo_xStageLoad8_Test" in [document.meta["name"] for document in selected] + + @pytest.mark.asyncio async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ @@ -334,6 +356,45 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_dbschema_retrieval_filters_unrequested_noisy_full_schema(): + class Retriever: + async def run(self, query_embedding, filters): + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "orders", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "orders"}, + ), + Document( + content=str( + { + "type": "TABLE", + "name": "orders_test_duplicate", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "orders_test_duplicate"}, + ), + ] + } + + documents = await dbschema_retrieval( + query="", + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=Retriever(), + ) + + assert [document.meta["name"] for document in documents] == ["orders"] + + @pytest.mark.asyncio async def test_dbschema_retrieval_does_not_load_full_schema_for_unmatched_question(): class Retriever: diff --git a/wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py b/wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py new file mode 100644 index 0000000000..87823b65a5 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py @@ -0,0 +1,72 @@ +from haystack import Document + +from src.pipelines.metadata_hygiene import ( + filter_business_documents, + filter_business_schema_contexts, + is_noisy_schema_context, + query_requests_noisy_metadata, +) + + +def test_filter_business_documents_excludes_unrequested_noisy_models(): + documents = [ + Document( + content="Customer order transaction data.", + meta={"name": "dbo_Orders", "type": "TABLE_DESCRIPTION"}, + ), + Document( + content="Raw temporary import rows.", + meta={"name": "tmp_orders_import", "type": "TABLE_DESCRIPTION"}, + ), + Document( + content="Duplicate backup copy of orders.", + meta={"name": "Orders_Backup_Copy", "type": "TABLE_DESCRIPTION"}, + ), + ] + + filtered = filter_business_documents("show orders by customer", documents) + + assert [document.meta["name"] for document in filtered] == ["dbo_Orders"] + + +def test_filter_business_documents_keeps_noisy_models_when_explicitly_requested(): + documents = [ + Document( + content="Customer order transaction data.", + meta={"name": "dbo_Orders", "type": "TABLE_DESCRIPTION"}, + ), + Document( + content="Raw temporary import rows.", + meta={"name": "tmp_orders_import", "type": "TABLE_DESCRIPTION"}, + ), + ] + + filtered = filter_business_documents("show temporary import rows", documents) + + assert [document.meta["name"] for document in filtered] == [ + "dbo_Orders", + "tmp_orders_import", + ] + + +def test_filter_business_schema_contexts_preserves_all_when_only_noisy_context_exists(): + contexts = ["CREATE TABLE stg_orders_load (id INT, order_id INT);"] + + assert filter_business_schema_contexts("show orders", contexts) == contexts + + +def test_filter_business_schema_contexts_removes_noisy_contexts_when_business_exists(): + contexts = [ + "CREATE TABLE orders (id INT, customer_id INT);", + "CREATE TABLE orders_test_duplicate (id INT, customer_id INT);", + "CREATE TABLE debug_order_log (id INT, message TEXT);", + ] + + assert filter_business_schema_contexts("show orders", contexts) == [ + "CREATE TABLE orders (id INT, customer_id INT);" + ] + + +def test_noisy_context_detection_and_explicit_query_terms(): + assert is_noisy_schema_context("CREATE TABLE dbo_xStageLoad8_Test (id INT);") + assert query_requests_noisy_metadata("compare staging load rows") From 0cadb0e8d23e8350b53fe7a523f0d82a00e4011a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 15 Jul 2026 22:39:25 +0530 Subject: [PATCH 0579/1087] Improve generated semantic description quality --- .../generation/semantics_description.py | 9 +- .../web/v1/services/semantics_description.py | 142 ++++++++++++++++-- .../services/test_semantics_description.py | 74 ++++++++- 3 files changed, 209 insertions(+), 16 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index ad72c659c5..08fa40d8c6 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -36,13 +36,14 @@ ``` Your task is to update this JSON structure by adding a `description` field inside both the `properties` attribute of each `column` and the `model` itself. -Each `description` should be derived from the user-provided dataset context, the model name, column names, data types, aliases, and any existing descriptions. +Each `description` should be derived from the user-provided dataset context, the full schema, relationships, model names, column names, data types, aliases, and any existing descriptions. Follow these steps: 1. **For the `model`**: Write a clear natural language business description of the model's purpose and what real-world records it represents. Insert this description in the `properties` field of the `model`. 2. **For each `column`**: Write a clear natural language business description of the column's meaning, not just its technical name. Each column's description should be added under its respective `properties` field in the format: `'description': 'business description'`. 3. Ensure that the output is a well-formatted JSON structure, preserving the input's original format and adding the appropriate `description` fields. 4. Avoid repeating technical table or column names as the whole description. Prefer business meaning such as identifiers, dates, amounts, statuses, dimensions, ownership, and operational usage. -5. Keep descriptions concise, factual, and useful for text-to-SQL retrieval. +5. Do not use generic boilerplate such as "stores the value", "contains records for", or "field from". Explain what the data means to a business user. +6. Make every model and column description unique, human-readable, concise, factual, and useful for text-to-SQL retrieval. ### Output Format: @@ -88,7 +89,9 @@ Picked models: {{ picked_models }} Localization Language: {{ language }} -Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt. Do not omit selected models or columns. +Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt and schema context. +Do not omit selected models or columns. Do not copy the table or column name as the description. +Use simple language that explains the business purpose, meaning, and analytical use of each field. """ diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index eb644ca1c6..370c3914f1 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,5 +1,6 @@ import asyncio import logging +import re from typing import Any, Dict, Literal, Optional import orjson @@ -85,10 +86,110 @@ def _text(self, value: Any) -> str: def _humanize_name(self, name: str) -> str: return " ".join( token - for token in name.replace(".", " ").replace("_", " ").split() + for token in re.sub( + r"(?<=[a-z0-9])(?=[A-Z])", + " ", + name.replace(".", " ").replace("_", " "), + ).split() if token.lower() not in {"dbo", "public"} ) or name + def _name_tokens(self, name: str) -> set[str]: + return { + token.lower() + for token in self._humanize_name(name).split() + if token and token.lower() not in {"x", "stage", "load", "tbl", "table"} + } + + def _table_context(self, model: dict[str, Any]) -> set[str]: + tokens = self._name_tokens(self._text(model.get("name", ""))) + for column in model.get("columns", []) or []: + tokens.update(self._name_tokens(self._text(column.get("name", "")))) + return tokens + + def _fallback_model_description(self, model: dict[str, Any]) -> str: + context = self._table_context(model) + if context & {"sales", "revenue", "order", "orders", "customer", "product"}: + return ( + "Captures commercial activity and related business dimensions for " + "sales reporting, performance analysis, and customer or product insights." + ) + if context & {"invoice", "payment", "price", "cost", "amount", "finance"}: + return ( + "Captures financial transactions and monetary measures used for " + "reconciliation, reporting, and performance analysis." + ) + if context & {"employee", "user", "person", "salesperson", "owner", "manager"}: + return ( + "Captures people and ownership attributes used to assign responsibility, " + "segment activity, and analyze performance." + ) + return ( + "Captures operational business records used for reporting, filtering, " + "trend analysis, and answering analytical questions." + ) + + def _fallback_column_description( + self, + model: dict[str, Any], + column: dict[str, Any], + ) -> str: + column_name = self._text(column.get("name", "")) + data_type = self._text(column.get("type", "")).lower() + tokens = self._name_tokens(column_name) + context = self._table_context(model) + + if tokens & {"division", "bu", "business", "unit", "department"}: + return "Organizational segment used to group records for ownership, reporting, and performance comparison." + if tokens & {"company", "entity", "organization", "org"}: + return "Legal or business entity associated with the record for company-level reporting and filtering." + if tokens & {"market", "region", "territory", "country", "state", "city", "location"}: + return "Geographic or market segment used to analyze activity by area and compare regional performance." + if tokens & {"product", "prod", "sku", "item", "material"}: + return "Product or item classification used to analyze sales, demand, and business activity by offering." + if tokens & {"type", "category", "class", "segment", "group"}: + return "Business classification used to segment records into meaningful reporting categories." + if tokens & {"customer", "client", "account"}: + return "Customer or account reference used to connect activity to the buyer or business relationship." + if tokens & {"salesperson", "seller", "rep", "owner", "manager", "person"}: + return "Responsible person or role associated with the record for ownership and performance analysis." + if tokens & {"status", "stage", "state"}: + return "Current business state used to track workflow progress, completion, or operational condition." + if tokens & {"date", "time", "day", "month", "year", "period", "created", "updated"}: + return "Time period used to sequence records, filter activity, and analyze trends over time." + if tokens & {"amount", "sales", "revenue", "cost", "price", "value", "total", "net", "gross"}: + return "Monetary measure used to calculate financial results, compare performance, and summarize business activity." + if tokens & {"quantity", "qty", "count", "units", "volume"}: + return "Quantity measure used to count activity, summarize volume, and compare operational scale." + if tokens & {"rate", "ratio", "percent", "percentage", "margin"}: + return "Calculated rate or percentage used to compare efficiency, contribution, or relative performance." + if tokens & {"id", "key", "code", "number", "no"}: + return "Identifier used to distinguish records and join this data with related business information." + if "date" in data_type or "time" in data_type: + return "Timestamp or calendar value used for time-based filtering, sequencing, and trend analysis." + if any(type_name in data_type for type_name in ("int", "float", "double", "decimal", "numeric", "number")): + return "Numeric business measure used for aggregation, comparison, and analytical calculations." + if context & {"sales", "order", "customer", "product"}: + return "Business attribute used to filter and explain commercial activity in reporting and analysis." + return "Business attribute used to categorize, filter, and explain records in analytical questions." + + def _is_low_quality_description(self, description: str, name: str) -> bool: + normalized = " ".join(description.lower().split()) + if not normalized: + return True + + name_text = self._humanize_name(name).lower() + low_quality_patterns = ( + "stores the", + "value used to describe or analyze", + "contains business records for", + "represents ", + "field from", + ) + if any(pattern in normalized for pattern in low_quality_patterns): + return True + return normalized in {name.lower(), name_text} + def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: output: dict[str, Any] = {} for model in chunk.get("mdl", {}).get("models", []): @@ -96,13 +197,10 @@ def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: if not model_name: continue - model_label = self._humanize_name(model_name) model_properties = self._properties(model) model_description = self._text(model_properties.get("description", "")) - if not model_description: - model_description = ( - f"Contains business records for {model_label}, used for reporting, analysis, and operational questions." - ) + if self._is_low_quality_description(model_description, model_name): + model_description = self._fallback_model_description(model) columns = [] for column in model.get("columns", []) or []: @@ -115,10 +213,10 @@ def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: column_description = self._text( column_properties.get("description", "") ) - if not column_description: - column_label = self._humanize_name(column_name) - column_description = ( - f"Stores the {column_label} value used to describe or analyze {model_label} records." + if self._is_low_quality_description(column_description, column_name): + column_description = self._fallback_column_description( + model, + column, ) columns.append( { @@ -153,10 +251,20 @@ def _complete_output_with_fallback( if not isinstance(properties, dict): properties = {} model_output["properties"] = properties - if not properties.get("description"): + if self._is_low_quality_description( + self._text(properties.get("description", "")), + model_name, + ): properties["description"] = self._text( model_output.get("description") ) or fallback_model["properties"]["description"] + if self._is_low_quality_description( + properties["description"], + model_name, + ): + properties["description"] = fallback_model["properties"][ + "description" + ] output_columns = { column.get("name"): column @@ -174,10 +282,20 @@ def _complete_output_with_fallback( if not isinstance(column_properties, dict): column_properties = {} output_column["properties"] = column_properties - if not column_properties.get("description"): + if self._is_low_quality_description( + self._text(column_properties.get("description", "")), + column_name, + ): column_properties["description"] = self._text( output_column.get("description") ) or fallback_column["properties"]["description"] + if self._is_low_quality_description( + column_properties["description"], + column_name, + ): + column_properties["description"] = fallback_column[ + "properties" + ]["description"] return completed diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 4b54a712fd..586c42f700 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -47,7 +47,7 @@ async def test_generate_semantics_description( "name": "column1", "type": "varchar", "properties": { - "description": "Stores the column1 value used to describe or analyze model1 records." + "description": "Business attribute used to categorize, filter, and explain records in analytical questions." }, } ], @@ -331,6 +331,78 @@ async def test_partial_llm_output_is_completed_for_all_selected_columns( assert response.response["customers"]["properties"]["description"] +@pytest.mark.asyncio +async def test_generic_llm_descriptions_are_replaced_with_business_descriptions( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Sales and operations reporting dataset", + selected_models=["dbo_xStageLoad2"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "dbo_xStageLoad2", + "columns": [ + {"name": "Division", "type": "varchar"}, + {"name": "SalesPerson", "type": "varchar"}, + {"name": "SalesAmount", "type": "float"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "dbo_xStageLoad2": { + "name": "dbo_xStageLoad2", + "columns": [ + { + "name": "Division", + "properties": { + "description": "Stores the Division value used to describe or analyze xStage records." + }, + }, + { + "name": "SalesPerson", + "properties": {"description": "SalesPerson"}, + }, + { + "name": "SalesAmount", + "properties": { + "description": "Stores the SalesAmount value." + }, + }, + ], + "properties": { + "description": "Contains business records for xStageLoad2." + }, + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + descriptions = [ + column["properties"]["description"] + for column in response.response["dbo_xStageLoad2"]["columns"] + ] + assert response.response["dbo_xStageLoad2"]["properties"]["description"].startswith( + "Captures commercial activity" + ) + assert descriptions == [ + "Organizational segment used to group records for ownership, reporting, and performance comparison.", + "Responsible person or role associated with the record for ownership and performance analysis.", + "Monetary measure used to calculate financial results, compare performance, and summarize business activity.", + ] + assert all("Stores the" not in description for description in descriptions) + + @pytest.mark.asyncio async def test_batch_processing_partial_failure( service: SemanticsDescription, From 7153ebf1e1c352ee173cf08d33c0485e80216709 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 15 Jul 2026 22:55:45 +0530 Subject: [PATCH 0580/1087] Revert "Add metadata hygiene for AI grounding" This reverts commit a153cc5888b8ab1a643abfc02ca99f9f731c0689. --- .../generation/followup_sql_generation.py | 4 - .../followup_sql_generation_reasoning.py | 2 - .../generation/question_recommendation.py | 3 - .../pipelines/generation/sql_generation.py | 4 - .../generation/sql_generation_reasoning.py | 2 - .../src/pipelines/metadata_hygiene.py | 130 ------------------ .../retrieval/db_schema_retrieval.py | 45 +++--- .../v1/services/question_recommendation.py | 8 +- .../retrieval/test_db_schema_retrieval.py | 61 -------- .../pytest/pipelines/test_metadata_hygiene.py | 72 ---------- 10 files changed, 26 insertions(+), 305 deletions(-) delete mode 100644 wren-ai-service/src/pipelines/metadata_hygiene.py delete mode 100644 wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py 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 f378f235c2..c9f8b23537 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -25,7 +25,6 @@ get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) -from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost @@ -131,7 +130,6 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - documents = filter_business_schema_contexts(query, documents or []) _prompt = prompt_builder.run( query=query, data_source=data_source, @@ -187,13 +185,11 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, documents: list[str], - query: str, data_source: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: - documents = filter_business_schema_contexts(query, documents or []) return await post_processor.run( generate_sql_in_followup.get("replies"), project_id=project_id, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 2140c6fc09..abbbb81d56 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -17,7 +17,6 @@ construct_instructions, sql_generation_reasoning_system_prompt, ) -from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_cost from src.web.v1.services import Configuration @@ -83,7 +82,6 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: - documents = filter_business_schema_contexts(query, documents or []) _prompt = prompt_builder.run( query=query, documents=documents, diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index a03d86a669..5a5b3e3e2d 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -12,7 +12,6 @@ from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines -from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_cost logger = logging.getLogger("wren-ai-service") @@ -184,8 +183,6 @@ def prompt( max_categories: int, prompt_builder: PromptBuilder, ) -> dict: - query_context = "\n".join(previous_questions or []) - documents = filter_business_schema_contexts(query_context, documents or []) _prompt = prompt_builder.run( documents=documents, previous_questions=previous_questions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index cd51ac7993..59bea2279c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -22,7 +22,6 @@ get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) -from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost @@ -120,7 +119,6 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - documents = filter_business_schema_contexts(query, documents or []) schema_context = "\n".join(documents or []).lower() has_pcb_context = any( term in schema_context @@ -186,14 +184,12 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, documents: list[str], - query: str, data_source: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, ) -> dict: - documents = filter_business_schema_contexts(query, documents or []) return await post_processor.run( generate_sql.get("replies"), project_id=project_id, 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 db0b953039..f91a4288e4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -15,7 +15,6 @@ construct_instructions, sql_generation_reasoning_system_prompt, ) -from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_cost from src.web.v1.services import Configuration @@ -67,7 +66,6 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: - documents = filter_business_schema_contexts(query, documents or []) _prompt = prompt_builder.run( query=query, documents=documents, diff --git a/wren-ai-service/src/pipelines/metadata_hygiene.py b/wren-ai-service/src/pipelines/metadata_hygiene.py deleted file mode 100644 index 25a26c0fd6..0000000000 --- a/wren-ai-service/src/pipelines/metadata_hygiene.py +++ /dev/null @@ -1,130 +0,0 @@ -import re -from typing import Any - -from haystack import Document - - -NOISY_METADATA_TERMS = ( - "archive", - "audit", - "backup", - "cache", - "copy", - "debug", - "dev", - "duplicate", - "etl", - "import", - "load", - "log", - "migration", - "raw", - "sample", - "scratch", - "sync", - "sys", - "technical", - "temp", - "test", - "tmp", -) - - -EXPLICIT_NOISY_METADATA_TERMS = NOISY_METADATA_TERMS + ( - "temporary", - "duplicates", - "logs", - "stage", - "staging", - "staged", - "stages", - "tests", -) - - -def normalize_metadata_token(value: Any) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - -def metadata_terms(value: Any) -> set[str]: - text = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", str(value or "")) - terms = { - normalize_metadata_token(token) - for token in re.split(r"[^A-Za-z0-9]+", text) - if token - } - return {term for term in terms if term} - - -def query_requests_noisy_metadata(query: str) -> bool: - query_terms = metadata_terms(query) - return bool(query_terms.intersection(EXPLICIT_NOISY_METADATA_TERMS)) - - -def is_noisy_metadata_name(value: Any) -> bool: - terms = metadata_terms(value) - if not terms: - return False - - compact_value = normalize_metadata_token(value) - if any( - compact_value.startswith(prefix) - for prefix in ("tmp", "temp", "test", "stg", "staging") - ): - return True - - return bool(terms.intersection(NOISY_METADATA_TERMS)) - - -def is_noisy_metadata_text(value: Any) -> bool: - terms = metadata_terms(value) - if terms.intersection(NOISY_METADATA_TERMS): - return True - - normalized = str(value or "").lower() - return any( - phrase in normalized - for phrase in ( - "raw load", - "load metadata", - "staging rows", - "temporary table", - "technical table", - ) - ) - - -def is_noisy_document(document: Document) -> bool: - name = document.meta.get("name", "") - description = document.meta.get("description", "") - return is_noisy_metadata_name(name) or is_noisy_metadata_text(description) - - -def filter_business_documents(query: str, documents: list[Document]) -> list[Document]: - if not documents or query_requests_noisy_metadata(query): - return documents - - filtered = [document for document in documents if not is_noisy_document(document)] - return filtered or documents - - -def _extract_context_name(context: str) -> str: - match = re.search( - r"\bCREATE\s+(?:TABLE|VIEW)\s+([^\s(]+)", - context or "", - flags=re.IGNORECASE, - ) - return match.group(1).strip("[]`\"") if match else "" - - -def is_noisy_schema_context(context: str) -> bool: - name = _extract_context_name(context) - return is_noisy_metadata_name(name) or is_noisy_metadata_text(context[:500]) - - -def filter_business_schema_contexts(query: str, contexts: list[str]) -> list[str]: - if not contexts or query_requests_noisy_metadata(query): - return contexts - - filtered = [context for context in contexts if not is_noisy_schema_context(context)] - return filtered or contexts diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 9b163d6712..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -21,12 +21,6 @@ get_engine_supported_data_type, normalize_data_type, ) -from src.pipelines.metadata_hygiene import ( - NOISY_METADATA_TERMS, - filter_business_documents, - metadata_terms, - query_requests_noisy_metadata, -) from src.utils import trace_cost if TYPE_CHECKING: from src.web.v1.services.ask import AskHistory @@ -190,6 +184,10 @@ def expand_business_terms_for_retrieval(query: str) -> str: return f"{query}\n" + "\n".join(expansions) +def _normalize_retrieval_token(value: str) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + + def _retrieval_terms(value: str) -> set[str]: stop_words = { "about", @@ -216,9 +214,9 @@ def _retrieval_terms(value: str) -> set[str]: "with", } terms = { - term - for term in metadata_terms(value) - if len(term) > 2 and term.lower() not in stop_words + _normalize_retrieval_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") + if len(token) > 2 and token.lower() not in stop_words } return {term for term in terms if term} @@ -245,8 +243,23 @@ def _source_shape_score(query: str, document: Document) -> int: source_terms = _retrieval_terms(source_text) score = 0 - if source_terms.intersection(NOISY_METADATA_TERMS) and not query_requests_noisy_metadata( - normalized_query + non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "stage", + "staging", + "temp", + "test", + "tmp", + ) + if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( + normalized_query, + non_production_terms, ): score -= 60 @@ -381,7 +394,6 @@ def _rerank_table_documents(query: str, documents: list[Document]) -> list[Docum if not documents: return documents - documents = filter_business_documents(query, documents) reranked = _score_table_documents(query, documents) if not reranked: return documents @@ -412,7 +424,6 @@ def _select_relevant_table_documents( if not documents or max_tables <= 0: return [] - documents = filter_business_documents(query, documents) reranked = _score_table_documents(query, documents) if not reranked: return documents[:max_tables] @@ -574,9 +585,6 @@ async def table_retrieval( query_embedding=embedding.get("embedding"), filters=base_filters, ) - results["documents"] = filter_business_documents( - query, results.get("documents") or [] - ) results["documents"] = _select_relevant_table_documents( query, results.get("documents") or [] ) @@ -643,10 +651,7 @@ async def dbschema_retrieval( return [] results = await dbschema_retriever.run(query_embedding=[], filters=filters) - documents = results.get("documents", []) - if not tables: - documents = filter_business_documents(query, documents) - return documents + return results.get("documents", []) @observe() diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 839fee9bbb..d93e456d2a 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -8,7 +8,6 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.metadata_hygiene import filter_business_schema_contexts from src.utils import trace_metadata from src.web.v1.services import BaseRequest, MetadataTraceable @@ -406,13 +405,8 @@ async def recommend(self, input: Request, **kwargs) -> Event: ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) - raw_table_ddls = [document.get("table_ddl") for document in documents] - table_ddls_for_recommendation = filter_business_schema_contexts( - "\n".join(input.previous_questions or []), - raw_table_ddls, - ) table_ddls = self._limit_text_items( - table_ddls_for_recommendation, + [document.get("table_ddl") for document in documents], max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 05b12ca408..7c7d46ce89 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -152,28 +152,6 @@ def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] -def test_select_relevant_table_documents_keeps_requested_test_candidate(): - documents = [ - Document( - content="Raw test load rows with order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ), - Document( - content="New order transaction records with market and customer details.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.4, - ), - ] - - selected = _select_relevant_table_documents( - "Show test load order distribution across markets.", - documents, - ) - - assert "dbo_xStageLoad8_Test" in [document.meta["name"] for document in selected] - - @pytest.mark.asyncio async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ @@ -356,45 +334,6 @@ async def run(self, query_embedding, filters): } -@pytest.mark.asyncio -async def test_dbschema_retrieval_filters_unrequested_noisy_full_schema(): - class Retriever: - async def run(self, query_embedding, filters): - return { - "documents": [ - Document( - content=str( - { - "type": "TABLE", - "name": "orders", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "orders"}, - ), - Document( - content=str( - { - "type": "TABLE", - "name": "orders_test_duplicate", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "orders_test_duplicate"}, - ), - ] - } - - documents = await dbschema_retrieval( - query="", - table_retrieval={"documents": []}, - project_id="project-1", - dbschema_retriever=Retriever(), - ) - - assert [document.meta["name"] for document in documents] == ["orders"] - - @pytest.mark.asyncio async def test_dbschema_retrieval_does_not_load_full_schema_for_unmatched_question(): class Retriever: diff --git a/wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py b/wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py deleted file mode 100644 index 87823b65a5..0000000000 --- a/wren-ai-service/tests/pytest/pipelines/test_metadata_hygiene.py +++ /dev/null @@ -1,72 +0,0 @@ -from haystack import Document - -from src.pipelines.metadata_hygiene import ( - filter_business_documents, - filter_business_schema_contexts, - is_noisy_schema_context, - query_requests_noisy_metadata, -) - - -def test_filter_business_documents_excludes_unrequested_noisy_models(): - documents = [ - Document( - content="Customer order transaction data.", - meta={"name": "dbo_Orders", "type": "TABLE_DESCRIPTION"}, - ), - Document( - content="Raw temporary import rows.", - meta={"name": "tmp_orders_import", "type": "TABLE_DESCRIPTION"}, - ), - Document( - content="Duplicate backup copy of orders.", - meta={"name": "Orders_Backup_Copy", "type": "TABLE_DESCRIPTION"}, - ), - ] - - filtered = filter_business_documents("show orders by customer", documents) - - assert [document.meta["name"] for document in filtered] == ["dbo_Orders"] - - -def test_filter_business_documents_keeps_noisy_models_when_explicitly_requested(): - documents = [ - Document( - content="Customer order transaction data.", - meta={"name": "dbo_Orders", "type": "TABLE_DESCRIPTION"}, - ), - Document( - content="Raw temporary import rows.", - meta={"name": "tmp_orders_import", "type": "TABLE_DESCRIPTION"}, - ), - ] - - filtered = filter_business_documents("show temporary import rows", documents) - - assert [document.meta["name"] for document in filtered] == [ - "dbo_Orders", - "tmp_orders_import", - ] - - -def test_filter_business_schema_contexts_preserves_all_when_only_noisy_context_exists(): - contexts = ["CREATE TABLE stg_orders_load (id INT, order_id INT);"] - - assert filter_business_schema_contexts("show orders", contexts) == contexts - - -def test_filter_business_schema_contexts_removes_noisy_contexts_when_business_exists(): - contexts = [ - "CREATE TABLE orders (id INT, customer_id INT);", - "CREATE TABLE orders_test_duplicate (id INT, customer_id INT);", - "CREATE TABLE debug_order_log (id INT, message TEXT);", - ] - - assert filter_business_schema_contexts("show orders", contexts) == [ - "CREATE TABLE orders (id INT, customer_id INT);" - ] - - -def test_noisy_context_detection_and_explicit_query_terms(): - assert is_noisy_schema_context("CREATE TABLE dbo_xStageLoad8_Test (id INT);") - assert query_requests_noisy_metadata("compare staging load rows") From 03e8bb250cb1a377b691918b60e3fa61621e5bb6 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 00:55:34 +0530 Subject: [PATCH 0581/1087] Fix ask pipeline schema retrieval stalls --- wren-ai-service/src/web/v1/services/ask.py | 184 +++++++++++++++++- .../pytest/services/test_ask_sales_sql.py | 37 ++++ 2 files changed, 219 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 13548f9620..c4a7e11b14 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1298,6 +1298,26 @@ def _find_any_temporal_schema_column(self, table: dict[str, Any]) -> str | None: return column_name return None + def _is_probable_explicit_table_token(self, token: str) -> bool: + normalized = re.sub(r"\s+", " ", str(token or "").strip().lower()) + if not normalized: + return False + if normalized in self._INTENT_STOPWORDS or normalized in { + "last", + "latest", + "recent", + "current", + "previous", + "next", + "month", + "year", + "quarter", + "week", + "day", + }: + return False + return True + def _quote_sql_identifier(self, identifier: str) -> str: return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' @@ -1493,6 +1513,103 @@ def _find_dimension_column_for_query( return column return candidate_columns[0] if candidate_columns else None + def _build_schema_ranked_measure_sql( + self, + query: str, + table_ddls: list[str], + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + if not re.search( + r"\b(?:top|highest|largest|biggest|best|lowest|smallest|bottom)\b", + normalized_query, + ): + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + query_tokens = self._intent_tokens(query) + if not query_tokens: + return None + + scored: list[tuple[int, dict[str, Any], str, str]] = [] + for table in tables: + text_columns: list[tuple[int, str]] = [] + numeric_columns: list[tuple[int, str]] = [] + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if not column_name: + continue + + column_tokens = self._schema_name_tokens(column_name) + query_overlap = column_tokens & query_tokens + + if self._is_numeric_schema_type(column_type): + if query_overlap: + numeric_columns.append( + (80 + 5 * len(query_overlap), column_name) + ) + continue + + if self._is_temporal_schema_type(column_type): + continue + + if query_overlap: + text_columns.append((60 + 5 * len(query_overlap), column_name)) + + if not text_columns or not numeric_columns: + continue + + dimension_score, dimension = sorted( + text_columns, + key=lambda item: item[0], + reverse=True, + )[0] + measure_score, measure = sorted( + numeric_columns, + key=lambda item: item[0], + reverse=True, + )[0] + table_score = dimension_score + measure_score + table_tokens = self._schema_name_tokens(str(table.get("name") or "")) + table_score += 5 * len(table_tokens & query_tokens) + scored.append((table_score, table, dimension, measure)) + + if not scored: + return None + + _, table, dimension, measure = sorted( + scored, + key=lambda item: item[0], + reverse=True, + )[0] + table_name = str(table.get("name") or "") + if not table_name: + return None + + limit = self._extract_requested_top_n(query, default_value=10) + table_ref = self._quote_sql_identifier(table_name) + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + measure_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" + metric_expr = f"SUM({measure_ref})" + direction = ( + "ASC" + if re.search(r"\b(?:lowest|smallest|least|bottom)\b", normalized_query) + else "DESC" + ) + return ( + f"SELECT TOP {limit} {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " + f"{metric_expr} AS {self._quote_sql_identifier('Total' + measure)} " + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL AND {measure_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f"ORDER BY {metric_expr} {direction}" + ) + def _find_temporal_column_for_query( self, query: str, table: dict[str, Any] ) -> str | None: @@ -1737,7 +1854,11 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - if table_name and table_name not in table_names: + if ( + table_name + and self._is_probable_explicit_table_token(table_name) + and table_name not in table_names + ): table_names.append(table_name) for match in re.finditer( r"\bin\s+(?:the\s+)?([A-Za-z_][A-Za-z0-9_.$]*)\s+table\b", @@ -1745,7 +1866,11 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: flags=re.IGNORECASE, ): table_name = match.group(1).strip(".,;:()[]{}") - if table_name and table_name not in table_names: + if ( + table_name + and self._is_probable_explicit_table_token(table_name) + and table_name not in table_names + ): table_names.append(table_name) for match in re.finditer( r"\bin\s+([A-Za-z_][A-Za-z0-9_.$]*)", @@ -1756,6 +1881,7 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: if ( table_name and ("." in table_name or "_" in table_name) + and self._is_probable_explicit_table_token(table_name) and table_name not in table_names ): table_names.append(table_name) @@ -1768,6 +1894,7 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: if ( table_name and ("." in table_name or "_" in table_name) + and self._is_probable_explicit_table_token(table_name) and table_name not in table_names ): table_names.append(table_name) @@ -1794,7 +1921,13 @@ def _explicit_table_name_candidates(self, table_name: str) -> list[str]: separator_normalized = re.sub(r"[.$]", "_", table_name) if separator_normalized not in candidates: candidates.append(separator_normalized) + if "_" in table_name and "." not in table_name: + dotted = table_name.replace("_", ".", 1) + if dotted not in candidates: + candidates.append(dotted) short_name = re.split(r"[.$]", table_name)[-1] + if "." not in table_name and "_" in table_name: + short_name = table_name.split("_", 1)[-1] if short_name and short_name not in candidates: candidates.append(short_name) return candidates @@ -5848,6 +5981,32 @@ async def ask( table_names, ) + if ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, + ): + ask_result = self._build_validated_ask_result_from_sql( + ranked_measure_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated ranked measure SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = ranked_measure_sql + if table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ): @@ -6474,6 +6633,27 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) + if not api_results and ( + ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, + ) + ): + logger.info( + "Using schema-grounded ranked measure SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + ranked_measure_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = ranked_measure_sql + error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." + if not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index c7ab4c9272..4a2b68027e 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -538,6 +538,12 @@ def test_extract_explicit_table_names_from_query(): assert service._extract_explicit_table_names_from_query( "Show the first 10 rows from tblNewOrders" ) == ["tblNewOrders"] + assert service._extract_explicit_table_names_from_query( + "Show the latest records from last month" + ) == [] + assert service._extract_explicit_table_names_from_query( + "Show all customers names" + ) == [] def test_extract_explicit_table_names_from_using_clause(): @@ -560,6 +566,37 @@ def test_extract_explicit_table_names_from_in_clause(): assert service._extract_explicit_table_names_from_query( "Which customers have the highest number of orders in market?" ) == [] + assert service._extract_explicit_table_names_from_query( + "Show top 10 customers by invoice amount in the current year" + ) == [] + + +def test_build_schema_ranked_measure_sql_uses_matching_dimension_and_measure(): + service = AskService.__new__(AskService) + + sql = service._build_schema_ranked_measure_sql( + "Show top 10 customers by invoice amount", + [ + """ + CREATE TABLE sales_fact ( + Customer_Name VARCHAR, + Product_Name VARCHAR, + Transaction_Amount FLOAT, + Invoice_ID VARCHAR + ); + """ + ], + ) + + assert sql == ( + 'SELECT TOP 10 "sales_fact"."Customer_Name" AS "Customer_Name", ' + 'SUM("sales_fact"."Transaction_Amount") AS "TotalTransaction_Amount" ' + 'FROM "sales_fact" ' + 'WHERE "sales_fact"."Customer_Name" IS NOT NULL ' + 'AND "sales_fact"."Transaction_Amount" IS NOT NULL ' + 'GROUP BY "sales_fact"."Customer_Name" ' + 'ORDER BY SUM("sales_fact"."Transaction_Amount") DESC' + ) def test_extract_explicit_table_names_from_repair_logs_phrase(): From 278c304a74cc6a2fcdd2cee603488b67f17d3bea Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 16 Jul 2026 02:14:49 +0530 Subject: [PATCH 0582/1087] Restore legacy ask retrieval flow --- .../retrieval/db_schema_retrieval.py | 514 +- wren-ai-service/src/web/v1/services/ask.py | 7235 +---------------- .../retrieval/test_db_schema_retrieval.py | 349 +- 3 files changed, 343 insertions(+), 7755 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..33c2617f55 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,6 +1,5 @@ import ast import logging -import re import sys from typing import TYPE_CHECKING, Any, Optional @@ -29,8 +28,6 @@ logger = logging.getLogger("wren-ai-service") -MAX_RELEVANT_TABLE_CANDIDATES = 5 - table_columns_selection_system_prompt = """ ### TASK ### @@ -128,413 +125,6 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline -def expand_business_terms_for_retrieval(query: str) -> str: - normalized = (query or "").lower() - expansions: list[str] = [] - - if any( - term in normalized - for term in ( - "amount", - "currency", - "currencies", - "customer", - "customers", - "invoice", - "invoices", - "market", - "markets", - "order", - "orders", - "product", - "products", - "category", - "categories", - "quantity", - "qty", - "region", - "regions", - "sales", - "salesperson", - "sales person", - "sold", - "value", - ) - ): - expansions.append( - "transaction purchase billing account geography area representative product item category sku quantity units sold amount value total metric money exchange currency" - ) - - if any( - term in normalized - for term in ("defect", "failure", "issue", "repair", "resolved", "status") - ): - expansions.append( - "issue defect category status resolved created updated date timestamp event" - ) - - if any(term in normalized for term in ("throughput", "production", "manufacturing")): - expansions.append( - "rate volume output capacity process unit group completed timestamp date" - ) - - if not expansions: - return query - - return f"{query}\n" + "\n".join(expansions) - - -def _normalize_retrieval_token(value: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - -def _retrieval_terms(value: str) -> set[str]: - stop_words = { - "about", - "across", - "and", - "are", - "ask", - "bar", - "chart", - "create", - "different", - "for", - "from", - "how", - "in", - "is", - "of", - "show", - "the", - "to", - "top", - "what", - "which", - "with", - } - terms = { - _normalize_retrieval_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") - if len(token) > 2 and token.lower() not in stop_words - } - return {term for term in terms if term} - - -def _query_mentions_any(query: str, terms: tuple[str, ...]) -> bool: - normalized = (query or "").lower() - return any(re.search(rf"\b{re.escape(term)}\b", normalized) for term in terms) - - -def _source_text(document: Document) -> str: - return " ".join( - str(part or "") - for part in ( - document.meta.get("name"), - document.meta.get("description"), - document.content, - ) - ).lower() - - -def _source_shape_score(query: str, document: Document) -> int: - normalized_query = (query or "").lower() - source_text = _source_text(document) - source_terms = _retrieval_terms(source_text) - - score = 0 - non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "stage", - "staging", - "temp", - "test", - "tmp", - ) - if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( - normalized_query, - non_production_terms, - ): - score -= 60 - - aggregation_terms = ( - "amount", - "average", - "avg", - "count", - "distribution", - "metric", - "revenue", - "sum", - "total", - "trend", - "value", - "volume", - ) - transaction_source_terms = ( - "activity", - "detail", - "event", - "fact", - "history", - "invoice", - "line", - "order", - "sale", - "sales", - "transaction", - ) - reference_source_terms = ( - "account", - "catalog", - "dimension", - "directory", - "entity", - "lookup", - "master", - "profile", - "reference", - ) - entity_listing_pattern = re.search( - r"\b(?:list|show|display|get|find)\b.*\b(?:accounts?|customers?|" - r"employees?|entities|items?|names?|products?|suppliers?|users?|vendors?)\b", - normalized_query, - ) - asks_for_aggregation = _query_mentions_any(normalized_query, aggregation_terms) or bool( - re.search(r"\b(?:by|per|each|top|bottom|rank|ranking)\b", normalized_query) - ) - asks_for_entity_listing = bool(entity_listing_pattern) and not asks_for_aggregation - - if asks_for_entity_listing: - if source_terms & set(reference_source_terms): - score += 35 - if source_terms & set(transaction_source_terms): - score -= 12 - elif asks_for_aggregation: - if source_terms & set(transaction_source_terms): - score += 25 - if source_terms & set(reference_source_terms): - score += 5 - - return score - - -def _document_relevance_score(document: Document, query_terms: set[str]) -> int: - if not query_terms: - return 0 - - document_terms = _retrieval_terms( - " ".join( - str(part or "") - for part in ( - document.meta.get("name"), - document.meta.get("description"), - document.content, - ) - ) - ) - if not document_terms: - return 0 - - score = 0 - for query_term in query_terms: - if query_term in document_terms: - score += 20 - continue - for document_term in document_terms: - if query_term in document_term or document_term in query_term: - score += 8 - break - return score - - -def _semantic_score(document: Document) -> float: - score = getattr(document, "score", None) - if isinstance(score, (int, float)): - return float(score) - score = document.meta.get("score") - if isinstance(score, (int, float)): - return float(score) - return 0.0 - - -def _score_table_documents( - query: str, documents: list[Document] -) -> list[tuple[float, int, Document, int, float]]: - if not documents: - return [] - - query_terms = _retrieval_terms(expand_business_terms_for_retrieval(query)) - if not query_terms: - return [ - (_semantic_score(document), -index, document, 0, _semantic_score(document)) - for index, document in enumerate(documents) - ] - - scored_documents: list[tuple[float, int, Document, int, float]] = [] - for index, document in enumerate(documents): - lexical_score = _document_relevance_score(document, query_terms) - semantic_score = _semantic_score(document) - source_shape_score = _source_shape_score(query, document) - combined_score = semantic_score + lexical_score + source_shape_score - scored_documents.append( - (combined_score, -index, document, lexical_score, semantic_score) - ) - - return sorted(scored_documents, key=lambda item: (item[0], item[1]), reverse=True) - - -def _rerank_table_documents(query: str, documents: list[Document]) -> list[Document]: - if not documents: - return documents - - reranked = _score_table_documents(query, documents) - if not reranked: - return documents - - logger.info( - "Top table candidates after retrieval rerank: %s", - [ - { - "name": document.meta.get("name"), - "semantic_score": round(semantic_score, 4), - "lexical_score": lexical_score, - "combined_score": round(combined_score, 4), - } - for combined_score, _index, document, lexical_score, semantic_score in reranked[ - :5 - ] - ], - ) - return [document for _score, _index, document, _lexical, _semantic in reranked] - - -def _select_relevant_table_documents( - query: str, - documents: list[Document], - *, - max_tables: int = MAX_RELEVANT_TABLE_CANDIDATES, -) -> list[Document]: - if not documents or max_tables <= 0: - return [] - - reranked = _score_table_documents(query, documents) - if not reranked: - return documents[:max_tables] - - candidate_pool = [item for item in reranked if item[3] > 0] or reranked - selected = [ - document - for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] - ] - if len(selected) < len(documents): - logger.info( - "Scoped table candidates for schema loading from %s to %s tables: %s", - len(documents), - len(selected), - [document.meta.get("name") for document in selected], - ) - return selected - - -def _is_project_wide_analysis_query(query: str) -> bool: - normalized = (query or "").lower() - if not normalized: - return False - - analysis_terms = { - "average", - "avg", - "bar chart", - "breakdown", - "chart", - "completed", - "compare", - "count", - "counts", - "distribution", - "group by", - "grouped", - "highest", - "line chart", - "lowest", - "maximum", - "minimum", - "monthly", - "most common", - "number of", - "pie chart", - "quarter", - "rank", - "ranking", - "recommend", - "recommended", - "show", - "status", - "sum", - "total", - "totals", - "top", - "trend", - "volume", - } - return any(term in normalized for term in analysis_terms) - - -def _dedupe_documents(documents: list[Document]) -> list[Document]: - deduped: list[Document] = [] - seen: set[tuple[str, str, str]] = set() - for document in documents: - key = ( - str(document.meta.get("name", "")), - str(document.meta.get("type", "")), - document.content, - ) - if key in seen: - continue - seen.add(key) - deduped.append(document) - return deduped - - -def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: - normalized: list[str] = [] - for table_name in table_names or []: - if not isinstance(table_name, str): - continue - table_name = table_name.strip() - if table_name and table_name not in normalized: - normalized.append(table_name) - return normalized - - -def _extract_table_names_from_table_retrieval( - table_retrieval: dict, explicit_tables: Optional[list[str]] = None -) -> list[str]: - table_names = _normalize_table_names(explicit_tables) - for document in table_retrieval.get("documents") or []: - if not isinstance(document, Document): - continue - table_name = document.meta.get("name") - if not isinstance(table_name, str): - try: - content = ast.literal_eval(document.content) - except (SyntaxError, ValueError): - content = {} - table_name = content.get("name") if isinstance(content, dict) else None - if isinstance(table_name, str): - table_name = table_name.strip() - if table_name and table_name not in table_names: - table_names.append(table_name) - return table_names - - @observe(capture_input=False, capture_output=False) async def embedding( query: str, @@ -553,7 +143,6 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query - query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -562,13 +151,9 @@ async def embedding( @observe(capture_input=False) async def table_retrieval( - query: str, - embedding: dict, - project_id: str, - tables: list[str], - table_retriever: Any, + embedding: dict, project_id: str, tables: list[str], table_retriever: Any ) -> dict: - base_filters = { + filters = { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, @@ -576,82 +161,61 @@ async def table_retrieval( } if project_id: - base_filters["conditions"].append( + filters["conditions"].append( {"field": "project_id", "operator": "==", "value": project_id} ) if embedding: - results = await table_retriever.run( + return await table_retriever.run( query_embedding=embedding.get("embedding"), - filters=base_filters, + filters=filters, ) - results["documents"] = _select_relevant_table_documents( - query, results.get("documents") or [] + elif tables: + filters["conditions"].append( + {"field": "name", "operator": "in", "value": tables} ) - return results - if tables: - logger.info("Loading explicit table descriptions: %s", tables) - explicit_filters = { - **base_filters, - "conditions": [ - *base_filters["conditions"], - {"field": "name", "operator": "in", "value": tables}, - ], - } - return await table_retriever.run(query_embedding=[], filters=explicit_filters) + return await table_retriever.run( + query_embedding=[], + filters=filters, + ) return {"documents": []} @observe(capture_input=False) async def dbschema_retrieval( - query: str, - table_retrieval: dict, - project_id: str, - dbschema_retriever: Any, - tables: Optional[list[str]] = None, + table_retrieval: dict, project_id: str, dbschema_retriever: Any ) -> list[Document]: - selected_table_names = _extract_table_names_from_table_retrieval( - table_retrieval, tables - ) + tables = table_retrieval.get("documents", []) + table_names = [] + for table in tables: + content = ast.literal_eval(table.content) + table_names.append(content["name"]) + + table_name_conditions = [ + {"field": "name", "operator": "==", "value": table_name} + for table_name in table_names + ] - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + if table_name_conditions: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } - if selected_table_names: - filters["conditions"].append( - {"field": "name", "operator": "in", "value": selected_table_names} - ) - logger.info( - "Loading selected deployed schema metadata for active project_id %s tables=%s", - project_id, - selected_table_names, - ) - elif not query: - logger.info( - "Loading complete deployed schema metadata for active project_id %s", - project_id, - ) - else: - logger.info( - "No relevant table-description candidates found for active project_id %s; " - "skipping full schema loading for query=%s", - project_id, - query, - ) - return [] + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results.get("documents", []) + return [] @observe() diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c4a7e11b14..c4f6e0bfa7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,34 +1,17 @@ import asyncio import logging -import re -from typing import Any, Dict, List, Literal, Optional +from typing import Dict, List, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - normalize_sql_direction_keywords, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, -) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent logger = logging.getLogger("wren-ai-service") -NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( - "No relevant data found in the active datasource for this question." -) -MAX_FORCED_EXPLICIT_TABLES = 5 - - -async def _return_value(value): - return value - class AskHistory(BaseModel): sql: str @@ -96,7 +79,7 @@ class _AskResultResponse(BaseModel): rephrased_question: Optional[str] = None intent_reasoning: Optional[str] = None sql_generation_reasoning: Optional[str] = None - type: Optional[Literal["GENERAL", "TEXT_TO_SQL", "MISLEADING_QUERY"]] = None + type: Optional[Literal["GENERAL", "TEXT_TO_SQL"]] = None retrieved_tables: Optional[List[str]] = None response: Optional[List[AskResult]] = None invalid_sql: Optional[str] = None @@ -116,39 +99,6 @@ class AskResultResponse(_AskResultResponse): class AskService: - _HISTORICAL_QUESTION_STOP_WORDS = { - "a", - "an", - "and", - "are", - "as", - "at", - "be", - "by", - "can", - "chart", - "create", - "each", - "for", - "from", - "give", - "graph", - "how", - "in", - "is", - "me", - "of", - "on", - "please", - "show", - "the", - "to", - "total", - "what", - "which", - "with", - } - def __init__( self, pipelines: Dict[str, BasicPipeline], @@ -158,9 +108,9 @@ def __init__( allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, - max_sql_correction_retries: int = 3, pipeline_timeout_seconds: int = 90, schema_retrieval_timeout_seconds: int = 180, + max_sql_correction_retries: int = 3, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -169,9 +119,6 @@ def __init__( self._ask_results: Dict[str, AskResultResponse] = TTLCache( maxsize=maxsize, ttl=ttl ) - self._general_streaming_results: Dict[str, str] = TTLCache( - maxsize=maxsize, ttl=ttl - ) self._allow_sql_generation_reasoning = allow_sql_generation_reasoning self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_intent_classification = allow_intent_classification @@ -191,6202 +138,129 @@ def _is_stopped(self, query_id: str, container: dict): return False - @classmethod - def _normalize_historical_question_text(cls, question: str | None) -> str: - return " ".join(re.findall(r"[a-z0-9]+", (question or "").lower())) - - @classmethod - def _historical_question_tokens(cls, question: str | None) -> set[str]: - normalized = cls._normalize_historical_question_text(question) - return { - token - for token in normalized.split() - if len(token) > 1 and token not in cls._HISTORICAL_QUESTION_STOP_WORDS - } - - @classmethod - def _is_reusable_historical_question( - cls, query: str | None, historical_question: str | None - ) -> bool: - normalized_query = cls._normalize_historical_question_text(query) - normalized_historical_question = cls._normalize_historical_question_text( - historical_question - ) - if not normalized_query or not normalized_historical_question: - return False - return normalized_query == normalized_historical_question - - @classmethod - def _should_use_histories_for_query(cls, query: str | None) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - contextual_prefixes = ( - "also ", - "and ", - "but ", - "for those ", - "for that ", - "for the same ", - "from that ", - "how about ", - "in that ", - "now ", - "same ", - "show more", - "show the same", - "then ", - "use that ", - "what about ", - "what if ", - ) - if normalized.startswith(contextual_prefixes): - return True - - contextual_patterns = ( - r"\b(previous|last|above|earlier|same|those|that|these|them|it|its|there)\b", - r"\b(add|break down|compare|filter|group|instead|only|sort|split)\b.+\b(by|to|with)\b", - r"\b(by|for|with)\s+(month|quarter|year|status|type|category|customer|market|region|country|division)\b", - ) - return any(re.search(pattern, normalized) for pattern in contextual_patterns) - - def _is_greeting_query(self, query: str) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - greeting_patterns = { - "hi", - "hello", - "hey", - "hii", - "hola", - "good morning", - "good afternoon", - "good evening", - "how are you", - "thanks", - "thank you", - } - return normalized in greeting_patterns - - def _is_data_analysis_query(self, query: str) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - analysis_terms = { - "amount", - "average", - "avg", - "bar chart", - "bottom", - "chart", - "common", - "compare", - "count", - "cost", - "claim", - "claims", - "currency", - "currencies", - "customer", - "customers", - "dashboard", - "debug", - "distribution", - "failure", - "fastest growing", - "growth", - "group", - "grouped", - "invoice", - "invoices", - "margin", - "market", - "markets", - "monthly", - "order", - "orders", - "pcb", - "performance", - "profit", - "product", - "products", - "product type", - "product types", - "quarter", - "quantity", - "rank", - "ranking", - "region", - "regions", - "repair", - "resolved", - "revenue", - "sale", - "sales", - "sales person", - "sales rep", - "salesperson", - "sla", - "top", - "trend", - "turnaround", - "value", - "volume", - "year", - "yearly", - } - return any(term in normalized for term in analysis_terms) - - def _needs_conversation_context(self, query: str) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - return any( - re.search(pattern, normalized) - for pattern in ( - r"\b(previous|last|above|earlier)\s+(query|question|answer|result|sql|chart)\b", - r"\b(same|that|those|them|it)\s+(table|query|question|result|chart|sql|period|filter)\b", - r"\b(use|using|based on|compare with|compared with)\s+(that|previous|last|above|earlier)\b", - r"\bwhat about\b", - r"\bhow about\b", - ) - ) - - def _should_reuse_historical_question_sql( - self, - query: str, - histories: list[AskHistory] | None, - ) -> bool: - return False - - def _rewrite_query_for_text_to_sql(self, query: str) -> str: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return query - - guidance: list[str] = [] - - if any( - term in normalized - for term in ("chart", "bar chart", "line chart", "pie chart", "graph") - ): - guidance.append( - "Return SQL only for the aggregated dataset required to build the requested chart." - ) - - if any( - term in normalized - for term in ( - "failure category", - "failure categories", - "common failure", - "common failures", - "failure code", - "top 10", - "most common", - ) - ): - guidance.append( - "Use an exposed failure category, failure name, or failure code field from the schema and return that dimension with a count metric." - ) - - if any( - term in normalized - for term in ("monthly", "last 12 months", "last month", "trend", "volume") - ): - guidance.append( - "Use a real timestamp column from the schema and aggregate results by calendar month when a monthly trend is requested." - ) - - if re.search( - r"\b(?:by|across|per|each|grouped by|group by)\s+[a-z][a-z0-9 _-]*", - normalized, - ) or "over time" in normalized: - guidance.append( - "Preserve explicit grouping dimensions requested by the question, such as market, region, currency, product, customer, status, type, or category, using only matching columns exposed in the provided schema." - ) - - if any(term in normalized for term in ("currency", "currencies", "fx")): - guidance.append( - "For currency questions, use an exposed currency, money, exchange, or FX code/name column from the schema and group by it." - ) - - if any( - term in normalized - for term in ( - "amount", - "cost", - "revenue", - "sales value", - "sum", - "total", - "value", - ) - ) and not any( - term in normalized - for term in ("count", "how many", "number of records", "record count") - ): - guidance.append( - "When the question asks for total, sum, amount, value, revenue, or cost, aggregate an exposed numeric measure with SUM; use COUNT only for record-count questions." - ) - - if not guidance: - return query - - return f"{query}\n\nSQL generation guidance:\n- " + "\n- ".join(guidance) - - def _schema_contains( - self, - table_ddls: list[str], - pattern: str, - table_names: Optional[list[str]] = None, - ) -> bool: - schema_text = "\n".join(table_ddls or []) - if table_names: - schema_text += "\n" + "\n".join(table_names) - return bool(re.search(pattern, schema_text, flags=re.IGNORECASE)) - - def _schema_has_table_column( - self, - table_ddls: list[str], - table_name: str, - column_name: str, - table_names: Optional[list[str]] = None, - ) -> bool: - table_pattern = rf"\b{re.escape(table_name)}\b" - column_pattern = rf"\b{re.escape(column_name)}\b" - - for ddl in table_ddls or []: - if re.search(table_pattern, ddl, flags=re.IGNORECASE) and re.search( - column_pattern, ddl, flags=re.IGNORECASE - ): - return True - - return False - - def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: - column_names: list[str] = [] - non_column_prefixes = ( - "create ", - "constraint ", - "foreign ", - "primary ", - "unique ", - "index ", - ")", - "/*", - "--", - ) - - for ddl in table_ddls: - if not isinstance(ddl, str): - continue - for line in ddl.splitlines(): - stripped = line.strip().rstrip(",") - if not stripped: - continue - if stripped.lower().startswith(non_column_prefixes): - continue - - column_match = re.match( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_]*))\s+", - stripped, - ) - if not column_match: - continue - - column_name = next( - value for value in column_match.groupdict().values() if value - ) - column_names.append(str(column_name).lower()) - - return column_names - - def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: - tables: list[dict[str, Any]] = [] - for ddl in table_ddls or []: - if not isinstance(ddl, str): - continue - table_match = re.search( - r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", - ddl, - flags=re.IGNORECASE, - ) - if not table_match: - continue - - table_name = next( - (value for value in table_match.groupdict().values() if value), - None, - ) - if not table_name: - continue - body_start = table_match.end() - depth = 1 - body_end = body_start - while body_end < len(ddl) and depth > 0: - if ddl[body_end] == "(": - depth += 1 - elif ddl[body_end] == ")": - depth -= 1 - body_end += 1 - - columns: list[dict[str, str]] = [] - for line in ddl[body_start : body_end - 1].splitlines(): - stripped = line.strip().rstrip(",") - if not stripped or stripped.startswith(("--", "/*")): - continue - if re.match( - r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", - stripped, - flags=re.IGNORECASE, - ): - continue - - column_match = re.match( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_$]*))" - r"\s+(?P[A-Za-z0-9_(),]+)", - stripped, - ) - if column_match: - column_name = next( - (value - for key, value in column_match.groupdict().items() - if key != "type" and value - ), - None, - ) - if not column_name: - continue - column_type = column_match.group("type") or "" - columns.append( - { - "name": str(column_name), - "type": str(column_type).lower(), - } - ) - - tables.append({"name": table_name, "columns": columns}) - - return tables - - _INTENT_STOPWORDS = { - "a", - "an", - "and", - "are", - "as", - "based", - "be", - "by", - "can", - "chart", - "correct", - "data", - "different", - "do", - "does", - "each", - "for", - "from", - "give", - "how", - "in", - "is", - "it", - "list", - "many", - "me", - "of", - "on", - "or", - "per", - "question", - "rate", - "records", - "reduce", - "show", - "system", - "taken", - "the", - "there", - "to", - "total", - "type", - "types", - "what", - "which", - "with", - } - - def _intent_tokens(self, text: str) -> set[str]: - tokens: set[str] = set() - for raw_token in re.findall(r"[A-Za-z][A-Za-z0-9_]*", text or ""): - split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) - for token in re.findall(r"[A-Za-z0-9]+", split_token.lower()): - if len(token) <= 2 or token in self._INTENT_STOPWORDS: - continue - tokens.add(token) - if token.endswith("ies") and len(token) > 4: - tokens.add(token[:-3] + "y") - elif token.endswith("s") and len(token) > 3: - tokens.add(token[:-1]) - return tokens - - def _schema_name_tokens(self, name: str) -> set[str]: - spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(name or "")) - return { - token - for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) - if len(token) > 1 - } - - def _table_for_sql_reference( - self, table_reference: str, valid_tables: dict[str, dict[str, Any]] - ) -> dict[str, Any] | None: - table_key = str(table_reference or "").lower() - if table_key in valid_tables: - return valid_tables[table_key] - suffix_key = table_key.split(".")[-1] - for valid_table_name, table in valid_tables.items(): - if valid_table_name.split(".")[-1] == suffix_key: - return table - return None - - def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return [] - - concept_groups: list[set[str]] = [] - if "product line" in normalized or "productline" in normalized: - concept_groups.append({"product", "prod", "line", "productline"}) - if "pcb" in normalized: - concept_groups.append( - { - "pcb", - "board", - "repair", - "repairs", - "debug", - "failure", - "failures", - } - ) - if "critical" in normalized: - concept_groups.append({"critical", "severity", "priority"}) - if "cost" in normalized: - concept_groups.append({"cost", "amount", "expense", "impact"}) - if "currency" in normalized or "currencies" in normalized: - concept_groups.append({"currency", "curr", "money", "fx", "exchange"}) - if "market" in normalized or "markets" in normalized: - concept_groups.append({"market", "region", "country", "territory"}) - if "region" in normalized or "regions" in normalized: - concept_groups.append({"region", "market", "area", "territory", "country"}) - if "quarterly" in normalized or "quarter" in normalized: - concept_groups.append({"quarter", "quarterly"}) - if "recurring" in normalized or "recurrence" in normalized: - concept_groups.append({"recurring", "recurrence", "occurrence", "occurrences", "count"}) - if "issue" in normalized or "issues" in normalized: - concept_groups.append({"issue", "issues", "failure", "failures", "problem", "defect"}) - - return concept_groups - - def _sql_covers_required_question_concepts( - self, - sql: str, - query: str | None, - referenced_column_tokens: set[str], - referenced_table_tokens: set[str], - ) -> bool: - sql_text = (sql or "").lower() - available_tokens = referenced_column_tokens | referenced_table_tokens - for concept_group in self._required_sql_concept_groups(query): - if concept_group & available_tokens: - continue - if any(token in sql_text for token in concept_group): - continue - logger.warning( - "Ignoring SQL because it does not cover required question concept. " - "query=%s required=%s referenced_column_tokens=%s referenced_table_tokens=%s sql=%s", - query, - sorted(concept_group), - sorted(referenced_column_tokens), - sorted(referenced_table_tokens), - sql, - ) - return False - return True - - def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return True - asks_for_measure_sum = any( - term in normalized_query - for term in ( - "amount", - "cost", - "revenue", - "sales value", - "sum", - "total", - "value", - ) - ) - asks_for_count = any( - term in normalized_query - for term in ( - "count", - "counts", - "how many", - "number of", - "record count", - "records", - "rows", - ) - ) - if not asks_for_measure_sum or asks_for_count: - return True - - normalized_sql = re.sub(r"\s+", " ", sql or "").lower() - if re.search(r"\b(sum|avg|min|max)\s*\(", normalized_sql): - return True - if re.search(r"\bcount\s*\(", normalized_sql): - logger.warning( - "Ignoring SQL because a measure-total question was answered with row counting. " - "query=%s sql=%s", - query, - sql, - ) - return False - return True - - def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not re.search( - r"\b(?:distinct|unique|no duplicate|no duplicates|without duplicates)\b", - normalized_query, - ): - return True - - normalized_sql = re.sub(r"\s+", " ", sql or "").strip() - if re.search(r"\bcount\s*\(\s*distinct\b", normalized_sql, flags=re.IGNORECASE): - return True - if re.search(r"\bGROUP\s+BY\b", normalized_sql, flags=re.IGNORECASE): - group_match = re.search( - r"\bGROUP\s+BY\b(?P.*?)(?:\bORDER\s+BY\b|\bHAVING\b|$)", - normalized_sql, - flags=re.IGNORECASE | re.DOTALL, - ) - if group_match: - group_items = [ - item.strip() - for item in re.split(r",(?![^()]*\))", group_match.group("group")) - if item.strip() - ] - if len(group_items) > 1: - logger.warning( - "Ignoring SQL because GROUP BY covers multiple columns and can still duplicate the requested entity. " - "query=%s sql=%s", - query, - sql, - ) - return False - return True - select_match = re.search( - r"\bSELECT\b(?P.*?)\bFROM\b", - sql or "", - flags=re.IGNORECASE | re.DOTALL, - ) - if not select_match: - return [] - - invalid: list[str] = [] - for match in re.finditer( - r'(?P(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$.]*))\s+AS\s+' - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_]*))', - select_match.group("select"), - flags=re.IGNORECASE, - ): - expression = match.group("expr").strip('"[]') - expression_key = expression.split(".")[-1].lower() - alias = match.group("quoted") or match.group("bracketed") or match.group("bare") or "" - alias_key = alias.lower() - if not expression_key or not alias_key: - continue - if alias_key in valid_columns or expression_key not in valid_columns: - continue - alias_terms = { - term for term in re.split(r"[^a-z0-9]+", alias_key) if term - } - if alias_terms & allowed_alias_terms: - continue - invalid.append(alias) - - return invalid - - def _unqualified_valid_sql_column_tokens( - self, sql: str, schema_tables: list[dict[str, Any]] - ) -> set[str]: - valid_columns = { - str(column.get("name") or "").lower(): str(column.get("name") or "") - for table in schema_tables - for column in table.get("columns", []) - if column.get("name") - } - if not valid_columns: - return set() - - sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") - identifier_pattern = re.compile( - r'"(?P[^"]+)"|(?P\b[A-Za-z_][A-Za-z0-9_]*\b)' - ) - tokens: set[str] = set() - for match in identifier_pattern.finditer(sql_without_strings): - identifier = match.group("quoted") or match.group("bare") or "" - identifier_key = identifier.lower() - if identifier_key not in valid_columns: - continue - - before = sql_without_strings[: match.start()].rstrip() - after = sql_without_strings[match.end() :].lstrip() - if before.endswith(".") or after.startswith("."): - continue - - previous_word_match = re.search(r"([A-Za-z_][A-Za-z0-9_]*)\s*$", before) - previous_word = ( - previous_word_match.group(1).lower() if previous_word_match else "" - ) - if previous_word == "as" and self._is_alias_identifier_position( - sql_without_strings, - match.start(), - ): - continue - - tokens.update(self._schema_name_tokens(valid_columns[identifier_key])) - - return tokens - - def _sql_matches_question_intent( - self, - sql: str, - query: str | None, - schema_tables: list[dict[str, Any]], - ) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - expects_dimension = bool( - re.search( - r"\b(?:by|per|each|which|different|type|types|category|" - r"categories|status|source|market|markets|region|regions|" - r"currency|currencies)\b", - normalized_query, - ) - ) - - question_tokens = self._intent_tokens(query or "") - required_concept_groups = self._required_sql_concept_groups(query) - if not question_tokens and not required_concept_groups: - return True - - valid_tables = { - str(table.get("name") or "").lower(): table - for table in schema_tables - if table.get("name") - } - if not valid_tables: - return True - - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_tables = [ - next(value for value in match.groupdict().values() if value) - for match in table_reference_pattern.finditer(sql) - ] - if not referenced_tables: - return True - - referenced_table_tokens = set().union( - *[self._schema_name_tokens(table) for table in referenced_tables] - ) - qualified_column_pattern = re.compile( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_$]*))", - flags=re.IGNORECASE, - ) - referenced_columns_by_table: dict[str, set[str]] = {} - for match in qualified_column_pattern.finditer(sql): - table_reference = ( - match.group("table_quoted") - or match.group("table_bracketed") - or match.group("table_bare") - or "" - ).lower() - column_reference = ( - match.group("column_quoted") - or match.group("column_bracketed") - or match.group("column_bare") - or "" - ) - referenced_columns_by_table.setdefault(table_reference, set()).add( - column_reference - ) - - all_referenced_column_tokens = set().union( - *[ - self._schema_name_tokens(column_name) - for columns in referenced_columns_by_table.values() - for column_name in columns - ] - ) if referenced_columns_by_table else set() - all_referenced_column_tokens.update( - self._unqualified_valid_sql_column_tokens(sql, schema_tables) - ) - if not self._sql_covers_required_question_concepts( - sql, - query, - all_referenced_column_tokens, - referenced_table_tokens, - ): - return False - if not self._sql_uses_required_measure_aggregation(sql, query): - return False - if not self._sql_satisfies_unique_entity_request(sql, query): - return False - - if not expects_dimension: - return True - - for table_reference in referenced_tables: - table = self._table_for_sql_reference(table_reference, valid_tables) - if not table: - continue - - columns = [ - column for column in table.get("columns", []) if column.get("name") - ] - intent_matching_columns = [ - str(column.get("name")) - for column in columns - if self._schema_name_tokens(str(column.get("name"))) & question_tokens - ] - if not intent_matching_columns: - continue - - table_key = str(table_reference or "").lower() - referenced_columns = referenced_columns_by_table.get( - table_key - ) or referenced_columns_by_table.get( - table_key.split(".")[-1], - set(), - ) - referenced_column_tokens = ( - set().union( - *[ - self._schema_name_tokens(column_name) - for column_name in referenced_columns - ] - ) - if referenced_columns - else all_referenced_column_tokens - ) - - if not referenced_column_tokens & question_tokens: - logger.warning( - "Ignoring SQL because selected columns do not match question intent. " - "query=%s table=%s matching_schema_columns=%s referenced_columns=%s sql=%s", - query, - table.get("name"), - intent_matching_columns, - sorted(referenced_columns), - sql, - ) - return False - - return True - - def _is_numeric_schema_type(self, column_type: str) -> bool: - return bool( - re.search( - r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|float|double|" - r"real|money|number)\b", - column_type, - flags=re.IGNORECASE, - ) - ) - - def _is_temporal_schema_type(self, column_type: str) -> bool: - return bool( - re.search( - r"\b(?:date|time|timestamp|datetime|smalldatetime)\b", - column_type, - flags=re.IGNORECASE, - ) - ) - - def _is_text_schema_type(self, column_type: str) -> bool: - return bool( - re.search( - r"\b(?:char|text|string|varchar|nvarchar|uuid|guid|json)\b", - column_type, - flags=re.IGNORECASE, - ) - ) - - def _find_schema_column( - self, - table: dict[str, Any], - candidates: tuple[str, ...], - numeric: bool | None = None, - temporal: bool | None = None, - ) -> str | None: - normalized_candidates = [ - re.sub(r"[^a-z0-9]", "", str(candidate).lower()) - for candidate in candidates - if candidate is not None - ] - if not normalized_candidates: - return None - scored: list[tuple[int, int, str]] = [] - for column in table.get("columns", []): - column_name = column.get("name") - if not column_name: - continue - column_name = str(column_name) - normalized_column = re.sub(r"[^a-z0-9]", "", column_name.lower()) - column_type = str(column.get("type") or "") - if numeric is True and not self._is_numeric_schema_type(column_type): - continue - if temporal is True and not self._is_temporal_schema_type(column_type): - continue - - for candidate_index, candidate in enumerate(normalized_candidates): - if normalized_column == candidate: - scored.append((100, candidate_index, column_name)) - elif candidate and candidate in normalized_column: - scored.append((60 + len(candidate), candidate_index, column_name)) - elif normalized_column and normalized_column in candidate: - scored.append( - (40 + len(normalized_column), candidate_index, column_name) - ) - - if not scored: - return None - - return sorted(scored, key=lambda item: (-item[0], item[1]))[0][2] - - def _find_first_schema_column( + @observe(name="Ask Question") + @trace_metadata + async def ask( self, - table: dict[str, Any], - candidates: tuple[str, ...], - *, - avoid: set[str] | None = None, - ) -> str | None: - avoid = {str(column).lower() for column in avoid or set()} - for candidate_group in candidates: - column = self._find_schema_column(table, (candidate_group,)) - if column and column.lower() not in avoid: - return column - return None - - def _find_any_temporal_schema_column(self, table: dict[str, Any]) -> str | None: - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if column_name and self._is_temporal_schema_type(column_type): - return column_name - return None - - def _is_probable_explicit_table_token(self, token: str) -> bool: - normalized = re.sub(r"\s+", " ", str(token or "").strip().lower()) - if not normalized: - return False - if normalized in self._INTENT_STOPWORDS or normalized in { - "last", - "latest", - "recent", - "current", - "previous", - "next", - "month", - "year", - "quarter", - "week", - "day", - }: - return False - return True - - def _quote_sql_identifier(self, identifier: str) -> str: - return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' - - def _normalize_schema_identifier_key(self, value: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - def _schema_identifier_alias_keys(self, value: str) -> set[str]: - raw_value = str(value or "").strip() - base_key = self._normalize_schema_identifier_key(raw_value) - separator_normalized_key = self._normalize_schema_identifier_key( - re.sub(r"[.$]", "_", raw_value) - ) - compact_parts_key = "".join( - self._normalize_schema_identifier_key(part) - for part in re.split(r"[.$_]+", raw_value) - if part - ) - return { - key - for key in (base_key, separator_normalized_key, compact_parts_key) - if key + ask_request: AskRequest, + **kwargs, + ): + trace_id = kwargs.get("trace_id") + results = { + "ask_result": {}, + "metadata": { + "type": "", + "error_type": "", + "error_message": "", + "request_from": ask_request.request_from, + }, } - def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: - return self._explicit_table_alias_keys( - self._extract_explicit_table_names_from_query(query or "") + query_id = ask_request.query_id + histories = ask_request.histories[: self._max_histories][ + ::-1 + ] # reverse the order of histories + rephrased_question = None + intent_reasoning = None + sql_generation_reasoning = None + sql_samples = [] + instructions = [] + api_results = [] + table_names = [] + error_message = None + invalid_sql = None + allow_sql_generation_reasoning = ( + self._allow_sql_generation_reasoning + and not ask_request.ignore_sql_generation_reasoning ) - - def _explicit_table_alias_keys(self, table_names: list[str]) -> set[str]: - keys: set[str] = set() - for table_name in table_names: - keys.update(self._schema_identifier_alias_keys(table_name)) - return keys - - def _filter_retrieval_metadata_for_explicit_query( - self, - query: str, - documents: list[dict], - explicit_table_names: Optional[list[str]] = None, - ) -> tuple[list[dict], list[str], list[str]]: - explicit_table_keys = ( - self._explicit_table_alias_keys(explicit_table_names) - if explicit_table_names - else self._explicit_table_alias_keys_from_query(query) + enable_column_pruning = ( + self._enable_column_pruning or ask_request.enable_column_pruning ) - if not explicit_table_keys: - table_names, table_ddls = self._metadata_from_documents(documents) - return documents, table_names, table_ddls + allow_sql_functions_retrieval = self._allow_sql_functions_retrieval + allow_sql_diagnosis = self._allow_sql_diagnosis + max_sql_correction_retries = self._max_sql_correction_retries + current_sql_correction_retries = 0 + use_dry_plan = ask_request.use_dry_plan + allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback - matched_documents: list[dict] = [] - for document in documents: - candidate_names = [] - if isinstance(table_name := document.get("table_name"), str): - candidate_names.append(table_name) - if isinstance(table_ddl := document.get("table_ddl"), str): - candidate_names.extend( - str(table.get("name") or "") - for table in self._parse_schema_tables([table_ddl]) - if table.get("name") - ) + try: + user_query = ask_request.query - candidate_keys: set[str] = set() - for candidate_name in candidate_names: - candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) - candidate_keys.update( - self._schema_identifier_alias_keys( - re.split(r"[.$_]", str(candidate_name or ""))[-1] - ) + # ask status can be understanding, searching, generating, finished, failed, stopped + # we will need to handle business logic for each status + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="understanding", + trace_id=trace_id, + is_followup=True if histories else False, ) - if explicit_table_keys.intersection(candidate_keys): - matched_documents.append(document) - - table_names, table_ddls = self._metadata_from_documents(matched_documents) - return matched_documents, table_names, table_ddls - - def _sql_references_explicit_table( - self, - sql: str, - query: str | None, - ) -> bool: - explicit_table_keys = self._explicit_table_alias_keys_from_query(query) - if not explicit_table_keys: - return True - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_tables = [ - next(value for value in match.groupdict().values() if value) - for match in table_reference_pattern.finditer(sql) - ] - - for table_reference in referenced_tables: - reference_keys = self._schema_identifier_alias_keys(table_reference) - reference_keys.update( - self._schema_identifier_alias_keys( - re.split(r"[.$_]", str(table_reference or ""))[-1] - ) - ) - if explicit_table_keys.intersection(reference_keys): - return True - - logger.warning( - "Ignoring SQL because it does not reference the explicitly requested table. " - "query=%s referenced_tables=%s sql=%s", - query, - referenced_tables, - sql, - ) - return False - - def _table_matches_query(self, table_name: str, query: str) -> bool: - query_keys = self._schema_identifier_alias_keys(query) - short_table = re.split(r"[.$_]", str(table_name or ""))[-1] - table_keys = self._schema_identifier_alias_keys(table_name) - table_keys.update(self._schema_identifier_alias_keys(short_table)) - return any( - table_key and any(table_key in query_key for query_key in query_keys) - for table_key in table_keys - ) - - def _find_best_schema_table_for_query( - self, query: str, tables: list[dict[str, Any]] - ) -> dict[str, Any] | None: - if not tables: - return None - - scored_tables: list[tuple[int, dict[str, Any]]] = [] - query_tokens = self._intent_tokens(query) - for table in tables: - table_name = str(table.get("name") or "") - if not table_name: - continue - - score = 0 - if self._table_matches_query(table_name, query): - score += 100 - - table_tokens = self._schema_name_tokens(table_name) - score += 8 * len(table_tokens & query_tokens) - - column_token_matches = 0 - for column in table.get("columns", []): - column_token_matches += len( - self._schema_name_tokens(str(column.get("name") or "")) - & query_tokens - ) - score += column_token_matches - - if score > 0: - scored_tables.append((score, table)) - - if scored_tables: - return sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] - if len(tables) == 1: - return tables[0] - return None - - def _query_mentions_column(self, query: str, column_name: str) -> bool: - normalized_query = self._normalize_schema_identifier_key(query) - normalized_column = self._normalize_schema_identifier_key(column_name) - if not normalized_column: - return False - if normalized_column in normalized_query: - return True - if normalized_column.endswith("y"): - return f"{normalized_column[:-1]}ies" in normalized_query - return f"{normalized_column}s" in normalized_query - - def _is_alias_identifier_position(self, sql: str, start: int) -> bool: - before = sql[:start].rstrip() - return bool(re.search(r"\bAS\s*$", before, flags=re.IGNORECASE)) - - def _find_dimension_column_for_query( - self, query: str, table: dict[str, Any] - ) -> str | None: - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - if column_name and self._query_mentions_column(query, column_name): - return column_name - - candidate_columns = [ - str(column.get("name")) - for column in table.get("columns", []) - if column.get("name") - and not self._is_temporal_schema_type(str(column.get("type") or "")) - ] - for candidate in ("name", "category", "type", "status", "code"): - column = self._find_schema_column(table, (candidate,)) - if column in candidate_columns: - return column - return candidate_columns[0] if candidate_columns else None - - def _build_schema_ranked_measure_sql( - self, - query: str, - table_ddls: list[str], - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - if not re.search( - r"\b(?:top|highest|largest|biggest|best|lowest|smallest|bottom)\b", - normalized_query, - ): - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - query_tokens = self._intent_tokens(query) - if not query_tokens: - return None - - scored: list[tuple[int, dict[str, Any], str, str]] = [] - for table in tables: - text_columns: list[tuple[int, str]] = [] - numeric_columns: list[tuple[int, str]] = [] - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if not column_name: - continue - - column_tokens = self._schema_name_tokens(column_name) - query_overlap = column_tokens & query_tokens - - if self._is_numeric_schema_type(column_type): - if query_overlap: - numeric_columns.append( - (80 + 5 * len(query_overlap), column_name) - ) - continue - - if self._is_temporal_schema_type(column_type): - continue - - if query_overlap: - text_columns.append((60 + 5 * len(query_overlap), column_name)) - - if not text_columns or not numeric_columns: - continue - - dimension_score, dimension = sorted( - text_columns, - key=lambda item: item[0], - reverse=True, - )[0] - measure_score, measure = sorted( - numeric_columns, - key=lambda item: item[0], - reverse=True, - )[0] - table_score = dimension_score + measure_score - table_tokens = self._schema_name_tokens(str(table.get("name") or "")) - table_score += 5 * len(table_tokens & query_tokens) - scored.append((table_score, table, dimension, measure)) - - if not scored: - return None - - _, table, dimension, measure = sorted( - scored, - key=lambda item: item[0], - reverse=True, - )[0] - table_name = str(table.get("name") or "") - if not table_name: - return None - - limit = self._extract_requested_top_n(query, default_value=10) - table_ref = self._quote_sql_identifier(table_name) - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - measure_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" - metric_expr = f"SUM({measure_ref})" - direction = ( - "ASC" - if re.search(r"\b(?:lowest|smallest|least|bottom)\b", normalized_query) - else "DESC" - ) - return ( - f"SELECT TOP {limit} {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " - f"{metric_expr} AS {self._quote_sql_identifier('Total' + measure)} " - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL AND {measure_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f"ORDER BY {metric_expr} {direction}" - ) - - def _find_temporal_column_for_query( - self, query: str, table: dict[str, Any] - ) -> str | None: - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if ( - column_name - and self._is_temporal_schema_type(column_type) - and self._query_mentions_column(query, column_name) - ): - return column_name - - return self._find_any_temporal_schema_column(table) - - def _build_schema_grounded_table_question_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - tables = self._parse_schema_tables(table_ddls) - table = self._find_best_schema_table_for_query(query, tables) - if not table: - return None - - table_name = str(table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - limit = self._extract_requested_top_n(query, default_value=10) - - wants_latest_records = any( - term in normalized - for term in ("latest", "recent", "newest", "last records", "latest records") - ) - if wants_latest_records: - date_column = self._find_temporal_column_for_query(query, table) - if not date_column: - return None - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT TOP {limit} * " - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"ORDER BY {date_ref} DESC" - ) - - wants_monthly_count = any( - term in normalized - for term in ("monthly", "by month", "per month", "month-wise") - ) and any(term in normalized for term in ("count", "records", "rows")) - if wants_monthly_count: - date_column = self._find_temporal_column_for_query(query, table) - if not date_column: - return None - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " - f"DATEPART(MONTH, {date_ref}) ASC" - ) - - wants_total_count = ( - re.search(r"\bhow many\b", normalized) - or "record count" in normalized - or "count of records" in normalized - or "number of records" in normalized - ) and not re.search(r"\b(?:by|per|each|distribution|highest|top)\b", normalized) - if wants_total_count: - return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' - - wants_grouped_count = ( - any( - term in normalized - for term in ( - "count", - "counts", - "record count", - "number of", - "how many", - ) - ) - and re.search(r"\b(?:by|per|each|grouped by|group by)\b", normalized) - ) - wants_ranked_count = any( - term in normalized for term in ("highest", "top", "most", "largest") - ) and any( - term in normalized - for term in ("count", "counts", "number of", "orders", "records", "rows") - ) - if wants_grouped_count or wants_ranked_count: - dimension_column = self._find_dimension_column_for_query(query, table) - if not dimension_column: - return None - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" - count_column = None - if any(term in normalized for term in ("order", "orders")): - count_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "OrderID", "id"), - ) - count_expression = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(count_column)})" - if count_column - else "COUNT(*)" - ) - top_clause = f"TOP {limit} " if wants_ranked_count else "" - nonblank_filter = ( - f"AND LTRIM(RTRIM({dimension_ref})) <> '' " - if wants_ranked_count - else "" - ) - return ( - f"SELECT {top_clause}{dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_column)}, " - f'{count_expression} AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"{nonblank_filter}" - f"GROUP BY {dimension_ref} " - f"ORDER BY {count_expression} DESC" - ) - - wants_distribution = any( - term in normalized - for term in ( - "distribution", - "highest occurrence", - "highest occurrences", - "most occurrence", - "most occurrences", - "occurrences", - "top", - "common", - ) - ) - if wants_distribution: - dimension_column = self._find_dimension_column_for_query(query, table) - if not dimension_column: - return None - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" - - occurrence_column = self._find_schema_column( - table, - ("occurrences", "occurrence", "count", "total_count", "record_count"), - numeric=True, - ) - if occurrence_column and self._query_mentions_column( - query, occurrence_column - ): - metric_ref = f"{table_ref}.{self._quote_sql_identifier(occurrence_column)}" - return ( - f"SELECT TOP {limit} {dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_column)}, " - f"{metric_ref} AS {self._quote_sql_identifier(occurrence_column)} " - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"ORDER BY {metric_ref} DESC" - ) - - return ( - f"SELECT TOP {limit} {dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_column)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - return None - - def _build_explicit_table_preview_sql( - self, query: str, table_ddls: list[str] - ) -> tuple[str, str] | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip()) - if not normalized_query: - return None - - if not re.search( - r"\b(?:first|top|sample|preview|show|list)\b", - normalized_query, - flags=re.IGNORECASE, - ): - return None - if not re.search( - r"\b(?:rows?|records?|data)\b", normalized_query, flags=re.IGNORECASE - ): - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - normalized_query_key = re.sub(r"[^a-z0-9]", "", normalized_query.lower()) - normalized_query_keys = self._schema_identifier_alias_keys(normalized_query) - scored_tables: list[tuple[int, str]] = [] - for table in tables: - table_name = table.get("name") - if not table_name: - continue - table_name = str(table_name) - normalized_table = re.sub(r"[^a-z0-9]", "", table_name.lower()) - if not normalized_table: - continue - table_keys = self._schema_identifier_alias_keys(table_name) - if normalized_table in normalized_query_key or any( - table_key in query_key - for table_key in table_keys - for query_key in normalized_query_keys - ): - scored_tables.append((100 + len(normalized_table), table_name)) - continue - - table_without_schema = re.split(r"[.$]", table_name)[-1] - normalized_short_name = re.sub( - r"[^a-z0-9]", "", table_without_schema.lower() - ) - if normalized_short_name and normalized_short_name in normalized_query_key: - scored_tables.append((80 + len(normalized_short_name), table_name)) - - if not scored_tables: - return None - - _, table_name = sorted(scored_tables, reverse=True)[0] - limit = self._extract_requested_top_n(query, default_value=10) - return ( - f"SELECT TOP {limit} * FROM {self._quote_sql_identifier(table_name)}", - table_name, - ) - - def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: - table_names: list[str] = [] - for match in re.finditer( - r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - for match in re.finditer( - r"\bin\s+(?:the\s+)?([A-Za-z_][A-Za-z0-9_.$]*)\s+table\b", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - for match in re.finditer( - r"\bin\s+([A-Za-z_][A-Za-z0-9_.$]*)", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and ("." in table_name or "_" in table_name) - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - for match in re.finditer( - r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and ("." in table_name or "_" in table_name) - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - if re.search( - r"\b(?:repair\s+logs?|repair\s+tickets?|board\s+models?)\b", - query or "", - flags=re.IGNORECASE, - ): - for table_name in ("repair_logs", "dbo_repair_logs"): - if table_name not in table_names: - table_names.append(table_name) - if re.search(r"\bticket\s+labels?\b", query or "", flags=re.IGNORECASE): - for table_name in ("ticket_labels", "dbo_ticket_labels"): - if table_name not in table_names: - table_names.append(table_name) - return table_names - - def _explicit_table_name_candidates(self, table_name: str) -> list[str]: - table_name = str(table_name or "").strip().strip(".,;:()[]{}") - if not table_name: - return [] - - candidates = [table_name] - separator_normalized = re.sub(r"[.$]", "_", table_name) - if separator_normalized not in candidates: - candidates.append(separator_normalized) - if "_" in table_name and "." not in table_name: - dotted = table_name.replace("_", ".", 1) - if dotted not in candidates: - candidates.append(dotted) - short_name = re.split(r"[.$]", table_name)[-1] - if "." not in table_name and "_" in table_name: - short_name = table_name.split("_", 1)[-1] - if short_name and short_name not in candidates: - candidates.append(short_name) - return candidates - - def _normalize_explicit_table_names( - self, table_names: Optional[list[str]] - ) -> list[str]: - normalized: list[str] = [] - for table_name in table_names or []: - for candidate in self._explicit_table_name_candidates(table_name): - if candidate not in normalized: - normalized.append(candidate) - return normalized - - def _build_direct_orders_sales_sql(self, query: str) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - is_sales_or_orders_query = any( - term in normalized - for term in ( - "sales", - "sale", - "order", - "orders", - "new order", - "new orders", - "market", - "salesperson", - "sales person", - ) - ) - if not is_sales_or_orders_query: - return None - - table_ref = '"dbo_tblSales"' - limit = self._extract_requested_top_n(query, default_value=10) - - if ( - ("salesperson" in normalized or "sales person" in normalized) - and ("order count" in normalized or "orders" in normalized or "count" in normalized) - ): - return ( - f'SELECT TOP {limit} {table_ref}."SalesPerson" AS "SalesPerson", ' - f'COUNT(*) AS "OrderCount" ' - f"FROM {table_ref} " - f'WHERE {table_ref}."SalesPerson" IS NOT NULL ' - f'GROUP BY {table_ref}."SalesPerson" ' - f"ORDER BY COUNT(*) DESC" - ) - - if "top" in normalized and "new order" in normalized: - date_filter = "" - if re.search(r"\b2026[\s-]*q1\b", normalized): - date_filter = ( - f'WHERE {table_ref}."OrdDate" >= \'2026-01-01 00:00:00\' ' - f'AND {table_ref}."OrdDate" < \'2026-04-01 00:00:00\' ' - ) - return ( - f'SELECT TOP {limit} {table_ref}."BU" AS "BU", ' - f'{table_ref}."Market" AS "Market", ' - f'{table_ref}."Customer" AS "Customer", ' - f'{table_ref}."ProdName" AS "ProdName", ' - f'{table_ref}."SalesValue" AS "SalesValue" ' - f"FROM {table_ref} " - f"{date_filter}" - f'ORDER BY {table_ref}."SalesValue" DESC' - ) - - if ( - "market" in normalized - and "growth" in normalized - and any(term in normalized for term in ("last year", "previous year")) - ): - return ( - f'SELECT {table_ref}."Market" AS "Market", ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "CurrentPeriodSales", ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "PreviousPeriodSales", ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) - ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "SalesGrowth" ' - f"FROM {table_ref} " - f'WHERE {table_ref}."Market" IS NOT NULL ' - f'GROUP BY {table_ref}."Market" ' - f'ORDER BY "SalesGrowth" DESC' - ) - - if ( - "distribution" in normalized - and "sales" in normalized - and ("market" in normalized or "by market" in normalized) - ): - return ( - f'SELECT {table_ref}."Market" AS "Market", ' - f'SUM({table_ref}."SalesValue") AS "TotalSalesValue" ' - f"FROM {table_ref} " - f'WHERE {table_ref}."Market" IS NOT NULL ' - f'GROUP BY {table_ref}."Market" ' - f'ORDER BY SUM({table_ref}."SalesValue") DESC' - ) - - return None - - def _extract_explicit_table_column_reference( - self, query: str - ) -> tuple[str, str] | None: - normalized_query = query or "" - reference_match = re.search( - r"\b(?P[A-Za-z_][A-Za-z0-9_]*)[._]" - r"(?P
[A-Za-z_][A-Za-z0-9_]*)[._]" - r"(?P[A-Za-z_][A-Za-z0-9_]*)\b", - normalized_query, - ) - if not reference_match: - return None - - schema = reference_match.group("schema") - table = reference_match.group("table") - column = reference_match.group("column") - table_name = f"{schema}_{table}" - return table_name, column - - def _build_explicit_group_count_sql(self, query: str) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - if not any( - term in normalized_query - for term in ( - "group by", - "grouped by", - "by ", - "pie chart", - "donut chart", - "bar chart", - "count", - "counts", - ) - ): - return None - - explicit_reference = self._extract_explicit_table_column_reference(query) - if not explicit_reference: - return None - - table_name, column = explicit_reference - table_ref = self._quote_sql_identifier(table_name) - column_ref = f"{table_ref}.{self._quote_sql_identifier(column)}" - return ( - f"SELECT {column_ref} AS {self._quote_sql_identifier(column)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"GROUP BY {column_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - def _build_date_filter(self, table_name: str, date_column: str, query: str) -> str: - date_ref = ( - f"{self._quote_sql_identifier(table_name)}." - f"{self._quote_sql_identifier(date_column)}" - ) - normalized_query = (query or "").lower() - if "this year" in normalized_query or "current year" in normalized_query: - return ( - f" WHERE {date_ref} >= '2026-01-01 00:00:00' " - f"AND {date_ref} < '2027-01-01 00:00:00'" - ) - year_match = re.search(r"\b(20\d{2})\b", normalized_query) - if year_match: - year = int(year_match.group(1)) - return ( - f" WHERE {date_ref} >= '{year}-01-01 00:00:00' " - f"AND {date_ref} < '{year + 1}-01-01 00:00:00'" - ) - return "" - - def _append_not_null_filters( - self, where_clause: str, column_refs: list[str] - ) -> str: - conditions = [f"{column_ref} IS NOT NULL" for column_ref in column_refs] - if not conditions: - return where_clause - - if where_clause.strip(): - return f"{where_clause.rstrip()} AND {' AND '.join(conditions)}" - return f" WHERE {' AND '.join(conditions)}" - - def _build_schema_literal_filter_conditions( - self, - query: str, - table: dict[str, Any], - table_ref: str, - ) -> list[str]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - conditions: list[str] = [] - - if "backlog" in normalized_query: - filter_column = self._find_schema_column( - table, - ( - "Category", - "OrderCategory", - "Order Category", - "Status", - "OrderStatus", - "Order Status", - "Stage", - "OrderStage", - "Order Stage", - ), - ) - if filter_column: - filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" - conditions.append(f"{filter_ref} = 'Backlog'") - - return conditions - - def _select_best_analytics_table( - self, - tables: list[dict[str, Any]], - required_dimensions: list[tuple[str, ...]], - measure_candidates: tuple[str, ...], - wants_date: bool = False, - allow_count_metric: bool = False, - query: str = "", - ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - scored: list[ - tuple[int, dict[str, Any], list[str], str | None, str | None] - ] = [] - for table in tables: - dimensions = [ - self._find_schema_column(table, candidates) - for candidates in required_dimensions - ] - if any(dimension is None for dimension in dimensions): - continue - - measure = self._find_schema_column( - table, measure_candidates, numeric=True - ) - if not measure and not allow_count_metric: - continue - date_column = self._find_schema_column( - table, - ( - "OrdDate", - "InvDate", - "OrderDate", - "NewOrderDate", - "Date", - "CreatedAt", - "created_at", - ), - temporal=True, - ) - if wants_date and not date_column: - continue - - score = 10 * len([dimension for dimension in dimensions if dimension]) - if measure: - score += 8 - elif allow_count_metric: - score += 2 - if date_column: - score += 4 - table_name = str(table.get("name") or "").lower() - if not table_name: - continue - if "sales" in table_name: - score += 5 - if "tblsales" in self._normalize_schema_token(table_name): - score += 25 - if "order" in table_name: - score += 4 - if "invoice" in table_name or "inv" in table_name: - score += 3 - if "stage" in table_name: - score -= 8 - if any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ): - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber"), - ) - if "order" in table_name: - score += 30 - if "neworder" in self._normalize_schema_token(table_name): - score += 15 - if order_column: - score += 12 - if "margin" in table_name and "margin" not in normalized_query: - score -= 12 - if "customer" in normalized_query: - if "customer" in table_name or "account" in table_name: - score += 16 - if self._find_schema_column( - table, - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - ), - ): - score += 10 - if any( - term in normalized_query for term in ("product", "products", "item") - ): - if "product" in table_name or "item" in table_name: - score += 16 - if any( - term in normalized_query - for term in ("sales", "revenue", "value", "amount") - ): - if "sales" in table_name: - score += 12 - if "invoice" in normalized_query and ( - "invoice" in table_name or "inv" in table_name - ): - score += 20 - - scored.append((score, table, dimensions, measure, date_column)) - - if not scored: - return None - - _, table, dimensions, measure, date_column = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - return ( - table, - [dimension for dimension in dimensions if dimension], - measure, - date_column, - ) - - def _build_schema_grounded_analytics_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - - if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): - return pcb_direct_sql - - if repair_failure_count_sql := self._build_repair_failure_count_sql( - query, table_ddls - ): - return repair_failure_count_sql - - if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( - query, table_ddls - ): - return monthly_repair_volume_sql - - is_sales_or_order_query = any( - term in normalized_query - for term in ( - "average order value", - "invoice", - "new order", - "order", - "orders", - "currency", - "currencies", - "market", - "markets", - "performance", - "product", - "products", - "quantity", - "qty", - "revenue", - "sale", - "sales", - "salesperson", - "sales person", - "sold", - ) - ) - if not is_sales_or_order_query: - if operational_sql := self._build_schema_grounded_operational_sql( - query, tables - ): - return operational_sql - - if conversion_sql := self._build_order_invoice_conversion_sql( - query, tables - ): - return conversion_sql - - if yoy_sql := self._build_yoy_sales_change_sql(query, tables): - return yoy_sql - - if contribution_sql := self._build_contribution_sql(query, tables): - return contribution_sql - - if not is_sales_or_order_query: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not any( - term in normalized_query - for term in ( - "amount", - "cost", - "expense", - "quantity", - "qty", - "revenue", - "sale", - "sales", - "sold", - "sum", - "total", - "value", - ) - ) - wants_average_metric = any( - term in normalized_query for term in ("average", "avg", "mean") - ) - - wants_monthly_count = ( - "monthly" in normalized_query - and any(term in normalized_query for term in ("count", "volume")) - and any(term in normalized_query for term in ("order", "orders")) - ) - if wants_monthly_count: - date_candidates = ( - ("InvDate", "InvoiceDate", "Invoice Date") - if "invdate" in compact_query or "invoice" in normalized_query - else ( - "OrdDate", - "OrderDate", - "NewOrderDate", - "InvDate", - "InvoiceDate", - "Date", - ) - ) - scored_tables: list[tuple[int, dict[str, Any], str]] = [] - for table in tables: - date_column = self._find_schema_column( - table, date_candidates, temporal=True - ) - if not date_column: - continue - table_name = str(table.get("name") or "") - score = 10 - if "sales" in table_name.lower() or "order" in table_name.lower(): - score += 5 - if self._find_schema_column( - table, ("OrdNo", "OrderNo", "OrderId", "InvoiceNo") - ): - score += 3 - scored_tables.append((score, table, date_column)) - - if scored_tables: - _, table, date_column = sorted( - scored_tables, key=lambda item: item[0], reverse=True - )[0] - table_name = table.get("name") - if table_name and date_column: - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - date_ref = ( - f"{table_ref}.{self._quote_sql_identifier(date_column)}" - ) - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "OrderCount" ' - f"FROM {table_ref}" - f"{self._build_date_filter(table_name, date_column, query)} " - f"GROUP BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref})" - ) - - dimension_candidates: list[tuple[str, ...]] = [] - if "salesperson" in normalized_query or "sales person" in normalized_query: - dimension_candidates.append( - ("SalesPerson", "Salesman", "Sales Rep", "SalesRep", "Rep", "Owner") - ) - if "business unit" in normalized_query or re.search(r"\bbu\b", normalized_query): - dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) - if "market" in normalized_query: - dimension_candidates.append(("Market", "MarketType", "MarketName", "Region", "Country")) - if "region" in normalized_query: - dimension_candidates.append(("Region", "Market", "Area", "Territory")) - if "currency" in normalized_query or "currencies" in normalized_query: - dimension_candidates.append( - ( - "Currency", - "CurrencyCode", - "Currency Code", - "Curr", - "CurrCode", - "MoneyCurrency", - "PaymentCurrency", - "FXCurrency", - ) - ) - if "country" in normalized_query or "countries" in normalized_query: - dimension_candidates.append(("Country", "CountryName", "Nation", "Market")) - if "division" in normalized_query: - dimension_candidates.append(("Division",)) - if ( - ( - "category" in normalized_query - or "categories" in normalized_query - or "prodcategory" in compact_query - or "productcategory" in compact_query - ) - and "product" in normalized_query - and "product type" not in normalized_query - and "prodtype" not in compact_query - and "producttype" not in compact_query - ): - dimension_candidates.append( - ( - "ProductCategory", - "Product Category", - "ProdCategory", - "Category", - "ProductType", - "Product Type", - "ProdType", - "ProdName", - "Product", - "ProductName", - ) - ) - elif ( - "product type" in normalized_query - or "prodtype" in normalized_query - or "producttype" in compact_query - or "prodtype" in compact_query - ): - dimension_candidates.append(("ProdType", "ProductType", "Product Type")) - elif "product" in normalized_query: - dimension_candidates.append( - ( - "ProdName", - "Product", - "ProductName", - "ProductDescription", - "Item", - "ItemName", - "ProdCode", - "ProductCode", - "PartNo", - "SKU", - ) - ) - if ( - "customer" in normalized_query - or "custname" in compact_query - or "custno" in compact_query - ): - dimension_candidates.append( - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - "Client", - "ClientName", - ) - ) - - measure_candidates = ( - "Qty", - "Quantity", - "QtySold", - "SoldQty", - "QuantitySold", - "UnitsSold", - "ItemQty", - "SalesQty", - "OrderQty", - "OrderQuantity", - "InvoiceQty", - "InvoiceQuantity", - ) if any(term in normalized_query for term in ("quantity", "qty")) else ( - "Sales", - "SalesValue", - "FXSalesValue", - "Revenue", - "NetSales", - "TotalSales", - "SalesAmount", - "SaleAmount", - "NewOrderValue", - "NewOrdersValue", - "InvoiceValue", - "InvoiceAmount", - "InvoiceAmt", - "OrderValue", - "TotalRevenue", - "Amount", - "Value", - "TotalOrderValue", - "Cost", - ) - if "invoice" in normalized_query: - measure_candidates = ( - "InvoiceValue", - "InvoiceAmount", - "InvoiceAmt", - "InvValue", - "InvAmount", - "SalesValue", - "FXSalesValue", - "Value", - "Amount", - ) - if wants_count_metric: - measure_candidates = () - wants_trend = ( - "trend" in normalized_query - or "line chart" in normalized_query - or "over time" in normalized_query - or "last 12 months" in normalized_query - or "by month" in normalized_query - or "monthly" in normalized_query - ) - wants_date_distribution = ( - any( - term in normalized_query - for term in ("distribution", "breakdown", "split") - ) - and any( - term in normalized_query - for term in ("date", "dates", "orddate", "order date", "order dates") - ) - ) - wants_order_count_metric = ( - any(term in normalized_query for term in ("order", "orders", "new order", "new orders")) - and any( - term in normalized_query - for term in ( - "count", - "counts", - "volume", - "how many", - "number of", - "monthly", - "over time", - "last 12 months", - "each", - "per ", - ) - ) - and not any( - term in normalized_query - for term in ("value", "amount", "revenue", "sales", "cost", "margin") - ) - ) - wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) - wants_top = wants_top or any( - term in normalized_query - for term in ( - "top ", - "best ", - "ranking", - "performance ranking", - ) - ) - wants_time_bucket = wants_trend or bool( - re.search(r"\bby\s+(?:month|year|quarter|date)\b", normalized_query) - ) - wants_detail_rows = ( - wants_top - and ("new order" in normalized_query or "orders" in normalized_query) - and any(term in normalized_query for term in ("including", "include")) - ) - mentions_date_column = any( - column_name in compact_query - for column_name in ( - "orddate", - "invdate", - "orderdate", - "invoicedate", - "createdat", - ) - ) - wants_date = ( - wants_trend - or wants_date_distribution - or mentions_date_column - or "this year" in normalized_query - or bool(re.search(r"\b20\d{2}\b", normalized_query)) - ) - wants_unique_customers_by_group = ( - any( - term in normalized_query - for term in ("unique customer", "unique customers") - ) - and "customer" in normalized_query - and "division" in normalized_query - and "market" in normalized_query - and any(term in normalized_query for term in ("highest", "top", "most")) - and any(term in normalized_query for term in ("each", "per ")) - ) - if wants_unique_customers_by_group: - selected = self._select_best_analytics_table( - tables, - [ - ("Market", "MarketType", "MarketName", "Region", "Country"), - ("Division",), - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - "Client", - "ClientName", - ), - ], - (), - wants_date=False, - allow_count_metric=True, - query=query, - ) - if selected: - table, dimensions, _measure, _date_column = selected - table_name = table.get("name") - if table_name and len(dimensions) >= 3: - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - market, division, customer = dimensions[:3] - market_ref = f"{table_ref}.{self._quote_sql_identifier(market)}" - division_ref = ( - f"{table_ref}.{self._quote_sql_identifier(division)}" - ) - customer_ref = ( - f"{table_ref}.{self._quote_sql_identifier(customer)}" - ) - where_clause = self._append_not_null_filters( - "", - [market_ref, division_ref, customer_ref], - ) - return ( - "WITH grouped_results AS (" - f"SELECT {market_ref} AS {self._quote_sql_identifier(market)}, " - f"{division_ref} AS {self._quote_sql_identifier(division)}, " - f"COUNT(DISTINCT {customer_ref}) AS \"UniqueCustomerCount\" " - f"FROM {table_ref}" - f"{where_clause} " - f"GROUP BY {market_ref}, {division_ref}" - "), ranked_results AS (" - f"SELECT {self._quote_sql_identifier(market)}, " - f"{self._quote_sql_identifier(division)}, " - "\"UniqueCustomerCount\", " - f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(market)} " - "ORDER BY \"UniqueCustomerCount\" DESC) AS \"rank\" " - "FROM grouped_results" - ") " - f"SELECT {self._quote_sql_identifier(market)}, " - f"{self._quote_sql_identifier(division)}, " - "\"UniqueCustomerCount\" " - "FROM ranked_results " - "WHERE \"rank\" = 1 " - "ORDER BY \"UniqueCustomerCount\" DESC" - ) - - if not dimension_candidates and wants_time_bucket: - selected = self._select_best_analytics_table( - tables, - [], - measure_candidates, - wants_date=True, - allow_count_metric=wants_count_metric, - query=query, - ) - if not selected: - return None - - table, _dimensions, measure, date_column = selected - table_name = table.get("name") - if not (table_name and date_column): - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - if wants_count_metric or not measure: - metric_expr = "COUNT(*)" - metric_alias = "RecordCount" - elif wants_average_metric: - metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Average{measure}" - else: - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Total{measure}" - - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " - f"FROM {table_ref}" - f"{self._build_date_filter(table_name, date_column, query)} " - f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" - ) - - if not dimension_candidates: - return None - - allow_count_metric = ( - wants_order_count_metric - or wants_count_metric - or ("performance" in normalized_query and wants_time_bucket) - ) - selected = self._select_best_analytics_table( - tables, - dimension_candidates, - measure_candidates, - wants_date=wants_date, - allow_count_metric=allow_count_metric, - query=query, - ) - if not selected: - return None - - table, dimensions, measure, date_column = selected - table_name = table.get("name") - if not table_name: - return None - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - dimension_refs = [ - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - for dimension in dimensions - ] - - if wants_detail_rows: - if not measure: - return None - metric_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" - limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) - limit = int(limit_match.group(1)) if limit_match else 20 - select_parts = [ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ] - select_parts.append( - f"{metric_ref} AS {self._quote_sql_identifier(measure)}" - ) - date_filter = ( - self._build_date_filter(table_name, date_column, query) - if date_column - else "" - ) - return ( - f"SELECT TOP {limit} {', '.join(select_parts)} " - f"FROM {table_ref}" - f"{self._append_not_null_filters(date_filter, dimension_refs)} " - f"ORDER BY {metric_ref} DESC" - ) - - if wants_date_distribution and date_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - select_parts = [ - f"DATEPART(YEAR, {date_ref}) AS \"year\"", - f"DATEPART(MONTH, {date_ref}) AS \"month\"", - *[ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ], - 'COUNT(*) AS "OrderCount"', - ] - group_parts = [ - f"DATEPART(YEAR, {date_ref})", - f"DATEPART(MONTH, {date_ref})", - *dimension_refs, - ] - where_clause = self._append_not_null_filters( - self._build_date_filter(table_name, date_column, query), - [date_ref, *dimension_refs], - ) - extra_conditions = self._build_schema_literal_filter_conditions( - query, - table, - table_ref, - ) - if extra_conditions: - where_clause = ( - f"{where_clause.rstrip()} AND {' AND '.join(extra_conditions)}" - if where_clause.strip() - else f" WHERE {' AND '.join(extra_conditions)}" - ) - return ( - f"SELECT {', '.join(select_parts)} FROM {table_ref}" - f"{where_clause} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}), COUNT(*) DESC" - ) - - wants_top_per_group = ( - len(dimensions) >= 2 - and any(term in normalized_query for term in ("highest", "top", "most")) - and any(term in normalized_query for term in ("each", "per ")) - ) - if wants_top_per_group: - partition_dimension = None - rank_dimension = None - if "market" in normalized_query: - partition_dimension = self._find_schema_column( - table, ("Market", "MarketType", "Region") - ) - if "region" in normalized_query and not partition_dimension: - partition_dimension = self._find_schema_column( - table, ("Region", "Market", "Area", "Territory") - ) - if "customer" in normalized_query: - rank_dimension = self._find_schema_column( - table, ("Customer", "CustName", "CustNo") - ) - if not partition_dimension: - partition_dimension = dimensions[0] - if not rank_dimension: - rank_dimension = next( - ( - dimension - for dimension in dimensions - if dimension != partition_dimension - ), - None, - ) - - if partition_dimension and rank_dimension: - partition_ref = ( - f"{table_ref}.{self._quote_sql_identifier(partition_dimension)}" - ) - rank_ref = f"{table_ref}.{self._quote_sql_identifier(rank_dimension)}" - if wants_order_count_metric: - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), - ) - metric_expr = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" - if order_column - else "COUNT(*)" - ) - metric_alias = "OrderCount" - else: - metric_expr = ( - f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - if measure - else "COUNT(*)" - ) - metric_alias = f"Total{measure}" if measure else "RecordCount" - where_clause = self._append_not_null_filters( - ( - self._build_date_filter(table_name, date_column, query) - if date_column - else "" - ), - [partition_ref, rank_ref], - ) - return ( - "WITH grouped_results AS (" - f"SELECT {partition_ref} AS {self._quote_sql_identifier(partition_dimension)}, " - f"{rank_ref} AS {self._quote_sql_identifier(rank_dimension)}, " - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " - f"FROM {table_ref}" - f"{where_clause} " - f"GROUP BY {partition_ref}, {rank_ref}" - "), ranked_results AS (" - f"SELECT {self._quote_sql_identifier(partition_dimension)}, " - f"{self._quote_sql_identifier(rank_dimension)}, " - f"{self._quote_sql_identifier(metric_alias)}, " - f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(partition_dimension)} " - f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC) AS \"rank\" " - "FROM grouped_results" - ") " - f"SELECT {self._quote_sql_identifier(partition_dimension)}, " - f"{self._quote_sql_identifier(rank_dimension)}, " - f"{self._quote_sql_identifier(metric_alias)} " - "FROM ranked_results " - "WHERE \"rank\" = 1 " - f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC" - ) - - if wants_trend and date_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - if wants_order_count_metric or wants_count_metric: - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), - ) - metric_expr = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" - if order_column - else "COUNT(*)" - ) - metric_alias = "OrderCount" - elif wants_average_metric: - metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Average{measure}" - else: - metric_expr = ( - f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - if measure - else "COUNT(*)" - ) - metric_alias = f"Total{measure}" if measure else "RecordCount" - select_parts = [ - f"DATEPART(YEAR, {date_ref}) AS \"year\"", - f"DATEPART(MONTH, {date_ref}) AS \"month\"", - *[ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ], - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}", - ] - group_parts = [ - f"DATEPART(YEAR, {date_ref})", - f"DATEPART(MONTH, {date_ref})", - *dimension_refs, - ] - return ( - f"SELECT {', '.join(select_parts)} FROM {table_ref}" - f"{self._append_not_null_filters(self._build_date_filter(table_name, date_column, query), dimension_refs)} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref})" - ) - - if wants_order_count_metric or wants_count_metric or not measure: - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), - ) - metric_expr = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" - if order_column and (wants_order_count_metric or wants_count_metric) - else "COUNT(*)" - ) - metric_alias = "OrderCount" if order_column else "RecordCount" - elif wants_average_metric: - metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Average{measure}" - else: - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Total{measure}" - limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) - limit = int(limit_match.group(1)) if limit_match else 10 - top_clause = f"TOP {limit} " if wants_top else "" - sort_direction = ( - "ASC" - if any( - term in normalized_query - for term in ( - "losing", - "lowest", - "least", - "bottom", - "declining", - "underperforming", - "smallest", - ) - ) - else "DESC" - ) - date_filter = ( - self._build_date_filter(table_name, date_column, query) - if date_column - else "" - ) - select_parts = [ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ] - select_parts.append( - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}" - ) - return ( - f"SELECT {top_clause}{', '.join(select_parts)} " - f"FROM {table_ref}" - f"{self._append_not_null_filters(date_filter, dimension_refs)} " - f"GROUP BY {', '.join(dimension_refs)} " - f"ORDER BY {metric_expr} {sort_direction}" - ) - - def _build_contribution_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not any(term in normalized_query for term in ("contribution", "pie chart")): - return None - - compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - dimension_candidates: tuple[str, ...] | None = None - if ( - "product type" in normalized_query - or "prodtype" in normalized_query - or "producttype" in compact_query - or "prodtype" in compact_query - ): - dimension_candidates = ("ProdType", "ProductType", "Product Type") - elif "market" in normalized_query: - dimension_candidates = ("Market", "MarketType") - elif "division" in normalized_query: - dimension_candidates = ("Division",) - elif "customer" in normalized_query: - dimension_candidates = ("Customer", "CustName", "CustNo") - - if not dimension_candidates: - return None - - selected = self._select_best_analytics_table( - tables, - [dimension_candidates], - ( - "SalesValue", - "FXSalesValue", - "OrderValue", - "NewOrderValue", - "Revenue", - "Value", - "Amount", - ), - wants_date=False, - query=query, - ) - if not selected: - return None - - table, dimensions, measure, _date_column = selected - table_name = table.get("name") - if not (table_name and dimensions and measure): - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - dimension = dimensions[0] - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - return ( - f"SELECT {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " - f"{metric_expr} AS \"Total{measure}\" " - f"FROM {table_ref} " - f"GROUP BY {dimension_ref} " - f"ORDER BY {metric_expr} DESC" - ) - - def _build_generic_categorical_count_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - if re.search( - r"\b(?:first|top|sample|preview|show|list)\b.*\b(?:rows?|records?|data)\b", - normalized_query, - ): - return None - - wants_categorical_summary = any( - term in normalized_query - for term in ( - "bar chart", - "by ", - "chart", - "count", - "distribution", - "donut chart", - "frequency", - "group by", - "grouped by", - "most often", - "often", - "pie chart", - "restored", - "status", - "type", - "category", - ) - ) - if not wants_categorical_summary: - return None - - query_key = self._normalize_schema_token(query) - query_terms = self._query_schema_terms(query) - scored: list[tuple[int, dict[str, Any], str]] = [] - low_value_column_patterns = ( - "id", - "no", - "number", - "date", - "time", - "description", - "comment", - "note", - "remark", - ) - - for table in tables: - table_name = str(table.get("name") or "") - normalized_table = self._normalize_schema_token(table_name) - normalized_short_table = self._normalize_schema_token( - re.split(r"[.$]", table_name)[-1] - ) - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if not column_name or not self._is_text_schema_type(column_type): - continue - normalized_column = self._normalize_schema_token(column_name) - if not normalized_column: - continue - - score = 0 - if normalized_table and normalized_table in query_key: - score += 120 - if normalized_short_table and normalized_short_table in query_key: - score += 100 - if normalized_column and normalized_column in query_key: - score += 180 - for term in query_terms: - if term == normalized_column: - score += 100 - elif term in normalized_column or normalized_column in term: - score += 45 - if term == normalized_table or term == normalized_short_table: - score += 40 - elif term in normalized_table or term in normalized_short_table: - score += 20 - if "status" in normalized_query and "status" in normalized_column: - score += 90 - if "category" in normalized_query and "category" in normalized_column: - score += 80 - if "type" in normalized_query and "type" in normalized_column: - score += 70 - if ( - "destination" in normalized_query - and "destination" in normalized_column - and ( - "database" in normalized_query - or "databases" in normalized_query - ) - and ( - "name" in normalized_column - or "phys" in normalized_column - or "db" in normalized_column - ) - ): - score += 140 - if any(pattern == normalized_column for pattern in low_value_column_patterns): - score -= 100 - elif any( - normalized_column.endswith(pattern) - for pattern in low_value_column_patterns - ): - score -= 35 - - if score > 0: - scored.append((score, table, column_name)) - - if not scored: - return None - - _, table, dimension = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - table_name = table.get("name") - if not table_name: - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - top_n = self._extract_requested_top_n(query, default_value=0) - top_clause = f"TOP {top_n} " if top_n else "" - return ( - f"SELECT {top_clause}{dimension_ref} AS {self._quote_sql_identifier(dimension)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - def _build_order_invoice_conversion_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not ( - "conversion" in normalized_query - and "order" in normalized_query - and "invoice" in normalized_query - ): - return None - - scored: list[tuple[int, dict[str, Any], str, str, str | None]] = [] - for table in tables: - order_column = self._find_schema_column( - table, ("OrdNo", "OrderNo", "OrderNumber", "NewOrderNo") - ) - invoice_column = self._find_schema_column( - table, ("InvoiceNo", "InvNo", "InvoiceNumber") - ) - date_column = self._find_schema_column( - table, - ("OrdDate", "InvDate", "OrderDate", "InvoiceDate", "Date"), - temporal=True, - ) - if not (order_column and invoice_column): - continue - - score = 20 - if date_column: - score += 5 - if "sales" in str(table.get("name") or "").lower(): - score += 5 - scored.append((score, table, order_column, invoice_column, date_column)) - - if not scored: - return None - - _, table, order_column, invoice_column, date_column = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - table_name = table.get("name") - if not table_name: - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - order_ref = f"{table_ref}.{self._quote_sql_identifier(order_column)}" - invoice_ref = f"{table_ref}.{self._quote_sql_identifier(invoice_column)}" - date_column = date_column or order_column - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f"COUNT(DISTINCT {order_ref}) AS \"OrderCount\", " - f"COUNT(DISTINCT {invoice_ref}) AS \"InvoiceCount\", " - f"(COUNT(DISTINCT {invoice_ref}) * 100.0 / " - f"NULLIF(COUNT(DISTINCT {order_ref}), 0)) AS \"ConversionRate\" " - f"FROM {table_ref} " - f"WHERE {order_ref} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" - ) - - def _build_yoy_sales_change_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not any(term in normalized_query for term in ("yoy", "year over year")): - return None - - required_dimensions: list[tuple[str, ...]] = [] - if "customer" in normalized_query: - required_dimensions.append(("Customer", "CustName", "CustNo")) - if "product" in normalized_query: - required_dimensions.append( - ("ProdName", "Product", "ProductName", "Item", "ProdCode") - ) - if "market" in normalized_query: - required_dimensions.append(("Market", "MarketType")) - - if not required_dimensions: - return None - - selected = self._select_best_analytics_table( - tables, - required_dimensions, - ( - "SalesValue", - "FXSalesValue", - "OrderValue", - "NewOrderValue", - "Revenue", - "Value", - "Amount", - ), - wants_date=False, - query=query, - ) - if not selected: - return None - - table, dimensions, measure, date_column = selected - table_name = table.get("name") - if not (table_name and measure): - return None - - year_column = self._find_schema_column( - table, ("YearInd", "Year", "OrderYear", "InvoiceYear"), numeric=True - ) - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - if year_column: - year_expr = f"{table_ref}.{self._quote_sql_identifier(year_column)}" - elif date_column: - year_expr = ( - f"DATEPART(YEAR, " - f"{table_ref}.{self._quote_sql_identifier(date_column)})" - ) - else: - return None - - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - dimension_refs = [ - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - for dimension in dimensions - ] - select_parts = [ - f"{year_expr} AS \"year\"", - *[ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ], - f"{metric_expr} AS \"Total{measure}\"", - ] - group_parts = [year_expr, *dimension_refs] - return ( - f"SELECT {', '.join(select_parts)} " - f"FROM {table_ref} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY {year_expr}, {metric_expr} DESC" - ) - - def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: - if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): - return max(1, min(int(match.group(1)), 100)) - if match := re.search( - r"\b(?:first|limit)\s+(\d+)\b", query or "", flags=re.IGNORECASE - ): - return max(1, min(int(match.group(1)), 100)) - if match := re.search(r"\b(\d+)\s+rows?\b", query or "", flags=re.IGNORECASE): - return max(1, min(int(match.group(1)), 100)) - return default_value - - def _build_manufacturing_throughput_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_throughput = "throughput" in normalized or ( - "repair" in normalized and "volume" in normalized - ) - wants_unit_breakdown = any( - term in normalized - for term in ( - "manufacturing unit", - "manufacturing units", - "business unit", - "business units", - "different unit", - "different units", - ) - ) - - if not (wants_throughput and wants_unit_breakdown): - return None - - tables = self._parse_schema_tables(table_ddls) - table = self._find_best_schema_table_for_query(query, tables) - if table: - unit_column = self._find_schema_column( - table, - ( - "BusinessUnit", - "business_unit", - "manufacturing_unit", - "manufacturingunit", - "unit", - "unit_name", - "BU", - "division", - ), - ) - if unit_column: - table_name = str(table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - unit_ref = f"{table_ref}.{self._quote_sql_identifier(unit_column)}" - timestamp_column = self._find_temporal_column_for_query(query, table) - - if timestamp_column and any( - term in normalized for term in ("trend", "monthly", "over time") - ): - timestamp_ref = ( - f"{table_ref}.{self._quote_sql_identifier(timestamp_column)}" - ) - return ( - f"SELECT {unit_ref} AS " - f"{self._quote_sql_identifier(unit_column)}, " - f"DATEPART(YEAR, {timestamp_ref}) AS \"year\", " - f"DATEPART(MONTH, {timestamp_ref}) AS \"month\", " - f'COUNT(*) AS "throughput" ' - f"FROM {table_ref} " - f"WHERE {unit_ref} IS NOT NULL " - f"AND {timestamp_ref} IS NOT NULL " - f"GROUP BY {unit_ref}, DATEPART(YEAR, {timestamp_ref}), " - f"DATEPART(MONTH, {timestamp_ref}) " - f"ORDER BY {unit_ref} ASC, DATEPART(YEAR, {timestamp_ref}) ASC, " - f"DATEPART(MONTH, {timestamp_ref}) ASC" - ) - - return ( - f"SELECT {unit_ref} AS " - f"{self._quote_sql_identifier(unit_column)}, " - f'COUNT(*) AS "throughput" ' - f"FROM {table_ref} " - f"WHERE {unit_ref} IS NOT NULL " - f"GROUP BY {unit_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - has_debug_entries = self._schema_contains( - table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names - ) - has_business_unit = self._schema_contains( - table_ddls, r"\bBusinessUnit\b", table_names=table_names - ) - if not (has_debug_entries and has_business_unit): - return None - - timestamp_column = None - for candidate in ("DateIn", "FailedAt"): - if self._schema_contains( - table_ddls, rf"\b{candidate}\b", table_names=table_names - ): - timestamp_column = candidate - break - - if timestamp_column and any( - term in normalized for term in ("trend", "monthly", "over time") - ): - timestamp_expression = f'"dbo_DebugEntries"."{timestamp_column}"' - return ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' - f'DATEPART(YEAR, {timestamp_expression}) AS "year", ' - f'DATEPART(MONTH, {timestamp_expression}) AS "month", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' - f'AND {timestamp_expression} IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' - f'DATEPART(YEAR, {timestamp_expression}), ' - f'DATEPART(MONTH, {timestamp_expression}) ' - 'ORDER BY "unit_name" ASC, "year" ASC, "month" ASC' - ) - - return ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries"."BusinessUnit" ' - 'ORDER BY "throughput" DESC' - ) - - def _build_audit_log_activity_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - if not ( - "audit" in normalized - and "log" in normalized - and any(term in normalized for term in ("activity", "over time", "trend")) - ): - return None - - table_name = "dbo_audit_log" - timestamp_column = "created_at" - if not self._schema_has_table_column( - table_ddls, - table_name, - timestamp_column, - table_names=table_names, - ): - return None - - dimension_column = None - condition_candidates = ( - "is_name_condition", - "name", - "action", - "entity_type", - ) - activity_candidates = ( - "action", - "entity_type", - "actor_name", - "actor_user_id", - "name", - ) - candidates = ( - condition_candidates - if "condition" in normalized - else activity_candidates - ) - for candidate in candidates: - if self._schema_has_table_column( - table_ddls, - table_name, - candidate, - table_names=table_names, - ): - dimension_column = candidate - break - - if not dimension_column: - return None - - timestamp_expression = f'"{table_name}"."{timestamp_column}"' - dimension_expression = f'"{table_name}"."{dimension_column}"' - return ( - f"SELECT DATEPART(YEAR, {timestamp_expression}) AS \"year\", " - f"DATEPART(MONTH, {timestamp_expression}) AS \"month\", " - f"{dimension_expression} AS \"{dimension_column}\", " - f'COUNT(*) AS "activity_count" ' - f'FROM "{table_name}" ' - f"WHERE {timestamp_expression} IS NOT NULL " - f"AND {dimension_expression} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {timestamp_expression}), " - f"DATEPART(MONTH, {timestamp_expression}), " - f"{dimension_expression} " - f"ORDER BY DATEPART(YEAR, {timestamp_expression}), " - f"DATEPART(MONTH, {timestamp_expression}), " - f'"activity_count" DESC' - ) - - def _build_repair_failure_count_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_failure_counts = ( - "failure" in normalized - and any( - term in normalized - for term in ( - "count", - "counts", - "category", - "code", - "grouped", - "common", - "most common", - "top", - ) - ) - and any(term in normalized for term in ("repair", "bar chart", "chart")) - ) - if not wants_failure_counts: - return None - - top_n = self._extract_requested_top_n(query) - - has_debug_fix_route = all( - ( - self._schema_has_table_column( - table_ddls, - "dbo_DebugEntries", - "DebugEntryId", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixLogs", - "DebugEntryId", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixLogs", - "FixId", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixes", - "Id", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixes", - "Description", - table_names=table_names, - ), - ) - ) - if has_debug_fix_route: - return ( - 'SELECT "dbo_DebugFixes"."Description" AS "failure_category", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_DebugEntries" ' - 'JOIN "dbo_DebugFixLogs" ' - 'ON "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" ' - 'JOIN "dbo_DebugFixes" ' - 'ON "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id" ' - 'WHERE "dbo_DebugFixes"."Description" IS NOT NULL ' - 'GROUP BY "dbo_DebugFixes"."Description" ' - 'ORDER BY "repair_count" DESC ' - f"LIMIT {top_n}" - ) - - has_debug_entries = self._schema_contains( - table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names - ) - has_failure_patterns = self._schema_contains( - table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names - ) - has_failure_sys = self._schema_contains( - table_ddls, r"\bFailureSys\b", table_names=table_names - ) - has_debug_entry_id = self._schema_contains( - table_ddls, r"\bDebugEntryId\b", table_names=table_names - ) - has_pattern_id = self._schema_contains( - table_ddls, r"\bid\b", table_names=table_names - ) - has_pattern_category = self._schema_contains( - table_ddls, r"\bcategory\b", table_names=table_names - ) - has_pattern_name = self._schema_contains( - table_ddls, r"\bname\b", table_names=table_names - ) - - if ( - has_debug_entries - and has_failure_patterns - and has_failure_sys - and has_debug_entry_id - and has_pattern_id - and (has_pattern_category or has_pattern_name) - ): - dimension_column = ( - "category" - if ("category" in normalized and has_pattern_category) - else ("name" if has_pattern_name else "category") - ) - return ( - f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' - f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' - f'FROM "dbo_DebugEntries" ' - f'JOIN "dbo_failure_patterns" ' - f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' - f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - has_repair_logs = self._schema_has_table_column( - table_ddls, - "dbo_repair_logs", - "failure_code", - table_names=table_names, - ) - if has_repair_logs: - return ( - f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_repair_logs" ' - f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' - f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - return None - - def _build_repair_sla_compliance_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_sla = "sla" in normalized and any( - term in normalized - for term in ("compliance", "dashboard", "chart", "repair", "repairs") - ) - if not wants_sla: - return None - - has_repair_status = self._schema_has_table_column( - table_ddls, - "dbo_repair_logs", - "status", - table_names=table_names, - ) - if has_repair_status: - return ( - 'SELECT "dbo_repair_logs"."status" AS "sla_status", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."status" IS NOT NULL ' - 'GROUP BY "dbo_repair_logs"."status" ' - 'ORDER BY "repair_count" DESC' - ) - - return None - - def _build_monthly_repair_volume_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_monthly_repairs = ( - "repair" in normalized - and any( - term in normalized - for term in ("monthly", "last 12 months", "trend", "volume") - ) - ) - if not wants_monthly_repairs: - return None - - tables = self._parse_schema_tables(table_ddls) - scored_tables: list[tuple[int, dict[str, Any], str]] = [] - for table in tables: - table_name = str(table.get("name") or "") - if table_names and table_name not in table_names: - continue - - date_column = self._find_schema_column( - table, - ( - "created_at", - "createdAt", - "created", - "DateIn", - "Date", - "repair_date", - "RepairDate", - "opened_at", - "started_at", - ), - temporal=True, - ) - if not date_column: - date_column = self._find_any_temporal_schema_column(table) - if not date_column: - continue - - normalized_table = self._normalize_schema_token(table_name) - score = 0 - if "repair" in normalized_table: - score += 30 - if "debugentries" in normalized_table or "debugentry" in normalized_table: - score += 25 - if "log" in normalized_table: - score += 10 - if self._find_schema_column( - table, - ("DebugEntryId", "RepairId", "repair_id", "id"), - ): - score += 5 - scored_tables.append((score, table, date_column)) - - if not scored_tables: - return None - - _score, table, date_column = sorted( - scored_tables, - key=lambda item: item[0], - reverse=True, - )[0] - table_name = str(table.get("name") or "") - if not table_name: - return None - - table_ref = self._quote_sql_identifier(table_name) - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "repair_count" ' - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " - f"DATEPART(MONTH, {date_ref}) ASC" - ) - - def _is_direct_heuristic_sql_query(self, query: str) -> bool: - return False - - def _build_heuristic_text_to_sql_fallback( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - if schema_grounded_sql := self._build_schema_grounded_sales_sql( - query, table_ddls - ): - return schema_grounded_sql - - if throughput_sql := self._build_manufacturing_throughput_sql( - query, table_ddls, table_names=table_names - ): - return throughput_sql - - if repair_failure_count_sql := self._build_repair_failure_count_sql( - query, table_ddls, table_names=table_names - ): - return repair_failure_count_sql - - if repair_sla_sql := self._build_repair_sla_compliance_sql( - query, table_ddls, table_names=table_names - ): - return repair_sla_sql - - if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( - query, table_ddls, table_names=table_names - ): - return monthly_repair_volume_sql - - wants_chart = any( - term in normalized for term in ("chart", "bar chart", "line chart", "graph") - ) - wants_failure_counts = any( - term in normalized - for term in ( - "failure", - "failure category", - "failure code", - "common pcb failures", - "common failures", - "most common", - "top 10", - "top ten", - ) - ) - wants_monthly_repairs = ( - "repair" in normalized - and any( - term in normalized - for term in ("monthly", "last 12 months", "trend", "volume") - ) - ) - - if wants_failure_counts and wants_chart: - top_n = self._extract_requested_top_n(query) - has_pattern_failure_sys = self._schema_contains( - table_ddls, r"\bFailuresys\b", table_names=table_names - ) - has_pattern_occurrences = self._schema_contains( - table_ddls, r"\boccurrences\b", table_names=table_names - ) - has_debug_entries = self._schema_contains( - table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names - ) - has_failure_patterns = self._schema_contains( - table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names - ) - has_failure_sys = self._schema_contains( - table_ddls, r"\bFailureSys\b", table_names=table_names - ) - has_debug_entry_id = self._schema_contains( - table_ddls, r"\bDebugEntryId\b", table_names=table_names - ) - has_pattern_id = self._schema_contains( - table_ddls, r"\bid\b", table_names=table_names - ) - has_pattern_category = self._schema_contains( - table_ddls, r"\bcategory\b", table_names=table_names - ) - has_pattern_name = self._schema_contains( - table_ddls, r"\bname\b", table_names=table_names - ) - - if has_failure_patterns and has_pattern_failure_sys and has_pattern_occurrences: - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'"dbo_failure_patterns"."occurrences" AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'AND "dbo_failure_patterns"."occurrences" IS NOT NULL ' - f'ORDER BY "dbo_failure_patterns"."occurrences" DESC ' - f'LIMIT {top_n}' - ) - - if has_failure_patterns and has_pattern_failure_sys: - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - if ( - has_debug_entries - and has_failure_patterns - and has_failure_sys - and has_debug_entry_id - and has_pattern_id - ): - dimension_column = ( - "category" - if ("category" in normalized and has_pattern_category) - else ("name" if has_pattern_name else "category") - ) - if dimension_column == "category" and not has_pattern_category: - dimension_column = "name" - - return ( - f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' - f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' - f'FROM "dbo_DebugEntries" ' - f'JOIN "dbo_failure_patterns" ' - f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' - f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - has_repair_logs = self._schema_contains( - table_ddls, r"\bdbo_repair_logs\b", table_names=table_names - ) - has_failure_code = self._schema_contains( - table_ddls, r"\bfailure_code\b", table_names=table_names - ) - if has_repair_logs and has_failure_code: - return ( - f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_repair_logs" ' - f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' - f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - if wants_failure_counts and wants_chart: - top_n = self._extract_requested_top_n(query) - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - return None - - def _is_schema_grounded_query( - self, query: str, db_schemas: Optional[list[str]] = None - ) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - explicit_schema_terms = ( - "table", - "column", - "schema", - "dataset", - "dbo.", - "select ", - " from ", - " join ", - " where ", - " group by ", - " order by ", - ) - if any(term in normalized for term in explicit_schema_terms): - return True - - identifier_tokens = re.findall(r"[a-zA-Z_][a-zA-Z0-9_\.]*", normalized) - if any("." in token for token in identifier_tokens): - return True - - for schema in db_schemas or []: - schema_text = schema.lower() - table_matches = re.findall( - r"create\s+table\s+([a-zA-Z0-9_\.\"]+)", schema_text - ) - column_matches = re.findall(r"\n\s*\"?([a-zA-Z_][a-zA-Z0-9_]*)\"?\s+", schema_text) - candidates = { - token.strip('"') - for token in table_matches + column_matches - if token and len(token.strip('"')) > 2 - } - if any(candidate in normalized for candidate in candidates): - return True - - return False - def _build_schema_grounded_operational_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - operational_terms = ( - "ticket", - "repair", - "failure", - "pcb", - "component", - "board", - "throughput", - "manufacturing", - "unit", - "knowledge", - "article", - "source", - "category", - "priority", - "status", - "open", - "closed", - "aging", - "workflow", - "time", - "duration", - "elapsed", - "estimated", - "volume", - "count", - ) - if not any(term in normalized_query for term in operational_terms): - return None - - scored_tables: list[tuple[int, dict[str, Any]]] = [] - for table in tables: - table_name = str(table.get("name") or "") - normalized_table = table_name.lower() - score = 0 - if any( - token in normalized_table - for token in ("ticket", "repair", "debug", "knowledge", "article") - ): - score += 10 - if "ticket" in normalized_query and "ticket" in normalized_table: - score += 8 - if "knowledge" in normalized_query and "knowledge" in normalized_table: - score += 8 - if "article" in normalized_query and "article" in normalized_table: - score += 5 - if "repair" in normalized_query and "repair" in normalized_table: - score += 5 - if "failure" in normalized_query and "failure" in normalized_table: - score += 8 - if any( - term in normalized_query - for term in ("business unit", "business units", "unit", "units") - ) and self._find_schema_column( - table, - ( - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "ManufacturingUnit", - "unit", - "BU", - "Division", - ), - ): - score += 15 - if any( - term in normalized_query - for term in ("product line", "product family", "product", "products") - ) and self._find_schema_column( - table, - ( - "Product_Family", - "ProductFamily", - "Product Family", - "ProductLine", - "Product_Line", - "Product", - "ProdType", - "Material", - ), - ): - score += 15 - if any(term in normalized_query for term in ("error", "failure")) and any( - self._find_schema_column(table, candidates) - for candidates in ( - ("failure_code", "FailureSys", "failure", "failure_type"), - ("category", "name", "description"), - ) - ): - score += 6 - if self._find_schema_column( - table, - ("created_at", "updated_at", "DateIn", "DateOut", "created", "date"), - temporal=True, - ): - score += 3 - if score: - scored_tables.append((score, table)) - - if not scored_tables: - return None - - table = sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] - table_name = str(table.get("name") or "") - if not table_name: - return None - - table_ref = self._quote_sql_identifier(table_name) - date_column = self._find_schema_column( - table, - ( - "created_at", - "created", - "DateIn", - "RepairDate", - "updated_at", - "DateOut", - "updated", - "date", - ), - temporal=True, - ) - - dimension_candidates: list[tuple[str, ...]] = [] - if any(term in normalized_query for term in ("failure", "failures", "error")): - dimension_candidates.append( - ( - "failure_code", - "FailureSys", - "failure", - "failure_type", - "failure_category", - "category", - "name", - "description", - ) - ) - if "manufacturing" in normalized_query or "unit" in normalized_query: - dimension_candidates.append( - ( - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "manufacturing unit", - "ManufacturingUnit", - "unit", - "BU", - "Division", - "assignee_user_id", - "created_by_user_id", - "org_id", - "status", - ) - ) - if any( - term in normalized_query - for term in ("product line", "product family", "product", "products") - ): - dimension_candidates.append( - ( - "Product_Family", - "ProductFamily", - "Product Family", - "ProductLine", - "Product_Line", - "Product", - "ProdType", - "Material", - ) - ) - if "component" in normalized_query: - dimension_candidates.append( - ("component", "component_type", "board_type", "title", "status") - ) - if "board" in normalized_query: - dimension_candidates.append(("board_type", "board", "title", "status")) - if "category" in normalized_query: - dimension_candidates.append(("category", "subcategory", "status", "priority")) - if "source" in normalized_query: - dimension_candidates.append(("source", "author", "category", "status")) - if "priority" in normalized_query: - dimension_candidates.append(("priority", "status")) - if ( - "status" in normalized_query - or "open" in normalized_query - or "closed" in normalized_query - ): - dimension_candidates.append(("status", "priority")) - if "assignee" in normalized_query: - dimension_candidates.append(("assignee_user_id", "created_by_user_id")) - if "workflow" in normalized_query: - dimension_candidates.append(("status", "priority", "assignee_user_id")) - - dimensions: list[str] = [] - for candidates in dimension_candidates: - dimension = self._find_schema_column(table, candidates) - if dimension and dimension not in dimensions: - dimensions.append(dimension) - - if not dimensions: - fallback_dimension = self._find_first_schema_column( - table, - ( - "status", - "priority", - "category", - "subcategory", - "author", - "assignee_user_id", - "created_by_user_id", - "org_id", - "title", - ), - ) - if fallback_dimension: - dimensions.append(fallback_dimension) - - wants_trend = any( - term in normalized_query - for term in ("trend", "monthly", "month", "line chart", "over time") - ) - wants_elapsed_time = ( - not wants_trend - and any( - term in normalized_query - for term in ( - "duration", - "elapsed", - "turnaround", - "estimated", - "time spent", - "time taken", - ) - ) - ) - wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) - limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) - limit = int(limit_match.group(1)) if limit_match else 10 - - if wants_elapsed_time: - temporal_columns = [ - str(column.get("name") or "") - for column in table.get("columns", []) - if column.get("name") - and self._is_temporal_schema_type(str(column.get("type") or "")) - ] - start_column = self._find_schema_column( - table, - ( - "created_at", - "created", - "DateIn", - "execution_date", - "opened_at", - "started_at", - "start_date", - "begin_date", - ), - temporal=True, - ) - end_column = self._find_schema_column( - table, - ( - "updated_at", - "updated", - "DateOut", - "closed_at", - "resolved_at", - "completed_at", - "finished_at", - "end_date", - ), - temporal=True, - ) - if not start_column and temporal_columns: - start_column = temporal_columns[0] - if not end_column: - for candidate in temporal_columns: - if candidate.lower() != str(start_column or "").lower(): - end_column = candidate - break - if start_column and end_column: - start_ref = f"{table_ref}.{self._quote_sql_identifier(start_column)}" - end_ref = f"{table_ref}.{self._quote_sql_identifier(end_column)}" - duration_expr = f"DATEDIFF('second', {start_ref}, {end_ref})" - if not dimensions: - fallback_dimension = self._find_first_schema_column( - table, - ( - "status", - "priority", - "assignee_user_id", - "created_by_user_id", - "org_id", - ), - ) - if fallback_dimension: - dimensions.append(fallback_dimension) - if dimensions: - dimension = dimensions[0] - dimension_ref = ( - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - ) - dimension_alias = ( - "workflow" if "workflow" in normalized_query else dimension - ) - return ( - f"SELECT {dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_alias)}, " - f'SUM({duration_expr}) AS "total_time_seconds" ' - f"FROM {table_ref} " - f"WHERE {start_ref} IS NOT NULL " - f"AND {end_ref} IS NOT NULL " - f"AND {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f'ORDER BY "total_time_seconds" DESC' - ) - return ( - f'SELECT SUM({duration_expr}) AS "total_time_seconds" ' - f"FROM {table_ref} " - f"WHERE {start_ref} IS NOT NULL " - f"AND {end_ref} IS NOT NULL" - ) - - if wants_trend and date_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - select_parts = [ - f"DATEPART(YEAR, {date_ref}) AS \"year\"", - f"DATEPART(MONTH, {date_ref}) AS \"month\"", - ] - group_parts = [ - f"DATEPART(YEAR, {date_ref})", - f"DATEPART(MONTH, {date_ref})", - ] - for dimension in dimensions[:2]: - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - select_parts.append( - f"{dimension_ref} AS {self._quote_sql_identifier(dimension)}" - ) - group_parts.append(dimension_ref) - select_parts.append('COUNT(*) AS "RecordCount"') - return ( - f"SELECT {', '.join(select_parts)} " - f"FROM {table_ref} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref})" - ) - - if dimensions: - top_clause = f"TOP {limit} " if wants_top else "" - dimension_refs = [ - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - for dimension in dimensions[:2] - ] - select_parts = [ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ] - select_parts.append('COUNT(*) AS "RecordCount"') - return ( - f"SELECT {top_clause}{', '.join(select_parts)} " - f"FROM {table_ref} " - f"GROUP BY {', '.join(dimension_refs)} " - f"ORDER BY COUNT(*) DESC" - ) - - return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' - - def _build_pcb_direct_question_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - repair_table = next( - ( - table - for table in tables - if str(table.get("name") or "").lower() == "dbo_repair_logs" - ), - None, - ) - ticket_label_table = next( - ( - table - for table in tables - if str(table.get("name") or "").lower() == "dbo_ticket_labels" - ), - None, - ) - limit = self._extract_requested_top_n(query, default_value=10) - - if ticket_label_table and "ticket" in normalized and "label" in normalized: - label_column = self._find_first_schema_column( - ticket_label_table, - ("name", "label", "title", "value", "id"), - ) - if label_column: - table_name = str(ticket_label_table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - label_ref = f"{table_ref}.{self._quote_sql_identifier(label_column)}" - return ( - f"SELECT TOP {limit} {label_ref} AS " - f"{self._quote_sql_identifier(label_column)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {label_ref} IS NOT NULL " - f"GROUP BY {label_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - if not repair_table: - return None - - table_name = str(repair_table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - board_model_column = self._find_schema_column( - repair_table, ("board_model", "boardModel", "board model", "product") - ) - failure_code_column = self._find_schema_column( - repair_table, ("failure_code", "failureCode", "failure code", "failure") - ) - created_at_column = self._find_schema_column( - repair_table, - ("created_at", "createdAt", "created", "date_received", "dateReceived"), - temporal=True, - ) - priority_column = self._find_schema_column(repair_table, ("priority",)) - status_column = self._find_schema_column(repair_table, ("status",)) - id_column = self._find_schema_column(repair_table, ("id", "repair_id")) - - asks_board_model_distribution = ( - "board model" in normalized - and any(term in normalized for term in ("distribution", "over time", "trend")) - ) - if asks_board_model_distribution and board_model_column and created_at_column: - board_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" - date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" - return ( - f"SELECT {board_ref} AS " - f"{self._quote_sql_identifier(board_model_column)}, " - f"DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {board_ref} IS NOT NULL " - f"AND {date_ref} IS NOT NULL " - f"GROUP BY {board_ref}, DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " - f"DATEPART(MONTH, {date_ref}) ASC, {board_ref} ASC" - ) - - asks_recurring_failures_by_product = ( - "recurring" in normalized - and "failure" in normalized - and ("product" in normalized or "pcb" in normalized) - ) - if ( - asks_recurring_failures_by_product - and board_model_column - and failure_code_column - ): - product_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" - failure_ref = f"{table_ref}.{self._quote_sql_identifier(failure_code_column)}" - return ( - f"SELECT {product_ref} AS " - f"{self._quote_sql_identifier(board_model_column)}, " - f"{failure_ref} AS " - f"{self._quote_sql_identifier(failure_code_column)}, " - f'COUNT(*) AS "failure_count" ' - f"FROM {table_ref} " - f"WHERE {product_ref} IS NOT NULL " - f"AND {failure_ref} IS NOT NULL " - f"GROUP BY {product_ref}, {failure_ref} " - f'ORDER BY "failure_count" DESC' - ) - - asks_highest_priority_repairs = ( - "repair" in normalized - and "priority" in normalized - and any(term in normalized for term in ("highest", "top", "high priority")) - ) - if asks_highest_priority_repairs and priority_column: - priority_ref = f"{table_ref}.{self._quote_sql_identifier(priority_column)}" - select_refs = [] - for column in ( - id_column, - board_model_column, - failure_code_column, - status_column, - priority_column, - created_at_column, - ): - if column and column not in select_refs: - select_refs.append(column) - select_sql = ", ".join( - f"{table_ref}.{self._quote_sql_identifier(column)} AS " - f"{self._quote_sql_identifier(column)}" - for column in select_refs - ) - return ( - f"SELECT TOP {limit} {select_sql} " - f"FROM {table_ref} " - f"WHERE {priority_ref} IS NOT NULL " - f"ORDER BY CASE LOWER({priority_ref}) " - f"WHEN 'critical' THEN 1 " - f"WHEN 'high' THEN 2 " - f"WHEN 'medium' THEN 3 " - f"WHEN 'low' THEN 4 " - f"ELSE 5 END" - ) - - asks_repair_ticket_distribution = ( - "repair" in normalized - and "ticket" in normalized - and ( - "distribution" in normalized - or "again distribution" in normalized - or "aging distribution" in normalized - ) - ) - if asks_repair_ticket_distribution: - if "aging" in normalized and created_at_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" - age_bucket = ( - f"CASE " - f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 7 THEN '0-7 days' " - f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 30 THEN '8-30 days' " - f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 90 THEN '31-90 days' " - f"ELSE '90+ days' END" - ) - return ( - f'SELECT {age_bucket} AS "age_bucket", ' - f'COUNT(*) AS "ticket_count" ' - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"GROUP BY {age_bucket} " - f'ORDER BY "ticket_count" DESC' - ) - distribution_column = status_column or priority_column or failure_code_column - if distribution_column: - dimension_ref = ( - f"{table_ref}.{self._quote_sql_identifier(distribution_column)}" - ) - return ( - f"SELECT {dimension_ref} AS " - f"{self._quote_sql_identifier(distribution_column)}, " - f'COUNT(*) AS "ticket_count" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f'ORDER BY "ticket_count" DESC' - ) - - return None - - def _get_unqueryable_metric_message( - self, query: str, table_ddls: list[str] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_schema = re.sub( - r"\s+", - " ", - " ".join(ddl for ddl in table_ddls if isinstance(ddl, str)).lower(), - ) - schema_column_names = self._extract_schema_column_names(table_ddls) - - if not normalized_query: - return None - - if "throughput" in normalized_query and any( - term in normalized_query for term in ("manufacturing", "unit", "units") - ): - unit_field_patterns = ( - r"\bbusiness[_ ]?unit\b", - r"\bmanufacturing[_ ]?unit\b", - r"\bunit[_ ]?name\b", - r"\bunit\b", - r"\bbu\b", - r"\bdivision\b", - ) - has_unit_field = any( - re.search(pattern, column_name) - for pattern in unit_field_patterns - for column_name in schema_column_names - ) - has_temporal_field = any( - self._is_temporal_schema_type(str(column.get("type") or "")) - for table in self._parse_schema_tables(table_ddls) - for column in table.get("columns", []) - ) - if not has_unit_field: - return ( - "The active datasource does not expose a manufacturing unit, " - "business unit, unit, BU, or division column. I cannot build " - "throughput trends across manufacturing units without a " - "queryable unit field." - ) - if "trend" in normalized_query and not has_temporal_field: - return ( - "The active datasource does not expose a queryable date or " - "timestamp column. I cannot build a throughput trend without " - "a first-class temporal field." - ) - - if any( - term in normalized_query - for term in ( - "monthly", - "trend", - "turnaround", - "time", - "duration", - "elapsed", - "latest", - "recent", - "newest", - "last records", - ) - ): - has_temporal_field = any( - self._is_temporal_schema_type(str(column.get("type") or "")) - for table in self._parse_schema_tables(table_ddls) - for column in table.get("columns", []) - ) - if not has_temporal_field: - return ( - "The active datasource does not expose a queryable date or " - "timestamp column. I cannot build a time-based analysis " - "without a first-class temporal field." - ) - - repair_cost_terms = ( - "repair cost", - "repair_cost", - "repaircost", - "cost", - "cost impact", - "cost_impact", - ) - if any(term in normalized_query for term in repair_cost_terms): - cost_field_patterns = ( - r"\brepair[_ ]?cost\b", - r"\bcost[_ ]?impact\b", - r"\bcost[_ ]?amount\b", - r"\btotal[_ ]?cost\b", - r"\bunit[_ ]?cost\b", - r"\bcost\b", - r"\bamount\b", - ) - has_cost_field = any( - re.search(pattern, column_name) - for pattern in cost_field_patterns - for column_name in schema_column_names - ) - - if not has_cost_field: - return ( - "The schema does not expose repair cost as a queryable " - "column. The MSSQL Wren/Ibis runtime cannot extract cost " - "from generic JSON/text fields such as data. Add repair " - "cost as a first-class column or calculated field, then " - "ask again." - ) - - first_pass_yield_terms = ( - "first pass yield", - "first-pass yield", - "first_pass_yield", - "fpy", - ) - if not any(term in normalized_query for term in first_pass_yield_terms): - return None - - required_field_patterns = ( - r"\bfirst[_ ]?pass[_ ]?yield\b", - r"\bfpy\b", - r"\battempt\b", - r"\battempt[_ ]?number\b", - r"\bfirst[_ ]?attempt\b", - r"\bpass[_ ]?fail\b", - r"\byield\b", - ) - has_required_field = any( - re.search(pattern, normalized_schema) - for pattern in required_field_patterns - ) - - if has_required_field: - return None - - return ( - "The schema does not expose first-pass yield, attempt number, " - "first-attempt result, or pass/fail fields as queryable columns. " - "I cannot calculate First Pass Yield from only generic JSON/text " - "fields such as data. Add those fields as first-class columns or " - "calculated fields, then ask again." - ) - - def _build_schema_grounded_sales_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_schema = "\n".join( - ddl for ddl in table_ddls or [] if isinstance(ddl, str) - ).lower() - if not normalized_query or not normalized_schema: - return None - - schema_key = self._normalize_schema_token(normalized_schema) - is_sales_specific_query = any( - term in normalized_query - for term in ("sale", "sales", "salesperson", "sales person") - ) - if is_sales_specific_query and "salesvalue" not in schema_key: - return None - - return self._build_schema_grounded_analytics_sql(query, table_ddls) - - async def _run_with_timeout( - self, - label: str, - coroutine, - timeout_seconds: Optional[int] = None, - ): - timeout = timeout_seconds or self._pipeline_timeout_seconds - try: - return await asyncio.wait_for( - coroutine, - timeout=timeout, - ) - except TimeoutError as exc: - raise TimeoutError(f"{label} timed out after {timeout} seconds") from exc - - def _should_retry_selected_schema_after_retrieval_timeout( - self, retrieval_table_names: Optional[list[str]] - ) -> bool: - return bool(retrieval_table_names) - - def _forced_explicit_table_names( - self, table_names: list[str], *, source: str = "request" - ) -> list[str]: - if not table_names: - return [] - if len(table_names) <= MAX_FORCED_EXPLICIT_TABLES: - return table_names - - logger.info( - "Treating broad %s explicit_tables list as retrieval candidates, not a forced schema scope: %s", - source, - table_names, - ) - return [] - - def _build_greeting_response(self, query: str) -> str: - return ( - f"Hi. I can help with questions about your active datasource and Wren AI.\n\n" - f"Try a data question like:\n" - f"- Show monthly trends for the last 12 months\n" - f"- Compare totals by category\n" - f"- Which records occur most often?\n\n" - f"If you want, ask a database question directly instead of `{query}`." - ) - - def _extract_pipeline_reply(self, result: dict, key: str) -> str: - payload = result.get(key) - if isinstance(payload, tuple): - payload = payload[0] - - if isinstance(payload, dict): - replies = payload.get("replies") or [] - if replies and isinstance(replies[0], str): - return replies[0] - - return "" - - def _extract_retrieval_documents(self, retrieval_result: dict) -> list[dict]: - construct_result = retrieval_result.get("construct_retrieval_results", {}) - documents = construct_result.get("retrieval_results", []) - if not isinstance(documents, list): - logger.warning("Schema retrieval returned invalid document payload") - return [] - - valid_documents = [] - for document in documents: - if not isinstance(document, dict): - logger.warning("Ignoring malformed retrieval document: %s", document) - continue - if not document.get("table_name") and not document.get("table_ddl"): - logger.warning("Ignoring retrieval document without table metadata") - continue - valid_documents.append(document) - - return valid_documents - - def _extract_retrieval_metadata( - self, retrieval_result: dict - ) -> tuple[list[dict], list[str], list[str]]: - documents = self._extract_retrieval_documents(retrieval_result) - return documents, *self._metadata_from_documents(documents) - - def _metadata_from_documents( - self, documents: list[dict] - ) -> tuple[list[str], list[str]]: - table_names = [ - table_name - for document in documents - if isinstance(table_name := document.get("table_name"), str) - and table_name.strip() - ] - table_ddls = [ - table_ddl - for document in documents - if isinstance(table_ddl := document.get("table_ddl"), str) - and table_ddl.strip() - ] - return table_names, table_ddls - - async def _complete_sql_generation_context( - self, - *, - query: str, - project_id: Optional[str], - documents: list[dict], - table_names: list[str], - table_ddls: list[str], - ) -> tuple[list[dict], list[str], list[str], dict]: - if not table_names or "db_schema_retrieval" not in self._pipelines: - return documents, table_names, table_ddls, {} - - selected_table_names = list(dict.fromkeys(table_names)) - try: - retrieval_result = await self._run_with_timeout( - "Complete selected schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=query, - tables=selected_table_names, - project_id=project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min(self._schema_retrieval_timeout_seconds, 30), - ) - except Exception as error: - logger.warning( - "Complete selected schema retrieval failed; using existing retrieval context. project_id=%s tables=%s error=%s", - project_id, - selected_table_names, - error, - ) - return documents, table_names, table_ddls, {} - - complete_documents, complete_table_names, complete_table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if not complete_documents: - logger.warning( - "Complete selected schema retrieval returned no documents; using existing retrieval context. project_id=%s tables=%s", - project_id, - selected_table_names, - ) - return documents, table_names, table_ddls, {} - - logger.info( - "Completed SQL generation context with full schemas for project_id %s tables=%s", - project_id, - complete_table_names, - ) - return ( - complete_documents, - complete_table_names, - complete_table_ddls, - retrieval_result.get("construct_retrieval_results", {}), - ) - - def _is_visualization_request(self, query: str) -> bool: - normalized = (query or "").lower() - return bool( - re.search( - r"\b(?:chart|graph|plot|visuali[sz]e|dashboard|bar|line|pie|donut|" - r"scatter|histogram|heatmap|trend|trends|distribution)\b", - normalized, - ) - ) - - def _get_metadata_question_kind(self, query: str) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").lower()).strip() - if not normalized: - return None - - if self._is_visualization_request(normalized): - return None - - if re.search(r"\b(?:row|rows|record|records)\s+count\b", normalized): - return None - - relationship_patterns = ( - r"\b(?:relationships?|relations?|joins?|foreign keys?|primary keys?)\b", - r"\b(?:how|what|which|show|list|describe)\b.*\b(?:tables?|models?)\b.*\b(?:connected|related|joined)\b", - ) - if any(re.search(pattern, normalized) for pattern in relationship_patterns): - return "relationships" - - table_count_patterns = ( - r"\b(?:how many|count|number of)\b.*\b(?:tables?|models?)\b", - r"\b(?:tables?|models?)\b.*\b(?:count|number)\b", - ) - if any(re.search(pattern, normalized) for pattern in table_count_patterns): - return "table_count" - - column_count_patterns = ( - r"\b(?:how many|count|number of)\b.*\b(?:columns?|fields?)\b", - r"\b(?:columns?|fields?)\b.*\b(?:count|number)\b", - ) - if any(re.search(pattern, normalized) for pattern in column_count_patterns): - return "column_count" - - schema_patterns = ( - r"\b(?:what|show|display|describe|list)\b.*\b(?:schema|metadata)\b", - r"\b(?:schema|metadata)\b.*\b(?:of|for|in)\b", - ) - if any(re.search(pattern, normalized) for pattern in schema_patterns): - return "schema" - - explicit_column_patterns = ( - r"\b(?:what|which|list|show|display|give|describe)\b.*\b(?:columns?|fields?)\b", - r"\b(?:columns?|fields?)\b.*\b(?:available|present|there|exist|schema|metadata)\b", - ) - if any(re.search(pattern, normalized) for pattern in explicit_column_patterns): - return "columns" - - table_patterns = ( - r"\b(?:what|which|list|show|display|give)\b.*\b(?:tables?|models?)\b", - r"\b(?:tables?|models?)\b.*\b(?:available|present|there|exist|in this datasource|in the datasource)\b", - r"\b(?:datasource|database|semantic layer|semantic model)\b.*\b(?:tables?|models?)\b", - ) - if any(re.search(pattern, normalized) for pattern in table_patterns): - return "tables" - - return None - - def _find_metadata_table_matches( - self, query: str, tables: list[dict[str, Any]] - ) -> list[dict[str, Any]]: - query_key = self._normalize_schema_token(query) - if not query_key: - return [] - - matches: list[tuple[int, dict[str, Any]]] = [] - for table in tables: - table_name = str(table.get("name") or "") - if not table_name: - continue - short_name = re.split(r"[.$]", table_name)[-1] - normalized_name = self._normalize_schema_token(table_name) - normalized_short_name = self._normalize_schema_token(short_name) - - score = 0 - if normalized_name and normalized_name in query_key: - score = 100 + len(normalized_name) - elif normalized_short_name and normalized_short_name in query_key: - score = 80 + len(normalized_short_name) - - if score: - matches.append((score, table)) - - return [ - table - for _, table in sorted(matches, key=lambda item: item[0], reverse=True) - ] - - def _format_metadata_table_list( - self, tables: list[dict[str, Any]], *, max_tables: int = 120 - ) -> str: - if not tables: - return "I couldn't find any deployed tables in the active datasource metadata." - - sorted_tables = sorted( - {str(table.get("name")) for table in tables if table.get("name")}, - key=str.lower, - ) - shown_tables = sorted_tables[:max_tables] - lines = [ - f"The active datasource has {len(sorted_tables)} deployed table" - f"{'' if len(sorted_tables) == 1 else 's'}:" - ] - lines.extend(f"- {table_name}" for table_name in shown_tables) - if len(sorted_tables) > max_tables: - lines.append( - f"- ...and {len(sorted_tables) - max_tables} more tables." - ) - return "\n".join(lines) - - def _format_metadata_columns( - self, - query: str, - tables: list[dict[str, Any]], - *, - max_tables: int = 25, - max_columns_per_table: int = 60, - ) -> str: - if not tables: - return "I couldn't find any deployed columns in the active datasource metadata." - - matched_tables = self._find_metadata_table_matches(query, tables) - selected_tables = matched_tables or sorted( - tables, key=lambda table: str(table.get("name") or "").lower() - ) - selected_tables = selected_tables[:max_tables] - - heading = ( - "Columns available in the matched deployed table" - if matched_tables and len(selected_tables) == 1 - else "Columns available in the active datasource metadata" - ) - lines = [f"{heading}:"] - for table in selected_tables: - table_name = str(table.get("name") or "unknown_table") - columns = [ - column - for column in table.get("columns", []) - if isinstance(column, dict) and column.get("name") - ] - if not columns: - lines.append(f"- {table_name}: no columns found") - continue - - column_parts = [] - for column in columns[:max_columns_per_table]: - column_name = str(column.get("name")) - column_type = str(column.get("type") or "").upper() - column_parts.append( - f"{column_name} ({column_type})" if column_type else column_name - ) - if len(columns) > max_columns_per_table: - column_parts.append( - f"...and {len(columns) - max_columns_per_table} more" - ) - lines.append(f"- {table_name}: {', '.join(column_parts)}") - - if len(tables) > max_tables and not matched_tables: - lines.append(f"- ...and {len(tables) - max_tables} more tables.") - - return "\n".join(lines) - - def _extract_metadata_relationships(self, table_ddls: list[str]) -> list[str]: - relationships: list[str] = [] - for ddl in table_ddls or []: - if not isinstance(ddl, str): - continue - table_match = re.search( - r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", - ddl, - flags=re.IGNORECASE, - ) - if not table_match: - continue - source_table = next( - (value for value in table_match.groupdict().values() if value), - "unknown_table", - ) - - for relationship_match in re.finditer( - r"FOREIGN\s+KEY\s*\((?P[^)]+)\)\s+REFERENCES\s+" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" - r"\s*\((?P[^)]+)\)", - ddl, - flags=re.IGNORECASE, - ): - target_table = next( - ( - value - for key, value in relationship_match.groupdict().items() - if key - in { - "quoted", - "bracketed", - "backticked", - "bare", - } - and value - ), - "unknown_table", - ) - source_columns = relationship_match.group("source_columns") - target_columns = relationship_match.group("target_columns") - relationships.append( - f"{source_table}({source_columns}) -> " - f"{target_table}({target_columns})" - ) - - return sorted(set(relationships), key=str.lower) - - def _format_metadata_relationships(self, table_ddls: list[str]) -> str: - relationships = self._extract_metadata_relationships(table_ddls) - if not relationships: - return ( - "I couldn't find explicit relationships or foreign keys in the " - "active datasource metadata." - ) - - lines = [ - f"The active datasource metadata has {len(relationships)} " - f"relationship{'' if len(relationships) == 1 else 's'}:" - ] - lines.extend(f"- {relationship}" for relationship in relationships[:120]) - if len(relationships) > 120: - lines.append(f"- ...and {len(relationships) - 120} more relationships.") - return "\n".join(lines) - - def _format_metadata_schema( - self, query: str, tables: list[dict[str, Any]], table_ddls: list[str] - ) -> str: - matched_tables = self._find_metadata_table_matches(query, tables) - selected_tables = matched_tables or sorted( - tables, key=lambda table: str(table.get("name") or "").lower() - ) - selected_tables = selected_tables[:20] - if not selected_tables: - return "I couldn't find schema details in the active datasource metadata." - - lines = ["Schema details from the active datasource metadata:"] - for table in selected_tables: - table_name = str(table.get("name") or "unknown_table") - columns = [ - column - for column in table.get("columns", []) - if isinstance(column, dict) and column.get("name") - ] - lines.append(f"- {table_name}") - if columns: - column_parts = [] - for column in columns[:60]: - column_name = str(column.get("name")) - column_type = str(column.get("type") or "").upper() - column_parts.append( - f"{column_name} ({column_type})" - if column_type - else column_name - ) - if len(columns) > 60: - column_parts.append(f"...and {len(columns) - 60} more") - lines.append(f" Columns: {', '.join(column_parts)}") - else: - lines.append(" Columns: no columns found") - - relationships = self._extract_metadata_relationships(table_ddls) - if relationships: - lines.append("Relationships:") - lines.extend(f"- {relationship}" for relationship in relationships[:40]) - if len(relationships) > 40: - lines.append(f"- ...and {len(relationships) - 40} more relationships.") - - return "\n".join(lines) - - def _format_metadata_table_count(self, tables: list[dict[str, Any]]) -> str: - table_names = {str(table.get("name")) for table in tables if table.get("name")} - return ( - f"The active datasource has {len(table_names)} deployed table" - f"{'' if len(table_names) == 1 else 's'}." - ) - - def _format_metadata_column_count( - self, query: str, tables: list[dict[str, Any]] - ) -> str: - matched_tables = self._find_metadata_table_matches(query, tables) - selected_tables = matched_tables or tables - total_columns = sum( - len( - [ - column - for column in table.get("columns", []) - if isinstance(column, dict) and column.get("name") - ] - ) - for table in selected_tables - ) - if matched_tables and len(selected_tables) == 1: - table_name = str(selected_tables[0].get("name") or "the matched table") - return f"{table_name} has {total_columns} deployed columns." - return ( - f"The active datasource metadata has {total_columns} deployed columns " - f"across {len(selected_tables)} table" - f"{'' if len(selected_tables) == 1 else 's'}." - ) - - def _build_metadata_response( - self, query: str, table_ddls: list[str], table_names: list[str] - ) -> str: - kind = self._get_metadata_question_kind(query) - parsed_tables = self._parse_schema_tables(table_ddls) - - if not parsed_tables and table_names: - parsed_tables = [ - {"name": table_name, "columns": []} for table_name in table_names - ] - - if kind == "schema": - return self._format_metadata_schema(query, parsed_tables, table_ddls) - if kind == "relationships": - return self._format_metadata_relationships(table_ddls) - if kind == "table_count": - return self._format_metadata_table_count(parsed_tables) - if kind == "column_count": - return self._format_metadata_column_count(query, parsed_tables) - if kind == "columns": - return self._format_metadata_columns(query, parsed_tables) - return self._format_metadata_table_list(parsed_tables) - - def _normalize_schema_token(self, value: str) -> str: - return re.sub(r"[^a-z0-9]", "", (value or "").lower()) - - def _query_schema_terms(self, query: str) -> set[str]: - normalized_query = (query or "").lower() - stop_words = { - "about", - "against", - "from", - "give", - "group", - "grouped", - "list", - "month", - "monthly", - "over", - "rows", - "show", - "table", - "tables", - "the", - "this", - "using", - "what", - "which", - "with", - "year", - } - terms = { - self._normalize_schema_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", normalized_query) - if len(token) > 2 and token not in stop_words - } - return {term for term in terms if term} - - def _prune_sql_generation_context( - self, - query: str, - documents: list[dict], - table_names: list[str], - table_ddls: list[str], - *, - max_tables: int = 8, - ) -> tuple[list[dict], list[str], list[str]]: - if len(table_ddls) <= max_tables: - return documents, table_names, table_ddls - - parsed_tables = self._parse_schema_tables(table_ddls) - if not parsed_tables: - return documents, table_names, table_ddls[:max_tables] - - query_key = self._normalize_schema_token(query) - query_terms = self._query_schema_terms(query) - explicit_tables = { - self._normalize_schema_token(table_name) - for table_name in self._extract_explicit_table_names_from_query(query) - } - - scored: list[tuple[int, int]] = [] - for index, table in enumerate(parsed_tables): - table_name = str(table.get("name") or "") - normalized_table = self._normalize_schema_token(table_name) - normalized_short_table = self._normalize_schema_token( - re.split(r"[.$]", table_name)[-1] - ) - column_terms = { - self._normalize_schema_token(str(column.get("name") or "")) - for column in table.get("columns", []) - if column.get("name") - } - - score = 0 - if normalized_table in explicit_tables or normalized_short_table in explicit_tables: - score += 1000 - if normalized_table and normalized_table in query_key: - score += 500 - if normalized_short_table and normalized_short_table in query_key: - score += 450 - for term in query_terms: - if not term: - continue - if term == normalized_table or term == normalized_short_table: - score += 80 - elif term in normalized_table or term in normalized_short_table: - score += 40 - for column_term in column_terms: - if term == column_term: - score += 60 - elif term in column_term or column_term in term: - score += 25 - - if score > 0: - scored.append((score, index)) - - if not scored: - return documents, table_names, table_ddls[:max_tables] - - sorted_scored_indexes = [ - index for _, index in sorted(scored, key=lambda item: item[0], reverse=True) - ] - core_limit = max(1, max_tables - 2) if max_tables > 2 else 1 - selected_indexes = sorted_scored_indexes[:core_limit] - selected_indexes = self._expand_pruned_context_with_related_tables( - selected_indexes, - parsed_tables, - table_ddls, - max_tables=max_tables, - ) - for index in sorted_scored_indexes: - if len(selected_indexes) >= max_tables: - break - if index not in selected_indexes: - selected_indexes.append(index) - selected_indexes = sorted(selected_indexes) - pruned_documents = [ - documents[index] for index in selected_indexes if index < len(documents) - ] - pruned_table_names = [ - table_names[index] for index in selected_indexes if index < len(table_names) - ] - pruned_table_ddls = [ - table_ddls[index] for index in selected_indexes if index < len(table_ddls) - ] - - logger.info( - "Pruned SQL generation context from %s to %s tables for query: %s", - len(table_ddls), - len(pruned_table_ddls), - query, - ) - return pruned_documents, pruned_table_names, pruned_table_ddls - - def _expand_pruned_context_with_related_tables( - self, - selected_indexes: list[int], - parsed_tables: list[dict[str, Any]], - table_ddls: list[str], - *, - max_tables: int, - ) -> list[int]: - if len(selected_indexes) >= max_tables: - return selected_indexes[:max_tables] - - selected: list[int] = list(dict.fromkeys(selected_indexes)) - selected_set = set(selected) - - def join_key_columns(table: dict[str, Any]) -> set[str]: - keys = set() - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - normalized = self._normalize_schema_token(column_name) - if not normalized: - continue - if ( - normalized == "id" - or normalized.endswith("id") - or normalized.endswith("no") - or normalized.endswith("number") - or normalized.endswith("code") - or normalized.endswith("key") - ): - keys.add(normalized) - return keys - - selected_table_names = { - self._normalize_schema_token( - str(parsed_tables[index].get("name") or "") - ) - for index in selected - if index < len(parsed_tables) - } - selected_join_keys: set[str] = set() - for index in selected: - if index < len(parsed_tables): - selected_join_keys.update(join_key_columns(parsed_tables[index])) - - candidates: list[tuple[int, int]] = [] - for index, table in enumerate(parsed_tables): - if index in selected_set: - continue - - table_name = str(table.get("name") or "") - normalized_table_name = self._normalize_schema_token(table_name) - ddl = table_ddls[index] if index < len(table_ddls) else "" - normalized_ddl = self._normalize_schema_token(ddl) - table_join_keys = join_key_columns(table) - - score = 0 - shared_keys = selected_join_keys & table_join_keys - if shared_keys: - score += 20 + 5 * len(shared_keys) - if normalized_table_name and any( - selected_table - and ( - selected_table in normalized_ddl - or normalized_table_name in selected_table - ) - for selected_table in selected_table_names - ): - score += 40 - if re.search(r"\b(?:foreign\s+key|references)\b", ddl, flags=re.IGNORECASE): - score += 15 - - if score > 0: - candidates.append((score, index)) - - for _, index in sorted(candidates, key=lambda item: item[0], reverse=True): - if len(selected) >= max_tables: - break - selected.append(index) - selected_set.add(index) - - return selected - - def _is_valid_select_sql(self, sql: Optional[str]) -> bool: - if not isinstance(sql, str): - return False - - normalized = re.sub(r"\s+", " ", sql.strip()) - if not normalized: - return False - - return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) - - def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: - if not self._is_valid_select_sql(sql): - return None - return AskResult(sql=sql.strip(), type="llm") - - def _build_validated_ask_result_from_sql( - self, - sql: Optional[str], - table_ddls: list[str], - query: str | None = None, - ) -> Optional[AskResult]: - if isinstance(sql, str): - sql = normalize_sql_direction_keywords(sql) - sql = normalize_sql_table_references_to_schema( - sql, - construct_valid_table_names(table_ddls), - ) - sql = normalize_sql_column_references_to_schema( - sql, - construct_valid_table_columns(table_ddls), - ) - ask_result = self._build_ask_result_from_sql(sql) - if not ask_result: - return None - - schema_tables = self._parse_schema_tables(table_ddls) - valid_tables = { - str(table.get("name") or "").lower(): table - for table in schema_tables - if table.get("name") - } - valid_table_suffixes = { - table_name.split(".")[-1].lower(): table - for table_name, table in valid_tables.items() - } - - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_tables = [ - next(value for value in match.groupdict().values() if value) - for match in table_reference_pattern.finditer(ask_result.sql) - ] - invalid_tables = [ - table - for table in referenced_tables - if table.lower() not in valid_tables - and table.lower().split(".")[-1] not in valid_table_suffixes - ] - - columns_by_table = { - table_name: { - str(column.get("name") or "").lower() - for column in table.get("columns", []) - if column.get("name") - } - for table_name, table in valid_tables.items() - } - columns_by_table.update( - { - table_name.split(".")[-1].lower(): columns - for table_name, columns in columns_by_table.items() - } - ) - - qualified_column_pattern = re.compile( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_$]*))", - flags=re.IGNORECASE, - ) - invalid_columns = [] - for match in qualified_column_pattern.finditer(ask_result.sql): - table_reference = ( - match.group("table_quoted") - or match.group("table_bracketed") - or match.group("table_bare") - or "" - ) - column_reference = ( - match.group("column_quoted") - or match.group("column_bracketed") - or match.group("column_bare") - or "" - ) - table_key = table_reference.lower() - column_key = column_reference.lower() - table_columns = columns_by_table.get(table_key) or columns_by_table.get( - table_key.split(".")[-1] - ) - if table_columns is not None and column_key not in table_columns: - invalid_columns.append(f"{table_reference}.{column_reference}") - - if invalid_tables or invalid_columns: - logger.warning( - "Ignoring heuristic SQL because it is not valid for active schema. " - "invalid_tables=%s invalid_columns=%s sql=%s", - invalid_tables, - invalid_columns, - ask_result.sql, - ) - return None - - invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( - ask_result.sql, - schema_tables, - ) - if invalid_unqualified_identifiers: - logger.warning( - "Ignoring SQL because it references unqualified fields outside the active schema. " - "invalid_identifiers=%s sql=%s", - invalid_unqualified_identifiers, - ask_result.sql, - ) - return None - - invalid_output_aliases = self._invalid_sql_output_aliases( - ask_result.sql, - schema_tables, - ) - if invalid_output_aliases: - logger.warning( - "Ignoring SQL because it aliases output fields to unavailable schema concepts. " - "invalid_aliases=%s sql=%s", - invalid_output_aliases, - ask_result.sql, - ) - return None - - if not self._sql_references_explicit_table(ask_result.sql, query): - return None - - if not self._sql_matches_question_intent( - ask_result.sql, - query, - schema_tables, - ): - return None - - return ask_result - - def _build_failed_text_to_sql_response( - self, - trace_id: Optional[str], - message: str, - *, - rephrased_question: Optional[str] = None, - intent_reasoning: Optional[str] = None, - retrieved_tables: Optional[list[str]] = None, - sql_generation_reasoning: Optional[str] = None, - invalid_sql: Optional[str] = None, - is_followup: bool = False, - code: Literal["NO_RELEVANT_DATA", "NO_RELEVANT_SQL", "OTHERS"] = "NO_RELEVANT_SQL", - ) -> AskResultResponse: - return AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError(code=code, message=message), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=retrieved_tables, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - trace_id=trace_id, - is_followup=is_followup, - ) - - def _build_no_relevant_active_datasource_response( - self, - trace_id: Optional[str], - *, - rephrased_question: Optional[str] = None, - intent_reasoning: Optional[str] = None, - retrieved_tables: Optional[list[str]] = None, - sql_generation_reasoning: Optional[str] = None, - is_followup: bool = False, - ) -> AskResultResponse: - return self._build_failed_text_to_sql_response( - trace_id, - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=retrieved_tables, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=None, - is_followup=is_followup, - code="NO_RELEVANT_DATA", - ) - - @observe(name="Ask Question") - @trace_metadata - async def ask( - self, - ask_request: AskRequest, - **kwargs, - ): - trace_id = kwargs.get("trace_id") - results = { - "ask_result": {}, - "metadata": { - "type": "", - "error_type": "", - "error_message": "", - "request_from": ask_request.request_from, - }, - } - - query_id = ask_request.query_id - if not query_id: - raise ValueError("query_id is required for ask service execution") - - user_query = (ask_request.query or "").strip() - if not user_query: - self._ask_results[query_id] = self._build_failed_text_to_sql_response( - trace_id, - "Question is required", - code="OTHERS", - ) - results["metadata"]["error_type"] = "OTHERS" - results["metadata"]["error_message"] = "Question is required" - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - logger.info(f"Ask pipeline started for query_id: {query_id}") - histories = ask_request.histories[: self._max_histories][ - ::-1 - ] # reverse the order of histories - if histories and not self._should_use_histories_for_query(user_query): - logger.info( - "Ignoring thread histories for independent question. query_id=%s query=%s", - query_id, - user_query, - ) - histories = [] - rephrased_question = None - intent_reasoning = None - sql_generation_reasoning = None - sql_samples = [] - instructions = [] - api_results = [] - documents = [] - table_names = [] - table_ddls = [] - _retrieval_result = {} - error_message = None - invalid_sql = None - allow_sql_generation_reasoning = ( - self._allow_sql_generation_reasoning - and not ask_request.ignore_sql_generation_reasoning - ) - enable_column_pruning = ( - self._enable_column_pruning or ask_request.enable_column_pruning - ) - allow_sql_functions_retrieval = self._allow_sql_functions_retrieval - allow_sql_diagnosis = self._allow_sql_diagnosis - allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - max_sql_correction_retries = self._max_sql_correction_retries - current_sql_correction_retries = 0 - use_dry_plan = ask_request.use_dry_plan - allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback - sql_knowledge = None - understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) - request_explicit_table_names = self._normalize_explicit_table_names( - ask_request.explicit_tables - ) - query_explicit_table_names = self._normalize_explicit_table_names( - self._extract_explicit_table_names_from_query(user_query) - ) - forced_request_explicit_table_names = self._forced_explicit_table_names( - request_explicit_table_names, - source="request", - ) - explicit_table_names = ( - forced_request_explicit_table_names or query_explicit_table_names - ) - retrieval_table_names = explicit_table_names or None - - try: - sql_user_query = user_query - - # ask status can be understanding, searching, generating, finished, failed, stopped - # we will need to handle business logic for each status - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="understanding", - trace_id=trace_id, - is_followup=True if histories else False, - ) - - if self._is_greeting_query(user_query): - self._general_streaming_results[query_id] = ( - self._build_greeting_response(user_query) - ) - - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - trace_id=trace_id, - is_followup=True if histories else False, - general_type="USER_GUIDE", - ) - results["metadata"]["type"] = "GENERAL" - return results - - metadata_question_kind = self._get_metadata_question_kind(user_query) - if metadata_question_kind: - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="GENERAL", - rephrased_question=user_query, - intent_reasoning=( - "Basic datasource metadata question detected; " - "retrieving deployed schema metadata directly." - ), - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - retrieval_result = await self._run_with_timeout( - "Metadata schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - metadata_answer = self._build_metadata_response( - user_query, table_ddls, table_names - ) - self._general_streaming_results[query_id] = metadata_answer - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=user_query, - intent_reasoning=( - "Answered from active datasource deployed metadata " - "without SQL generation." - ), - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - results["metadata"]["type"] = "GENERAL" - results["metadata"]["metadata_question_kind"] = ( - metadata_question_kind - ) - results["metadata"]["retrieved_table_count"] = len(documents) - return results - - if explicit_table_names: - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - rephrased_question=user_query, - intent_reasoning="Explicit table name detected; retrieving that deployed schema directly.", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - if not documents and not request_explicit_table_names: - logger.info( - "Explicit table retrieval did not return requested active-schema table; " - "loading full active schema. query_id=%s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval for explicit table", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - all_documents, _, _ = self._extract_retrieval_metadata( - retrieval_result - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - all_documents, - explicit_table_names, - ) - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - logger.info( - "Retrieved explicit tables for query_id %s: %s", - query_id, - table_names, - ) - - if ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, - ): - ask_result = self._build_validated_ask_result_from_sql( - ranked_measure_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated ranked measure SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = ranked_measure_sql - - if table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ): - ask_result = self._build_validated_ask_result_from_sql( - table_question_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table question matched deployed schema.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = table_question_sql - - if explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ): - explicit_sql, explicit_table_name = explicit_table_preview - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table preview request matched deployed schema.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = explicit_sql - - if documents and ( - deterministic_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls - ) - ): - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = deterministic_sql - - if not documents: - error_message = ( - "The requested table was not found in the deployed schema: " - + ", ".join(explicit_table_names) - ) - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_DATA", - message=error_message, - ), - rephrased_question=user_query, - intent_reasoning="Explicit table request did not match any deployed schema table.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = error_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - rephrased_question = user_query - intent_reasoning = ( - "Explicit table request matched deployed schema; generating SQL against retrieved schema." - ) - sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - - if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - logger.info( - "Retrieved tables for direct heuristic query_id %s: %s", - query_id, - table_names, - ) - - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using direct heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - if ask_result := self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ): - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - - if explicit_group_count_sql := self._build_explicit_group_count_sql( - user_query - ): - invalid_sql = explicit_group_count_sql - rephrased_question = user_query - logger.info( - "Deferring explicit grouped count SQL until active schema validation for query_id %s", - query_id, - ) - - historical_question_result = [] - should_skip_pre_sql_retrieval = self._is_data_analysis_query( - user_query - ) - if should_skip_pre_sql_retrieval: - rephrased_question = user_query - intent_reasoning = ( - "Detected a deployed-data analytics question; skipping " - "intent classification and using SQL generation." - ) - sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - logger.info( - "Skipping pre-SQL retrieval for analytics query_id %s: %s", - query_id, - user_query, - ) - - if ( - not api_results - and not should_skip_pre_sql_retrieval - and self._should_reuse_historical_question_sql( - user_query, histories - ) - ): - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - - try: - historical_question = await self._run_with_timeout( - "Historical question retrieval", - self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ), - timeout_seconds=min(understanding_timeout_seconds, 10), - ) - - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - except TimeoutError as exc: - logger.warning( - "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) + historical_question = await self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ) - valid_historical_results = [] - for result in historical_question_result: - historical_question_text = result.get("question") - if not self._is_reusable_historical_question( - user_query, historical_question_text - ): - logger.info( - "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", - query_id, - user_query, - historical_question_text, - ) - continue + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] - sql_statement = result.get("statement") - if not self._is_valid_select_sql(sql_statement): - logger.warning( - "Ignoring historical question without valid SQL for query_id %s", - query_id, - ) - continue - valid_historical_results.append( + if historical_question_result: + api_results = [ AskResult( **{ - "sql": sql_statement.strip(), + "sql": result.get("statement"), "type": "view" if result.get("viewId") else "llm", "viewId": result.get("viewId"), } ) - ) - - if valid_historical_results: - api_results = valid_historical_results + for result in historical_question_result + ] sql_generation_reasoning = "" - elif not api_results and not should_skip_pre_sql_retrieval: - original_user_query = user_query + else: # Run both pipeline operations concurrently - try: - sql_samples_task, instructions_task = await self._run_with_timeout( - "SQL pair and instruction retrieval", - asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", - ), - ), - timeout_seconds=understanding_timeout_seconds, - ) + sql_samples_task, instructions_task = await asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), + ) - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] - ) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) - except TimeoutError as exc: - logger.warning( - "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) - sql_samples = [] - instructions = [] + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] + ) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) if self._allow_intent_classification: - try: - intent_classification_result = ( - await self._run_with_timeout( - "Intent classification", - self._pipelines["intent_classification"].run( - query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, - ), - timeout_seconds=understanding_timeout_seconds, - ) - ).get("post_process", {}) - except TimeoutError as exc: - logger.warning( - "Intent classification timed out; continuing with TEXT_TO_SQL. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, + intent_classification_result = ( + await self._pipelines["intent_classification"].run( + query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, ) - intent_classification_result = { - "intent": "TEXT_TO_SQL", - "rephrased_question": user_query, - "reasoning": "Intent classification timed out; using SQL generation.", - "db_schemas": [], - } + ).get("post_process", {}) intent = intent_classification_result.get("intent") rephrased_question = intent_classification_result.get( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") - retrieved_db_schemas = intent_classification_result.get( - "db_schemas" - ) or [] - is_original_analytics_query = self._is_data_analysis_query( - original_user_query - ) - is_schema_grounded_query = self._is_schema_grounded_query( - original_user_query, retrieved_db_schemas - ) or self._is_schema_grounded_query( - rephrased_question or "", retrieved_db_schemas - ) - - if intent in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"} and ( - is_original_analytics_query - or is_schema_grounded_query - or self._is_data_analysis_query(rephrased_question or "") - ): - logger.info( - "Overriding intent %s to TEXT_TO_SQL for schema/data query: %s", - intent, - user_query, - ) - intent = "TEXT_TO_SQL" - if is_original_analytics_query: - if rephrased_question and rephrased_question != user_query: - logger.info( - "Ignoring rephrased analytics query from intent classification. original=%s rephrased=%s", - original_user_query, - rephrased_question, - ) - user_query = original_user_query - rephrased_question = original_user_query - elif rephrased_question: + if rephrased_question: user_query = rephrased_question - sql_user_query = ( - self._rewrite_query_for_text_to_sql(user_query) - if self._is_data_analysis_query(user_query) - else user_query - ) - if intent == "MISLEADING_QUERY": - general_result = await self._run_with_timeout( - "Misleading assistance", + asyncio.create_task( self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, @@ -6394,18 +268,14 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "misleading_assistance" ) ) self._ask_results[query_id] = AskResultResponse( status="finished", - type="MISLEADING_QUERY", + type="GENERAL", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, @@ -6415,8 +285,7 @@ async def ask( results["metadata"]["type"] = "MISLEADING_QUERY" return results elif intent == "GENERAL": - general_result = await self._run_with_timeout( - "Data assistance", + asyncio.create_task( self._pipelines["data_assistance"].run( query=user_query, histories=histories, @@ -6424,12 +293,8 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "data_assistance" ) ) @@ -6445,17 +310,12 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results elif intent == "USER_GUIDE": - general_result = await self._run_with_timeout( - "User guide assistance", + asyncio.create_task( self._pipelines["user_guide_assistance"].run( query=user_query, language=ask_request.configurations.language, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "user_guide_assistance" ) ) @@ -6479,11 +339,7 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - if ( - not self._is_stopped(query_id, self._ask_results) - and not api_results - and not documents - ): + if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -6493,530 +349,39 @@ async def ask( is_followup=True if histories else False, ) - try: - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) - ), - ), - timeout_seconds=self._schema_retrieval_timeout_seconds, - ) - except TimeoutError as error: - if not self._should_retry_selected_schema_after_retrieval_timeout( - retrieval_table_names - ): - logger.warning( - "Schema retrieval timed out for data query; not loading full project schema. " - "query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - error, - ) - retrieval_result = {"construct_retrieval_results": {}} - else: - logger.warning( - "Schema retrieval timed out; retrying only explicit selected schemas. " - "query_id=%s project_id=%s tables=%s error=%s", - query_id, - ask_request.project_id, - retrieval_table_names, - error, - ) - retrieval_result = await self._run_with_timeout( - "Selected schema fallback retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 30, - ), - ) + retrieval_result = await self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=ask_request.explicit_tables, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=enable_column_pruning, + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - if not documents: - if explicit_table_names: - logger.info( - "Retrying schema retrieval for explicit tables query_id %s: %s", - query_id, - explicit_table_names, - ) - retrieval_result = await self._run_with_timeout( - "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=enable_column_pruning, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 20, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - if ( - not documents - and self._get_metadata_question_kind(user_query) - and not request_explicit_table_names - ): - logger.info( - "Query-based schema retrieval returned no tables for data question; " - "retrying full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - logger.info( - "Retrieved tables for query_id %s: %s", query_id, table_names - ) - - if not api_results and ( - ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, - ) - ): - logger.info( - "Using schema-grounded ranked measure SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - ranked_measure_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = ranked_measure_sql - error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." - - if not api_results and ( - table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded table question SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - table_question_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = table_question_sql - error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - - if not api_results and ( - explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - ): - explicit_sql, explicit_table_name = explicit_table_preview - logger.info( - "Using explicit table preview SQL for query_id %s and table %s", - query_id, - explicit_table_name, - ) - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = explicit_sql - error_message = "Explicit table preview SQL was not valid for the active datasource schema." - - if not api_results and ( - audit_log_activity_sql := self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ) - ): - logger.info( - "Using schema-grounded audit log activity SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - audit_log_activity_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = audit_log_activity_sql - error_message = ( - "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." - ) - - if ( - not api_results - and self._is_data_analysis_query(user_query) - and ( - schema_grounded_sql := self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - ) - ): - logger.info( - "Using generic schema-grounded analytics SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - schema_grounded_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = schema_grounded_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and any( - term in user_query.lower() - for term in ( - "pcb", - "repair", - "failure", - "business unit", - "business units", - "product line", - "product family", - ) - ): - operational_sql = self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - if operational_sql: - logger.info( - "Using schema-grounded operational SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - operational_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = operational_sql - error_message = ( - "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and ( - deterministic_sales_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded CWSales SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sales_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = deterministic_sales_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - - should_retry_full_schema = ( - not api_results - and self._get_metadata_question_kind(user_query) - and "db_schema_retrieval" in self._pipelines - and not request_explicit_table_names - and not table_names - ) - if should_retry_full_schema: - logger.info( - "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retry", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 30, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - full_documents, full_table_names, full_table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - full_documents, full_table_names, full_table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - full_documents, - explicit_table_names, - ) - ) - if full_documents: - documents, table_names, table_ddls = ( - full_documents, - full_table_names, - full_table_ddls, - ) - logger.info( - "Using full active deployed schema retry for query_id %s: %s", - query_id, - table_names, - ) - - full_schema_preview = self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - full_schema_sql_candidates = ( - self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ), - full_schema_preview[0] if full_schema_preview else None, - self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ), - self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ), - self._build_schema_grounded_sales_sql( - user_query, table_ddls - ), - ) - for full_schema_sql in full_schema_sql_candidates: - if not full_schema_sql: - continue - ask_result = self._build_validated_ask_result_from_sql( - full_schema_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - break - invalid_sql = full_schema_sql - error_message = ( - "Full-schema grounded SQL was not valid for the active datasource schema and question intent." - ) + documents = _retrieval_result.get("retrieval_results", []) + table_names = [document.get("table_name") for document in documents] + table_ddls = [document.get("table_ddl") for document in documents] - if not api_results and ( - unqueryable_metric_message := self._get_unqueryable_metric_message( - user_query, table_ddls - ) - ): - logger.info( - "ask pipeline - NO_RELEVANT_SQL due to unqueryable metric: %s", - user_query, - ) + if not documents: + logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="failed", type="TEXT_TO_SQL", error=AskError( - code="NO_RELEVANT_SQL", - message=unqueryable_metric_message, + code="NO_RELEVANT_DATA", + message="No relevant data", ), rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, - retrieved_tables=table_names, trace_id=trace_id, is_followup=True if histories else False, ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = unqueryable_metric_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - if not documents: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - is_followup=True if histories else False, - ) - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - is_followup=True if histories else False, - ) - ) results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) results["metadata"]["type"] = "TEXT_TO_SQL" return results - if documents and not api_results: - documents, table_names, table_ddls = self._prune_sql_generation_context( - sql_user_query, - documents, - table_names, - table_ddls, - ) - ( - documents, - table_names, - table_ddls, - completed_retrieval_result, - ) = await self._complete_sql_generation_context( - query=sql_user_query, - project_id=ask_request.project_id, - documents=documents, - table_names=table_names, - table_ddls=table_ddls, - ) - if completed_retrieval_result: - _retrieval_result = completed_retrieval_result - - sql_generation_histories = histories - if self._is_data_analysis_query( - sql_user_query - ) and not self._needs_conversation_context(sql_user_query): - sql_generation_histories = [] - allow_sql_generation_reasoning = False - allow_sql_knowledge_retrieval = False - max_sql_correction_retries = min(max_sql_correction_retries, 1) - logger.info( - "Using fast standalone SQL generation path for query_id %s", - query_id, - ) - if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -7032,53 +397,29 @@ async def ask( is_followup=True if histories else False, ) - if sql_generation_histories: - try: - sql_generation_reasoning = ( - await self._run_with_timeout( - "Follow-up SQL generation reasoning", - self._pipelines[ - "followup_sql_generation_reasoning" - ].run( - query=sql_user_query, - contexts=table_ddls, - histories=sql_generation_histories, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), - ) - ).get("post_process", {}) - except Exception as reasoning_error: - logger.warning( - "Follow-up SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", - query_id, - reasoning_error, - ) - sql_generation_reasoning = "" + if histories: + sql_generation_reasoning = ( + await self._pipelines["followup_sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ) + ).get("post_process", {}) else: - try: - sql_generation_reasoning = ( - await self._run_with_timeout( - "SQL generation reasoning", - self._pipelines["sql_generation_reasoning"].run( - query=sql_user_query, - contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), - ) - ).get("post_process", {}) - except Exception as reasoning_error: - logger.warning( - "SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", - query_id, - reasoning_error, + sql_generation_reasoning = ( + await self._pipelines["sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, ) - sql_generation_reasoning = "" + ).get("post_process", {}) self._ask_results[query_id] = AskResultResponse( status="planning", @@ -7103,34 +444,14 @@ async def ask( is_followup=True if histories else False, ) - try: - sql_functions, sql_knowledge = await self._run_with_timeout( - "SQL helper retrieval", - asyncio.gather( - ( - self._pipelines["sql_functions_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_functions_retrieval - else _return_value([]) - ), - ( - self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_knowledge_retrieval - else _return_value(None) - ), - ), - timeout_seconds=min(self._pipeline_timeout_seconds, 10), + if allow_sql_functions_retrieval: + sql_functions = await self._pipelines[ + "sql_functions_retrieval" + ].run( + project_id=ask_request.project_id, ) - except TimeoutError as helper_timeout: - logger.warning( - "SQL helper retrieval timed out for query_id %s; continuing with schema only: %s", - query_id, - helper_timeout, - ) - sql_functions, sql_knowledge = [], None + else: + sql_functions = [] has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -7138,92 +459,63 @@ async def ask( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - try: - if sql_generation_histories: - text_to_sql_generation_results = await self._run_with_timeout( - "Follow-up SQL generation", - self._pipelines["followup_sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=sql_generation_histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), - ) - else: - text_to_sql_generation_results = await self._run_with_timeout( - "SQL generation", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), - ) - except TimeoutError as generation_timeout: - logger.warning( - "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", - query_id, - generation_timeout, + if histories: + text_to_sql_generation_results = await self._pipelines[ + "followup_sql_generation" + ].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + ) + else: + text_to_sql_generation_results = await self._pipelines[ + "sql_generation" + ].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, ) - text_to_sql_generation_results = { - "post_process": { - "valid_generation_result": None, - "invalid_generation_result": None, - } - } - error_message = str(generation_timeout) if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - if ask_result := self._build_validated_ask_result_from_sql( - sql_valid_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - else: - invalid_sql = sql_valid_result.get("sql") - error_message = ( - "SQL generation did not produce SQL that matches the active datasource schema and question intent." + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } ) + ] elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] in { - "TIME_OUT", - "UNSUPPORTED_SQL", - }: - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) + if failed_dry_run_result["type"] == "TIME_OUT": break original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] - sql_diagnosis_reasoning = None current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -7238,67 +530,52 @@ async def ask( ) if allow_sql_diagnosis: - sql_diagnosis_results = await self._run_with_timeout( - "SQL diagnosis", - self._pipelines["sql_diagnosis"].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_request.configurations.language, - ), + sql_diagnosis_results = await self._pipelines[ + "sql_diagnosis" + ].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_request.configurations.language, ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") - correction_error_message = error_message - if sql_diagnosis_reasoning: - correction_error_message = ( - f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" - ) - - sql_correction_results = await self._run_with_timeout( - "SQL correction", - self._pipelines["sql_correction"].run( - contexts=table_ddls, - instructions=instructions, - invalid_generation_result={ - "original_sql": original_sql, - "sql": invalid_sql, - "error": correction_error_message, - }, - project_id=ask_request.project_id, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - query=sql_user_query, - ), + sql_correction_results = await self._pipelines[ + "sql_correction" + ].run( + contexts=table_ddls, + instructions=instructions, + invalid_generation_result={ + "sql": original_sql, + "error": sql_diagnosis_reasoning + if allow_sql_diagnosis + else error_message, + }, + project_id=ask_request.project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_functions=sql_functions, ) if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - if ask_result := self._build_validated_ask_result_from_sql( - valid_generation_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - break - invalid_sql = valid_generation_result.get("sql") - error_message = ( - "SQL correction did not produce SQL that matches the active datasource schema and question intent." - ) + api_results = [ + AskResult( + **{ + "sql": valid_generation_result.get("sql"), + "type": "llm", + } + ) + ] + break failed_dry_run_result = sql_correction_results["post_process"][ "invalid_generation_result" ] - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) if api_results: if not self._is_stopped(query_id, self._ask_results): @@ -7316,64 +593,25 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - else: - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - is_followup=True if histories else False, - ) - ) - if error_message or invalid_sql: - logger.info( - "Suppressed technical SQL failure for query_id %s. " - "error=%s invalid_sql=%s", - query_id, - error_message, - invalid_sql, + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_SQL", + message=error_message or "No relevant SQL", + ), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + trace_id=trace_id, + is_followup=True if histories else False, ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = error_message results["metadata"]["type"] = "TEXT_TO_SQL" return results @@ -7427,13 +665,6 @@ async def get_ask_streaming_result( self, query_id: str, ): - if general_response := self._general_streaming_results.get(query_id): - event = SSEEvent( - data=SSEEvent.SSEEventMessage(message=general_response), - ) - yield event.serialize() - return - if self._ask_results.get(query_id): _pipeline_name = "" if self._ask_results.get(query_id).type == "GENERAL": diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7c7d46ce89..a6c6241fa7 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,189 +2,44 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( - _is_project_wide_analysis_query, - _rerank_table_documents, - _select_relevant_table_documents, check_using_db_schemas_without_pruning, dbschema_retrieval, - expand_business_terms_for_retrieval, + embedding, table_retrieval, ) -def test_project_wide_analysis_query_includes_broad_ranking_questions(): - assert _is_project_wide_analysis_query( - "Which projects have the highest number of completed questions?" - ) - - -def test_project_wide_analysis_query_ignores_empty_query(): - assert not _is_project_wide_analysis_query("") - - -def test_expand_business_terms_for_retrieval_adds_generic_sales_order_terms(): - query = "Show top customers by invoice amount" - - expanded_query = expand_business_terms_for_retrieval(query) - - assert query in expanded_query - assert "transaction purchase billing account geography" in expanded_query - assert "money exchange currency" in expanded_query - - -def test_expand_business_terms_for_retrieval_adds_generic_currency_market_terms(): - query = "Show invoice distribution by currency across markets" - - expanded_query = expand_business_terms_for_retrieval(query) - - assert query in expanded_query - assert "money exchange currency" in expanded_query - - -def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): - query = "Explain what this workspace does" - - assert expand_business_terms_for_retrieval(query) == query - - -def test_rerank_table_documents_prefers_question_relevant_table_text(): - generic_stage = Document( - content="Generic imported staging records with product labels.", - meta={"type": "TABLE_DESCRIPTION", "name": "generic_stage_load"}, - score=0.99, - ) - order_region_table = Document( - content="Business transactions grouped by customer geography and amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "business_transactions"}, - score=0.01, - ) - - documents = _rerank_table_documents( - "Show order distribution across regions.", - [generic_stage, order_region_table], - ) - - assert documents[0].meta["name"] == "business_transactions" - - -def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): - test_load = Document( - content="Raw test load rows for order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ) - order_market_table = Document( - content="New order transaction records with market and customer fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.45, - ) - - documents = _rerank_table_documents( - "Show order distribution across markets.", - [test_load, order_market_table], - ) - - assert documents[0].meta["name"] == "dbo_xStageNewOrders" - - -def test_select_relevant_table_documents_limits_weak_extra_candidates(): - documents = [ - Document( - content="Invoice transactions with product, customer, currency, and amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "invoices"}, - score=0.92, - ), - Document( - content="Product catalog with product names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "products"}, - score=0.86, - ), - Document( - content="Customer account master data.", - meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, - score=0.82, - ), - Document( - content="Exchange rate lookup by currency.", - meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, - score=0.78, - ), - Document( - content="Sales regions and market hierarchy.", - meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, - score=0.74, - ), - Document( - content="Raw staging audit rows with load metadata.", - meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, - score=0.99, - ), - ] - - selected = _select_relevant_table_documents( - "Show invoice distribution by currency across markets", - documents, - ) - - assert 1 <= len(selected) <= 5 - assert "staging_audit" not in [document.meta["name"] for document in selected] +@pytest.mark.asyncio +async def test_embedding_skips_vector_lookup_for_explicit_tables(): + class Embedder: + def __init__(self): + self.called = False + async def run(self, query): + self.called = True + return {"embedding": [0.1]} -def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): - documents = [ - Document( - content="Raw test load rows with order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ), - Document( - content="New order transaction records with market and customer details.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.4, - ), - ] + embedder = Embedder() - selected = _select_relevant_table_documents( - "Show order distribution across markets.", - documents, + result = await embedding( + query="show rows", + embedder=embedder, + histories=[], + tables=["orders"], ) - assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] + assert result == {} + assert not embedder.called @pytest.mark.asyncio -async def test_table_retrieval_caps_embedding_results_before_schema_loading(): +async def test_table_retrieval_returns_embedding_results_without_reranking_or_capping(): documents = [ Document( - content="Raw staging audit rows with load metadata.", - meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, - score=0.99, - ), - Document( - content="Invoice sales transactions with product categories and sales value.", - meta={"type": "TABLE_DESCRIPTION", "name": "sales_invoices"}, - score=0.8, - ), - Document( - content="Product catalog with product names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "products"}, - score=0.7, - ), - Document( - content="Customer account master data.", - meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, - score=0.6, - ), - Document( - content="Sales regions and market hierarchy.", - meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, - score=0.5, - ), - Document( - content="Exchange rate lookup by currency.", - meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, - score=0.4, - ), + content=str({"name": f"table_{index}"}), + meta={"type": "TABLE_DESCRIPTION", "name": f"table_{index}"}, + ) + for index in range(8) ] class Retriever: @@ -192,60 +47,17 @@ async def run(self, query_embedding, filters): return {"documents": documents} result = await table_retrieval( - query="What is the distribution of sales across product categories?", embedding={"embedding": [0.1, 0.2]}, project_id="project-1", tables=[], table_retriever=Retriever(), ) - selected_names = [document.meta["name"] for document in result["documents"]] - assert 1 <= len(selected_names) <= 5 - assert "staging_audit" not in selected_names - - -def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): - transaction_source = Document( - content="Invoice transaction fact rows with customer id and invoice amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "invoice_fact"}, - score=0.95, - ) - reference_source = Document( - content="Customer master reference directory with customer names and accounts.", - meta={"type": "TABLE_DESCRIPTION", "name": "customer_master"}, - score=0.7, - ) - - documents = _rerank_table_documents( - "List customer names without duplicates.", - [transaction_source, reference_source], - ) - - assert documents[0].meta["name"] == "customer_master" - - -def test_rerank_table_documents_prefers_transaction_source_for_metric_question(): - reference_source = Document( - content="Product catalog reference table with names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "product_master"}, - score=0.95, - ) - transaction_source = Document( - content="Sales transaction fact table with product, amount, and revenue.", - meta={"type": "TABLE_DESCRIPTION", "name": "sales_fact"}, - score=0.7, - ) - - documents = _rerank_table_documents( - "Show total sales amount by product.", - [reference_source, transaction_source], - ) - - assert documents[0].meta["name"] == "sales_fact" + assert result["documents"] == documents @pytest.mark.asyncio -async def test_table_retrieval_fetches_explicit_table_descriptions(): +async def test_table_retrieval_fetches_explicit_table_descriptions_with_project_scope(): class Retriever: def __init__(self): self.filters = None @@ -257,7 +69,6 @@ async def run(self, query_embedding, filters): retriever = Retriever() await table_retrieval( - query="show rows", embedding={}, project_id="project-1", tables=["orders"], @@ -275,7 +86,30 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_dbschema_retrieval_loads_selected_active_project_schema(): +async def test_table_retrieval_without_embedding_or_explicit_tables_returns_empty(): + class Retriever: + def __init__(self): + self.called = False + + async def run(self, query_embedding, filters): + self.called = True + return {"documents": []} + + retriever = Retriever() + + result = await table_retrieval( + embedding={}, + project_id="project-1", + tables=[], + table_retriever=retriever, + ) + + assert result == {"documents": []} + assert not retriever.called + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_loads_schema_for_retrieved_tables_with_project_scope(): class Retriever: def __init__(self): self.filters = None @@ -293,24 +127,13 @@ async def run(self, query_embedding, filters): } ), meta={"type": "TABLE_SCHEMA", "name": "orders"}, - ), - Document( - content=str( - { - "type": "TABLE", - "name": "customers", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "customers"}, - ), + ) ] } retriever = Retriever() documents = await dbschema_retrieval( - query="total orders", table_retrieval={ "documents": [ Document( @@ -323,19 +146,24 @@ async def run(self, query_embedding, filters): dbschema_retriever=retriever, ) - assert [document.meta["name"] for document in documents] == ["orders", "customers"] + assert [document.meta["name"] for document in documents] == ["orders"] assert retriever.filters == { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": "orders"} + ], + }, {"field": "project_id", "operator": "==", "value": "project-1"}, - {"field": "name", "operator": "in", "value": ["orders"]}, ], } @pytest.mark.asyncio -async def test_dbschema_retrieval_does_not_load_full_schema_for_unmatched_question(): +async def test_dbschema_retrieval_returns_empty_when_no_tables_are_retrieved(): class Retriever: def __init__(self): self.called = False @@ -347,7 +175,6 @@ async def run(self, query_embedding, filters): retriever = Retriever() documents = await dbschema_retrieval( - query="show top customers by invoice amount", table_retrieval={"documents": []}, project_id="project-1", dbschema_retriever=retriever, @@ -364,20 +191,20 @@ def encode(self, value): result = check_using_db_schemas_without_pruning( construct_db_schemas=[ - { - "type": "TABLE", - "name": "orders", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "amount", - "data_type": "DOUBLE", - "comment": "", - "is_primary_key": False, - } - ], - "properties": {}, + { + "type": "TABLE", + "name": "orders", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, "primaryKey": "", } ], @@ -389,37 +216,3 @@ def encode(self, value): assert result["db_schemas"] == [] assert result["tokens"] > 0 - - -@pytest.mark.asyncio -async def test_dbschema_retrieval_uses_explicit_tables_as_scope(): - class Retriever: - def __init__(self): - self.filters = None - - async def run(self, query_embedding, filters): - self.filters = filters - return {"documents": []} - - retriever = Retriever() - - await dbschema_retrieval( - query="show failed repairs", - table_retrieval={"documents": []}, - project_id="project-1", - dbschema_retriever=retriever, - tables=["dbo.failure_patterns", "dbo_failure_patterns"], - ) - - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - { - "field": "name", - "operator": "in", - "value": ["dbo.failure_patterns", "dbo_failure_patterns"], - }, - ], - } From 932cade7a4394c74ab7eebd58a0d9de5cc5ec440 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 16 Jul 2026 03:09:46 +0530 Subject: [PATCH 0583/1087] Restore ask greeting compatibility --- wren-ai-service/src/web/v1/services/ask.py | 27 ++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c4f6e0bfa7..158daaf624 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,5 +1,6 @@ import asyncio import logging +import re from typing import Dict, List, Literal, Optional from cachetools import TTLCache @@ -119,6 +120,9 @@ def __init__( self._ask_results: Dict[str, AskResultResponse] = TTLCache( maxsize=maxsize, ttl=ttl ) + self._general_streaming_results: Dict[str, str] = TTLCache( + maxsize=maxsize, ttl=ttl + ) self._allow_sql_generation_reasoning = allow_sql_generation_reasoning self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_intent_classification = allow_intent_classification @@ -138,6 +142,20 @@ def _is_stopped(self, query_id: str, container: dict): return False + def _is_greeting_query(self, query: str | None) -> bool: + normalized = " ".join(re.findall(r"[a-z]+", (query or "").lower())) + return normalized in { + "hi", + "hello", + "hey", + "good morning", + "good afternoon", + "good evening", + } + + def _build_greeting_response(self, query: str | None) -> str: + return "Hello! Ask me a question about your data, and I will help generate an answer." + @observe(name="Ask Question") @trace_metadata async def ask( @@ -665,6 +683,15 @@ async def get_ask_streaming_result( self, query_id: str, ): + if query_id in self._general_streaming_results: + event = SSEEvent( + data=SSEEvent.SSEEventMessage( + message=self._general_streaming_results[query_id], + ), + ) + yield event.serialize() + return + if self._ask_results.get(query_id): _pipeline_name = "" if self._ask_results.get(query_id).type == "GENERAL": From af62997d867f307206f78381592155d255f10901 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 16 Jul 2026 03:17:47 +0530 Subject: [PATCH 0584/1087] Undo last ask retrieval changes --- .../retrieval/db_schema_retrieval.py | 514 +- wren-ai-service/src/web/v1/services/ask.py | 7224 ++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 349 +- 3 files changed, 7736 insertions(+), 351 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 33c2617f55..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,5 +1,6 @@ import ast import logging +import re import sys from typing import TYPE_CHECKING, Any, Optional @@ -28,6 +29,8 @@ logger = logging.getLogger("wren-ai-service") +MAX_RELEVANT_TABLE_CANDIDATES = 5 + table_columns_selection_system_prompt = """ ### TASK ### @@ -125,6 +128,413 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline +def expand_business_terms_for_retrieval(query: str) -> str: + normalized = (query or "").lower() + expansions: list[str] = [] + + if any( + term in normalized + for term in ( + "amount", + "currency", + "currencies", + "customer", + "customers", + "invoice", + "invoices", + "market", + "markets", + "order", + "orders", + "product", + "products", + "category", + "categories", + "quantity", + "qty", + "region", + "regions", + "sales", + "salesperson", + "sales person", + "sold", + "value", + ) + ): + expansions.append( + "transaction purchase billing account geography area representative product item category sku quantity units sold amount value total metric money exchange currency" + ) + + if any( + term in normalized + for term in ("defect", "failure", "issue", "repair", "resolved", "status") + ): + expansions.append( + "issue defect category status resolved created updated date timestamp event" + ) + + if any(term in normalized for term in ("throughput", "production", "manufacturing")): + expansions.append( + "rate volume output capacity process unit group completed timestamp date" + ) + + if not expansions: + return query + + return f"{query}\n" + "\n".join(expansions) + + +def _normalize_retrieval_token(value: str) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + + +def _retrieval_terms(value: str) -> set[str]: + stop_words = { + "about", + "across", + "and", + "are", + "ask", + "bar", + "chart", + "create", + "different", + "for", + "from", + "how", + "in", + "is", + "of", + "show", + "the", + "to", + "top", + "what", + "which", + "with", + } + terms = { + _normalize_retrieval_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") + if len(token) > 2 and token.lower() not in stop_words + } + return {term for term in terms if term} + + +def _query_mentions_any(query: str, terms: tuple[str, ...]) -> bool: + normalized = (query or "").lower() + return any(re.search(rf"\b{re.escape(term)}\b", normalized) for term in terms) + + +def _source_text(document: Document) -> str: + return " ".join( + str(part or "") + for part in ( + document.meta.get("name"), + document.meta.get("description"), + document.content, + ) + ).lower() + + +def _source_shape_score(query: str, document: Document) -> int: + normalized_query = (query or "").lower() + source_text = _source_text(document) + source_terms = _retrieval_terms(source_text) + + score = 0 + non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "stage", + "staging", + "temp", + "test", + "tmp", + ) + if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( + normalized_query, + non_production_terms, + ): + score -= 60 + + aggregation_terms = ( + "amount", + "average", + "avg", + "count", + "distribution", + "metric", + "revenue", + "sum", + "total", + "trend", + "value", + "volume", + ) + transaction_source_terms = ( + "activity", + "detail", + "event", + "fact", + "history", + "invoice", + "line", + "order", + "sale", + "sales", + "transaction", + ) + reference_source_terms = ( + "account", + "catalog", + "dimension", + "directory", + "entity", + "lookup", + "master", + "profile", + "reference", + ) + entity_listing_pattern = re.search( + r"\b(?:list|show|display|get|find)\b.*\b(?:accounts?|customers?|" + r"employees?|entities|items?|names?|products?|suppliers?|users?|vendors?)\b", + normalized_query, + ) + asks_for_aggregation = _query_mentions_any(normalized_query, aggregation_terms) or bool( + re.search(r"\b(?:by|per|each|top|bottom|rank|ranking)\b", normalized_query) + ) + asks_for_entity_listing = bool(entity_listing_pattern) and not asks_for_aggregation + + if asks_for_entity_listing: + if source_terms & set(reference_source_terms): + score += 35 + if source_terms & set(transaction_source_terms): + score -= 12 + elif asks_for_aggregation: + if source_terms & set(transaction_source_terms): + score += 25 + if source_terms & set(reference_source_terms): + score += 5 + + return score + + +def _document_relevance_score(document: Document, query_terms: set[str]) -> int: + if not query_terms: + return 0 + + document_terms = _retrieval_terms( + " ".join( + str(part or "") + for part in ( + document.meta.get("name"), + document.meta.get("description"), + document.content, + ) + ) + ) + if not document_terms: + return 0 + + score = 0 + for query_term in query_terms: + if query_term in document_terms: + score += 20 + continue + for document_term in document_terms: + if query_term in document_term or document_term in query_term: + score += 8 + break + return score + + +def _semantic_score(document: Document) -> float: + score = getattr(document, "score", None) + if isinstance(score, (int, float)): + return float(score) + score = document.meta.get("score") + if isinstance(score, (int, float)): + return float(score) + return 0.0 + + +def _score_table_documents( + query: str, documents: list[Document] +) -> list[tuple[float, int, Document, int, float]]: + if not documents: + return [] + + query_terms = _retrieval_terms(expand_business_terms_for_retrieval(query)) + if not query_terms: + return [ + (_semantic_score(document), -index, document, 0, _semantic_score(document)) + for index, document in enumerate(documents) + ] + + scored_documents: list[tuple[float, int, Document, int, float]] = [] + for index, document in enumerate(documents): + lexical_score = _document_relevance_score(document, query_terms) + semantic_score = _semantic_score(document) + source_shape_score = _source_shape_score(query, document) + combined_score = semantic_score + lexical_score + source_shape_score + scored_documents.append( + (combined_score, -index, document, lexical_score, semantic_score) + ) + + return sorted(scored_documents, key=lambda item: (item[0], item[1]), reverse=True) + + +def _rerank_table_documents(query: str, documents: list[Document]) -> list[Document]: + if not documents: + return documents + + reranked = _score_table_documents(query, documents) + if not reranked: + return documents + + logger.info( + "Top table candidates after retrieval rerank: %s", + [ + { + "name": document.meta.get("name"), + "semantic_score": round(semantic_score, 4), + "lexical_score": lexical_score, + "combined_score": round(combined_score, 4), + } + for combined_score, _index, document, lexical_score, semantic_score in reranked[ + :5 + ] + ], + ) + return [document for _score, _index, document, _lexical, _semantic in reranked] + + +def _select_relevant_table_documents( + query: str, + documents: list[Document], + *, + max_tables: int = MAX_RELEVANT_TABLE_CANDIDATES, +) -> list[Document]: + if not documents or max_tables <= 0: + return [] + + reranked = _score_table_documents(query, documents) + if not reranked: + return documents[:max_tables] + + candidate_pool = [item for item in reranked if item[3] > 0] or reranked + selected = [ + document + for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] + ] + if len(selected) < len(documents): + logger.info( + "Scoped table candidates for schema loading from %s to %s tables: %s", + len(documents), + len(selected), + [document.meta.get("name") for document in selected], + ) + return selected + + +def _is_project_wide_analysis_query(query: str) -> bool: + normalized = (query or "").lower() + if not normalized: + return False + + analysis_terms = { + "average", + "avg", + "bar chart", + "breakdown", + "chart", + "completed", + "compare", + "count", + "counts", + "distribution", + "group by", + "grouped", + "highest", + "line chart", + "lowest", + "maximum", + "minimum", + "monthly", + "most common", + "number of", + "pie chart", + "quarter", + "rank", + "ranking", + "recommend", + "recommended", + "show", + "status", + "sum", + "total", + "totals", + "top", + "trend", + "volume", + } + return any(term in normalized for term in analysis_terms) + + +def _dedupe_documents(documents: list[Document]) -> list[Document]: + deduped: list[Document] = [] + seen: set[tuple[str, str, str]] = set() + for document in documents: + key = ( + str(document.meta.get("name", "")), + str(document.meta.get("type", "")), + document.content, + ) + if key in seen: + continue + seen.add(key) + deduped.append(document) + return deduped + + +def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: + normalized: list[str] = [] + for table_name in table_names or []: + if not isinstance(table_name, str): + continue + table_name = table_name.strip() + if table_name and table_name not in normalized: + normalized.append(table_name) + return normalized + + +def _extract_table_names_from_table_retrieval( + table_retrieval: dict, explicit_tables: Optional[list[str]] = None +) -> list[str]: + table_names = _normalize_table_names(explicit_tables) + for document in table_retrieval.get("documents") or []: + if not isinstance(document, Document): + continue + table_name = document.meta.get("name") + if not isinstance(table_name, str): + try: + content = ast.literal_eval(document.content) + except (SyntaxError, ValueError): + content = {} + table_name = content.get("name") if isinstance(content, dict) else None + if isinstance(table_name, str): + table_name = table_name.strip() + if table_name and table_name not in table_names: + table_names.append(table_name) + return table_names + + @observe(capture_input=False, capture_output=False) async def embedding( query: str, @@ -143,6 +553,7 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query + query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -151,9 +562,13 @@ async def embedding( @observe(capture_input=False) async def table_retrieval( - embedding: dict, project_id: str, tables: list[str], table_retriever: Any + query: str, + embedding: dict, + project_id: str, + tables: list[str], + table_retriever: Any, ) -> dict: - filters = { + base_filters = { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, @@ -161,61 +576,82 @@ async def table_retrieval( } if project_id: - filters["conditions"].append( + base_filters["conditions"].append( {"field": "project_id", "operator": "==", "value": project_id} ) if embedding: - return await table_retriever.run( + results = await table_retriever.run( query_embedding=embedding.get("embedding"), - filters=filters, + filters=base_filters, ) - elif tables: - filters["conditions"].append( - {"field": "name", "operator": "in", "value": tables} + results["documents"] = _select_relevant_table_documents( + query, results.get("documents") or [] ) + return results - return await table_retriever.run( - query_embedding=[], - filters=filters, - ) + if tables: + logger.info("Loading explicit table descriptions: %s", tables) + explicit_filters = { + **base_filters, + "conditions": [ + *base_filters["conditions"], + {"field": "name", "operator": "in", "value": tables}, + ], + } + return await table_retriever.run(query_embedding=[], filters=explicit_filters) return {"documents": []} @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, project_id: str, dbschema_retriever: Any + query: str, + table_retrieval: dict, + project_id: str, + dbschema_retriever: Any, + tables: Optional[list[str]] = None, ) -> list[Document]: - tables = table_retrieval.get("documents", []) - table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) - - table_name_conditions = [ - {"field": "name", "operator": "==", "value": table_name} - for table_name in table_names - ] - - if table_name_conditions: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } + selected_table_names = _extract_table_names_from_table_retrieval( + table_retrieval, tables + ) - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] + if selected_table_names: + filters["conditions"].append( + {"field": "name", "operator": "in", "value": selected_table_names} + ) + logger.info( + "Loading selected deployed schema metadata for active project_id %s tables=%s", + project_id, + selected_table_names, + ) + elif not query: + logger.info( + "Loading complete deployed schema metadata for active project_id %s", + project_id, + ) + else: + logger.info( + "No relevant table-description candidates found for active project_id %s; " + "skipping full schema loading for query=%s", + project_id, + query, + ) + return [] - return [] + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results.get("documents", []) @observe() diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 158daaf624..c4a7e11b14 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,18 +1,34 @@ import asyncio import logging import re -from typing import Dict, List, Literal, Optional +from typing import Any, Dict, List, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + normalize_sql_direction_keywords, + normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, +) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent logger = logging.getLogger("wren-ai-service") +NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( + "No relevant data found in the active datasource for this question." +) +MAX_FORCED_EXPLICIT_TABLES = 5 + + +async def _return_value(value): + return value + class AskHistory(BaseModel): sql: str @@ -80,7 +96,7 @@ class _AskResultResponse(BaseModel): rephrased_question: Optional[str] = None intent_reasoning: Optional[str] = None sql_generation_reasoning: Optional[str] = None - type: Optional[Literal["GENERAL", "TEXT_TO_SQL"]] = None + type: Optional[Literal["GENERAL", "TEXT_TO_SQL", "MISLEADING_QUERY"]] = None retrieved_tables: Optional[List[str]] = None response: Optional[List[AskResult]] = None invalid_sql: Optional[str] = None @@ -100,6 +116,39 @@ class AskResultResponse(_AskResultResponse): class AskService: + _HISTORICAL_QUESTION_STOP_WORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "can", + "chart", + "create", + "each", + "for", + "from", + "give", + "graph", + "how", + "in", + "is", + "me", + "of", + "on", + "please", + "show", + "the", + "to", + "total", + "what", + "which", + "with", + } + def __init__( self, pipelines: Dict[str, BasicPipeline], @@ -109,9 +158,9 @@ def __init__( allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, + max_sql_correction_retries: int = 3, pipeline_timeout_seconds: int = 90, schema_retrieval_timeout_seconds: int = 180, - max_sql_correction_retries: int = 3, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -142,143 +191,6202 @@ def _is_stopped(self, query_id: str, container: dict): return False - def _is_greeting_query(self, query: str | None) -> bool: - normalized = " ".join(re.findall(r"[a-z]+", (query or "").lower())) - return normalized in { + @classmethod + def _normalize_historical_question_text(cls, question: str | None) -> str: + return " ".join(re.findall(r"[a-z0-9]+", (question or "").lower())) + + @classmethod + def _historical_question_tokens(cls, question: str | None) -> set[str]: + normalized = cls._normalize_historical_question_text(question) + return { + token + for token in normalized.split() + if len(token) > 1 and token not in cls._HISTORICAL_QUESTION_STOP_WORDS + } + + @classmethod + def _is_reusable_historical_question( + cls, query: str | None, historical_question: str | None + ) -> bool: + normalized_query = cls._normalize_historical_question_text(query) + normalized_historical_question = cls._normalize_historical_question_text( + historical_question + ) + if not normalized_query or not normalized_historical_question: + return False + return normalized_query == normalized_historical_question + + @classmethod + def _should_use_histories_for_query(cls, query: str | None) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + contextual_prefixes = ( + "also ", + "and ", + "but ", + "for those ", + "for that ", + "for the same ", + "from that ", + "how about ", + "in that ", + "now ", + "same ", + "show more", + "show the same", + "then ", + "use that ", + "what about ", + "what if ", + ) + if normalized.startswith(contextual_prefixes): + return True + + contextual_patterns = ( + r"\b(previous|last|above|earlier|same|those|that|these|them|it|its|there)\b", + r"\b(add|break down|compare|filter|group|instead|only|sort|split)\b.+\b(by|to|with)\b", + r"\b(by|for|with)\s+(month|quarter|year|status|type|category|customer|market|region|country|division)\b", + ) + return any(re.search(pattern, normalized) for pattern in contextual_patterns) + + def _is_greeting_query(self, query: str) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + greeting_patterns = { "hi", "hello", "hey", + "hii", + "hola", "good morning", "good afternoon", "good evening", + "how are you", + "thanks", + "thank you", } + return normalized in greeting_patterns - def _build_greeting_response(self, query: str | None) -> str: - return "Hello! Ask me a question about your data, and I will help generate an answer." + def _is_data_analysis_query(self, query: str) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False - @observe(name="Ask Question") - @trace_metadata - async def ask( + analysis_terms = { + "amount", + "average", + "avg", + "bar chart", + "bottom", + "chart", + "common", + "compare", + "count", + "cost", + "claim", + "claims", + "currency", + "currencies", + "customer", + "customers", + "dashboard", + "debug", + "distribution", + "failure", + "fastest growing", + "growth", + "group", + "grouped", + "invoice", + "invoices", + "margin", + "market", + "markets", + "monthly", + "order", + "orders", + "pcb", + "performance", + "profit", + "product", + "products", + "product type", + "product types", + "quarter", + "quantity", + "rank", + "ranking", + "region", + "regions", + "repair", + "resolved", + "revenue", + "sale", + "sales", + "sales person", + "sales rep", + "salesperson", + "sla", + "top", + "trend", + "turnaround", + "value", + "volume", + "year", + "yearly", + } + return any(term in normalized for term in analysis_terms) + + def _needs_conversation_context(self, query: str) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + return any( + re.search(pattern, normalized) + for pattern in ( + r"\b(previous|last|above|earlier)\s+(query|question|answer|result|sql|chart)\b", + r"\b(same|that|those|them|it)\s+(table|query|question|result|chart|sql|period|filter)\b", + r"\b(use|using|based on|compare with|compared with)\s+(that|previous|last|above|earlier)\b", + r"\bwhat about\b", + r"\bhow about\b", + ) + ) + + def _should_reuse_historical_question_sql( self, - ask_request: AskRequest, - **kwargs, - ): - trace_id = kwargs.get("trace_id") - results = { - "ask_result": {}, - "metadata": { - "type": "", - "error_type": "", - "error_message": "", - "request_from": ask_request.request_from, - }, + query: str, + histories: list[AskHistory] | None, + ) -> bool: + return False + + def _rewrite_query_for_text_to_sql(self, query: str) -> str: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return query + + guidance: list[str] = [] + + if any( + term in normalized + for term in ("chart", "bar chart", "line chart", "pie chart", "graph") + ): + guidance.append( + "Return SQL only for the aggregated dataset required to build the requested chart." + ) + + if any( + term in normalized + for term in ( + "failure category", + "failure categories", + "common failure", + "common failures", + "failure code", + "top 10", + "most common", + ) + ): + guidance.append( + "Use an exposed failure category, failure name, or failure code field from the schema and return that dimension with a count metric." + ) + + if any( + term in normalized + for term in ("monthly", "last 12 months", "last month", "trend", "volume") + ): + guidance.append( + "Use a real timestamp column from the schema and aggregate results by calendar month when a monthly trend is requested." + ) + + if re.search( + r"\b(?:by|across|per|each|grouped by|group by)\s+[a-z][a-z0-9 _-]*", + normalized, + ) or "over time" in normalized: + guidance.append( + "Preserve explicit grouping dimensions requested by the question, such as market, region, currency, product, customer, status, type, or category, using only matching columns exposed in the provided schema." + ) + + if any(term in normalized for term in ("currency", "currencies", "fx")): + guidance.append( + "For currency questions, use an exposed currency, money, exchange, or FX code/name column from the schema and group by it." + ) + + if any( + term in normalized + for term in ( + "amount", + "cost", + "revenue", + "sales value", + "sum", + "total", + "value", + ) + ) and not any( + term in normalized + for term in ("count", "how many", "number of records", "record count") + ): + guidance.append( + "When the question asks for total, sum, amount, value, revenue, or cost, aggregate an exposed numeric measure with SUM; use COUNT only for record-count questions." + ) + + if not guidance: + return query + + return f"{query}\n\nSQL generation guidance:\n- " + "\n- ".join(guidance) + + def _schema_contains( + self, + table_ddls: list[str], + pattern: str, + table_names: Optional[list[str]] = None, + ) -> bool: + schema_text = "\n".join(table_ddls or []) + if table_names: + schema_text += "\n" + "\n".join(table_names) + return bool(re.search(pattern, schema_text, flags=re.IGNORECASE)) + + def _schema_has_table_column( + self, + table_ddls: list[str], + table_name: str, + column_name: str, + table_names: Optional[list[str]] = None, + ) -> bool: + table_pattern = rf"\b{re.escape(table_name)}\b" + column_pattern = rf"\b{re.escape(column_name)}\b" + + for ddl in table_ddls or []: + if re.search(table_pattern, ddl, flags=re.IGNORECASE) and re.search( + column_pattern, ddl, flags=re.IGNORECASE + ): + return True + + return False + + def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: + column_names: list[str] = [] + non_column_prefixes = ( + "create ", + "constraint ", + "foreign ", + "primary ", + "unique ", + "index ", + ")", + "/*", + "--", + ) + + for ddl in table_ddls: + if not isinstance(ddl, str): + continue + for line in ddl.splitlines(): + stripped = line.strip().rstrip(",") + if not stripped: + continue + if stripped.lower().startswith(non_column_prefixes): + continue + + column_match = re.match( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_]*))\s+", + stripped, + ) + if not column_match: + continue + + column_name = next( + value for value in column_match.groupdict().values() if value + ) + column_names.append(str(column_name).lower()) + + return column_names + + def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: + tables: list[dict[str, Any]] = [] + for ddl in table_ddls or []: + if not isinstance(ddl, str): + continue + table_match = re.search( + r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", + ddl, + flags=re.IGNORECASE, + ) + if not table_match: + continue + + table_name = next( + (value for value in table_match.groupdict().values() if value), + None, + ) + if not table_name: + continue + body_start = table_match.end() + depth = 1 + body_end = body_start + while body_end < len(ddl) and depth > 0: + if ddl[body_end] == "(": + depth += 1 + elif ddl[body_end] == ")": + depth -= 1 + body_end += 1 + + columns: list[dict[str, str]] = [] + for line in ddl[body_start : body_end - 1].splitlines(): + stripped = line.strip().rstrip(",") + if not stripped or stripped.startswith(("--", "/*")): + continue + if re.match( + r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", + stripped, + flags=re.IGNORECASE, + ): + continue + + column_match = re.match( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_$]*))" + r"\s+(?P[A-Za-z0-9_(),]+)", + stripped, + ) + if column_match: + column_name = next( + (value + for key, value in column_match.groupdict().items() + if key != "type" and value + ), + None, + ) + if not column_name: + continue + column_type = column_match.group("type") or "" + columns.append( + { + "name": str(column_name), + "type": str(column_type).lower(), + } + ) + + tables.append({"name": table_name, "columns": columns}) + + return tables + + _INTENT_STOPWORDS = { + "a", + "an", + "and", + "are", + "as", + "based", + "be", + "by", + "can", + "chart", + "correct", + "data", + "different", + "do", + "does", + "each", + "for", + "from", + "give", + "how", + "in", + "is", + "it", + "list", + "many", + "me", + "of", + "on", + "or", + "per", + "question", + "rate", + "records", + "reduce", + "show", + "system", + "taken", + "the", + "there", + "to", + "total", + "type", + "types", + "what", + "which", + "with", + } + + def _intent_tokens(self, text: str) -> set[str]: + tokens: set[str] = set() + for raw_token in re.findall(r"[A-Za-z][A-Za-z0-9_]*", text or ""): + split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) + for token in re.findall(r"[A-Za-z0-9]+", split_token.lower()): + if len(token) <= 2 or token in self._INTENT_STOPWORDS: + continue + tokens.add(token) + if token.endswith("ies") and len(token) > 4: + tokens.add(token[:-3] + "y") + elif token.endswith("s") and len(token) > 3: + tokens.add(token[:-1]) + return tokens + + def _schema_name_tokens(self, name: str) -> set[str]: + spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(name or "")) + return { + token + for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) + if len(token) > 1 } - query_id = ask_request.query_id - histories = ask_request.histories[: self._max_histories][ - ::-1 - ] # reverse the order of histories - rephrased_question = None - intent_reasoning = None - sql_generation_reasoning = None - sql_samples = [] - instructions = [] - api_results = [] - table_names = [] - error_message = None - invalid_sql = None - allow_sql_generation_reasoning = ( - self._allow_sql_generation_reasoning - and not ask_request.ignore_sql_generation_reasoning + def _table_for_sql_reference( + self, table_reference: str, valid_tables: dict[str, dict[str, Any]] + ) -> dict[str, Any] | None: + table_key = str(table_reference or "").lower() + if table_key in valid_tables: + return valid_tables[table_key] + suffix_key = table_key.split(".")[-1] + for valid_table_name, table in valid_tables.items(): + if valid_table_name.split(".")[-1] == suffix_key: + return table + return None + + def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return [] + + concept_groups: list[set[str]] = [] + if "product line" in normalized or "productline" in normalized: + concept_groups.append({"product", "prod", "line", "productline"}) + if "pcb" in normalized: + concept_groups.append( + { + "pcb", + "board", + "repair", + "repairs", + "debug", + "failure", + "failures", + } + ) + if "critical" in normalized: + concept_groups.append({"critical", "severity", "priority"}) + if "cost" in normalized: + concept_groups.append({"cost", "amount", "expense", "impact"}) + if "currency" in normalized or "currencies" in normalized: + concept_groups.append({"currency", "curr", "money", "fx", "exchange"}) + if "market" in normalized or "markets" in normalized: + concept_groups.append({"market", "region", "country", "territory"}) + if "region" in normalized or "regions" in normalized: + concept_groups.append({"region", "market", "area", "territory", "country"}) + if "quarterly" in normalized or "quarter" in normalized: + concept_groups.append({"quarter", "quarterly"}) + if "recurring" in normalized or "recurrence" in normalized: + concept_groups.append({"recurring", "recurrence", "occurrence", "occurrences", "count"}) + if "issue" in normalized or "issues" in normalized: + concept_groups.append({"issue", "issues", "failure", "failures", "problem", "defect"}) + + return concept_groups + + def _sql_covers_required_question_concepts( + self, + sql: str, + query: str | None, + referenced_column_tokens: set[str], + referenced_table_tokens: set[str], + ) -> bool: + sql_text = (sql or "").lower() + available_tokens = referenced_column_tokens | referenced_table_tokens + for concept_group in self._required_sql_concept_groups(query): + if concept_group & available_tokens: + continue + if any(token in sql_text for token in concept_group): + continue + logger.warning( + "Ignoring SQL because it does not cover required question concept. " + "query=%s required=%s referenced_column_tokens=%s referenced_table_tokens=%s sql=%s", + query, + sorted(concept_group), + sorted(referenced_column_tokens), + sorted(referenced_table_tokens), + sql, + ) + return False + return True + + def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return True + asks_for_measure_sum = any( + term in normalized_query + for term in ( + "amount", + "cost", + "revenue", + "sales value", + "sum", + "total", + "value", + ) ) - enable_column_pruning = ( - self._enable_column_pruning or ask_request.enable_column_pruning + asks_for_count = any( + term in normalized_query + for term in ( + "count", + "counts", + "how many", + "number of", + "record count", + "records", + "rows", + ) ) - allow_sql_functions_retrieval = self._allow_sql_functions_retrieval - allow_sql_diagnosis = self._allow_sql_diagnosis - max_sql_correction_retries = self._max_sql_correction_retries - current_sql_correction_retries = 0 - use_dry_plan = ask_request.use_dry_plan - allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback + if not asks_for_measure_sum or asks_for_count: + return True - try: - user_query = ask_request.query + normalized_sql = re.sub(r"\s+", " ", sql or "").lower() + if re.search(r"\b(sum|avg|min|max)\s*\(", normalized_sql): + return True + if re.search(r"\bcount\s*\(", normalized_sql): + logger.warning( + "Ignoring SQL because a measure-total question was answered with row counting. " + "query=%s sql=%s", + query, + sql, + ) + return False + return True - # ask status can be understanding, searching, generating, finished, failed, stopped - # we will need to handle business logic for each status - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="understanding", - trace_id=trace_id, - is_followup=True if histories else False, + def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not re.search( + r"\b(?:distinct|unique|no duplicate|no duplicates|without duplicates)\b", + normalized_query, + ): + return True + + normalized_sql = re.sub(r"\s+", " ", sql or "").strip() + if re.search(r"\bcount\s*\(\s*distinct\b", normalized_sql, flags=re.IGNORECASE): + return True + if re.search(r"\bGROUP\s+BY\b", normalized_sql, flags=re.IGNORECASE): + group_match = re.search( + r"\bGROUP\s+BY\b(?P.*?)(?:\bORDER\s+BY\b|\bHAVING\b|$)", + normalized_sql, + flags=re.IGNORECASE | re.DOTALL, + ) + if group_match: + group_items = [ + item.strip() + for item in re.split(r",(?![^()]*\))", group_match.group("group")) + if item.strip() + ] + if len(group_items) > 1: + logger.warning( + "Ignoring SQL because GROUP BY covers multiple columns and can still duplicate the requested entity. " + "query=%s sql=%s", + query, + sql, + ) + return False + return True + select_match = re.search( + r"\bSELECT\b(?P.*?)\bFROM\b", + sql or "", + flags=re.IGNORECASE | re.DOTALL, + ) + if not select_match: + return [] + + invalid: list[str] = [] + for match in re.finditer( + r'(?P(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$.]*))\s+AS\s+' + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_]*))', + select_match.group("select"), + flags=re.IGNORECASE, + ): + expression = match.group("expr").strip('"[]') + expression_key = expression.split(".")[-1].lower() + alias = match.group("quoted") or match.group("bracketed") or match.group("bare") or "" + alias_key = alias.lower() + if not expression_key or not alias_key: + continue + if alias_key in valid_columns or expression_key not in valid_columns: + continue + alias_terms = { + term for term in re.split(r"[^a-z0-9]+", alias_key) if term + } + if alias_terms & allowed_alias_terms: + continue + invalid.append(alias) + + return invalid + + def _unqualified_valid_sql_column_tokens( + self, sql: str, schema_tables: list[dict[str, Any]] + ) -> set[str]: + valid_columns = { + str(column.get("name") or "").lower(): str(column.get("name") or "") + for table in schema_tables + for column in table.get("columns", []) + if column.get("name") + } + if not valid_columns: + return set() + + sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") + identifier_pattern = re.compile( + r'"(?P[^"]+)"|(?P\b[A-Za-z_][A-Za-z0-9_]*\b)' + ) + tokens: set[str] = set() + for match in identifier_pattern.finditer(sql_without_strings): + identifier = match.group("quoted") or match.group("bare") or "" + identifier_key = identifier.lower() + if identifier_key not in valid_columns: + continue + + before = sql_without_strings[: match.start()].rstrip() + after = sql_without_strings[match.end() :].lstrip() + if before.endswith(".") or after.startswith("."): + continue + + previous_word_match = re.search(r"([A-Za-z_][A-Za-z0-9_]*)\s*$", before) + previous_word = ( + previous_word_match.group(1).lower() if previous_word_match else "" + ) + if previous_word == "as" and self._is_alias_identifier_position( + sql_without_strings, + match.start(), + ): + continue + + tokens.update(self._schema_name_tokens(valid_columns[identifier_key])) + + return tokens + + def _sql_matches_question_intent( + self, + sql: str, + query: str | None, + schema_tables: list[dict[str, Any]], + ) -> bool: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + expects_dimension = bool( + re.search( + r"\b(?:by|per|each|which|different|type|types|category|" + r"categories|status|source|market|markets|region|regions|" + r"currency|currencies)\b", + normalized_query, + ) + ) + + question_tokens = self._intent_tokens(query or "") + required_concept_groups = self._required_sql_concept_groups(query) + if not question_tokens and not required_concept_groups: + return True + + valid_tables = { + str(table.get("name") or "").lower(): table + for table in schema_tables + if table.get("name") + } + if not valid_tables: + return True + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, + ) + referenced_tables = [ + next(value for value in match.groupdict().values() if value) + for match in table_reference_pattern.finditer(sql) + ] + if not referenced_tables: + return True + + referenced_table_tokens = set().union( + *[self._schema_name_tokens(table) for table in referenced_tables] + ) + qualified_column_pattern = re.compile( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_$]*))", + flags=re.IGNORECASE, + ) + referenced_columns_by_table: dict[str, set[str]] = {} + for match in qualified_column_pattern.finditer(sql): + table_reference = ( + match.group("table_quoted") + or match.group("table_bracketed") + or match.group("table_bare") + or "" + ).lower() + column_reference = ( + match.group("column_quoted") + or match.group("column_bracketed") + or match.group("column_bare") + or "" + ) + referenced_columns_by_table.setdefault(table_reference, set()).add( + column_reference + ) + + all_referenced_column_tokens = set().union( + *[ + self._schema_name_tokens(column_name) + for columns in referenced_columns_by_table.values() + for column_name in columns + ] + ) if referenced_columns_by_table else set() + all_referenced_column_tokens.update( + self._unqualified_valid_sql_column_tokens(sql, schema_tables) + ) + if not self._sql_covers_required_question_concepts( + sql, + query, + all_referenced_column_tokens, + referenced_table_tokens, + ): + return False + if not self._sql_uses_required_measure_aggregation(sql, query): + return False + if not self._sql_satisfies_unique_entity_request(sql, query): + return False + + if not expects_dimension: + return True + + for table_reference in referenced_tables: + table = self._table_for_sql_reference(table_reference, valid_tables) + if not table: + continue + + columns = [ + column for column in table.get("columns", []) if column.get("name") + ] + intent_matching_columns = [ + str(column.get("name")) + for column in columns + if self._schema_name_tokens(str(column.get("name"))) & question_tokens + ] + if not intent_matching_columns: + continue + + table_key = str(table_reference or "").lower() + referenced_columns = referenced_columns_by_table.get( + table_key + ) or referenced_columns_by_table.get( + table_key.split(".")[-1], + set(), + ) + referenced_column_tokens = ( + set().union( + *[ + self._schema_name_tokens(column_name) + for column_name in referenced_columns + ] + ) + if referenced_columns + else all_referenced_column_tokens + ) + + if not referenced_column_tokens & question_tokens: + logger.warning( + "Ignoring SQL because selected columns do not match question intent. " + "query=%s table=%s matching_schema_columns=%s referenced_columns=%s sql=%s", + query, + table.get("name"), + intent_matching_columns, + sorted(referenced_columns), + sql, ) + return False - historical_question = await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, + return True + + def _is_numeric_schema_type(self, column_type: str) -> bool: + return bool( + re.search( + r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|float|double|" + r"real|money|number)\b", + column_type, + flags=re.IGNORECASE, + ) + ) + + def _is_temporal_schema_type(self, column_type: str) -> bool: + return bool( + re.search( + r"\b(?:date|time|timestamp|datetime|smalldatetime)\b", + column_type, + flags=re.IGNORECASE, + ) + ) + + def _is_text_schema_type(self, column_type: str) -> bool: + return bool( + re.search( + r"\b(?:char|text|string|varchar|nvarchar|uuid|guid|json)\b", + column_type, + flags=re.IGNORECASE, + ) + ) + + def _find_schema_column( + self, + table: dict[str, Any], + candidates: tuple[str, ...], + numeric: bool | None = None, + temporal: bool | None = None, + ) -> str | None: + normalized_candidates = [ + re.sub(r"[^a-z0-9]", "", str(candidate).lower()) + for candidate in candidates + if candidate is not None + ] + if not normalized_candidates: + return None + scored: list[tuple[int, int, str]] = [] + for column in table.get("columns", []): + column_name = column.get("name") + if not column_name: + continue + column_name = str(column_name) + normalized_column = re.sub(r"[^a-z0-9]", "", column_name.lower()) + column_type = str(column.get("type") or "") + if numeric is True and not self._is_numeric_schema_type(column_type): + continue + if temporal is True and not self._is_temporal_schema_type(column_type): + continue + + for candidate_index, candidate in enumerate(normalized_candidates): + if normalized_column == candidate: + scored.append((100, candidate_index, column_name)) + elif candidate and candidate in normalized_column: + scored.append((60 + len(candidate), candidate_index, column_name)) + elif normalized_column and normalized_column in candidate: + scored.append( + (40 + len(normalized_column), candidate_index, column_name) + ) + + if not scored: + return None + + return sorted(scored, key=lambda item: (-item[0], item[1]))[0][2] + + def _find_first_schema_column( + self, + table: dict[str, Any], + candidates: tuple[str, ...], + *, + avoid: set[str] | None = None, + ) -> str | None: + avoid = {str(column).lower() for column in avoid or set()} + for candidate_group in candidates: + column = self._find_schema_column(table, (candidate_group,)) + if column and column.lower() not in avoid: + return column + return None + + def _find_any_temporal_schema_column(self, table: dict[str, Any]) -> str | None: + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if column_name and self._is_temporal_schema_type(column_type): + return column_name + return None + + def _is_probable_explicit_table_token(self, token: str) -> bool: + normalized = re.sub(r"\s+", " ", str(token or "").strip().lower()) + if not normalized: + return False + if normalized in self._INTENT_STOPWORDS or normalized in { + "last", + "latest", + "recent", + "current", + "previous", + "next", + "month", + "year", + "quarter", + "week", + "day", + }: + return False + return True + + def _quote_sql_identifier(self, identifier: str) -> str: + return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' + + def _normalize_schema_identifier_key(self, value: str) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + + def _schema_identifier_alias_keys(self, value: str) -> set[str]: + raw_value = str(value or "").strip() + base_key = self._normalize_schema_identifier_key(raw_value) + separator_normalized_key = self._normalize_schema_identifier_key( + re.sub(r"[.$]", "_", raw_value) + ) + compact_parts_key = "".join( + self._normalize_schema_identifier_key(part) + for part in re.split(r"[.$_]+", raw_value) + if part + ) + return { + key + for key in (base_key, separator_normalized_key, compact_parts_key) + if key + } + + def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: + return self._explicit_table_alias_keys( + self._extract_explicit_table_names_from_query(query or "") + ) + + def _explicit_table_alias_keys(self, table_names: list[str]) -> set[str]: + keys: set[str] = set() + for table_name in table_names: + keys.update(self._schema_identifier_alias_keys(table_name)) + return keys + + def _filter_retrieval_metadata_for_explicit_query( + self, + query: str, + documents: list[dict], + explicit_table_names: Optional[list[str]] = None, + ) -> tuple[list[dict], list[str], list[str]]: + explicit_table_keys = ( + self._explicit_table_alias_keys(explicit_table_names) + if explicit_table_names + else self._explicit_table_alias_keys_from_query(query) + ) + if not explicit_table_keys: + table_names, table_ddls = self._metadata_from_documents(documents) + return documents, table_names, table_ddls + + matched_documents: list[dict] = [] + for document in documents: + candidate_names = [] + if isinstance(table_name := document.get("table_name"), str): + candidate_names.append(table_name) + if isinstance(table_ddl := document.get("table_ddl"), str): + candidate_names.extend( + str(table.get("name") or "") + for table in self._parse_schema_tables([table_ddl]) + if table.get("name") + ) + + candidate_keys: set[str] = set() + for candidate_name in candidate_names: + candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) + candidate_keys.update( + self._schema_identifier_alias_keys( + re.split(r"[.$_]", str(candidate_name or ""))[-1] + ) + ) + if explicit_table_keys.intersection(candidate_keys): + matched_documents.append(document) + + table_names, table_ddls = self._metadata_from_documents(matched_documents) + return matched_documents, table_names, table_ddls + + def _sql_references_explicit_table( + self, + sql: str, + query: str | None, + ) -> bool: + explicit_table_keys = self._explicit_table_alias_keys_from_query(query) + if not explicit_table_keys: + return True + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, + ) + referenced_tables = [ + next(value for value in match.groupdict().values() if value) + for match in table_reference_pattern.finditer(sql) + ] + + for table_reference in referenced_tables: + reference_keys = self._schema_identifier_alias_keys(table_reference) + reference_keys.update( + self._schema_identifier_alias_keys( + re.split(r"[.$_]", str(table_reference or ""))[-1] ) + ) + if explicit_table_keys.intersection(reference_keys): + return True + + logger.warning( + "Ignoring SQL because it does not reference the explicitly requested table. " + "query=%s referenced_tables=%s sql=%s", + query, + referenced_tables, + sql, + ) + return False + + def _table_matches_query(self, table_name: str, query: str) -> bool: + query_keys = self._schema_identifier_alias_keys(query) + short_table = re.split(r"[.$_]", str(table_name or ""))[-1] + table_keys = self._schema_identifier_alias_keys(table_name) + table_keys.update(self._schema_identifier_alias_keys(short_table)) + return any( + table_key and any(table_key in query_key for query_key in query_keys) + for table_key in table_keys + ) + + def _find_best_schema_table_for_query( + self, query: str, tables: list[dict[str, Any]] + ) -> dict[str, Any] | None: + if not tables: + return None + + scored_tables: list[tuple[int, dict[str, Any]]] = [] + query_tokens = self._intent_tokens(query) + for table in tables: + table_name = str(table.get("name") or "") + if not table_name: + continue + + score = 0 + if self._table_matches_query(table_name, query): + score += 100 + + table_tokens = self._schema_name_tokens(table_name) + score += 8 * len(table_tokens & query_tokens) + + column_token_matches = 0 + for column in table.get("columns", []): + column_token_matches += len( + self._schema_name_tokens(str(column.get("name") or "")) + & query_tokens + ) + score += column_token_matches + + if score > 0: + scored_tables.append((score, table)) + + if scored_tables: + return sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] + if len(tables) == 1: + return tables[0] + return None + + def _query_mentions_column(self, query: str, column_name: str) -> bool: + normalized_query = self._normalize_schema_identifier_key(query) + normalized_column = self._normalize_schema_identifier_key(column_name) + if not normalized_column: + return False + if normalized_column in normalized_query: + return True + if normalized_column.endswith("y"): + return f"{normalized_column[:-1]}ies" in normalized_query + return f"{normalized_column}s" in normalized_query + + def _is_alias_identifier_position(self, sql: str, start: int) -> bool: + before = sql[:start].rstrip() + return bool(re.search(r"\bAS\s*$", before, flags=re.IGNORECASE)) + + def _find_dimension_column_for_query( + self, query: str, table: dict[str, Any] + ) -> str | None: + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + if column_name and self._query_mentions_column(query, column_name): + return column_name + + candidate_columns = [ + str(column.get("name")) + for column in table.get("columns", []) + if column.get("name") + and not self._is_temporal_schema_type(str(column.get("type") or "")) + ] + for candidate in ("name", "category", "type", "status", "code"): + column = self._find_schema_column(table, (candidate,)) + if column in candidate_columns: + return column + return candidate_columns[0] if candidate_columns else None + + def _build_schema_ranked_measure_sql( + self, + query: str, + table_ddls: list[str], + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + if not re.search( + r"\b(?:top|highest|largest|biggest|best|lowest|smallest|bottom)\b", + normalized_query, + ): + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + query_tokens = self._intent_tokens(query) + if not query_tokens: + return None + + scored: list[tuple[int, dict[str, Any], str, str]] = [] + for table in tables: + text_columns: list[tuple[int, str]] = [] + numeric_columns: list[tuple[int, str]] = [] + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if not column_name: + continue + + column_tokens = self._schema_name_tokens(column_name) + query_overlap = column_tokens & query_tokens + + if self._is_numeric_schema_type(column_type): + if query_overlap: + numeric_columns.append( + (80 + 5 * len(query_overlap), column_name) + ) + continue + + if self._is_temporal_schema_type(column_type): + continue + + if query_overlap: + text_columns.append((60 + 5 * len(query_overlap), column_name)) + + if not text_columns or not numeric_columns: + continue + + dimension_score, dimension = sorted( + text_columns, + key=lambda item: item[0], + reverse=True, + )[0] + measure_score, measure = sorted( + numeric_columns, + key=lambda item: item[0], + reverse=True, + )[0] + table_score = dimension_score + measure_score + table_tokens = self._schema_name_tokens(str(table.get("name") or "")) + table_score += 5 * len(table_tokens & query_tokens) + scored.append((table_score, table, dimension, measure)) + + if not scored: + return None + + _, table, dimension, measure = sorted( + scored, + key=lambda item: item[0], + reverse=True, + )[0] + table_name = str(table.get("name") or "") + if not table_name: + return None + + limit = self._extract_requested_top_n(query, default_value=10) + table_ref = self._quote_sql_identifier(table_name) + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + measure_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" + metric_expr = f"SUM({measure_ref})" + direction = ( + "ASC" + if re.search(r"\b(?:lowest|smallest|least|bottom)\b", normalized_query) + else "DESC" + ) + return ( + f"SELECT TOP {limit} {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " + f"{metric_expr} AS {self._quote_sql_identifier('Total' + measure)} " + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL AND {measure_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f"ORDER BY {metric_expr} {direction}" + ) + + def _find_temporal_column_for_query( + self, query: str, table: dict[str, Any] + ) -> str | None: + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if ( + column_name + and self._is_temporal_schema_type(column_type) + and self._query_mentions_column(query, column_name) + ): + return column_name + + return self._find_any_temporal_schema_column(table) + + def _build_schema_grounded_table_question_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + tables = self._parse_schema_tables(table_ddls) + table = self._find_best_schema_table_for_query(query, tables) + if not table: + return None + + table_name = str(table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + limit = self._extract_requested_top_n(query, default_value=10) + + wants_latest_records = any( + term in normalized + for term in ("latest", "recent", "newest", "last records", "latest records") + ) + if wants_latest_records: + date_column = self._find_temporal_column_for_query(query, table) + if not date_column: + return None + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT TOP {limit} * " + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"ORDER BY {date_ref} DESC" + ) + + wants_monthly_count = any( + term in normalized + for term in ("monthly", "by month", "per month", "month-wise") + ) and any(term in normalized for term in ("count", "records", "rows")) + if wants_monthly_count: + date_column = self._find_temporal_column_for_query(query, table) + if not date_column: + return None + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " + f"DATEPART(MONTH, {date_ref}) ASC" + ) + + wants_total_count = ( + re.search(r"\bhow many\b", normalized) + or "record count" in normalized + or "count of records" in normalized + or "number of records" in normalized + ) and not re.search(r"\b(?:by|per|each|distribution|highest|top)\b", normalized) + if wants_total_count: + return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' + + wants_grouped_count = ( + any( + term in normalized + for term in ( + "count", + "counts", + "record count", + "number of", + "how many", + ) + ) + and re.search(r"\b(?:by|per|each|grouped by|group by)\b", normalized) + ) + wants_ranked_count = any( + term in normalized for term in ("highest", "top", "most", "largest") + ) and any( + term in normalized + for term in ("count", "counts", "number of", "orders", "records", "rows") + ) + if wants_grouped_count or wants_ranked_count: + dimension_column = self._find_dimension_column_for_query(query, table) + if not dimension_column: + return None + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" + count_column = None + if any(term in normalized for term in ("order", "orders")): + count_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "OrderID", "id"), + ) + count_expression = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(count_column)})" + if count_column + else "COUNT(*)" + ) + top_clause = f"TOP {limit} " if wants_ranked_count else "" + nonblank_filter = ( + f"AND LTRIM(RTRIM({dimension_ref})) <> '' " + if wants_ranked_count + else "" + ) + return ( + f"SELECT {top_clause}{dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_column)}, " + f'{count_expression} AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"{nonblank_filter}" + f"GROUP BY {dimension_ref} " + f"ORDER BY {count_expression} DESC" + ) + + wants_distribution = any( + term in normalized + for term in ( + "distribution", + "highest occurrence", + "highest occurrences", + "most occurrence", + "most occurrences", + "occurrences", + "top", + "common", + ) + ) + if wants_distribution: + dimension_column = self._find_dimension_column_for_query(query, table) + if not dimension_column: + return None + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" + + occurrence_column = self._find_schema_column( + table, + ("occurrences", "occurrence", "count", "total_count", "record_count"), + numeric=True, + ) + if occurrence_column and self._query_mentions_column( + query, occurrence_column + ): + metric_ref = f"{table_ref}.{self._quote_sql_identifier(occurrence_column)}" + return ( + f"SELECT TOP {limit} {dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_column)}, " + f"{metric_ref} AS {self._quote_sql_identifier(occurrence_column)} " + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"ORDER BY {metric_ref} DESC" + ) + + return ( + f"SELECT TOP {limit} {dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_column)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f"ORDER BY COUNT(*) DESC" + ) + + return None + + def _build_explicit_table_preview_sql( + self, query: str, table_ddls: list[str] + ) -> tuple[str, str] | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) + if not normalized_query: + return None + + if not re.search( + r"\b(?:first|top|sample|preview|show|list)\b", + normalized_query, + flags=re.IGNORECASE, + ): + return None + if not re.search( + r"\b(?:rows?|records?|data)\b", normalized_query, flags=re.IGNORECASE + ): + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + normalized_query_key = re.sub(r"[^a-z0-9]", "", normalized_query.lower()) + normalized_query_keys = self._schema_identifier_alias_keys(normalized_query) + scored_tables: list[tuple[int, str]] = [] + for table in tables: + table_name = table.get("name") + if not table_name: + continue + table_name = str(table_name) + normalized_table = re.sub(r"[^a-z0-9]", "", table_name.lower()) + if not normalized_table: + continue + table_keys = self._schema_identifier_alias_keys(table_name) + if normalized_table in normalized_query_key or any( + table_key in query_key + for table_key in table_keys + for query_key in normalized_query_keys + ): + scored_tables.append((100 + len(normalized_table), table_name)) + continue + + table_without_schema = re.split(r"[.$]", table_name)[-1] + normalized_short_name = re.sub( + r"[^a-z0-9]", "", table_without_schema.lower() + ) + if normalized_short_name and normalized_short_name in normalized_query_key: + scored_tables.append((80 + len(normalized_short_name), table_name)) + + if not scored_tables: + return None + + _, table_name = sorted(scored_tables, reverse=True)[0] + limit = self._extract_requested_top_n(query, default_value=10) + return ( + f"SELECT TOP {limit} * FROM {self._quote_sql_identifier(table_name)}", + table_name, + ) + + def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: + table_names: list[str] = [] + for match in re.finditer( + r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if ( + table_name + and self._is_probable_explicit_table_token(table_name) + and table_name not in table_names + ): + table_names.append(table_name) + for match in re.finditer( + r"\bin\s+(?:the\s+)?([A-Za-z_][A-Za-z0-9_.$]*)\s+table\b", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if ( + table_name + and self._is_probable_explicit_table_token(table_name) + and table_name not in table_names + ): + table_names.append(table_name) + for match in re.finditer( + r"\bin\s+([A-Za-z_][A-Za-z0-9_.$]*)", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if ( + table_name + and ("." in table_name or "_" in table_name) + and self._is_probable_explicit_table_token(table_name) + and table_name not in table_names + ): + table_names.append(table_name) + for match in re.finditer( + r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", + query or "", + flags=re.IGNORECASE, + ): + table_name = match.group(1).strip(".,;:()[]{}") + if ( + table_name + and ("." in table_name or "_" in table_name) + and self._is_probable_explicit_table_token(table_name) + and table_name not in table_names + ): + table_names.append(table_name) + if re.search( + r"\b(?:repair\s+logs?|repair\s+tickets?|board\s+models?)\b", + query or "", + flags=re.IGNORECASE, + ): + for table_name in ("repair_logs", "dbo_repair_logs"): + if table_name not in table_names: + table_names.append(table_name) + if re.search(r"\bticket\s+labels?\b", query or "", flags=re.IGNORECASE): + for table_name in ("ticket_labels", "dbo_ticket_labels"): + if table_name not in table_names: + table_names.append(table_name) + return table_names + + def _explicit_table_name_candidates(self, table_name: str) -> list[str]: + table_name = str(table_name or "").strip().strip(".,;:()[]{}") + if not table_name: + return [] + + candidates = [table_name] + separator_normalized = re.sub(r"[.$]", "_", table_name) + if separator_normalized not in candidates: + candidates.append(separator_normalized) + if "_" in table_name and "." not in table_name: + dotted = table_name.replace("_", ".", 1) + if dotted not in candidates: + candidates.append(dotted) + short_name = re.split(r"[.$]", table_name)[-1] + if "." not in table_name and "_" in table_name: + short_name = table_name.split("_", 1)[-1] + if short_name and short_name not in candidates: + candidates.append(short_name) + return candidates + + def _normalize_explicit_table_names( + self, table_names: Optional[list[str]] + ) -> list[str]: + normalized: list[str] = [] + for table_name in table_names or []: + for candidate in self._explicit_table_name_candidates(table_name): + if candidate not in normalized: + normalized.append(candidate) + return normalized + + def _build_direct_orders_sales_sql(self, query: str) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + is_sales_or_orders_query = any( + term in normalized + for term in ( + "sales", + "sale", + "order", + "orders", + "new order", + "new orders", + "market", + "salesperson", + "sales person", + ) + ) + if not is_sales_or_orders_query: + return None + + table_ref = '"dbo_tblSales"' + limit = self._extract_requested_top_n(query, default_value=10) + + if ( + ("salesperson" in normalized or "sales person" in normalized) + and ("order count" in normalized or "orders" in normalized or "count" in normalized) + ): + return ( + f'SELECT TOP {limit} {table_ref}."SalesPerson" AS "SalesPerson", ' + f'COUNT(*) AS "OrderCount" ' + f"FROM {table_ref} " + f'WHERE {table_ref}."SalesPerson" IS NOT NULL ' + f'GROUP BY {table_ref}."SalesPerson" ' + f"ORDER BY COUNT(*) DESC" + ) + + if "top" in normalized and "new order" in normalized: + date_filter = "" + if re.search(r"\b2026[\s-]*q1\b", normalized): + date_filter = ( + f'WHERE {table_ref}."OrdDate" >= \'2026-01-01 00:00:00\' ' + f'AND {table_ref}."OrdDate" < \'2026-04-01 00:00:00\' ' + ) + return ( + f'SELECT TOP {limit} {table_ref}."BU" AS "BU", ' + f'{table_ref}."Market" AS "Market", ' + f'{table_ref}."Customer" AS "Customer", ' + f'{table_ref}."ProdName" AS "ProdName", ' + f'{table_ref}."SalesValue" AS "SalesValue" ' + f"FROM {table_ref} " + f"{date_filter}" + f'ORDER BY {table_ref}."SalesValue" DESC' + ) + + if ( + "market" in normalized + and "growth" in normalized + and any(term in normalized for term in ("last year", "previous year")) + ): + return ( + f'SELECT {table_ref}."Market" AS "Market", ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "CurrentPeriodSales", ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "PreviousPeriodSales", ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) - ' + f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " + f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " + f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "SalesGrowth" ' + f"FROM {table_ref} " + f'WHERE {table_ref}."Market" IS NOT NULL ' + f'GROUP BY {table_ref}."Market" ' + f'ORDER BY "SalesGrowth" DESC' + ) + + if ( + "distribution" in normalized + and "sales" in normalized + and ("market" in normalized or "by market" in normalized) + ): + return ( + f'SELECT {table_ref}."Market" AS "Market", ' + f'SUM({table_ref}."SalesValue") AS "TotalSalesValue" ' + f"FROM {table_ref} " + f'WHERE {table_ref}."Market" IS NOT NULL ' + f'GROUP BY {table_ref}."Market" ' + f'ORDER BY SUM({table_ref}."SalesValue") DESC' + ) + + return None + + def _extract_explicit_table_column_reference( + self, query: str + ) -> tuple[str, str] | None: + normalized_query = query or "" + reference_match = re.search( + r"\b(?P[A-Za-z_][A-Za-z0-9_]*)[._]" + r"(?P
[A-Za-z_][A-Za-z0-9_]*)[._]" + r"(?P[A-Za-z_][A-Za-z0-9_]*)\b", + normalized_query, + ) + if not reference_match: + return None + + schema = reference_match.group("schema") + table = reference_match.group("table") + column = reference_match.group("column") + table_name = f"{schema}_{table}" + return table_name, column + + def _build_explicit_group_count_sql(self, query: str) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + if not any( + term in normalized_query + for term in ( + "group by", + "grouped by", + "by ", + "pie chart", + "donut chart", + "bar chart", + "count", + "counts", + ) + ): + return None + + explicit_reference = self._extract_explicit_table_column_reference(query) + if not explicit_reference: + return None + + table_name, column = explicit_reference + table_ref = self._quote_sql_identifier(table_name) + column_ref = f"{table_ref}.{self._quote_sql_identifier(column)}" + return ( + f"SELECT {column_ref} AS {self._quote_sql_identifier(column)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"GROUP BY {column_ref} " + f"ORDER BY COUNT(*) DESC" + ) + + def _build_date_filter(self, table_name: str, date_column: str, query: str) -> str: + date_ref = ( + f"{self._quote_sql_identifier(table_name)}." + f"{self._quote_sql_identifier(date_column)}" + ) + normalized_query = (query or "").lower() + if "this year" in normalized_query or "current year" in normalized_query: + return ( + f" WHERE {date_ref} >= '2026-01-01 00:00:00' " + f"AND {date_ref} < '2027-01-01 00:00:00'" + ) + year_match = re.search(r"\b(20\d{2})\b", normalized_query) + if year_match: + year = int(year_match.group(1)) + return ( + f" WHERE {date_ref} >= '{year}-01-01 00:00:00' " + f"AND {date_ref} < '{year + 1}-01-01 00:00:00'" + ) + return "" + + def _append_not_null_filters( + self, where_clause: str, column_refs: list[str] + ) -> str: + conditions = [f"{column_ref} IS NOT NULL" for column_ref in column_refs] + if not conditions: + return where_clause + + if where_clause.strip(): + return f"{where_clause.rstrip()} AND {' AND '.join(conditions)}" + return f" WHERE {' AND '.join(conditions)}" + + def _build_schema_literal_filter_conditions( + self, + query: str, + table: dict[str, Any], + table_ref: str, + ) -> list[str]: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + conditions: list[str] = [] + + if "backlog" in normalized_query: + filter_column = self._find_schema_column( + table, + ( + "Category", + "OrderCategory", + "Order Category", + "Status", + "OrderStatus", + "Order Status", + "Stage", + "OrderStage", + "Order Stage", + ), + ) + if filter_column: + filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" + conditions.append(f"{filter_ref} = 'Backlog'") + + return conditions + + def _select_best_analytics_table( + self, + tables: list[dict[str, Any]], + required_dimensions: list[tuple[str, ...]], + measure_candidates: tuple[str, ...], + wants_date: bool = False, + allow_count_metric: bool = False, + query: str = "", + ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + scored: list[ + tuple[int, dict[str, Any], list[str], str | None, str | None] + ] = [] + for table in tables: + dimensions = [ + self._find_schema_column(table, candidates) + for candidates in required_dimensions + ] + if any(dimension is None for dimension in dimensions): + continue + + measure = self._find_schema_column( + table, measure_candidates, numeric=True + ) + if not measure and not allow_count_metric: + continue + date_column = self._find_schema_column( + table, + ( + "OrdDate", + "InvDate", + "OrderDate", + "NewOrderDate", + "Date", + "CreatedAt", + "created_at", + ), + temporal=True, + ) + if wants_date and not date_column: + continue + + score = 10 * len([dimension for dimension in dimensions if dimension]) + if measure: + score += 8 + elif allow_count_metric: + score += 2 + if date_column: + score += 4 + table_name = str(table.get("name") or "").lower() + if not table_name: + continue + if "sales" in table_name: + score += 5 + if "tblsales" in self._normalize_schema_token(table_name): + score += 25 + if "order" in table_name: + score += 4 + if "invoice" in table_name or "inv" in table_name: + score += 3 + if "stage" in table_name: + score -= 8 + if any( + term in normalized_query + for term in ("order", "orders", "new order", "new orders") + ): + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber"), + ) + if "order" in table_name: + score += 30 + if "neworder" in self._normalize_schema_token(table_name): + score += 15 + if order_column: + score += 12 + if "margin" in table_name and "margin" not in normalized_query: + score -= 12 + if "customer" in normalized_query: + if "customer" in table_name or "account" in table_name: + score += 16 + if self._find_schema_column( + table, + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + ), + ): + score += 10 + if any( + term in normalized_query for term in ("product", "products", "item") + ): + if "product" in table_name or "item" in table_name: + score += 16 + if any( + term in normalized_query + for term in ("sales", "revenue", "value", "amount") + ): + if "sales" in table_name: + score += 12 + if "invoice" in normalized_query and ( + "invoice" in table_name or "inv" in table_name + ): + score += 20 + + scored.append((score, table, dimensions, measure, date_column)) + + if not scored: + return None + + _, table, dimensions, measure, date_column = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + return ( + table, + [dimension for dimension in dimensions if dimension], + measure, + date_column, + ) + + def _build_schema_grounded_analytics_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + + if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): + return pcb_direct_sql + + if repair_failure_count_sql := self._build_repair_failure_count_sql( + query, table_ddls + ): + return repair_failure_count_sql + + if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( + query, table_ddls + ): + return monthly_repair_volume_sql + + is_sales_or_order_query = any( + term in normalized_query + for term in ( + "average order value", + "invoice", + "new order", + "order", + "orders", + "currency", + "currencies", + "market", + "markets", + "performance", + "product", + "products", + "quantity", + "qty", + "revenue", + "sale", + "sales", + "salesperson", + "sales person", + "sold", + ) + ) + if not is_sales_or_order_query: + if operational_sql := self._build_schema_grounded_operational_sql( + query, tables + ): + return operational_sql + + if conversion_sql := self._build_order_invoice_conversion_sql( + query, tables + ): + return conversion_sql + + if yoy_sql := self._build_yoy_sales_change_sql(query, tables): + return yoy_sql + + if contribution_sql := self._build_contribution_sql(query, tables): + return contribution_sql + + if not is_sales_or_order_query: + if categorical_count_sql := self._build_generic_categorical_count_sql( + query, tables + ): + return categorical_count_sql + + wants_count_metric = any( + term in normalized_query + for term in ("count", "counts", "volume", "how many", "distribution") + ) and not any( + term in normalized_query + for term in ( + "amount", + "cost", + "expense", + "quantity", + "qty", + "revenue", + "sale", + "sales", + "sold", + "sum", + "total", + "value", + ) + ) + wants_average_metric = any( + term in normalized_query for term in ("average", "avg", "mean") + ) + + wants_monthly_count = ( + "monthly" in normalized_query + and any(term in normalized_query for term in ("count", "volume")) + and any(term in normalized_query for term in ("order", "orders")) + ) + if wants_monthly_count: + date_candidates = ( + ("InvDate", "InvoiceDate", "Invoice Date") + if "invdate" in compact_query or "invoice" in normalized_query + else ( + "OrdDate", + "OrderDate", + "NewOrderDate", + "InvDate", + "InvoiceDate", + "Date", + ) + ) + scored_tables: list[tuple[int, dict[str, Any], str]] = [] + for table in tables: + date_column = self._find_schema_column( + table, date_candidates, temporal=True + ) + if not date_column: + continue + table_name = str(table.get("name") or "") + score = 10 + if "sales" in table_name.lower() or "order" in table_name.lower(): + score += 5 + if self._find_schema_column( + table, ("OrdNo", "OrderNo", "OrderId", "InvoiceNo") + ): + score += 3 + scored_tables.append((score, table, date_column)) + + if scored_tables: + _, table, date_column = sorted( + scored_tables, key=lambda item: item[0], reverse=True + )[0] + table_name = table.get("name") + if table_name and date_column: + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + date_ref = ( + f"{table_ref}.{self._quote_sql_identifier(date_column)}" + ) + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "OrderCount" ' + f"FROM {table_ref}" + f"{self._build_date_filter(table_name, date_column, query)} " + f"GROUP BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref})" + ) + + dimension_candidates: list[tuple[str, ...]] = [] + if "salesperson" in normalized_query or "sales person" in normalized_query: + dimension_candidates.append( + ("SalesPerson", "Salesman", "Sales Rep", "SalesRep", "Rep", "Owner") + ) + if "business unit" in normalized_query or re.search(r"\bbu\b", normalized_query): + dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) + if "market" in normalized_query: + dimension_candidates.append(("Market", "MarketType", "MarketName", "Region", "Country")) + if "region" in normalized_query: + dimension_candidates.append(("Region", "Market", "Area", "Territory")) + if "currency" in normalized_query or "currencies" in normalized_query: + dimension_candidates.append( + ( + "Currency", + "CurrencyCode", + "Currency Code", + "Curr", + "CurrCode", + "MoneyCurrency", + "PaymentCurrency", + "FXCurrency", + ) + ) + if "country" in normalized_query or "countries" in normalized_query: + dimension_candidates.append(("Country", "CountryName", "Nation", "Market")) + if "division" in normalized_query: + dimension_candidates.append(("Division",)) + if ( + ( + "category" in normalized_query + or "categories" in normalized_query + or "prodcategory" in compact_query + or "productcategory" in compact_query + ) + and "product" in normalized_query + and "product type" not in normalized_query + and "prodtype" not in compact_query + and "producttype" not in compact_query + ): + dimension_candidates.append( + ( + "ProductCategory", + "Product Category", + "ProdCategory", + "Category", + "ProductType", + "Product Type", + "ProdType", + "ProdName", + "Product", + "ProductName", + ) + ) + elif ( + "product type" in normalized_query + or "prodtype" in normalized_query + or "producttype" in compact_query + or "prodtype" in compact_query + ): + dimension_candidates.append(("ProdType", "ProductType", "Product Type")) + elif "product" in normalized_query: + dimension_candidates.append( + ( + "ProdName", + "Product", + "ProductName", + "ProductDescription", + "Item", + "ItemName", + "ProdCode", + "ProductCode", + "PartNo", + "SKU", + ) + ) + if ( + "customer" in normalized_query + or "custname" in compact_query + or "custno" in compact_query + ): + dimension_candidates.append( + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + "Client", + "ClientName", + ) + ) + + measure_candidates = ( + "Qty", + "Quantity", + "QtySold", + "SoldQty", + "QuantitySold", + "UnitsSold", + "ItemQty", + "SalesQty", + "OrderQty", + "OrderQuantity", + "InvoiceQty", + "InvoiceQuantity", + ) if any(term in normalized_query for term in ("quantity", "qty")) else ( + "Sales", + "SalesValue", + "FXSalesValue", + "Revenue", + "NetSales", + "TotalSales", + "SalesAmount", + "SaleAmount", + "NewOrderValue", + "NewOrdersValue", + "InvoiceValue", + "InvoiceAmount", + "InvoiceAmt", + "OrderValue", + "TotalRevenue", + "Amount", + "Value", + "TotalOrderValue", + "Cost", + ) + if "invoice" in normalized_query: + measure_candidates = ( + "InvoiceValue", + "InvoiceAmount", + "InvoiceAmt", + "InvValue", + "InvAmount", + "SalesValue", + "FXSalesValue", + "Value", + "Amount", + ) + if wants_count_metric: + measure_candidates = () + wants_trend = ( + "trend" in normalized_query + or "line chart" in normalized_query + or "over time" in normalized_query + or "last 12 months" in normalized_query + or "by month" in normalized_query + or "monthly" in normalized_query + ) + wants_date_distribution = ( + any( + term in normalized_query + for term in ("distribution", "breakdown", "split") + ) + and any( + term in normalized_query + for term in ("date", "dates", "orddate", "order date", "order dates") + ) + ) + wants_order_count_metric = ( + any(term in normalized_query for term in ("order", "orders", "new order", "new orders")) + and any( + term in normalized_query + for term in ( + "count", + "counts", + "volume", + "how many", + "number of", + "monthly", + "over time", + "last 12 months", + "each", + "per ", + ) + ) + and not any( + term in normalized_query + for term in ("value", "amount", "revenue", "sales", "cost", "margin") + ) + ) + wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) + wants_top = wants_top or any( + term in normalized_query + for term in ( + "top ", + "best ", + "ranking", + "performance ranking", + ) + ) + wants_time_bucket = wants_trend or bool( + re.search(r"\bby\s+(?:month|year|quarter|date)\b", normalized_query) + ) + wants_detail_rows = ( + wants_top + and ("new order" in normalized_query or "orders" in normalized_query) + and any(term in normalized_query for term in ("including", "include")) + ) + mentions_date_column = any( + column_name in compact_query + for column_name in ( + "orddate", + "invdate", + "orderdate", + "invoicedate", + "createdat", + ) + ) + wants_date = ( + wants_trend + or wants_date_distribution + or mentions_date_column + or "this year" in normalized_query + or bool(re.search(r"\b20\d{2}\b", normalized_query)) + ) + wants_unique_customers_by_group = ( + any( + term in normalized_query + for term in ("unique customer", "unique customers") + ) + and "customer" in normalized_query + and "division" in normalized_query + and "market" in normalized_query + and any(term in normalized_query for term in ("highest", "top", "most")) + and any(term in normalized_query for term in ("each", "per ")) + ) + if wants_unique_customers_by_group: + selected = self._select_best_analytics_table( + tables, + [ + ("Market", "MarketType", "MarketName", "Region", "Country"), + ("Division",), + ( + "Customer", + "CustomerName", + "CustName", + "CustNo", + "CustomerNo", + "CustomerCode", + "Account", + "AccountName", + "Client", + "ClientName", + ), + ], + (), + wants_date=False, + allow_count_metric=True, + query=query, + ) + if selected: + table, dimensions, _measure, _date_column = selected + table_name = table.get("name") + if table_name and len(dimensions) >= 3: + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + market, division, customer = dimensions[:3] + market_ref = f"{table_ref}.{self._quote_sql_identifier(market)}" + division_ref = ( + f"{table_ref}.{self._quote_sql_identifier(division)}" + ) + customer_ref = ( + f"{table_ref}.{self._quote_sql_identifier(customer)}" + ) + where_clause = self._append_not_null_filters( + "", + [market_ref, division_ref, customer_ref], + ) + return ( + "WITH grouped_results AS (" + f"SELECT {market_ref} AS {self._quote_sql_identifier(market)}, " + f"{division_ref} AS {self._quote_sql_identifier(division)}, " + f"COUNT(DISTINCT {customer_ref}) AS \"UniqueCustomerCount\" " + f"FROM {table_ref}" + f"{where_clause} " + f"GROUP BY {market_ref}, {division_ref}" + "), ranked_results AS (" + f"SELECT {self._quote_sql_identifier(market)}, " + f"{self._quote_sql_identifier(division)}, " + "\"UniqueCustomerCount\", " + f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(market)} " + "ORDER BY \"UniqueCustomerCount\" DESC) AS \"rank\" " + "FROM grouped_results" + ") " + f"SELECT {self._quote_sql_identifier(market)}, " + f"{self._quote_sql_identifier(division)}, " + "\"UniqueCustomerCount\" " + "FROM ranked_results " + "WHERE \"rank\" = 1 " + "ORDER BY \"UniqueCustomerCount\" DESC" + ) + + if not dimension_candidates and wants_time_bucket: + selected = self._select_best_analytics_table( + tables, + [], + measure_candidates, + wants_date=True, + allow_count_metric=wants_count_metric, + query=query, + ) + if not selected: + return None + + table, _dimensions, measure, date_column = selected + table_name = table.get("name") + if not (table_name and date_column): + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + if wants_count_metric or not measure: + metric_expr = "COUNT(*)" + metric_alias = "RecordCount" + elif wants_average_metric: + metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Average{measure}" + else: + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Total{measure}" + + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " + f"FROM {table_ref}" + f"{self._build_date_filter(table_name, date_column, query)} " + f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" + ) + + if not dimension_candidates: + return None + + allow_count_metric = ( + wants_order_count_metric + or wants_count_metric + or ("performance" in normalized_query and wants_time_bucket) + ) + selected = self._select_best_analytics_table( + tables, + dimension_candidates, + measure_candidates, + wants_date=wants_date, + allow_count_metric=allow_count_metric, + query=query, + ) + if not selected: + return None + + table, dimensions, measure, date_column = selected + table_name = table.get("name") + if not table_name: + return None + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + dimension_refs = [ + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + for dimension in dimensions + ] + + if wants_detail_rows: + if not measure: + return None + metric_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" + limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) + limit = int(limit_match.group(1)) if limit_match else 20 + select_parts = [ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ] + select_parts.append( + f"{metric_ref} AS {self._quote_sql_identifier(measure)}" + ) + date_filter = ( + self._build_date_filter(table_name, date_column, query) + if date_column + else "" + ) + return ( + f"SELECT TOP {limit} {', '.join(select_parts)} " + f"FROM {table_ref}" + f"{self._append_not_null_filters(date_filter, dimension_refs)} " + f"ORDER BY {metric_ref} DESC" + ) + + if wants_date_distribution and date_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + select_parts = [ + f"DATEPART(YEAR, {date_ref}) AS \"year\"", + f"DATEPART(MONTH, {date_ref}) AS \"month\"", + *[ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ], + 'COUNT(*) AS "OrderCount"', + ] + group_parts = [ + f"DATEPART(YEAR, {date_ref})", + f"DATEPART(MONTH, {date_ref})", + *dimension_refs, + ] + where_clause = self._append_not_null_filters( + self._build_date_filter(table_name, date_column, query), + [date_ref, *dimension_refs], + ) + extra_conditions = self._build_schema_literal_filter_conditions( + query, + table, + table_ref, + ) + if extra_conditions: + where_clause = ( + f"{where_clause.rstrip()} AND {' AND '.join(extra_conditions)}" + if where_clause.strip() + else f" WHERE {' AND '.join(extra_conditions)}" + ) + return ( + f"SELECT {', '.join(select_parts)} FROM {table_ref}" + f"{where_clause} " + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}), COUNT(*) DESC" + ) + + wants_top_per_group = ( + len(dimensions) >= 2 + and any(term in normalized_query for term in ("highest", "top", "most")) + and any(term in normalized_query for term in ("each", "per ")) + ) + if wants_top_per_group: + partition_dimension = None + rank_dimension = None + if "market" in normalized_query: + partition_dimension = self._find_schema_column( + table, ("Market", "MarketType", "Region") + ) + if "region" in normalized_query and not partition_dimension: + partition_dimension = self._find_schema_column( + table, ("Region", "Market", "Area", "Territory") + ) + if "customer" in normalized_query: + rank_dimension = self._find_schema_column( + table, ("Customer", "CustName", "CustNo") + ) + if not partition_dimension: + partition_dimension = dimensions[0] + if not rank_dimension: + rank_dimension = next( + ( + dimension + for dimension in dimensions + if dimension != partition_dimension + ), + None, + ) + + if partition_dimension and rank_dimension: + partition_ref = ( + f"{table_ref}.{self._quote_sql_identifier(partition_dimension)}" + ) + rank_ref = f"{table_ref}.{self._quote_sql_identifier(rank_dimension)}" + if wants_order_count_metric: + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), + ) + metric_expr = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" + if order_column + else "COUNT(*)" + ) + metric_alias = "OrderCount" + else: + metric_expr = ( + f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + if measure + else "COUNT(*)" + ) + metric_alias = f"Total{measure}" if measure else "RecordCount" + where_clause = self._append_not_null_filters( + ( + self._build_date_filter(table_name, date_column, query) + if date_column + else "" + ), + [partition_ref, rank_ref], + ) + return ( + "WITH grouped_results AS (" + f"SELECT {partition_ref} AS {self._quote_sql_identifier(partition_dimension)}, " + f"{rank_ref} AS {self._quote_sql_identifier(rank_dimension)}, " + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " + f"FROM {table_ref}" + f"{where_clause} " + f"GROUP BY {partition_ref}, {rank_ref}" + "), ranked_results AS (" + f"SELECT {self._quote_sql_identifier(partition_dimension)}, " + f"{self._quote_sql_identifier(rank_dimension)}, " + f"{self._quote_sql_identifier(metric_alias)}, " + f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(partition_dimension)} " + f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC) AS \"rank\" " + "FROM grouped_results" + ") " + f"SELECT {self._quote_sql_identifier(partition_dimension)}, " + f"{self._quote_sql_identifier(rank_dimension)}, " + f"{self._quote_sql_identifier(metric_alias)} " + "FROM ranked_results " + "WHERE \"rank\" = 1 " + f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC" + ) + + if wants_trend and date_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + if wants_order_count_metric or wants_count_metric: + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), + ) + metric_expr = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" + if order_column + else "COUNT(*)" + ) + metric_alias = "OrderCount" + elif wants_average_metric: + metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Average{measure}" + else: + metric_expr = ( + f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + if measure + else "COUNT(*)" + ) + metric_alias = f"Total{measure}" if measure else "RecordCount" + select_parts = [ + f"DATEPART(YEAR, {date_ref}) AS \"year\"", + f"DATEPART(MONTH, {date_ref}) AS \"month\"", + *[ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ], + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}", + ] + group_parts = [ + f"DATEPART(YEAR, {date_ref})", + f"DATEPART(MONTH, {date_ref})", + *dimension_refs, + ] + return ( + f"SELECT {', '.join(select_parts)} FROM {table_ref}" + f"{self._append_not_null_filters(self._build_date_filter(table_name, date_column, query), dimension_refs)} " + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref})" + ) + + if wants_order_count_metric or wants_count_metric or not measure: + order_column = self._find_schema_column( + table, + ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), + ) + metric_expr = ( + f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" + if order_column and (wants_order_count_metric or wants_count_metric) + else "COUNT(*)" + ) + metric_alias = "OrderCount" if order_column else "RecordCount" + elif wants_average_metric: + metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Average{measure}" + else: + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + metric_alias = f"Total{measure}" + limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) + limit = int(limit_match.group(1)) if limit_match else 10 + top_clause = f"TOP {limit} " if wants_top else "" + sort_direction = ( + "ASC" + if any( + term in normalized_query + for term in ( + "losing", + "lowest", + "least", + "bottom", + "declining", + "underperforming", + "smallest", + ) + ) + else "DESC" + ) + date_filter = ( + self._build_date_filter(table_name, date_column, query) + if date_column + else "" + ) + select_parts = [ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ] + select_parts.append( + f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}" + ) + return ( + f"SELECT {top_clause}{', '.join(select_parts)} " + f"FROM {table_ref}" + f"{self._append_not_null_filters(date_filter, dimension_refs)} " + f"GROUP BY {', '.join(dimension_refs)} " + f"ORDER BY {metric_expr} {sort_direction}" + ) + + def _build_contribution_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not any(term in normalized_query for term in ("contribution", "pie chart")): + return None + + compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) + dimension_candidates: tuple[str, ...] | None = None + if ( + "product type" in normalized_query + or "prodtype" in normalized_query + or "producttype" in compact_query + or "prodtype" in compact_query + ): + dimension_candidates = ("ProdType", "ProductType", "Product Type") + elif "market" in normalized_query: + dimension_candidates = ("Market", "MarketType") + elif "division" in normalized_query: + dimension_candidates = ("Division",) + elif "customer" in normalized_query: + dimension_candidates = ("Customer", "CustName", "CustNo") + + if not dimension_candidates: + return None + + selected = self._select_best_analytics_table( + tables, + [dimension_candidates], + ( + "SalesValue", + "FXSalesValue", + "OrderValue", + "NewOrderValue", + "Revenue", + "Value", + "Amount", + ), + wants_date=False, + query=query, + ) + if not selected: + return None + + table, dimensions, measure, _date_column = selected + table_name = table.get("name") + if not (table_name and dimensions and measure): + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + dimension = dimensions[0] + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + return ( + f"SELECT {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " + f"{metric_expr} AS \"Total{measure}\" " + f"FROM {table_ref} " + f"GROUP BY {dimension_ref} " + f"ORDER BY {metric_expr} DESC" + ) + + def _build_generic_categorical_count_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + if re.search( + r"\b(?:first|top|sample|preview|show|list)\b.*\b(?:rows?|records?|data)\b", + normalized_query, + ): + return None + + wants_categorical_summary = any( + term in normalized_query + for term in ( + "bar chart", + "by ", + "chart", + "count", + "distribution", + "donut chart", + "frequency", + "group by", + "grouped by", + "most often", + "often", + "pie chart", + "restored", + "status", + "type", + "category", + ) + ) + if not wants_categorical_summary: + return None + + query_key = self._normalize_schema_token(query) + query_terms = self._query_schema_terms(query) + scored: list[tuple[int, dict[str, Any], str]] = [] + low_value_column_patterns = ( + "id", + "no", + "number", + "date", + "time", + "description", + "comment", + "note", + "remark", + ) + + for table in tables: + table_name = str(table.get("name") or "") + normalized_table = self._normalize_schema_token(table_name) + normalized_short_table = self._normalize_schema_token( + re.split(r"[.$]", table_name)[-1] + ) + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + column_type = str(column.get("type") or "") + if not column_name or not self._is_text_schema_type(column_type): + continue + normalized_column = self._normalize_schema_token(column_name) + if not normalized_column: + continue + + score = 0 + if normalized_table and normalized_table in query_key: + score += 120 + if normalized_short_table and normalized_short_table in query_key: + score += 100 + if normalized_column and normalized_column in query_key: + score += 180 + for term in query_terms: + if term == normalized_column: + score += 100 + elif term in normalized_column or normalized_column in term: + score += 45 + if term == normalized_table or term == normalized_short_table: + score += 40 + elif term in normalized_table or term in normalized_short_table: + score += 20 + if "status" in normalized_query and "status" in normalized_column: + score += 90 + if "category" in normalized_query and "category" in normalized_column: + score += 80 + if "type" in normalized_query and "type" in normalized_column: + score += 70 + if ( + "destination" in normalized_query + and "destination" in normalized_column + and ( + "database" in normalized_query + or "databases" in normalized_query + ) + and ( + "name" in normalized_column + or "phys" in normalized_column + or "db" in normalized_column + ) + ): + score += 140 + if any(pattern == normalized_column for pattern in low_value_column_patterns): + score -= 100 + elif any( + normalized_column.endswith(pattern) + for pattern in low_value_column_patterns + ): + score -= 35 + + if score > 0: + scored.append((score, table, column_name)) + + if not scored: + return None + + _, table, dimension = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + table_name = table.get("name") + if not table_name: + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + top_n = self._extract_requested_top_n(query, default_value=0) + top_clause = f"TOP {top_n} " if top_n else "" + return ( + f"SELECT {top_clause}{dimension_ref} AS {self._quote_sql_identifier(dimension)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f"ORDER BY COUNT(*) DESC" + ) + + def _build_order_invoice_conversion_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not ( + "conversion" in normalized_query + and "order" in normalized_query + and "invoice" in normalized_query + ): + return None + + scored: list[tuple[int, dict[str, Any], str, str, str | None]] = [] + for table in tables: + order_column = self._find_schema_column( + table, ("OrdNo", "OrderNo", "OrderNumber", "NewOrderNo") + ) + invoice_column = self._find_schema_column( + table, ("InvoiceNo", "InvNo", "InvoiceNumber") + ) + date_column = self._find_schema_column( + table, + ("OrdDate", "InvDate", "OrderDate", "InvoiceDate", "Date"), + temporal=True, + ) + if not (order_column and invoice_column): + continue + + score = 20 + if date_column: + score += 5 + if "sales" in str(table.get("name") or "").lower(): + score += 5 + scored.append((score, table, order_column, invoice_column, date_column)) + + if not scored: + return None + + _, table, order_column, invoice_column, date_column = sorted( + scored, key=lambda item: item[0], reverse=True + )[0] + table_name = table.get("name") + if not table_name: + return None + + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + order_ref = f"{table_ref}.{self._quote_sql_identifier(order_column)}" + invoice_ref = f"{table_ref}.{self._quote_sql_identifier(invoice_column)}" + date_column = date_column or order_column + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f"COUNT(DISTINCT {order_ref}) AS \"OrderCount\", " + f"COUNT(DISTINCT {invoice_ref}) AS \"InvoiceCount\", " + f"(COUNT(DISTINCT {invoice_ref}) * 100.0 / " + f"NULLIF(COUNT(DISTINCT {order_ref}), 0)) AS \"ConversionRate\" " + f"FROM {table_ref} " + f"WHERE {order_ref} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" + ) + + def _build_yoy_sales_change_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not any(term in normalized_query for term in ("yoy", "year over year")): + return None + + required_dimensions: list[tuple[str, ...]] = [] + if "customer" in normalized_query: + required_dimensions.append(("Customer", "CustName", "CustNo")) + if "product" in normalized_query: + required_dimensions.append( + ("ProdName", "Product", "ProductName", "Item", "ProdCode") + ) + if "market" in normalized_query: + required_dimensions.append(("Market", "MarketType")) + + if not required_dimensions: + return None + + selected = self._select_best_analytics_table( + tables, + required_dimensions, + ( + "SalesValue", + "FXSalesValue", + "OrderValue", + "NewOrderValue", + "Revenue", + "Value", + "Amount", + ), + wants_date=False, + query=query, + ) + if not selected: + return None + + table, dimensions, measure, date_column = selected + table_name = table.get("name") + if not (table_name and measure): + return None + + year_column = self._find_schema_column( + table, ("YearInd", "Year", "OrderYear", "InvoiceYear"), numeric=True + ) + table_name = str(table_name) + table_ref = self._quote_sql_identifier(table_name) + if year_column: + year_expr = f"{table_ref}.{self._quote_sql_identifier(year_column)}" + elif date_column: + year_expr = ( + f"DATEPART(YEAR, " + f"{table_ref}.{self._quote_sql_identifier(date_column)})" + ) + else: + return None + + metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" + dimension_refs = [ + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + for dimension in dimensions + ] + select_parts = [ + f"{year_expr} AS \"year\"", + *[ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ], + f"{metric_expr} AS \"Total{measure}\"", + ] + group_parts = [year_expr, *dimension_refs] + return ( + f"SELECT {', '.join(select_parts)} " + f"FROM {table_ref} " + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY {year_expr}, {metric_expr} DESC" + ) + + def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: + if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): + return max(1, min(int(match.group(1)), 100)) + if match := re.search( + r"\b(?:first|limit)\s+(\d+)\b", query or "", flags=re.IGNORECASE + ): + return max(1, min(int(match.group(1)), 100)) + if match := re.search(r"\b(\d+)\s+rows?\b", query or "", flags=re.IGNORECASE): + return max(1, min(int(match.group(1)), 100)) + return default_value + + def _build_manufacturing_throughput_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_throughput = "throughput" in normalized or ( + "repair" in normalized and "volume" in normalized + ) + wants_unit_breakdown = any( + term in normalized + for term in ( + "manufacturing unit", + "manufacturing units", + "business unit", + "business units", + "different unit", + "different units", + ) + ) + + if not (wants_throughput and wants_unit_breakdown): + return None + + tables = self._parse_schema_tables(table_ddls) + table = self._find_best_schema_table_for_query(query, tables) + if table: + unit_column = self._find_schema_column( + table, + ( + "BusinessUnit", + "business_unit", + "manufacturing_unit", + "manufacturingunit", + "unit", + "unit_name", + "BU", + "division", + ), + ) + if unit_column: + table_name = str(table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + unit_ref = f"{table_ref}.{self._quote_sql_identifier(unit_column)}" + timestamp_column = self._find_temporal_column_for_query(query, table) + + if timestamp_column and any( + term in normalized for term in ("trend", "monthly", "over time") + ): + timestamp_ref = ( + f"{table_ref}.{self._quote_sql_identifier(timestamp_column)}" + ) + return ( + f"SELECT {unit_ref} AS " + f"{self._quote_sql_identifier(unit_column)}, " + f"DATEPART(YEAR, {timestamp_ref}) AS \"year\", " + f"DATEPART(MONTH, {timestamp_ref}) AS \"month\", " + f'COUNT(*) AS "throughput" ' + f"FROM {table_ref} " + f"WHERE {unit_ref} IS NOT NULL " + f"AND {timestamp_ref} IS NOT NULL " + f"GROUP BY {unit_ref}, DATEPART(YEAR, {timestamp_ref}), " + f"DATEPART(MONTH, {timestamp_ref}) " + f"ORDER BY {unit_ref} ASC, DATEPART(YEAR, {timestamp_ref}) ASC, " + f"DATEPART(MONTH, {timestamp_ref}) ASC" + ) + + return ( + f"SELECT {unit_ref} AS " + f"{self._quote_sql_identifier(unit_column)}, " + f'COUNT(*) AS "throughput" ' + f"FROM {table_ref} " + f"WHERE {unit_ref} IS NOT NULL " + f"GROUP BY {unit_ref} " + f"ORDER BY COUNT(*) DESC" + ) + + has_debug_entries = self._schema_contains( + table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names + ) + has_business_unit = self._schema_contains( + table_ddls, r"\bBusinessUnit\b", table_names=table_names + ) + if not (has_debug_entries and has_business_unit): + return None + + timestamp_column = None + for candidate in ("DateIn", "FailedAt"): + if self._schema_contains( + table_ddls, rf"\b{candidate}\b", table_names=table_names + ): + timestamp_column = candidate + break + + if timestamp_column and any( + term in normalized for term in ("trend", "monthly", "over time") + ): + timestamp_expression = f'"dbo_DebugEntries"."{timestamp_column}"' + return ( + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' + f'DATEPART(YEAR, {timestamp_expression}) AS "year", ' + f'DATEPART(MONTH, {timestamp_expression}) AS "month", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries" ' + 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' + f'AND {timestamp_expression} IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' + f'DATEPART(YEAR, {timestamp_expression}), ' + f'DATEPART(MONTH, {timestamp_expression}) ' + 'ORDER BY "unit_name" ASC, "year" ASC, "month" ASC' + ) + + return ( + 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' + 'COUNT(*) AS "throughput" ' + 'FROM "dbo_DebugEntries" ' + 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' + 'GROUP BY "dbo_DebugEntries"."BusinessUnit" ' + 'ORDER BY "throughput" DESC' + ) + + def _build_audit_log_activity_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + if not ( + "audit" in normalized + and "log" in normalized + and any(term in normalized for term in ("activity", "over time", "trend")) + ): + return None + + table_name = "dbo_audit_log" + timestamp_column = "created_at" + if not self._schema_has_table_column( + table_ddls, + table_name, + timestamp_column, + table_names=table_names, + ): + return None + + dimension_column = None + condition_candidates = ( + "is_name_condition", + "name", + "action", + "entity_type", + ) + activity_candidates = ( + "action", + "entity_type", + "actor_name", + "actor_user_id", + "name", + ) + candidates = ( + condition_candidates + if "condition" in normalized + else activity_candidates + ) + for candidate in candidates: + if self._schema_has_table_column( + table_ddls, + table_name, + candidate, + table_names=table_names, + ): + dimension_column = candidate + break + + if not dimension_column: + return None + + timestamp_expression = f'"{table_name}"."{timestamp_column}"' + dimension_expression = f'"{table_name}"."{dimension_column}"' + return ( + f"SELECT DATEPART(YEAR, {timestamp_expression}) AS \"year\", " + f"DATEPART(MONTH, {timestamp_expression}) AS \"month\", " + f"{dimension_expression} AS \"{dimension_column}\", " + f'COUNT(*) AS "activity_count" ' + f'FROM "{table_name}" ' + f"WHERE {timestamp_expression} IS NOT NULL " + f"AND {dimension_expression} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {timestamp_expression}), " + f"DATEPART(MONTH, {timestamp_expression}), " + f"{dimension_expression} " + f"ORDER BY DATEPART(YEAR, {timestamp_expression}), " + f"DATEPART(MONTH, {timestamp_expression}), " + f'"activity_count" DESC' + ) + + def _build_repair_failure_count_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_failure_counts = ( + "failure" in normalized + and any( + term in normalized + for term in ( + "count", + "counts", + "category", + "code", + "grouped", + "common", + "most common", + "top", + ) + ) + and any(term in normalized for term in ("repair", "bar chart", "chart")) + ) + if not wants_failure_counts: + return None + + top_n = self._extract_requested_top_n(query) + + has_debug_fix_route = all( + ( + self._schema_has_table_column( + table_ddls, + "dbo_DebugEntries", + "DebugEntryId", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixLogs", + "DebugEntryId", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixLogs", + "FixId", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixes", + "Id", + table_names=table_names, + ), + self._schema_has_table_column( + table_ddls, + "dbo_DebugFixes", + "Description", + table_names=table_names, + ), + ) + ) + if has_debug_fix_route: + return ( + 'SELECT "dbo_DebugFixes"."Description" AS "failure_category", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_DebugEntries" ' + 'JOIN "dbo_DebugFixLogs" ' + 'ON "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" ' + 'JOIN "dbo_DebugFixes" ' + 'ON "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id" ' + 'WHERE "dbo_DebugFixes"."Description" IS NOT NULL ' + 'GROUP BY "dbo_DebugFixes"."Description" ' + 'ORDER BY "repair_count" DESC ' + f"LIMIT {top_n}" + ) + + has_debug_entries = self._schema_contains( + table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names + ) + has_failure_patterns = self._schema_contains( + table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names + ) + has_failure_sys = self._schema_contains( + table_ddls, r"\bFailureSys\b", table_names=table_names + ) + has_debug_entry_id = self._schema_contains( + table_ddls, r"\bDebugEntryId\b", table_names=table_names + ) + has_pattern_id = self._schema_contains( + table_ddls, r"\bid\b", table_names=table_names + ) + has_pattern_category = self._schema_contains( + table_ddls, r"\bcategory\b", table_names=table_names + ) + has_pattern_name = self._schema_contains( + table_ddls, r"\bname\b", table_names=table_names + ) + + if ( + has_debug_entries + and has_failure_patterns + and has_failure_sys + and has_debug_entry_id + and has_pattern_id + and (has_pattern_category or has_pattern_name) + ): + dimension_column = ( + "category" + if ("category" in normalized and has_pattern_category) + else ("name" if has_pattern_name else "category") + ) + return ( + f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' + f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' + f'FROM "dbo_DebugEntries" ' + f'JOIN "dbo_failure_patterns" ' + f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' + f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + has_repair_logs = self._schema_has_table_column( + table_ddls, + "dbo_repair_logs", + "failure_code", + table_names=table_names, + ) + if has_repair_logs: + return ( + f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_repair_logs" ' + f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' + f'GROUP BY "dbo_repair_logs"."failure_code" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + return None + + def _build_repair_sla_compliance_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_sla = "sla" in normalized and any( + term in normalized + for term in ("compliance", "dashboard", "chart", "repair", "repairs") + ) + if not wants_sla: + return None + + has_repair_status = self._schema_has_table_column( + table_ddls, + "dbo_repair_logs", + "status", + table_names=table_names, + ) + if has_repair_status: + return ( + 'SELECT "dbo_repair_logs"."status" AS "sla_status", ' + 'COUNT(*) AS "repair_count" ' + 'FROM "dbo_repair_logs" ' + 'WHERE "dbo_repair_logs"."status" IS NOT NULL ' + 'GROUP BY "dbo_repair_logs"."status" ' + 'ORDER BY "repair_count" DESC' + ) + + return None + + def _build_monthly_repair_volume_sql( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + wants_monthly_repairs = ( + "repair" in normalized + and any( + term in normalized + for term in ("monthly", "last 12 months", "trend", "volume") + ) + ) + if not wants_monthly_repairs: + return None + + tables = self._parse_schema_tables(table_ddls) + scored_tables: list[tuple[int, dict[str, Any], str]] = [] + for table in tables: + table_name = str(table.get("name") or "") + if table_names and table_name not in table_names: + continue + + date_column = self._find_schema_column( + table, + ( + "created_at", + "createdAt", + "created", + "DateIn", + "Date", + "repair_date", + "RepairDate", + "opened_at", + "started_at", + ), + temporal=True, + ) + if not date_column: + date_column = self._find_any_temporal_schema_column(table) + if not date_column: + continue + + normalized_table = self._normalize_schema_token(table_name) + score = 0 + if "repair" in normalized_table: + score += 30 + if "debugentries" in normalized_table or "debugentry" in normalized_table: + score += 25 + if "log" in normalized_table: + score += 10 + if self._find_schema_column( + table, + ("DebugEntryId", "RepairId", "repair_id", "id"), + ): + score += 5 + scored_tables.append((score, table, date_column)) + + if not scored_tables: + return None + + _score, table, date_column = sorted( + scored_tables, + key=lambda item: item[0], + reverse=True, + )[0] + table_name = str(table.get("name") or "") + if not table_name: + return None + + table_ref = self._quote_sql_identifier(table_name) + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + return ( + f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "repair_count" ' + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"GROUP BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " + f"DATEPART(MONTH, {date_ref}) ASC" + ) + + def _is_direct_heuristic_sql_query(self, query: str) -> bool: + return False + + def _build_heuristic_text_to_sql_fallback( + self, + query: str, + table_ddls: list[str], + table_names: Optional[list[str]] = None, + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + if schema_grounded_sql := self._build_schema_grounded_sales_sql( + query, table_ddls + ): + return schema_grounded_sql + + if throughput_sql := self._build_manufacturing_throughput_sql( + query, table_ddls, table_names=table_names + ): + return throughput_sql + + if repair_failure_count_sql := self._build_repair_failure_count_sql( + query, table_ddls, table_names=table_names + ): + return repair_failure_count_sql + + if repair_sla_sql := self._build_repair_sla_compliance_sql( + query, table_ddls, table_names=table_names + ): + return repair_sla_sql + + if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( + query, table_ddls, table_names=table_names + ): + return monthly_repair_volume_sql + + wants_chart = any( + term in normalized for term in ("chart", "bar chart", "line chart", "graph") + ) + wants_failure_counts = any( + term in normalized + for term in ( + "failure", + "failure category", + "failure code", + "common pcb failures", + "common failures", + "most common", + "top 10", + "top ten", + ) + ) + wants_monthly_repairs = ( + "repair" in normalized + and any( + term in normalized + for term in ("monthly", "last 12 months", "trend", "volume") + ) + ) + + if wants_failure_counts and wants_chart: + top_n = self._extract_requested_top_n(query) + has_pattern_failure_sys = self._schema_contains( + table_ddls, r"\bFailuresys\b", table_names=table_names + ) + has_pattern_occurrences = self._schema_contains( + table_ddls, r"\boccurrences\b", table_names=table_names + ) + has_debug_entries = self._schema_contains( + table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names + ) + has_failure_patterns = self._schema_contains( + table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names + ) + has_failure_sys = self._schema_contains( + table_ddls, r"\bFailureSys\b", table_names=table_names + ) + has_debug_entry_id = self._schema_contains( + table_ddls, r"\bDebugEntryId\b", table_names=table_names + ) + has_pattern_id = self._schema_contains( + table_ddls, r"\bid\b", table_names=table_names + ) + has_pattern_category = self._schema_contains( + table_ddls, r"\bcategory\b", table_names=table_names + ) + has_pattern_name = self._schema_contains( + table_ddls, r"\bname\b", table_names=table_names + ) + + if has_failure_patterns and has_pattern_failure_sys and has_pattern_occurrences: + return ( + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'"dbo_failure_patterns"."occurrences" AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'AND "dbo_failure_patterns"."occurrences" IS NOT NULL ' + f'ORDER BY "dbo_failure_patterns"."occurrences" DESC ' + f'LIMIT {top_n}' + ) + + if has_failure_patterns and has_pattern_failure_sys: + return ( + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."Failuresys" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + if ( + has_debug_entries + and has_failure_patterns + and has_failure_sys + and has_debug_entry_id + and has_pattern_id + ): + dimension_column = ( + "category" + if ("category" in normalized and has_pattern_category) + else ("name" if has_pattern_name else "category") + ) + if dimension_column == "category" and not has_pattern_category: + dimension_column = "name" + + return ( + f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' + f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' + f'FROM "dbo_DebugEntries" ' + f'JOIN "dbo_failure_patterns" ' + f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' + f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + has_repair_logs = self._schema_contains( + table_ddls, r"\bdbo_repair_logs\b", table_names=table_names + ) + has_failure_code = self._schema_contains( + table_ddls, r"\bfailure_code\b", table_names=table_names + ) + if has_repair_logs and has_failure_code: + return ( + f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_repair_logs" ' + f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' + f'GROUP BY "dbo_repair_logs"."failure_code" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + if wants_failure_counts and wants_chart: + top_n = self._extract_requested_top_n(query) + return ( + f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' + f'COUNT(*) AS "repair_count" ' + f'FROM "dbo_failure_patterns" ' + f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' + f'GROUP BY "dbo_failure_patterns"."Failuresys" ' + f'ORDER BY "repair_count" DESC ' + f'LIMIT {top_n}' + ) + + return None + + def _is_schema_grounded_query( + self, query: str, db_schemas: Optional[list[str]] = None + ) -> bool: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return False + + explicit_schema_terms = ( + "table", + "column", + "schema", + "dataset", + "dbo.", + "select ", + " from ", + " join ", + " where ", + " group by ", + " order by ", + ) + if any(term in normalized for term in explicit_schema_terms): + return True + + identifier_tokens = re.findall(r"[a-zA-Z_][a-zA-Z0-9_\.]*", normalized) + if any("." in token for token in identifier_tokens): + return True + + for schema in db_schemas or []: + schema_text = schema.lower() + table_matches = re.findall( + r"create\s+table\s+([a-zA-Z0-9_\.\"]+)", schema_text + ) + column_matches = re.findall(r"\n\s*\"?([a-zA-Z_][a-zA-Z0-9_]*)\"?\s+", schema_text) + candidates = { + token.strip('"') + for token in table_matches + column_matches + if token and len(token.strip('"')) > 2 + } + if any(candidate in normalized for candidate in candidates): + return True + + return False + def _build_schema_grounded_operational_sql( + self, query: str, tables: list[dict[str, Any]] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized_query: + return None + + operational_terms = ( + "ticket", + "repair", + "failure", + "pcb", + "component", + "board", + "throughput", + "manufacturing", + "unit", + "knowledge", + "article", + "source", + "category", + "priority", + "status", + "open", + "closed", + "aging", + "workflow", + "time", + "duration", + "elapsed", + "estimated", + "volume", + "count", + ) + if not any(term in normalized_query for term in operational_terms): + return None + + scored_tables: list[tuple[int, dict[str, Any]]] = [] + for table in tables: + table_name = str(table.get("name") or "") + normalized_table = table_name.lower() + score = 0 + if any( + token in normalized_table + for token in ("ticket", "repair", "debug", "knowledge", "article") + ): + score += 10 + if "ticket" in normalized_query and "ticket" in normalized_table: + score += 8 + if "knowledge" in normalized_query and "knowledge" in normalized_table: + score += 8 + if "article" in normalized_query and "article" in normalized_table: + score += 5 + if "repair" in normalized_query and "repair" in normalized_table: + score += 5 + if "failure" in normalized_query and "failure" in normalized_table: + score += 8 + if any( + term in normalized_query + for term in ("business unit", "business units", "unit", "units") + ) and self._find_schema_column( + table, + ( + "BusinessUnit", + "Business_Unit", + "Business Unit", + "manufacturing_unit", + "ManufacturingUnit", + "unit", + "BU", + "Division", + ), + ): + score += 15 + if any( + term in normalized_query + for term in ("product line", "product family", "product", "products") + ) and self._find_schema_column( + table, + ( + "Product_Family", + "ProductFamily", + "Product Family", + "ProductLine", + "Product_Line", + "Product", + "ProdType", + "Material", + ), + ): + score += 15 + if any(term in normalized_query for term in ("error", "failure")) and any( + self._find_schema_column(table, candidates) + for candidates in ( + ("failure_code", "FailureSys", "failure", "failure_type"), + ("category", "name", "description"), + ) + ): + score += 6 + if self._find_schema_column( + table, + ("created_at", "updated_at", "DateIn", "DateOut", "created", "date"), + temporal=True, + ): + score += 3 + if score: + scored_tables.append((score, table)) + + if not scored_tables: + return None + + table = sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] + table_name = str(table.get("name") or "") + if not table_name: + return None + + table_ref = self._quote_sql_identifier(table_name) + date_column = self._find_schema_column( + table, + ( + "created_at", + "created", + "DateIn", + "RepairDate", + "updated_at", + "DateOut", + "updated", + "date", + ), + temporal=True, + ) + + dimension_candidates: list[tuple[str, ...]] = [] + if any(term in normalized_query for term in ("failure", "failures", "error")): + dimension_candidates.append( + ( + "failure_code", + "FailureSys", + "failure", + "failure_type", + "failure_category", + "category", + "name", + "description", + ) + ) + if "manufacturing" in normalized_query or "unit" in normalized_query: + dimension_candidates.append( + ( + "BusinessUnit", + "Business_Unit", + "Business Unit", + "manufacturing_unit", + "manufacturing unit", + "ManufacturingUnit", + "unit", + "BU", + "Division", + "assignee_user_id", + "created_by_user_id", + "org_id", + "status", + ) + ) + if any( + term in normalized_query + for term in ("product line", "product family", "product", "products") + ): + dimension_candidates.append( + ( + "Product_Family", + "ProductFamily", + "Product Family", + "ProductLine", + "Product_Line", + "Product", + "ProdType", + "Material", + ) + ) + if "component" in normalized_query: + dimension_candidates.append( + ("component", "component_type", "board_type", "title", "status") + ) + if "board" in normalized_query: + dimension_candidates.append(("board_type", "board", "title", "status")) + if "category" in normalized_query: + dimension_candidates.append(("category", "subcategory", "status", "priority")) + if "source" in normalized_query: + dimension_candidates.append(("source", "author", "category", "status")) + if "priority" in normalized_query: + dimension_candidates.append(("priority", "status")) + if ( + "status" in normalized_query + or "open" in normalized_query + or "closed" in normalized_query + ): + dimension_candidates.append(("status", "priority")) + if "assignee" in normalized_query: + dimension_candidates.append(("assignee_user_id", "created_by_user_id")) + if "workflow" in normalized_query: + dimension_candidates.append(("status", "priority", "assignee_user_id")) + + dimensions: list[str] = [] + for candidates in dimension_candidates: + dimension = self._find_schema_column(table, candidates) + if dimension and dimension not in dimensions: + dimensions.append(dimension) + + if not dimensions: + fallback_dimension = self._find_first_schema_column( + table, + ( + "status", + "priority", + "category", + "subcategory", + "author", + "assignee_user_id", + "created_by_user_id", + "org_id", + "title", + ), + ) + if fallback_dimension: + dimensions.append(fallback_dimension) + + wants_trend = any( + term in normalized_query + for term in ("trend", "monthly", "month", "line chart", "over time") + ) + wants_elapsed_time = ( + not wants_trend + and any( + term in normalized_query + for term in ( + "duration", + "elapsed", + "turnaround", + "estimated", + "time spent", + "time taken", + ) + ) + ) + wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) + limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) + limit = int(limit_match.group(1)) if limit_match else 10 + + if wants_elapsed_time: + temporal_columns = [ + str(column.get("name") or "") + for column in table.get("columns", []) + if column.get("name") + and self._is_temporal_schema_type(str(column.get("type") or "")) + ] + start_column = self._find_schema_column( + table, + ( + "created_at", + "created", + "DateIn", + "execution_date", + "opened_at", + "started_at", + "start_date", + "begin_date", + ), + temporal=True, + ) + end_column = self._find_schema_column( + table, + ( + "updated_at", + "updated", + "DateOut", + "closed_at", + "resolved_at", + "completed_at", + "finished_at", + "end_date", + ), + temporal=True, + ) + if not start_column and temporal_columns: + start_column = temporal_columns[0] + if not end_column: + for candidate in temporal_columns: + if candidate.lower() != str(start_column or "").lower(): + end_column = candidate + break + if start_column and end_column: + start_ref = f"{table_ref}.{self._quote_sql_identifier(start_column)}" + end_ref = f"{table_ref}.{self._quote_sql_identifier(end_column)}" + duration_expr = f"DATEDIFF('second', {start_ref}, {end_ref})" + if not dimensions: + fallback_dimension = self._find_first_schema_column( + table, + ( + "status", + "priority", + "assignee_user_id", + "created_by_user_id", + "org_id", + ), + ) + if fallback_dimension: + dimensions.append(fallback_dimension) + if dimensions: + dimension = dimensions[0] + dimension_ref = ( + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + ) + dimension_alias = ( + "workflow" if "workflow" in normalized_query else dimension + ) + return ( + f"SELECT {dimension_ref} AS " + f"{self._quote_sql_identifier(dimension_alias)}, " + f'SUM({duration_expr}) AS "total_time_seconds" ' + f"FROM {table_ref} " + f"WHERE {start_ref} IS NOT NULL " + f"AND {end_ref} IS NOT NULL " + f"AND {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f'ORDER BY "total_time_seconds" DESC' + ) + return ( + f'SELECT SUM({duration_expr}) AS "total_time_seconds" ' + f"FROM {table_ref} " + f"WHERE {start_ref} IS NOT NULL " + f"AND {end_ref} IS NOT NULL" + ) + + if wants_trend and date_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" + select_parts = [ + f"DATEPART(YEAR, {date_ref}) AS \"year\"", + f"DATEPART(MONTH, {date_ref}) AS \"month\"", + ] + group_parts = [ + f"DATEPART(YEAR, {date_ref})", + f"DATEPART(MONTH, {date_ref})", + ] + for dimension in dimensions[:2]: + dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" + select_parts.append( + f"{dimension_ref} AS {self._quote_sql_identifier(dimension)}" + ) + group_parts.append(dimension_ref) + select_parts.append('COUNT(*) AS "RecordCount"') + return ( + f"SELECT {', '.join(select_parts)} " + f"FROM {table_ref} " + f"GROUP BY {', '.join(group_parts)} " + f"ORDER BY DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref})" + ) + + if dimensions: + top_clause = f"TOP {limit} " if wants_top else "" + dimension_refs = [ + f"{table_ref}.{self._quote_sql_identifier(dimension)}" + for dimension in dimensions[:2] + ] + select_parts = [ + f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" + for index, dimension_ref in enumerate(dimension_refs) + ] + select_parts.append('COUNT(*) AS "RecordCount"') + return ( + f"SELECT {top_clause}{', '.join(select_parts)} " + f"FROM {table_ref} " + f"GROUP BY {', '.join(dimension_refs)} " + f"ORDER BY COUNT(*) DESC" + ) + + return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' + + def _build_pcb_direct_question_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) + if not normalized: + return None + + tables = self._parse_schema_tables(table_ddls) + if not tables: + return None + + repair_table = next( + ( + table + for table in tables + if str(table.get("name") or "").lower() == "dbo_repair_logs" + ), + None, + ) + ticket_label_table = next( + ( + table + for table in tables + if str(table.get("name") or "").lower() == "dbo_ticket_labels" + ), + None, + ) + limit = self._extract_requested_top_n(query, default_value=10) + + if ticket_label_table and "ticket" in normalized and "label" in normalized: + label_column = self._find_first_schema_column( + ticket_label_table, + ("name", "label", "title", "value", "id"), + ) + if label_column: + table_name = str(ticket_label_table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + label_ref = f"{table_ref}.{self._quote_sql_identifier(label_column)}" + return ( + f"SELECT TOP {limit} {label_ref} AS " + f"{self._quote_sql_identifier(label_column)}, " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {label_ref} IS NOT NULL " + f"GROUP BY {label_ref} " + f"ORDER BY COUNT(*) DESC" + ) + + if not repair_table: + return None + + table_name = str(repair_table.get("name") or "") + table_ref = self._quote_sql_identifier(table_name) + board_model_column = self._find_schema_column( + repair_table, ("board_model", "boardModel", "board model", "product") + ) + failure_code_column = self._find_schema_column( + repair_table, ("failure_code", "failureCode", "failure code", "failure") + ) + created_at_column = self._find_schema_column( + repair_table, + ("created_at", "createdAt", "created", "date_received", "dateReceived"), + temporal=True, + ) + priority_column = self._find_schema_column(repair_table, ("priority",)) + status_column = self._find_schema_column(repair_table, ("status",)) + id_column = self._find_schema_column(repair_table, ("id", "repair_id")) + + asks_board_model_distribution = ( + "board model" in normalized + and any(term in normalized for term in ("distribution", "over time", "trend")) + ) + if asks_board_model_distribution and board_model_column and created_at_column: + board_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" + date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" + return ( + f"SELECT {board_ref} AS " + f"{self._quote_sql_identifier(board_model_column)}, " + f"DATEPART(YEAR, {date_ref}) AS \"year\", " + f"DATEPART(MONTH, {date_ref}) AS \"month\", " + f'COUNT(*) AS "RecordCount" ' + f"FROM {table_ref} " + f"WHERE {board_ref} IS NOT NULL " + f"AND {date_ref} IS NOT NULL " + f"GROUP BY {board_ref}, DATEPART(YEAR, {date_ref}), " + f"DATEPART(MONTH, {date_ref}) " + f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " + f"DATEPART(MONTH, {date_ref}) ASC, {board_ref} ASC" + ) + + asks_recurring_failures_by_product = ( + "recurring" in normalized + and "failure" in normalized + and ("product" in normalized or "pcb" in normalized) + ) + if ( + asks_recurring_failures_by_product + and board_model_column + and failure_code_column + ): + product_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" + failure_ref = f"{table_ref}.{self._quote_sql_identifier(failure_code_column)}" + return ( + f"SELECT {product_ref} AS " + f"{self._quote_sql_identifier(board_model_column)}, " + f"{failure_ref} AS " + f"{self._quote_sql_identifier(failure_code_column)}, " + f'COUNT(*) AS "failure_count" ' + f"FROM {table_ref} " + f"WHERE {product_ref} IS NOT NULL " + f"AND {failure_ref} IS NOT NULL " + f"GROUP BY {product_ref}, {failure_ref} " + f'ORDER BY "failure_count" DESC' + ) + + asks_highest_priority_repairs = ( + "repair" in normalized + and "priority" in normalized + and any(term in normalized for term in ("highest", "top", "high priority")) + ) + if asks_highest_priority_repairs and priority_column: + priority_ref = f"{table_ref}.{self._quote_sql_identifier(priority_column)}" + select_refs = [] + for column in ( + id_column, + board_model_column, + failure_code_column, + status_column, + priority_column, + created_at_column, + ): + if column and column not in select_refs: + select_refs.append(column) + select_sql = ", ".join( + f"{table_ref}.{self._quote_sql_identifier(column)} AS " + f"{self._quote_sql_identifier(column)}" + for column in select_refs + ) + return ( + f"SELECT TOP {limit} {select_sql} " + f"FROM {table_ref} " + f"WHERE {priority_ref} IS NOT NULL " + f"ORDER BY CASE LOWER({priority_ref}) " + f"WHEN 'critical' THEN 1 " + f"WHEN 'high' THEN 2 " + f"WHEN 'medium' THEN 3 " + f"WHEN 'low' THEN 4 " + f"ELSE 5 END" + ) + + asks_repair_ticket_distribution = ( + "repair" in normalized + and "ticket" in normalized + and ( + "distribution" in normalized + or "again distribution" in normalized + or "aging distribution" in normalized + ) + ) + if asks_repair_ticket_distribution: + if "aging" in normalized and created_at_column: + date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" + age_bucket = ( + f"CASE " + f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 7 THEN '0-7 days' " + f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 30 THEN '8-30 days' " + f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 90 THEN '31-90 days' " + f"ELSE '90+ days' END" + ) + return ( + f'SELECT {age_bucket} AS "age_bucket", ' + f'COUNT(*) AS "ticket_count" ' + f"FROM {table_ref} " + f"WHERE {date_ref} IS NOT NULL " + f"GROUP BY {age_bucket} " + f'ORDER BY "ticket_count" DESC' + ) + distribution_column = status_column or priority_column or failure_code_column + if distribution_column: + dimension_ref = ( + f"{table_ref}.{self._quote_sql_identifier(distribution_column)}" + ) + return ( + f"SELECT {dimension_ref} AS " + f"{self._quote_sql_identifier(distribution_column)}, " + f'COUNT(*) AS "ticket_count" ' + f"FROM {table_ref} " + f"WHERE {dimension_ref} IS NOT NULL " + f"GROUP BY {dimension_ref} " + f'ORDER BY "ticket_count" DESC' + ) + + return None + + def _get_unqueryable_metric_message( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + normalized_schema = re.sub( + r"\s+", + " ", + " ".join(ddl for ddl in table_ddls if isinstance(ddl, str)).lower(), + ) + schema_column_names = self._extract_schema_column_names(table_ddls) + + if not normalized_query: + return None + + if "throughput" in normalized_query and any( + term in normalized_query for term in ("manufacturing", "unit", "units") + ): + unit_field_patterns = ( + r"\bbusiness[_ ]?unit\b", + r"\bmanufacturing[_ ]?unit\b", + r"\bunit[_ ]?name\b", + r"\bunit\b", + r"\bbu\b", + r"\bdivision\b", + ) + has_unit_field = any( + re.search(pattern, column_name) + for pattern in unit_field_patterns + for column_name in schema_column_names + ) + has_temporal_field = any( + self._is_temporal_schema_type(str(column.get("type") or "")) + for table in self._parse_schema_tables(table_ddls) + for column in table.get("columns", []) + ) + if not has_unit_field: + return ( + "The active datasource does not expose a manufacturing unit, " + "business unit, unit, BU, or division column. I cannot build " + "throughput trends across manufacturing units without a " + "queryable unit field." + ) + if "trend" in normalized_query and not has_temporal_field: + return ( + "The active datasource does not expose a queryable date or " + "timestamp column. I cannot build a throughput trend without " + "a first-class temporal field." + ) + + if any( + term in normalized_query + for term in ( + "monthly", + "trend", + "turnaround", + "time", + "duration", + "elapsed", + "latest", + "recent", + "newest", + "last records", + ) + ): + has_temporal_field = any( + self._is_temporal_schema_type(str(column.get("type") or "")) + for table in self._parse_schema_tables(table_ddls) + for column in table.get("columns", []) + ) + if not has_temporal_field: + return ( + "The active datasource does not expose a queryable date or " + "timestamp column. I cannot build a time-based analysis " + "without a first-class temporal field." + ) + + repair_cost_terms = ( + "repair cost", + "repair_cost", + "repaircost", + "cost", + "cost impact", + "cost_impact", + ) + if any(term in normalized_query for term in repair_cost_terms): + cost_field_patterns = ( + r"\brepair[_ ]?cost\b", + r"\bcost[_ ]?impact\b", + r"\bcost[_ ]?amount\b", + r"\btotal[_ ]?cost\b", + r"\bunit[_ ]?cost\b", + r"\bcost\b", + r"\bamount\b", + ) + has_cost_field = any( + re.search(pattern, column_name) + for pattern in cost_field_patterns + for column_name in schema_column_names + ) + + if not has_cost_field: + return ( + "The schema does not expose repair cost as a queryable " + "column. The MSSQL Wren/Ibis runtime cannot extract cost " + "from generic JSON/text fields such as data. Add repair " + "cost as a first-class column or calculated field, then " + "ask again." + ) + + first_pass_yield_terms = ( + "first pass yield", + "first-pass yield", + "first_pass_yield", + "fpy", + ) + if not any(term in normalized_query for term in first_pass_yield_terms): + return None + + required_field_patterns = ( + r"\bfirst[_ ]?pass[_ ]?yield\b", + r"\bfpy\b", + r"\battempt\b", + r"\battempt[_ ]?number\b", + r"\bfirst[_ ]?attempt\b", + r"\bpass[_ ]?fail\b", + r"\byield\b", + ) + has_required_field = any( + re.search(pattern, normalized_schema) + for pattern in required_field_patterns + ) + + if has_required_field: + return None + + return ( + "The schema does not expose first-pass yield, attempt number, " + "first-attempt result, or pass/fail fields as queryable columns. " + "I cannot calculate First Pass Yield from only generic JSON/text " + "fields such as data. Add those fields as first-class columns or " + "calculated fields, then ask again." + ) + + def _build_schema_grounded_sales_sql( + self, query: str, table_ddls: list[str] + ) -> str | None: + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) + normalized_schema = "\n".join( + ddl for ddl in table_ddls or [] if isinstance(ddl, str) + ).lower() + if not normalized_query or not normalized_schema: + return None + + schema_key = self._normalize_schema_token(normalized_schema) + is_sales_specific_query = any( + term in normalized_query + for term in ("sale", "sales", "salesperson", "sales person") + ) + if is_sales_specific_query and "salesvalue" not in schema_key: + return None + + return self._build_schema_grounded_analytics_sql(query, table_ddls) + + async def _run_with_timeout( + self, + label: str, + coroutine, + timeout_seconds: Optional[int] = None, + ): + timeout = timeout_seconds or self._pipeline_timeout_seconds + try: + return await asyncio.wait_for( + coroutine, + timeout=timeout, + ) + except TimeoutError as exc: + raise TimeoutError(f"{label} timed out after {timeout} seconds") from exc + + def _should_retry_selected_schema_after_retrieval_timeout( + self, retrieval_table_names: Optional[list[str]] + ) -> bool: + return bool(retrieval_table_names) + + def _forced_explicit_table_names( + self, table_names: list[str], *, source: str = "request" + ) -> list[str]: + if not table_names: + return [] + if len(table_names) <= MAX_FORCED_EXPLICIT_TABLES: + return table_names + + logger.info( + "Treating broad %s explicit_tables list as retrieval candidates, not a forced schema scope: %s", + source, + table_names, + ) + return [] + + def _build_greeting_response(self, query: str) -> str: + return ( + f"Hi. I can help with questions about your active datasource and Wren AI.\n\n" + f"Try a data question like:\n" + f"- Show monthly trends for the last 12 months\n" + f"- Compare totals by category\n" + f"- Which records occur most often?\n\n" + f"If you want, ask a database question directly instead of `{query}`." + ) + + def _extract_pipeline_reply(self, result: dict, key: str) -> str: + payload = result.get(key) + if isinstance(payload, tuple): + payload = payload[0] + + if isinstance(payload, dict): + replies = payload.get("replies") or [] + if replies and isinstance(replies[0], str): + return replies[0] + + return "" + + def _extract_retrieval_documents(self, retrieval_result: dict) -> list[dict]: + construct_result = retrieval_result.get("construct_retrieval_results", {}) + documents = construct_result.get("retrieval_results", []) + if not isinstance(documents, list): + logger.warning("Schema retrieval returned invalid document payload") + return [] + + valid_documents = [] + for document in documents: + if not isinstance(document, dict): + logger.warning("Ignoring malformed retrieval document: %s", document) + continue + if not document.get("table_name") and not document.get("table_ddl"): + logger.warning("Ignoring retrieval document without table metadata") + continue + valid_documents.append(document) + + return valid_documents + + def _extract_retrieval_metadata( + self, retrieval_result: dict + ) -> tuple[list[dict], list[str], list[str]]: + documents = self._extract_retrieval_documents(retrieval_result) + return documents, *self._metadata_from_documents(documents) + + def _metadata_from_documents( + self, documents: list[dict] + ) -> tuple[list[str], list[str]]: + table_names = [ + table_name + for document in documents + if isinstance(table_name := document.get("table_name"), str) + and table_name.strip() + ] + table_ddls = [ + table_ddl + for document in documents + if isinstance(table_ddl := document.get("table_ddl"), str) + and table_ddl.strip() + ] + return table_names, table_ddls + + async def _complete_sql_generation_context( + self, + *, + query: str, + project_id: Optional[str], + documents: list[dict], + table_names: list[str], + table_ddls: list[str], + ) -> tuple[list[dict], list[str], list[str], dict]: + if not table_names or "db_schema_retrieval" not in self._pipelines: + return documents, table_names, table_ddls, {} + + selected_table_names = list(dict.fromkeys(table_names)) + try: + retrieval_result = await self._run_with_timeout( + "Complete selected schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=query, + tables=selected_table_names, + project_id=project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min(self._schema_retrieval_timeout_seconds, 30), + ) + except Exception as error: + logger.warning( + "Complete selected schema retrieval failed; using existing retrieval context. project_id=%s tables=%s error=%s", + project_id, + selected_table_names, + error, + ) + return documents, table_names, table_ddls, {} + + complete_documents, complete_table_names, complete_table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if not complete_documents: + logger.warning( + "Complete selected schema retrieval returned no documents; using existing retrieval context. project_id=%s tables=%s", + project_id, + selected_table_names, + ) + return documents, table_names, table_ddls, {} + + logger.info( + "Completed SQL generation context with full schemas for project_id %s tables=%s", + project_id, + complete_table_names, + ) + return ( + complete_documents, + complete_table_names, + complete_table_ddls, + retrieval_result.get("construct_retrieval_results", {}), + ) + + def _is_visualization_request(self, query: str) -> bool: + normalized = (query or "").lower() + return bool( + re.search( + r"\b(?:chart|graph|plot|visuali[sz]e|dashboard|bar|line|pie|donut|" + r"scatter|histogram|heatmap|trend|trends|distribution)\b", + normalized, + ) + ) + + def _get_metadata_question_kind(self, query: str) -> str | None: + normalized = re.sub(r"\s+", " ", (query or "").lower()).strip() + if not normalized: + return None + + if self._is_visualization_request(normalized): + return None + + if re.search(r"\b(?:row|rows|record|records)\s+count\b", normalized): + return None + + relationship_patterns = ( + r"\b(?:relationships?|relations?|joins?|foreign keys?|primary keys?)\b", + r"\b(?:how|what|which|show|list|describe)\b.*\b(?:tables?|models?)\b.*\b(?:connected|related|joined)\b", + ) + if any(re.search(pattern, normalized) for pattern in relationship_patterns): + return "relationships" + + table_count_patterns = ( + r"\b(?:how many|count|number of)\b.*\b(?:tables?|models?)\b", + r"\b(?:tables?|models?)\b.*\b(?:count|number)\b", + ) + if any(re.search(pattern, normalized) for pattern in table_count_patterns): + return "table_count" + + column_count_patterns = ( + r"\b(?:how many|count|number of)\b.*\b(?:columns?|fields?)\b", + r"\b(?:columns?|fields?)\b.*\b(?:count|number)\b", + ) + if any(re.search(pattern, normalized) for pattern in column_count_patterns): + return "column_count" + + schema_patterns = ( + r"\b(?:what|show|display|describe|list)\b.*\b(?:schema|metadata)\b", + r"\b(?:schema|metadata)\b.*\b(?:of|for|in)\b", + ) + if any(re.search(pattern, normalized) for pattern in schema_patterns): + return "schema" + + explicit_column_patterns = ( + r"\b(?:what|which|list|show|display|give|describe)\b.*\b(?:columns?|fields?)\b", + r"\b(?:columns?|fields?)\b.*\b(?:available|present|there|exist|schema|metadata)\b", + ) + if any(re.search(pattern, normalized) for pattern in explicit_column_patterns): + return "columns" + + table_patterns = ( + r"\b(?:what|which|list|show|display|give)\b.*\b(?:tables?|models?)\b", + r"\b(?:tables?|models?)\b.*\b(?:available|present|there|exist|in this datasource|in the datasource)\b", + r"\b(?:datasource|database|semantic layer|semantic model)\b.*\b(?:tables?|models?)\b", + ) + if any(re.search(pattern, normalized) for pattern in table_patterns): + return "tables" + + return None + + def _find_metadata_table_matches( + self, query: str, tables: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + query_key = self._normalize_schema_token(query) + if not query_key: + return [] + + matches: list[tuple[int, dict[str, Any]]] = [] + for table in tables: + table_name = str(table.get("name") or "") + if not table_name: + continue + short_name = re.split(r"[.$]", table_name)[-1] + normalized_name = self._normalize_schema_token(table_name) + normalized_short_name = self._normalize_schema_token(short_name) + + score = 0 + if normalized_name and normalized_name in query_key: + score = 100 + len(normalized_name) + elif normalized_short_name and normalized_short_name in query_key: + score = 80 + len(normalized_short_name) + + if score: + matches.append((score, table)) + + return [ + table + for _, table in sorted(matches, key=lambda item: item[0], reverse=True) + ] + + def _format_metadata_table_list( + self, tables: list[dict[str, Any]], *, max_tables: int = 120 + ) -> str: + if not tables: + return "I couldn't find any deployed tables in the active datasource metadata." + + sorted_tables = sorted( + {str(table.get("name")) for table in tables if table.get("name")}, + key=str.lower, + ) + shown_tables = sorted_tables[:max_tables] + lines = [ + f"The active datasource has {len(sorted_tables)} deployed table" + f"{'' if len(sorted_tables) == 1 else 's'}:" + ] + lines.extend(f"- {table_name}" for table_name in shown_tables) + if len(sorted_tables) > max_tables: + lines.append( + f"- ...and {len(sorted_tables) - max_tables} more tables." + ) + return "\n".join(lines) + + def _format_metadata_columns( + self, + query: str, + tables: list[dict[str, Any]], + *, + max_tables: int = 25, + max_columns_per_table: int = 60, + ) -> str: + if not tables: + return "I couldn't find any deployed columns in the active datasource metadata." + + matched_tables = self._find_metadata_table_matches(query, tables) + selected_tables = matched_tables or sorted( + tables, key=lambda table: str(table.get("name") or "").lower() + ) + selected_tables = selected_tables[:max_tables] + + heading = ( + "Columns available in the matched deployed table" + if matched_tables and len(selected_tables) == 1 + else "Columns available in the active datasource metadata" + ) + lines = [f"{heading}:"] + for table in selected_tables: + table_name = str(table.get("name") or "unknown_table") + columns = [ + column + for column in table.get("columns", []) + if isinstance(column, dict) and column.get("name") + ] + if not columns: + lines.append(f"- {table_name}: no columns found") + continue + + column_parts = [] + for column in columns[:max_columns_per_table]: + column_name = str(column.get("name")) + column_type = str(column.get("type") or "").upper() + column_parts.append( + f"{column_name} ({column_type})" if column_type else column_name + ) + if len(columns) > max_columns_per_table: + column_parts.append( + f"...and {len(columns) - max_columns_per_table} more" + ) + lines.append(f"- {table_name}: {', '.join(column_parts)}") + + if len(tables) > max_tables and not matched_tables: + lines.append(f"- ...and {len(tables) - max_tables} more tables.") + + return "\n".join(lines) + + def _extract_metadata_relationships(self, table_ddls: list[str]) -> list[str]: + relationships: list[str] = [] + for ddl in table_ddls or []: + if not isinstance(ddl, str): + continue + table_match = re.search( + r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", + ddl, + flags=re.IGNORECASE, + ) + if not table_match: + continue + source_table = next( + (value for value in table_match.groupdict().values() if value), + "unknown_table", + ) + + for relationship_match in re.finditer( + r"FOREIGN\s+KEY\s*\((?P[^)]+)\)\s+REFERENCES\s+" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" + r"\s*\((?P[^)]+)\)", + ddl, + flags=re.IGNORECASE, + ): + target_table = next( + ( + value + for key, value in relationship_match.groupdict().items() + if key + in { + "quoted", + "bracketed", + "backticked", + "bare", + } + and value + ), + "unknown_table", + ) + source_columns = relationship_match.group("source_columns") + target_columns = relationship_match.group("target_columns") + relationships.append( + f"{source_table}({source_columns}) -> " + f"{target_table}({target_columns})" + ) + + return sorted(set(relationships), key=str.lower) + + def _format_metadata_relationships(self, table_ddls: list[str]) -> str: + relationships = self._extract_metadata_relationships(table_ddls) + if not relationships: + return ( + "I couldn't find explicit relationships or foreign keys in the " + "active datasource metadata." + ) + + lines = [ + f"The active datasource metadata has {len(relationships)} " + f"relationship{'' if len(relationships) == 1 else 's'}:" + ] + lines.extend(f"- {relationship}" for relationship in relationships[:120]) + if len(relationships) > 120: + lines.append(f"- ...and {len(relationships) - 120} more relationships.") + return "\n".join(lines) + + def _format_metadata_schema( + self, query: str, tables: list[dict[str, Any]], table_ddls: list[str] + ) -> str: + matched_tables = self._find_metadata_table_matches(query, tables) + selected_tables = matched_tables or sorted( + tables, key=lambda table: str(table.get("name") or "").lower() + ) + selected_tables = selected_tables[:20] + if not selected_tables: + return "I couldn't find schema details in the active datasource metadata." + + lines = ["Schema details from the active datasource metadata:"] + for table in selected_tables: + table_name = str(table.get("name") or "unknown_table") + columns = [ + column + for column in table.get("columns", []) + if isinstance(column, dict) and column.get("name") + ] + lines.append(f"- {table_name}") + if columns: + column_parts = [] + for column in columns[:60]: + column_name = str(column.get("name")) + column_type = str(column.get("type") or "").upper() + column_parts.append( + f"{column_name} ({column_type})" + if column_type + else column_name + ) + if len(columns) > 60: + column_parts.append(f"...and {len(columns) - 60} more") + lines.append(f" Columns: {', '.join(column_parts)}") + else: + lines.append(" Columns: no columns found") + + relationships = self._extract_metadata_relationships(table_ddls) + if relationships: + lines.append("Relationships:") + lines.extend(f"- {relationship}" for relationship in relationships[:40]) + if len(relationships) > 40: + lines.append(f"- ...and {len(relationships) - 40} more relationships.") + + return "\n".join(lines) + + def _format_metadata_table_count(self, tables: list[dict[str, Any]]) -> str: + table_names = {str(table.get("name")) for table in tables if table.get("name")} + return ( + f"The active datasource has {len(table_names)} deployed table" + f"{'' if len(table_names) == 1 else 's'}." + ) + + def _format_metadata_column_count( + self, query: str, tables: list[dict[str, Any]] + ) -> str: + matched_tables = self._find_metadata_table_matches(query, tables) + selected_tables = matched_tables or tables + total_columns = sum( + len( + [ + column + for column in table.get("columns", []) + if isinstance(column, dict) and column.get("name") + ] + ) + for table in selected_tables + ) + if matched_tables and len(selected_tables) == 1: + table_name = str(selected_tables[0].get("name") or "the matched table") + return f"{table_name} has {total_columns} deployed columns." + return ( + f"The active datasource metadata has {total_columns} deployed columns " + f"across {len(selected_tables)} table" + f"{'' if len(selected_tables) == 1 else 's'}." + ) + + def _build_metadata_response( + self, query: str, table_ddls: list[str], table_names: list[str] + ) -> str: + kind = self._get_metadata_question_kind(query) + parsed_tables = self._parse_schema_tables(table_ddls) + + if not parsed_tables and table_names: + parsed_tables = [ + {"name": table_name, "columns": []} for table_name in table_names + ] + + if kind == "schema": + return self._format_metadata_schema(query, parsed_tables, table_ddls) + if kind == "relationships": + return self._format_metadata_relationships(table_ddls) + if kind == "table_count": + return self._format_metadata_table_count(parsed_tables) + if kind == "column_count": + return self._format_metadata_column_count(query, parsed_tables) + if kind == "columns": + return self._format_metadata_columns(query, parsed_tables) + return self._format_metadata_table_list(parsed_tables) + + def _normalize_schema_token(self, value: str) -> str: + return re.sub(r"[^a-z0-9]", "", (value or "").lower()) + + def _query_schema_terms(self, query: str) -> set[str]: + normalized_query = (query or "").lower() + stop_words = { + "about", + "against", + "from", + "give", + "group", + "grouped", + "list", + "month", + "monthly", + "over", + "rows", + "show", + "table", + "tables", + "the", + "this", + "using", + "what", + "which", + "with", + "year", + } + terms = { + self._normalize_schema_token(token) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", normalized_query) + if len(token) > 2 and token not in stop_words + } + return {term for term in terms if term} + + def _prune_sql_generation_context( + self, + query: str, + documents: list[dict], + table_names: list[str], + table_ddls: list[str], + *, + max_tables: int = 8, + ) -> tuple[list[dict], list[str], list[str]]: + if len(table_ddls) <= max_tables: + return documents, table_names, table_ddls + + parsed_tables = self._parse_schema_tables(table_ddls) + if not parsed_tables: + return documents, table_names, table_ddls[:max_tables] + + query_key = self._normalize_schema_token(query) + query_terms = self._query_schema_terms(query) + explicit_tables = { + self._normalize_schema_token(table_name) + for table_name in self._extract_explicit_table_names_from_query(query) + } + + scored: list[tuple[int, int]] = [] + for index, table in enumerate(parsed_tables): + table_name = str(table.get("name") or "") + normalized_table = self._normalize_schema_token(table_name) + normalized_short_table = self._normalize_schema_token( + re.split(r"[.$]", table_name)[-1] + ) + column_terms = { + self._normalize_schema_token(str(column.get("name") or "")) + for column in table.get("columns", []) + if column.get("name") + } + + score = 0 + if normalized_table in explicit_tables or normalized_short_table in explicit_tables: + score += 1000 + if normalized_table and normalized_table in query_key: + score += 500 + if normalized_short_table and normalized_short_table in query_key: + score += 450 + for term in query_terms: + if not term: + continue + if term == normalized_table or term == normalized_short_table: + score += 80 + elif term in normalized_table or term in normalized_short_table: + score += 40 + for column_term in column_terms: + if term == column_term: + score += 60 + elif term in column_term or column_term in term: + score += 25 + + if score > 0: + scored.append((score, index)) + + if not scored: + return documents, table_names, table_ddls[:max_tables] + + sorted_scored_indexes = [ + index for _, index in sorted(scored, key=lambda item: item[0], reverse=True) + ] + core_limit = max(1, max_tables - 2) if max_tables > 2 else 1 + selected_indexes = sorted_scored_indexes[:core_limit] + selected_indexes = self._expand_pruned_context_with_related_tables( + selected_indexes, + parsed_tables, + table_ddls, + max_tables=max_tables, + ) + for index in sorted_scored_indexes: + if len(selected_indexes) >= max_tables: + break + if index not in selected_indexes: + selected_indexes.append(index) + selected_indexes = sorted(selected_indexes) + pruned_documents = [ + documents[index] for index in selected_indexes if index < len(documents) + ] + pruned_table_names = [ + table_names[index] for index in selected_indexes if index < len(table_names) + ] + pruned_table_ddls = [ + table_ddls[index] for index in selected_indexes if index < len(table_ddls) + ] + + logger.info( + "Pruned SQL generation context from %s to %s tables for query: %s", + len(table_ddls), + len(pruned_table_ddls), + query, + ) + return pruned_documents, pruned_table_names, pruned_table_ddls + + def _expand_pruned_context_with_related_tables( + self, + selected_indexes: list[int], + parsed_tables: list[dict[str, Any]], + table_ddls: list[str], + *, + max_tables: int, + ) -> list[int]: + if len(selected_indexes) >= max_tables: + return selected_indexes[:max_tables] + + selected: list[int] = list(dict.fromkeys(selected_indexes)) + selected_set = set(selected) + + def join_key_columns(table: dict[str, Any]) -> set[str]: + keys = set() + for column in table.get("columns", []): + column_name = str(column.get("name") or "") + normalized = self._normalize_schema_token(column_name) + if not normalized: + continue + if ( + normalized == "id" + or normalized.endswith("id") + or normalized.endswith("no") + or normalized.endswith("number") + or normalized.endswith("code") + or normalized.endswith("key") + ): + keys.add(normalized) + return keys + + selected_table_names = { + self._normalize_schema_token( + str(parsed_tables[index].get("name") or "") + ) + for index in selected + if index < len(parsed_tables) + } + selected_join_keys: set[str] = set() + for index in selected: + if index < len(parsed_tables): + selected_join_keys.update(join_key_columns(parsed_tables[index])) + + candidates: list[tuple[int, int]] = [] + for index, table in enumerate(parsed_tables): + if index in selected_set: + continue + + table_name = str(table.get("name") or "") + normalized_table_name = self._normalize_schema_token(table_name) + ddl = table_ddls[index] if index < len(table_ddls) else "" + normalized_ddl = self._normalize_schema_token(ddl) + table_join_keys = join_key_columns(table) + + score = 0 + shared_keys = selected_join_keys & table_join_keys + if shared_keys: + score += 20 + 5 * len(shared_keys) + if normalized_table_name and any( + selected_table + and ( + selected_table in normalized_ddl + or normalized_table_name in selected_table + ) + for selected_table in selected_table_names + ): + score += 40 + if re.search(r"\b(?:foreign\s+key|references)\b", ddl, flags=re.IGNORECASE): + score += 15 + + if score > 0: + candidates.append((score, index)) + + for _, index in sorted(candidates, key=lambda item: item[0], reverse=True): + if len(selected) >= max_tables: + break + selected.append(index) + selected_set.add(index) + + return selected + + def _is_valid_select_sql(self, sql: Optional[str]) -> bool: + if not isinstance(sql, str): + return False + + normalized = re.sub(r"\s+", " ", sql.strip()) + if not normalized: + return False + + return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) + + def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: + if not self._is_valid_select_sql(sql): + return None + return AskResult(sql=sql.strip(), type="llm") + + def _build_validated_ask_result_from_sql( + self, + sql: Optional[str], + table_ddls: list[str], + query: str | None = None, + ) -> Optional[AskResult]: + if isinstance(sql, str): + sql = normalize_sql_direction_keywords(sql) + sql = normalize_sql_table_references_to_schema( + sql, + construct_valid_table_names(table_ddls), + ) + sql = normalize_sql_column_references_to_schema( + sql, + construct_valid_table_columns(table_ddls), + ) + ask_result = self._build_ask_result_from_sql(sql) + if not ask_result: + return None + + schema_tables = self._parse_schema_tables(table_ddls) + valid_tables = { + str(table.get("name") or "").lower(): table + for table in schema_tables + if table.get("name") + } + valid_table_suffixes = { + table_name.split(".")[-1].lower(): table + for table_name, table in valid_tables.items() + } + + table_reference_pattern = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", + flags=re.IGNORECASE, + ) + referenced_tables = [ + next(value for value in match.groupdict().values() if value) + for match in table_reference_pattern.finditer(ask_result.sql) + ] + invalid_tables = [ + table + for table in referenced_tables + if table.lower() not in valid_tables + and table.lower().split(".")[-1] not in valid_table_suffixes + ] + + columns_by_table = { + table_name: { + str(column.get("name") or "").lower() + for column in table.get("columns", []) + if column.get("name") + } + for table_name, table in valid_tables.items() + } + columns_by_table.update( + { + table_name.split(".")[-1].lower(): columns + for table_name, columns in columns_by_table.items() + } + ) + + qualified_column_pattern = re.compile( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"(?P[A-Za-z_][A-Za-z0-9_$]*))", + flags=re.IGNORECASE, + ) + invalid_columns = [] + for match in qualified_column_pattern.finditer(ask_result.sql): + table_reference = ( + match.group("table_quoted") + or match.group("table_bracketed") + or match.group("table_bare") + or "" + ) + column_reference = ( + match.group("column_quoted") + or match.group("column_bracketed") + or match.group("column_bare") + or "" + ) + table_key = table_reference.lower() + column_key = column_reference.lower() + table_columns = columns_by_table.get(table_key) or columns_by_table.get( + table_key.split(".")[-1] + ) + if table_columns is not None and column_key not in table_columns: + invalid_columns.append(f"{table_reference}.{column_reference}") + + if invalid_tables or invalid_columns: + logger.warning( + "Ignoring heuristic SQL because it is not valid for active schema. " + "invalid_tables=%s invalid_columns=%s sql=%s", + invalid_tables, + invalid_columns, + ask_result.sql, + ) + return None + + invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( + ask_result.sql, + schema_tables, + ) + if invalid_unqualified_identifiers: + logger.warning( + "Ignoring SQL because it references unqualified fields outside the active schema. " + "invalid_identifiers=%s sql=%s", + invalid_unqualified_identifiers, + ask_result.sql, + ) + return None + + invalid_output_aliases = self._invalid_sql_output_aliases( + ask_result.sql, + schema_tables, + ) + if invalid_output_aliases: + logger.warning( + "Ignoring SQL because it aliases output fields to unavailable schema concepts. " + "invalid_aliases=%s sql=%s", + invalid_output_aliases, + ask_result.sql, + ) + return None + + if not self._sql_references_explicit_table(ask_result.sql, query): + return None + + if not self._sql_matches_question_intent( + ask_result.sql, + query, + schema_tables, + ): + return None + + return ask_result + + def _build_failed_text_to_sql_response( + self, + trace_id: Optional[str], + message: str, + *, + rephrased_question: Optional[str] = None, + intent_reasoning: Optional[str] = None, + retrieved_tables: Optional[list[str]] = None, + sql_generation_reasoning: Optional[str] = None, + invalid_sql: Optional[str] = None, + is_followup: bool = False, + code: Literal["NO_RELEVANT_DATA", "NO_RELEVANT_SQL", "OTHERS"] = "NO_RELEVANT_SQL", + ) -> AskResultResponse: + return AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError(code=code, message=message), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=retrieved_tables, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + trace_id=trace_id, + is_followup=is_followup, + ) + + def _build_no_relevant_active_datasource_response( + self, + trace_id: Optional[str], + *, + rephrased_question: Optional[str] = None, + intent_reasoning: Optional[str] = None, + retrieved_tables: Optional[list[str]] = None, + sql_generation_reasoning: Optional[str] = None, + is_followup: bool = False, + ) -> AskResultResponse: + return self._build_failed_text_to_sql_response( + trace_id, + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=retrieved_tables, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=None, + is_followup=is_followup, + code="NO_RELEVANT_DATA", + ) + + @observe(name="Ask Question") + @trace_metadata + async def ask( + self, + ask_request: AskRequest, + **kwargs, + ): + trace_id = kwargs.get("trace_id") + results = { + "ask_result": {}, + "metadata": { + "type": "", + "error_type": "", + "error_message": "", + "request_from": ask_request.request_from, + }, + } + + query_id = ask_request.query_id + if not query_id: + raise ValueError("query_id is required for ask service execution") + + user_query = (ask_request.query or "").strip() + if not user_query: + self._ask_results[query_id] = self._build_failed_text_to_sql_response( + trace_id, + "Question is required", + code="OTHERS", + ) + results["metadata"]["error_type"] = "OTHERS" + results["metadata"]["error_message"] = "Question is required" + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + logger.info(f"Ask pipeline started for query_id: {query_id}") + histories = ask_request.histories[: self._max_histories][ + ::-1 + ] # reverse the order of histories + if histories and not self._should_use_histories_for_query(user_query): + logger.info( + "Ignoring thread histories for independent question. query_id=%s query=%s", + query_id, + user_query, + ) + histories = [] + rephrased_question = None + intent_reasoning = None + sql_generation_reasoning = None + sql_samples = [] + instructions = [] + api_results = [] + documents = [] + table_names = [] + table_ddls = [] + _retrieval_result = {} + error_message = None + invalid_sql = None + allow_sql_generation_reasoning = ( + self._allow_sql_generation_reasoning + and not ask_request.ignore_sql_generation_reasoning + ) + enable_column_pruning = ( + self._enable_column_pruning or ask_request.enable_column_pruning + ) + allow_sql_functions_retrieval = self._allow_sql_functions_retrieval + allow_sql_diagnosis = self._allow_sql_diagnosis + allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval + max_sql_correction_retries = self._max_sql_correction_retries + current_sql_correction_retries = 0 + use_dry_plan = ask_request.use_dry_plan + allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback + sql_knowledge = None + understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) + request_explicit_table_names = self._normalize_explicit_table_names( + ask_request.explicit_tables + ) + query_explicit_table_names = self._normalize_explicit_table_names( + self._extract_explicit_table_names_from_query(user_query) + ) + forced_request_explicit_table_names = self._forced_explicit_table_names( + request_explicit_table_names, + source="request", + ) + explicit_table_names = ( + forced_request_explicit_table_names or query_explicit_table_names + ) + retrieval_table_names = explicit_table_names or None + + try: + sql_user_query = user_query + + # ask status can be understanding, searching, generating, finished, failed, stopped + # we will need to handle business logic for each status + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="understanding", + trace_id=trace_id, + is_followup=True if histories else False, + ) + + if self._is_greeting_query(user_query): + self._general_streaming_results[query_id] = ( + self._build_greeting_response(user_query) + ) + + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + trace_id=trace_id, + is_followup=True if histories else False, + general_type="USER_GUIDE", + ) + results["metadata"]["type"] = "GENERAL" + return results + + metadata_question_kind = self._get_metadata_question_kind(user_query) + if metadata_question_kind: + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="GENERAL", + rephrased_question=user_query, + intent_reasoning=( + "Basic datasource metadata question detected; " + "retrieving deployed schema metadata directly." + ), + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + retrieval_result = await self._run_with_timeout( + "Metadata schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + metadata_answer = self._build_metadata_response( + user_query, table_ddls, table_names + ) + self._general_streaming_results[query_id] = metadata_answer + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=user_query, + intent_reasoning=( + "Answered from active datasource deployed metadata " + "without SQL generation." + ), + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + results["metadata"]["type"] = "GENERAL" + results["metadata"]["metadata_question_kind"] = ( + metadata_question_kind + ) + results["metadata"]["retrieved_table_count"] = len(documents) + return results + + if explicit_table_names: + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=user_query, + intent_reasoning="Explicit table name detected; retrieving that deployed schema directly.", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Explicit table schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) + if not documents and not request_explicit_table_names: + logger.info( + "Explicit table retrieval did not return requested active-schema table; " + "loading full active schema. query_id=%s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval for explicit table", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + all_documents, _, _ = self._extract_retrieval_metadata( + retrieval_result + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + all_documents, + explicit_table_names, + ) + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + logger.info( + "Retrieved explicit tables for query_id %s: %s", + query_id, + table_names, + ) + + if ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, + ): + ask_result = self._build_validated_ask_result_from_sql( + ranked_measure_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated ranked measure SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = ranked_measure_sql + + if table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ): + ask_result = self._build_validated_ask_result_from_sql( + table_question_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table question matched deployed schema.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = table_question_sql + + if explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ): + explicit_sql, explicit_table_name = explicit_table_preview + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table preview request matched deployed schema.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = explicit_sql + + if documents and ( + deterministic_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ) + ): + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = deterministic_sql + + if not documents: + error_message = ( + "The requested table was not found in the deployed schema: " + + ", ".join(explicit_table_names) + ) + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_DATA", + message=error_message, + ), + rephrased_question=user_query, + intent_reasoning="Explicit table request did not match any deployed schema table.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = error_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + rephrased_question = user_query + intent_reasoning = ( + "Explicit table request matched deployed schema; generating SQL against retrieved schema." + ) + sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + + if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + logger.info( + "Retrieved tables for direct heuristic query_id %s: %s", + query_id, + table_names, + ) + + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using direct heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + if ask_result := self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ): + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + + if explicit_group_count_sql := self._build_explicit_group_count_sql( + user_query + ): + invalid_sql = explicit_group_count_sql + rephrased_question = user_query + logger.info( + "Deferring explicit grouped count SQL until active schema validation for query_id %s", + query_id, + ) + + historical_question_result = [] + should_skip_pre_sql_retrieval = self._is_data_analysis_query( + user_query + ) + if should_skip_pre_sql_retrieval: + rephrased_question = user_query + intent_reasoning = ( + "Detected a deployed-data analytics question; skipping " + "intent classification and using SQL generation." + ) + sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + logger.info( + "Skipping pre-SQL retrieval for analytics query_id %s: %s", + query_id, + user_query, + ) + + if ( + not api_results + and not should_skip_pre_sql_retrieval + and self._should_reuse_historical_question_sql( + user_query, histories + ) + ): + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + + try: + historical_question = await self._run_with_timeout( + "Historical question retrieval", + self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ), + timeout_seconds=min(understanding_timeout_seconds, 10), + ) - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] + except TimeoutError as exc: + logger.warning( + "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) + + valid_historical_results = [] + for result in historical_question_result: + historical_question_text = result.get("question") + if not self._is_reusable_historical_question( + user_query, historical_question_text + ): + logger.info( + "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", + query_id, + user_query, + historical_question_text, + ) + continue - if historical_question_result: - api_results = [ + sql_statement = result.get("statement") + if not self._is_valid_select_sql(sql_statement): + logger.warning( + "Ignoring historical question without valid SQL for query_id %s", + query_id, + ) + continue + valid_historical_results.append( AskResult( **{ - "sql": result.get("statement"), + "sql": sql_statement.strip(), "type": "view" if result.get("viewId") else "llm", "viewId": result.get("viewId"), } ) - for result in historical_question_result - ] + ) + + if valid_historical_results: + api_results = valid_historical_results sql_generation_reasoning = "" - else: + elif not api_results and not should_skip_pre_sql_retrieval: + original_user_query = user_query # Run both pipeline operations concurrently - sql_samples_task, instructions_task = await asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", - ), - ) + try: + sql_samples_task, instructions_task = await self._run_with_timeout( + "SQL pair and instruction retrieval", + asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), + ), + timeout_seconds=understanding_timeout_seconds, + ) - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] - ) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] + ) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) + except TimeoutError as exc: + logger.warning( + "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) + sql_samples = [] + instructions = [] if self._allow_intent_classification: - intent_classification_result = ( - await self._pipelines["intent_classification"].run( - query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, + try: + intent_classification_result = ( + await self._run_with_timeout( + "Intent classification", + self._pipelines["intent_classification"].run( + query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, + ), + timeout_seconds=understanding_timeout_seconds, + ) + ).get("post_process", {}) + except TimeoutError as exc: + logger.warning( + "Intent classification timed out; continuing with TEXT_TO_SQL. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, ) - ).get("post_process", {}) + intent_classification_result = { + "intent": "TEXT_TO_SQL", + "rephrased_question": user_query, + "reasoning": "Intent classification timed out; using SQL generation.", + "db_schemas": [], + } intent = intent_classification_result.get("intent") rephrased_question = intent_classification_result.get( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") + retrieved_db_schemas = intent_classification_result.get( + "db_schemas" + ) or [] + is_original_analytics_query = self._is_data_analysis_query( + original_user_query + ) + is_schema_grounded_query = self._is_schema_grounded_query( + original_user_query, retrieved_db_schemas + ) or self._is_schema_grounded_query( + rephrased_question or "", retrieved_db_schemas + ) + + if intent in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"} and ( + is_original_analytics_query + or is_schema_grounded_query + or self._is_data_analysis_query(rephrased_question or "") + ): + logger.info( + "Overriding intent %s to TEXT_TO_SQL for schema/data query: %s", + intent, + user_query, + ) + intent = "TEXT_TO_SQL" - if rephrased_question: + if is_original_analytics_query: + if rephrased_question and rephrased_question != user_query: + logger.info( + "Ignoring rephrased analytics query from intent classification. original=%s rephrased=%s", + original_user_query, + rephrased_question, + ) + user_query = original_user_query + rephrased_question = original_user_query + elif rephrased_question: user_query = rephrased_question + sql_user_query = ( + self._rewrite_query_for_text_to_sql(user_query) + if self._is_data_analysis_query(user_query) + else user_query + ) + if intent == "MISLEADING_QUERY": - asyncio.create_task( + general_result = await self._run_with_timeout( + "Misleading assistance", self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, @@ -286,14 +6394,18 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "misleading_assistance" ) ) self._ask_results[query_id] = AskResultResponse( status="finished", - type="GENERAL", + type="MISLEADING_QUERY", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, @@ -303,7 +6415,8 @@ async def ask( results["metadata"]["type"] = "MISLEADING_QUERY" return results elif intent == "GENERAL": - asyncio.create_task( + general_result = await self._run_with_timeout( + "Data assistance", self._pipelines["data_assistance"].run( query=user_query, histories=histories, @@ -311,8 +6424,12 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "data_assistance" ) ) @@ -328,12 +6445,17 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results elif intent == "USER_GUIDE": - asyncio.create_task( + general_result = await self._run_with_timeout( + "User guide assistance", self._pipelines["user_guide_assistance"].run( query=user_query, language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "user_guide_assistance" ) ) @@ -357,7 +6479,11 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - if not self._is_stopped(query_id, self._ask_results) and not api_results: + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and not documents + ): self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -367,39 +6493,530 @@ async def ask( is_followup=True if histories else False, ) - retrieval_result = await self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=ask_request.explicit_tables, - histories=histories, - project_id=ask_request.project_id, - enable_column_pruning=enable_column_pruning, - ) + try: + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), + ), + timeout_seconds=self._schema_retrieval_timeout_seconds, + ) + except TimeoutError as error: + if not self._should_retry_selected_schema_after_retrieval_timeout( + retrieval_table_names + ): + logger.warning( + "Schema retrieval timed out for data query; not loading full project schema. " + "query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + error, + ) + retrieval_result = {"construct_retrieval_results": {}} + else: + logger.warning( + "Schema retrieval timed out; retrying only explicit selected schemas. " + "query_id=%s project_id=%s tables=%s error=%s", + query_id, + ask_request.project_id, + retrieval_table_names, + error, + ) + retrieval_result = await self._run_with_timeout( + "Selected schema fallback retrieval", + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 30, + ), + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - documents = _retrieval_result.get("retrieval_results", []) - table_names = [document.get("table_name") for document in documents] - table_ddls = [document.get("table_ddl") for document in documents] - + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) if not documents: - logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") + if explicit_table_names: + logger.info( + "Retrying schema retrieval for explicit tables query_id %s: %s", + query_id, + explicit_table_names, + ) + retrieval_result = await self._run_with_timeout( + "Explicit table schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=enable_column_pruning, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 20, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) + if ( + not documents + and self._get_metadata_question_kind(user_query) + and not request_explicit_table_names + ): + logger.info( + "Query-based schema retrieval returned no tables for data question; " + "retrying full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) + logger.info( + "Retrieved tables for query_id %s: %s", query_id, table_names + ) + + if not api_results and ( + ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, + ) + ): + logger.info( + "Using schema-grounded ranked measure SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + ranked_measure_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = ranked_measure_sql + error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." + + if not api_results and ( + table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ) + ): + logger.info( + "Using schema-grounded table question SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + table_question_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = table_question_sql + error_message = "Schema-grounded table SQL was not valid for the active datasource schema." + + if not api_results and ( + explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + ): + explicit_sql, explicit_table_name = explicit_table_preview + logger.info( + "Using explicit table preview SQL for query_id %s and table %s", + query_id, + explicit_table_name, + ) + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = explicit_sql + error_message = "Explicit table preview SQL was not valid for the active datasource schema." + + if not api_results and ( + audit_log_activity_sql := self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ) + ): + logger.info( + "Using schema-grounded audit log activity SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + audit_log_activity_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = audit_log_activity_sql + error_message = ( + "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." + ) + + if ( + not api_results + and self._is_data_analysis_query(user_query) + and ( + schema_grounded_sql := self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + ) + ): + logger.info( + "Using generic schema-grounded analytics SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + schema_grounded_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = schema_grounded_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and any( + term in user_query.lower() + for term in ( + "pcb", + "repair", + "failure", + "business unit", + "business units", + "product line", + "product family", + ) + ): + operational_sql = self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + if operational_sql: + logger.info( + "Using schema-grounded operational SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + operational_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = operational_sql + error_message = ( + "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and ( + deterministic_sales_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ) + ): + logger.info( + "Using schema-grounded CWSales SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sales_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = deterministic_sales_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." + ) + + should_retry_full_schema = ( + not api_results + and self._get_metadata_question_kind(user_query) + and "db_schema_retrieval" in self._pipelines + and not request_explicit_table_names + and not table_names + ) + if should_retry_full_schema: + logger.info( + "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retry", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 30, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + full_documents, full_table_names, full_table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + full_documents, full_table_names, full_table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + full_documents, + explicit_table_names, + ) + ) + if full_documents: + documents, table_names, table_ddls = ( + full_documents, + full_table_names, + full_table_ddls, + ) + logger.info( + "Using full active deployed schema retry for query_id %s: %s", + query_id, + table_names, + ) + + full_schema_preview = self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + full_schema_sql_candidates = ( + self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ), + full_schema_preview[0] if full_schema_preview else None, + self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ), + self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ), + self._build_schema_grounded_sales_sql( + user_query, table_ddls + ), + ) + for full_schema_sql in full_schema_sql_candidates: + if not full_schema_sql: + continue + ask_result = self._build_validated_ask_result_from_sql( + full_schema_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + break + invalid_sql = full_schema_sql + error_message = ( + "Full-schema grounded SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and ( + unqueryable_metric_message := self._get_unqueryable_metric_message( + user_query, table_ddls + ) + ): + logger.info( + "ask pipeline - NO_RELEVANT_SQL due to unqueryable metric: %s", + user_query, + ) if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="failed", type="TEXT_TO_SQL", error=AskError( - code="NO_RELEVANT_DATA", - message="No relevant data", + code="NO_RELEVANT_SQL", + message=unqueryable_metric_message, ), rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, + retrieved_tables=table_names, trace_id=trace_id, is_followup=True if histories else False, ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = unqueryable_metric_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + if not documents: + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", + query_id, + user_query, + ) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + is_followup=True if histories else False, + ) + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + is_followup=True if histories else False, + ) + ) results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) results["metadata"]["type"] = "TEXT_TO_SQL" return results + if documents and not api_results: + documents, table_names, table_ddls = self._prune_sql_generation_context( + sql_user_query, + documents, + table_names, + table_ddls, + ) + ( + documents, + table_names, + table_ddls, + completed_retrieval_result, + ) = await self._complete_sql_generation_context( + query=sql_user_query, + project_id=ask_request.project_id, + documents=documents, + table_names=table_names, + table_ddls=table_ddls, + ) + if completed_retrieval_result: + _retrieval_result = completed_retrieval_result + + sql_generation_histories = histories + if self._is_data_analysis_query( + sql_user_query + ) and not self._needs_conversation_context(sql_user_query): + sql_generation_histories = [] + allow_sql_generation_reasoning = False + allow_sql_knowledge_retrieval = False + max_sql_correction_retries = min(max_sql_correction_retries, 1) + logger.info( + "Using fast standalone SQL generation path for query_id %s", + query_id, + ) + if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -415,29 +7032,53 @@ async def ask( is_followup=True if histories else False, ) - if histories: - sql_generation_reasoning = ( - await self._pipelines["followup_sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ) - ).get("post_process", {}) + if sql_generation_histories: + try: + sql_generation_reasoning = ( + await self._run_with_timeout( + "Follow-up SQL generation reasoning", + self._pipelines[ + "followup_sql_generation_reasoning" + ].run( + query=sql_user_query, + contexts=table_ddls, + histories=sql_generation_histories, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), + ) + ).get("post_process", {}) + except Exception as reasoning_error: + logger.warning( + "Follow-up SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", + query_id, + reasoning_error, + ) + sql_generation_reasoning = "" else: - sql_generation_reasoning = ( - await self._pipelines["sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, + try: + sql_generation_reasoning = ( + await self._run_with_timeout( + "SQL generation reasoning", + self._pipelines["sql_generation_reasoning"].run( + query=sql_user_query, + contexts=table_ddls, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), + ) + ).get("post_process", {}) + except Exception as reasoning_error: + logger.warning( + "SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", + query_id, + reasoning_error, ) - ).get("post_process", {}) + sql_generation_reasoning = "" self._ask_results[query_id] = AskResultResponse( status="planning", @@ -462,14 +7103,34 @@ async def ask( is_followup=True if histories else False, ) - if allow_sql_functions_retrieval: - sql_functions = await self._pipelines[ - "sql_functions_retrieval" - ].run( - project_id=ask_request.project_id, + try: + sql_functions, sql_knowledge = await self._run_with_timeout( + "SQL helper retrieval", + asyncio.gather( + ( + self._pipelines["sql_functions_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_functions_retrieval + else _return_value([]) + ), + ( + self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_knowledge_retrieval + else _return_value(None) + ), + ), + timeout_seconds=min(self._pipeline_timeout_seconds, 10), ) - else: - sql_functions = [] + except TimeoutError as helper_timeout: + logger.warning( + "SQL helper retrieval timed out for query_id %s; continuing with schema only: %s", + query_id, + helper_timeout, + ) + sql_functions, sql_knowledge = [], None has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -477,63 +7138,92 @@ async def ask( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - if histories: - text_to_sql_generation_results = await self._pipelines[ - "followup_sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - ) - else: - text_to_sql_generation_results = await self._pipelines[ - "sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, + try: + if sql_generation_histories: + text_to_sql_generation_results = await self._run_with_timeout( + "Follow-up SQL generation", + self._pipelines["followup_sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=sql_generation_histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), + ) + else: + text_to_sql_generation_results = await self._run_with_timeout( + "SQL generation", + self._pipelines["sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), + ) + except TimeoutError as generation_timeout: + logger.warning( + "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + query_id, + generation_timeout, ) + text_to_sql_generation_results = { + "post_process": { + "valid_generation_result": None, + "invalid_generation_result": None, + } + } + error_message = str(generation_timeout) if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - api_results = [ - AskResult( - **{ - "sql": sql_valid_result.get("sql"), - "type": "llm", - } + if ask_result := self._build_validated_ask_result_from_sql( + sql_valid_result.get("sql"), + table_ddls, + sql_user_query, + ): + api_results = [ask_result] + else: + invalid_sql = sql_valid_result.get("sql") + error_message = ( + "SQL generation did not produce SQL that matches the active datasource schema and question intent." ) - ] elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + if failed_dry_run_result["type"] in { + "TIME_OUT", + "UNSUPPORTED_SQL", + }: + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) break original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + sql_diagnosis_reasoning = None current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -548,52 +7238,67 @@ async def ask( ) if allow_sql_diagnosis: - sql_diagnosis_results = await self._pipelines[ - "sql_diagnosis" - ].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_request.configurations.language, + sql_diagnosis_results = await self._run_with_timeout( + "SQL diagnosis", + self._pipelines["sql_diagnosis"].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_request.configurations.language, + ), ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") - sql_correction_results = await self._pipelines[ - "sql_correction" - ].run( - contexts=table_ddls, - instructions=instructions, - invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, - }, - project_id=ask_request.project_id, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_functions=sql_functions, + correction_error_message = error_message + if sql_diagnosis_reasoning: + correction_error_message = ( + f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" + ) + + sql_correction_results = await self._run_with_timeout( + "SQL correction", + self._pipelines["sql_correction"].run( + contexts=table_ddls, + instructions=instructions, + invalid_generation_result={ + "original_sql": original_sql, + "sql": invalid_sql, + "error": correction_error_message, + }, + project_id=ask_request.project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + query=sql_user_query, + ), ) if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - api_results = [ - AskResult( - **{ - "sql": valid_generation_result.get("sql"), - "type": "llm", - } - ) - ] - break + if ask_result := self._build_validated_ask_result_from_sql( + valid_generation_result.get("sql"), + table_ddls, + sql_user_query, + ): + api_results = [ask_result] + break + invalid_sql = valid_generation_result.get("sql") + error_message = ( + "SQL correction did not produce SQL that matches the active datasource schema and question intent." + ) failed_dry_run_result = sql_correction_results["post_process"][ "invalid_generation_result" ] + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) if api_results: if not self._is_stopped(query_id, self._ask_results): @@ -611,25 +7316,64 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + else: + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_SQL", - message=error_message or "No relevant SQL", - ), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - trace_id=trace_id, - is_followup=True if histories else False, + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + is_followup=True if histories else False, + ) + ) + if error_message or invalid_sql: + logger.info( + "Suppressed technical SQL failure for query_id %s. " + "error=%s invalid_sql=%s", + query_id, + error_message, + invalid_sql, ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = error_message + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) results["metadata"]["type"] = "TEXT_TO_SQL" return results @@ -683,11 +7427,9 @@ async def get_ask_streaming_result( self, query_id: str, ): - if query_id in self._general_streaming_results: + if general_response := self._general_streaming_results.get(query_id): event = SSEEvent( - data=SSEEvent.SSEEventMessage( - message=self._general_streaming_results[query_id], - ), + data=SSEEvent.SSEEventMessage(message=general_response), ) yield event.serialize() return diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index a6c6241fa7..7c7d46ce89 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,44 +2,189 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( + _is_project_wide_analysis_query, + _rerank_table_documents, + _select_relevant_table_documents, check_using_db_schemas_without_pruning, dbschema_retrieval, - embedding, + expand_business_terms_for_retrieval, table_retrieval, ) -@pytest.mark.asyncio -async def test_embedding_skips_vector_lookup_for_explicit_tables(): - class Embedder: - def __init__(self): - self.called = False +def test_project_wide_analysis_query_includes_broad_ranking_questions(): + assert _is_project_wide_analysis_query( + "Which projects have the highest number of completed questions?" + ) - async def run(self, query): - self.called = True - return {"embedding": [0.1]} - embedder = Embedder() +def test_project_wide_analysis_query_ignores_empty_query(): + assert not _is_project_wide_analysis_query("") - result = await embedding( - query="show rows", - embedder=embedder, - histories=[], - tables=["orders"], + +def test_expand_business_terms_for_retrieval_adds_generic_sales_order_terms(): + query = "Show top customers by invoice amount" + + expanded_query = expand_business_terms_for_retrieval(query) + + assert query in expanded_query + assert "transaction purchase billing account geography" in expanded_query + assert "money exchange currency" in expanded_query + + +def test_expand_business_terms_for_retrieval_adds_generic_currency_market_terms(): + query = "Show invoice distribution by currency across markets" + + expanded_query = expand_business_terms_for_retrieval(query) + + assert query in expanded_query + assert "money exchange currency" in expanded_query + + +def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): + query = "Explain what this workspace does" + + assert expand_business_terms_for_retrieval(query) == query + + +def test_rerank_table_documents_prefers_question_relevant_table_text(): + generic_stage = Document( + content="Generic imported staging records with product labels.", + meta={"type": "TABLE_DESCRIPTION", "name": "generic_stage_load"}, + score=0.99, + ) + order_region_table = Document( + content="Business transactions grouped by customer geography and amount.", + meta={"type": "TABLE_DESCRIPTION", "name": "business_transactions"}, + score=0.01, + ) + + documents = _rerank_table_documents( + "Show order distribution across regions.", + [generic_stage, order_region_table], + ) + + assert documents[0].meta["name"] == "business_transactions" + + +def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): + test_load = Document( + content="Raw test load rows for order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ) + order_market_table = Document( + content="New order transaction records with market and customer fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.45, + ) + + documents = _rerank_table_documents( + "Show order distribution across markets.", + [test_load, order_market_table], ) - assert result == {} - assert not embedder.called + assert documents[0].meta["name"] == "dbo_xStageNewOrders" + + +def test_select_relevant_table_documents_limits_weak_extra_candidates(): + documents = [ + Document( + content="Invoice transactions with product, customer, currency, and amount.", + meta={"type": "TABLE_DESCRIPTION", "name": "invoices"}, + score=0.92, + ), + Document( + content="Product catalog with product names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "products"}, + score=0.86, + ), + Document( + content="Customer account master data.", + meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, + score=0.82, + ), + Document( + content="Exchange rate lookup by currency.", + meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, + score=0.78, + ), + Document( + content="Sales regions and market hierarchy.", + meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, + score=0.74, + ), + Document( + content="Raw staging audit rows with load metadata.", + meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, + score=0.99, + ), + ] + + selected = _select_relevant_table_documents( + "Show invoice distribution by currency across markets", + documents, + ) + + assert 1 <= len(selected) <= 5 + assert "staging_audit" not in [document.meta["name"] for document in selected] + + +def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): + documents = [ + Document( + content="Raw test load rows with order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ), + Document( + content="New order transaction records with market and customer details.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.4, + ), + ] + + selected = _select_relevant_table_documents( + "Show order distribution across markets.", + documents, + ) + + assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] @pytest.mark.asyncio -async def test_table_retrieval_returns_embedding_results_without_reranking_or_capping(): +async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ Document( - content=str({"name": f"table_{index}"}), - meta={"type": "TABLE_DESCRIPTION", "name": f"table_{index}"}, - ) - for index in range(8) + content="Raw staging audit rows with load metadata.", + meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, + score=0.99, + ), + Document( + content="Invoice sales transactions with product categories and sales value.", + meta={"type": "TABLE_DESCRIPTION", "name": "sales_invoices"}, + score=0.8, + ), + Document( + content="Product catalog with product names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "products"}, + score=0.7, + ), + Document( + content="Customer account master data.", + meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, + score=0.6, + ), + Document( + content="Sales regions and market hierarchy.", + meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, + score=0.5, + ), + Document( + content="Exchange rate lookup by currency.", + meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, + score=0.4, + ), ] class Retriever: @@ -47,17 +192,60 @@ async def run(self, query_embedding, filters): return {"documents": documents} result = await table_retrieval( + query="What is the distribution of sales across product categories?", embedding={"embedding": [0.1, 0.2]}, project_id="project-1", tables=[], table_retriever=Retriever(), ) - assert result["documents"] == documents + selected_names = [document.meta["name"] for document in result["documents"]] + assert 1 <= len(selected_names) <= 5 + assert "staging_audit" not in selected_names + + +def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): + transaction_source = Document( + content="Invoice transaction fact rows with customer id and invoice amount.", + meta={"type": "TABLE_DESCRIPTION", "name": "invoice_fact"}, + score=0.95, + ) + reference_source = Document( + content="Customer master reference directory with customer names and accounts.", + meta={"type": "TABLE_DESCRIPTION", "name": "customer_master"}, + score=0.7, + ) + + documents = _rerank_table_documents( + "List customer names without duplicates.", + [transaction_source, reference_source], + ) + + assert documents[0].meta["name"] == "customer_master" + + +def test_rerank_table_documents_prefers_transaction_source_for_metric_question(): + reference_source = Document( + content="Product catalog reference table with names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "product_master"}, + score=0.95, + ) + transaction_source = Document( + content="Sales transaction fact table with product, amount, and revenue.", + meta={"type": "TABLE_DESCRIPTION", "name": "sales_fact"}, + score=0.7, + ) + + documents = _rerank_table_documents( + "Show total sales amount by product.", + [reference_source, transaction_source], + ) + + assert documents[0].meta["name"] == "sales_fact" @pytest.mark.asyncio -async def test_table_retrieval_fetches_explicit_table_descriptions_with_project_scope(): +async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: def __init__(self): self.filters = None @@ -69,6 +257,7 @@ async def run(self, query_embedding, filters): retriever = Retriever() await table_retrieval( + query="show rows", embedding={}, project_id="project-1", tables=["orders"], @@ -86,30 +275,7 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_table_retrieval_without_embedding_or_explicit_tables_returns_empty(): - class Retriever: - def __init__(self): - self.called = False - - async def run(self, query_embedding, filters): - self.called = True - return {"documents": []} - - retriever = Retriever() - - result = await table_retrieval( - embedding={}, - project_id="project-1", - tables=[], - table_retriever=retriever, - ) - - assert result == {"documents": []} - assert not retriever.called - - -@pytest.mark.asyncio -async def test_dbschema_retrieval_loads_schema_for_retrieved_tables_with_project_scope(): +async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: def __init__(self): self.filters = None @@ -127,13 +293,24 @@ async def run(self, query_embedding, filters): } ), meta={"type": "TABLE_SCHEMA", "name": "orders"}, - ) + ), + Document( + content=str( + { + "type": "TABLE", + "name": "customers", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "customers"}, + ), ] } retriever = Retriever() documents = await dbschema_retrieval( + query="total orders", table_retrieval={ "documents": [ Document( @@ -146,24 +323,19 @@ async def run(self, query_embedding, filters): dbschema_retriever=retriever, ) - assert [document.meta["name"] for document in documents] == ["orders"] + assert [document.meta["name"] for document in documents] == ["orders", "customers"] assert retriever.filters == { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - { - "operator": "OR", - "conditions": [ - {"field": "name", "operator": "==", "value": "orders"} - ], - }, {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "name", "operator": "in", "value": ["orders"]}, ], } @pytest.mark.asyncio -async def test_dbschema_retrieval_returns_empty_when_no_tables_are_retrieved(): +async def test_dbschema_retrieval_does_not_load_full_schema_for_unmatched_question(): class Retriever: def __init__(self): self.called = False @@ -175,6 +347,7 @@ async def run(self, query_embedding, filters): retriever = Retriever() documents = await dbschema_retrieval( + query="show top customers by invoice amount", table_retrieval={"documents": []}, project_id="project-1", dbschema_retriever=retriever, @@ -191,20 +364,20 @@ def encode(self, value): result = check_using_db_schemas_without_pruning( construct_db_schemas=[ - { - "type": "TABLE", - "name": "orders", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "amount", - "data_type": "DOUBLE", - "comment": "", - "is_primary_key": False, - } - ], - "properties": {}, + { + "type": "TABLE", + "name": "orders", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, "primaryKey": "", } ], @@ -216,3 +389,37 @@ def encode(self, value): assert result["db_schemas"] == [] assert result["tokens"] > 0 + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_uses_explicit_tables_as_scope(): + class Retriever: + def __init__(self): + self.filters = None + + async def run(self, query_embedding, filters): + self.filters = filters + return {"documents": []} + + retriever = Retriever() + + await dbschema_retrieval( + query="show failed repairs", + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=retriever, + tables=["dbo.failure_patterns", "dbo_failure_patterns"], + ) + + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + { + "field": "name", + "operator": "in", + "value": ["dbo.failure_patterns", "dbo_failure_patterns"], + }, + ], + } From bf25e1362766eba148f9c348c86957e201ddb266 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 16 Jul 2026 03:33:41 +0530 Subject: [PATCH 0585/1087] Restore scoped ask schema retrieval flow --- .../retrieval/db_schema_retrieval.py | 7 +- wren-ai-service/src/web/v1/services/ask.py | 66 ++-- .../retrieval/test_db_schema_retrieval.py | 323 +++++------------- 3 files changed, 124 insertions(+), 272 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..32cdc0072b 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -553,7 +553,6 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query - query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -581,14 +580,10 @@ async def table_retrieval( ) if embedding: - results = await table_retriever.run( + return await table_retriever.run( query_embedding=embedding.get("embedding"), filters=base_filters, ) - results["documents"] = _select_relevant_table_documents( - query, results.get("documents") or [] - ) - return results if tables: logger.info("Loading explicit table descriptions: %s", tables) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c4a7e11b14..d6188cb085 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -161,6 +161,7 @@ def __init__( max_sql_correction_retries: int = 3, pipeline_timeout_seconds: int = 90, schema_retrieval_timeout_seconds: int = 180, + allow_schema_sql_shortcuts: bool = False, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -180,6 +181,7 @@ def __init__( self._enable_column_pruning = enable_column_pruning self._pipeline_timeout_seconds = pipeline_timeout_seconds self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds + self._allow_schema_sql_shortcuts = allow_schema_sql_shortcuts self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries @@ -5981,9 +5983,11 @@ async def ask( table_names, ) - if ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, + if self._allow_schema_sql_shortcuts and ( + ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, + ) ): ask_result = self._build_validated_ask_result_from_sql( ranked_measure_sql, @@ -6007,8 +6011,10 @@ async def ask( return results invalid_sql = ranked_measure_sql - if table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls + if self._allow_schema_sql_shortcuts and ( + table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ) ): ask_result = self._build_validated_ask_result_from_sql( table_question_sql, @@ -6032,8 +6038,10 @@ async def ask( return results invalid_sql = table_question_sql - if explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls + if self._allow_schema_sql_shortcuts and ( + explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ) ): explicit_sql, explicit_table_name = explicit_table_preview if explicit_table_name not in table_names: @@ -6060,7 +6068,7 @@ async def ask( return results invalid_sql = explicit_sql - if documents and ( + if self._allow_schema_sql_shortcuts and documents and ( deterministic_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6116,7 +6124,11 @@ async def ask( ) sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): + if ( + self._allow_schema_sql_shortcuts + and not explicit_table_names + and self._is_direct_heuristic_sql_query(user_query) + ): self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -6171,8 +6183,10 @@ async def ask( invalid_sql = heuristic_sql error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - if explicit_group_count_sql := self._build_explicit_group_count_sql( - user_query + if self._allow_schema_sql_shortcuts and ( + explicit_group_count_sql := self._build_explicit_group_count_sql( + user_query + ) ): invalid_sql = explicit_group_count_sql rephrased_question = user_query @@ -6633,7 +6647,7 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( ranked_measure_sql := self._build_schema_ranked_measure_sql( user_query, table_ddls, @@ -6654,7 +6668,7 @@ async def ask( invalid_sql = ranked_measure_sql error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6674,7 +6688,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6698,7 +6712,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6721,7 +6735,8 @@ async def ask( ) if ( - not api_results + self._allow_schema_sql_shortcuts + and not api_results and self._is_data_analysis_query(user_query) and ( schema_grounded_sql := self._build_schema_grounded_analytics_sql( @@ -6746,7 +6761,7 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if not api_results and any( + if self._allow_schema_sql_shortcuts and not api_results and any( term in user_query.lower() for term in ( "pcb", @@ -6779,7 +6794,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6802,7 +6817,8 @@ async def ask( ) should_retry_full_schema = ( - not api_results + self._allow_schema_sql_shortcuts + and not api_results and self._get_metadata_question_kind(user_query) and "db_schema_retrieval" in self._pipelines and not request_explicit_table_names @@ -6916,8 +6932,10 @@ async def ask( return results if not documents: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names + if self._allow_schema_sql_shortcuts and ( + heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ) ): logger.info( "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", @@ -7316,8 +7334,10 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names + if self._allow_schema_sql_shortcuts and ( + heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ) ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7c7d46ce89..77591333fa 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,158 +2,54 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( - _is_project_wide_analysis_query, - _rerank_table_documents, - _select_relevant_table_documents, check_using_db_schemas_without_pruning, dbschema_retrieval, - expand_business_terms_for_retrieval, + embedding, table_retrieval, ) -def test_project_wide_analysis_query_includes_broad_ranking_questions(): - assert _is_project_wide_analysis_query( - "Which projects have the highest number of completed questions?" - ) - - -def test_project_wide_analysis_query_ignores_empty_query(): - assert not _is_project_wide_analysis_query("") - - -def test_expand_business_terms_for_retrieval_adds_generic_sales_order_terms(): - query = "Show top customers by invoice amount" - - expanded_query = expand_business_terms_for_retrieval(query) - - assert query in expanded_query - assert "transaction purchase billing account geography" in expanded_query - assert "money exchange currency" in expanded_query - - -def test_expand_business_terms_for_retrieval_adds_generic_currency_market_terms(): - query = "Show invoice distribution by currency across markets" +class RecordingEmbedder: + def __init__(self): + self.query = None - expanded_query = expand_business_terms_for_retrieval(query) + async def run(self, query): + self.query = query + return {"embedding": [0.1, 0.2]} - assert query in expanded_query - assert "money exchange currency" in expanded_query - - -def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): - query = "Explain what this workspace does" - - assert expand_business_terms_for_retrieval(query) == query - - -def test_rerank_table_documents_prefers_question_relevant_table_text(): - generic_stage = Document( - content="Generic imported staging records with product labels.", - meta={"type": "TABLE_DESCRIPTION", "name": "generic_stage_load"}, - score=0.99, - ) - order_region_table = Document( - content="Business transactions grouped by customer geography and amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "business_transactions"}, - score=0.01, - ) - - documents = _rerank_table_documents( - "Show order distribution across regions.", - [generic_stage, order_region_table], - ) - assert documents[0].meta["name"] == "business_transactions" - - -def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): - test_load = Document( - content="Raw test load rows for order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ) - order_market_table = Document( - content="New order transaction records with market and customer fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.45, - ) - - documents = _rerank_table_documents( - "Show order distribution across markets.", - [test_load, order_market_table], - ) - - assert documents[0].meta["name"] == "dbo_xStageNewOrders" - - -def test_select_relevant_table_documents_limits_weak_extra_candidates(): - documents = [ - Document( - content="Invoice transactions with product, customer, currency, and amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "invoices"}, - score=0.92, - ), - Document( - content="Product catalog with product names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "products"}, - score=0.86, - ), - Document( - content="Customer account master data.", - meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, - score=0.82, - ), - Document( - content="Exchange rate lookup by currency.", - meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, - score=0.78, - ), - Document( - content="Sales regions and market hierarchy.", - meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, - score=0.74, - ), - Document( - content="Raw staging audit rows with load metadata.", - meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, - score=0.99, - ), - ] +@pytest.mark.asyncio +async def test_embedding_uses_question_and_histories_without_query_expansion(): + embedder = RecordingEmbedder() - selected = _select_relevant_table_documents( - "Show invoice distribution by currency across markets", - documents, + result = await embedding( + query="Show top customers by invoice amount", + embedder=embedder, + histories=[], ) - assert 1 <= len(selected) <= 5 - assert "staging_audit" not in [document.meta["name"] for document in selected] + assert result == {"embedding": [0.1, 0.2]} + assert embedder.query == "\nShow top customers by invoice amount" + assert "transaction purchase billing" not in embedder.query -def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): - documents = [ - Document( - content="Raw test load rows with order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ), - Document( - content="New order transaction records with market and customer details.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.4, - ), - ] +@pytest.mark.asyncio +async def test_embedding_skips_vector_lookup_for_explicit_tables(): + embedder = RecordingEmbedder() - selected = _select_relevant_table_documents( - "Show order distribution across markets.", - documents, + result = await embedding( + query="Show rows", + embedder=embedder, + histories=[], + tables=["orders"], ) - assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] + assert result == {} + assert embedder.query is None @pytest.mark.asyncio -async def test_table_retrieval_caps_embedding_results_before_schema_loading(): +async def test_table_retrieval_uses_vector_retriever_results_without_local_rerank(): documents = [ Document( content="Raw staging audit rows with load metadata.", @@ -165,83 +61,34 @@ async def test_table_retrieval_caps_embedding_results_before_schema_loading(): meta={"type": "TABLE_DESCRIPTION", "name": "sales_invoices"}, score=0.8, ), - Document( - content="Product catalog with product names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "products"}, - score=0.7, - ), - Document( - content="Customer account master data.", - meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, - score=0.6, - ), - Document( - content="Sales regions and market hierarchy.", - meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, - score=0.5, - ), - Document( - content="Exchange rate lookup by currency.", - meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, - score=0.4, - ), ] class Retriever: + def __init__(self): + self.filters = None + async def run(self, query_embedding, filters): + self.filters = filters return {"documents": documents} + retriever = Retriever() + result = await table_retrieval( query="What is the distribution of sales across product categories?", embedding={"embedding": [0.1, 0.2]}, project_id="project-1", tables=[], - table_retriever=Retriever(), - ) - - selected_names = [document.meta["name"] for document in result["documents"]] - assert 1 <= len(selected_names) <= 5 - assert "staging_audit" not in selected_names - - -def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): - transaction_source = Document( - content="Invoice transaction fact rows with customer id and invoice amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "invoice_fact"}, - score=0.95, - ) - reference_source = Document( - content="Customer master reference directory with customer names and accounts.", - meta={"type": "TABLE_DESCRIPTION", "name": "customer_master"}, - score=0.7, - ) - - documents = _rerank_table_documents( - "List customer names without duplicates.", - [transaction_source, reference_source], - ) - - assert documents[0].meta["name"] == "customer_master" - - -def test_rerank_table_documents_prefers_transaction_source_for_metric_question(): - reference_source = Document( - content="Product catalog reference table with names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "product_master"}, - score=0.95, - ) - transaction_source = Document( - content="Sales transaction fact table with product, amount, and revenue.", - meta={"type": "TABLE_DESCRIPTION", "name": "sales_fact"}, - score=0.7, - ) - - documents = _rerank_table_documents( - "Show total sales amount by product.", - [reference_source, transaction_source], + table_retriever=retriever, ) - assert documents[0].meta["name"] == "sales_fact" + assert result["documents"] == documents + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } @pytest.mark.asyncio @@ -275,7 +122,7 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_dbschema_retrieval_loads_selected_active_project_schema(): +async def test_dbschema_retrieval_loads_only_selected_active_project_schema(): class Retriever: def __init__(self): self.filters = None @@ -293,17 +140,7 @@ async def run(self, query_embedding, filters): } ), meta={"type": "TABLE_SCHEMA", "name": "orders"}, - ), - Document( - content=str( - { - "type": "TABLE", - "name": "customers", - "columns": [], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "customers"}, - ), + ) ] } @@ -323,7 +160,7 @@ async def run(self, query_embedding, filters): dbschema_retriever=retriever, ) - assert [document.meta["name"] for document in documents] == ["orders", "customers"] + assert [document.meta["name"] for document in documents] == ["orders"] assert retriever.filters == { "operator": "AND", "conditions": [ @@ -357,40 +194,6 @@ async def run(self, query_embedding, filters): assert not retriever.called -def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): - class Encoding: - def encode(self, value): - return value.split() - - result = check_using_db_schemas_without_pruning( - construct_db_schemas=[ - { - "type": "TABLE", - "name": "orders", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "amount", - "data_type": "DOUBLE", - "comment": "", - "is_primary_key": False, - } - ], - "properties": {}, - "primaryKey": "", - } - ], - dbschema_retrieval=[], - encoding=Encoding(), - enable_column_pruning=True, - context_window_size=1000, - ) - - assert result["db_schemas"] == [] - assert result["tokens"] > 0 - - @pytest.mark.asyncio async def test_dbschema_retrieval_uses_explicit_tables_as_scope(): class Retriever: @@ -423,3 +226,37 @@ async def run(self, query_embedding, filters): }, ], } + + +def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "orders", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=True, + context_window_size=1000, + ) + + assert result["db_schemas"] == [] + assert result["tokens"] > 0 From 789c5d8db3834e88c2915882cac66be7933cb7d3 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 03:39:01 +0530 Subject: [PATCH 0586/1087] Align calculated field expressions with docs --- .../src/apollo/client/graphql/__types__.ts | 2 - wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 26 ++++++- .../apollo/server/mdl/test/mdlBuilder.test.ts | 70 ++++++++++++++++++- wren-ui/src/apollo/server/schema.ts | 2 - 4 files changed, 94 insertions(+), 6 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 72135acbd4..f3da8f04fd 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -604,9 +604,7 @@ export enum ExpressionName { AVG = 'AVG', CBRT = 'CBRT', CEIL = 'CEIL', - CEILING = 'CEILING', COUNT = 'COUNT', - COUNT_IF = 'COUNT_IF', EXP = 'EXP', FLOOR = 'FLOOR', LENGTH = 'LENGTH', diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index bc166e7f55..86e0a6c87a 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -23,6 +23,24 @@ const logger = getLogger('MDLBuilder'); logger.level = 'debug'; const config = getConfig(); +const DOCUMENTED_CALCULATED_FIELD_FUNCTIONS = new Map([ + ['ABS', 'abs'], + ['AVG', 'avg'], + ['COUNT', 'count'], + ['MAX', 'max'], + ['MIN', 'min'], + ['SUM', 'sum'], + ['CBRT', 'cbrt'], + ['CEIL', 'ceil'], + ['EXP', 'exp'], + ['FLOOR', 'floor'], + ['LN', 'ln'], + ['LOG10', 'log10'], + ['ROUND', 'round'], + ['SIGN', 'sign'], + ['LENGTH', 'length'], + ['REVERSE', 'reverse'], +]); export interface MDLBuilderBuildFromOptions { project: Project; @@ -534,7 +552,13 @@ export class MDLBuilder implements IMDLBuilder { if (fieldExpression.length !== lineage.length) { return null; } - return `${column.aggregation}(${fieldExpression.join('.')})`; + const functionName = DOCUMENTED_CALCULATED_FIELD_FUNCTIONS.get( + String(column.aggregation).toUpperCase(), + ); + if (!functionName) { + return null; + } + return `${functionName}(${fieldExpression.join('.')})`; } protected getRelationCondition(relation: RelationInfo): string { diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 97c1c45b0f..5f41ee3f9c 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -852,7 +852,7 @@ describe('MDLBuilder', () => { displayName: 'total_payment', referenceName: 'total_payment', sourceColumnName: 'total_payment', - aggregation: 'sum', + aggregation: 'SUM', lineage: JSON.stringify([1, 2, 8]), customExpression: null, type: 'FLOAT', @@ -993,6 +993,74 @@ describe('MDLBuilder', () => { ); }); + it('should skip calculated fields with expressions outside the documented function list.', () => { + const models = [ + { + id: 1, + projectId: 1, + displayName: 'orders', + sourceTableName: 'orders', + referenceName: 'orders', + refSql: 'SELECT * FROM orders', + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const columns = [ + { + id: 1, + modelId: 1, + isCalculated: false, + displayName: 'id', + referenceName: 'id', + sourceColumnName: 'id', + aggregation: null, + lineage: null, + customExpression: null, + type: 'INTEGER', + notNull: true, + isPk: true, + properties: null, + }, + { + id: 2, + modelId: 1, + isCalculated: true, + displayName: 'unsupported_count_if', + referenceName: 'unsupported_count_if', + sourceColumnName: 'unsupported_count_if', + aggregation: 'COUNT_IF', + lineage: JSON.stringify([1]), + customExpression: null, + type: 'BIGINT', + notNull: false, + isPk: false, + properties: null, + }, + ] as ModelColumn[]; + const builderOptions = { + project: { + schema: 'public', + catalog: 'wrenai', + }, + models, + columns, + relations: [], + relatedModels: models, + relatedColumns: columns, + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + const ordersModel = manifest.models.find((m) => m.name === 'orders'); + expect( + ordersModel.columns.find((c) => c.name === 'unsupported_count_if'), + ).toBeUndefined(); + }); + it.each(Object.values(DataSourceName))( `should return correct data source type`, (type) => { diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 1c3dad1cc7..21c9f214ea 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -110,13 +110,11 @@ export const typeDefs = gql` ABS AVG COUNT - COUNT_IF MAX MIN SUM CBRT CEIL - CEILING EXP FLOOR LN From f0b14100df0b9b24447a18b1014097bf6c5f9476 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 16 Jul 2026 13:46:30 +0530 Subject: [PATCH 0587/1087] Revert "Restore scoped ask schema retrieval flow" This reverts commit bf25e1362766eba148f9c348c86957e201ddb266. --- .../retrieval/db_schema_retrieval.py | 7 +- wren-ai-service/src/web/v1/services/ask.py | 66 ++-- .../retrieval/test_db_schema_retrieval.py | 323 +++++++++++++----- 3 files changed, 272 insertions(+), 124 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 32cdc0072b..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -553,6 +553,7 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query + query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -580,10 +581,14 @@ async def table_retrieval( ) if embedding: - return await table_retriever.run( + results = await table_retriever.run( query_embedding=embedding.get("embedding"), filters=base_filters, ) + results["documents"] = _select_relevant_table_documents( + query, results.get("documents") or [] + ) + return results if tables: logger.info("Loading explicit table descriptions: %s", tables) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d6188cb085..c4a7e11b14 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -161,7 +161,6 @@ def __init__( max_sql_correction_retries: int = 3, pipeline_timeout_seconds: int = 90, schema_retrieval_timeout_seconds: int = 180, - allow_schema_sql_shortcuts: bool = False, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -181,7 +180,6 @@ def __init__( self._enable_column_pruning = enable_column_pruning self._pipeline_timeout_seconds = pipeline_timeout_seconds self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds - self._allow_schema_sql_shortcuts = allow_schema_sql_shortcuts self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries @@ -5983,11 +5981,9 @@ async def ask( table_names, ) - if self._allow_schema_sql_shortcuts and ( - ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, - ) + if ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, ): ask_result = self._build_validated_ask_result_from_sql( ranked_measure_sql, @@ -6011,10 +6007,8 @@ async def ask( return results invalid_sql = ranked_measure_sql - if self._allow_schema_sql_shortcuts and ( - table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ) + if table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls ): ask_result = self._build_validated_ask_result_from_sql( table_question_sql, @@ -6038,10 +6032,8 @@ async def ask( return results invalid_sql = table_question_sql - if self._allow_schema_sql_shortcuts and ( - explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ) + if explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls ): explicit_sql, explicit_table_name = explicit_table_preview if explicit_table_name not in table_names: @@ -6068,7 +6060,7 @@ async def ask( return results invalid_sql = explicit_sql - if self._allow_schema_sql_shortcuts and documents and ( + if documents and ( deterministic_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6124,11 +6116,7 @@ async def ask( ) sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - if ( - self._allow_schema_sql_shortcuts - and not explicit_table_names - and self._is_direct_heuristic_sql_query(user_query) - ): + if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -6183,10 +6171,8 @@ async def ask( invalid_sql = heuristic_sql error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - if self._allow_schema_sql_shortcuts and ( - explicit_group_count_sql := self._build_explicit_group_count_sql( - user_query - ) + if explicit_group_count_sql := self._build_explicit_group_count_sql( + user_query ): invalid_sql = explicit_group_count_sql rephrased_question = user_query @@ -6647,7 +6633,7 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( ranked_measure_sql := self._build_schema_ranked_measure_sql( user_query, table_ddls, @@ -6668,7 +6654,7 @@ async def ask( invalid_sql = ranked_measure_sql error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6688,7 +6674,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6712,7 +6698,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6735,8 +6721,7 @@ async def ask( ) if ( - self._allow_schema_sql_shortcuts - and not api_results + not api_results and self._is_data_analysis_query(user_query) and ( schema_grounded_sql := self._build_schema_grounded_analytics_sql( @@ -6761,7 +6746,7 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if self._allow_schema_sql_shortcuts and not api_results and any( + if not api_results and any( term in user_query.lower() for term in ( "pcb", @@ -6794,7 +6779,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6817,8 +6802,7 @@ async def ask( ) should_retry_full_schema = ( - self._allow_schema_sql_shortcuts - and not api_results + not api_results and self._get_metadata_question_kind(user_query) and "db_schema_retrieval" in self._pipelines and not request_explicit_table_names @@ -6932,10 +6916,8 @@ async def ask( return results if not documents: - if self._allow_schema_sql_shortcuts and ( - heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ) + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names ): logger.info( "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", @@ -7334,10 +7316,8 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if self._allow_schema_sql_shortcuts and ( - heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ) + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 77591333fa..7c7d46ce89 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,54 +2,158 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( + _is_project_wide_analysis_query, + _rerank_table_documents, + _select_relevant_table_documents, check_using_db_schemas_without_pruning, dbschema_retrieval, - embedding, + expand_business_terms_for_retrieval, table_retrieval, ) -class RecordingEmbedder: - def __init__(self): - self.query = None +def test_project_wide_analysis_query_includes_broad_ranking_questions(): + assert _is_project_wide_analysis_query( + "Which projects have the highest number of completed questions?" + ) - async def run(self, query): - self.query = query - return {"embedding": [0.1, 0.2]} +def test_project_wide_analysis_query_ignores_empty_query(): + assert not _is_project_wide_analysis_query("") -@pytest.mark.asyncio -async def test_embedding_uses_question_and_histories_without_query_expansion(): - embedder = RecordingEmbedder() - result = await embedding( - query="Show top customers by invoice amount", - embedder=embedder, - histories=[], +def test_expand_business_terms_for_retrieval_adds_generic_sales_order_terms(): + query = "Show top customers by invoice amount" + + expanded_query = expand_business_terms_for_retrieval(query) + + assert query in expanded_query + assert "transaction purchase billing account geography" in expanded_query + assert "money exchange currency" in expanded_query + + +def test_expand_business_terms_for_retrieval_adds_generic_currency_market_terms(): + query = "Show invoice distribution by currency across markets" + + expanded_query = expand_business_terms_for_retrieval(query) + + assert query in expanded_query + assert "money exchange currency" in expanded_query + + +def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): + query = "Explain what this workspace does" + + assert expand_business_terms_for_retrieval(query) == query + + +def test_rerank_table_documents_prefers_question_relevant_table_text(): + generic_stage = Document( + content="Generic imported staging records with product labels.", + meta={"type": "TABLE_DESCRIPTION", "name": "generic_stage_load"}, + score=0.99, + ) + order_region_table = Document( + content="Business transactions grouped by customer geography and amount.", + meta={"type": "TABLE_DESCRIPTION", "name": "business_transactions"}, + score=0.01, ) - assert result == {"embedding": [0.1, 0.2]} - assert embedder.query == "\nShow top customers by invoice amount" - assert "transaction purchase billing" not in embedder.query + documents = _rerank_table_documents( + "Show order distribution across regions.", + [generic_stage, order_region_table], + ) + assert documents[0].meta["name"] == "business_transactions" -@pytest.mark.asyncio -async def test_embedding_skips_vector_lookup_for_explicit_tables(): - embedder = RecordingEmbedder() - result = await embedding( - query="Show rows", - embedder=embedder, - histories=[], - tables=["orders"], +def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): + test_load = Document( + content="Raw test load rows for order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ) + order_market_table = Document( + content="New order transaction records with market and customer fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.45, ) - assert result == {} - assert embedder.query is None + documents = _rerank_table_documents( + "Show order distribution across markets.", + [test_load, order_market_table], + ) + + assert documents[0].meta["name"] == "dbo_xStageNewOrders" + + +def test_select_relevant_table_documents_limits_weak_extra_candidates(): + documents = [ + Document( + content="Invoice transactions with product, customer, currency, and amount.", + meta={"type": "TABLE_DESCRIPTION", "name": "invoices"}, + score=0.92, + ), + Document( + content="Product catalog with product names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "products"}, + score=0.86, + ), + Document( + content="Customer account master data.", + meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, + score=0.82, + ), + Document( + content="Exchange rate lookup by currency.", + meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, + score=0.78, + ), + Document( + content="Sales regions and market hierarchy.", + meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, + score=0.74, + ), + Document( + content="Raw staging audit rows with load metadata.", + meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, + score=0.99, + ), + ] + + selected = _select_relevant_table_documents( + "Show invoice distribution by currency across markets", + documents, + ) + + assert 1 <= len(selected) <= 5 + assert "staging_audit" not in [document.meta["name"] for document in selected] + + +def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): + documents = [ + Document( + content="Raw test load rows with order market fields.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, + score=0.99, + ), + Document( + content="New order transaction records with market and customer details.", + meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, + score=0.4, + ), + ] + + selected = _select_relevant_table_documents( + "Show order distribution across markets.", + documents, + ) + + assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] @pytest.mark.asyncio -async def test_table_retrieval_uses_vector_retriever_results_without_local_rerank(): +async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ Document( content="Raw staging audit rows with load metadata.", @@ -61,34 +165,83 @@ async def test_table_retrieval_uses_vector_retriever_results_without_local_reran meta={"type": "TABLE_DESCRIPTION", "name": "sales_invoices"}, score=0.8, ), + Document( + content="Product catalog with product names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "products"}, + score=0.7, + ), + Document( + content="Customer account master data.", + meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, + score=0.6, + ), + Document( + content="Sales regions and market hierarchy.", + meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, + score=0.5, + ), + Document( + content="Exchange rate lookup by currency.", + meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, + score=0.4, + ), ] class Retriever: - def __init__(self): - self.filters = None - async def run(self, query_embedding, filters): - self.filters = filters return {"documents": documents} - retriever = Retriever() - result = await table_retrieval( query="What is the distribution of sales across product categories?", embedding={"embedding": [0.1, 0.2]}, project_id="project-1", tables=[], - table_retriever=retriever, + table_retriever=Retriever(), ) - assert result["documents"] == documents - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - ], - } + selected_names = [document.meta["name"] for document in result["documents"]] + assert 1 <= len(selected_names) <= 5 + assert "staging_audit" not in selected_names + + +def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): + transaction_source = Document( + content="Invoice transaction fact rows with customer id and invoice amount.", + meta={"type": "TABLE_DESCRIPTION", "name": "invoice_fact"}, + score=0.95, + ) + reference_source = Document( + content="Customer master reference directory with customer names and accounts.", + meta={"type": "TABLE_DESCRIPTION", "name": "customer_master"}, + score=0.7, + ) + + documents = _rerank_table_documents( + "List customer names without duplicates.", + [transaction_source, reference_source], + ) + + assert documents[0].meta["name"] == "customer_master" + + +def test_rerank_table_documents_prefers_transaction_source_for_metric_question(): + reference_source = Document( + content="Product catalog reference table with names and categories.", + meta={"type": "TABLE_DESCRIPTION", "name": "product_master"}, + score=0.95, + ) + transaction_source = Document( + content="Sales transaction fact table with product, amount, and revenue.", + meta={"type": "TABLE_DESCRIPTION", "name": "sales_fact"}, + score=0.7, + ) + + documents = _rerank_table_documents( + "Show total sales amount by product.", + [reference_source, transaction_source], + ) + + assert documents[0].meta["name"] == "sales_fact" @pytest.mark.asyncio @@ -122,7 +275,7 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_dbschema_retrieval_loads_only_selected_active_project_schema(): +async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: def __init__(self): self.filters = None @@ -140,7 +293,17 @@ async def run(self, query_embedding, filters): } ), meta={"type": "TABLE_SCHEMA", "name": "orders"}, - ) + ), + Document( + content=str( + { + "type": "TABLE", + "name": "customers", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "customers"}, + ), ] } @@ -160,7 +323,7 @@ async def run(self, query_embedding, filters): dbschema_retriever=retriever, ) - assert [document.meta["name"] for document in documents] == ["orders"] + assert [document.meta["name"] for document in documents] == ["orders", "customers"] assert retriever.filters == { "operator": "AND", "conditions": [ @@ -194,6 +357,40 @@ async def run(self, query_embedding, filters): assert not retriever.called +def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "orders", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=True, + context_window_size=1000, + ) + + assert result["db_schemas"] == [] + assert result["tokens"] > 0 + + @pytest.mark.asyncio async def test_dbschema_retrieval_uses_explicit_tables_as_scope(): class Retriever: @@ -226,37 +423,3 @@ async def run(self, query_embedding, filters): }, ], } - - -def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): - class Encoding: - def encode(self, value): - return value.split() - - result = check_using_db_schemas_without_pruning( - construct_db_schemas=[ - { - "type": "TABLE", - "name": "orders", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "amount", - "data_type": "DOUBLE", - "comment": "", - "is_primary_key": False, - } - ], - "properties": {}, - "primaryKey": "", - } - ], - dbschema_retrieval=[], - encoding=Encoding(), - enable_column_pruning=True, - context_window_size=1000, - ) - - assert result["db_schemas"] == [] - assert result["tokens"] > 0 From 49eb8081a95a900eb61897ece2d6225ca30e86ea Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 14:27:43 +0530 Subject: [PATCH 0588/1087] Use semantic metadata for table retrieval intent --- wren-ai-service/src/web/v1/services/ask.py | 94 +------------------ .../pytest/services/test_ask_sales_sql.py | 16 ++-- 2 files changed, 9 insertions(+), 101 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c4a7e11b14..33877768bb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -362,80 +362,7 @@ def _should_reuse_historical_question_sql( return False def _rewrite_query_for_text_to_sql(self, query: str) -> str: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return query - - guidance: list[str] = [] - - if any( - term in normalized - for term in ("chart", "bar chart", "line chart", "pie chart", "graph") - ): - guidance.append( - "Return SQL only for the aggregated dataset required to build the requested chart." - ) - - if any( - term in normalized - for term in ( - "failure category", - "failure categories", - "common failure", - "common failures", - "failure code", - "top 10", - "most common", - ) - ): - guidance.append( - "Use an exposed failure category, failure name, or failure code field from the schema and return that dimension with a count metric." - ) - - if any( - term in normalized - for term in ("monthly", "last 12 months", "last month", "trend", "volume") - ): - guidance.append( - "Use a real timestamp column from the schema and aggregate results by calendar month when a monthly trend is requested." - ) - - if re.search( - r"\b(?:by|across|per|each|grouped by|group by)\s+[a-z][a-z0-9 _-]*", - normalized, - ) or "over time" in normalized: - guidance.append( - "Preserve explicit grouping dimensions requested by the question, such as market, region, currency, product, customer, status, type, or category, using only matching columns exposed in the provided schema." - ) - - if any(term in normalized for term in ("currency", "currencies", "fx")): - guidance.append( - "For currency questions, use an exposed currency, money, exchange, or FX code/name column from the schema and group by it." - ) - - if any( - term in normalized - for term in ( - "amount", - "cost", - "revenue", - "sales value", - "sum", - "total", - "value", - ) - ) and not any( - term in normalized - for term in ("count", "how many", "number of records", "record count") - ): - guidance.append( - "When the question asks for total, sum, amount, value, revenue, or cost, aggregate an exposed numeric measure with SUM; use COUNT only for record-count questions." - ) - - if not guidance: - return query - - return f"{query}\n\nSQL generation guidance:\n- " + "\n- ".join(guidance) + return query def _schema_contains( self, @@ -1898,18 +1825,6 @@ def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: and table_name not in table_names ): table_names.append(table_name) - if re.search( - r"\b(?:repair\s+logs?|repair\s+tickets?|board\s+models?)\b", - query or "", - flags=re.IGNORECASE, - ): - for table_name in ("repair_logs", "dbo_repair_logs"): - if table_name not in table_names: - table_names.append(table_name) - if re.search(r"\bticket\s+labels?\b", query or "", flags=re.IGNORECASE): - for table_name in ("ticket_labels", "dbo_ticket_labels"): - if table_name not in table_names: - table_names.append(table_name) return table_names def _explicit_table_name_candidates(self, table_name: str) -> list[str]: @@ -5814,16 +5729,11 @@ async def ask( request_explicit_table_names = self._normalize_explicit_table_names( ask_request.explicit_tables ) - query_explicit_table_names = self._normalize_explicit_table_names( - self._extract_explicit_table_names_from_query(user_query) - ) forced_request_explicit_table_names = self._forced_explicit_table_names( request_explicit_table_names, source="request", ) - explicit_table_names = ( - forced_request_explicit_table_names or query_explicit_table_names - ) + explicit_table_names = forced_request_explicit_table_names retrieval_table_names = explicit_table_names or None try: diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 4a2b68027e..9593fca054 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -184,15 +184,13 @@ def test_data_query_timeout_retry_does_not_allow_full_project_schema(): assert service._should_retry_selected_schema_after_retrieval_timeout(["orders"]) -def test_rewrite_query_for_text_to_sql_guides_total_amount_to_sum(): +def test_rewrite_query_for_text_to_sql_preserves_user_question(): service = AskService.__new__(AskService) - rewritten = service._rewrite_query_for_text_to_sql( - "Show total invoice amount by currency." - ) + query = "Show total invoice amount by currency." + rewritten = service._rewrite_query_for_text_to_sql(query) - assert "aggregate an exposed numeric measure with SUM" in rewritten - assert "use COUNT only for record-count questions" in rewritten + assert rewritten == query def test_validated_sql_rejects_count_for_total_amount_question(): @@ -604,7 +602,7 @@ def test_extract_explicit_table_names_from_repair_logs_phrase(): assert service._extract_explicit_table_names_from_query( "Count repair logs by failure_code in repair logs." - ) == ["repair_logs", "dbo_repair_logs"] + ) == [] def test_extract_explicit_table_names_from_pcb_repair_phrases(): @@ -612,10 +610,10 @@ def test_extract_explicit_table_names_from_pcb_repair_phrases(): assert service._extract_explicit_table_names_from_query( "How many different board models are present in the dbo.repair_logs table?" - ) == ["dbo.repair_logs", "repair_logs", "dbo_repair_logs"] + ) == ["dbo.repair_logs"] assert service._extract_explicit_table_names_from_query( "Display top 10 ticket labels." - ) == ["ticket_labels", "dbo_ticket_labels"] + ) == [] def test_explicit_table_name_candidates_include_dotted_and_short_forms(): From 43ffcfe5d86ead9ef78bb0d78e504a3a7b0eaed2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 16 Jul 2026 19:05:26 +0530 Subject: [PATCH 0589/1087] Restore legacy ask retrieval flow --- .../retrieval/db_schema_retrieval.py | 34 +++- wren-ai-service/src/web/v1/services/ask.py | 174 +++++++++++------- .../retrieval/test_db_schema_retrieval.py | 28 ++- .../pytest/services/test_ask_sales_sql.py | 25 +++ 4 files changed, 181 insertions(+), 80 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..e5091fe9cf 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -428,7 +428,32 @@ def _select_relevant_table_documents( if not reranked: return documents[:max_tables] - candidate_pool = [item for item in reranked if item[3] > 0] or reranked + non_production_terms = ( + "archive", + "backup", + "copy", + "dev", + "development", + "duplicate", + "sample", + "stage", + "staging", + "temp", + "test", + "tmp", + ) + query_mentions_non_production = _query_mentions_any(query, non_production_terms) + production_candidates = [ + item + for item in reranked + if query_mentions_non_production + or not (_retrieval_terms(_source_text(item[2])) & set(non_production_terms)) + ] + candidate_pool = ( + [item for item in production_candidates if item[3] > 0] + or production_candidates + or reranked + ) selected = [ document for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] @@ -553,7 +578,6 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query - query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -581,14 +605,10 @@ async def table_retrieval( ) if embedding: - results = await table_retriever.run( + return await table_retriever.run( query_embedding=embedding.get("embedding"), filters=base_filters, ) - results["documents"] = _select_relevant_table_documents( - query, results.get("documents") or [] - ) - return results if tables: logger.info("Loading explicit table descriptions: %s", tables) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 33877768bb..b633ac9d7b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -161,6 +161,7 @@ def __init__( max_sql_correction_retries: int = 3, pipeline_timeout_seconds: int = 90, schema_retrieval_timeout_seconds: int = 180, + allow_schema_sql_shortcuts: bool = False, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -178,6 +179,7 @@ def __init__( self._allow_sql_diagnosis = allow_sql_diagnosis self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval self._enable_column_pruning = enable_column_pruning + self._allow_schema_sql_shortcuts = allow_schema_sql_shortcuts self._pipeline_timeout_seconds = pipeline_timeout_seconds self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds self._max_histories = max_histories @@ -5471,6 +5473,7 @@ def _build_validated_ask_result_from_sql( sql: Optional[str], table_ddls: list[str], query: str | None = None, + strict_semantic_validation: bool = True, ) -> Optional[AskResult]: if isinstance(sql, str): sql = normalize_sql_direction_keywords(sql) @@ -5567,41 +5570,42 @@ def _build_validated_ask_result_from_sql( ) return None - invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( - ask_result.sql, - schema_tables, - ) - if invalid_unqualified_identifiers: - logger.warning( - "Ignoring SQL because it references unqualified fields outside the active schema. " - "invalid_identifiers=%s sql=%s", - invalid_unqualified_identifiers, - ask_result.sql, - ) + if not self._sql_references_explicit_table(ask_result.sql, query): return None - invalid_output_aliases = self._invalid_sql_output_aliases( - ask_result.sql, - schema_tables, - ) - if invalid_output_aliases: - logger.warning( - "Ignoring SQL because it aliases output fields to unavailable schema concepts. " - "invalid_aliases=%s sql=%s", - invalid_output_aliases, + if strict_semantic_validation: + invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( ask_result.sql, + schema_tables, ) - return None + if invalid_unqualified_identifiers: + logger.warning( + "Ignoring SQL because it references unqualified fields outside the active schema. " + "invalid_identifiers=%s sql=%s", + invalid_unqualified_identifiers, + ask_result.sql, + ) + return None - if not self._sql_references_explicit_table(ask_result.sql, query): - return None + invalid_output_aliases = self._invalid_sql_output_aliases( + ask_result.sql, + schema_tables, + ) + if invalid_output_aliases: + logger.warning( + "Ignoring SQL because it aliases output fields to unavailable schema concepts. " + "invalid_aliases=%s sql=%s", + invalid_output_aliases, + ask_result.sql, + ) + return None - if not self._sql_matches_question_intent( - ask_result.sql, - query, - schema_tables, - ): - return None + if not self._sql_matches_question_intent( + ask_result.sql, + query, + schema_tables, + ): + return None return ask_result @@ -6092,9 +6096,7 @@ async def ask( ) historical_question_result = [] - should_skip_pre_sql_retrieval = self._is_data_analysis_query( - user_query - ) + should_skip_pre_sql_retrieval = False if should_skip_pre_sql_retrieval: rephrased_question = user_query intent_reasoning = ( @@ -6543,7 +6545,7 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( ranked_measure_sql := self._build_schema_ranked_measure_sql( user_query, table_ddls, @@ -6564,7 +6566,7 @@ async def ask( invalid_sql = ranked_measure_sql error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6584,7 +6586,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6608,7 +6610,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6631,6 +6633,8 @@ async def ask( ) if ( + self._allow_schema_sql_shortcuts + and not api_results and self._is_data_analysis_query(user_query) and ( @@ -6656,7 +6660,7 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if not api_results and any( + if self._allow_schema_sql_shortcuts and not api_results and any( term in user_query.lower() for term in ( "pcb", @@ -6689,7 +6693,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if not api_results and ( + if self._allow_schema_sql_shortcuts and not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6763,24 +6767,26 @@ async def ask( table_names, ) - full_schema_preview = self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - full_schema_sql_candidates = ( - self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ), - full_schema_preview[0] if full_schema_preview else None, - self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ), - self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ), - self._build_schema_grounded_sales_sql( + full_schema_sql_candidates = () + if self._allow_schema_sql_shortcuts: + full_schema_preview = self._build_explicit_table_preview_sql( user_query, table_ddls - ), - ) + ) + full_schema_sql_candidates = ( + self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ), + full_schema_preview[0] if full_schema_preview else None, + self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ), + self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ), + self._build_schema_grounded_sales_sql( + user_query, table_ddls + ), + ) for full_schema_sql in full_schema_sql_candidates: if not full_schema_sql: continue @@ -6825,7 +6831,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if not documents: + if self._allow_schema_sql_shortcuts and not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): @@ -7109,6 +7115,7 @@ async def ask( sql_valid_result.get("sql"), table_ddls, sql_user_query, + strict_semantic_validation=False, ): api_results = [ask_result] else: @@ -7194,6 +7201,7 @@ async def ask( valid_generation_result.get("sql"), table_ddls, sql_user_query, + strict_semantic_validation=False, ): api_results = [ask_result] break @@ -7226,8 +7234,10 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names + if self._allow_schema_sql_shortcuts and ( + heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ) ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", @@ -7261,17 +7271,41 @@ async def ask( return results logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") + failure_message = ( + error_message + or "SQL generation could not produce a valid query for the retrieved schema." + ) + failure_code = ( + "NO_RELEVANT_SQL" + if documents or table_names + else "NO_RELEVANT_DATA" + ) if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - is_followup=True if histories else False, + if failure_code == "NO_RELEVANT_SQL": + self._ask_results[query_id] = ( + self._build_failed_text_to_sql_response( + trace_id, + failure_message, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + is_followup=True if histories else False, + code="NO_RELEVANT_SQL", + ) + ) + else: + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + is_followup=True if histories else False, + ) ) - ) if error_message or invalid_sql: logger.info( "Suppressed technical SQL failure for query_id %s. " @@ -7280,9 +7314,11 @@ async def ask( error_message, invalid_sql, ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_type"] = failure_code results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + failure_message + if failure_code == "NO_RELEVANT_SQL" + else NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE ) results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7c7d46ce89..5662058dd8 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -7,6 +7,7 @@ _select_relevant_table_documents, check_using_db_schemas_without_pruning, dbschema_retrieval, + embedding, expand_business_terms_for_retrieval, table_retrieval, ) @@ -47,6 +48,27 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query +@pytest.mark.asyncio +async def test_embedding_uses_original_query_without_business_term_expansion(): + class Embedder: + def __init__(self): + self.query = None + + async def run(self, query): + self.query = query + return {"embedding": [0.1, 0.2]} + + embedder = Embedder() + + await embedding( + query="Show top customers by invoice amount", + embedder=embedder, + histories=[], + ) + + assert embedder.query == "\nShow top customers by invoice amount" + + def test_rerank_table_documents_prefers_question_relevant_table_text(): generic_stage = Document( content="Generic imported staging records with product labels.", @@ -153,7 +175,7 @@ def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): @pytest.mark.asyncio -async def test_table_retrieval_caps_embedding_results_before_schema_loading(): +async def test_table_retrieval_returns_embedding_results_without_local_reranking(): documents = [ Document( content="Raw staging audit rows with load metadata.", @@ -199,9 +221,7 @@ async def run(self, query_embedding, filters): table_retriever=Retriever(), ) - selected_names = [document.meta["name"] for document in result["documents"]] - assert 1 <= len(selected_names) <= 5 - assert "staging_audit" not in selected_names + assert result["documents"] == documents def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index 9593fca054..a994519419 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -217,6 +217,31 @@ def test_validated_sql_rejects_count_for_total_amount_question(): assert result is None +def test_validated_llm_sql_accepts_schema_valid_count_for_total_amount_question(): + service = AskService.__new__(AskService) + + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_Invoices"."Currency" AS "Currency", ' + 'COUNT(*) AS "RecordCount" ' + 'FROM "dbo_Invoices" ' + 'GROUP BY "dbo_Invoices"."Currency"' + ), + [ + """ + CREATE TABLE dbo_Invoices ( + Currency VARCHAR, + InvoiceAmount DOUBLE + ); + """ + ], + "Show total invoice amount by currency.", + strict_semantic_validation=False, + ) + + assert result is not None + + def test_validated_sql_rejects_detail_rows_for_customer_order_count_question(): service = AskService.__new__(AskService) From 0a5faf93e2ba1f23b1baad7b3e7663f9b27bfb3b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 16 Jul 2026 19:30:42 +0530 Subject: [PATCH 0590/1087] Revert "Restore legacy ask retrieval flow" This reverts commit 43ffcfe5d86ead9ef78bb0d78e504a3a7b0eaed2. --- .../retrieval/db_schema_retrieval.py | 34 +--- wren-ai-service/src/web/v1/services/ask.py | 174 +++++++----------- .../retrieval/test_db_schema_retrieval.py | 28 +-- .../pytest/services/test_ask_sales_sql.py | 25 --- 4 files changed, 80 insertions(+), 181 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e5091fe9cf..b58e771490 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -428,32 +428,7 @@ def _select_relevant_table_documents( if not reranked: return documents[:max_tables] - non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "stage", - "staging", - "temp", - "test", - "tmp", - ) - query_mentions_non_production = _query_mentions_any(query, non_production_terms) - production_candidates = [ - item - for item in reranked - if query_mentions_non_production - or not (_retrieval_terms(_source_text(item[2])) & set(non_production_terms)) - ] - candidate_pool = ( - [item for item in production_candidates if item[3] > 0] - or production_candidates - or reranked - ) + candidate_pool = [item for item in reranked if item[3] > 0] or reranked selected = [ document for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] @@ -578,6 +553,7 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query + query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -605,10 +581,14 @@ async def table_retrieval( ) if embedding: - return await table_retriever.run( + results = await table_retriever.run( query_embedding=embedding.get("embedding"), filters=base_filters, ) + results["documents"] = _select_relevant_table_documents( + query, results.get("documents") or [] + ) + return results if tables: logger.info("Loading explicit table descriptions: %s", tables) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b633ac9d7b..33877768bb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -161,7 +161,6 @@ def __init__( max_sql_correction_retries: int = 3, pipeline_timeout_seconds: int = 90, schema_retrieval_timeout_seconds: int = 180, - allow_schema_sql_shortcuts: bool = False, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -179,7 +178,6 @@ def __init__( self._allow_sql_diagnosis = allow_sql_diagnosis self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval self._enable_column_pruning = enable_column_pruning - self._allow_schema_sql_shortcuts = allow_schema_sql_shortcuts self._pipeline_timeout_seconds = pipeline_timeout_seconds self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds self._max_histories = max_histories @@ -5473,7 +5471,6 @@ def _build_validated_ask_result_from_sql( sql: Optional[str], table_ddls: list[str], query: str | None = None, - strict_semantic_validation: bool = True, ) -> Optional[AskResult]: if isinstance(sql, str): sql = normalize_sql_direction_keywords(sql) @@ -5570,42 +5567,41 @@ def _build_validated_ask_result_from_sql( ) return None - if not self._sql_references_explicit_table(ask_result.sql, query): - return None - - if strict_semantic_validation: - invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( + invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( + ask_result.sql, + schema_tables, + ) + if invalid_unqualified_identifiers: + logger.warning( + "Ignoring SQL because it references unqualified fields outside the active schema. " + "invalid_identifiers=%s sql=%s", + invalid_unqualified_identifiers, ask_result.sql, - schema_tables, ) - if invalid_unqualified_identifiers: - logger.warning( - "Ignoring SQL because it references unqualified fields outside the active schema. " - "invalid_identifiers=%s sql=%s", - invalid_unqualified_identifiers, - ask_result.sql, - ) - return None + return None - invalid_output_aliases = self._invalid_sql_output_aliases( + invalid_output_aliases = self._invalid_sql_output_aliases( + ask_result.sql, + schema_tables, + ) + if invalid_output_aliases: + logger.warning( + "Ignoring SQL because it aliases output fields to unavailable schema concepts. " + "invalid_aliases=%s sql=%s", + invalid_output_aliases, ask_result.sql, - schema_tables, ) - if invalid_output_aliases: - logger.warning( - "Ignoring SQL because it aliases output fields to unavailable schema concepts. " - "invalid_aliases=%s sql=%s", - invalid_output_aliases, - ask_result.sql, - ) - return None + return None - if not self._sql_matches_question_intent( - ask_result.sql, - query, - schema_tables, - ): - return None + if not self._sql_references_explicit_table(ask_result.sql, query): + return None + + if not self._sql_matches_question_intent( + ask_result.sql, + query, + schema_tables, + ): + return None return ask_result @@ -6096,7 +6092,9 @@ async def ask( ) historical_question_result = [] - should_skip_pre_sql_retrieval = False + should_skip_pre_sql_retrieval = self._is_data_analysis_query( + user_query + ) if should_skip_pre_sql_retrieval: rephrased_question = user_query intent_reasoning = ( @@ -6545,7 +6543,7 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( ranked_measure_sql := self._build_schema_ranked_measure_sql( user_query, table_ddls, @@ -6566,7 +6564,7 @@ async def ask( invalid_sql = ranked_measure_sql error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( table_question_sql := self._build_schema_grounded_table_question_sql( user_query, table_ddls ) @@ -6586,7 +6584,7 @@ async def ask( invalid_sql = table_question_sql error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( explicit_table_preview := self._build_explicit_table_preview_sql( user_query, table_ddls ) @@ -6610,7 +6608,7 @@ async def ask( invalid_sql = explicit_sql error_message = "Explicit table preview SQL was not valid for the active datasource schema." - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( audit_log_activity_sql := self._build_audit_log_activity_sql( user_query, table_ddls, table_names=table_names ) @@ -6633,8 +6631,6 @@ async def ask( ) if ( - self._allow_schema_sql_shortcuts - and not api_results and self._is_data_analysis_query(user_query) and ( @@ -6660,7 +6656,7 @@ async def ask( "Schema-grounded SQL was not valid for the active datasource schema and question intent." ) - if self._allow_schema_sql_shortcuts and not api_results and any( + if not api_results and any( term in user_query.lower() for term in ( "pcb", @@ -6693,7 +6689,7 @@ async def ask( "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." ) - if self._allow_schema_sql_shortcuts and not api_results and ( + if not api_results and ( deterministic_sales_sql := self._build_schema_grounded_sales_sql( user_query, table_ddls ) @@ -6767,26 +6763,24 @@ async def ask( table_names, ) - full_schema_sql_candidates = () - if self._allow_schema_sql_shortcuts: - full_schema_preview = self._build_explicit_table_preview_sql( + full_schema_preview = self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + full_schema_sql_candidates = ( + self._build_schema_grounded_table_question_sql( user_query, table_ddls - ) - full_schema_sql_candidates = ( - self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ), - full_schema_preview[0] if full_schema_preview else None, - self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ), - self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ), - self._build_schema_grounded_sales_sql( - user_query, table_ddls - ), - ) + ), + full_schema_preview[0] if full_schema_preview else None, + self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ), + self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ), + self._build_schema_grounded_sales_sql( + user_query, table_ddls + ), + ) for full_schema_sql in full_schema_sql_candidates: if not full_schema_sql: continue @@ -6831,7 +6825,7 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results - if self._allow_schema_sql_shortcuts and not documents: + if not documents: if heuristic_sql := self._build_heuristic_text_to_sql_fallback( user_query, table_ddls, table_names=table_names ): @@ -7115,7 +7109,6 @@ async def ask( sql_valid_result.get("sql"), table_ddls, sql_user_query, - strict_semantic_validation=False, ): api_results = [ask_result] else: @@ -7201,7 +7194,6 @@ async def ask( valid_generation_result.get("sql"), table_ddls, sql_user_query, - strict_semantic_validation=False, ): api_results = [ask_result] break @@ -7234,10 +7226,8 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if self._allow_schema_sql_shortcuts and ( - heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ) + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names ): logger.info( "Using heuristic text-to-sql fallback for query_id %s: %s", @@ -7271,41 +7261,17 @@ async def ask( return results logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") - failure_message = ( - error_message - or "SQL generation could not produce a valid query for the retrieved schema." - ) - failure_code = ( - "NO_RELEVANT_SQL" - if documents or table_names - else "NO_RELEVANT_DATA" - ) if not self._is_stopped(query_id, self._ask_results): - if failure_code == "NO_RELEVANT_SQL": - self._ask_results[query_id] = ( - self._build_failed_text_to_sql_response( - trace_id, - failure_message, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - is_followup=True if histories else False, - code="NO_RELEVANT_SQL", - ) - ) - else: - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - is_followup=True if histories else False, - ) + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + is_followup=True if histories else False, ) + ) if error_message or invalid_sql: logger.info( "Suppressed technical SQL failure for query_id %s. " @@ -7314,11 +7280,9 @@ async def ask( error_message, invalid_sql, ) - results["metadata"]["error_type"] = failure_code + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" results["metadata"]["error_message"] = ( - failure_message - if failure_code == "NO_RELEVANT_SQL" - else NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE ) results["metadata"]["type"] = "TEXT_TO_SQL" diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 5662058dd8..7c7d46ce89 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -7,7 +7,6 @@ _select_relevant_table_documents, check_using_db_schemas_without_pruning, dbschema_retrieval, - embedding, expand_business_terms_for_retrieval, table_retrieval, ) @@ -48,27 +47,6 @@ def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): assert expand_business_terms_for_retrieval(query) == query -@pytest.mark.asyncio -async def test_embedding_uses_original_query_without_business_term_expansion(): - class Embedder: - def __init__(self): - self.query = None - - async def run(self, query): - self.query = query - return {"embedding": [0.1, 0.2]} - - embedder = Embedder() - - await embedding( - query="Show top customers by invoice amount", - embedder=embedder, - histories=[], - ) - - assert embedder.query == "\nShow top customers by invoice amount" - - def test_rerank_table_documents_prefers_question_relevant_table_text(): generic_stage = Document( content="Generic imported staging records with product labels.", @@ -175,7 +153,7 @@ def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): @pytest.mark.asyncio -async def test_table_retrieval_returns_embedding_results_without_local_reranking(): +async def test_table_retrieval_caps_embedding_results_before_schema_loading(): documents = [ Document( content="Raw staging audit rows with load metadata.", @@ -221,7 +199,9 @@ async def run(self, query_embedding, filters): table_retriever=Retriever(), ) - assert result["documents"] == documents + selected_names = [document.meta["name"] for document in result["documents"]] + assert 1 <= len(selected_names) <= 5 + assert "staging_audit" not in selected_names def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py index a994519419..9593fca054 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py @@ -217,31 +217,6 @@ def test_validated_sql_rejects_count_for_total_amount_question(): assert result is None -def test_validated_llm_sql_accepts_schema_valid_count_for_total_amount_question(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_Invoices"."Currency" AS "Currency", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_Invoices" ' - 'GROUP BY "dbo_Invoices"."Currency"' - ), - [ - """ - CREATE TABLE dbo_Invoices ( - Currency VARCHAR, - InvoiceAmount DOUBLE - ); - """ - ], - "Show total invoice amount by currency.", - strict_semantic_validation=False, - ) - - assert result is not None - - def test_validated_sql_rejects_detail_rows_for_customer_order_count_question(): service = AskService.__new__(AskService) From b94909e46775027249c7684c25344041159f2b15 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 18:29:38 +0530 Subject: [PATCH 0591/1087] Update db_schema_retrieval.py --- .../retrieval/db_schema_retrieval.py | 541 ++---------------- 1 file changed, 45 insertions(+), 496 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index b58e771490..2ddd931fe7 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,8 +1,7 @@ import ast import logging -import re import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional import orjson import tiktoken @@ -19,18 +18,12 @@ build_table_ddl, clean_up_new_lines, get_engine_supported_data_type, - normalize_data_type, ) from src.utils import trace_cost -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") -MAX_RELEVANT_TABLE_CANDIDATES = 5 - table_columns_selection_system_prompt = """ ### TASK ### @@ -108,9 +101,9 @@ def _build_metric_ddl(content: dict) -> str: columns_ddl = [ - f"{column['comment']}{column['name']} {get_engine_supported_data_type(normalize_data_type(column.get('data_type')))}" + f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" for column in content["columns"] - if normalize_data_type(column.get("data_type")).lower() + if column["data_type"].lower() != "unknown" # quick fix: filtering out UNKNOWN column type ] @@ -128,424 +121,8 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline -def expand_business_terms_for_retrieval(query: str) -> str: - normalized = (query or "").lower() - expansions: list[str] = [] - - if any( - term in normalized - for term in ( - "amount", - "currency", - "currencies", - "customer", - "customers", - "invoice", - "invoices", - "market", - "markets", - "order", - "orders", - "product", - "products", - "category", - "categories", - "quantity", - "qty", - "region", - "regions", - "sales", - "salesperson", - "sales person", - "sold", - "value", - ) - ): - expansions.append( - "transaction purchase billing account geography area representative product item category sku quantity units sold amount value total metric money exchange currency" - ) - - if any( - term in normalized - for term in ("defect", "failure", "issue", "repair", "resolved", "status") - ): - expansions.append( - "issue defect category status resolved created updated date timestamp event" - ) - - if any(term in normalized for term in ("throughput", "production", "manufacturing")): - expansions.append( - "rate volume output capacity process unit group completed timestamp date" - ) - - if not expansions: - return query - - return f"{query}\n" + "\n".join(expansions) - - -def _normalize_retrieval_token(value: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - -def _retrieval_terms(value: str) -> set[str]: - stop_words = { - "about", - "across", - "and", - "are", - "ask", - "bar", - "chart", - "create", - "different", - "for", - "from", - "how", - "in", - "is", - "of", - "show", - "the", - "to", - "top", - "what", - "which", - "with", - } - terms = { - _normalize_retrieval_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", value or "") - if len(token) > 2 and token.lower() not in stop_words - } - return {term for term in terms if term} - - -def _query_mentions_any(query: str, terms: tuple[str, ...]) -> bool: - normalized = (query or "").lower() - return any(re.search(rf"\b{re.escape(term)}\b", normalized) for term in terms) - - -def _source_text(document: Document) -> str: - return " ".join( - str(part or "") - for part in ( - document.meta.get("name"), - document.meta.get("description"), - document.content, - ) - ).lower() - - -def _source_shape_score(query: str, document: Document) -> int: - normalized_query = (query or "").lower() - source_text = _source_text(document) - source_terms = _retrieval_terms(source_text) - - score = 0 - non_production_terms = ( - "archive", - "backup", - "copy", - "dev", - "development", - "duplicate", - "sample", - "stage", - "staging", - "temp", - "test", - "tmp", - ) - if any(term in source_terms for term in non_production_terms) and not _query_mentions_any( - normalized_query, - non_production_terms, - ): - score -= 60 - - aggregation_terms = ( - "amount", - "average", - "avg", - "count", - "distribution", - "metric", - "revenue", - "sum", - "total", - "trend", - "value", - "volume", - ) - transaction_source_terms = ( - "activity", - "detail", - "event", - "fact", - "history", - "invoice", - "line", - "order", - "sale", - "sales", - "transaction", - ) - reference_source_terms = ( - "account", - "catalog", - "dimension", - "directory", - "entity", - "lookup", - "master", - "profile", - "reference", - ) - entity_listing_pattern = re.search( - r"\b(?:list|show|display|get|find)\b.*\b(?:accounts?|customers?|" - r"employees?|entities|items?|names?|products?|suppliers?|users?|vendors?)\b", - normalized_query, - ) - asks_for_aggregation = _query_mentions_any(normalized_query, aggregation_terms) or bool( - re.search(r"\b(?:by|per|each|top|bottom|rank|ranking)\b", normalized_query) - ) - asks_for_entity_listing = bool(entity_listing_pattern) and not asks_for_aggregation - - if asks_for_entity_listing: - if source_terms & set(reference_source_terms): - score += 35 - if source_terms & set(transaction_source_terms): - score -= 12 - elif asks_for_aggregation: - if source_terms & set(transaction_source_terms): - score += 25 - if source_terms & set(reference_source_terms): - score += 5 - - return score - - -def _document_relevance_score(document: Document, query_terms: set[str]) -> int: - if not query_terms: - return 0 - - document_terms = _retrieval_terms( - " ".join( - str(part or "") - for part in ( - document.meta.get("name"), - document.meta.get("description"), - document.content, - ) - ) - ) - if not document_terms: - return 0 - - score = 0 - for query_term in query_terms: - if query_term in document_terms: - score += 20 - continue - for document_term in document_terms: - if query_term in document_term or document_term in query_term: - score += 8 - break - return score - - -def _semantic_score(document: Document) -> float: - score = getattr(document, "score", None) - if isinstance(score, (int, float)): - return float(score) - score = document.meta.get("score") - if isinstance(score, (int, float)): - return float(score) - return 0.0 - - -def _score_table_documents( - query: str, documents: list[Document] -) -> list[tuple[float, int, Document, int, float]]: - if not documents: - return [] - - query_terms = _retrieval_terms(expand_business_terms_for_retrieval(query)) - if not query_terms: - return [ - (_semantic_score(document), -index, document, 0, _semantic_score(document)) - for index, document in enumerate(documents) - ] - - scored_documents: list[tuple[float, int, Document, int, float]] = [] - for index, document in enumerate(documents): - lexical_score = _document_relevance_score(document, query_terms) - semantic_score = _semantic_score(document) - source_shape_score = _source_shape_score(query, document) - combined_score = semantic_score + lexical_score + source_shape_score - scored_documents.append( - (combined_score, -index, document, lexical_score, semantic_score) - ) - - return sorted(scored_documents, key=lambda item: (item[0], item[1]), reverse=True) - - -def _rerank_table_documents(query: str, documents: list[Document]) -> list[Document]: - if not documents: - return documents - - reranked = _score_table_documents(query, documents) - if not reranked: - return documents - - logger.info( - "Top table candidates after retrieval rerank: %s", - [ - { - "name": document.meta.get("name"), - "semantic_score": round(semantic_score, 4), - "lexical_score": lexical_score, - "combined_score": round(combined_score, 4), - } - for combined_score, _index, document, lexical_score, semantic_score in reranked[ - :5 - ] - ], - ) - return [document for _score, _index, document, _lexical, _semantic in reranked] - - -def _select_relevant_table_documents( - query: str, - documents: list[Document], - *, - max_tables: int = MAX_RELEVANT_TABLE_CANDIDATES, -) -> list[Document]: - if not documents or max_tables <= 0: - return [] - - reranked = _score_table_documents(query, documents) - if not reranked: - return documents[:max_tables] - - candidate_pool = [item for item in reranked if item[3] > 0] or reranked - selected = [ - document - for _score, _index, document, _lexical, _semantic in candidate_pool[:max_tables] - ] - if len(selected) < len(documents): - logger.info( - "Scoped table candidates for schema loading from %s to %s tables: %s", - len(documents), - len(selected), - [document.meta.get("name") for document in selected], - ) - return selected - - -def _is_project_wide_analysis_query(query: str) -> bool: - normalized = (query or "").lower() - if not normalized: - return False - - analysis_terms = { - "average", - "avg", - "bar chart", - "breakdown", - "chart", - "completed", - "compare", - "count", - "counts", - "distribution", - "group by", - "grouped", - "highest", - "line chart", - "lowest", - "maximum", - "minimum", - "monthly", - "most common", - "number of", - "pie chart", - "quarter", - "rank", - "ranking", - "recommend", - "recommended", - "show", - "status", - "sum", - "total", - "totals", - "top", - "trend", - "volume", - } - return any(term in normalized for term in analysis_terms) - - -def _dedupe_documents(documents: list[Document]) -> list[Document]: - deduped: list[Document] = [] - seen: set[tuple[str, str, str]] = set() - for document in documents: - key = ( - str(document.meta.get("name", "")), - str(document.meta.get("type", "")), - document.content, - ) - if key in seen: - continue - seen.add(key) - deduped.append(document) - return deduped - - -def _normalize_table_names(table_names: Optional[list[str]]) -> list[str]: - normalized: list[str] = [] - for table_name in table_names or []: - if not isinstance(table_name, str): - continue - table_name = table_name.strip() - if table_name and table_name not in normalized: - normalized.append(table_name) - return normalized - - -def _extract_table_names_from_table_retrieval( - table_retrieval: dict, explicit_tables: Optional[list[str]] = None -) -> list[str]: - table_names = _normalize_table_names(explicit_tables) - for document in table_retrieval.get("documents") or []: - if not isinstance(document, Document): - continue - table_name = document.meta.get("name") - if not isinstance(table_name, str): - try: - content = ast.literal_eval(document.content) - except (SyntaxError, ValueError): - content = {} - table_name = content.get("name") if isinstance(content, dict) else None - if isinstance(table_name, str): - table_name = table_name.strip() - if table_name and table_name not in table_names: - table_names.append(table_name) - return table_names - - @observe(capture_input=False, capture_output=False) -async def embedding( - query: str, - embedder: Any, - histories: list[AskHistory], - tables: Optional[list[str]] = None, -) -> dict: - if tables: - logger.info("Skipping embedding retrieval for explicit tables: %s", tables) - return {} - +async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: if histories: previous_query_summaries = [history.question for history in histories] @@ -553,7 +130,6 @@ async def embedding( previous_query_summaries = [] query = "\n".join(previous_query_summaries) + "\n" + query - query = expand_business_terms_for_retrieval(query) return await embedder.run(query) else: @@ -562,13 +138,9 @@ async def embedding( @observe(capture_input=False) async def table_retrieval( - query: str, - embedding: dict, - project_id: str, - tables: list[str], - table_retriever: Any, + embedding: dict, project_id: str, tables: list[str], table_retriever: Any ) -> dict: - base_filters = { + filters = { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, @@ -576,82 +148,59 @@ async def table_retrieval( } if project_id: - base_filters["conditions"].append( + filters["conditions"].append( {"field": "project_id", "operator": "==", "value": project_id} ) if embedding: - results = await table_retriever.run( + return await table_retriever.run( query_embedding=embedding.get("embedding"), - filters=base_filters, + filters=filters, ) - results["documents"] = _select_relevant_table_documents( - query, results.get("documents") or [] + else: + filters["conditions"].append( + {"field": "name", "operator": "in", "value": tables} ) - return results - if tables: - logger.info("Loading explicit table descriptions: %s", tables) - explicit_filters = { - **base_filters, - "conditions": [ - *base_filters["conditions"], - {"field": "name", "operator": "in", "value": tables}, - ], - } - return await table_retriever.run(query_embedding=[], filters=explicit_filters) - - return {"documents": []} + return await table_retriever.run( + query_embedding=[], + filters=filters, + ) @observe(capture_input=False) async def dbschema_retrieval( - query: str, - table_retrieval: dict, - project_id: str, - dbschema_retriever: Any, - tables: Optional[list[str]] = None, + table_retrieval: dict, project_id: str, dbschema_retriever: Any ) -> list[Document]: - selected_table_names = _extract_table_names_from_table_retrieval( - table_retrieval, tables - ) + tables = table_retrieval.get("documents", []) + table_names = [] + for table in tables: + content = ast.literal_eval(table.content) + table_names.append(content["name"]) + + table_name_conditions = [ + {"field": "name", "operator": "==", "value": table_name} + for table_name in table_names + ] - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + if table_name_conditions: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } - if selected_table_names: - filters["conditions"].append( - {"field": "name", "operator": "in", "value": selected_table_names} - ) - logger.info( - "Loading selected deployed schema metadata for active project_id %s tables=%s", - project_id, - selected_table_names, - ) - elif not query: - logger.info( - "Loading complete deployed schema metadata for active project_id %s", - project_id, - ) - else: - logger.info( - "No relevant table-description candidates found for active project_id %s; " - "skipping full schema loading for query=%s", - project_id, - query, - ) - return [] + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results.get("documents", []) + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] + + return [] @observe() @@ -967,4 +516,4 @@ async def run( **self._components, **self._configs, }, - ) + ) \ No newline at end of file From d31228bae7a39856eedfd564f56dcdd56ee886df Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 19:52:01 +0530 Subject: [PATCH 0592/1087] Restore legacy ask retrieval flow --- wren-ai-service/src/web/v1/services/ask.py | 1583 +++----------------- 1 file changed, 237 insertions(+), 1346 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 33877768bb..bafbe31820 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -5653,14 +5653,11 @@ def _build_no_relevant_active_datasource_response( code="NO_RELEVANT_DATA", ) - @observe(name="Ask Question") - @trace_metadata - async def ask( + async def _ask_with_legacy_retrieval_flow( self, ask_request: AskRequest, - **kwargs, + trace_id: Optional[str], ): - trace_id = kwargs.get("trace_id") results = { "ask_result": {}, "metadata": { @@ -5672,42 +5669,15 @@ async def ask( } query_id = ask_request.query_id - if not query_id: - raise ValueError("query_id is required for ask service execution") - - user_query = (ask_request.query or "").strip() - if not user_query: - self._ask_results[query_id] = self._build_failed_text_to_sql_response( - trace_id, - "Question is required", - code="OTHERS", - ) - results["metadata"]["error_type"] = "OTHERS" - results["metadata"]["error_message"] = "Question is required" - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - logger.info(f"Ask pipeline started for query_id: {query_id}") - histories = ask_request.histories[: self._max_histories][ - ::-1 - ] # reverse the order of histories - if histories and not self._should_use_histories_for_query(user_query): - logger.info( - "Ignoring thread histories for independent question. query_id=%s query=%s", - query_id, - user_query, - ) - histories = [] + histories = ask_request.histories[: self._max_histories][::-1] rephrased_question = None intent_reasoning = None sql_generation_reasoning = None sql_samples = [] instructions = [] api_results = [] - documents = [] table_names = [] table_ddls = [] - _retrieval_result = {} error_message = None invalid_sql = None allow_sql_generation_reasoning = ( @@ -5725,22 +5695,10 @@ async def ask( use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None - understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) - request_explicit_table_names = self._normalize_explicit_table_names( - ask_request.explicit_tables - ) - forced_request_explicit_table_names = self._forced_explicit_table_names( - request_explicit_table_names, - source="request", - ) - explicit_table_names = forced_request_explicit_table_names - retrieval_table_names = explicit_table_names or None try: - sql_user_query = user_query + user_query = ask_request.query - # ask status can be understanding, searching, generating, finished, failed, stopped - # we will need to handle business logic for each status if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="understanding", @@ -5752,7 +5710,6 @@ async def ask( self._general_streaming_results[query_id] = ( self._build_greeting_response(user_query) ) - self._ask_results[query_id] = AskResultResponse( status="finished", type="GENERAL", @@ -5763,540 +5720,66 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results - metadata_question_kind = self._get_metadata_question_kind(user_query) - if metadata_question_kind: - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="GENERAL", - rephrased_question=user_query, - intent_reasoning=( - "Basic datasource metadata question detected; " - "retrieving deployed schema metadata directly." - ), - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - retrieval_result = await self._run_with_timeout( - "Metadata schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - metadata_answer = self._build_metadata_response( - user_query, table_ddls, table_names - ) - self._general_streaming_results[query_id] = metadata_answer - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=user_query, - intent_reasoning=( - "Answered from active datasource deployed metadata " - "without SQL generation." - ), - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - results["metadata"]["type"] = "GENERAL" - results["metadata"]["metadata_question_kind"] = ( - metadata_question_kind - ) - results["metadata"]["retrieved_table_count"] = len(documents) - return results + historical_question = await self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ) + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] - if explicit_table_names: - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - rephrased_question=user_query, - intent_reasoning="Explicit table name detected; retrieving that deployed schema directly.", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( + if historical_question_result: + api_results = [ + AskResult( + sql=result.get("statement"), + type="view" if result.get("viewId") else "llm", + viewId=result.get("viewId"), + ) + for result in historical_question_result + ] + sql_generation_reasoning = "" + else: + sql_samples_task, instructions_task = await asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( query=user_query, - tables=explicit_table_names, project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - if not documents and not request_explicit_table_names: - logger.info( - "Explicit table retrieval did not return requested active-schema table; " - "loading full active schema. query_id=%s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval for explicit table", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - all_documents, _, _ = self._extract_retrieval_metadata( - retrieval_result - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - all_documents, - explicit_table_names, - ) - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - logger.info( - "Retrieved explicit tables for query_id %s: %s", - query_id, - table_names, - ) - - if ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, - ): - ask_result = self._build_validated_ask_result_from_sql( - ranked_measure_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated ranked measure SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = ranked_measure_sql - - if table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ): - ask_result = self._build_validated_ask_result_from_sql( - table_question_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table question matched deployed schema.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = table_question_sql - - if explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ): - explicit_sql, explicit_table_name = explicit_table_preview - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table preview request matched deployed schema.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = explicit_sql - - if documents and ( - deterministic_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls - ) - ): - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = deterministic_sql - - if not documents: - error_message = ( - "The requested table was not found in the deployed schema: " - + ", ".join(explicit_table_names) - ) - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_DATA", - message=error_message, - ), - rephrased_question=user_query, - intent_reasoning="Explicit table request did not match any deployed schema table.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = error_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - rephrased_question = user_query - intent_reasoning = ( - "Explicit table request matched deployed schema; generating SQL against retrieved schema." - ) - sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - - if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( + self._pipelines["instructions_retrieval"].run( query=user_query, - histories=[], project_id=ask_request.project_id, - enable_column_pruning=False, + scope="sql", ), ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - logger.info( - "Retrieved tables for direct heuristic query_id %s: %s", - query_id, - table_names, - ) - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using direct heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - if ask_result := self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ): - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - - if explicit_group_count_sql := self._build_explicit_group_count_sql( - user_query - ): - invalid_sql = explicit_group_count_sql - rephrased_question = user_query - logger.info( - "Deferring explicit grouped count SQL until active schema validation for query_id %s", - query_id, - ) - - historical_question_result = [] - should_skip_pre_sql_retrieval = self._is_data_analysis_query( - user_query - ) - if should_skip_pre_sql_retrieval: - rephrased_question = user_query - intent_reasoning = ( - "Detected a deployed-data analytics question; skipping " - "intent classification and using SQL generation." + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] ) - sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - logger.info( - "Skipping pre-SQL retrieval for analytics query_id %s: %s", - query_id, - user_query, + instructions = instructions_task["formatted_output"].get( + "documents", [] ) - if ( - not api_results - and not should_skip_pre_sql_retrieval - and self._should_reuse_historical_question_sql( - user_query, histories - ) - ): - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - - try: - historical_question = await self._run_with_timeout( - "Historical question retrieval", - self._pipelines["historical_question"].run( + if self._allow_intent_classification: + intent_classification_result = ( + await self._pipelines["intent_classification"].run( query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, project_id=ask_request.project_id, - ), - timeout_seconds=min(understanding_timeout_seconds, 10), - ) - - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - except TimeoutError as exc: - logger.warning( - "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) - - valid_historical_results = [] - for result in historical_question_result: - historical_question_text = result.get("question") - if not self._is_reusable_historical_question( - user_query, historical_question_text - ): - logger.info( - "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", - query_id, - user_query, - historical_question_text, - ) - continue - - sql_statement = result.get("statement") - if not self._is_valid_select_sql(sql_statement): - logger.warning( - "Ignoring historical question without valid SQL for query_id %s", - query_id, - ) - continue - valid_historical_results.append( - AskResult( - **{ - "sql": sql_statement.strip(), - "type": "view" if result.get("viewId") else "llm", - "viewId": result.get("viewId"), - } - ) - ) - - if valid_historical_results: - api_results = valid_historical_results - sql_generation_reasoning = "" - elif not api_results and not should_skip_pre_sql_retrieval: - original_user_query = user_query - # Run both pipeline operations concurrently - try: - sql_samples_task, instructions_task = await self._run_with_timeout( - "SQL pair and instruction retrieval", - asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", - ), - ), - timeout_seconds=understanding_timeout_seconds, - ) - - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] - ) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) - except TimeoutError as exc: - logger.warning( - "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) - sql_samples = [] - instructions = [] - - if self._allow_intent_classification: - try: - intent_classification_result = ( - await self._run_with_timeout( - "Intent classification", - self._pipelines["intent_classification"].run( - query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, - ), - timeout_seconds=understanding_timeout_seconds, - ) - ).get("post_process", {}) - except TimeoutError as exc: - logger.warning( - "Intent classification timed out; continuing with TEXT_TO_SQL. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, + configuration=ask_request.configurations, ) - intent_classification_result = { - "intent": "TEXT_TO_SQL", - "rephrased_question": user_query, - "reasoning": "Intent classification timed out; using SQL generation.", - "db_schemas": [], - } + ).get("post_process", {}) intent = intent_classification_result.get("intent") rephrased_question = intent_classification_result.get( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") - retrieved_db_schemas = intent_classification_result.get( - "db_schemas" - ) or [] - is_original_analytics_query = self._is_data_analysis_query( - original_user_query - ) - is_schema_grounded_query = self._is_schema_grounded_query( - original_user_query, retrieved_db_schemas - ) or self._is_schema_grounded_query( - rephrased_question or "", retrieved_db_schemas - ) - if intent in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"} and ( - is_original_analytics_query - or is_schema_grounded_query - or self._is_data_analysis_query(rephrased_question or "") - ): - logger.info( - "Overriding intent %s to TEXT_TO_SQL for schema/data query: %s", - intent, - user_query, - ) - intent = "TEXT_TO_SQL" - - if is_original_analytics_query: - if rephrased_question and rephrased_question != user_query: - logger.info( - "Ignoring rephrased analytics query from intent classification. original=%s rephrased=%s", - original_user_query, - rephrased_question, - ) - user_query = original_user_query - rephrased_question = original_user_query - elif rephrased_question: + if rephrased_question: user_query = rephrased_question - sql_user_query = ( - self._rewrite_query_for_text_to_sql(user_query) - if self._is_data_analysis_query(user_query) - else user_query - ) - if intent == "MISLEADING_QUERY": - general_result = await self._run_with_timeout( - "Misleading assistance", + asyncio.create_task( self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, @@ -6304,18 +5787,13 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "misleading_assistance" ) ) - self._ask_results[query_id] = AskResultResponse( status="finished", - type="MISLEADING_QUERY", + type="GENERAL", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, @@ -6324,9 +5802,9 @@ async def ask( ) results["metadata"]["type"] = "MISLEADING_QUERY" return results - elif intent == "GENERAL": - general_result = await self._run_with_timeout( - "Data assistance", + + if intent == "GENERAL": + asyncio.create_task( self._pipelines["data_assistance"].run( query=user_query, histories=histories, @@ -6334,15 +5812,10 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "data_assistance" ) ) - self._ask_results[query_id] = AskResultResponse( status="finished", type="GENERAL", @@ -6354,21 +5827,16 @@ async def ask( ) results["metadata"]["type"] = "GENERAL" return results - elif intent == "USER_GUIDE": - general_result = await self._run_with_timeout( - "User guide assistance", + + if intent == "USER_GUIDE": + asyncio.create_task( self._pipelines["user_guide_assistance"].run( query=user_query, language=ask_request.configurations.language, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "user_guide_assistance" ) ) - self._ask_results[query_id] = AskResultResponse( status="finished", type="GENERAL", @@ -6380,20 +5848,17 @@ async def ask( ) results["metadata"]["type"] = "GENERAL" return results - else: - self._ask_results[query_id] = AskResultResponse( - status="understanding", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - if ( - not self._is_stopped(query_id, self._ask_results) - and not api_results - and not documents - ): + + self._ask_results[query_id] = AskResultResponse( + status="understanding", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + + if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -6403,530 +5868,38 @@ async def ask( is_followup=True if histories else False, ) - try: - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) - ), - ), - timeout_seconds=self._schema_retrieval_timeout_seconds, - ) - except TimeoutError as error: - if not self._should_retry_selected_schema_after_retrieval_timeout( - retrieval_table_names - ): - logger.warning( - "Schema retrieval timed out for data query; not loading full project schema. " - "query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - error, - ) - retrieval_result = {"construct_retrieval_results": {}} - else: - logger.warning( - "Schema retrieval timed out; retrying only explicit selected schemas. " - "query_id=%s project_id=%s tables=%s error=%s", - query_id, - ask_request.project_id, - retrieval_table_names, - error, - ) - retrieval_result = await self._run_with_timeout( - "Selected schema fallback retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 30, - ), - ) + retrieval_result = await self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=enable_column_pruning, + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - if not documents: - if explicit_table_names: - logger.info( - "Retrying schema retrieval for explicit tables query_id %s: %s", - query_id, - explicit_table_names, - ) - retrieval_result = await self._run_with_timeout( - "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=enable_column_pruning, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 20, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - if ( - not documents - and self._get_metadata_question_kind(user_query) - and not request_explicit_table_names - ): - logger.info( - "Query-based schema retrieval returned no tables for data question; " - "retrying full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - logger.info( - "Retrieved tables for query_id %s: %s", query_id, table_names - ) - - if not api_results and ( - ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, - ) - ): - logger.info( - "Using schema-grounded ranked measure SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - ranked_measure_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = ranked_measure_sql - error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." - - if not api_results and ( - table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded table question SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - table_question_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = table_question_sql - error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - - if not api_results and ( - explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - ): - explicit_sql, explicit_table_name = explicit_table_preview - logger.info( - "Using explicit table preview SQL for query_id %s and table %s", - query_id, - explicit_table_name, - ) - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = explicit_sql - error_message = "Explicit table preview SQL was not valid for the active datasource schema." - - if not api_results and ( - audit_log_activity_sql := self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ) - ): - logger.info( - "Using schema-grounded audit log activity SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - audit_log_activity_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = audit_log_activity_sql - error_message = ( - "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." - ) + documents = _retrieval_result.get("retrieval_results", []) + table_names = [document.get("table_name") for document in documents] + table_ddls = [document.get("table_ddl") for document in documents] - if ( - not api_results - and self._is_data_analysis_query(user_query) - and ( - schema_grounded_sql := self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - ) - ): - logger.info( - "Using generic schema-grounded analytics SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - schema_grounded_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = schema_grounded_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and any( - term in user_query.lower() - for term in ( - "pcb", - "repair", - "failure", - "business unit", - "business units", - "product line", - "product family", - ) - ): - operational_sql = self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - if operational_sql: - logger.info( - "Using schema-grounded operational SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - operational_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = operational_sql - error_message = ( - "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and ( - deterministic_sales_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded CWSales SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sales_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = deterministic_sales_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - - should_retry_full_schema = ( - not api_results - and self._get_metadata_question_kind(user_query) - and "db_schema_retrieval" in self._pipelines - and not request_explicit_table_names - and not table_names - ) - if should_retry_full_schema: - logger.info( - "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retry", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 30, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - full_documents, full_table_names, full_table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - full_documents, full_table_names, full_table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - full_documents, - explicit_table_names, - ) - ) - if full_documents: - documents, table_names, table_ddls = ( - full_documents, - full_table_names, - full_table_ddls, - ) - logger.info( - "Using full active deployed schema retry for query_id %s: %s", - query_id, - table_names, - ) - - full_schema_preview = self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - full_schema_sql_candidates = ( - self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ), - full_schema_preview[0] if full_schema_preview else None, - self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ), - self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ), - self._build_schema_grounded_sales_sql( - user_query, table_ddls - ), - ) - for full_schema_sql in full_schema_sql_candidates: - if not full_schema_sql: - continue - ask_result = self._build_validated_ask_result_from_sql( - full_schema_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - break - invalid_sql = full_schema_sql - error_message = ( - "Full-schema grounded SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and ( - unqueryable_metric_message := self._get_unqueryable_metric_message( - user_query, table_ddls - ) - ): - logger.info( - "ask pipeline - NO_RELEVANT_SQL due to unqueryable metric: %s", - user_query, - ) + if not documents: + logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="failed", type="TEXT_TO_SQL", error=AskError( - code="NO_RELEVANT_SQL", - message=unqueryable_metric_message, + code="NO_RELEVANT_DATA", + message="No relevant data", ), rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, - retrieved_tables=table_names, trace_id=trace_id, is_followup=True if histories else False, ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = unqueryable_metric_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - if not documents: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - is_followup=True if histories else False, - ) - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - is_followup=True if histories else False, - ) - ) results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) results["metadata"]["type"] = "TEXT_TO_SQL" return results - if documents and not api_results: - documents, table_names, table_ddls = self._prune_sql_generation_context( - sql_user_query, - documents, - table_names, - table_ddls, - ) - ( - documents, - table_names, - table_ddls, - completed_retrieval_result, - ) = await self._complete_sql_generation_context( - query=sql_user_query, - project_id=ask_request.project_id, - documents=documents, - table_names=table_names, - table_ddls=table_ddls, - ) - if completed_retrieval_result: - _retrieval_result = completed_retrieval_result - - sql_generation_histories = histories - if self._is_data_analysis_query( - sql_user_query - ) and not self._needs_conversation_context(sql_user_query): - sql_generation_histories = [] - allow_sql_generation_reasoning = False - allow_sql_knowledge_retrieval = False - max_sql_correction_retries = min(max_sql_correction_retries, 1) - logger.info( - "Using fast standalone SQL generation path for query_id %s", - query_id, - ) - if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -6942,53 +5915,29 @@ async def ask( is_followup=True if histories else False, ) - if sql_generation_histories: - try: - sql_generation_reasoning = ( - await self._run_with_timeout( - "Follow-up SQL generation reasoning", - self._pipelines[ - "followup_sql_generation_reasoning" - ].run( - query=sql_user_query, - contexts=table_ddls, - histories=sql_generation_histories, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), - ) - ).get("post_process", {}) - except Exception as reasoning_error: - logger.warning( - "Follow-up SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", - query_id, - reasoning_error, + if histories: + sql_generation_reasoning = ( + await self._pipelines["followup_sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, ) - sql_generation_reasoning = "" + ).get("post_process", {}) else: - try: - sql_generation_reasoning = ( - await self._run_with_timeout( - "SQL generation reasoning", - self._pipelines["sql_generation_reasoning"].run( - query=sql_user_query, - contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), - ) - ).get("post_process", {}) - except Exception as reasoning_error: - logger.warning( - "SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", - query_id, - reasoning_error, + sql_generation_reasoning = ( + await self._pipelines["sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, ) - sql_generation_reasoning = "" + ).get("post_process", {}) self._ask_results[query_id] = AskResultResponse( status="planning", @@ -7013,34 +5962,22 @@ async def ask( is_followup=True if histories else False, ) - try: - sql_functions, sql_knowledge = await self._run_with_timeout( - "SQL helper retrieval", - asyncio.gather( - ( - self._pipelines["sql_functions_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_functions_retrieval - else _return_value([]) - ), - ( - self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_knowledge_retrieval - else _return_value(None) - ), - ), - timeout_seconds=min(self._pipeline_timeout_seconds, 10), - ) - except TimeoutError as helper_timeout: - logger.warning( - "SQL helper retrieval timed out for query_id %s; continuing with schema only: %s", - query_id, - helper_timeout, - ) - sql_functions, sql_knowledge = [], None + sql_functions, sql_knowledge = await asyncio.gather( + ( + self._pipelines["sql_functions_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_functions_retrieval + else _return_value([]) + ), + ( + self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_knowledge_retrieval + else _return_value(None) + ), + ) has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -7048,92 +5985,63 @@ async def ask( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - try: - if sql_generation_histories: - text_to_sql_generation_results = await self._run_with_timeout( - "Follow-up SQL generation", - self._pipelines["followup_sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=sql_generation_histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), - ) - else: - text_to_sql_generation_results = await self._run_with_timeout( - "SQL generation", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), - ) - except TimeoutError as generation_timeout: - logger.warning( - "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", - query_id, - generation_timeout, + if histories: + text_to_sql_generation_results = await self._pipelines[ + "followup_sql_generation" + ].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ) + else: + text_to_sql_generation_results = await self._pipelines[ + "sql_generation" + ].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, ) - text_to_sql_generation_results = { - "post_process": { - "valid_generation_result": None, - "invalid_generation_result": None, - } - } - error_message = str(generation_timeout) if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - if ask_result := self._build_validated_ask_result_from_sql( - sql_valid_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - else: - invalid_sql = sql_valid_result.get("sql") - error_message = ( - "SQL generation did not produce SQL that matches the active datasource schema and question intent." + api_results = [ + AskResult( + sql=sql_valid_result.get("sql"), + type="llm", ) + ] elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] in { - "TIME_OUT", - "UNSUPPORTED_SQL", - }: - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) + if failed_dry_run_result["type"] == "TIME_OUT": break original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] - sql_diagnosis_reasoning = None current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -7147,68 +6055,53 @@ async def ask( is_followup=True if histories else False, ) + sql_diagnosis_reasoning = None if allow_sql_diagnosis: - sql_diagnosis_results = await self._run_with_timeout( - "SQL diagnosis", - self._pipelines["sql_diagnosis"].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_request.configurations.language, - ), + sql_diagnosis_results = await self._pipelines[ + "sql_diagnosis" + ].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_request.configurations.language, ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") - correction_error_message = error_message - if sql_diagnosis_reasoning: - correction_error_message = ( - f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" - ) - - sql_correction_results = await self._run_with_timeout( - "SQL correction", - self._pipelines["sql_correction"].run( - contexts=table_ddls, - instructions=instructions, - invalid_generation_result={ - "original_sql": original_sql, - "sql": invalid_sql, - "error": correction_error_message, - }, - project_id=ask_request.project_id, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - query=sql_user_query, - ), + sql_correction_results = await self._pipelines[ + "sql_correction" + ].run( + contexts=table_ddls, + instructions=instructions, + invalid_generation_result={ + "sql": original_sql, + "error": sql_diagnosis_reasoning + if allow_sql_diagnosis + else error_message, + }, + project_id=ask_request.project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, ) if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - if ask_result := self._build_validated_ask_result_from_sql( - valid_generation_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - break - invalid_sql = valid_generation_result.get("sql") - error_message = ( - "SQL correction did not produce SQL that matches the active datasource schema and question intent." - ) + api_results = [ + AskResult( + sql=valid_generation_result.get("sql"), + type="llm", + ) + ] + break failed_dry_run_result = sql_correction_results["post_process"][ "invalid_generation_result" ] - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) if api_results: if not self._is_stopped(query_id, self._ask_results): @@ -7226,64 +6119,25 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - else: - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - is_followup=True if histories else False, - ) - ) - if error_message or invalid_sql: - logger.info( - "Suppressed technical SQL failure for query_id %s. " - "error=%s invalid_sql=%s", - query_id, - error_message, - invalid_sql, + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_SQL", + message=error_message or "No relevant SQL", + ), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + trace_id=trace_id, + is_followup=True if histories else False, ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = error_message results["metadata"]["type"] = "TEXT_TO_SQL" return results @@ -7306,6 +6160,43 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + @observe(name="Ask Question") + @trace_metadata + async def ask( + self, + ask_request: AskRequest, + **kwargs, + ): + trace_id = kwargs.get("trace_id") + results = { + "ask_result": {}, + "metadata": { + "type": "", + "error_type": "", + "error_message": "", + "request_from": ask_request.request_from, + }, + } + + query_id = ask_request.query_id + if not query_id: + raise ValueError("query_id is required for ask service execution") + + user_query = (ask_request.query or "").strip() + if not user_query: + self._ask_results[query_id] = self._build_failed_text_to_sql_response( + trace_id, + "Question is required", + code="OTHERS", + ) + results["metadata"]["error_type"] = "OTHERS" + results["metadata"]["error_message"] = "Question is required" + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + logger.info(f"Ask pipeline started for query_id: {query_id}") + return await self._ask_with_legacy_retrieval_flow(ask_request, trace_id) + def stop_ask( self, stop_ask_request: StopAskRequest, From ba5d8ce4ab940ab9249bbf3a99e06d50e2df9095 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 20:06:45 +0530 Subject: [PATCH 0593/1087] Fix retrieval circular import --- .../pipelines/retrieval/db_schema_retrieval.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 2ddd931fe7..e6f8ad435a 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,7 @@ import ast import logging import sys -from typing import Any, Optional +from typing import Any, Optional, Protocol import orjson import tiktoken @@ -20,7 +20,11 @@ get_engine_supported_data_type, ) from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory + + +class AskHistoryLike(Protocol): + question: str + sql: str logger = logging.getLogger("wren-ai-service") @@ -122,7 +126,7 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline @observe(capture_input=False, capture_output=False) -async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: +async def embedding(query: str, embedder: Any, histories: list[AskHistoryLike]) -> dict: if query: if histories: previous_query_summaries = [history.question for history in histories] @@ -305,7 +309,7 @@ def prompt( construct_db_schemas: list[dict], prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, - histories: list[AskHistory], + histories: list[AskHistoryLike], ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -501,7 +505,7 @@ async def run( query: str = "", tables: Optional[list[str]] = None, project_id: Optional[str] = None, - histories: Optional[list[AskHistory]] = None, + histories: Optional[list[AskHistoryLike]] = None, enable_column_pruning: bool = False, ): logger.info("Ask Retrieval pipeline is running...") @@ -516,4 +520,4 @@ async def run( **self._components, **self._configs, }, - ) \ No newline at end of file + ) From d3ca5455db56d23f54ee243d6615ab2c65f4de45 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 20:28:52 +0530 Subject: [PATCH 0594/1087] Revert "Fix retrieval circular import" This reverts commit ba5d8ce4ab940ab9249bbf3a99e06d50e2df9095. --- .../pipelines/retrieval/db_schema_retrieval.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e6f8ad435a..2ddd931fe7 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,7 @@ import ast import logging import sys -from typing import Any, Optional, Protocol +from typing import Any, Optional import orjson import tiktoken @@ -20,11 +20,7 @@ get_engine_supported_data_type, ) from src.utils import trace_cost - - -class AskHistoryLike(Protocol): - question: str - sql: str +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -126,7 +122,7 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline @observe(capture_input=False, capture_output=False) -async def embedding(query: str, embedder: Any, histories: list[AskHistoryLike]) -> dict: +async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: if histories: previous_query_summaries = [history.question for history in histories] @@ -309,7 +305,7 @@ def prompt( construct_db_schemas: list[dict], prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, - histories: list[AskHistoryLike], + histories: list[AskHistory], ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -505,7 +501,7 @@ async def run( query: str = "", tables: Optional[list[str]] = None, project_id: Optional[str] = None, - histories: Optional[list[AskHistoryLike]] = None, + histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, ): logger.info("Ask Retrieval pipeline is running...") @@ -520,4 +516,4 @@ async def run( **self._components, **self._configs, }, - ) + ) \ No newline at end of file From af7478b8e4dd9ed802855d169e765d1d94ee3b22 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 20:28:52 +0530 Subject: [PATCH 0595/1087] Revert "Restore legacy ask retrieval flow" This reverts commit d31228bae7a39856eedfd564f56dcdd56ee886df. --- wren-ai-service/src/web/v1/services/ask.py | 1583 +++++++++++++++++--- 1 file changed, 1346 insertions(+), 237 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bafbe31820..33877768bb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -5653,11 +5653,14 @@ def _build_no_relevant_active_datasource_response( code="NO_RELEVANT_DATA", ) - async def _ask_with_legacy_retrieval_flow( + @observe(name="Ask Question") + @trace_metadata + async def ask( self, ask_request: AskRequest, - trace_id: Optional[str], + **kwargs, ): + trace_id = kwargs.get("trace_id") results = { "ask_result": {}, "metadata": { @@ -5669,15 +5672,42 @@ async def _ask_with_legacy_retrieval_flow( } query_id = ask_request.query_id - histories = ask_request.histories[: self._max_histories][::-1] + if not query_id: + raise ValueError("query_id is required for ask service execution") + + user_query = (ask_request.query or "").strip() + if not user_query: + self._ask_results[query_id] = self._build_failed_text_to_sql_response( + trace_id, + "Question is required", + code="OTHERS", + ) + results["metadata"]["error_type"] = "OTHERS" + results["metadata"]["error_message"] = "Question is required" + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + logger.info(f"Ask pipeline started for query_id: {query_id}") + histories = ask_request.histories[: self._max_histories][ + ::-1 + ] # reverse the order of histories + if histories and not self._should_use_histories_for_query(user_query): + logger.info( + "Ignoring thread histories for independent question. query_id=%s query=%s", + query_id, + user_query, + ) + histories = [] rephrased_question = None intent_reasoning = None sql_generation_reasoning = None sql_samples = [] instructions = [] api_results = [] + documents = [] table_names = [] table_ddls = [] + _retrieval_result = {} error_message = None invalid_sql = None allow_sql_generation_reasoning = ( @@ -5695,10 +5725,22 @@ async def _ask_with_legacy_retrieval_flow( use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None + understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) + request_explicit_table_names = self._normalize_explicit_table_names( + ask_request.explicit_tables + ) + forced_request_explicit_table_names = self._forced_explicit_table_names( + request_explicit_table_names, + source="request", + ) + explicit_table_names = forced_request_explicit_table_names + retrieval_table_names = explicit_table_names or None try: - user_query = ask_request.query + sql_user_query = user_query + # ask status can be understanding, searching, generating, finished, failed, stopped + # we will need to handle business logic for each status if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="understanding", @@ -5710,6 +5752,7 @@ async def _ask_with_legacy_retrieval_flow( self._general_streaming_results[query_id] = ( self._build_greeting_response(user_query) ) + self._ask_results[query_id] = AskResultResponse( status="finished", type="GENERAL", @@ -5720,66 +5763,540 @@ async def _ask_with_legacy_retrieval_flow( results["metadata"]["type"] = "GENERAL" return results - historical_question = await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ) - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] + metadata_question_kind = self._get_metadata_question_kind(user_query) + if metadata_question_kind: + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="GENERAL", + rephrased_question=user_query, + intent_reasoning=( + "Basic datasource metadata question detected; " + "retrieving deployed schema metadata directly." + ), + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + retrieval_result = await self._run_with_timeout( + "Metadata schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + metadata_answer = self._build_metadata_response( + user_query, table_ddls, table_names + ) + self._general_streaming_results[query_id] = metadata_answer + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=user_query, + intent_reasoning=( + "Answered from active datasource deployed metadata " + "without SQL generation." + ), + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + results["metadata"]["type"] = "GENERAL" + results["metadata"]["metadata_question_kind"] = ( + metadata_question_kind + ) + results["metadata"]["retrieved_table_count"] = len(documents) + return results - if historical_question_result: - api_results = [ - AskResult( - sql=result.get("statement"), - type="view" if result.get("viewId") else "llm", - viewId=result.get("viewId"), - ) - for result in historical_question_result - ] - sql_generation_reasoning = "" - else: - sql_samples_task, instructions_task = await asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( + if explicit_table_names: + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=user_query, + intent_reasoning="Explicit table name detected; retrieving that deployed schema directly.", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Explicit table schema retrieval", + self._pipelines["db_schema_retrieval"].run( query=user_query, + tables=explicit_table_names, project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, ), - self._pipelines["instructions_retrieval"].run( + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) + if not documents and not request_explicit_table_names: + logger.info( + "Explicit table retrieval did not return requested active-schema table; " + "loading full active schema. query_id=%s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval for explicit table", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + all_documents, _, _ = self._extract_retrieval_metadata( + retrieval_result + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + all_documents, + explicit_table_names, + ) + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + logger.info( + "Retrieved explicit tables for query_id %s: %s", + query_id, + table_names, + ) + + if ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, + ): + ask_result = self._build_validated_ask_result_from_sql( + ranked_measure_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated ranked measure SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = ranked_measure_sql + + if table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ): + ask_result = self._build_validated_ask_result_from_sql( + table_question_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table question matched deployed schema.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = table_question_sql + + if explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ): + explicit_sql, explicit_table_name = explicit_table_preview + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table preview request matched deployed schema.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = explicit_sql + + if documents and ( + deterministic_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ) + ): + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = deterministic_sql + + if not documents: + error_message = ( + "The requested table was not found in the deployed schema: " + + ", ".join(explicit_table_names) + ) + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_DATA", + message=error_message, + ), + rephrased_question=user_query, + intent_reasoning="Explicit table request did not match any deployed schema table.", + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = error_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + rephrased_question = user_query + intent_reasoning = ( + "Explicit table request matched deployed schema; generating SQL against retrieved schema." + ) + sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + + if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + trace_id=trace_id, + is_followup=True if histories else False, + ) + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( query=user_query, + histories=[], project_id=ask_request.project_id, - scope="sql", + enable_column_pruning=False, ), ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + logger.info( + "Retrieved tables for direct heuristic query_id %s: %s", + query_id, + table_names, + ) - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using direct heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + if ask_result := self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ): + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=user_query, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + + if explicit_group_count_sql := self._build_explicit_group_count_sql( + user_query + ): + invalid_sql = explicit_group_count_sql + rephrased_question = user_query + logger.info( + "Deferring explicit grouped count SQL until active schema validation for query_id %s", + query_id, + ) + + historical_question_result = [] + should_skip_pre_sql_retrieval = self._is_data_analysis_query( + user_query + ) + if should_skip_pre_sql_retrieval: + rephrased_question = user_query + intent_reasoning = ( + "Detected a deployed-data analytics question; skipping " + "intent classification and using SQL generation." ) - instructions = instructions_task["formatted_output"].get( - "documents", [] + sql_user_query = self._rewrite_query_for_text_to_sql(user_query) + logger.info( + "Skipping pre-SQL retrieval for analytics query_id %s: %s", + query_id, + user_query, ) - if self._allow_intent_classification: - intent_classification_result = ( - await self._pipelines["intent_classification"].run( + if ( + not api_results + and not should_skip_pre_sql_retrieval + and self._should_reuse_historical_question_sql( + user_query, histories + ) + ): + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + + try: + historical_question = await self._run_with_timeout( + "Historical question retrieval", + self._pipelines["historical_question"].run( query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, project_id=ask_request.project_id, - configuration=ask_request.configurations, + ), + timeout_seconds=min(understanding_timeout_seconds, 10), + ) + + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] + except TimeoutError as exc: + logger.warning( + "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) + + valid_historical_results = [] + for result in historical_question_result: + historical_question_text = result.get("question") + if not self._is_reusable_historical_question( + user_query, historical_question_text + ): + logger.info( + "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", + query_id, + user_query, + historical_question_text, + ) + continue + + sql_statement = result.get("statement") + if not self._is_valid_select_sql(sql_statement): + logger.warning( + "Ignoring historical question without valid SQL for query_id %s", + query_id, + ) + continue + valid_historical_results.append( + AskResult( + **{ + "sql": sql_statement.strip(), + "type": "view" if result.get("viewId") else "llm", + "viewId": result.get("viewId"), + } + ) + ) + + if valid_historical_results: + api_results = valid_historical_results + sql_generation_reasoning = "" + elif not api_results and not should_skip_pre_sql_retrieval: + original_user_query = user_query + # Run both pipeline operations concurrently + try: + sql_samples_task, instructions_task = await self._run_with_timeout( + "SQL pair and instruction retrieval", + asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), + ), + timeout_seconds=understanding_timeout_seconds, + ) + + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] + ) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) + except TimeoutError as exc: + logger.warning( + "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, + ) + sql_samples = [] + instructions = [] + + if self._allow_intent_classification: + try: + intent_classification_result = ( + await self._run_with_timeout( + "Intent classification", + self._pipelines["intent_classification"].run( + query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, + ), + timeout_seconds=understanding_timeout_seconds, + ) + ).get("post_process", {}) + except TimeoutError as exc: + logger.warning( + "Intent classification timed out; continuing with TEXT_TO_SQL. query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + exc, ) - ).get("post_process", {}) + intent_classification_result = { + "intent": "TEXT_TO_SQL", + "rephrased_question": user_query, + "reasoning": "Intent classification timed out; using SQL generation.", + "db_schemas": [], + } intent = intent_classification_result.get("intent") rephrased_question = intent_classification_result.get( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") + retrieved_db_schemas = intent_classification_result.get( + "db_schemas" + ) or [] + is_original_analytics_query = self._is_data_analysis_query( + original_user_query + ) + is_schema_grounded_query = self._is_schema_grounded_query( + original_user_query, retrieved_db_schemas + ) or self._is_schema_grounded_query( + rephrased_question or "", retrieved_db_schemas + ) - if rephrased_question: + if intent in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"} and ( + is_original_analytics_query + or is_schema_grounded_query + or self._is_data_analysis_query(rephrased_question or "") + ): + logger.info( + "Overriding intent %s to TEXT_TO_SQL for schema/data query: %s", + intent, + user_query, + ) + intent = "TEXT_TO_SQL" + + if is_original_analytics_query: + if rephrased_question and rephrased_question != user_query: + logger.info( + "Ignoring rephrased analytics query from intent classification. original=%s rephrased=%s", + original_user_query, + rephrased_question, + ) + user_query = original_user_query + rephrased_question = original_user_query + elif rephrased_question: user_query = rephrased_question + sql_user_query = ( + self._rewrite_query_for_text_to_sql(user_query) + if self._is_data_analysis_query(user_query) + else user_query + ) + if intent == "MISLEADING_QUERY": - asyncio.create_task( + general_result = await self._run_with_timeout( + "Misleading assistance", self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, @@ -5787,13 +6304,18 @@ async def _ask_with_legacy_retrieval_flow( "db_schemas" ), language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "misleading_assistance" ) ) + self._ask_results[query_id] = AskResultResponse( status="finished", - type="GENERAL", + type="MISLEADING_QUERY", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, @@ -5802,9 +6324,9 @@ async def _ask_with_legacy_retrieval_flow( ) results["metadata"]["type"] = "MISLEADING_QUERY" return results - - if intent == "GENERAL": - asyncio.create_task( + elif intent == "GENERAL": + general_result = await self._run_with_timeout( + "Data assistance", self._pipelines["data_assistance"].run( query=user_query, histories=histories, @@ -5812,10 +6334,15 @@ async def _ask_with_legacy_retrieval_flow( "db_schemas" ), language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "data_assistance" ) ) + self._ask_results[query_id] = AskResultResponse( status="finished", type="GENERAL", @@ -5827,16 +6354,21 @@ async def _ask_with_legacy_retrieval_flow( ) results["metadata"]["type"] = "GENERAL" return results - - if intent == "USER_GUIDE": - asyncio.create_task( + elif intent == "USER_GUIDE": + general_result = await self._run_with_timeout( + "User guide assistance", self._pipelines["user_guide_assistance"].run( query=user_query, language=ask_request.configurations.language, - query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, + ), + ) + self._general_streaming_results[query_id] = ( + self._extract_pipeline_reply( + general_result, "user_guide_assistance" ) ) + self._ask_results[query_id] = AskResultResponse( status="finished", type="GENERAL", @@ -5848,17 +6380,20 @@ async def _ask_with_legacy_retrieval_flow( ) results["metadata"]["type"] = "GENERAL" return results - - self._ask_results[query_id] = AskResultResponse( - status="understanding", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - - if not self._is_stopped(query_id, self._ask_results) and not api_results: + else: + self._ask_results[query_id] = AskResultResponse( + status="understanding", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and not documents + ): self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -5868,38 +6403,530 @@ async def _ask_with_legacy_retrieval_flow( is_followup=True if histories else False, ) - retrieval_result = await self._pipelines["db_schema_retrieval"].run( - query=user_query, - histories=histories, - project_id=ask_request.project_id, - enable_column_pruning=enable_column_pruning, - ) + try: + retrieval_result = await self._run_with_timeout( + "Schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=( + enable_column_pruning + and not self._is_data_analysis_query(user_query) + ), + ), + timeout_seconds=self._schema_retrieval_timeout_seconds, + ) + except TimeoutError as error: + if not self._should_retry_selected_schema_after_retrieval_timeout( + retrieval_table_names + ): + logger.warning( + "Schema retrieval timed out for data query; not loading full project schema. " + "query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + error, + ) + retrieval_result = {"construct_retrieval_results": {}} + else: + logger.warning( + "Schema retrieval timed out; retrying only explicit selected schemas. " + "query_id=%s project_id=%s tables=%s error=%s", + query_id, + ask_request.project_id, + retrieval_table_names, + error, + ) + retrieval_result = await self._run_with_timeout( + "Selected schema fallback retrieval", + self._pipelines["db_schema_retrieval"].run( + query=sql_user_query, + tables=retrieval_table_names, + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 30, + ), + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - documents = _retrieval_result.get("retrieval_results", []) - table_names = [document.get("table_name") for document in documents] - table_ddls = [document.get("table_ddl") for document in documents] - + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) if not documents: - logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") + if explicit_table_names: + logger.info( + "Retrying schema retrieval for explicit tables query_id %s: %s", + query_id, + explicit_table_names, + ) + retrieval_result = await self._run_with_timeout( + "Explicit table schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=explicit_table_names, + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=enable_column_pruning, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 20, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) + if ( + not documents + and self._get_metadata_question_kind(user_query) + and not request_explicit_table_names + ): + logger.info( + "Query-based schema retrieval returned no tables for data question; " + "retrying full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + documents, table_names, table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + documents, + explicit_table_names, + ) + ) + logger.info( + "Retrieved tables for query_id %s: %s", query_id, table_names + ) + + if not api_results and ( + ranked_measure_sql := self._build_schema_ranked_measure_sql( + user_query, + table_ddls, + ) + ): + logger.info( + "Using schema-grounded ranked measure SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + ranked_measure_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = ranked_measure_sql + error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." + + if not api_results and ( + table_question_sql := self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ) + ): + logger.info( + "Using schema-grounded table question SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + table_question_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = table_question_sql + error_message = "Schema-grounded table SQL was not valid for the active datasource schema." + + if not api_results and ( + explicit_table_preview := self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + ): + explicit_sql, explicit_table_name = explicit_table_preview + logger.info( + "Using explicit table preview SQL for query_id %s and table %s", + query_id, + explicit_table_name, + ) + if explicit_table_name not in table_names: + table_names.append(explicit_table_name) + ask_result = self._build_validated_ask_result_from_sql( + explicit_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = explicit_sql + error_message = "Explicit table preview SQL was not valid for the active datasource schema." + + if not api_results and ( + audit_log_activity_sql := self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ) + ): + logger.info( + "Using schema-grounded audit log activity SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + audit_log_activity_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = audit_log_activity_sql + error_message = ( + "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." + ) + + if ( + not api_results + and self._is_data_analysis_query(user_query) + and ( + schema_grounded_sql := self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + ) + ): + logger.info( + "Using generic schema-grounded analytics SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + schema_grounded_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = schema_grounded_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and any( + term in user_query.lower() + for term in ( + "pcb", + "repair", + "failure", + "business unit", + "business units", + "product line", + "product family", + ) + ): + operational_sql = self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ) + if operational_sql: + logger.info( + "Using schema-grounded operational SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + operational_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = operational_sql + error_message = ( + "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and ( + deterministic_sales_sql := self._build_schema_grounded_sales_sql( + user_query, table_ddls + ) + ): + logger.info( + "Using schema-grounded CWSales SQL for query_id %s", + query_id, + ) + ask_result = self._build_validated_ask_result_from_sql( + deterministic_sales_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + else: + invalid_sql = deterministic_sales_sql + error_message = ( + "Schema-grounded SQL was not valid for the active datasource schema and question intent." + ) + + should_retry_full_schema = ( + not api_results + and self._get_metadata_question_kind(user_query) + and "db_schema_retrieval" in self._pipelines + and not request_explicit_table_names + and not table_names + ) + if should_retry_full_schema: + logger.info( + "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", + query_id, + ) + retrieval_result = await self._run_with_timeout( + "Full active schema retry", + self._pipelines["db_schema_retrieval"].run( + query="", + histories=[], + project_id=ask_request.project_id, + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 30, + ), + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + full_documents, full_table_names, full_table_ddls = ( + self._extract_retrieval_metadata(retrieval_result) + ) + if explicit_table_names: + full_documents, full_table_names, full_table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + full_documents, + explicit_table_names, + ) + ) + if full_documents: + documents, table_names, table_ddls = ( + full_documents, + full_table_names, + full_table_ddls, + ) + logger.info( + "Using full active deployed schema retry for query_id %s: %s", + query_id, + table_names, + ) + + full_schema_preview = self._build_explicit_table_preview_sql( + user_query, table_ddls + ) + full_schema_sql_candidates = ( + self._build_schema_grounded_table_question_sql( + user_query, table_ddls + ), + full_schema_preview[0] if full_schema_preview else None, + self._build_audit_log_activity_sql( + user_query, table_ddls, table_names=table_names + ), + self._build_schema_grounded_analytics_sql( + user_query, table_ddls + ), + self._build_schema_grounded_sales_sql( + user_query, table_ddls + ), + ) + for full_schema_sql in full_schema_sql_candidates: + if not full_schema_sql: + continue + ask_result = self._build_validated_ask_result_from_sql( + full_schema_sql, + table_ddls, + user_query, + ) + if ask_result: + api_results = [ask_result] + break + invalid_sql = full_schema_sql + error_message = ( + "Full-schema grounded SQL was not valid for the active datasource schema and question intent." + ) + + if not api_results and ( + unqueryable_metric_message := self._get_unqueryable_metric_message( + user_query, table_ddls + ) + ): + logger.info( + "ask pipeline - NO_RELEVANT_SQL due to unqueryable metric: %s", + user_query, + ) if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="failed", type="TEXT_TO_SQL", error=AskError( - code="NO_RELEVANT_DATA", - message="No relevant data", + code="NO_RELEVANT_SQL", + message=unqueryable_metric_message, ), rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, + retrieved_tables=table_names, trace_id=trace_id, is_followup=True if histories else False, ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = unqueryable_metric_message + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + if not documents: + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", + query_id, + user_query, + ) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + is_followup=True if histories else False, + ) + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + is_followup=True if histories else False, + ) + ) results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) results["metadata"]["type"] = "TEXT_TO_SQL" return results + if documents and not api_results: + documents, table_names, table_ddls = self._prune_sql_generation_context( + sql_user_query, + documents, + table_names, + table_ddls, + ) + ( + documents, + table_names, + table_ddls, + completed_retrieval_result, + ) = await self._complete_sql_generation_context( + query=sql_user_query, + project_id=ask_request.project_id, + documents=documents, + table_names=table_names, + table_ddls=table_ddls, + ) + if completed_retrieval_result: + _retrieval_result = completed_retrieval_result + + sql_generation_histories = histories + if self._is_data_analysis_query( + sql_user_query + ) and not self._needs_conversation_context(sql_user_query): + sql_generation_histories = [] + allow_sql_generation_reasoning = False + allow_sql_knowledge_retrieval = False + max_sql_correction_retries = min(max_sql_correction_retries, 1) + logger.info( + "Using fast standalone SQL generation path for query_id %s", + query_id, + ) + if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -5915,29 +6942,53 @@ async def _ask_with_legacy_retrieval_flow( is_followup=True if histories else False, ) - if histories: - sql_generation_reasoning = ( - await self._pipelines["followup_sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, + if sql_generation_histories: + try: + sql_generation_reasoning = ( + await self._run_with_timeout( + "Follow-up SQL generation reasoning", + self._pipelines[ + "followup_sql_generation_reasoning" + ].run( + query=sql_user_query, + contexts=table_ddls, + histories=sql_generation_histories, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), + ) + ).get("post_process", {}) + except Exception as reasoning_error: + logger.warning( + "Follow-up SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", + query_id, + reasoning_error, ) - ).get("post_process", {}) + sql_generation_reasoning = "" else: - sql_generation_reasoning = ( - await self._pipelines["sql_generation_reasoning"].run( - query=user_query, - contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, + try: + sql_generation_reasoning = ( + await self._run_with_timeout( + "SQL generation reasoning", + self._pipelines["sql_generation_reasoning"].run( + query=sql_user_query, + contexts=table_ddls, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, + ), + ) + ).get("post_process", {}) + except Exception as reasoning_error: + logger.warning( + "SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", + query_id, + reasoning_error, ) - ).get("post_process", {}) + sql_generation_reasoning = "" self._ask_results[query_id] = AskResultResponse( status="planning", @@ -5962,22 +7013,34 @@ async def _ask_with_legacy_retrieval_flow( is_followup=True if histories else False, ) - sql_functions, sql_knowledge = await asyncio.gather( - ( - self._pipelines["sql_functions_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_functions_retrieval - else _return_value([]) - ), - ( - self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_knowledge_retrieval - else _return_value(None) - ), - ) + try: + sql_functions, sql_knowledge = await self._run_with_timeout( + "SQL helper retrieval", + asyncio.gather( + ( + self._pipelines["sql_functions_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_functions_retrieval + else _return_value([]) + ), + ( + self._pipelines["sql_knowledge_retrieval"].run( + project_id=ask_request.project_id, + ) + if allow_sql_knowledge_retrieval + else _return_value(None) + ), + ), + timeout_seconds=min(self._pipeline_timeout_seconds, 10), + ) + except TimeoutError as helper_timeout: + logger.warning( + "SQL helper retrieval timed out for query_id %s; continuing with schema only: %s", + query_id, + helper_timeout, + ) + sql_functions, sql_knowledge = [], None has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -5985,63 +7048,92 @@ async def _ask_with_legacy_retrieval_flow( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - if histories: - text_to_sql_generation_results = await self._pipelines[ - "followup_sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ) - else: - text_to_sql_generation_results = await self._pipelines[ - "sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, + try: + if sql_generation_histories: + text_to_sql_generation_results = await self._run_with_timeout( + "Follow-up SQL generation", + self._pipelines["followup_sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=sql_generation_histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), + ) + else: + text_to_sql_generation_results = await self._run_with_timeout( + "SQL generation", + self._pipelines["sql_generation"].run( + query=sql_user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ), + ) + except TimeoutError as generation_timeout: + logger.warning( + "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + query_id, + generation_timeout, ) + text_to_sql_generation_results = { + "post_process": { + "valid_generation_result": None, + "invalid_generation_result": None, + } + } + error_message = str(generation_timeout) if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - api_results = [ - AskResult( - sql=sql_valid_result.get("sql"), - type="llm", + if ask_result := self._build_validated_ask_result_from_sql( + sql_valid_result.get("sql"), + table_ddls, + sql_user_query, + ): + api_results = [ask_result] + else: + invalid_sql = sql_valid_result.get("sql") + error_message = ( + "SQL generation did not produce SQL that matches the active datasource schema and question intent." ) - ] elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + if failed_dry_run_result["type"] in { + "TIME_OUT", + "UNSUPPORTED_SQL", + }: + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) break original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + sql_diagnosis_reasoning = None current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -6055,53 +7147,68 @@ async def _ask_with_legacy_retrieval_flow( is_followup=True if histories else False, ) - sql_diagnosis_reasoning = None if allow_sql_diagnosis: - sql_diagnosis_results = await self._pipelines[ - "sql_diagnosis" - ].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_request.configurations.language, + sql_diagnosis_results = await self._run_with_timeout( + "SQL diagnosis", + self._pipelines["sql_diagnosis"].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_request.configurations.language, + ), ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") - sql_correction_results = await self._pipelines[ - "sql_correction" - ].run( - contexts=table_ddls, - instructions=instructions, - invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, - }, - project_id=ask_request.project_id, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, + correction_error_message = error_message + if sql_diagnosis_reasoning: + correction_error_message = ( + f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" + ) + + sql_correction_results = await self._run_with_timeout( + "SQL correction", + self._pipelines["sql_correction"].run( + contexts=table_ddls, + instructions=instructions, + invalid_generation_result={ + "original_sql": original_sql, + "sql": invalid_sql, + "error": correction_error_message, + }, + project_id=ask_request.project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + query=sql_user_query, + ), ) if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - api_results = [ - AskResult( - sql=valid_generation_result.get("sql"), - type="llm", - ) - ] - break + if ask_result := self._build_validated_ask_result_from_sql( + valid_generation_result.get("sql"), + table_ddls, + sql_user_query, + ): + api_results = [ask_result] + break + invalid_sql = valid_generation_result.get("sql") + error_message = ( + "SQL correction did not produce SQL that matches the active datasource schema and question intent." + ) failed_dry_run_result = sql_correction_results["post_process"][ "invalid_generation_result" ] + invalid_sql = failed_dry_run_result.get("sql", invalid_sql) + error_message = failed_dry_run_result.get( + "error", error_message + ) if api_results: if not self._is_stopped(query_id, self._ask_results): @@ -6119,25 +7226,64 @@ async def _ask_with_legacy_retrieval_flow( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: + if heuristic_sql := self._build_heuristic_text_to_sql_fallback( + user_query, table_ddls, table_names=table_names + ): + logger.info( + "Using heuristic text-to-sql fallback for query_id %s: %s", + query_id, + user_query, + ) + ask_result = self._build_validated_ask_result_from_sql( + heuristic_sql, + table_ddls, + user_query, + ) + if not ask_result: + invalid_sql = heuristic_sql + error_message = "Heuristic SQL fallback was not valid for the active datasource schema." + else: + api_results = [ask_result] + if not self._is_stopped(query_id, self._ask_results): + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="TEXT_TO_SQL", + response=api_results, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + results["ask_result"] = api_results + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_SQL", - message=error_message or "No relevant SQL", - ), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - trace_id=trace_id, - is_followup=True if histories else False, + self._ask_results[query_id] = ( + self._build_no_relevant_active_datasource_response( + trace_id, + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + is_followup=True if histories else False, + ) ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = error_message + if error_message or invalid_sql: + logger.info( + "Suppressed technical SQL failure for query_id %s. " + "error=%s invalid_sql=%s", + query_id, + error_message, + invalid_sql, + ) + results["metadata"]["error_type"] = "NO_RELEVANT_DATA" + results["metadata"]["error_message"] = ( + NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE + ) results["metadata"]["type"] = "TEXT_TO_SQL" return results @@ -6160,43 +7306,6 @@ async def _ask_with_legacy_retrieval_flow( results["metadata"]["type"] = "TEXT_TO_SQL" return results - @observe(name="Ask Question") - @trace_metadata - async def ask( - self, - ask_request: AskRequest, - **kwargs, - ): - trace_id = kwargs.get("trace_id") - results = { - "ask_result": {}, - "metadata": { - "type": "", - "error_type": "", - "error_message": "", - "request_from": ask_request.request_from, - }, - } - - query_id = ask_request.query_id - if not query_id: - raise ValueError("query_id is required for ask service execution") - - user_query = (ask_request.query or "").strip() - if not user_query: - self._ask_results[query_id] = self._build_failed_text_to_sql_response( - trace_id, - "Question is required", - code="OTHERS", - ) - results["metadata"]["error_type"] = "OTHERS" - results["metadata"]["error_message"] = "Question is required" - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - logger.info(f"Ask pipeline started for query_id: {query_id}") - return await self._ask_with_legacy_retrieval_flow(ask_request, trace_id) - def stop_ask( self, stop_ask_request: StopAskRequest, From 3d2547921e2f5ab670fbb3aee7071e8c3e5e48e9 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 21:04:09 +0530 Subject: [PATCH 0596/1087] Fix retrieval circular import --- .../pipelines/retrieval/db_schema_retrieval.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 2ddd931fe7..e6f8ad435a 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,7 @@ import ast import logging import sys -from typing import Any, Optional +from typing import Any, Optional, Protocol import orjson import tiktoken @@ -20,7 +20,11 @@ get_engine_supported_data_type, ) from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory + + +class AskHistoryLike(Protocol): + question: str + sql: str logger = logging.getLogger("wren-ai-service") @@ -122,7 +126,7 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline @observe(capture_input=False, capture_output=False) -async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: +async def embedding(query: str, embedder: Any, histories: list[AskHistoryLike]) -> dict: if query: if histories: previous_query_summaries = [history.question for history in histories] @@ -305,7 +309,7 @@ def prompt( construct_db_schemas: list[dict], prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, - histories: list[AskHistory], + histories: list[AskHistoryLike], ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -501,7 +505,7 @@ async def run( query: str = "", tables: Optional[list[str]] = None, project_id: Optional[str] = None, - histories: Optional[list[AskHistory]] = None, + histories: Optional[list[AskHistoryLike]] = None, enable_column_pruning: bool = False, ): logger.info("Ask Retrieval pipeline is running...") @@ -516,4 +520,4 @@ async def run( **self._components, **self._configs, }, - ) \ No newline at end of file + ) From 544e8a0bdb332d0ac64f0407357c99d52b3cb32a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 16 Jul 2026 22:50:57 +0530 Subject: [PATCH 0597/1087] Disable heuristic SQL overrides in ask pipeline --- wren-ai-service/src/web/v1/services/ask.py | 66 ++++++++++------------ 1 file changed, 29 insertions(+), 37 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 33877768bb..b8077c0f98 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -269,6 +269,8 @@ def _is_greeting_query(self, query: str) -> bool: return normalized in greeting_patterns def _is_data_analysis_query(self, query: str) -> bool: + return False + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return False @@ -997,6 +999,8 @@ def _sql_matches_question_intent( query: str | None, schema_tables: list[dict[str, Any]], ) -> bool: + return True + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) expects_dimension = bool( re.search( @@ -1325,6 +1329,8 @@ def _sql_references_explicit_table( sql: str, query: str | None, ) -> bool: + return True + explicit_table_keys = self._explicit_table_alias_keys_from_query(query) if not explicit_table_keys: return True @@ -1445,6 +1451,8 @@ def _build_schema_ranked_measure_sql( query: str, table_ddls: list[str], ) -> str | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized_query: return None @@ -1555,6 +1563,8 @@ def _find_temporal_column_for_query( def _build_schema_grounded_table_question_sql( self, query: str, table_ddls: list[str] ) -> str | None: + return None + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None @@ -1717,6 +1727,8 @@ def _build_schema_grounded_table_question_sql( def _build_explicit_table_preview_sql( self, query: str, table_ddls: list[str] ) -> tuple[str, str] | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip()) if not normalized_query: return None @@ -1974,6 +1986,8 @@ def _extract_explicit_table_column_reference( return table_name, column def _build_explicit_group_count_sql(self, query: str) -> str | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized_query: return None @@ -2198,6 +2212,8 @@ def _select_best_analytics_table( def _build_schema_grounded_analytics_sql( self, query: str, table_ddls: list[str] ) -> str | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized_query: return None @@ -3453,6 +3469,8 @@ def _build_audit_log_activity_sql( table_ddls: list[str], table_names: Optional[list[str]] = None, ) -> str | None: + return None + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None @@ -3803,6 +3821,8 @@ def _build_heuristic_text_to_sql_fallback( table_ddls: list[str], table_names: Optional[list[str]] = None, ) -> str | None: + return None + normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized: return None @@ -4011,6 +4031,8 @@ def _is_schema_grounded_query( def _build_schema_grounded_operational_sql( self, query: str, tables: list[dict[str, Any]] ) -> str | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) if not normalized_query: return None @@ -4584,6 +4606,8 @@ def _build_pcb_direct_question_sql( def _get_unqueryable_metric_message( self, query: str, table_ddls: list[str] ) -> str | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) normalized_schema = re.sub( r"\s+", @@ -4727,6 +4751,8 @@ def _get_unqueryable_metric_message( def _build_schema_grounded_sales_sql( self, query: str, table_ddls: list[str] ) -> str | None: + return None + normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) normalized_schema = "\n".join( ddl for ddl in table_ddls or [] if isinstance(ddl, str) @@ -5281,6 +5307,8 @@ def _prune_sql_generation_context( *, max_tables: int = 8, ) -> tuple[list[dict], list[str], list[str]]: + return documents, table_names, table_ddls + if len(table_ddls) <= max_tables: return documents, table_names, table_ddls @@ -5559,7 +5587,7 @@ def _build_validated_ask_result_from_sql( if invalid_tables or invalid_columns: logger.warning( - "Ignoring heuristic SQL because it is not valid for active schema. " + "Ignoring generated SQL because it is not valid for active schema. " "invalid_tables=%s invalid_columns=%s sql=%s", invalid_tables, invalid_columns, @@ -5567,42 +5595,6 @@ def _build_validated_ask_result_from_sql( ) return None - invalid_unqualified_identifiers = self._invalid_unqualified_sql_identifiers( - ask_result.sql, - schema_tables, - ) - if invalid_unqualified_identifiers: - logger.warning( - "Ignoring SQL because it references unqualified fields outside the active schema. " - "invalid_identifiers=%s sql=%s", - invalid_unqualified_identifiers, - ask_result.sql, - ) - return None - - invalid_output_aliases = self._invalid_sql_output_aliases( - ask_result.sql, - schema_tables, - ) - if invalid_output_aliases: - logger.warning( - "Ignoring SQL because it aliases output fields to unavailable schema concepts. " - "invalid_aliases=%s sql=%s", - invalid_output_aliases, - ask_result.sql, - ) - return None - - if not self._sql_references_explicit_table(ask_result.sql, query): - return None - - if not self._sql_matches_question_intent( - ask_result.sql, - query, - schema_tables, - ): - return None - return ask_result def _build_failed_text_to_sql_response( From 7f246890e1a6188e10c55f6ee23bbd98b182a86c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 00:43:12 +0530 Subject: [PATCH 0598/1087] Refactor ask service to LLM orchestration --- wren-ai-service/src/web/v1/services/ask.py | 5436 +------------------- 1 file changed, 174 insertions(+), 5262 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b8077c0f98..d9d6b42729 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -247,7 +247,6 @@ def _should_use_histories_for_query(cls, query: str | None) -> bool: contextual_patterns = ( r"\b(previous|last|above|earlier|same|those|that|these|them|it|its|there)\b", r"\b(add|break down|compare|filter|group|instead|only|sort|split)\b.+\b(by|to|with)\b", - r"\b(by|for|with)\s+(month|quarter|year|status|type|category|customer|market|region|country|division)\b", ) return any(re.search(pattern, normalized) for pattern in contextual_patterns) @@ -268,94 +267,6 @@ def _is_greeting_query(self, query: str) -> bool: } return normalized in greeting_patterns - def _is_data_analysis_query(self, query: str) -> bool: - return False - - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - analysis_terms = { - "amount", - "average", - "avg", - "bar chart", - "bottom", - "chart", - "common", - "compare", - "count", - "cost", - "claim", - "claims", - "currency", - "currencies", - "customer", - "customers", - "dashboard", - "debug", - "distribution", - "failure", - "fastest growing", - "growth", - "group", - "grouped", - "invoice", - "invoices", - "margin", - "market", - "markets", - "monthly", - "order", - "orders", - "pcb", - "performance", - "profit", - "product", - "products", - "product type", - "product types", - "quarter", - "quantity", - "rank", - "ranking", - "region", - "regions", - "repair", - "resolved", - "revenue", - "sale", - "sales", - "sales person", - "sales rep", - "salesperson", - "sla", - "top", - "trend", - "turnaround", - "value", - "volume", - "year", - "yearly", - } - return any(term in normalized for term in analysis_terms) - - def _needs_conversation_context(self, query: str) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - return any( - re.search(pattern, normalized) - for pattern in ( - r"\b(previous|last|above|earlier)\s+(query|question|answer|result|sql|chart)\b", - r"\b(same|that|those|them|it)\s+(table|query|question|result|chart|sql|period|filter)\b", - r"\b(use|using|based on|compare with|compared with)\s+(that|previous|last|above|earlier)\b", - r"\bwhat about\b", - r"\bhow about\b", - ) - ) - def _should_reuse_historical_question_sql( self, query: str, @@ -363,4427 +274,163 @@ def _should_reuse_historical_question_sql( ) -> bool: return False - def _rewrite_query_for_text_to_sql(self, query: str) -> str: - return query - - def _schema_contains( - self, - table_ddls: list[str], - pattern: str, - table_names: Optional[list[str]] = None, - ) -> bool: - schema_text = "\n".join(table_ddls or []) - if table_names: - schema_text += "\n" + "\n".join(table_names) - return bool(re.search(pattern, schema_text, flags=re.IGNORECASE)) - - def _schema_has_table_column( - self, - table_ddls: list[str], - table_name: str, - column_name: str, - table_names: Optional[list[str]] = None, - ) -> bool: - table_pattern = rf"\b{re.escape(table_name)}\b" - column_pattern = rf"\b{re.escape(column_name)}\b" - - for ddl in table_ddls or []: - if re.search(table_pattern, ddl, flags=re.IGNORECASE) and re.search( - column_pattern, ddl, flags=re.IGNORECASE - ): - return True - - return False - - def _extract_schema_column_names(self, table_ddls: list[str]) -> list[str]: - column_names: list[str] = [] - non_column_prefixes = ( - "create ", - "constraint ", - "foreign ", - "primary ", - "unique ", - "index ", - ")", - "/*", - "--", - ) - - for ddl in table_ddls: - if not isinstance(ddl, str): - continue - for line in ddl.splitlines(): - stripped = line.strip().rstrip(",") - if not stripped: - continue - if stripped.lower().startswith(non_column_prefixes): - continue - - column_match = re.match( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_]*))\s+", - stripped, - ) - if not column_match: - continue - - column_name = next( - value for value in column_match.groupdict().values() if value - ) - column_names.append(str(column_name).lower()) - - return column_names - def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: tables: list[dict[str, Any]] = [] for ddl in table_ddls or []: if not isinstance(ddl, str): - continue - table_match = re.search( - r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", - ddl, - flags=re.IGNORECASE, - ) - if not table_match: - continue - - table_name = next( - (value for value in table_match.groupdict().values() if value), - None, - ) - if not table_name: - continue - body_start = table_match.end() - depth = 1 - body_end = body_start - while body_end < len(ddl) and depth > 0: - if ddl[body_end] == "(": - depth += 1 - elif ddl[body_end] == ")": - depth -= 1 - body_end += 1 - - columns: list[dict[str, str]] = [] - for line in ddl[body_start : body_end - 1].splitlines(): - stripped = line.strip().rstrip(",") - if not stripped or stripped.startswith(("--", "/*")): - continue - if re.match( - r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", - stripped, - flags=re.IGNORECASE, - ): - continue - - column_match = re.match( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_$]*))" - r"\s+(?P[A-Za-z0-9_(),]+)", - stripped, - ) - if column_match: - column_name = next( - (value - for key, value in column_match.groupdict().items() - if key != "type" and value - ), - None, - ) - if not column_name: - continue - column_type = column_match.group("type") or "" - columns.append( - { - "name": str(column_name), - "type": str(column_type).lower(), - } - ) - - tables.append({"name": table_name, "columns": columns}) - - return tables - - _INTENT_STOPWORDS = { - "a", - "an", - "and", - "are", - "as", - "based", - "be", - "by", - "can", - "chart", - "correct", - "data", - "different", - "do", - "does", - "each", - "for", - "from", - "give", - "how", - "in", - "is", - "it", - "list", - "many", - "me", - "of", - "on", - "or", - "per", - "question", - "rate", - "records", - "reduce", - "show", - "system", - "taken", - "the", - "there", - "to", - "total", - "type", - "types", - "what", - "which", - "with", - } - - def _intent_tokens(self, text: str) -> set[str]: - tokens: set[str] = set() - for raw_token in re.findall(r"[A-Za-z][A-Za-z0-9_]*", text or ""): - split_token = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", raw_token) - for token in re.findall(r"[A-Za-z0-9]+", split_token.lower()): - if len(token) <= 2 or token in self._INTENT_STOPWORDS: - continue - tokens.add(token) - if token.endswith("ies") and len(token) > 4: - tokens.add(token[:-3] + "y") - elif token.endswith("s") and len(token) > 3: - tokens.add(token[:-1]) - return tokens - - def _schema_name_tokens(self, name: str) -> set[str]: - spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(name or "")) - return { - token - for token in re.findall(r"[A-Za-z0-9]+", spaced.lower()) - if len(token) > 1 - } - - def _table_for_sql_reference( - self, table_reference: str, valid_tables: dict[str, dict[str, Any]] - ) -> dict[str, Any] | None: - table_key = str(table_reference or "").lower() - if table_key in valid_tables: - return valid_tables[table_key] - suffix_key = table_key.split(".")[-1] - for valid_table_name, table in valid_tables.items(): - if valid_table_name.split(".")[-1] == suffix_key: - return table - return None - - def _required_sql_concept_groups(self, query: str | None) -> list[set[str]]: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return [] - - concept_groups: list[set[str]] = [] - if "product line" in normalized or "productline" in normalized: - concept_groups.append({"product", "prod", "line", "productline"}) - if "pcb" in normalized: - concept_groups.append( - { - "pcb", - "board", - "repair", - "repairs", - "debug", - "failure", - "failures", - } - ) - if "critical" in normalized: - concept_groups.append({"critical", "severity", "priority"}) - if "cost" in normalized: - concept_groups.append({"cost", "amount", "expense", "impact"}) - if "currency" in normalized or "currencies" in normalized: - concept_groups.append({"currency", "curr", "money", "fx", "exchange"}) - if "market" in normalized or "markets" in normalized: - concept_groups.append({"market", "region", "country", "territory"}) - if "region" in normalized or "regions" in normalized: - concept_groups.append({"region", "market", "area", "territory", "country"}) - if "quarterly" in normalized or "quarter" in normalized: - concept_groups.append({"quarter", "quarterly"}) - if "recurring" in normalized or "recurrence" in normalized: - concept_groups.append({"recurring", "recurrence", "occurrence", "occurrences", "count"}) - if "issue" in normalized or "issues" in normalized: - concept_groups.append({"issue", "issues", "failure", "failures", "problem", "defect"}) - - return concept_groups - - def _sql_covers_required_question_concepts( - self, - sql: str, - query: str | None, - referenced_column_tokens: set[str], - referenced_table_tokens: set[str], - ) -> bool: - sql_text = (sql or "").lower() - available_tokens = referenced_column_tokens | referenced_table_tokens - for concept_group in self._required_sql_concept_groups(query): - if concept_group & available_tokens: - continue - if any(token in sql_text for token in concept_group): - continue - logger.warning( - "Ignoring SQL because it does not cover required question concept. " - "query=%s required=%s referenced_column_tokens=%s referenced_table_tokens=%s sql=%s", - query, - sorted(concept_group), - sorted(referenced_column_tokens), - sorted(referenced_table_tokens), - sql, - ) - return False - return True - - def _sql_uses_required_measure_aggregation(self, sql: str, query: str | None) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return True - asks_for_measure_sum = any( - term in normalized_query - for term in ( - "amount", - "cost", - "revenue", - "sales value", - "sum", - "total", - "value", - ) - ) - asks_for_count = any( - term in normalized_query - for term in ( - "count", - "counts", - "how many", - "number of", - "record count", - "records", - "rows", - ) - ) - if not asks_for_measure_sum or asks_for_count: - return True - - normalized_sql = re.sub(r"\s+", " ", sql or "").lower() - if re.search(r"\b(sum|avg|min|max)\s*\(", normalized_sql): - return True - if re.search(r"\bcount\s*\(", normalized_sql): - logger.warning( - "Ignoring SQL because a measure-total question was answered with row counting. " - "query=%s sql=%s", - query, - sql, - ) - return False - return True - - def _sql_satisfies_unique_entity_request(self, sql: str, query: str | None) -> bool: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not re.search( - r"\b(?:distinct|unique|no duplicate|no duplicates|without duplicates)\b", - normalized_query, - ): - return True - - normalized_sql = re.sub(r"\s+", " ", sql or "").strip() - if re.search(r"\bcount\s*\(\s*distinct\b", normalized_sql, flags=re.IGNORECASE): - return True - if re.search(r"\bGROUP\s+BY\b", normalized_sql, flags=re.IGNORECASE): - group_match = re.search( - r"\bGROUP\s+BY\b(?P.*?)(?:\bORDER\s+BY\b|\bHAVING\b|$)", - normalized_sql, - flags=re.IGNORECASE | re.DOTALL, - ) - if group_match: - group_items = [ - item.strip() - for item in re.split(r",(?![^()]*\))", group_match.group("group")) - if item.strip() - ] - if len(group_items) > 1: - logger.warning( - "Ignoring SQL because GROUP BY covers multiple columns and can still duplicate the requested entity. " - "query=%s sql=%s", - query, - sql, - ) - return False - return True - select_match = re.search( - r"\bSELECT\b(?P.*?)\bFROM\b", - sql or "", - flags=re.IGNORECASE | re.DOTALL, - ) - if not select_match: - return [] - - invalid: list[str] = [] - for match in re.finditer( - r'(?P(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$.]*))\s+AS\s+' - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_]*))', - select_match.group("select"), - flags=re.IGNORECASE, - ): - expression = match.group("expr").strip('"[]') - expression_key = expression.split(".")[-1].lower() - alias = match.group("quoted") or match.group("bracketed") or match.group("bare") or "" - alias_key = alias.lower() - if not expression_key or not alias_key: - continue - if alias_key in valid_columns or expression_key not in valid_columns: - continue - alias_terms = { - term for term in re.split(r"[^a-z0-9]+", alias_key) if term - } - if alias_terms & allowed_alias_terms: - continue - invalid.append(alias) - - return invalid - - def _unqualified_valid_sql_column_tokens( - self, sql: str, schema_tables: list[dict[str, Any]] - ) -> set[str]: - valid_columns = { - str(column.get("name") or "").lower(): str(column.get("name") or "") - for table in schema_tables - for column in table.get("columns", []) - if column.get("name") - } - if not valid_columns: - return set() - - sql_without_strings = re.sub(r"'(?:''|[^'])*'", "''", sql or "") - identifier_pattern = re.compile( - r'"(?P[^"]+)"|(?P\b[A-Za-z_][A-Za-z0-9_]*\b)' - ) - tokens: set[str] = set() - for match in identifier_pattern.finditer(sql_without_strings): - identifier = match.group("quoted") or match.group("bare") or "" - identifier_key = identifier.lower() - if identifier_key not in valid_columns: - continue - - before = sql_without_strings[: match.start()].rstrip() - after = sql_without_strings[match.end() :].lstrip() - if before.endswith(".") or after.startswith("."): - continue - - previous_word_match = re.search(r"([A-Za-z_][A-Za-z0-9_]*)\s*$", before) - previous_word = ( - previous_word_match.group(1).lower() if previous_word_match else "" - ) - if previous_word == "as" and self._is_alias_identifier_position( - sql_without_strings, - match.start(), - ): - continue - - tokens.update(self._schema_name_tokens(valid_columns[identifier_key])) - - return tokens - - def _sql_matches_question_intent( - self, - sql: str, - query: str | None, - schema_tables: list[dict[str, Any]], - ) -> bool: - return True - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - expects_dimension = bool( - re.search( - r"\b(?:by|per|each|which|different|type|types|category|" - r"categories|status|source|market|markets|region|regions|" - r"currency|currencies)\b", - normalized_query, - ) - ) - - question_tokens = self._intent_tokens(query or "") - required_concept_groups = self._required_sql_concept_groups(query) - if not question_tokens and not required_concept_groups: - return True - - valid_tables = { - str(table.get("name") or "").lower(): table - for table in schema_tables - if table.get("name") - } - if not valid_tables: - return True - - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_tables = [ - next(value for value in match.groupdict().values() if value) - for match in table_reference_pattern.finditer(sql) - ] - if not referenced_tables: - return True - - referenced_table_tokens = set().union( - *[self._schema_name_tokens(table) for table in referenced_tables] - ) - qualified_column_pattern = re.compile( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_$]*))", - flags=re.IGNORECASE, - ) - referenced_columns_by_table: dict[str, set[str]] = {} - for match in qualified_column_pattern.finditer(sql): - table_reference = ( - match.group("table_quoted") - or match.group("table_bracketed") - or match.group("table_bare") - or "" - ).lower() - column_reference = ( - match.group("column_quoted") - or match.group("column_bracketed") - or match.group("column_bare") - or "" - ) - referenced_columns_by_table.setdefault(table_reference, set()).add( - column_reference - ) - - all_referenced_column_tokens = set().union( - *[ - self._schema_name_tokens(column_name) - for columns in referenced_columns_by_table.values() - for column_name in columns - ] - ) if referenced_columns_by_table else set() - all_referenced_column_tokens.update( - self._unqualified_valid_sql_column_tokens(sql, schema_tables) - ) - if not self._sql_covers_required_question_concepts( - sql, - query, - all_referenced_column_tokens, - referenced_table_tokens, - ): - return False - if not self._sql_uses_required_measure_aggregation(sql, query): - return False - if not self._sql_satisfies_unique_entity_request(sql, query): - return False - - if not expects_dimension: - return True - - for table_reference in referenced_tables: - table = self._table_for_sql_reference(table_reference, valid_tables) - if not table: - continue - - columns = [ - column for column in table.get("columns", []) if column.get("name") - ] - intent_matching_columns = [ - str(column.get("name")) - for column in columns - if self._schema_name_tokens(str(column.get("name"))) & question_tokens - ] - if not intent_matching_columns: - continue - - table_key = str(table_reference or "").lower() - referenced_columns = referenced_columns_by_table.get( - table_key - ) or referenced_columns_by_table.get( - table_key.split(".")[-1], - set(), - ) - referenced_column_tokens = ( - set().union( - *[ - self._schema_name_tokens(column_name) - for column_name in referenced_columns - ] - ) - if referenced_columns - else all_referenced_column_tokens - ) - - if not referenced_column_tokens & question_tokens: - logger.warning( - "Ignoring SQL because selected columns do not match question intent. " - "query=%s table=%s matching_schema_columns=%s referenced_columns=%s sql=%s", - query, - table.get("name"), - intent_matching_columns, - sorted(referenced_columns), - sql, - ) - return False - - return True - - def _is_numeric_schema_type(self, column_type: str) -> bool: - return bool( - re.search( - r"\b(?:int|integer|bigint|smallint|tinyint|decimal|numeric|float|double|" - r"real|money|number)\b", - column_type, - flags=re.IGNORECASE, - ) - ) - - def _is_temporal_schema_type(self, column_type: str) -> bool: - return bool( - re.search( - r"\b(?:date|time|timestamp|datetime|smalldatetime)\b", - column_type, - flags=re.IGNORECASE, - ) - ) - - def _is_text_schema_type(self, column_type: str) -> bool: - return bool( - re.search( - r"\b(?:char|text|string|varchar|nvarchar|uuid|guid|json)\b", - column_type, - flags=re.IGNORECASE, - ) - ) - - def _find_schema_column( - self, - table: dict[str, Any], - candidates: tuple[str, ...], - numeric: bool | None = None, - temporal: bool | None = None, - ) -> str | None: - normalized_candidates = [ - re.sub(r"[^a-z0-9]", "", str(candidate).lower()) - for candidate in candidates - if candidate is not None - ] - if not normalized_candidates: - return None - scored: list[tuple[int, int, str]] = [] - for column in table.get("columns", []): - column_name = column.get("name") - if not column_name: - continue - column_name = str(column_name) - normalized_column = re.sub(r"[^a-z0-9]", "", column_name.lower()) - column_type = str(column.get("type") or "") - if numeric is True and not self._is_numeric_schema_type(column_type): - continue - if temporal is True and not self._is_temporal_schema_type(column_type): - continue - - for candidate_index, candidate in enumerate(normalized_candidates): - if normalized_column == candidate: - scored.append((100, candidate_index, column_name)) - elif candidate and candidate in normalized_column: - scored.append((60 + len(candidate), candidate_index, column_name)) - elif normalized_column and normalized_column in candidate: - scored.append( - (40 + len(normalized_column), candidate_index, column_name) - ) - - if not scored: - return None - - return sorted(scored, key=lambda item: (-item[0], item[1]))[0][2] - - def _find_first_schema_column( - self, - table: dict[str, Any], - candidates: tuple[str, ...], - *, - avoid: set[str] | None = None, - ) -> str | None: - avoid = {str(column).lower() for column in avoid or set()} - for candidate_group in candidates: - column = self._find_schema_column(table, (candidate_group,)) - if column and column.lower() not in avoid: - return column - return None - - def _find_any_temporal_schema_column(self, table: dict[str, Any]) -> str | None: - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if column_name and self._is_temporal_schema_type(column_type): - return column_name - return None - - def _is_probable_explicit_table_token(self, token: str) -> bool: - normalized = re.sub(r"\s+", " ", str(token or "").strip().lower()) - if not normalized: - return False - if normalized in self._INTENT_STOPWORDS or normalized in { - "last", - "latest", - "recent", - "current", - "previous", - "next", - "month", - "year", - "quarter", - "week", - "day", - }: - return False - return True - - def _quote_sql_identifier(self, identifier: str) -> str: - return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' - - def _normalize_schema_identifier_key(self, value: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - def _schema_identifier_alias_keys(self, value: str) -> set[str]: - raw_value = str(value or "").strip() - base_key = self._normalize_schema_identifier_key(raw_value) - separator_normalized_key = self._normalize_schema_identifier_key( - re.sub(r"[.$]", "_", raw_value) - ) - compact_parts_key = "".join( - self._normalize_schema_identifier_key(part) - for part in re.split(r"[.$_]+", raw_value) - if part - ) - return { - key - for key in (base_key, separator_normalized_key, compact_parts_key) - if key - } - - def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: - return self._explicit_table_alias_keys( - self._extract_explicit_table_names_from_query(query or "") - ) - - def _explicit_table_alias_keys(self, table_names: list[str]) -> set[str]: - keys: set[str] = set() - for table_name in table_names: - keys.update(self._schema_identifier_alias_keys(table_name)) - return keys - - def _filter_retrieval_metadata_for_explicit_query( - self, - query: str, - documents: list[dict], - explicit_table_names: Optional[list[str]] = None, - ) -> tuple[list[dict], list[str], list[str]]: - explicit_table_keys = ( - self._explicit_table_alias_keys(explicit_table_names) - if explicit_table_names - else self._explicit_table_alias_keys_from_query(query) - ) - if not explicit_table_keys: - table_names, table_ddls = self._metadata_from_documents(documents) - return documents, table_names, table_ddls - - matched_documents: list[dict] = [] - for document in documents: - candidate_names = [] - if isinstance(table_name := document.get("table_name"), str): - candidate_names.append(table_name) - if isinstance(table_ddl := document.get("table_ddl"), str): - candidate_names.extend( - str(table.get("name") or "") - for table in self._parse_schema_tables([table_ddl]) - if table.get("name") - ) - - candidate_keys: set[str] = set() - for candidate_name in candidate_names: - candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) - candidate_keys.update( - self._schema_identifier_alias_keys( - re.split(r"[.$_]", str(candidate_name or ""))[-1] - ) - ) - if explicit_table_keys.intersection(candidate_keys): - matched_documents.append(document) - - table_names, table_ddls = self._metadata_from_documents(matched_documents) - return matched_documents, table_names, table_ddls - - def _sql_references_explicit_table( - self, - sql: str, - query: str | None, - ) -> bool: - return True - - explicit_table_keys = self._explicit_table_alias_keys_from_query(query) - if not explicit_table_keys: - return True - - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_tables = [ - next(value for value in match.groupdict().values() if value) - for match in table_reference_pattern.finditer(sql) - ] - - for table_reference in referenced_tables: - reference_keys = self._schema_identifier_alias_keys(table_reference) - reference_keys.update( - self._schema_identifier_alias_keys( - re.split(r"[.$_]", str(table_reference or ""))[-1] - ) - ) - if explicit_table_keys.intersection(reference_keys): - return True - - logger.warning( - "Ignoring SQL because it does not reference the explicitly requested table. " - "query=%s referenced_tables=%s sql=%s", - query, - referenced_tables, - sql, - ) - return False - - def _table_matches_query(self, table_name: str, query: str) -> bool: - query_keys = self._schema_identifier_alias_keys(query) - short_table = re.split(r"[.$_]", str(table_name or ""))[-1] - table_keys = self._schema_identifier_alias_keys(table_name) - table_keys.update(self._schema_identifier_alias_keys(short_table)) - return any( - table_key and any(table_key in query_key for query_key in query_keys) - for table_key in table_keys - ) - - def _find_best_schema_table_for_query( - self, query: str, tables: list[dict[str, Any]] - ) -> dict[str, Any] | None: - if not tables: - return None - - scored_tables: list[tuple[int, dict[str, Any]]] = [] - query_tokens = self._intent_tokens(query) - for table in tables: - table_name = str(table.get("name") or "") - if not table_name: - continue - - score = 0 - if self._table_matches_query(table_name, query): - score += 100 - - table_tokens = self._schema_name_tokens(table_name) - score += 8 * len(table_tokens & query_tokens) - - column_token_matches = 0 - for column in table.get("columns", []): - column_token_matches += len( - self._schema_name_tokens(str(column.get("name") or "")) - & query_tokens - ) - score += column_token_matches - - if score > 0: - scored_tables.append((score, table)) - - if scored_tables: - return sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] - if len(tables) == 1: - return tables[0] - return None - - def _query_mentions_column(self, query: str, column_name: str) -> bool: - normalized_query = self._normalize_schema_identifier_key(query) - normalized_column = self._normalize_schema_identifier_key(column_name) - if not normalized_column: - return False - if normalized_column in normalized_query: - return True - if normalized_column.endswith("y"): - return f"{normalized_column[:-1]}ies" in normalized_query - return f"{normalized_column}s" in normalized_query - - def _is_alias_identifier_position(self, sql: str, start: int) -> bool: - before = sql[:start].rstrip() - return bool(re.search(r"\bAS\s*$", before, flags=re.IGNORECASE)) - - def _find_dimension_column_for_query( - self, query: str, table: dict[str, Any] - ) -> str | None: - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - if column_name and self._query_mentions_column(query, column_name): - return column_name - - candidate_columns = [ - str(column.get("name")) - for column in table.get("columns", []) - if column.get("name") - and not self._is_temporal_schema_type(str(column.get("type") or "")) - ] - for candidate in ("name", "category", "type", "status", "code"): - column = self._find_schema_column(table, (candidate,)) - if column in candidate_columns: - return column - return candidate_columns[0] if candidate_columns else None - - def _build_schema_ranked_measure_sql( - self, - query: str, - table_ddls: list[str], - ) -> str | None: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - if not re.search( - r"\b(?:top|highest|largest|biggest|best|lowest|smallest|bottom)\b", - normalized_query, - ): - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - query_tokens = self._intent_tokens(query) - if not query_tokens: - return None - - scored: list[tuple[int, dict[str, Any], str, str]] = [] - for table in tables: - text_columns: list[tuple[int, str]] = [] - numeric_columns: list[tuple[int, str]] = [] - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if not column_name: - continue - - column_tokens = self._schema_name_tokens(column_name) - query_overlap = column_tokens & query_tokens - - if self._is_numeric_schema_type(column_type): - if query_overlap: - numeric_columns.append( - (80 + 5 * len(query_overlap), column_name) - ) - continue - - if self._is_temporal_schema_type(column_type): - continue - - if query_overlap: - text_columns.append((60 + 5 * len(query_overlap), column_name)) - - if not text_columns or not numeric_columns: - continue - - dimension_score, dimension = sorted( - text_columns, - key=lambda item: item[0], - reverse=True, - )[0] - measure_score, measure = sorted( - numeric_columns, - key=lambda item: item[0], - reverse=True, - )[0] - table_score = dimension_score + measure_score - table_tokens = self._schema_name_tokens(str(table.get("name") or "")) - table_score += 5 * len(table_tokens & query_tokens) - scored.append((table_score, table, dimension, measure)) - - if not scored: - return None - - _, table, dimension, measure = sorted( - scored, - key=lambda item: item[0], - reverse=True, - )[0] - table_name = str(table.get("name") or "") - if not table_name: - return None - - limit = self._extract_requested_top_n(query, default_value=10) - table_ref = self._quote_sql_identifier(table_name) - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - measure_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" - metric_expr = f"SUM({measure_ref})" - direction = ( - "ASC" - if re.search(r"\b(?:lowest|smallest|least|bottom)\b", normalized_query) - else "DESC" - ) - return ( - f"SELECT TOP {limit} {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " - f"{metric_expr} AS {self._quote_sql_identifier('Total' + measure)} " - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL AND {measure_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f"ORDER BY {metric_expr} {direction}" - ) - - def _find_temporal_column_for_query( - self, query: str, table: dict[str, Any] - ) -> str | None: - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if ( - column_name - and self._is_temporal_schema_type(column_type) - and self._query_mentions_column(query, column_name) - ): - return column_name - - return self._find_any_temporal_schema_column(table) - - def _build_schema_grounded_table_question_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - return None - - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - tables = self._parse_schema_tables(table_ddls) - table = self._find_best_schema_table_for_query(query, tables) - if not table: - return None - - table_name = str(table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - limit = self._extract_requested_top_n(query, default_value=10) - - wants_latest_records = any( - term in normalized - for term in ("latest", "recent", "newest", "last records", "latest records") - ) - if wants_latest_records: - date_column = self._find_temporal_column_for_query(query, table) - if not date_column: - return None - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT TOP {limit} * " - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"ORDER BY {date_ref} DESC" - ) - - wants_monthly_count = any( - term in normalized - for term in ("monthly", "by month", "per month", "month-wise") - ) and any(term in normalized for term in ("count", "records", "rows")) - if wants_monthly_count: - date_column = self._find_temporal_column_for_query(query, table) - if not date_column: - return None - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " - f"DATEPART(MONTH, {date_ref}) ASC" - ) - - wants_total_count = ( - re.search(r"\bhow many\b", normalized) - or "record count" in normalized - or "count of records" in normalized - or "number of records" in normalized - ) and not re.search(r"\b(?:by|per|each|distribution|highest|top)\b", normalized) - if wants_total_count: - return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' - - wants_grouped_count = ( - any( - term in normalized - for term in ( - "count", - "counts", - "record count", - "number of", - "how many", - ) - ) - and re.search(r"\b(?:by|per|each|grouped by|group by)\b", normalized) - ) - wants_ranked_count = any( - term in normalized for term in ("highest", "top", "most", "largest") - ) and any( - term in normalized - for term in ("count", "counts", "number of", "orders", "records", "rows") - ) - if wants_grouped_count or wants_ranked_count: - dimension_column = self._find_dimension_column_for_query(query, table) - if not dimension_column: - return None - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" - count_column = None - if any(term in normalized for term in ("order", "orders")): - count_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "OrderID", "id"), - ) - count_expression = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(count_column)})" - if count_column - else "COUNT(*)" - ) - top_clause = f"TOP {limit} " if wants_ranked_count else "" - nonblank_filter = ( - f"AND LTRIM(RTRIM({dimension_ref})) <> '' " - if wants_ranked_count - else "" - ) - return ( - f"SELECT {top_clause}{dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_column)}, " - f'{count_expression} AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"{nonblank_filter}" - f"GROUP BY {dimension_ref} " - f"ORDER BY {count_expression} DESC" - ) - - wants_distribution = any( - term in normalized - for term in ( - "distribution", - "highest occurrence", - "highest occurrences", - "most occurrence", - "most occurrences", - "occurrences", - "top", - "common", - ) - ) - if wants_distribution: - dimension_column = self._find_dimension_column_for_query(query, table) - if not dimension_column: - return None - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension_column)}" - - occurrence_column = self._find_schema_column( - table, - ("occurrences", "occurrence", "count", "total_count", "record_count"), - numeric=True, - ) - if occurrence_column and self._query_mentions_column( - query, occurrence_column - ): - metric_ref = f"{table_ref}.{self._quote_sql_identifier(occurrence_column)}" - return ( - f"SELECT TOP {limit} {dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_column)}, " - f"{metric_ref} AS {self._quote_sql_identifier(occurrence_column)} " - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"ORDER BY {metric_ref} DESC" - ) - - return ( - f"SELECT TOP {limit} {dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_column)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - return None - - def _build_explicit_table_preview_sql( - self, query: str, table_ddls: list[str] - ) -> tuple[str, str] | None: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip()) - if not normalized_query: - return None - - if not re.search( - r"\b(?:first|top|sample|preview|show|list)\b", - normalized_query, - flags=re.IGNORECASE, - ): - return None - if not re.search( - r"\b(?:rows?|records?|data)\b", normalized_query, flags=re.IGNORECASE - ): - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - normalized_query_key = re.sub(r"[^a-z0-9]", "", normalized_query.lower()) - normalized_query_keys = self._schema_identifier_alias_keys(normalized_query) - scored_tables: list[tuple[int, str]] = [] - for table in tables: - table_name = table.get("name") - if not table_name: - continue - table_name = str(table_name) - normalized_table = re.sub(r"[^a-z0-9]", "", table_name.lower()) - if not normalized_table: - continue - table_keys = self._schema_identifier_alias_keys(table_name) - if normalized_table in normalized_query_key or any( - table_key in query_key - for table_key in table_keys - for query_key in normalized_query_keys - ): - scored_tables.append((100 + len(normalized_table), table_name)) - continue - - table_without_schema = re.split(r"[.$]", table_name)[-1] - normalized_short_name = re.sub( - r"[^a-z0-9]", "", table_without_schema.lower() - ) - if normalized_short_name and normalized_short_name in normalized_query_key: - scored_tables.append((80 + len(normalized_short_name), table_name)) - - if not scored_tables: - return None - - _, table_name = sorted(scored_tables, reverse=True)[0] - limit = self._extract_requested_top_n(query, default_value=10) - return ( - f"SELECT TOP {limit} * FROM {self._quote_sql_identifier(table_name)}", - table_name, - ) - - def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: - table_names: list[str] = [] - for match in re.finditer( - r"\b(?:from|table|model)\s+([A-Za-z_][A-Za-z0-9_.$]*)", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - for match in re.finditer( - r"\bin\s+(?:the\s+)?([A-Za-z_][A-Za-z0-9_.$]*)\s+table\b", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - for match in re.finditer( - r"\bin\s+([A-Za-z_][A-Za-z0-9_.$]*)", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and ("." in table_name or "_" in table_name) - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - for match in re.finditer( - r"\busing\s+([A-Za-z_][A-Za-z0-9_.$]*)", - query or "", - flags=re.IGNORECASE, - ): - table_name = match.group(1).strip(".,;:()[]{}") - if ( - table_name - and ("." in table_name or "_" in table_name) - and self._is_probable_explicit_table_token(table_name) - and table_name not in table_names - ): - table_names.append(table_name) - return table_names - - def _explicit_table_name_candidates(self, table_name: str) -> list[str]: - table_name = str(table_name or "").strip().strip(".,;:()[]{}") - if not table_name: - return [] - - candidates = [table_name] - separator_normalized = re.sub(r"[.$]", "_", table_name) - if separator_normalized not in candidates: - candidates.append(separator_normalized) - if "_" in table_name and "." not in table_name: - dotted = table_name.replace("_", ".", 1) - if dotted not in candidates: - candidates.append(dotted) - short_name = re.split(r"[.$]", table_name)[-1] - if "." not in table_name and "_" in table_name: - short_name = table_name.split("_", 1)[-1] - if short_name and short_name not in candidates: - candidates.append(short_name) - return candidates - - def _normalize_explicit_table_names( - self, table_names: Optional[list[str]] - ) -> list[str]: - normalized: list[str] = [] - for table_name in table_names or []: - for candidate in self._explicit_table_name_candidates(table_name): - if candidate not in normalized: - normalized.append(candidate) - return normalized - - def _build_direct_orders_sales_sql(self, query: str) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - is_sales_or_orders_query = any( - term in normalized - for term in ( - "sales", - "sale", - "order", - "orders", - "new order", - "new orders", - "market", - "salesperson", - "sales person", - ) - ) - if not is_sales_or_orders_query: - return None - - table_ref = '"dbo_tblSales"' - limit = self._extract_requested_top_n(query, default_value=10) - - if ( - ("salesperson" in normalized or "sales person" in normalized) - and ("order count" in normalized or "orders" in normalized or "count" in normalized) - ): - return ( - f'SELECT TOP {limit} {table_ref}."SalesPerson" AS "SalesPerson", ' - f'COUNT(*) AS "OrderCount" ' - f"FROM {table_ref} " - f'WHERE {table_ref}."SalesPerson" IS NOT NULL ' - f'GROUP BY {table_ref}."SalesPerson" ' - f"ORDER BY COUNT(*) DESC" - ) - - if "top" in normalized and "new order" in normalized: - date_filter = "" - if re.search(r"\b2026[\s-]*q1\b", normalized): - date_filter = ( - f'WHERE {table_ref}."OrdDate" >= \'2026-01-01 00:00:00\' ' - f'AND {table_ref}."OrdDate" < \'2026-04-01 00:00:00\' ' - ) - return ( - f'SELECT TOP {limit} {table_ref}."BU" AS "BU", ' - f'{table_ref}."Market" AS "Market", ' - f'{table_ref}."Customer" AS "Customer", ' - f'{table_ref}."ProdName" AS "ProdName", ' - f'{table_ref}."SalesValue" AS "SalesValue" ' - f"FROM {table_ref} " - f"{date_filter}" - f'ORDER BY {table_ref}."SalesValue" DESC' - ) - - if ( - "market" in normalized - and "growth" in normalized - and any(term in normalized for term in ("last year", "previous year")) - ): - return ( - f'SELECT {table_ref}."Market" AS "Market", ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "CurrentPeriodSales", ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "PreviousPeriodSales", ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2026-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2026-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) - ' - f"SUM(CASE WHEN {table_ref}.\"OrdDate\" >= '2025-01-01 00:00:00' " - f"AND {table_ref}.\"OrdDate\" < '2025-07-01 00:00:00' " - f'THEN {table_ref}."SalesValue" ELSE 0 END) AS "SalesGrowth" ' - f"FROM {table_ref} " - f'WHERE {table_ref}."Market" IS NOT NULL ' - f'GROUP BY {table_ref}."Market" ' - f'ORDER BY "SalesGrowth" DESC' - ) - - if ( - "distribution" in normalized - and "sales" in normalized - and ("market" in normalized or "by market" in normalized) - ): - return ( - f'SELECT {table_ref}."Market" AS "Market", ' - f'SUM({table_ref}."SalesValue") AS "TotalSalesValue" ' - f"FROM {table_ref} " - f'WHERE {table_ref}."Market" IS NOT NULL ' - f'GROUP BY {table_ref}."Market" ' - f'ORDER BY SUM({table_ref}."SalesValue") DESC' - ) - - return None - - def _extract_explicit_table_column_reference( - self, query: str - ) -> tuple[str, str] | None: - normalized_query = query or "" - reference_match = re.search( - r"\b(?P[A-Za-z_][A-Za-z0-9_]*)[._]" - r"(?P
[A-Za-z_][A-Za-z0-9_]*)[._]" - r"(?P[A-Za-z_][A-Za-z0-9_]*)\b", - normalized_query, - ) - if not reference_match: - return None - - schema = reference_match.group("schema") - table = reference_match.group("table") - column = reference_match.group("column") - table_name = f"{schema}_{table}" - return table_name, column - - def _build_explicit_group_count_sql(self, query: str) -> str | None: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - if not any( - term in normalized_query - for term in ( - "group by", - "grouped by", - "by ", - "pie chart", - "donut chart", - "bar chart", - "count", - "counts", - ) - ): - return None - - explicit_reference = self._extract_explicit_table_column_reference(query) - if not explicit_reference: - return None - - table_name, column = explicit_reference - table_ref = self._quote_sql_identifier(table_name) - column_ref = f"{table_ref}.{self._quote_sql_identifier(column)}" - return ( - f"SELECT {column_ref} AS {self._quote_sql_identifier(column)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"GROUP BY {column_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - def _build_date_filter(self, table_name: str, date_column: str, query: str) -> str: - date_ref = ( - f"{self._quote_sql_identifier(table_name)}." - f"{self._quote_sql_identifier(date_column)}" - ) - normalized_query = (query or "").lower() - if "this year" in normalized_query or "current year" in normalized_query: - return ( - f" WHERE {date_ref} >= '2026-01-01 00:00:00' " - f"AND {date_ref} < '2027-01-01 00:00:00'" - ) - year_match = re.search(r"\b(20\d{2})\b", normalized_query) - if year_match: - year = int(year_match.group(1)) - return ( - f" WHERE {date_ref} >= '{year}-01-01 00:00:00' " - f"AND {date_ref} < '{year + 1}-01-01 00:00:00'" - ) - return "" - - def _append_not_null_filters( - self, where_clause: str, column_refs: list[str] - ) -> str: - conditions = [f"{column_ref} IS NOT NULL" for column_ref in column_refs] - if not conditions: - return where_clause - - if where_clause.strip(): - return f"{where_clause.rstrip()} AND {' AND '.join(conditions)}" - return f" WHERE {' AND '.join(conditions)}" - - def _build_schema_literal_filter_conditions( - self, - query: str, - table: dict[str, Any], - table_ref: str, - ) -> list[str]: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - conditions: list[str] = [] - - if "backlog" in normalized_query: - filter_column = self._find_schema_column( - table, - ( - "Category", - "OrderCategory", - "Order Category", - "Status", - "OrderStatus", - "Order Status", - "Stage", - "OrderStage", - "Order Stage", - ), - ) - if filter_column: - filter_ref = f"{table_ref}.{self._quote_sql_identifier(filter_column)}" - conditions.append(f"{filter_ref} = 'Backlog'") - - return conditions - - def _select_best_analytics_table( - self, - tables: list[dict[str, Any]], - required_dimensions: list[tuple[str, ...]], - measure_candidates: tuple[str, ...], - wants_date: bool = False, - allow_count_metric: bool = False, - query: str = "", - ) -> tuple[dict[str, Any], list[str], str | None, str | None] | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - scored: list[ - tuple[int, dict[str, Any], list[str], str | None, str | None] - ] = [] - for table in tables: - dimensions = [ - self._find_schema_column(table, candidates) - for candidates in required_dimensions - ] - if any(dimension is None for dimension in dimensions): - continue - - measure = self._find_schema_column( - table, measure_candidates, numeric=True - ) - if not measure and not allow_count_metric: - continue - date_column = self._find_schema_column( - table, - ( - "OrdDate", - "InvDate", - "OrderDate", - "NewOrderDate", - "Date", - "CreatedAt", - "created_at", - ), - temporal=True, - ) - if wants_date and not date_column: - continue - - score = 10 * len([dimension for dimension in dimensions if dimension]) - if measure: - score += 8 - elif allow_count_metric: - score += 2 - if date_column: - score += 4 - table_name = str(table.get("name") or "").lower() - if not table_name: - continue - if "sales" in table_name: - score += 5 - if "tblsales" in self._normalize_schema_token(table_name): - score += 25 - if "order" in table_name: - score += 4 - if "invoice" in table_name or "inv" in table_name: - score += 3 - if "stage" in table_name: - score -= 8 - if any( - term in normalized_query - for term in ("order", "orders", "new order", "new orders") - ): - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "OrderNumber"), - ) - if "order" in table_name: - score += 30 - if "neworder" in self._normalize_schema_token(table_name): - score += 15 - if order_column: - score += 12 - if "margin" in table_name and "margin" not in normalized_query: - score -= 12 - if "customer" in normalized_query: - if "customer" in table_name or "account" in table_name: - score += 16 - if self._find_schema_column( - table, - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - ), - ): - score += 10 - if any( - term in normalized_query for term in ("product", "products", "item") - ): - if "product" in table_name or "item" in table_name: - score += 16 - if any( - term in normalized_query - for term in ("sales", "revenue", "value", "amount") - ): - if "sales" in table_name: - score += 12 - if "invoice" in normalized_query and ( - "invoice" in table_name or "inv" in table_name - ): - score += 20 - - scored.append((score, table, dimensions, measure, date_column)) - - if not scored: - return None - - _, table, dimensions, measure, date_column = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - return ( - table, - [dimension for dimension in dimensions if dimension], - measure, - date_column, - ) - - def _build_schema_grounded_analytics_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - - if pcb_direct_sql := self._build_pcb_direct_question_sql(query, table_ddls): - return pcb_direct_sql - - if repair_failure_count_sql := self._build_repair_failure_count_sql( - query, table_ddls - ): - return repair_failure_count_sql - - if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( - query, table_ddls - ): - return monthly_repair_volume_sql - - is_sales_or_order_query = any( - term in normalized_query - for term in ( - "average order value", - "invoice", - "new order", - "order", - "orders", - "currency", - "currencies", - "market", - "markets", - "performance", - "product", - "products", - "quantity", - "qty", - "revenue", - "sale", - "sales", - "salesperson", - "sales person", - "sold", - ) - ) - if not is_sales_or_order_query: - if operational_sql := self._build_schema_grounded_operational_sql( - query, tables - ): - return operational_sql - - if conversion_sql := self._build_order_invoice_conversion_sql( - query, tables - ): - return conversion_sql - - if yoy_sql := self._build_yoy_sales_change_sql(query, tables): - return yoy_sql - - if contribution_sql := self._build_contribution_sql(query, tables): - return contribution_sql - - if not is_sales_or_order_query: - if categorical_count_sql := self._build_generic_categorical_count_sql( - query, tables - ): - return categorical_count_sql - - wants_count_metric = any( - term in normalized_query - for term in ("count", "counts", "volume", "how many", "distribution") - ) and not any( - term in normalized_query - for term in ( - "amount", - "cost", - "expense", - "quantity", - "qty", - "revenue", - "sale", - "sales", - "sold", - "sum", - "total", - "value", - ) - ) - wants_average_metric = any( - term in normalized_query for term in ("average", "avg", "mean") - ) - - wants_monthly_count = ( - "monthly" in normalized_query - and any(term in normalized_query for term in ("count", "volume")) - and any(term in normalized_query for term in ("order", "orders")) - ) - if wants_monthly_count: - date_candidates = ( - ("InvDate", "InvoiceDate", "Invoice Date") - if "invdate" in compact_query or "invoice" in normalized_query - else ( - "OrdDate", - "OrderDate", - "NewOrderDate", - "InvDate", - "InvoiceDate", - "Date", - ) - ) - scored_tables: list[tuple[int, dict[str, Any], str]] = [] - for table in tables: - date_column = self._find_schema_column( - table, date_candidates, temporal=True - ) - if not date_column: - continue - table_name = str(table.get("name") or "") - score = 10 - if "sales" in table_name.lower() or "order" in table_name.lower(): - score += 5 - if self._find_schema_column( - table, ("OrdNo", "OrderNo", "OrderId", "InvoiceNo") - ): - score += 3 - scored_tables.append((score, table, date_column)) - - if scored_tables: - _, table, date_column = sorted( - scored_tables, key=lambda item: item[0], reverse=True - )[0] - table_name = table.get("name") - if table_name and date_column: - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - date_ref = ( - f"{table_ref}.{self._quote_sql_identifier(date_column)}" - ) - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "OrderCount" ' - f"FROM {table_ref}" - f"{self._build_date_filter(table_name, date_column, query)} " - f"GROUP BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref})" - ) - - dimension_candidates: list[tuple[str, ...]] = [] - if "salesperson" in normalized_query or "sales person" in normalized_query: - dimension_candidates.append( - ("SalesPerson", "Salesman", "Sales Rep", "SalesRep", "Rep", "Owner") - ) - if "business unit" in normalized_query or re.search(r"\bbu\b", normalized_query): - dimension_candidates.append(("BusinessUnit", "Business Unit", "BU")) - if "market" in normalized_query: - dimension_candidates.append(("Market", "MarketType", "MarketName", "Region", "Country")) - if "region" in normalized_query: - dimension_candidates.append(("Region", "Market", "Area", "Territory")) - if "currency" in normalized_query or "currencies" in normalized_query: - dimension_candidates.append( - ( - "Currency", - "CurrencyCode", - "Currency Code", - "Curr", - "CurrCode", - "MoneyCurrency", - "PaymentCurrency", - "FXCurrency", - ) - ) - if "country" in normalized_query or "countries" in normalized_query: - dimension_candidates.append(("Country", "CountryName", "Nation", "Market")) - if "division" in normalized_query: - dimension_candidates.append(("Division",)) - if ( - ( - "category" in normalized_query - or "categories" in normalized_query - or "prodcategory" in compact_query - or "productcategory" in compact_query - ) - and "product" in normalized_query - and "product type" not in normalized_query - and "prodtype" not in compact_query - and "producttype" not in compact_query - ): - dimension_candidates.append( - ( - "ProductCategory", - "Product Category", - "ProdCategory", - "Category", - "ProductType", - "Product Type", - "ProdType", - "ProdName", - "Product", - "ProductName", - ) - ) - elif ( - "product type" in normalized_query - or "prodtype" in normalized_query - or "producttype" in compact_query - or "prodtype" in compact_query - ): - dimension_candidates.append(("ProdType", "ProductType", "Product Type")) - elif "product" in normalized_query: - dimension_candidates.append( - ( - "ProdName", - "Product", - "ProductName", - "ProductDescription", - "Item", - "ItemName", - "ProdCode", - "ProductCode", - "PartNo", - "SKU", - ) - ) - if ( - "customer" in normalized_query - or "custname" in compact_query - or "custno" in compact_query - ): - dimension_candidates.append( - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - "Client", - "ClientName", - ) - ) - - measure_candidates = ( - "Qty", - "Quantity", - "QtySold", - "SoldQty", - "QuantitySold", - "UnitsSold", - "ItemQty", - "SalesQty", - "OrderQty", - "OrderQuantity", - "InvoiceQty", - "InvoiceQuantity", - ) if any(term in normalized_query for term in ("quantity", "qty")) else ( - "Sales", - "SalesValue", - "FXSalesValue", - "Revenue", - "NetSales", - "TotalSales", - "SalesAmount", - "SaleAmount", - "NewOrderValue", - "NewOrdersValue", - "InvoiceValue", - "InvoiceAmount", - "InvoiceAmt", - "OrderValue", - "TotalRevenue", - "Amount", - "Value", - "TotalOrderValue", - "Cost", - ) - if "invoice" in normalized_query: - measure_candidates = ( - "InvoiceValue", - "InvoiceAmount", - "InvoiceAmt", - "InvValue", - "InvAmount", - "SalesValue", - "FXSalesValue", - "Value", - "Amount", - ) - if wants_count_metric: - measure_candidates = () - wants_trend = ( - "trend" in normalized_query - or "line chart" in normalized_query - or "over time" in normalized_query - or "last 12 months" in normalized_query - or "by month" in normalized_query - or "monthly" in normalized_query - ) - wants_date_distribution = ( - any( - term in normalized_query - for term in ("distribution", "breakdown", "split") - ) - and any( - term in normalized_query - for term in ("date", "dates", "orddate", "order date", "order dates") - ) - ) - wants_order_count_metric = ( - any(term in normalized_query for term in ("order", "orders", "new order", "new orders")) - and any( - term in normalized_query - for term in ( - "count", - "counts", - "volume", - "how many", - "number of", - "monthly", - "over time", - "last 12 months", - "each", - "per ", - ) - ) - and not any( - term in normalized_query - for term in ("value", "amount", "revenue", "sales", "cost", "margin") - ) - ) - wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) - wants_top = wants_top or any( - term in normalized_query - for term in ( - "top ", - "best ", - "ranking", - "performance ranking", - ) - ) - wants_time_bucket = wants_trend or bool( - re.search(r"\bby\s+(?:month|year|quarter|date)\b", normalized_query) - ) - wants_detail_rows = ( - wants_top - and ("new order" in normalized_query or "orders" in normalized_query) - and any(term in normalized_query for term in ("including", "include")) - ) - mentions_date_column = any( - column_name in compact_query - for column_name in ( - "orddate", - "invdate", - "orderdate", - "invoicedate", - "createdat", - ) - ) - wants_date = ( - wants_trend - or wants_date_distribution - or mentions_date_column - or "this year" in normalized_query - or bool(re.search(r"\b20\d{2}\b", normalized_query)) - ) - wants_unique_customers_by_group = ( - any( - term in normalized_query - for term in ("unique customer", "unique customers") - ) - and "customer" in normalized_query - and "division" in normalized_query - and "market" in normalized_query - and any(term in normalized_query for term in ("highest", "top", "most")) - and any(term in normalized_query for term in ("each", "per ")) - ) - if wants_unique_customers_by_group: - selected = self._select_best_analytics_table( - tables, - [ - ("Market", "MarketType", "MarketName", "Region", "Country"), - ("Division",), - ( - "Customer", - "CustomerName", - "CustName", - "CustNo", - "CustomerNo", - "CustomerCode", - "Account", - "AccountName", - "Client", - "ClientName", - ), - ], - (), - wants_date=False, - allow_count_metric=True, - query=query, - ) - if selected: - table, dimensions, _measure, _date_column = selected - table_name = table.get("name") - if table_name and len(dimensions) >= 3: - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - market, division, customer = dimensions[:3] - market_ref = f"{table_ref}.{self._quote_sql_identifier(market)}" - division_ref = ( - f"{table_ref}.{self._quote_sql_identifier(division)}" - ) - customer_ref = ( - f"{table_ref}.{self._quote_sql_identifier(customer)}" - ) - where_clause = self._append_not_null_filters( - "", - [market_ref, division_ref, customer_ref], - ) - return ( - "WITH grouped_results AS (" - f"SELECT {market_ref} AS {self._quote_sql_identifier(market)}, " - f"{division_ref} AS {self._quote_sql_identifier(division)}, " - f"COUNT(DISTINCT {customer_ref}) AS \"UniqueCustomerCount\" " - f"FROM {table_ref}" - f"{where_clause} " - f"GROUP BY {market_ref}, {division_ref}" - "), ranked_results AS (" - f"SELECT {self._quote_sql_identifier(market)}, " - f"{self._quote_sql_identifier(division)}, " - "\"UniqueCustomerCount\", " - f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(market)} " - "ORDER BY \"UniqueCustomerCount\" DESC) AS \"rank\" " - "FROM grouped_results" - ") " - f"SELECT {self._quote_sql_identifier(market)}, " - f"{self._quote_sql_identifier(division)}, " - "\"UniqueCustomerCount\" " - "FROM ranked_results " - "WHERE \"rank\" = 1 " - "ORDER BY \"UniqueCustomerCount\" DESC" - ) - - if not dimension_candidates and wants_time_bucket: - selected = self._select_best_analytics_table( - tables, - [], - measure_candidates, - wants_date=True, - allow_count_metric=wants_count_metric, - query=query, - ) - if not selected: - return None - - table, _dimensions, measure, date_column = selected - table_name = table.get("name") - if not (table_name and date_column): - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - if wants_count_metric or not measure: - metric_expr = "COUNT(*)" - metric_alias = "RecordCount" - elif wants_average_metric: - metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Average{measure}" - else: - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Total{measure}" - - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " - f"FROM {table_ref}" - f"{self._build_date_filter(table_name, date_column, query)} " - f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" - ) - - if not dimension_candidates: - return None - - allow_count_metric = ( - wants_order_count_metric - or wants_count_metric - or ("performance" in normalized_query and wants_time_bucket) - ) - selected = self._select_best_analytics_table( - tables, - dimension_candidates, - measure_candidates, - wants_date=wants_date, - allow_count_metric=allow_count_metric, - query=query, - ) - if not selected: - return None - - table, dimensions, measure, date_column = selected - table_name = table.get("name") - if not table_name: - return None - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - dimension_refs = [ - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - for dimension in dimensions - ] - - if wants_detail_rows: - if not measure: - return None - metric_ref = f"{table_ref}.{self._quote_sql_identifier(measure)}" - limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) - limit = int(limit_match.group(1)) if limit_match else 20 - select_parts = [ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ] - select_parts.append( - f"{metric_ref} AS {self._quote_sql_identifier(measure)}" - ) - date_filter = ( - self._build_date_filter(table_name, date_column, query) - if date_column - else "" - ) - return ( - f"SELECT TOP {limit} {', '.join(select_parts)} " - f"FROM {table_ref}" - f"{self._append_not_null_filters(date_filter, dimension_refs)} " - f"ORDER BY {metric_ref} DESC" - ) - - if wants_date_distribution and date_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - select_parts = [ - f"DATEPART(YEAR, {date_ref}) AS \"year\"", - f"DATEPART(MONTH, {date_ref}) AS \"month\"", - *[ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ], - 'COUNT(*) AS "OrderCount"', - ] - group_parts = [ - f"DATEPART(YEAR, {date_ref})", - f"DATEPART(MONTH, {date_ref})", - *dimension_refs, - ] - where_clause = self._append_not_null_filters( - self._build_date_filter(table_name, date_column, query), - [date_ref, *dimension_refs], - ) - extra_conditions = self._build_schema_literal_filter_conditions( - query, - table, - table_ref, - ) - if extra_conditions: - where_clause = ( - f"{where_clause.rstrip()} AND {' AND '.join(extra_conditions)}" - if where_clause.strip() - else f" WHERE {' AND '.join(extra_conditions)}" - ) - return ( - f"SELECT {', '.join(select_parts)} FROM {table_ref}" - f"{where_clause} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}), COUNT(*) DESC" - ) - - wants_top_per_group = ( - len(dimensions) >= 2 - and any(term in normalized_query for term in ("highest", "top", "most")) - and any(term in normalized_query for term in ("each", "per ")) - ) - if wants_top_per_group: - partition_dimension = None - rank_dimension = None - if "market" in normalized_query: - partition_dimension = self._find_schema_column( - table, ("Market", "MarketType", "Region") - ) - if "region" in normalized_query and not partition_dimension: - partition_dimension = self._find_schema_column( - table, ("Region", "Market", "Area", "Territory") - ) - if "customer" in normalized_query: - rank_dimension = self._find_schema_column( - table, ("Customer", "CustName", "CustNo") - ) - if not partition_dimension: - partition_dimension = dimensions[0] - if not rank_dimension: - rank_dimension = next( - ( - dimension - for dimension in dimensions - if dimension != partition_dimension - ), - None, - ) - - if partition_dimension and rank_dimension: - partition_ref = ( - f"{table_ref}.{self._quote_sql_identifier(partition_dimension)}" - ) - rank_ref = f"{table_ref}.{self._quote_sql_identifier(rank_dimension)}" - if wants_order_count_metric: - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), - ) - metric_expr = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" - if order_column - else "COUNT(*)" - ) - metric_alias = "OrderCount" - else: - metric_expr = ( - f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - if measure - else "COUNT(*)" - ) - metric_alias = f"Total{measure}" if measure else "RecordCount" - where_clause = self._append_not_null_filters( - ( - self._build_date_filter(table_name, date_column, query) - if date_column - else "" - ), - [partition_ref, rank_ref], - ) - return ( - "WITH grouped_results AS (" - f"SELECT {partition_ref} AS {self._quote_sql_identifier(partition_dimension)}, " - f"{rank_ref} AS {self._quote_sql_identifier(rank_dimension)}, " - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)} " - f"FROM {table_ref}" - f"{where_clause} " - f"GROUP BY {partition_ref}, {rank_ref}" - "), ranked_results AS (" - f"SELECT {self._quote_sql_identifier(partition_dimension)}, " - f"{self._quote_sql_identifier(rank_dimension)}, " - f"{self._quote_sql_identifier(metric_alias)}, " - f"ROW_NUMBER() OVER (PARTITION BY {self._quote_sql_identifier(partition_dimension)} " - f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC) AS \"rank\" " - "FROM grouped_results" - ") " - f"SELECT {self._quote_sql_identifier(partition_dimension)}, " - f"{self._quote_sql_identifier(rank_dimension)}, " - f"{self._quote_sql_identifier(metric_alias)} " - "FROM ranked_results " - "WHERE \"rank\" = 1 " - f"ORDER BY {self._quote_sql_identifier(metric_alias)} DESC" - ) - - if wants_trend and date_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - if wants_order_count_metric or wants_count_metric: - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), - ) - metric_expr = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" - if order_column - else "COUNT(*)" - ) - metric_alias = "OrderCount" - elif wants_average_metric: - metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Average{measure}" - else: - metric_expr = ( - f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - if measure - else "COUNT(*)" - ) - metric_alias = f"Total{measure}" if measure else "RecordCount" - select_parts = [ - f"DATEPART(YEAR, {date_ref}) AS \"year\"", - f"DATEPART(MONTH, {date_ref}) AS \"month\"", - *[ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ], - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}", - ] - group_parts = [ - f"DATEPART(YEAR, {date_ref})", - f"DATEPART(MONTH, {date_ref})", - *dimension_refs, - ] - return ( - f"SELECT {', '.join(select_parts)} FROM {table_ref}" - f"{self._append_not_null_filters(self._build_date_filter(table_name, date_column, query), dimension_refs)} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref})" - ) - - if wants_order_count_metric or wants_count_metric or not measure: - order_column = self._find_schema_column( - table, - ("OrdNo", "OrderNo", "OrderId", "NewOrderId", "InvoiceNo"), - ) - metric_expr = ( - f"COUNT(DISTINCT {table_ref}.{self._quote_sql_identifier(order_column)})" - if order_column and (wants_order_count_metric or wants_count_metric) - else "COUNT(*)" - ) - metric_alias = "OrderCount" if order_column else "RecordCount" - elif wants_average_metric: - metric_expr = f"AVG({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Average{measure}" - else: - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - metric_alias = f"Total{measure}" - limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) - limit = int(limit_match.group(1)) if limit_match else 10 - top_clause = f"TOP {limit} " if wants_top else "" - sort_direction = ( - "ASC" - if any( - term in normalized_query - for term in ( - "losing", - "lowest", - "least", - "bottom", - "declining", - "underperforming", - "smallest", - ) - ) - else "DESC" - ) - date_filter = ( - self._build_date_filter(table_name, date_column, query) - if date_column - else "" - ) - select_parts = [ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ] - select_parts.append( - f"{metric_expr} AS {self._quote_sql_identifier(metric_alias)}" - ) - return ( - f"SELECT {top_clause}{', '.join(select_parts)} " - f"FROM {table_ref}" - f"{self._append_not_null_filters(date_filter, dimension_refs)} " - f"GROUP BY {', '.join(dimension_refs)} " - f"ORDER BY {metric_expr} {sort_direction}" - ) - - def _build_contribution_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not any(term in normalized_query for term in ("contribution", "pie chart")): - return None - - compact_query = re.sub(r"[^a-z0-9]", "", normalized_query) - dimension_candidates: tuple[str, ...] | None = None - if ( - "product type" in normalized_query - or "prodtype" in normalized_query - or "producttype" in compact_query - or "prodtype" in compact_query - ): - dimension_candidates = ("ProdType", "ProductType", "Product Type") - elif "market" in normalized_query: - dimension_candidates = ("Market", "MarketType") - elif "division" in normalized_query: - dimension_candidates = ("Division",) - elif "customer" in normalized_query: - dimension_candidates = ("Customer", "CustName", "CustNo") - - if not dimension_candidates: - return None - - selected = self._select_best_analytics_table( - tables, - [dimension_candidates], - ( - "SalesValue", - "FXSalesValue", - "OrderValue", - "NewOrderValue", - "Revenue", - "Value", - "Amount", - ), - wants_date=False, - query=query, - ) - if not selected: - return None - - table, dimensions, measure, _date_column = selected - table_name = table.get("name") - if not (table_name and dimensions and measure): - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - dimension = dimensions[0] - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - return ( - f"SELECT {dimension_ref} AS {self._quote_sql_identifier(dimension)}, " - f"{metric_expr} AS \"Total{measure}\" " - f"FROM {table_ref} " - f"GROUP BY {dimension_ref} " - f"ORDER BY {metric_expr} DESC" - ) - - def _build_generic_categorical_count_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - if re.search( - r"\b(?:first|top|sample|preview|show|list)\b.*\b(?:rows?|records?|data)\b", - normalized_query, - ): - return None - - wants_categorical_summary = any( - term in normalized_query - for term in ( - "bar chart", - "by ", - "chart", - "count", - "distribution", - "donut chart", - "frequency", - "group by", - "grouped by", - "most often", - "often", - "pie chart", - "restored", - "status", - "type", - "category", - ) - ) - if not wants_categorical_summary: - return None - - query_key = self._normalize_schema_token(query) - query_terms = self._query_schema_terms(query) - scored: list[tuple[int, dict[str, Any], str]] = [] - low_value_column_patterns = ( - "id", - "no", - "number", - "date", - "time", - "description", - "comment", - "note", - "remark", - ) - - for table in tables: - table_name = str(table.get("name") or "") - normalized_table = self._normalize_schema_token(table_name) - normalized_short_table = self._normalize_schema_token( - re.split(r"[.$]", table_name)[-1] - ) - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - column_type = str(column.get("type") or "") - if not column_name or not self._is_text_schema_type(column_type): - continue - normalized_column = self._normalize_schema_token(column_name) - if not normalized_column: - continue - - score = 0 - if normalized_table and normalized_table in query_key: - score += 120 - if normalized_short_table and normalized_short_table in query_key: - score += 100 - if normalized_column and normalized_column in query_key: - score += 180 - for term in query_terms: - if term == normalized_column: - score += 100 - elif term in normalized_column or normalized_column in term: - score += 45 - if term == normalized_table or term == normalized_short_table: - score += 40 - elif term in normalized_table or term in normalized_short_table: - score += 20 - if "status" in normalized_query and "status" in normalized_column: - score += 90 - if "category" in normalized_query and "category" in normalized_column: - score += 80 - if "type" in normalized_query and "type" in normalized_column: - score += 70 - if ( - "destination" in normalized_query - and "destination" in normalized_column - and ( - "database" in normalized_query - or "databases" in normalized_query - ) - and ( - "name" in normalized_column - or "phys" in normalized_column - or "db" in normalized_column - ) - ): - score += 140 - if any(pattern == normalized_column for pattern in low_value_column_patterns): - score -= 100 - elif any( - normalized_column.endswith(pattern) - for pattern in low_value_column_patterns - ): - score -= 35 - - if score > 0: - scored.append((score, table, column_name)) - - if not scored: - return None - - _, table, dimension = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - table_name = table.get("name") - if not table_name: - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - top_n = self._extract_requested_top_n(query, default_value=0) - top_clause = f"TOP {top_n} " if top_n else "" - return ( - f"SELECT {top_clause}{dimension_ref} AS {self._quote_sql_identifier(dimension)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - def _build_order_invoice_conversion_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not ( - "conversion" in normalized_query - and "order" in normalized_query - and "invoice" in normalized_query - ): - return None - - scored: list[tuple[int, dict[str, Any], str, str, str | None]] = [] - for table in tables: - order_column = self._find_schema_column( - table, ("OrdNo", "OrderNo", "OrderNumber", "NewOrderNo") - ) - invoice_column = self._find_schema_column( - table, ("InvoiceNo", "InvNo", "InvoiceNumber") - ) - date_column = self._find_schema_column( - table, - ("OrdDate", "InvDate", "OrderDate", "InvoiceDate", "Date"), - temporal=True, - ) - if not (order_column and invoice_column): - continue - - score = 20 - if date_column: - score += 5 - if "sales" in str(table.get("name") or "").lower(): - score += 5 - scored.append((score, table, order_column, invoice_column, date_column)) - - if not scored: - return None - - _, table, order_column, invoice_column, date_column = sorted( - scored, key=lambda item: item[0], reverse=True - )[0] - table_name = table.get("name") - if not table_name: - return None - - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - order_ref = f"{table_ref}.{self._quote_sql_identifier(order_column)}" - invoice_ref = f"{table_ref}.{self._quote_sql_identifier(invoice_column)}" - date_column = date_column or order_column - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f"COUNT(DISTINCT {order_ref}) AS \"OrderCount\", " - f"COUNT(DISTINCT {invoice_ref}) AS \"InvoiceCount\", " - f"(COUNT(DISTINCT {invoice_ref}) * 100.0 / " - f"NULLIF(COUNT(DISTINCT {order_ref}), 0)) AS \"ConversionRate\" " - f"FROM {table_ref} " - f"WHERE {order_ref} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}), DATEPART(MONTH, {date_ref})" - ) - - def _build_yoy_sales_change_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not any(term in normalized_query for term in ("yoy", "year over year")): - return None - - required_dimensions: list[tuple[str, ...]] = [] - if "customer" in normalized_query: - required_dimensions.append(("Customer", "CustName", "CustNo")) - if "product" in normalized_query: - required_dimensions.append( - ("ProdName", "Product", "ProductName", "Item", "ProdCode") - ) - if "market" in normalized_query: - required_dimensions.append(("Market", "MarketType")) - - if not required_dimensions: - return None - - selected = self._select_best_analytics_table( - tables, - required_dimensions, - ( - "SalesValue", - "FXSalesValue", - "OrderValue", - "NewOrderValue", - "Revenue", - "Value", - "Amount", - ), - wants_date=False, - query=query, - ) - if not selected: - return None - - table, dimensions, measure, date_column = selected - table_name = table.get("name") - if not (table_name and measure): - return None - - year_column = self._find_schema_column( - table, ("YearInd", "Year", "OrderYear", "InvoiceYear"), numeric=True - ) - table_name = str(table_name) - table_ref = self._quote_sql_identifier(table_name) - if year_column: - year_expr = f"{table_ref}.{self._quote_sql_identifier(year_column)}" - elif date_column: - year_expr = ( - f"DATEPART(YEAR, " - f"{table_ref}.{self._quote_sql_identifier(date_column)})" - ) - else: - return None - - metric_expr = f"SUM({table_ref}.{self._quote_sql_identifier(measure)})" - dimension_refs = [ - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - for dimension in dimensions - ] - select_parts = [ - f"{year_expr} AS \"year\"", - *[ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ], - f"{metric_expr} AS \"Total{measure}\"", - ] - group_parts = [year_expr, *dimension_refs] - return ( - f"SELECT {', '.join(select_parts)} " - f"FROM {table_ref} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY {year_expr}, {metric_expr} DESC" - ) - - def _extract_requested_top_n(self, query: str, default_value: int = 10) -> int: - if match := re.search(r"\btop\s+(\d+)\b", query or "", flags=re.IGNORECASE): - return max(1, min(int(match.group(1)), 100)) - if match := re.search( - r"\b(?:first|limit)\s+(\d+)\b", query or "", flags=re.IGNORECASE - ): - return max(1, min(int(match.group(1)), 100)) - if match := re.search(r"\b(\d+)\s+rows?\b", query or "", flags=re.IGNORECASE): - return max(1, min(int(match.group(1)), 100)) - return default_value - - def _build_manufacturing_throughput_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_throughput = "throughput" in normalized or ( - "repair" in normalized and "volume" in normalized - ) - wants_unit_breakdown = any( - term in normalized - for term in ( - "manufacturing unit", - "manufacturing units", - "business unit", - "business units", - "different unit", - "different units", - ) - ) - - if not (wants_throughput and wants_unit_breakdown): - return None - - tables = self._parse_schema_tables(table_ddls) - table = self._find_best_schema_table_for_query(query, tables) - if table: - unit_column = self._find_schema_column( - table, - ( - "BusinessUnit", - "business_unit", - "manufacturing_unit", - "manufacturingunit", - "unit", - "unit_name", - "BU", - "division", - ), - ) - if unit_column: - table_name = str(table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - unit_ref = f"{table_ref}.{self._quote_sql_identifier(unit_column)}" - timestamp_column = self._find_temporal_column_for_query(query, table) - - if timestamp_column and any( - term in normalized for term in ("trend", "monthly", "over time") - ): - timestamp_ref = ( - f"{table_ref}.{self._quote_sql_identifier(timestamp_column)}" - ) - return ( - f"SELECT {unit_ref} AS " - f"{self._quote_sql_identifier(unit_column)}, " - f"DATEPART(YEAR, {timestamp_ref}) AS \"year\", " - f"DATEPART(MONTH, {timestamp_ref}) AS \"month\", " - f'COUNT(*) AS "throughput" ' - f"FROM {table_ref} " - f"WHERE {unit_ref} IS NOT NULL " - f"AND {timestamp_ref} IS NOT NULL " - f"GROUP BY {unit_ref}, DATEPART(YEAR, {timestamp_ref}), " - f"DATEPART(MONTH, {timestamp_ref}) " - f"ORDER BY {unit_ref} ASC, DATEPART(YEAR, {timestamp_ref}) ASC, " - f"DATEPART(MONTH, {timestamp_ref}) ASC" - ) - - return ( - f"SELECT {unit_ref} AS " - f"{self._quote_sql_identifier(unit_column)}, " - f'COUNT(*) AS "throughput" ' - f"FROM {table_ref} " - f"WHERE {unit_ref} IS NOT NULL " - f"GROUP BY {unit_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - has_debug_entries = self._schema_contains( - table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names - ) - has_business_unit = self._schema_contains( - table_ddls, r"\bBusinessUnit\b", table_names=table_names - ) - if not (has_debug_entries and has_business_unit): - return None - - timestamp_column = None - for candidate in ("DateIn", "FailedAt"): - if self._schema_contains( - table_ddls, rf"\b{candidate}\b", table_names=table_names - ): - timestamp_column = candidate - break - - if timestamp_column and any( - term in normalized for term in ("trend", "monthly", "over time") - ): - timestamp_expression = f'"dbo_DebugEntries"."{timestamp_column}"' - return ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' - f'DATEPART(YEAR, {timestamp_expression}) AS "year", ' - f'DATEPART(MONTH, {timestamp_expression}) AS "month", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' - f'AND {timestamp_expression} IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries"."BusinessUnit", ' - f'DATEPART(YEAR, {timestamp_expression}), ' - f'DATEPART(MONTH, {timestamp_expression}) ' - 'ORDER BY "unit_name" ASC, "year" ASC, "month" ASC' - ) - - return ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'WHERE "dbo_DebugEntries"."BusinessUnit" IS NOT NULL ' - 'GROUP BY "dbo_DebugEntries"."BusinessUnit" ' - 'ORDER BY "throughput" DESC' - ) - - def _build_audit_log_activity_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - return None - - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - if not ( - "audit" in normalized - and "log" in normalized - and any(term in normalized for term in ("activity", "over time", "trend")) - ): - return None - - table_name = "dbo_audit_log" - timestamp_column = "created_at" - if not self._schema_has_table_column( - table_ddls, - table_name, - timestamp_column, - table_names=table_names, - ): - return None - - dimension_column = None - condition_candidates = ( - "is_name_condition", - "name", - "action", - "entity_type", - ) - activity_candidates = ( - "action", - "entity_type", - "actor_name", - "actor_user_id", - "name", - ) - candidates = ( - condition_candidates - if "condition" in normalized - else activity_candidates - ) - for candidate in candidates: - if self._schema_has_table_column( - table_ddls, - table_name, - candidate, - table_names=table_names, - ): - dimension_column = candidate - break - - if not dimension_column: - return None - - timestamp_expression = f'"{table_name}"."{timestamp_column}"' - dimension_expression = f'"{table_name}"."{dimension_column}"' - return ( - f"SELECT DATEPART(YEAR, {timestamp_expression}) AS \"year\", " - f"DATEPART(MONTH, {timestamp_expression}) AS \"month\", " - f"{dimension_expression} AS \"{dimension_column}\", " - f'COUNT(*) AS "activity_count" ' - f'FROM "{table_name}" ' - f"WHERE {timestamp_expression} IS NOT NULL " - f"AND {dimension_expression} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {timestamp_expression}), " - f"DATEPART(MONTH, {timestamp_expression}), " - f"{dimension_expression} " - f"ORDER BY DATEPART(YEAR, {timestamp_expression}), " - f"DATEPART(MONTH, {timestamp_expression}), " - f'"activity_count" DESC' - ) - - def _build_repair_failure_count_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_failure_counts = ( - "failure" in normalized - and any( - term in normalized - for term in ( - "count", - "counts", - "category", - "code", - "grouped", - "common", - "most common", - "top", - ) - ) - and any(term in normalized for term in ("repair", "bar chart", "chart")) - ) - if not wants_failure_counts: - return None - - top_n = self._extract_requested_top_n(query) - - has_debug_fix_route = all( - ( - self._schema_has_table_column( - table_ddls, - "dbo_DebugEntries", - "DebugEntryId", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixLogs", - "DebugEntryId", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixLogs", - "FixId", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixes", - "Id", - table_names=table_names, - ), - self._schema_has_table_column( - table_ddls, - "dbo_DebugFixes", - "Description", - table_names=table_names, - ), - ) - ) - if has_debug_fix_route: - return ( - 'SELECT "dbo_DebugFixes"."Description" AS "failure_category", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_DebugEntries" ' - 'JOIN "dbo_DebugFixLogs" ' - 'ON "dbo_DebugEntries"."DebugEntryId" = "dbo_DebugFixLogs"."DebugEntryId" ' - 'JOIN "dbo_DebugFixes" ' - 'ON "dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id" ' - 'WHERE "dbo_DebugFixes"."Description" IS NOT NULL ' - 'GROUP BY "dbo_DebugFixes"."Description" ' - 'ORDER BY "repair_count" DESC ' - f"LIMIT {top_n}" - ) - - has_debug_entries = self._schema_contains( - table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names - ) - has_failure_patterns = self._schema_contains( - table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names - ) - has_failure_sys = self._schema_contains( - table_ddls, r"\bFailureSys\b", table_names=table_names - ) - has_debug_entry_id = self._schema_contains( - table_ddls, r"\bDebugEntryId\b", table_names=table_names - ) - has_pattern_id = self._schema_contains( - table_ddls, r"\bid\b", table_names=table_names - ) - has_pattern_category = self._schema_contains( - table_ddls, r"\bcategory\b", table_names=table_names - ) - has_pattern_name = self._schema_contains( - table_ddls, r"\bname\b", table_names=table_names - ) - - if ( - has_debug_entries - and has_failure_patterns - and has_failure_sys - and has_debug_entry_id - and has_pattern_id - and (has_pattern_category or has_pattern_name) - ): - dimension_column = ( - "category" - if ("category" in normalized and has_pattern_category) - else ("name" if has_pattern_name else "category") - ) - return ( - f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' - f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' - f'FROM "dbo_DebugEntries" ' - f'JOIN "dbo_failure_patterns" ' - f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' - f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - has_repair_logs = self._schema_has_table_column( - table_ddls, - "dbo_repair_logs", - "failure_code", - table_names=table_names, - ) - if has_repair_logs: - return ( - f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_repair_logs" ' - f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' - f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - return None - - def _build_repair_sla_compliance_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_sla = "sla" in normalized and any( - term in normalized - for term in ("compliance", "dashboard", "chart", "repair", "repairs") - ) - if not wants_sla: - return None - - has_repair_status = self._schema_has_table_column( - table_ddls, - "dbo_repair_logs", - "status", - table_names=table_names, - ) - if has_repair_status: - return ( - 'SELECT "dbo_repair_logs"."status" AS "sla_status", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."status" IS NOT NULL ' - 'GROUP BY "dbo_repair_logs"."status" ' - 'ORDER BY "repair_count" DESC' - ) - - return None - - def _build_monthly_repair_volume_sql( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - wants_monthly_repairs = ( - "repair" in normalized - and any( - term in normalized - for term in ("monthly", "last 12 months", "trend", "volume") - ) - ) - if not wants_monthly_repairs: - return None - - tables = self._parse_schema_tables(table_ddls) - scored_tables: list[tuple[int, dict[str, Any], str]] = [] - for table in tables: - table_name = str(table.get("name") or "") - if table_names and table_name not in table_names: - continue - - date_column = self._find_schema_column( - table, - ( - "created_at", - "createdAt", - "created", - "DateIn", - "Date", - "repair_date", - "RepairDate", - "opened_at", - "started_at", - ), - temporal=True, - ) - if not date_column: - date_column = self._find_any_temporal_schema_column(table) - if not date_column: - continue - - normalized_table = self._normalize_schema_token(table_name) - score = 0 - if "repair" in normalized_table: - score += 30 - if "debugentries" in normalized_table or "debugentry" in normalized_table: - score += 25 - if "log" in normalized_table: - score += 10 - if self._find_schema_column( - table, - ("DebugEntryId", "RepairId", "repair_id", "id"), - ): - score += 5 - scored_tables.append((score, table, date_column)) - - if not scored_tables: - return None - - _score, table, date_column = sorted( - scored_tables, - key=lambda item: item[0], - reverse=True, - )[0] - table_name = str(table.get("name") or "") - if not table_name: - return None - - table_ref = self._quote_sql_identifier(table_name) - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - return ( - f"SELECT DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "repair_count" ' - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"GROUP BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " - f"DATEPART(MONTH, {date_ref}) ASC" - ) - - def _is_direct_heuristic_sql_query(self, query: str) -> bool: - return False - - def _build_heuristic_text_to_sql_fallback( - self, - query: str, - table_ddls: list[str], - table_names: Optional[list[str]] = None, - ) -> str | None: - return None - - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - if schema_grounded_sql := self._build_schema_grounded_sales_sql( - query, table_ddls - ): - return schema_grounded_sql - - if throughput_sql := self._build_manufacturing_throughput_sql( - query, table_ddls, table_names=table_names - ): - return throughput_sql - - if repair_failure_count_sql := self._build_repair_failure_count_sql( - query, table_ddls, table_names=table_names - ): - return repair_failure_count_sql - - if repair_sla_sql := self._build_repair_sla_compliance_sql( - query, table_ddls, table_names=table_names - ): - return repair_sla_sql - - if monthly_repair_volume_sql := self._build_monthly_repair_volume_sql( - query, table_ddls, table_names=table_names - ): - return monthly_repair_volume_sql - - wants_chart = any( - term in normalized for term in ("chart", "bar chart", "line chart", "graph") - ) - wants_failure_counts = any( - term in normalized - for term in ( - "failure", - "failure category", - "failure code", - "common pcb failures", - "common failures", - "most common", - "top 10", - "top ten", - ) - ) - wants_monthly_repairs = ( - "repair" in normalized - and any( - term in normalized - for term in ("monthly", "last 12 months", "trend", "volume") - ) - ) - - if wants_failure_counts and wants_chart: - top_n = self._extract_requested_top_n(query) - has_pattern_failure_sys = self._schema_contains( - table_ddls, r"\bFailuresys\b", table_names=table_names - ) - has_pattern_occurrences = self._schema_contains( - table_ddls, r"\boccurrences\b", table_names=table_names - ) - has_debug_entries = self._schema_contains( - table_ddls, r"\bdbo_DebugEntries\b", table_names=table_names - ) - has_failure_patterns = self._schema_contains( - table_ddls, r"\bdbo_failure_patterns\b", table_names=table_names - ) - has_failure_sys = self._schema_contains( - table_ddls, r"\bFailureSys\b", table_names=table_names - ) - has_debug_entry_id = self._schema_contains( - table_ddls, r"\bDebugEntryId\b", table_names=table_names - ) - has_pattern_id = self._schema_contains( - table_ddls, r"\bid\b", table_names=table_names - ) - has_pattern_category = self._schema_contains( - table_ddls, r"\bcategory\b", table_names=table_names - ) - has_pattern_name = self._schema_contains( - table_ddls, r"\bname\b", table_names=table_names - ) - - if has_failure_patterns and has_pattern_failure_sys and has_pattern_occurrences: - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'"dbo_failure_patterns"."occurrences" AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'AND "dbo_failure_patterns"."occurrences" IS NOT NULL ' - f'ORDER BY "dbo_failure_patterns"."occurrences" DESC ' - f'LIMIT {top_n}' - ) - - if has_failure_patterns and has_pattern_failure_sys: - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - if ( - has_debug_entries - and has_failure_patterns - and has_failure_sys - and has_debug_entry_id - and has_pattern_id - ): - dimension_column = ( - "category" - if ("category" in normalized and has_pattern_category) - else ("name" if has_pattern_name else "category") - ) - if dimension_column == "category" and not has_pattern_category: - dimension_column = "name" - - return ( - f'SELECT "dbo_failure_patterns"."{dimension_column}" AS "failure_category", ' - f'COUNT("dbo_DebugEntries"."DebugEntryId") AS "repair_count" ' - f'FROM "dbo_DebugEntries" ' - f'JOIN "dbo_failure_patterns" ' - f'ON "dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id" ' - f'WHERE "dbo_failure_patterns"."{dimension_column}" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."{dimension_column}" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - has_repair_logs = self._schema_contains( - table_ddls, r"\bdbo_repair_logs\b", table_names=table_names - ) - has_failure_code = self._schema_contains( - table_ddls, r"\bfailure_code\b", table_names=table_names - ) - if has_repair_logs and has_failure_code: - return ( - f'SELECT "dbo_repair_logs"."failure_code" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_repair_logs" ' - f'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' - f'GROUP BY "dbo_repair_logs"."failure_code" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - if wants_failure_counts and wants_chart: - top_n = self._extract_requested_top_n(query) - return ( - f'SELECT "dbo_failure_patterns"."Failuresys" AS "failure_category", ' - f'COUNT(*) AS "repair_count" ' - f'FROM "dbo_failure_patterns" ' - f'WHERE "dbo_failure_patterns"."Failuresys" IS NOT NULL ' - f'GROUP BY "dbo_failure_patterns"."Failuresys" ' - f'ORDER BY "repair_count" DESC ' - f'LIMIT {top_n}' - ) - - return None - - def _is_schema_grounded_query( - self, query: str, db_schemas: Optional[list[str]] = None - ) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - explicit_schema_terms = ( - "table", - "column", - "schema", - "dataset", - "dbo.", - "select ", - " from ", - " join ", - " where ", - " group by ", - " order by ", - ) - if any(term in normalized for term in explicit_schema_terms): - return True - - identifier_tokens = re.findall(r"[a-zA-Z_][a-zA-Z0-9_\.]*", normalized) - if any("." in token for token in identifier_tokens): - return True - - for schema in db_schemas or []: - schema_text = schema.lower() - table_matches = re.findall( - r"create\s+table\s+([a-zA-Z0-9_\.\"]+)", schema_text - ) - column_matches = re.findall(r"\n\s*\"?([a-zA-Z_][a-zA-Z0-9_]*)\"?\s+", schema_text) - candidates = { - token.strip('"') - for token in table_matches + column_matches - if token and len(token.strip('"')) > 2 - } - if any(candidate in normalized for candidate in candidates): - return True - - return False - def _build_schema_grounded_operational_sql( - self, query: str, tables: list[dict[str, Any]] - ) -> str | None: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized_query: - return None - - operational_terms = ( - "ticket", - "repair", - "failure", - "pcb", - "component", - "board", - "throughput", - "manufacturing", - "unit", - "knowledge", - "article", - "source", - "category", - "priority", - "status", - "open", - "closed", - "aging", - "workflow", - "time", - "duration", - "elapsed", - "estimated", - "volume", - "count", - ) - if not any(term in normalized_query for term in operational_terms): - return None - - scored_tables: list[tuple[int, dict[str, Any]]] = [] - for table in tables: - table_name = str(table.get("name") or "") - normalized_table = table_name.lower() - score = 0 - if any( - token in normalized_table - for token in ("ticket", "repair", "debug", "knowledge", "article") - ): - score += 10 - if "ticket" in normalized_query and "ticket" in normalized_table: - score += 8 - if "knowledge" in normalized_query and "knowledge" in normalized_table: - score += 8 - if "article" in normalized_query and "article" in normalized_table: - score += 5 - if "repair" in normalized_query and "repair" in normalized_table: - score += 5 - if "failure" in normalized_query and "failure" in normalized_table: - score += 8 - if any( - term in normalized_query - for term in ("business unit", "business units", "unit", "units") - ) and self._find_schema_column( - table, - ( - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "ManufacturingUnit", - "unit", - "BU", - "Division", - ), - ): - score += 15 - if any( - term in normalized_query - for term in ("product line", "product family", "product", "products") - ) and self._find_schema_column( - table, - ( - "Product_Family", - "ProductFamily", - "Product Family", - "ProductLine", - "Product_Line", - "Product", - "ProdType", - "Material", - ), - ): - score += 15 - if any(term in normalized_query for term in ("error", "failure")) and any( - self._find_schema_column(table, candidates) - for candidates in ( - ("failure_code", "FailureSys", "failure", "failure_type"), - ("category", "name", "description"), - ) - ): - score += 6 - if self._find_schema_column( - table, - ("created_at", "updated_at", "DateIn", "DateOut", "created", "date"), - temporal=True, - ): - score += 3 - if score: - scored_tables.append((score, table)) - - if not scored_tables: - return None - - table = sorted(scored_tables, key=lambda item: item[0], reverse=True)[0][1] - table_name = str(table.get("name") or "") - if not table_name: - return None - - table_ref = self._quote_sql_identifier(table_name) - date_column = self._find_schema_column( - table, - ( - "created_at", - "created", - "DateIn", - "RepairDate", - "updated_at", - "DateOut", - "updated", - "date", - ), - temporal=True, - ) - - dimension_candidates: list[tuple[str, ...]] = [] - if any(term in normalized_query for term in ("failure", "failures", "error")): - dimension_candidates.append( - ( - "failure_code", - "FailureSys", - "failure", - "failure_type", - "failure_category", - "category", - "name", - "description", - ) - ) - if "manufacturing" in normalized_query or "unit" in normalized_query: - dimension_candidates.append( - ( - "BusinessUnit", - "Business_Unit", - "Business Unit", - "manufacturing_unit", - "manufacturing unit", - "ManufacturingUnit", - "unit", - "BU", - "Division", - "assignee_user_id", - "created_by_user_id", - "org_id", - "status", - ) - ) - if any( - term in normalized_query - for term in ("product line", "product family", "product", "products") - ): - dimension_candidates.append( - ( - "Product_Family", - "ProductFamily", - "Product Family", - "ProductLine", - "Product_Line", - "Product", - "ProdType", - "Material", - ) - ) - if "component" in normalized_query: - dimension_candidates.append( - ("component", "component_type", "board_type", "title", "status") - ) - if "board" in normalized_query: - dimension_candidates.append(("board_type", "board", "title", "status")) - if "category" in normalized_query: - dimension_candidates.append(("category", "subcategory", "status", "priority")) - if "source" in normalized_query: - dimension_candidates.append(("source", "author", "category", "status")) - if "priority" in normalized_query: - dimension_candidates.append(("priority", "status")) - if ( - "status" in normalized_query - or "open" in normalized_query - or "closed" in normalized_query - ): - dimension_candidates.append(("status", "priority")) - if "assignee" in normalized_query: - dimension_candidates.append(("assignee_user_id", "created_by_user_id")) - if "workflow" in normalized_query: - dimension_candidates.append(("status", "priority", "assignee_user_id")) - - dimensions: list[str] = [] - for candidates in dimension_candidates: - dimension = self._find_schema_column(table, candidates) - if dimension and dimension not in dimensions: - dimensions.append(dimension) - - if not dimensions: - fallback_dimension = self._find_first_schema_column( - table, - ( - "status", - "priority", - "category", - "subcategory", - "author", - "assignee_user_id", - "created_by_user_id", - "org_id", - "title", - ), - ) - if fallback_dimension: - dimensions.append(fallback_dimension) - - wants_trend = any( - term in normalized_query - for term in ("trend", "monthly", "month", "line chart", "over time") - ) - wants_elapsed_time = ( - not wants_trend - and any( - term in normalized_query - for term in ( - "duration", - "elapsed", - "turnaround", - "estimated", - "time spent", - "time taken", - ) - ) - ) - wants_top = bool(re.search(r"\btop\s+\d+\b", normalized_query)) - limit_match = re.search(r"\btop\s+(\d+)\b", normalized_query) - limit = int(limit_match.group(1)) if limit_match else 10 - - if wants_elapsed_time: - temporal_columns = [ - str(column.get("name") or "") - for column in table.get("columns", []) - if column.get("name") - and self._is_temporal_schema_type(str(column.get("type") or "")) - ] - start_column = self._find_schema_column( - table, - ( - "created_at", - "created", - "DateIn", - "execution_date", - "opened_at", - "started_at", - "start_date", - "begin_date", - ), - temporal=True, - ) - end_column = self._find_schema_column( - table, - ( - "updated_at", - "updated", - "DateOut", - "closed_at", - "resolved_at", - "completed_at", - "finished_at", - "end_date", - ), - temporal=True, - ) - if not start_column and temporal_columns: - start_column = temporal_columns[0] - if not end_column: - for candidate in temporal_columns: - if candidate.lower() != str(start_column or "").lower(): - end_column = candidate - break - if start_column and end_column: - start_ref = f"{table_ref}.{self._quote_sql_identifier(start_column)}" - end_ref = f"{table_ref}.{self._quote_sql_identifier(end_column)}" - duration_expr = f"DATEDIFF('second', {start_ref}, {end_ref})" - if not dimensions: - fallback_dimension = self._find_first_schema_column( - table, - ( - "status", - "priority", - "assignee_user_id", - "created_by_user_id", - "org_id", - ), - ) - if fallback_dimension: - dimensions.append(fallback_dimension) - if dimensions: - dimension = dimensions[0] - dimension_ref = ( - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - ) - dimension_alias = ( - "workflow" if "workflow" in normalized_query else dimension - ) - return ( - f"SELECT {dimension_ref} AS " - f"{self._quote_sql_identifier(dimension_alias)}, " - f'SUM({duration_expr}) AS "total_time_seconds" ' - f"FROM {table_ref} " - f"WHERE {start_ref} IS NOT NULL " - f"AND {end_ref} IS NOT NULL " - f"AND {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f'ORDER BY "total_time_seconds" DESC' - ) - return ( - f'SELECT SUM({duration_expr}) AS "total_time_seconds" ' - f"FROM {table_ref} " - f"WHERE {start_ref} IS NOT NULL " - f"AND {end_ref} IS NOT NULL" - ) - - if wants_trend and date_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(date_column)}" - select_parts = [ - f"DATEPART(YEAR, {date_ref}) AS \"year\"", - f"DATEPART(MONTH, {date_ref}) AS \"month\"", - ] - group_parts = [ - f"DATEPART(YEAR, {date_ref})", - f"DATEPART(MONTH, {date_ref})", - ] - for dimension in dimensions[:2]: - dimension_ref = f"{table_ref}.{self._quote_sql_identifier(dimension)}" - select_parts.append( - f"{dimension_ref} AS {self._quote_sql_identifier(dimension)}" - ) - group_parts.append(dimension_ref) - select_parts.append('COUNT(*) AS "RecordCount"') - return ( - f"SELECT {', '.join(select_parts)} " - f"FROM {table_ref} " - f"GROUP BY {', '.join(group_parts)} " - f"ORDER BY DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref})" - ) - - if dimensions: - top_clause = f"TOP {limit} " if wants_top else "" - dimension_refs = [ - f"{table_ref}.{self._quote_sql_identifier(dimension)}" - for dimension in dimensions[:2] - ] - select_parts = [ - f"{dimension_ref} AS {self._quote_sql_identifier(dimensions[index])}" - for index, dimension_ref in enumerate(dimension_refs) - ] - select_parts.append('COUNT(*) AS "RecordCount"') - return ( - f"SELECT {top_clause}{', '.join(select_parts)} " - f"FROM {table_ref} " - f"GROUP BY {', '.join(dimension_refs)} " - f"ORDER BY COUNT(*) DESC" - ) - - return f'SELECT COUNT(*) AS "RecordCount" FROM {table_ref}' - - def _build_pcb_direct_question_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return None - - tables = self._parse_schema_tables(table_ddls) - if not tables: - return None - - repair_table = next( - ( - table - for table in tables - if str(table.get("name") or "").lower() == "dbo_repair_logs" - ), - None, - ) - ticket_label_table = next( - ( - table - for table in tables - if str(table.get("name") or "").lower() == "dbo_ticket_labels" - ), - None, - ) - limit = self._extract_requested_top_n(query, default_value=10) - - if ticket_label_table and "ticket" in normalized and "label" in normalized: - label_column = self._find_first_schema_column( - ticket_label_table, - ("name", "label", "title", "value", "id"), - ) - if label_column: - table_name = str(ticket_label_table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - label_ref = f"{table_ref}.{self._quote_sql_identifier(label_column)}" - return ( - f"SELECT TOP {limit} {label_ref} AS " - f"{self._quote_sql_identifier(label_column)}, " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {label_ref} IS NOT NULL " - f"GROUP BY {label_ref} " - f"ORDER BY COUNT(*) DESC" - ) - - if not repair_table: - return None - - table_name = str(repair_table.get("name") or "") - table_ref = self._quote_sql_identifier(table_name) - board_model_column = self._find_schema_column( - repair_table, ("board_model", "boardModel", "board model", "product") - ) - failure_code_column = self._find_schema_column( - repair_table, ("failure_code", "failureCode", "failure code", "failure") - ) - created_at_column = self._find_schema_column( - repair_table, - ("created_at", "createdAt", "created", "date_received", "dateReceived"), - temporal=True, - ) - priority_column = self._find_schema_column(repair_table, ("priority",)) - status_column = self._find_schema_column(repair_table, ("status",)) - id_column = self._find_schema_column(repair_table, ("id", "repair_id")) - - asks_board_model_distribution = ( - "board model" in normalized - and any(term in normalized for term in ("distribution", "over time", "trend")) - ) - if asks_board_model_distribution and board_model_column and created_at_column: - board_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" - date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" - return ( - f"SELECT {board_ref} AS " - f"{self._quote_sql_identifier(board_model_column)}, " - f"DATEPART(YEAR, {date_ref}) AS \"year\", " - f"DATEPART(MONTH, {date_ref}) AS \"month\", " - f'COUNT(*) AS "RecordCount" ' - f"FROM {table_ref} " - f"WHERE {board_ref} IS NOT NULL " - f"AND {date_ref} IS NOT NULL " - f"GROUP BY {board_ref}, DATEPART(YEAR, {date_ref}), " - f"DATEPART(MONTH, {date_ref}) " - f"ORDER BY DATEPART(YEAR, {date_ref}) ASC, " - f"DATEPART(MONTH, {date_ref}) ASC, {board_ref} ASC" - ) - - asks_recurring_failures_by_product = ( - "recurring" in normalized - and "failure" in normalized - and ("product" in normalized or "pcb" in normalized) - ) - if ( - asks_recurring_failures_by_product - and board_model_column - and failure_code_column - ): - product_ref = f"{table_ref}.{self._quote_sql_identifier(board_model_column)}" - failure_ref = f"{table_ref}.{self._quote_sql_identifier(failure_code_column)}" - return ( - f"SELECT {product_ref} AS " - f"{self._quote_sql_identifier(board_model_column)}, " - f"{failure_ref} AS " - f"{self._quote_sql_identifier(failure_code_column)}, " - f'COUNT(*) AS "failure_count" ' - f"FROM {table_ref} " - f"WHERE {product_ref} IS NOT NULL " - f"AND {failure_ref} IS NOT NULL " - f"GROUP BY {product_ref}, {failure_ref} " - f'ORDER BY "failure_count" DESC' - ) - - asks_highest_priority_repairs = ( - "repair" in normalized - and "priority" in normalized - and any(term in normalized for term in ("highest", "top", "high priority")) - ) - if asks_highest_priority_repairs and priority_column: - priority_ref = f"{table_ref}.{self._quote_sql_identifier(priority_column)}" - select_refs = [] - for column in ( - id_column, - board_model_column, - failure_code_column, - status_column, - priority_column, - created_at_column, - ): - if column and column not in select_refs: - select_refs.append(column) - select_sql = ", ".join( - f"{table_ref}.{self._quote_sql_identifier(column)} AS " - f"{self._quote_sql_identifier(column)}" - for column in select_refs - ) - return ( - f"SELECT TOP {limit} {select_sql} " - f"FROM {table_ref} " - f"WHERE {priority_ref} IS NOT NULL " - f"ORDER BY CASE LOWER({priority_ref}) " - f"WHEN 'critical' THEN 1 " - f"WHEN 'high' THEN 2 " - f"WHEN 'medium' THEN 3 " - f"WHEN 'low' THEN 4 " - f"ELSE 5 END" + continue + table_match = re.search( + r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' + r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" + r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", + ddl, + flags=re.IGNORECASE, ) + if not table_match: + continue - asks_repair_ticket_distribution = ( - "repair" in normalized - and "ticket" in normalized - and ( - "distribution" in normalized - or "again distribution" in normalized - or "aging distribution" in normalized + table_name = next( + (value for value in table_match.groupdict().values() if value), + None, ) - ) - if asks_repair_ticket_distribution: - if "aging" in normalized and created_at_column: - date_ref = f"{table_ref}.{self._quote_sql_identifier(created_at_column)}" - age_bucket = ( - f"CASE " - f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 7 THEN '0-7 days' " - f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 30 THEN '8-30 days' " - f"WHEN DATEDIFF('day', {date_ref}, CURRENT_TIMESTAMP) <= 90 THEN '31-90 days' " - f"ELSE '90+ days' END" - ) - return ( - f'SELECT {age_bucket} AS "age_bucket", ' - f'COUNT(*) AS "ticket_count" ' - f"FROM {table_ref} " - f"WHERE {date_ref} IS NOT NULL " - f"GROUP BY {age_bucket} " - f'ORDER BY "ticket_count" DESC' - ) - distribution_column = status_column or priority_column or failure_code_column - if distribution_column: - dimension_ref = ( - f"{table_ref}.{self._quote_sql_identifier(distribution_column)}" - ) - return ( - f"SELECT {dimension_ref} AS " - f"{self._quote_sql_identifier(distribution_column)}, " - f'COUNT(*) AS "ticket_count" ' - f"FROM {table_ref} " - f"WHERE {dimension_ref} IS NOT NULL " - f"GROUP BY {dimension_ref} " - f'ORDER BY "ticket_count" DESC' - ) - - return None - - def _get_unqueryable_metric_message( - self, query: str, table_ddls: list[str] - ) -> str | None: - return None - - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_schema = re.sub( - r"\s+", - " ", - " ".join(ddl for ddl in table_ddls if isinstance(ddl, str)).lower(), - ) - schema_column_names = self._extract_schema_column_names(table_ddls) + if not table_name: + continue + body_start = table_match.end() + depth = 1 + body_end = body_start + while body_end < len(ddl) and depth > 0: + if ddl[body_end] == "(": + depth += 1 + elif ddl[body_end] == ")": + depth -= 1 + body_end += 1 - if not normalized_query: - return None + columns: list[dict[str, str]] = [] + for line in ddl[body_start : body_end - 1].splitlines(): + stripped = line.strip().rstrip(",") + if not stripped or stripped.startswith(("--", "/*")): + continue + if re.match( + r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", + stripped, + flags=re.IGNORECASE, + ): + continue - if "throughput" in normalized_query and any( - term in normalized_query for term in ("manufacturing", "unit", "units") - ): - unit_field_patterns = ( - r"\bbusiness[_ ]?unit\b", - r"\bmanufacturing[_ ]?unit\b", - r"\bunit[_ ]?name\b", - r"\bunit\b", - r"\bbu\b", - r"\bdivision\b", - ) - has_unit_field = any( - re.search(pattern, column_name) - for pattern in unit_field_patterns - for column_name in schema_column_names - ) - has_temporal_field = any( - self._is_temporal_schema_type(str(column.get("type") or "")) - for table in self._parse_schema_tables(table_ddls) - for column in table.get("columns", []) - ) - if not has_unit_field: - return ( - "The active datasource does not expose a manufacturing unit, " - "business unit, unit, BU, or division column. I cannot build " - "throughput trends across manufacturing units without a " - "queryable unit field." - ) - if "trend" in normalized_query and not has_temporal_field: - return ( - "The active datasource does not expose a queryable date or " - "timestamp column. I cannot build a throughput trend without " - "a first-class temporal field." + column_match = re.match( + r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' + r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_$]*))" + r"\s+(?P[A-Za-z0-9_(),]+)", + stripped, ) + if column_match: + column_name = next( + (value + for key, value in column_match.groupdict().items() + if key != "type" and value + ), + None, + ) + if not column_name: + continue + column_type = column_match.group("type") or "" + columns.append( + { + "name": str(column_name), + "type": str(column_type).lower(), + } + ) - if any( - term in normalized_query - for term in ( - "monthly", - "trend", - "turnaround", - "time", - "duration", - "elapsed", - "latest", - "recent", - "newest", - "last records", - ) - ): - has_temporal_field = any( - self._is_temporal_schema_type(str(column.get("type") or "")) - for table in self._parse_schema_tables(table_ddls) - for column in table.get("columns", []) - ) - if not has_temporal_field: - return ( - "The active datasource does not expose a queryable date or " - "timestamp column. I cannot build a time-based analysis " - "without a first-class temporal field." - ) + tables.append({"name": table_name, "columns": columns}) - repair_cost_terms = ( - "repair cost", - "repair_cost", - "repaircost", - "cost", - "cost impact", - "cost_impact", - ) - if any(term in normalized_query for term in repair_cost_terms): - cost_field_patterns = ( - r"\brepair[_ ]?cost\b", - r"\bcost[_ ]?impact\b", - r"\bcost[_ ]?amount\b", - r"\btotal[_ ]?cost\b", - r"\bunit[_ ]?cost\b", - r"\bcost\b", - r"\bamount\b", - ) - has_cost_field = any( - re.search(pattern, column_name) - for pattern in cost_field_patterns - for column_name in schema_column_names - ) + return tables - if not has_cost_field: - return ( - "The schema does not expose repair cost as a queryable " - "column. The MSSQL Wren/Ibis runtime cannot extract cost " - "from generic JSON/text fields such as data. Add repair " - "cost as a first-class column or calculated field, then " - "ask again." - ) + def _normalize_schema_identifier_key(self, value: str) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - first_pass_yield_terms = ( - "first pass yield", - "first-pass yield", - "first_pass_yield", - "fpy", + def _schema_identifier_alias_keys(self, value: str) -> set[str]: + raw_value = str(value or "") + base_key = self._normalize_schema_identifier_key(raw_value) + separator_normalized_key = self._normalize_schema_identifier_key( + re.sub(r"[._\-]+", " ", raw_value) ) - if not any(term in normalized_query for term in first_pass_yield_terms): - return None + part_keys = { + self._normalize_schema_identifier_key(part) + for part in re.split(r"[._\-]+", raw_value) + if part + } + return {key for key in {base_key, separator_normalized_key, *part_keys} if key} - required_field_patterns = ( - r"\bfirst[_ ]?pass[_ ]?yield\b", - r"\bfpy\b", - r"\battempt\b", - r"\battempt[_ ]?number\b", - r"\bfirst[_ ]?attempt\b", - r"\bpass[_ ]?fail\b", - r"\byield\b", - ) - has_required_field = any( - re.search(pattern, normalized_schema) - for pattern in required_field_patterns - ) + def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: + names: list[str] = [] + for quoted in re.findall(r"[`\"\[]([^`\"\]]+)[`\"\]]", query or ""): + if quoted and quoted not in names: + names.append(quoted) + for token in re.findall(r"\b[A-Za-z_][A-Za-z0-9_$]*(?:\.[A-Za-z_][A-Za-z0-9_$]*)+\b", query or ""): + if token and token not in names: + names.append(token) + return names - if has_required_field: - return None + def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: + keys: set[str] = set() + for table_name in self._extract_explicit_table_names_from_query(query or ""): + keys.update(self._schema_identifier_alias_keys(table_name)) + return keys - return ( - "The schema does not expose first-pass yield, attempt number, " - "first-attempt result, or pass/fail fields as queryable columns. " - "I cannot calculate First Pass Yield from only generic JSON/text " - "fields such as data. Add those fields as first-class columns or " - "calculated fields, then ask again." + def _explicit_table_alias_keys(self, table_names: list[str]) -> set[str]: + keys: set[str] = set() + for table_name in table_names: + keys.update(self._schema_identifier_alias_keys(table_name)) + return keys + + def _filter_retrieval_metadata_for_explicit_query( + self, + query: str, + documents: list[dict], + explicit_table_names: list[str] | None = None, + ) -> tuple[list[dict], list[str], list[str]]: + explicit_keys = ( + self._explicit_table_alias_keys(explicit_table_names or []) + if explicit_table_names + else self._explicit_table_alias_keys_from_query(query) ) + if not explicit_keys: + return self._metadata_from_documents(documents) - def _build_schema_grounded_sales_sql( - self, query: str, table_ddls: list[str] - ) -> str | None: - return None + filtered_documents: list[dict] = [] + for document in documents or []: + metadata = document.get("metadata") or {} + table_name = metadata.get("table_name") or document.get("table_name") + table_ddl = metadata.get("table_ddl") or document.get("table_ddl") + candidate_names = [table_name] + if table_ddl: + candidate_names.extend( + table.get("name") + for table in self._parse_schema_tables([table_ddl]) + if table.get("name") + ) - normalized_query = re.sub(r"\s+", " ", (query or "").strip().lower()) - normalized_schema = "\n".join( - ddl for ddl in table_ddls or [] if isinstance(ddl, str) - ).lower() - if not normalized_query or not normalized_schema: - return None + candidate_keys: set[str] = set() + for candidate_name in candidate_names: + candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) - schema_key = self._normalize_schema_token(normalized_schema) - is_sales_specific_query = any( - term in normalized_query - for term in ("sale", "sales", "salesperson", "sales person") - ) - if is_sales_specific_query and "salesvalue" not in schema_key: - return None + if candidate_keys & explicit_keys: + filtered_documents.append(document) - return self._build_schema_grounded_analytics_sql(query, table_ddls) + if not filtered_documents: + return [], [], [] + return self._metadata_from_documents(filtered_documents) - async def _run_with_timeout( + def _normalize_explicit_table_names( self, - label: str, - coroutine, - timeout_seconds: Optional[int] = None, - ): - timeout = timeout_seconds or self._pipeline_timeout_seconds - try: - return await asyncio.wait_for( - coroutine, - timeout=timeout, - ) - except TimeoutError as exc: - raise TimeoutError(f"{label} timed out after {timeout} seconds") from exc + table_names: list[str] | None, + ) -> list[str]: + normalized: list[str] = [] + for table_name in table_names or []: + candidate = str(table_name or "").strip() + if candidate and candidate not in normalized: + normalized.append(candidate) + return normalized def _should_retry_selected_schema_after_retrieval_timeout( self, retrieval_table_names: Optional[list[str]] @@ -5266,38 +913,6 @@ def _build_metadata_response( def _normalize_schema_token(self, value: str) -> str: return re.sub(r"[^a-z0-9]", "", (value or "").lower()) - def _query_schema_terms(self, query: str) -> set[str]: - normalized_query = (query or "").lower() - stop_words = { - "about", - "against", - "from", - "give", - "group", - "grouped", - "list", - "month", - "monthly", - "over", - "rows", - "show", - "table", - "tables", - "the", - "this", - "using", - "what", - "which", - "with", - "year", - } - terms = { - self._normalize_schema_token(token) - for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", normalized_query) - if len(token) > 2 and token not in stop_words - } - return {term for term in terms if term} - def _prune_sql_generation_context( self, query: str, @@ -5308,177 +923,6 @@ def _prune_sql_generation_context( max_tables: int = 8, ) -> tuple[list[dict], list[str], list[str]]: return documents, table_names, table_ddls - - if len(table_ddls) <= max_tables: - return documents, table_names, table_ddls - - parsed_tables = self._parse_schema_tables(table_ddls) - if not parsed_tables: - return documents, table_names, table_ddls[:max_tables] - - query_key = self._normalize_schema_token(query) - query_terms = self._query_schema_terms(query) - explicit_tables = { - self._normalize_schema_token(table_name) - for table_name in self._extract_explicit_table_names_from_query(query) - } - - scored: list[tuple[int, int]] = [] - for index, table in enumerate(parsed_tables): - table_name = str(table.get("name") or "") - normalized_table = self._normalize_schema_token(table_name) - normalized_short_table = self._normalize_schema_token( - re.split(r"[.$]", table_name)[-1] - ) - column_terms = { - self._normalize_schema_token(str(column.get("name") or "")) - for column in table.get("columns", []) - if column.get("name") - } - - score = 0 - if normalized_table in explicit_tables or normalized_short_table in explicit_tables: - score += 1000 - if normalized_table and normalized_table in query_key: - score += 500 - if normalized_short_table and normalized_short_table in query_key: - score += 450 - for term in query_terms: - if not term: - continue - if term == normalized_table or term == normalized_short_table: - score += 80 - elif term in normalized_table or term in normalized_short_table: - score += 40 - for column_term in column_terms: - if term == column_term: - score += 60 - elif term in column_term or column_term in term: - score += 25 - - if score > 0: - scored.append((score, index)) - - if not scored: - return documents, table_names, table_ddls[:max_tables] - - sorted_scored_indexes = [ - index for _, index in sorted(scored, key=lambda item: item[0], reverse=True) - ] - core_limit = max(1, max_tables - 2) if max_tables > 2 else 1 - selected_indexes = sorted_scored_indexes[:core_limit] - selected_indexes = self._expand_pruned_context_with_related_tables( - selected_indexes, - parsed_tables, - table_ddls, - max_tables=max_tables, - ) - for index in sorted_scored_indexes: - if len(selected_indexes) >= max_tables: - break - if index not in selected_indexes: - selected_indexes.append(index) - selected_indexes = sorted(selected_indexes) - pruned_documents = [ - documents[index] for index in selected_indexes if index < len(documents) - ] - pruned_table_names = [ - table_names[index] for index in selected_indexes if index < len(table_names) - ] - pruned_table_ddls = [ - table_ddls[index] for index in selected_indexes if index < len(table_ddls) - ] - - logger.info( - "Pruned SQL generation context from %s to %s tables for query: %s", - len(table_ddls), - len(pruned_table_ddls), - query, - ) - return pruned_documents, pruned_table_names, pruned_table_ddls - - def _expand_pruned_context_with_related_tables( - self, - selected_indexes: list[int], - parsed_tables: list[dict[str, Any]], - table_ddls: list[str], - *, - max_tables: int, - ) -> list[int]: - if len(selected_indexes) >= max_tables: - return selected_indexes[:max_tables] - - selected: list[int] = list(dict.fromkeys(selected_indexes)) - selected_set = set(selected) - - def join_key_columns(table: dict[str, Any]) -> set[str]: - keys = set() - for column in table.get("columns", []): - column_name = str(column.get("name") or "") - normalized = self._normalize_schema_token(column_name) - if not normalized: - continue - if ( - normalized == "id" - or normalized.endswith("id") - or normalized.endswith("no") - or normalized.endswith("number") - or normalized.endswith("code") - or normalized.endswith("key") - ): - keys.add(normalized) - return keys - - selected_table_names = { - self._normalize_schema_token( - str(parsed_tables[index].get("name") or "") - ) - for index in selected - if index < len(parsed_tables) - } - selected_join_keys: set[str] = set() - for index in selected: - if index < len(parsed_tables): - selected_join_keys.update(join_key_columns(parsed_tables[index])) - - candidates: list[tuple[int, int]] = [] - for index, table in enumerate(parsed_tables): - if index in selected_set: - continue - - table_name = str(table.get("name") or "") - normalized_table_name = self._normalize_schema_token(table_name) - ddl = table_ddls[index] if index < len(table_ddls) else "" - normalized_ddl = self._normalize_schema_token(ddl) - table_join_keys = join_key_columns(table) - - score = 0 - shared_keys = selected_join_keys & table_join_keys - if shared_keys: - score += 20 + 5 * len(shared_keys) - if normalized_table_name and any( - selected_table - and ( - selected_table in normalized_ddl - or normalized_table_name in selected_table - ) - for selected_table in selected_table_names - ): - score += 40 - if re.search(r"\b(?:foreign\s+key|references)\b", ddl, flags=re.IGNORECASE): - score += 15 - - if score > 0: - candidates.append((score, index)) - - for _, index in sorted(candidates, key=lambda item: item[0], reverse=True): - if len(selected) >= max_tables: - break - selected.append(index) - selected_set.add(index) - - return selected - def _is_valid_select_sql(self, sql: Optional[str]) -> bool: if not isinstance(sql, str): return False @@ -5844,150 +1288,44 @@ async def ask( explicit_table_names, ) ) - if not documents and not request_explicit_table_names: - logger.info( - "Explicit table retrieval did not return requested active-schema table; " - "loading full active schema. query_id=%s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval for explicit table", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - all_documents, _, _ = self._extract_retrieval_metadata( - retrieval_result - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - all_documents, - explicit_table_names, - ) - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - logger.info( - "Retrieved explicit tables for query_id %s: %s", - query_id, - table_names, - ) - - if ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, - ): - ask_result = self._build_validated_ask_result_from_sql( - ranked_measure_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated ranked measure SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = ranked_measure_sql - - if table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ): - ask_result = self._build_validated_ask_result_from_sql( - table_question_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table question matched deployed schema.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = table_question_sql - - if explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ): - explicit_sql, explicit_table_name = explicit_table_preview - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table preview request matched deployed schema.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = explicit_sql - - if documents and ( - deterministic_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls + if not documents and not request_explicit_table_names: + logger.info( + "Explicit table retrieval did not return requested active-schema table; " + "loading full active schema. query_id=%s", + query_id, ) - ): - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sql, - table_ddls, - user_query, + retrieval_result = await self._run_with_timeout( + "Full active schema retrieval for explicit table", + self._pipelines["db_schema_retrieval"].run( + query="", + project_id=ask_request.project_id, + histories=[], + enable_column_pruning=False, + ), + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + self._pipeline_timeout_seconds, + 20, + ), ) - if ask_result: - api_results = [ask_result] - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - intent_reasoning="Explicit table request matched deployed schema and generated SQL locally.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, + all_documents, _, _ = self._extract_retrieval_metadata( + retrieval_result + ) + documents, table_names, table_ddls = ( + self._filter_retrieval_metadata_for_explicit_query( + user_query, + all_documents, + explicit_table_names, ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = deterministic_sql + ) + _retrieval_result = retrieval_result.get( + "construct_retrieval_results", {} + ) + logger.info( + "Retrieved explicit tables for query_id %s: %s", + query_id, + table_names, + ) if not documents: error_message = ( @@ -6016,93 +1354,11 @@ async def ask( intent_reasoning = ( "Explicit table request matched deployed schema; generating SQL against retrieved schema." ) - sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - - if not explicit_table_names and self._is_direct_heuristic_sql_query(user_query): - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - logger.info( - "Retrieved tables for direct heuristic query_id %s: %s", - query_id, - table_names, - ) - - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using direct heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - if ask_result := self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ): - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=user_query, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - - if explicit_group_count_sql := self._build_explicit_group_count_sql( - user_query - ): - invalid_sql = explicit_group_count_sql - rephrased_question = user_query - logger.info( - "Deferring explicit grouped count SQL until active schema validation for query_id %s", - query_id, - ) + sql_user_query = user_query historical_question_result = [] - should_skip_pre_sql_retrieval = self._is_data_analysis_query( - user_query - ) - if should_skip_pre_sql_retrieval: - rephrased_question = user_query - intent_reasoning = ( - "Detected a deployed-data analytics question; skipping " - "intent classification and using SQL generation." - ) - sql_user_query = self._rewrite_query_for_text_to_sql(user_query) - logger.info( - "Skipping pre-SQL retrieval for analytics query_id %s: %s", - query_id, - user_query, - ) - if ( not api_results - and not should_skip_pre_sql_retrieval and self._should_reuse_historical_question_sql( user_query, histories ) @@ -6173,7 +1429,7 @@ async def ask( if valid_historical_results: api_results = valid_historical_results sql_generation_reasoning = "" - elif not api_results and not should_skip_pre_sql_retrieval: + elif not api_results: original_user_query = user_query # Run both pipeline operations concurrently try: @@ -6247,9 +1503,6 @@ async def ask( retrieved_db_schemas = intent_classification_result.get( "db_schemas" ) or [] - is_original_analytics_query = self._is_data_analysis_query( - original_user_query - ) is_schema_grounded_query = self._is_schema_grounded_query( original_user_query, retrieved_db_schemas ) or self._is_schema_grounded_query( @@ -6257,9 +1510,7 @@ async def ask( ) if intent in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"} and ( - is_original_analytics_query - or is_schema_grounded_query - or self._is_data_analysis_query(rephrased_question or "") + is_schema_grounded_query ): logger.info( "Overriding intent %s to TEXT_TO_SQL for schema/data query: %s", @@ -6268,23 +1519,10 @@ async def ask( ) intent = "TEXT_TO_SQL" - if is_original_analytics_query: - if rephrased_question and rephrased_question != user_query: - logger.info( - "Ignoring rephrased analytics query from intent classification. original=%s rephrased=%s", - original_user_query, - rephrased_question, - ) - user_query = original_user_query - rephrased_question = original_user_query - elif rephrased_question: + if rephrased_question: user_query = rephrased_question - sql_user_query = ( - self._rewrite_query_for_text_to_sql(user_query) - if self._is_data_analysis_query(user_query) - else user_query - ) + sql_user_query = user_query if intent == "MISLEADING_QUERY": general_result = await self._run_with_timeout( @@ -6403,10 +1641,7 @@ async def ask( tables=retrieval_table_names, histories=[], project_id=ask_request.project_id, - enable_column_pruning=( - enable_column_pruning - and not self._is_data_analysis_query(user_query) - ), + enable_column_pruning=enable_column_pruning, ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) @@ -6535,174 +1770,6 @@ async def ask( "Retrieved tables for query_id %s: %s", query_id, table_names ) - if not api_results and ( - ranked_measure_sql := self._build_schema_ranked_measure_sql( - user_query, - table_ddls, - ) - ): - logger.info( - "Using schema-grounded ranked measure SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - ranked_measure_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = ranked_measure_sql - error_message = "Schema-grounded ranked measure SQL was not valid for the active datasource schema." - - if not api_results and ( - table_question_sql := self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded table question SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - table_question_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = table_question_sql - error_message = "Schema-grounded table SQL was not valid for the active datasource schema." - - if not api_results and ( - explicit_table_preview := self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - ): - explicit_sql, explicit_table_name = explicit_table_preview - logger.info( - "Using explicit table preview SQL for query_id %s and table %s", - query_id, - explicit_table_name, - ) - if explicit_table_name not in table_names: - table_names.append(explicit_table_name) - ask_result = self._build_validated_ask_result_from_sql( - explicit_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = explicit_sql - error_message = "Explicit table preview SQL was not valid for the active datasource schema." - - if not api_results and ( - audit_log_activity_sql := self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ) - ): - logger.info( - "Using schema-grounded audit log activity SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - audit_log_activity_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = audit_log_activity_sql - error_message = ( - "Schema-grounded audit SQL was not valid for the active datasource schema and question intent." - ) - - if ( - not api_results - and self._is_data_analysis_query(user_query) - and ( - schema_grounded_sql := self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - ) - ): - logger.info( - "Using generic schema-grounded analytics SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - schema_grounded_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = schema_grounded_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and any( - term in user_query.lower() - for term in ( - "pcb", - "repair", - "failure", - "business unit", - "business units", - "product line", - "product family", - ) - ): - operational_sql = self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ) - if operational_sql: - logger.info( - "Using schema-grounded operational SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - operational_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = operational_sql - error_message = ( - "Schema-grounded operational SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and ( - deterministic_sales_sql := self._build_schema_grounded_sales_sql( - user_query, table_ddls - ) - ): - logger.info( - "Using schema-grounded CWSales SQL for query_id %s", - query_id, - ) - ask_result = self._build_validated_ask_result_from_sql( - deterministic_sales_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - else: - invalid_sql = deterministic_sales_sql - error_message = ( - "Schema-grounded SQL was not valid for the active datasource schema and question intent." - ) - should_retry_full_schema = ( not api_results and self._get_metadata_question_kind(user_query) @@ -6755,117 +1822,7 @@ async def ask( table_names, ) - full_schema_preview = self._build_explicit_table_preview_sql( - user_query, table_ddls - ) - full_schema_sql_candidates = ( - self._build_schema_grounded_table_question_sql( - user_query, table_ddls - ), - full_schema_preview[0] if full_schema_preview else None, - self._build_audit_log_activity_sql( - user_query, table_ddls, table_names=table_names - ), - self._build_schema_grounded_analytics_sql( - user_query, table_ddls - ), - self._build_schema_grounded_sales_sql( - user_query, table_ddls - ), - ) - for full_schema_sql in full_schema_sql_candidates: - if not full_schema_sql: - continue - ask_result = self._build_validated_ask_result_from_sql( - full_schema_sql, - table_ddls, - user_query, - ) - if ask_result: - api_results = [ask_result] - break - invalid_sql = full_schema_sql - error_message = ( - "Full-schema grounded SQL was not valid for the active datasource schema and question intent." - ) - - if not api_results and ( - unqueryable_metric_message := self._get_unqueryable_metric_message( - user_query, table_ddls - ) - ): - logger.info( - "ask pipeline - NO_RELEVANT_SQL due to unqueryable metric: %s", - user_query, - ) - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_SQL", - message=unqueryable_metric_message, - ), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_SQL" - results["metadata"]["error_message"] = unqueryable_metric_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - if not documents: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql fallback before retrieval failure for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - is_followup=True if histories else False, - ) - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( @@ -6907,17 +1864,6 @@ async def ask( _retrieval_result = completed_retrieval_result sql_generation_histories = histories - if self._is_data_analysis_query( - sql_user_query - ) and not self._needs_conversation_context(sql_user_query): - sql_generation_histories = [] - allow_sql_generation_reasoning = False - allow_sql_knowledge_retrieval = False - max_sql_correction_retries = min(max_sql_correction_retries, 1) - logger.info( - "Using fast standalone SQL generation path for query_id %s", - query_id, - ) if ( not self._is_stopped(query_id, self._ask_results) @@ -7082,7 +2028,7 @@ async def ask( ) except TimeoutError as generation_timeout: logger.warning( - "SQL generation timed out for query_id %s; trying schema-grounded fallback: %s", + "SQL generation timed out for query_id %s: %s", query_id, generation_timeout, ) @@ -7218,40 +2164,6 @@ async def ask( results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" else: - if heuristic_sql := self._build_heuristic_text_to_sql_fallback( - user_query, table_ddls, table_names=table_names - ): - logger.info( - "Using heuristic text-to-sql fallback for query_id %s: %s", - query_id, - user_query, - ) - ask_result = self._build_validated_ask_result_from_sql( - heuristic_sql, - table_ddls, - user_query, - ) - if not ask_result: - invalid_sql = heuristic_sql - error_message = "Heuristic SQL fallback was not valid for the active datasource schema." - else: - api_results = [ask_result] - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="TEXT_TO_SQL", - response=api_results, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["ask_result"] = api_results - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = ( From 2f7029d51406860d7528d5ca85e9bccf240b9afd Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 01:18:26 +0530 Subject: [PATCH 0599/1087] Fix ask service orchestration cleanup --- wren-ai-service/src/web/v1/services/ask.py | 156 ++------------------- 1 file changed, 14 insertions(+), 142 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d9d6b42729..4c711c471c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -116,39 +116,6 @@ class AskResultResponse(_AskResultResponse): class AskService: - _HISTORICAL_QUESTION_STOP_WORDS = { - "a", - "an", - "and", - "are", - "as", - "at", - "be", - "by", - "can", - "chart", - "create", - "each", - "for", - "from", - "give", - "graph", - "how", - "in", - "is", - "me", - "of", - "on", - "please", - "show", - "the", - "to", - "total", - "what", - "which", - "with", - } - def __init__( self, pipelines: Dict[str, BasicPipeline], @@ -191,64 +158,19 @@ def _is_stopped(self, query_id: str, container: dict): return False - @classmethod - def _normalize_historical_question_text(cls, question: str | None) -> str: - return " ".join(re.findall(r"[a-z0-9]+", (question or "").lower())) - - @classmethod - def _historical_question_tokens(cls, question: str | None) -> set[str]: - normalized = cls._normalize_historical_question_text(question) - return { - token - for token in normalized.split() - if len(token) > 1 and token not in cls._HISTORICAL_QUESTION_STOP_WORDS - } - - @classmethod - def _is_reusable_historical_question( - cls, query: str | None, historical_question: str | None - ) -> bool: - normalized_query = cls._normalize_historical_question_text(query) - normalized_historical_question = cls._normalize_historical_question_text( - historical_question - ) - if not normalized_query or not normalized_historical_question: - return False - return normalized_query == normalized_historical_question - - @classmethod - def _should_use_histories_for_query(cls, query: str | None) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - if not normalized: - return False - - contextual_prefixes = ( - "also ", - "and ", - "but ", - "for those ", - "for that ", - "for the same ", - "from that ", - "how about ", - "in that ", - "now ", - "same ", - "show more", - "show the same", - "then ", - "use that ", - "what about ", - "what if ", - ) - if normalized.startswith(contextual_prefixes): - return True - - contextual_patterns = ( - r"\b(previous|last|above|earlier|same|those|that|these|them|it|its|there)\b", - r"\b(add|break down|compare|filter|group|instead|only|sort|split)\b.+\b(by|to|with)\b", - ) - return any(re.search(pattern, normalized) for pattern in contextual_patterns) + async def _run_with_timeout( + self, + label: str, + awaitable, + *, + timeout_seconds: Optional[int] = None, + ): + timeout = timeout_seconds or self._pipeline_timeout_seconds + try: + return await asyncio.wait_for(awaitable, timeout=timeout) + except TimeoutError: + logger.warning("%s timed out after %s seconds", label, timeout) + raise def _is_greeting_query(self, query: str) -> bool: normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) @@ -267,13 +189,6 @@ def _is_greeting_query(self, query: str) -> bool: } return normalized in greeting_patterns - def _should_reuse_historical_question_sql( - self, - query: str, - histories: list[AskHistory] | None, - ) -> bool: - return False - def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: tables: list[dict[str, Any]] = [] for ddl in table_ddls or []: @@ -1127,13 +1042,6 @@ async def ask( histories = ask_request.histories[: self._max_histories][ ::-1 ] # reverse the order of histories - if histories and not self._should_use_histories_for_query(user_query): - logger.info( - "Ignoring thread histories for independent question. query_id=%s query=%s", - query_id, - user_query, - ) - histories = [] rephrased_question = None intent_reasoning = None sql_generation_reasoning = None @@ -1357,12 +1265,7 @@ async def ask( sql_user_query = user_query historical_question_result = [] - if ( - not api_results - and self._should_reuse_historical_question_sql( - user_query, histories - ) - ): + if not api_results: if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( status="searching", @@ -1397,18 +1300,6 @@ async def ask( valid_historical_results = [] for result in historical_question_result: - historical_question_text = result.get("question") - if not self._is_reusable_historical_question( - user_query, historical_question_text - ): - logger.info( - "Ignoring historical SQL for materially different question. query_id=%s query=%s historical_question=%s", - query_id, - user_query, - historical_question_text, - ) - continue - sql_statement = result.get("statement") if not self._is_valid_select_sql(sql_statement): logger.warning( @@ -1430,7 +1321,6 @@ async def ask( api_results = valid_historical_results sql_generation_reasoning = "" elif not api_results: - original_user_query = user_query # Run both pipeline operations concurrently try: sql_samples_task, instructions_task = await self._run_with_timeout( @@ -1500,24 +1390,6 @@ async def ask( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") - retrieved_db_schemas = intent_classification_result.get( - "db_schemas" - ) or [] - is_schema_grounded_query = self._is_schema_grounded_query( - original_user_query, retrieved_db_schemas - ) or self._is_schema_grounded_query( - rephrased_question or "", retrieved_db_schemas - ) - - if intent in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"} and ( - is_schema_grounded_query - ): - logger.info( - "Overriding intent %s to TEXT_TO_SQL for schema/data query: %s", - intent, - user_query, - ) - intent = "TEXT_TO_SQL" if rephrased_question: user_query = rephrased_question From 53b39077a70be804e9d6663c8740827933e03792 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 01:18:39 +0530 Subject: [PATCH 0600/1087] Update SQL generation pipeline --- .../pipelines/generation/sql_generation.py | 65 +++---------------- 1 file changed, 10 insertions(+), 55 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 59bea2279c..28f33d138e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -12,14 +12,12 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, - construct_valid_table_columns, - construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, - get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -30,24 +28,11 @@ sql_generation_user_prompt_template = """ -### TARGET DATA SOURCE ### -{{ data_source }} - -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -interpreting intent and generating SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -### VALID TABLE NAMES ### -Only use these exact table names from the schema. Do not invent, rename, singularize, -pluralize, or add catalog/schema prefixes unless the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -87,13 +72,6 @@ ### QUESTION ### User's Question: {{ query }} -### INTENT AND SCHEMA GROUNDING ### -Interpret the user's business terms by matching them to explicit tables, columns, -metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Do not answer with -general guidance when the question can be answered with SQL over the active metadata. -Never reuse table or column names from SQL SAMPLES unless those exact names also -appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. - {% if sql_generation_reasoning %} ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -109,7 +87,6 @@ def prompt( query: str, documents: list[str], prompt_builder: PromptBuilder, - data_source: str, sql_generation_reasoning: str | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -119,25 +96,9 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - schema_context = "\n".join(documents or []).lower() - has_pcb_context = any( - term in schema_context - for term in ( - "dbo_debugentries", - "debugentryid", - "failure_patterns", - "repair_logs", - "failedat", - "failuresys", - "workorder", - ) - ) _prompt = prompt_builder.run( query=query, - data_source=data_source, documents=documents, - has_pcb_context=has_pcb_context, - valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -148,9 +109,7 @@ def prompt( else "" ), metric_instructions=( - get_metric_instructions(sql_knowledge, data_source=data_source) - if has_metric - else "" + get_metric_instructions(sql_knowledge) if has_metric else "" ), json_field_instructions=( get_json_field_instructions(sql_knowledge) if has_json_field else "" @@ -167,13 +126,9 @@ async def generate_sql( prompt: dict, generator: Any, generator_name: str, - data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_generation_system_prompt( - sql_knowledge, - data_source=data_source, - ) + current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -183,7 +138,6 @@ async def generate_sql( async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, - documents: list[str], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -197,8 +151,6 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) @@ -220,7 +172,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_generation_system_prompt(None), - generation_kwargs=get_sql_generation_model_kwargs(llm_provider), + generation_kwargs=SQL_GENERATION_MODEL_KWARGS, ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( @@ -253,7 +205,10 @@ async def run( ): logger.info("SQL Generation pipeline is running...") - metadata = await retrieve_metadata(project_id or "", self._retriever) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], @@ -275,4 +230,4 @@ async def run( "sql_knowledge": sql_knowledge, **self._components, }, - ) + ) \ No newline at end of file From 8f154ee289fc1332e1d0d82cbb412c3117930045 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 01:40:52 +0530 Subject: [PATCH 0601/1087] Fix ask general assistance streaming state --- wren-ai-service/src/web/v1/services/ask.py | 29 ++++++---------------- 1 file changed, 7 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 4c711c471c..32647a953d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1397,8 +1397,7 @@ async def ask( sql_user_query = user_query if intent == "MISLEADING_QUERY": - general_result = await self._run_with_timeout( - "Misleading assistance", + asyncio.create_task( self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, @@ -1406,18 +1405,14 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, + query_id=query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "misleading_assistance" ) ) self._ask_results[query_id] = AskResultResponse( status="finished", - type="MISLEADING_QUERY", + type="GENERAL", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, @@ -1427,8 +1422,7 @@ async def ask( results["metadata"]["type"] = "MISLEADING_QUERY" return results elif intent == "GENERAL": - general_result = await self._run_with_timeout( - "Data assistance", + asyncio.create_task( self._pipelines["data_assistance"].run( query=user_query, histories=histories, @@ -1436,12 +1430,8 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, + query_id=query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "data_assistance" ) ) @@ -1457,17 +1447,12 @@ async def ask( results["metadata"]["type"] = "GENERAL" return results elif intent == "USER_GUIDE": - general_result = await self._run_with_timeout( - "User guide assistance", + asyncio.create_task( self._pipelines["user_guide_assistance"].run( query=user_query, language=ask_request.configurations.language, + query_id=query_id, custom_instruction=ask_request.custom_instruction, - ), - ) - self._general_streaming_results[query_id] = ( - self._extract_pipeline_reply( - general_result, "user_guide_assistance" ) ) From 699faf9cb3f594f063d26c9317c55bc1f9083b35 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 17 Jul 2026 13:04:49 +0530 Subject: [PATCH 0602/1087] Fix table validation for EXTRACT expressions --- .../src/pipelines/generation/utils/sql.py | 7 +++++++ .../pytest/pipelines/generation/test_sql_utils.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b60124120d..6f7e8178bb 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -2786,9 +2786,16 @@ def _quote_table_reference(table_reference: str) -> str: ) +def _is_extract_argument_from_keyword(sql: str, from_keyword_start: int) -> bool: + prefix = sql[:from_keyword_start] + return bool(re.search(r"\bEXTRACT\s*\([^)]*$", prefix, flags=re.IGNORECASE)) + + def extract_sql_table_references(sql: str) -> list[str]: references = [] for match in _SQL_TABLE_REFERENCE_PATTERN.finditer(sql): + if _is_extract_argument_from_keyword(sql, match.start()): + continue table_reference = match.group("table") if table_reference.startswith("("): continue diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6dc0b45205..0ee2519757 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -285,6 +285,21 @@ def test_schema_validation_ignores_null_table_metadata(): ) == [] +def test_schema_validation_does_not_treat_extract_from_expression_as_table(): + sql = """ + SELECT + "Market", + SUM(CASE WHEN EXTRACT(YEAR FROM "InvDate") = EXTRACT(YEAR FROM GETDATE()) - 1 + THEN "SalesValue" + ELSE 0 + END) AS "LastYearSales" + FROM "dbo_tblSales" + GROUP BY "Market" + """ + + assert find_invalid_table_references(sql, ["dbo_tblSales"]) == [] + + def test_schema_validation_ignores_null_column_metadata(): assert find_invalid_column_references( 'SELECT "dbo_tblSales"."Market" FROM "dbo_tblSales"', From 8a00c7d54f4c04dbe3ce52f45e53ca74eb08ea4a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 13:20:33 +0530 Subject: [PATCH 0603/1087] Keep ask flow semantic retrieval scoped --- .../src/pipelines/generation/utils/sql.py | 98 +---------- .../src/pipelines/sql_normalizer.py | 23 --- wren-ai-service/src/web/v1/services/ask.py | 155 ++++-------------- .../apollo/server/utils/mssqlSqlNormalizer.ts | 70 +------- wren-ui/src/hooks/useAskPrompt.tsx | 5 - 5 files changed, 36 insertions(+), 315 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 6f7e8178bb..6b69e452c6 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -501,30 +501,6 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: if match := timestamp_column_pattern.search(sql): return match.group(0) - table_pattern = re.compile( - r'\bFROM\s+("[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)', - re.IGNORECASE, - ) - if match := table_pattern.search(sql): - table_name = match.group(1) - raw_table_name = str(table_name or "").strip('"[]') - normalized_table_name = raw_table_name.lower() - quoted_table_name = f'"{raw_table_name}"' - if normalized_table_name == "dbo_debugentries": - return f'{quoted_table_name}."DateIn"' - if "report" in normalized_table_name: - return f'{quoted_table_name}."generated_at"' - if any( - token in normalized_table_name - for token in ("knowledge", "kb_article", "kb_articles", "article") - ): - return f'{quoted_table_name}."created_at"' - if any( - token in normalized_table_name - for token in ("repair", "ticket", "event", "log") - ): - return f'{quoted_table_name}."created_at"' - return None @@ -1428,10 +1404,6 @@ def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: normalized = _rewrite_mssql_limit_clause(normalized) normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) - normalized = _rewrite_mssql_invented_date_identifiers(normalized) - normalized = _rewrite_mssql_invented_report_fields(normalized) - normalized = _rewrite_mssql_invented_ticket_metrics(normalized) - normalized = _rewrite_mssql_invented_knowledge_article_fields(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) normalized = _rewrite_mssql_datepart_alias_references(normalized) @@ -1522,7 +1494,6 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> normalized = _rewrite_mssql_to_date_buckets(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) - normalized = _rewrite_mssql_invented_date_identifiers(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) @@ -2661,74 +2632,7 @@ def _sql_identifier_alias_candidates(identifier: str) -> set[str]: return {candidate for candidate in candidates if candidate} -_SEMANTIC_COLUMN_ALIASES: dict[str, tuple[str, ...]] = { - "source": ( - "source", - "source_ticket_id", - "source_repair_ids", - "category", - "subcategory", - "author", - "created_by_user_id", - "status", - ), - "leadsource": ("source", "source_ticket_id", "source_repair_ids", "category"), - "articletype": ("article_type", "category", "subcategory", "type", "status"), - "articleid": ("article_id", "id"), - "knowledgearticleid": ("knowledge_article_id", "id"), - "type": ("type", "category", "subcategory", "status"), - "category": ("category", "subcategory", "status", "priority"), - "subcategory": ("subcategory", "category", "status", "priority"), - "author": ("author", "created_by_user_id", "owner", "assignee_user_id"), - "createdby": ("created_by", "created_by_user_id", "author"), - "createdbyuser": ("created_by_user", "created_by_user_id", "author"), - "createdbyuserid": ("created_by_user_id", "author"), - "lastupdatedate": ( - "last_update_date", - "last_updated_at", - "updated_at", - "DateOut", - "DateIn", - "FailedAt", - "ModifiedAt", - "CreatedAt", - "created_at", - ), - "lastupdate": ( - "last_update_date", - "last_updated_at", - "updated_at", - "DateOut", - "DateIn", - "FailedAt", - "ModifiedAt", - "CreatedAt", - "created_at", - ), - "updateddate": ( - "updated_at", - "last_update_date", - "last_updated_at", - "DateOut", - "DateIn", - "FailedAt", - "ModifiedAt", - "CreatedAt", - "created_at", - ), - "invoicequantity": ("Qty", "Quantity", "InvoiceQty", "InvoiceCount"), - "otddate": ("InvDate", "OrdDate", "OrderDate", "InvoiceDate", "Date"), - "period": ("timeid", "TimeID", "TimeId", "YearInd", "Year", "Date"), - "periodid": ("timeid", "TimeID", "TimeId"), - "timeid": ("timeid", "TimeID", "TimeId"), - "customer": ("account", "Customer", "CustName", "CustNo", "customerpo"), - "customers": ("account", "Customer", "CustName", "CustNo", "customerpo"), - "customername": ("account", "Customer", "CustName", "CustNo", "customerpo"), - "customerid": ("account", "Customer", "CustNo", "customerpo"), - "customeraccount": ("account", "Customer", "CustName", "CustNo"), - "customerregion": ("Country", "Market", "Region", "CustomerRegion"), - "fixlogid": ("DebugEntryId", "FixId", "RepairItem", "id"), -} +_SEMANTIC_COLUMN_ALIASES: dict[str, tuple[str, ...]] = {} def _find_semantic_column_alias( diff --git a/wren-ai-service/src/pipelines/sql_normalizer.py b/wren-ai-service/src/pipelines/sql_normalizer.py index 7107446686..ac61a57741 100644 --- a/wren-ai-service/src/pipelines/sql_normalizer.py +++ b/wren-ai-service/src/pipelines/sql_normalizer.py @@ -309,21 +309,6 @@ def _infer_mssql_timestamp_expression(sql: str) -> str | None: if match := timestamp_column_pattern.search(sql): return match.group(0) - table_pattern = re.compile( - r'\bFROM\s+("[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)', - re.IGNORECASE, - ) - if match := table_pattern.search(sql): - table_name = match.group(1) - normalized_table_name = str(table_name or "").strip('"[]').lower() - if "report" in normalized_table_name: - return f'{table_name}."generated_at"' - if any( - token in normalized_table_name - for token in ("repair", "ticket", "debug", "event", "log") - ): - return f'{table_name}."created_at"' - return None @@ -495,17 +480,9 @@ def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> flags=re.IGNORECASE, ) normalized = _replace_relative_getdate_calls(normalized, now) - normalized = re.sub( - r"\bAS\s+failure\s+category\b", - 'AS "failure_category"', - normalized, - flags=re.IGNORECASE, - ) normalized = _rewrite_mssql_to_unixtime(normalized) normalized = _rewrite_mssql_timestamp_subtraction(normalized) normalized = _rewrite_mssql_timestamp_casts(normalized) - normalized = _rewrite_mssql_invented_date_identifiers(normalized) - normalized = _rewrite_mssql_invented_failure_category(normalized) normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) normalized = _rewrite_mssql_bucket_functions(normalized) normalized = _rewrite_temporal_bucket_functions(normalized) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 32647a953d..e0124ad35c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -126,8 +126,8 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, - pipeline_timeout_seconds: int = 90, - schema_retrieval_timeout_seconds: int = 180, + pipeline_timeout_seconds: int = 45, + schema_retrieval_timeout_seconds: int = 25, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -835,9 +835,22 @@ def _prune_sql_generation_context( table_names: list[str], table_ddls: list[str], *, - max_tables: int = 8, + max_tables: int = 10, ) -> tuple[list[dict], list[str], list[str]]: - return documents, table_names, table_ddls + if len(documents) <= max_tables: + return documents, table_names, table_ddls + + pruned_documents = documents[:max_tables] + pruned_table_names = table_names[:max_tables] + pruned_table_ddls = table_ddls[:max_tables] + logger.info( + "Pruned SQL generation context from %s to %s tables for query: %s", + len(documents), + len(pruned_documents), + query, + ) + return pruned_documents, pruned_table_names, pruned_table_ddls + def _is_valid_select_sql(self, sql: Optional[str]) -> bool: if not isinstance(sql, str): return False @@ -1069,7 +1082,10 @@ async def ask( use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None - understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 20) + understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 12) + planning_timeout_seconds = min(self._pipeline_timeout_seconds, 15) + generation_timeout_seconds = min(self._pipeline_timeout_seconds, 30) + correction_timeout_seconds = min(self._pipeline_timeout_seconds, 15) request_explicit_table_names = self._normalize_explicit_table_names( ask_request.explicit_tables ) @@ -1196,36 +1212,6 @@ async def ask( explicit_table_names, ) ) - if not documents and not request_explicit_table_names: - logger.info( - "Explicit table retrieval did not return requested active-schema table; " - "loading full active schema. query_id=%s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval for explicit table", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - all_documents, _, _ = self._extract_retrieval_metadata( - retrieval_result - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - all_documents, - explicit_table_names, - ) - ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) @@ -1268,8 +1254,7 @@ async def ask( if not api_results: if not self._is_stopped(query_id, self._ask_results): self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", + status="understanding", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, @@ -1585,100 +1570,10 @@ async def ask( explicit_table_names, ) ) - if ( - not documents - and self._get_metadata_question_kind(user_query) - and not request_explicit_table_names - ): - logger.info( - "Query-based schema retrieval returned no tables for data question; " - "retrying full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) - should_retry_full_schema = ( - not api_results - and self._get_metadata_question_kind(user_query) - and "db_schema_retrieval" in self._pipelines - and not request_explicit_table_names - and not table_names - ) - if should_retry_full_schema: - logger.info( - "No grounded SQL from retrieved schema; retrying with full active deployed schema for query_id %s", - query_id, - ) - retrieval_result = await self._run_with_timeout( - "Full active schema retry", - self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 30, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - full_documents, full_table_names, full_table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if explicit_table_names: - full_documents, full_table_names, full_table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - full_documents, - explicit_table_names, - ) - ) - if full_documents: - documents, table_names, table_ddls = ( - full_documents, - full_table_names, - full_table_ddls, - ) - logger.info( - "Using full active deployed schema retry for query_id %s: %s", - query_id, - table_names, - ) - if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): @@ -1753,6 +1648,7 @@ async def ask( configuration=ask_request.configurations, query_id=query_id, ), + timeout_seconds=planning_timeout_seconds, ) ).get("post_process", {}) except Exception as reasoning_error: @@ -1775,6 +1671,7 @@ async def ask( configuration=ask_request.configurations, query_id=query_id, ), + timeout_seconds=planning_timeout_seconds, ) ).get("post_process", {}) except Exception as reasoning_error: @@ -1863,6 +1760,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, ), + timeout_seconds=generation_timeout_seconds, ) else: text_to_sql_generation_results = await self._run_with_timeout( @@ -1882,6 +1780,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, ), + timeout_seconds=generation_timeout_seconds, ) except TimeoutError as generation_timeout: logger.warning( @@ -1952,6 +1851,7 @@ async def ask( error_message=error_message, language=ask_request.configurations.language, ), + timeout_seconds=correction_timeout_seconds, ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" @@ -1980,6 +1880,7 @@ async def ask( sql_knowledge=sql_knowledge, query=sql_user_query, ), + timeout_seconds=correction_timeout_seconds, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts index b2c7a1fdb4..1b004f9406 100644 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts @@ -50,32 +50,14 @@ const inferMssqlTimestampExpression = (sql: string): string => { return qualifiedTimestamp[0]; } - const fromTable = sql.match(/\bFROM\s+(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))/i); - if (fromTable) { - const tableName = fromTable[1] || fromTable[2] || fromTable[3]; - if (tableName.toLowerCase() === 'dbo_debugentries') { - return `"${tableName}"."DateIn"`; - } - if (tableName.toLowerCase() === 'dbo_reports') { - return `"${tableName}"."generated_at"`; - } - if ( - tableName.toLowerCase() === 'dbo_knowledge_articles' || - tableName.toLowerCase() === 'dbo_kb_articles' - ) { - return `"${tableName}"."created_at"`; - } - if (tableName.toLowerCase() === 'dbo_repair_logs') { - return `"${tableName}"."created_at"`; - } - return `"${tableName}"."created_at"`; - } - - return '"created_at"'; + return ''; }; const replaceInventedDateFields = (sql: string): string => { const timestampExpression = inferMssqlTimestampExpression(sql); + if (!timestampExpression) { + return sql; + } const inventedDateFields = [ 'RepairDate', 'repairDate', @@ -120,41 +102,11 @@ const rewriteMssqlDatepartFunctions = (sql: string): string => `EXTRACT(${String(part).toUpperCase()} FROM ${String(expression).trim()})`, ); -const replaceCwSalesAliases = (sql: string): string => { - if (!/\bdbo_(?:qSales1|tblSalesHistory|tblSales)\b/i.test(sql)) { - return sql; - } - - const salesTable = String.raw`(?:"dbo_(?:qSales1|tblSalesHistory|tblSales)"|\[dbo_(?:qSales1|tblSalesHistory|tblSales)\]|dbo_(?:qSales1|tblSalesHistory|tblSales))`; - const otdDate = String.raw`(?:"OTD_Date"|"OTDDate"|\[OTD_Date\]|\[OTDDate\]|OTD_Date|OTDDate)`; - const fixLogId = String.raw`(?:"FixLogId"|"FixLogID"|\[FixLogId\]|\[FixLogID\]|FixLogId|FixLogID)`; - - sql = sql.replace( - new RegExp(String.raw`(${salesTable})\s*\.\s*${otdDate}`, 'gi'), - '$1."InvDate"', - ); - sql = sql.replace( - new RegExp(String.raw`(${salesTable})\s*\.\s*${fixLogId}`, 'gi'), - '$1."InvoiceNo"', - ); - - const tableReferences = sql.match(new RegExp(salesTable, 'gi')) || []; - if (new Set(tableReferences.map((table) => table.toLowerCase())).size === 1) { - sql = sql.replace( - new RegExp(String.raw`(? { const timestampExpression = inferMssqlTimestampExpression(sql); + if (!timestampExpression) { + return sql; + } const bucketExpressions: Record = { YEAR: `EXTRACT(YEAR FROM ${timestampExpression})`, MONTH: `EXTRACT(MONTH FROM ${timestampExpression})`, @@ -524,16 +476,8 @@ export const normalizeMssqlGeneratedSqlFields = ( sql = normalizeMssqlGeneratedSqlSyntax(sql); sql = rewriteMssqlDatepartFunctions(sql); sql = replaceRelativeCurrentDateCalls(sql); - sql = replaceCwSalesAliases(sql); sql = replaceInventedDateFields(sql); - sql = replaceRepairLogThroughputShape(sql); - sql = replaceTicketCycleTurnaroundShape(sql); - sql = replacePcbThroughputFields(sql); - sql = replaceInventedFailureCategory(sql); - sql = replaceInventedReportFields(sql); - sql = replaceInventedKnowledgeArticleFields(sql); sql = replaceInventedTimeBuckets(sql); - sql = replaceBadFailurePatternJoins(sql); sql = rewriteMssqlDatepartFunctions(sql); sql = quoteMssqlDboModelReferences(sql); return normalizeMssqlGeneratedSqlSyntax(sql); diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index 9270483192..47253f4dac 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -135,11 +135,6 @@ const handleUpdateRerunAskingTaskCache = ( if (result?.thread) { const task = cloneDeep(askingTask); - // bypass understanding status to thread response - if (task.status === AskingTaskStatus.UNDERSTANDING) { - task.status = AskingTaskStatus.SEARCHING; - task.type = AskingTaskType.TEXT_TO_SQL; - } client.cache.updateQuery( { query: THREAD, From 259b7c230727c323bbd98a671cd0a7e9eb5b9a05 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 13:53:55 +0530 Subject: [PATCH 0604/1087] Route general data asks through SQL flow --- wren-ai-service/src/web/v1/services/ask.py | 25 +++++++--------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e0124ad35c..cd3c965e47 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1407,30 +1407,21 @@ async def ask( results["metadata"]["type"] = "MISLEADING_QUERY" return results elif intent == "GENERAL": - asyncio.create_task( - self._pipelines["data_assistance"].run( - query=user_query, - histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=query_id, - custom_instruction=ask_request.custom_instruction, - ) + intent_reasoning = ( + f"{intent_reasoning or ''}\n" + "Classifier returned GENERAL, but this ask flow " + "treats non-schema, non-guide questions as data " + "retrieval requests so they continue through " + "semantic retrieval and SQL generation." ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", + status="understanding", + type="TEXT_TO_SQL", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", ) - results["metadata"]["type"] = "GENERAL" - return results elif intent == "USER_GUIDE": asyncio.create_task( self._pipelines["user_guide_assistance"].run( From 27cc41a88fb1bc377f2033363d7f3b3710c4cbf3 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 15:58:04 +0530 Subject: [PATCH 0605/1087] Fix schema retrieval and SQL date normalization --- .../src/pipelines/generation/utils/sql.py | 50 +++++++++++++------ .../retrieval/db_schema_retrieval.py | 21 ++++---- wren-ai-service/src/web/v1/services/chart.py | 9 +++- 3 files changed, 55 insertions(+), 25 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 6b69e452c6..0f3c16c00c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -199,27 +199,41 @@ def replace_day_offset(match: re.Match[str]) -> str: def _replace_relative_current_date_calls(sql: str, now: datetime) -> str: + def timestamp_for_interval(amount: int, unit: str, direction: int = -1) -> str: + normalized_unit = unit.lower() + signed_amount = amount * direction + if normalized_unit.startswith("month"): + return _format_timestamp_literal(_add_months(now, signed_amount)) + if normalized_unit.startswith("year"): + return _format_timestamp_literal(_add_months(now, signed_amount * 12)) + if normalized_unit.startswith("day"): + return _format_timestamp_literal(now + timedelta(days=signed_amount)) + if normalized_unit.startswith("hour"): + return _format_timestamp_literal(now + timedelta(hours=signed_amount)) + if normalized_unit.startswith("minute"): + return _format_timestamp_literal(now + timedelta(minutes=signed_amount)) + if normalized_unit.startswith("second"): + return _format_timestamp_literal(now + timedelta(seconds=signed_amount)) + return "" + def replace_date_sub_interval(match: re.Match[str]) -> str: amount = int(match.group("amount")) unit = match.group("unit").lower() - if unit.startswith("month"): - return _format_timestamp_literal(_add_months(now, -amount)) - if unit.startswith("year"): - return _format_timestamp_literal(_add_months(now, -amount * 12)) - if unit.startswith("day"): - return _format_timestamp_literal(now - timedelta(days=amount)) - return match.group(0) + return timestamp_for_interval(amount, unit) or match.group(0) def replace_date_sub_unit_amount(match: re.Match[str]) -> str: unit = match.group("unit").lower() amount = int(match.group("amount")) - if unit.startswith("month"): - return _format_timestamp_literal(_add_months(now, -amount)) - if unit.startswith("year"): - return _format_timestamp_literal(_add_months(now, -amount * 12)) - if unit.startswith("day"): - return _format_timestamp_literal(now - timedelta(days=amount)) - return match.group(0) + return timestamp_for_interval(amount, unit) or match.group(0) + + def replace_timestamp_interval(match: re.Match[str]) -> str: + direction = -1 if match.group("operator") == "-" else 1 + return ( + timestamp_for_interval( + int(match.group("amount")), match.group("unit"), direction + ) + or match.group(0) + ) sql = re.sub( r"\bDATE_SUB\(\s*CURRENT_DATE(?:\(\))?\s*,\s*INTERVAL\s+(?P\d+)\s+(?PYEAR|MONTH|DAY)S?\s*\)", @@ -234,7 +248,13 @@ def replace_date_sub_unit_amount(match: re.Match[str]) -> str: flags=re.IGNORECASE, ) sql = re.sub( - r"\bCURRENT_DATE(?:\(\))?\b", + r"\b(?:CURRENT_DATE|CURRENT_TIMESTAMP|NOW)(?:\(\))?\s*(?P[+-])\s*INTERVAL\s+'?(?P\d+)\s+(?PYEAR|MONTH|DAY|HOUR|MINUTE|SECOND)S?'?", + replace_timestamp_interval, + sql, + flags=re.IGNORECASE, + ) + sql = re.sub( + r"\b(?:CURRENT_DATE|CURRENT_TIMESTAMP|NOW)(?:\(\))?\b", _format_timestamp_literal(now), sql, flags=re.IGNORECASE, diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e6f8ad435a..3879e3b63e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -142,7 +142,10 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistoryLike]) @observe(capture_input=False) async def table_retrieval( - embedding: dict, project_id: str, tables: list[str], table_retriever: Any + embedding: dict, + project_id: str, + tables: list[str] | None, + table_retriever: Any, ) -> dict: filters = { "operator": "AND", @@ -161,15 +164,15 @@ async def table_retrieval( query_embedding=embedding.get("embedding"), filters=filters, ) - else: - filters["conditions"].append( - {"field": "name", "operator": "in", "value": tables} - ) - return await table_retriever.run( - query_embedding=[], - filters=filters, - ) + if not tables: + return {"documents": []} + + filters["conditions"].append({"field": "name", "operator": "in", "value": tables}) + return await table_retriever.run( + query_embedding=[], + filters=filters, + ) @observe(capture_input=False) diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index efb3656512..b9152f7d83 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -9,6 +9,7 @@ from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, construct_valid_table_names, + extract_sql_table_references, find_invalid_column_references, find_invalid_table_references, normalize_sql_column_references_to_schema, @@ -93,14 +94,19 @@ def _is_stopped(self, query_id: str): return False async def _load_active_schema_contexts( - self, project_id: Optional[str] + self, project_id: Optional[str], sql: str ) -> list[str]: retrieval_pipeline = self._pipelines.get("db_schema_retrieval") if not retrieval_pipeline: return [] + table_references = extract_sql_table_references(sql) + if not table_references: + return [] + retrieval_result = await retrieval_pipeline.run( query="", + tables=table_references, histories=[], project_id=project_id, enable_column_pruning=False, @@ -155,6 +161,7 @@ async def chart( execute_sql_error_message = None schema_contexts = await self._load_active_schema_contexts( chart_request.project_id, + chart_request.sql, ) normalized_sql = self._normalize_and_validate_sql( chart_request.sql, From 76f17cc2fa4f9c1aa129b3fb5c72fa81d2fddeb9 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 17 Jul 2026 23:01:29 +0530 Subject: [PATCH 0606/1087] Scope SQL generation to retrieved schema --- .../pipelines/generation/sql_generation.py | 24 +++++++- wren-ai-service/src/web/v1/services/ask.py | 60 +++++++++++++++---- 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 28f33d138e..03cbc6c275 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -15,6 +15,8 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_valid_table_columns, + construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -33,6 +35,16 @@ {{ document }} {% endfor %} +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only use these exact table names from DATABASE SCHEMA. Do not invent, rename, +singularize, pluralize, or add catalog/schema prefixes unless the table name is +shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -77,6 +89,13 @@ {{ sql_generation_reasoning }} {% endif %} +### INTENT AND SCHEMA GROUNDING ### +Interpret the user's business terms by matching them to explicit tables, columns, +metrics, views, and relationships in DATABASE SCHEMA. Never reuse table or column +names from SQL SAMPLES unless those exact names also appear in DATABASE SCHEMA or +VALID TABLE NAMES for the active datasource. Do not invent tables, columns, joins, +metrics, or relationships. + Let's think step by step. """ @@ -99,6 +118,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, + valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -151,6 +171,8 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) @@ -230,4 +252,4 @@ async def run( "sql_knowledge": sql_knowledge, **self._components, }, - ) \ No newline at end of file + ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cd3c965e47..c19d7770a9 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -11,6 +11,7 @@ from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, construct_valid_table_names, + extract_sql_table_references, normalize_sql_direction_keywords, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, @@ -851,6 +852,40 @@ def _prune_sql_generation_context( ) return pruned_documents, pruned_table_names, pruned_table_ddls + def _filter_sql_samples_for_retrieved_schema( + self, sql_samples: list[dict], table_ddls: list[str] + ) -> list[dict]: + valid_table_names = { + table_name.lower() + for table_name in construct_valid_table_names(table_ddls) + if table_name + } + if not valid_table_names: + return [] + + filtered_samples = [] + for sample in sql_samples or []: + sql = sample.get("sql") or sample.get("statement") or "" + table_references = { + table_reference.lower() + for table_reference in extract_sql_table_references(sql) + if table_reference + } + if not table_references: + filtered_samples.append(sample) + continue + + if table_references.issubset(valid_table_names): + filtered_samples.append(sample) + + if len(filtered_samples) != len(sql_samples or []): + logger.info( + "Filtered SQL samples from %s to %s using retrieved schema context.", + len(sql_samples or []), + len(filtered_samples), + ) + return filtered_samples + def _is_valid_select_sql(self, sql: Optional[str]) -> bool: if not isinstance(sql, str): return False @@ -1083,6 +1118,9 @@ async def ask( allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 12) + optional_context_timeout_seconds = min(self._pipeline_timeout_seconds, 3) + intent_timeout_seconds = min(self._pipeline_timeout_seconds, 6) + schema_timeout_seconds = min(self._schema_retrieval_timeout_seconds, 12) planning_timeout_seconds = min(self._pipeline_timeout_seconds, 15) generation_timeout_seconds = min(self._pipeline_timeout_seconds, 30) correction_timeout_seconds = min(self._pipeline_timeout_seconds, 15) @@ -1268,7 +1306,7 @@ async def ask( query=user_query, project_id=ask_request.project_id, ), - timeout_seconds=min(understanding_timeout_seconds, 10), + timeout_seconds=optional_context_timeout_seconds, ) # we only return top 1 result @@ -1321,7 +1359,7 @@ async def ask( scope="sql", ), ), - timeout_seconds=understanding_timeout_seconds, + timeout_seconds=optional_context_timeout_seconds, ) # Extract results from completed tasks @@ -1354,7 +1392,7 @@ async def ask( project_id=ask_request.project_id, configuration=ask_request.configurations, ), - timeout_seconds=understanding_timeout_seconds, + timeout_seconds=intent_timeout_seconds, ) ).get("post_process", {}) except TimeoutError as exc: @@ -1476,7 +1514,7 @@ async def ask( project_id=ask_request.project_id, enable_column_pruning=enable_column_pruning, ), - timeout_seconds=self._schema_retrieval_timeout_seconds, + timeout_seconds=schema_timeout_seconds, ) except TimeoutError as error: if not self._should_retry_selected_schema_after_retrieval_timeout( @@ -1508,10 +1546,7 @@ async def ask( project_id=ask_request.project_id, enable_column_pruning=False, ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 30, - ), + timeout_seconds=schema_timeout_seconds, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -1543,10 +1578,7 @@ async def ask( histories=[], enable_column_pruning=enable_column_pruning, ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 20, - ), + timeout_seconds=schema_timeout_seconds, ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -1605,6 +1637,10 @@ async def ask( ) if completed_retrieval_result: _retrieval_result = completed_retrieval_result + sql_samples = self._filter_sql_samples_for_retrieved_schema( + sql_samples, + table_ddls, + ) sql_generation_histories = histories From faf230f8d2d63f1ff0fc47c072837152e30541d0 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 18 Jul 2026 00:15:08 +0530 Subject: [PATCH 0607/1087] Revert "Scope SQL generation to retrieved schema" This reverts commit 76f17cc2fa4f9c1aa129b3fb5c72fa81d2fddeb9. --- .../pipelines/generation/sql_generation.py | 24 +------- wren-ai-service/src/web/v1/services/ask.py | 60 ++++--------------- 2 files changed, 13 insertions(+), 71 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 03cbc6c275..28f33d138e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -15,8 +15,6 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, - construct_valid_table_columns, - construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -35,16 +33,6 @@ {{ document }} {% endfor %} -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only use these exact table names from DATABASE SCHEMA. Do not invent, rename, -singularize, pluralize, or add catalog/schema prefixes unless the table name is -shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -89,13 +77,6 @@ {{ sql_generation_reasoning }} {% endif %} -### INTENT AND SCHEMA GROUNDING ### -Interpret the user's business terms by matching them to explicit tables, columns, -metrics, views, and relationships in DATABASE SCHEMA. Never reuse table or column -names from SQL SAMPLES unless those exact names also appear in DATABASE SCHEMA or -VALID TABLE NAMES for the active datasource. Do not invent tables, columns, joins, -metrics, or relationships. - Let's think step by step. """ @@ -118,7 +99,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -171,8 +151,6 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) @@ -252,4 +230,4 @@ async def run( "sql_knowledge": sql_knowledge, **self._components, }, - ) + ) \ No newline at end of file diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c19d7770a9..cd3c965e47 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -11,7 +11,6 @@ from src.pipelines.generation.utils.sql import ( construct_valid_table_columns, construct_valid_table_names, - extract_sql_table_references, normalize_sql_direction_keywords, normalize_sql_column_references_to_schema, normalize_sql_table_references_to_schema, @@ -852,40 +851,6 @@ def _prune_sql_generation_context( ) return pruned_documents, pruned_table_names, pruned_table_ddls - def _filter_sql_samples_for_retrieved_schema( - self, sql_samples: list[dict], table_ddls: list[str] - ) -> list[dict]: - valid_table_names = { - table_name.lower() - for table_name in construct_valid_table_names(table_ddls) - if table_name - } - if not valid_table_names: - return [] - - filtered_samples = [] - for sample in sql_samples or []: - sql = sample.get("sql") or sample.get("statement") or "" - table_references = { - table_reference.lower() - for table_reference in extract_sql_table_references(sql) - if table_reference - } - if not table_references: - filtered_samples.append(sample) - continue - - if table_references.issubset(valid_table_names): - filtered_samples.append(sample) - - if len(filtered_samples) != len(sql_samples or []): - logger.info( - "Filtered SQL samples from %s to %s using retrieved schema context.", - len(sql_samples or []), - len(filtered_samples), - ) - return filtered_samples - def _is_valid_select_sql(self, sql: Optional[str]) -> bool: if not isinstance(sql, str): return False @@ -1118,9 +1083,6 @@ async def ask( allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 12) - optional_context_timeout_seconds = min(self._pipeline_timeout_seconds, 3) - intent_timeout_seconds = min(self._pipeline_timeout_seconds, 6) - schema_timeout_seconds = min(self._schema_retrieval_timeout_seconds, 12) planning_timeout_seconds = min(self._pipeline_timeout_seconds, 15) generation_timeout_seconds = min(self._pipeline_timeout_seconds, 30) correction_timeout_seconds = min(self._pipeline_timeout_seconds, 15) @@ -1306,7 +1268,7 @@ async def ask( query=user_query, project_id=ask_request.project_id, ), - timeout_seconds=optional_context_timeout_seconds, + timeout_seconds=min(understanding_timeout_seconds, 10), ) # we only return top 1 result @@ -1359,7 +1321,7 @@ async def ask( scope="sql", ), ), - timeout_seconds=optional_context_timeout_seconds, + timeout_seconds=understanding_timeout_seconds, ) # Extract results from completed tasks @@ -1392,7 +1354,7 @@ async def ask( project_id=ask_request.project_id, configuration=ask_request.configurations, ), - timeout_seconds=intent_timeout_seconds, + timeout_seconds=understanding_timeout_seconds, ) ).get("post_process", {}) except TimeoutError as exc: @@ -1514,7 +1476,7 @@ async def ask( project_id=ask_request.project_id, enable_column_pruning=enable_column_pruning, ), - timeout_seconds=schema_timeout_seconds, + timeout_seconds=self._schema_retrieval_timeout_seconds, ) except TimeoutError as error: if not self._should_retry_selected_schema_after_retrieval_timeout( @@ -1546,7 +1508,10 @@ async def ask( project_id=ask_request.project_id, enable_column_pruning=False, ), - timeout_seconds=schema_timeout_seconds, + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 30, + ), ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -1578,7 +1543,10 @@ async def ask( histories=[], enable_column_pruning=enable_column_pruning, ), - timeout_seconds=schema_timeout_seconds, + timeout_seconds=min( + self._schema_retrieval_timeout_seconds, + 20, + ), ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} @@ -1637,10 +1605,6 @@ async def ask( ) if completed_retrieval_result: _retrieval_result = completed_retrieval_result - sql_samples = self._filter_sql_samples_for_retrieved_schema( - sql_samples, - table_ddls, - ) sql_generation_histories = histories From 8d1a63aeb4617f50d6dfe404c2909fb19c538579 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 18 Jul 2026 02:11:57 +0530 Subject: [PATCH 0608/1087] Fix semantic schema selection for ask retrieval --- .../retrieval/db_schema_retrieval.py | 13 +- .../retrieval/test_db_schema_retrieval.py | 341 ++++-------------- 2 files changed, 87 insertions(+), 267 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 3879e3b63e..ceeedeb70d 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -240,6 +240,8 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: @observe(capture_input=False) def check_using_db_schemas_without_pruning( + query: str, + tables: list[str] | None, construct_db_schemas: list[dict], dbschema_retrieval: list[Document], encoding: tiktoken.Encoding, @@ -288,7 +290,16 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) - if _token_count > context_window_size or enable_column_pruning: + should_select_tables_for_question = ( + bool((query or "").strip()) + and not tables + and len(retrieval_results) > 1 + ) + if ( + _token_count > context_window_size + or enable_column_pruning + or should_select_tables_for_question + ): return { "db_schemas": [], "tokens": _token_count, diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7c7d46ce89..fe93aa8779 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -2,248 +2,12 @@ from haystack import Document from src.pipelines.retrieval.db_schema_retrieval import ( - _is_project_wide_analysis_query, - _rerank_table_documents, - _select_relevant_table_documents, check_using_db_schemas_without_pruning, dbschema_retrieval, - expand_business_terms_for_retrieval, table_retrieval, ) -def test_project_wide_analysis_query_includes_broad_ranking_questions(): - assert _is_project_wide_analysis_query( - "Which projects have the highest number of completed questions?" - ) - - -def test_project_wide_analysis_query_ignores_empty_query(): - assert not _is_project_wide_analysis_query("") - - -def test_expand_business_terms_for_retrieval_adds_generic_sales_order_terms(): - query = "Show top customers by invoice amount" - - expanded_query = expand_business_terms_for_retrieval(query) - - assert query in expanded_query - assert "transaction purchase billing account geography" in expanded_query - assert "money exchange currency" in expanded_query - - -def test_expand_business_terms_for_retrieval_adds_generic_currency_market_terms(): - query = "Show invoice distribution by currency across markets" - - expanded_query = expand_business_terms_for_retrieval(query) - - assert query in expanded_query - assert "money exchange currency" in expanded_query - - -def test_expand_business_terms_for_retrieval_leaves_query_unchanged(): - query = "Explain what this workspace does" - - assert expand_business_terms_for_retrieval(query) == query - - -def test_rerank_table_documents_prefers_question_relevant_table_text(): - generic_stage = Document( - content="Generic imported staging records with product labels.", - meta={"type": "TABLE_DESCRIPTION", "name": "generic_stage_load"}, - score=0.99, - ) - order_region_table = Document( - content="Business transactions grouped by customer geography and amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "business_transactions"}, - score=0.01, - ) - - documents = _rerank_table_documents( - "Show order distribution across regions.", - [generic_stage, order_region_table], - ) - - assert documents[0].meta["name"] == "business_transactions" - - -def test_rerank_table_documents_penalizes_test_sources_even_when_name_uses_underscores(): - test_load = Document( - content="Raw test load rows for order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ) - order_market_table = Document( - content="New order transaction records with market and customer fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.45, - ) - - documents = _rerank_table_documents( - "Show order distribution across markets.", - [test_load, order_market_table], - ) - - assert documents[0].meta["name"] == "dbo_xStageNewOrders" - - -def test_select_relevant_table_documents_limits_weak_extra_candidates(): - documents = [ - Document( - content="Invoice transactions with product, customer, currency, and amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "invoices"}, - score=0.92, - ), - Document( - content="Product catalog with product names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "products"}, - score=0.86, - ), - Document( - content="Customer account master data.", - meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, - score=0.82, - ), - Document( - content="Exchange rate lookup by currency.", - meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, - score=0.78, - ), - Document( - content="Sales regions and market hierarchy.", - meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, - score=0.74, - ), - Document( - content="Raw staging audit rows with load metadata.", - meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, - score=0.99, - ), - ] - - selected = _select_relevant_table_documents( - "Show invoice distribution by currency across markets", - documents, - ) - - assert 1 <= len(selected) <= 5 - assert "staging_audit" not in [document.meta["name"] for document in selected] - - -def test_select_relevant_table_documents_excludes_unrequested_test_candidate(): - documents = [ - Document( - content="Raw test load rows with order market fields.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageLoad8_Test"}, - score=0.99, - ), - Document( - content="New order transaction records with market and customer details.", - meta={"type": "TABLE_DESCRIPTION", "name": "dbo_xStageNewOrders"}, - score=0.4, - ), - ] - - selected = _select_relevant_table_documents( - "Show order distribution across markets.", - documents, - ) - - assert [document.meta["name"] for document in selected] == ["dbo_xStageNewOrders"] - - -@pytest.mark.asyncio -async def test_table_retrieval_caps_embedding_results_before_schema_loading(): - documents = [ - Document( - content="Raw staging audit rows with load metadata.", - meta={"type": "TABLE_DESCRIPTION", "name": "staging_audit"}, - score=0.99, - ), - Document( - content="Invoice sales transactions with product categories and sales value.", - meta={"type": "TABLE_DESCRIPTION", "name": "sales_invoices"}, - score=0.8, - ), - Document( - content="Product catalog with product names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "products"}, - score=0.7, - ), - Document( - content="Customer account master data.", - meta={"type": "TABLE_DESCRIPTION", "name": "customers"}, - score=0.6, - ), - Document( - content="Sales regions and market hierarchy.", - meta={"type": "TABLE_DESCRIPTION", "name": "regions"}, - score=0.5, - ), - Document( - content="Exchange rate lookup by currency.", - meta={"type": "TABLE_DESCRIPTION", "name": "currency_rates"}, - score=0.4, - ), - ] - - class Retriever: - async def run(self, query_embedding, filters): - return {"documents": documents} - - result = await table_retrieval( - query="What is the distribution of sales across product categories?", - embedding={"embedding": [0.1, 0.2]}, - project_id="project-1", - tables=[], - table_retriever=Retriever(), - ) - - selected_names = [document.meta["name"] for document in result["documents"]] - assert 1 <= len(selected_names) <= 5 - assert "staging_audit" not in selected_names - - -def test_rerank_table_documents_prefers_reference_source_for_entity_listing(): - transaction_source = Document( - content="Invoice transaction fact rows with customer id and invoice amount.", - meta={"type": "TABLE_DESCRIPTION", "name": "invoice_fact"}, - score=0.95, - ) - reference_source = Document( - content="Customer master reference directory with customer names and accounts.", - meta={"type": "TABLE_DESCRIPTION", "name": "customer_master"}, - score=0.7, - ) - - documents = _rerank_table_documents( - "List customer names without duplicates.", - [transaction_source, reference_source], - ) - - assert documents[0].meta["name"] == "customer_master" - - -def test_rerank_table_documents_prefers_transaction_source_for_metric_question(): - reference_source = Document( - content="Product catalog reference table with names and categories.", - meta={"type": "TABLE_DESCRIPTION", "name": "product_master"}, - score=0.95, - ) - transaction_source = Document( - content="Sales transaction fact table with product, amount, and revenue.", - meta={"type": "TABLE_DESCRIPTION", "name": "sales_fact"}, - score=0.7, - ) - - documents = _rerank_table_documents( - "Show total sales amount by product.", - [reference_source, transaction_source], - ) - - assert documents[0].meta["name"] == "sales_fact" - - @pytest.mark.asyncio async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: @@ -257,7 +21,6 @@ async def run(self, query_embedding, filters): retriever = Retriever() await table_retrieval( - query="show rows", embedding={}, project_id="project-1", tables=["orders"], @@ -310,7 +73,6 @@ async def run(self, query_embedding, filters): retriever = Retriever() documents = await dbschema_retrieval( - query="total orders", table_retrieval={ "documents": [ Document( @@ -328,8 +90,13 @@ async def run(self, query_embedding, filters): "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": "orders"}, + ], + }, {"field": "project_id", "operator": "==", "value": "project-1"}, - {"field": "name", "operator": "in", "value": ["orders"]}, ], } @@ -347,7 +114,6 @@ async def run(self, query_embedding, filters): retriever = Retriever() documents = await dbschema_retrieval( - query="show top customers by invoice amount", table_retrieval={"documents": []}, project_id="project-1", dbschema_retriever=retriever, @@ -363,6 +129,8 @@ def encode(self, value): return value.split() result = check_using_db_schemas_without_pruning( + query="show top customers by invoice amount", + tables=None, construct_db_schemas=[ { "type": "TABLE", @@ -391,35 +159,76 @@ def encode(self, value): assert result["tokens"] > 0 -@pytest.mark.asyncio -async def test_dbschema_retrieval_uses_explicit_tables_as_scope(): - class Retriever: - def __init__(self): - self.filters = None - - async def run(self, query_embedding, filters): - self.filters = filters - return {"documents": []} +def test_check_using_db_schemas_without_pruning_selects_tables_for_question(): + class Encoding: + def encode(self, value): + return value.split() - retriever = Retriever() + def table_schema(name): + return { + "type": "TABLE", + "name": name, + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } - await dbschema_retrieval( - query="show failed repairs", - table_retrieval={"documents": []}, - project_id="project-1", - dbschema_retriever=retriever, - tables=["dbo.failure_patterns", "dbo_failure_patterns"], + result = check_using_db_schemas_without_pruning( + query="compare recent activity by account", + tables=None, + construct_db_schemas=[ + table_schema("activity"), + table_schema("account"), + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, ) - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, + assert result["db_schemas"] == [] + assert result["tokens"] > 0 + + +def test_check_using_db_schemas_without_pruning_keeps_explicit_table_fast_path(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + query="show records from activity", + tables=["activity"], + construct_db_schemas=[ { - "field": "name", - "operator": "in", - "value": ["dbo.failure_patterns", "dbo_failure_patterns"], - }, + "type": "TABLE", + "name": "activity", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } ], - } + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + assert [schema["table_name"] for schema in result["db_schemas"]] == ["activity"] From 55219abe524e379ca5898a5df6a31f1dd5a5c7d4 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 00:22:50 +0530 Subject: [PATCH 0609/1087] Refactor SQL utilities and simplify rules --- .../src/pipelines/generation/utils/sql.py | 2691 +---------------- 1 file changed, 84 insertions(+), 2607 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 0f3c16c00c..45d3011bf4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,6 +1,5 @@ import logging import re -from datetime import datetime, timedelta from typing import Any, Dict, List import aiohttp @@ -13,1520 +12,60 @@ Engine, clean_generation_result, ) -from src.core.provider import LLMProvider from src.pipelines.retrieval.sql_knowledge import SqlKnowledge logger = logging.getLogger("wren-ai-service") -def _parse_sql_json_payload(payload: str) -> str | None: - try: - parsed_payload = orjson.loads(payload) - except orjson.JSONDecodeError: - return None - - if isinstance(parsed_payload, dict) and isinstance(parsed_payload.get("sql"), str): - return parsed_payload["sql"] - - return None - - -def _extract_json_object_with_sql(result: str) -> str | None: - sql_key_match = re.search(r'"sql"\s*:', result, flags=re.IGNORECASE) - if not sql_key_match: - return None +def _extract_ddl_columns(ddl: str) -> list[str]: + if not isinstance(ddl, str): + return [] - start = result.rfind("{", 0, sql_key_match.start()) - if start == -1: - return None + column_section = re.search(r"\((.*)\)", ddl, flags=re.DOTALL) + if not column_section: + return [] - depth = 0 - in_string = False - escape_next = False - for index, char in enumerate(result[start:], start=start): - if escape_next: - escape_next = False - continue - if char == "\\" and in_string: - escape_next = True + columns = [] + for raw_line in re.split(r",\s*(?:\n|(?=[A-Za-z_\"`\[]))", column_section.group(1)): + line = re.sub(r"/\*.*?\*/", "", raw_line).strip() + line = re.sub(r"^--.*$", "", line).strip().rstrip(",") + if not line: continue - if char == '"': - in_string = not in_string - continue - if in_string: - continue - if char == "{": - depth += 1 - elif char == "}": - depth -= 1 - if depth == 0: - return result[start : index + 1] - - return None - - -def _extract_select_statement(result: str) -> str | None: - sql_match = re.search(r"\b(?:WITH|SELECT)\b", result, flags=re.IGNORECASE) - if not sql_match: - return None - - statement = result[sql_match.start() :].strip() - semicolon_index = statement.find(";") - if semicolon_index >= 0: - statement = statement[:semicolon_index] - - return statement - - -def extract_sql_generation_result(result: str) -> str: - fenced_blocks = re.findall( - r"```(?:json|sql)?\s*(.*?)```", result, flags=re.IGNORECASE | re.DOTALL - ) - for block in fenced_blocks: - if sql := _parse_sql_json_payload(block.strip()): - return clean_generation_result(sql) - if sql := _extract_select_statement(block): - return clean_generation_result(sql) - - cleaned_result = clean_generation_result(result) - if sql := _parse_sql_json_payload(cleaned_result): - return clean_generation_result(sql) - - if json_payload := _extract_json_object_with_sql(result): - if sql := _parse_sql_json_payload(json_payload): - return clean_generation_result(sql) - - if sql := _extract_select_statement(result): - return clean_generation_result(sql) - - return cleaned_result - - -def is_select_statement(sql: str) -> bool: - return bool(re.match(r"^\s*(?:WITH|SELECT)\b", sql, flags=re.IGNORECASE)) - - -def normalize_data_source(data_source: str | None) -> str: - normalized = (data_source or "").strip().upper().replace("-", "_").replace( - " ", "_" - ) - if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: - return "MSSQL" - return normalized - - -def _format_timestamp_literal(value: datetime) -> str: - return value.strftime("'%Y-%m-%d %H:%M:%S'") - - -def _add_months(value: datetime, months: int) -> datetime: - month_index = value.month - 1 + months - year = value.year + month_index // 12 - month = month_index % 12 + 1 - day = min( - value.day, - [ - 31, - 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28, - 31, - 30, - 31, - 30, - 31, - 31, - 30, - 31, - 30, - 31, - ][month - 1], - ) - return value.replace(year=year, month=month, day=day) - - -def _start_of_month(value: datetime) -> datetime: - return value.replace(day=1, hour=0, minute=0, second=0, microsecond=0) - - -def _replace_relative_getdate_calls(sql: str, now: datetime) -> str: - def replace_month_offset(match: re.Match[str]) -> str: - months = int(match.group(1)) - return _format_timestamp_literal(_add_months(now, months)) - - def replace_year_offset(match: re.Match[str]) -> str: - years = int(match.group(1)) - return _format_timestamp_literal(_add_months(now, years * 12)) - - def replace_day_offset(match: re.Match[str]) -> str: - days = int(match.group(1)) - return _format_timestamp_literal(now + timedelta(days=days)) - - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_month_offset, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*year\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_year_offset, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*day\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_day_offset, - sql, - flags=re.IGNORECASE, - ) - - current_month_start = _format_timestamp_literal(_start_of_month(now)) - previous_month_start = _format_timestamp_literal( - _start_of_month(_add_months(now, -1)) - ) - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*,\s*0\s*\)", - current_month_start, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*-\s*1\s*,\s*0\s*\)", - previous_month_start, - sql, - flags=re.IGNORECASE, - ) - return sql - - -def _replace_relative_current_date_calls(sql: str, now: datetime) -> str: - def timestamp_for_interval(amount: int, unit: str, direction: int = -1) -> str: - normalized_unit = unit.lower() - signed_amount = amount * direction - if normalized_unit.startswith("month"): - return _format_timestamp_literal(_add_months(now, signed_amount)) - if normalized_unit.startswith("year"): - return _format_timestamp_literal(_add_months(now, signed_amount * 12)) - if normalized_unit.startswith("day"): - return _format_timestamp_literal(now + timedelta(days=signed_amount)) - if normalized_unit.startswith("hour"): - return _format_timestamp_literal(now + timedelta(hours=signed_amount)) - if normalized_unit.startswith("minute"): - return _format_timestamp_literal(now + timedelta(minutes=signed_amount)) - if normalized_unit.startswith("second"): - return _format_timestamp_literal(now + timedelta(seconds=signed_amount)) - return "" - - def replace_date_sub_interval(match: re.Match[str]) -> str: - amount = int(match.group("amount")) - unit = match.group("unit").lower() - return timestamp_for_interval(amount, unit) or match.group(0) - - def replace_date_sub_unit_amount(match: re.Match[str]) -> str: - unit = match.group("unit").lower() - amount = int(match.group("amount")) - return timestamp_for_interval(amount, unit) or match.group(0) - - def replace_timestamp_interval(match: re.Match[str]) -> str: - direction = -1 if match.group("operator") == "-" else 1 - return ( - timestamp_for_interval( - int(match.group("amount")), match.group("unit"), direction - ) - or match.group(0) - ) - - sql = re.sub( - r"\bDATE_SUB\(\s*CURRENT_DATE(?:\(\))?\s*,\s*INTERVAL\s+(?P\d+)\s+(?PYEAR|MONTH|DAY)S?\s*\)", - replace_date_sub_interval, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"\bDATE_SUB\(\s*'?(?PYEAR|MONTH|DAY)'?\s*,\s*(?P\d+)\s*,\s*CURRENT_DATE(?:\(\))?\s*\)", - replace_date_sub_unit_amount, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"\b(?:CURRENT_DATE|CURRENT_TIMESTAMP|NOW)(?:\(\))?\s*(?P[+-])\s*INTERVAL\s+'?(?P\d+)\s+(?PYEAR|MONTH|DAY|HOUR|MINUTE|SECOND)S?'?", - replace_timestamp_interval, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"\b(?:CURRENT_DATE|CURRENT_TIMESTAMP|NOW)(?:\(\))?\b", - _format_timestamp_literal(now), - sql, - flags=re.IGNORECASE, - ) - return sql - - -def _rewrite_mssql_bucket_functions(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - - sql = re.sub( - rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: ( - f"(EXTRACT(YEAR FROM {m.group(1)}) * 100 + EXTRACT(MONTH FROM {m.group(1)}))" - ), - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - sql, - flags=re.IGNORECASE, - ) - return sql - - -def _rewrite_temporal_bucket_functions(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - replacements = [ - ( - re.compile( - rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ] - - rewritten = sql - for pattern, replacement in replacements: - rewritten = pattern.sub(replacement, rewritten) - - return rewritten - - -def _qualify_mssql_temporal_expression(expression: str, sql: str) -> str: - expression = expression.strip() - if "." in expression or expression.startswith(('"', "[")): - return expression - - if re.search(r"\bdbo_DebugEntries\b", sql, flags=re.IGNORECASE) and re.fullmatch( - r"(?:DateIn|DateOut|FailedAt)", expression, flags=re.IGNORECASE - ): - canonical_columns = { - "datein": "DateIn", - "dateout": "DateOut", - "failedat": "FailedAt", - } - return f'"dbo_DebugEntries"."{canonical_columns[expression.lower()]}"' - - return f'"{expression}"' - - -def _rewrite_mssql_to_date_buckets(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - - def make_day_bucket(expression: str) -> str: - timestamp_expression = _qualify_mssql_temporal_expression(expression, sql) - return ( - f"(EXTRACT(YEAR FROM {timestamp_expression}) * 10000 + " - f"EXTRACT(MONTH FROM {timestamp_expression}) * 100 + " - f"EXTRACT(DAY FROM {timestamp_expression}))" - ) - - rewritten = re.sub( - rf"\bTO_DATE\(\s*{expression_pattern}\s*,\s*'YYYY-MM-DD'\s*\)", - lambda match: make_day_bucket(match.group(1)), - sql, - flags=re.IGNORECASE, - ) - rewritten = re.sub( - rf"\bDATE\(\s*{expression_pattern}\s*\)", - lambda match: make_day_bucket(match.group(1)), - rewritten, - flags=re.IGNORECASE, - ) - return rewritten - - -def _rewrite_mssql_timestamp_casts(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - timestamp_function_pattern = re.compile( - rf"\bTO_TIMESTAMP(?:_(?:MILLIS|SECONDS|MICROS|NANOS))?\(\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ) - timestamp_cast_pattern = re.compile( - r"CAST\(\s*((?:[^()]|\([^()]*\))+?)\s+AS\s+TIMESTAMP\s*\)", - re.IGNORECASE, - ) - - rewritten = timestamp_function_pattern.sub( - lambda m: f"CAST({m.group(1)} AS DATETIME)", sql - ) - rewritten = timestamp_cast_pattern.sub( - lambda m: f"CAST({m.group(1)} AS DATETIME)", rewritten - ) - return rewritten - - -def _rewrite_mssql_to_unixtime(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - to_unixtime_pattern = re.compile( - rf"\bTO_UNIXTIME\(\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ) - - return to_unixtime_pattern.sub(lambda m: m.group(1), sql) - - -def _rewrite_mssql_timestamp_subtraction(sql: str) -> str: - expression_pattern = r"((?:[^(),+\-]|\([^()]*\))+?)" - timestamp_subtraction_pattern = re.compile( - rf"{expression_pattern}\s*-\s*{expression_pattern}\s+AS\s+(\"[^\"]+\")", - re.IGNORECASE, - ) - - def replace_subtraction(match: re.Match[str]) -> str: - left = match.group(1).strip() - right = match.group(2).strip() - alias = match.group(3) - alias_text = str(alias or "").strip('"').lower() - - if not any(token in alias_text for token in ("duration", "turnaround")): - return match.group(0) - - return f"DATEDIFF('second', {right}, {left}) AS {alias}" - - return timestamp_subtraction_pattern.sub(replace_subtraction, sql) - - -def _infer_mssql_timestamp_expression(sql: str) -> str | None: - timestamp_column_pattern = re.compile( - r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at|DateIn|DateOut|FailedAt)")', - re.IGNORECASE, - ) - if match := timestamp_column_pattern.search(sql): - return match.group(0) - - return None - - -def _rewrite_mssql_aggregate_qualified_temporal_columns(sql: str) -> str: - table_references = extract_sql_table_references(sql) - if len(table_references) != 1: - return sql - - table_name = table_references[0] - if not table_name: - return sql - - table_ref = _quote_sql_identifier(table_name) - aggregate_qualifier_pattern = ( - r'(?:"(?:SUM|COUNT|AVG|MIN|MAX)"|\[(?:SUM|COUNT|AVG|MIN|MAX)\]|' - r'\b(?:SUM|COUNT|AVG|MIN|MAX)\b)' - ) - temporal_column_pattern = ( - r'(?:"(?Pcreated_at|updated_at|generated_at|opened_at|closed_at|' - r'completed_at|resolved_at|DateIn|DateOut|FailedAt)"|' - r'\[(?Pcreated_at|updated_at|generated_at|opened_at|closed_at|' - r'completed_at|resolved_at|DateIn|DateOut|FailedAt)\]|' - r'(?Pcreated_at|updated_at|generated_at|opened_at|closed_at|' - r'completed_at|resolved_at|DateIn|DateOut|FailedAt))' - ) - pattern = re.compile( - rf"{aggregate_qualifier_pattern}\s*\.\s*{temporal_column_pattern}", - re.IGNORECASE, - ) - - def replace_reference(match: re.Match[str]) -> str: - column = ( - match.group("quoted") - or match.group("bracketed") - or match.group("bare") - ) - return f"{table_ref}.{_quote_sql_identifier(str(column))}" - - return pattern.sub(replace_reference, sql) - - -def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: - timestamp_expression = _infer_mssql_timestamp_expression(sql) - if not timestamp_expression: - return sql - - qualified_invented_date_identifier_pattern = re.compile( - r'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' - r'(?:"(?:RepairDate|repair_date|repairDate|date|month_date|event_date)"|\[(?:RepairDate|repair_date|repairDate|date|month_date|event_date)\]|(?:RepairDate|repair_date|repairDate|month_date|event_date))', - re.IGNORECASE, - ) - invented_date_identifier_pattern = re.compile( - r'(? str: - sales_table_identifier = ( - r'(?:"dbo_(?:qSales1|tblSalesHistory|tblSales)"' - r"|\[dbo_(?:qSales1|tblSalesHistory|tblSales)\]" - r"|dbo_(?:qSales1|tblSalesHistory|tblSales))" - ) - - def replace_qualified_column(column_pattern: str, canonical_column: str) -> None: - nonlocal sql - sql = re.sub( - rf"(?P
{sales_table_identifier})\s*\.\s*{column_pattern}", - lambda match: f'{match.group("table")}."{canonical_column}"', - sql, - flags=re.IGNORECASE, - ) - - otd_date_pattern = ( - r'(?:"OTD_Date"|"OTDDate"|\[OTD_Date\]|\[OTDDate\]|OTD_Date|OTDDate)' - ) - fix_log_id_pattern = ( - r'(?:"FixLogId"|"FixLogID"|\[FixLogId\]|\[FixLogID\]|FixLogId|FixLogID)' - ) - replace_qualified_column(otd_date_pattern, "InvDate") - replace_qualified_column(fix_log_id_pattern, "InvoiceNo") - - if len(extract_sql_table_references(sql)) == 1 and re.search( - rf"\bFROM\s+{sales_table_identifier}\b", sql, flags=re.IGNORECASE - ): - sql = re.sub( - rf"(? str: - rewritten = sql - - if re.search(r"\bdbo_repair_logs\b", rewritten, flags=re.IGNORECASE): - invented_failure_pattern_id_pattern = re.compile( - r'(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)', - re.IGNORECASE, - ) - rewritten = invented_failure_pattern_id_pattern.sub( - '"dbo_repair_logs"."failure_code"', rewritten - ) - - if re.search(r"\bdbo_DebugEntries\b", rewritten, flags=re.IGNORECASE) and re.search( - r"\bdbo_failure_patterns\b", rewritten, flags=re.IGNORECASE - ): - invented_debug_failure_pattern_pattern = re.compile( - r'(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)', - re.IGNORECASE, - ) - rewritten = invented_debug_failure_pattern_pattern.sub( - '"dbo_DebugEntries"."FailureSys"', rewritten - ) - - debug_id_identifier = ( - r'(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)' - r'\s*\.\s*(?:"DebugEntryId"|\[DebugEntryId\]|DebugEntryId)' - ) - failure_pattern_id_identifier = ( - r'(?:"dbo_failure_patterns"|\[dbo_failure_patterns\]|dbo_failure_patterns)' - r'\s*\.\s*(?:"id"|\[id\]|id)' - ) - debug_id_to_failure_pattern_pattern = re.compile( - rf"{debug_id_identifier}\s*=\s*{failure_pattern_id_identifier}", - re.IGNORECASE, - ) - failure_pattern_to_debug_id_pattern = re.compile( - rf"{failure_pattern_id_identifier}\s*=\s*{debug_id_identifier}", - re.IGNORECASE, - ) - rewritten = debug_id_to_failure_pattern_pattern.sub( - '"dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id"', - rewritten, - ) - rewritten = failure_pattern_to_debug_id_pattern.sub( - '"dbo_failure_patterns"."id" = "dbo_DebugEntries"."FailureSys"', - rewritten, - ) - - return rewritten - - -def _rewrite_mssql_invented_pcb_throughput_identifiers(sql: str) -> str: - rewritten = sql - manufacturing_unit_pattern = ( - r'(?:"ManufacturingUnit"|"Manufacturing_Unit"|"manufacturing_unit"|' - r'\[ManufacturingUnit\]|\[Manufacturing_Unit\]|\[manufacturing_unit\]|' - r'ManufacturingUnit|Manufacturing_Unit|manufacturing_unit)' - ) - - if re.search(r"\bdbo_DebugEntries\b", rewritten, flags=re.IGNORECASE): - debug_table_pattern = ( - r'(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)\s*\.\s*' - ) - rewritten = re.sub( - rf"{debug_table_pattern}{manufacturing_unit_pattern}", - '"dbo_DebugEntries"."BusinessUnit"', - rewritten, - flags=re.IGNORECASE, - ) - - if re.search(r"\bdbo_repair_logs\b", rewritten, flags=re.IGNORECASE) and re.search( - manufacturing_unit_pattern, rewritten, flags=re.IGNORECASE - ): - rewritten = re.sub( - r'(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)', - '"dbo_DebugEntries"', - rewritten, - flags=re.IGNORECASE, - ) - rewritten = re.sub( - rf'"dbo_DebugEntries"\s*\.\s*{manufacturing_unit_pattern}', - '"dbo_DebugEntries"."BusinessUnit"', - rewritten, - flags=re.IGNORECASE, - ) - rewritten = re.sub( - r'"dbo_DebugEntries"\s*\.\s*(?:"id"|\[id\]|id)', - '"dbo_DebugEntries"."DebugEntryId"', - rewritten, - flags=re.IGNORECASE, - ) - rewritten = re.sub( - r'"dbo_DebugEntries"\s*\.\s*(?:"created_at"|"updated_at"|\[created_at\]|\[updated_at\]|created_at|updated_at)', - '"dbo_DebugEntries"."DateIn"', - rewritten, - flags=re.IGNORECASE, - ) - - return rewritten - - -def _rewrite_mssql_repair_log_throughput_shape(sql: str) -> str: - if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\bavg_turnaround_time\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\brepair_count\b|\bthroughput\b", sql, flags=re.IGNORECASE): - return sql - - return ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' - 'COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'GROUP BY "dbo_DebugEntries"."BusinessUnit" ' - 'ORDER BY "throughput" DESC' - ) - - -def _rewrite_mssql_repair_log_turnaround_trend_shape(sql: str) -> str: - if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\bavg_turnaround_time\b|\bturnaround\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\bMONTH\b|DATEPART\(\s*'?\s*MONTH", sql, flags=re.IGNORECASE): - return sql - - return ( - "SELECT EXTRACT(YEAR FROM \"dbo_repair_logs\".\"created_at\") AS \"year\", " - "EXTRACT(MONTH FROM \"dbo_repair_logs\".\"created_at\") AS \"month\", " - 'AVG(DATEDIFF(\'second\', "dbo_repair_logs"."created_at", ' - '"dbo_repair_logs"."updated_at")) AS "avg_turnaround_seconds" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' - 'AND "dbo_repair_logs"."updated_at" IS NOT NULL ' - "GROUP BY EXTRACT(YEAR FROM \"dbo_repair_logs\".\"created_at\"), " - "EXTRACT(MONTH FROM \"dbo_repair_logs\".\"created_at\") " - "ORDER BY EXTRACT(YEAR FROM \"dbo_repair_logs\".\"created_at\") ASC, " - "EXTRACT(MONTH FROM \"dbo_repair_logs\".\"created_at\") ASC" - ) - - -def _rewrite_mssql_ticket_cycle_turnaround_shape(sql: str) -> str: - if not re.search(r"\bdbo_ticket_cycles\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\bturnaround_time\b|\bavg_turnaround_time\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\bMONTH\b|DATEPART\(\s*'?\s*MONTH", sql, flags=re.IGNORECASE): - return sql - - return ( - 'SELECT EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") AS "year", ' - 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") AS "month", ' - 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", ' - '"dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds" ' - 'FROM "dbo_ticket_cycles" ' - 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL ' - 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL ' - 'GROUP BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at"), ' - 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") ' - 'ORDER BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") ASC, ' - 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") ASC' - ) - - -def _rewrite_mssql_invented_failure_category(sql: str) -> str: - if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\bfailure[_\s]+category\b", sql, flags=re.IGNORECASE): - return sql - - failure_code_expression = '"dbo_repair_logs"."failure_code"' - rewritten = re.sub( - r'(?P\bSELECT\s+|,\s*)(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)(?P\s*(?:,|\bFROM\b))', - rf'\g{failure_code_expression} AS "failure_category"\g', - sql, - flags=re.IGNORECASE, - ) - rewritten = re.sub( - r"\bAS\s+failure\s+category\b", - 'AS "failure_category"', - rewritten, - flags=re.IGNORECASE, - ) - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match[str]) -> str: - body = re.sub( - r'(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)', - failure_code_expression, - match.group("body"), - flags=re.IGNORECASE, - ) - return f"{match.group(1)}{body}" - - return clause_pattern.sub(replace_clause, rewritten) - - -def _rewrite_mssql_invented_report_fields(sql: str) -> str: - if not re.search(r"\bdbo_reports\b", sql, flags=re.IGNORECASE): - return sql - - report_table = r'(?:"dbo_reports"|\[dbo_reports\]|dbo_reports)' - rewritten = re.sub( - rf"(?:(?:{report_table})\s*\.\s*)?(?:\"filters\"|\[filters\]|\bfilters\b)", - '"dbo_reports"."data"', - sql, - flags=re.IGNORECASE, - ) - rewritten = re.sub( - rf"(?:(?:{report_table})\s*\.\s*)?(?:\"report_size\"|\"file_size\"|\[report_size\]|\[file_size\]|\breport_size\b|\bfile_size\b)", - '"dbo_reports"."size_bytes"', - rewritten, - flags=re.IGNORECASE, - ) - return rewritten - - -def _rewrite_mssql_invented_ticket_metrics(sql: str) -> str: - if not re.search(r"\bdbo_tickets\b", sql, flags=re.IGNORECASE): - return sql - - if not re.search( - r"\b(?:token_cost|average_token_cost|avg_token_cost|cost)\b", - sql, - flags=re.IGNORECASE, - ): - return sql - - dimension = '"dbo_tickets"."status"' - alias = "status" - if re.search(r"\bpriority\b", sql, flags=re.IGNORECASE): - dimension = '"dbo_tickets"."priority"' - alias = "priority" - - return ( - f'SELECT {dimension} AS "{alias}", ' - 'COUNT("dbo_tickets"."id") AS "ticket_count" ' - 'FROM "dbo_tickets" ' - f"GROUP BY {dimension} " - 'ORDER BY "ticket_count" DESC' - ) - - -def _rewrite_mssql_invented_knowledge_article_fields(sql: str) -> str: - if not re.search( - r"\b(?:dbo_knowledge_articles|dbo_kb_articles)\b", sql, flags=re.IGNORECASE - ): - return sql - - rewritten = sql - table_replacements = { - "dbo_knowledge_articles": { - "article_id": '"id"', - "knowledge_article_id": '"id"', - "article_content": '"content"', - "article_text": '"content"', - "article_body": '"content"', - "effectiveness_score": '"helpful"', - "created_by": '"author"', - "created_by_user": '"author"', - "created_by_user_id": '"author"', - "author_id": '"author"', - }, - "dbo_kb_articles": { - "article_id": '"id"', - "knowledge_article_id": '"id"', - "article_content": '"content"', - "article_text": '"content"', - "article_body": '"content"', - "category": '"category"', - "section": '"category"', - "article_section": '"category"', - "created_by": '"created_by_user_id"', - "created_by_user": '"created_by_user_id"', - "author": '"created_by_user_id"', - "author_id": '"created_by_user_id"', - }, - } - - for table_name, field_replacements in table_replacements.items(): - article_table = ( - rf'(?:"{table_name}"|\[{table_name}\]|{table_name})' - ) - for invented_field, replacement_field in field_replacements.items(): - rewritten = re.sub( - rf"(?P
{article_table})\s*\.\s*(?:\"{invented_field}\"|\[{invented_field}\]|\b{invented_field}\b)", - rf"\g
.{replacement_field}", - rewritten, - flags=re.IGNORECASE, - ) - - if re.search(r"\bdbo_knowledge_articles\b", rewritten, flags=re.IGNORECASE): - unqualified_replacements = table_replacements["dbo_knowledge_articles"] - elif re.search(r"\bdbo_kb_articles\b", rewritten, flags=re.IGNORECASE): - unqualified_replacements = table_replacements["dbo_kb_articles"] - else: - unqualified_replacements = {} - - for invented_field, replacement_field in unqualified_replacements.items(): - rewritten = re.sub( - rf'(? bool: - if re.search(r"(?:->>|->)", sql): - return True - - unsupported_json_functions = ( - "JSON_VALUE", - "JSON_QUERY", - "JSON_EXTRACT", - "JSON_EXTRACT_SCALAR", - "JSON_EXTRACT_ARRAY", - "LAX_BOOL", - "LAX_FLOAT64", - "LAX_INT64", - "LAX_STRING", - ) - function_pattern = r"\b(?:{})\s*\(".format("|".join(unsupported_json_functions)) - return bool(re.search(function_pattern, sql, flags=re.IGNORECASE)) - - -def _rewrite_mssql_bare_time_bucket_identifiers(sql: str) -> str: - timestamp_expression = _infer_mssql_timestamp_expression(sql) - if not timestamp_expression: - return sql - - bucket_expressions = { - "year": f"EXTRACT(YEAR FROM {timestamp_expression})", - "month": f"EXTRACT(MONTH FROM {timestamp_expression})", - "day": f"EXTRACT(DAY FROM {timestamp_expression})", - } - rewritten = sql - - for bucket, expression in bucket_expressions.items(): - qualified_bucket_pattern = re.compile( - rf'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' - rf'(?:"{bucket}"|\[{bucket}\]|{bucket})', - re.IGNORECASE, - ) - select_pattern = re.compile( - r"\bSELECT\b(?P.*?)(?=\bFROM\b)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_select(match: re.Match[str]) -> str: - body = match.group("body") - items = _split_top_level_select_items(body) - if not items: - return match.group(0) - - rebuilt: list[str] = [] - changed = False - bucket_select_item_pattern = re.compile( - rf"^(?P(?:\"{bucket}\"|\[{bucket}\]|{bucket}|{qualified_bucket_pattern.pattern}))" - rf"(?:\s+(?:AS\s+)?(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*))?$", - re.IGNORECASE, - ) - for item in items: - item_match = bucket_select_item_pattern.fullmatch(item.strip()) - if item_match: - alias = item_match.group("alias") or f'"{bucket}"' - rebuilt.append(f"{expression} AS {alias}") - changed = True - else: - rebuilt.append(item) - - if not changed: - return match.group(0) - - return "SELECT " + ", ".join(rebuilt) + " " - - rewritten = select_pattern.sub(replace_select, rewritten) - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match[str]) -> str: - body = match.group("body") - for bucket, expression in bucket_expressions.items(): - qualified_bucket_pattern = re.compile( - rf'(?:(?:"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*\.\s*)' - rf'(?:"{bucket}"|\[{bucket}\]|{bucket})', - re.IGNORECASE, - ) - body = qualified_bucket_pattern.sub(expression, body) - body = re.sub( - rf'"{bucket}"', - expression, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"\[{bucket}\]", - expression, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"(? str: - limit_match = re.search(r"\s+LIMIT\s+(\d+)\s*;?\s*$", sql, flags=re.IGNORECASE) - if not limit_match: - return sql - - limit = limit_match.group(1) - without_limit = sql[: limit_match.start()].rstrip() - if re.search( - r"\bSELECT\s+(?:DISTINCT\s+)?TOP\s+(?:\(\s*)?\d+", - without_limit, - flags=re.IGNORECASE, - ): - return without_limit - - if re.match(r"\s*SELECT\s+DISTINCT\b", without_limit, flags=re.IGNORECASE): - return re.sub( - r"\bSELECT\s+DISTINCT\b", - f"SELECT DISTINCT TOP {limit}", - without_limit, - count=1, - flags=re.IGNORECASE, - ) - - if re.match(r"\s*SELECT\b", without_limit, flags=re.IGNORECASE): - return re.sub( - r"\bSELECT\b", - f"SELECT TOP {limit}", - without_limit, - count=1, - flags=re.IGNORECASE, - ) - - return without_limit - - -def _unwrap_simple_mssql_where_parentheses(sql: str) -> str: - return re.sub( - r"\bWHERE\s*\(\s*([^()]+?)\s*\)(?=\s*(?:GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|$))", - r"WHERE \1", - sql, - flags=re.IGNORECASE | re.DOTALL, - ) - - -def _rewrite_mssql_datepart_alias_references(sql: str) -> str: - datepart_alias_pattern = re.compile( - r"\b((?:DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)|(?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\)|EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\)))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", - re.IGNORECASE, - ) - aliases: dict[str, str] = {} - - for match in datepart_alias_pattern.finditer(sql): - expression = match.group(1) - alias = match.group(5) or match.group(6) or match.group(7) - if alias: - aliases[str(alias).lower()] = expression - - if not aliases: - return sql - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match) -> str: - body = match.group("body") - placeholders: dict[str, str] = {} - for alias, expression in aliases.items(): - placeholder = f"__WREN_MSSQL_DATEPART_ALIAS_{len(placeholders)}__" - placeholders[placeholder] = expression - body = re.sub( - rf'"{re.escape(alias)}"', - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"\[{re.escape(alias)}\]", - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"(? list[str]: - items: list[str] = [] - current: list[str] = [] - depth = 0 - in_single_quote = False - in_double_quote = False - - for char in select_body: - if char == "'" and not in_double_quote: - in_single_quote = not in_single_quote - elif char == '"' and not in_single_quote: - in_double_quote = not in_double_quote - elif not in_single_quote and not in_double_quote: - if char == "(": - depth += 1 - elif char == ")" and depth > 0: - depth -= 1 - elif char == "," and depth == 0: - items.append("".join(current).strip()) - current = [] - continue - current.append(char) - - if current: - items.append("".join(current).strip()) - - return items - - -def _is_word_at(sql: str, index: int, word: str) -> bool: - end = index + len(word) - if sql[index:end].upper() != word: - return False - before = sql[index - 1] if index > 0 else "" - after = sql[end] if end < len(sql) else "" - return not (before.isalnum() or before == "_") and not ( - after.isalnum() or after == "_" - ) - - -def _find_select_list_spans(sql: str) -> list[tuple[int, int]]: - spans: list[tuple[int, int]] = [] - depth = 0 - in_single_quote = False - in_double_quote = False - in_bracket = False - index = 0 - - while index < len(sql): - char = sql[index] - if char == "'" and not in_double_quote and not in_bracket: - in_single_quote = not in_single_quote - elif char == '"' and not in_single_quote and not in_bracket: - in_double_quote = not in_double_quote - elif char == "[" and not in_single_quote and not in_double_quote: - in_bracket = True - elif char == "]" and in_bracket: - in_bracket = False - elif not in_single_quote and not in_double_quote and not in_bracket: - if char == "(": - depth += 1 - elif char == ")" and depth > 0: - depth -= 1 - elif _is_word_at(sql, index, "SELECT"): - select_depth = depth - select_body_start = index + len("SELECT") - cursor = select_body_start - cursor_depth = depth - cursor_in_single_quote = False - cursor_in_double_quote = False - cursor_in_bracket = False - - while cursor < len(sql): - cursor_char = sql[cursor] - if ( - cursor_char == "'" - and not cursor_in_double_quote - and not cursor_in_bracket - ): - cursor_in_single_quote = not cursor_in_single_quote - elif ( - cursor_char == '"' - and not cursor_in_single_quote - and not cursor_in_bracket - ): - cursor_in_double_quote = not cursor_in_double_quote - elif ( - cursor_char == "[" - and not cursor_in_single_quote - and not cursor_in_double_quote - ): - cursor_in_bracket = True - elif cursor_char == "]" and cursor_in_bracket: - cursor_in_bracket = False - elif ( - not cursor_in_single_quote - and not cursor_in_double_quote - and not cursor_in_bracket - ): - if cursor_char == "(": - cursor_depth += 1 - elif cursor_char == ")" and cursor_depth > 0: - cursor_depth -= 1 - elif cursor_depth == select_depth and _is_word_at( - sql, cursor, "FROM" - ): - spans.append((select_body_start, cursor)) - break - cursor += 1 - index = cursor - index += 1 - - return spans - - -def _strip_projection_alias(item: str) -> str: - alias_match = re.search( - r"\s+(?:AS\s+)?(?:\"[^\"]+\"|\[[^\]]+\]|`[^`]+`|[A-Za-z_][A-Za-z0-9_]*)\s*$", - item, - flags=re.IGNORECASE, - ) - if not alias_match: - return item.strip() - - expression = item[: alias_match.start()].strip() - if not expression: - return item.strip() - return expression - -def _simple_projection_key(item: str) -> str | None: - cleaned = re.sub(r"^\s*DISTINCT\s+", "", item, flags=re.IGNORECASE).strip() - cleaned = _strip_projection_alias(cleaned) - identifier_pattern = ( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_]*|\*))" - ) - identifier_match_pattern = ( - r'(?:"[^"]+"|\[[^\]]+\]|`[^`]+`|[A-Za-z_][A-Za-z0-9_]*|\*)' - ) - qualified_identifier_pattern = rf"^\s*{identifier_match_pattern}(?:\s*\.\s*{identifier_match_pattern})*\s*$" - if not re.match(qualified_identifier_pattern, cleaned): - return None - - parts = re.findall(identifier_pattern, cleaned) - if not parts: - return None - last_part = next((value for value in parts[-1] if value), "") - return str(last_part).lower() - - -def _dedupe_duplicate_simple_select_items(sql: str) -> str: - spans = _find_select_list_spans(sql) - if not spans: - return sql - - normalized = sql - for start, end in reversed(spans): - body = normalized[start:end] - items = _split_top_level_select_items(body) - if len(items) < 2: + first_token = line.split()[0].strip('"`[]') + if first_token.upper() in { + "CONSTRAINT", + "FOREIGN", + "PRIMARY", + "UNIQUE", + "KEY", + }: continue + columns.append(first_token) - seen_simple_projection_keys: set[str] = set() - deduped_items: list[str] = [] - changed = False - for item in items: - key = _simple_projection_key(item) - if key and key in seen_simple_projection_keys: - changed = True - logger.debug( - 'Removing duplicate simple projection "%s" from generated SQL', - item, - ) - continue - if key: - seen_simple_projection_keys.add(key) - deduped_items.append(item) - - if changed: - normalized = ( - normalized[:start] - + " " - + ", ".join(deduped_items) - + " " - + normalized[end:] - ) - - return normalized + return columns -def _rewrite_mssql_temporal_bucket_alias_references(sql: str) -> str: - select_match = re.search( - r"\bSELECT\b(?P.*?)(?=\bFROM\b)", - sql, - flags=re.IGNORECASE | re.DOTALL, - ) - if not select_match: - return sql - - aliases: dict[str, str] = {} - for item in _split_top_level_select_items(select_match.group("body")): - if not re.search( - r"\b(?:DATEPART|YEAR|MONTH|DAY|EXTRACT)\s*\(", item, flags=re.IGNORECASE - ): - continue +def format_retrieved_schema_manifest( + documents: list[str] | None, + allowed_table_names: list[str] | None, +) -> list[dict[str, Any]]: + table_names = allowed_table_names or [] + ddls = documents or [] - alias_match = re.search( - r"\s+(?:AS\s+)?(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))\s*$", - item, - flags=re.IGNORECASE, - ) - if not alias_match: + manifest = [] + for index, table_name in enumerate(table_names): + if not isinstance(table_name, str) or not table_name.strip(): continue - - alias = alias_match.group(1) or alias_match.group(2) or alias_match.group(3) - expression = item[: alias_match.start()].strip() - if alias: - aliases[str(alias).lower()] = expression - - if not aliases: - return sql - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match[str]) -> str: - body = match.group("body") - placeholders: dict[str, str] = {} - for alias, expression in aliases.items(): - placeholder = f"__WREN_MSSQL_TEMPORAL_BUCKET_ALIAS_{len(placeholders)}__" - placeholders[placeholder] = expression - body = re.sub( - rf'"{re.escape(alias)}"', - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"\[{re.escape(alias)}\]", - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"(? bool: - return False - - -def _rewrite_known_schema_hallucinations(sql: str, now: datetime) -> str: - normalized = _replace_relative_current_date_calls(sql, now) - normalized = _unwrap_simple_mssql_where_parentheses(normalized) - normalized = _rewrite_mssql_limit_clause(normalized) - normalized = _rewrite_mssql_to_date_buckets(normalized) - normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) - normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) - normalized = _rewrite_temporal_bucket_functions(normalized) - normalized = _rewrite_mssql_datepart_alias_references(normalized) - normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) - return normalized - - -def _rewrite_mssql_limit_clause(sql: str) -> str: - limit_match = re.search(r"\s+LIMIT\s+(\d+)\s*;?\s*$", sql, flags=re.IGNORECASE) - if not limit_match: - return sql - - limit = limit_match.group(1) - without_limit = sql[: limit_match.start()].rstrip() - if re.search(r"\bUNION(?:\s+ALL)?\b", without_limit, flags=re.IGNORECASE): - return without_limit - - if re.search( - r"\bSELECT\s+(?:DISTINCT\s+)?TOP\s*\(?\s*\d+\s*\)?", - without_limit, - flags=re.IGNORECASE, - ): - return without_limit - - return re.sub( - r"\bSELECT\s+(DISTINCT\s+)?", - lambda match: f"{match.group(0)}TOP {limit} ", - without_limit, - count=1, - flags=re.IGNORECASE, - ) - - -def _normalize_identifier_quote_syntax(sql: str) -> str: - normalized = re.sub( - r"`([^`]+)`", - lambda match: _quote_sql_identifier(match.group(1)), - sql, - ) - normalized = re.sub( - r"\[([^\]]+)\]", - lambda match: _quote_sql_identifier(match.group(1)), - normalized, - ) - normalized = re.sub( - r'"{2,}([A-Za-z_][A-Za-z0-9_$]*)"{2,}', - lambda match: _quote_sql_identifier(match.group(1)), - normalized, - ) - return normalized - - -def normalize_sql_direction_keywords(sql: str) -> str: - parts = re.split(r'("(?:[^"]|"")*"|\'(?:\'\'|[^\'])*\')', sql or "") - for index in range(0, len(parts), 2): - parts[index] = re.sub( - r"\b(asc|desc)\b", - lambda match: match.group(1).upper(), - parts[index], - flags=re.IGNORECASE, + ddl = ddls[index] if index < len(ddls) and isinstance(ddls[index], str) else "" + manifest.append( + { + "table_name": table_name.strip(), + "columns": _extract_ddl_columns(ddl), + } ) - return "".join(parts) - -def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: - normalized = _normalize_identifier_quote_syntax( - normalize_sql_direction_keywords(sql) - ) - normalized_data_source = normalize_data_source(data_source) - - if normalized_data_source == "MSSQL": - now = datetime.now() - normalized = re.sub( - r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE - ) - normalized = _unwrap_simple_mssql_where_parentheses(normalized) - normalized = _rewrite_mssql_limit_clause(normalized) - normalized = re.sub( - r"CAST\(\s*('(?:[^']|'')*')\s+AS\s+DATETIME(?:2|OFFSET)\s*\)", - r"\1", - normalized, - flags=re.IGNORECASE, - ) - normalized = _replace_relative_getdate_calls(normalized, now) - normalized = _replace_relative_current_date_calls(normalized, now) - normalized = _rewrite_mssql_to_unixtime(normalized) - normalized = _rewrite_mssql_timestamp_subtraction(normalized) - normalized = _rewrite_mssql_to_date_buckets(normalized) - normalized = _rewrite_mssql_timestamp_casts(normalized) - normalized = _rewrite_mssql_aggregate_qualified_temporal_columns(normalized) - normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) - normalized = _rewrite_mssql_bucket_functions(normalized) - normalized = _rewrite_temporal_bucket_functions(normalized) - normalized = _rewrite_mssql_datepart_alias_references(normalized) - normalized = _rewrite_mssql_temporal_bucket_alias_references(normalized) - normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) - normalized = _rewrite_mssql_limit_clause(normalized) - elif _references_known_hallucination_prone_schema(normalized): - normalized = _rewrite_known_schema_hallucinations(normalized, datetime.now()) - - normalized = _dedupe_duplicate_simple_select_items(normalized) - - return re.sub(r"\s+", " ", normalized).strip() + return manifest @component @@ -1546,102 +85,15 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - valid_table_names: list[str] | None = None, - valid_table_columns: dict[str, list[str]] | None = None, ) -> dict: try: - cleaned_generation_result = extract_sql_generation_result(replies[0]) - - cleaned_generation_result = normalize_generation_result_sql( - cleaned_generation_result, data_source=data_source - ) - cleaned_generation_result = normalize_sql_table_references_to_schema( - cleaned_generation_result, - valid_table_names or [], - ) - cleaned_generation_result = normalize_sql_column_references_to_schema( - cleaned_generation_result, - valid_table_columns or {}, - ) + cleaned_generation_result = clean_generation_result(replies[0]) - if not is_select_statement(cleaned_generation_result): - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "DRY_RUN", - "error": "Generated response did not contain a SQL SELECT statement.", - "correlation_id": "", - }, - } - - invalid_table_references = find_invalid_table_references( - cleaned_generation_result, - valid_table_names or [], - ) - if invalid_table_references: - valid_table_list = ", ".join(valid_table_names or []) - invalid_table_list = ", ".join(invalid_table_references) - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_VALIDATION", - "error": ( - "Generated SQL references table(s) not present in the " - f"active datasource metadata: {invalid_table_list}. " - "Use only these valid table names exactly as shown: " - f"{valid_table_list}" - ), - "correlation_id": "", - }, - } - - invalid_column_references = find_invalid_column_references( - cleaned_generation_result, - valid_table_columns or {}, - ) - if invalid_column_references: - invalid_column_list = ", ".join(invalid_column_references) - valid_column_list = format_valid_table_columns(valid_table_columns or {}) - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_VALIDATION", - "error": ( - "Generated SQL references column(s) not present in the " - f"active datasource metadata: {invalid_column_list}. " - "Use only these valid table columns exactly as shown: " - f"{valid_column_list}" - ), - "correlation_id": "", - }, - } - - if normalize_data_source( - data_source - ) == "MSSQL" and contains_unsupported_mssql_json_access( - cleaned_generation_result - ): - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "UNSUPPORTED_SQL", - "error": ( - "Generated SQL uses JSON extraction, but the MSSQL " - "Wren/Ibis runtime does not support JSON operators " - "or JSON extraction functions. Use only first-class " - "columns exposed in the schema." - ), - "correlation_id": "", - }, - } + # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' + if cleaned_generation_result.startswith("{"): + cleaned_generation_result = orjson.loads(cleaned_generation_result)[ + "sql" + ] ( valid_generation_result, @@ -1678,9 +130,6 @@ async def _classify_generation_result( ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} - generation_result = normalize_generation_result_sql( - generation_result, data_source=data_source - ) use_dry_run = not allow_data_preview async with aiohttp.ClientSession() as session: @@ -1722,12 +171,8 @@ async def _classify_generation_result( } else: error_message = addition.get("error_message", "") - normalized_error_sql = normalize_generation_result_sql( - generation_result, - data_source=data_source, - ) invalid_generation_result = { - "sql": normalized_error_sql, + "sql": addition.get("error_sql", generation_result), "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") @@ -1756,12 +201,8 @@ async def _classify_generation_result( if error_message == "" else "PREVIEW_FAILED" ) - normalized_error_sql = normalize_generation_result_sql( - generation_result, - data_source=data_source, - ) invalid_generation_result = { - "sql": normalized_error_sql, + "sql": addition.get("error_sql", generation_result), "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") @@ -1777,6 +218,7 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. +- NEVER invent, assume, or rename tables and columns. Generate SQL only from tables and columns present in the provided database schema. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. @@ -1795,16 +237,17 @@ async def _classify_generation_result( - Use "lower(.) = lower()" when: - The user requests an exact, specific value. - There is no ambiguity or pattern in the value. -- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to a temporal type first before using it in the query. - - For engines that list these functions in SQL FUNCTIONS, use TO_TIMESTAMP_MILLIS("") if the timestamp_column is in milliseconds. - - For engines that list these functions in SQL FUNCTIONS, use TO_TIMESTAMP_SECONDS("") if the timestamp_column is in seconds. - - For engines that list these functions in SQL FUNCTIONS, use TO_TIMESTAMP_MICROS("") if the timestamp_column is in microseconds. -- When you need to cast a date/time related field, CAST it to a temporal type that is supported by the target data source and consistent with the SQL FUNCTIONS section. - - example 1: CAST(properties_closedate AS TIMESTAMP) - - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP) +- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query + - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds + - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds + - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds +- ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query + - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) + - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) + - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) - If the user asks for a specific date, please give the date range in SQL query - example: "What is the total revenue for the month of 2024-11-01?" - - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP) AND CAST(r.PurchaseTimestamp AS TIMESTAMP) < CAST('2024-11-02 00:00:00' AS TIMESTAMP)" + - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. @@ -1827,43 +270,8 @@ async def _classify_generation_result( - DON'T USE "TO_CHAR" function in the generated SQL query. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- Never invent foreign key columns or relationship fields from table names unless that exact column appears in the DATABASE SCHEMA. Join only on explicit schema columns or explicit relationships. -- Never invent time bucket columns such as "MONTH", "YEAR", "DAY", "month", "year", or "date" unless that exact column appears in the DATABASE SCHEMA. For monthly, yearly, or daily trends, apply a supported date/time bucket function from SQL FUNCTIONS to a real timestamp column from the selected table. -- Every generated SQL query must be grounded only in the connected datasource metadata, deployed semantic model definitions, relationships, and DATABASE SCHEMA shown in the prompt. Do not use table names, column names, join paths, JSON keys, or business dimensions that are not explicitly present in that context. -- For trend questions, choose an explicit timestamp/date column from the active schema and bucket it with supported SQL FUNCTIONS. Do not select, group by, or order by invented time bucket columns unless they explicitly appear in the schema. -- For grouped count questions, choose an explicit dimension column from the active schema. Do not invent category/status/type columns unless they explicitly appear in the schema. -- For top/bottom N questions, return exactly the business columns needed to answer the question. For example, "top 10 common failures" should return the failure field and the failure count. -- For top/bottom N questions, prefer ORDER BY on the metric plus a row limit instead of adding ranking helper columns. -- Do not include helper ranking columns such as "rank", "row_number", or "dense_rank" in the final SELECT unless the user explicitly asks to see ranks. -- If a ranking helper is required internally, compute it in a subquery/CTE and filter on it, but omit it from the final SELECT unless explicitly requested. -""" - -_MSSQL_TEXT_TO_SQL_RULES = """ -### MSSQL-SPECIFIC RULES ### -- The target database is MSSQL. -- Prefer parser-safe date bucket syntax such as EXTRACT(YEAR FROM "created_at") and EXTRACT(MONTH FROM "created_at"). -- DO NOT use PostgreSQL-style or Trino-style date syntax such as DATE_TRUNC, DATETRUNC, INTERVAL, CURRENT_DATE, TIMESTAMP WITH TIME ZONE, TO_CHAR, TO_UNIXTIME, TO_TIMESTAMP, TO_TIMESTAMP_MILLIS, TO_TIMESTAMP_SECONDS, TO_TIMESTAMP_MICROS, TO_TIMESTAMP_NANOS, or :: casts. -- DO NOT use JSON extraction functions or operators such as JSON_VALUE, JSON_QUERY, JSON_EXTRACT, JSON_EXTRACT_SCALAR, JSON_EXTRACT_ARRAY, json_value, json_extract, ->, or ->>. The MSSQL Wren/Ibis runtime does not support them. -- If a table has a generic JSON/text column such as "data", do not assume keys inside it are queryable. Only use fields that are exposed as first-class columns in the DATABASE SCHEMA. -- If a requested metric is only present inside a JSON/text column and is not exposed as a first-class column or calculated field, do not generate SQL that extracts it from JSON. -- Never invent JSON-derived columns unless they are explicitly listed as columns in the DATABASE SCHEMA. -- For trend or volume questions, use explicit timestamp/date columns only when those exact columns appear in the selected table schema. -- For grouped count questions, use explicit exposed dimension fields and only join tables when an explicit join key or relationship exists in the DATABASE SCHEMA. -- DO NOT use DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET unless the SQL FUNCTIONS section explicitly proves they are supported by the target runtime. -- Do not subtract timestamp/date columns directly. If a duration or turnaround column exists in the schema, select that column directly. If only start/end timestamps exist and the SQL FUNCTIONS section lists DATEDIFF, use DATEDIFF('second', , ) for duration in seconds. -- Resolve relative time phrases such as "last 12 months", "last month", or "this year" into absolute ISO timestamp boundaries using the current time context. Prefer closed-open literal ranges over runtime date arithmetic. -- For month bucketing, prefer separate year/month fields: - - EXTRACT(YEAR FROM ) AS "year" - - EXTRACT(MONTH FROM ) AS "month" - Then GROUP BY and ORDER BY the same year/month expressions. -- Do not GROUP BY or ORDER BY quoted year/month aliases such as "YEAR" or "MONTH"; repeat the EXTRACT(...) expression instead. -- For year bucketing, prefer EXTRACT(YEAR FROM ). -- For top/bottom N questions in MSSQL, prefer SELECT TOP (N) with ORDER BY over DENSE_RANK/ROW_NUMBER when the user did not explicitly request ranks. -- For filtering a specific year such as 2025, prefer a closed-open range: - - >= '2025-01-01 00:00:00' - - AND < '2026-01-01 00:00:00' -- When a temporal cast is required, use CAST( AS DATETIME), or keep literal timestamps as plain ISO strings if the column is already datetime-like. -- Keep MSSQL date logic simple and planner-safe. Never emit DATEADD/DATEDIFF fallback expressions unless the SQL FUNCTIONS section explicitly requires them. +- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. +- For the ranking problem, you must add the ranking column to the final SELECT clause. """ @@ -1992,9 +400,17 @@ async def _classify_generation_result( 1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. 2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. 3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. -So utilize those metric components in the SQL generation process to give an answer using the date functions that are valid for the target data source and listed in the SQL FUNCTIONS section. - -For example, the SQL should filter PurchaseTimestamp to the previous calendar month using the dialect-appropriate month-boundary functions from the SQL FUNCTIONS section. +So utilize those metric components in the SQL generation process to give an answer like this: + +SQL Query: +SELECT + CustomerId, + PriceSum AS TotalRevenue +FROM + Revenue +WHERE + PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND + PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -2068,9 +484,9 @@ async def _classify_generation_result( 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. 2. Explicitly state the following information in the reasoning plan: if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; -if the user uses a relative timeframe and Current Time is provided in the input, you will resolve it into an absolute time frame in the SQL query using exact dates rather than relative date arithmetic. -3. For top/bottom N questions, plan to order by the relevant metric and limit the result to N rows. Do not add a rank column unless the user explicitly asks to see ranks. -4. For questions like "top 10 common failures", the final table should contain the grouped business field and its count/metric, not helper ranking columns. +otherwise, you will put the relative timeframe in the SQL query. +3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. +4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. @@ -2081,6 +497,8 @@ async def _classify_generation_result( 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. 14. ONLY SHOWING the reasoning plan in bullet points. +15. Never include SQL code, table aliases, or assumed table/column names in the reasoning plan. +16. Only mention a table or column when the exact name appears in the DATABASE SCHEMA. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -2097,24 +515,13 @@ def _extract_from_sql_knowledge( return value if value and value.strip() else default_value -def _append_data_source_rules(base_rules: str, data_source: str | None = None) -> str: - normalized_data_source = normalize_data_source(data_source) - if normalized_data_source == "MSSQL": - return f"{base_rules}\n\n{_MSSQL_TEXT_TO_SQL_RULES}" - return base_rules - - -def get_text_to_sql_rules( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: +def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: if sql_knowledge is not None: - base_rules = _extract_from_sql_knowledge( + return _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) - return _append_data_source_rules(base_rules, data_source) - return _append_data_source_rules(_DEFAULT_TEXT_TO_SQL_RULES, data_source) + return _DEFAULT_TEXT_TO_SQL_RULES def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: @@ -2128,46 +535,16 @@ def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) return _DEFAULT_CALCULATED_FIELD_INSTRUCTIONS -def get_metric_instructions( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: - instructions = _DEFAULT_METRIC_INSTRUCTIONS +def get_metric_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: if sql_knowledge is not None: - instructions = _extract_from_sql_knowledge( + return _extract_from_sql_knowledge( sql_knowledge, "metric_instructions", _DEFAULT_METRIC_INSTRUCTIONS ) - if normalize_data_source(data_source) == "MSSQL": - instructions += """ + return _DEFAULT_METRIC_INSTRUCTIONS -#### MSSQL Metric Notes #### -- Resolve relative metric time windows into absolute ISO date ranges whenever current time context is available. -- Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE in MSSQL metric queries unless the SQL FUNCTIONS section explicitly shows they are supported by the target engine. -- For month trend metrics, prefer EXTRACT(YEAR FROM ) and EXTRACT(MONTH FROM ) as separate grouped columns. -""" - - return instructions - - -_MSSQL_JSON_FIELD_INSTRUCTIONS = """ -#### MSSQL JSON Field Instructions #### -- The target runtime cannot execute JSON extraction from generic JSON/text columns. -- Do not use JSON operators or functions such as ->, ->>, JSON_VALUE, JSON_QUERY, - JSON_EXTRACT, JSON_EXTRACT_SCALAR, LAX_STRING, LAX_INT64, LAX_FLOAT64, or LAX_BOOL. -- If the requested value is only inside a generic JSON/text column such as "data", - do not infer or extract it. Use only first-class columns and calculated fields - that are explicitly exposed in the DATABASE SCHEMA. -""" - - -def get_json_field_instructions( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: - if normalize_data_source(data_source) == "MSSQL": - return _MSSQL_JSON_FIELD_INSTRUCTIONS +def get_json_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: if sql_knowledge is not None: return _extract_from_sql_knowledge( sql_knowledge, "json_field_instructions", _DEFAULT_JSON_FIELD_INSTRUCTIONS @@ -2176,33 +553,21 @@ def get_json_field_instructions( return _DEFAULT_JSON_FIELD_INSTRUCTIONS -def get_sql_generation_system_prompt( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: - text_to_sql_rules = get_text_to_sql_rules( - sql_knowledge, - data_source=data_source, - ) +def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) -> str: + text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" You are a helpful assistant that converts natural language queries into ANSI SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. +Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query from the provided database schema. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +4. If the section of REASONING PLAN is available in user's input, treat it only as high-level guidance. Ignore any table, column, alias, filter, or SQL fragment from the reasoning plan that is not explicitly present in the DATABASE SCHEMA. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. -6. YOU MUST ONLY use table names and column names that are explicitly present in the ACTIVE DATASOURCE METADATA, DATABASE SCHEMA, or VALID TABLE NAMES sections. -7. SQL SAMPLES are examples of style only. NEVER reuse a sample table or column name unless that exact table or column also appears in the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES sections. -8. NEVER invent generic table names from the user's business terms unless that exact table name is present in the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES sections. -9. Map business concepts to the closest explicit tables, columns, metrics, views, and relationships from the active metadata. Do not create a new table or column name from the business concept. -10. Before applying SUM, AVG, MIN, MAX, or arithmetic to a column, verify that the chosen column is numeric in the active metadata. Do not aggregate text/string columns as numeric values. -11. Do not prefix table names with catalog or schema names unless the active metadata, DATABASE SCHEMA, or VALID TABLE NAMES section shows the table name with that exact prefix. {text_to_sql_rules} @@ -2230,45 +595,6 @@ class SqlGenerationResult(BaseModel): } -def get_sql_generation_model_kwargs(llm_provider: LLMProvider) -> dict: - model_kwargs = llm_provider.get_model_kwargs() or {} - response_format = model_kwargs.get("response_format", {}) - - if isinstance(response_format, dict) and response_format.get("type") == "text": - return {} - - return SQL_GENERATION_MODEL_KWARGS - - -_SCHEMA_IDENTIFIER_PATTERN = ( - r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)' -) -_SCHEMA_TABLE_REFERENCE_PATTERN = ( - rf"{_SCHEMA_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SCHEMA_IDENTIFIER_PATTERN})*" -) -_SEMANTIC_TABLE_NAME_KEYS = ( - "name", - "referenceName", - "sourceTableName", - "tableName", - "table", -) -_SEMANTIC_COLUMN_CONTAINER_KEYS = ( - "columns", - "fields", - "calculatedFields", - "dimensions", - "measures", -) -_SEMANTIC_COLUMN_NAME_KEYS = ( - "name", - "referenceName", - "sourceColumnName", - "columnName", - "fieldName", -) - - def construct_instructions( instructions: list[dict] | None = None, ): @@ -2281,855 +607,6 @@ def construct_instructions( return _instructions -def _parse_semantic_metadata_content(content: str) -> Any | None: - content = content.strip() - if not content: - return None - - if content.startswith("```"): - content = re.sub(r"^```(?:json|mdl)?\s*", "", content, flags=re.IGNORECASE) - content = re.sub(r"\s*```$", "", content) - - candidates = [content] - object_start = content.find("{") - object_end = content.rfind("}") - if object_start >= 0 and object_end > object_start: - candidates.append(content[object_start : object_end + 1]) - array_start = content.find("[") - array_end = content.rfind("]") - if array_start >= 0 and array_end > array_start: - candidates.append(content[array_start : array_end + 1]) - - for candidate in candidates: - try: - return orjson.loads(candidate) - except orjson.JSONDecodeError: - continue - - return None - - -def _semantic_name_values(metadata: dict[str, Any], keys: tuple[str, ...]) -> list[str]: - values = [] - for key in keys: - value = metadata.get(key) - if isinstance(value, str) and value.strip(): - values.append(value.strip()) - return values - - -def _semantic_column_names(metadata: Any) -> set[str]: - columns: set[str] = set() - if isinstance(metadata, dict): - for column_name in _semantic_name_values(metadata, _SEMANTIC_COLUMN_NAME_KEYS): - columns.add(column_name) - for key in _SEMANTIC_COLUMN_CONTAINER_KEYS: - value = metadata.get(key) - if value is not None: - columns.update(_semantic_column_names(value)) - elif isinstance(metadata, list): - for item in metadata: - columns.update(_semantic_column_names(item)) - - return columns - - -def _construct_semantic_table_columns(content: str) -> dict[str, set[str]]: - parsed_content = _parse_semantic_metadata_content(content) - if parsed_content is None: - return {} - - table_columns: dict[str, set[str]] = {} - - def collect(metadata: Any) -> None: - if isinstance(metadata, list): - for item in metadata: - collect(item) - return - - if not isinstance(metadata, dict): - return - - column_containers = [ - metadata.get(key) - for key in _SEMANTIC_COLUMN_CONTAINER_KEYS - if metadata.get(key) is not None - ] - if column_containers: - columns = set() - for container in column_containers: - columns.update(_semantic_column_names(container)) - - if columns: - for table_reference in _semantic_name_values( - metadata, _SEMANTIC_TABLE_NAME_KEYS - ): - for table_name in _table_reference_suffixes(table_reference): - table_columns.setdefault(table_name, set()).update(columns) - - for value in metadata.values(): - collect(value) - - collect(parsed_content) - return table_columns - - -def construct_valid_table_names(documents: list[Any] | None = None) -> list[str]: - table_names: set[str] = set() - for document in documents or []: - content = getattr(document, "content", document) - if not isinstance(content, str): - continue - - for table_name in _construct_semantic_table_columns(content): - table_names.add(table_name) - - for match in re.finditer( - rf"\bCREATE\s+TABLE\s+(?P
{_SCHEMA_TABLE_REFERENCE_PATTERN})", - content, - flags=re.IGNORECASE, - ): - for table_name in _table_reference_suffixes(match.group("table")): - table_names.add(table_name) - - for table_reference in extract_sql_table_references(content): - for table_name in _table_reference_suffixes(table_reference): - table_names.add(table_name) - - return sorted(table_names) - - -def construct_valid_table_columns( - documents: list[Any] | None = None, -) -> dict[str, list[str]]: - table_columns: dict[str, set[str]] = {} - for document in documents or []: - content = getattr(document, "content", document) - if not isinstance(content, str): - continue - - for table_name, columns in _construct_semantic_table_columns( - content - ).items(): - table_columns.setdefault(table_name, set()).update(columns) - - for table_match in re.finditer( - rf"\bCREATE\s+TABLE\s+(?P
{_SCHEMA_TABLE_REFERENCE_PATTERN})\s*\(", - content, - flags=re.IGNORECASE, - ): - table_names = _table_reference_suffixes(table_match.group("table")) - body_start = table_match.end() - depth = 1 - body_end = body_start - while body_end < len(content) and depth > 0: - char = content[body_end] - if char == "(": - depth += 1 - elif char == ")": - depth -= 1 - body_end += 1 - - table_body = content[body_start : body_end - 1] - columns_by_table = [ - table_columns.setdefault(table_name, set()) - for table_name in table_names - ] - for line in table_body.splitlines(): - line = line.strip() - if not line or line.startswith("--"): - continue - line = line.rstrip(",") - if re.match( - r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", - line, - flags=re.IGNORECASE, - ): - continue - - column_match = re.match( - r"([`\"\[]?)(?P[A-Za-z_][A-Za-z0-9_$]*)\1\s+", - line, - ) - if column_match: - for columns in columns_by_table: - columns.add(column_match.group("column")) - - return { - table_name: sorted(columns) - for table_name, columns in sorted(table_columns.items()) - } - - -_SQL_IDENTIFIER_PATTERN = ( - r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)' -) -_SQL_TABLE_REFERENCE_PATTERN = re.compile( - rf"\b(?:FROM|JOIN)\s+" - rf"(?P
{_SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SQL_IDENTIFIER_PATTERN})*)", - flags=re.IGNORECASE, -) -_SQL_TABLE_WITH_ALIAS_PATTERN = re.compile( - rf"\b(?:FROM|JOIN)\s+" - rf"(?P
{_SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*{_SQL_IDENTIFIER_PATTERN})*)" - rf"(?:\s+(?:AS\s+)?(?P{_SQL_IDENTIFIER_PATTERN}))?", - flags=re.IGNORECASE, -) -_SQL_CTE_PATTERN = re.compile( - rf"(?:\bWITH\b|,)\s*(?P{_SQL_IDENTIFIER_PATTERN})\s+AS\s*\(", - flags=re.IGNORECASE, -) -_SQL_QUALIFIED_COLUMN_PATTERN = re.compile( - rf"(?P{_SQL_IDENTIFIER_PATTERN})\s*\.\s*(?P{_SQL_IDENTIFIER_PATTERN})", - flags=re.IGNORECASE, -) -_SQL_RESERVED_ALIASES = { - "where", - "join", - "left", - "right", - "inner", - "outer", - "full", - "cross", - "on", - "group", - "order", - "having", - "limit", - "union", -} -_SQL_NON_COLUMN_IDENTIFIERS = { - *_SQL_RESERVED_ALIASES, - "and", - "as", - "asc", - "between", - "by", - "case", - "cast", - "count", - "date", - "dateadd", - "datediff", - "datepart", - "day", - "desc", - "distinct", - "else", - "end", - "false", - "from", - "getdate", - "hour", - "in", - "is", - "like", - "max", - "min", - "month", - "not", - "null", - "or", - "select", - "sum", - "then", - "top", - "true", - "when", - "year", -} - - -def _normalize_sql_identifier(identifier: str) -> str: - identifier = identifier.strip() - if ( - (identifier.startswith('"') and identifier.endswith('"')) - or (identifier.startswith("`") and identifier.endswith("`")) - or (identifier.startswith("[") and identifier.endswith("]")) - ): - return identifier[1:-1] - return identifier - - -def _quote_sql_identifier(identifier: str) -> str: - return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' - - -def _compact_sql_identifier(identifier: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(identifier or "").lower()) - - -def _strip_sql_literals(sql: str) -> str: - without_strings = re.sub(r"'(?:''|[^'])*'", " ", sql) - return re.sub(r"\b\d+(?:\.\d+)?\b", " ", without_strings) - - -def _extract_clause_bodies( - sql: str, clause: str, terminators: list[str] -) -> list[str]: - terminator_pattern = "|".join(rf"\b{terminator}\b" for terminator in terminators) - pattern = re.compile( - rf"\b{clause}\b(?P.*?)(?={terminator_pattern}|$)", - flags=re.IGNORECASE | re.DOTALL, - ) - return [match.group("body") or "" for match in pattern.finditer(sql)] - - -def _find_unqualified_column_candidates(sql: str) -> set[str]: - bodies = [ - *_extract_clause_bodies( - sql, - r"WHERE", - [r"GROUP\s+BY", r"ORDER\s+BY", "HAVING", "LIMIT", "FETCH", "UNION"], - ), - *_extract_clause_bodies( - sql, - r"HAVING", - [r"GROUP\s+BY", r"ORDER\s+BY", "LIMIT", "FETCH", "UNION"], - ), - *_extract_clause_bodies( - sql, - r"ON", - [ - "WHERE", - r"GROUP\s+BY", - r"ORDER\s+BY", - "HAVING", - "JOIN", - "LIMIT", - "FETCH", - "UNION", - ], - ), - *_extract_clause_bodies( - sql, - r"GROUP\s+BY", - [r"ORDER\s+BY", "HAVING", "LIMIT", "FETCH", "UNION"], - ), - ] - candidates: set[str] = set() - identifier_pattern = re.compile(_SQL_IDENTIFIER_PATTERN, flags=re.IGNORECASE) - - for body in bodies: - searchable_body = _strip_sql_literals(body) - for match in identifier_pattern.finditer(searchable_body): - token = match.group(0) - before = searchable_body[: match.start()].rstrip() - after = searchable_body[match.end() :].lstrip() - identifier = _normalize_sql_identifier(token) - normalized = identifier.lower() - if ( - not identifier - or normalized in _SQL_NON_COLUMN_IDENTIFIERS - or before.endswith(".") - or after.startswith(".") - or after.startswith("(") - ): - continue - candidates.add(identifier) - - function_argument_pattern = re.compile( - rf"\b[A-Za-z_][A-Za-z0-9_$]*\s*\(\s*" - rf"(?:DISTINCT\s+)?(?P{_SQL_IDENTIFIER_PATTERN})" - rf"(?:\s*\.\s*(?P{_SQL_IDENTIFIER_PATTERN}))?", - flags=re.IGNORECASE, - ) - for match in function_argument_pattern.finditer(_strip_sql_literals(sql)): - column = _normalize_sql_identifier(match.group("column") or match.group("arg")) - if column and column != "*" and column.lower() not in _SQL_NON_COLUMN_IDENTIFIERS: - candidates.add(column) - - return candidates - - -def _sql_identifier_alias_candidates(identifier: str) -> set[str]: - normalized = _normalize_sql_identifier(identifier) - candidates = {normalized} - acronym_split = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", normalized) - snake_case = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", acronym_split) - candidates.add(snake_case) - return {candidate for candidate in candidates if candidate} - - -_SEMANTIC_COLUMN_ALIASES: dict[str, tuple[str, ...]] = {} - - -def _find_semantic_column_alias( - requested_column: str, - canonical_columns: dict[str, str], -) -> str | None: - for alias in _SEMANTIC_COLUMN_ALIASES.get( - _compact_sql_identifier(requested_column), () - ): - canonical = canonical_columns.get(_compact_sql_identifier(alias)) - if canonical: - return canonical - return None - - -def _split_table_reference(table_reference: str) -> list[str]: - stripped = table_reference.strip() - if not stripped: - return [] - - is_multipart_quoted_reference = bool( - re.search(r'"\s*\.\s*"|\]\s*\.\s*\[|`\s*\.\s*`', stripped) - ) - if ( - not is_multipart_quoted_reference - and ( - (stripped.startswith('"') and stripped.endswith('"')) - or (stripped.startswith("[") and stripped.endswith("]")) - or (stripped.startswith("`") and stripped.endswith("`")) - ) - ): - normalized_identifier = _normalize_sql_identifier(stripped) - if "." in normalized_identifier: - return [ - part - for part in re.split(r"\s*\.\s*", normalized_identifier) - if part.strip() - ] - - return [ - _normalize_sql_identifier(part) - for part in re.split(r"\s*\.\s*", stripped) - if part.strip() - ] - - -def _table_reference_suffixes(table_reference: str) -> list[str]: - parts = _split_table_reference(table_reference) - return [".".join(parts[index:]) for index in range(len(parts))] - - -def _quote_table_reference(table_reference: str) -> str: - return ".".join( - _quote_sql_identifier(part) for part in _split_table_reference(table_reference) - ) - - -def _is_extract_argument_from_keyword(sql: str, from_keyword_start: int) -> bool: - prefix = sql[:from_keyword_start] - return bool(re.search(r"\bEXTRACT\s*\([^)]*$", prefix, flags=re.IGNORECASE)) - - -def extract_sql_table_references(sql: str) -> list[str]: - references = [] - for match in _SQL_TABLE_REFERENCE_PATTERN.finditer(sql): - if _is_extract_argument_from_keyword(sql, match.start()): - continue - table_reference = match.group("table") - if table_reference.startswith("("): - continue - references.append(".".join(_split_table_reference(table_reference))) - return references - - -def extract_cte_names(sql: str) -> set[str]: - return { - _normalize_sql_identifier(match.group("cte")).lower() - for match in _SQL_CTE_PATTERN.finditer(sql) - } - - -def find_invalid_table_references(sql: str, valid_table_names: list[str]) -> list[str]: - if not valid_table_names: - return [] - - valid_tables = { - str(table_name).lower() - for table_name in valid_table_names - if table_name is not None - } - if not valid_tables: - return [] - cte_names = extract_cte_names(sql) - invalid_references = [] - - for table_reference in extract_sql_table_references(sql): - normalized_reference = table_reference.lower() - reference_suffixes = { - suffix.lower() for suffix in _table_reference_suffixes(table_reference) - } - if ( - normalized_reference in valid_tables - or normalized_reference in cte_names - or reference_suffixes.intersection(valid_tables) - or reference_suffixes.intersection(cte_names) - ): - continue - invalid_references.append(table_reference) - - return sorted(set(invalid_references)) - - -def _find_schema_table_alias( - requested_table: str, valid_table_names: list[str] -) -> str | None: - requested_suffixes = _table_reference_suffixes(requested_table) - requested_candidates = [ - suffix for suffix in requested_suffixes if suffix and len(suffix) > 2 - ] - if not requested_candidates: - return None - - valid_candidates = [ - str(table_name) - for table_name in valid_table_names or [] - if table_name is not None and str(table_name).strip() - ] - valid_by_lower = {table.lower(): table for table in valid_candidates} - for candidate in requested_candidates: - exact = valid_by_lower.get(candidate.lower()) - if exact: - return exact - - scored: list[tuple[int, int, str]] = [] - for valid_table in valid_candidates: - valid_suffixes = _table_reference_suffixes(valid_table) - valid_keys = [valid_table, *valid_suffixes] - for requested_key in requested_candidates: - requested_compact = _compact_sql_identifier(requested_key) - if not requested_compact: - continue - for valid_key in valid_keys: - valid_compact = _compact_sql_identifier(valid_key) - if not valid_compact: - continue - score = 0 - if requested_compact == valid_compact: - score = 1000 + len(valid_compact) - elif valid_compact.endswith(requested_compact): - score = 800 + len(requested_compact) - elif requested_compact.endswith(valid_compact): - score = 700 + len(valid_compact) - elif ( - len(requested_compact) >= 6 - and valid_compact.startswith(requested_compact) - ): - score = 600 + len(requested_compact) - elif ( - len(valid_compact) >= 6 - and requested_compact.startswith(valid_compact) - ): - score = 500 + len(valid_compact) - if score: - scored.append((score, len(valid_compact), valid_table)) - - if not scored: - return None - - scored.sort(reverse=True) - best_score = scored[0][0] - best_tables = {table for score, _, table in scored if score == best_score} - if len(best_tables) != 1: - return None - return scored[0][2] - - -def normalize_sql_table_references_to_schema( - sql: str, valid_table_names: list[str] -) -> str: - if not sql or not valid_table_names: - return sql - - replacements: dict[str, str] = {} - for table_reference in extract_sql_table_references(sql): - canonical_table = _find_schema_table_alias(table_reference, valid_table_names) - if not canonical_table or canonical_table == table_reference: - continue - replacements[table_reference] = canonical_table - - if not replacements: - return sql - - normalized_sql = sql - for requested_table, canonical_table in sorted( - replacements.items(), key=lambda item: len(item[0]), reverse=True - ): - requested_parts = _split_table_reference(requested_table) - if not requested_parts: - continue - quoted_requested = r"\s*\.\s*".join( - re.escape(_quote_sql_identifier(part)) for part in requested_parts - ) - single_quoted_requested = re.escape(_quote_sql_identifier(requested_table)) - bracketed_requested = re.escape(f"[{requested_table}]") - backticked_requested = re.escape(f"`{requested_table}`") - bare_requested = r"\s*\.\s*".join( - re.escape(part) for part in requested_parts - ) - table_pattern = re.compile( - rf"(? dict[str, str]: - valid_tables = { - str(table_name).lower(): table_name - for table_name in valid_table_columns - if table_name is not None - } - cte_names = extract_cte_names(sql) - aliases: dict[str, str] = {} - - for table_name in valid_table_columns: - if table_name is None: - continue - aliases[str(table_name).lower()] = table_name - - for match in _SQL_TABLE_WITH_ALIAS_PATTERN.finditer(sql): - table_reference = ".".join(_split_table_reference(match.group("table"))) - normalized_table = table_reference.lower() - matched_table = valid_tables.get(normalized_table) - if not matched_table: - for suffix in _table_reference_suffixes(table_reference): - matched_table = valid_tables.get(suffix.lower()) - if matched_table: - break - if not matched_table or normalized_table in cte_names: - continue - - alias = match.group("alias") - if not alias: - continue - - normalized_alias = _normalize_sql_identifier(alias or "").lower() - if normalized_alias in _SQL_RESERVED_ALIASES: - continue - aliases[normalized_alias] = matched_table - - return aliases - - -def normalize_sql_column_references_to_schema( - sql: str, valid_table_columns: dict[str, list[str]] -) -> str: - if not valid_table_columns: - return sql - - aliases = _extract_table_aliases(sql, valid_table_columns) - if not aliases: - return sql - - canonical_columns_by_table: dict[str, dict[str, str]] = {} - for table_name, columns in valid_table_columns.items(): - if table_name is None: - continue - compact_columns = { - _compact_sql_identifier(column): str(column) - for column in columns - if column is not None - } - compact_columns = { - compact: column - for compact, column in compact_columns.items() - if compact - } - canonical_columns_by_table[str(table_name)] = compact_columns - - def replace_column_reference(match: re.Match[str]) -> str: - qualifier = match.group("qualifier") - column = match.group("column") - normalized_qualifier = _normalize_sql_identifier(qualifier).lower() - table_name = aliases.get(normalized_qualifier) - if not table_name: - return match.group(0) - - canonical_columns = canonical_columns_by_table.get(str(table_name), {}) - normalized_column = _normalize_sql_identifier(column) - compact_column = _compact_sql_identifier(normalized_column) - canonical_column = canonical_columns.get( - compact_column - ) or _find_semantic_column_alias(normalized_column, canonical_columns) - if not canonical_column or canonical_column == normalized_column: - return match.group(0) - - return f"{qualifier}.{_quote_sql_identifier(canonical_column)}" - - normalized_sql = _SQL_QUALIFIED_COLUMN_PATTERN.sub(replace_column_reference, sql) - - referenced_tables = { - aliases.get(table_reference.lower()) - for table_reference in extract_sql_table_references(normalized_sql) - } - referenced_tables = {table for table in referenced_tables if table} - if len(referenced_tables) != 1: - return normalized_sql - - table_name = next(iter(referenced_tables)) - canonical_columns = canonical_columns_by_table.get(str(table_name), {}) - if not canonical_columns: - return normalized_sql - - valid_compact_columns = set(canonical_columns) - - def replace_unqualified_identifier(match: re.Match[str]) -> str: - identifier = next( - value - for value in ( - match.group("quoted"), - match.group("bracketed"), - match.group("bare"), - ) - if value - ) - compact_identifier = _compact_sql_identifier(identifier) - canonical_column = canonical_columns.get( - compact_identifier - ) or _find_semantic_column_alias(identifier, canonical_columns) - if not canonical_column: - return match.group(0) - - if canonical_column == identifier: - return match.group(0) - - return _quote_sql_identifier(canonical_column) - - schema_candidates = { - candidate - for column in canonical_columns.values() - if column - and _normalize_sql_identifier(column).lower() not in _SQL_RESERVED_ALIASES - for candidate in _sql_identifier_alias_candidates(column) - } - unqualified_candidates = sorted( - schema_candidates - | { - alias_key - for alias_key in _SEMANTIC_COLUMN_ALIASES - if alias_key - and _normalize_sql_identifier(alias_key).lower() - not in _SQL_RESERVED_ALIASES - } - | { - alias - for aliases in _SEMANTIC_COLUMN_ALIASES.values() - for alias in aliases - if alias - and _normalize_sql_identifier(alias).lower() not in _SQL_RESERVED_ALIASES - }, - key=len, - reverse=True, - ) - if not unqualified_candidates: - return normalized_sql - - candidate_pattern = "|".join( - re.escape(candidate) for candidate in unqualified_candidates - ) - unqualified_identifier_pattern = re.compile( - rf'(?{candidate_pattern})"' - rf"|\[(?P{candidate_pattern})\]" - rf"|(?P{candidate_pattern}))(?!\w)", - flags=re.IGNORECASE, - ) - return unqualified_identifier_pattern.sub( - replace_unqualified_identifier, - normalized_sql, - ) - - -def find_invalid_column_references( - sql: str, valid_table_columns: dict[str, list[str]] -) -> list[str]: - if not valid_table_columns: - return [] - - aliases = _extract_table_aliases(sql, valid_table_columns) - cte_names = extract_cte_names(sql) - invalid_references = [] - - for match in _SQL_QUALIFIED_COLUMN_PATTERN.finditer(sql): - qualifier = _normalize_sql_identifier(match.group("qualifier")) - column = _normalize_sql_identifier(match.group("column")) - normalized_qualifier = qualifier.lower() - - if normalized_qualifier in cte_names: - continue - - table_name = aliases.get(normalized_qualifier) - if not table_name: - continue - - valid_columns = { - str(col).lower() - for col in valid_table_columns.get(table_name, []) - if col is not None - } - if column.lower() not in valid_columns: - invalid_references.append(f"{qualifier}.{column}") - - referenced_tables = { - aliases.get(table_reference.lower()) - for table_reference in extract_sql_table_references(sql) - } - referenced_tables = {table for table in referenced_tables if table} - if len(referenced_tables) == 1: - table_name = next(iter(referenced_tables)) - valid_columns = { - str(col).lower() - for col in valid_table_columns.get(table_name, []) - if col is not None - } - valid_compact_columns = { - _compact_sql_identifier(col) - for col in valid_table_columns.get(table_name, []) - if col is not None - } - for start, end in _find_select_list_spans(sql): - for item in _split_top_level_select_items(sql[start:end]): - expression = _strip_projection_alias( - re.sub(r"^\s*DISTINCT\s+", "", item, flags=re.IGNORECASE) - ) - if not re.fullmatch(_SQL_IDENTIFIER_PATTERN, expression.strip()): - continue - column = _normalize_sql_identifier(expression) - if column == "*": - continue - if ( - column.lower() not in valid_columns - and _compact_sql_identifier(column) not in valid_compact_columns - ): - invalid_references.append(column) - table_aliases = { - alias.lower() - for alias, alias_table in aliases.items() - if alias_table == table_name - } - for column in _find_unqualified_column_candidates(sql): - normalized_column = column.lower() - if ( - normalized_column in table_aliases - or normalized_column in valid_columns - or _compact_sql_identifier(column) in valid_compact_columns - ): - continue - invalid_references.append(column) - - return sorted(set(invalid_references)) - - -def format_valid_table_columns(valid_table_columns: dict[str, list[str]]) -> str: - return "; ".join( - f"{table}: {', '.join(columns)}" - for table, columns in sorted(valid_table_columns.items()) - ) - - def construct_ask_history_messages( histories: list[Any] | list[dict], ) -> list[ChatMessage]: From efc9ed651d1669b7663f4c2a28d8dbc2686ac8af Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 02:01:52 +0530 Subject: [PATCH 0610/1087] Fix SQL generation correction flow --- .../generation/followup_sql_generation.py | 3 - .../pipelines/generation/sql_correction.py | 3 - .../pipelines/generation/sql_regeneration.py | 3 - .../retrieval/db_schema_retrieval.py | 13 +-- .../retrieval/test_db_schema_retrieval.py | 15 ++-- .../apollo/server/services/askingService.ts | 52 +++++++++--- .../services/tests/askingService.test.ts | 80 +++++++++++++++++++ 7 files changed, 127 insertions(+), 42 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 c9f8b23537..0f5da7a28d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -17,7 +17,6 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, - construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, @@ -196,8 +195,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index daae37e02e..5c7452870d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,7 +15,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_valid_table_columns, construct_valid_table_names, get_sql_generation_model_kwargs, get_text_to_sql_rules, @@ -179,8 +178,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 8b3eaafcb1..74afbb276a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,7 +14,6 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, - construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, @@ -182,8 +181,6 @@ async def post_process( regenerate_sql.get("replies"), project_id=project_id, data_source=data_source, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ceeedeb70d..3879e3b63e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -240,8 +240,6 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: @observe(capture_input=False) def check_using_db_schemas_without_pruning( - query: str, - tables: list[str] | None, construct_db_schemas: list[dict], dbschema_retrieval: list[Document], encoding: tiktoken.Encoding, @@ -290,16 +288,7 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) - should_select_tables_for_question = ( - bool((query or "").strip()) - and not tables - and len(retrieval_results) > 1 - ) - if ( - _token_count > context_window_size - or enable_column_pruning - or should_select_tables_for_question - ): + if _token_count > context_window_size or enable_column_pruning: return { "db_schemas": [], "tokens": _token_count, diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index fe93aa8779..17dac743c6 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -129,8 +129,6 @@ def encode(self, value): return value.split() result = check_using_db_schemas_without_pruning( - query="show top customers by invoice amount", - tables=None, construct_db_schemas=[ { "type": "TABLE", @@ -159,7 +157,7 @@ def encode(self, value): assert result["tokens"] > 0 -def test_check_using_db_schemas_without_pruning_selects_tables_for_question(): +def test_check_using_db_schemas_without_pruning_keeps_semantic_retrieval_context(): class Encoding: def encode(self, value): return value.split() @@ -183,8 +181,6 @@ def table_schema(name): } result = check_using_db_schemas_without_pruning( - query="compare recent activity by account", - tables=None, construct_db_schemas=[ table_schema("activity"), table_schema("account"), @@ -195,18 +191,19 @@ def table_schema(name): context_window_size=1000, ) - assert result["db_schemas"] == [] + assert [schema["table_name"] for schema in result["db_schemas"]] == [ + "activity", + "account", + ] assert result["tokens"] > 0 -def test_check_using_db_schemas_without_pruning_keeps_explicit_table_fast_path(): +def test_check_using_db_schemas_without_pruning_keeps_selected_table_context(): class Encoding: def encode(self, value): return value.split() result = check_using_db_schemas_without_pruning( - query="show records from activity", - tables=["activity"], construct_db_schemas=[ { "type": "TABLE", diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index f5b1115cae..55718b1174 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1,6 +1,8 @@ import { IWrenAIAdaptor } from '@server/adaptors/wrenAIAdaptor'; import { AskResultStatus, + AskResultType, + AskCandidateType, RecommendationQuestionsResult, RecommendationQuestionsInput, RecommendationQuestion, @@ -17,7 +19,6 @@ import { IThreadResponseRepository, ThreadResponse, ThreadResponseAnswerDetail, - ThreadResponseAdjustmentType, } from '../repositories/threadResponseRepository'; import { getLogger } from '@server/utils'; import { isEmpty, isNil } from 'lodash'; @@ -1409,18 +1410,45 @@ export class AskingService implements IAskingService { } await this.ensureThreadInCurrentProject(response.threadId); - return await this.threadResponseRepository.createOne({ - sql: input.sql, - threadId: response.threadId, - question: response.question, - adjustment: { - type: ThreadResponseAdjustmentType.APPLY_SQL, - payload: { - originalThreadResponseId: response.id, - sql: input.sql, - }, - }, + const project = await this.getProjectForThreadResponse(response); + const deployment = await this.deployService.getLastDeployment(project.id); + await this.queryService.preview(input.sql, { + project, + manifest: deployment.manifest, + modelingOnly: false, + limit: 1, + cacheEnabled: false, }); + + const updatedResponse = await this.threadResponseRepository.updateOne( + response.id, + { + sql: input.sql, + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, + }, + ); + + if (response.askingTaskId) { + await this.askingTaskRepository.updateOne(response.askingTaskId, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FINISHED, + response: [ + { + type: AskCandidateType.LLM, + sql: input.sql, + }, + ], + error: null, + invalidSql: null, + }, + }); + } + + return updatedResponse; } public async adjustThreadResponseAnswer( diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index d33c9e166d..aca8d01df5 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -231,4 +231,84 @@ describe('AskingService', () => { }); }); }); + + describe('SQL correction application', () => { + test('updates the original response and clears stale downstream state', async () => { + const service = Object.create(AskingService.prototype) as any; + service.threadResponseRepository = { + findOneBy: jest.fn().mockResolvedValue({ + id: 11, + threadId: 7, + askingTaskId: 42, + question: 'Show the corrected result', + sql: 'SELECT 0', + answerDetail: { status: 'FINISHED', content: 'stale answer' }, + chartDetail: { status: 'FINISHED', chartSchema: { stale: true } }, + }), + updateOne: jest.fn().mockResolvedValue({ + id: 11, + threadId: 7, + askingTaskId: 42, + question: 'Show the corrected result', + sql: 'SELECT 1', + }), + createOne: jest.fn(), + }; + service.threadRepository = { + findOneBy: jest.fn().mockResolvedValue({ id: 7, projectId: 1 }), + }; + service.projectService = { + getCurrentProject: jest.fn().mockResolvedValue({ id: 1 }), + getProjectById: jest.fn().mockResolvedValue({ id: 1 }), + }; + service.deployService = { + getLastDeployment: jest.fn().mockResolvedValue({ manifest: {} }), + }; + service.queryService = { + preview: jest.fn().mockResolvedValue({ columns: [], data: [] }), + }; + service.askingTaskRepository = { + updateOne: jest.fn().mockResolvedValue(undefined), + }; + + const response = await service.adjustThreadResponseWithSQL(11, { + sql: 'SELECT 1', + }); + + expect(response.sql).toBe('SELECT 1'); + expect(service.threadResponseRepository.createOne).not.toHaveBeenCalled(); + expect(service.queryService.preview).toHaveBeenCalledWith('SELECT 1', { + project: { id: 1 }, + manifest: {}, + modelingOnly: false, + limit: 1, + cacheEnabled: false, + }); + expect(service.threadResponseRepository.updateOne).toHaveBeenCalledWith( + 11, + expect.objectContaining({ + sql: 'SELECT 1', + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, + }), + ); + expect(service.askingTaskRepository.updateOne).toHaveBeenCalledWith( + 42, + expect.objectContaining({ + detail: expect.objectContaining({ + status: 'FINISHED', + response: [ + { + type: 'LLM', + sql: 'SELECT 1', + }, + ], + invalidSql: null, + }), + }), + ); + }); + }); }); From a8449782bcc390de42c06b238846ce0d5d1b50d8 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 02:56:51 +0530 Subject: [PATCH 0611/1087] Fix backend SQL flow imports --- .../generation/followup_sql_generation.py | 6 ++-- .../pipelines/generation/sql_correction.py | 6 ++-- .../pipelines/generation/sql_regeneration.py | 5 ++-- .../src/pipelines/generation/utils/sql.py | 10 +++++-- wren-ai-service/src/web/v1/services/ask.py | 17 ----------- wren-ai-service/src/web/v1/services/chart.py | 30 +------------------ .../src/web/v1/services/sql_answer.py | 24 +-------------- 7 files changed, 16 insertions(+), 82 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 0f5da7a28d..d73b71d7a3 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -14,14 +14,13 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, - construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, - get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -133,7 +132,6 @@ def prompt( query=query, data_source=data_source, documents=documents, - valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -216,7 +214,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_generation_system_prompt(None), - generation_kwargs=get_sql_generation_model_kwargs(llm_provider), + generation_kwargs=SQL_GENERATION_MODEL_KWARGS, ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 5c7452870d..0fef7ec436 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -13,10 +13,9 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, - construct_valid_table_names, - get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -134,7 +133,6 @@ def prompt( query=query, data_source=data_source, documents=documents, - valid_table_names=construct_valid_table_names(documents), invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, @@ -199,7 +197,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_correction_system_prompt(None), - generation_kwargs=get_sql_generation_model_kwargs(llm_provider), + generation_kwargs=SQL_GENERATION_MODEL_KWARGS, ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 74afbb276a..498e718f61 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -12,13 +12,12 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, - construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, - get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -197,7 +196,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_regeneration_system_prompt(None), - generation_kwargs=get_sql_generation_model_kwargs(llm_provider), + generation_kwargs=SQL_GENERATION_MODEL_KWARGS, ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 45d3011bf4..5b1004ff20 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -515,7 +515,10 @@ def _extract_from_sql_knowledge( return value if value and value.strip() else default_value -def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: +def get_text_to_sql_rules( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: if sql_knowledge is not None: return _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES @@ -535,7 +538,10 @@ def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) return _DEFAULT_CALCULATED_FIELD_INSTRUCTIONS -def get_metric_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: +def get_metric_instructions( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: if sql_knowledge is not None: return _extract_from_sql_knowledge( sql_knowledge, "metric_instructions", _DEFAULT_METRIC_INSTRUCTIONS diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cd3c965e47..95eac8c695 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -8,13 +8,6 @@ from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - normalize_sql_direction_keywords, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, -) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -872,16 +865,6 @@ def _build_validated_ask_result_from_sql( table_ddls: list[str], query: str | None = None, ) -> Optional[AskResult]: - if isinstance(sql, str): - sql = normalize_sql_direction_keywords(sql) - sql = normalize_sql_table_references_to_schema( - sql, - construct_valid_table_names(table_ddls), - ) - sql = normalize_sql_column_references_to_schema( - sql, - construct_valid_table_columns(table_ddls), - ) ask_result = self._build_ask_result_from_sql(sql) if not ask_result: return None diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index b9152f7d83..62586c9f1a 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,15 +6,6 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - extract_sql_table_references, - find_invalid_column_references, - find_invalid_table_references, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, -) from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest @@ -100,13 +91,8 @@ async def _load_active_schema_contexts( if not retrieval_pipeline: return [] - table_references = extract_sql_table_references(sql) - if not table_references: - return [] - retrieval_result = await retrieval_pipeline.run( query="", - tables=table_references, histories=[], project_id=project_id, enable_column_pruning=False, @@ -123,21 +109,7 @@ async def _load_active_schema_contexts( def _normalize_and_validate_sql( self, sql: str, schema_contexts: list[str] ) -> str | None: - valid_table_names = construct_valid_table_names(schema_contexts) - valid_table_columns = construct_valid_table_columns(schema_contexts) - normalized_sql = normalize_sql_table_references_to_schema( - sql, - valid_table_names, - ) - normalized_sql = normalize_sql_column_references_to_schema( - normalized_sql, - valid_table_columns, - ) - if find_invalid_table_references(normalized_sql, valid_table_names): - return None - if find_invalid_column_references(normalized_sql, valid_table_columns): - return None - return normalized_sql + return sql @observe(name="Generate Chart") @trace_metadata diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index 7442f2f884..dea25902a9 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -7,14 +7,6 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - find_invalid_column_references, - find_invalid_table_references, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, -) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -90,21 +82,7 @@ async def _load_active_schema_contexts( def _normalize_and_validate_sql( self, sql: str, schema_contexts: list[str] ) -> str | None: - valid_table_names = construct_valid_table_names(schema_contexts) - valid_table_columns = construct_valid_table_columns(schema_contexts) - normalized_sql = normalize_sql_table_references_to_schema( - sql, - valid_table_names, - ) - normalized_sql = normalize_sql_column_references_to_schema( - normalized_sql, - valid_table_columns, - ) - if find_invalid_table_references(normalized_sql, valid_table_names): - return None - if find_invalid_column_references(normalized_sql, valid_table_columns): - return None - return normalized_sql + return sql @observe(name="SQL Answer") @trace_metadata From d93a7c4fcd3c2a835007c2039824eb88f9956e4e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:08:15 +0530 Subject: [PATCH 0612/1087] Accept SQL generation table limit config --- wren-ai-service/src/web/v1/services/ask.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 95eac8c695..2a7e82c53b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -119,6 +119,7 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, + max_sql_generation_tables: int = 10, pipeline_timeout_seconds: int = 45, schema_retrieval_timeout_seconds: int = 25, max_histories: int = 5, @@ -142,6 +143,7 @@ def __init__( self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries + self._max_sql_generation_tables = max_sql_generation_tables def _is_stopped(self, query_id: str, container: dict): if ( @@ -1573,6 +1575,7 @@ async def ask( documents, table_names, table_ddls, + max_tables=self._max_sql_generation_tables, ) ( documents, From 0146f27987e13a411679f7b2226d38ffa7dc3472 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:17:27 +0530 Subject: [PATCH 0613/1087] Accept explicit table limit config --- wren-ai-service/src/web/v1/services/ask.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2a7e82c53b..9f228affac 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -120,6 +120,7 @@ def __init__( enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, max_sql_generation_tables: int = 10, + max_forced_explicit_tables: int = MAX_FORCED_EXPLICIT_TABLES, pipeline_timeout_seconds: int = 45, schema_retrieval_timeout_seconds: int = 25, max_histories: int = 5, @@ -144,6 +145,7 @@ def __init__( self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries self._max_sql_generation_tables = max_sql_generation_tables + self._max_forced_explicit_tables = max_forced_explicit_tables def _is_stopped(self, query_id: str, container: dict): if ( @@ -352,7 +354,7 @@ def _forced_explicit_table_names( ) -> list[str]: if not table_names: return [] - if len(table_names) <= MAX_FORCED_EXPLICIT_TABLES: + if len(table_names) <= self._max_forced_explicit_tables: return table_names logger.info( From 9e199ddf8f2bf795a6b08dc619a5d01c4a2175cb Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:34:53 +0530 Subject: [PATCH 0614/1087] Tolerate extra AskService config keys --- wren-ai-service/src/web/v1/services/ask.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9f228affac..960d177bf9 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -126,7 +126,13 @@ def __init__( max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, + **kwargs: Any, ): + if kwargs: + logger.info( + "Ignoring unsupported AskService config keys: %s", + sorted(kwargs), + ) self._pipelines = pipelines self._ask_results: Dict[str, AskResultResponse] = TTLCache( maxsize=maxsize, ttl=ttl From 2f39c29f36bcafb96855ba9672e80bcfc5bb9917 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:47:50 +0530 Subject: [PATCH 0615/1087] Revert "Tolerate extra AskService config keys" This reverts commit 9e199ddf8f2bf795a6b08dc619a5d01c4a2175cb. --- wren-ai-service/src/web/v1/services/ask.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 960d177bf9..9f228affac 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -126,13 +126,7 @@ def __init__( max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, - **kwargs: Any, ): - if kwargs: - logger.info( - "Ignoring unsupported AskService config keys: %s", - sorted(kwargs), - ) self._pipelines = pipelines self._ask_results: Dict[str, AskResultResponse] = TTLCache( maxsize=maxsize, ttl=ttl From 0966f2c41433fd20e4127065c7b423182d5045b5 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:47:51 +0530 Subject: [PATCH 0616/1087] Revert "Accept explicit table limit config" This reverts commit 0146f27987e13a411679f7b2226d38ffa7dc3472. --- wren-ai-service/src/web/v1/services/ask.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 9f228affac..2a7e82c53b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -120,7 +120,6 @@ def __init__( enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, max_sql_generation_tables: int = 10, - max_forced_explicit_tables: int = MAX_FORCED_EXPLICIT_TABLES, pipeline_timeout_seconds: int = 45, schema_retrieval_timeout_seconds: int = 25, max_histories: int = 5, @@ -145,7 +144,6 @@ def __init__( self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries self._max_sql_generation_tables = max_sql_generation_tables - self._max_forced_explicit_tables = max_forced_explicit_tables def _is_stopped(self, query_id: str, container: dict): if ( @@ -354,7 +352,7 @@ def _forced_explicit_table_names( ) -> list[str]: if not table_names: return [] - if len(table_names) <= self._max_forced_explicit_tables: + if len(table_names) <= MAX_FORCED_EXPLICIT_TABLES: return table_names logger.info( From d5e7a9e49c63ab4c79134aed727dcd7d8f2a466f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:47:51 +0530 Subject: [PATCH 0617/1087] Revert "Accept SQL generation table limit config" This reverts commit d93a7c4fcd3c2a835007c2039824eb88f9956e4e. --- wren-ai-service/src/web/v1/services/ask.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2a7e82c53b..95eac8c695 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -119,7 +119,6 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, - max_sql_generation_tables: int = 10, pipeline_timeout_seconds: int = 45, schema_retrieval_timeout_seconds: int = 25, max_histories: int = 5, @@ -143,7 +142,6 @@ def __init__( self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries - self._max_sql_generation_tables = max_sql_generation_tables def _is_stopped(self, query_id: str, container: dict): if ( @@ -1575,7 +1573,6 @@ async def ask( documents, table_names, table_ddls, - max_tables=self._max_sql_generation_tables, ) ( documents, From 3063ad81c11f96ce46e36855bf650cafb71b6f72 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:47:51 +0530 Subject: [PATCH 0618/1087] Revert "Fix backend SQL flow imports" This reverts commit a8449782bcc390de42c06b238846ce0d5d1b50d8. --- .../generation/followup_sql_generation.py | 6 ++-- .../pipelines/generation/sql_correction.py | 6 ++-- .../pipelines/generation/sql_regeneration.py | 5 ++-- .../src/pipelines/generation/utils/sql.py | 10 ++----- wren-ai-service/src/web/v1/services/ask.py | 17 +++++++++++ wren-ai-service/src/web/v1/services/chart.py | 30 ++++++++++++++++++- .../src/web/v1/services/sql_answer.py | 24 ++++++++++++++- 7 files changed, 82 insertions(+), 16 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 d73b71d7a3..0f5da7a28d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -14,13 +14,14 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, + get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -132,6 +133,7 @@ def prompt( query=query, data_source=data_source, documents=documents, + valid_table_names=construct_valid_table_names(documents), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -214,7 +216,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_generation_system_prompt(None), - generation_kwargs=SQL_GENERATION_MODEL_KWARGS, + generation_kwargs=get_sql_generation_model_kwargs(llm_provider), ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 0fef7ec436..5c7452870d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -13,9 +13,10 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_valid_table_names, + get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -133,6 +134,7 @@ def prompt( query=query, data_source=data_source, documents=documents, + valid_table_names=construct_valid_table_names(documents), invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, @@ -197,7 +199,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_correction_system_prompt(None), - generation_kwargs=SQL_GENERATION_MODEL_KWARGS, + generation_kwargs=get_sql_generation_model_kwargs(llm_provider), ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 498e718f61..74afbb276a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -12,12 +12,13 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, + get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -196,7 +197,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_regeneration_system_prompt(None), - generation_kwargs=SQL_GENERATION_MODEL_KWARGS, + generation_kwargs=get_sql_generation_model_kwargs(llm_provider), ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5b1004ff20..45d3011bf4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -515,10 +515,7 @@ def _extract_from_sql_knowledge( return value if value and value.strip() else default_value -def get_text_to_sql_rules( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: +def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: if sql_knowledge is not None: return _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES @@ -538,10 +535,7 @@ def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) return _DEFAULT_CALCULATED_FIELD_INSTRUCTIONS -def get_metric_instructions( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: +def get_metric_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: if sql_knowledge is not None: return _extract_from_sql_knowledge( sql_knowledge, "metric_instructions", _DEFAULT_METRIC_INSTRUCTIONS diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 95eac8c695..cd3c965e47 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -8,6 +8,13 @@ from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + normalize_sql_direction_keywords, + normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, +) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -865,6 +872,16 @@ def _build_validated_ask_result_from_sql( table_ddls: list[str], query: str | None = None, ) -> Optional[AskResult]: + if isinstance(sql, str): + sql = normalize_sql_direction_keywords(sql) + sql = normalize_sql_table_references_to_schema( + sql, + construct_valid_table_names(table_ddls), + ) + sql = normalize_sql_column_references_to_schema( + sql, + construct_valid_table_columns(table_ddls), + ) ask_result = self._build_ask_result_from_sql(sql) if not ask_result: return None diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index 62586c9f1a..b9152f7d83 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,6 +6,15 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + extract_sql_table_references, + find_invalid_column_references, + find_invalid_table_references, + normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, +) from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest @@ -91,8 +100,13 @@ async def _load_active_schema_contexts( if not retrieval_pipeline: return [] + table_references = extract_sql_table_references(sql) + if not table_references: + return [] + retrieval_result = await retrieval_pipeline.run( query="", + tables=table_references, histories=[], project_id=project_id, enable_column_pruning=False, @@ -109,7 +123,21 @@ async def _load_active_schema_contexts( def _normalize_and_validate_sql( self, sql: str, schema_contexts: list[str] ) -> str | None: - return sql + valid_table_names = construct_valid_table_names(schema_contexts) + valid_table_columns = construct_valid_table_columns(schema_contexts) + normalized_sql = normalize_sql_table_references_to_schema( + sql, + valid_table_names, + ) + normalized_sql = normalize_sql_column_references_to_schema( + normalized_sql, + valid_table_columns, + ) + if find_invalid_table_references(normalized_sql, valid_table_names): + return None + if find_invalid_column_references(normalized_sql, valid_table_columns): + return None + return normalized_sql @observe(name="Generate Chart") @trace_metadata diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index dea25902a9..7442f2f884 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -7,6 +7,14 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.sql import ( + construct_valid_table_columns, + construct_valid_table_names, + find_invalid_column_references, + find_invalid_table_references, + normalize_sql_column_references_to_schema, + normalize_sql_table_references_to_schema, +) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -82,7 +90,21 @@ async def _load_active_schema_contexts( def _normalize_and_validate_sql( self, sql: str, schema_contexts: list[str] ) -> str | None: - return sql + valid_table_names = construct_valid_table_names(schema_contexts) + valid_table_columns = construct_valid_table_columns(schema_contexts) + normalized_sql = normalize_sql_table_references_to_schema( + sql, + valid_table_names, + ) + normalized_sql = normalize_sql_column_references_to_schema( + normalized_sql, + valid_table_columns, + ) + if find_invalid_table_references(normalized_sql, valid_table_names): + return None + if find_invalid_column_references(normalized_sql, valid_table_columns): + return None + return normalized_sql @observe(name="SQL Answer") @trace_metadata From ef389e1aca31d7b5238be3de55fdd4964608f9d3 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 03:47:51 +0530 Subject: [PATCH 0619/1087] Revert "Fix SQL generation correction flow" This reverts commit efc9ed651d1669b7663f4c2a28d8dbc2686ac8af. --- .../generation/followup_sql_generation.py | 3 + .../pipelines/generation/sql_correction.py | 3 + .../pipelines/generation/sql_regeneration.py | 3 + .../retrieval/db_schema_retrieval.py | 13 ++- .../retrieval/test_db_schema_retrieval.py | 15 ++-- .../apollo/server/services/askingService.ts | 52 +++--------- .../services/tests/askingService.test.ts | 80 ------------------- 7 files changed, 42 insertions(+), 127 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 0f5da7a28d..c9f8b23537 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -17,6 +17,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, @@ -195,6 +196,8 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 5c7452870d..daae37e02e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_columns, construct_valid_table_names, get_sql_generation_model_kwargs, get_text_to_sql_rules, @@ -178,6 +179,8 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 74afbb276a..8b3eaafcb1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, construct_instructions, + construct_valid_table_columns, construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, @@ -181,6 +182,8 @@ async def post_process( regenerate_sql.get("replies"), project_id=project_id, data_source=data_source, + valid_table_names=construct_valid_table_names(documents), + valid_table_columns=construct_valid_table_columns(documents), ) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 3879e3b63e..ceeedeb70d 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -240,6 +240,8 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: @observe(capture_input=False) def check_using_db_schemas_without_pruning( + query: str, + tables: list[str] | None, construct_db_schemas: list[dict], dbschema_retrieval: list[Document], encoding: tiktoken.Encoding, @@ -288,7 +290,16 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) - if _token_count > context_window_size or enable_column_pruning: + should_select_tables_for_question = ( + bool((query or "").strip()) + and not tables + and len(retrieval_results) > 1 + ) + if ( + _token_count > context_window_size + or enable_column_pruning + or should_select_tables_for_question + ): return { "db_schemas": [], "tokens": _token_count, diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 17dac743c6..fe93aa8779 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -129,6 +129,8 @@ def encode(self, value): return value.split() result = check_using_db_schemas_without_pruning( + query="show top customers by invoice amount", + tables=None, construct_db_schemas=[ { "type": "TABLE", @@ -157,7 +159,7 @@ def encode(self, value): assert result["tokens"] > 0 -def test_check_using_db_schemas_without_pruning_keeps_semantic_retrieval_context(): +def test_check_using_db_schemas_without_pruning_selects_tables_for_question(): class Encoding: def encode(self, value): return value.split() @@ -181,6 +183,8 @@ def table_schema(name): } result = check_using_db_schemas_without_pruning( + query="compare recent activity by account", + tables=None, construct_db_schemas=[ table_schema("activity"), table_schema("account"), @@ -191,19 +195,18 @@ def table_schema(name): context_window_size=1000, ) - assert [schema["table_name"] for schema in result["db_schemas"]] == [ - "activity", - "account", - ] + assert result["db_schemas"] == [] assert result["tokens"] > 0 -def test_check_using_db_schemas_without_pruning_keeps_selected_table_context(): +def test_check_using_db_schemas_without_pruning_keeps_explicit_table_fast_path(): class Encoding: def encode(self, value): return value.split() result = check_using_db_schemas_without_pruning( + query="show records from activity", + tables=["activity"], construct_db_schemas=[ { "type": "TABLE", diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 55718b1174..f5b1115cae 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1,8 +1,6 @@ import { IWrenAIAdaptor } from '@server/adaptors/wrenAIAdaptor'; import { AskResultStatus, - AskResultType, - AskCandidateType, RecommendationQuestionsResult, RecommendationQuestionsInput, RecommendationQuestion, @@ -19,6 +17,7 @@ import { IThreadResponseRepository, ThreadResponse, ThreadResponseAnswerDetail, + ThreadResponseAdjustmentType, } from '../repositories/threadResponseRepository'; import { getLogger } from '@server/utils'; import { isEmpty, isNil } from 'lodash'; @@ -1410,45 +1409,18 @@ export class AskingService implements IAskingService { } await this.ensureThreadInCurrentProject(response.threadId); - const project = await this.getProjectForThreadResponse(response); - const deployment = await this.deployService.getLastDeployment(project.id); - await this.queryService.preview(input.sql, { - project, - manifest: deployment.manifest, - modelingOnly: false, - limit: 1, - cacheEnabled: false, - }); - - const updatedResponse = await this.threadResponseRepository.updateOne( - response.id, - { - sql: input.sql, - viewId: null, - answerDetail: null, - breakdownDetail: null, - chartDetail: null, - }, - ); - - if (response.askingTaskId) { - await this.askingTaskRepository.updateOne(response.askingTaskId, { - detail: { - type: AskResultType.TEXT_TO_SQL, - status: AskResultStatus.FINISHED, - response: [ - { - type: AskCandidateType.LLM, - sql: input.sql, - }, - ], - error: null, - invalidSql: null, + return await this.threadResponseRepository.createOne({ + sql: input.sql, + threadId: response.threadId, + question: response.question, + adjustment: { + type: ThreadResponseAdjustmentType.APPLY_SQL, + payload: { + originalThreadResponseId: response.id, + sql: input.sql, }, - }); - } - - return updatedResponse; + }, + }); } public async adjustThreadResponseAnswer( diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index aca8d01df5..d33c9e166d 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -231,84 +231,4 @@ describe('AskingService', () => { }); }); }); - - describe('SQL correction application', () => { - test('updates the original response and clears stale downstream state', async () => { - const service = Object.create(AskingService.prototype) as any; - service.threadResponseRepository = { - findOneBy: jest.fn().mockResolvedValue({ - id: 11, - threadId: 7, - askingTaskId: 42, - question: 'Show the corrected result', - sql: 'SELECT 0', - answerDetail: { status: 'FINISHED', content: 'stale answer' }, - chartDetail: { status: 'FINISHED', chartSchema: { stale: true } }, - }), - updateOne: jest.fn().mockResolvedValue({ - id: 11, - threadId: 7, - askingTaskId: 42, - question: 'Show the corrected result', - sql: 'SELECT 1', - }), - createOne: jest.fn(), - }; - service.threadRepository = { - findOneBy: jest.fn().mockResolvedValue({ id: 7, projectId: 1 }), - }; - service.projectService = { - getCurrentProject: jest.fn().mockResolvedValue({ id: 1 }), - getProjectById: jest.fn().mockResolvedValue({ id: 1 }), - }; - service.deployService = { - getLastDeployment: jest.fn().mockResolvedValue({ manifest: {} }), - }; - service.queryService = { - preview: jest.fn().mockResolvedValue({ columns: [], data: [] }), - }; - service.askingTaskRepository = { - updateOne: jest.fn().mockResolvedValue(undefined), - }; - - const response = await service.adjustThreadResponseWithSQL(11, { - sql: 'SELECT 1', - }); - - expect(response.sql).toBe('SELECT 1'); - expect(service.threadResponseRepository.createOne).not.toHaveBeenCalled(); - expect(service.queryService.preview).toHaveBeenCalledWith('SELECT 1', { - project: { id: 1 }, - manifest: {}, - modelingOnly: false, - limit: 1, - cacheEnabled: false, - }); - expect(service.threadResponseRepository.updateOne).toHaveBeenCalledWith( - 11, - expect.objectContaining({ - sql: 'SELECT 1', - viewId: null, - answerDetail: null, - breakdownDetail: null, - chartDetail: null, - }), - ); - expect(service.askingTaskRepository.updateOne).toHaveBeenCalledWith( - 42, - expect.objectContaining({ - detail: expect.objectContaining({ - status: 'FINISHED', - response: [ - { - type: 'LLM', - sql: 'SELECT 1', - }, - ], - invalidSql: null, - }), - }), - ); - }); - }); }); From bbf0bcd9c6c49807809fc1fbfa79d50c40112ec6 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 04:21:36 +0530 Subject: [PATCH 0620/1087] Fix reverted SQL utility imports --- .../generation/followup_sql_generation.py | 14 +-- .../pipelines/generation/sql_correction.py | 15 +-- .../pipelines/generation/sql_regeneration.py | 17 +-- wren-ai-service/src/web/v1/services/ask.py | 104 +----------------- wren-ai-service/src/web/v1/services/chart.py | 30 +---- .../src/web/v1/services/sql_answer.py | 24 +--- 6 files changed, 15 insertions(+), 189 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 c9f8b23537..5f769b974a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -14,15 +14,13 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, - construct_valid_table_columns, - construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, - get_sql_generation_model_kwargs, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -134,7 +132,7 @@ def prompt( query=query, data_source=data_source, documents=documents, - valid_table_names=construct_valid_table_names(documents), + valid_table_names=[], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -145,9 +143,7 @@ def prompt( else "" ), metric_instructions=( - get_metric_instructions(sql_knowledge, data_source=data_source) - if has_metric - else "" + get_metric_instructions(sql_knowledge) if has_metric else "" ), json_field_instructions=( get_json_field_instructions(sql_knowledge) if has_json_field else "" @@ -196,8 +192,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) @@ -219,7 +213,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_generation_system_prompt(None), - generation_kwargs=get_sql_generation_model_kwargs(llm_provider), + generation_kwargs=SQL_GENERATION_MODEL_KWARGS, ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index daae37e02e..284cf09a6d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -13,11 +13,9 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, - construct_valid_table_columns, - construct_valid_table_names, - get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -31,10 +29,7 @@ def get_sql_correction_system_prompt( sql_knowledge: SqlKnowledge | None = None, data_source: str | None = None, ) -> str: - text_to_sql_rules = get_text_to_sql_rules( - sql_knowledge, - data_source=data_source, - ) + text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" ### TASK ### @@ -135,7 +130,7 @@ def prompt( query=query, data_source=data_source, documents=documents, - valid_table_names=construct_valid_table_names(documents), + valid_table_names=[], invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, @@ -179,8 +174,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) @@ -202,7 +195,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_correction_system_prompt(None), - generation_kwargs=get_sql_generation_model_kwargs(llm_provider), + generation_kwargs=SQL_GENERATION_MODEL_KWARGS, ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 8b3eaafcb1..a9b93bc942 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -12,14 +12,12 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, - construct_valid_table_columns, - construct_valid_table_names, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, - get_sql_generation_model_kwargs, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -33,10 +31,7 @@ def get_sql_regeneration_system_prompt( sql_knowledge: SqlKnowledge | None = None, data_source: str | None = None, ) -> str: - text_to_sql_rules = get_text_to_sql_rules( - sql_knowledge, - data_source=data_source, - ) + text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" ### TASK ### @@ -139,9 +134,7 @@ def prompt( else "" ), metric_instructions=( - get_metric_instructions(sql_knowledge, data_source=data_source) - if has_metric - else "" + get_metric_instructions(sql_knowledge) if has_metric else "" ), json_field_instructions=( get_json_field_instructions(sql_knowledge) if has_json_field else "" @@ -182,8 +175,6 @@ async def post_process( regenerate_sql.get("replies"), project_id=project_id, data_source=data_source, - valid_table_names=construct_valid_table_names(documents), - valid_table_columns=construct_valid_table_columns(documents), ) @@ -200,7 +191,7 @@ def __init__( self._components = { "generator": llm_provider.get_generator( system_prompt=get_sql_regeneration_system_prompt(None), - generation_kwargs=get_sql_generation_model_kwargs(llm_provider), + generation_kwargs=SQL_GENERATION_MODEL_KWARGS, ), "generator_name": llm_provider.get_model(), "prompt_builder": PromptBuilder( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cd3c965e47..404a27341c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -8,13 +8,6 @@ from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - normalize_sql_direction_keywords, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, -) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -872,102 +865,7 @@ def _build_validated_ask_result_from_sql( table_ddls: list[str], query: str | None = None, ) -> Optional[AskResult]: - if isinstance(sql, str): - sql = normalize_sql_direction_keywords(sql) - sql = normalize_sql_table_references_to_schema( - sql, - construct_valid_table_names(table_ddls), - ) - sql = normalize_sql_column_references_to_schema( - sql, - construct_valid_table_columns(table_ddls), - ) - ask_result = self._build_ask_result_from_sql(sql) - if not ask_result: - return None - - schema_tables = self._parse_schema_tables(table_ddls) - valid_tables = { - str(table.get("name") or "").lower(): table - for table in schema_tables - if table.get("name") - } - valid_table_suffixes = { - table_name.split(".")[-1].lower(): table - for table_name, table in valid_tables.items() - } - - table_reference_pattern = re.compile( - r'\b(?:FROM|JOIN)\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|(?P[A-Za-z_][A-Za-z0-9_.$]*))", - flags=re.IGNORECASE, - ) - referenced_tables = [ - next(value for value in match.groupdict().values() if value) - for match in table_reference_pattern.finditer(ask_result.sql) - ] - invalid_tables = [ - table - for table in referenced_tables - if table.lower() not in valid_tables - and table.lower().split(".")[-1] not in valid_table_suffixes - ] - - columns_by_table = { - table_name: { - str(column.get("name") or "").lower() - for column in table.get("columns", []) - if column.get("name") - } - for table_name, table in valid_tables.items() - } - columns_by_table.update( - { - table_name.split(".")[-1].lower(): columns - for table_name, columns in columns_by_table.items() - } - ) - - qualified_column_pattern = re.compile( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\.\s*" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"(?P[A-Za-z_][A-Za-z0-9_$]*))", - flags=re.IGNORECASE, - ) - invalid_columns = [] - for match in qualified_column_pattern.finditer(ask_result.sql): - table_reference = ( - match.group("table_quoted") - or match.group("table_bracketed") - or match.group("table_bare") - or "" - ) - column_reference = ( - match.group("column_quoted") - or match.group("column_bracketed") - or match.group("column_bare") - or "" - ) - table_key = table_reference.lower() - column_key = column_reference.lower() - table_columns = columns_by_table.get(table_key) or columns_by_table.get( - table_key.split(".")[-1] - ) - if table_columns is not None and column_key not in table_columns: - invalid_columns.append(f"{table_reference}.{column_reference}") - - if invalid_tables or invalid_columns: - logger.warning( - "Ignoring generated SQL because it is not valid for active schema. " - "invalid_tables=%s invalid_columns=%s sql=%s", - invalid_tables, - invalid_columns, - ask_result.sql, - ) - return None - - return ask_result + return self._build_ask_result_from_sql(sql) def _build_failed_text_to_sql_response( self, diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index b9152f7d83..62586c9f1a 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,15 +6,6 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - extract_sql_table_references, - find_invalid_column_references, - find_invalid_table_references, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, -) from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest @@ -100,13 +91,8 @@ async def _load_active_schema_contexts( if not retrieval_pipeline: return [] - table_references = extract_sql_table_references(sql) - if not table_references: - return [] - retrieval_result = await retrieval_pipeline.run( query="", - tables=table_references, histories=[], project_id=project_id, enable_column_pruning=False, @@ -123,21 +109,7 @@ async def _load_active_schema_contexts( def _normalize_and_validate_sql( self, sql: str, schema_contexts: list[str] ) -> str | None: - valid_table_names = construct_valid_table_names(schema_contexts) - valid_table_columns = construct_valid_table_columns(schema_contexts) - normalized_sql = normalize_sql_table_references_to_schema( - sql, - valid_table_names, - ) - normalized_sql = normalize_sql_column_references_to_schema( - normalized_sql, - valid_table_columns, - ) - if find_invalid_table_references(normalized_sql, valid_table_names): - return None - if find_invalid_column_references(normalized_sql, valid_table_columns): - return None - return normalized_sql + return sql @observe(name="Generate Chart") @trace_metadata diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index 7442f2f884..dea25902a9 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -7,14 +7,6 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.sql import ( - construct_valid_table_columns, - construct_valid_table_names, - find_invalid_column_references, - find_invalid_table_references, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, -) from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -90,21 +82,7 @@ async def _load_active_schema_contexts( def _normalize_and_validate_sql( self, sql: str, schema_contexts: list[str] ) -> str | None: - valid_table_names = construct_valid_table_names(schema_contexts) - valid_table_columns = construct_valid_table_columns(schema_contexts) - normalized_sql = normalize_sql_table_references_to_schema( - sql, - valid_table_names, - ) - normalized_sql = normalize_sql_column_references_to_schema( - normalized_sql, - valid_table_columns, - ) - if find_invalid_table_references(normalized_sql, valid_table_names): - return None - if find_invalid_column_references(normalized_sql, valid_table_columns): - return None - return normalized_sql + return sql @observe(name="SQL Answer") @trace_metadata From 9ea0de804ea7f07377f947c4fa6fb370a8044e6d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 11:21:25 +0530 Subject: [PATCH 0621/1087] Restore semantic schema retrieval flow --- .../retrieval/db_schema_retrieval.py | 11 +- wren-ai-service/src/web/v1/services/ask.py | 279 +----------------- 2 files changed, 11 insertions(+), 279 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ceeedeb70d..6f1e3949ad 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -290,16 +290,7 @@ def check_using_db_schemas_without_pruning( retrieval_result["table_ddl"] for retrieval_result in retrieval_results ] _token_count = len(encoding.encode(" ".join(table_ddls))) - should_select_tables_for_question = ( - bool((query or "").strip()) - and not tables - and len(retrieval_results) > 1 - ) - if ( - _token_count > context_window_size - or enable_column_pruning - or should_select_tables_for_question - ): + if _token_count > context_window_size or enable_column_pruning: return { "db_schemas": [], "tokens": _token_count, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 404a27341c..8eab218197 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -16,8 +16,6 @@ NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( "No relevant data found in the active datasource for this question." ) -MAX_FORCED_EXPLICIT_TABLES = 5 - async def _return_value(value): return value @@ -253,113 +251,6 @@ def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: return tables - def _normalize_schema_identifier_key(self, value: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) - - def _schema_identifier_alias_keys(self, value: str) -> set[str]: - raw_value = str(value or "") - base_key = self._normalize_schema_identifier_key(raw_value) - separator_normalized_key = self._normalize_schema_identifier_key( - re.sub(r"[._\-]+", " ", raw_value) - ) - part_keys = { - self._normalize_schema_identifier_key(part) - for part in re.split(r"[._\-]+", raw_value) - if part - } - return {key for key in {base_key, separator_normalized_key, *part_keys} if key} - - def _extract_explicit_table_names_from_query(self, query: str) -> list[str]: - names: list[str] = [] - for quoted in re.findall(r"[`\"\[]([^`\"\]]+)[`\"\]]", query or ""): - if quoted and quoted not in names: - names.append(quoted) - for token in re.findall(r"\b[A-Za-z_][A-Za-z0-9_$]*(?:\.[A-Za-z_][A-Za-z0-9_$]*)+\b", query or ""): - if token and token not in names: - names.append(token) - return names - - def _explicit_table_alias_keys_from_query(self, query: str | None) -> set[str]: - keys: set[str] = set() - for table_name in self._extract_explicit_table_names_from_query(query or ""): - keys.update(self._schema_identifier_alias_keys(table_name)) - return keys - - def _explicit_table_alias_keys(self, table_names: list[str]) -> set[str]: - keys: set[str] = set() - for table_name in table_names: - keys.update(self._schema_identifier_alias_keys(table_name)) - return keys - - def _filter_retrieval_metadata_for_explicit_query( - self, - query: str, - documents: list[dict], - explicit_table_names: list[str] | None = None, - ) -> tuple[list[dict], list[str], list[str]]: - explicit_keys = ( - self._explicit_table_alias_keys(explicit_table_names or []) - if explicit_table_names - else self._explicit_table_alias_keys_from_query(query) - ) - if not explicit_keys: - return self._metadata_from_documents(documents) - - filtered_documents: list[dict] = [] - for document in documents or []: - metadata = document.get("metadata") or {} - table_name = metadata.get("table_name") or document.get("table_name") - table_ddl = metadata.get("table_ddl") or document.get("table_ddl") - candidate_names = [table_name] - if table_ddl: - candidate_names.extend( - table.get("name") - for table in self._parse_schema_tables([table_ddl]) - if table.get("name") - ) - - candidate_keys: set[str] = set() - for candidate_name in candidate_names: - candidate_keys.update(self._schema_identifier_alias_keys(candidate_name)) - - if candidate_keys & explicit_keys: - filtered_documents.append(document) - - if not filtered_documents: - return [], [], [] - return self._metadata_from_documents(filtered_documents) - - def _normalize_explicit_table_names( - self, - table_names: list[str] | None, - ) -> list[str]: - normalized: list[str] = [] - for table_name in table_names or []: - candidate = str(table_name or "").strip() - if candidate and candidate not in normalized: - normalized.append(candidate) - return normalized - - def _should_retry_selected_schema_after_retrieval_timeout( - self, retrieval_table_names: Optional[list[str]] - ) -> bool: - return bool(retrieval_table_names) - - def _forced_explicit_table_names( - self, table_names: list[str], *, source: str = "request" - ) -> list[str]: - if not table_names: - return [] - if len(table_names) <= MAX_FORCED_EXPLICIT_TABLES: - return table_names - - logger.info( - "Treating broad %s explicit_tables list as retrieval candidates, not a forced schema scope: %s", - source, - table_names, - ) - return [] - def _build_greeting_response(self, query: str) -> str: return ( f"Hi. I can help with questions about your active datasource and Wren AI.\n\n" @@ -984,16 +875,6 @@ async def ask( planning_timeout_seconds = min(self._pipeline_timeout_seconds, 15) generation_timeout_seconds = min(self._pipeline_timeout_seconds, 30) correction_timeout_seconds = min(self._pipeline_timeout_seconds, 15) - request_explicit_table_names = self._normalize_explicit_table_names( - ask_request.explicit_tables - ) - forced_request_explicit_table_names = self._forced_explicit_table_names( - request_explicit_table_names, - source="request", - ) - explicit_table_names = forced_request_explicit_table_names - retrieval_table_names = explicit_table_names or None - try: sql_user_query = user_query @@ -1076,78 +957,6 @@ async def ask( results["metadata"]["retrieved_table_count"] = len(documents) return results - if explicit_table_names: - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="TEXT_TO_SQL", - rephrased_question=user_query, - intent_reasoning="Explicit table name detected; retrieving that deployed schema directly.", - trace_id=trace_id, - is_followup=True if histories else False, - ) - retrieval_result = await self._run_with_timeout( - "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - logger.info( - "Retrieved explicit tables for query_id %s: %s", - query_id, - table_names, - ) - - if not documents: - error_message = ( - "The requested table was not found in the deployed schema: " - + ", ".join(explicit_table_names) - ) - self._ask_results[query_id] = AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError( - code="NO_RELEVANT_DATA", - message=error_message, - ), - rephrased_question=user_query, - intent_reasoning="Explicit table request did not match any deployed schema table.", - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = error_message - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - rephrased_question = user_query - intent_reasoning = ( - "Explicit table request matched deployed schema; generating SQL against retrieved schema." - ) - sql_user_query = user_query - historical_question_result = [] if not api_results: if not self._is_stopped(query_id, self._ask_results): @@ -1369,96 +1178,28 @@ async def ask( "Schema retrieval", self._pipelines["db_schema_retrieval"].run( query=sql_user_query, - tables=retrieval_table_names, - histories=[], + tables=None, + histories=histories, project_id=ask_request.project_id, enable_column_pruning=enable_column_pruning, ), timeout_seconds=self._schema_retrieval_timeout_seconds, ) except TimeoutError as error: - if not self._should_retry_selected_schema_after_retrieval_timeout( - retrieval_table_names - ): - logger.warning( - "Schema retrieval timed out for data query; not loading full project schema. " - "query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - error, - ) - retrieval_result = {"construct_retrieval_results": {}} - else: - logger.warning( - "Schema retrieval timed out; retrying only explicit selected schemas. " - "query_id=%s project_id=%s tables=%s error=%s", - query_id, - ask_request.project_id, - retrieval_table_names, - error, - ) - retrieval_result = await self._run_with_timeout( - "Selected schema fallback retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=retrieval_table_names, - histories=[], - project_id=ask_request.project_id, - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 30, - ), - ) + logger.warning( + "Schema retrieval timed out for data query; not loading full project schema. " + "query_id=%s project_id=%s error=%s", + query_id, + ask_request.project_id, + error, + ) + retrieval_result = {"construct_retrieval_results": {}} _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) documents, table_names, table_ddls = ( self._extract_retrieval_metadata(retrieval_result) ) - if explicit_table_names: - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) - if not documents: - if explicit_table_names: - logger.info( - "Retrying schema retrieval for explicit tables query_id %s: %s", - query_id, - explicit_table_names, - ) - retrieval_result = await self._run_with_timeout( - "Explicit table schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=user_query, - tables=explicit_table_names, - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=enable_column_pruning, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - 20, - ), - ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - documents, table_names, table_ddls = ( - self._filter_retrieval_metadata_for_explicit_query( - user_query, - documents, - explicit_table_names, - ) - ) logger.info( "Retrieved tables for query_id %s: %s", query_id, table_names ) From 66efba364468b8ca674c9edf37e1cf74dd94fef0 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 13:15:31 +0530 Subject: [PATCH 0622/1087] Accept SQL generation table limit config --- wren-ai-service/src/config.py | 1 + wren-ai-service/src/globals.py | 1 + wren-ai-service/src/web/v1/services/ask.py | 3 +++ 3 files changed, 5 insertions(+) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 46f53cb250..7fa5cc4d74 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -44,6 +44,7 @@ class Settings(BaseSettings): allow_sql_knowledge_retrieval: bool = Field(default=True) max_histories: int = Field(default=5) max_sql_correction_retries: int = Field(default=3) + max_sql_generation_tables: int = Field(default=10) pipeline_timeout_seconds: int = Field(default=90) schema_retrieval_timeout_seconds: int = Field(default=600) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index f2a0db24a6..3756110078 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -163,6 +163,7 @@ def create_service_container( max_histories=settings.max_histories, enable_column_pruning=settings.enable_column_pruning, max_sql_correction_retries=settings.max_sql_correction_retries, + max_sql_generation_tables=settings.max_sql_generation_tables, pipeline_timeout_seconds=settings.pipeline_timeout_seconds, schema_retrieval_timeout_seconds=settings.schema_retrieval_timeout_seconds, **query_cache, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8eab218197..c25f71d26c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -120,6 +120,7 @@ def __init__( pipeline_timeout_seconds: int = 45, schema_retrieval_timeout_seconds: int = 25, max_histories: int = 5, + max_sql_generation_tables: int = 10, maxsize: int = 1_000_000, ttl: int = 120, ): @@ -139,6 +140,7 @@ def __init__( self._pipeline_timeout_seconds = pipeline_timeout_seconds self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds self._max_histories = max_histories + self._max_sql_generation_tables = max_sql_generation_tables self._max_sql_correction_retries = max_sql_correction_retries def _is_stopped(self, query_id: str, container: dict): @@ -1229,6 +1231,7 @@ async def ask( documents, table_names, table_ddls, + max_tables=self._max_sql_generation_tables, ) ( documents, From 29d51b14032c30bcf3acaee8672977264385e77c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 14:32:45 +0530 Subject: [PATCH 0623/1087] Fix metadata-grounded SQL column selection --- .../src/pipelines/generation/sql_answer.py | 4 + .../src/pipelines/generation/utils/sql.py | 10 +- wren-ai-service/src/web/v1/services/ask.py | 190 +++++++++++++++++- .../services/test_metadata_grounding.py | 57 ++++++ 4 files changed, 259 insertions(+), 2 deletions(-) create mode 100644 wren-ai-service/tests/pytest/services/test_metadata_grounding.py diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 7ed59d5b63..4f180c9e25 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -35,6 +35,9 @@ 9. Always produce a narrative answer. Never return an empty response. 10. If the user asks for a chart or trend, still summarize the result in words and mention the chart-ready fields. 11. If the data contains only raw rows or a single column, summarize what those rows show, mention the visible date/category range when possible, and state that the result table contains the detailed rows. +12. If Data rows are present, answer from those rows only. Never say you do not have access to the database, system, records, or source data after rows are provided. +13. Do not give generic instructions about how the user can find the data when SQL results are present. Summarize the returned rows instead. +14. If Data rows are empty, say the SQL ran but returned no rows and briefly mention the selected columns, filters, or grouping visible in the SQL. ### OUTPUT FORMAT @@ -64,6 +67,7 @@ Custom Instruction: {{ custom_instruction }} Please think step by step and answer the user's question. +If rows are present in Data, summarize those rows directly and do not claim that the data is unavailable. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 45d3011bf4..bbb4b40d0c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -219,6 +219,13 @@ async def _classify_generation_result( - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. - NEVER invent, assume, or rename tables and columns. Generate SQL only from tables and columns present in the provided database schema. +- Treat the retrieved schema manifest and DATABASE SCHEMA as the exact identifier allowlist. Every SELECT, WHERE, JOIN, GROUP BY, HAVING, and ORDER BY table/column must exist there. +- Select columns by business meaning, not by name similarity alone. Match the user's entities, metrics, dimensions, filters, and dates to column names, descriptions, aliases, data types, user instructions, and SQL samples. +- Prefer semantically described business/canonical models, metrics, and views over staging, temp, test, raw, backup, load, or legacy tables when the metadata indicates that distinction. +- For value, amount, total, rate, count, average, or KPI questions, use numeric measures or numeric columns whose description/alias matches the requested metric. Do not SUM or AVG string columns. +- For comparison or "by" questions, use categorical/date dimension columns for grouping and numeric measures for aggregation. +- Do not use technical audit or ingestion columns (created_at, updated_at, loaded_at, inserted_at, file_date, batch_id, row_id, ingestion timestamps, etc.) unless the user explicitly asks about sync, load, audit, or ingestion. +- If multiple columns are equally plausible and no metadata/rule/sample disambiguates them, do not guess; return SQL only when the chosen columns are grounded by the metadata. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. @@ -567,7 +574,8 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. 4. If the section of REASONING PLAN is available in user's input, treat it only as high-level guidance. Ignore any table, column, alias, filter, or SQL fragment from the reasoning plan that is not explicitly present in the DATABASE SCHEMA. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +5. Before finalizing, validate the SQL against the retrieved metadata: every table, column, join, filter, GROUP BY, HAVING, and ORDER BY identifier must exist in the provided schema, and selected columns must match the user's business intent. +6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c25f71d26c..ebf400b745 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -752,13 +752,201 @@ def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: return None return AskResult(sql=sql.strip(), type="llm") + def _extract_sql_table_aliases( + self, sql: str, schema_tables: dict[str, dict[str, Any]] + ) -> dict[str, str]: + aliases: dict[str, str] = {} + reserved_aliases = { + "on", + "where", + "group", + "order", + "having", + "limit", + "join", + "left", + "right", + "inner", + "outer", + "full", + "cross", + } + pattern = re.compile( + r"\b(?:FROM|JOIN)\s+([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*)" + r"(?:\s+(?:AS\s+)?([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*))?", + flags=re.IGNORECASE, + ) + + for match in pattern.finditer(sql): + table_ref = match.group(1).strip().strip('"`[]').lower() + matched_table = next( + ( + table_name + for table_name in schema_tables + if self._normalize_schema_token(table_ref) + == self._normalize_schema_token(table_name) + ), + table_ref, + ) + aliases[table_ref] = matched_table + + alias = match.group(2) + if alias: + normalized_alias = alias.strip().strip('"`[]').lower() + if normalized_alias not in reserved_aliases: + aliases[normalized_alias] = matched_table + + return aliases + + def _find_invalid_sql_identifiers( + self, sql: str, table_ddls: list[str] + ) -> list[str]: + parsed_tables = self._parse_schema_tables(table_ddls) + schema_tables = { + str(table.get("name") or "").lower(): table + for table in parsed_tables + if table.get("name") + } + if not schema_tables: + return [] + + aliases = self._extract_sql_table_aliases(sql, schema_tables) + invalid_identifiers: set[str] = set() + + for table_ref in re.findall( + r"\b(?:FROM|JOIN)\s+([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*)", + sql, + flags=re.IGNORECASE, + ): + normalized_table_ref = table_ref.strip().strip('"`[]').lower() + if normalized_table_ref not in aliases: + invalid_identifiers.add(normalized_table_ref) + + for match in re.finditer( + r"([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*)\s*\.\s*([A-Za-z_\"`\[][A-Za-z0-9_$\"`\[\]]*)", + sql, + ): + table_or_alias = match.group(1).strip().strip('"`[]').lower() + column_name = match.group(2).strip().strip('"`[]').lower() + table_name = aliases.get(table_or_alias, table_or_alias) + table = schema_tables.get(table_name) + if not table: + invalid_identifiers.add(table_or_alias) + continue + + schema_columns = { + str(column.get("name") or "").lower() + for column in table.get("columns", []) + if isinstance(column, dict) and column.get("name") + } + if column_name not in schema_columns: + invalid_identifiers.add(f"{table_or_alias}.{column_name}") + + referenced_tables = set(aliases.values()) + if len(referenced_tables) == 1: + table_name = next(iter(referenced_tables)) + table = schema_tables.get(table_name) + schema_columns = { + str(column.get("name") or "").lower() + for column in (table or {}).get("columns", []) + if isinstance(column, dict) and column.get("name") + } + sql_without_qualified_refs = re.sub( + r"[A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*\s*\.\s*[A-Za-z_\"`\[][A-Za-z0-9_$\"`\[\]]*", + " ", + sql, + ) + sql_without_literals = re.sub( + r"'(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"", + " ", + sql_without_qualified_refs, + ) + keywords = { + "as", + "and", + "or", + "not", + "null", + "is", + "in", + "like", + "between", + "case", + "when", + "then", + "else", + "end", + "select", + "from", + "join", + "on", + "where", + "group", + "by", + "order", + "having", + "limit", + "distinct", + "asc", + "desc", + "with", + "over", + "partition", + "sum", + "avg", + "count", + "min", + "max", + "cast", + "datepart", + "dateadd", + "datediff", + "lower", + "upper", + "coalesce", + "round", + "rank", + "dense_rank", + "row_number", + } + for candidate in re.findall( + r"\b[A-Za-z_][A-Za-z0-9_$]*\b", sql_without_literals + ): + normalized = candidate.lower() + if ( + normalized in keywords + or normalized in aliases + or normalized in schema_tables + or normalized in schema_columns + ): + continue + invalid_identifiers.add(f"{table_name}.{normalized}") + + return sorted(invalid_identifiers) + def _build_validated_ask_result_from_sql( self, sql: Optional[str], table_ddls: list[str], query: str | None = None, ) -> Optional[AskResult]: - return self._build_ask_result_from_sql(sql) + ask_result = self._build_ask_result_from_sql(sql) + if not ask_result: + return None + + invalid_identifiers = self._find_invalid_sql_identifiers( + ask_result.sql, table_ddls + ) + if invalid_identifiers: + logger.warning( + "Generated SQL references identifiers outside deployed metadata. query=%s invalid_identifiers=%s sql=%s", + query, + invalid_identifiers, + ask_result.sql, + ) + return None + + return ask_result def _build_failed_text_to_sql_response( self, diff --git a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py new file mode 100644 index 0000000000..978f5a0f2e --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py @@ -0,0 +1,57 @@ +from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt +from src.pipelines.generation.utils.sql import get_sql_generation_system_prompt +from src.web.v1.services.ask import AskService + + +def test_schema_grounding_rejects_unknown_column_after_table_match(): + service = AskService.__new__(AskService) + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_orders"."CustomerName", ' + 'SUM("dbo_orders"."HallucinatedValue") AS "TotalValue" ' + 'FROM "dbo_orders" GROUP BY "dbo_orders"."CustomerName"' + ), + [ + """ + CREATE TABLE dbo_orders ( + CustomerName VARCHAR, + OrderValue DOUBLE + ); + """ + ], + "compare order value by customer", + ) + + assert result is None + + +def test_schema_grounding_accepts_valid_columns(): + service = AskService.__new__(AskService) + result = service._build_validated_ask_result_from_sql( + ( + 'SELECT "dbo_orders"."CustomerName", ' + 'SUM("dbo_orders"."OrderValue") AS "TotalValue" ' + 'FROM "dbo_orders" GROUP BY "dbo_orders"."CustomerName"' + ), + [ + """ + CREATE TABLE dbo_orders ( + CustomerName VARCHAR, + OrderValue DOUBLE + ); + """ + ], + "compare order value by customer", + ) + + assert result is not None + + +def test_prompts_enforce_metadata_grounding_and_result_grounded_answers(): + sql_prompt = get_sql_generation_system_prompt() + + assert "exact identifier allowlist" in sql_prompt + assert "Select columns by business meaning" in sql_prompt + assert "Do not SUM or AVG string columns" in sql_prompt + assert "Never say you do not have access" in sql_to_answer_system_prompt + assert "If Data rows are empty" in sql_to_answer_system_prompt From a631afcf294dacc7d6840f9d6e75684410d702c6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 14:45:58 +0530 Subject: [PATCH 0624/1087] Restore legacy SQL validation flow --- .../generation/followup_sql_generation.py | 5 +- wren-ai-service/src/web/v1/services/ask.py | 190 +----------------- .../services/test_metadata_grounding.py | 65 +++--- 3 files changed, 35 insertions(+), 225 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 5f769b974a..43a6f55d7b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -165,10 +165,7 @@ async def generate_sql_in_followup( sql_knowledge: SqlKnowledge | None = None, ) -> dict: history_messages = construct_ask_history_messages(histories) - current_system_prompt = get_sql_generation_system_prompt( - sql_knowledge, - data_source=data_source, - ) + current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) return await generator( prompt=prompt.get("prompt"), history_messages=history_messages, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ebf400b745..c25f71d26c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -752,201 +752,13 @@ def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: return None return AskResult(sql=sql.strip(), type="llm") - def _extract_sql_table_aliases( - self, sql: str, schema_tables: dict[str, dict[str, Any]] - ) -> dict[str, str]: - aliases: dict[str, str] = {} - reserved_aliases = { - "on", - "where", - "group", - "order", - "having", - "limit", - "join", - "left", - "right", - "inner", - "outer", - "full", - "cross", - } - pattern = re.compile( - r"\b(?:FROM|JOIN)\s+([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*)" - r"(?:\s+(?:AS\s+)?([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*))?", - flags=re.IGNORECASE, - ) - - for match in pattern.finditer(sql): - table_ref = match.group(1).strip().strip('"`[]').lower() - matched_table = next( - ( - table_name - for table_name in schema_tables - if self._normalize_schema_token(table_ref) - == self._normalize_schema_token(table_name) - ), - table_ref, - ) - aliases[table_ref] = matched_table - - alias = match.group(2) - if alias: - normalized_alias = alias.strip().strip('"`[]').lower() - if normalized_alias not in reserved_aliases: - aliases[normalized_alias] = matched_table - - return aliases - - def _find_invalid_sql_identifiers( - self, sql: str, table_ddls: list[str] - ) -> list[str]: - parsed_tables = self._parse_schema_tables(table_ddls) - schema_tables = { - str(table.get("name") or "").lower(): table - for table in parsed_tables - if table.get("name") - } - if not schema_tables: - return [] - - aliases = self._extract_sql_table_aliases(sql, schema_tables) - invalid_identifiers: set[str] = set() - - for table_ref in re.findall( - r"\b(?:FROM|JOIN)\s+([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*)", - sql, - flags=re.IGNORECASE, - ): - normalized_table_ref = table_ref.strip().strip('"`[]').lower() - if normalized_table_ref not in aliases: - invalid_identifiers.add(normalized_table_ref) - - for match in re.finditer( - r"([A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*)\s*\.\s*([A-Za-z_\"`\[][A-Za-z0-9_$\"`\[\]]*)", - sql, - ): - table_or_alias = match.group(1).strip().strip('"`[]').lower() - column_name = match.group(2).strip().strip('"`[]').lower() - table_name = aliases.get(table_or_alias, table_or_alias) - table = schema_tables.get(table_name) - if not table: - invalid_identifiers.add(table_or_alias) - continue - - schema_columns = { - str(column.get("name") or "").lower() - for column in table.get("columns", []) - if isinstance(column, dict) and column.get("name") - } - if column_name not in schema_columns: - invalid_identifiers.add(f"{table_or_alias}.{column_name}") - - referenced_tables = set(aliases.values()) - if len(referenced_tables) == 1: - table_name = next(iter(referenced_tables)) - table = schema_tables.get(table_name) - schema_columns = { - str(column.get("name") or "").lower() - for column in (table or {}).get("columns", []) - if isinstance(column, dict) and column.get("name") - } - sql_without_qualified_refs = re.sub( - r"[A-Za-z_\"`\[][A-Za-z0-9_.$\"`\[\]]*\s*\.\s*[A-Za-z_\"`\[][A-Za-z0-9_$\"`\[\]]*", - " ", - sql, - ) - sql_without_literals = re.sub( - r"'(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"", - " ", - sql_without_qualified_refs, - ) - keywords = { - "as", - "and", - "or", - "not", - "null", - "is", - "in", - "like", - "between", - "case", - "when", - "then", - "else", - "end", - "select", - "from", - "join", - "on", - "where", - "group", - "by", - "order", - "having", - "limit", - "distinct", - "asc", - "desc", - "with", - "over", - "partition", - "sum", - "avg", - "count", - "min", - "max", - "cast", - "datepart", - "dateadd", - "datediff", - "lower", - "upper", - "coalesce", - "round", - "rank", - "dense_rank", - "row_number", - } - for candidate in re.findall( - r"\b[A-Za-z_][A-Za-z0-9_$]*\b", sql_without_literals - ): - normalized = candidate.lower() - if ( - normalized in keywords - or normalized in aliases - or normalized in schema_tables - or normalized in schema_columns - ): - continue - invalid_identifiers.add(f"{table_name}.{normalized}") - - return sorted(invalid_identifiers) - def _build_validated_ask_result_from_sql( self, sql: Optional[str], table_ddls: list[str], query: str | None = None, ) -> Optional[AskResult]: - ask_result = self._build_ask_result_from_sql(sql) - if not ask_result: - return None - - invalid_identifiers = self._find_invalid_sql_identifiers( - ask_result.sql, table_ddls - ) - if invalid_identifiers: - logger.warning( - "Generated SQL references identifiers outside deployed metadata. query=%s invalid_identifiers=%s sql=%s", - query, - invalid_identifiers, - ask_result.sql, - ) - return None - - return ask_result + return self._build_ask_result_from_sql(sql) def _build_failed_text_to_sql_response( self, diff --git a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py index 978f5a0f2e..b30e76d0ce 100644 --- a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py +++ b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py @@ -1,50 +1,30 @@ +from src.pipelines.generation.followup_sql_generation import generate_sql_in_followup from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt from src.pipelines.generation.utils.sql import get_sql_generation_system_prompt from src.web.v1.services.ask import AskService -def test_schema_grounding_rejects_unknown_column_after_table_match(): +def test_build_validated_ask_result_keeps_legacy_select_flow(): service = AskService.__new__(AskService) result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_orders"."CustomerName", ' - 'SUM("dbo_orders"."HallucinatedValue") AS "TotalValue" ' - 'FROM "dbo_orders" GROUP BY "dbo_orders"."CustomerName"' - ), - [ - """ - CREATE TABLE dbo_orders ( - CustomerName VARCHAR, - OrderValue DOUBLE - ); - """ - ], - "compare order value by customer", + 'SELECT "dbo_orders"."CustomerName" FROM "dbo_orders"', + [], + "show customers", ) - assert result is None + assert result is not None + assert result.sql == 'SELECT "dbo_orders"."CustomerName" FROM "dbo_orders"' -def test_schema_grounding_accepts_valid_columns(): +def test_build_validated_ask_result_rejects_non_select_sql(): service = AskService.__new__(AskService) result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_orders"."CustomerName", ' - 'SUM("dbo_orders"."OrderValue") AS "TotalValue" ' - 'FROM "dbo_orders" GROUP BY "dbo_orders"."CustomerName"' - ), - [ - """ - CREATE TABLE dbo_orders ( - CustomerName VARCHAR, - OrderValue DOUBLE - ); - """ - ], - "compare order value by customer", + 'DELETE FROM "dbo_orders"', + [], + "delete customers", ) - assert result is not None + assert result is None def test_prompts_enforce_metadata_grounding_and_result_grounded_answers(): @@ -55,3 +35,24 @@ def test_prompts_enforce_metadata_grounding_and_result_grounded_answers(): assert "Do not SUM or AVG string columns" in sql_prompt assert "Never say you do not have access" in sql_to_answer_system_prompt assert "If Data rows are empty" in sql_to_answer_system_prompt + + +async def test_followup_sql_generation_uses_current_system_prompt_signature(): + calls = {} + + class Generator: + async def __call__(self, **kwargs): + calls.update(kwargs) + return {"replies": ['{"sql": "SELECT 1"}']} + + result = await generate_sql_in_followup( + prompt={"prompt": "prompt"}, + generator=Generator(), + histories=[], + generator_name="test", + data_source="mssql", + sql_knowledge=None, + ) + + assert result[1] == "test" + assert "current_system_prompt" in calls From 643880737fc3bbe494a60eb66d60dffb511e053b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 15:02:06 +0530 Subject: [PATCH 0625/1087] Improve natural language schema retrieval flow --- .../generation/intent_classification.py | 66 ++++++++++++-- .../retrieval/db_schema_retrieval.py | 89 +++++++++++++++---- .../services/test_metadata_grounding.py | 33 +++++++ 3 files changed, 162 insertions(+), 26 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 8fe3ada814..4e993a52bb 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,5 +1,6 @@ import ast import logging +import re import sys from typing import TYPE_CHECKING, Any, Literal, Optional @@ -36,7 +37,8 @@ - **Rephrase Question:** Rewrite follow-up questions into full standalone questions using prior conversation context. - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. -- **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. +- **Natural-language data questions:** A user does not need to mention exact table or column names. If the question asks for data, counts, totals, averages, trends, comparisons, rankings, distinct values, frequencies, distributions, filters, lists, or chart-ready analysis that could be answered from the provided schema, classify it as `TEXT_TO_SQL`. +- **Vague Queries:** If the question is vague and cannot be connected to any data concept, table, column, metric, or SQL sample from the schema, classify it as `MISLEADING_QUERY`. - **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. @@ -46,19 +48,22 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes references to specific tables, columns, or data details. -- The question includes **complete information** with specific tables, columns, or data values needed for execution. +- The question (or related previous query) includes references to business entities, measures, dimensions, dates, filters, or data details, even when exact table or column names are not mentioned. +- The question includes enough business intent to retrieve matching tables and columns from the schema. - The question provides **all necessary parameters** to generate executable SQL. **Requirements:** -- Must have complete filter criteria, specific values, or clear references to previous context. -- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. -- Reference phrases from the user's inputs that clearly relate to the schema. +- Do not require users to provide table names or column names. +- Use the provided schema to map natural-language concepts to tables and columns. +- Reference phrases from the user's inputs that clearly describe the requested data operation. **Examples:** - "What is the total sales for last quarter?" - "Show me all customers who purchased product X." - "List the top 10 products by revenue." +- "What are the unique values and frequencies for business unit?" +- "What is the average length of emails?" +- "Which markets have the highest growth rate this year compared to last year?" @@ -84,6 +89,7 @@ **When to Use:** - The user's inputs pertains to Wren AI's features, usage, or capabilities. - The query relates directly to content in the user guide. +- Do not use this category for business/data questions just because the exact table name is missing. **Examples:** - "What can Wren AI do?" @@ -96,7 +102,7 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. +- The user's inputs lacks any data-analysis intent that can be mapped to the provided schema. - It appears off-topic or is simply a casual conversation starter. **Requirements:** @@ -299,13 +305,55 @@ async def classify_intent(prompt: dict, generator: Any, generator_name: str) -> return await generator(prompt=prompt.get("prompt")), generator_name +def _looks_like_data_question(query: str) -> bool: + return bool( + re.search( + r"\b(" + r"show|list|find|get|give|fetch|compare|calculate|count|sum|total|" + r"average|avg|mean|min|max|median|distinct|unique|frequency|" + r"frequencies|distribution|trend|growth|rate|ratio|percentage|" + r"top|bottom|highest|lowest|rank|group|breakdown|by|where|filter|" + r"chart|graph|plot|values|records|rows" + r")\b", + query or "", + flags=re.IGNORECASE, + ) + ) + + +def _looks_like_user_guide_question(query: str) -> bool: + return bool( + re.search( + r"\b(" + r"wren|project|workspace|connect|connection|database connection|" + r"setup|configure|configuration|delete|reset|invite|permission|" + r"role|user guide|documentation|how do i|how can i" + r")\b", + query or "", + flags=re.IGNORECASE, + ) + ) + + @observe(capture_input=False) -def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict: +def post_process( + classify_intent: dict, construct_db_schemas: list[str], query: str +) -> dict: try: results = orjson.loads(classify_intent.get("replies")[0]) + intent = results["results"] + if ( + construct_db_schemas + and intent in {"MISLEADING_QUERY", "USER_GUIDE"} + and _looks_like_data_question(query) + and not ( + intent == "USER_GUIDE" and _looks_like_user_guide_question(query) + ) + ): + intent = "TEXT_TO_SQL" return { "rephrased_question": results["rephrased_question"], - "intent": results["results"], + "intent": intent, "reasoning": results["reasoning"], "db_schemas": construct_db_schemas, } diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6f1e3949ad..3a35a2c436 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -43,6 +43,12 @@ class AskHistoryLike(Protocol): 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +8. The user does not need to mention exact table or column names. Map natural-language business terms to tables and columns using table names, column names, comments, aliases, descriptions, data types, primary keys, foreign keys, relationships, metrics, and views. +9. Prefer tables, views, or metrics whose descriptions and columns directly match the requested business entities, measures, dimensions, dates, filters, rankings, frequencies, or comparisons. +10. For metric questions such as total, count, average, rate, value, amount, growth, frequency, or distribution, include numeric/measure columns and the grouping/filter/date columns needed to answer the question. +11. For "unique values", "frequency", "distribution", or "count by" questions, include the categorical column being counted and any requested filter/date columns. +12. Prefer canonical/business-ready models, metrics, and views over staging, raw, temporary, backup, load, or test tables when the metadata indicates that distinction. +13. If no table clearly matches, return an empty results list instead of selecting unrelated tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -86,6 +92,7 @@ class AskHistoryLike(Protocol): - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. - Match Table names with the definition in the "Create Table" statement. +- Do not choose a column only because it has a similar name; choose it only when its metadata meaning and data type fit the user's question. Good luck! @@ -354,30 +361,78 @@ def construct_retrieval_results( construct_db_schemas: list[dict], dbschema_retrieval: list[Document], ) -> dict[str, Any]: + fallback_retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] + fallback_result = { + "retrieval_results": fallback_retrieval_results, + "has_calculated_field": check_using_db_schemas_without_pruning[ + "has_calculated_field" + ], + "has_metric": check_using_db_schemas_without_pruning["has_metric"], + "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], + } + if filter_columns_in_tables: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] + try: + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + )["results"] + except Exception: + logger.warning( + "Column pruning returned invalid output; using full retrieved schemas." + ) + return fallback_result + + if not columns_and_tables_needed: + logger.info( + "Column pruning selected no tables; using full retrieved schemas." + ) + return fallback_result # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format reformated_json = {} for table in columns_and_tables_needed: - reformated_json[table["table_name"]] = table["table_contents"] + table_name = table.get("table_name") + table_contents = table.get("table_contents") or {} + if table_name: + reformated_json[table_name] = table_contents columns_and_tables_needed = reformated_json + if not columns_and_tables_needed: + logger.info( + "Column pruning selected no usable tables; using full retrieved schemas." + ) + return fallback_result + tables = set(columns_and_tables_needed.keys()) retrieval_results = [] has_calculated_field = False has_metric = False has_json_field = False + valid_table_names = { + table_schema["name"] for table_schema in construct_db_schemas + } for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: + requested_columns = set( + columns_and_tables_needed[table_schema["name"]].get("columns", []) + ) + valid_columns = { + column.get("name") + for column in table_schema.get("columns", []) + if isinstance(column, dict) and column.get("name") + } + selected_columns = requested_columns & valid_columns + if not selected_columns: + logger.info( + "Column pruning selected no valid columns for table %s; using full retrieved schemas.", + table_schema["name"], + ) + return fallback_result + ddl, _has_calculated_field, _has_json_field = build_table_ddl( table_schema, - columns=set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ), + columns=selected_columns, tables=tables, ) if _has_calculated_field: @@ -395,6 +450,7 @@ def construct_retrieval_results( for document in dbschema_retrieval: if document.meta["name"] in columns_and_tables_needed: content = ast.literal_eval(document.content) + valid_table_names.add(content["name"]) if content["type"] == "METRIC": retrieval_results.append( @@ -412,6 +468,14 @@ def construct_retrieval_results( } ) + unknown_tables = tables - valid_table_names + if unknown_tables or not retrieval_results: + logger.info( + "Column pruning selected unknown or unusable tables %s; using full retrieved schemas.", + sorted(unknown_tables), + ) + return fallback_result + return { "retrieval_results": retrieval_results, "has_calculated_field": has_calculated_field, @@ -419,16 +483,7 @@ def construct_retrieval_results( "has_json_field": has_json_field, } else: - retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] - - return { - "retrieval_results": retrieval_results, - "has_calculated_field": check_using_db_schemas_without_pruning[ - "has_calculated_field" - ], - "has_metric": check_using_db_schemas_without_pruning["has_metric"], - "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], - } + return fallback_result ## End of Pipeline diff --git a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py index b30e76d0ce..6534dfc683 100644 --- a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py +++ b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py @@ -1,4 +1,5 @@ from src.pipelines.generation.followup_sql_generation import generate_sql_in_followup +from src.pipelines.generation.intent_classification import post_process from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt from src.pipelines.generation.utils.sql import get_sql_generation_system_prompt from src.web.v1.services.ask import AskService @@ -56,3 +57,35 @@ async def __call__(self, **kwargs): assert result[1] == "test" assert "current_system_prompt" in calls + + +def test_analytic_question_with_schema_stays_text_to_sql_without_table_name(): + result = post_process( + classify_intent={ + "replies": [ + '{"rephrased_question":"What is the average length of emails?",' + '"reasoning":"Misclassified as guide.",' + '"results":"USER_GUIDE"}' + ] + }, + construct_db_schemas=["CREATE TABLE users (email VARCHAR);"], + query="What is the average length of emails?", + ) + + assert result["intent"] == "TEXT_TO_SQL" + + +def test_user_guide_question_stays_user_guide(): + result = post_process( + classify_intent={ + "replies": [ + '{"rephrased_question":"How can I connect to a database?",' + '"reasoning":"Wren setup question.",' + '"results":"USER_GUIDE"}' + ] + }, + construct_db_schemas=["CREATE TABLE users (email VARCHAR);"], + query="How can I connect to a database?", + ) + + assert result["intent"] == "USER_GUIDE" From 3be8d23f18ae5ffd35bb085261d45899c1ec3c6a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 15:21:17 +0530 Subject: [PATCH 0626/1087] Pass retrieved table allowlist to SQL prompts --- .../generation/followup_sql_generation.py | 5 ++++- .../followup_sql_generation_reasoning.py | 14 ++++++++++++++ .../src/pipelines/generation/sql_correction.py | 5 ++++- .../src/pipelines/generation/sql_generation.py | 16 +++++++++++++++- .../generation/sql_generation_reasoning.py | 14 ++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 5 +++++ .../pytest/services/test_metadata_grounding.py | 3 +++ 7 files changed, 59 insertions(+), 3 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 43a6f55d7b..b013f9b25c 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -127,12 +127,13 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, data_source=data_source, documents=documents, - valid_table_names=[], + valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -240,6 +241,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -250,6 +252,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index abbbb81d56..4d4402ea69 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -37,6 +37,16 @@ {{ document }} {% endfor %} +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only mention these exact table names in the reasoning plan. Do not invent, rename, +singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless +the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -81,10 +91,12 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + valid_table_names=valid_table_names or [], histories=histories, sql_samples=sql_samples, instructions=construct_instructions( @@ -185,6 +197,7 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, + valid_table_names: Optional[list[str]] = None, ): logger.info("Followup SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -192,6 +205,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "histories": histories, "sql_samples": sql_samples or [], "instructions": instructions or [], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 284cf09a6d..e61f873292 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -125,12 +125,13 @@ def prompt( query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, data_source=data_source, documents=documents, - valid_table_names=[], + valid_table_names=valid_table_names or [], invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, @@ -220,6 +221,7 @@ async def run( allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, query: str | None = None, + valid_table_names: list[str] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -230,6 +232,7 @@ async def run( inputs={ "invalid_generation_result": invalid_generation_result, "documents": contexts, + "valid_table_names": valid_table_names or [], "query": query, "instructions": instructions, "sql_functions": sql_functions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 28f33d138e..d2c5ac4724 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,6 +28,16 @@ sql_generation_user_prompt_template = """ +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only use these exact table names from the retrieved schema. Do not invent, rename, +singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless +the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -95,10 +105,12 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -202,6 +214,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -215,6 +228,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, @@ -230,4 +244,4 @@ async def run( "sql_knowledge": sql_knowledge, **self._components, }, - ) \ No newline at end of file + ) 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 f91a4288e4..fc52e5a4d5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -30,6 +30,16 @@ {{ document }} {% endfor %} +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only mention these exact table names in the reasoning plan. Do not invent, rename, +singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless +the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -65,10 +75,12 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + valid_table_names=valid_table_names or [], sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, @@ -163,6 +175,7 @@ async def run( instructions: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, + valid_table_names: Optional[list[str]] = None, ): logger.info("SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -170,6 +183,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "sql_samples": sql_samples or [], "instructions": instructions or [], "configuration": configuration, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c25f71d26c..258944a5fb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1275,6 +1275,7 @@ async def ask( ].run( query=sql_user_query, contexts=table_ddls, + valid_table_names=table_names, histories=sql_generation_histories, sql_samples=sql_samples, instructions=instructions, @@ -1299,6 +1300,7 @@ async def ask( self._pipelines["sql_generation_reasoning"].run( query=sql_user_query, contexts=table_ddls, + valid_table_names=table_names, sql_samples=sql_samples, instructions=instructions, configuration=ask_request.configurations, @@ -1380,6 +1382,7 @@ async def ask( self._pipelines["followup_sql_generation"].run( query=sql_user_query, contexts=table_ddls, + valid_table_names=table_names, sql_generation_reasoning=sql_generation_reasoning, histories=sql_generation_histories, project_id=ask_request.project_id, @@ -1401,6 +1404,7 @@ async def ask( self._pipelines["sql_generation"].run( query=sql_user_query, contexts=table_ddls, + valid_table_names=table_names, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, sql_samples=sql_samples, @@ -1500,6 +1504,7 @@ async def ask( "SQL correction", self._pipelines["sql_correction"].run( contexts=table_ddls, + valid_table_names=table_names, instructions=instructions, invalid_generation_result={ "original_sql": original_sql, diff --git a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py index 6534dfc683..96fb49b77a 100644 --- a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py +++ b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py @@ -1,4 +1,5 @@ from src.pipelines.generation.followup_sql_generation import generate_sql_in_followup +from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template from src.pipelines.generation.intent_classification import post_process from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt from src.pipelines.generation.utils.sql import get_sql_generation_system_prompt @@ -36,6 +37,8 @@ def test_prompts_enforce_metadata_grounding_and_result_grounded_answers(): assert "Do not SUM or AVG string columns" in sql_prompt assert "Never say you do not have access" in sql_to_answer_system_prompt assert "If Data rows are empty" in sql_to_answer_system_prompt + assert "VALID TABLE NAMES" in sql_generation_user_prompt_template + assert "Do not invent, rename" in sql_generation_user_prompt_template async def test_followup_sql_generation_uses_current_system_prompt_signature(): From bbc317abdcdac467f1dacf19bb217fc4deafbc44 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 15:21:35 +0530 Subject: [PATCH 0627/1087] Restore legacy ask pipeline flow --- wren-ai-service/src/config.py | 5 +- wren-ai-service/src/globals.py | 6 - .../generation/intent_classification.py | 73 +- .../retrieval/db_schema_retrieval.py | 126 +- wren-ai-service/src/web/v1/services/ask.py | 1398 +++-------------- 5 files changed, 254 insertions(+), 1354 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 7fa5cc4d74..c5acf4ae47 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -41,12 +41,9 @@ class Settings(BaseSettings): allow_sql_generation_reasoning: bool = Field(default=True) allow_sql_functions_retrieval: bool = Field(default=True) allow_sql_diagnosis: bool = Field(default=True) - allow_sql_knowledge_retrieval: bool = Field(default=True) + allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) max_sql_correction_retries: int = Field(default=3) - max_sql_generation_tables: int = Field(default=10) - pipeline_timeout_seconds: int = Field(default=90) - schema_retrieval_timeout_seconds: int = Field(default=600) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 3756110078..9343344616 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -90,7 +90,6 @@ def create_service_container( **pipe_components["semantics_description"], ) }, - generation_timeout_seconds=settings.pipeline_timeout_seconds, **query_cache, ), semantics_preparation_service=services.SemanticsPreparationService( @@ -163,9 +162,6 @@ def create_service_container( max_histories=settings.max_histories, enable_column_pruning=settings.enable_column_pruning, max_sql_correction_retries=settings.max_sql_correction_retries, - max_sql_generation_tables=settings.max_sql_generation_tables, - pipeline_timeout_seconds=settings.pipeline_timeout_seconds, - schema_retrieval_timeout_seconds=settings.schema_retrieval_timeout_seconds, **query_cache, ), ask_feedback_service=services.AskFeedbackService( @@ -191,7 +187,6 @@ def create_service_container( chart_service=services.ChartService( pipelines={ "sql_executor": _sql_executor_pipeline, - "db_schema_retrieval": _db_schema_retrieval_pipeline, "chart_generation": generation.ChartGeneration( **pipe_components["chart_generation"], ), @@ -209,7 +204,6 @@ def create_service_container( ), sql_answer_service=services.SqlAnswerService( pipelines={ - "db_schema_retrieval": _db_schema_retrieval_pipeline, "preprocess_sql_data": retrieval.PreprocessSqlData( **pipe_components["preprocess_sql_data"], ), diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4e993a52bb..4d6cd313cd 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,8 +1,7 @@ import ast import logging -import re import sys -from typing import TYPE_CHECKING, Any, Literal, Optional +from typing import Any, Literal, Optional import orjson from hamilton import base @@ -18,10 +17,7 @@ from src.pipelines.generation.utils.sql import construct_instructions from src.utils import trace_cost from src.web.v1.services import Configuration -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -37,8 +33,7 @@ - **Rephrase Question:** Rewrite follow-up questions into full standalone questions using prior conversation context. - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. -- **Natural-language data questions:** A user does not need to mention exact table or column names. If the question asks for data, counts, totals, averages, trends, comparisons, rankings, distinct values, frequencies, distributions, filters, lists, or chart-ready analysis that could be answered from the provided schema, classify it as `TEXT_TO_SQL`. -- **Vague Queries:** If the question is vague and cannot be connected to any data concept, table, column, metric, or SQL sample from the schema, classify it as `MISLEADING_QUERY`. +- **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. - **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. @@ -48,22 +43,19 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes references to business entities, measures, dimensions, dates, filters, or data details, even when exact table or column names are not mentioned. -- The question includes enough business intent to retrieve matching tables and columns from the schema. +- The question (or related previous query) includes references to specific tables, columns, or data details. +- The question includes **complete information** with specific tables, columns, or data values needed for execution. - The question provides **all necessary parameters** to generate executable SQL. **Requirements:** -- Do not require users to provide table names or column names. -- Use the provided schema to map natural-language concepts to tables and columns. -- Reference phrases from the user's inputs that clearly describe the requested data operation. +- Must have complete filter criteria, specific values, or clear references to previous context. +- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. +- Reference phrases from the user's inputs that clearly relate to the schema. **Examples:** - "What is the total sales for last quarter?" - "Show me all customers who purchased product X." - "List the top 10 products by revenue." -- "What are the unique values and frequencies for business unit?" -- "What is the average length of emails?" -- "Which markets have the highest growth rate this year compared to last year?" @@ -89,7 +81,6 @@ **When to Use:** - The user's inputs pertains to Wren AI's features, usage, or capabilities. - The query relates directly to content in the user guide. -- Do not use this category for business/data questions just because the exact table name is missing. **Examples:** - "What can Wren AI do?" @@ -102,7 +93,7 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs lacks any data-analysis intent that can be mapped to the provided schema. +- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. - It appears off-topic or is simply a casual conversation starter. **Requirements:** @@ -305,55 +296,13 @@ async def classify_intent(prompt: dict, generator: Any, generator_name: str) -> return await generator(prompt=prompt.get("prompt")), generator_name -def _looks_like_data_question(query: str) -> bool: - return bool( - re.search( - r"\b(" - r"show|list|find|get|give|fetch|compare|calculate|count|sum|total|" - r"average|avg|mean|min|max|median|distinct|unique|frequency|" - r"frequencies|distribution|trend|growth|rate|ratio|percentage|" - r"top|bottom|highest|lowest|rank|group|breakdown|by|where|filter|" - r"chart|graph|plot|values|records|rows" - r")\b", - query or "", - flags=re.IGNORECASE, - ) - ) - - -def _looks_like_user_guide_question(query: str) -> bool: - return bool( - re.search( - r"\b(" - r"wren|project|workspace|connect|connection|database connection|" - r"setup|configure|configuration|delete|reset|invite|permission|" - r"role|user guide|documentation|how do i|how can i" - r")\b", - query or "", - flags=re.IGNORECASE, - ) - ) - - @observe(capture_input=False) -def post_process( - classify_intent: dict, construct_db_schemas: list[str], query: str -) -> dict: +def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict: try: results = orjson.loads(classify_intent.get("replies")[0]) - intent = results["results"] - if ( - construct_db_schemas - and intent in {"MISLEADING_QUERY", "USER_GUIDE"} - and _looks_like_data_question(query) - and not ( - intent == "USER_GUIDE" and _looks_like_user_guide_question(query) - ) - ): - intent = "TEXT_TO_SQL" return { "rephrased_question": results["rephrased_question"], - "intent": intent, + "intent": results["results"], "reasoning": results["reasoning"], "db_schemas": construct_db_schemas, } diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 3a35a2c436..6c8dd7bbe3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,7 @@ import ast import logging import sys -from typing import Any, Optional, Protocol +from typing import Any, Optional import orjson import tiktoken @@ -20,11 +20,7 @@ get_engine_supported_data_type, ) from src.utils import trace_cost - - -class AskHistoryLike(Protocol): - question: str - sql: str +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -43,12 +39,6 @@ class AskHistoryLike(Protocol): 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -8. The user does not need to mention exact table or column names. Map natural-language business terms to tables and columns using table names, column names, comments, aliases, descriptions, data types, primary keys, foreign keys, relationships, metrics, and views. -9. Prefer tables, views, or metrics whose descriptions and columns directly match the requested business entities, measures, dimensions, dates, filters, rankings, frequencies, or comparisons. -10. For metric questions such as total, count, average, rate, value, amount, growth, frequency, or distribution, include numeric/measure columns and the grouping/filter/date columns needed to answer the question. -11. For "unique values", "frequency", "distribution", or "count by" questions, include the categorical column being counted and any requested filter/date columns. -12. Prefer canonical/business-ready models, metrics, and views over staging, raw, temporary, backup, load, or test tables when the metadata indicates that distinction. -13. If no table clearly matches, return an empty results list instead of selecting unrelated tables. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -92,7 +82,6 @@ class AskHistoryLike(Protocol): - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. - Match Table names with the definition in the "Create Table" statement. -- Do not choose a column only because it has a similar name; choose it only when its metadata meaning and data type fit the user's question. Good luck! @@ -133,7 +122,7 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline @observe(capture_input=False, capture_output=False) -async def embedding(query: str, embedder: Any, histories: list[AskHistoryLike]) -> dict: +async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: if histories: previous_query_summaries = [history.question for history in histories] @@ -149,10 +138,7 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistoryLike]) @observe(capture_input=False) async def table_retrieval( - embedding: dict, - project_id: str, - tables: list[str] | None, - table_retriever: Any, + embedding: dict, project_id: str, tables: list[str], table_retriever: Any ) -> dict: filters = { "operator": "AND", @@ -171,15 +157,15 @@ async def table_retrieval( query_embedding=embedding.get("embedding"), filters=filters, ) + else: + filters["conditions"].append( + {"field": "name", "operator": "in", "value": tables} + ) - if not tables: - return {"documents": []} - - filters["conditions"].append({"field": "name", "operator": "in", "value": tables}) - return await table_retriever.run( - query_embedding=[], - filters=filters, - ) + return await table_retriever.run( + query_embedding=[], + filters=filters, + ) @observe(capture_input=False) @@ -247,8 +233,6 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: @observe(capture_input=False) def check_using_db_schemas_without_pruning( - query: str, - tables: list[str] | None, construct_db_schemas: list[dict], dbschema_retrieval: list[Document], encoding: tiktoken.Encoding, @@ -321,7 +305,7 @@ def prompt( construct_db_schemas: list[dict], prompt_builder: PromptBuilder, check_using_db_schemas_without_pruning: dict, - histories: list[AskHistoryLike], + histories: list[AskHistory], ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ @@ -361,78 +345,30 @@ def construct_retrieval_results( construct_db_schemas: list[dict], dbschema_retrieval: list[Document], ) -> dict[str, Any]: - fallback_retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] - fallback_result = { - "retrieval_results": fallback_retrieval_results, - "has_calculated_field": check_using_db_schemas_without_pruning[ - "has_calculated_field" - ], - "has_metric": check_using_db_schemas_without_pruning["has_metric"], - "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], - } - if filter_columns_in_tables: - try: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] - except Exception: - logger.warning( - "Column pruning returned invalid output; using full retrieved schemas." - ) - return fallback_result - - if not columns_and_tables_needed: - logger.info( - "Column pruning selected no tables; using full retrieved schemas." - ) - return fallback_result + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + )["results"] # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format reformated_json = {} for table in columns_and_tables_needed: - table_name = table.get("table_name") - table_contents = table.get("table_contents") or {} - if table_name: - reformated_json[table_name] = table_contents + reformated_json[table["table_name"]] = table["table_contents"] columns_and_tables_needed = reformated_json - if not columns_and_tables_needed: - logger.info( - "Column pruning selected no usable tables; using full retrieved schemas." - ) - return fallback_result - tables = set(columns_and_tables_needed.keys()) retrieval_results = [] has_calculated_field = False has_metric = False has_json_field = False - valid_table_names = { - table_schema["name"] for table_schema in construct_db_schemas - } for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - requested_columns = set( - columns_and_tables_needed[table_schema["name"]].get("columns", []) - ) - valid_columns = { - column.get("name") - for column in table_schema.get("columns", []) - if isinstance(column, dict) and column.get("name") - } - selected_columns = requested_columns & valid_columns - if not selected_columns: - logger.info( - "Column pruning selected no valid columns for table %s; using full retrieved schemas.", - table_schema["name"], - ) - return fallback_result - ddl, _has_calculated_field, _has_json_field = build_table_ddl( table_schema, - columns=selected_columns, + columns=set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ), tables=tables, ) if _has_calculated_field: @@ -450,7 +386,6 @@ def construct_retrieval_results( for document in dbschema_retrieval: if document.meta["name"] in columns_and_tables_needed: content = ast.literal_eval(document.content) - valid_table_names.add(content["name"]) if content["type"] == "METRIC": retrieval_results.append( @@ -468,14 +403,6 @@ def construct_retrieval_results( } ) - unknown_tables = tables - valid_table_names - if unknown_tables or not retrieval_results: - logger.info( - "Column pruning selected unknown or unusable tables %s; using full retrieved schemas.", - sorted(unknown_tables), - ) - return fallback_result - return { "retrieval_results": retrieval_results, "has_calculated_field": has_calculated_field, @@ -483,7 +410,16 @@ def construct_retrieval_results( "has_json_field": has_json_field, } else: - return fallback_result + retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] + + return { + "retrieval_results": retrieval_results, + "has_calculated_field": check_using_db_schemas_without_pruning[ + "has_calculated_field" + ], + "has_metric": check_using_db_schemas_without_pruning["has_metric"], + "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], + } ## End of Pipeline @@ -565,7 +501,7 @@ async def run( query: str = "", tables: Optional[list[str]] = None, project_id: Optional[str] = None, - histories: Optional[list[AskHistoryLike]] = None, + histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, ): logger.info("Ask Retrieval pipeline is running...") diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 258944a5fb..aa26fa3f81 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,7 +1,6 @@ import asyncio import logging -import re -from typing import Any, Dict, List, Literal, Optional +from typing import Dict, List, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe @@ -13,13 +12,6 @@ logger = logging.getLogger("wren-ai-service") -NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( - "No relevant data found in the active datasource for this question." -) - -async def _return_value(value): - return value - class AskHistory(BaseModel): sql: str @@ -38,10 +30,6 @@ class AskRequest(BaseRequest): use_dry_plan: bool = False allow_dry_plan_fallback: bool = True custom_instruction: Optional[str] = None - explicit_tables: Optional[list[str]] = Field( - default=None, - validation_alias=AliasChoices("explicit_tables", "explicitTables"), - ) class AskResponse(BaseModel): @@ -87,7 +75,7 @@ class _AskResultResponse(BaseModel): rephrased_question: Optional[str] = None intent_reasoning: Optional[str] = None sql_generation_reasoning: Optional[str] = None - type: Optional[Literal["GENERAL", "TEXT_TO_SQL", "MISLEADING_QUERY"]] = None + type: Optional[Literal["GENERAL", "TEXT_TO_SQL"]] = None retrieved_tables: Optional[List[str]] = None response: Optional[List[AskResult]] = None invalid_sql: Optional[str] = None @@ -117,10 +105,7 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, - pipeline_timeout_seconds: int = 45, - schema_retrieval_timeout_seconds: int = 25, max_histories: int = 5, - max_sql_generation_tables: int = 10, maxsize: int = 1_000_000, ttl: int = 120, ): @@ -128,19 +113,13 @@ def __init__( self._ask_results: Dict[str, AskResultResponse] = TTLCache( maxsize=maxsize, ttl=ttl ) - self._general_streaming_results: Dict[str, str] = TTLCache( - maxsize=maxsize, ttl=ttl - ) self._allow_sql_generation_reasoning = allow_sql_generation_reasoning self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_intent_classification = allow_intent_classification self._allow_sql_diagnosis = allow_sql_diagnosis self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval self._enable_column_pruning = enable_column_pruning - self._pipeline_timeout_seconds = pipeline_timeout_seconds - self._schema_retrieval_timeout_seconds = schema_retrieval_timeout_seconds self._max_histories = max_histories - self._max_sql_generation_tables = max_sql_generation_tables self._max_sql_correction_retries = max_sql_correction_retries def _is_stopped(self, query_id: str, container: dict): @@ -151,663 +130,6 @@ def _is_stopped(self, query_id: str, container: dict): return False - async def _run_with_timeout( - self, - label: str, - awaitable, - *, - timeout_seconds: Optional[int] = None, - ): - timeout = timeout_seconds or self._pipeline_timeout_seconds - try: - return await asyncio.wait_for(awaitable, timeout=timeout) - except TimeoutError: - logger.warning("%s timed out after %s seconds", label, timeout) - raise - - def _is_greeting_query(self, query: str) -> bool: - normalized = re.sub(r"\s+", " ", (query or "").strip().lower()) - greeting_patterns = { - "hi", - "hello", - "hey", - "hii", - "hola", - "good morning", - "good afternoon", - "good evening", - "how are you", - "thanks", - "thank you", - } - return normalized in greeting_patterns - - def _parse_schema_tables(self, table_ddls: list[str]) -> list[dict[str, Any]]: - tables: list[dict[str, Any]] = [] - for ddl in table_ddls or []: - if not isinstance(ddl, str): - continue - table_match = re.search( - r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", - ddl, - flags=re.IGNORECASE, - ) - if not table_match: - continue - - table_name = next( - (value for value in table_match.groupdict().values() if value), - None, - ) - if not table_name: - continue - body_start = table_match.end() - depth = 1 - body_end = body_start - while body_end < len(ddl) and depth > 0: - if ddl[body_end] == "(": - depth += 1 - elif ddl[body_end] == ")": - depth -= 1 - body_end += 1 - - columns: list[dict[str, str]] = [] - for line in ddl[body_start : body_end - 1].splitlines(): - stripped = line.strip().rstrip(",") - if not stripped or stripped.startswith(("--", "/*")): - continue - if re.match( - r"^(?:PRIMARY|FOREIGN|CONSTRAINT|UNIQUE|KEY)\b", - stripped, - flags=re.IGNORECASE, - ): - continue - - column_match = re.match( - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_$]*))" - r"\s+(?P[A-Za-z0-9_(),]+)", - stripped, - ) - if column_match: - column_name = next( - (value - for key, value in column_match.groupdict().items() - if key != "type" and value - ), - None, - ) - if not column_name: - continue - column_type = column_match.group("type") or "" - columns.append( - { - "name": str(column_name), - "type": str(column_type).lower(), - } - ) - - tables.append({"name": table_name, "columns": columns}) - - return tables - - def _build_greeting_response(self, query: str) -> str: - return ( - f"Hi. I can help with questions about your active datasource and Wren AI.\n\n" - f"Try a data question like:\n" - f"- Show monthly trends for the last 12 months\n" - f"- Compare totals by category\n" - f"- Which records occur most often?\n\n" - f"If you want, ask a database question directly instead of `{query}`." - ) - - def _extract_pipeline_reply(self, result: dict, key: str) -> str: - payload = result.get(key) - if isinstance(payload, tuple): - payload = payload[0] - - if isinstance(payload, dict): - replies = payload.get("replies") or [] - if replies and isinstance(replies[0], str): - return replies[0] - - return "" - - def _extract_retrieval_documents(self, retrieval_result: dict) -> list[dict]: - construct_result = retrieval_result.get("construct_retrieval_results", {}) - documents = construct_result.get("retrieval_results", []) - if not isinstance(documents, list): - logger.warning("Schema retrieval returned invalid document payload") - return [] - - valid_documents = [] - for document in documents: - if not isinstance(document, dict): - logger.warning("Ignoring malformed retrieval document: %s", document) - continue - if not document.get("table_name") and not document.get("table_ddl"): - logger.warning("Ignoring retrieval document without table metadata") - continue - valid_documents.append(document) - - return valid_documents - - def _extract_retrieval_metadata( - self, retrieval_result: dict - ) -> tuple[list[dict], list[str], list[str]]: - documents = self._extract_retrieval_documents(retrieval_result) - return documents, *self._metadata_from_documents(documents) - - def _metadata_from_documents( - self, documents: list[dict] - ) -> tuple[list[str], list[str]]: - table_names = [ - table_name - for document in documents - if isinstance(table_name := document.get("table_name"), str) - and table_name.strip() - ] - table_ddls = [ - table_ddl - for document in documents - if isinstance(table_ddl := document.get("table_ddl"), str) - and table_ddl.strip() - ] - return table_names, table_ddls - - async def _complete_sql_generation_context( - self, - *, - query: str, - project_id: Optional[str], - documents: list[dict], - table_names: list[str], - table_ddls: list[str], - ) -> tuple[list[dict], list[str], list[str], dict]: - if not table_names or "db_schema_retrieval" not in self._pipelines: - return documents, table_names, table_ddls, {} - - selected_table_names = list(dict.fromkeys(table_names)) - try: - retrieval_result = await self._run_with_timeout( - "Complete selected schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=query, - tables=selected_table_names, - project_id=project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min(self._schema_retrieval_timeout_seconds, 30), - ) - except Exception as error: - logger.warning( - "Complete selected schema retrieval failed; using existing retrieval context. project_id=%s tables=%s error=%s", - project_id, - selected_table_names, - error, - ) - return documents, table_names, table_ddls, {} - - complete_documents, complete_table_names, complete_table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - if not complete_documents: - logger.warning( - "Complete selected schema retrieval returned no documents; using existing retrieval context. project_id=%s tables=%s", - project_id, - selected_table_names, - ) - return documents, table_names, table_ddls, {} - - logger.info( - "Completed SQL generation context with full schemas for project_id %s tables=%s", - project_id, - complete_table_names, - ) - return ( - complete_documents, - complete_table_names, - complete_table_ddls, - retrieval_result.get("construct_retrieval_results", {}), - ) - - def _is_visualization_request(self, query: str) -> bool: - normalized = (query or "").lower() - return bool( - re.search( - r"\b(?:chart|graph|plot|visuali[sz]e|dashboard|bar|line|pie|donut|" - r"scatter|histogram|heatmap|trend|trends|distribution)\b", - normalized, - ) - ) - - def _get_metadata_question_kind(self, query: str) -> str | None: - normalized = re.sub(r"\s+", " ", (query or "").lower()).strip() - if not normalized: - return None - - if self._is_visualization_request(normalized): - return None - - if re.search(r"\b(?:row|rows|record|records)\s+count\b", normalized): - return None - - relationship_patterns = ( - r"\b(?:relationships?|relations?|joins?|foreign keys?|primary keys?)\b", - r"\b(?:how|what|which|show|list|describe)\b.*\b(?:tables?|models?)\b.*\b(?:connected|related|joined)\b", - ) - if any(re.search(pattern, normalized) for pattern in relationship_patterns): - return "relationships" - - table_count_patterns = ( - r"\b(?:how many|count|number of)\b.*\b(?:tables?|models?)\b", - r"\b(?:tables?|models?)\b.*\b(?:count|number)\b", - ) - if any(re.search(pattern, normalized) for pattern in table_count_patterns): - return "table_count" - - column_count_patterns = ( - r"\b(?:how many|count|number of)\b.*\b(?:columns?|fields?)\b", - r"\b(?:columns?|fields?)\b.*\b(?:count|number)\b", - ) - if any(re.search(pattern, normalized) for pattern in column_count_patterns): - return "column_count" - - schema_patterns = ( - r"\b(?:what|show|display|describe|list)\b.*\b(?:schema|metadata)\b", - r"\b(?:schema|metadata)\b.*\b(?:of|for|in)\b", - ) - if any(re.search(pattern, normalized) for pattern in schema_patterns): - return "schema" - - explicit_column_patterns = ( - r"\b(?:what|which|list|show|display|give|describe)\b.*\b(?:columns?|fields?)\b", - r"\b(?:columns?|fields?)\b.*\b(?:available|present|there|exist|schema|metadata)\b", - ) - if any(re.search(pattern, normalized) for pattern in explicit_column_patterns): - return "columns" - - table_patterns = ( - r"\b(?:what|which|list|show|display|give)\b.*\b(?:tables?|models?)\b", - r"\b(?:tables?|models?)\b.*\b(?:available|present|there|exist|in this datasource|in the datasource)\b", - r"\b(?:datasource|database|semantic layer|semantic model)\b.*\b(?:tables?|models?)\b", - ) - if any(re.search(pattern, normalized) for pattern in table_patterns): - return "tables" - - return None - - def _find_metadata_table_matches( - self, query: str, tables: list[dict[str, Any]] - ) -> list[dict[str, Any]]: - query_key = self._normalize_schema_token(query) - if not query_key: - return [] - - matches: list[tuple[int, dict[str, Any]]] = [] - for table in tables: - table_name = str(table.get("name") or "") - if not table_name: - continue - short_name = re.split(r"[.$]", table_name)[-1] - normalized_name = self._normalize_schema_token(table_name) - normalized_short_name = self._normalize_schema_token(short_name) - - score = 0 - if normalized_name and normalized_name in query_key: - score = 100 + len(normalized_name) - elif normalized_short_name and normalized_short_name in query_key: - score = 80 + len(normalized_short_name) - - if score: - matches.append((score, table)) - - return [ - table - for _, table in sorted(matches, key=lambda item: item[0], reverse=True) - ] - - def _format_metadata_table_list( - self, tables: list[dict[str, Any]], *, max_tables: int = 120 - ) -> str: - if not tables: - return "I couldn't find any deployed tables in the active datasource metadata." - - sorted_tables = sorted( - {str(table.get("name")) for table in tables if table.get("name")}, - key=str.lower, - ) - shown_tables = sorted_tables[:max_tables] - lines = [ - f"The active datasource has {len(sorted_tables)} deployed table" - f"{'' if len(sorted_tables) == 1 else 's'}:" - ] - lines.extend(f"- {table_name}" for table_name in shown_tables) - if len(sorted_tables) > max_tables: - lines.append( - f"- ...and {len(sorted_tables) - max_tables} more tables." - ) - return "\n".join(lines) - - def _format_metadata_columns( - self, - query: str, - tables: list[dict[str, Any]], - *, - max_tables: int = 25, - max_columns_per_table: int = 60, - ) -> str: - if not tables: - return "I couldn't find any deployed columns in the active datasource metadata." - - matched_tables = self._find_metadata_table_matches(query, tables) - selected_tables = matched_tables or sorted( - tables, key=lambda table: str(table.get("name") or "").lower() - ) - selected_tables = selected_tables[:max_tables] - - heading = ( - "Columns available in the matched deployed table" - if matched_tables and len(selected_tables) == 1 - else "Columns available in the active datasource metadata" - ) - lines = [f"{heading}:"] - for table in selected_tables: - table_name = str(table.get("name") or "unknown_table") - columns = [ - column - for column in table.get("columns", []) - if isinstance(column, dict) and column.get("name") - ] - if not columns: - lines.append(f"- {table_name}: no columns found") - continue - - column_parts = [] - for column in columns[:max_columns_per_table]: - column_name = str(column.get("name")) - column_type = str(column.get("type") or "").upper() - column_parts.append( - f"{column_name} ({column_type})" if column_type else column_name - ) - if len(columns) > max_columns_per_table: - column_parts.append( - f"...and {len(columns) - max_columns_per_table} more" - ) - lines.append(f"- {table_name}: {', '.join(column_parts)}") - - if len(tables) > max_tables and not matched_tables: - lines.append(f"- ...and {len(tables) - max_tables} more tables.") - - return "\n".join(lines) - - def _extract_metadata_relationships(self, table_ddls: list[str]) -> list[str]: - relationships: list[str] = [] - for ddl in table_ddls or []: - if not isinstance(ddl, str): - continue - table_match = re.search( - r'\bCREATE\s+TABLE\s+(?:"(?P[^"]+)"|' - r"\[(?P[^\]]+)\]|`(?P[^`]+)`|" - r"(?P[A-Za-z_][A-Za-z0-9_.$]*))\s*\(", - ddl, - flags=re.IGNORECASE, - ) - if not table_match: - continue - source_table = next( - (value for value in table_match.groupdict().values() if value), - "unknown_table", - ) - - for relationship_match in re.finditer( - r"FOREIGN\s+KEY\s*\((?P[^)]+)\)\s+REFERENCES\s+" - r'(?:"(?P[^"]+)"|\[(?P[^\]]+)\]|' - r"`(?P[^`]+)`|(?P[A-Za-z_][A-Za-z0-9_.$]*))" - r"\s*\((?P[^)]+)\)", - ddl, - flags=re.IGNORECASE, - ): - target_table = next( - ( - value - for key, value in relationship_match.groupdict().items() - if key - in { - "quoted", - "bracketed", - "backticked", - "bare", - } - and value - ), - "unknown_table", - ) - source_columns = relationship_match.group("source_columns") - target_columns = relationship_match.group("target_columns") - relationships.append( - f"{source_table}({source_columns}) -> " - f"{target_table}({target_columns})" - ) - - return sorted(set(relationships), key=str.lower) - - def _format_metadata_relationships(self, table_ddls: list[str]) -> str: - relationships = self._extract_metadata_relationships(table_ddls) - if not relationships: - return ( - "I couldn't find explicit relationships or foreign keys in the " - "active datasource metadata." - ) - - lines = [ - f"The active datasource metadata has {len(relationships)} " - f"relationship{'' if len(relationships) == 1 else 's'}:" - ] - lines.extend(f"- {relationship}" for relationship in relationships[:120]) - if len(relationships) > 120: - lines.append(f"- ...and {len(relationships) - 120} more relationships.") - return "\n".join(lines) - - def _format_metadata_schema( - self, query: str, tables: list[dict[str, Any]], table_ddls: list[str] - ) -> str: - matched_tables = self._find_metadata_table_matches(query, tables) - selected_tables = matched_tables or sorted( - tables, key=lambda table: str(table.get("name") or "").lower() - ) - selected_tables = selected_tables[:20] - if not selected_tables: - return "I couldn't find schema details in the active datasource metadata." - - lines = ["Schema details from the active datasource metadata:"] - for table in selected_tables: - table_name = str(table.get("name") or "unknown_table") - columns = [ - column - for column in table.get("columns", []) - if isinstance(column, dict) and column.get("name") - ] - lines.append(f"- {table_name}") - if columns: - column_parts = [] - for column in columns[:60]: - column_name = str(column.get("name")) - column_type = str(column.get("type") or "").upper() - column_parts.append( - f"{column_name} ({column_type})" - if column_type - else column_name - ) - if len(columns) > 60: - column_parts.append(f"...and {len(columns) - 60} more") - lines.append(f" Columns: {', '.join(column_parts)}") - else: - lines.append(" Columns: no columns found") - - relationships = self._extract_metadata_relationships(table_ddls) - if relationships: - lines.append("Relationships:") - lines.extend(f"- {relationship}" for relationship in relationships[:40]) - if len(relationships) > 40: - lines.append(f"- ...and {len(relationships) - 40} more relationships.") - - return "\n".join(lines) - - def _format_metadata_table_count(self, tables: list[dict[str, Any]]) -> str: - table_names = {str(table.get("name")) for table in tables if table.get("name")} - return ( - f"The active datasource has {len(table_names)} deployed table" - f"{'' if len(table_names) == 1 else 's'}." - ) - - def _format_metadata_column_count( - self, query: str, tables: list[dict[str, Any]] - ) -> str: - matched_tables = self._find_metadata_table_matches(query, tables) - selected_tables = matched_tables or tables - total_columns = sum( - len( - [ - column - for column in table.get("columns", []) - if isinstance(column, dict) and column.get("name") - ] - ) - for table in selected_tables - ) - if matched_tables and len(selected_tables) == 1: - table_name = str(selected_tables[0].get("name") or "the matched table") - return f"{table_name} has {total_columns} deployed columns." - return ( - f"The active datasource metadata has {total_columns} deployed columns " - f"across {len(selected_tables)} table" - f"{'' if len(selected_tables) == 1 else 's'}." - ) - - def _build_metadata_response( - self, query: str, table_ddls: list[str], table_names: list[str] - ) -> str: - kind = self._get_metadata_question_kind(query) - parsed_tables = self._parse_schema_tables(table_ddls) - - if not parsed_tables and table_names: - parsed_tables = [ - {"name": table_name, "columns": []} for table_name in table_names - ] - - if kind == "schema": - return self._format_metadata_schema(query, parsed_tables, table_ddls) - if kind == "relationships": - return self._format_metadata_relationships(table_ddls) - if kind == "table_count": - return self._format_metadata_table_count(parsed_tables) - if kind == "column_count": - return self._format_metadata_column_count(query, parsed_tables) - if kind == "columns": - return self._format_metadata_columns(query, parsed_tables) - return self._format_metadata_table_list(parsed_tables) - - def _normalize_schema_token(self, value: str) -> str: - return re.sub(r"[^a-z0-9]", "", (value or "").lower()) - - def _prune_sql_generation_context( - self, - query: str, - documents: list[dict], - table_names: list[str], - table_ddls: list[str], - *, - max_tables: int = 10, - ) -> tuple[list[dict], list[str], list[str]]: - if len(documents) <= max_tables: - return documents, table_names, table_ddls - - pruned_documents = documents[:max_tables] - pruned_table_names = table_names[:max_tables] - pruned_table_ddls = table_ddls[:max_tables] - logger.info( - "Pruned SQL generation context from %s to %s tables for query: %s", - len(documents), - len(pruned_documents), - query, - ) - return pruned_documents, pruned_table_names, pruned_table_ddls - - def _is_valid_select_sql(self, sql: Optional[str]) -> bool: - if not isinstance(sql, str): - return False - - normalized = re.sub(r"\s+", " ", sql.strip()) - if not normalized: - return False - - return bool(re.match(r"^(?:WITH|SELECT)\b", normalized, flags=re.IGNORECASE)) - - def _build_ask_result_from_sql(self, sql: Optional[str]) -> Optional[AskResult]: - if not self._is_valid_select_sql(sql): - return None - return AskResult(sql=sql.strip(), type="llm") - - def _build_validated_ask_result_from_sql( - self, - sql: Optional[str], - table_ddls: list[str], - query: str | None = None, - ) -> Optional[AskResult]: - return self._build_ask_result_from_sql(sql) - - def _build_failed_text_to_sql_response( - self, - trace_id: Optional[str], - message: str, - *, - rephrased_question: Optional[str] = None, - intent_reasoning: Optional[str] = None, - retrieved_tables: Optional[list[str]] = None, - sql_generation_reasoning: Optional[str] = None, - invalid_sql: Optional[str] = None, - is_followup: bool = False, - code: Literal["NO_RELEVANT_DATA", "NO_RELEVANT_SQL", "OTHERS"] = "NO_RELEVANT_SQL", - ) -> AskResultResponse: - return AskResultResponse( - status="failed", - type="TEXT_TO_SQL", - error=AskError(code=code, message=message), - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=retrieved_tables, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=invalid_sql, - trace_id=trace_id, - is_followup=is_followup, - ) - - def _build_no_relevant_active_datasource_response( - self, - trace_id: Optional[str], - *, - rephrased_question: Optional[str] = None, - intent_reasoning: Optional[str] = None, - retrieved_tables: Optional[list[str]] = None, - sql_generation_reasoning: Optional[str] = None, - is_followup: bool = False, - ) -> AskResultResponse: - return self._build_failed_text_to_sql_response( - trace_id, - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=retrieved_tables, - sql_generation_reasoning=sql_generation_reasoning, - invalid_sql=None, - is_followup=is_followup, - code="NO_RELEVANT_DATA", - ) - @observe(name="Ask Question") @trace_metadata async def ask( @@ -827,22 +149,6 @@ async def ask( } query_id = ask_request.query_id - if not query_id: - raise ValueError("query_id is required for ask service execution") - - user_query = (ask_request.query or "").strip() - if not user_query: - self._ask_results[query_id] = self._build_failed_text_to_sql_response( - trace_id, - "Question is required", - code="OTHERS", - ) - results["metadata"]["error_type"] = "OTHERS" - results["metadata"]["error_message"] = "Question is required" - results["metadata"]["type"] = "TEXT_TO_SQL" - return results - - logger.info(f"Ask pipeline started for query_id: {query_id}") histories = ask_request.histories[: self._max_histories][ ::-1 ] # reverse the order of histories @@ -852,10 +158,7 @@ async def ask( sql_samples = [] instructions = [] api_results = [] - documents = [] table_names = [] - table_ddls = [] - _retrieval_result = {} error_message = None invalid_sql = None allow_sql_generation_reasoning = ( @@ -873,12 +176,9 @@ async def ask( use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback sql_knowledge = None - understanding_timeout_seconds = min(self._pipeline_timeout_seconds, 12) - planning_timeout_seconds = min(self._pipeline_timeout_seconds, 15) - generation_timeout_seconds = min(self._pipeline_timeout_seconds, 30) - correction_timeout_seconds = min(self._pipeline_timeout_seconds, 15) + try: - sql_user_query = user_query + user_query = ask_request.query # ask status can be understanding, searching, generating, finished, failed, stopped # we will need to handle business logic for each status @@ -889,196 +189,61 @@ async def ask( is_followup=True if histories else False, ) - if self._is_greeting_query(user_query): - self._general_streaming_results[query_id] = ( - self._build_greeting_response(user_query) - ) - - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - trace_id=trace_id, - is_followup=True if histories else False, - general_type="USER_GUIDE", - ) - results["metadata"]["type"] = "GENERAL" - return results - - metadata_question_kind = self._get_metadata_question_kind(user_query) - if metadata_question_kind: - self._ask_results[query_id] = AskResultResponse( - status="searching", - type="GENERAL", - rephrased_question=user_query, - intent_reasoning=( - "Basic datasource metadata question detected; " - "retrieving deployed schema metadata directly." - ), - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - retrieval_result = await self._run_with_timeout( - "Metadata schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query="", - project_id=ask_request.project_id, - histories=[], - enable_column_pruning=False, - ), - timeout_seconds=min( - self._schema_retrieval_timeout_seconds, - self._pipeline_timeout_seconds, - 20, - ), - ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - metadata_answer = self._build_metadata_response( - user_query, table_ddls, table_names - ) - self._general_streaming_results[query_id] = metadata_answer - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=user_query, - intent_reasoning=( - "Answered from active datasource deployed metadata " - "without SQL generation." - ), - retrieved_tables=table_names, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - results["metadata"]["type"] = "GENERAL" - results["metadata"]["metadata_question_kind"] = ( - metadata_question_kind - ) - results["metadata"]["retrieved_table_count"] = len(documents) - return results - - historical_question_result = [] - if not api_results: - if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = AskResultResponse( - status="understanding", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) - - try: - historical_question = await self._run_with_timeout( - "Historical question retrieval", - self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ), - timeout_seconds=min(understanding_timeout_seconds, 10), - ) + historical_question = await self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ) - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - except TimeoutError as exc: - logger.warning( - "Historical question retrieval timed out; continuing without history match. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] - valid_historical_results = [] - for result in historical_question_result: - sql_statement = result.get("statement") - if not self._is_valid_select_sql(sql_statement): - logger.warning( - "Ignoring historical question without valid SQL for query_id %s", - query_id, - ) - continue - valid_historical_results.append( + if historical_question_result: + api_results = [ AskResult( **{ - "sql": sql_statement.strip(), + "sql": result.get("statement"), "type": "view" if result.get("viewId") else "llm", "viewId": result.get("viewId"), } ) - ) - - if valid_historical_results: - api_results = valid_historical_results + for result in historical_question_result + ] sql_generation_reasoning = "" - elif not api_results: + else: # Run both pipeline operations concurrently - try: - sql_samples_task, instructions_task = await self._run_with_timeout( - "SQL pair and instruction retrieval", - asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", - ), - ), - timeout_seconds=understanding_timeout_seconds, - ) + sql_samples_task, instructions_task = await asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), + ) - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] - ) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) - except TimeoutError as exc: - logger.warning( - "SQL pair and instruction retrieval timed out; continuing without optional examples. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, - ) - sql_samples = [] - instructions = [] + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] + ) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) if self._allow_intent_classification: - try: - intent_classification_result = ( - await self._run_with_timeout( - "Intent classification", - self._pipelines["intent_classification"].run( - query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, - ), - timeout_seconds=understanding_timeout_seconds, - ) - ).get("post_process", {}) - except TimeoutError as exc: - logger.warning( - "Intent classification timed out; continuing with TEXT_TO_SQL. query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - exc, + intent_classification_result = ( + await self._pipelines["intent_classification"].run( + query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, ) - intent_classification_result = { - "intent": "TEXT_TO_SQL", - "rephrased_question": user_query, - "reasoning": "Intent classification timed out; using SQL generation.", - "db_schemas": [], - } + ).get("post_process", {}) intent = intent_classification_result.get("intent") rephrased_question = intent_classification_result.get( "rephrased_question" @@ -1088,8 +253,6 @@ async def ask( if rephrased_question: user_query = rephrased_question - sql_user_query = user_query - if intent == "MISLEADING_QUERY": asyncio.create_task( self._pipelines["misleading_assistance"].run( @@ -1099,7 +262,7 @@ async def ask( "db_schemas" ), language=ask_request.configurations.language, - query_id=query_id, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, ) ) @@ -1116,27 +279,36 @@ async def ask( results["metadata"]["type"] = "MISLEADING_QUERY" return results elif intent == "GENERAL": - intent_reasoning = ( - f"{intent_reasoning or ''}\n" - "Classifier returned GENERAL, but this ask flow " - "treats non-schema, non-guide questions as data " - "retrieval requests so they continue through " - "semantic retrieval and SQL generation." + asyncio.create_task( + self._pipelines["data_assistance"].run( + query=user_query, + histories=histories, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) ) + self._ask_results[query_id] = AskResultResponse( - status="understanding", - type="TEXT_TO_SQL", + status="finished", + type="GENERAL", rephrased_question=rephrased_question, intent_reasoning=intent_reasoning, trace_id=trace_id, is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", ) + results["metadata"]["type"] = "GENERAL" + return results elif intent == "USER_GUIDE": asyncio.create_task( self._pipelines["user_guide_assistance"].run( query=user_query, language=ask_request.configurations.language, - query_id=query_id, + query_id=ask_request.query_id, custom_instruction=ask_request.custom_instruction, ) ) @@ -1161,11 +333,7 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - if ( - not self._is_stopped(query_id, self._ask_results) - and not api_results - and not documents - ): + if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -1175,81 +343,38 @@ async def ask( is_followup=True if histories else False, ) - try: - retrieval_result = await self._run_with_timeout( - "Schema retrieval", - self._pipelines["db_schema_retrieval"].run( - query=sql_user_query, - tables=None, - histories=histories, - project_id=ask_request.project_id, - enable_column_pruning=enable_column_pruning, - ), - timeout_seconds=self._schema_retrieval_timeout_seconds, - ) - except TimeoutError as error: - logger.warning( - "Schema retrieval timed out for data query; not loading full project schema. " - "query_id=%s project_id=%s error=%s", - query_id, - ask_request.project_id, - error, - ) - retrieval_result = {"construct_retrieval_results": {}} + retrieval_result = await self._pipelines["db_schema_retrieval"].run( + query=user_query, + histories=histories, + project_id=ask_request.project_id, + enable_column_pruning=enable_column_pruning, + ) _retrieval_result = retrieval_result.get( "construct_retrieval_results", {} ) - documents, table_names, table_ddls = ( - self._extract_retrieval_metadata(retrieval_result) - ) - logger.info( - "Retrieved tables for query_id %s: %s", query_id, table_names - ) + documents = _retrieval_result.get("retrieval_results", []) + table_names = [document.get("table_name") for document in documents] + table_ddls = [document.get("table_ddl") for document in documents] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - is_followup=True if histories else False, - ) + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_DATA", + message="No relevant data", + ), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, ) results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) results["metadata"]["type"] = "TEXT_TO_SQL" return results - if documents and not api_results: - documents, table_names, table_ddls = self._prune_sql_generation_context( - sql_user_query, - documents, - table_names, - table_ddls, - max_tables=self._max_sql_generation_tables, - ) - ( - documents, - table_names, - table_ddls, - completed_retrieval_result, - ) = await self._complete_sql_generation_context( - query=sql_user_query, - project_id=ask_request.project_id, - documents=documents, - table_names=table_names, - table_ddls=table_ddls, - ) - if completed_retrieval_result: - _retrieval_result = completed_retrieval_result - - sql_generation_histories = histories - if ( not self._is_stopped(query_id, self._ask_results) and not api_results @@ -1265,57 +390,29 @@ async def ask( is_followup=True if histories else False, ) - if sql_generation_histories: - try: - sql_generation_reasoning = ( - await self._run_with_timeout( - "Follow-up SQL generation reasoning", - self._pipelines[ - "followup_sql_generation_reasoning" - ].run( - query=sql_user_query, - contexts=table_ddls, - valid_table_names=table_names, - histories=sql_generation_histories, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), - timeout_seconds=planning_timeout_seconds, - ) - ).get("post_process", {}) - except Exception as reasoning_error: - logger.warning( - "Follow-up SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", - query_id, - reasoning_error, + if histories: + sql_generation_reasoning = ( + await self._pipelines["followup_sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, ) - sql_generation_reasoning = "" + ).get("post_process", {}) else: - try: - sql_generation_reasoning = ( - await self._run_with_timeout( - "SQL generation reasoning", - self._pipelines["sql_generation_reasoning"].run( - query=sql_user_query, - contexts=table_ddls, - valid_table_names=table_names, - sql_samples=sql_samples, - instructions=instructions, - configuration=ask_request.configurations, - query_id=query_id, - ), - timeout_seconds=planning_timeout_seconds, - ) - ).get("post_process", {}) - except Exception as reasoning_error: - logger.warning( - "SQL generation reasoning failed for query_id %s; continuing without reasoning: %s", - query_id, - reasoning_error, + sql_generation_reasoning = ( + await self._pipelines["sql_generation_reasoning"].run( + query=user_query, + contexts=table_ddls, + sql_samples=sql_samples, + instructions=instructions, + configuration=ask_request.configurations, + query_id=query_id, ) - sql_generation_reasoning = "" + ).get("post_process", {}) self._ask_results[query_id] = AskResultResponse( status="planning", @@ -1340,34 +437,21 @@ async def ask( is_followup=True if histories else False, ) - try: - sql_functions, sql_knowledge = await self._run_with_timeout( - "SQL helper retrieval", - asyncio.gather( - ( - self._pipelines["sql_functions_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_functions_retrieval - else _return_value([]) - ), - ( - self._pipelines["sql_knowledge_retrieval"].run( - project_id=ask_request.project_id, - ) - if allow_sql_knowledge_retrieval - else _return_value(None) - ), - ), - timeout_seconds=min(self._pipeline_timeout_seconds, 10), + if allow_sql_functions_retrieval: + sql_functions = await self._pipelines[ + "sql_functions_retrieval" + ].run( + project_id=ask_request.project_id, ) - except TimeoutError as helper_timeout: - logger.warning( - "SQL helper retrieval timed out for query_id %s; continuing with schema only: %s", - query_id, - helper_timeout, + else: + sql_functions = [] + + if allow_sql_knowledge_retrieval: + sql_knowledge = await self._pipelines[ + "sql_knowledge_retrieval" + ].run( + project_id=ask_request.project_id, ) - sql_functions, sql_knowledge = [], None has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -1375,96 +459,65 @@ async def ask( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - try: - if sql_generation_histories: - text_to_sql_generation_results = await self._run_with_timeout( - "Follow-up SQL generation", - self._pipelines["followup_sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - valid_table_names=table_names, - sql_generation_reasoning=sql_generation_reasoning, - histories=sql_generation_histories, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), - timeout_seconds=generation_timeout_seconds, - ) - else: - text_to_sql_generation_results = await self._run_with_timeout( - "SQL generation", - self._pipelines["sql_generation"].run( - query=sql_user_query, - contexts=table_ddls, - valid_table_names=table_names, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - ), - timeout_seconds=generation_timeout_seconds, - ) - except TimeoutError as generation_timeout: - logger.warning( - "SQL generation timed out for query_id %s: %s", - query_id, - generation_timeout, + if histories: + text_to_sql_generation_results = await self._pipelines[ + "followup_sql_generation" + ].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + ) + else: + text_to_sql_generation_results = await self._pipelines[ + "sql_generation" + ].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, ) - text_to_sql_generation_results = { - "post_process": { - "valid_generation_result": None, - "invalid_generation_result": None, - } - } - error_message = str(generation_timeout) if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" ]: - if ask_result := self._build_validated_ask_result_from_sql( - sql_valid_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - else: - invalid_sql = sql_valid_result.get("sql") - error_message = ( - "SQL generation did not produce SQL that matches the active datasource schema and question intent." + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } ) + ] elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] in { - "TIME_OUT", - "UNSUPPORTED_SQL", - }: - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) + if failed_dry_run_result["type"] == "TIME_OUT": break original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] - sql_diagnosis_reasoning = None current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -1479,70 +532,53 @@ async def ask( ) if allow_sql_diagnosis: - sql_diagnosis_results = await self._run_with_timeout( - "SQL diagnosis", - self._pipelines["sql_diagnosis"].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_request.configurations.language, - ), - timeout_seconds=correction_timeout_seconds, + sql_diagnosis_results = await self._pipelines[ + "sql_diagnosis" + ].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_request.configurations.language, ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") - correction_error_message = error_message - if sql_diagnosis_reasoning: - correction_error_message = ( - f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" - ) - - sql_correction_results = await self._run_with_timeout( - "SQL correction", - self._pipelines["sql_correction"].run( - contexts=table_ddls, - valid_table_names=table_names, - instructions=instructions, - invalid_generation_result={ - "original_sql": original_sql, - "sql": invalid_sql, - "error": correction_error_message, - }, - project_id=ask_request.project_id, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - query=sql_user_query, - ), - timeout_seconds=correction_timeout_seconds, + sql_correction_results = await self._pipelines[ + "sql_correction" + ].run( + contexts=table_ddls, + instructions=instructions, + invalid_generation_result={ + "sql": original_sql, + "error": sql_diagnosis_reasoning + if allow_sql_diagnosis + else error_message, + }, + project_id=ask_request.project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, ) if valid_generation_result := sql_correction_results[ "post_process" ]["valid_generation_result"]: - if ask_result := self._build_validated_ask_result_from_sql( - valid_generation_result.get("sql"), - table_ddls, - sql_user_query, - ): - api_results = [ask_result] - break - invalid_sql = valid_generation_result.get("sql") - error_message = ( - "SQL correction did not produce SQL that matches the active datasource schema and question intent." - ) + api_results = [ + AskResult( + **{ + "sql": valid_generation_result.get("sql"), + "type": "llm", + } + ) + ] + break failed_dry_run_result = sql_correction_results["post_process"][ "invalid_generation_result" ] - invalid_sql = failed_dry_run_result.get("sql", invalid_sql) - error_message = failed_dry_run_result.get( - "error", error_message - ) if api_results: if not self._is_stopped(query_id, self._ask_results): @@ -1562,28 +598,23 @@ async def ask( else: logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): - self._ask_results[query_id] = ( - self._build_no_relevant_active_datasource_response( - trace_id, - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - retrieved_tables=table_names, - sql_generation_reasoning=sql_generation_reasoning, - is_followup=True if histories else False, - ) - ) - if error_message or invalid_sql: - logger.info( - "Suppressed technical SQL failure for query_id %s. " - "error=%s invalid_sql=%s", - query_id, - error_message, - invalid_sql, + self._ask_results[query_id] = AskResultResponse( + status="failed", + type="TEXT_TO_SQL", + error=AskError( + code="NO_RELEVANT_SQL", + message=error_message or "No relevant SQL", + ), + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + retrieved_tables=table_names, + sql_generation_reasoning=sql_generation_reasoning, + invalid_sql=invalid_sql, + trace_id=trace_id, + is_followup=True if histories else False, ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) + results["metadata"]["error_type"] = "NO_RELEVANT_SQL" + results["metadata"]["error_message"] = error_message results["metadata"]["type"] = "TEXT_TO_SQL" return results @@ -1637,13 +668,6 @@ async def get_ask_streaming_result( self, query_id: str, ): - if general_response := self._general_streaming_results.get(query_id): - event = SSEEvent( - data=SSEEvent.SSEEventMessage(message=general_response), - ) - yield event.serialize() - return - if self._ask_results.get(query_id): _pipeline_name = "" if self._ask_results.get(query_id).type == "GENERAL": From d1dc6042a9d6f75e4e0a48b846eda64db2b69886 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 15:32:31 +0530 Subject: [PATCH 0628/1087] Restore legacy ask SQL flow --- .../generation/followup_sql_generation.py | 46 ++------- .../followup_sql_generation_reasoning.py | 29 +----- .../pipelines/generation/sql_correction.py | 62 ++---------- .../pipelines/generation/sql_generation.py | 14 --- .../generation/sql_generation_reasoning.py | 19 +--- .../src/pipelines/generation/utils/sql.py | 76 ++------------- wren-ai-service/src/web/v1/services/ask.py | 3 + .../services/test_metadata_grounding.py | 94 ------------------- 8 files changed, 29 insertions(+), 314 deletions(-) delete mode 100644 wren-ai-service/tests/pytest/services/test_metadata_grounding.py 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 b013f9b25c..35cfb8fccf 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -1,8 +1,6 @@ -from __future__ import annotations - import logging import sys -from typing import TYPE_CHECKING, Any +from typing import Any from hamilton import base from hamilton.async_driver import AsyncDriver @@ -26,11 +24,7 @@ from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -40,24 +34,11 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. -### TARGET DATA SOURCE ### -{{ data_source }} - -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -interpreting intent and generating SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -### VALID TABLE NAMES ### -Only use these exact table names from the schema. Do not invent, rename, singularize, -pluralize, or add catalog/schema prefixes unless the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -97,14 +78,6 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -### INTENT AND SCHEMA GROUNDING ### -Interpret the user's business terms by matching them to explicit tables, columns, -metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Never reuse table -or column names from SQL SAMPLES or chat history unless those exact names also -appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. -Only apply aggregate functions to columns whose active metadata type supports that -operation. - ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -119,7 +92,6 @@ def prompt( documents: list[str], sql_generation_reasoning: str, prompt_builder: PromptBuilder, - data_source: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -127,13 +99,10 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, - data_source=data_source, documents=documents, - valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -162,7 +131,6 @@ async def generate_sql_in_followup( generator: Any, histories: list[AskHistory], generator_name: str, - data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: history_messages = construct_ask_history_messages(histories) @@ -178,7 +146,6 @@ async def generate_sql_in_followup( async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, - documents: list[str], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -241,18 +208,19 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") - metadata = await retrieve_metadata(project_id or "", self._retriever) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 4d4402ea69..42b28c5b8f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -1,9 +1,7 @@ -from __future__ import annotations - import asyncio import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -19,34 +17,17 @@ ) from src.utils import trace_cost from src.web.v1.services import Configuration - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") sql_generation_reasoning_user_prompt_template = """ -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -planning SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only mention these exact table names in the reasoning plan. Do not invent, rename, -singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless -the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -91,12 +72,10 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - valid_table_names=valid_table_names or [], histories=histories, sql_samples=sql_samples, instructions=construct_instructions( @@ -197,7 +176,6 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, - valid_table_names: Optional[list[str]] = None, ): logger.info("Followup SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -205,7 +183,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "histories": histories, "sql_samples": sql_samples or [], "instructions": instructions or [], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index e61f873292..973b8c69a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -25,10 +25,7 @@ logger = logging.getLogger("wren-ai-service") -def get_sql_correction_system_prompt( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: +def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) -> str: text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" @@ -55,30 +52,13 @@ def get_sql_correction_system_prompt( sql_correction_user_prompt_template = """ -### TARGET DATA SOURCE ### -{{ data_source }} - {% if documents %} -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -correcting SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} {% endif %} -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only use these exact table names from the schema. If the invalid SQL references a -table not listed here, replace it with the closest listed table only when the schema -clearly supports the user's request. Do not invent, rename, singularize, pluralize, -or add catalog/schema prefixes unless the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - {% if sql_functions %} ### SQL FUNCTIONS ### {% for function in sql_functions %} @@ -94,23 +74,9 @@ def get_sql_correction_system_prompt( {% endif %} ### QUESTION ### -{% if query %} -User's Question: {{ query }} -{% endif %} -{% if invalid_generation_result.original_sql %} -Original SQL: {{ invalid_generation_result.original_sql }} -{% endif %} -Invalid SQL: {{ invalid_generation_result.sql }} +SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} -### CORRECTION GROUNDING ### -Use ACTIVE DATASOURCE METADATA and VALID TABLE NAMES as the source of truth. If the -invalid SQL references a table or column not listed above, replace it only when the -active datasource metadata clearly contains an equivalent object that supports the -user's request. Do not invent tables, columns, joins, metrics, or relationships. -Only apply aggregate functions to columns whose active metadata type supports that -operation. - Let's think step by step. """ @@ -121,17 +87,11 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, - data_source: str, - query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( - query=query, - data_source=data_source, documents=documents, - valid_table_names=valid_table_names or [], invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, @@ -147,13 +107,9 @@ async def generate_sql_correction( prompt: dict, generator: Any, generator_name: str, - data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_correction_system_prompt( - sql_knowledge, - data_source=data_source, - ) + current_system_prompt = get_sql_correction_system_prompt(sql_knowledge) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -163,7 +119,6 @@ async def generate_sql_correction( async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, - documents: List[Document], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -220,20 +175,19 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - query: str | None = None, - valid_table_names: list[str] | None = None, ): logger.info("SQLCorrection pipeline is running...") - metadata = await retrieve_metadata(project_id or "", self._retriever) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, "documents": contexts, - "valid_table_names": valid_table_names or [], - "query": query, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index d2c5ac4724..1ee4952b3e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,16 +28,6 @@ sql_generation_user_prompt_template = """ -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only use these exact table names from the retrieved schema. Do not invent, rename, -singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless -the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -105,12 +95,10 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -214,7 +202,6 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -228,7 +215,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, 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 fc52e5a4d5..00b731cb2c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -22,24 +22,11 @@ sql_generation_reasoning_user_prompt_template = """ -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -planning SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only mention these exact table names in the reasoning plan. Do not invent, rename, -singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless -the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -75,12 +62,10 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - valid_table_names=valid_table_names or [], sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, @@ -175,7 +160,6 @@ async def run( instructions: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, - valid_table_names: Optional[list[str]] = None, ): logger.info("SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -183,7 +167,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "sql_samples": sql_samples or [], "instructions": instructions or [], "configuration": configuration, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index bbb4b40d0c..088282574e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,4 @@ import logging -import re from typing import Any, Dict, List import aiohttp @@ -13,61 +12,11 @@ clean_generation_result, ) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") -def _extract_ddl_columns(ddl: str) -> list[str]: - if not isinstance(ddl, str): - return [] - - column_section = re.search(r"\((.*)\)", ddl, flags=re.DOTALL) - if not column_section: - return [] - - columns = [] - for raw_line in re.split(r",\s*(?:\n|(?=[A-Za-z_\"`\[]))", column_section.group(1)): - line = re.sub(r"/\*.*?\*/", "", raw_line).strip() - line = re.sub(r"^--.*$", "", line).strip().rstrip(",") - if not line: - continue - - first_token = line.split()[0].strip('"`[]') - if first_token.upper() in { - "CONSTRAINT", - "FOREIGN", - "PRIMARY", - "UNIQUE", - "KEY", - }: - continue - columns.append(first_token) - - return columns - - -def format_retrieved_schema_manifest( - documents: list[str] | None, - allowed_table_names: list[str] | None, -) -> list[dict[str, Any]]: - table_names = allowed_table_names or [] - ddls = documents or [] - - manifest = [] - for index, table_name in enumerate(table_names): - if not isinstance(table_name, str) or not table_name.strip(): - continue - ddl = ddls[index] if index < len(ddls) and isinstance(ddls[index], str) else "" - manifest.append( - { - "table_name": table_name.strip(), - "columns": _extract_ddl_columns(ddl), - } - ) - - return manifest - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -218,14 +167,6 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. -- NEVER invent, assume, or rename tables and columns. Generate SQL only from tables and columns present in the provided database schema. -- Treat the retrieved schema manifest and DATABASE SCHEMA as the exact identifier allowlist. Every SELECT, WHERE, JOIN, GROUP BY, HAVING, and ORDER BY table/column must exist there. -- Select columns by business meaning, not by name similarity alone. Match the user's entities, metrics, dimensions, filters, and dates to column names, descriptions, aliases, data types, user instructions, and SQL samples. -- Prefer semantically described business/canonical models, metrics, and views over staging, temp, test, raw, backup, load, or legacy tables when the metadata indicates that distinction. -- For value, amount, total, rate, count, average, or KPI questions, use numeric measures or numeric columns whose description/alias matches the requested metric. Do not SUM or AVG string columns. -- For comparison or "by" questions, use categorical/date dimension columns for grouping and numeric measures for aggregation. -- Do not use technical audit or ingestion columns (created_at, updated_at, loaded_at, inserted_at, file_date, batch_id, row_id, ingestion timestamps, etc.) unless the user explicitly asks about sync, load, audit, or ingestion. -- If multiple columns are equally plausible and no metadata/rule/sample disambiguates them, do not guess; return SQL only when the chosen columns are grounded by the metadata. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. @@ -489,8 +430,8 @@ async def _classify_generation_result( ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; +2. Explicitly state the following information in the reasoning plan: +if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; otherwise, you will put the relative timeframe in the SQL query. 3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. @@ -504,8 +445,6 @@ async def _classify_generation_result( 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. 14. ONLY SHOWING the reasoning plan in bullet points. -15. Never include SQL code, table aliases, or assumed table/column names in the reasoning plan. -16. Only mention a table or column when the exact name appears in the DATABASE SCHEMA. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -566,16 +505,15 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" You are a helpful assistant that converts natural language queries into ANSI SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query from the provided database schema. +Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. If the section of REASONING PLAN is available in user's input, treat it only as high-level guidance. Ignore any table, column, alias, filter, or SQL fragment from the reasoning plan that is not explicitly present in the DATABASE SCHEMA. -5. Before finalizing, validate the SQL against the retrieved metadata: every table, column, join, filter, GROUP BY, HAVING, and ORDER BY identifier must exist in the provided schema, and selected columns must match the user's business intent. -6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} @@ -616,7 +554,7 @@ def construct_instructions( def construct_ask_history_messages( - histories: list[Any] | list[dict], + histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: messages = [] for history in histories: diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index aa26fa3f81..844330eac7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -105,6 +105,9 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, + max_sql_generation_tables: int = 0, + pipeline_timeout_seconds: int = 0, + schema_retrieval_timeout_seconds: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, diff --git a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py b/wren-ai-service/tests/pytest/services/test_metadata_grounding.py deleted file mode 100644 index 96fb49b77a..0000000000 --- a/wren-ai-service/tests/pytest/services/test_metadata_grounding.py +++ /dev/null @@ -1,94 +0,0 @@ -from src.pipelines.generation.followup_sql_generation import generate_sql_in_followup -from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template -from src.pipelines.generation.intent_classification import post_process -from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt -from src.pipelines.generation.utils.sql import get_sql_generation_system_prompt -from src.web.v1.services.ask import AskService - - -def test_build_validated_ask_result_keeps_legacy_select_flow(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - 'SELECT "dbo_orders"."CustomerName" FROM "dbo_orders"', - [], - "show customers", - ) - - assert result is not None - assert result.sql == 'SELECT "dbo_orders"."CustomerName" FROM "dbo_orders"' - - -def test_build_validated_ask_result_rejects_non_select_sql(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - 'DELETE FROM "dbo_orders"', - [], - "delete customers", - ) - - assert result is None - - -def test_prompts_enforce_metadata_grounding_and_result_grounded_answers(): - sql_prompt = get_sql_generation_system_prompt() - - assert "exact identifier allowlist" in sql_prompt - assert "Select columns by business meaning" in sql_prompt - assert "Do not SUM or AVG string columns" in sql_prompt - assert "Never say you do not have access" in sql_to_answer_system_prompt - assert "If Data rows are empty" in sql_to_answer_system_prompt - assert "VALID TABLE NAMES" in sql_generation_user_prompt_template - assert "Do not invent, rename" in sql_generation_user_prompt_template - - -async def test_followup_sql_generation_uses_current_system_prompt_signature(): - calls = {} - - class Generator: - async def __call__(self, **kwargs): - calls.update(kwargs) - return {"replies": ['{"sql": "SELECT 1"}']} - - result = await generate_sql_in_followup( - prompt={"prompt": "prompt"}, - generator=Generator(), - histories=[], - generator_name="test", - data_source="mssql", - sql_knowledge=None, - ) - - assert result[1] == "test" - assert "current_system_prompt" in calls - - -def test_analytic_question_with_schema_stays_text_to_sql_without_table_name(): - result = post_process( - classify_intent={ - "replies": [ - '{"rephrased_question":"What is the average length of emails?",' - '"reasoning":"Misclassified as guide.",' - '"results":"USER_GUIDE"}' - ] - }, - construct_db_schemas=["CREATE TABLE users (email VARCHAR);"], - query="What is the average length of emails?", - ) - - assert result["intent"] == "TEXT_TO_SQL" - - -def test_user_guide_question_stays_user_guide(): - result = post_process( - classify_intent={ - "replies": [ - '{"rephrased_question":"How can I connect to a database?",' - '"reasoning":"Wren setup question.",' - '"results":"USER_GUIDE"}' - ] - }, - construct_db_schemas=["CREATE TABLE users (email VARCHAR);"], - query="How can I connect to a database?", - ) - - assert result["intent"] == "USER_GUIDE" From 6b031bb93a9ff82440f95575030818e8c36d8674 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 15:39:28 +0530 Subject: [PATCH 0629/1087] Restore legacy generation prompts --- .../generation/followup_sql_generation.py | 46 +++++++++-- .../followup_sql_generation_reasoning.py | 29 ++++++- .../pipelines/generation/sql_correction.py | 62 +++++++++++++-- .../pipelines/generation/sql_generation.py | 14 ++++ .../generation/sql_generation_reasoning.py | 19 ++++- .../src/pipelines/generation/utils/sql.py | 76 +++++++++++++++++-- 6 files changed, 220 insertions(+), 26 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 35cfb8fccf..b013f9b25c 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import logging import sys -from typing import Any +from typing import TYPE_CHECKING, Any from hamilton import base from hamilton.async_driver import AsyncDriver @@ -24,7 +26,11 @@ from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") @@ -34,11 +40,24 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. -### DATABASE SCHEMA ### +### TARGET DATA SOURCE ### +{{ data_source }} + +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +interpreting intent and generating SQL. {% for document in documents %} {{ document }} {% endfor %} +### VALID TABLE NAMES ### +Only use these exact table names from the schema. Do not invent, rename, singularize, +pluralize, or add catalog/schema prefixes unless the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -78,6 +97,14 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +### INTENT AND SCHEMA GROUNDING ### +Interpret the user's business terms by matching them to explicit tables, columns, +metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Never reuse table +or column names from SQL SAMPLES or chat history unless those exact names also +appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. +Only apply aggregate functions to columns whose active metadata type supports that +operation. + ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -92,6 +119,7 @@ def prompt( documents: list[str], sql_generation_reasoning: str, prompt_builder: PromptBuilder, + data_source: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -99,10 +127,13 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, + data_source=data_source, documents=documents, + valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -131,6 +162,7 @@ async def generate_sql_in_followup( generator: Any, histories: list[AskHistory], generator_name: str, + data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: history_messages = construct_ask_history_messages(histories) @@ -146,6 +178,7 @@ async def generate_sql_in_followup( async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, + documents: list[str], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -208,19 +241,18 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = await retrieve_metadata(project_id or "", self._retriever) return await self._pipe.execute( ["post_process"], inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 42b28c5b8f..4d4402ea69 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import asyncio import logging import sys -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -17,17 +19,34 @@ ) from src.utils import trace_cost from src.web.v1.services import Configuration -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory +else: + AskHistory = Any logger = logging.getLogger("wren-ai-service") sql_generation_reasoning_user_prompt_template = """ -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +planning SQL. {% for document in documents %} {{ document }} {% endfor %} +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only mention these exact table names in the reasoning plan. Do not invent, rename, +singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless +the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -72,10 +91,12 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + valid_table_names=valid_table_names or [], histories=histories, sql_samples=sql_samples, instructions=construct_instructions( @@ -176,6 +197,7 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, + valid_table_names: Optional[list[str]] = None, ): logger.info("Followup SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -183,6 +205,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "histories": histories, "sql_samples": sql_samples or [], "instructions": instructions or [], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..e61f873292 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -25,7 +25,10 @@ logger = logging.getLogger("wren-ai-service") -def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) -> str: +def get_sql_correction_system_prompt( + sql_knowledge: SqlKnowledge | None = None, + data_source: str | None = None, +) -> str: text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" @@ -52,13 +55,30 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) sql_correction_user_prompt_template = """ +### TARGET DATA SOURCE ### +{{ data_source }} + {% if documents %} -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +correcting SQL. {% for document in documents %} {{ document }} {% endfor %} {% endif %} +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only use these exact table names from the schema. If the invalid SQL references a +table not listed here, replace it with the closest listed table only when the schema +clearly supports the user's request. Do not invent, rename, singularize, pluralize, +or add catalog/schema prefixes unless the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + {% if sql_functions %} ### SQL FUNCTIONS ### {% for function in sql_functions %} @@ -74,9 +94,23 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -SQL: {{ invalid_generation_result.sql }} +{% if query %} +User's Question: {{ query }} +{% endif %} +{% if invalid_generation_result.original_sql %} +Original SQL: {{ invalid_generation_result.original_sql }} +{% endif %} +Invalid SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} +### CORRECTION GROUNDING ### +Use ACTIVE DATASOURCE METADATA and VALID TABLE NAMES as the source of truth. If the +invalid SQL references a table or column not listed above, replace it only when the +active datasource metadata clearly contains an equivalent object that supports the +user's request. Do not invent tables, columns, joins, metrics, or relationships. +Only apply aggregate functions to columns whose active metadata type supports that +operation. + Let's think step by step. """ @@ -87,11 +121,17 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, + data_source: str, + query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( + query=query, + data_source=data_source, documents=documents, + valid_table_names=valid_table_names or [], invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, @@ -107,9 +147,13 @@ async def generate_sql_correction( prompt: dict, generator: Any, generator_name: str, + data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_correction_system_prompt(sql_knowledge) + current_system_prompt = get_sql_correction_system_prompt( + sql_knowledge, + data_source=data_source, + ) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -119,6 +163,7 @@ async def generate_sql_correction( async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, + documents: List[Document], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -175,19 +220,20 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + query: str | None = None, + valid_table_names: list[str] | None = None, ): logger.info("SQLCorrection pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = await retrieve_metadata(project_id or "", self._retriever) return await self._pipe.execute( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, "documents": contexts, + "valid_table_names": valid_table_names or [], + "query": query, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..d2c5ac4724 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,6 +28,16 @@ sql_generation_user_prompt_template = """ +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only use these exact table names from the retrieved schema. Do not invent, rename, +singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless +the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -95,10 +105,12 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -202,6 +214,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + valid_table_names: list[str] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -215,6 +228,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, 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 00b731cb2c..fc52e5a4d5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -22,11 +22,24 @@ sql_generation_reasoning_user_prompt_template = """ -### DATABASE SCHEMA ### +### ACTIVE DATASOURCE METADATA ### +This is the complete deployed metadata for the active datasource, including schema, +tables, columns, metrics, views, and relationships. Use only this metadata when +planning SQL. {% for document in documents %} {{ document }} {% endfor %} +{% if valid_table_names %} +### VALID TABLE NAMES ### +Only mention these exact table names in the reasoning plan. Do not invent, rename, +singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless +the table name is shown that way here. +{% for table_name in valid_table_names %} +- {{ table_name }} +{% endfor %} +{% endif %} + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -62,10 +75,12 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), + valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + valid_table_names=valid_table_names or [], sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, @@ -160,6 +175,7 @@ async def run( instructions: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, + valid_table_names: Optional[list[str]] = None, ): logger.info("SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -167,6 +183,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "valid_table_names": valid_table_names or [], "sql_samples": sql_samples or [], "instructions": instructions or [], "configuration": configuration, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..bbb4b40d0c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, List import aiohttp @@ -12,11 +13,61 @@ clean_generation_result, ) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge -from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") +def _extract_ddl_columns(ddl: str) -> list[str]: + if not isinstance(ddl, str): + return [] + + column_section = re.search(r"\((.*)\)", ddl, flags=re.DOTALL) + if not column_section: + return [] + + columns = [] + for raw_line in re.split(r",\s*(?:\n|(?=[A-Za-z_\"`\[]))", column_section.group(1)): + line = re.sub(r"/\*.*?\*/", "", raw_line).strip() + line = re.sub(r"^--.*$", "", line).strip().rstrip(",") + if not line: + continue + + first_token = line.split()[0].strip('"`[]') + if first_token.upper() in { + "CONSTRAINT", + "FOREIGN", + "PRIMARY", + "UNIQUE", + "KEY", + }: + continue + columns.append(first_token) + + return columns + + +def format_retrieved_schema_manifest( + documents: list[str] | None, + allowed_table_names: list[str] | None, +) -> list[dict[str, Any]]: + table_names = allowed_table_names or [] + ddls = documents or [] + + manifest = [] + for index, table_name in enumerate(table_names): + if not isinstance(table_name, str) or not table_name.strip(): + continue + ddl = ddls[index] if index < len(ddls) and isinstance(ddls[index], str) else "" + manifest.append( + { + "table_name": table_name.strip(), + "columns": _extract_ddl_columns(ddl), + } + ) + + return manifest + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -167,6 +218,14 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. +- NEVER invent, assume, or rename tables and columns. Generate SQL only from tables and columns present in the provided database schema. +- Treat the retrieved schema manifest and DATABASE SCHEMA as the exact identifier allowlist. Every SELECT, WHERE, JOIN, GROUP BY, HAVING, and ORDER BY table/column must exist there. +- Select columns by business meaning, not by name similarity alone. Match the user's entities, metrics, dimensions, filters, and dates to column names, descriptions, aliases, data types, user instructions, and SQL samples. +- Prefer semantically described business/canonical models, metrics, and views over staging, temp, test, raw, backup, load, or legacy tables when the metadata indicates that distinction. +- For value, amount, total, rate, count, average, or KPI questions, use numeric measures or numeric columns whose description/alias matches the requested metric. Do not SUM or AVG string columns. +- For comparison or "by" questions, use categorical/date dimension columns for grouping and numeric measures for aggregation. +- Do not use technical audit or ingestion columns (created_at, updated_at, loaded_at, inserted_at, file_date, batch_id, row_id, ingestion timestamps, etc.) unless the user explicitly asks about sync, load, audit, or ingestion. +- If multiple columns are equally plausible and no metadata/rule/sample disambiguates them, do not guess; return SQL only when the chosen columns are grounded by the metadata. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. @@ -430,8 +489,8 @@ async def _classify_generation_result( ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; +2. Explicitly state the following information in the reasoning plan: +if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; otherwise, you will put the relative timeframe in the SQL query. 3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. @@ -445,6 +504,8 @@ async def _classify_generation_result( 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. 14. ONLY SHOWING the reasoning plan in bullet points. +15. Never include SQL code, table aliases, or assumed table/column names in the reasoning plan. +16. Only mention a table or column when the exact name appears in the DATABASE SCHEMA. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -505,15 +566,16 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" You are a helpful assistant that converts natural language queries into ANSI SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. +Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query from the provided database schema. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +4. If the section of REASONING PLAN is available in user's input, treat it only as high-level guidance. Ignore any table, column, alias, filter, or SQL fragment from the reasoning plan that is not explicitly present in the DATABASE SCHEMA. +5. Before finalizing, validate the SQL against the retrieved metadata: every table, column, join, filter, GROUP BY, HAVING, and ORDER BY identifier must exist in the provided schema, and selected columns must match the user's business intent. +6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} @@ -554,7 +616,7 @@ def construct_instructions( def construct_ask_history_messages( - histories: list[AskHistory] | list[dict], + histories: list[Any] | list[dict], ) -> list[ChatMessage]: messages = [] for history in histories: From 3bdba9aa7bbd5d3680becc8846d7d3781fde8a95 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 15:58:18 +0530 Subject: [PATCH 0630/1087] Remove preview manifest SQL validator --- .../apollo/server/services/queryService.ts | 548 +----------------- 1 file changed, 1 insertion(+), 547 deletions(-) diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 4ae56a81e8..e4853589ac 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -1,5 +1,5 @@ import { DataSourceName } from '@server/types'; -import { ColumnMDL, Manifest, TableReference } from '@server/mdl/type'; +import { Manifest, TableReference } from '@server/mdl/type'; import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; import { SupportedDataSource, @@ -344,9 +344,6 @@ const quoteTableReference = (tableReference: string) => const escapeRegExp = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); -const compactSqlIdentifier = (identifier: string) => - normalizeSqlIdentifier(identifier).replace(/[^A-Za-z0-9]/g, '').toLowerCase(); - const extractSqlTableReferences = (sql: string) => { const references: string[] = []; const tablePattern = new RegExp( @@ -360,505 +357,6 @@ const extractSqlTableReferences = (sql: string) => { return references; }; -const addTableReferenceName = ( - names: Set, - parts: Array, -) => { - const normalizedParts = parts - .filter((part): part is string => Boolean(part)) - .map((part) => part.toLowerCase()); - if (!normalizedParts.length) { - return; - } - - for (let index = 0; index < normalizedParts.length; index += 1) { - names.add(normalizedParts.slice(index).join('.')); - } -}; - -const extractCteNames = (sql: string) => { - const cteNames = new Set(); - const ctePattern = new RegExp( - String.raw`(?:\bWITH\b|,)\s*(${SQL_IDENTIFIER_PATTERN})\s+AS\s*\(`, - 'gi', - ); - let match: RegExpExecArray | null; - while ((match = ctePattern.exec(sql))) { - cteNames.add(normalizeSqlIdentifier(match[1]).toLowerCase()); - } - return cteNames; -}; - -const getManifestQueryableNames = (manifest?: Manifest) => { - const names = new Set(); - for (const model of manifest?.models || []) { - if (model.name) names.add(model.name.toLowerCase()); - if (model.tableReference?.table) { - addTableReferenceName(names, [ - model.tableReference.catalog, - model.tableReference.schema, - model.tableReference.table, - ]); - } - if (model.refSql) { - for (const reference of extractSqlTableReferences(model.refSql)) { - addTableReferenceName(names, splitTableReference(reference)); - } - } - } - for (const view of manifest?.views || []) { - if (view.name) names.add(view.name.toLowerCase()); - } - return names; -}; - -interface ManifestModelSchema { - name: string; - columns: Map; -} - -const addManifestModelSchemaAlias = ( - schemas: Map, - alias: string | undefined, - schema: ManifestModelSchema, -) => { - if (!alias) { - return; - } - schemas.set(alias.toLowerCase(), schema); - schemas.set(compactSqlIdentifier(alias), schema); -}; - -const getManifestModelSchemas = (manifest?: Manifest) => { - const schemas = new Map(); - - for (const model of manifest?.models || []) { - if (!model.name || !model.columns?.length) { - continue; - } - - const columns = new Map(); - model.columns - .filter((column) => column?.name) - .forEach((column) => { - columns.set(column.name.toLowerCase(), column); - columns.set(compactSqlIdentifier(column.name), column); - }); - - const schema: ManifestModelSchema = { - name: model.name, - columns, - }; - - addManifestModelSchemaAlias(schemas, model.name, schema); - addManifestModelSchemaAlias(schemas, model.tableReference?.table, schema); - - const referenceParts = [ - model.tableReference?.catalog, - model.tableReference?.schema, - model.tableReference?.table, - ].filter(Boolean); - if (referenceParts.length) { - addManifestModelSchemaAlias(schemas, referenceParts.join('.'), schema); - } - } - - return schemas; -}; - -const extractSqlTableAliases = ( - sql: string, - schemas: Map, -) => { - const aliases = new Map(); - const tablePattern = new RegExp( - String.raw`\b(?:FROM|JOIN)\s+(${SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*${SQL_IDENTIFIER_PATTERN})*)(?:\s+(?:AS\s+)?(${SQL_IDENTIFIER_PATTERN}))?`, - 'gi', - ); - const clauseWords = new Set([ - 'where', - 'join', - 'inner', - 'left', - 'right', - 'full', - 'cross', - 'group', - 'order', - 'having', - 'limit', - 'offset', - 'fetch', - 'union', - 'on', - ]); - - let match: RegExpExecArray | null; - while ((match = tablePattern.exec(sql))) { - const reference = splitTableReference(match[1]).join('.'); - const lastPart = splitTableReference(reference).pop(); - const schema = - schemas.get(reference.toLowerCase()) || - schemas.get(compactSqlIdentifier(reference)) || - (lastPart - ? schemas.get(lastPart.toLowerCase()) || - schemas.get(compactSqlIdentifier(lastPart)) - : undefined); - - if (!schema) { - continue; - } - - aliases.set(reference.toLowerCase(), schema); - if (lastPart) { - aliases.set(lastPart.toLowerCase(), schema); - } - - const alias = match[2] ? normalizeSqlIdentifier(match[2]) : null; - if (alias && !clauseWords.has(alias.toLowerCase())) { - aliases.set(alias.toLowerCase(), schema); - } - } - - return aliases; -}; - -const isNumericColumnType = (type?: string) => - !!type && - /(?:int|integer|bigint|smallint|tinyint|float|double|decimal|numeric|number|real|money)/i.test( - type, - ); - -const isDefined = (value: T | undefined | null): value is T => - value !== undefined && value !== null; - -const SQL_NON_COLUMN_IDENTIFIERS = new Set([ - 'and', - 'as', - 'asc', - 'between', - 'by', - 'case', - 'cast', - 'count', - 'date', - 'dateadd', - 'datediff', - 'datepart', - 'day', - 'desc', - 'distinct', - 'else', - 'end', - 'false', - 'from', - 'getdate', - 'group', - 'having', - 'hour', - 'in', - 'is', - 'join', - 'left', - 'like', - 'limit', - 'max', - 'min', - 'month', - 'not', - 'null', - 'on', - 'or', - 'order', - 'right', - 'select', - 'sum', - 'then', - 'top', - 'true', - 'when', - 'where', - 'year', -]); - -const splitTopLevelSqlList = (body: string) => { - const items: string[] = []; - let current = ''; - let depth = 0; - let inSingleQuote = false; - let inDoubleQuote = false; - let inBracket = false; - - for (const char of body) { - if (char === "'" && !inDoubleQuote && !inBracket) { - inSingleQuote = !inSingleQuote; - } else if (char === '"' && !inSingleQuote && !inBracket) { - inDoubleQuote = !inDoubleQuote; - } else if (char === '[' && !inSingleQuote && !inDoubleQuote) { - inBracket = true; - } else if (char === ']' && inBracket) { - inBracket = false; - } else if (!inSingleQuote && !inDoubleQuote && !inBracket) { - if (char === '(') depth += 1; - if (char === ')' && depth > 0) depth -= 1; - if (char === ',' && depth === 0) { - items.push(current.trim()); - current = ''; - continue; - } - } - current += char; - } - - if (current.trim()) { - items.push(current.trim()); - } - return items; -}; - -const stripProjectionAlias = (item: string) => { - const aliasMatch = item.match( - /\s+(?:AS\s+)?(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_]*)\s*$/i, - ); - if (!aliasMatch || aliasMatch.index === undefined) { - return item.trim(); - } - - const expression = item.slice(0, aliasMatch.index).trim(); - return expression || item.trim(); -}; - -const extractSimpleProjectionColumns = (sql: string) => { - const columns: string[] = []; - const selectPattern = /\bSELECT\b(?.*?)(?=\bFROM\b)/gis; - let match: RegExpExecArray | null; - while ((match = selectPattern.exec(sql))) { - const body = match.groups?.body || ''; - splitTopLevelSqlList(body).forEach((item) => { - const expression = stripProjectionAlias( - item.replace(/^\s*DISTINCT\s+/i, ''), - ); - const identifierMatch = expression.match( - new RegExp(String.raw`^${SQL_IDENTIFIER_PATTERN}$`, 'i'), - ); - if (identifierMatch && normalizeSqlIdentifier(expression) !== '*') { - columns.push(normalizeSqlIdentifier(expression)); - } - }); - } - return columns; -}; - -const stripSqlLiterals = (sql: string) => - sql - .replace(/'(?:''|[^'])*'/g, ' ') - .replace(/\b\d+(?:\.\d+)?\b/g, ' '); - -const extractClauseBody = ( - sql: string, - clause: string, - terminators: string[], -) => { - const pattern = new RegExp( - String.raw`\b${clause}\b(?.*?)(?=${ - terminators.map((word) => String.raw`\b${word}\b`).join('|') - }|$)`, - 'gis', - ); - const bodies: string[] = []; - let match: RegExpExecArray | null; - while ((match = pattern.exec(sql))) { - bodies.push(match.groups?.body || ''); - } - return bodies; -}; - -const extractPotentialUnqualifiedColumnReferences = (sql: string) => { - const bodies = [ - ...extractClauseBody(sql, 'WHERE', [ - 'GROUP\\s+BY', - 'ORDER\\s+BY', - 'HAVING', - 'LIMIT', - 'FETCH', - 'UNION', - ]), - ...extractClauseBody(sql, 'HAVING', [ - 'GROUP\\s+BY', - 'ORDER\\s+BY', - 'LIMIT', - 'FETCH', - 'UNION', - ]), - ...extractClauseBody(sql, 'ON', [ - 'WHERE', - 'GROUP\\s+BY', - 'ORDER\\s+BY', - 'HAVING', - 'JOIN', - 'LIMIT', - 'FETCH', - 'UNION', - ]), - ...extractClauseBody(sql, 'GROUP\\s+BY', [ - 'ORDER\\s+BY', - 'HAVING', - 'LIMIT', - 'FETCH', - 'UNION', - ]), - ]; - const identifiers = new Set(); - const identifierPattern = new RegExp(SQL_IDENTIFIER_PATTERN, 'gi'); - - for (const body of bodies) { - const searchableBody = stripSqlLiterals(body); - let match: RegExpExecArray | null; - while ((match = identifierPattern.exec(searchableBody))) { - const token = match[0]; - const before = searchableBody.slice(0, match.index).trimEnd(); - const after = searchableBody.slice(match.index + token.length).trimStart(); - const identifier = normalizeSqlIdentifier(token); - const normalized = identifier.toLowerCase(); - if ( - !identifier || - SQL_NON_COLUMN_IDENTIFIERS.has(normalized) || - before.endsWith('.') || - after.startsWith('.') || - after.startsWith('(') - ) { - continue; - } - identifiers.add(identifier); - } - } - - const functionArgumentPattern = new RegExp( - String.raw`\b[A-Za-z_][A-Za-z0-9_$]*\s*\(\s*(?:DISTINCT\s+)?(${SQL_IDENTIFIER_PATTERN})(?:\s*\.\s*(${SQL_IDENTIFIER_PATTERN}))?`, - 'gi', - ); - let match: RegExpExecArray | null; - while ((match = functionArgumentPattern.exec(stripSqlLiterals(sql)))) { - const column = normalizeSqlIdentifier(match[2] || match[1]); - if ( - column && - column !== '*' && - !SQL_NON_COLUMN_IDENTIFIERS.has(column.toLowerCase()) - ) { - identifiers.add(column); - } - } - - return [...identifiers]; -}; - -const findSqlReferenceValidationErrors = ( - sql: string, - manifest?: Manifest, -) => { - const schemas = getManifestModelSchemas(manifest); - if (!schemas.size) { - return []; - } - - const aliases = extractSqlTableAliases(sql, schemas); - const errors: string[] = []; - const qualifiedColumnPattern = new RegExp( - String.raw`(${SQL_IDENTIFIER_PATTERN})\s*\.\s*(${SQL_IDENTIFIER_PATTERN})`, - 'gi', - ); - - let match: RegExpExecArray | null; - while ((match = qualifiedColumnPattern.exec(sql))) { - const qualifier = normalizeSqlIdentifier(match[1]); - const column = normalizeSqlIdentifier(match[2]); - const schema = aliases.get(qualifier.toLowerCase()); - - if (!schema || column === '*') { - continue; - } - - if ( - !schema.columns.has(column.toLowerCase()) && - !schema.columns.has(compactSqlIdentifier(column)) - ) { - errors.push(`${qualifier}.${column}`); - } - } - - const aggregatePattern = new RegExp( - String.raw`\b(AVG|SUM)\s*\(\s*(?:DISTINCT\s+)?(${SQL_IDENTIFIER_PATTERN})(?:\s*\.\s*(${SQL_IDENTIFIER_PATTERN}))?\s*\)`, - 'gi', - ); - while ((match = aggregatePattern.exec(sql))) { - const functionName = match[1].toUpperCase(); - const qualifier = match[3] ? normalizeSqlIdentifier(match[2]) : null; - const column = normalizeSqlIdentifier(match[3] || match[2]); - if (column === '*') { - continue; - } - - const candidateSchemas = qualifier - ? [aliases.get(qualifier.toLowerCase())].filter(isDefined) - : [...new Set(aliases.values())]; - const matchingColumns = candidateSchemas - .map( - (schema) => - schema?.columns.get(column.toLowerCase()) || - schema?.columns.get(compactSqlIdentifier(column)), - ) - .filter(isDefined); - - if (!matchingColumns.length) { - errors.push(qualifier ? `${qualifier}.${column}` : column); - continue; - } - - if ( - matchingColumns.some( - (columnSchema) => !isNumericColumnType(columnSchema.type), - ) - ) { - errors.push( - `${functionName}(${qualifier ? `${qualifier}.` : ''}${column}) uses a non-numeric column`, - ); - } - } - - const activeSchemas = [...new Set(aliases.values())]; - if (activeSchemas.length === 1) { - const [schema] = activeSchemas; - extractSimpleProjectionColumns(sql).forEach((column) => { - if ( - !schema.columns.has(column.toLowerCase()) && - !schema.columns.has(compactSqlIdentifier(column)) - ) { - errors.push(column); - } - }); - - const tableAliases = new Set( - [...aliases.entries()] - .filter(([, aliasSchema]) => aliasSchema === schema) - .map(([alias]) => alias.toLowerCase()), - ); - const validCompactColumns = new Set([...schema.columns.keys()]); - extractPotentialUnqualifiedColumnReferences(sql).forEach((column) => { - const normalizedColumn = column.toLowerCase(); - if ( - tableAliases.has(normalizedColumn) || - schema.columns.has(normalizedColumn) || - validCompactColumns.has(compactSqlIdentifier(column)) - ) { - return; - } - errors.push(column); - }); - } - - return [...new Set(errors)]; -}; - const addModelReferenceAlias = ( aliases: Map, parts: Array, @@ -963,49 +461,6 @@ const normalizeSqlReferencesToManifest = (sql: string, manifest?: Manifest) => { return normalizedSql; }; -const validateSqlReferencesManifest = (sql: string, manifest?: Manifest) => { - const validNames = getManifestQueryableNames(manifest); - if (!validNames.size) { - return; - } - - const cteNames = extractCteNames(sql); - const invalidReferences = extractSqlTableReferences(sql).filter((reference) => { - const normalized = reference.toLowerCase(); - const parts = splitTableReference(reference); - const lastPart = parts[parts.length - 1]?.toLowerCase(); - const firstPart = parts[0]?.toLowerCase(); - const suffixes = parts.map((_, index) => - parts.slice(index).join('.').toLowerCase(), - ); - return ( - !validNames.has(normalized) && - !cteNames.has(normalized) && - !suffixes.some((suffix) => validNames.has(suffix)) && - !(parts.length === 2 && firstPart && validNames.has(firstPart)) && - (!lastPart || !validNames.has(lastPart)) - ); - }); - - if (invalidReferences.length) { - throw new Error( - `Generated SQL references table(s) not present in the active datasource metadata: ${[ - ...new Set(invalidReferences), - ].join(', ')}`, - ); - } - - const invalidColumnReferences = findSqlReferenceValidationErrors( - sql, - manifest, - ); - if (invalidColumnReferences.length) { - throw new Error( - `Generated SQL references column(s) or expressions not valid for the active datasource metadata: ${invalidColumnReferences.join(', ')}`, - ); - } -}; - export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; @@ -1045,7 +500,6 @@ export class QueryService implements IQueryService { dataSource, limit, ); - validateSqlReferencesManifest(normalizedPreview.sql, mdl); if (this.useEngine(dataSource)) { if (dryRun) { logger.debug('Using wren engine to dry run'); From 467de749122c6581f65d68a6d42ae76650f97b1a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 16:07:57 +0530 Subject: [PATCH 0631/1087] Route analytical questions to text-to-sql --- .../generation/intent_classification.py | 20 ++++++++++--------- wren-ai-service/src/web/v1/services/ask.py | 10 ++++++++++ 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..908850151b 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -42,15 +42,16 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. -- The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes references to specific tables, columns, or data details. -- The question includes **complete information** with specific tables, columns, or data values needed for execution. -- The question provides **all necessary parameters** to generate executable SQL. +- The user's inputs ask to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data. +- The user's inputs are related to the database schema and require an SQL query. +- The question can be answered by selecting relevant tables and columns from the provided schema, even if the user does not mention exact physical table or column names. +- The question includes enough business meaning, dimensions, metrics, filters, or time criteria to attempt SQL generation from the schema. +- Natural-language analytical questions should be classified as `TEXT_TO_SQL` when they can reasonably be grounded in the schema. **Requirements:** -- Must have complete filter criteria, specific values, or clear references to previous context. -- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. -- Reference phrases from the user's inputs that clearly relate to the schema. +- Do not require the user to explicitly name a table or column. +- Use the schema to determine whether the user's business terms can map to available tables or columns. +- Reference phrases from the user's inputs that clearly indicate a data retrieval or analysis request. **Examples:** - "What is the total sales for last quarter?" @@ -63,7 +64,8 @@ - The user seeks general information about the database schema or its overall capabilities. - The query references **missing information** (e.g., "the following items" without listing them). - The query contains **placeholder references** that cannot be resolved from context. -- The query is **incomplete for SQL generation** despite mentioning database concepts. +- The query is **incomplete for SQL generation** because required values or references are missing, not merely because exact table or column names are absent. +- The user is asking for explanation, guidance, or clarification rather than asking to retrieve or analyze rows from the data. **Requirements:** - Incorporate phrases from the user's inputs that indicate incompleteness or lack of relevance to the database schema. @@ -93,7 +95,7 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. +- The user's inputs cannot be interpreted as a database question or data analysis request. - It appears off-topic or is simply a casual conversation starter. **Requirements:** diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 844330eac7..d4cbe0c1a7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -256,6 +256,16 @@ async def ask( if rephrased_question: user_query = rephrased_question + if ( + intent == "GENERAL" + and not intent_classification_result.get("db_schemas") + ): + logger.info( + "Intent classification returned GENERAL without schema context; continuing Text-to-SQL retrieval for query_id %s", + query_id, + ) + intent = "TEXT_TO_SQL" + if intent == "MISLEADING_QUERY": asyncio.create_task( self._pipelines["misleading_assistance"].run( From 07c1d8f33cc9a5227cccf92e41788db818db8aca Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 16:33:25 +0530 Subject: [PATCH 0632/1087] Continue text-to-sql for analytical queries --- .../generation/intent_classification.py | 3 ++- wren-ai-service/src/web/v1/services/ask.py | 27 +++++++++++++++++-- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 908850151b..ebefab117d 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -33,7 +33,8 @@ - **Rephrase Question:** Rewrite follow-up questions into full standalone questions using prior conversation context. - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. -- **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. +- **Vague Queries:** If the question is not a data retrieval or analysis request and does not relate to the schema, classify it as `MISLEADING_QUERY`. +- **Natural Language Data Queries:** Do not require exact physical table or column names. If the user asks to show, list, find, compare, count, aggregate, rank, filter, sort, or analyze data, classify it as `TEXT_TO_SQL` when schema context may answer it. - **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index d4cbe0c1a7..48968c9a5e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,5 +1,6 @@ import asyncio import logging +import re from typing import Dict, List, Literal, Optional from cachetools import TTLCache @@ -13,6 +14,18 @@ logger = logging.getLogger("wren-ai-service") +DATA_QUERY_PATTERN = re.compile( + r"\b(" + r"show|list|find|get|give|display|retrieve|fetch|compare|count|sum|total|" + r"average|avg|min|max|rank|top|bottom|highest|lowest|latest|earliest|" + r"newest|oldest|sort|order|group|filter|where|between|starts|ends|" + r"contains|duplicate|unique|distinct|percentage|percent|trend|growth|" + r"breakdown|by|per" + r")\b", + re.IGNORECASE, +) + + class AskHistory(BaseModel): sql: str question: str @@ -125,6 +138,12 @@ def __init__( self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries + def _should_continue_text_to_sql(self, query: str, intent: Optional[str]) -> bool: + if intent not in {"GENERAL", "MISLEADING_QUERY"}: + return False + + return bool(DATA_QUERY_PATTERN.search(query or "")) + def _is_stopped(self, query_id: str, container: dict): if ( result := container.get(query_id) @@ -256,12 +275,16 @@ async def ask( if rephrased_question: user_query = rephrased_question - if ( + if self._should_continue_text_to_sql( + user_query, + intent, + ) or ( intent == "GENERAL" and not intent_classification_result.get("db_schemas") ): logger.info( - "Intent classification returned GENERAL without schema context; continuing Text-to-SQL retrieval for query_id %s", + "Intent classification returned %s for an analytical query; continuing Text-to-SQL retrieval for query_id %s", + intent, query_id, ) intent = "TEXT_TO_SQL" From d8024b0d58eaf438369148b745d19010b3822cb3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 16:45:28 +0530 Subject: [PATCH 0633/1087] Prevent user guide routing for data queries --- .../src/pipelines/generation/intent_classification.py | 2 ++ wren-ai-service/src/web/v1/services/ask.py | 8 ++++---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index ebefab117d..e858b46aaa 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -35,6 +35,7 @@ - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. - **Vague Queries:** If the question is not a data retrieval or analysis request and does not relate to the schema, classify it as `MISLEADING_QUERY`. - **Natural Language Data Queries:** Do not require exact physical table or column names. If the user asks to show, list, find, compare, count, aggregate, rank, filter, sort, or analyze data, classify it as `TEXT_TO_SQL` when schema context may answer it. +- **User Guide Boundary:** Do not classify a data retrieval, filtering, date, metric, aggregation, or row-listing question as `USER_GUIDE` just because it contains words like "how", "show", or "filter". - **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. @@ -84,6 +85,7 @@ **When to Use:** - The user's inputs pertains to Wren AI's features, usage, or capabilities. - The query relates directly to content in the user guide. +- The query asks how to use Wren AI itself, not how to retrieve or analyze rows from the connected data. **Examples:** - "What can Wren AI do?" diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 48968c9a5e..6bf806392d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -15,13 +15,13 @@ DATA_QUERY_PATTERN = re.compile( - r"\b(" + r"(\bhow\s+many\b|\bhow\s+much\b|\bwhat\s+(?:is|are|was|were)\b|\bwhich\b|\b(" r"show|list|find|get|give|display|retrieve|fetch|compare|count|sum|total|" r"average|avg|min|max|rank|top|bottom|highest|lowest|latest|earliest|" r"newest|oldest|sort|order|group|filter|where|between|starts|ends|" r"contains|duplicate|unique|distinct|percentage|percent|trend|growth|" - r"breakdown|by|per" - r")\b", + r"breakdown|placed|created|became|active|inactive|expired|valid|by|per" + r")\b)", re.IGNORECASE, ) @@ -139,7 +139,7 @@ def __init__( self._max_sql_correction_retries = max_sql_correction_retries def _should_continue_text_to_sql(self, query: str, intent: Optional[str]) -> bool: - if intent not in {"GENERAL", "MISLEADING_QUERY"}: + if intent not in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"}: return False return bool(DATA_QUERY_PATTERN.search(query or "")) From 61fe2e5c2d69360d8ce876343f21a66903fa1677 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 17:22:10 +0530 Subject: [PATCH 0634/1087] Ground SQL generation with retrieved schema --- .../pipelines/generation/sql_correction.py | 3 +++ .../pipelines/generation/sql_generation.py | 8 ++++++++ wren-ai-service/src/web/v1/services/ask.py | 19 +++++++++++++++++-- 3 files changed, 28 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index e61f873292..1074b16482 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -101,6 +101,9 @@ def get_sql_correction_system_prompt( Original SQL: {{ invalid_generation_result.original_sql }} {% endif %} Invalid SQL: {{ invalid_generation_result.sql }} +{% if invalid_generation_result.executed_sql %} +Engine SQL that failed validation/execution: {{ invalid_generation_result.executed_sql }} +{% endif %} Error Message: {{ invalid_generation_result.error }} ### CORRECTION GROUNDING ### diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index d2c5ac4724..90b15d2fac 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -82,6 +82,14 @@ ### QUESTION ### User's Question: {{ query }} +### INTENT AND SCHEMA GROUNDING ### +Interpret the user's business terms by matching them to explicit tables, columns, +metrics, views, and relationships in DATABASE SCHEMA. Never copy table or column +names from SQL SAMPLES, REASONING PLAN, or prior assumptions unless those exact +names also appear in DATABASE SCHEMA or VALID TABLE NAMES. If a required metric, +dimension, filter, date, or join key cannot be grounded by the retrieved metadata, +do not replace it with a generic placeholder. + {% if sql_generation_reasoning %} ### REASONING PLAN ### {{ sql_generation_reasoning }} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6bf806392d..0d518eba8b 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -389,8 +389,16 @@ async def ask( "construct_retrieval_results", {} ) documents = _retrieval_result.get("retrieval_results", []) - table_names = [document.get("table_name") for document in documents] - table_ddls = [document.get("table_ddl") for document in documents] + table_names = [ + document.get("table_name") + for document in documents + if document.get("table_name") + ] + table_ddls = [ + document.get("table_ddl") + for document in documents + if document.get("table_ddl") + ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -436,6 +444,7 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, + valid_table_names=table_names, ) ).get("post_process", {}) else: @@ -447,6 +456,7 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, + valid_table_names=table_names, ) ).get("post_process", {}) @@ -513,6 +523,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + valid_table_names=table_names, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -531,6 +542,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + valid_table_names=table_names, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -588,6 +600,7 @@ async def ask( instructions=instructions, invalid_generation_result={ "sql": original_sql, + "executed_sql": invalid_sql, "error": sql_diagnosis_reasoning if allow_sql_diagnosis else error_message, @@ -597,6 +610,8 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + query=user_query, + valid_table_names=table_names, ) if valid_generation_result := sql_correction_results[ From be7971c0b97dc75a9ca3c727f35818ae88504b7c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 17:39:48 +0530 Subject: [PATCH 0635/1087] Prevent reasoning from polluting SQL generation --- wren-ai-service/src/pipelines/generation/utils/sql.py | 3 +++ .../src/pipelines/retrieval/db_schema_retrieval.py | 5 +++++ wren-ai-service/src/web/v1/services/ask.py | 4 ++-- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index bbb4b40d0c..305473255d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -506,6 +506,9 @@ async def _classify_generation_result( 14. ONLY SHOWING the reasoning plan in bullet points. 15. Never include SQL code, table aliases, or assumed table/column names in the reasoning plan. 16. Only mention a table or column when the exact name appears in the DATABASE SCHEMA. +17. Map the user's business wording to schema meaning, not exact word matches only. Use table/column descriptions, aliases, data types, relationships, metrics, calculated fields, and SQL samples to identify grounded concepts. +18. If a business concept cannot be grounded to an exact table/column from the DATABASE SCHEMA, say that the concept is not grounded in the retrieved metadata. Do not create substitute table names, column names, or example SQL. +19. Do not write phrases like "Here is a possible SQL query", "replace with actual table", or any placeholder table/column names. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..30928321ad 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -39,6 +39,9 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +8. Match the user's business wording to schema meaning, not exact word matches only. Use table names, column names, comments, aliases, descriptions, data types, relationships, metrics, calculated fields, and SQL samples when available. +9. Do not invent a table or column because the user used a business word. Select only tables and columns that are explicitly present in the provided schema. +10. If a requested concept cannot be grounded to a schema table or column, omit that table/column instead of creating a placeholder. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -82,6 +85,8 @@ - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. - Match Table names with the definition in the "Create Table" statement. +- Never return placeholder, example, or inferred names unless those exact names are + present in the provided schema. Good luck! diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0d518eba8b..682a98e69c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -511,7 +511,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning="", histories=histories, project_id=ask_request.project_id, sql_samples=sql_samples, @@ -531,7 +531,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, project_id=ask_request.project_id, sql_samples=sql_samples, instructions=instructions, From 61f6bd72c5887018923399b1f91d5435291fa90e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 17:49:54 +0530 Subject: [PATCH 0636/1087] Restore legacy SQL generation flow --- .../generation/followup_sql_generation.py | 46 ++--------- .../followup_sql_generation_reasoning.py | 29 +------ .../pipelines/generation/sql_correction.py | 65 ++------------- .../pipelines/generation/sql_generation.py | 22 ------ .../generation/sql_generation_reasoning.py | 19 +---- .../src/pipelines/generation/utils/sql.py | 79 ++----------------- .../retrieval/db_schema_retrieval.py | 5 -- wren-ai-service/src/web/v1/services/ask.py | 23 +----- 8 files changed, 30 insertions(+), 258 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 b013f9b25c..35cfb8fccf 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -1,8 +1,6 @@ -from __future__ import annotations - import logging import sys -from typing import TYPE_CHECKING, Any +from typing import Any from hamilton import base from hamilton.async_driver import AsyncDriver @@ -26,11 +24,7 @@ from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.utils import trace_cost - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -40,24 +34,11 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. -### TARGET DATA SOURCE ### -{{ data_source }} - -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -interpreting intent and generating SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -### VALID TABLE NAMES ### -Only use these exact table names from the schema. Do not invent, rename, singularize, -pluralize, or add catalog/schema prefixes unless the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -97,14 +78,6 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -### INTENT AND SCHEMA GROUNDING ### -Interpret the user's business terms by matching them to explicit tables, columns, -metrics, views, and relationships in ACTIVE DATASOURCE METADATA. Never reuse table -or column names from SQL SAMPLES or chat history unless those exact names also -appear in ACTIVE DATASOURCE METADATA or VALID TABLE NAMES for the active datasource. -Only apply aggregate functions to columns whose active metadata type supports that -operation. - ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -119,7 +92,6 @@ def prompt( documents: list[str], sql_generation_reasoning: str, prompt_builder: PromptBuilder, - data_source: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -127,13 +99,10 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, - data_source=data_source, documents=documents, - valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -162,7 +131,6 @@ async def generate_sql_in_followup( generator: Any, histories: list[AskHistory], generator_name: str, - data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: history_messages = construct_ask_history_messages(histories) @@ -178,7 +146,6 @@ async def generate_sql_in_followup( async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, - documents: list[str], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -241,18 +208,19 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") - metadata = await retrieve_metadata(project_id or "", self._retriever) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 4d4402ea69..42b28c5b8f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -1,9 +1,7 @@ -from __future__ import annotations - import asyncio import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -19,34 +17,17 @@ ) from src.utils import trace_cost from src.web.v1.services import Configuration - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") sql_generation_reasoning_user_prompt_template = """ -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -planning SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only mention these exact table names in the reasoning plan. Do not invent, rename, -singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless -the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -91,12 +72,10 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - valid_table_names=valid_table_names or [], histories=histories, sql_samples=sql_samples, instructions=construct_instructions( @@ -197,7 +176,6 @@ async def run( instructions: Optional[list[dict]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, - valid_table_names: Optional[list[str]] = None, ): logger.info("Followup SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -205,7 +183,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "histories": histories, "sql_samples": sql_samples or [], "instructions": instructions or [], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 1074b16482..973b8c69a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -25,10 +25,7 @@ logger = logging.getLogger("wren-ai-service") -def get_sql_correction_system_prompt( - sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, -) -> str: +def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) -> str: text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" @@ -55,30 +52,13 @@ def get_sql_correction_system_prompt( sql_correction_user_prompt_template = """ -### TARGET DATA SOURCE ### -{{ data_source }} - {% if documents %} -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -correcting SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} {% endif %} -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only use these exact table names from the schema. If the invalid SQL references a -table not listed here, replace it with the closest listed table only when the schema -clearly supports the user's request. Do not invent, rename, singularize, pluralize, -or add catalog/schema prefixes unless the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - {% if sql_functions %} ### SQL FUNCTIONS ### {% for function in sql_functions %} @@ -94,26 +74,9 @@ def get_sql_correction_system_prompt( {% endif %} ### QUESTION ### -{% if query %} -User's Question: {{ query }} -{% endif %} -{% if invalid_generation_result.original_sql %} -Original SQL: {{ invalid_generation_result.original_sql }} -{% endif %} -Invalid SQL: {{ invalid_generation_result.sql }} -{% if invalid_generation_result.executed_sql %} -Engine SQL that failed validation/execution: {{ invalid_generation_result.executed_sql }} -{% endif %} +SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} -### CORRECTION GROUNDING ### -Use ACTIVE DATASOURCE METADATA and VALID TABLE NAMES as the source of truth. If the -invalid SQL references a table or column not listed above, replace it only when the -active datasource metadata clearly contains an equivalent object that supports the -user's request. Do not invent tables, columns, joins, metrics, or relationships. -Only apply aggregate functions to columns whose active metadata type supports that -operation. - Let's think step by step. """ @@ -124,17 +87,11 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, - data_source: str, - query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( - query=query, - data_source=data_source, documents=documents, - valid_table_names=valid_table_names or [], invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, @@ -150,13 +107,9 @@ async def generate_sql_correction( prompt: dict, generator: Any, generator_name: str, - data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_correction_system_prompt( - sql_knowledge, - data_source=data_source, - ) + current_system_prompt = get_sql_correction_system_prompt(sql_knowledge) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -166,7 +119,6 @@ async def generate_sql_correction( async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, - documents: List[Document], data_source: str, project_id: str | None = None, use_dry_plan: bool = False, @@ -223,20 +175,19 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - query: str | None = None, - valid_table_names: list[str] | None = None, ): logger.info("SQLCorrection pipeline is running...") - metadata = await retrieve_metadata(project_id or "", self._retriever) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, "documents": contexts, - "valid_table_names": valid_table_names or [], - "query": query, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 90b15d2fac..1ee4952b3e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,16 +28,6 @@ sql_generation_user_prompt_template = """ -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only use these exact table names from the retrieved schema. Do not invent, rename, -singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless -the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -82,14 +72,6 @@ ### QUESTION ### User's Question: {{ query }} -### INTENT AND SCHEMA GROUNDING ### -Interpret the user's business terms by matching them to explicit tables, columns, -metrics, views, and relationships in DATABASE SCHEMA. Never copy table or column -names from SQL SAMPLES, REASONING PLAN, or prior assumptions unless those exact -names also appear in DATABASE SCHEMA or VALID TABLE NAMES. If a required metric, -dimension, filter, date, or join key cannot be grounded by the retrieved metadata, -do not replace it with a generic placeholder. - {% if sql_generation_reasoning %} ### REASONING PLAN ### {{ sql_generation_reasoning }} @@ -113,12 +95,10 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - valid_table_names=valid_table_names or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -222,7 +202,6 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - valid_table_names: list[str] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -236,7 +215,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, 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 fc52e5a4d5..00b731cb2c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -22,24 +22,11 @@ sql_generation_reasoning_user_prompt_template = """ -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource, including schema, -tables, columns, metrics, views, and relationships. Use only this metadata when -planning SQL. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -{% if valid_table_names %} -### VALID TABLE NAMES ### -Only mention these exact table names in the reasoning plan. Do not invent, rename, -singularize, pluralize, normalize, abbreviate, or add catalog/schema prefixes unless -the table name is shown that way here. -{% for table_name in valid_table_names %} -- {{ table_name }} -{% endfor %} -{% endif %} - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -75,12 +62,10 @@ def prompt( instructions: list[dict], prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), - valid_table_names: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - valid_table_names=valid_table_names or [], sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, @@ -175,7 +160,6 @@ async def run( instructions: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, - valid_table_names: Optional[list[str]] = None, ): logger.info("SQL Generation Reasoning pipeline is running...") return await self._pipe.execute( @@ -183,7 +167,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "valid_table_names": valid_table_names or [], "sql_samples": sql_samples or [], "instructions": instructions or [], "configuration": configuration, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 305473255d..088282574e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,4 @@ import logging -import re from typing import Any, Dict, List import aiohttp @@ -13,61 +12,11 @@ clean_generation_result, ) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") -def _extract_ddl_columns(ddl: str) -> list[str]: - if not isinstance(ddl, str): - return [] - - column_section = re.search(r"\((.*)\)", ddl, flags=re.DOTALL) - if not column_section: - return [] - - columns = [] - for raw_line in re.split(r",\s*(?:\n|(?=[A-Za-z_\"`\[]))", column_section.group(1)): - line = re.sub(r"/\*.*?\*/", "", raw_line).strip() - line = re.sub(r"^--.*$", "", line).strip().rstrip(",") - if not line: - continue - - first_token = line.split()[0].strip('"`[]') - if first_token.upper() in { - "CONSTRAINT", - "FOREIGN", - "PRIMARY", - "UNIQUE", - "KEY", - }: - continue - columns.append(first_token) - - return columns - - -def format_retrieved_schema_manifest( - documents: list[str] | None, - allowed_table_names: list[str] | None, -) -> list[dict[str, Any]]: - table_names = allowed_table_names or [] - ddls = documents or [] - - manifest = [] - for index, table_name in enumerate(table_names): - if not isinstance(table_name, str) or not table_name.strip(): - continue - ddl = ddls[index] if index < len(ddls) and isinstance(ddls[index], str) else "" - manifest.append( - { - "table_name": table_name.strip(), - "columns": _extract_ddl_columns(ddl), - } - ) - - return manifest - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -218,14 +167,6 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. -- NEVER invent, assume, or rename tables and columns. Generate SQL only from tables and columns present in the provided database schema. -- Treat the retrieved schema manifest and DATABASE SCHEMA as the exact identifier allowlist. Every SELECT, WHERE, JOIN, GROUP BY, HAVING, and ORDER BY table/column must exist there. -- Select columns by business meaning, not by name similarity alone. Match the user's entities, metrics, dimensions, filters, and dates to column names, descriptions, aliases, data types, user instructions, and SQL samples. -- Prefer semantically described business/canonical models, metrics, and views over staging, temp, test, raw, backup, load, or legacy tables when the metadata indicates that distinction. -- For value, amount, total, rate, count, average, or KPI questions, use numeric measures or numeric columns whose description/alias matches the requested metric. Do not SUM or AVG string columns. -- For comparison or "by" questions, use categorical/date dimension columns for grouping and numeric measures for aggregation. -- Do not use technical audit or ingestion columns (created_at, updated_at, loaded_at, inserted_at, file_date, batch_id, row_id, ingestion timestamps, etc.) unless the user explicitly asks about sync, load, audit, or ingestion. -- If multiple columns are equally plausible and no metadata/rule/sample disambiguates them, do not guess; return SQL only when the chosen columns are grounded by the metadata. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. @@ -489,8 +430,8 @@ async def _classify_generation_result( ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; +2. Explicitly state the following information in the reasoning plan: +if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; otherwise, you will put the relative timeframe in the SQL query. 3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. @@ -504,11 +445,6 @@ async def _classify_generation_result( 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. 14. ONLY SHOWING the reasoning plan in bullet points. -15. Never include SQL code, table aliases, or assumed table/column names in the reasoning plan. -16. Only mention a table or column when the exact name appears in the DATABASE SCHEMA. -17. Map the user's business wording to schema meaning, not exact word matches only. Use table/column descriptions, aliases, data types, relationships, metrics, calculated fields, and SQL samples to identify grounded concepts. -18. If a business concept cannot be grounded to an exact table/column from the DATABASE SCHEMA, say that the concept is not grounded in the retrieved metadata. Do not create substitute table names, column names, or example SQL. -19. Do not write phrases like "Here is a possible SQL query", "replace with actual table", or any placeholder table/column names. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -569,16 +505,15 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" You are a helpful assistant that converts natural language queries into ANSI SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query from the provided database schema. +Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. If the section of REASONING PLAN is available in user's input, treat it only as high-level guidance. Ignore any table, column, alias, filter, or SQL fragment from the reasoning plan that is not explicitly present in the DATABASE SCHEMA. -5. Before finalizing, validate the SQL against the retrieved metadata: every table, column, join, filter, GROUP BY, HAVING, and ORDER BY identifier must exist in the provided schema, and selected columns must match the user's business intent. -6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} @@ -619,7 +554,7 @@ def construct_instructions( def construct_ask_history_messages( - histories: list[Any] | list[dict], + histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: messages = [] for history in histories: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 30928321ad..6c8dd7bbe3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -39,9 +39,6 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -8. Match the user's business wording to schema meaning, not exact word matches only. Use table names, column names, comments, aliases, descriptions, data types, relationships, metrics, calculated fields, and SQL samples when available. -9. Do not invent a table or column because the user used a business word. Select only tables and columns that are explicitly present in the provided schema. -10. If a requested concept cannot be grounded to a schema table or column, omit that table/column instead of creating a placeholder. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -85,8 +82,6 @@ - Use table name used in the "Create Table" statement, don't use "alias". - Match Column names with the definition in the "Create Table" statement. - Match Table names with the definition in the "Create Table" statement. -- Never return placeholder, example, or inferred names unless those exact names are - present in the provided schema. Good luck! diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 682a98e69c..6bf806392d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -389,16 +389,8 @@ async def ask( "construct_retrieval_results", {} ) documents = _retrieval_result.get("retrieval_results", []) - table_names = [ - document.get("table_name") - for document in documents - if document.get("table_name") - ] - table_ddls = [ - document.get("table_ddl") - for document in documents - if document.get("table_ddl") - ] + table_names = [document.get("table_name") for document in documents] + table_ddls = [document.get("table_ddl") for document in documents] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -444,7 +436,6 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, - valid_table_names=table_names, ) ).get("post_process", {}) else: @@ -456,7 +447,6 @@ async def ask( instructions=instructions, configuration=ask_request.configurations, query_id=query_id, - valid_table_names=table_names, ) ).get("post_process", {}) @@ -511,7 +501,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning="", + sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, sql_samples=sql_samples, @@ -523,7 +513,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - valid_table_names=table_names, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -531,7 +520,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, sql_samples=sql_samples, instructions=instructions, @@ -542,7 +531,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - valid_table_names=table_names, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -600,7 +588,6 @@ async def ask( instructions=instructions, invalid_generation_result={ "sql": original_sql, - "executed_sql": invalid_sql, "error": sql_diagnosis_reasoning if allow_sql_diagnosis else error_message, @@ -610,8 +597,6 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - query=user_query, - valid_table_names=table_names, ) if valid_generation_result := sql_correction_results[ From a30af6bc9ce60fa8163fef80c2622aa6583bfd1a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 18:13:48 +0530 Subject: [PATCH 0637/1087] Enrich table retrieval index with column semantics --- .../pipelines/indexing/table_description.py | 59 +++++++++++++++++-- 1 file changed, 54 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index cb39b097a3..d1d0a6aa6a 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -14,12 +14,18 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.indexing import AsyncDocumentWriter, DocumentCleaner, MDLValidator +from src.pipelines.indexing import ( + AsyncDocumentWriter, + DocumentCleaner, + MDLValidator, + clean_display_name, +) logger = logging.getLogger("wren-ai-service") MAX_TABLE_DESCRIPTION_COLUMNS = 200 MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH = 4000 +MAX_TABLE_DESCRIPTION_COLUMN_DESCRIPTION_LENGTH = 500 @component @@ -37,6 +43,18 @@ def _truncate_description(self, description: Any) -> str: + "..." ) + def _truncate_column_description(self, description: Any) -> str: + normalized_description = self._normalize_text(description) + if len(normalized_description) <= MAX_TABLE_DESCRIPTION_COLUMN_DESCRIPTION_LENGTH: + return normalized_description + + return ( + normalized_description[ + :MAX_TABLE_DESCRIPTION_COLUMN_DESCRIPTION_LENGTH + ].rstrip() + + "..." + ) + def _format_columns(self, columns: List[Any]) -> str: normalized_columns = [self._normalize_text(column) for column in columns] if len(normalized_columns) <= MAX_TABLE_DESCRIPTION_COLUMNS: @@ -48,6 +66,37 @@ def _format_columns(self, columns: List[Any]) -> str: ] return ", ".join(truncated_columns) + def _properties(self, payload: Dict[str, Any]) -> Dict[str, Any]: + properties = payload.get("properties") + return properties if isinstance(properties, dict) else {} + + def _display_name(self, properties: Dict[str, Any]) -> str: + return clean_display_name( + self._normalize_text(properties.get("displayName", "")) + ) + + def _column_text(self, column: Dict[str, Any]) -> str: + properties = self._properties(column) + parts = [self._normalize_text(column.get("name", ""))] + + display_name = self._display_name(properties) + if display_name: + parts.append(f"alias: {display_name}") + + data_type = self._normalize_text( + column.get("type", column.get("data_type", "")) + ) + if data_type: + parts.append(f"type: {data_type}") + + description = self._truncate_column_description( + properties.get("description", "") + ) + if description: + parts.append(f"description: {description}") + + return " | ".join(part for part in parts if part) + @component.output_types(documents=List[Document]) def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None): def _additional_meta() -> Dict[str, Any]: @@ -78,15 +127,13 @@ def _additional_meta() -> Dict[str, Any]: def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[str]: def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: - properties = payload.get("properties") - if not isinstance(properties, dict): - properties = {} + properties = self._properties(payload) return { "mdl_type": mdl_type, "name": payload.get("name"), "columns": [ - column.get("name", "") + self._column_text(column) for column in payload.get("columns", []) if isinstance(column, dict) ], @@ -102,6 +149,8 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: return [ { "name": resource["name"], + "type": resource["mdl_type"], + "alias": self._display_name(resource["properties"]), "description": self._truncate_description( resource["properties"].get("description", "") ), From aa13116ef4c1a17a054aa6656f46d2203bde338c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 19:09:28 +0530 Subject: [PATCH 0638/1087] Ground SQL generation to retrieved schema --- .../generation/followup_sql_generation.py | 7 ++++++ .../followup_sql_generation_reasoning.py | 6 +++++ .../pipelines/generation/sql_correction.py | 8 +++++++ .../pipelines/generation/sql_generation.py | 7 ++++++ .../generation/sql_generation_reasoning.py | 6 +++++ .../src/pipelines/generation/utils/sql.py | 24 ++++++++++++------- 6 files changed, 50 insertions(+), 8 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 35cfb8fccf..06fb69824b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -39,6 +39,13 @@ {{ document }} {% endfor %} +### SCHEMA GROUNDING ### +- Use only tables and columns that appear in DATABASE SCHEMA. +- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. +- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. +- If the REASONING PLAN, SQL SAMPLES, USER INSTRUCTIONS, or QUERY HISTORY mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. +- Before returning the SQL, verify every table and column name exists exactly in DATABASE SCHEMA. + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 42b28c5b8f..2b82a199cb 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -28,6 +28,12 @@ {{ document }} {% endfor %} +### SCHEMA GROUNDING ### +- Mention only tables and columns that appear in DATABASE SCHEMA. +- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. +- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. +- If SQL SAMPLES, USER INSTRUCTIONS, or QUERY HISTORY mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..7f3b1ff75f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -36,6 +36,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). 2. Then, generate the syntactically correct ANSI SQL query to correct the error. +3. If the error says a table or column does not exist, replace it only with an exact table or column identifier from DATABASE SCHEMA. Do not create a normalized, friendly, translated, or guessed identifier. +4. The corrected SQL must use only tables and columns that appear in DATABASE SCHEMA. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -57,6 +59,12 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% for document in documents %} {{ document }} {% endfor %} + +### SCHEMA GROUNDING ### +- Use only tables and columns that appear in DATABASE SCHEMA. +- Resolve invalid identifiers by matching the user's wording and error message to existing schema comments, aliases, descriptions, and column names. +- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. +- Before returning the corrected SQL, verify every table and column name exists exactly in DATABASE SCHEMA. {% endif %} {% if sql_functions %} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..3efa510d12 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -33,6 +33,13 @@ {{ document }} {% endfor %} +### SCHEMA GROUNDING ### +- Use only tables and columns that appear in DATABASE SCHEMA. +- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. +- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. +- If the REASONING PLAN, SQL SAMPLES, or USER INSTRUCTIONS mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. +- Before returning the SQL, verify every table and column name exists exactly in DATABASE SCHEMA. + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} 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 00b731cb2c..a08a94ac3d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -27,6 +27,12 @@ {{ document }} {% endfor %} +### SCHEMA GROUNDING ### +- Mention only tables and columns that appear in DATABASE SCHEMA. +- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. +- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. +- If SQL SAMPLES or USER INSTRUCTIONS mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..086a81bdd5 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -167,6 +167,9 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. +- Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. Never invent a normalized, friendly, translated, or guessed identifier. +- If a user uses business wording that does not exactly match a column name, map it only to an existing table or column by using the schema comments, aliases, descriptions, and available column names. +- When the schema exposes raw or generated names, use those exact names in SQL. Do not replace them with natural-language names. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. @@ -437,14 +440,17 @@ async def _classify_generation_result( 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. -7. Give a step by step reasoning plan in order to answer user's question. -8. The reasoning plan should be in the language same as the language user provided in the input. -9. Don't include SQL in the reasoning plan. -10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. -11. Do not include ```markdown or ``` in the answer. -12. A table name in the reasoning plan must be in this format: `table: `. -13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +7. Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. +8. Do not mention a table or column in the reasoning plan unless it appears in the DATABASE SCHEMA section. +9. If the user's wording is different from the schema names, map the wording to existing schema names by using comments, aliases, descriptions, and available column names. Do not invent normalized or friendly names. +10. Give a step by step reasoning plan in order to answer user's question. +11. The reasoning plan should be in the language same as the language user provided in the input. +12. Don't include SQL in the reasoning plan. +13. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. +14. Do not include ```markdown or ``` in the answer. +15. A table name in the reasoning plan must be in this format: `table: `. +16. A column name in the reasoning plan must be in this format: `column: .`. +17. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -514,6 +520,8 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. 4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +6. DATABASE SCHEMA is more authoritative than the reasoning plan, SQL samples, and user wording. If any of those mention a table or column that is not present in DATABASE SCHEMA, do not use it. +7. For every table and column in the final SQL, verify that the exact identifier appears in DATABASE SCHEMA before returning the SQL. {text_to_sql_rules} From 1f6cc7d925732ab850f607a1864a2d21e7caabc5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 19:29:28 +0530 Subject: [PATCH 0639/1087] Add schema identifier grounding to SQL prompts --- .../generation/followup_sql_generation.py | 8 ++ .../followup_sql_generation_reasoning.py | 8 ++ .../pipelines/generation/sql_correction.py | 8 ++ .../pipelines/generation/sql_generation.py | 8 ++ .../generation/sql_generation_reasoning.py | 8 ++ .../src/pipelines/generation/utils/sql.py | 123 ++++++++++++++++++ 6 files changed, 163 insertions(+) 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 06fb69824b..0ddd1270d2 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -15,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, + construct_schema_grounding_context, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -45,6 +46,11 @@ - Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. - If the REASONING PLAN, SQL SAMPLES, USER INSTRUCTIONS, or QUERY HISTORY mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. - Before returning the SQL, verify every table and column name exists exactly in DATABASE SCHEMA. +{% if schema_grounding_context %} + +### VALID SCHEMA IDENTIFIERS ### +{{ schema_grounding_context }} +{% endif %} {% if calculated_field_instructions %} {{ calculated_field_instructions }} @@ -107,9 +113,11 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: + schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, + schema_grounding_context=schema_grounding_context, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 2b82a199cb..e5844d1253 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -12,6 +12,7 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( + construct_schema_grounding_context, construct_instructions, sql_generation_reasoning_system_prompt, ) @@ -33,6 +34,11 @@ - Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. - Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. - If SQL SAMPLES, USER INSTRUCTIONS, or QUERY HISTORY mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. +{% if schema_grounding_context %} + +### VALID SCHEMA IDENTIFIERS ### +{{ schema_grounding_context }} +{% endif %} {% if sql_samples %} ### SQL SAMPLES ### @@ -79,9 +85,11 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: + schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, + schema_grounding_context=schema_grounding_context, histories=histories, sql_samples=sql_samples, instructions=construct_instructions( diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 7f3b1ff75f..034a45a5f0 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + construct_schema_grounding_context, construct_instructions, get_text_to_sql_rules, ) @@ -65,6 +66,11 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) - Resolve invalid identifiers by matching the user's wording and error message to existing schema comments, aliases, descriptions, and column names. - Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. - Before returning the corrected SQL, verify every table and column name exists exactly in DATABASE SCHEMA. +{% if schema_grounding_context %} + +### VALID SCHEMA IDENTIFIERS ### +{{ schema_grounding_context }} +{% endif %} {% endif %} {% if sql_functions %} @@ -98,8 +104,10 @@ def prompt( instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: + schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( documents=documents, + schema_grounding_context=schema_grounding_context, invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 3efa510d12..df087f378a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + construct_schema_grounding_context, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -39,6 +40,11 @@ - Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. - If the REASONING PLAN, SQL SAMPLES, or USER INSTRUCTIONS mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. - Before returning the SQL, verify every table and column name exists exactly in DATABASE SCHEMA. +{% if schema_grounding_context %} + +### VALID SCHEMA IDENTIFIERS ### +{{ schema_grounding_context }} +{% endif %} {% if calculated_field_instructions %} {{ calculated_field_instructions }} @@ -103,9 +109,11 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: + schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, + schema_grounding_context=schema_grounding_context, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, 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 a08a94ac3d..3e0f561cbd 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -12,6 +12,7 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( + construct_schema_grounding_context, construct_instructions, sql_generation_reasoning_system_prompt, ) @@ -32,6 +33,11 @@ - Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. - Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. - If SQL SAMPLES or USER INSTRUCTIONS mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. +{% if schema_grounding_context %} + +### VALID SCHEMA IDENTIFIERS ### +{{ schema_grounding_context }} +{% endif %} {% if sql_samples %} ### SQL SAMPLES ### @@ -69,9 +75,11 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: + schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, + schema_grounding_context=schema_grounding_context, sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 086a81bdd5..d1c7dd6c6d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, List import aiohttp @@ -17,6 +18,127 @@ logger = logging.getLogger("wren-ai-service") +_SQL_IDENTIFIER_PATTERN = r'"([^"]+)"|`([^`]+)`|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_$]*)' +_CREATE_TABLE_PATTERN = re.compile( + rf"\bCREATE\s+(?:TABLE|VIEW)\s+{_SQL_IDENTIFIER_PATTERN}", + re.IGNORECASE, +) +_COLUMN_PATTERN = re.compile(rf"^\s*{_SQL_IDENTIFIER_PATTERN}\s+([A-Za-z][\w()]+)") +_COLUMN_COMMENT_PATTERN = re.compile( + rf"--\s*(\{{.*?\}})\s*\n\s*{_SQL_IDENTIFIER_PATTERN}\s+([A-Za-z][\w()]+)", + re.DOTALL, +) +_SKIPPED_COLUMN_TOKENS = { + "CONSTRAINT", + "FOREIGN", + "PRIMARY", + "UNIQUE", + "CHECK", + "KEY", + "REFERENCES", +} +MAX_SCHEMA_GROUNDING_COLUMNS_PER_TABLE = 160 +MAX_SCHEMA_GROUNDING_CONTEXT_LENGTH = 20000 + + +def _first_identifier(match: re.Match) -> str: + for group in match.groups(): + if group: + return group.strip() + return "" + + +def _truncate_schema_text(value: Any, limit: int = 120) -> str: + text = " ".join(str(value or "").split()) + if len(text) <= limit: + return text + return text[:limit].rstrip() + "..." + + +def construct_schema_grounding_context(documents: list[str] | None) -> str: + if not documents: + return "" + + sections: list[str] = [] + + for document in documents: + table_match = _CREATE_TABLE_PATTERN.search(document or "") + table_name = _first_identifier(table_match) if table_match else "" + if not table_name: + continue + + comment_metadata_by_column: dict[str, dict[str, Any]] = {} + for comment_match in _COLUMN_COMMENT_PATTERN.finditer(document): + column_name = next( + ( + group.strip() + for group in comment_match.groups()[1:-1] + if group and group.strip() + ), + "", + ) + try: + metadata = orjson.loads(comment_match.group(1)) + except Exception: + metadata = {} + comment_metadata_by_column[column_name] = metadata + + columns: list[str] = [] + seen_columns: set[str] = set() + for line in (document or "").splitlines(): + if len(columns) >= MAX_SCHEMA_GROUNDING_COLUMNS_PER_TABLE: + break + + match = _COLUMN_PATTERN.match(line) + if not match: + continue + + column_name = _first_identifier(match) + if ( + not column_name + or column_name.upper() in _SKIPPED_COLUMN_TOKENS + or column_name in seen_columns + ): + continue + + seen_columns.add(column_name) + metadata = comment_metadata_by_column.get(column_name, {}) + details = [column_name] + + data_type = _truncate_schema_text(match.group(5)) + if data_type: + details.append(f"type={data_type}") + + alias = _truncate_schema_text(metadata.get("alias")) + if alias and alias != column_name: + details.append(f"alias={alias}") + + description = _truncate_schema_text(metadata.get("description")) + if description: + details.append(f"description={description}") + + columns.append("; ".join(details)) + + if columns: + if len(seen_columns) >= MAX_SCHEMA_GROUNDING_COLUMNS_PER_TABLE: + columns.append("...") + + sections.append( + f"- table `{table_name}` valid columns: {', '.join(columns)}" + ) + else: + sections.append(f"- table `{table_name}`") + + schema_grounding_context = "\n".join(sections) + if len(schema_grounding_context) <= MAX_SCHEMA_GROUNDING_CONTEXT_LENGTH: + return schema_grounding_context + + return ( + schema_grounding_context[:MAX_SCHEMA_GROUNDING_CONTEXT_LENGTH].rstrip() + + "\n..." + ) + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -196,6 +318,7 @@ async def _classify_generation_result( - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) +- For date parts such as year, month, day, quarter, or week, use `EXTRACT( FROM CAST( AS TIMESTAMP WITH TIME ZONE))`. Do not use `YEAR()`, `MONTH()`, `DAY()`, or `DATEPART()` unless the SQL FUNCTIONS section explicitly says that function is supported. - If the user asks for a specific date, please give the date range in SQL query - example: "What is the total revenue for the month of 2024-11-01?" - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" From 90755b0d2798d32153af305200d3c846b1abb487 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 20:18:55 +0530 Subject: [PATCH 0640/1087] Refactor schema grounding and SQL adjustment handling --- .../generation/followup_sql_generation.py | 28 ++-- .../followup_sql_generation_reasoning.py | 27 ++-- .../pipelines/generation/sql_correction.py | 31 ++--- .../pipelines/generation/sql_generation.py | 28 ++-- .../generation/sql_generation_reasoning.py | 27 ++-- .../src/pipelines/generation/utils/sql.py | 122 ------------------ .../apollo/server/services/askingService.ts | 52 ++++++-- .../pages/home/promptThread/ChartAnswer.tsx | 2 +- .../home/promptThread/TextBasedAnswer.tsx | 2 +- .../home/promptThread/ViewSQLTabContent.tsx | 2 +- 10 files changed, 85 insertions(+), 236 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 0ddd1270d2..9624e941b6 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -15,7 +15,6 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, - construct_schema_grounding_context, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -35,23 +34,6 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - -### SCHEMA GROUNDING ### -- Use only tables and columns that appear in DATABASE SCHEMA. -- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. -- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. -- If the REASONING PLAN, SQL SAMPLES, USER INSTRUCTIONS, or QUERY HISTORY mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. -- Before returning the SQL, verify every table and column name exists exactly in DATABASE SCHEMA. -{% if schema_grounding_context %} - -### VALID SCHEMA IDENTIFIERS ### -{{ schema_grounding_context }} -{% endif %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -95,6 +77,14 @@ {{ sql_generation_reasoning }} Let's think step by step. + +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + +### FINAL SQL INSTRUCTION ### +Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. """ @@ -113,11 +103,9 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, - schema_grounding_context=schema_grounding_context, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index e5844d1253..a7c9b7e15d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -12,7 +12,6 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( - construct_schema_grounding_context, construct_instructions, sql_generation_reasoning_system_prompt, ) @@ -24,22 +23,6 @@ sql_generation_reasoning_user_prompt_template = """ -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - -### SCHEMA GROUNDING ### -- Mention only tables and columns that appear in DATABASE SCHEMA. -- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. -- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. -- If SQL SAMPLES, USER INSTRUCTIONS, or QUERY HISTORY mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. -{% if schema_grounding_context %} - -### VALID SCHEMA IDENTIFIERS ### -{{ schema_grounding_context }} -{% endif %} - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -71,6 +54,14 @@ Current Time: {{ current_time }} Let's think step by step. + +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + +### FINAL REASONING INSTRUCTION ### +Mention only the exact table and column identifiers in DATABASE SCHEMA above. """ @@ -85,11 +76,9 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: - schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, - schema_grounding_context=schema_grounding_context, histories=histories, sql_samples=sql_samples, instructions=construct_instructions( diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 034a45a5f0..44a1da94a2 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,7 +15,6 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, - construct_schema_grounding_context, construct_instructions, get_text_to_sql_rules, ) @@ -55,24 +54,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) sql_correction_user_prompt_template = """ -{% if documents %} -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - -### SCHEMA GROUNDING ### -- Use only tables and columns that appear in DATABASE SCHEMA. -- Resolve invalid identifiers by matching the user's wording and error message to existing schema comments, aliases, descriptions, and column names. -- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. -- Before returning the corrected SQL, verify every table and column name exists exactly in DATABASE SCHEMA. -{% if schema_grounding_context %} - -### VALID SCHEMA IDENTIFIERS ### -{{ schema_grounding_context }} -{% endif %} -{% endif %} - {% if sql_functions %} ### SQL FUNCTIONS ### {% for function in sql_functions %} @@ -92,6 +73,16 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Error Message: {{ invalid_generation_result.error }} Let's think step by step. + +{% if documents %} +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + +### FINAL CORRECTION INSTRUCTION ### +Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat the invalid SQL and error message as guidance only. +{% endif %} """ @@ -104,10 +95,8 @@ def prompt( instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: - schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( documents=documents, - schema_grounding_context=schema_grounding_context, invalid_generation_result=invalid_generation_result, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index df087f378a..df3cbde82e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,7 +14,6 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, - construct_schema_grounding_context, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -29,23 +28,6 @@ sql_generation_user_prompt_template = """ -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - -### SCHEMA GROUNDING ### -- Use only tables and columns that appear in DATABASE SCHEMA. -- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. -- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. -- If the REASONING PLAN, SQL SAMPLES, or USER INSTRUCTIONS mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. -- Before returning the SQL, verify every table and column name exists exactly in DATABASE SCHEMA. -{% if schema_grounding_context %} - -### VALID SCHEMA IDENTIFIERS ### -{{ schema_grounding_context }} -{% endif %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -91,6 +73,14 @@ {% endif %} Let's think step by step. + +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + +### FINAL SQL INSTRUCTION ### +Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. """ @@ -109,11 +99,9 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, - schema_grounding_context=schema_grounding_context, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, 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 3e0f561cbd..110cda207a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -12,7 +12,6 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( - construct_schema_grounding_context, construct_instructions, sql_generation_reasoning_system_prompt, ) @@ -23,22 +22,6 @@ sql_generation_reasoning_user_prompt_template = """ -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - -### SCHEMA GROUNDING ### -- Mention only tables and columns that appear in DATABASE SCHEMA. -- Resolve user business wording by matching it to existing schema comments, aliases, descriptions, and column names. -- Do not create normalized or friendly identifiers that are not present in DATABASE SCHEMA. -- If SQL SAMPLES or USER INSTRUCTIONS mention identifiers that are not present in DATABASE SCHEMA, ignore those identifiers. -{% if schema_grounding_context %} - -### VALID SCHEMA IDENTIFIERS ### -{{ schema_grounding_context }} -{% endif %} - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -62,6 +45,14 @@ Current Time: {{ current_time }} Let's think step by step. + +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + +### FINAL REASONING INSTRUCTION ### +Mention only the exact table and column identifiers in DATABASE SCHEMA above. """ @@ -75,11 +66,9 @@ def prompt( prompt_builder: PromptBuilder, configuration: Configuration | None = Configuration(), ) -> dict: - schema_grounding_context = construct_schema_grounding_context(documents) _prompt = prompt_builder.run( query=query, documents=documents, - schema_grounding_context=schema_grounding_context, sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index d1c7dd6c6d..a572b29456 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,4 @@ import logging -import re from typing import Any, Dict, List import aiohttp @@ -18,127 +17,6 @@ logger = logging.getLogger("wren-ai-service") -_SQL_IDENTIFIER_PATTERN = r'"([^"]+)"|`([^`]+)`|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_$]*)' -_CREATE_TABLE_PATTERN = re.compile( - rf"\bCREATE\s+(?:TABLE|VIEW)\s+{_SQL_IDENTIFIER_PATTERN}", - re.IGNORECASE, -) -_COLUMN_PATTERN = re.compile(rf"^\s*{_SQL_IDENTIFIER_PATTERN}\s+([A-Za-z][\w()]+)") -_COLUMN_COMMENT_PATTERN = re.compile( - rf"--\s*(\{{.*?\}})\s*\n\s*{_SQL_IDENTIFIER_PATTERN}\s+([A-Za-z][\w()]+)", - re.DOTALL, -) -_SKIPPED_COLUMN_TOKENS = { - "CONSTRAINT", - "FOREIGN", - "PRIMARY", - "UNIQUE", - "CHECK", - "KEY", - "REFERENCES", -} -MAX_SCHEMA_GROUNDING_COLUMNS_PER_TABLE = 160 -MAX_SCHEMA_GROUNDING_CONTEXT_LENGTH = 20000 - - -def _first_identifier(match: re.Match) -> str: - for group in match.groups(): - if group: - return group.strip() - return "" - - -def _truncate_schema_text(value: Any, limit: int = 120) -> str: - text = " ".join(str(value or "").split()) - if len(text) <= limit: - return text - return text[:limit].rstrip() + "..." - - -def construct_schema_grounding_context(documents: list[str] | None) -> str: - if not documents: - return "" - - sections: list[str] = [] - - for document in documents: - table_match = _CREATE_TABLE_PATTERN.search(document or "") - table_name = _first_identifier(table_match) if table_match else "" - if not table_name: - continue - - comment_metadata_by_column: dict[str, dict[str, Any]] = {} - for comment_match in _COLUMN_COMMENT_PATTERN.finditer(document): - column_name = next( - ( - group.strip() - for group in comment_match.groups()[1:-1] - if group and group.strip() - ), - "", - ) - try: - metadata = orjson.loads(comment_match.group(1)) - except Exception: - metadata = {} - comment_metadata_by_column[column_name] = metadata - - columns: list[str] = [] - seen_columns: set[str] = set() - for line in (document or "").splitlines(): - if len(columns) >= MAX_SCHEMA_GROUNDING_COLUMNS_PER_TABLE: - break - - match = _COLUMN_PATTERN.match(line) - if not match: - continue - - column_name = _first_identifier(match) - if ( - not column_name - or column_name.upper() in _SKIPPED_COLUMN_TOKENS - or column_name in seen_columns - ): - continue - - seen_columns.add(column_name) - metadata = comment_metadata_by_column.get(column_name, {}) - details = [column_name] - - data_type = _truncate_schema_text(match.group(5)) - if data_type: - details.append(f"type={data_type}") - - alias = _truncate_schema_text(metadata.get("alias")) - if alias and alias != column_name: - details.append(f"alias={alias}") - - description = _truncate_schema_text(metadata.get("description")) - if description: - details.append(f"description={description}") - - columns.append("; ".join(details)) - - if columns: - if len(seen_columns) >= MAX_SCHEMA_GROUNDING_COLUMNS_PER_TABLE: - columns.append("...") - - sections.append( - f"- table `{table_name}` valid columns: {', '.join(columns)}" - ) - else: - sections.append(f"- table `{table_name}`") - - schema_grounding_context = "\n".join(sections) - if len(schema_grounding_context) <= MAX_SCHEMA_GROUNDING_CONTEXT_LENGTH: - return schema_grounding_context - - return ( - schema_grounding_context[:MAX_SCHEMA_GROUNDING_CONTEXT_LENGTH].rstrip() - + "\n..." - ) - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index f5b1115cae..55718b1174 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1,6 +1,8 @@ import { IWrenAIAdaptor } from '@server/adaptors/wrenAIAdaptor'; import { AskResultStatus, + AskResultType, + AskCandidateType, RecommendationQuestionsResult, RecommendationQuestionsInput, RecommendationQuestion, @@ -17,7 +19,6 @@ import { IThreadResponseRepository, ThreadResponse, ThreadResponseAnswerDetail, - ThreadResponseAdjustmentType, } from '../repositories/threadResponseRepository'; import { getLogger } from '@server/utils'; import { isEmpty, isNil } from 'lodash'; @@ -1409,18 +1410,45 @@ export class AskingService implements IAskingService { } await this.ensureThreadInCurrentProject(response.threadId); - return await this.threadResponseRepository.createOne({ - sql: input.sql, - threadId: response.threadId, - question: response.question, - adjustment: { - type: ThreadResponseAdjustmentType.APPLY_SQL, - payload: { - originalThreadResponseId: response.id, - sql: input.sql, - }, - }, + const project = await this.getProjectForThreadResponse(response); + const deployment = await this.deployService.getLastDeployment(project.id); + await this.queryService.preview(input.sql, { + project, + manifest: deployment.manifest, + modelingOnly: false, + limit: 1, + cacheEnabled: false, }); + + const updatedResponse = await this.threadResponseRepository.updateOne( + response.id, + { + sql: input.sql, + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, + }, + ); + + if (response.askingTaskId) { + await this.askingTaskRepository.updateOne(response.askingTaskId, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FINISHED, + response: [ + { + type: AskCandidateType.LLM, + sql: input.sql, + }, + ], + error: null, + invalidSql: null, + }, + }); + } + + return updatedResponse; } public async adjustThreadResponseAnswer( diff --git a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx index e5ef14bb01..2db457ba35 100644 --- a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx @@ -124,7 +124,7 @@ export default function ChartAnswer(props: AnswerResultProps) { previewData({ variables: { where: { responseId: threadResponse.id } }, }); - }, [previewData, status, threadResponse.id]); + }, [previewData, status, threadResponse.id, threadResponse.sql]); const chartSpec = useMemo(() => { if ( diff --git a/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx index 37126e24c2..78d740d022 100644 --- a/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx @@ -135,7 +135,7 @@ export default function TextBasedAnswer(props: AnswerResultProps) { onInitPreviewDone(); } - }, [isLastThreadResponse, allowPreviewData]); + }, [isLastThreadResponse, allowPreviewData, threadResponse.sql]); const loading = !getIsLoadingFinished(status); diff --git a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx index 135f8e2e5f..1cee69660a 100644 --- a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx @@ -71,7 +71,7 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { if (isLastThreadResponse) { autoTriggerPreviewDataButton(); } - }, [isLastThreadResponse]); + }, [isLastThreadResponse, threadResponse.sql]); const { id, sql } = threadResponse; From 4d5fe09780008c575d52ba9e312436d4360b4490 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 22:07:10 +0530 Subject: [PATCH 0641/1087] Use TYPE_CHECKING and future annotations for AskHistory --- wren-ai-service/src/pipelines/generation/utils/sql.py | 10 +++++++--- 1 file changed, 7 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 a572b29456..c02bd752be 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,7 @@ +from __future__ import annotations + import logging -from typing import Any, Dict, List +from typing import TYPE_CHECKING, Any, Dict, List import aiohttp import orjson @@ -12,7 +14,9 @@ clean_generation_result, ) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge -from src.web.v1.services.ask import AskHistory + +if TYPE_CHECKING: + from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -563,7 +567,7 @@ def construct_instructions( def construct_ask_history_messages( - histories: list[AskHistory] | list[dict], + histories: list["AskHistory"] | list[dict], ) -> list[ChatMessage]: messages = [] for history in histories: From 925f003e1c0dfd7d5cd9940b5298fec7cc886741 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Fri, 24 Jul 2026 23:12:44 +0530 Subject: [PATCH 0642/1087] Remove unsupported ask greeting shortcut --- wren-ai-service/src/web/v1/routers/ask.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/wren-ai-service/src/web/v1/routers/ask.py b/wren-ai-service/src/web/v1/routers/ask.py index 8e63f711d1..595e8fe16e 100644 --- a/wren-ai-service/src/web/v1/routers/ask.py +++ b/wren-ai-service/src/web/v1/routers/ask.py @@ -36,16 +36,6 @@ async def ask( status="understanding", ) - if ask_service._is_greeting_query(ask_request.query): - ask_service._general_streaming_results[query_id] = ( - ask_service._build_greeting_response(ask_request.query) - ) - ask_service._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - ) - return AskResponse(query_id=query_id) - task = asyncio.create_task( ask_service.ask( ask_request, From 40213337b7a22f5e1f46a359b8849827931c317f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 24 Jul 2026 23:55:17 +0530 Subject: [PATCH 0643/1087] Reorder prompts and enhance SQL safety rules --- .../src/pipelines/generation/sql_correction.py | 18 +++++++++--------- .../src/pipelines/generation/sql_generation.py | 10 +++++----- .../src/pipelines/generation/utils/sql.py | 2 +- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 44a1da94a2..52262fbf97 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -54,6 +54,13 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) sql_correction_user_prompt_template = """ +{% if documents %} +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} +{% endif %} + {% if sql_functions %} ### SQL FUNCTIONS ### {% for function in sql_functions %} @@ -72,17 +79,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} -Let's think step by step. - -{% if documents %} -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - ### FINAL CORRECTION INSTRUCTION ### Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat the invalid SQL and error message as guidance only. -{% endif %} + +Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index df3cbde82e..cee5020ef3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,6 +28,11 @@ sql_generation_user_prompt_template = """ +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -74,11 +79,6 @@ Let's think step by step. -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - ### FINAL SQL INSTRUCTION ### Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c02bd752be..ff52bd8107 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -169,7 +169,7 @@ async def _classify_generation_result( _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### -- ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. +- ONLY USE SELECT statements. Do not generate ALTER, CREATE, DROP, INSERT, UPDATE, DELETE, MERGE, TRUNCATE, GRANT, REVOKE, or any other statement that can change the database or schema. - ONLY USE the tables and columns mentioned in the database schema. - Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. Never invent a normalized, friendly, translated, or guessed identifier. - If a user uses business wording that does not exactly match a column name, map it only to an existing table or column by using the schema comments, aliases, descriptions, and available column names. From 3d50dcb7d7f70832b9334b009215bb0771115eff Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 01:02:58 +0530 Subject: [PATCH 0644/1087] Restore legacy SQL prompt context flow --- .../generation/followup_sql_generation.py | 12 ++++++------ .../followup_sql_generation_reasoning.py | 12 ++++++------ .../src/pipelines/generation/sql_correction.py | 2 +- .../src/pipelines/generation/sql_generation.py | 2 +- .../generation/sql_generation_reasoning.py | 12 ++++++------ .../src/pipelines/generation/utils/sql.py | 14 ++++++++++---- 6 files changed, 30 insertions(+), 24 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 9624e941b6..1cbf000588 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -34,6 +34,11 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -78,13 +83,8 @@ Let's think step by step. -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - ### FINAL SQL INSTRUCTION ### -Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. +Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. Do not convert underscores to dots or add database/schema prefixes. """ diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index a7c9b7e15d..34c983f748 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -23,6 +23,11 @@ sql_generation_reasoning_user_prompt_template = """ +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -55,13 +60,8 @@ Let's think step by step. -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - ### FINAL REASONING INSTRUCTION ### -Mention only the exact table and column identifiers in DATABASE SCHEMA above. +Mention only the exact table and column identifiers in DATABASE SCHEMA above. Do not convert underscores to dots or add database/schema prefixes. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 52262fbf97..35d963dae0 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -80,7 +80,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Error Message: {{ invalid_generation_result.error }} ### FINAL CORRECTION INSTRUCTION ### -Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat the invalid SQL and error message as guidance only. +Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat the invalid SQL and error message as guidance only. Do not convert underscores to dots or add database/schema prefixes. Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index cee5020ef3..f63632fe77 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -80,7 +80,7 @@ Let's think step by step. ### FINAL SQL INSTRUCTION ### -Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. +Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. Do not convert underscores to dots or add database/schema prefixes. """ 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 110cda207a..3df34c1ea3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -22,6 +22,11 @@ sql_generation_reasoning_user_prompt_template = """ +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document }} +{% endfor %} + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} @@ -46,13 +51,8 @@ Let's think step by step. -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document }} -{% endfor %} - ### FINAL REASONING INSTRUCTION ### -Mention only the exact table and column identifiers in DATABASE SCHEMA above. +Mention only the exact table and column identifiers in DATABASE SCHEMA above. Do not convert underscores to dots or add database/schema prefixes. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ff52bd8107..f883ca2848 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -172,6 +172,8 @@ async def _classify_generation_result( - ONLY USE SELECT statements. Do not generate ALTER, CREATE, DROP, INSERT, UPDATE, DELETE, MERGE, TRUNCATE, GRANT, REVOKE, or any other statement that can change the database or schema. - ONLY USE the tables and columns mentioned in the database schema. - Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. Never invent a normalized, friendly, translated, or guessed identifier. +- Do not use names from policies, instructions, SQL functions, SQL samples, query history, reasoning text, error messages, or the user's wording as table or column identifiers unless they appear exactly in DATABASE SCHEMA. +- Policy documents and policy names are not database tables or columns unless they appear exactly in DATABASE SCHEMA. - If a user uses business wording that does not exactly match a column name, map it only to an existing table or column by using the schema comments, aliases, descriptions, and available column names. - When the schema exposes raw or generated names, use those exact names in SQL. Do not replace them with natural-language names. - ONLY USE "*" if the user query asks for all the columns of a table. @@ -474,9 +476,13 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: if sql_knowledge is not None: - return _extract_from_sql_knowledge( - sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES - ) + value = getattr(sql_knowledge, "text_to_sql_rule", "") + if value and value.strip(): + return ( + f"{_DEFAULT_TEXT_TO_SQL_RULES}\n\n" + "### PROJECT SQL RULES ###\n" + f"{value.strip()}" + ) return _DEFAULT_TEXT_TO_SQL_RULES @@ -523,7 +529,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +4. YOU MUST USE the reasoning plan only as analytical guidance if the section of REASONING PLAN is available in user's input. Do not copy table or column names from it unless they appear exactly in DATABASE SCHEMA. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. 6. DATABASE SCHEMA is more authoritative than the reasoning plan, SQL samples, and user wording. If any of those mention a table or column that is not present in DATABASE SCHEMA, do not use it. 7. For every table and column in the final SQL, verify that the exact identifier appears in DATABASE SCHEMA before returning the SQL. From c80c157b04b15cbaaac02029744023fcebed59a7 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 01:30:53 +0530 Subject: [PATCH 0645/1087] Restore legacy SQL prompt endings --- .../src/pipelines/generation/followup_sql_generation.py | 3 --- .../pipelines/generation/followup_sql_generation_reasoning.py | 3 --- wren-ai-service/src/pipelines/generation/sql_correction.py | 3 --- wren-ai-service/src/pipelines/generation/sql_generation.py | 3 --- .../src/pipelines/generation/sql_generation_reasoning.py | 3 --- 5 files changed, 15 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 1cbf000588..35cfb8fccf 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -82,9 +82,6 @@ {{ sql_generation_reasoning }} Let's think step by step. - -### FINAL SQL INSTRUCTION ### -Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. Do not convert underscores to dots or add database/schema prefixes. """ diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 34c983f748..42b28c5b8f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -59,9 +59,6 @@ Current Time: {{ current_time }} Let's think step by step. - -### FINAL REASONING INSTRUCTION ### -Mention only the exact table and column identifiers in DATABASE SCHEMA above. Do not convert underscores to dots or add database/schema prefixes. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 35d963dae0..ed1f322dd3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -79,9 +79,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} -### FINAL CORRECTION INSTRUCTION ### -Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat the invalid SQL and error message as guidance only. Do not convert underscores to dots or add database/schema prefixes. - Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index f63632fe77..1ee4952b3e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -78,9 +78,6 @@ {% endif %} Let's think step by step. - -### FINAL SQL INSTRUCTION ### -Use only the exact table and column identifiers in DATABASE SCHEMA above. Treat all other sections as guidance only. Do not convert underscores to dots or add database/schema prefixes. """ 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 3df34c1ea3..00b731cb2c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -50,9 +50,6 @@ Current Time: {{ current_time }} Let's think step by step. - -### FINAL REASONING INSTRUCTION ### -Mention only the exact table and column identifiers in DATABASE SCHEMA above. Do not convert underscores to dots or add database/schema prefixes. """ From be326f352e06e8b4e7dc296858e0143977cd3c97 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 02:08:56 +0530 Subject: [PATCH 0646/1087] Tighten legacy SQL prompt grounding --- .../src/pipelines/generation/utils/sql.py | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index f883ca2848..01929d4049 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -176,10 +176,12 @@ async def _classify_generation_result( - Policy documents and policy names are not database tables or columns unless they appear exactly in DATABASE SCHEMA. - If a user uses business wording that does not exactly match a column name, map it only to an existing table or column by using the schema comments, aliases, descriptions, and available column names. - When the schema exposes raw or generated names, use those exact names in SQL. Do not replace them with natural-language names. -- ONLY USE "*" if the user query asks for all the columns of a table. +- Do not use "*" unless the user explicitly asks for every column or every field from a table. For ordinary requests such as recent records, customers, transactions, comparisons, summaries, or aggregations, select only the columns needed to answer the question. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! +- Use multiple retrieved tables when the question requires them, and join only with relationships that are present in the DATABASE SCHEMA section. Do not infer joins from similar names, business wording, SQL samples, query history, or reasoning text. +- Do not invent measure, dimension, or filter columns. If the question refers to a business concept, map it only to an exact table or column that appears in DATABASE SCHEMA. - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. @@ -450,14 +452,16 @@ async def _classify_generation_result( 7. Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. 8. Do not mention a table or column in the reasoning plan unless it appears in the DATABASE SCHEMA section. 9. If the user's wording is different from the schema names, map the wording to existing schema names by using comments, aliases, descriptions, and available column names. Do not invent normalized or friendly names. -10. Give a step by step reasoning plan in order to answer user's question. -11. The reasoning plan should be in the language same as the language user provided in the input. -12. Don't include SQL in the reasoning plan. -13. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. -14. Do not include ```markdown or ``` in the answer. -15. A table name in the reasoning plan must be in this format: `table: `. -16. A column name in the reasoning plan must be in this format: `column: .`. -17. ONLY SHOWING the reasoning plan in bullet points. +10. Do not enumerate the full schema. Mention only the retrieved tables, columns, and relationships that are relevant to answering the user's question. +11. When the question requires multiple tables, use only relationships that are present in the DATABASE SCHEMA section; do not infer relationships from similar names, business wording, SQL samples, query history, or reasoning text. +12. Give a step by step reasoning plan in order to answer user's question. +13. The reasoning plan should be in the language same as the language user provided in the input. +14. Don't include SQL in the reasoning plan. +15. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. +16. Do not include ```markdown or ``` in the answer. +17. A table name in the reasoning plan must be in this format: `table: `. +18. A column name in the reasoning plan must be in this format: `column: .`. +19. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format From 10a8c56de66b950d886ec9a90d808326b4b82118 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 02:33:42 +0530 Subject: [PATCH 0647/1087] Restore correction question context --- .../src/pipelines/generation/sql_correction.py | 14 ++++++++++---- .../src/web/v1/routers/sql_corrections.py | 1 + wren-ai-service/src/web/v1/services/ask.py | 1 + .../src/web/v1/services/ask_feedback.py | 1 + .../src/web/v1/services/sql_corrections.py | 3 +++ 5 files changed, 16 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index ed1f322dd3..54c79830a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -35,9 +35,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### SQL CORRECTION INSTRUCTIONS ### 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). -2. Then, generate the syntactically correct ANSI SQL query to correct the error. -3. If the error says a table or column does not exist, replace it only with an exact table or column identifier from DATABASE SCHEMA. Do not create a normalized, friendly, translated, or guessed identifier. -4. The corrected SQL must use only tables and columns that appear in DATABASE SCHEMA. +2. Then, generate the syntactically correct ANSI SQL query that answers the user's question. +3. Use the failed SQL and error message as diagnostic context only. Regenerate the corrected SQL from the DATABASE SCHEMA instead of patching invalid table or column names from the failed SQL or error message. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -76,7 +75,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -SQL: {{ invalid_generation_result.sql }} +{% if query %} +User's Question: {{ query }} +{% endif %} +Failed SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} Let's think step by step. @@ -89,10 +91,12 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, + query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( + query=query, documents=documents, invalid_generation_result=invalid_generation_result, instructions=construct_instructions( @@ -171,6 +175,7 @@ async def run( self, contexts: List[Document], invalid_generation_result: Dict[str, str], + query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, @@ -189,6 +194,7 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, + "query": query, "documents": contexts, "instructions": instructions, "sql_functions": sql_functions, diff --git a/wren-ai-service/src/web/v1/routers/sql_corrections.py b/wren-ai-service/src/web/v1/routers/sql_corrections.py index b74be58bcf..9ccc7be79b 100644 --- a/wren-ai-service/src/web/v1/routers/sql_corrections.py +++ b/wren-ai-service/src/web/v1/routers/sql_corrections.py @@ -19,6 +19,7 @@ class PostRequest(BaseRequest): sql: str error: str + query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = False allow_dry_plan_fallback: bool = True diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6bf806392d..78e05a889e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -585,6 +585,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, + query=user_query, instructions=instructions, invalid_generation_result={ "sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 25044de18c..85ed45ff79 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -239,6 +239,7 @@ async def ask_feedback( "sql_correction" ].run( contexts=table_ddls, + query=ask_feedback_request.question, instructions=instructions, invalid_generation_result={ "original_sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 86d0f55301..3634199927 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -60,6 +60,7 @@ class CorrectionRequest(BaseRequest): event_id: str sql: str error: str + query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = False allow_dry_plan_fallback: bool = True @@ -76,6 +77,7 @@ async def correct( event_id = request.event_id sql = request.sql error = request.error + query = request.query project_id = request.project_id retrieved_tables = request.retrieved_tables use_dry_plan = request.use_dry_plan @@ -114,6 +116,7 @@ async def correct( res = await self._pipelines["sql_correction"].run( contexts=table_ddls, + query=query, invalid_generation_result=_invalid, project_id=project_id, use_dry_plan=use_dry_plan, From a9622f1adfb8a8f4e1efbf60521bc23beb604ff7 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 03:00:30 +0530 Subject: [PATCH 0648/1087] Restore schema-only SQL prompt flow --- .../pipelines/generation/sql_correction.py | 1 + .../src/pipelines/generation/utils/sql.py | 36 ++++++++++--------- wren-ai-service/src/web/v1/services/ask.py | 33 ----------------- 3 files changed, 20 insertions(+), 50 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 54c79830a7..f4b0ad64d6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -37,6 +37,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). 2. Then, generate the syntactically correct ANSI SQL query that answers the user's question. 3. Use the failed SQL and error message as diagnostic context only. Regenerate the corrected SQL from the DATABASE SCHEMA instead of patching invalid table or column names from the failed SQL or error message. +4. Do not copy table or column identifiers from the failed SQL, error message, reasoning text, SQL samples, query history, or user wording unless they appear exactly in DATABASE SCHEMA. ### SQL RULES ### Make sure you follow the SQL Rules strictly. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 01929d4049..1b503f0c88 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -171,17 +171,18 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements. Do not generate ALTER, CREATE, DROP, INSERT, UPDATE, DELETE, MERGE, TRUNCATE, GRANT, REVOKE, or any other statement that can change the database or schema. - ONLY USE the tables and columns mentioned in the database schema. -- Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. Never invent a normalized, friendly, translated, or guessed identifier. +- Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. Never invent, rename, translate, normalize, prefix, suffix, or guess an identifier. - Do not use names from policies, instructions, SQL functions, SQL samples, query history, reasoning text, error messages, or the user's wording as table or column identifiers unless they appear exactly in DATABASE SCHEMA. - Policy documents and policy names are not database tables or columns unless they appear exactly in DATABASE SCHEMA. -- If a user uses business wording that does not exactly match a column name, map it only to an existing table or column by using the schema comments, aliases, descriptions, and available column names. -- When the schema exposes raw or generated names, use those exact names in SQL. Do not replace them with natural-language names. +- Use schema comments, aliases, and descriptions only to choose among identifiers that already appear exactly in DATABASE SCHEMA. Do not derive new table or column names from them. +- When the schema exposes raw or generated names, use those exact names in SQL. Do not replace them with natural-language names or inferred business names. +- If the user's requested concept is not represented by any exact table or column in DATABASE SCHEMA, do not assume a fallback table or column. - Do not use "*" unless the user explicitly asks for every column or every field from a table. For ordinary requests such as recent records, customers, transactions, comparisons, summaries, or aggregations, select only the columns needed to answer the question. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! - Use multiple retrieved tables when the question requires them, and join only with relationships that are present in the DATABASE SCHEMA section. Do not infer joins from similar names, business wording, SQL samples, query history, or reasoning text. -- Do not invent measure, dimension, or filter columns. If the question refers to a business concept, map it only to an exact table or column that appears in DATABASE SCHEMA. +- Do not invent measure, dimension, or filter columns. If the question refers to a business concept, use only an exact table or column that appears in DATABASE SCHEMA. - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. @@ -447,21 +448,22 @@ async def _classify_generation_result( otherwise, you will put the relative timeframe in the SQL query. 3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. -5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. -6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. +5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan, but do not treat them as schema. +6. If SQL SAMPLES section is provided, use them only as examples of SQL structure. Do not treat their identifiers as schema unless they appear exactly in DATABASE SCHEMA. 7. Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. 8. Do not mention a table or column in the reasoning plan unless it appears in the DATABASE SCHEMA section. -9. If the user's wording is different from the schema names, map the wording to existing schema names by using comments, aliases, descriptions, and available column names. Do not invent normalized or friendly names. +9. If the user's wording is different from the schema names, use comments, aliases, and descriptions only to choose among existing schema names. Do not invent normalized, friendly, or inferred names. 10. Do not enumerate the full schema. Mention only the retrieved tables, columns, and relationships that are relevant to answering the user's question. 11. When the question requires multiple tables, use only relationships that are present in the DATABASE SCHEMA section; do not infer relationships from similar names, business wording, SQL samples, query history, or reasoning text. -12. Give a step by step reasoning plan in order to answer user's question. -13. The reasoning plan should be in the language same as the language user provided in the input. -14. Don't include SQL in the reasoning plan. -15. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. -16. Do not include ```markdown or ``` in the answer. -17. A table name in the reasoning plan must be in this format: `table: `. -18. A column name in the reasoning plan must be in this format: `column: .`. -19. ONLY SHOWING the reasoning plan in bullet points. +12. If the required table or column is not present in DATABASE SCHEMA, state that the exact identifier is unavailable; do not assume a replacement. +13. Give a step by step reasoning plan in order to answer user's question. +14. The reasoning plan should be in the language same as the language user provided in the input. +15. Don't include SQL in the reasoning plan. +16. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. +17. Do not include ```markdown or ``` in the answer. +18. A table name in the reasoning plan must be in this format: `table: `. +19. A column name in the reasoning plan must be in this format: `column: .`. +20. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -530,9 +532,9 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### GENERAL RULES ### -1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. +1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input, unless they conflict with DATABASE SCHEMA. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. +3. YOU MUST REFER to sql samples only for SQL structure and patterns if the section of SQL SAMPLES is available in user's input. Do not copy identifiers from samples unless they appear exactly in DATABASE SCHEMA. 4. YOU MUST USE the reasoning plan only as analytical guidance if the section of REASONING PLAN is available in user's input. Do not copy table or column names from it unless they appear exactly in DATABASE SCHEMA. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. 6. DATABASE SCHEMA is more authoritative than the reasoning plan, SQL samples, and user wording. If any of those mention a table or column that is not present in DATABASE SCHEMA, do not use it. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 78e05a889e..7ce6ea2e15 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,6 +1,5 @@ import asyncio import logging -import re from typing import Dict, List, Literal, Optional from cachetools import TTLCache @@ -14,18 +13,6 @@ logger = logging.getLogger("wren-ai-service") -DATA_QUERY_PATTERN = re.compile( - r"(\bhow\s+many\b|\bhow\s+much\b|\bwhat\s+(?:is|are|was|were)\b|\bwhich\b|\b(" - r"show|list|find|get|give|display|retrieve|fetch|compare|count|sum|total|" - r"average|avg|min|max|rank|top|bottom|highest|lowest|latest|earliest|" - r"newest|oldest|sort|order|group|filter|where|between|starts|ends|" - r"contains|duplicate|unique|distinct|percentage|percent|trend|growth|" - r"breakdown|placed|created|became|active|inactive|expired|valid|by|per" - r")\b)", - re.IGNORECASE, -) - - class AskHistory(BaseModel): sql: str question: str @@ -138,12 +125,6 @@ def __init__( self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries - def _should_continue_text_to_sql(self, query: str, intent: Optional[str]) -> bool: - if intent not in {"GENERAL", "MISLEADING_QUERY", "USER_GUIDE"}: - return False - - return bool(DATA_QUERY_PATTERN.search(query or "")) - def _is_stopped(self, query_id: str, container: dict): if ( result := container.get(query_id) @@ -275,20 +256,6 @@ async def ask( if rephrased_question: user_query = rephrased_question - if self._should_continue_text_to_sql( - user_query, - intent, - ) or ( - intent == "GENERAL" - and not intent_classification_result.get("db_schemas") - ): - logger.info( - "Intent classification returned %s for an analytical query; continuing Text-to-SQL retrieval for query_id %s", - intent, - query_id, - ) - intent = "TEXT_TO_SQL" - if intent == "MISLEADING_QUERY": asyncio.create_task( self._pipelines["misleading_assistance"].run( From ad99e4d78f0e78a822dc01e6be44bc1d80981f83 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 03:09:39 +0530 Subject: [PATCH 0649/1087] Route data requests to SQL flow --- .../src/pipelines/generation/intent_classification.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index e858b46aaa..da52f2eb34 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -34,9 +34,10 @@ - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. - **Vague Queries:** If the question is not a data retrieval or analysis request and does not relate to the schema, classify it as `MISLEADING_QUERY`. -- **Natural Language Data Queries:** Do not require exact physical table or column names. If the user asks to show, list, find, compare, count, aggregate, rank, filter, sort, or analyze data, classify it as `TEXT_TO_SQL` when schema context may answer it. +- **Natural Language Data Queries:** Do not require exact physical table or column names. If the user asks to retrieve or analyze data and the schema context may answer it, classify it as `TEXT_TO_SQL`. +- **Intent Precedence:** Prefer `TEXT_TO_SQL` over `GENERAL` for data retrieval or analysis requests. Lack of exact table or column names is not a reason to choose `GENERAL`. - **User Guide Boundary:** Do not classify a data retrieval, filtering, date, metric, aggregation, or row-listing question as `USER_GUIDE` just because it contains words like "how", "show", or "filter". -- **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. +- **Incomplete Queries:** If the question is related to the database schema but references unresolved placeholders (e.g., "the following", "these", "those", "the previous ones") without providing them or prior context, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. ### Intent Definitions ### @@ -49,10 +50,12 @@ - The question can be answered by selecting relevant tables and columns from the provided schema, even if the user does not mention exact physical table or column names. - The question includes enough business meaning, dimensions, metrics, filters, or time criteria to attempt SQL generation from the schema. - Natural-language analytical questions should be classified as `TEXT_TO_SQL` when they can reasonably be grounded in the schema. +- Broad row-listing requests are `TEXT_TO_SQL` if related schema is present. **Requirements:** - Do not require the user to explicitly name a table or column. - Use the schema to determine whether the user's business terms can map to available tables or columns. +- Do not generate SQL in this step. Only classify intent and rephrase the question. - Reference phrases from the user's inputs that clearly indicate a data retrieval or analysis request. **Examples:** @@ -68,6 +71,7 @@ - The query contains **placeholder references** that cannot be resolved from context. - The query is **incomplete for SQL generation** because required values or references are missing, not merely because exact table or column names are absent. - The user is asking for explanation, guidance, or clarification rather than asking to retrieve or analyze rows from the data. +- Do not classify row-listing, filtering, ordering, metric, date, or aggregation requests as `GENERAL` merely because the user used business wording. **Requirements:** - Incorporate phrases from the user's inputs that indicate incompleteness or lack of relevance to the database schema. From cbd0cc23166730a601973ce9413f43dd0c60eda7 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 15:14:41 +0530 Subject: [PATCH 0650/1087] Restore legacy post SQL flows --- .../pipelines/generation/chart_generation.py | 22 +- .../generation/question_recommendation.py | 21 +- .../src/pipelines/generation/sql_answer.py | 21 - .../pipelines/generation/sql_correction.py | 12 +- .../src/pipelines/generation/utils/chart.py | 742 +----------------- .../src/pipelines/generation/utils/sql.py | 21 +- .../src/pipelines/indexing/db_schema.py | 157 +--- .../src/web/v1/routers/sql_corrections.py | 1 - wren-ai-service/src/web/v1/services/ask.py | 1 - .../src/web/v1/services/ask_feedback.py | 1 - wren-ai-service/src/web/v1/services/chart.py | 86 +- .../v1/services/question_recommendation.py | 253 +----- .../src/web/v1/services/sql_answer.py | 56 +- .../src/web/v1/services/sql_corrections.py | 3 - 14 files changed, 86 insertions(+), 1311 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/chart_generation.py b/wren-ai-service/src/pipelines/generation/chart_generation.py index 1b0c159655..6daca5ec17 100644 --- a/wren-ai-service/src/pipelines/generation/chart_generation.py +++ b/wren-ai-service/src/pipelines/generation/chart_generation.py @@ -42,16 +42,6 @@ """ chart_generation_user_prompt_template = """ -{% if documents %} -### ACTIVE DATASOURCE METADATA ### -This is the complete deployed metadata for the active datasource. Use it to -understand the schema, tables, columns, metrics, views, and relationships behind -the SQL before choosing a chart. -{% for document in documents %} - {{ document }} -{% endfor %} -{% endif %} - ### INPUT ### Question: {{ query }} SQL: {{ sql }} @@ -80,7 +70,6 @@ def prompt( language: str, custom_instruction: str, prompt_builder: PromptBuilder, - documents: list[str] | None = None, ) -> dict: sample_data = preprocess_data.get("sample_data") sample_column_values = preprocess_data.get("sample_column_values") @@ -92,7 +81,6 @@ def prompt( sample_column_values=sample_column_values, language=language, custom_instruction=custom_instruction, - documents=documents or [], ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -109,14 +97,12 @@ def post_process( vega_schema: Dict[str, Any], remove_data_from_chart_schema: bool, preprocess_data: dict, - query: str, post_processor: ChartGenerationPostProcessor, ) -> dict: return post_processor.run( generate_chart.get("replies"), vega_schema, preprocess_data["sample_data"], - query, remove_data_from_chart_schema, ) @@ -152,11 +138,7 @@ def __init__( "post_processor": ChartGenerationPostProcessor(), } - with open( - "src/pipelines/generation/utils/vega-lite-schema-v5.json", - "r", - encoding="utf-8", - ) as f: + with open("src/pipelines/generation/utils/vega-lite-schema-v5.json", "r") as f: _vega_schema = orjson.loads(f.read()) self._configs = { @@ -176,7 +158,6 @@ async def run( language: str, remove_data_from_chart_schema: bool = True, custom_instruction: Optional[str] = None, - contexts: Optional[list[str]] = None, ) -> dict: logger.info("Chart Generation pipeline is running...") return await self._pipe.execute( @@ -188,7 +169,6 @@ async def run( "language": language, "remove_data_from_chart_schema": remove_data_from_chart_schema, "custom_instruction": custom_instruction or "", - "documents": contexts or [], **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index 5a5b3e3e2d..a6e7c17b02 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -16,13 +16,6 @@ logger = logging.getLogger("wren-ai-service") -DEFAULT_QUESTION_CATEGORIES = [ - "Descriptive Questions", - "Segmentation Questions", - "Comparative Questions", - "Data Quality/Accuracy Questions", -] - system_prompt = """ You are an expert in data analysis and SQL query generation. Given a data model specification, optionally a user's question, and a list of categories, your task is to generate insightful, specific questions that can be answered using the provided data model. Each question should be accompanied by a brief explanation of its relevance or importance. @@ -73,9 +66,6 @@ 5. **General Guidelines for All Questions:** - Ensure questions can be answered using the data model. - - Use only the tables, fields, and relationships that are explicitly present in the provided database schema. - - Do not invent tables, columns, business entities, or time dimensions that are not present in the schema. - - Keep the question grounded in the deployed database domain shown by the schema context. - Mix simple and complex questions. - Avoid open-ended questions - each should have a definite answer. - Incorporate time-based analysis where relevant. @@ -160,10 +150,7 @@ {% endif %} {% if documents %} -### ACTIVE DATASOURCE METADATA ### -Use only this latest deployed metadata from the active datasource when generating -recommended questions. Do not reuse tables, columns, or business terms from prior -questions unless they are answerable from this metadata. +### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} @@ -183,6 +170,12 @@ def prompt( max_categories: int, prompt_builder: PromptBuilder, ) -> dict: + """ + If previous_questions is provided, the MDL is omitted to allow the LLM to focus on + generating recommendations based on the question history. This helps provide more + contextually relevant questions that build on previous questions. + """ + _prompt = prompt_builder.run( documents=documents, previous_questions=previous_questions, diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 4f180c9e25..81289081b5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -32,12 +32,6 @@ 6. Answer must be in the same language user specified. 7. Do not include ```markdown or ``` in the answer. 8. If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. -9. Always produce a narrative answer. Never return an empty response. -10. If the user asks for a chart or trend, still summarize the result in words and mention the chart-ready fields. -11. If the data contains only raw rows or a single column, summarize what those rows show, mention the visible date/category range when possible, and state that the result table contains the detailed rows. -12. If Data rows are present, answer from those rows only. Never say you do not have access to the database, system, records, or source data after rows are provided. -13. Do not give generic instructions about how the user can find the data when SQL results are present. Summarize the returned rows instead. -14. If Data rows are empty, say the SQL ran but returned no rows and briefly mention the selected columns, filters, or grouping visible in the SQL. ### OUTPUT FORMAT @@ -45,16 +39,6 @@ """ sql_to_answer_user_prompt_template = """ -{% if documents %} -### Active Datasource Metadata ### -This is the complete deployed metadata for the active datasource. Use it to -understand the schema, tables, columns, metrics, views, and relationships behind -the SQL before answering. -{% for document in documents %} - {{ document }} -{% endfor %} -{% endif %} - ### Inputs ### User's question: {{ query }} SQL: {{ sql }} @@ -67,7 +51,6 @@ Custom Instruction: {{ custom_instruction }} Please think step by step and answer the user's question. -If rows are present in Data, summarize those rows directly and do not claim that the data is unavailable. """ @@ -81,7 +64,6 @@ def prompt( current_time: str, custom_instruction: str, prompt_builder: PromptBuilder, - documents: list[str] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, @@ -90,7 +72,6 @@ def prompt( language=language, current_time=current_time, custom_instruction=custom_instruction, - documents=documents or [], ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -175,7 +156,6 @@ async def run( current_time: str = Configuration().show_current_time(), query_id: Optional[str] = None, custom_instruction: Optional[str] = None, - contexts: Optional[list[str]] = None, ) -> dict: logger.info("Sql_Answer Generation pipeline is running...") return await self._pipe.execute( @@ -188,7 +168,6 @@ async def run( "current_time": current_time, "query_id": query_id, "custom_instruction": custom_instruction or "", - "documents": contexts or [], **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f4b0ad64d6..601e75d2ae 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -35,9 +35,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### SQL CORRECTION INSTRUCTIONS ### 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). -2. Then, generate the syntactically correct ANSI SQL query that answers the user's question. +2. Then, generate the syntactically correct ANSI SQL query to correct the error. 3. Use the failed SQL and error message as diagnostic context only. Regenerate the corrected SQL from the DATABASE SCHEMA instead of patching invalid table or column names from the failed SQL or error message. -4. Do not copy table or column identifiers from the failed SQL, error message, reasoning text, SQL samples, query history, or user wording unless they appear exactly in DATABASE SCHEMA. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -76,10 +75,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -{% if query %} -User's Question: {{ query }} -{% endif %} -Failed SQL: {{ invalid_generation_result.sql }} +SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} Let's think step by step. @@ -92,12 +88,10 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, - query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( - query=query, documents=documents, invalid_generation_result=invalid_generation_result, instructions=construct_instructions( @@ -176,7 +170,6 @@ async def run( self, contexts: List[Document], invalid_generation_result: Dict[str, str], - query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, @@ -195,7 +188,6 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, - "query": query, "documents": contexts, "instructions": instructions, "sql_functions": sql_functions, diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 673e8be576..5d06b949a9 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -1,6 +1,4 @@ import logging -import re -from copy import deepcopy from typing import Any, Dict, Literal, Optional import orjson @@ -13,665 +11,6 @@ logger = logging.getLogger("wren-ai-service") -def _humanize_title(name: str | None) -> str: - text = str(name or "") - text = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", " ", text) - text = re.sub(r"[_\s]+", " ", text) - return text.strip().title() - - -def _detect_requested_chart_type(query: str | None) -> str: - normalized = (query or "").lower() - checks = [ - ("grouped_bar", ["grouped bar"]), - ("stacked_bar", ["stacked bar"]), - ("multi_line", ["multi line", "multi-line"]), - ("line", ["line chart", "line graph", "line plot"]), - ("bar", ["bar chart", "bar graph", "column chart"]), - ("bar", ["waterfall", "waterfall chart"]), - ("pie", ["pie chart", "donut chart", "doughnut chart"]), - ("area", ["area chart", "area graph"]), - ] - for chart_type, patterns in checks: - if any(pattern in normalized for pattern in patterns): - return chart_type - if re.search(r"\b(chart|graph|plot|visuali[sz](?:e|ation)?)\b", normalized): - return "bar" - return "" - - -def _safe_column_names(columns: list[Any]) -> list[str]: - return [str(column) for column in columns if column is not None and str(column)] - - -def _normalize_identifier(value: str | None) -> str: - return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) - - -def _identifier_tokens(value: str | None) -> set[str]: - text = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", " ", str(value or "")) - return { - token - for token in re.split(r"[^a-zA-Z0-9]+", text.lower()) - if len(token) > 1 - } - - -def _query_relevant_columns(query: str | None, columns: list[str]) -> list[str]: - normalized_query = str(query or "").lower() - compact_query = _normalize_identifier(normalized_query) - query_tokens = _identifier_tokens(normalized_query) - - scored_columns: list[tuple[int, int, str]] = [] - for index, column in enumerate(columns): - tokens = _identifier_tokens(column) - compact_column = _normalize_identifier(column) - score = 0 - - if compact_column and compact_column in compact_query: - score += 100 - - matched_tokens = tokens.intersection(query_tokens) - score += len(matched_tokens) * 20 - - # Prefer multi-word business dimensions that the user explicitly asks - # for, e.g. "Business Unit" matching BusinessUnit. - if tokens and tokens.issubset(query_tokens): - score += 40 - - if score: - scored_columns.append((-score, index, column)) - - return [column for _, _, column in sorted(scored_columns)] - - -def _query_grouping_columns(query: str | None, columns: list[str]) -> set[str]: - normalized_query = str(query or "").lower() - grouping_terms = set() - for match in re.finditer( - r"\b(?:by|per|each|grouped by|group by)\s+([a-zA-Z0-9_ ]+)", - normalized_query, - ): - phrase = match.group(1) - phrase = re.split( - r"\b(?:and|with|over|for|where|order|sort|top|last|using)\b", - phrase, - maxsplit=1, - )[0] - grouping_terms.update(_identifier_tokens(phrase)) - - grouped_columns = set() - for column in columns: - tokens = _identifier_tokens(column) - if tokens and tokens.intersection(grouping_terms): - grouped_columns.add(column) - - return grouped_columns - - -def _select_measure_column(query: str | None, quantitative: list[str]) -> str | None: - if not quantitative: - return None - - grouping_columns = _query_grouping_columns(query, quantitative) - measure_candidates = [ - column for column in quantitative if column not in grouping_columns - ] or quantitative - - relevant = _query_relevant_columns(query, measure_candidates) - if relevant: - return relevant[0] - - metric_keywords = ( - "count", - "total", - "sum", - "amount", - "value", - "volume", - "order", - "sales", - "revenue", - "quantity", - "workload", - "throughput", - "failure", - "repair", - ) - for column in measure_candidates: - normalized = str(column).lower() - if any(keyword in normalized for keyword in metric_keywords): - return column - - return measure_candidates[0] - - -def _select_axis_columns( - query: str | None, - chart_type: str, - quantitative: list[str], - temporal: list[str], - nominal: list[str], - columns: list[str], -) -> tuple[str | None, str | None, list[str]]: - dimensions = _select_dimension_columns( - query, chart_type, nominal, temporal, columns - ) - measure = _select_measure_column(query, quantitative) - - if dimensions: - return dimensions[0], measure, dimensions - - # Numeric-only SQL results are still chartable, but x/y must never point to - # the same column. Use the first non-measure numeric column as the category - # axis and the selected measure as y. - if measure and len(quantitative) > 1: - x_field = next( - (column for column in quantitative if column != measure), - quantitative[0], - ) - return x_field, measure, [x_field] - - return (columns[0] if columns else None), measure, dimensions - - -def _count_axis_title(query: str | None) -> str: - normalized = str(query or "").lower() - if "new order" in normalized: - return "New Orders Count" - if "order" in normalized: - return "Order Count" - if "repair" in normalized: - return "Repair Count" - if "failure" in normalized: - return "Failure Count" - if "ticket" in normalized: - return "Ticket Count" - return "Count" - - -def _wants_time_axis(query: str | None, chart_type: str) -> bool: - normalized = str(query or "").lower() - time_pattern = ( - r"\b(trend|over time|timeline|monthly|weekly|daily|yearly|" - r"by month|by week|by day|by year)\b" - ) - return chart_type in {"line", "area", "multi_line"} or bool( - re.search(time_pattern, normalized) - ) - - -def _select_dimension_columns( - query: str | None, - chart_type: str, - nominal: list[str], - temporal: list[str], - columns: list[str], -) -> list[str]: - relevant = _query_relevant_columns(query, columns) - relevant_nominal = [column for column in relevant if column in nominal] - relevant_temporal = [column for column in relevant if column in temporal] - - if _wants_time_axis(query, chart_type): - ordered = relevant_temporal + [ - column for column in temporal if column not in relevant_temporal - ] - ordered += relevant_nominal + [ - column for column in nominal if column not in relevant_nominal - ] - return ordered - - ordered = relevant_nominal + [ - column for column in nominal if column not in relevant_nominal - ] - ordered += relevant_temporal + [ - column for column in temporal if column not in relevant_temporal - ] - if len(ordered) > 1: - normalized_query = str(query or "").lower() - for column in list(ordered): - tokens = _identifier_tokens(column) - if any( - re.search(rf"\b(?:each|per)\s+{re.escape(token)}s?\b", normalized_query) - for token in tokens - ): - ordered.remove(column) - ordered.insert(1, column) - break - return ordered - - -def _refine_chart_type_for_columns( - query: str | None, - chart_type: str, - sample_data: list[dict], -) -> str: - if chart_type != "bar" or not sample_data: - return chart_type - - columns = _safe_column_names(list(sample_data[0].keys())) - inferred = _infer_column_types(sample_data) - dimensions = _select_dimension_columns( - query, chart_type, inferred["nominal"], inferred["temporal"], columns - ) - if len([column for column in dimensions if column in inferred["nominal"]]) > 1: - return "grouped_bar" - - return chart_type - - -def _match_column_name(field: str | None, columns: list[str]) -> str: - if field is None: - return "" - field = str(field) - columns = _safe_column_names(columns) - if field in columns: - return field - - lowered = {column.lower(): column for column in columns} - normalized = lowered.get(field.lower()) - if normalized: - return normalized - - compact = re.sub(r"[\s_]+", "", field.lower()) - for column in columns: - if re.sub(r"[\s_]+", "", column.lower()) == compact: - return column - - return field - - -def _normalize_chart_schema_fields(chart_schema: dict, columns: list[str]) -> dict: - normalized = deepcopy(chart_schema) - encoding = normalized.get("encoding", {}) - - for key in ("x", "y", "x2", "y2", "color", "xOffset", "theta"): - axis = encoding.get(key) - if isinstance(axis, dict) and axis.get("field"): - axis["field"] = _match_column_name(axis["field"], columns) - - for transform in normalized.get("transform", []) or []: - if isinstance(transform, dict) and isinstance(transform.get("fold"), list): - transform["fold"] = [ - _match_column_name(field, columns) - for field in transform["fold"] - if field is not None - ] - - return normalized - - -def _infer_column_types(sample_data: list[dict]) -> dict[str, list[str]]: - if not sample_data: - return {"quantitative": [], "temporal": [], "nominal": []} - - df = pd.DataFrame(sample_data) - quantitative: list[str] = [] - temporal: list[str] = [] - nominal: list[str] = [] - - for column in df.columns: - if column is None or not str(column): - continue - - values = df[column].dropna() - if values.empty: - continue - - column_name = str(column).lower() - numeric_values = pd.to_numeric( - values.astype(str).str.replace(",", "", regex=False), - errors="coerce", - ) - is_temporal_name = bool( - re.search(r"(date|time|month|year|day|created|updated)", column_name) - ) - string_values = values.astype(str) - looks_temporal = string_values.str.match( - r"^\d{4}[-/]\d{1,2}([-/]\d{1,2})?" - ).all() - temporal_values = ( - pd.to_datetime(values, errors="coerce") - if is_temporal_name or looks_temporal - else None - ) - - if numeric_values.notna().all() and not is_temporal_name: - quantitative.append(str(column)) - elif is_temporal_name or ( - temporal_values is not None and temporal_values.notna().all() - ): - temporal.append(str(column)) - else: - nominal.append(str(column)) - - return { - "quantitative": quantitative, - "temporal": temporal, - "nominal": nominal, - } - - -def _chart_mark_type(chart_schema: dict) -> str: - mark = chart_schema.get("mark", {}) - if isinstance(mark, str): - return mark - if isinstance(mark, dict): - return str(mark.get("type") or "") - return "" - - -def _chart_type_from_schema(chart_schema: dict, default: str = "") -> str: - mark_type = _chart_mark_type(chart_schema) - encoding = chart_schema.get("encoding", {}) - - if mark_type == "arc": - return "pie" - if mark_type == "area": - return "area" - if mark_type == "line": - return "multi_line" if chart_schema.get("transform") else "line" - if mark_type == "bar": - if encoding.get("xOffset"): - return "grouped_bar" - if isinstance(encoding.get("y"), dict) and encoding["y"].get("stack"): - return "stacked_bar" - return "bar" - - return default - - -def _is_quantitative_encoding(axis: Any) -> bool: - return isinstance(axis, dict) and axis.get("type") == "quantitative" - - -def _is_categorical_encoding(axis: Any) -> bool: - return isinstance(axis, dict) and axis.get("type") in { - "nominal", - "ordinal", - "temporal", - } - - -def _fallback_chart_type( - requested_chart_type: str, - quantitative: list[str], - temporal: list[str], - nominal: list[str], -) -> str: - chart_type = requested_chart_type or "bar" - - if chart_type == "pie": - return "pie" if nominal else "" - - if chart_type in {"line", "area", "multi_line"}: - return chart_type if temporal or nominal or quantitative else "" - - if chart_type in {"grouped_bar", "stacked_bar"}: - return chart_type if len(nominal) > 1 else "" - - return "bar" if nominal or temporal or quantitative else "" - - -def _build_fallback_chart_schema( - query: str | None, - chart_type: str, - sample_data: list[dict], -) -> dict: - if not sample_data: - return {} - - columns = _safe_column_names(list(sample_data[0].keys())) - if not columns: - return {} - inferred = _infer_column_types(sample_data) - quantitative = inferred["quantitative"] - temporal = inferred["temporal"] - nominal = inferred["nominal"] - chart_type = _fallback_chart_type(chart_type, quantitative, temporal, nominal) - if not chart_type: - return {} - - x_field, measure, dimensions = _select_axis_columns( - query, chart_type, quantitative, temporal, nominal, columns - ) - - title = _humanize_title(query or "Chart") - - def axis(field: str, field_type: str) -> dict: - base = {"field": field, "type": field_type, "title": _humanize_title(field)} - if field_type == "temporal": - base["timeUnit"] = "yearmonth" - return base - - count_axis = { - "aggregate": "count", - "type": "quantitative", - "title": _count_axis_title(query), - } - - if chart_type == "pie": - color_field = x_field or columns[0] - theta_axis = axis(measure, "quantitative") if measure else count_axis - return { - "title": title, - "mark": {"type": "arc"}, - "encoding": { - "theta": theta_axis, - "color": axis(color_field, "nominal"), - }, - } - - if chart_type in {"line", "area", "multi_line"}: - if not measure: - y_encoding = count_axis - else: - y_encoding = axis(measure, "quantitative") - if {"year", "month"}.issubset({str(c).lower() for c in columns}): - month_field = next(c for c in columns if str(c).lower() == "month") - encoding = { - "x": axis(month_field, "ordinal"), - "y": y_encoding, - } - years = [c for c in columns if str(c).lower() == "year"] - if years: - encoding["color"] = axis(years[0], "nominal") - return { - "title": title, - "mark": {"type": "area" if chart_type == "area" else "line"}, - "encoding": encoding, - } - - x_field = x_field or columns[0] - x_type = "temporal" if x_field in temporal else "ordinal" - encoding = { - "x": axis(x_field, x_type), - "y": y_encoding, - } - series_field = next( - (column for column in dimensions[1:] if column in nominal), - None, - ) - if series_field: - encoding["color"] = axis(series_field, "nominal") - return { - "title": title, - "mark": {"type": "area" if chart_type == "area" else "line"}, - "encoding": encoding, - } - - x_field = x_field or columns[0] - x_type = ( - "nominal" - if x_field in nominal - else ("temporal" if x_field in temporal else "ordinal") - ) - y_encoding = axis(measure, "quantitative") if measure else count_axis - encoding = { - "x": axis(x_field, x_type), - "y": y_encoding, - } - comparison_field = next( - (column for column in dimensions[1:] if column in nominal), - None, - ) - if comparison_field: - encoding["color"] = axis(comparison_field, "nominal") - nominal_dimension_count = len([c for c in dimensions if c in nominal]) - if chart_type == "grouped_bar" or nominal_dimension_count > 1: - encoding["xOffset"] = axis(comparison_field, "nominal") - elif x_field in nominal: - encoding["color"] = axis(x_field, "nominal") - - mark = {"type": "bar"} - if chart_type == "stacked_bar": - encoding["y"]["stack"] = "zero" - - return { - "title": title, - "mark": mark, - "encoding": encoding, - } - - -def build_fallback_chart_result( - query: str | None, - data: Dict[str, Any], - remove_data_from_chart_schema: bool = True, -) -> dict: - processed = ChartDataPreprocessor().run(data) - sample_data = processed.get("sample_data", []) - chart_type = _detect_requested_chart_type(query) or "bar" - chart_type = _refine_chart_type_for_columns(query, chart_type, sample_data) - chart_schema = _build_fallback_chart_schema(query, chart_type, sample_data) - if not chart_schema: - return { - "chart_schema": {}, - "reasoning": "", - "chart_type": "", - } - chart_type = _chart_type_from_schema(chart_schema, chart_type) - - chart_schema["$schema"] = "https://vega.github.io/schema/vega-lite/v5.json" - chart_schema["data"] = {"values": sample_data} - if remove_data_from_chart_schema: - chart_schema["data"]["values"] = [] - - return { - "chart_schema": chart_schema, - "reasoning": "Generated from the SQL result columns and requested chart type.", - "chart_type": chart_type, - } - - -def _is_schema_compatible_with_sample_data( - chart_schema: dict, - sample_data: list[dict], -) -> bool: - if not chart_schema or not sample_data: - return False - - columns = set(_safe_column_names(list(sample_data[0].keys()))) - inferred = _infer_column_types(sample_data) - quantitative = set(inferred["quantitative"]) - temporal = set(inferred["temporal"]) - nominal = set(inferred["nominal"]) - categorical = nominal | temporal - encoding = chart_schema.get("encoding", {}) - for key in ("x", "y", "x2", "y2", "color", "xOffset", "theta"): - axis = encoding.get(key) - if isinstance(axis, dict) and axis.get("aggregate"): - return False - field = axis.get("field") if isinstance(axis, dict) else None - if field and str(field) not in columns: - return False - - if ( - field - and _is_quantitative_encoding(axis) - and str(field) not in quantitative - ): - return False - - if field and key in {"color", "xOffset"} and str(field) not in categorical: - return False - - for transform in chart_schema.get("transform", []) or []: - if isinstance(transform, dict): - for field in transform.get("fold", []) or []: - if field is not None and str(field) not in columns: - return False - - mark_type = _chart_mark_type(chart_schema) - x_axis = encoding.get("x") - y_axis = encoding.get("y") - theta_axis = encoding.get("theta") - has_quantitative_measure = any( - _is_quantitative_encoding(axis) - for axis in (x_axis, y_axis, theta_axis) - ) - - if mark_type == "arc": - return _is_quantitative_encoding(theta_axis) and _is_categorical_encoding( - encoding.get("color") - ) - - if mark_type in {"bar", "line", "area"}: - return has_quantitative_measure and ( - _is_categorical_encoding(x_axis) or _is_categorical_encoding(y_axis) - ) - - return True - - -def _needs_deterministic_bar_fallback( - chart_schema: dict, - chart_type: str, - sample_data: list[dict], -) -> bool: - if chart_type not in {"bar", "grouped_bar", "stacked_bar"}: - return False - if not sample_data: - return False - - encoding = chart_schema.get("encoding", {}) if chart_schema else {} - - # Reject range-style bar encodings for simple grouped-count datasets. - for key in ("x2", "y2"): - axis = encoding.get(key) - if isinstance(axis, dict) and axis.get("field"): - return True - - for axis_name in ("x", "y", "color", "xOffset", "theta"): - axis = encoding.get(axis_name) - if not isinstance(axis, dict): - continue - field = axis.get("field", "") - if isinstance(field, str) and ( - field.endswith("_start") or field.endswith("_end") - ): - return True - - inferred = _infer_column_types(sample_data) - quantitative = inferred["quantitative"] - nominal = inferred["nominal"] - - # For the common case "category + count", prefer a deterministic bar spec - # if the model did not produce a usable quantitative y axis. - if len(quantitative) == 1 and len(nominal) >= 1: - y_axis = encoding.get("y") - x_axis = encoding.get("x") - if not isinstance(y_axis, dict) or y_axis.get("field") not in quantitative: - return True - if not isinstance(x_axis, dict) or x_axis.get("field") not in nominal: - return True - - if len(quantitative) == 0 and len(nominal) >= 1: - return True - - return False - - chart_generation_instructions = """ ### INSTRUCTIONS ### @@ -699,7 +38,6 @@ def _needs_deterministic_bar_fallback( - Default time unit is "yearmonth". - For each axis, generate the corresponding human-readable title based on the language provided by the user. - Make sure all of the fields(x, y, xOffset, color, etc.) in the encoding section of the chart schema are present in the column names of the data. -- Do not use Vega-Lite aggregate count or calculate new measures in the chart schema. The SQL must return the metric column, and the chart must encode that returned metric field. ### GUIDELINES TO PLOT CHART ### @@ -918,13 +256,10 @@ def run( sample_data_count: int = 15, sample_column_size: int = 5, ): - columns = [] - for index, column in enumerate(data.get("columns", [])): - if isinstance(column, dict): - column_name = str(column.get("name") or "").strip() - else: - column_name = str(column or "").strip() - columns.append(column_name or f"column_{index + 1}") + columns = [ + column.get("name", "") if isinstance(column, dict) else column + for column in data.get("columns", []) + ] data = data.get("data", []) df = pd.DataFrame(data, columns=columns) @@ -932,7 +267,10 @@ def run( col: list(df[col].unique())[:sample_column_size] for col in df.columns } - sample_data = df.head(sample_data_count).to_dict(orient="records") + if len(df) > sample_data_count: + sample_data = df.sample(n=sample_data_count).to_dict(orient="records") + else: + sample_data = df.to_dict(orient="records") return { "sample_data": sample_data, @@ -950,43 +288,17 @@ def run( replies: str, vega_schema: Dict[str, Any], sample_data: list[dict], - query: Optional[str] = None, remove_data_from_chart_schema: Optional[bool] = True, ): try: generation_result = orjson.loads(replies[0]) reasoning = generation_result.get("reasoning", "") - requested_chart_type = _detect_requested_chart_type(query) - chart_type = requested_chart_type or generation_result.get("chart_type", "") + chart_type = generation_result.get("chart_type", "") if chart_schema := generation_result.get("chart_schema", {}): # sometimes the chart_schema is still in string format if isinstance(chart_schema, str): chart_schema = orjson.loads(chart_schema) - chart_schema = _normalize_chart_schema_fields( - chart_schema, list(sample_data[0].keys()) if sample_data else [] - ) - - if ( - not _is_schema_compatible_with_sample_data(chart_schema, sample_data) - or _needs_deterministic_bar_fallback( - chart_schema, chart_type or "", sample_data - ) - ): - chart_schema = _build_fallback_chart_schema( - query, chart_type or "bar", sample_data - ) - chart_type = _chart_type_from_schema(chart_schema, chart_type) - - if not chart_schema: - return { - "results": { - "chart_schema": {}, - "reasoning": reasoning, - "chart_type": "", - } - } - chart_schema[ "$schema" ] = "https://vega.github.io/schema/vega-lite/v5.json" @@ -1005,55 +317,31 @@ def run( } } - fallback_schema = _build_fallback_chart_schema( - query, chart_type or "bar", sample_data - ) - fallback_chart_type = _chart_type_from_schema(fallback_schema, chart_type) - if not fallback_schema: - fallback_chart_type = "" - return { "results": { - "chart_schema": fallback_schema, + "chart_schema": {}, "reasoning": reasoning, - "chart_type": fallback_chart_type, + "chart_type": chart_type, } } except ValidationError as e: logger.exception(f"Vega-lite schema is not valid: {e}") - fallback_schema = _build_fallback_chart_schema( - query, - _detect_requested_chart_type(query) or "", - sample_data, - ) return { "results": { - "chart_schema": fallback_schema, + "chart_schema": {}, "reasoning": "", - "chart_type": _chart_type_from_schema( - fallback_schema, _detect_requested_chart_type(query) or "" - ) - if fallback_schema - else "", + "chart_type": "", } } except Exception as e: logger.exception(f"JSON deserialization failed: {e}") - fallback_chart_type = _detect_requested_chart_type(query) or "" - fallback_schema = _build_fallback_chart_schema( - query, fallback_chart_type, sample_data - ) return { "results": { - "chart_schema": fallback_schema, + "chart_schema": {}, "reasoning": "", - "chart_type": _chart_type_from_schema( - fallback_schema, fallback_chart_type - ) - if fallback_schema - else "", + "chart_type": "", } } diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1b503f0c88..18c9010cf3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -176,7 +176,6 @@ async def _classify_generation_result( - Policy documents and policy names are not database tables or columns unless they appear exactly in DATABASE SCHEMA. - Use schema comments, aliases, and descriptions only to choose among identifiers that already appear exactly in DATABASE SCHEMA. Do not derive new table or column names from them. - When the schema exposes raw or generated names, use those exact names in SQL. Do not replace them with natural-language names or inferred business names. -- If the user's requested concept is not represented by any exact table or column in DATABASE SCHEMA, do not assume a fallback table or column. - Do not use "*" unless the user explicitly asks for every column or every field from a table. For ordinary requests such as recent records, customers, transactions, comparisons, summaries, or aggregations, select only the columns needed to answer the question. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. @@ -455,15 +454,14 @@ async def _classify_generation_result( 9. If the user's wording is different from the schema names, use comments, aliases, and descriptions only to choose among existing schema names. Do not invent normalized, friendly, or inferred names. 10. Do not enumerate the full schema. Mention only the retrieved tables, columns, and relationships that are relevant to answering the user's question. 11. When the question requires multiple tables, use only relationships that are present in the DATABASE SCHEMA section; do not infer relationships from similar names, business wording, SQL samples, query history, or reasoning text. -12. If the required table or column is not present in DATABASE SCHEMA, state that the exact identifier is unavailable; do not assume a replacement. -13. Give a step by step reasoning plan in order to answer user's question. -14. The reasoning plan should be in the language same as the language user provided in the input. -15. Don't include SQL in the reasoning plan. -16. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. -17. Do not include ```markdown or ``` in the answer. -18. A table name in the reasoning plan must be in this format: `table: `. -19. A column name in the reasoning plan must be in this format: `column: .`. -20. ONLY SHOWING the reasoning plan in bullet points. +12. Give a step by step reasoning plan in order to answer user's question. +13. The reasoning plan should be in the language same as the language user provided in the input. +14. Don't include SQL in the reasoning plan. +15. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. +16. Do not include ```markdown or ``` in the answer. +17. A table name in the reasoning plan must be in this format: `table: `. +18. A column name in the reasoning plan must be in this format: `column: .`. +19. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -532,13 +530,12 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### GENERAL RULES ### -1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input, unless they conflict with DATABASE SCHEMA. +1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to sql samples only for SQL structure and patterns if the section of SQL SAMPLES is available in user's input. Do not copy identifiers from samples unless they appear exactly in DATABASE SCHEMA. 4. YOU MUST USE the reasoning plan only as analytical guidance if the section of REASONING PLAN is available in user's input. Do not copy table or column names from it unless they appear exactly in DATABASE SCHEMA. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. 6. DATABASE SCHEMA is more authoritative than the reasoning plan, SQL samples, and user wording. If any of those mention a table or column that is not present in DATABASE SCHEMA, do not use it. -7. For every table and column in the final SQL, verify that the exact identifier appears in DATABASE SCHEMA before returning the SQL. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 8bdf0f17f1..394d087b46 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -25,94 +25,9 @@ logger = logging.getLogger("wren-ai-service") -MAX_DB_SCHEMA_COMMENT_LENGTH = 4000 -MAX_DB_SCHEMA_STATEMENT_LENGTH = 8000 -MAX_DB_SCHEMA_DOCUMENT_LENGTH = 12000 - @component class DDLChunker: - def _normalize_text(self, value: Any) -> str: - return "" if value is None else str(value) - - def _truncate_text(self, text: Any, max_length: int) -> str: - normalized_text = self._normalize_text(text) - if len(normalized_text) <= max_length: - return normalized_text - - return normalized_text[:max_length].rstrip() + "..." - - def _serialize_table_columns_payload(self, columns: List[dict]) -> str: - return str({"type": "TABLE_COLUMNS", "columns": columns}) - - def _fit_table_column_command(self, command: dict) -> dict: - if ( - len(self._serialize_table_columns_payload([command])) - <= MAX_DB_SCHEMA_DOCUMENT_LENGTH - ): - return command - - fitted_command = {**command} - comment = fitted_command.get("comment", "") - if comment: - base_length = len( - self._serialize_table_columns_payload( - [{**fitted_command, "comment": ""}] - ) - ) - available_comment_length = max( - MAX_DB_SCHEMA_COMMENT_LENGTH // 8, - MAX_DB_SCHEMA_DOCUMENT_LENGTH - base_length - 3, - ) - fitted_command["comment"] = self._truncate_text( - comment, - available_comment_length, - ) - - if ( - len(self._serialize_table_columns_payload([fitted_command])) - <= MAX_DB_SCHEMA_DOCUMENT_LENGTH - ): - return fitted_command - - fitted_command["comment"] = "" - return fitted_command - - def _build_table_column_payloads( - self, - model_name: str, - commands: List[dict], - column_batch_size: int, - ) -> List[Dict[str, str]]: - batches: List[List[dict]] = [] - current_batch: List[dict] = [] - effective_batch_size = max(column_batch_size, 1) - - for command in [self._fit_table_column_command(command) for command in commands]: - candidate_batch = current_batch + [command] - candidate_payload = self._serialize_table_columns_payload(candidate_batch) - - if current_batch and ( - len(current_batch) >= effective_batch_size - or len(candidate_payload) > MAX_DB_SCHEMA_DOCUMENT_LENGTH - ): - batches.append(current_batch) - current_batch = [command] - continue - - current_batch = candidate_batch - - if current_batch: - batches.append(current_batch) - - return [ - { - "name": model_name, - "payload": self._serialize_table_columns_payload(batch), - } - for batch in batches - ] - @component.output_types(documents=List[Document]) async def run( self, @@ -178,13 +93,9 @@ async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: for column in model.get("columns", []) if column.get("isHidden") is not True ] - properties = model.get("properties") - if not isinstance(properties, dict): - properties = {} - return { "name": model.get("name", ""), - "properties": properties, + "properties": model.get("properties", {}), "columns": columns, "primaryKey": model.get("primaryKey", ""), } @@ -220,23 +131,12 @@ def _convert_models_and_relationships( ) -> List[Dict[str, str]]: def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) - if not isinstance(properties, dict): - properties = {} - display_name = properties.get("displayName", "") model_properties = { - "alias": clean_display_name( - "" if display_name is None else str(display_name) - ), - "description": self._truncate_text( - properties.get("description", ""), - MAX_DB_SCHEMA_COMMENT_LENGTH, - ), + "alias": clean_display_name(properties.get("displayName", "")), + "description": properties.get("description", ""), } - comment = self._truncate_text( - f"\n/* {str(model_properties)} */\n", - MAX_DB_SCHEMA_COMMENT_LENGTH, - ) + comment = f"\n/* {str(model_properties)} */\n" table_name = model["name"] payload = { @@ -255,14 +155,10 @@ def _column_command(column: Dict[str, Any], model: Dict[str, Any]) -> dict: for helper in helper.COLUMN_COMMENT_HELPERS.values() if helper.condition(column) ] - comment = self._truncate_text( - "".join(comments), - MAX_DB_SCHEMA_COMMENT_LENGTH, - ) return { "type": "COLUMN", - "comment": comment, + "comment": "".join(comments), "name": column["name"], "data_type": column["type"], "is_primary_key": column["name"] == model["primaryKey"], @@ -297,10 +193,7 @@ def _relationship_command( return { "type": "FOREIGN_KEY", - "comment": self._truncate_text( - f'-- {{"condition": {condition}, "joinType": {join_type}}}\n ', - MAX_DB_SCHEMA_COMMENT_LENGTH, - ), + "comment": f'-- {{"condition": {condition}, "joinType": {join_type}}}\n ', "constraint": fk_constraint, "tables": models, } @@ -317,11 +210,18 @@ def _column_batch( filtered = [command for command in commands if command is not None] - return self._build_table_column_payloads( - model["name"], - filtered, - column_batch_size, - ) + return [ + { + "name": model["name"], + "payload": str( + { + "type": "TABLE_COLUMNS", + "columns": filtered[i : i + column_batch_size], + } + ), + } + for i in range(0, len(filtered), column_batch_size) + ] # A map to store model primary keys for foreign key relationships primary_keys_map = {model["name"]: model["primaryKey"] for model in models} @@ -335,21 +235,13 @@ def _column_batch( def _convert_views(self, views: List[Dict[str, Any]]) -> List[Dict[str, str]]: def _payload(view: Dict[str, Any]) -> dict: - properties = view.get("properties") - if not isinstance(properties, dict): - properties = {} - return { "type": "VIEW", - "comment": self._truncate_text( - f"/* {properties} */\n" if properties else "", - MAX_DB_SCHEMA_COMMENT_LENGTH, - ), + "comment": f"/* {view['properties']} */\n" + if "properties" in view + else "", "name": view["name"], - "statement": self._truncate_text( - view["statement"], - MAX_DB_SCHEMA_STATEMENT_LENGTH, - ), + "statement": view["statement"], } return [ @@ -360,7 +252,7 @@ def _convert_metrics(self, metrics: List[Dict[str, Any]]) -> List[Dict[str, str] def _create_column(name: str, data_type: str, comment: str) -> dict: return { "type": "COLUMN", - "comment": self._truncate_text(comment, MAX_DB_SCHEMA_COMMENT_LENGTH), + "comment": comment, "name": name, "data_type": data_type, } @@ -423,9 +315,6 @@ async def chunk( @observe(capture_input=False, capture_output=False) async def embedding(chunk: Dict[str, Any], embedder: Any) -> Dict[str, Any]: - if not chunk["documents"]: - return chunk - return await embedder.run(documents=chunk["documents"]) diff --git a/wren-ai-service/src/web/v1/routers/sql_corrections.py b/wren-ai-service/src/web/v1/routers/sql_corrections.py index 9ccc7be79b..b74be58bcf 100644 --- a/wren-ai-service/src/web/v1/routers/sql_corrections.py +++ b/wren-ai-service/src/web/v1/routers/sql_corrections.py @@ -19,7 +19,6 @@ class PostRequest(BaseRequest): sql: str error: str - query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = False allow_dry_plan_fallback: bool = True diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7ce6ea2e15..844330eac7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -552,7 +552,6 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - query=user_query, instructions=instructions, invalid_generation_result={ "sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 85ed45ff79..25044de18c 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -239,7 +239,6 @@ async def ask_feedback( "sql_correction" ].run( contexts=table_ddls, - query=ask_feedback_request.question, instructions=instructions, invalid_generation_result={ "original_sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index 62586c9f1a..ed3f47f1bc 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,16 +6,11 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline -from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest logger = logging.getLogger("wren-ai-service") -NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( - "No relevant data found in the active datasource for this question." -) - # POST /v1/charts class ChartRequest(BaseRequest): @@ -84,33 +79,6 @@ def _is_stopped(self, query_id: str): return False - async def _load_active_schema_contexts( - self, project_id: Optional[str], sql: str - ) -> list[str]: - retrieval_pipeline = self._pipelines.get("db_schema_retrieval") - if not retrieval_pipeline: - return [] - - retrieval_result = await retrieval_pipeline.run( - query="", - histories=[], - project_id=project_id, - enable_column_pruning=False, - ) - documents = retrieval_result.get("construct_retrieval_results", {}).get( - "retrieval_results", [] - ) - return [ - document["table_ddl"] - for document in documents - if isinstance(document, dict) and document.get("table_ddl") - ] - - def _normalize_and_validate_sql( - self, sql: str, schema_contexts: list[str] - ) -> str | None: - return sql - @observe(name="Generate Chart") @trace_metadata async def chart( @@ -131,28 +99,6 @@ async def chart( try: query_id = chart_request.query_id execute_sql_error_message = None - schema_contexts = await self._load_active_schema_contexts( - chart_request.project_id, - chart_request.sql, - ) - normalized_sql = self._normalize_and_validate_sql( - chart_request.sql, - schema_contexts, - ) - if not normalized_sql: - self._chart_results[query_id] = ChartResultResponse( - status="failed", - error=ChartError( - code="OTHERS", - message=NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, - ), - trace_id=trace_id, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) - return results if not chart_request.data: self._chart_results[query_id] = ChartResultResponse( @@ -162,7 +108,7 @@ async def chart( execute_sql_result = ( await self._pipelines["sql_executor"].run( - sql=normalized_sql, + sql=chart_request.sql, project_id=chart_request.project_id, ) )["execute_sql"] @@ -180,19 +126,12 @@ async def chart( status="failed", error=ChartError( code="OTHERS", - message=NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, + message=execute_sql_error_message, ), trace_id=trace_id, ) - logger.info( - "Suppressed chart SQL execution failure for query_id %s: %s", - query_id, - execute_sql_error_message, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) + results["metadata"]["error_type"] = "OTHERS" + results["metadata"]["error_message"] = execute_sql_error_message return results self._chart_results[query_id] = ChartResultResponse( @@ -200,28 +139,13 @@ async def chart( trace_id=trace_id, ) - deterministic_chart_result = build_fallback_chart_result( - chart_request.query, - sql_data, - chart_request.remove_data_from_chart_schema, - ) - if deterministic_chart_result.get("chart_schema"): - self._chart_results[query_id] = ChartResultResponse( - status="finished", - response=ChartResult(**deterministic_chart_result), - trace_id=trace_id, - ) - results["chart_result"] = deterministic_chart_result - return results - chart_generation_result = await self._pipelines["chart_generation"].run( query=chart_request.query, - sql=normalized_sql, + sql=chart_request.sql, data=sql_data, language=chart_request.configurations.language, remove_data_from_chart_schema=chart_request.remove_data_from_chart_schema, custom_instruction=chart_request.custom_instruction, - contexts=schema_contexts, ) chart_result = chart_generation_result["post_process"]["results"] diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index d93e456d2a..6033237a45 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -13,31 +13,6 @@ logger = logging.getLogger("wren-ai-service") -DEFAULT_RECOMMENDATION_CONTEXT_ITEMS = 8 -DEFAULT_RECOMMENDATION_CONTEXT_CHARS = 12000 -DEFAULT_VALIDATION_CONTEXT_ITEMS = 4 -DEFAULT_VALIDATION_CONTEXT_CHARS = 6000 -STRICT_VALIDATION_CONTEXT_ITEMS = 2 -STRICT_VALIDATION_CONTEXT_CHARS = 2500 -DEFAULT_VALIDATION_SQL_SAMPLE_ITEMS = 2 -DEFAULT_VALIDATION_SQL_SAMPLE_CHARS = 2000 -STRICT_VALIDATION_SQL_SAMPLE_ITEMS = 0 -STRICT_VALIDATION_SQL_SAMPLE_CHARS = 0 -DEFAULT_VALIDATION_INSTRUCTION_ITEMS = 6 -DEFAULT_VALIDATION_INSTRUCTION_CHARS = 1500 -STRICT_VALIDATION_INSTRUCTION_ITEMS = 2 -STRICT_VALIDATION_INSTRUCTION_CHARS = 500 -DEFAULT_VALIDATION_SQL_FUNCTION_ITEMS = 12 -DEFAULT_VALIDATION_SQL_FUNCTION_CHARS = 2500 -STRICT_VALIDATION_SQL_FUNCTION_ITEMS = 0 -STRICT_VALIDATION_SQL_FUNCTION_CHARS = 0 -DEFAULT_QUESTION_CATEGORIES = [ - "Descriptive Questions", - "Segmentation Questions", - "Comparative Questions", - "Data Quality/Accuracy Questions", -] - class QuestionRecommendation: class Error(BaseModel): @@ -67,112 +42,6 @@ def __init__( self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval - def _truncate_text(self, value: str, max_chars: int) -> str: - if max_chars <= 0: - return "" - if len(value) <= max_chars: - return value - return value[:max_chars].rstrip() + "..." - - def _limit_text_items( - self, items: list[str], max_items: int, max_chars: int - ) -> list[str]: - if max_items <= 0 or max_chars <= 0: - return [] - - limited_items: list[str] = [] - remaining_chars = max_chars - - for item in items: - if len(limited_items) >= max_items or remaining_chars <= 0: - break - - truncated = self._truncate_text(item, remaining_chars) - if not truncated: - break - - limited_items.append(truncated) - remaining_chars -= len(truncated) - - return limited_items - - def _limit_dict_items( - self, - items: list[dict], - key: str, - max_items: int, - max_chars: int, - ) -> list[dict]: - if max_items <= 0 or max_chars <= 0: - return [] - - limited_items: list[dict] = [] - remaining_chars = max_chars - - for item in items: - if len(limited_items) >= max_items or remaining_chars <= 0: - break - - value = str(item.get(key, "")) - truncated_value = self._truncate_text(value, remaining_chars) - if not truncated_value: - break - - next_item = {**item, key: truncated_value} - limited_items.append(next_item) - remaining_chars -= len(truncated_value) - - return limited_items - - def _limit_sql_functions( - self, functions: list, max_items: int, max_chars: int - ) -> list: - serialized_functions = [str(function) for function in functions] - limited_values = self._limit_text_items( - serialized_functions, - max_items=max_items, - max_chars=max_chars, - ) - limited_count = len(limited_values) - return functions[:limited_count] - - def _is_context_size_error(self, error: Exception) -> bool: - error_message = str(error).lower() - return any( - phrase in error_message - for phrase in [ - "context size has been exceeded", - "too large to process", - "maximum context length", - "prompt is too long", - ] - ) - - def _get_target_categories( - self, - requested_categories: list[str], - max_categories: int, - ) -> list[str]: - categories = requested_categories or DEFAULT_QUESTION_CATEGORIES - return categories[:max_categories] - - def _get_underfilled_categories( - self, - response_questions: dict[str, list[dict]], - requested_categories: list[str], - max_categories: int, - max_questions: int, - ) -> list[str]: - target_categories = self._get_target_categories( - requested_categories=requested_categories, - max_categories=max_categories, - ) - return [ - category - for category in target_categories - if len(response_questions.get(category, [])) < max_questions - ] - def _handle_exception( self, event_id: str, @@ -252,83 +121,19 @@ async def _instructions_retrieval() -> list[dict]: else: sql_knowledge = None - validation_attempts = [ - { - "contexts": self._limit_text_items( - table_ddls, - max_items=DEFAULT_VALIDATION_CONTEXT_ITEMS, - max_chars=DEFAULT_VALIDATION_CONTEXT_CHARS, - ), - "sql_samples": self._limit_dict_items( - sql_samples, - key="sql", - max_items=DEFAULT_VALIDATION_SQL_SAMPLE_ITEMS, - max_chars=DEFAULT_VALIDATION_SQL_SAMPLE_CHARS, - ), - "instructions": self._limit_dict_items( - instructions, - key="instruction", - max_items=DEFAULT_VALIDATION_INSTRUCTION_ITEMS, - max_chars=DEFAULT_VALIDATION_INSTRUCTION_CHARS, - ), - "sql_functions": self._limit_sql_functions( - sql_functions, - max_items=DEFAULT_VALIDATION_SQL_FUNCTION_ITEMS, - max_chars=DEFAULT_VALIDATION_SQL_FUNCTION_CHARS, - ), - }, - { - "contexts": self._limit_text_items( - table_ddls, - max_items=STRICT_VALIDATION_CONTEXT_ITEMS, - max_chars=STRICT_VALIDATION_CONTEXT_CHARS, - ), - "sql_samples": self._limit_dict_items( - sql_samples, - key="sql", - max_items=STRICT_VALIDATION_SQL_SAMPLE_ITEMS, - max_chars=STRICT_VALIDATION_SQL_SAMPLE_CHARS, - ), - "instructions": self._limit_dict_items( - instructions, - key="instruction", - max_items=STRICT_VALIDATION_INSTRUCTION_ITEMS, - max_chars=STRICT_VALIDATION_INSTRUCTION_CHARS, - ), - "sql_functions": self._limit_sql_functions( - sql_functions, - max_items=STRICT_VALIDATION_SQL_FUNCTION_ITEMS, - max_chars=STRICT_VALIDATION_SQL_FUNCTION_CHARS, - ), - }, - ] - - generated_sql = None - for attempt_index, attempt in enumerate(validation_attempts): - try: - generated_sql = await self._pipelines["sql_generation"].run( - query=candidate["question"], - contexts=attempt["contexts"], - project_id=project_id, - sql_samples=attempt["sql_samples"], - instructions=attempt["instructions"], - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=attempt["sql_functions"], - allow_data_preview=allow_data_preview, - sql_knowledge=sql_knowledge, - ) - break - except Exception as error: - is_last_attempt = attempt_index == len(validation_attempts) - 1 - if is_last_attempt or not self._is_context_size_error(error): - raise - - logger.warning( - "Request %s: SQL validation prompt exceeded context window; retrying with reduced context", - request_id, - ) + generated_sql = await self._pipelines["sql_generation"].run( + query=candidate["question"], + contexts=table_ddls, + project_id=project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + allow_data_preview=allow_data_preview, + sql_knowledge=sql_knowledge, + ) post_process = generated_sql["post_process"] @@ -364,7 +169,6 @@ class Request(BaseRequest): event_id: str mdl: str previous_questions: list[str] = [] - categories: list[str] = [] max_questions: int = 5 max_categories: int = 3 regenerate: bool = False @@ -396,28 +200,18 @@ async def recommend(self, input: Request, **kwargs) -> Event: trace_id = kwargs.get("trace_id") try: - orjson.loads(input.mdl) + mdl = orjson.loads(input.mdl) retrieval_result = await self._pipelines["db_schema_retrieval"].run( - query="", - histories=[], + tables=[model["name"] for model in mdl["models"]], project_id=input.project_id, - enable_column_pruning=False, ) _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) - table_ddls = self._limit_text_items( - [document.get("table_ddl") for document in documents], - max_items=DEFAULT_RECOMMENDATION_CONTEXT_ITEMS, - max_chars=DEFAULT_RECOMMENDATION_CONTEXT_CHARS, - ) + table_ddls = [document.get("table_ddl") for document in documents] request = { "contexts": table_ddls, "previous_questions": input.previous_questions, - "categories": self._get_target_categories( - requested_categories=input.categories, - max_categories=input.max_categories, - ), "language": input.configurations.language, "max_questions": input.max_questions, "max_categories": input.max_categories, @@ -432,18 +226,17 @@ async def recommend(self, input: Request, **kwargs) -> Event: resource.trace_id = trace_id response = resource.response - categories = self._get_underfilled_categories( - response_questions=response["questions"], - requested_categories=request["categories"], - max_categories=input.max_categories, - max_questions=input.max_questions, - ) - need_regenerate = bool(categories) and input.regenerate + categories_count = { + category: input.max_questions - len(questions) + for category, questions in response["questions"].items() + if len(questions) < input.max_questions + } + categories = list(categories_count.keys()) + need_regenerate = len(categories) > 0 and input.regenerate resource.status = "generating" if need_regenerate else "finished" if resource.status == "finished": - resource.request_from = input.request_from return resource.with_metadata() await self._recommend( diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index dea25902a9..c907387300 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -12,10 +12,6 @@ logger = logging.getLogger("wren-ai-service") -NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE = ( - "No relevant data found in the active datasource for this question." -) - # POST /v1/sql-answers class SqlAnswerRequest(BaseRequest): @@ -57,33 +53,6 @@ def __init__( maxsize=maxsize, ttl=ttl ) - async def _load_active_schema_contexts( - self, project_id: Optional[str] - ) -> list[str]: - retrieval_pipeline = self._pipelines.get("db_schema_retrieval") - if not retrieval_pipeline: - return [] - - retrieval_result = await retrieval_pipeline.run( - query="", - histories=[], - project_id=project_id, - enable_column_pruning=False, - ) - documents = retrieval_result.get("construct_retrieval_results", {}).get( - "retrieval_results", [] - ) - return [ - document["table_ddl"] - for document in documents - if isinstance(document, dict) and document.get("table_ddl") - ] - - def _normalize_and_validate_sql( - self, sql: str, schema_contexts: list[str] - ) -> str | None: - return sql - @observe(name="SQL Answer") @trace_metadata async def sql_answer( @@ -110,28 +79,6 @@ async def sql_answer( trace_id=trace_id, ) - schema_contexts = await self._load_active_schema_contexts( - sql_answer_request.project_id, - ) - normalized_sql = self._normalize_and_validate_sql( - sql_answer_request.sql, - schema_contexts, - ) - if not normalized_sql: - self._sql_answer_results[query_id] = SqlAnswerResultResponse( - status="failed", - error=SqlAnswerResultResponse.SqlAnswerError( - code="OTHERS", - message=NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE, - ), - trace_id=trace_id, - ) - results["metadata"]["error_type"] = "NO_RELEVANT_DATA" - results["metadata"]["error_message"] = ( - NO_RELEVANT_ACTIVE_DATASOURCE_MESSAGE - ) - return results - preprocessed_sql_data = self._pipelines["preprocess_sql_data"].run( sql_data=sql_answer_request.sql_data, )["preprocess"] @@ -149,13 +96,12 @@ async def sql_answer( asyncio.create_task( self._pipelines["sql_answer"].run( query=sql_answer_request.query, - sql=normalized_sql, + sql=sql_answer_request.sql, sql_data=preprocessed_sql_data.get("sql_data", {}), language=sql_answer_request.configurations.language, current_time=sql_answer_request.configurations.show_current_time(), query_id=query_id, custom_instruction=sql_answer_request.custom_instruction, - contexts=schema_contexts, ) ) diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 3634199927..86d0f55301 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -60,7 +60,6 @@ class CorrectionRequest(BaseRequest): event_id: str sql: str error: str - query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = False allow_dry_plan_fallback: bool = True @@ -77,7 +76,6 @@ async def correct( event_id = request.event_id sql = request.sql error = request.error - query = request.query project_id = request.project_id retrieved_tables = request.retrieved_tables use_dry_plan = request.use_dry_plan @@ -116,7 +114,6 @@ async def correct( res = await self._pipelines["sql_correction"].run( contexts=table_ddls, - query=query, invalid_generation_result=_invalid, project_id=project_id, use_dry_plan=use_dry_plan, From 1b4c5cb1e37f6da8e873717b6ae75745340aa3bf Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 18:30:11 +0530 Subject: [PATCH 0651/1087] Fix ask pipeline SQL routing --- .../pipelines/generation/data_assistance.py | 7 +-- .../generation/intent_classification.py | 28 ++++----- .../generation/misleading_assistance.py | 7 +-- .../pipelines/generation/sql_correction.py | 1 - .../src/pipelines/generation/sql_diagnosis.py | 12 +--- .../src/pipelines/generation/utils/sql.py | 62 ++++++------------- wren-ai-service/src/web/v1/services/ask.py | 3 - 7 files changed, 37 insertions(+), 83 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index 09ea192ad0..51b91197f9 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -1,7 +1,7 @@ import asyncio import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -12,10 +12,7 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index da52f2eb34..90cbba6310 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -33,11 +33,10 @@ - **Rephrase Question:** Rewrite follow-up questions into full standalone questions using prior conversation context. - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. -- **Vague Queries:** If the question is not a data retrieval or analysis request and does not relate to the schema, classify it as `MISLEADING_QUERY`. -- **Natural Language Data Queries:** Do not require exact physical table or column names. If the user asks to retrieve or analyze data and the schema context may answer it, classify it as `TEXT_TO_SQL`. -- **Intent Precedence:** Prefer `TEXT_TO_SQL` over `GENERAL` for data retrieval or analysis requests. Lack of exact table or column names is not a reason to choose `GENERAL`. -- **User Guide Boundary:** Do not classify a data retrieval, filtering, date, metric, aggregation, or row-listing question as `USER_GUIDE` just because it contains words like "how", "show", or "filter". -- **Incomplete Queries:** If the question is related to the database schema but references unresolved placeholders (e.g., "the following", "these", "those", "the previous ones") without providing them or prior context, classify as `GENERAL`. +- **Data Retrieval Requests:** If the user asks to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data from the connected database, classify it as `TEXT_TO_SQL`. +- **Database Schema Exploration:** If the user asks about available tables, columns, relationships, schema meaning, or what can be asked, classify it as `GENERAL`. +- **Out-of-Scope Queries:** If the question is unrelated to the database schema or data retrieval, classify it as `MISLEADING_QUERY`. +- **Incomplete Queries:** If the question references unresolved placeholders (e.g., "the following", "these", "those") without providing them or prior context, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. ### Intent Definitions ### @@ -45,18 +44,15 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. +- The user's inputs are related to the database schema and requires an SQL query. - The user's inputs ask to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data. -- The user's inputs are related to the database schema and require an SQL query. - The question can be answered by selecting relevant tables and columns from the provided schema, even if the user does not mention exact physical table or column names. - The question includes enough business meaning, dimensions, metrics, filters, or time criteria to attempt SQL generation from the schema. -- Natural-language analytical questions should be classified as `TEXT_TO_SQL` when they can reasonably be grounded in the schema. -- Broad row-listing requests are `TEXT_TO_SQL` if related schema is present. **Requirements:** - Do not require the user to explicitly name a table or column. -- Use the schema to determine whether the user's business terms can map to available tables or columns. -- Do not generate SQL in this step. Only classify intent and rephrase the question. -- Reference phrases from the user's inputs that clearly indicate a data retrieval or analysis request. +- Use the provided schema context to decide whether the user's business terms can be answered by SQL. +- Reference phrases from the user's inputs that clearly indicate a data retrieval request. **Examples:** - "What is the total sales for last quarter?" @@ -67,11 +63,10 @@ **When to Use:** - The user seeks general information about the database schema or its overall capabilities. +- The user asks about available tables, columns, relationships, schema meaning, or what questions can be asked. - The query references **missing information** (e.g., "the following items" without listing them). - The query contains **placeholder references** that cannot be resolved from context. -- The query is **incomplete for SQL generation** because required values or references are missing, not merely because exact table or column names are absent. -- The user is asking for explanation, guidance, or clarification rather than asking to retrieve or analyze rows from the data. -- Do not classify row-listing, filtering, ordering, metric, date, or aggregation requests as `GENERAL` merely because the user used business wording. +- The query is asking for explanation or guidance rather than retrieval, filtering, ordering, aggregation, or analysis of rows. **Requirements:** - Incorporate phrases from the user's inputs that indicate incompleteness or lack of relevance to the database schema. @@ -80,6 +75,8 @@ **Examples:** - "What is the dataset about?" - "Tell me more about the database." +- "Explain the customer table to me." +- "What tables do I have?" - "How can I analyze customer behavior with this data?" - "Show me orders for these products" (without specifying which products) - "Filter by the criteria I mentioned" (without previous context defining criteria) @@ -89,7 +86,6 @@ **When to Use:** - The user's inputs pertains to Wren AI's features, usage, or capabilities. - The query relates directly to content in the user guide. -- The query asks how to use Wren AI itself, not how to retrieve or analyze rows from the connected data. **Examples:** - "What can Wren AI do?" @@ -102,7 +98,7 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs cannot be interpreted as a database question or data analysis request. +- The user's inputs cannot be interpreted as a database question or data retrieval request. - It appears off-topic or is simply a casual conversation starter. **Requirements:** diff --git a/wren-ai-service/src/pipelines/generation/misleading_assistance.py b/wren-ai-service/src/pipelines/generation/misleading_assistance.py index 3053ab53e0..a35738ecf5 100644 --- a/wren-ai-service/src/pipelines/generation/misleading_assistance.py +++ b/wren-ai-service/src/pipelines/generation/misleading_assistance.py @@ -1,7 +1,7 @@ import asyncio import logging import sys -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from hamilton import base from hamilton.async_driver import AsyncDriver @@ -12,10 +12,7 @@ from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.utils import trace_cost -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory -else: - AskHistory = Any +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 601e75d2ae..973b8c69a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -36,7 +36,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). 2. Then, generate the syntactically correct ANSI SQL query to correct the error. -3. Use the failed SQL and error message as diagnostic context only. Regenerate the corrected SQL from the DATABASE SCHEMA instead of patching invalid table or column names from the failed SQL or error message. ### SQL RULES ### Make sure you follow the SQL Rules strictly. diff --git a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py index 129536cabe..3f22b9d512 100644 --- a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py +++ b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py @@ -92,17 +92,7 @@ async def generate_sql_diagnosis( async def post_process( generate_sql_diagnosis: dict, ) -> str: - reply = (generate_sql_diagnosis.get("replies") or [""])[0] - try: - parsed_reply = orjson.loads(reply) - except orjson.JSONDecodeError: - return {"reasoning": reply.strip()} - - if isinstance(parsed_reply, dict): - reasoning = parsed_reply.get("reasoning", "") - return {"reasoning": reasoning if isinstance(reasoning, str) else ""} - - return {"reasoning": str(parsed_reply)} + return orjson.loads(generate_sql_diagnosis.get("replies")[0]) ## End of Pipeline diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 18c9010cf3..088282574e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,7 +1,5 @@ -from __future__ import annotations - import logging -from typing import TYPE_CHECKING, Any, Dict, List +from typing import Any, Dict, List import aiohttp import orjson @@ -14,9 +12,7 @@ clean_generation_result, ) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge - -if TYPE_CHECKING: - from src.web.v1.services.ask import AskHistory +from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -169,19 +165,12 @@ async def _classify_generation_result( _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### -- ONLY USE SELECT statements. Do not generate ALTER, CREATE, DROP, INSERT, UPDATE, DELETE, MERGE, TRUNCATE, GRANT, REVOKE, or any other statement that can change the database or schema. +- ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. -- Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. Never invent, rename, translate, normalize, prefix, suffix, or guess an identifier. -- Do not use names from policies, instructions, SQL functions, SQL samples, query history, reasoning text, error messages, or the user's wording as table or column identifiers unless they appear exactly in DATABASE SCHEMA. -- Policy documents and policy names are not database tables or columns unless they appear exactly in DATABASE SCHEMA. -- Use schema comments, aliases, and descriptions only to choose among identifiers that already appear exactly in DATABASE SCHEMA. Do not derive new table or column names from them. -- When the schema exposes raw or generated names, use those exact names in SQL. Do not replace them with natural-language names or inferred business names. -- Do not use "*" unless the user explicitly asks for every column or every field from a table. For ordinary requests such as recent records, customers, transactions, comparisons, summaries, or aggregations, select only the columns needed to answer the question. +- ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! -- Use multiple retrieved tables when the question requires them, and join only with relationships that are present in the DATABASE SCHEMA section. Do not infer joins from similar names, business wording, SQL samples, query history, or reasoning text. -- Do not invent measure, dimension, or filter columns. If the question refers to a business concept, use only an exact table or column that appears in DATABASE SCHEMA. - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. @@ -204,7 +193,6 @@ async def _classify_generation_result( - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) -- For date parts such as year, month, day, quarter, or week, use `EXTRACT( FROM CAST( AS TIMESTAMP WITH TIME ZONE))`. Do not use `YEAR()`, `MONTH()`, `DAY()`, or `DATEPART()` unless the SQL FUNCTIONS section explicitly says that function is supported. - If the user asks for a specific date, please give the date range in SQL query - example: "What is the total revenue for the month of 2024-11-01?" - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" @@ -447,21 +435,16 @@ async def _classify_generation_result( otherwise, you will put the relative timeframe in the SQL query. 3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. -5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan, but do not treat them as schema. -6. If SQL SAMPLES section is provided, use them only as examples of SQL structure. Do not treat their identifiers as schema unless they appear exactly in DATABASE SCHEMA. -7. Treat the DATABASE SCHEMA section as the only authoritative source for table and column identifiers. -8. Do not mention a table or column in the reasoning plan unless it appears in the DATABASE SCHEMA section. -9. If the user's wording is different from the schema names, use comments, aliases, and descriptions only to choose among existing schema names. Do not invent normalized, friendly, or inferred names. -10. Do not enumerate the full schema. Mention only the retrieved tables, columns, and relationships that are relevant to answering the user's question. -11. When the question requires multiple tables, use only relationships that are present in the DATABASE SCHEMA section; do not infer relationships from similar names, business wording, SQL samples, query history, or reasoning text. -12. Give a step by step reasoning plan in order to answer user's question. -13. The reasoning plan should be in the language same as the language user provided in the input. -14. Don't include SQL in the reasoning plan. -15. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. -16. Do not include ```markdown or ``` in the answer. -17. A table name in the reasoning plan must be in this format: `table: `. -18. A column name in the reasoning plan must be in this format: `column: .`. -19. ONLY SHOWING the reasoning plan in bullet points. +5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. +6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. +7. Give a step by step reasoning plan in order to answer user's question. +8. The reasoning plan should be in the language same as the language user provided in the input. +9. Don't include SQL in the reasoning plan. +10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. +11. Do not include ```markdown or ``` in the answer. +12. A table name in the reasoning plan must be in this format: `table: `. +13. A column name in the reasoning plan must be in this format: `column: .`. +14. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -480,13 +463,9 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: if sql_knowledge is not None: - value = getattr(sql_knowledge, "text_to_sql_rule", "") - if value and value.strip(): - return ( - f"{_DEFAULT_TEXT_TO_SQL_RULES}\n\n" - "### PROJECT SQL RULES ###\n" - f"{value.strip()}" - ) + return _extract_from_sql_knowledge( + sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES + ) return _DEFAULT_TEXT_TO_SQL_RULES @@ -532,10 +511,9 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -3. YOU MUST REFER to sql samples only for SQL structure and patterns if the section of SQL SAMPLES is available in user's input. Do not copy identifiers from samples unless they appear exactly in DATABASE SCHEMA. -4. YOU MUST USE the reasoning plan only as analytical guidance if the section of REASONING PLAN is available in user's input. Do not copy table or column names from it unless they appear exactly in DATABASE SCHEMA. +3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. +4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. -6. DATABASE SCHEMA is more authoritative than the reasoning plan, SQL samples, and user wording. If any of those mention a table or column that is not present in DATABASE SCHEMA, do not use it. {text_to_sql_rules} @@ -576,7 +554,7 @@ def construct_instructions( def construct_ask_history_messages( - histories: list["AskHistory"] | list[dict], + histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: messages = [] for history in histories: diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 844330eac7..aa26fa3f81 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -105,9 +105,6 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, - max_sql_generation_tables: int = 0, - pipeline_timeout_seconds: int = 0, - schema_retrieval_timeout_seconds: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, From e7d58f7fdd66384a606f816a5b355da0f6e12247 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sat, 25 Jul 2026 18:48:11 +0530 Subject: [PATCH 0652/1087] Fix AskService startup kwargs --- wren-ai-service/src/web/v1/services/ask.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index aa26fa3f81..844330eac7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -105,6 +105,9 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, + max_sql_generation_tables: int = 0, + pipeline_timeout_seconds: int = 0, + schema_retrieval_timeout_seconds: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, From ba8c4cc65f4edc691c6627bb6db45d4a76e0a933 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 19:06:58 +0530 Subject: [PATCH 0653/1087] Remove hardcoded SQL normalization paths --- .../generation/followup_sql_generation.py | 7 +- .../pipelines/generation/sql_correction.py | 7 +- .../pipelines/generation/sql_generation.py | 7 +- .../src/pipelines/generation/utils/sql.py | 4 + .../src/pipelines/retrieval/sql_executor.py | 3 - .../src/pipelines/sql_normalizer.py | 491 --------------- .../src/apollo/server/adaptors/ibisAdaptor.ts | 5 +- .../apollo/server/services/queryService.ts | 427 +------------ .../apollo/server/utils/mssqlSqlNormalizer.ts | 562 ------------------ .../server/utils/recommendationQuestions.ts | 13 +- .../utils/tests/mssqlSqlNormalizer.test.ts | 258 -------- 11 files changed, 33 insertions(+), 1751 deletions(-) delete mode 100644 wren-ai-service/src/pipelines/sql_normalizer.py delete mode 100644 wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts delete mode 100644 wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts 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 35cfb8fccf..d82d6e40b6 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -211,10 +211,9 @@ async def run( ): logger.info("Follow-Up SQL Generation pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = ( + await retrieve_metadata(project_id, self._retriever) if project_id else {} + ) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..d9caf87a84 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -178,10 +178,9 @@ async def run( ): logger.info("SQLCorrection pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = ( + await retrieve_metadata(project_id, self._retriever) if project_id else {} + ) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..d0b921f70c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -205,10 +205,9 @@ async def run( ): logger.info("SQL Generation pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) - else: - metadata = {} + metadata = ( + await retrieve_metadata(project_id, self._retriever) if project_id else {} + ) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..ab22d2fde8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -553,6 +553,10 @@ def construct_instructions( return _instructions +def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: + return " ".join(sql.replace('\\"', '"').split()) + + def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: diff --git a/wren-ai-service/src/pipelines/retrieval/sql_executor.py b/wren-ai-service/src/pipelines/retrieval/sql_executor.py index 129c4992c3..97d0bf8bc5 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_executor.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_executor.py @@ -29,9 +29,6 @@ async def run( project_id: str | None = None, limit: int = 500, ): - from src.pipelines.generation.utils.sql import normalize_generation_result_sql - - sql = normalize_generation_result_sql(sql, data_source=self._data_source) async with aiohttp.ClientSession() as session: _, data, addition = await self._engine.execute_sql( sql, diff --git a/wren-ai-service/src/pipelines/sql_normalizer.py b/wren-ai-service/src/pipelines/sql_normalizer.py deleted file mode 100644 index ac61a57741..0000000000 --- a/wren-ai-service/src/pipelines/sql_normalizer.py +++ /dev/null @@ -1,491 +0,0 @@ -import re -from datetime import datetime, timedelta - - -def normalize_data_source(data_source: str | None) -> str: - normalized = (data_source or "").strip().upper().replace("-", "_").replace( - " ", "_" - ) - if normalized in {"SQLSERVER", "SQL_SERVER", "MS_SQL", "MSSQLSERVER"}: - return "MSSQL" - return normalized - - -def _format_timestamp_literal(value: datetime) -> str: - return value.strftime("'%Y-%m-%d %H:%M:%S'") - - -def _add_months(value: datetime, months: int) -> datetime: - month_index = value.month - 1 + months - year = value.year + month_index // 12 - month = month_index % 12 + 1 - day = min( - value.day, - [ - 31, - 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28, - 31, - 30, - 31, - 30, - 31, - 31, - 30, - 31, - 30, - 31, - ][month - 1], - ) - return value.replace(year=year, month=month, day=day) - - -def _start_of_month(value: datetime) -> datetime: - return value.replace(day=1, hour=0, minute=0, second=0, microsecond=0) - - -def _replace_relative_getdate_calls(sql: str, now: datetime) -> str: - def replace_month_offset(match: re.Match[str]) -> str: - months = int(match.group(1)) - return _format_timestamp_literal(_add_months(now, months)) - - def replace_year_offset(match: re.Match[str]) -> str: - years = int(match.group(1)) - return _format_timestamp_literal(_add_months(now, years * 12)) - - def replace_day_offset(match: re.Match[str]) -> str: - days = int(match.group(1)) - return _format_timestamp_literal(now + timedelta(days=days)) - - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_month_offset, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*year\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_year_offset, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*day\s*,\s*([+-]?\d+)\s*,\s*GETDATE\(\)\s*\)", - replace_day_offset, - sql, - flags=re.IGNORECASE, - ) - - current_month_start = _format_timestamp_literal(_start_of_month(now)) - previous_month_start = _format_timestamp_literal( - _start_of_month(_add_months(now, -1)) - ) - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*,\s*0\s*\)", - current_month_start, - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - r"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*GETDATE\(\)\s*\)\s*-\s*1\s*,\s*0\s*\)", - previous_month_start, - sql, - flags=re.IGNORECASE, - ) - return sql - - -def _rewrite_mssql_bucket_functions(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - - sql = re.sub( - rf"DATEADD\(\s*month\s*,\s*DATEDIFF\(\s*month\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: ( - f"(EXTRACT(YEAR FROM {m.group(1)}) * 100 + EXTRACT(MONTH FROM {m.group(1)}))" - ), - sql, - flags=re.IGNORECASE, - ) - sql = re.sub( - rf"DATEADD\(\s*year\s*,\s*DATEDIFF\(\s*year\s*,\s*0\s*,\s*{expression_pattern}\s*\)\s*,\s*0\s*\)", - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - sql, - flags=re.IGNORECASE, - ) - return sql - - -def _rewrite_temporal_bucket_functions(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - replacements = [ - ( - re.compile( - rf"DATEPART\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*DAY\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile(rf"YEAR\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile(rf"MONTH\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile(rf"DAY\(\s*{expression_pattern}\s*\)", re.IGNORECASE), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATETRUNC\(\s*MONTH\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATETRUNC\(\s*YEAR\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_TRUNC\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_TRUNC\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*YEAR\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*MONTH\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATE_PART\(\s*'?\s*DAY\s*'?\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*YEAR\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*MONTH\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"EXTRACT\(\s*DAY\s+FROM\s+{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'YEAR'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(YEAR FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'MONTH'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(MONTH FROM {m.group(1)})", - ), - ( - re.compile( - rf"DATEPART\(\s*'DAY'\s*,\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ), - lambda m: f"EXTRACT(DAY FROM {m.group(1)})", - ), - ] - - rewritten = sql - for pattern, replacement in replacements: - rewritten = pattern.sub(replacement, rewritten) - - return rewritten - - -def _rewrite_mssql_timestamp_casts(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - timestamp_function_pattern = re.compile( - rf"\bTO_TIMESTAMP(?:_(?:MILLIS|SECONDS|MICROS|NANOS))?\(\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ) - timestamp_cast_pattern = re.compile( - r"CAST\(\s*((?:[^()]|\([^()]*\))+?)\s+AS\s+TIMESTAMP\s*\)", - re.IGNORECASE, - ) - - rewritten = timestamp_function_pattern.sub( - lambda m: f"CAST({m.group(1)} AS DATETIME)", sql - ) - rewritten = timestamp_cast_pattern.sub( - lambda m: f"CAST({m.group(1)} AS DATETIME)", rewritten - ) - return rewritten - - -def _rewrite_mssql_to_unixtime(sql: str) -> str: - expression_pattern = r"((?:[^(),]|\([^()]*\))+?)" - to_unixtime_pattern = re.compile( - rf"\bTO_UNIXTIME\(\s*{expression_pattern}\s*\)", - re.IGNORECASE, - ) - - return to_unixtime_pattern.sub(lambda m: m.group(1), sql) - - -def _rewrite_mssql_timestamp_subtraction(sql: str) -> str: - expression_pattern = r"((?:[^(),+\-]|\([^()]*\))+?)" - timestamp_subtraction_pattern = re.compile( - rf"{expression_pattern}\s*-\s*{expression_pattern}\s+AS\s+(\"[^\"]+\")", - re.IGNORECASE, - ) - - def replace_subtraction(match: re.Match[str]) -> str: - left = match.group(1).strip() - right = match.group(2).strip() - alias = match.group(3) - alias_text = str(alias or "").strip('"').lower() - - if not any(token in alias_text for token in ("duration", "turnaround")): - return match.group(0) - - return f"DATEDIFF('second', {right}, {left}) AS {alias}" - - return timestamp_subtraction_pattern.sub(replace_subtraction, sql) - - -def _infer_mssql_timestamp_expression(sql: str) -> str | None: - timestamp_column_pattern = re.compile( - r'(?:(?:"[^"]+"\.)?"(?:created_at|updated_at|generated_at|opened_at|closed_at|completed_at|resolved_at)")', - re.IGNORECASE, - ) - if match := timestamp_column_pattern.search(sql): - return match.group(0) - - return None - - -def _rewrite_mssql_invented_date_identifiers(sql: str) -> str: - timestamp_expression = _infer_mssql_timestamp_expression(sql) - if not timestamp_expression: - return sql - - invented_date_identifier_pattern = re.compile( - r'(? str: - timestamp_expression = _infer_mssql_timestamp_expression(sql) - if not timestamp_expression: - return sql - - bucket_expressions = { - "year": f"EXTRACT(YEAR FROM {timestamp_expression})", - "month": f"EXTRACT(MONTH FROM {timestamp_expression})", - "day": f"EXTRACT(DAY FROM {timestamp_expression})", - } - rewritten = sql - - for bucket, expression in bucket_expressions.items(): - select_identifier_pattern = re.compile( - rf'(?P\bSELECT\s+|,\s*)"{bucket}"(?P\s*(?:,|\bFROM\b))', - re.IGNORECASE, - ) - - def replace_select_identifier(match: re.Match[str]) -> str: - prefix = match.group("prefix") - suffix = match.group("suffix") - return f'{prefix}{expression} AS "{bucket}"{suffix}' - - rewritten = select_identifier_pattern.sub( - replace_select_identifier, rewritten - ) - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match[str]) -> str: - body = match.group("body") - for bucket, expression in bucket_expressions.items(): - body = re.sub( - rf'"{bucket}"', - expression, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"\[{bucket}\]", - expression, - body, - flags=re.IGNORECASE, - ) - return f"{match.group(1)}{body}" - - return clause_pattern.sub(replace_clause, rewritten) - - -def _rewrite_mssql_invented_failure_category(sql: str) -> str: - if not re.search(r"\bdbo_repair_logs\b", sql, flags=re.IGNORECASE): - return sql - if not re.search(r"\bfailure[_\s]+category\b", sql, flags=re.IGNORECASE): - return sql - - failure_code_expression = '"dbo_repair_logs"."failure_code"' - rewritten = re.sub( - r'(?P\bSELECT\s+|,\s*)(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)(?P\s*(?:,|\bFROM\b))', - rf'\g{failure_code_expression} AS "failure_category"\g', - sql, - flags=re.IGNORECASE, - ) - rewritten = re.sub( - r"\bAS\s+failure\s+category\b", - 'AS "failure_category"', - rewritten, - flags=re.IGNORECASE, - ) - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match[str]) -> str: - body = re.sub( - r'(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)', - failure_code_expression, - match.group("body"), - flags=re.IGNORECASE, - ) - return f"{match.group(1)}{body}" - - return clause_pattern.sub(replace_clause, rewritten) - - -def _rewrite_mssql_datepart_alias_references(sql: str) -> str: - datepart_alias_pattern = re.compile( - r"\b((?:DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)|(?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\)|EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\)))\s+AS\s+(?:\"([^\"]+)\"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))", - re.IGNORECASE, - ) - aliases: dict[str, str] = {} - - for match in datepart_alias_pattern.finditer(sql): - expression = match.group(1) - alias = match.group(5) or match.group(6) or match.group(7) - if alias: - aliases[str(alias).lower()] = expression - - if not aliases: - return sql - - clause_pattern = re.compile( - r"\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?P.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - - def replace_clause(match: re.Match) -> str: - body = match.group("body") - placeholders: dict[str, str] = {} - for alias, expression in aliases.items(): - placeholder = f"__WREN_MSSQL_DATEPART_ALIAS_{len(placeholders)}__" - placeholders[placeholder] = expression - body = re.sub( - rf'"{re.escape(alias)}"', - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"\[{re.escape(alias)}\]", - placeholder, - body, - flags=re.IGNORECASE, - ) - body = re.sub( - rf"(? str: - normalized = sql - - if normalize_data_source(data_source) == "MSSQL": - now = datetime.now() - normalized = re.sub( - r"\s+NULLS\s+(?:LAST|FIRST)\b", "", normalized, flags=re.IGNORECASE - ) - normalized = re.sub( - r"CAST\(\s*('(?:[^']|'')*')\s+AS\s+DATETIME(?:2|OFFSET)\s*\)", - r"\1", - normalized, - flags=re.IGNORECASE, - ) - normalized = _replace_relative_getdate_calls(normalized, now) - normalized = _rewrite_mssql_to_unixtime(normalized) - normalized = _rewrite_mssql_timestamp_subtraction(normalized) - normalized = _rewrite_mssql_timestamp_casts(normalized) - normalized = _rewrite_mssql_bare_time_bucket_identifiers(normalized) - normalized = _rewrite_mssql_bucket_functions(normalized) - normalized = _rewrite_temporal_bucket_functions(normalized) - normalized = _rewrite_mssql_datepart_alias_references(normalized) - - return re.sub(r"\s+", " ", normalized).strip() diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index a6adb47b1d..3b16be0c34 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -6,7 +6,6 @@ import { Manifest } from '@server/mdl/type'; import * as Errors from '@server/utils/error'; import { getConfig } from '@server/config'; import { toDockerHost } from '@server/utils'; -import { normalizeMssqlSqlForIbis } from '@server/utils/mssqlSqlNormalizer'; import { CompactColumn, CompactTable, @@ -271,7 +270,7 @@ export class IbisAdaptor implements IIbisAdaptor { public async getNativeSql(options: IbisDryPlanOptions): Promise { const { dataSource, mdl, sql } = options; const body = { - sql: normalizeMssqlSqlForIbis(sql, dataSource), + sql, manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), }; try { @@ -295,7 +294,6 @@ export class IbisAdaptor implements IIbisAdaptor { options: IbisQueryOptions, ): Promise { const { dataSource, mdl } = options; - query = normalizeMssqlSqlForIbis(query, dataSource); const connectionInfo = this.updateConnectionInfo(options.connectionInfo); const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo); const queryString = this.buildQueryString(options); @@ -340,7 +338,6 @@ export class IbisAdaptor implements IIbisAdaptor { options: IbisQueryOptions, ): Promise { const { dataSource, mdl } = options; - query = normalizeMssqlSqlForIbis(query, dataSource); const connectionInfo = this.updateConnectionInfo(options.connectionInfo); const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo); const body = { diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index e4853589ac..62434e4ab7 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -1,5 +1,5 @@ import { DataSourceName } from '@server/types'; -import { Manifest, TableReference } from '@server/mdl/type'; +import { Manifest } from '@server/mdl/type'; import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; import { SupportedDataSource, @@ -9,7 +9,6 @@ import { IbisResponse, } from '../adaptors/ibisAdaptor'; import { getLogger } from '@server/utils'; -import { normalizeMssqlSqlForIbis } from '@server/utils/mssqlSqlNormalizer'; import { Project } from '../repositories'; import { PostHogTelemetry, TelemetryEvent } from '../telemetry/telemetry'; @@ -77,390 +76,6 @@ export interface IQueryService { ): Promise; } -const normalizePreviewSqlForIbis = ( - sql: string, - dataSource: DataSourceName, - limit?: number, -): { sql: string; limit?: number } => { - sql = normalizeMssqlSqlForIbis(sql, dataSource); - - if (dataSource !== DataSourceName.MSSQL) { - return { - sql: normalizeNonMssqlGeneratedSqlSyntax(sql, dataSource), - limit, - }; - } - - const topMatch = sql.match(/^\s*SELECT\s+(DISTINCT\s+)?TOP\s*\(?\s*(\d+)\s*\)?\s+/i); - if (!topMatch) { - return { sql, limit }; - } - - const distinctClause = topMatch[1] || ''; - const topLimit = Number(topMatch[2]); - const normalizedSql = sql.replace( - /^\s*SELECT\s+(DISTINCT\s+)?TOP\s*\(?\s*\d+\s*\)?\s+/i, - `SELECT ${distinctClause}`, - ); - - return { - sql: normalizedSql, - limit: - limit && limit > 0 ? Math.min(limit, topLimit) : topLimit, - }; -}; - -const normalizeDeployedManifestForDatasource = ( - manifest: Manifest, - project: Project, -): Manifest => { - if (project.type === DataSourceName.MSSQL || !manifest?.models?.length) { - return manifest; - } - - const fallbackCatalog = manifest.catalog || project.catalog || null; - const fallbackSchema = manifest.schema || project.schema || null; - - return { - ...manifest, - models: manifest.models.map((model) => { - const tableReferenceResult = normalizeTableReference( - model.tableReference, - fallbackSchema, - ); - const synthesizedTableReference = - tableReferenceResult.tableReference || - buildTableReferenceFromDboModelName( - model.name, - fallbackCatalog, - fallbackSchema, - ) || - buildTableReferenceFromDboRefSql( - model.refSql, - fallbackCatalog, - fallbackSchema, - ); - - if (!synthesizedTableReference) { - return model; - } - - const normalizedModel = { - ...model, - tableReference: synthesizedTableReference, - }; - - if ( - tableReferenceResult.changed || - isDboPrefixedModelName(model.name) || - containsDboPhysicalReference(model.refSql) - ) { - delete normalizedModel.refSql; - } - - return normalizedModel; - }), - }; -}; - -const normalizeNonMssqlGeneratedSqlSyntax = ( - sql: string, - dataSource: DataSourceName, -): string => { - if (dataSource === DataSourceName.MSSQL) { - return sql; - } - - return rewriteGeneratedDateDiff(sql); -}; - -const rewriteGeneratedDateDiff = (sql: string): string => { - const dateDiffPattern = - /\bdate_?diff\s*\(\s*'?([A-Za-z]+)'?\s*,\s*([^,()]+(?:\([^)]*\))?[^,()]*)\s*,\s*([^()]+(?:\([^)]*\))?[^()]*)\)/gi; - - return sql.replace( - dateDiffPattern, - (_match, unit: string, startExpression: string, endExpression: string) => { - const normalizedUnit = unit.toLowerCase(); - const start = startExpression.trim(); - const end = endExpression.trim(); - - if (['day', 'dd', 'd'].includes(normalizedUnit)) { - return `EXTRACT(DAY FROM (${end} - ${start}))`; - } - - if (['month', 'mm', 'm'].includes(normalizedUnit)) { - return `((EXTRACT(YEAR FROM ${end}) - EXTRACT(YEAR FROM ${start})) * 12 + (EXTRACT(MONTH FROM ${end}) - EXTRACT(MONTH FROM ${start})))`; - } - - if (['year', 'yy', 'yyyy'].includes(normalizedUnit)) { - return `(EXTRACT(YEAR FROM ${end}) - EXTRACT(YEAR FROM ${start}))`; - } - - return `EXTRACT(DAY FROM (${end} - ${start}))`; - }, - ); -}; - -const normalizeTableReference = ( - tableReference: TableReference | undefined, - fallbackSchema: string | null, -): { tableReference?: TableReference; changed: boolean } => { - if (!tableReference?.table) { - return { tableReference, changed: false }; - } - - const normalizedTableName = normalizeDboPrefixedTableName( - tableReference.table, - ); - const shouldReplaceDboSchema = - tableReference.schema?.toLowerCase() === 'dbo' && fallbackSchema; - - if ( - normalizedTableName === tableReference.table && - !shouldReplaceDboSchema - ) { - return { tableReference, changed: false }; - } - - return { - tableReference: { - ...tableReference, - schema: shouldReplaceDboSchema ? fallbackSchema : tableReference.schema, - table: normalizedTableName, - }, - changed: true, - }; -}; - -const normalizeDboPrefixedTableName = (tableName: string): string => { - const match = tableName.match(/^dbo_(.+)$/i); - return match ? match[1] : tableName; -}; - -const buildTableReferenceFromDboModelName = ( - modelName: string | undefined, - fallbackCatalog: string | null, - fallbackSchema: string | null, -): TableReference | undefined => { - if (!modelName || !isDboPrefixedModelName(modelName)) { - return undefined; - } - - return { - catalog: fallbackCatalog, - schema: fallbackSchema, - table: normalizeDboPrefixedTableName(modelName), - }; -}; - -const buildTableReferenceFromDboRefSql = ( - refSql: string | undefined, - fallbackCatalog: string | null, - fallbackSchema: string | null, -): TableReference | undefined => { - if (!refSql || !containsDboPhysicalReference(refSql)) { - return undefined; - } - - const tableName = - extractDboPrefixedTableName(refSql) || extractDboSchemaTableName(refSql); - if (!tableName) { - return undefined; - } - - return { - catalog: fallbackCatalog, - schema: fallbackSchema, - table: normalizeDboPrefixedTableName(tableName), - }; -}; - -const isDboPrefixedModelName = (modelName: string | undefined): boolean => - !!modelName && /^dbo_.+/i.test(modelName); - -const containsDboPhysicalReference = (sql: string | undefined): boolean => - !!sql && /(?:^|[.\s"])(?:dbo_[\w]+|dbo\.[\w"]+)/i.test(sql); - -const extractDboPrefixedTableName = (sql: string): string | undefined => { - const match = sql.match(/\bdbo_([A-Za-z0-9_]+)\b/i); - return match ? `dbo_${match[1]}` : undefined; -}; - -const extractDboSchemaTableName = (sql: string): string | undefined => { - const match = sql.match(/\bdbo\.("?)([A-Za-z0-9_]+)\1/i); - return match ? match[2] : undefined; -}; - -const SQL_IDENTIFIER_PATTERN = - String.raw`(?:"[^"]+"|` + - '`[^`]+`' + - String.raw`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)`; - -const normalizeSqlIdentifier = (identifier: string) => { - const trimmed = identifier.trim(); - if ( - (trimmed.startsWith('"') && trimmed.endsWith('"')) || - (trimmed.startsWith('`') && trimmed.endsWith('`')) || - (trimmed.startsWith('[') && trimmed.endsWith(']')) - ) { - return trimmed.slice(1, -1); - } - return trimmed; -}; - -const splitTableReference = (tableReference: string) => { - const trimmed = tableReference.trim(); - if (!trimmed) { - return []; - } - - const isMultipartQuotedReference = - /"\s*\.\s*"|\]\s*\.\s*\[|`\s*\.\s*`/.test(trimmed); - if ( - !isMultipartQuotedReference && - ((trimmed.startsWith('"') && trimmed.endsWith('"')) || - (trimmed.startsWith('`') && trimmed.endsWith('`')) || - (trimmed.startsWith('[') && trimmed.endsWith(']'))) - ) { - const normalized = normalizeSqlIdentifier(trimmed); - if (normalized.includes('.')) { - return normalized.split(/\s*\.\s*/).filter(Boolean); - } - } - - return trimmed - .split(/\s*\.\s*/) - .map(normalizeSqlIdentifier) - .filter(Boolean); -}; - -const quoteSqlIdentifier = (identifier: string) => - `"${identifier.replace(/"/g, '""')}"`; - -const quoteTableReference = (tableReference: string) => - splitTableReference(tableReference).map(quoteSqlIdentifier).join('.'); - -const escapeRegExp = (value: string) => - value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); - -const extractSqlTableReferences = (sql: string) => { - const references: string[] = []; - const tablePattern = new RegExp( - String.raw`\b(?:FROM|JOIN)\s+(${SQL_IDENTIFIER_PATTERN}(?:\s*\.\s*${SQL_IDENTIFIER_PATTERN})*)`, - 'gi', - ); - let match: RegExpExecArray | null; - while ((match = tablePattern.exec(sql))) { - references.push(splitTableReference(match[1]).join('.')); - } - return references; -}; - -const addModelReferenceAlias = ( - aliases: Map, - parts: Array, - modelName?: string, -) => { - if (!modelName) { - return; - } - - const normalizedParts = parts - .filter((part): part is string => Boolean(part)) - .map((part) => part.toLowerCase()); - if (!normalizedParts.length) { - return; - } - - for (let index = 0; index < normalizedParts.length; index += 1) { - aliases.set(normalizedParts.slice(index).join('.'), modelName); - } -}; - -const getManifestTableReferenceAliases = (manifest?: Manifest) => { - const aliases = new Map(); - for (const model of manifest?.models || []) { - if (!model.name) { - continue; - } - - aliases.set(model.name.toLowerCase(), model.name); - const dboModelMatch = model.name.match(/^dbo_(.+)$/i); - if (dboModelMatch?.[1]) { - aliases.set(`dbo.${dboModelMatch[1]}`.toLowerCase(), model.name); - } - const basePatternMatch = model.name.match(/^(.+)_patterns$/i); - if (basePatternMatch?.[1]) { - aliases.set(basePatternMatch[1].toLowerCase(), model.name); - const dboBasePatternMatch = basePatternMatch[1].match(/^dbo_(.+)$/i); - if (dboBasePatternMatch?.[1]) { - aliases.set(`dbo.${dboBasePatternMatch[1]}`.toLowerCase(), model.name); - } - } - if (model.tableReference?.table) { - addModelReferenceAlias( - aliases, - [ - model.tableReference.catalog, - model.tableReference.schema, - model.tableReference.table, - ], - model.name, - ); - } - } - return aliases; -}; - -const tableReferencePatternFor = (tableReference: string) => { - const parts = splitTableReference(tableReference); - if (!parts.length) { - return undefined; - } - - const multipartQuoted = parts - .map(quoteSqlIdentifier) - .map(escapeRegExp) - .join(String.raw`\s*\.\s*`); - const bare = parts.map(escapeRegExp).join(String.raw`\s*\.\s*`); - const singleQuoted = escapeRegExp(quoteSqlIdentifier(tableReference)); - const bracketed = escapeRegExp(`[${tableReference}]`); - const backticked = escapeRegExp(`\`${tableReference}\``); - return new RegExp( - String.raw`(? { - const aliases = getManifestTableReferenceAliases(manifest); - if (!aliases.size) { - return sql; - } - - let normalizedSql = sql; - const replacements = new Map(); - for (const reference of extractSqlTableReferences(sql)) { - const canonicalName = aliases.get(reference.toLowerCase()); - if (canonicalName && canonicalName.toLowerCase() !== reference.toLowerCase()) { - replacements.set(reference, canonicalName); - } - } - - for (const [reference, canonicalName] of [...replacements.entries()].sort( - ([left], [right]) => right.length - left.length, - )) { - const pattern = tableReferencePatternFor(reference); - if (!pattern) { - continue; - } - normalizedSql = normalizedSql.replace(pattern, quoteTableReference(canonicalName)); - } - - return normalizedSql; -}; - export class QueryService implements IQueryService { private readonly ibisAdaptor: IIbisAdaptor; private readonly wrenEngineAdaptor: IWrenEngineAdaptor; @@ -486,34 +101,27 @@ export class QueryService implements IQueryService { ): Promise { const { project, - manifest: rawMdl, + manifest: mdl, limit, dryRun, refresh, cacheEnabled, } = options; - const mdl = normalizeDeployedManifestForDatasource(rawMdl, project); const { type: dataSource, connectionInfo } = project; - const manifestNormalizedSql = normalizeSqlReferencesToManifest(sql, mdl); - const normalizedPreview = normalizePreviewSqlForIbis( - manifestNormalizedSql, - dataSource, - limit, - ); if (this.useEngine(dataSource)) { if (dryRun) { logger.debug('Using wren engine to dry run'); - await this.wrenEngineAdaptor.dryRun(normalizedPreview.sql, { + await this.wrenEngineAdaptor.dryRun(sql, { manifest: mdl, - limit: normalizedPreview.limit, + limit, }); return true; } else { logger.debug('Using wren engine to preview'); const data = await this.wrenEngineAdaptor.previewData( - normalizedPreview.sql, + sql, mdl, - normalizedPreview.limit, + limit, ); return data as PreviewDataResponse; } @@ -522,18 +130,18 @@ export class QueryService implements IQueryService { logger.debug('Use ibis adaptor to preview'); if (dryRun) { return await this.ibisDryRun( - normalizedPreview.sql, + sql, dataSource, connectionInfo, mdl, ); } else { return await this.ibisQuery( - normalizedPreview.sql, + sql, dataSource, connectionInfo, mdl, - normalizedPreview.limit, + limit, refresh, cacheEnabled, ); @@ -563,12 +171,11 @@ export class QueryService implements IQueryService { parameters: Record, ): Promise { const { type: dataSource, connectionInfo } = project; - const mdl = normalizeDeployedManifestForDatasource(manifest, project); const res = await this.ibisAdaptor.validate( dataSource, rule, connectionInfo, - mdl, + manifest, parameters, ); return res; @@ -596,22 +203,21 @@ export class QueryService implements IQueryService { connectionInfo: any, mdl: Manifest, ): Promise { - const normalizedQuery = normalizePreviewSqlForIbis(sql, dataSource).sql; const event = TelemetryEvent.IBIS_DRY_RUN; try { - const res = await this.ibisAdaptor.dryRun(normalizedQuery, { + const res = await this.ibisAdaptor.dryRun(sql, { dataSource, connectionInfo, mdl, }); - this.sendIbisEvent(event, res, { dataSource, sql: normalizedQuery }); + this.sendIbisEvent(event, res, { dataSource, sql }); return { correlationId: res.correlationId, }; } catch (err: any) { this.sendIbisFailedEvent(event, err, { dataSource, - sql: normalizedQuery, + sql, }); throw err; } @@ -626,20 +232,19 @@ export class QueryService implements IQueryService { refresh?: boolean, cacheEnabled?: boolean, ): Promise { - const normalizedPreview = normalizePreviewSqlForIbis(sql, dataSource, limit); const event = TelemetryEvent.IBIS_QUERY; try { - const res = await this.ibisAdaptor.query(normalizedPreview.sql, { + const res = await this.ibisAdaptor.query(sql, { dataSource, connectionInfo, mdl, - limit: normalizedPreview.limit, + limit, refresh, cacheEnabled, }); this.sendIbisEvent(event, res, { dataSource, - sql: normalizedPreview.sql, + sql, }); const data = this.transformDataType(res); return { diff --git a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts b/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts deleted file mode 100644 index 1b004f9406..0000000000 --- a/wren-ui/src/apollo/server/utils/mssqlSqlNormalizer.ts +++ /dev/null @@ -1,562 +0,0 @@ -import { DataSourceName } from '@server/types'; - -const escapeRegex = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); - -const formatTimestampLiteral = (date: Date) => { - const pad = (value: number) => String(value).padStart(2, '0'); - return `'${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())} ${pad(date.getHours())}:${pad(date.getMinutes())}:${pad(date.getSeconds())}'`; -}; - -const addMonths = (date: Date, months: number) => { - const next = new Date(date); - next.setMonth(next.getMonth() + months); - return next; -}; - -const replaceRelativeCurrentDateCalls = (sql: string): string => { - const now = new Date(); - const relativeLiteral = (unit: string, amount: number) => { - const normalizedUnit = unit.toLowerCase(); - if (normalizedUnit.startsWith('month')) { - return formatTimestampLiteral(addMonths(now, -amount)); - } - if (normalizedUnit.startsWith('year')) { - return formatTimestampLiteral(addMonths(now, -amount * 12)); - } - if (normalizedUnit.startsWith('day')) { - const next = new Date(now); - next.setDate(next.getDate() - amount); - return formatTimestampLiteral(next); - } - return null; - }; - - sql = sql.replace( - /\bDATE_SUB\(\s*CURRENT_DATE(?:\(\))?\s*,\s*INTERVAL\s+(\d+)\s+(YEAR|MONTH|DAY)S?\s*\)/gi, - (match, amount, unit) => relativeLiteral(unit, Number(amount)) || match, - ); - sql = sql.replace( - /\bDATE_SUB\(\s*'?(YEAR|MONTH|DAY)'?\s*,\s*(\d+)\s*,\s*CURRENT_DATE(?:\(\))?\s*\)/gi, - (match, unit, amount) => relativeLiteral(unit, Number(amount)) || match, - ); - return sql.replace(/\bCURRENT_DATE(?:\(\))?\b/gi, formatTimestampLiteral(now)); -}; - -const inferMssqlTimestampExpression = (sql: string): string => { - const qualifiedTimestamp = sql.match( - /"([^"]+)"\."(created_at|updated_at|generated_at|created_date|date|DateIn|DateOut|FailedAt)"/i, - ); - if (qualifiedTimestamp) { - return qualifiedTimestamp[0]; - } - - return ''; -}; - -const replaceInventedDateFields = (sql: string): string => { - const timestampExpression = inferMssqlTimestampExpression(sql); - if (!timestampExpression) { - return sql; - } - const inventedDateFields = [ - 'RepairDate', - 'repairDate', - 'repair_date', - 'Repair_Date', - 'EventDate', - 'event_date', - 'last_update_date', - 'last_updated_date', - 'last_updated_at', - 'lastUpdateDate', - 'lastUpdatedDate', - 'updated_date', - 'updated_at', - 'Date', - 'date', - ]; - - inventedDateFields.forEach((field) => { - const escaped = escapeRegex(field); - sql = sql.replace( - new RegExp(String.raw`(?:"[^"]+"\.)"${escaped}"`, 'gi'), - timestampExpression, - ); - sql = sql.replace(new RegExp(String.raw`"${escaped}"`, 'gi'), timestampExpression); - sql = sql.replace(new RegExp(String.raw`\[${escaped}\]`, 'gi'), timestampExpression); - if (field.toLowerCase() !== 'date') { - sql = sql.replace( - new RegExp(String.raw`(? - sql.replace( - /\bDATEPART\(\s*'?\s*(YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\)/gi, - (_match, part, expression) => - `EXTRACT(${String(part).toUpperCase()} FROM ${String(expression).trim()})`, - ); - -const replaceInventedTimeBuckets = (sql: string): string => { - const timestampExpression = inferMssqlTimestampExpression(sql); - if (!timestampExpression) { - return sql; - } - const bucketExpressions: Record = { - YEAR: `EXTRACT(YEAR FROM ${timestampExpression})`, - MONTH: `EXTRACT(MONTH FROM ${timestampExpression})`, - DAY: `EXTRACT(DAY FROM ${timestampExpression})`, - }; - - sql = sql.replace(/\bSELECT\b(?.*?)(?=\bFROM\b)/is, (match, _body, _offset, _source, groups) => { - let body = groups?.body || ''; - Object.entries(bucketExpressions).forEach(([bucket, expression]) => { - const alias = bucket.toLowerCase(); - body = body.replace( - new RegExp( - String.raw`(^|,)\s*(?:(?:"[^"]+"\.)"?${bucket}"?|(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})|\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b|"${bucket}"|\[${bucket}\]|\b${bucket}\b)(?:\s+(?:AS\s+)?(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*)))?(?=\s*(?:,|$))`, - 'gi', - ), - (_match, prefix, quotedAlias, bracketAlias, bareAlias) => { - const selectedAlias = quotedAlias - ? `"${quotedAlias}"` - : bracketAlias - ? `[${bracketAlias}]` - : bareAlias || `"${alias}"`; - return `${prefix} ${expression} AS ${selectedAlias}`; - }, - ); - }); - return `SELECT${body}`; - }); - - Object.entries(bucketExpressions).forEach(([bucket, expression]) => { - sql = sql.replace( - new RegExp(String.raw`(?:"[^"]+"\.)"?${bucket}"?`, 'gi'), - expression, - ); - sql = sql.replace( - new RegExp(String.raw`(?:\[[^\]]+\]\.)(?:\[${bucket}\]|${bucket})`, 'gi'), - expression, - ); - sql = sql.replace( - new RegExp(String.raw`\b[A-Za-z_][A-Za-z0-9_]*\.${bucket}\b`, 'gi'), - expression, - ); - }); - - const clausePattern = - /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; - sql = sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { - let body = groups?.body || ''; - Object.entries(bucketExpressions).forEach(([bucket, expression]) => { - body = body.replace(new RegExp(String.raw`"${bucket}"`, 'gi'), expression); - body = body.replace(new RegExp(String.raw`\[${bucket}\]`, 'gi'), expression); - body = body.replace( - new RegExp(String.raw`(? { - const limitMatch = sql.match(/\s+LIMIT\s+(\d+)\s*;?\s*$/i); - if (!limitMatch || limitMatch.index === undefined) { - return sql; - } - - const limit = limitMatch[1]; - const withoutLimit = sql.slice(0, limitMatch.index).trimEnd(); - if (/\bSELECT\s+(?:DISTINCT\s+)?TOP\s+(?:\(\s*)?\d+/i.test(withoutLimit)) { - return withoutLimit; - } - - if (/^\s*SELECT\s+DISTINCT\b/i.test(withoutLimit)) { - return withoutLimit.replace(/\bSELECT\s+DISTINCT\b/i, `SELECT DISTINCT TOP ${limit}`); - } - - if (/^\s*SELECT\b/i.test(withoutLimit)) { - return withoutLimit.replace(/\bSELECT\b/i, `SELECT TOP ${limit}`); - } - - return withoutLimit; -}; - -const unwrapSimpleMssqlWhereParentheses = (sql: string): string => - sql.replace( - /\bWHERE\s*\(\s*([^()]+?)\s*\)(?=\s*(?:GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|$))/gis, - 'WHERE $1', - ); - -const normalizeMssqlGeneratedSqlSyntax = (sql: string): string => { - sql = sql.replace(/\s+NULLS\s+(?:LAST|FIRST)\b/gi, ''); - sql = unwrapSimpleMssqlWhereParentheses(sql); - return rewriteMssqlLimitClause(sql); -}; - -const quoteMssqlDboModelReferences = (sql: string): string => { - sql = sql.replace( - /(?:"[^"]+"\.){1,2}"(dbo_[A-Za-z0-9_]+)"/gi, - '"$1"', - ); - sql = sql.replace( - /\b[A-Za-z_][A-Za-z0-9_]*\.[A-Za-z_][A-Za-z0-9_]*\.(dbo_[A-Za-z0-9_]+)\b/g, - '"$1"', - ); - sql = sql.replace( - /\[[^\]]+\]\.\[[^\]]+\]\.\[(dbo_[A-Za-z0-9_]+)\]/gi, - '"$1"', - ); - sql = sql.replace(/\bdbo\.([A-Za-z0-9_]+)\b/g, '"dbo_$1"'); - - const quotedModels = new Set(); - sql.replace(/"dbo_[A-Za-z0-9_]+"/gi, (match) => { - quotedModels.add(match.slice(1, -1)); - return match; - }); - - const bareDboModel = /\bdbo_[A-Za-z0-9_]+\b/g; - return sql.replace(bareDboModel, (modelName, offset, source) => { - if (source[offset - 1] === '"' || source[offset + modelName.length] === '"') { - return modelName; - } - - if (!quotedModels.has(modelName)) { - quotedModels.add(modelName); - } - - return `"${modelName}"`; - }); -}; - -const replaceBadFailurePatternJoins = (sql: string): string => { - if ( - !/\bdbo_DebugEntries\b/i.test(sql) || - !/\bdbo_failure_patterns\b/i.test(sql) - ) { - return sql; - } - - const debugTable = String.raw`(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)`; - const failurePatternTable = String.raw`(?:"dbo_failure_patterns"|\[dbo_failure_patterns\]|dbo_failure_patterns)`; - const debugEntryId = String.raw`${debugTable}\s*\.\s*(?:"DebugEntryId"|\[DebugEntryId\]|DebugEntryId)`; - const debugFailureSys = '"dbo_DebugEntries"."FailureSys"'; - const failurePatternId = String.raw`${failurePatternTable}\s*\.\s*(?:"id"|\[id\]|id)`; - const normalizedFailurePatternId = '"dbo_failure_patterns"."id"'; - - sql = sql.replace( - new RegExp(String.raw`${debugEntryId}\s*=\s*${failurePatternId}`, 'gi'), - `${debugFailureSys} = ${normalizedFailurePatternId}`, - ); - sql = sql.replace( - new RegExp(String.raw`${failurePatternId}\s*=\s*${debugEntryId}`, 'gi'), - `${normalizedFailurePatternId} = ${debugFailureSys}`, - ); - sql = sql.replace( - new RegExp( - String.raw`${debugTable}\s*\.\s*(?:"FailurePatternID"|"FailurePatternId"|\[FailurePatternID\]|\[FailurePatternId\]|FailurePatternID|FailurePatternId)`, - 'gi', - ), - debugFailureSys, - ); - - return sql; -}; - -const replacePcbThroughputFields = (sql: string): string => { - const manufacturingUnitField = - String.raw`(?:"ManufacturingUnit"|"Manufacturing_Unit"|"manufacturing_unit"|\[ManufacturingUnit\]|\[Manufacturing_Unit\]|\[manufacturing_unit\]|ManufacturingUnit|Manufacturing_Unit|manufacturing_unit)`; - - if (/\bdbo_DebugEntries\b/i.test(sql)) { - const debugTable = String.raw`(?:"dbo_DebugEntries"|\[dbo_DebugEntries\]|dbo_DebugEntries)`; - sql = sql.replace( - new RegExp(String.raw`${debugTable}\s*\.\s*${manufacturingUnitField}`, 'gi'), - '"dbo_DebugEntries"."BusinessUnit"', - ); - } - - if (/\bdbo_repair_logs\b/i.test(sql) && new RegExp(manufacturingUnitField, 'i').test(sql)) { - sql = sql.replace( - /(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)/gi, - '"dbo_DebugEntries"', - ); - sql = sql.replace( - new RegExp(String.raw`"dbo_DebugEntries"\s*\.\s*${manufacturingUnitField}`, 'gi'), - '"dbo_DebugEntries"."BusinessUnit"', - ); - sql = sql.replace( - /"dbo_DebugEntries"\s*\.\s*(?:"id"|\[id\]|id)/gi, - '"dbo_DebugEntries"."DebugEntryId"', - ); - sql = sql.replace( - /"dbo_DebugEntries"\s*\.\s*(?:"created_at"|"updated_at"|\[created_at\]|\[updated_at\]|created_at|updated_at)/gi, - '"dbo_DebugEntries"."DateIn"', - ); - } - - return sql; -}; - -const replaceRepairLogThroughputShape = (sql: string): string => { - if ( - !/\bdbo_repair_logs\b/i.test(sql) || - !/\bavg_turnaround_time\b/i.test(sql) || - !/\b(?:repair_count|throughput)\b/i.test(sql) - ) { - return sql; - } - - return [ - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name",', - 'COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput"', - 'FROM "dbo_DebugEntries"', - 'GROUP BY "dbo_DebugEntries"."BusinessUnit"', - 'ORDER BY "throughput" DESC', - ].join(' '); -}; - -const replaceTicketCycleTurnaroundShape = (sql: string): string => { - if ( - !/\bdbo_ticket_cycles\b/i.test(sql) || - !/\b(?:turnaround_time|avg_turnaround_time)\b/i.test(sql) || - !/\bMONTH\b|DATEPART\(\s*'?\s*MONTH/i.test(sql) - ) { - return sql; - } - - return [ - 'SELECT EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") AS "year",', - 'EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") AS "month",', - 'AVG(DATEDIFF(\'second\', "dbo_ticket_cycles"."start_date", "dbo_ticket_cycles"."end_date")) AS "avg_turnaround_seconds"', - 'FROM "dbo_ticket_cycles"', - 'WHERE "dbo_ticket_cycles"."start_date" IS NOT NULL', - 'AND "dbo_ticket_cycles"."end_date" IS NOT NULL', - 'GROUP BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at"), EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at")', - 'ORDER BY EXTRACT(YEAR FROM "dbo_ticket_cycles"."created_at") ASC, EXTRACT(MONTH FROM "dbo_ticket_cycles"."created_at") ASC', - ].join(' '); -}; - -const replaceInventedFailureCategory = (sql: string): string => { - if (!/\bdbo_repair_logs\b/i.test(sql) || !/\bfailure[_\s]+category\b/i.test(sql)) { - return sql; - } - - const failureCodeExpression = '"dbo_repair_logs"."failure_code"'; - sql = sql.replace( - /\bSELECT\b(?.*?)(?=\bFROM\b)/is, - (match, _body, _offset, _source, groups) => { - let body = groups?.body || ''; - body = body.replace( - /(^|,)\s*(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)(?=\s*(?:,|$))/gi, - `$1 ${failureCodeExpression} AS "failure_category"`, - ); - return `SELECT${body}`; - }, - ); - sql = sql.replace(/\bAS\s+failure\s+category\b/gi, 'AS "failure_category"'); - - const clausePattern = - /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; - return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { - const body = (groups?.body || '').replace( - /(?:(?:"dbo_repair_logs"|\[dbo_repair_logs\]|dbo_repair_logs)\s*\.\s*)?(?:"failure[_\s]+category"|\[failure[_\s]+category\]|failure[_\s]+category)/gi, - failureCodeExpression, - ); - return `${clause}${body}`; - }); -}; - -const replaceInventedReportFields = (sql: string): string => { - if (!/\bdbo_reports\b/i.test(sql)) { - return sql; - } - - const reportTable = String.raw`(?:"dbo_reports"|\[dbo_reports\]|dbo_reports)`; - sql = sql.replace( - new RegExp(String.raw`(?:(?:${reportTable})\s*\.\s*)?(?:"filters"|\[filters\]|\bfilters\b)`, 'gi'), - '"dbo_reports"."data"', - ); - sql = sql.replace( - new RegExp( - String.raw`(?:(?:${reportTable})\s*\.\s*)?(?:"report_size"|"file_size"|\[report_size\]|\[file_size\]|\breport_size\b|\bfile_size\b)`, - 'gi', - ), - '"dbo_reports"."size_bytes"', - ); - return sql; -}; - -const replaceInventedKnowledgeArticleFields = (sql: string): string => { - if (!/\b(?:dbo_knowledge_articles|dbo_kb_articles)\b/i.test(sql)) { - return sql; - } - - const replacementsByTable: Record> = { - dbo_knowledge_articles: { - effectiveness_score: '"helpful"', - created_by: '"author"', - created_by_user: '"author"', - created_by_user_id: '"author"', - author_id: '"author"', - }, - dbo_kb_articles: { - category: '"category"', - section: '"category"', - article_section: '"category"', - created_by: '"created_by_user_id"', - created_by_user: '"created_by_user_id"', - author: '"created_by_user_id"', - author_id: '"created_by_user_id"', - }, - }; - - Object.entries(replacementsByTable).forEach(([tableName, replacements]) => { - const tablePattern = String.raw`(?:"${tableName}"|\[${tableName}\]|${tableName})`; - Object.entries(replacements).forEach(([inventedField, replacementField]) => { - const escapedField = escapeRegex(inventedField); - sql = sql.replace( - new RegExp( - String.raw`(${tablePattern})\s*\.\s*(?:"${escapedField}"|\[${escapedField}\]|\b${escapedField}\b)`, - 'gi', - ), - `$1.${replacementField}`, - ); - }); - }); - - const activeTable = /\bdbo_knowledge_articles\b/i.test(sql) - ? 'dbo_knowledge_articles' - : /\bdbo_kb_articles\b/i.test(sql) - ? 'dbo_kb_articles' - : null; - - if (activeTable) { - Object.entries(replacementsByTable[activeTable]).forEach( - ([inventedField, replacementField]) => { - const escapedField = escapeRegex(inventedField); - sql = sql.replace( - new RegExp(String.raw`(? { - sql = sql.replace(/\\"/g, '"'); - sql = quoteMssqlDboModelReferences(sql); - - if (dataSource !== DataSourceName.MSSQL) { - return sql; - } - - sql = normalizeMssqlGeneratedSqlSyntax(sql); - sql = rewriteMssqlDatepartFunctions(sql); - sql = replaceRelativeCurrentDateCalls(sql); - sql = replaceInventedDateFields(sql); - sql = replaceInventedTimeBuckets(sql); - sql = rewriteMssqlDatepartFunctions(sql); - sql = quoteMssqlDboModelReferences(sql); - return normalizeMssqlGeneratedSqlSyntax(sql); -}; - -export const rewriteMssqlDatepartAliasReferences = ( - sql: string, - dataSource: DataSourceName, -): string => { - if (dataSource !== DataSourceName.MSSQL) { - return sql; - } - - sql = sql.replace(/\\"/g, '"'); - - const aliases: Record = {}; - const aliasTargetPattern = - String.raw`(?:"([^"]+)"|\[([^\]]+)\]|([A-Za-z_][A-Za-z0-9_]*))`; - const aliasPatterns = [ - new RegExp( - String.raw`\b(DATEPART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b((?:YEAR|MONTH|DAY)\(\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b(DATE_PART\(\s*'?\s*(?:YEAR|MONTH|DAY)\s*'?\s*,\s*((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - new RegExp( - String.raw`\b(EXTRACT\(\s*(?:YEAR|MONTH|DAY)\s+FROM\s+((?:[^()]|\([^()]*\))+?)\s*\))\s+AS\s+${aliasTargetPattern}`, - 'gi', - ), - ]; - - aliasPatterns.forEach((aliasPattern) => { - for (const match of sql.matchAll(aliasPattern)) { - const expression = match[1]; - const alias = match[3] || match[4] || match[5]; - aliases[alias.toLowerCase()] = expression; - } - }); - - if (!Object.keys(aliases).length) { - return sql; - } - - const clausePattern = - /\b(GROUP\s+BY|ORDER\s+BY|HAVING)\b(?.*?)(?=\b(?:ORDER\s+BY|GROUP\s+BY|HAVING|LIMIT|OFFSET|FETCH|UNION|WHERE)\b|$)/gis; - - return sql.replace(clausePattern, (match, clause, _body, _offset, _source, groups) => { - let body = groups?.body || ''; - const placeholders: Record = {}; - - Object.entries(aliases).forEach(([alias, expression]) => { - const placeholder = `__WREN_MSSQL_DATEPART_ALIAS_${Object.keys(placeholders).length}__`; - placeholders[placeholder] = expression; - const escapedAlias = escapeRegex(alias); - - body = body.replace(new RegExp(`"${escapedAlias}"`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\\\+"${escapedAlias}\\\\+"`, 'gi'), placeholder); - body = body.replace(new RegExp(`\\[${escapedAlias}\\]`, 'gi'), placeholder); - }); - - Object.entries(placeholders).forEach(([placeholder, expression]) => { - body = body.replaceAll(placeholder, expression); - }); - - return `${clause}${body}`; - }); -}; - -export const normalizeMssqlSqlForIbis = ( - sql: string, - dataSource: DataSourceName, -): string => { - sql = normalizeMssqlGeneratedSqlFields(sql, dataSource); - sql = rewriteMssqlDatepartAliasReferences(sql, dataSource); - return normalizeMssqlGeneratedSqlFields(sql, dataSource); -}; diff --git a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts index 691db14ec7..b2db036db5 100644 --- a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts +++ b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts @@ -12,22 +12,15 @@ const quoteIdentifier = (identifier: unknown) => { }; const isMetricColumn = (column: Partial) => { - const name = String(column.name || '').toLowerCase(); const type = String(column.type || '').toLowerCase(); - return ( - /int|float|double|decimal|numeric|number|real|money/.test(type) || - /amount|value|total|count|qty|quantity|price|cost|revenue|sales|margin|rate/.test( - name, - ) + return /int|bigint|smallint|tinyint|float|double|decimal|numeric|number|real|money/.test( + type, ); }; const isDateColumn = (column: Partial) => { - const name = String(column.name || '').toLowerCase(); const type = String(column.type || '').toLowerCase(); - return ( - /date|time|timestamp/.test(type) || /date|time|created|updated/.test(name) - ); + return /date|time|timestamp/.test(type); }; const isDimensionColumn = (column: Partial) => { diff --git a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts b/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts deleted file mode 100644 index fc5990f17e..0000000000 --- a/wren-ui/src/apollo/server/utils/tests/mssqlSqlNormalizer.test.ts +++ /dev/null @@ -1,258 +0,0 @@ -import { DataSourceName } from '../../types'; -import { normalizeMssqlSqlForIbis } from '../mssqlSqlNormalizer'; - -describe('mssqlSqlNormalizer', () => { - it('rewrites CWSales OTD date aliases before MSSQL preview execution', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT - DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date") AS "year", - DATEPART(MONTH, "dbo_tblSalesHistory"."OTD_Date") AS "month", - "dbo_tblSalesHistory"."MarketType" AS "MarketType", - SUM("dbo_tblSalesHistory"."Qty") AS "TotalQty" - FROM "dbo_tblSalesHistory" - GROUP BY - DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date"), - DATEPART(MONTH, "dbo_tblSalesHistory"."OTD_Date"), - "dbo_tblSalesHistory"."MarketType" - ORDER BY - DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date"), - DATEPART(MONTH, "dbo_tblSalesHistory"."OTD_Date") - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('OTD_Date'); - expect(normalized).toContain('"dbo_tblSalesHistory"."InvDate"'); - }); - - it('rewrites CWSales FixLogId aliases before MSSQL preview execution', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT - "SalesPerson", - Country, - COUNT("dbo_qSales1"."FixLogId") AS NumberOfInvoices - FROM "dbo_qSales1" - GROUP BY "SalesPerson", Country - ORDER BY NumberOfInvoices DESC - LIMIT 1 - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('FixLogId'); - expect(normalized).toContain('"dbo_qSales1"."InvoiceNo"'); - }); - - it('does not rewrite CWSales aliases for non-MSSQL datasources', () => { - const sql = 'SELECT "dbo_tblSalesHistory"."OTD_Date" FROM "dbo_tblSalesHistory"'; - - expect(normalizeMssqlSqlForIbis(sql, DataSourceName.POSTGRES)).toBe(sql); - }); - - it('rewrites aliased repair log time buckets', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT YEAR AS year, MONTH AS month, COUNT(*) AS repair_count - FROM dbo_repair_logs - GROUP BY YEAR, MONTH - ORDER BY YEAR ASC, MONTH ASC - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('YEAR AS year'); - expect(normalized).not.toContain('MONTH AS month'); - expect(normalized).toContain( - 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS year', - ); - expect(normalized).toContain( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS month', - ); - expect(normalized).toContain( - 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at")', - ); - }); - - it('rewrites quoted debug entry year aliases', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT - "YEAR" AS "YEAR", - "dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit", - COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" - FROM "dbo_DebugEntries" - GROUP BY "YEAR", "dbo_DebugEntries"."BusinessUnit" - ORDER BY "YEAR" ASC - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('"YEAR" AS "YEAR"'); - expect(normalized).not.toContain('GROUP BY "YEAR"'); - expect(normalized).toContain( - 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "YEAR"', - ); - expect(normalized).toContain( - 'GROUP BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn")', - ); - }); - - it('rewrites knowledge article time buckets', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT "YEAR", COUNT("dbo_knowledge_articles"."id") AS "article_count" - FROM "dbo_knowledge_articles" - GROUP BY "YEAR" - ORDER BY "YEAR" ASC - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('SELECT "YEAR"'); - expect(normalized).not.toContain('GROUP BY "YEAR"'); - expect(normalized).toContain( - 'DATEPART(YEAR, "dbo_knowledge_articles"."created_at") AS "year"', - ); - }); - - it('rewrites stale last_update_date references for debug entry trends', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT - DATEPART(YEAR, last_update_date) AS "year", - DATEPART(MONTH, last_update_date) AS "month", - "dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit", - COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" - FROM "dbo_DebugEntries" - GROUP BY - DATEPART(YEAR, last_update_date), - DATEPART(MONTH, last_update_date), - "dbo_DebugEntries"."BusinessUnit" - ORDER BY - DATEPART(YEAR, last_update_date), - DATEPART(MONTH, last_update_date) - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('last_update_date'); - expect(normalized).toContain('DATEPART(YEAR, "dbo_DebugEntries"."DateIn")'); - expect(normalized).toContain('DATEPART(MONTH, "dbo_DebugEntries"."DateIn")'); - }); - - it('rewrites hallucinated knowledge article fields', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT - AVG("dbo_knowledge_articles"."effectiveness_score") AS "avg_effectiveness", - "dbo_knowledge_articles"."created_by" AS "created_by" - FROM "dbo_knowledge_articles" - GROUP BY "dbo_knowledge_articles"."created_by" - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('effectiveness_score'); - expect(normalized).not.toContain('"dbo_knowledge_articles"."created_by"'); - expect(normalized).toContain( - 'AVG("dbo_knowledge_articles"."helpful") AS "avg_effectiveness"', - ); - expect(normalized).toContain('"dbo_knowledge_articles"."author" AS "author"'); - expect(normalized).toContain('GROUP BY "dbo_knowledge_articles"."author"'); - }); - - it('rewrites hallucinated kb article creator fields', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT created_by, COUNT(*) AS article_count - FROM dbo_kb_articles - GROUP BY created_by - ORDER BY article_count DESC - `, - DataSourceName.MSSQL, - ); - - expect(normalized).not.toContain('created_by,'); - expect(normalized).not.toContain('GROUP BY created_by'); - expect(normalized).toContain('"created_by_user_id"'); - }); - - it('keeps quoted dbo-prefixed model names for ibis model resolution', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT "dbo_search_queries"."org_id", COUNT(*) AS completed_questions - FROM "dbo_search_queries" - WHERE "dbo_search_queries"."result_count" > 0 - GROUP BY "dbo_search_queries"."org_id" - `, - DataSourceName.MSSQL, - ); - - expect(normalized).toContain('FROM "dbo_search_queries"'); - expect(normalized).toContain('"dbo_search_queries"."org_id"'); - expect(normalized).not.toContain('FROM dbo_search_queries'); - }); - - it('quotes unquoted dbo-prefixed model names for ibis model resolution', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT status, COUNT(*) AS count_of_questions - FROM dbo_tickets - GROUP BY status - `, - DataSourceName.MSSQL, - ); - - expect(normalized).toContain('FROM "dbo_tickets"'); - expect(normalized).not.toContain('FROM dbo_tickets'); - }); - - it('quotes dbo-prefixed model names for non-MSSQL project contexts', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT org_id, COUNT(*) AS num_questions - FROM dbo_search_queries - GROUP BY org_id - `, - DataSourceName.POSTGRES, - ); - - expect(normalized).toContain('FROM "dbo_search_queries"'); - expect(normalized).not.toContain('FROM dbo_search_queries'); - }); - - it('collapses fully qualified dbo-prefixed model names for switched projects', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT wrenai.public.dbo_search_queries.org_id, COUNT(*) AS num_questions - FROM wrenai.public.dbo_search_queries - GROUP BY wrenai.public.dbo_search_queries.org_id - `, - DataSourceName.POSTGRES, - ); - - expect(normalized).toContain('FROM "dbo_search_queries"'); - expect(normalized).toContain('"dbo_search_queries".org_id'); - expect(normalized).not.toContain('wrenai.public.dbo_search_queries'); - }); - - it('normalizes schema-qualified dbo references back to model names', () => { - const normalized = normalizeMssqlSqlForIbis( - ` - SELECT dbo.search_queries.org_id, COUNT(*) AS completed_questions - FROM dbo.search_queries - INNER JOIN dbo.organizations ON dbo.search_queries.org_id = dbo.organizations.id - GROUP BY dbo.organizations.name - `, - DataSourceName.MSSQL, - ); - - expect(normalized).toContain('FROM "dbo_search_queries"'); - expect(normalized).toContain('INNER JOIN "dbo_organizations"'); - expect(normalized).toContain('"dbo_search_queries".org_id'); - expect(normalized).not.toContain('dbo.search_queries'); - expect(normalized).not.toContain('dbo.organizations'); - }); -}); From 7190bc6e54c540c1943dba29c2917e6c37976f83 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 19:22:55 +0530 Subject: [PATCH 0654/1087] Fix ask routing and Wren SQL previews --- .../generation/intent_classification.py | 47 ++++++++++++++++++- .../apollo/server/services/queryService.ts | 2 +- .../server/utils/recommendationQuestions.ts | 2 +- 3 files changed, 47 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 90cbba6310..4dda8d63aa 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,5 +1,6 @@ import ast import logging +import re import sys from typing import Any, Literal, Optional @@ -302,12 +303,54 @@ async def classify_intent(prompt: dict, generator: Any, generator_name: str) -> @observe(capture_input=False) -def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict: +def post_process( + query: str, classify_intent: dict, construct_db_schemas: list[str] +) -> dict: try: results = orjson.loads(classify_intent.get("replies")[0]) + intent = results["results"] + query_text = " ".join( + [ + str(query or ""), + str(results.get("rephrased_question") or ""), + ] + ).lower() + data_request_patterns = ( + r"\bshow\b", + r"\blist\b", + r"\bdisplay\b", + r"\bfetch\b", + r"\bfind\b", + r"\bcompare\b", + r"\bcount\b", + r"\bsum\b", + r"\btotal\b", + r"\baverage\b", + r"\bavg\b", + r"\bmin\b", + r"\bmax\b", + r"\btop\b", + r"\bbottom\b", + r"\brecent\b", + r"\blatest\b", + r"\bfilter\b", + r"\bgroup\b", + r"\border(?:s|ed|ing)?\b", + r"\bsort\b", + r"\brank\b", + r"\btrend\b", + ) + if ( + intent in {"GENERAL", "USER_GUIDE"} + and construct_db_schemas + and any( + re.search(pattern, query_text) for pattern in data_request_patterns + ) + ): + intent = "TEXT_TO_SQL" return { "rephrased_question": results["rephrased_question"], - "intent": results["results"], + "intent": intent, "reasoning": results["reasoning"], "db_schemas": construct_db_schemas, } diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 62434e4ab7..ee7b341fdd 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -258,7 +258,7 @@ export class QueryService implements IQueryService { } catch (err: any) { this.sendIbisFailedEvent(event, err, { dataSource, - sql: normalizedPreview.sql, + sql, }); throw err; } diff --git a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts index b2db036db5..4b145ed362 100644 --- a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts +++ b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts @@ -140,7 +140,7 @@ export const buildFastRecommendationQuestions = ( addQuestion({ category: label, question: `Show the first 10 rows from ${label}.`, - sql: `SELECT TOP 10 ${previewColumns.join(', ')} FROM ${modelRef}`, + sql: `SELECT ${previewColumns.join(', ')} FROM ${modelRef} LIMIT 10`, }); } From de4f7b20f4a9f1044c1fcc16d0cdfbc0e3252c06 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 19:24:33 +0530 Subject: [PATCH 0655/1087] Remove ask intent heuristic override --- .../generation/intent_classification.py | 47 +------------------ 1 file changed, 2 insertions(+), 45 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4dda8d63aa..90cbba6310 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,6 +1,5 @@ import ast import logging -import re import sys from typing import Any, Literal, Optional @@ -303,54 +302,12 @@ async def classify_intent(prompt: dict, generator: Any, generator_name: str) -> @observe(capture_input=False) -def post_process( - query: str, classify_intent: dict, construct_db_schemas: list[str] -) -> dict: +def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict: try: results = orjson.loads(classify_intent.get("replies")[0]) - intent = results["results"] - query_text = " ".join( - [ - str(query or ""), - str(results.get("rephrased_question") or ""), - ] - ).lower() - data_request_patterns = ( - r"\bshow\b", - r"\blist\b", - r"\bdisplay\b", - r"\bfetch\b", - r"\bfind\b", - r"\bcompare\b", - r"\bcount\b", - r"\bsum\b", - r"\btotal\b", - r"\baverage\b", - r"\bavg\b", - r"\bmin\b", - r"\bmax\b", - r"\btop\b", - r"\bbottom\b", - r"\brecent\b", - r"\blatest\b", - r"\bfilter\b", - r"\bgroup\b", - r"\border(?:s|ed|ing)?\b", - r"\bsort\b", - r"\brank\b", - r"\btrend\b", - ) - if ( - intent in {"GENERAL", "USER_GUIDE"} - and construct_db_schemas - and any( - re.search(pattern, query_text) for pattern in data_request_patterns - ) - ): - intent = "TEXT_TO_SQL" return { "rephrased_question": results["rephrased_question"], - "intent": intent, + "intent": results["results"], "reasoning": results["reasoning"], "db_schemas": construct_db_schemas, } From a0070cae614feba8c9ff7b80cf1e824948014153 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 19:45:23 +0530 Subject: [PATCH 0656/1087] Route asks through SQL generation by default --- deployment/kustomizations/base/cm.yaml | 2 +- docker/config.example.yaml | 4 +-- wren-ai-service/src/config.py | 2 +- .../generation/intent_classification.py | 25 ++++++++----------- .../tools/config/config.example.yaml | 2 +- wren-ai-service/tools/config/config.full.yaml | 2 +- 6 files changed, 16 insertions(+), 21 deletions(-) diff --git a/deployment/kustomizations/base/cm.yaml b/deployment/kustomizations/base/cm.yaml index 0bbe6d5e1b..9f7a5bef25 100644 --- a/deployment/kustomizations/base/cm.yaml +++ b/deployment/kustomizations/base/cm.yaml @@ -222,7 +222,7 @@ data: column_indexing_batch_size: 50 table_retrieval_size: 10 table_column_retrieval_size: 100 - allow_intent_classification: true + allow_intent_classification: false allow_sql_generation_reasoning: true allow_sql_functions_retrieval: true enable_column_pruning: false diff --git a/docker/config.example.yaml b/docker/config.example.yaml index 2a22e788e9..aa81c74c60 100644 --- a/docker/config.example.yaml +++ b/docker/config.example.yaml @@ -175,7 +175,7 @@ settings: column_indexing_batch_size: 50 table_retrieval_size: 10 table_column_retrieval_size: 100 - allow_intent_classification: true + allow_intent_classification: false allow_sql_generation_reasoning: true allow_sql_functions_retrieval: true enable_column_pruning: false @@ -190,4 +190,4 @@ settings: 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 + instructions_top_k: 10 diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index c5acf4ae47..a5c6aed24b 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -37,7 +37,7 @@ class Settings(BaseSettings): instructions_top_k: int = Field(default=10) # generation config - allow_intent_classification: bool = Field(default=True) + allow_intent_classification: bool = Field(default=False) allow_sql_generation_reasoning: bool = Field(default=True) allow_sql_functions_retrieval: bool = Field(default=True) allow_sql_diagnosis: bool = Field(default=True) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 90cbba6310..4d6cd313cd 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -33,10 +33,8 @@ - **Rephrase Question:** Rewrite follow-up questions into full standalone questions using prior conversation context. - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. -- **Data Retrieval Requests:** If the user asks to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data from the connected database, classify it as `TEXT_TO_SQL`. -- **Database Schema Exploration:** If the user asks about available tables, columns, relationships, schema meaning, or what can be asked, classify it as `GENERAL`. -- **Out-of-Scope Queries:** If the question is unrelated to the database schema or data retrieval, classify it as `MISLEADING_QUERY`. -- **Incomplete Queries:** If the question references unresolved placeholders (e.g., "the following", "these", "those") without providing them or prior context, classify as `GENERAL`. +- **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. +- **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. ### Intent Definitions ### @@ -45,14 +43,14 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The user's inputs ask to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data. -- The question can be answered by selecting relevant tables and columns from the provided schema, even if the user does not mention exact physical table or column names. -- The question includes enough business meaning, dimensions, metrics, filters, or time criteria to attempt SQL generation from the schema. +- The question (or related previous query) includes references to specific tables, columns, or data details. +- The question includes **complete information** with specific tables, columns, or data values needed for execution. +- The question provides **all necessary parameters** to generate executable SQL. **Requirements:** -- Do not require the user to explicitly name a table or column. -- Use the provided schema context to decide whether the user's business terms can be answered by SQL. -- Reference phrases from the user's inputs that clearly indicate a data retrieval request. +- Must have complete filter criteria, specific values, or clear references to previous context. +- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. +- Reference phrases from the user's inputs that clearly relate to the schema. **Examples:** - "What is the total sales for last quarter?" @@ -63,10 +61,9 @@ **When to Use:** - The user seeks general information about the database schema or its overall capabilities. -- The user asks about available tables, columns, relationships, schema meaning, or what questions can be asked. - The query references **missing information** (e.g., "the following items" without listing them). - The query contains **placeholder references** that cannot be resolved from context. -- The query is asking for explanation or guidance rather than retrieval, filtering, ordering, aggregation, or analysis of rows. +- The query is **incomplete for SQL generation** despite mentioning database concepts. **Requirements:** - Incorporate phrases from the user's inputs that indicate incompleteness or lack of relevance to the database schema. @@ -75,8 +72,6 @@ **Examples:** - "What is the dataset about?" - "Tell me more about the database." -- "Explain the customer table to me." -- "What tables do I have?" - "How can I analyze customer behavior with this data?" - "Show me orders for these products" (without specifying which products) - "Filter by the criteria I mentioned" (without previous context defining criteria) @@ -98,7 +93,7 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs cannot be interpreted as a database question or data retrieval request. +- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. - It appears off-topic or is simply a casual conversation starter. **Requirements:** diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index dd3904529b..b7675000e0 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -188,7 +188,7 @@ settings: column_indexing_batch_size: 50 table_retrieval_size: 10 table_column_retrieval_size: 100 - allow_intent_classification: true + allow_intent_classification: false allow_sql_generation_reasoning: true allow_sql_functions_retrieval: true enable_column_pruning: false diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 8fc543403e..40657375ea 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -186,7 +186,7 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 query_cache_maxsize: 1000 - allow_intent_classification: true + allow_intent_classification: false allow_sql_generation_reasoning: true allow_sql_functions_retrieval: true enable_column_pruning: false From bf1351f1d3aa56c541f5b5a9d95c67c39489e049 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 19:50:54 +0530 Subject: [PATCH 0657/1087] Add ask timeout settings --- wren-ai-service/src/config.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index a5c6aed24b..bc97d8cc71 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -44,6 +44,9 @@ class Settings(BaseSettings): allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) max_sql_correction_retries: int = Field(default=3) + max_sql_generation_tables: int = Field(default=0) + pipeline_timeout_seconds: int = Field(default=0) + schema_retrieval_timeout_seconds: int = Field(default=0) # engine config engine_timeout: float = Field(default=30.0) From 0c4ba51ed5d83a29e59b3f325ffd4e5a075ff6c2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 20:43:02 +0530 Subject: [PATCH 0658/1087] Preserve ranked schema retrieval order --- .../retrieval/db_schema_retrieval.py | 32 +++++++++++++++---- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..28df670a57 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -120,6 +120,14 @@ def _build_view_ddl(content: dict) -> str: ) +def _table_names_from_retrieval(table_retrieval: dict) -> list[str]: + table_names = [] + for table in table_retrieval.get("documents", []): + content = ast.literal_eval(table.content) + table_names.append(content["name"]) + return table_names + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -172,11 +180,7 @@ async def table_retrieval( async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any ) -> list[Document]: - tables = table_retrieval.get("documents", []) - table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) + table_names = _table_names_from_retrieval(table_retrieval) table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} @@ -204,7 +208,9 @@ async def dbschema_retrieval( @observe() -def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: +def construct_db_schemas( + dbschema_retrieval: list[Document], table_retrieval: dict +) -> list[dict]: db_schemas = {} for document in dbschema_retrieval: content = ast.literal_eval(document.content) @@ -228,7 +234,19 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: # remove incomplete schemas db_schemas = {k: v for k, v in db_schemas.items() if "type" in v and "columns" in v} - return list(db_schemas.values()) + ranked_table_names = _table_names_from_retrieval(table_retrieval) + ranked_schemas = [ + db_schemas[table_name] + for table_name in ranked_table_names + if table_name in db_schemas + ] + remaining_schemas = [ + table_schema + for table_name, table_schema in db_schemas.items() + if table_name not in set(ranked_table_names) + ] + + return ranked_schemas + remaining_schemas @observe(capture_input=False) From 9e6f96dbd1f7706cabe7d3e5d9b5deca5235bbfa Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 21:49:58 +0530 Subject: [PATCH 0659/1087] Restore metadata-grounded ask flow --- .../generation/followup_sql_generation.py | 4 + .../followup_sql_generation_reasoning.py | 4 + .../pipelines/generation/sql_correction.py | 4 + .../pipelines/generation/sql_generation.py | 4 + .../generation/sql_generation_reasoning.py | 4 + .../pipelines/generation/sql_regeneration.py | 4 + .../src/pipelines/generation/utils/sql.py | 33 +- .../pipelines/generation/test_sql_utils.py | 1593 +---------------- .../textBasedAnswerBackgroundTracker.ts | 100 +- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 37 +- .../apollo/server/mdl/test/mdlBuilder.test.ts | 18 +- .../apollo/server/services/askingService.ts | 87 - .../apollo/server/services/projectService.ts | 56 - .../server/utils/recommendationQuestions.ts | 148 -- 14 files changed, 78 insertions(+), 2018 deletions(-) delete mode 100644 wren-ui/src/apollo/server/utils/recommendationQuestions.ts 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 d82d6e40b6..fcfbd227ca 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -39,6 +39,10 @@ {{ document }} {% endfor %} +Use this DATABASE SCHEMA as the complete allowed identifier set for this query. +Only generate SQL with table, column, schema, model, and datasource names present above. +Do not infer identifiers from the follow-up question, previous SQL, summary, SQL samples, user instructions, or prior examples. + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 42b28c5b8f..6c9b89db8a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -28,6 +28,10 @@ {{ document }} {% endfor %} +Use this DATABASE SCHEMA as the complete allowed identifier set for the reasoning plan. +Only refer to table, column, schema, model, and datasource names present above. +Do not infer identifiers from the follow-up question, previous SQL, SQL samples, user instructions, or prior examples. + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index d9caf87a84..999b11c1d8 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -57,6 +57,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% for document in documents %} {{ document }} {% endfor %} + +Use this DATABASE SCHEMA as the complete allowed identifier set for the corrected SQL. +Only correct SQL with table, column, schema, model, and datasource names present above. +Do not infer identifiers from the original SQL, error message, user instructions, or prior examples. {% endif %} {% if sql_functions %} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index d0b921f70c..40097aa44e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -33,6 +33,10 @@ {{ document }} {% endfor %} +Use this DATABASE SCHEMA as the complete allowed identifier set for this query. +Only generate SQL with table, column, schema, model, and datasource names present above. +Do not infer identifiers from the question, SQL samples, user instructions, or prior examples. + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} 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 00b731cb2c..f9f9af9f4f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -27,6 +27,10 @@ {{ document }} {% endfor %} +Use this DATABASE SCHEMA as the complete allowed identifier set for the reasoning plan. +Only refer to table, column, schema, model, and datasource names present above. +Do not infer identifiers from the question, SQL samples, user instructions, or prior examples. + {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index a9b93bc942..2d273a9aea 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -60,6 +60,10 @@ def get_sql_regeneration_system_prompt( {{ document }} {% endfor %} +Use this DATABASE SCHEMA as the complete allowed identifier set for the regenerated SQL. +Only regenerate SQL with table, column, schema, model, and datasource names present above. +Do not infer identifiers from the original SQL, reasoning, SQL samples, user instructions, or prior examples. + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ab22d2fde8..eccf4b8a31 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -437,14 +437,16 @@ async def _classify_generation_result( 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. -7. Give a step by step reasoning plan in order to answer user's question. -8. The reasoning plan should be in the language same as the language user provided in the input. -9. Don't include SQL in the reasoning plan. -10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. -11. Do not include ```markdown or ``` in the answer. -12. A table name in the reasoning plan must be in this format: `table: `. -13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +7. Use DATABASE SCHEMA as the complete and only source of valid table, column, schema, model, and datasource names. +8. Do not introduce, infer, or copy any identifier from the question, SQL samples, user instructions, or query history unless it also appears in DATABASE SCHEMA. +9. Give a step by step reasoning plan in order to answer user's question. +10. The reasoning plan should be in the language same as the language user provided in the input. +11. Don't include SQL in the reasoning plan. +12. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. +13. Do not include ```markdown or ``` in the answer. +14. A table name in the reasoning plan must be in this format: `table: `. +15. A column name in the reasoning plan must be in this format: `column: .`. +16. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -510,10 +512,13 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. -2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +2. The DATABASE SCHEMA section is the complete and only source of valid table, column, schema, model, and datasource names for this request. +3. YOU MUST NOT introduce, infer, copy, or repair any table, column, schema, model, or datasource name that is absent from DATABASE SCHEMA. +4. SQL SAMPLES and USER INSTRUCTIONS are usage guidance only. Do not copy identifiers from them unless those identifiers also appear in DATABASE SCHEMA. +5. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. +6. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. +7. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +8. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} @@ -553,10 +558,6 @@ def construct_instructions( return _instructions -def normalize_generation_result_sql(sql: str, data_source: str | None = None) -> str: - return " ".join(sql.replace('\\"', '"').split()) - - def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 0ee2519757..51f9d4f2e4 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,1592 +1,49 @@ from src.pipelines.generation.utils.sql import ( - contains_unsupported_mssql_json_access, - construct_valid_table_columns, - construct_valid_table_names, - extract_sql_generation_result, - find_invalid_column_references, - find_invalid_table_references, + construct_instructions, get_json_field_instructions, get_metric_instructions, - normalize_data_source, - normalize_generation_result_sql, - normalize_sql_direction_keywords, - normalize_sql_column_references_to_schema, - normalize_sql_table_references_to_schema, get_sql_generation_system_prompt, get_text_to_sql_rules, ) -def test_construct_valid_table_names_from_schema_documents(): - documents = [ - 'CREATE TABLE repair_logs ("id" INTEGER);', - '/* comment */ CREATE TABLE "employees" ("emp_no" INTEGER);', - ] +class _SqlKnowledge: + text_to_sql_rule = "Use the supplied model context only." + metric_instructions = "Use the supplied metric definitions only." + json_field_instructions = "Use the supplied JSON field definitions only." - assert construct_valid_table_names(documents) == ["employees", "repair_logs"] +def test_construct_instructions_uses_instruction_text(): + assert construct_instructions( + [{"instruction": "First rule."}, {"instruction": "Second rule."}] + ) == ["First rule.", "Second rule."] -def test_construct_valid_table_names_includes_ref_sql_source_tables(): - documents = [ - ''' - CREATE TABLE repair_logs ("id" INTEGER); - refSql: SELECT "created_at", "warning_signals" FROM "wrenai"."public"."dbo_repair_logs" - ''', - ] - assert construct_valid_table_names(documents) == [ - "dbo_repair_logs", - "public.dbo_repair_logs", - "repair_logs", - "wrenai.public.dbo_repair_logs", - ] +def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): + rules = get_text_to_sql_rules() + assert "ONLY USE the tables and columns mentioned in the database schema" in rules + assert 'ONLY USE "*" if the user query asks for all the columns' in rules -def test_schema_validation_allows_qualified_suffix_table_references(): - assert find_invalid_table_references( - 'SELECT * FROM "wrenai"."public"."dbo_repair_logs"', - ["dbo_repair_logs"], - ) == [] - -def test_column_validation_allows_qualified_suffix_table_references(): - sql = ( - 'SELECT "wrenai"."public"."dbo_repair_logs"."warning_signals" ' - 'FROM "wrenai"."public"."dbo_repair_logs"' - ) - - assert find_invalid_column_references( - sql, - {"dbo_repair_logs": ["warning_signals"]}, - ) == [] - - -def test_column_validation_rejects_invalid_unqualified_projection_for_single_table(): - assert find_invalid_column_references( - 'SELECT warning_signals FROM "dbo_repair_logs"', - {"dbo_repair_logs": ["id", "created_at", "status"]}, - ) == ["warning_signals"] - - -def test_column_validation_rejects_invalid_unqualified_projection_alias(): - assert find_invalid_column_references( - 'SELECT categories AS category_count FROM "dbo_repair_logs"', - {"dbo_repair_logs": ["id", "created_at", "status"]}, - ) == ["categories"] - - -def test_column_validation_rejects_invalid_unqualified_filter_column(): - assert find_invalid_column_references( - 'SELECT id FROM "policies" WHERE policy_category_id = 1', - {"policies": ["id", "policy_name"]}, - ) == ["policy_category_id"] - - -def test_column_validation_rejects_invalid_unqualified_function_argument(): - assert find_invalid_column_references( - 'SELECT COUNT(policy_category_id) FROM "policies"', - {"policies": ["id", "policy_name"]}, - ) == ["policy_category_id"] - - -def test_column_validation_allows_valid_unqualified_projection_for_single_table(): - assert find_invalid_column_references( - 'SELECT status AS repair_status FROM "dbo_repair_logs"', - {"dbo_repair_logs": ["id", "created_at", "status"]}, - ) == [] - - -def test_normalize_sql_table_references_to_schema_maps_unique_prefix_table(): - sql = ( - 'SELECT "public"."dbo_failure"."created_at" ' - 'FROM "public"."dbo_failure"' - ) - - normalized = normalize_sql_table_references_to_schema( - sql, - ["dbo_failure_patterns"], - ) - - assert normalized == ( - 'SELECT "dbo_failure_patterns"."created_at" ' - 'FROM "dbo_failure_patterns"' - ) - - -def test_normalize_sql_table_references_to_schema_maps_single_quoted_full_table(): - sql = ( - 'SELECT "wrenai.public.dbo_failure"."created_at" ' - 'FROM "wrenai.public.dbo_failure"' - ) - - normalized = normalize_sql_table_references_to_schema( - sql, - ["dbo_failure_patterns"], - ) - - assert normalized == ( - 'SELECT "dbo_failure_patterns"."created_at" ' - 'FROM "dbo_failure_patterns"' - ) - - -def test_normalize_sql_table_references_to_schema_maps_multipart_quoted_table(): - sql = ( - 'SELECT "wrenai"."public"."dbo_failure"."created_at" ' - 'FROM "wrenai"."public"."dbo_failure"' - ) - - normalized = normalize_sql_table_references_to_schema( - sql, - ["dbo_failure_patterns"], - ) - - assert normalized == ( - 'SELECT "dbo_failure_patterns"."created_at" ' - 'FROM "dbo_failure_patterns"' - ) - - -def test_construct_valid_table_columns_adds_qualified_suffix_tables(): - documents = [ - 'CREATE TABLE "wrenai"."public"."dbo_repair_logs" ("warning_signals" INTEGER);', - ] - - assert construct_valid_table_columns(documents) == { - "dbo_repair_logs": ["warning_signals"], - "public.dbo_repair_logs": ["warning_signals"], - "wrenai.public.dbo_repair_logs": ["warning_signals"], - } - - -def test_construct_valid_table_columns_from_semantic_metadata_document(): - documents = [ - """ - { - "models": [ - { - "name": "dbo_new_orders", - "referenceName": "sales.dbo_new_orders", - "columns": [ - {"name": "business"}, - {"name": "market"}, - {"name": "customer_name"}, - {"name": "product_name"}, - {"name": "order_value"} - ], - "calculatedFields": [ - {"name": "order_month"} - ] - } - ] - } - """, - ] - - assert construct_valid_table_names(documents) == [ - "dbo_new_orders", - "sales.dbo_new_orders", - ] - assert construct_valid_table_columns(documents) == { - "dbo_new_orders": [ - "business", - "customer_name", - "market", - "order_month", - "order_value", - "product_name", - ], - "sales.dbo_new_orders": [ - "business", - "customer_name", - "market", - "order_month", - "order_value", - "product_name", - ], - } - - -def test_column_validation_uses_semantic_metadata_document_columns(): - documents = [ - """ - { - "models": [ - { - "name": "dbo_new_orders", - "columns": [ - {"name": "customer_name"}, - {"name": "product_name"}, - {"name": "order_value"} - ] - } - ] - } - """, - ] - valid_table_columns = construct_valid_table_columns(documents) - - assert find_invalid_column_references( - 'SELECT "dbo_new_orders"."customer_name", ' - '"dbo_new_orders"."product_name", ' - '"dbo_new_orders"."order_value" ' - 'FROM "dbo_new_orders"', - valid_table_columns, - ) == [] - assert find_invalid_column_references( - 'SELECT "dbo_new_orders"."missing_value" FROM "dbo_new_orders"', - valid_table_columns, - ) == ["dbo_new_orders.missing_value"] - - -def test_normalize_generation_result_sql_standardizes_identifier_quotes(): - sql = ( - 'SELECT COUNT(*) AS `num_tags`, SUM(`tokenCost`) AS `popularity` ' - 'FROM `dbo_kb_articles` WHERE (""""category"""" = \'Ticket Sourcing\')' - ) - - normalized = normalize_generation_result_sql(sql, data_source="mssql") - - assert "`" not in normalized - assert '""""category""""' not in normalized - assert '"num_tags"' in normalized - assert '"tokenCost"' in normalized - assert '"category" = \'Ticket Sourcing\'' in normalized - - -def test_normalize_sql_direction_keywords_preserves_quoted_text(): - sql = ( - 'SELECT "Description" AS "Desc", \'asc desc\' AS "label" ' - 'FROM "dbo_tblSales" ORDER BY SUM("SalesValue") Desc' - ) - - assert normalize_sql_direction_keywords(sql) == ( - 'SELECT "Description" AS "Desc", \'asc desc\' AS "label" ' - 'FROM "dbo_tblSales" ORDER BY SUM("SalesValue") DESC' - ) - - -def test_normalize_generation_result_sql_uppercases_direction_keywords_before_execution(): - sql = ( - 'SELECT "ProdName", SUM("SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" GROUP BY "ProdName" ' - 'ORDER BY SUM("SalesValue") Desc' - ) - - assert normalize_generation_result_sql(sql).endswith( - 'ORDER BY SUM("SalesValue") DESC' - ) - - -def test_schema_validation_ignores_null_table_metadata(): - assert find_invalid_table_references( - 'SELECT * FROM "dbo_tblSales"', - [None, "dbo_tblSales"], - ) == [] - - -def test_schema_validation_does_not_treat_extract_from_expression_as_table(): - sql = """ - SELECT - "Market", - SUM(CASE WHEN EXTRACT(YEAR FROM "InvDate") = EXTRACT(YEAR FROM GETDATE()) - 1 - THEN "SalesValue" - ELSE 0 - END) AS "LastYearSales" - FROM "dbo_tblSales" - GROUP BY "Market" - """ - - assert find_invalid_table_references(sql, ["dbo_tblSales"]) == [] - - -def test_schema_validation_ignores_null_column_metadata(): - assert find_invalid_column_references( - 'SELECT "dbo_tblSales"."Market" FROM "dbo_tblSales"', - {"dbo_tblSales": [None, "Market"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_uses_exact_schema_names(): - sql = ( - 'SELECT "dbo_xStageLoad8_Test"."PH-BU", "dbo_xStageLoad8_Test"."P-M" ' - 'FROM "dbo_xStageLoad8_Test"' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_xStageLoad8_Test": ["PH_BU", "P_M"]}, - ) - - assert '"dbo_xStageLoad8_Test"."PH_BU"' in normalized - assert '"dbo_xStageLoad8_Test"."P_M"' in normalized - assert "PH-BU" not in normalized - assert "P-M" not in normalized - - -def test_normalize_sql_column_references_to_schema_maps_underscore_to_camel_columns(): - sql = ( - 'SELECT "dbo_tblSalesHistory"."OTD_Date", SUM("dbo_tblSalesHistory"."Sales_Value") ' - 'FROM "dbo_tblSalesHistory" ' - 'GROUP BY "dbo_tblSalesHistory"."OTD_Date"' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, - ) - - assert '"dbo_tblSalesHistory"."OTDDate"' in normalized - assert '"dbo_tblSalesHistory"."SalesValue"' in normalized - assert "OTD_Date" not in normalized - assert "Sales_Value" not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_maps_unqualified_underscore_to_camel_columns(): - sql = ( - 'SELECT DATEPART(YEAR, "OTD_Date") AS "Year", SUM(Sales_Value) ' - 'FROM "dbo_tblSalesHistory" ' - 'GROUP BY DATEPART(YEAR, "OTD_Date")' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, - ) - - assert 'DATEPART(YEAR, "OTDDate") AS "Year"' in normalized - assert 'SUM("SalesValue")' in normalized - assert "OTD_Date" not in normalized - assert "Sales_Value" not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_tblSalesHistory": ["OTDDate", "SalesValue"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_maps_otd_date_to_invoice_date(): - sql = ( - 'SELECT DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date") AS "Year", ' - 'SUM("dbo_tblSalesHistory"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSalesHistory" ' - 'GROUP BY DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date")' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_tblSalesHistory": ["InvDate", "SalesValue"]}, - ) - - assert 'DATEPART(YEAR, "dbo_tblSalesHistory"."InvDate") AS "Year"' in normalized - assert "OTD_Date" not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_tblSalesHistory": ["InvDate", "SalesValue"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_maps_sales_business_aliases(): - sql = ( - 'SELECT "dbo_qSales1"."Customer_Region", ' - 'SUM("dbo_qSales1"."InvoiceQuantity") AS "InvoiceQuantity" ' - 'FROM "dbo_qSales1" ' - 'GROUP BY "dbo_qSales1"."Customer_Region"' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_qSales1": ["Country", "Market", "Qty"]}, - ) - - assert '"dbo_qSales1"."Country"' in normalized - assert '"dbo_qSales1"."Qty"' in normalized - assert "Customer_Region" not in normalized - assert "InvoiceQuantity" not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_qSales1": ["Country", "Market", "Qty"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_maps_period_to_timeid(): - sql = 'SELECT "dbo_tblFactSales"."Period" FROM "dbo_tblFactSales"' - - normalized = normalize_sql_column_references_to_schema( - sql, - { - "dbo_tblFactSales": [ - "account", - "customerpo", - "timeid", - "amount", - ] - }, - ) - - assert normalized == 'SELECT "dbo_tblFactSales"."timeid" FROM "dbo_tblFactSales"' - - -def test_normalize_sql_column_references_to_schema_maps_customer_to_account(): - sql = 'SELECT "dbo_tblFactSales"."customer" FROM "dbo_tblFactSales"' - - normalized = normalize_sql_column_references_to_schema( - sql, - { - "dbo_tblFactSales": [ - "account", - "customerpo", - "timeid", - "amount", - ] - }, - ) - - assert normalized == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' - - -def test_normalize_sql_column_references_to_schema_maps_unqualified_customer_to_account(): - sql = 'SELECT "customer" FROM "dbo_tblFactSales"' - - normalized = normalize_sql_column_references_to_schema( - sql, - { - "dbo_tblFactSales": [ - "account", - "customerpo", - "timeid", - "amount", - ] - }, - ) - - assert normalized == 'SELECT "account" FROM "dbo_tblFactSales"' - - -def test_normalize_sql_column_references_to_schema_maps_debug_business_aliases(): - sql = ( - 'SELECT COUNT("FixLogId") AS "FixLogCount" ' - 'FROM "dbo_DebugEntries"' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_DebugEntries": ["DebugEntryId", "FixId"]}, - ) - - assert 'COUNT("DebugEntryId") AS "FixLogCount"' in normalized - assert "FixLogId" not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_DebugEntries": ["DebugEntryId", "FixId"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_maps_last_update_date_alias(): - sql = ( - 'SELECT DATEPART(YEAR, last_update_date) AS "year", ' - 'DATEPART(MONTH, last_update_date) AS "month", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_DebugEntries" ' - 'GROUP BY DATEPART(YEAR, last_update_date), ' - 'DATEPART(MONTH, last_update_date)' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_DebugEntries": ["DebugEntryId", "BusinessUnit", "DateIn", "FailedAt"]}, - ) - - assert "last_update_date" not in normalized - assert 'DATEPART(YEAR, "DateIn") AS "year"' in normalized - assert 'DATEPART(MONTH, "DateIn") AS "month"' in normalized - assert find_invalid_column_references( - normalized, - {"dbo_DebugEntries": ["DebugEntryId", "BusinessUnit", "DateIn", "FailedAt"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_keeps_unknown_columns_invalid(): - sql = 'SELECT "dbo_qSales1"."UnitPrice" FROM "dbo_qSales1"' - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_qSales1": ["SalesValue", "Cost"]}, - ) - - assert normalized == sql - assert find_invalid_column_references( - normalized, - {"dbo_qSales1": ["SalesValue", "Cost"]}, - ) == ["dbo_qSales1.UnitPrice"] - - -def test_normalize_sql_column_references_to_schema_maps_kb_article_aliases(): - sql = ( - 'SELECT "dbo_kb_articles"."article_type", COUNT(*) AS "RecordCount" ' - 'FROM "dbo_kb_articles" ' - 'GROUP BY "dbo_kb_articles"."article_type"' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_kb_articles": ["id", "category", "source_ticket_id"]}, - ) - - assert '"dbo_kb_articles"."category"' in normalized - assert "article_type" not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_kb_articles": ["id", "category", "source_ticket_id"]}, - ) == [] - - -def test_normalize_sql_column_references_to_schema_maps_unqualified_source(): - sql = ( - 'SELECT source, COUNT(*) AS "RecordCount" ' - 'FROM "dbo_kb_articles" ' - 'GROUP BY source ' - 'ORDER BY COUNT(*) DESC' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - {"dbo_kb_articles": ["id", "category", "source_ticket_id"]}, - ) - - assert '"source_ticket_id"' in normalized - assert " source" not in normalized - assert "GROUP BY source" not in normalized - - -def test_normalize_sql_column_references_to_schema_maps_unqualified_article_id(): - sql = ( - 'SELECT COUNT(article_id) AS "article_count" ' - 'FROM "dbo_knowledge_articles"' - ) - - normalized = normalize_sql_column_references_to_schema( - sql, - { - "dbo_knowledge_articles": [ - "id", - "org_id", - "title", - "category", - "subcategory", - "content", - "author", - "tags", - "views", - "helpful", - "data", - "created_at", - "updated_at", - ] - }, - ) - - assert 'COUNT("id") AS "article_count"' in normalized - assert "article_id" not in normalized - assert find_invalid_column_references( - normalized, - {"dbo_knowledge_articles": ["id", "org_id", "title"]}, - ) == [] - - -def test_sql_generation_system_prompt_rejects_stale_sales_sample_schema(): - prompt = get_sql_generation_system_prompt() - - assert "SQL SAMPLES are examples of style only" in prompt - assert "sales performance" in prompt - assert "Do not SUM or AVG string columns" in prompt - - -def test_extract_sql_generation_result_from_json_payload(): - result = '{"sql": "SELECT COUNT(*) AS repair_count FROM repairs;"}' - - assert ( - extract_sql_generation_result(result) - == "SELECT COUNT(*) AS repair_count FROM repairs" - ) - - -def test_extract_sql_generation_result_from_prose_wrapped_sql(): - result = ( - "The SQL query is: SELECT DATEPART(YEAR, created_at) AS year, " - "COUNT(*) AS repair_count FROM repairs GROUP BY DATEPART(YEAR, created_at);" - ) - - assert extract_sql_generation_result(result) == ( - "SELECT DATEPART(YEAR, created_at) AS year, COUNT(*) AS repair_count " - "FROM repairs GROUP BY DATEPART(YEAR, created_at)" - ) - - -def test_extract_sql_generation_result_from_fenced_sql(): - result = """ - Here is the query: - ```sql - SELECT id FROM repairs; - ``` - """ - - assert extract_sql_generation_result(result) == "SELECT id FROM repairs" - - -def test_extract_sql_generation_result_from_prose_wrapped_json(): - result = 'Here is the result:\n{"sql": "SELECT id FROM repairs;"}' - - assert extract_sql_generation_result(result) == "SELECT id FROM repairs" - - -def test_get_text_to_sql_rules_adds_mssql_specific_constraints(): - rules = get_text_to_sql_rules(data_source="MSSQL") - - assert "The target database is MSSQL." in rules - assert "DATEPART(YEAR, )" in rules - assert "DATEADD, DATEDIFF, DATETIME2, or DATETIMEOFFSET" in rules - assert "TO_UNIXTIME" in rules - assert "Do not subtract timestamp/date columns directly" in rules - assert "TO_TIMESTAMP_MILLIS" in rules - assert "DO NOT use PostgreSQL-style or Trino-style date syntax" in rules - assert "DO NOT use JSON extraction functions or operators" in rules - assert "JSON_VALUE" in rules - assert "JSON_EXTRACT" in rules - assert "->>" in rules - assert "do not assume keys inside it are queryable" in rules - assert "Never invent JSON-derived columns" in rules - assert "Resolve relative time phrases" in rules - assert "Do not include helper ranking columns" in rules - assert "prefer SELECT TOP (N)" in rules - assert "connected datasource metadata" in rules - assert '"dbo_DebugFixes"."Description"' in rules - assert '"dbo_DebugFixLogs"."FixId"' in rules - assert "repair SLA compliance" in rules - assert '"dbo_repair_logs"."status"' in rules - assert "FailurePatternID" in rules - assert "failure_code" in rules - assert "CURRENT_DATE - INTERVAL '1 month'" not in rules - - -def test_get_json_field_instructions_for_mssql_disables_json_extraction(): - instructions = get_json_field_instructions(data_source="MSSQL") - - assert "cannot execute JSON extraction" in instructions - assert "JSON_VALUE" in instructions - assert "->>" in instructions - assert "Use only first-class columns" in instructions - assert "LAX_STRING(JSON_QUERY" not in instructions - - -def test_contains_unsupported_mssql_json_access_detects_json_syntax(): - assert contains_unsupported_mssql_json_access( - 'SELECT "data" ->> \'AttemptNumber\' FROM "dbo_repair_logs"' - ) - assert contains_unsupported_mssql_json_access( - 'SELECT JSON_VALUE("data", \'$.AttemptNumber\') FROM "dbo_repair_logs"' - ) - assert not contains_unsupported_mssql_json_access( - 'SELECT "created_at", "status" FROM "dbo_repair_logs"' - ) - - -def test_get_metric_instructions_for_mssql_avoids_date_trunc_example(): - instructions = get_metric_instructions(data_source="MSSQL") - - assert "DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month')" not in instructions - assert "Do not use DATE_TRUNC, DATETRUNC, DATEADD, DATEDIFF, INTERVAL, CURRENT_DATE, or TIMESTAMP WITH TIME ZONE" in instructions - assert "DATEPART(YEAR, )" in instructions - - -def test_get_sql_generation_system_prompt_uses_data_source_specific_rules(): - prompt = get_sql_generation_system_prompt(data_source="MSSQL") - - assert "The target database is MSSQL." in prompt - assert "DATEPART(YEAR, )" in prompt - assert "deployed semantic model definitions" in prompt - assert '"dbo_DebugFixes"."Description"' in prompt - assert "repair SLA compliance" in prompt - - -def test_normalize_generation_result_sql_rewrites_common_mssql_time_patterns(): - sql = """ - SELECT - DATEPART(YEAR, "created_at") AS "year", - DATEPART(MONTH, "created_at") AS "month", - COUNT("id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY DATEPART(YEAR, "created_at"), DATEPART(MONTH, "created_at") - ORDER BY "year" ASC NULLS LAST, "month" ASC NULLS LAST - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert 'DATEPART(\'YEAR\', "created_at")' not in normalized - assert 'DATEPART(\'MONTH\', "created_at")' not in normalized - assert "NULLS LAST" not in normalized - assert 'DATEPART(YEAR, "created_at")' in normalized - assert 'DATEPART(MONTH, "created_at")' in normalized - - -def test_normalize_generation_result_sql_rewrites_cwsales_otd_date_for_mssql(): - sql = ( - 'SELECT DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date") AS "Year", ' - 'SUM("dbo_tblSalesHistory"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSalesHistory" ' - 'GROUP BY DATEPART(YEAR, "dbo_tblSalesHistory"."OTD_Date")' - ) - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "OTD_Date" not in normalized - assert 'DATEPART(YEAR, "dbo_tblSalesHistory"."InvDate") AS "Year"' in normalized - assert 'GROUP BY DATEPART(YEAR, "dbo_tblSalesHistory"."InvDate")' in normalized - - -def test_normalize_generation_result_sql_rewrites_cwsales_fix_log_id_for_mssql(): - sql = ( - 'SELECT COUNT("dbo_qSales1"."FixLogId") AS "NumberOfInvoices" ' - 'FROM "dbo_qSales1"' - ) - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "FixLogId" not in normalized - assert 'COUNT("dbo_qSales1"."InvoiceNo") AS "NumberOfInvoices"' in normalized - - -def test_normalize_generation_result_sql_rewrites_common_mssql_dateadd_patterns(): - sql = """ - SELECT - DATEADD(month, DATEDIFF(month, 0, "created_at"), 0) AS "month_start", - COUNT("id") AS "repair_count" - FROM "dbo_repair_logs" - WHERE "created_at" >= DATEADD(month, -12, GETDATE()) - AND "created_at" < DATEADD(month, DATEDIFF(month, 0, GETDATE()), 0) - GROUP BY DATEADD(month, DATEDIFF(month, 0, "created_at"), 0) - ORDER BY "month_start" ASC NULLS LAST - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "DATEADD(" not in normalized - assert "DATEDIFF(" not in normalized - assert "NULLS LAST" not in normalized - assert 'DATEPART(YEAR, "created_at")' in normalized - assert 'DATEPART(MONTH, "created_at")' in normalized +def test_get_text_to_sql_rules_uses_sql_knowledge_override(): + assert get_text_to_sql_rules(_SqlKnowledge()) == _SqlKnowledge.text_to_sql_rule -def test_normalize_data_source_maps_sql_server_aliases_to_mssql(): - assert normalize_data_source("sqlserver") == "MSSQL" - assert normalize_data_source("SQL Server") == "MSSQL" +def test_get_metric_instructions_uses_sql_knowledge_override(): + assert get_metric_instructions(_SqlKnowledge()) == _SqlKnowledge.metric_instructions -def test_normalize_generation_result_sql_rewrites_nested_temporal_patterns_for_mssql(): - sql = """ - SELECT - DATE_PART('YEAR', CAST("created_at" AS DATETIME)) AS "year", - DATE_TRUNC('MONTH', CAST("created_at" AS DATETIME)) AS "month_bucket", - EXTRACT(DAY FROM CAST("created_at" AS DATETIME)) AS "day_of_month" - FROM "dbo_repair_logs" - ORDER BY "month_bucket" ASC NULLS LAST - """ - - normalized = normalize_generation_result_sql(sql, data_source="sqlserver") - - assert "DATE_PART(" not in normalized - assert "DATE_TRUNC(" not in normalized - assert "EXTRACT(" not in normalized - assert "NULLS LAST" not in normalized - assert 'DATEPART(YEAR, CAST("created_at" AS DATETIME))' in normalized - assert 'DATEPART(MONTH, CAST("created_at" AS DATETIME))' in normalized - assert 'DATEPART(DAY, CAST("created_at" AS DATETIME))' in normalized - - -def test_normalize_generation_result_sql_rewrites_to_timestamp_for_mssql(): - sql = """ - SELECT - DATEPART(YEAR, TO_TIMESTAMP("created_at")) AS "year", - COUNT("id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY DATEPART(YEAR, TO_TIMESTAMP("created_at")) - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "TO_TIMESTAMP(" not in normalized - assert 'DATEPART(YEAR, CAST("created_at" AS DATETIME))' in normalized - - -def test_normalize_generation_result_sql_rewrites_to_timestamp_variants_for_mssql(): - sql = """ - SELECT - TO_TIMESTAMP_MILLIS("created_at_ms") AS "created_at", - TO_TIMESTAMP_SECONDS("closed_at_sec") AS "closed_at" - FROM "dbo_repair_logs" - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "TO_TIMESTAMP_MILLIS(" not in normalized - assert "TO_TIMESTAMP_SECONDS(" not in normalized - assert 'CAST("created_at_ms" AS DATETIME)' in normalized - assert 'CAST("closed_at_sec" AS DATETIME)' in normalized - - -def test_normalize_generation_result_sql_rewrites_mssql_datepart_alias_references(): - sql = """ - SELECT - DATEPART(YEAR, "created_at") AS "YEAR", - DATEPART(MONTH, "created_at") AS "MONTH", - COUNT("id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY "YEAR", "MONTH" - ORDER BY "YEAR" ASC, "MONTH" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert 'GROUP BY "YEAR"' not in normalized - assert 'ORDER BY "YEAR"' not in normalized - assert 'DATEPART(YEAR, "created_at") AS "YEAR"' in normalized - assert ( - 'GROUP BY DATEPART(YEAR, "created_at"), DATEPART(MONTH, "created_at")' - in normalized - ) - assert ( - 'ORDER BY DATEPART(YEAR, "created_at") ASC, DATEPART(MONTH, "created_at") ASC' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_unquoted_mssql_datepart_alias_references(): - sql = """ - SELECT - DATEPART(YEAR, "created_at") AS YEAR, - DATEPART(MONTH, "created_at") AS MONTH, - COUNT("id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY YEAR, MONTH - ORDER BY YEAR ASC, MONTH ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "GROUP BY YEAR" not in normalized - assert "ORDER BY YEAR" not in normalized - assert 'DATEPART(YEAR, "created_at") AS YEAR' in normalized - assert ( - 'GROUP BY DATEPART(YEAR, "created_at"), DATEPART(MONTH, "created_at")' - in normalized - ) - assert ( - 'ORDER BY DATEPART(YEAR, "created_at") ASC, DATEPART(MONTH, "created_at") ASC' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_invented_repair_date_for_mssql(): - sql = """ - SELECT - DATEPART(YEAR, "RepairDate") AS "YEAR", - DATEPART(MONTH, "RepairDate") AS "MONTH", - COUNT("dbo_repair_logs"."id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY DATEPART(YEAR, "RepairDate"), DATEPART(MONTH, "RepairDate") - ORDER BY "YEAR" ASC, "MONTH" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert '"RepairDate"' not in normalized - assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "YEAR"' in normalized - assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "MONTH"' in normalized - assert ( - 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at")' - in normalized - ) - assert ( - 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_qualified_invented_repair_date_for_mssql(): - sql = """ - SELECT - DATEPART(YEAR, "dbo_repair_logs"."RepairDate") AS "YEAR", - DATEPART(MONTH, "dbo_repair_logs"."RepairDate") AS "MONTH", - COUNT("dbo_repair_logs"."id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY DATEPART(YEAR, "dbo_repair_logs"."RepairDate"), - DATEPART(MONTH, "dbo_repair_logs"."RepairDate") - ORDER BY "YEAR" ASC, "MONTH" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "RepairDate" not in normalized - assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "YEAR"' in normalized - assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "MONTH"' in normalized - - -def test_normalize_generation_result_sql_rewrites_invented_repair_failure_pattern_id_for_mssql(): - sql = """ - SELECT - "dbo_failure_patterns"."category" AS "failure_category", - COUNT("dbo_repair_logs"."id") AS "repair_count" - FROM "dbo_repair_logs" - JOIN "dbo_failure_patterns" - ON "dbo_repair_logs"."FailurePatternID" = "dbo_failure_patterns"."id" - GROUP BY "dbo_failure_patterns"."category" - ORDER BY "repair_count" DESC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "FailurePatternID" not in normalized - assert ( - '"dbo_repair_logs"."failure_code" = "dbo_failure_patterns"."id"' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_debug_entry_failure_pattern_join_for_mssql(): - sql = """ - SELECT TOP 10 - "dbo_failure_patterns"."category" AS "FailureCategory", - COUNT_BIG(1) AS "FailureCount" - FROM "dbo_DebugEntries" - INNER JOIN "dbo_failure_patterns" - ON "dbo_DebugEntries"."DebugEntryId" = "dbo_failure_patterns"."id" - GROUP BY "dbo_failure_patterns"."category" - ORDER BY "FailureCount" DESC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert '"dbo_DebugEntries"."DebugEntryId" = "dbo_failure_patterns"."id"' not in normalized - assert ( - '"dbo_DebugEntries"."FailureSys" = "dbo_failure_patterns"."id"' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_pcb_throughput_repair_log_fields_for_mssql(): - sql = """ - SELECT - "dbo_repair_logs"."ManufacturingUnit" AS "manufacturing_unit", - "dbo_repair_logs"."MONTH", - COUNT("dbo_repair_logs"."id") AS "throughput" - FROM "dbo_repair_logs" - GROUP BY "dbo_repair_logs"."ManufacturingUnit", "dbo_repair_logs"."MONTH" - ORDER BY "dbo_repair_logs"."MONTH" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "dbo_repair_logs" not in normalized - assert "ManufacturingUnit" not in normalized - assert '"dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit"' in normalized - assert '"dbo_DebugEntries"."DebugEntryId"' in normalized - assert 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month"' in normalized - assert 'GROUP BY "dbo_DebugEntries"."BusinessUnit"' in normalized - assert 'ORDER BY DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_repair_log_turnaround_throughput_shape_for_mssql(): - sql = """ - SELECT - board_model AS unit_name, - COUNT(*) AS repair_count, - AVG((DATEPART(DAY, updated_at) - DATEPART(DAY, created_at))) AS avg_turnaround_time - FROM dbo_repair_logs - GROUP BY board_model - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "DATEPART(DAY" not in normalized - assert "avg_turnaround_time" not in normalized - assert "dbo_repair_logs" not in normalized - assert ( - 'SELECT "dbo_DebugEntries"."BusinessUnit" AS "unit_name", ' - 'COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput"' - in normalized - ) - assert 'FROM "dbo_DebugEntries"' in normalized - - -def test_normalize_generation_result_sql_rewrites_repair_log_turnaround_month_trend_shape_for_mssql(): - sql = """ - SELECT - MONTH, - AVG(avg_turnaround_time) AS avg_turnaround_time - FROM dbo_repair_logs - GROUP BY MONTH - ORDER BY MONTH ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "avg_turnaround_time" not in normalized - assert "GROUP BY MONTH" not in normalized - assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year"' in normalized - assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' in normalized - assert ( - 'AVG(DATEDIFF(\'second\', "dbo_repair_logs"."created_at", ' - '"dbo_repair_logs"."updated_at")) AS "avg_turnaround_seconds"' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_bare_month_field_for_mssql(): - sql = """ - SELECT - "MONTH", - COUNT("dbo_repair_logs"."id") AS "repair_count" - FROM "dbo_repair_logs" - WHERE "dbo_repair_logs"."created_at" >= '2025-05-01 00:00:00' - GROUP BY "MONTH" - ORDER BY "MONTH" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert 'SELECT "MONTH"' not in normalized - assert 'GROUP BY "MONTH"' not in normalized - assert 'ORDER BY "MONTH"' not in normalized - assert ( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' - in normalized - ) - assert 'GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized - assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_qualified_month_field_for_mssql(): - sql = """ - SELECT - "dbo_repair_logs"."MONTH", - COUNT("dbo_repair_logs"."id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY "dbo_repair_logs"."MONTH" - ORDER BY "dbo_repair_logs"."MONTH" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert '"dbo_repair_logs"."MONTH"' not in normalized +def test_get_json_field_instructions_uses_sql_knowledge_override(): assert ( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' - in normalized + get_json_field_instructions(_SqlKnowledge()) + == _SqlKnowledge.json_field_instructions ) - assert 'GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized - assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_aggregate_qualified_temporal_column_for_mssql(): - sql = """ - SELECT - DATEPART(YEAR, "SUM"."created_at") AS "year", - DATEPART(MONTH, "SUM"."created_at") AS "month", - COUNT(*) AS "ticket_count" - FROM "dbo_tickets" - GROUP BY DATEPART(YEAR, "SUM"."created_at"), DATEPART(MONTH, "SUM"."created_at") - ORDER BY DATEPART(YEAR, "SUM"."created_at"), DATEPART(MONTH, "SUM"."created_at") - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert '"SUM"."created_at"' not in normalized - assert 'DATEPART(YEAR, "dbo_tickets"."created_at") AS "year"' in normalized - assert 'DATEPART(MONTH, "dbo_tickets"."created_at") AS "month"' in normalized - assert 'GROUP BY DATEPART(YEAR, "dbo_tickets"."created_at")' in normalized - assert 'ORDER BY DATEPART(YEAR, "dbo_tickets"."created_at")' in normalized - - -def test_normalize_generation_result_sql_rewrites_unquoted_qualified_month_field_for_mssql(): - sql = """ - SELECT - dbo_repair_logs.MONTH, - COUNT(dbo_repair_logs.id) AS "repair_count" - FROM dbo_repair_logs - GROUP BY dbo_repair_logs.MONTH - ORDER BY dbo_repair_logs.MONTH ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "dbo_repair_logs.MONTH" not in normalized - assert ( - 'DATEPART(MONTH, dbo_repair_logs."created_at") AS "month"' - in normalized - ) - assert 'GROUP BY DATEPART(MONTH, dbo_repair_logs."created_at")' in normalized - assert 'ORDER BY DATEPART(MONTH, dbo_repair_logs."created_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_bare_unquoted_month_field_for_mssql(): - sql = """ - SELECT - MONTH, - COUNT(*) AS repair_volume - FROM dbo_repair_logs - GROUP BY MONTH - ORDER BY MONTH ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "SELECT MONTH" not in normalized - assert "GROUP BY MONTH" not in normalized - assert "ORDER BY MONTH" not in normalized - assert ( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' - in normalized - ) - assert 'GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized - assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_bare_month_field_for_local_file(): - sql = """ - SELECT - MONTH, - COUNT(*) AS repair_volume - FROM dbo_repair_logs - GROUP BY MONTH - ORDER BY MONTH ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="local_file") - - assert "SELECT MONTH" not in normalized - assert "GROUP BY MONTH" not in normalized - assert "ORDER BY MONTH" not in normalized - assert ( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' - in normalized - ) - assert 'GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized - assert 'ORDER BY DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_bare_month_field_for_sqlite(): - sql = """ - SELECT - "MONTH", - COUNT("dbo_repair_logs"."id") AS "repair_count" - FROM "dbo_repair_logs" - GROUP BY "MONTH" - ORDER BY "MONTH" ASC - """ - normalized = normalize_generation_result_sql(sql, data_source="sqlite") - assert 'SELECT "MONTH"' not in normalized - assert 'GROUP BY "MONTH"' not in normalized - assert 'ORDER BY "MONTH"' not in normalized - assert ( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_bare_year_for_report_charts(): - sql = """ - SELECT - "YEAR", - COUNT("dbo_reports"."id") AS "report_count" - FROM "dbo_reports" - GROUP BY "YEAR" - ORDER BY "YEAR" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert 'SELECT "YEAR"' not in normalized - assert 'GROUP BY "YEAR"' not in normalized - assert 'ORDER BY "YEAR"' not in normalized - assert 'DATEPART(YEAR, "dbo_reports"."generated_at") AS "year"' in normalized - assert 'GROUP BY DATEPART(YEAR, "dbo_reports"."generated_at")' in normalized - assert 'ORDER BY DATEPART(YEAR, "dbo_reports"."generated_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_timestamp_casts_for_mssql(): - sql = """ - SELECT COUNT("id") - FROM "dbo_repair_logs" - WHERE CAST("created_at" AS TIMESTAMP) >= CAST('2026-01-01 00:00:00' AS TIMESTAMP) - """ - - normalized = normalize_generation_result_sql(sql, data_source="sqlserver") - - assert " AS TIMESTAMP" not in normalized - assert 'CAST("created_at" AS DATETIME)' in normalized - assert "CAST('2026-01-01 00:00:00' AS DATETIME)" in normalized - - -def test_normalize_generation_result_sql_rewrites_to_date_bucket_for_mssql(): - sql = """ - SELECT - TO_DATE(DateIn, 'YYYY-MM-DD') EntryDate, - COUNT(*) Throughput - FROM dbo_DebugEntries - GROUP BY EntryDate - ORDER BY EntryDate ASC NULLS LAST - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "TO_DATE(" not in normalized - assert "NULLS LAST" not in normalized - assert "GROUP BY EntryDate" not in normalized - assert "ORDER BY EntryDate" not in normalized - assert 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn")' in normalized - assert 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn")' in normalized - assert 'DATEPART(DAY, "dbo_DebugEntries"."DateIn")' in normalized - - -def test_normalize_generation_result_sql_rewrites_date_function_for_mssql(): - sql = """ - SELECT - DATE(DateIn) AS EntryDate, - COUNT(*) AS Throughput - FROM dbo_DebugEntries - GROUP BY EntryDate - ORDER BY EntryDate ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "DATE(DateIn)" not in normalized - assert "GROUP BY EntryDate" not in normalized - assert 'DATEPART(DAY, "dbo_DebugEntries"."DateIn")' in normalized - - -def test_normalize_generation_result_sql_rewrites_date_sub_for_mssql(): - sql = """ - SELECT - DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "MONTH", - COUNT(*) AS "repair_volume" - FROM "dbo_repair_logs" - WHERE "dbo_repair_logs"."created_at" >= DATE_SUB(CURRENT_DATE, INTERVAL 12 MONTH) - GROUP BY DATEPART(MONTH, "dbo_repair_logs"."created_at") - ORDER BY "MONTH" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "DATE_SUB(" not in normalized - assert "CURRENT_DATE" not in normalized - assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at")' in normalized - - -def test_normalize_generation_result_sql_rewrites_repair_log_failure_category_for_mssql(): - sql = """ - SELECT - failure_category, - COUNT(*) AS repair_count - FROM dbo_repair_logs - GROUP BY failure_category - ORDER BY repair_count DESC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "failure_category," not in normalized - assert "GROUP BY failure_category" not in normalized - assert '"dbo_repair_logs"."failure_code" AS "failure_category"' in normalized - assert 'GROUP BY "dbo_repair_logs"."failure_code"' in normalized - - -def test_normalize_generation_result_sql_rewrites_report_hallucinated_fields_for_mssql(): - sql = """ - SELECT - COUNT(*) total_reports, - SUM(CASE WHEN (filters LIKE '%raw%data%file%') THEN 1 ELSE 0 END) raw_data_files_included, - AVG((CASE WHEN (filters LIKE '%raw%data%file%') THEN file_size ELSE null END)) avg_file_size_with_raw_data - FROM dbo_reports - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "filters" not in normalized - assert "THEN file_size" not in normalized - assert '"dbo_reports"."data" LIKE' in normalized - assert '"dbo_reports"."size_bytes"' in normalized - - -def test_normalize_generation_result_sql_rewrites_ticket_token_cost_for_mssql(): - sql = """ - SELECT - "status", - AVG(token_cost) average_token_cost - FROM "dbo_tickets" - GROUP BY "status" - ORDER BY average_token_cost DESC NULLS LAST - LIMIT 1 - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "token_cost" not in normalized - assert 'SELECT "dbo_tickets"."status" AS "status"' in normalized - assert 'COUNT("dbo_tickets"."id") AS "ticket_count"' in normalized - assert 'FROM "dbo_tickets"' in normalized - assert 'GROUP BY "dbo_tickets"."status"' in normalized - - -def test_normalize_generation_result_sql_strips_to_unixtime_for_mssql(): - sql = """ - SELECT - TO_UNIXTIME(CAST("created_at" AS TIMESTAMP)) AS "created_at_unix", - AVG("repair_cost") AS "avg_repair_cost" - FROM "dbo_repair_logs" - GROUP BY TO_UNIXTIME(CAST("created_at" AS TIMESTAMP)) - ORDER BY "created_at_unix" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "TO_UNIXTIME(" not in normalized - assert 'CAST("created_at" AS DATETIME) AS "created_at_unix"' in normalized - assert 'GROUP BY CAST("created_at" AS DATETIME)' in normalized - - -def test_normalize_generation_result_sql_rewrites_timestamp_subtraction_for_mssql(): - sql = """ - SELECT - "updated_at" - "created_at" AS "turnaround_seconds", - "repair_cost" - FROM "dbo_repair_logs" - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert '"updated_at" - "created_at"' not in normalized - assert ( - 'DATEDIFF(\'second\', "created_at", "updated_at") AS "turnaround_seconds"' - in normalized - ) - -def test_normalize_generation_result_sql_rewrites_mssql_time_buckets_and_ordering(): - sql = """ - SELECT - YEAR, - MONTH, - COUNT(*) AS repair_count - FROM dbo_repair_logs - GROUP BY YEAR, MONTH - ORDER BY YEAR ASC NULLS LAST, MONTH ASC NULLS LAST - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "NULLS LAST" not in normalized - assert "SELECT YEAR" not in normalized - assert "GROUP BY YEAR" not in normalized - assert "ORDER BY YEAR" not in normalized - assert ( - 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year"' - in normalized - ) - assert ( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' - in normalized - ) - assert 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at")' in normalized - assert 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_aliased_mssql_time_buckets(): - sql = """ - SELECT - YEAR AS year, - MONTH AS month, - COUNT(*) AS repair_count - FROM dbo_repair_logs - GROUP BY YEAR, MONTH - ORDER BY YEAR ASC, MONTH ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "YEAR AS year" not in normalized - assert "MONTH AS month" not in normalized - assert ( - 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS year' - in normalized - ) - assert ( - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS month' - in normalized - ) - assert 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at")' in normalized - assert 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_debug_entry_quoted_year_alias(): - sql = """ - SELECT - "YEAR" AS "YEAR", - "dbo_DebugEntries"."BusinessUnit" AS "manufacturing_unit", - COUNT("dbo_DebugEntries"."DebugEntryId") AS "throughput" - FROM "dbo_DebugEntries" - GROUP BY "YEAR", "dbo_DebugEntries"."BusinessUnit" - ORDER BY "YEAR" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert '"YEAR" AS "YEAR"' not in normalized - assert 'GROUP BY "YEAR"' not in normalized - assert 'ORDER BY "YEAR"' not in normalized - assert 'DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "YEAR"' in normalized - assert 'GROUP BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn")' in normalized - assert 'ORDER BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC' in normalized - - -def test_normalize_generation_result_sql_rewrites_mssql_limit_and_where_parentheses(): - sql = """ - SELECT model_id, COUNT(*) AS ticket_count - FROM dbo_tickets - WHERE (source = 'AI') - GROUP BY model_id - ORDER BY ticket_count DESC NULLS LAST - LIMIT 1 - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert normalized.startswith("SELECT TOP 1") - assert "WHERE (source = 'AI')" not in normalized - assert "WHERE source = 'AI'" in normalized - assert "NULLS LAST" not in normalized - assert "LIMIT 1" not in normalized - - -def test_normalize_generation_result_sql_removes_mssql_limit_when_top_exists(): - sql = """ - SELECT TOP 5 - id - FROM dbo_tickets - ORDER BY id DESC - LIMIT 1 - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "LIMIT" not in normalized - assert "SELECT TOP 5 id" in normalized - - -def test_normalize_generation_result_sql_rewrites_knowledge_article_time_buckets_for_mssql(): - sql = """ - SELECT - "YEAR", - COUNT("dbo_knowledge_articles"."id") AS "article_count" - FROM "dbo_knowledge_articles" - GROUP BY "YEAR" - ORDER BY "YEAR" ASC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert 'SELECT "YEAR"' not in normalized - assert 'GROUP BY "YEAR"' not in normalized - assert 'ORDER BY "YEAR"' not in normalized - assert ( - 'DATEPART(YEAR, "dbo_knowledge_articles"."created_at") AS "year"' - in normalized - ) - assert ( - 'GROUP BY DATEPART(YEAR, "dbo_knowledge_articles"."created_at")' - in normalized - ) - - -def test_normalize_generation_result_sql_rewrites_knowledge_article_hallucinated_fields_for_mssql(): - sql = """ - SELECT - AVG("dbo_knowledge_articles"."effectiveness_score") AS "avg_effectiveness", - "dbo_knowledge_articles"."created_by" AS "created_by" - FROM "dbo_knowledge_articles" - GROUP BY "dbo_knowledge_articles"."created_by" - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "effectiveness_score" not in normalized - assert '"dbo_knowledge_articles"."created_by"' not in normalized - assert 'AVG("dbo_knowledge_articles"."helpful") AS "avg_effectiveness"' in normalized - assert '"dbo_knowledge_articles"."author" AS "author"' in normalized - assert 'GROUP BY "dbo_knowledge_articles"."author"' in normalized - - -def test_normalize_generation_result_sql_rewrites_knowledge_article_id_for_mssql(): - sql = """ - SELECT - COUNT("dbo_knowledge_articles"."article_id") AS "article_count" - FROM "dbo_knowledge_articles" - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "article_id" not in normalized - assert 'COUNT("dbo_knowledge_articles"."id") AS "article_count"' in normalized - - -def test_normalize_generation_result_sql_rewrites_article_content_for_mssql(): - sql = """ - SELECT - LENGTH("dbo_kb_articles"."article_text") AS "article_length" - FROM "dbo_kb_articles" - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "article_text" not in normalized - assert 'LENGTH("dbo_kb_articles"."content") AS "article_length"' in normalized - - -def test_normalize_generation_result_sql_keeps_union_limit_planner_safe_for_mssql(): - sql = """ - SELECT - LENGTH(article_text) AS article_length - FROM "dbo_kb_articles" - UNION ALL SELECT - LENGTH(article_text) AS article_length - FROM "dbo_knowledge_articles" - LIMIT 1 - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") - - assert "LIMIT" not in normalized - assert "TOP 1" not in normalized - assert 'LENGTH("content") AS article_length' in normalized - assert "UNION ALL SELECT" in normalized - - -def test_normalize_generation_result_sql_rewrites_kb_article_created_by_for_mssql(): - sql = """ - SELECT - created_by, - COUNT(*) AS article_count - FROM dbo_kb_articles - GROUP BY created_by - ORDER BY article_count DESC - """ - - normalized = normalize_generation_result_sql(sql, data_source="MSSQL") +def test_sql_generation_system_prompt_grounding_contract(): + prompt = get_sql_generation_system_prompt() - assert "created_by," not in normalized - assert "GROUP BY created_by" not in normalized - assert '"created_by_user_id"' in normalized + assert "DATABASE SCHEMA section is the complete and only source" in prompt + assert "MUST NOT introduce, infer, copy, or repair" in prompt + assert "Do not copy identifiers" in prompt diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index 561e798cc7..bb20cc8c9d 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -21,85 +21,6 @@ import { getLogger } from '@server/utils'; const logger = getLogger('TextBasedAnswerBackgroundTracker'); logger.level = 'debug'; -const ANSWER_PREVIEW_LIMIT = 500; - -const formatValue = (value: unknown) => { - if (value === null || value === undefined || value === '') { - return '(blank)'; - } - if (typeof value === 'number') { - return Number.isInteger(value) ? value.toString() : value.toLocaleString(); - } - return String(value); -}; - -const buildFastAnswer = ( - question: string, - data: PreviewDataResponse, -): string | null => { - const rows = data?.data || []; - const columns = data?.columns || []; - if (!columns.length) { - return null; - } - if (!rows.length) { - return 'No rows were returned for this question.'; - } - - const columnNames = columns.map((column) => column.name); - const rowCount = rows.length; - const sampleRows = rows.slice(0, Math.min(rowCount, 10)); - const hasMetricColumn = columns.some((column) => - /count|total|sum|amount|value|revenue|sales|qty|quantity|rate|percent/i.test( - column.name, - ), - ); - const questionPrefix = question ? `For "${question}", ` : ''; - - if (columns.length === 1) { - const values = sampleRows.map((row) => formatValue(row[0])).join(', '); - const columnName = columnNames[0]; - const isCountColumn = /count|recordcount|rowcount/i.test(columnName); - if (isCountColumn && Number(sampleRows[0]?.[0]) === 0) { - return `${questionPrefix}the active datasource returned 0 matching records.`; - } - return `${questionPrefix}the query returned ${rowCount} row${ - rowCount === 1 ? '' : 's' - }. Values: ${values}.`; - } - - if (columns.length === 2 && hasMetricColumn) { - const [labelColumn, metricColumn] = columnNames; - const lines = sampleRows.map( - (row, index) => - `${index + 1}. ${formatValue(row[0])}: ${formatValue(row[1])}`, - ); - return [ - `${questionPrefix}the top ${sampleRows.length} results by ${metricColumn} are:`, - ...lines, - `Columns used: ${labelColumn}, ${metricColumn}.`, - ].join('\n'); - } - - const preview = sampleRows - .map((row, index) => { - const values = columnNames - .map((columnName, columnIndex) => { - return `${columnName}: ${formatValue(row[columnIndex])}`; - }) - .join(', '); - return `${index + 1}. ${values}`; - }) - .join('\n'); - - return [ - `${questionPrefix}the query returned ${rowCount} row${ - rowCount === 1 ? '' : 's' - }. Showing the first ${sampleRows.length}:`, - preview, - ].join('\n'); -}; - export class TextBasedAnswerBackgroundTracker { // tasks is a kv pair of task id and thread response private tasks: Record = {}; @@ -184,7 +105,7 @@ export class TextBasedAnswerBackgroundTracker { project, manifest: mdl, modelingOnly: false, - limit: ANSWER_PREVIEW_LIMIT, + limit: 500, cacheEnabled: false, })) as PreviewDataResponse; } catch (error) { @@ -202,25 +123,6 @@ export class TextBasedAnswerBackgroundTracker { throw error; } - const fastAnswer = buildFastAnswer(threadResponse.question, data); - if (fastAnswer) { - const finishedDetail = { - ...threadResponse.answerDetail, - status: ThreadResponseAnswerStatus.FINISHED, - content: fastAnswer, - numRowsUsedInLLM: Math.min( - data?.data?.length || 0, - ANSWER_PREVIEW_LIMIT, - ), - }; - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: finishedDetail, - }); - threadResponse.answerDetail = finishedDetail; - delete this.tasks[threadResponse.id]; - return; - } - const response = await this.wrenAIAdaptor.createTextBasedAnswer({ query: threadResponse.question, sql: threadResponse.sql, diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 86e0a6c87a..5a39c65933 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -637,13 +637,10 @@ export class MDLBuilder implements IMDLBuilder { /^([^.]+)\.([^.]+)\.([^.]+)$/, ); if (catalogQualifiedMatch) { - const normalizedTableName = this.normalizeDboPrefixedTableName( - catalogQualifiedMatch[3], - ); return { catalog: catalogQualifiedMatch[1], - schema: normalizedTableName.schema || catalogQualifiedMatch[2], - table: normalizedTableName.table, + schema: catalogQualifiedMatch[2], + table: catalogQualifiedMatch[3], }; } @@ -656,38 +653,8 @@ export class MDLBuilder implements IMDLBuilder { }; } - const underscoreQualifiedMatch = sourceTableName.match(/^(dbo)_(.+)$/i); - if (underscoreQualifiedMatch) { - return { - catalog: null, - schema: - this.project.type === DataSourceName.MSSQL - ? underscoreQualifiedMatch[1] - : null, - table: underscoreQualifiedMatch[2], - }; - } - return null; } - - private normalizeDboPrefixedTableName(tableName: string): { - schema: string | null; - table: string; - } { - const underscoreQualifiedMatch = tableName.match(/^(dbo)_(.+)$/i); - if (!underscoreQualifiedMatch) { - return { schema: null, table: tableName }; - } - - return { - schema: - this.project.type === DataSourceName.MSSQL - ? underscoreQualifiedMatch[1] - : null, - table: underscoreQualifiedMatch[2], - }; - } private hasDuplicateSourceColumns(modelId: number): boolean { const sourceColumnNames = new Set(); for (const column of this.columns.filter( diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 5f41ee3f9c..930a963254 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -585,7 +585,7 @@ describe('MDLBuilder', () => { expect(manifest.models[0].refSql).toBeUndefined(); }); - it('should split mssql schema-normalized source table names in rust engine table reference fallback.', () => { + it('should preserve unqualified source table names in rust engine table reference fallback.', () => { const project = { id: 1, type: DataSourceName.MSSQL, @@ -625,13 +625,13 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: null, - schema: 'dbo', - table: 'tickets', + schema: null, + table: 'dbo_tickets', }); expect(manifest.models[0].refSql).toBeUndefined(); }); - it('should strip dbo-prefixed source table names for non-mssql projects.', () => { + it('should preserve dbo-prefixed source table names for non-mssql projects.', () => { const project = { id: 1, type: DataSourceName.POSTGRES, @@ -672,12 +672,12 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: null, schema: null, - table: 'search_queries', + table: 'dbo_search_queries', }); expect(manifest.models[0].refSql).toBeUndefined(); }); - it('should strip dbo-prefixed property table names before project schema fallback.', () => { + it('should preserve dbo-prefixed property table names before project schema fallback.', () => { const project = { id: 1, type: DataSourceName.POSTGRES, @@ -721,12 +721,12 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: null, schema: 'public', - table: 'search_queries', + table: 'dbo_search_queries', }); expect(manifest.models[0].refSql).toBeUndefined(); }); - it('should strip dbo-prefixed table names from catalog-qualified non-mssql table references.', () => { + it('should preserve catalog-qualified non-mssql table references.', () => { const project = { id: 1, type: DataSourceName.POSTGRES, @@ -769,7 +769,7 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: 'wrenai', schema: 'public', - table: 'search_queries', + table: 'dbo_search_queries', }); expect(manifest.models[0].refSql).toBeUndefined(); }); diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 55718b1174..4993320ec0 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -5,7 +5,6 @@ import { AskCandidateType, RecommendationQuestionsResult, RecommendationQuestionsInput, - RecommendationQuestion, WrenAIError, RecommendationQuestionStatus, ChartStatus, @@ -43,7 +42,6 @@ import { TrackedAdjustmentResult, } from '../backgrounds'; import { getConfig } from '@server/config'; -import { buildFastRecommendationQuestions } from '../utils/recommendationQuestions'; import { TextBasedAnswerBackgroundTracker } from '../backgrounds/textBasedAnswerBackgroundTracker'; import { IAskingTaskTracker, TrackedAskingResult } from './askingTaskTracker'; @@ -447,10 +445,6 @@ export class AskingService implements IAskingService { private askingTaskRepository: IAskingTaskRepository; private adjustmentBackgroundTracker: AdjustmentBackgroundTaskTracker; private instantRecommendationJobs = new Map>(); - private instantRecommendationResults = new Map< - string, - RecommendationQuestionsResult - >(); private threadRecommendationJobs = new Map>(); private initialized = false; @@ -611,26 +605,6 @@ export class AskingService implements IAskingService { .sort((a, b) => b.id - a.id) .slice(0, 5); const questions = slicedThreadResponses.map(({ question }) => question); - const fastQuestions = await this.filterExecutableRecommendationQuestions( - buildFastRecommendationQuestions( - manifest, - this.getThreadRecommendationQuestionsConfig(project).maxQuestions, - questions, - ), - project, - manifest, - this.getThreadRecommendationQuestionsConfig(project).maxQuestions, - ); - if (fastQuestions.length) { - await this.threadRepository.updateOne(threadId, { - queryId: `fast-thread-${threadId}-${Date.now()}`, - questionsStatus: RecommendationQuestionStatus.FINISHED, - questions: fastQuestions, - questionsError: null, - }); - return; - } - const recommendQuestionData: RecommendationQuestionsInput = { manifest, projectId: project.id.toString(), @@ -1243,31 +1217,6 @@ export class AskingService implements IAskingService { currentProject.id, ); - const fastQuestions = await this.filterExecutableRecommendationQuestions( - buildFastRecommendationQuestions( - manifest, - this.getThreadRecommendationQuestionsConfig(currentProject).maxQuestions, - input.previousQuestions || [], - ), - currentProject, - manifest, - this.getThreadRecommendationQuestionsConfig(currentProject).maxQuestions, - ); - if (fastQuestions.length) { - const queryId = `fast-instant-${currentProject.id}-${Date.now()}`; - this.instantRecommendationResults.set(queryId, { - status: RecommendationQuestionStatus.FINISHED, - type: null, - response: { questions: fastQuestions }, - error: null, - }); - setTimeout( - () => this.instantRecommendationResults.delete(queryId), - 5 * 60 * 1000, - ); - return { id: queryId }; - } - const response = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, projectId: currentProject.id.toString(), @@ -1280,47 +1229,11 @@ export class AskingService implements IAskingService { public async getInstantRecommendedQuestions( queryId: string, ): Promise { - const localResult = this.instantRecommendationResults.get(queryId); - if (localResult) { - return localResult; - } - const response = await this.wrenAIAdaptor.getRecommendationQuestionsResult(queryId); return response; } - private async filterExecutableRecommendationQuestions( - questions: RecommendationQuestion[], - project: Project, - manifest: any, - maxQuestions: number, - ): Promise { - const validQuestions: RecommendationQuestion[] = []; - for (const question of questions) { - try { - const result = (await this.queryService.preview(question.sql, { - project, - manifest, - modelingOnly: false, - limit: 1, - cacheEnabled: false, - })) as PreviewDataResponse; - if (result?.data?.length) { - validQuestions.push(question); - if (validQuestions.length >= maxQuestions) { - break; - } - } - } catch (error) { - logger.warn( - `Skipping recommended question because SQL preview failed: ${question.question}. ${error}`, - ); - } - } - return validQuestions; - } - public async deleteAllByProjectId(projectId: number): Promise { // delete all threads await this.threadRepository.deleteAllBy({ projectId }); diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index f27ca78255..7d63f0f528 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -27,8 +27,6 @@ import { IMDLService } from './mdlService'; import { ProjectRecommendQuestionBackgroundTracker } from '../backgrounds'; import { ITelemetry } from '../telemetry/telemetry'; import { getConfig } from '../config'; -import { buildFastRecommendationQuestions } from '../utils/recommendationQuestions'; -import { IQueryService, PreviewDataResponse } from './queryService'; const config = getConfig(); @@ -95,7 +93,6 @@ export class ProjectService implements IProjectService { private projectRepository: IProjectRepository; private metadataService: IDataSourceMetadataService; private mdlService: IMDLService; - private queryService: IQueryService; private wrenAIAdaptor: IWrenAIAdaptor; private projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; private projectRecommendationJobs = new Map>(); @@ -103,7 +100,6 @@ export class ProjectService implements IProjectService { projectRepository, metadataService, mdlService, - queryService, wrenAIAdaptor, telemetry, projectRecommendQuestionBackgroundTracker, @@ -111,7 +107,6 @@ export class ProjectService implements IProjectService { projectRepository: IProjectRepository; metadataService: IDataSourceMetadataService; mdlService: IMDLService; - queryService: IQueryService; wrenAIAdaptor: IWrenAIAdaptor; telemetry: ITelemetry; projectRecommendQuestionBackgroundTracker?: ProjectRecommendQuestionBackgroundTracker; @@ -119,7 +114,6 @@ export class ProjectService implements IProjectService { this.projectRepository = projectRepository; this.metadataService = metadataService; this.mdlService = mdlService; - this.queryService = queryService; this.wrenAIAdaptor = wrenAIAdaptor; this.projectRecommendQuestionBackgroundTracker = projectRecommendQuestionBackgroundTracker ?? @@ -180,25 +174,6 @@ export class ProjectService implements IProjectService { project: Project, ): Promise { const { manifest } = await this.mdlService.makeModelMDL(project); - const fastQuestions = await this.filterExecutableRecommendationQuestions( - buildFastRecommendationQuestions( - manifest, - this.getProjectRecommendationQuestionsConfig(project).maxQuestions, - ), - project, - manifest, - this.getProjectRecommendationQuestionsConfig(project).maxQuestions, - ); - if (fastQuestions.length) { - await this.projectRepository.updateOne(project.id, { - queryId: `fast-project-${project.id}-${Date.now()}`, - questionsStatus: RecommendationQuestionStatus.FINISHED, - questions: fastQuestions, - questionsError: null, - }); - return; - } - const recommendQuestionResult = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, @@ -369,37 +344,6 @@ export class ProjectService implements IProjectService { ); } - private async filterExecutableRecommendationQuestions( - questions: RecommendationQuestion[], - project: Project, - manifest: any, - maxQuestions: number, - ): Promise { - const validQuestions: RecommendationQuestion[] = []; - for (const question of questions) { - try { - const result = (await this.queryService.preview(question.sql, { - project, - manifest, - modelingOnly: false, - limit: 1, - cacheEnabled: false, - })) as PreviewDataResponse; - if (result?.data?.length) { - validQuestions.push(question); - if (validQuestions.length >= maxQuestions) { - break; - } - } - } catch (error) { - logger.warn( - `Skipping project recommended question because SQL preview failed: ${question.question}. ${error}`, - ); - } - } - return validQuestions; - } - private getProjectRecommendationQuestionsConfig(project: Project) { return { maxCategories: config.projectRecommendationQuestionMaxCategories, diff --git a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts b/wren-ui/src/apollo/server/utils/recommendationQuestions.ts deleted file mode 100644 index 4b145ed362..0000000000 --- a/wren-ui/src/apollo/server/utils/recommendationQuestions.ts +++ /dev/null @@ -1,148 +0,0 @@ -import { RecommendationQuestion } from '@server/models/adaptor'; -import { Manifest, ModelMDL, ColumnMDL } from '@server/mdl/type'; - -const MAX_MODEL_COUNT = 6; -type RecommendationModel = Partial & { - name: string; - columns?: Partial[]; -}; - -const quoteIdentifier = (identifier: unknown) => { - return `"${String(identifier || '').replace(/"/g, '""')}"`; -}; - -const isMetricColumn = (column: Partial) => { - const type = String(column.type || '').toLowerCase(); - return /int|bigint|smallint|tinyint|float|double|decimal|numeric|number|real|money/.test( - type, - ); -}; - -const isDateColumn = (column: Partial) => { - const type = String(column.type || '').toLowerCase(); - return /date|time|timestamp/.test(type); -}; - -const isDimensionColumn = (column: Partial) => { - const name = String(column.name || '').toLowerCase(); - if ( - !column.name || - column.isCalculated || - isMetricColumn(column) || - isDateColumn(column) - ) { - return false; - } - return !/(^id$|_id$|uuid|guid|password|token|secret|json|payload)/.test(name); -}; - -const displayName = (model: RecommendationModel) => { - return model.properties?.displayName || model.name; -}; - -const firstUsableModels = (manifest: Manifest) => { - return (manifest.models || []) - .filter((model): model is RecommendationModel => Boolean(model?.name)) - .filter((model) => (model.columns || []).some((column) => column.name)) - .slice(0, MAX_MODEL_COUNT); -}; - -export const buildFastRecommendationQuestions = ( - manifest: Manifest, - maxQuestions = 5, - previousQuestions: string[] = [], -): RecommendationQuestion[] => { - const candidateLimit = Math.max(maxQuestions * 3, maxQuestions); - const seen = new Set( - previousQuestions.map((question) => question.trim().toLowerCase()), - ); - const questions: RecommendationQuestion[] = []; - const addQuestion = (question: RecommendationQuestion) => { - if (questions.length >= candidateLimit) { - return; - } - const key = question.question.trim().toLowerCase(); - if (seen.has(key)) { - return; - } - seen.add(key); - questions.push(question); - }; - - for (const model of firstUsableModels(manifest)) { - const modelRef = quoteIdentifier(model.name); - const columns = model.columns || []; - const dimensions = columns.filter(isDimensionColumn); - const metrics = columns.filter(isMetricColumn); - const dates = columns.filter(isDateColumn); - const label = displayName(model); - - addQuestion({ - category: label, - question: `How many records are in ${label}?`, - sql: `SELECT COUNT(*) AS "RecordCount" FROM ${modelRef}`, - }); - - if (dimensions[0]) { - const column = dimensions[0]; - const columnRef = `${modelRef}.${quoteIdentifier(column.name)}`; - addQuestion({ - category: label, - question: `What is the distribution of ${column.name} in ${label}?`, - sql: - `SELECT ${columnRef} AS ${quoteIdentifier(column.name)}, ` + - `COUNT(*) AS "RecordCount" FROM ${modelRef} ` + - `GROUP BY ${columnRef} ORDER BY COUNT(*) DESC`, - }); - } - - if (dimensions[0] && metrics[0]) { - const dimension = dimensions[0]; - const metric = metrics[0]; - const dimensionRef = `${modelRef}.${quoteIdentifier(dimension.name)}`; - const metricRef = `${modelRef}.${quoteIdentifier(metric.name)}`; - addQuestion({ - category: label, - question: `Which ${dimension.name} values have the highest ${metric.name} in ${label}?`, - sql: - `SELECT ${dimensionRef} AS ${quoteIdentifier(dimension.name)}, ` + - `SUM(${metricRef}) AS ${quoteIdentifier(`Total${metric.name}`)} ` + - `FROM ${modelRef} GROUP BY ${dimensionRef} ` + - `ORDER BY SUM(${metricRef}) DESC`, - }); - } - - if (dates[0]) { - const date = dates[0]; - const dateRef = `${modelRef}.${quoteIdentifier(date.name)}`; - addQuestion({ - category: label, - question: `Show monthly record count by ${date.name} in ${label}.`, - sql: - `SELECT DATEPART(YEAR, ${dateRef}) AS "year", ` + - `DATEPART(MONTH, ${dateRef}) AS "month", ` + - `COUNT(*) AS "RecordCount" FROM ${modelRef} ` + - `GROUP BY DATEPART(YEAR, ${dateRef}), DATEPART(MONTH, ${dateRef}) ` + - `ORDER BY DATEPART(YEAR, ${dateRef}), DATEPART(MONTH, ${dateRef})`, - }); - } - - const previewColumns = columns - .filter((column) => column.name) - .slice(0, 8) - .map( - (column) => - `${modelRef}.${quoteIdentifier(column.name)} AS ${quoteIdentifier( - column.name, - )}`, - ); - - addQuestion({ - category: label, - question: `Show the first 10 rows from ${label}.`, - sql: `SELECT ${previewColumns.join(', ')} FROM ${modelRef} LIMIT 10`, - }); - } - - return questions; -}; From a0205dc86b3b92195841eb4adeb9ca672fdb440d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 22:37:58 +0530 Subject: [PATCH 0660/1087] Use retrieved schema only for executable SQL --- .../generation/question_recommendation.py | 4 + .../src/pipelines/generation/utils/sql.py | 2 + wren-ai-service/src/web/v1/services/ask.py | 228 ++++++++---------- .../src/web/v1/services/ask_feedback.py | 34 +-- .../v1/services/question_recommendation.py | 27 +-- 5 files changed, 118 insertions(+), 177 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index a6e7c17b02..0d8abab7a7 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -66,6 +66,8 @@ 5. **General Guidelines for All Questions:** - Ensure questions can be answered using the data model. + - Use only the models, columns, relationships, and business concepts present in the DATABASE SCHEMA. + - Do not invent generic business entities, tables, columns, metrics, or topics that are absent from the DATABASE SCHEMA. - Mix simple and complex questions. - Avoid open-ended questions - each should have a definite answer. - Incorporate time-based analysis where relevant. @@ -156,6 +158,8 @@ {% endfor %} {% endif %} +Use the DATABASE SCHEMA as the only source for question topics. Generated questions must be answerable from the provided schema without inventing tables, columns, metrics, or entities. + Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based on the provided data model. Both the questions and category names should be translated into {{language}}{% if user_question %} and be related to the user's question{% endif %}. The output format should maintain the structure but with localized text. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index eccf4b8a31..865f4fd0b3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -166,8 +166,10 @@ async def _classify_generation_result( _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. +- The CREATE TABLE and CREATE VIEW statements in DATABASE SCHEMA describe metadata only. NEVER output CREATE, ALTER, DROP, INSERT, UPDATE, DELETE, MERGE, or any other DDL/DML statement. - ONLY USE the tables and columns mentioned in the database schema. - ONLY USE "*" if the user query asks for all the columns of a table. +- ONLY SELECT columns and aggregations that are required to answer the user's question. Do not include unrelated columns. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 844330eac7..6f6ab5bafb 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -192,30 +192,7 @@ async def ask( is_followup=True if histories else False, ) - historical_question = await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ) - - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - - if historical_question_result: - api_results = [ - AskResult( - **{ - "sql": result.get("statement"), - "type": "view" if result.get("viewId") else "llm", - "viewId": result.get("viewId"), - } - ) - for result in historical_question_result - ] - sql_generation_reasoning = "" - else: - # Run both pipeline operations concurrently + if self._allow_intent_classification: sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( query=user_query, @@ -236,106 +213,105 @@ async def ask( "documents", [] ) - if self._allow_intent_classification: - intent_classification_result = ( - await self._pipelines["intent_classification"].run( + intent_classification_result = ( + await self._pipelines["intent_classification"].run( + query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, + ) + ).get("post_process", {}) + intent = intent_classification_result.get("intent") + rephrased_question = intent_classification_result.get( + "rephrased_question" + ) + intent_reasoning = intent_classification_result.get("reasoning") + + if rephrased_question: + user_query = rephrased_question + + if intent == "MISLEADING_QUERY": + asyncio.create_task( + self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, ) - ).get("post_process", {}) - intent = intent_classification_result.get("intent") - rephrased_question = intent_classification_result.get( - "rephrased_question" ) - intent_reasoning = intent_classification_result.get("reasoning") - - if rephrased_question: - user_query = rephrased_question - - if intent == "MISLEADING_QUERY": - asyncio.create_task( - self._pipelines["misleading_assistance"].run( - query=user_query, - histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, - ) - ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="MISLEADING_QUERY", - ) - results["metadata"]["type"] = "MISLEADING_QUERY" - return results - elif intent == "GENERAL": - asyncio.create_task( - self._pipelines["data_assistance"].run( - query=user_query, - histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="MISLEADING_QUERY", + ) + results["metadata"]["type"] = "MISLEADING_QUERY" + return results + elif intent == "GENERAL": + asyncio.create_task( + self._pipelines["data_assistance"].run( + query=user_query, + histories=histories, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, ) + ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - results["metadata"]["type"] = "GENERAL" - return results - elif intent == "USER_GUIDE": - asyncio.create_task( - self._pipelines["user_guide_assistance"].run( - query=user_query, - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + results["metadata"]["type"] = "GENERAL" + return results + elif intent == "USER_GUIDE": + asyncio.create_task( + self._pipelines["user_guide_assistance"].run( + query=user_query, + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, ) + ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="USER_GUIDE", - ) - results["metadata"]["type"] = "GENERAL" - return results - else: - self._ask_results[query_id] = AskResultResponse( - status="understanding", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="USER_GUIDE", + ) + results["metadata"]["type"] = "GENERAL" + return results + else: + self._ask_results[query_id] = AskResultResponse( + status="understanding", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="searching", @@ -399,8 +375,8 @@ async def ask( query=user_query, contexts=table_ddls, histories=histories, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], configuration=ask_request.configurations, query_id=query_id, ) @@ -410,8 +386,8 @@ async def ask( await self._pipelines["sql_generation_reasoning"].run( query=user_query, contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], configuration=ask_request.configurations, query_id=query_id, ) @@ -468,11 +444,11 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, histories=histories, project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -487,10 +463,10 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -552,7 +528,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - instructions=instructions, + instructions=[], invalid_generation_result={ "sql": original_sql, "error": sql_diagnosis_reasoning diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 25044de18c..b976ebbb28 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -1,4 +1,3 @@ -import asyncio import logging from typing import Dict, List, Literal, Optional @@ -114,24 +113,9 @@ async def ask_feedback( trace_id=trace_id, ) - ( - retrieval_task, - sql_samples_task, - instructions_task, - ) = await asyncio.gather( - self._pipelines["db_schema_retrieval"].run( - tables=ask_feedback_request.tables, - project_id=ask_feedback_request.project_id, - ), - self._pipelines["sql_pairs_retrieval"].run( - query=ask_feedback_request.question, - project_id=ask_feedback_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=ask_feedback_request.question, - project_id=ask_feedback_request.project_id, - scope="sql", - ), + retrieval_task = await self._pipelines["db_schema_retrieval"].run( + tables=ask_feedback_request.tables, + project_id=ask_feedback_request.project_id, ) if allow_sql_functions_retrieval: @@ -161,10 +145,6 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - 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_feedback_results): self._ask_feedback_results[query_id] = AskFeedbackResultResponse( @@ -176,11 +156,11 @@ async def ask_feedback( "sql_regeneration" ].run( contexts=table_ddls, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + sql_generation_reasoning=None, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -239,7 +219,7 @@ async def ask_feedback( "sql_correction" ].run( contexts=table_ddls, - instructions=instructions, + instructions=[], invalid_generation_result={ "original_sql": original_sql, "sql": invalid_sql, diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 6033237a45..0fb8e5711b 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -82,29 +82,8 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: has_json_field = _retrieval_result.get("has_json_field", False) return table_ddls, has_calculated_field, has_metric, has_json_field - async def _sql_pairs_retrieval() -> list[dict]: - sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( - query=candidate["question"], - project_id=project_id, - ) - sql_samples = sql_pairs_result["formatted_output"].get("documents", []) - return sql_samples - - async def _instructions_retrieval() -> list[dict]: - result = await self._pipelines["instructions_retrieval"].run( - query=candidate["question"], - project_id=project_id, - scope="sql", - ) - instructions = result["formatted_output"].get("instructions", []) - return instructions - try: - _document, sql_samples, instructions = await asyncio.gather( - _document_retrieval(), - _sql_pairs_retrieval(), - _instructions_retrieval(), - ) + _document = await _document_retrieval() table_ddls, has_calculated_field, has_metric, has_json_field = _document if self._allow_sql_functions_retrieval: @@ -125,8 +104,8 @@ async def _instructions_retrieval() -> list[dict]: query=candidate["question"], contexts=table_ddls, project_id=project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, From 97ab2d9950df9d9a6b572e7d0c115062dc579e4b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 22:42:53 +0530 Subject: [PATCH 0661/1087] Revert "Use retrieved schema only for executable SQL" This reverts commit a0205dc86b3b92195841eb4adeb9ca672fdb440d. --- .../generation/question_recommendation.py | 4 - .../src/pipelines/generation/utils/sql.py | 2 - wren-ai-service/src/web/v1/services/ask.py | 228 ++++++++++-------- .../src/web/v1/services/ask_feedback.py | 34 ++- .../v1/services/question_recommendation.py | 27 ++- 5 files changed, 177 insertions(+), 118 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index 0d8abab7a7..a6e7c17b02 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -66,8 +66,6 @@ 5. **General Guidelines for All Questions:** - Ensure questions can be answered using the data model. - - Use only the models, columns, relationships, and business concepts present in the DATABASE SCHEMA. - - Do not invent generic business entities, tables, columns, metrics, or topics that are absent from the DATABASE SCHEMA. - Mix simple and complex questions. - Avoid open-ended questions - each should have a definite answer. - Incorporate time-based analysis where relevant. @@ -158,8 +156,6 @@ {% endfor %} {% endif %} -Use the DATABASE SCHEMA as the only source for question topics. Generated questions must be answerable from the provided schema without inventing tables, columns, metrics, or entities. - Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based on the provided data model. Both the questions and category names should be translated into {{language}}{% if user_question %} and be related to the user's question{% endif %}. The output format should maintain the structure but with localized text. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 865f4fd0b3..eccf4b8a31 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -166,10 +166,8 @@ async def _classify_generation_result( _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. -- The CREATE TABLE and CREATE VIEW statements in DATABASE SCHEMA describe metadata only. NEVER output CREATE, ALTER, DROP, INSERT, UPDATE, DELETE, MERGE, or any other DDL/DML statement. - ONLY USE the tables and columns mentioned in the database schema. - ONLY USE "*" if the user query asks for all the columns of a table. -- ONLY SELECT columns and aggregations that are required to answer the user's question. Do not include unrelated columns. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6f6ab5bafb..844330eac7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -192,7 +192,30 @@ async def ask( is_followup=True if histories else False, ) - if self._allow_intent_classification: + historical_question = await self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ) + + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] + + if historical_question_result: + api_results = [ + AskResult( + **{ + "sql": result.get("statement"), + "type": "view" if result.get("viewId") else "llm", + "viewId": result.get("viewId"), + } + ) + for result in historical_question_result + ] + sql_generation_reasoning = "" + else: + # Run both pipeline operations concurrently sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( query=user_query, @@ -213,105 +236,106 @@ async def ask( "documents", [] ) - intent_classification_result = ( - await self._pipelines["intent_classification"].run( - query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, - ) - ).get("post_process", {}) - intent = intent_classification_result.get("intent") - rephrased_question = intent_classification_result.get( - "rephrased_question" - ) - intent_reasoning = intent_classification_result.get("reasoning") - - if rephrased_question: - user_query = rephrased_question - - if intent == "MISLEADING_QUERY": - asyncio.create_task( - self._pipelines["misleading_assistance"].run( + if self._allow_intent_classification: + intent_classification_result = ( + await self._pipelines["intent_classification"].run( query=user_query, histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, ) + ).get("post_process", {}) + intent = intent_classification_result.get("intent") + rephrased_question = intent_classification_result.get( + "rephrased_question" ) + intent_reasoning = intent_classification_result.get("reasoning") + + if rephrased_question: + user_query = rephrased_question + + if intent == "MISLEADING_QUERY": + asyncio.create_task( + self._pipelines["misleading_assistance"].run( + query=user_query, + histories=histories, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) + ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="MISLEADING_QUERY", - ) - results["metadata"]["type"] = "MISLEADING_QUERY" - return results - elif intent == "GENERAL": - asyncio.create_task( - self._pipelines["data_assistance"].run( - query=user_query, - histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="MISLEADING_QUERY", + ) + results["metadata"]["type"] = "MISLEADING_QUERY" + return results + elif intent == "GENERAL": + asyncio.create_task( + self._pipelines["data_assistance"].run( + query=user_query, + histories=histories, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) ) - ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - results["metadata"]["type"] = "GENERAL" - return results - elif intent == "USER_GUIDE": - asyncio.create_task( - self._pipelines["user_guide_assistance"].run( - query=user_query, - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + results["metadata"]["type"] = "GENERAL" + return results + elif intent == "USER_GUIDE": + asyncio.create_task( + self._pipelines["user_guide_assistance"].run( + query=user_query, + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) ) - ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="USER_GUIDE", - ) - results["metadata"]["type"] = "GENERAL" - return results - else: - self._ask_results[query_id] = AskResultResponse( - status="understanding", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="USER_GUIDE", + ) + results["metadata"]["type"] = "GENERAL" + return results + else: + self._ask_results[query_id] = AskResultResponse( + status="understanding", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="searching", @@ -375,8 +399,8 @@ async def ask( query=user_query, contexts=table_ddls, histories=histories, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, configuration=ask_request.configurations, query_id=query_id, ) @@ -386,8 +410,8 @@ async def ask( await self._pipelines["sql_generation_reasoning"].run( query=user_query, contexts=table_ddls, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, configuration=ask_request.configurations, query_id=query_id, ) @@ -444,11 +468,11 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -463,10 +487,10 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -528,7 +552,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - instructions=[], + instructions=instructions, invalid_generation_result={ "sql": original_sql, "error": sql_diagnosis_reasoning diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index b976ebbb28..25044de18c 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -1,3 +1,4 @@ +import asyncio import logging from typing import Dict, List, Literal, Optional @@ -113,9 +114,24 @@ async def ask_feedback( trace_id=trace_id, ) - retrieval_task = await self._pipelines["db_schema_retrieval"].run( - tables=ask_feedback_request.tables, - project_id=ask_feedback_request.project_id, + ( + retrieval_task, + sql_samples_task, + instructions_task, + ) = await asyncio.gather( + self._pipelines["db_schema_retrieval"].run( + tables=ask_feedback_request.tables, + project_id=ask_feedback_request.project_id, + ), + self._pipelines["sql_pairs_retrieval"].run( + query=ask_feedback_request.question, + project_id=ask_feedback_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=ask_feedback_request.question, + project_id=ask_feedback_request.project_id, + scope="sql", + ), ) if allow_sql_functions_retrieval: @@ -145,6 +161,10 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + 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_feedback_results): self._ask_feedback_results[query_id] = AskFeedbackResultResponse( @@ -156,11 +176,11 @@ async def ask_feedback( "sql_regeneration" ].run( contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -219,7 +239,7 @@ async def ask_feedback( "sql_correction" ].run( contexts=table_ddls, - instructions=[], + instructions=instructions, invalid_generation_result={ "original_sql": original_sql, "sql": invalid_sql, diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 0fb8e5711b..6033237a45 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -82,8 +82,29 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: has_json_field = _retrieval_result.get("has_json_field", False) return table_ddls, has_calculated_field, has_metric, has_json_field + async def _sql_pairs_retrieval() -> list[dict]: + sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( + query=candidate["question"], + project_id=project_id, + ) + sql_samples = sql_pairs_result["formatted_output"].get("documents", []) + return sql_samples + + async def _instructions_retrieval() -> list[dict]: + result = await self._pipelines["instructions_retrieval"].run( + query=candidate["question"], + project_id=project_id, + scope="sql", + ) + instructions = result["formatted_output"].get("instructions", []) + return instructions + try: - _document = await _document_retrieval() + _document, sql_samples, instructions = await asyncio.gather( + _document_retrieval(), + _sql_pairs_retrieval(), + _instructions_retrieval(), + ) table_ddls, has_calculated_field, has_metric, has_json_field = _document if self._allow_sql_functions_retrieval: @@ -104,8 +125,8 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: query=candidate["question"], contexts=table_ddls, project_id=project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, From 2f4e373a391f1e9ff671a5651fcbb11267f31011 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 23:12:21 +0530 Subject: [PATCH 0662/1087] Align ask SQL flow with legacy grounding --- docker/config.example.yaml | 2 +- wren-ai-service/src/config.py | 5 +- .../retrieval/db_schema_retrieval.py | 32 +-- wren-ai-service/src/web/v1/services/ask.py | 251 +++++++++--------- 4 files changed, 128 insertions(+), 162 deletions(-) diff --git a/docker/config.example.yaml b/docker/config.example.yaml index aa81c74c60..a22444f84c 100644 --- a/docker/config.example.yaml +++ b/docker/config.example.yaml @@ -175,7 +175,7 @@ settings: column_indexing_batch_size: 50 table_retrieval_size: 10 table_column_retrieval_size: 100 - allow_intent_classification: false + allow_intent_classification: true allow_sql_generation_reasoning: true allow_sql_functions_retrieval: true enable_column_pruning: false diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index bc97d8cc71..c5acf4ae47 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -37,16 +37,13 @@ class Settings(BaseSettings): instructions_top_k: int = Field(default=10) # generation config - allow_intent_classification: bool = Field(default=False) + allow_intent_classification: bool = Field(default=True) allow_sql_generation_reasoning: bool = Field(default=True) allow_sql_functions_retrieval: bool = Field(default=True) allow_sql_diagnosis: bool = Field(default=True) allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) max_sql_correction_retries: int = Field(default=3) - max_sql_generation_tables: int = Field(default=0) - pipeline_timeout_seconds: int = Field(default=0) - schema_retrieval_timeout_seconds: int = Field(default=0) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 28df670a57..6c8dd7bbe3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -120,14 +120,6 @@ def _build_view_ddl(content: dict) -> str: ) -def _table_names_from_retrieval(table_retrieval: dict) -> list[str]: - table_names = [] - for table in table_retrieval.get("documents", []): - content = ast.literal_eval(table.content) - table_names.append(content["name"]) - return table_names - - ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -180,7 +172,11 @@ async def table_retrieval( async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any ) -> list[Document]: - table_names = _table_names_from_retrieval(table_retrieval) + tables = table_retrieval.get("documents", []) + table_names = [] + for table in tables: + content = ast.literal_eval(table.content) + table_names.append(content["name"]) table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} @@ -208,9 +204,7 @@ async def dbschema_retrieval( @observe() -def construct_db_schemas( - dbschema_retrieval: list[Document], table_retrieval: dict -) -> list[dict]: +def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: db_schemas = {} for document in dbschema_retrieval: content = ast.literal_eval(document.content) @@ -234,19 +228,7 @@ def construct_db_schemas( # remove incomplete schemas db_schemas = {k: v for k, v in db_schemas.items() if "type" in v and "columns" in v} - ranked_table_names = _table_names_from_retrieval(table_retrieval) - ranked_schemas = [ - db_schemas[table_name] - for table_name in ranked_table_names - if table_name in db_schemas - ] - remaining_schemas = [ - table_schema - for table_name, table_schema in db_schemas.items() - if table_name not in set(ranked_table_names) - ] - - return ranked_schemas + remaining_schemas + return list(db_schemas.values()) @observe(capture_input=False) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 844330eac7..fc85ff4d60 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -105,9 +105,6 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 3, - max_sql_generation_tables: int = 0, - pipeline_timeout_seconds: int = 0, - schema_retrieval_timeout_seconds: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -203,139 +200,129 @@ async def ask( ).get("documents", [])[:1] if historical_question_result: - api_results = [ - AskResult( - **{ - "sql": result.get("statement"), - "type": "view" if result.get("viewId") else "llm", - "viewId": result.get("viewId"), - } - ) - for result in historical_question_result - ] sql_generation_reasoning = "" - else: - # Run both pipeline operations concurrently - sql_samples_task, instructions_task = await asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( + + # Run both pipeline operations concurrently. These are still used by + # the legacy intent path, but executable SQL generation below is + # grounded only by the retrieved schema for the current deployment. + sql_samples_task, instructions_task = await asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), + ) + + # Extract results from completed tasks + sql_samples = sql_samples_task["formatted_output"].get("documents", []) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) + + if self._allow_intent_classification: + intent_classification_result = ( + await self._pipelines["intent_classification"].run( query=user_query, + histories=histories, + sql_samples=sql_samples, + instructions=instructions, project_id=ask_request.project_id, - scope="sql", - ), + configuration=ask_request.configurations, + ) + ).get("post_process", {}) + intent = intent_classification_result.get("intent") + rephrased_question = intent_classification_result.get( + "rephrased_question" ) + intent_reasoning = intent_classification_result.get("reasoning") - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get( - "documents", [] - ) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) + if rephrased_question: + user_query = rephrased_question - if self._allow_intent_classification: - intent_classification_result = ( - await self._pipelines["intent_classification"].run( + if intent == "MISLEADING_QUERY": + asyncio.create_task( + self._pipelines["misleading_assistance"].run( query=user_query, histories=histories, - sql_samples=sql_samples, - instructions=instructions, - project_id=ask_request.project_id, - configuration=ask_request.configurations, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, ) - ).get("post_process", {}) - intent = intent_classification_result.get("intent") - rephrased_question = intent_classification_result.get( - "rephrased_question" ) - intent_reasoning = intent_classification_result.get("reasoning") - - if rephrased_question: - user_query = rephrased_question - - if intent == "MISLEADING_QUERY": - asyncio.create_task( - self._pipelines["misleading_assistance"].run( - query=user_query, - histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, - ) - ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="MISLEADING_QUERY", - ) - results["metadata"]["type"] = "MISLEADING_QUERY" - return results - elif intent == "GENERAL": - asyncio.create_task( - self._pipelines["data_assistance"].run( - query=user_query, - histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="MISLEADING_QUERY", + ) + results["metadata"]["type"] = "MISLEADING_QUERY" + return results + elif intent == "GENERAL": + asyncio.create_task( + self._pipelines["data_assistance"].run( + query=user_query, + histories=histories, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, ) + ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - results["metadata"]["type"] = "GENERAL" - return results - elif intent == "USER_GUIDE": - asyncio.create_task( - self._pipelines["user_guide_assistance"].run( - query=user_query, - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + results["metadata"]["type"] = "GENERAL" + return results + elif intent == "USER_GUIDE": + asyncio.create_task( + self._pipelines["user_guide_assistance"].run( + query=user_query, + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, ) + ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="USER_GUIDE", - ) - results["metadata"]["type"] = "GENERAL" - return results - else: - self._ask_results[query_id] = AskResultResponse( - status="understanding", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="USER_GUIDE", + ) + results["metadata"]["type"] = "GENERAL" + return results + else: + self._ask_results[query_id] = AskResultResponse( + status="understanding", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="searching", @@ -399,8 +386,8 @@ async def ask( query=user_query, contexts=table_ddls, histories=histories, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], configuration=ask_request.configurations, query_id=query_id, ) @@ -410,8 +397,8 @@ async def ask( await self._pipelines["sql_generation_reasoning"].run( query=user_query, contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], configuration=ask_request.configurations, query_id=query_id, ) @@ -468,11 +455,11 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, histories=histories, project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -487,10 +474,10 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -552,7 +539,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - instructions=instructions, + instructions=[], invalid_generation_result={ "sql": original_sql, "error": sql_diagnosis_reasoning From 52c0d51bd5f23cfb2a60ca60d71ae317a90b732c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sat, 25 Jul 2026 23:33:07 +0530 Subject: [PATCH 0663/1087] Remove generic identifiers from SQL prompts --- .../src/pipelines/generation/utils/sql.py | 139 ------------------ 1 file changed, 139 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index eccf4b8a31..9eb56d606a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -176,7 +176,6 @@ async def _classify_generation_result( - Put double quotes around column and table names. - Put single quotes around string literals. - Never quote numeric literals. - For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992; - YOU MUST USE "lower(.) like lower()" function or "lower(.) = lower()" function for case-insensitive comparison! - Use "lower(.) LIKE lower()" when: - The user requests a pattern or partial match. @@ -186,30 +185,12 @@ async def _classify_generation_result( - The user requests an exact, specific value. - There is no ambiguity or pattern in the value. - If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query - - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds - - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds - - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds - ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query - - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) - If the user asks for a specific date, please give the date range in SQL query - - example: "What is the total revenue for the month of 2024-11-01?" - - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. - Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. - - EXAMPLE - DATABASE SCHEMA - /* {"alias":"_orders","description":"A model representing the orders data."} */ - CREATE TABLE orders ( - -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} - ApprovedTimestamp TIMESTAMP - } - - SQL - SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -230,44 +211,6 @@ async def _classify_generation_result( First, provide a brief explanation of what each field represents in the context of the schema, including how each field is computed using the relationships between models. Then, during the following tasks, if the user queries pertain to any calculated fields defined in the database schema, ensure to utilize those calculated fields appropriately in the output SQL queries. The goal is to accurately reflect the intent of the question in the SQL syntax, leveraging the pre-computed logic embedded within the calculated fields. - -EXAMPLES: -The given schema is created by the SQL command: - -CREATE TABLE orders ( - OrderId VARCHAR PRIMARY KEY, - CustomerId VARCHAR, - -- This column is a Calculated Field - -- column expression: avg(reviews.Score) - Rating DOUBLE, - -- This column is a Calculated Field - -- column expression: count(reviews.Id) - ReviewCount BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) - Size BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) > 1 - Large BOOLEAN, - FOREIGN KEY (CustomerId) REFERENCES customers(Id) -); - -Interpret the columns that are marked as Calculated Fields in the schema: -Rating (DOUBLE) - Calculated as the average score (avg) of the Score field from the reviews table where the reviews are associated with the order. This field represents the overall customer satisfaction rating for the order based on review scores. -ReviewCount (BIGINT) - Calculated by counting (count) the number of entries in the reviews table associated with this order. It measures the volume of customer feedback received for the order. -Size (BIGINT) - Represents the total number of items in the order, calculated by counting the number of item entries (ItemNumber) in the order_items table linked to this order. This field is useful for understanding the scale or size of an order. -Large (BOOLEAN) - A boolean value calculated to check if the number of items in the order exceeds one (count(order_items.ItemNumber) > 1). It indicates whether the order is considered large in terms of item quantity. - -And if the user input queries like these: -1. "How many large orders have been placed by customer with ID 'C1234'?" -2. "What is the average customer rating for orders that were rated by more than 10 reviewers?" - -For the first query: -First try to intepret the user query, the user wants to know the average rating for orders which have attracted significant review activity, specifically those with more than 10 reviews. -Then, according to the above intepretation about the given schema, the term 'Rating' is predefined in the Calculated Field of the 'orders' model. And, the number of reviews is also predefined in the 'ReviewCount' Calculated Field. -So utilize those Calculated Fields in the SQL generation process to give an answer like this: - -SQL Query: SELECT AVG(Rating) FROM orders WHERE ReviewCount > 10 """ _DEFAULT_METRIC_INSTRUCTIONS = """ @@ -297,68 +240,6 @@ async def _classify_generation_result( If the given schema contains the structures marked as 'metric', you should first interpret the metric schema based on the above definition. Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. - -EXAMPLES: -The given schema is created by the SQL command: - -/* This table is a metric */ -/* Metric Base Object: orders */ -CREATE TABLE Revenue ( - -- This column is a dimension - PurchaseTimestamp TIMESTAMP, - -- This column is a dimension - CustomerId VARCHAR, - -- This column is a dimension - Status VARCHAR, - -- This column is a measure - -- expression: sum(order_items.Price) - PriceSum DOUBLE, - -- This column is a measure - -- expression: count(OrderId) - NumberOfOrders BIGINT -); - -Interpret the metric with the understanding of the metric structure: -1. Base Object: orders -This is the primary data source for the metric. -The orders table provides the underlying data from which dimensions and measures are derived. -It is the foundation upon which the metric is built, though it itself is not directly used in queries against the Revenue table. -It shows the reference between the 'Revenue' metric and the 'orders' model. For the user queries pretain to the 'Revenue' of 'orders', the metric should be utilize in the sql generation process. -2. Dimensions -The metric contains the columns marked as 'dimension'. They can be interpreted as below: -- PurchaseTimestamp (TIMESTAMP) - Acts as a temporal dimension, allowing analysis of revenue over time. This can be used to observe trends, seasonal variations, or performance over specific periods. -- CustomerId (VARCHAR) - A key dimension for customer segmentation, it enables the analysis of revenue generated from individual customers or customer groups. -- Status (VARCHAR) - Reflects the current state of an order (e.g., pending, completed, cancelled). This dimension is crucial for analyses that differentiate performance based on order status. -3. Measures -The metric contains the columns marked as 'measure'. They can be interpreted as below: -- PriceSum (DOUBLE) - A financial measure calculated as sum(order_items.Price), representing the total revenue generated from orders. This measure is vital for tracking overall sales performance and is the primary output of interest in many financial and business analyses. -- NumberOfOrders (BIGINT) - A count measure that provides the total number of orders. This is essential for operational metrics, such as assessing the volume of business activity and evaluating the efficiency of sales processes. - -Now, if the user input queries like this: -Question: "What was the total revenue from each customer last month?" - -First try to intepret the user query, the user asks for a breakdown of the total revenue generated by each customer in the previous calendar month. -The user is specifically interested in understanding how much each customer contributed to the total sales during this period. -To answer this question, it is suitable to use the following components from the metric: -1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. -2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. -3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. -So utilize those metric components in the SQL generation process to give an answer like this: - -SQL Query: -SELECT - CustomerId, - PriceSum AS TotalRevenue -FROM - Revenue -WHERE - PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND - PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -369,31 +250,11 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - For Example: - DATA SCHEMA: - `/* {"alias":"users","description":"A model representing the users data."} */ - CREATE TABLE users ( - -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} - address JSON - )` - To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. - - For Example: - DATA SCHEMA - `/* {"alias":"my_table","description":"A test my_table"} */ - CREATE TABLE my_table ( - -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} - elements JSON - )` - To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". - DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". """ From 77b06246c01e42f1d037462b55064e5276a7c9c6 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 26 Jul 2026 00:06:15 +0530 Subject: [PATCH 0664/1087] Restore SQL reasoning into generation flow --- wren-ai-service/src/web/v1/services/ask.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index fc85ff4d60..3f6b38b4dd 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -455,7 +455,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, sql_samples=[], @@ -474,7 +474,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, sql_samples=[], instructions=[], From 098b5188c0b7d97eca4f255f8f6862ea222fe358 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 26 Jul 2026 01:02:29 +0530 Subject: [PATCH 0665/1087] Restore deployed SQL context in ask flow --- wren-ai-service/src/web/v1/services/ask.py | 24 +++++++++++----------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 3f6b38b4dd..bc66f5be90 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -202,9 +202,9 @@ async def ask( if historical_question_result: sql_generation_reasoning = "" - # Run both pipeline operations concurrently. These are still used by - # the legacy intent path, but executable SQL generation below is - # grounded only by the retrieved schema for the current deployment. + # Run both pipeline operations concurrently. The retrieved samples + # and instructions are passed through the same reasoning, + # generation, and correction flow as legacy/v1. sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( query=user_query, @@ -386,8 +386,8 @@ async def ask( query=user_query, contexts=table_ddls, histories=histories, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, configuration=ask_request.configurations, query_id=query_id, ) @@ -397,8 +397,8 @@ async def ask( await self._pipelines["sql_generation_reasoning"].run( query=user_query, contexts=table_ddls, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, configuration=ask_request.configurations, query_id=query_id, ) @@ -458,8 +458,8 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -476,8 +476,8 @@ async def ask( contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -539,7 +539,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - instructions=[], + instructions=instructions, invalid_generation_result={ "sql": original_sql, "error": sql_diagnosis_reasoning From c3cde7cfcc7daf8aca31cb777fd33d42ad8e363d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 26 Jul 2026 19:55:44 +0530 Subject: [PATCH 0666/1087] Fix ask reuse and MDL table references --- wren-ai-service/src/web/v1/services/ask.py | 224 +++++++++--------- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 70 +----- .../apollo/server/mdl/test/mdlBuilder.test.ts | 188 +-------------- 3 files changed, 121 insertions(+), 361 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bc66f5be90..03efe00799 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -200,129 +200,137 @@ async def ask( ).get("documents", [])[:1] if historical_question_result: + api_results = [ + AskResult( + **{ + "sql": result.get("statement"), + "type": "view" if result.get("viewId") else "llm", + "viewId": result.get("viewId"), + } + ) + for result in historical_question_result + ] sql_generation_reasoning = "" - - # Run both pipeline operations concurrently. The retrieved samples - # and instructions are passed through the same reasoning, - # generation, and correction flow as legacy/v1. - sql_samples_task, instructions_task = await asyncio.gather( - self._pipelines["sql_pairs_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - ), - self._pipelines["instructions_retrieval"].run( - query=user_query, - project_id=ask_request.project_id, - scope="sql", - ), - ) - - # Extract results from completed tasks - sql_samples = sql_samples_task["formatted_output"].get("documents", []) - instructions = instructions_task["formatted_output"].get( - "documents", [] - ) - - if self._allow_intent_classification: - intent_classification_result = ( - await self._pipelines["intent_classification"].run( + else: + sql_samples_task, instructions_task = await asyncio.gather( + self._pipelines["sql_pairs_retrieval"].run( query=user_query, - histories=histories, - sql_samples=sql_samples, - instructions=instructions, project_id=ask_request.project_id, - configuration=ask_request.configurations, - ) - ).get("post_process", {}) - intent = intent_classification_result.get("intent") - rephrased_question = intent_classification_result.get( - "rephrased_question" + ), + self._pipelines["instructions_retrieval"].run( + query=user_query, + project_id=ask_request.project_id, + scope="sql", + ), ) - intent_reasoning = intent_classification_result.get("reasoning") - if rephrased_question: - user_query = rephrased_question + sql_samples = sql_samples_task["formatted_output"].get( + "documents", [] + ) + instructions = instructions_task["formatted_output"].get( + "documents", [] + ) - if intent == "MISLEADING_QUERY": - asyncio.create_task( - self._pipelines["misleading_assistance"].run( + if self._allow_intent_classification: + intent_classification_result = ( + await self._pipelines["intent_classification"].run( query=user_query, histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, + sql_samples=sql_samples, + instructions=instructions, + project_id=ask_request.project_id, + configuration=ask_request.configurations, ) + ).get("post_process", {}) + intent = intent_classification_result.get("intent") + rephrased_question = intent_classification_result.get( + "rephrased_question" ) + intent_reasoning = intent_classification_result.get("reasoning") + + if rephrased_question: + user_query = rephrased_question + + if intent == "MISLEADING_QUERY": + asyncio.create_task( + self._pipelines["misleading_assistance"].run( + query=user_query, + histories=histories, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) + ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="MISLEADING_QUERY", - ) - results["metadata"]["type"] = "MISLEADING_QUERY" - return results - elif intent == "GENERAL": - asyncio.create_task( - self._pipelines["data_assistance"].run( - query=user_query, - histories=histories, - db_schemas=intent_classification_result.get( - "db_schemas" - ), - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="MISLEADING_QUERY", + ) + results["metadata"]["type"] = "MISLEADING_QUERY" + return results + elif intent == "GENERAL": + asyncio.create_task( + self._pipelines["data_assistance"].run( + query=user_query, + histories=histories, + db_schemas=intent_classification_result.get( + "db_schemas" + ), + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) ) - ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="DATA_ASSISTANCE", - ) - results["metadata"]["type"] = "GENERAL" - return results - elif intent == "USER_GUIDE": - asyncio.create_task( - self._pipelines["user_guide_assistance"].run( - query=user_query, - language=ask_request.configurations.language, - query_id=ask_request.query_id, - custom_instruction=ask_request.custom_instruction, + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="DATA_ASSISTANCE", + ) + results["metadata"]["type"] = "GENERAL" + return results + elif intent == "USER_GUIDE": + asyncio.create_task( + self._pipelines["user_guide_assistance"].run( + query=user_query, + language=ask_request.configurations.language, + query_id=ask_request.query_id, + custom_instruction=ask_request.custom_instruction, + ) ) - ) - self._ask_results[query_id] = AskResultResponse( - status="finished", - type="GENERAL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - general_type="USER_GUIDE", - ) - results["metadata"]["type"] = "GENERAL" - return results - else: - self._ask_results[query_id] = AskResultResponse( - status="understanding", - type="TEXT_TO_SQL", - rephrased_question=rephrased_question, - intent_reasoning=intent_reasoning, - trace_id=trace_id, - is_followup=True if histories else False, - ) + self._ask_results[query_id] = AskResultResponse( + status="finished", + type="GENERAL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + general_type="USER_GUIDE", + ) + results["metadata"]["type"] = "GENERAL" + return results + else: + self._ask_results[query_id] = AskResultResponse( + status="understanding", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) if not self._is_stopped(query_id, self._ask_results) and not api_results: self._ask_results[query_id] = AskResultResponse( status="searching", diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 5a39c65933..bf94b7fb67 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -583,78 +583,16 @@ export class MDLBuilder implements IMDLBuilder { model.properties && typeof model.properties === 'string' ? this.parseProperties(model.properties) : {}; - const propertyTableReference = - typeof modelProps.table === 'string' - ? this.buildTableReferenceFromTableName(modelProps.table) - : null; - const fallbackTableReference = this.buildFallbackTableReference(model); - const table = - propertyTableReference?.table || - modelProps.table || - fallbackTableReference?.table; - if (!table) { + if (!modelProps.table) { return null; } return { - catalog: - propertyTableReference?.catalog || - modelProps.catalog || - fallbackTableReference?.catalog || - null, - schema: - propertyTableReference?.schema || - modelProps.schema || - fallbackTableReference?.schema || - null, - table, + catalog: modelProps.catalog || null, + schema: modelProps.schema || null, + table: modelProps.table, }; } - private buildFallbackTableReference(model: Model): TableReference | null { - if (!this.useRustWrenEngine() || !model.sourceTableName) { - return null; - } - - const sourceTableName = model.sourceTableName.trim(); - const normalizedTableReference = - this.buildTableReferenceFromTableName(sourceTableName); - if (normalizedTableReference) { - return normalizedTableReference; - } - - return { - catalog: null, - schema: null, - table: sourceTableName, - }; - } - - private buildTableReferenceFromTableName( - tableName: string, - ): TableReference | null { - const sourceTableName = tableName.trim(); - const catalogQualifiedMatch = sourceTableName.match( - /^([^.]+)\.([^.]+)\.([^.]+)$/, - ); - if (catalogQualifiedMatch) { - return { - catalog: catalogQualifiedMatch[1], - schema: catalogQualifiedMatch[2], - table: catalogQualifiedMatch[3], - }; - } - - const dotQualifiedMatch = sourceTableName.match(/^([^.]+)\.([^.]+)$/); - if (dotQualifiedMatch) { - return { - catalog: null, - schema: dotQualifiedMatch[1], - table: dotQualifiedMatch[2], - }; - } - - return null; - } private hasDuplicateSourceColumns(modelId: number): boolean { const sourceColumnNames = new Set(); for (const column of this.columns.filter( diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 930a963254..3eeadea7f5 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -539,145 +539,7 @@ describe('MDLBuilder', () => { expect(manifest.views).toEqual(expectedViews); }); - it('should use source table name as rust engine table reference fallback.', () => { - const project = { - id: 1, - type: DataSourceName.MSSQL, - displayName: 'my project', - connectionInfo: {}, - catalog: 'wrenai', - schema: 'public', - sampleDataset: null, - } as Project; - const models = [ - { - id: 1, - projectId: 1, - displayName: 'Search Queries', - sourceTableName: 'dbo.search_queries', - referenceName: 'dbo_search_queries', - refSql: 'SELECT * FROM dbo.search_queries', - cached: false, - refreshTime: null, - properties: null, - }, - ] as Model[]; - const builderOptions = { - project, - models, - columns: [], - nestedColumns: [], - relations: [], - views: [], - relatedModels: [], - relatedColumns: [], - relatedRelations: [], - } as MDLBuilderBuildFromOptions; - mdlBuilder = new MDLBuilder(builderOptions); - - const manifest = mdlBuilder.build(); - - expect(manifest.models[0].tableReference).toEqual({ - catalog: null, - schema: 'dbo', - table: 'search_queries', - }); - expect(manifest.models[0].refSql).toBeUndefined(); - }); - - it('should preserve unqualified source table names in rust engine table reference fallback.', () => { - const project = { - id: 1, - type: DataSourceName.MSSQL, - displayName: 'my project', - connectionInfo: {}, - catalog: 'wrenai', - schema: 'public', - sampleDataset: null, - } as Project; - const models = [ - { - id: 1, - projectId: 1, - displayName: 'Tickets', - sourceTableName: 'dbo_tickets', - referenceName: 'dbo_tickets', - refSql: 'SELECT * FROM dbo.tickets', - cached: false, - refreshTime: null, - properties: null, - }, - ] as Model[]; - const builderOptions = { - project, - models, - columns: [], - nestedColumns: [], - relations: [], - views: [], - relatedModels: [], - relatedColumns: [], - relatedRelations: [], - } as MDLBuilderBuildFromOptions; - mdlBuilder = new MDLBuilder(builderOptions); - - const manifest = mdlBuilder.build(); - - expect(manifest.models[0].tableReference).toEqual({ - catalog: null, - schema: null, - table: 'dbo_tickets', - }); - expect(manifest.models[0].refSql).toBeUndefined(); - }); - - it('should preserve dbo-prefixed source table names for non-mssql projects.', () => { - const project = { - id: 1, - type: DataSourceName.POSTGRES, - displayName: 'wren ai project', - connectionInfo: {}, - catalog: 'wrenai', - schema: 'public', - sampleDataset: null, - } as Project; - const models = [ - { - id: 1, - projectId: 1, - displayName: 'Search Queries', - sourceTableName: 'dbo_search_queries', - referenceName: 'dbo_search_queries', - refSql: 'SELECT * FROM dbo.search_queries', - cached: false, - refreshTime: null, - properties: null, - }, - ] as Model[]; - const builderOptions = { - project, - models, - columns: [], - nestedColumns: [], - relations: [], - views: [], - relatedModels: [], - relatedColumns: [], - relatedRelations: [], - } as MDLBuilderBuildFromOptions; - mdlBuilder = new MDLBuilder(builderOptions); - - const manifest = mdlBuilder.build(); - - expect(manifest.models[0].tableReference).toEqual({ - catalog: null, - schema: null, - table: 'dbo_search_queries', - }); - expect(manifest.models[0].refSql).toBeUndefined(); - }); - - it('should preserve dbo-prefixed property table names before project schema fallback.', () => { + it('should preserve dbo-prefixed property table names.', () => { const project = { id: 1, type: DataSourceName.POSTGRES, @@ -726,54 +588,6 @@ describe('MDLBuilder', () => { expect(manifest.models[0].refSql).toBeUndefined(); }); - it('should preserve catalog-qualified non-mssql table references.', () => { - const project = { - id: 1, - type: DataSourceName.POSTGRES, - displayName: 'wren ai project', - connectionInfo: {}, - catalog: 'wrenai', - schema: 'public', - sampleDataset: null, - } as Project; - const models = [ - { - id: 1, - projectId: 1, - displayName: 'Search Queries', - sourceTableName: 'wrenai.public.dbo_search_queries', - referenceName: 'dbo_search_queries', - refSql: 'SELECT * FROM wrenai.public.dbo_search_queries', - cached: false, - refreshTime: null, - properties: JSON.stringify({ - table: 'wrenai.public.dbo_search_queries', - }), - }, - ] as Model[]; - const builderOptions = { - project, - models, - columns: [], - nestedColumns: [], - relations: [], - views: [], - relatedModels: [], - relatedColumns: [], - relatedRelations: [], - } as MDLBuilderBuildFromOptions; - mdlBuilder = new MDLBuilder(builderOptions); - - const manifest = mdlBuilder.build(); - - expect(manifest.models[0].tableReference).toEqual({ - catalog: 'wrenai', - schema: 'public', - table: 'dbo_search_queries', - }); - expect(manifest.models[0].refSql).toBeUndefined(); - }); - it('should return correct expression in calculated field.', () => { const models = [ // customer model From 89a538a4fd307f0d0218e3c7366c84a75e74c248 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 26 Jul 2026 20:46:01 +0530 Subject: [PATCH 0667/1087] Preserve model refSql without tableReference --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 6 +-- .../apollo/server/mdl/test/mdlBuilder.test.ts | 44 +++++++++++++++++++ 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index bf94b7fb67..b91a4cc103 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -147,11 +147,7 @@ export class MDLBuilder implements IMDLBuilder { columns: [], tableReference, // can only have one of refSql or tableReference - refSql: this.useRustWrenEngine() - ? null - : tableReference - ? null - : model.refSql, + refSql: tableReference ? null : model.refSql, cached: model.cached ? true : false, refreshTime: model.refreshTime, properties: { diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 3eeadea7f5..1cba2a12cf 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -588,6 +588,50 @@ describe('MDLBuilder', () => { expect(manifest.models[0].refSql).toBeUndefined(); }); + it('should preserve refSql when a model has no tableReference.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Semantic Model', + sourceTableName: 'physical_table', + referenceName: 'semantic_model', + refSql: 'SELECT * FROM physical_schema.physical_table', + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toBeFalsy(); + expect(manifest.models[0].refSql).toEqual( + 'SELECT * FROM physical_schema.physical_table', + ); + }); + it('should return correct expression in calculated field.', () => { const models = [ // customer model From ba1451000d29d3b1d75c7fdc9bb5af14b8fee1a7 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 26 Jul 2026 21:31:23 +0530 Subject: [PATCH 0668/1087] Fix ask SQL source mapping --- .../generation/followup_sql_generation.py | 11 +- .../followup_sql_generation_reasoning.py | 4 - .../pipelines/generation/sql_correction.py | 11 +- .../pipelines/generation/sql_generation.py | 11 +- .../generation/sql_generation_reasoning.py | 4 - .../pipelines/generation/sql_regeneration.py | 21 +-- .../src/pipelines/generation/utils/sql.py | 168 ++++++++++++++++-- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 5 +- .../apollo/server/mdl/test/mdlBuilder.test.ts | 61 ++++++- 9 files changed, 221 insertions(+), 75 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 fcfbd227ca..35cfb8fccf 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -39,10 +39,6 @@ {{ document }} {% endfor %} -Use this DATABASE SCHEMA as the complete allowed identifier set for this query. -Only generate SQL with table, column, schema, model, and datasource names present above. -Do not infer identifiers from the follow-up question, previous SQL, summary, SQL samples, user instructions, or prior examples. - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -215,9 +211,10 @@ async def run( ): logger.info("Follow-Up SQL Generation pipeline is running...") - metadata = ( - await retrieve_metadata(project_id, self._retriever) if project_id else {} - ) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 6c9b89db8a..42b28c5b8f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -28,10 +28,6 @@ {{ document }} {% endfor %} -Use this DATABASE SCHEMA as the complete allowed identifier set for the reasoning plan. -Only refer to table, column, schema, model, and datasource names present above. -Do not infer identifiers from the follow-up question, previous SQL, SQL samples, user instructions, or prior examples. - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 999b11c1d8..973b8c69a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -57,10 +57,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% for document in documents %} {{ document }} {% endfor %} - -Use this DATABASE SCHEMA as the complete allowed identifier set for the corrected SQL. -Only correct SQL with table, column, schema, model, and datasource names present above. -Do not infer identifiers from the original SQL, error message, user instructions, or prior examples. {% endif %} {% if sql_functions %} @@ -182,9 +178,10 @@ async def run( ): logger.info("SQLCorrection pipeline is running...") - metadata = ( - await retrieve_metadata(project_id, self._retriever) if project_id else {} - ) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 40097aa44e..1ee4952b3e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -33,10 +33,6 @@ {{ document }} {% endfor %} -Use this DATABASE SCHEMA as the complete allowed identifier set for this query. -Only generate SQL with table, column, schema, model, and datasource names present above. -Do not infer identifiers from the question, SQL samples, user instructions, or prior examples. - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -209,9 +205,10 @@ async def run( ): logger.info("SQL Generation pipeline is running...") - metadata = ( - await retrieve_metadata(project_id, self._retriever) if project_id else {} - ) + if use_dry_plan: + metadata = await retrieve_metadata(project_id or "", self._retriever) + else: + metadata = {} return await self._pipe.execute( ["post_process"], 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 f9f9af9f4f..00b731cb2c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -27,10 +27,6 @@ {{ document }} {% endfor %} -Use this DATABASE SCHEMA as the complete allowed identifier set for the reasoning plan. -Only refer to table, column, schema, model, and datasource names present above. -Do not infer identifiers from the question, SQL samples, user instructions, or prior examples. - {% if sql_samples %} ### SQL SAMPLES ### {% for sql_sample in sql_samples %} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 2d273a9aea..4b7284aa26 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -29,7 +29,6 @@ def get_sql_regeneration_system_prompt( sql_knowledge: SqlKnowledge | None = None, - data_source: str | None = None, ) -> str: text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) @@ -52,18 +51,11 @@ def get_sql_regeneration_system_prompt( sql_regeneration_user_prompt_template = """ -### TARGET DATA SOURCE ### -{{ data_source }} - ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} {% endfor %} -Use this DATABASE SCHEMA as the complete allowed identifier set for the regenerated SQL. -Only regenerate SQL with table, column, schema, model, and datasource names present above. -Do not infer identifiers from the original SQL, reasoning, SQL samples, user instructions, or prior examples. - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -115,7 +107,6 @@ def prompt( sql_generation_reasoning: str, sql: str, prompt_builder: PromptBuilder, - data_source: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -126,7 +117,6 @@ def prompt( ) -> dict: _prompt = prompt_builder.run( sql=sql, - data_source=data_source, documents=documents, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( @@ -155,13 +145,9 @@ async def regenerate_sql( prompt: dict, generator: Any, generator_name: str, - data_source: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - current_system_prompt = get_sql_regeneration_system_prompt( - sql_knowledge, - data_source=data_source, - ) + current_system_prompt = get_sql_regeneration_system_prompt(sql_knowledge) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt ), generator_name @@ -171,14 +157,11 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, - documents: list[str], - data_source: str, project_id: str | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, - data_source=data_source, ) @@ -214,7 +197,6 @@ async def run( contexts: list[str], sql_generation_reasoning: str, sql: str, - data_source: str = "local_file", sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, @@ -240,7 +222,6 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, - "data_source": data_source, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 9eb56d606a..088282574e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -176,6 +176,7 @@ async def _classify_generation_result( - Put double quotes around column and table names. - Put single quotes around string literals. - Never quote numeric literals. + For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992; - YOU MUST USE "lower(.) like lower()" function or "lower(.) = lower()" function for case-insensitive comparison! - Use "lower(.) LIKE lower()" when: - The user requests a pattern or partial match. @@ -185,12 +186,30 @@ async def _classify_generation_result( - The user requests an exact, specific value. - There is no ambiguity or pattern in the value. - If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query + - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds + - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds + - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds - ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query + - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) + - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) + - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) - If the user asks for a specific date, please give the date range in SQL query + - example: "What is the total revenue for the month of 2024-11-01?" + - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. - Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. + - EXAMPLE + DATABASE SCHEMA + /* {"alias":"_orders","description":"A model representing the orders data."} */ + CREATE TABLE orders ( + -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} + ApprovedTimestamp TIMESTAMP + } + + SQL + SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -211,6 +230,44 @@ async def _classify_generation_result( First, provide a brief explanation of what each field represents in the context of the schema, including how each field is computed using the relationships between models. Then, during the following tasks, if the user queries pertain to any calculated fields defined in the database schema, ensure to utilize those calculated fields appropriately in the output SQL queries. The goal is to accurately reflect the intent of the question in the SQL syntax, leveraging the pre-computed logic embedded within the calculated fields. + +EXAMPLES: +The given schema is created by the SQL command: + +CREATE TABLE orders ( + OrderId VARCHAR PRIMARY KEY, + CustomerId VARCHAR, + -- This column is a Calculated Field + -- column expression: avg(reviews.Score) + Rating DOUBLE, + -- This column is a Calculated Field + -- column expression: count(reviews.Id) + ReviewCount BIGINT, + -- This column is a Calculated Field + -- column expression: count(order_items.ItemNumber) + Size BIGINT, + -- This column is a Calculated Field + -- column expression: count(order_items.ItemNumber) > 1 + Large BOOLEAN, + FOREIGN KEY (CustomerId) REFERENCES customers(Id) +); + +Interpret the columns that are marked as Calculated Fields in the schema: +Rating (DOUBLE) - Calculated as the average score (avg) of the Score field from the reviews table where the reviews are associated with the order. This field represents the overall customer satisfaction rating for the order based on review scores. +ReviewCount (BIGINT) - Calculated by counting (count) the number of entries in the reviews table associated with this order. It measures the volume of customer feedback received for the order. +Size (BIGINT) - Represents the total number of items in the order, calculated by counting the number of item entries (ItemNumber) in the order_items table linked to this order. This field is useful for understanding the scale or size of an order. +Large (BOOLEAN) - A boolean value calculated to check if the number of items in the order exceeds one (count(order_items.ItemNumber) > 1). It indicates whether the order is considered large in terms of item quantity. + +And if the user input queries like these: +1. "How many large orders have been placed by customer with ID 'C1234'?" +2. "What is the average customer rating for orders that were rated by more than 10 reviewers?" + +For the first query: +First try to intepret the user query, the user wants to know the average rating for orders which have attracted significant review activity, specifically those with more than 10 reviews. +Then, according to the above intepretation about the given schema, the term 'Rating' is predefined in the Calculated Field of the 'orders' model. And, the number of reviews is also predefined in the 'ReviewCount' Calculated Field. +So utilize those Calculated Fields in the SQL generation process to give an answer like this: + +SQL Query: SELECT AVG(Rating) FROM orders WHERE ReviewCount > 10 """ _DEFAULT_METRIC_INSTRUCTIONS = """ @@ -240,6 +297,68 @@ async def _classify_generation_result( If the given schema contains the structures marked as 'metric', you should first interpret the metric schema based on the above definition. Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. + +EXAMPLES: +The given schema is created by the SQL command: + +/* This table is a metric */ +/* Metric Base Object: orders */ +CREATE TABLE Revenue ( + -- This column is a dimension + PurchaseTimestamp TIMESTAMP, + -- This column is a dimension + CustomerId VARCHAR, + -- This column is a dimension + Status VARCHAR, + -- This column is a measure + -- expression: sum(order_items.Price) + PriceSum DOUBLE, + -- This column is a measure + -- expression: count(OrderId) + NumberOfOrders BIGINT +); + +Interpret the metric with the understanding of the metric structure: +1. Base Object: orders +This is the primary data source for the metric. +The orders table provides the underlying data from which dimensions and measures are derived. +It is the foundation upon which the metric is built, though it itself is not directly used in queries against the Revenue table. +It shows the reference between the 'Revenue' metric and the 'orders' model. For the user queries pretain to the 'Revenue' of 'orders', the metric should be utilize in the sql generation process. +2. Dimensions +The metric contains the columns marked as 'dimension'. They can be interpreted as below: +- PurchaseTimestamp (TIMESTAMP) + Acts as a temporal dimension, allowing analysis of revenue over time. This can be used to observe trends, seasonal variations, or performance over specific periods. +- CustomerId (VARCHAR) + A key dimension for customer segmentation, it enables the analysis of revenue generated from individual customers or customer groups. +- Status (VARCHAR) + Reflects the current state of an order (e.g., pending, completed, cancelled). This dimension is crucial for analyses that differentiate performance based on order status. +3. Measures +The metric contains the columns marked as 'measure'. They can be interpreted as below: +- PriceSum (DOUBLE) + A financial measure calculated as sum(order_items.Price), representing the total revenue generated from orders. This measure is vital for tracking overall sales performance and is the primary output of interest in many financial and business analyses. +- NumberOfOrders (BIGINT) + A count measure that provides the total number of orders. This is essential for operational metrics, such as assessing the volume of business activity and evaluating the efficiency of sales processes. + +Now, if the user input queries like this: +Question: "What was the total revenue from each customer last month?" + +First try to intepret the user query, the user asks for a breakdown of the total revenue generated by each customer in the previous calendar month. +The user is specifically interested in understanding how much each customer contributed to the total sales during this period. +To answer this question, it is suitable to use the following components from the metric: +1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. +2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. +3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. +So utilize those metric components in the SQL generation process to give an answer like this: + +SQL Query: +SELECT + CustomerId, + PriceSum AS TotalRevenue +FROM + Revenue +WHERE + PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND + PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -250,11 +369,31 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields + - For Example: + DATA SCHEMA: + `/* {"alias":"users","description":"A model representing the users data."} */ + CREATE TABLE users ( + -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} + address JSON + )` + To get the city of address in user table use SQL: + `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` + - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. + - For Example: + DATA SCHEMA + `/* {"alias":"my_table","description":"A test my_table"} */ + CREATE TABLE my_table ( + -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} + elements JSON + )` + To get the number of elements in my_table table use SQL: + `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) + - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". - DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". """ @@ -298,16 +437,14 @@ async def _classify_generation_result( 4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. -7. Use DATABASE SCHEMA as the complete and only source of valid table, column, schema, model, and datasource names. -8. Do not introduce, infer, or copy any identifier from the question, SQL samples, user instructions, or query history unless it also appears in DATABASE SCHEMA. -9. Give a step by step reasoning plan in order to answer user's question. -10. The reasoning plan should be in the language same as the language user provided in the input. -11. Don't include SQL in the reasoning plan. -12. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. -13. Do not include ```markdown or ``` in the answer. -14. A table name in the reasoning plan must be in this format: `table: `. -15. A column name in the reasoning plan must be in this format: `column: .`. -16. ONLY SHOWING the reasoning plan in bullet points. +7. Give a step by step reasoning plan in order to answer user's question. +8. The reasoning plan should be in the language same as the language user provided in the input. +9. Don't include SQL in the reasoning plan. +10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. +11. Do not include ```markdown or ``` in the answer. +12. A table name in the reasoning plan must be in this format: `table: `. +13. A column name in the reasoning plan must be in this format: `column: .`. +14. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -373,13 +510,10 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. -2. The DATABASE SCHEMA section is the complete and only source of valid table, column, schema, model, and datasource names for this request. -3. YOU MUST NOT introduce, infer, copy, or repair any table, column, schema, model, or datasource name that is absent from DATABASE SCHEMA. -4. SQL SAMPLES and USER INSTRUCTIONS are usage guidance only. Do not copy identifiers from them unless those identifiers also appear in DATABASE SCHEMA. -5. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -6. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -7. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. -8. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. +3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. +4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index b91a4cc103..2c391dbc80 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -141,13 +141,14 @@ export class MDLBuilder implements IMDLBuilder { if (model.displayName) { properties.displayName = model.displayName; } - const tableReference = this.buildTableReference(model); + const refSql = model.refSql || null; + const tableReference = refSql ? null : this.buildTableReference(model); const modelMdl = { name: model.referenceName, columns: [], tableReference, // can only have one of refSql or tableReference - refSql: tableReference ? null : model.refSql, + refSql, cached: model.cached ? true : false, refreshTime: model.refreshTime, properties: { diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 1cba2a12cf..a5713e0d48 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -539,7 +539,7 @@ describe('MDLBuilder', () => { expect(manifest.views).toEqual(expectedViews); }); - it('should preserve dbo-prefixed property table names.', () => { + it('should use tableReference when a model has no refSql.', () => { const project = { id: 1, type: DataSourceName.POSTGRES, @@ -554,14 +554,14 @@ describe('MDLBuilder', () => { id: 1, projectId: 1, displayName: 'Search Queries', - sourceTableName: 'dbo_search_queries', - referenceName: 'dbo_search_queries', - refSql: 'SELECT * FROM dbo.search_queries', + sourceTableName: 'search_queries', + referenceName: 'search_queries', + refSql: null, cached: false, refreshTime: null, properties: JSON.stringify({ schema: 'public', - table: 'dbo_search_queries', + table: 'search_queries', }), }, ] as Model[]; @@ -583,9 +583,56 @@ describe('MDLBuilder', () => { expect(manifest.models[0].tableReference).toEqual({ catalog: null, schema: 'public', - table: 'dbo_search_queries', + table: 'search_queries', }); - expect(manifest.models[0].refSql).toBeUndefined(); + expect(manifest.models[0].refSql).toBeFalsy(); + }); + + it('should prefer refSql over tableReference when both are present.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Semantic Model', + sourceTableName: 'physical_table', + referenceName: 'semantic_model', + refSql: 'SELECT * FROM physical_schema.physical_table', + cached: false, + refreshTime: null, + properties: JSON.stringify({ + schema: 'public', + table: 'semantic_model', + }), + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toBeFalsy(); + expect(manifest.models[0].refSql).toEqual( + 'SELECT * FROM physical_schema.physical_table', + ); }); it('should preserve refSql when a model has no tableReference.', () => { From f79610a416f597020b6ea20ede34752c70b8fc50 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 26 Jul 2026 22:17:26 +0530 Subject: [PATCH 0669/1087] Derive model source SQL from table references --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 46 +++--------- .../apollo/server/mdl/test/mdlBuilder.test.ts | 74 +++++++++++++++++++ 2 files changed, 86 insertions(+), 34 deletions(-) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 2c391dbc80..efe098ec1d 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -141,8 +141,17 @@ export class MDLBuilder implements IMDLBuilder { if (model.displayName) { properties.displayName = model.displayName; } - const refSql = model.refSql || null; - const tableReference = refSql ? null : this.buildTableReference(model); + const sourceTableReference = this.buildTableReference(model); + const refSql = + model.refSql || + (sourceTableReference + ? this.buildTableReferenceSql( + model.id, + { name: model.referenceName, columns: [] }, + sourceTableReference, + ) + : null); + const tableReference = refSql ? null : sourceTableReference; const modelMdl = { name: model.referenceName, columns: [], @@ -158,21 +167,6 @@ export class MDLBuilder implements IMDLBuilder { primaryKey: '', // will be modified in addColumn } as ModelMDL; - if (tableReference && this.hasDuplicateSourceColumns(model.id)) { - const refSql = this.buildDedupedTableReferenceSql( - model.id, - modelMdl, - tableReference, - ); - if (refSql) { - logger.debug( - `Using deduped explicit projection for model "${model.referenceName}" because its source table contains duplicate column names.`, - ); - modelMdl.tableReference = null; - modelMdl.refSql = refSql; - } - } - return modelMdl; }); } @@ -590,23 +584,7 @@ export class MDLBuilder implements IMDLBuilder { }; } - private hasDuplicateSourceColumns(modelId: number): boolean { - const sourceColumnNames = new Set(); - for (const column of this.columns.filter( - ({ isCalculated, modelId: columnModelId }) => - !isCalculated && columnModelId === modelId, - )) { - const sourceColumnName = ( - column.sourceColumnName || column.referenceName - ).toLowerCase(); - if (sourceColumnNames.has(sourceColumnName)) { - return true; - } - sourceColumnNames.add(sourceColumnName); - } - return false; - } - private buildDedupedTableReferenceSql( + private buildTableReferenceSql( modelId: number, model: Partial, tableReference: TableReference, diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index a5713e0d48..8a950f1931 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -635,6 +635,80 @@ describe('MDLBuilder', () => { ); }); + it('should build refSql from tableReference metadata when a model has columns but no refSql.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Imported Model', + sourceTableName: 'imported_model', + referenceName: 'imported_model', + refSql: null, + cached: false, + refreshTime: null, + properties: JSON.stringify({ + catalog: 'physical_catalog', + schema: 'physical_schema', + table: 'physical_table', + }), + }, + ] as Model[]; + const columns = [ + { + id: 1, + modelId: 1, + isCalculated: false, + displayName: 'Source Name', + referenceName: 'SourceName', + sourceColumnName: 'Source Name', + type: 'VARCHAR', + notNull: false, + isPk: false, + properties: null, + }, + { + id: 2, + modelId: 1, + isCalculated: false, + displayName: 'Source Type', + referenceName: 'SourceType', + sourceColumnName: 'Source Type', + type: 'VARCHAR', + notNull: false, + isPk: false, + properties: null, + }, + ] as ModelColumn[]; + const builderOptions = { + project, + models, + columns, + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toBeFalsy(); + expect(manifest.models[0].refSql).toEqual( + 'SELECT "Source Name" AS "SourceName", "Source Type" AS "SourceType" FROM "physical_catalog"."physical_schema"."physical_table"', + ); + }); + it('should preserve refSql when a model has no tableReference.', () => { const project = { id: 1, From ec8afe48572513ec480de428e6db4f6164d23f90 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 26 Jul 2026 22:01:11 +0530 Subject: [PATCH 0670/1087] Ground executable SQL on retrieved DDL --- wren-ai-service/src/web/v1/services/ask.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 03efe00799..2f3b1602fe 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -212,6 +212,8 @@ async def ask( ] sql_generation_reasoning = "" else: + # Run both retrievals for intent classification. Executable SQL + # generation is grounded only by the schema DDL retrieved below. sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( query=user_query, @@ -394,8 +396,8 @@ async def ask( query=user_query, contexts=table_ddls, histories=histories, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], configuration=ask_request.configurations, query_id=query_id, ) @@ -405,8 +407,8 @@ async def ask( await self._pipelines["sql_generation_reasoning"].run( query=user_query, contexts=table_ddls, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], configuration=ask_request.configurations, query_id=query_id, ) @@ -466,8 +468,8 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -484,8 +486,8 @@ async def ask( contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, - sql_samples=sql_samples, - instructions=instructions, + sql_samples=[], + instructions=[], has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -547,7 +549,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - instructions=instructions, + instructions=[], invalid_generation_result={ "sql": original_sql, "error": sql_diagnosis_reasoning From 8752c9c41c3b08a78d42766dbdd39e20ba362a8d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 26 Jul 2026 23:40:50 +0530 Subject: [PATCH 0671/1087] Restore legacy semantics preparation flow --- .../web/v1/services/semantics_preparation.py | 162 +----------------- 1 file changed, 3 insertions(+), 159 deletions(-) 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 23dd0994d1..2ff6215cbe 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -1,10 +1,9 @@ import asyncio import logging -from typing import Any, Dict, Literal, Optional +from typing import Dict, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe -import orjson from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline @@ -56,154 +55,6 @@ def __init__( str, SemanticsPreparationStatusResponse ] = TTLCache(maxsize=maxsize, ttl=ttl) - def _parse_mdl(self, mdl: str) -> dict[str, Any]: - parsed = orjson.loads(mdl) - parsed.setdefault("models", []) - parsed.setdefault("views", []) - parsed.setdefault("metrics", []) - parsed.setdefault("relationships", []) - self._normalize_mdl_metadata(parsed) - return parsed - - def _normalize_text(self, value: Any) -> str: - return "" if value is None else str(value) - - def _normalize_properties(self, payload: dict[str, Any]) -> None: - properties = payload.get("properties") - if not isinstance(properties, dict): - properties = {} - payload["properties"] = properties - - for key in ("description", "displayName"): - if key in properties: - properties[key] = self._normalize_text(properties[key]) - - def _normalize_resource_metadata(self, payload: dict[str, Any]) -> None: - self._normalize_properties(payload) - columns = payload.get("columns", []) - if not isinstance(columns, list): - payload["columns"] = [] - return - - for column in columns: - if not isinstance(column, dict): - continue - self._normalize_properties(column) - - def _normalize_mdl_metadata(self, mdl: dict[str, Any]) -> None: - for collection in ("models", "views", "metrics"): - resources = mdl.get(collection, []) - if not isinstance(resources, list): - mdl[collection] = [] - continue - - for resource in resources: - if isinstance(resource, dict): - self._normalize_resource_metadata(resource) - - def _validate_mdl_integrity(self, mdl: dict[str, Any]) -> None: - model_names = set() - for model in mdl["models"]: - model_name = model.get("name") - if not model_name: - raise ValueError("MDL contains a model without a name") - - normalized_model_name = str(model_name).lower() - if normalized_model_name in model_names: - raise ValueError(f'MDL contains duplicate model name "{model_name}"') - model_names.add(normalized_model_name) - - column_names = set() - for column in model.get("columns", []): - column_name = column.get("name") - if not column_name: - raise ValueError( - f'MDL model "{model_name}" contains a column without a name' - ) - - normalized_column_name = str(column_name).lower() - if normalized_column_name in column_names: - raise ValueError( - f'MDL model "{model_name}" contains duplicate column name "{column_name}"' - ) - column_names.add(normalized_column_name) - - for relationship in mdl["relationships"]: - for model_name in relationship.get("models", []): - if not model_name: - raise ValueError( - f'MDL relationship "{relationship.get("name", "")}" references an empty model name' - ) - if str(model_name).lower() not in model_names: - raise ValueError( - f'MDL relationship "{relationship.get("name", "")}" references missing model "{model_name}"' - ) - - def _project_filter( - self, project_id: Optional[str], *conditions: dict[str, Any] - ) -> dict[str, Any] | None: - all_conditions = list(conditions) - if project_id: - all_conditions.append( - {"field": "project_id", "operator": "==", "value": project_id} - ) - if not all_conditions: - return None - return {"operator": "AND", "conditions": all_conditions} - - async def _count_indexed_documents( - self, - pipeline_name: str, - project_id: Optional[str], - *conditions: dict[str, Any], - ) -> int: - pipeline = self._pipelines[pipeline_name] - writer = pipeline._components["writer"] - return await writer.document_store.count_documents( - filters=self._project_filter(project_id, *conditions) - ) - - async def _validate_index_integrity( - self, mdl: dict[str, Any], project_id: Optional[str] - ) -> None: - resource_count = ( - len(mdl["models"]) + len(mdl["views"]) + len(mdl["metrics"]) - ) - expected_schema_documents = len(mdl["views"]) + len(mdl["metrics"]) - for model in mdl["models"]: - expected_schema_documents += 1 - if model.get("columns") or mdl["relationships"]: - expected_schema_documents += 1 - - schema_count, table_description_count, project_meta_count = await asyncio.gather( - self._count_indexed_documents( - "db_schema", - project_id, - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ), - self._count_indexed_documents( - "table_description", - project_id, - {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, - ), - self._count_indexed_documents("project_meta", project_id), - ) - - if schema_count < expected_schema_documents: - raise ValueError( - "Incomplete DB schema index: " - f"expected at least {expected_schema_documents} documents, found {schema_count}" - ) - - if table_description_count < resource_count: - raise ValueError( - "Incomplete table-description index: " - f"expected at least {resource_count} documents, found {table_description_count}" - ) - - if project_meta_count < 1: - raise ValueError("Project metadata was not indexed") - @observe(name="Prepare Semantics") @trace_metadata async def prepare_semantics( @@ -220,13 +71,10 @@ async def prepare_semantics( } try: - mdl = self._parse_mdl(prepare_semantics_request.mdl) - self._validate_mdl_integrity(mdl) - normalized_mdl = orjson.dumps(mdl).decode("utf-8") - logger.info(f"MDL: {normalized_mdl}") + logger.info(f"MDL: {prepare_semantics_request.mdl}") input = { - "mdl_str": normalized_mdl, + "mdl_str": prepare_semantics_request.mdl, "project_id": prepare_semantics_request.project_id, } @@ -242,10 +90,6 @@ async def prepare_semantics( ] await asyncio.gather(*tasks) - await self._validate_index_integrity( - mdl, - prepare_semantics_request.project_id, - ) self._prepare_semantics_statuses[ prepare_semantics_request.mdl_hash From d441bfba96552387a0ef5235a710c0b42a6618af Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 00:38:38 +0530 Subject: [PATCH 0672/1087] Restore legacy ask SQL grounding flow --- wren-ai-service/src/web/v1/services/ask.py | 21 +++---- wren-ui/src/apollo/server/config.ts | 5 +- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 69 ++------------------- 3 files changed, 20 insertions(+), 75 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2f3b1602fe..7a20e792c0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -212,8 +212,7 @@ async def ask( ] sql_generation_reasoning = "" else: - # Run both retrievals for intent classification. Executable SQL - # generation is grounded only by the schema DDL retrieved below. + # Run both pipeline operations concurrently sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( query=user_query, @@ -396,8 +395,8 @@ async def ask( query=user_query, contexts=table_ddls, histories=histories, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, configuration=ask_request.configurations, query_id=query_id, ) @@ -407,8 +406,8 @@ async def ask( await self._pipelines["sql_generation_reasoning"].run( query=user_query, contexts=table_ddls, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, configuration=ask_request.configurations, query_id=query_id, ) @@ -468,8 +467,8 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -486,8 +485,8 @@ async def ask( contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, - sql_samples=[], - instructions=[], + sql_samples=sql_samples, + instructions=instructions, has_calculated_field=has_calculated_field, has_metric=has_metric, has_json_field=has_json_field, @@ -549,7 +548,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - instructions=[], + instructions=instructions, invalid_generation_result={ "sql": original_sql, "error": sql_diagnosis_reasoning diff --git a/wren-ui/src/apollo/server/config.ts b/wren-ui/src/apollo/server/config.ts index 576f380f0b..558bea67c0 100644 --- a/wren-ui/src/apollo/server/config.ts +++ b/wren-ui/src/apollo/server/config.ts @@ -182,5 +182,8 @@ const config = { }; export function getConfig(): IConfig { - return { ...defaultConfig, ...pickBy(config) }; + return { + ...defaultConfig, + ...pickBy(config, (value) => value !== undefined && value !== null), + }; } diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index efe098ec1d..08aea09c54 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -141,23 +141,17 @@ export class MDLBuilder implements IMDLBuilder { if (model.displayName) { properties.displayName = model.displayName; } - const sourceTableReference = this.buildTableReference(model); - const refSql = - model.refSql || - (sourceTableReference - ? this.buildTableReferenceSql( - model.id, - { name: model.referenceName, columns: [] }, - sourceTableReference, - ) - : null); - const tableReference = refSql ? null : sourceTableReference; + const tableReference = this.buildTableReference(model); const modelMdl = { name: model.referenceName, columns: [], tableReference, // can only have one of refSql or tableReference - refSql, + refSql: this.useRustWrenEngine() + ? null + : tableReference + ? null + : model.refSql, cached: model.cached ? true : false, refreshTime: model.refreshTime, properties: { @@ -584,57 +578,6 @@ export class MDLBuilder implements IMDLBuilder { }; } - private buildTableReferenceSql( - modelId: number, - model: Partial, - tableReference: TableReference, - ): string | null { - const sourceColumnNames = new Map(); - const projections: string[] = []; - - this.columns - .filter( - ({ isCalculated, modelId: columnModelId }) => - !isCalculated && columnModelId === modelId, - ) - .forEach((column) => { - const sourceColumnName = column.sourceColumnName || column.referenceName; - const normalizedSourceColumnName = sourceColumnName.toLowerCase(); - const existingColumnName = sourceColumnNames.get( - normalizedSourceColumnName, - ); - - if (existingColumnName) { - this.columnNameAliases.set(column.id, existingColumnName); - return; - } - - const columnName = this.getManifestColumnName(column, model); - sourceColumnNames.set(normalizedSourceColumnName, columnName); - const sourceExpression = this.quoteSqlIdentifier(sourceColumnName); - projections.push( - sourceColumnName === columnName - ? sourceExpression - : `${sourceExpression} AS ${this.quoteSqlIdentifier(columnName)}`, - ); - }); - - if (!projections.length) { - return null; - } - - const tableParts = [ - tableReference.catalog, - tableReference.schema, - tableReference.table, - ].filter((part): part is string => Boolean(part)); - return `SELECT ${projections.join(', ')} FROM ${tableParts - .map((part) => this.quoteSqlIdentifier(part)) - .join('.')}`; - } - private quoteSqlIdentifier(identifier: string): string { - return `"${identifier.replace(/"/g, '""')}"`; - } private parseLineage(lineage?: string): number[] { if (!lineage) { return []; From cebcc9b5b63bb439624c362cbce30a4d4a73bc49 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 02:53:50 +0530 Subject: [PATCH 0673/1087] Improve modeling semantics generation and save flow --- .../generation/semantics_description.py | 10 +- .../src/pipelines/generation/utils/sql.py | 6 +- .../web/v1/routers/semantics_description.py | 3 +- .../web/v1/services/semantics_description.py | 359 ++---------------- .../apollo/server/adaptors/wrenAIAdaptor.ts | 3 + wren-ui/src/apollo/server/models/model.ts | 36 +- wren-ui/src/apollo/server/resolvers.ts | 1 + .../apollo/server/resolvers/modelResolver.ts | 103 +++++ wren-ui/src/apollo/server/schema.ts | 7 + .../apollo/server/services/modelService.ts | 8 +- wren-ui/src/pages/modeling.tsx | 61 +-- 11 files changed, 217 insertions(+), 380 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 08fa40d8c6..93b684ea1a 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -36,7 +36,7 @@ ``` Your task is to update this JSON structure by adding a `description` field inside both the `properties` attribute of each `column` and the `model` itself. -Each `description` should be derived from the user-provided dataset context, the full schema, relationships, model names, column names, data types, aliases, and any existing descriptions. +Each `description` should be derived from the user-provided dataset context, the full schema, relationships, model names, column names, data types, aliases, existing descriptions, and provided sample rows. Follow these steps: 1. **For the `model`**: Write a clear natural language business description of the model's purpose and what real-world records it represents. Insert this description in the `properties` field of the `model`. 2. **For each `column`**: Write a clear natural language business description of the column's meaning, not just its technical name. Each column's description should be added under its respective `properties` field in the format: `'description': 'business description'`. @@ -89,7 +89,9 @@ Picked models: {{ picked_models }} Localization Language: {{ language }} -Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt and schema context. +Sample data for picked models: {{ data_samples }} + +Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt, schema context, and sample data. Do not omit selected models or columns. Do not copy the table or column name as the description. Use simple language that explains the business purpose, meaning, and analytical use of each field. """ @@ -151,11 +153,13 @@ def prompt( user_prompt: str, prompt_builder: PromptBuilder, language: str, + data_samples: dict[str, Any], ) -> dict: _prompt = prompt_builder.run( picked_models=picked_models, user_prompt=user_prompt, language=language, + data_samples=data_samples, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -266,6 +270,7 @@ async def run( selected_models: list[str], mdl: dict, language: str = "en", + data_samples: dict[str, Any] | None = None, ) -> dict: logger.info("Semantics Description Generation pipeline is running...") return await self._pipe.execute( @@ -275,6 +280,7 @@ async def run( "selected_models": selected_models, "mdl": mdl, "language": language, + "data_samples": data_samples or {}, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..0e2a2225e2 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -198,8 +198,10 @@ async def _classify_generation_result( - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. -- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. +- Table/column aliases in schema comments are display labels only. Never use an alias as an executable table or column identifier. +- Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. +- You may use aliases from schema comments only after AS in the final SELECT clause. +- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section only for final SELECT output labels. - EXAMPLE DATABASE SCHEMA /* {"alias":"_orders","description":"A model representing the orders data."} */ diff --git a/wren-ai-service/src/web/v1/routers/semantics_description.py b/wren-ai-service/src/web/v1/routers/semantics_description.py index 3e36d299bf..3118e59036 100644 --- a/wren-ai-service/src/web/v1/routers/semantics_description.py +++ b/wren-ai-service/src/web/v1/routers/semantics_description.py @@ -3,7 +3,7 @@ from typing import Literal, Optional from fastapi import APIRouter, BackgroundTasks, Depends -from pydantic import BaseModel +from pydantic import BaseModel, Field from src.globals import ( ServiceContainer, @@ -20,6 +20,7 @@ class PostRequest(BaseRequest): selected_models: list[str] user_prompt: str mdl: str + data_samples: dict = Field(default_factory=dict) class PostResponse(BaseModel): diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 370c3914f1..afd5ab812b 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,12 +1,11 @@ import asyncio import logging -import re from typing import Any, Dict, Literal, Optional import orjson from cachetools import TTLCache from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import BaseModel, Field from src.core.pipeline import BasicPipeline from src.utils import trace_metadata @@ -14,11 +13,6 @@ logger = logging.getLogger("wren-ai-service") -MAX_UI_WAIT_SECONDS = 180 -SEMANTICS_STATUS_TTL_BUFFER_SECONDS = 300 -SEMANTICS_MODEL_CHUNK_SIZE = 1000 -SEMANTICS_MAX_CONCURRENT_LLM_CALLS = 6 - class SemanticsDescription: class Resource(BaseModel, MetadataTraceable): @@ -38,20 +32,9 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, - generation_timeout_seconds: int = 90, ): self._pipelines = pipelines - self._generation_timeout_seconds = min( - generation_timeout_seconds, - MAX_UI_WAIT_SECONDS - 30, - ) - self._cache: Dict[str, self.Resource] = TTLCache( - maxsize=maxsize, - ttl=max( - ttl, - self._generation_timeout_seconds + SEMANTICS_STATUS_TTL_BUFFER_SECONDS, - ), - ) + self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) def _handle_exception( self, @@ -75,316 +58,45 @@ class GenerateRequest(BaseRequest): selected_models: list[str] user_prompt: str mdl: str - - def _properties(self, payload: dict[str, Any]) -> dict[str, Any]: - properties = payload.get("properties") - return properties if isinstance(properties, dict) else {} - - def _text(self, value: Any) -> str: - return "" if value is None else str(value) - - def _humanize_name(self, name: str) -> str: - return " ".join( - token - for token in re.sub( - r"(?<=[a-z0-9])(?=[A-Z])", - " ", - name.replace(".", " ").replace("_", " "), - ).split() - if token.lower() not in {"dbo", "public"} - ) or name - - def _name_tokens(self, name: str) -> set[str]: - return { - token.lower() - for token in self._humanize_name(name).split() - if token and token.lower() not in {"x", "stage", "load", "tbl", "table"} - } - - def _table_context(self, model: dict[str, Any]) -> set[str]: - tokens = self._name_tokens(self._text(model.get("name", ""))) - for column in model.get("columns", []) or []: - tokens.update(self._name_tokens(self._text(column.get("name", "")))) - return tokens - - def _fallback_model_description(self, model: dict[str, Any]) -> str: - context = self._table_context(model) - if context & {"sales", "revenue", "order", "orders", "customer", "product"}: - return ( - "Captures commercial activity and related business dimensions for " - "sales reporting, performance analysis, and customer or product insights." - ) - if context & {"invoice", "payment", "price", "cost", "amount", "finance"}: - return ( - "Captures financial transactions and monetary measures used for " - "reconciliation, reporting, and performance analysis." - ) - if context & {"employee", "user", "person", "salesperson", "owner", "manager"}: - return ( - "Captures people and ownership attributes used to assign responsibility, " - "segment activity, and analyze performance." - ) - return ( - "Captures operational business records used for reporting, filtering, " - "trend analysis, and answering analytical questions." - ) - - def _fallback_column_description( - self, - model: dict[str, Any], - column: dict[str, Any], - ) -> str: - column_name = self._text(column.get("name", "")) - data_type = self._text(column.get("type", "")).lower() - tokens = self._name_tokens(column_name) - context = self._table_context(model) - - if tokens & {"division", "bu", "business", "unit", "department"}: - return "Organizational segment used to group records for ownership, reporting, and performance comparison." - if tokens & {"company", "entity", "organization", "org"}: - return "Legal or business entity associated with the record for company-level reporting and filtering." - if tokens & {"market", "region", "territory", "country", "state", "city", "location"}: - return "Geographic or market segment used to analyze activity by area and compare regional performance." - if tokens & {"product", "prod", "sku", "item", "material"}: - return "Product or item classification used to analyze sales, demand, and business activity by offering." - if tokens & {"type", "category", "class", "segment", "group"}: - return "Business classification used to segment records into meaningful reporting categories." - if tokens & {"customer", "client", "account"}: - return "Customer or account reference used to connect activity to the buyer or business relationship." - if tokens & {"salesperson", "seller", "rep", "owner", "manager", "person"}: - return "Responsible person or role associated with the record for ownership and performance analysis." - if tokens & {"status", "stage", "state"}: - return "Current business state used to track workflow progress, completion, or operational condition." - if tokens & {"date", "time", "day", "month", "year", "period", "created", "updated"}: - return "Time period used to sequence records, filter activity, and analyze trends over time." - if tokens & {"amount", "sales", "revenue", "cost", "price", "value", "total", "net", "gross"}: - return "Monetary measure used to calculate financial results, compare performance, and summarize business activity." - if tokens & {"quantity", "qty", "count", "units", "volume"}: - return "Quantity measure used to count activity, summarize volume, and compare operational scale." - if tokens & {"rate", "ratio", "percent", "percentage", "margin"}: - return "Calculated rate or percentage used to compare efficiency, contribution, or relative performance." - if tokens & {"id", "key", "code", "number", "no"}: - return "Identifier used to distinguish records and join this data with related business information." - if "date" in data_type or "time" in data_type: - return "Timestamp or calendar value used for time-based filtering, sequencing, and trend analysis." - if any(type_name in data_type for type_name in ("int", "float", "double", "decimal", "numeric", "number")): - return "Numeric business measure used for aggregation, comparison, and analytical calculations." - if context & {"sales", "order", "customer", "product"}: - return "Business attribute used to filter and explain commercial activity in reporting and analysis." - return "Business attribute used to categorize, filter, and explain records in analytical questions." - - def _is_low_quality_description(self, description: str, name: str) -> bool: - normalized = " ".join(description.lower().split()) - if not normalized: - return True - - name_text = self._humanize_name(name).lower() - low_quality_patterns = ( - "stores the", - "value used to describe or analyze", - "contains business records for", - "represents ", - "field from", - ) - if any(pattern in normalized for pattern in low_quality_patterns): - return True - return normalized in {name.lower(), name_text} - - def _fallback_output(self, chunk: dict[str, Any]) -> dict[str, Any]: - output: dict[str, Any] = {} - for model in chunk.get("mdl", {}).get("models", []): - model_name = self._text(model.get("name", "")) - if not model_name: - continue - - model_properties = self._properties(model) - model_description = self._text(model_properties.get("description", "")) - if self._is_low_quality_description(model_description, model_name): - model_description = self._fallback_model_description(model) - - columns = [] - for column in model.get("columns", []) or []: - if column.get("relationship"): - continue - column_name = self._text(column.get("name", "")) - if not column_name: - continue - column_properties = self._properties(column) - column_description = self._text( - column_properties.get("description", "") - ) - if self._is_low_quality_description(column_description, column_name): - column_description = self._fallback_column_description( - model, - column, - ) - columns.append( - { - "name": column_name, - "type": self._text(column.get("type", "")), - "properties": {"description": column_description}, - } - ) - - output[model_name] = { - "name": model_name, - "columns": columns, - "properties": {"description": model_description}, - } - return output - - def _complete_output_with_fallback( - self, - output: dict[str, Any], - chunk: dict[str, Any], - ) -> dict[str, Any]: - fallback = self._fallback_output(chunk) - completed = dict(output) - - for model_name, fallback_model in fallback.items(): - model_output = completed.get(model_name) - if not isinstance(model_output, dict): - completed[model_name] = fallback_model - continue - - properties = model_output.get("properties") - if not isinstance(properties, dict): - properties = {} - model_output["properties"] = properties - if self._is_low_quality_description( - self._text(properties.get("description", "")), - model_name, - ): - properties["description"] = self._text( - model_output.get("description") - ) or fallback_model["properties"]["description"] - if self._is_low_quality_description( - properties["description"], - model_name, - ): - properties["description"] = fallback_model["properties"][ - "description" - ] - - output_columns = { - column.get("name"): column - for column in model_output.get("columns", []) - if isinstance(column, dict) and column.get("name") - } - for fallback_column in fallback_model.get("columns", []): - column_name = fallback_column.get("name") - output_column = output_columns.get(column_name) - if not output_column: - model_output.setdefault("columns", []).append(fallback_column) - continue - - column_properties = output_column.get("properties") - if not isinstance(column_properties, dict): - column_properties = {} - output_column["properties"] = column_properties - if self._is_low_quality_description( - self._text(column_properties.get("description", "")), - column_name, - ): - column_properties["description"] = self._text( - output_column.get("description") - ) or fallback_column["properties"]["description"] - if self._is_low_quality_description( - column_properties["description"], - column_name, - ): - column_properties["description"] = fallback_column[ - "properties" - ]["description"] - - return completed + data_samples: dict[str, Any] = Field(default_factory=dict) def _chunking( - self, - mdl_dict: dict, - request: GenerateRequest, - chunk_size: int = SEMANTICS_MODEL_CHUNK_SIZE, + self, mdl_dict: dict, request: GenerateRequest, chunk_size: int = 50 ) -> list[dict]: template = { "user_prompt": request.user_prompt, "language": request.configurations.language, + "data_samples": request.data_samples, } - chunks: list[dict[str, Any]] = [] - selected_models = set(request.selected_models) - current_models: list[dict[str, Any]] = [] - current_column_count = 0 - - def _flush_current_models(): - nonlocal current_models, current_column_count - if not current_models: - return - chunks.append({"models": current_models}) - current_models = [] - current_column_count = 0 - - def _append_model(model: dict[str, Any]): - nonlocal current_models, current_column_count - column_count = len(model.get("columns") or []) - if current_models and current_column_count + column_count > chunk_size: - _flush_current_models() - current_models.append(model) - current_column_count += column_count - - for model in mdl_dict.get("models", []): - model_name = model.get("name") - if model_name not in selected_models: - continue - - columns = model.get("columns") or [] - if not columns: - _append_model({**model, "columns": []}) - continue - - for i in range(0, len(columns), chunk_size): - _append_model({**model, "columns": columns[i : i + chunk_size]}) - - _flush_current_models() + chunks = [ + { + **model, + "columns": model.get("columns", [])[i : i + chunk_size], + } + for model in mdl_dict.get("models", []) + if model.get("name") in request.selected_models + for i in range(0, len(model.get("columns", [])), chunk_size) + ] return [ { **template, - "mdl": chunk, - "selected_models": [model["name"] for model in chunk["models"]], + "mdl": {"models": [chunk]}, + "selected_models": [chunk["name"]], + "data_samples": { + chunk["name"]: request.data_samples.get(chunk["name"]) + } + if chunk["name"] in request.data_samples + else {}, } for chunk in chunks ] async def _generate_task(self, request_id: str, chunk: dict): - try: - logger.info( - "Calling configured LLM for semantics descriptions. " - "models=%s timeout_seconds=%s", - chunk.get("selected_models", []), - self._generation_timeout_seconds, - ) - resp = await asyncio.wait_for( - self._pipelines["semantics_description"].run(**chunk), - timeout=self._generation_timeout_seconds, - ) - output = resp.get("output") - if not output: - logger.warning( - "Configured LLM returned empty semantics output; " - "returning metadata-based fallback descriptions." - ) - output = self._fallback_output(chunk) - else: - output = self._complete_output_with_fallback(output, chunk) - except TimeoutError: - logger.warning( - "Semantics description LLM call timed out after %s seconds; " - "returning metadata-based fallback descriptions.", - self._generation_timeout_seconds, - ) - output = self._fallback_output(chunk) - - if not isinstance(output, dict): + resp = await self._pipelines["semantics_description"].run(**chunk) + output = resp.get("output") + if not isinstance(output, dict) or not output: raise ValueError("Semantics description pipeline returned no output") current = self[request_id] @@ -398,15 +110,6 @@ async def _generate_task(self, request_id: str, chunk: dict): current.response[key].setdefault("columns", []) current.response[key]["columns"].extend(output[key].get("columns", [])) - async def _generate_task_with_semaphore( - self, - semaphore: asyncio.Semaphore, - request_id: str, - chunk: dict, - ): - async with semaphore: - await self._generate_task(request_id, chunk) - @observe(name="Generate Semantics Description") @trace_metadata async def generate(self, request: GenerateRequest, **kwargs) -> Resource: @@ -421,17 +124,9 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: raise ValueError( "No selected models matched the current semantic model metadata" ) - semaphore = asyncio.Semaphore(SEMANTICS_MAX_CONCURRENT_LLM_CALLS) - await asyncio.gather( - *[ - self._generate_task_with_semaphore( - semaphore, - request.id, - chunk, - ) - for chunk in chunks - ] - ) + tasks = [self._generate_task(request.id, chunk) for chunk in chunks] + + await asyncio.gather(*tasks) self[request.id].status = "finished" self[request.id].trace_id = trace_id diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 5acc9af772..a9eeb7f2f6 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -88,6 +88,7 @@ export interface IWrenAIAdaptor { selectedModels: string[]; userPrompt: string; projectId: number; + dataSamples?: Record; }): Promise; getSemanticsDescriptionResult(queryId: string): Promise; generateRelationshipRecommendations(input: { @@ -432,6 +433,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { selectedModels: string[]; userPrompt: string; projectId: number; + dataSamples?: Record; }): Promise { try { const res = await axios.post( @@ -441,6 +443,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { selected_models: input.selectedModels, user_prompt: input.userPrompt, project_id: String(input.projectId), + data_samples: input.dataSamples || {}, }, ); return { queryId: res.data.id }; diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index fd763d3de6..c6403db9b9 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -11,44 +11,50 @@ export interface UpdateModelData { export interface NestedColumnMetadataInput { id: number; - displayName: string; - description: string; + displayName?: string; + description?: string; } export interface ColumnMetadataInput { id: number; - displayName: string; - description: string; + displayName?: string; + description?: string; } export interface CalculatedFieldMetadataInput { id: number; - description: string; + description?: string; } export interface RelationshipMetadataInput { id: number; - description: string; + description?: string; } export interface ViewColumnMetadataInput { referenceName: string; - description: string; + description?: string; } export interface UpdateModelMetadataInput { - displayName: string; - description: string; + displayName?: string; + description?: string; + columns?: Array; + nestedColumns?: Array; + calculatedFields?: Array; + relationships?: Array; +} + +export interface SaveModelingSemanticInput { + modelId: number; + description?: string; columns: Array; - nestedColumns: Array; - calculatedFields: Array; - relationships: Array; } export interface UpdateViewMetadataInput { - displayName: string; - description: string; - columns: Array; + displayName?: string; + description?: string; + columns?: Array; } export enum ExpressionName { diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 0227636968..95c9afd840 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -152,6 +152,7 @@ const resolvers = { validateView: modelResolver.validateView, updateViewMetadata: modelResolver.updateViewMetadata, generateModelingSemantics: modelResolver.generateModelingSemantics, + saveModelingSemantics: modelResolver.saveModelingSemantics, generateModelingRelationships: modelResolver.generateModelingRelationships, saveModelingRelationships: modelResolver.saveModelingRelationships, diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index a0a43f42ff..f3f13ba5e4 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -2,6 +2,7 @@ import { CreateModelData, UpdateModelData, UpdateModelMetadataInput, + SaveModelingSemanticInput, CreateCalculatedFieldData, UpdateCalculatedFieldData, UpdateViewMetadataInput, @@ -61,6 +62,7 @@ export class ModelResolver { this.generateModelingSemantics = this.generateModelingSemantics.bind(this); this.getModelingSemanticsResult = this.getModelingSemanticsResult.bind(this); + this.saveModelingSemantics = this.saveModelingSemantics.bind(this); this.generateModelingRelationships = this.generateModelingRelationships.bind(this); this.getModelingRelationshipsResult = @@ -463,14 +465,63 @@ export class ModelResolver { ) { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const dataSamples = await this.collectModelingDataSamples( + args.data.selectedModels, + project, + manifest, + ctx, + ); return await ctx.wrenAIAdaptor.generateSemanticsDescription({ manifest, selectedModels: args.data.selectedModels, userPrompt: args.data.userPrompt, projectId: project.id, + dataSamples, }); } + private async collectModelingDataSamples( + selectedModels: string[], + project: Project, + manifest: any, + ctx: IContext, + ): Promise> { + const samples: Record = {}; + const selectedModelNames = new Set(selectedModels); + const models = (manifest.models || []).filter((model) => + selectedModelNames.has(model.name), + ); + + await Promise.all( + models.map(async (model) => { + try { + const preview = (await ctx.queryService.preview( + `SELECT * FROM "${model.name}"`, + { + project, + modelingOnly: false, + manifest, + limit: 5, + refresh: true, + cacheEnabled: false, + }, + )) as PreviewDataResponse; + + samples[model.name] = { + columns: preview.columns || [], + rows: (preview.data || []).slice(0, 5), + }; + } catch (err: any) { + logger.warn( + `Failed to collect semantic sample data for model "${model.name}": ${err.message}`, + ); + } + }), + ); + + return samples; + } + public async getModelingSemanticsResult( _root: any, args: { queryId: string }, @@ -479,6 +530,58 @@ export class ModelResolver { return await ctx.wrenAIAdaptor.getSemanticsDescriptionResult(args.queryId); } + public async saveModelingSemantics( + _root: any, + args: { data: SaveModelingSemanticInput[] }, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const models = await ctx.modelRepository.findAllBy({ + projectId: project.id, + }); + const modelById = new Map(models.map((model) => [model.id, model])); + + await Promise.all( + (args.data || []).map(async (item) => { + const model = modelById.get(item.modelId); + if (!model) { + throw new Error(`Model not found: ${item.modelId}`); + } + + await this.handleUpdateModelMetadata( + { + displayName: undefined, + description: item.description, + columns: [], + nestedColumns: [], + calculatedFields: [], + relationships: [], + }, + model, + ctx, + item.modelId, + ); + + if (!isEmpty(item.columns)) { + await this.handleUpdateColumnMetadata( + { + displayName: undefined, + description: undefined, + columns: item.columns, + nestedColumns: [], + calculatedFields: [], + relationships: [], + }, + ctx, + ); + } + }), + ); + + this.markProjectDirty(project.id); + return { savedCount: args.data?.length || 0 }; + } + public async generateModelingRelationships( _root: any, _args: any, diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 21c9f214ea..37e2549949 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -379,6 +379,12 @@ export const typeDefs = gql` userPrompt: String! } + input SaveModelingSemanticInput { + modelId: Int! + description: String + columns: [UpdateColumnMetadataInput!]! + } + type NestedFieldInfo { id: Int! displayName: String! @@ -1338,6 +1344,7 @@ export const typeDefs = gql` data: UpdateViewMetadataInput! ): Boolean! generateModelingSemantics(data: GenerateModelingSemanticsInput!): JSON! + saveModelingSemantics(data: [SaveModelingSemanticInput!]!): JSON! generateModelingRelationships: JSON! saveModelingRelationships(data: [ModelingRelationshipInput!]!): JSON! diff --git a/wren-ui/src/apollo/server/services/modelService.ts b/wren-ui/src/apollo/server/services/modelService.ts index c0b8bb5083..bd1235945f 100644 --- a/wren-ui/src/apollo/server/services/modelService.ts +++ b/wren-ui/src/apollo/server/services/modelService.ts @@ -260,7 +260,7 @@ export class ModelService implements IModelService { projectId: id, }); - await Promise.all([ + await Promise.all( tables.map(async (table) => { const model = models.find((m) => m.sourceTableName === table.tableName); if (!model) { @@ -275,7 +275,7 @@ export class ModelService implements IModelService { properties: JSON.stringify(properties), }); }), - ]); + ); } public async batchUpdateColumnProperties(tables: SampleDatasetTable[]) { @@ -298,7 +298,7 @@ export class ModelService implements IModelService { return acc; }, []); - await Promise.all([ + await Promise.all( transformedColumns.map(async (column) => { if (!column.properties) { return; @@ -327,7 +327,7 @@ export class ModelService implements IModelService { properties: JSON.stringify(properties), }); }), - ]); + ); } public generateReferenceName(data: GenerateReferenceNameData): string { diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 97dcebd732..9a075d4fb9 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -94,6 +94,12 @@ const SAVE_MODELING_RELATIONSHIPS = gql` } `; +const SAVE_MODELING_SEMANTICS = gql` + mutation SaveModelingSemantics($data: [SaveModelingSemanticInput!]!) { + saveModelingSemantics(data: $data) + } +`; + const Diagram = dynamic(() => import('@/components/diagram'), { ssr: false }); // https://github.com/vercel/next.js/issues/4957#issuecomment-413841689 const ForwardDiagram = forwardRef(function ForwardDiagram(props: any, ref) { @@ -400,6 +406,7 @@ export default function Modeling() { const [generateModelingRelationships] = useMutation( GENERATE_MODELING_RELATIONSHIPS, ); + const [saveModelingSemantics] = useMutation(SAVE_MODELING_SEMANTICS); const [saveModelingRelationships] = useMutation(SAVE_MODELING_RELATIONSHIPS); const diagramData = useMemo(() => { @@ -854,34 +861,40 @@ export default function Modeling() { if (!diagramData) return; setAssistantLoading(true); if (assistantMode === 'semantics') { - for (const model of semanticResult) { + const data = semanticResult.flatMap((model) => { const diagramModel = diagramData.models.find( (item) => item.referenceName === model.name, ); - if (!diagramModel) continue; - await updateModelMetadata({ - variables: { - where: { id: diagramModel.modelId }, - data: { - description: model.description, - columns: (model.columns || []) - .map((column) => { - const field = diagramModel.fields.find( - (item) => item.referenceName === column.name, - ); - return field - ? { - id: field.columnId, - displayName: field.displayName, - description: column.description, - } - : null; - }) - .filter(Boolean), - }, - }, - }); + if (!diagramModel) return []; + return { + modelId: diagramModel.modelId, + description: model.description, + columns: (model.columns || []) + .map((column) => { + const field = diagramModel.fields.find( + (item) => item.referenceName === column.name, + ); + return field + ? { + id: field.columnId, + displayName: field.displayName, + description: column.description, + } + : null; + }) + .filter(Boolean), + }; + }); + + if (!data.length) { + throw new Error('No semantic descriptions to save.'); } + + await saveModelingSemantics({ + variables: { data }, + refetchQueries, + awaitRefetchQueries: true, + }); } if (assistantMode === 'relationships') { const res = await saveModelingRelationships({ From b9cb1615d89f8045c1e9f5f28dd8510f88058047 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 14:18:35 +0530 Subject: [PATCH 0674/1087] Restore legacy table retrieval grounding --- .../generation/semantics_description.py | 10 +- .../pipelines/indexing/table_description.py | 96 +------------------ .../web/v1/routers/semantics_description.py | 3 +- .../web/v1/services/semantics_description.py | 17 +--- .../apollo/server/adaptors/wrenAIAdaptor.ts | 3 - .../apollo/server/resolvers/modelResolver.ts | 49 ---------- 6 files changed, 13 insertions(+), 165 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 93b684ea1a..25e824370d 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -36,7 +36,7 @@ ``` Your task is to update this JSON structure by adding a `description` field inside both the `properties` attribute of each `column` and the `model` itself. -Each `description` should be derived from the user-provided dataset context, the full schema, relationships, model names, column names, data types, aliases, existing descriptions, and provided sample rows. +Each `description` should be derived from the user-provided dataset context, the full schema, relationships, model names, column names, data types, aliases, and existing descriptions. Follow these steps: 1. **For the `model`**: Write a clear natural language business description of the model's purpose and what real-world records it represents. Insert this description in the `properties` field of the `model`. 2. **For each `column`**: Write a clear natural language business description of the column's meaning, not just its technical name. Each column's description should be added under its respective `properties` field in the format: `'description': 'business description'`. @@ -89,9 +89,7 @@ Picked models: {{ picked_models }} Localization Language: {{ language }} -Sample data for picked models: {{ data_samples }} - -Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt, schema context, and sample data. +Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt and schema context. Do not omit selected models or columns. Do not copy the table or column name as the description. Use simple language that explains the business purpose, meaning, and analytical use of each field. """ @@ -153,13 +151,11 @@ def prompt( user_prompt: str, prompt_builder: PromptBuilder, language: str, - data_samples: dict[str, Any], ) -> dict: _prompt = prompt_builder.run( picked_models=picked_models, user_prompt=user_prompt, language=language, - data_samples=data_samples, ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -270,7 +266,6 @@ async def run( selected_models: list[str], mdl: dict, language: str = "en", - data_samples: dict[str, Any] | None = None, ) -> dict: logger.info("Semantics Description Generation pipeline is running...") return await self._pipe.execute( @@ -280,7 +275,6 @@ async def run( "selected_models": selected_models, "mdl": mdl, "language": language, - "data_samples": data_samples or {}, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index d1d0a6aa6a..cc24d3331b 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -14,89 +14,13 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.indexing import ( - AsyncDocumentWriter, - DocumentCleaner, - MDLValidator, - clean_display_name, -) +from src.pipelines.indexing import AsyncDocumentWriter, DocumentCleaner, MDLValidator logger = logging.getLogger("wren-ai-service") -MAX_TABLE_DESCRIPTION_COLUMNS = 200 -MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH = 4000 -MAX_TABLE_DESCRIPTION_COLUMN_DESCRIPTION_LENGTH = 500 - @component class TableDescriptionChunker: - def _normalize_text(self, value: Any) -> str: - return "" if value is None else str(value) - - def _truncate_description(self, description: Any) -> str: - normalized_description = self._normalize_text(description) - if len(normalized_description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: - return normalized_description - - return ( - normalized_description[:MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH].rstrip() - + "..." - ) - - def _truncate_column_description(self, description: Any) -> str: - normalized_description = self._normalize_text(description) - if len(normalized_description) <= MAX_TABLE_DESCRIPTION_COLUMN_DESCRIPTION_LENGTH: - return normalized_description - - return ( - normalized_description[ - :MAX_TABLE_DESCRIPTION_COLUMN_DESCRIPTION_LENGTH - ].rstrip() - + "..." - ) - - def _format_columns(self, columns: List[Any]) -> str: - normalized_columns = [self._normalize_text(column) for column in columns] - if len(normalized_columns) <= MAX_TABLE_DESCRIPTION_COLUMNS: - return ", ".join(normalized_columns) - - remaining_columns = len(normalized_columns) - MAX_TABLE_DESCRIPTION_COLUMNS - truncated_columns = normalized_columns[:MAX_TABLE_DESCRIPTION_COLUMNS] + [ - f"... (+{remaining_columns} more columns)" - ] - return ", ".join(truncated_columns) - - def _properties(self, payload: Dict[str, Any]) -> Dict[str, Any]: - properties = payload.get("properties") - return properties if isinstance(properties, dict) else {} - - def _display_name(self, properties: Dict[str, Any]) -> str: - return clean_display_name( - self._normalize_text(properties.get("displayName", "")) - ) - - def _column_text(self, column: Dict[str, Any]) -> str: - properties = self._properties(column) - parts = [self._normalize_text(column.get("name", ""))] - - display_name = self._display_name(properties) - if display_name: - parts.append(f"alias: {display_name}") - - data_type = self._normalize_text( - column.get("type", column.get("data_type", "")) - ) - if data_type: - parts.append(f"type: {data_type}") - - description = self._truncate_column_description( - properties.get("description", "") - ) - if description: - parts.append(f"description: {description}") - - return " | ".join(part for part in parts if part) - @component.output_types(documents=List[Document]) def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None): def _additional_meta() -> Dict[str, Any]: @@ -127,17 +51,11 @@ def _additional_meta() -> Dict[str, Any]: def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[str]: def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: - properties = self._properties(payload) - return { "mdl_type": mdl_type, "name": payload.get("name"), - "columns": [ - self._column_text(column) - for column in payload.get("columns", []) - if isinstance(column, dict) - ], - "properties": properties, + "columns": [column["name"] for column in payload.get("columns", [])], + "properties": payload.get("properties", {}), } resources = ( @@ -149,12 +67,8 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: return [ { "name": resource["name"], - "type": resource["mdl_type"], - "alias": self._display_name(resource["properties"]), - "description": self._truncate_description( - resource["properties"].get("description", "") - ), - "columns": self._format_columns(resource["columns"]), + "description": resource["properties"].get("description", ""), + "columns": ", ".join(resource["columns"]), } for resource in resources if resource["name"] is not None diff --git a/wren-ai-service/src/web/v1/routers/semantics_description.py b/wren-ai-service/src/web/v1/routers/semantics_description.py index 3118e59036..3e36d299bf 100644 --- a/wren-ai-service/src/web/v1/routers/semantics_description.py +++ b/wren-ai-service/src/web/v1/routers/semantics_description.py @@ -3,7 +3,7 @@ from typing import Literal, Optional from fastapi import APIRouter, BackgroundTasks, Depends -from pydantic import BaseModel, Field +from pydantic import BaseModel from src.globals import ( ServiceContainer, @@ -20,7 +20,6 @@ class PostRequest(BaseRequest): selected_models: list[str] user_prompt: str mdl: str - data_samples: dict = Field(default_factory=dict) class PostResponse(BaseModel): diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index afd5ab812b..67d282591c 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,11 +1,11 @@ import asyncio import logging -from typing import Any, Dict, Literal, Optional +from typing import Dict, Literal, Optional import orjson from cachetools import TTLCache from langfuse.decorators import observe -from pydantic import BaseModel, Field +from pydantic import BaseModel from src.core.pipeline import BasicPipeline from src.utils import trace_metadata @@ -58,7 +58,6 @@ class GenerateRequest(BaseRequest): selected_models: list[str] user_prompt: str mdl: str - data_samples: dict[str, Any] = Field(default_factory=dict) def _chunking( self, mdl_dict: dict, request: GenerateRequest, chunk_size: int = 50 @@ -66,7 +65,6 @@ def _chunking( template = { "user_prompt": request.user_prompt, "language": request.configurations.language, - "data_samples": request.data_samples, } chunks = [ @@ -84,20 +82,15 @@ def _chunking( **template, "mdl": {"models": [chunk]}, "selected_models": [chunk["name"]], - "data_samples": { - chunk["name"]: request.data_samples.get(chunk["name"]) - } - if chunk["name"] in request.data_samples - else {}, } for chunk in chunks ] async def _generate_task(self, request_id: str, chunk: dict): resp = await self._pipelines["semantics_description"].run(**chunk) - output = resp.get("output") - if not isinstance(output, dict) or not output: - raise ValueError("Semantics description pipeline returned no output") + output = resp.get("output") or {} + if not isinstance(output, dict): + raise ValueError("Semantics description pipeline returned invalid output") current = self[request_id] current.response = current.response or {} diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index a9eeb7f2f6..5acc9af772 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -88,7 +88,6 @@ export interface IWrenAIAdaptor { selectedModels: string[]; userPrompt: string; projectId: number; - dataSamples?: Record; }): Promise; getSemanticsDescriptionResult(queryId: string): Promise; generateRelationshipRecommendations(input: { @@ -433,7 +432,6 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { selectedModels: string[]; userPrompt: string; projectId: number; - dataSamples?: Record; }): Promise { try { const res = await axios.post( @@ -443,7 +441,6 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { selected_models: input.selectedModels, user_prompt: input.userPrompt, project_id: String(input.projectId), - data_samples: input.dataSamples || {}, }, ); return { queryId: res.data.id }; diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index f3f13ba5e4..3cd42c4505 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -465,63 +465,14 @@ export class ModelResolver { ) { const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const dataSamples = await this.collectModelingDataSamples( - args.data.selectedModels, - project, - manifest, - ctx, - ); return await ctx.wrenAIAdaptor.generateSemanticsDescription({ manifest, selectedModels: args.data.selectedModels, userPrompt: args.data.userPrompt, projectId: project.id, - dataSamples, }); } - private async collectModelingDataSamples( - selectedModels: string[], - project: Project, - manifest: any, - ctx: IContext, - ): Promise> { - const samples: Record = {}; - const selectedModelNames = new Set(selectedModels); - const models = (manifest.models || []).filter((model) => - selectedModelNames.has(model.name), - ); - - await Promise.all( - models.map(async (model) => { - try { - const preview = (await ctx.queryService.preview( - `SELECT * FROM "${model.name}"`, - { - project, - modelingOnly: false, - manifest, - limit: 5, - refresh: true, - cacheEnabled: false, - }, - )) as PreviewDataResponse; - - samples[model.name] = { - columns: preview.columns || [], - rows: (preview.data || []).slice(0, 5), - }; - } catch (err: any) { - logger.warn( - `Failed to collect semantic sample data for model "${model.name}": ${err.message}`, - ); - } - }), - ); - - return samples; - } - public async getModelingSemanticsResult( _root: any, args: { queryId: string }, From 0f24e30a52a1d213bdd82f2206b8db9266e039e6 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 15:07:54 +0530 Subject: [PATCH 0675/1087] Route data asks to SQL generation --- .../generation/intent_classification.py | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..e2569f803a 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -33,8 +33,10 @@ - **Rephrase Question:** Rewrite follow-up questions into full standalone questions using prior conversation context. - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. -- **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. -- **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. +- **Data Retrieval Requests:** If the user asks to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data from the connected database, classify it as `TEXT_TO_SQL`. +- **Database Schema Exploration:** If the user asks about available tables, columns, relationships, schema meaning, or what can be asked, classify it as `GENERAL`. +- **Out-of-Scope Queries:** If the question is unrelated to the database schema or data retrieval, classify it as `MISLEADING_QUERY`. +- **Incomplete Queries:** If the question references unresolved placeholders (e.g., "the following", "these", "those") without providing them or prior context, classify it as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. ### Intent Definitions ### @@ -43,14 +45,14 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes references to specific tables, columns, or data details. -- The question includes **complete information** with specific tables, columns, or data values needed for execution. -- The question provides **all necessary parameters** to generate executable SQL. +- The user's inputs ask to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data. +- The question can be answered by selecting relevant tables and columns from the provided schema, even if the user does not mention exact physical table or column names. +- The question includes enough business meaning, dimensions, metrics, filters, or time criteria to attempt SQL generation from the schema. **Requirements:** -- Must have complete filter criteria, specific values, or clear references to previous context. -- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. -- Reference phrases from the user's inputs that clearly relate to the schema. +- Do not require the user to explicitly name a table or column. +- Use the provided schema context to decide whether the user's business terms can be answered by SQL. +- Reference phrases from the user's inputs that clearly indicate a data retrieval request. **Examples:** - "What is the total sales for last quarter?" @@ -61,9 +63,10 @@ **When to Use:** - The user seeks general information about the database schema or its overall capabilities. +- The user asks about available tables, columns, relationships, schema meaning, or what questions can be asked. - The query references **missing information** (e.g., "the following items" without listing them). - The query contains **placeholder references** that cannot be resolved from context. -- The query is **incomplete for SQL generation** despite mentioning database concepts. +- The query is asking for explanation or guidance rather than retrieval, filtering, ordering, aggregation, or analysis of rows. **Requirements:** - Incorporate phrases from the user's inputs that indicate incompleteness or lack of relevance to the database schema. @@ -72,6 +75,8 @@ **Examples:** - "What is the dataset about?" - "Tell me more about the database." +- "Explain the customer table to me." +- "What tables do I have?" - "How can I analyze customer behavior with this data?" - "Show me orders for these products" (without specifying which products) - "Filter by the criteria I mentioned" (without previous context defining criteria) @@ -93,7 +98,7 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. +- The user's inputs cannot be interpreted as a database question or data retrieval request. - It appears off-topic or is simply a casual conversation starter. **Requirements:** From 600105861765df604199ba7b88e06213a1741d41 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 15:32:35 +0530 Subject: [PATCH 0676/1087] Revert "Route data asks to SQL generation" This reverts commit 0f24e30a52a1d213bdd82f2206b8db9266e039e6. --- .../generation/intent_classification.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index e2569f803a..4d6cd313cd 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -33,10 +33,8 @@ - **Rephrase Question:** Rewrite follow-up questions into full standalone questions using prior conversation context. - **Concise Reasoning:** The reasoning must be clear, concise, and limited to 20 words. - **Language Consistency:** Use the same language as specified in the user's output language for the rephrased question and reasoning. -- **Data Retrieval Requests:** If the user asks to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data from the connected database, classify it as `TEXT_TO_SQL`. -- **Database Schema Exploration:** If the user asks about available tables, columns, relationships, schema meaning, or what can be asked, classify it as `GENERAL`. -- **Out-of-Scope Queries:** If the question is unrelated to the database schema or data retrieval, classify it as `MISLEADING_QUERY`. -- **Incomplete Queries:** If the question references unresolved placeholders (e.g., "the following", "these", "those") without providing them or prior context, classify it as `GENERAL`. +- **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. +- **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. ### Intent Definitions ### @@ -45,14 +43,14 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The user's inputs ask to retrieve, list, show, compare, count, aggregate, rank, filter, group, sort, or analyze data. -- The question can be answered by selecting relevant tables and columns from the provided schema, even if the user does not mention exact physical table or column names. -- The question includes enough business meaning, dimensions, metrics, filters, or time criteria to attempt SQL generation from the schema. +- The question (or related previous query) includes references to specific tables, columns, or data details. +- The question includes **complete information** with specific tables, columns, or data values needed for execution. +- The question provides **all necessary parameters** to generate executable SQL. **Requirements:** -- Do not require the user to explicitly name a table or column. -- Use the provided schema context to decide whether the user's business terms can be answered by SQL. -- Reference phrases from the user's inputs that clearly indicate a data retrieval request. +- Must have complete filter criteria, specific values, or clear references to previous context. +- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. +- Reference phrases from the user's inputs that clearly relate to the schema. **Examples:** - "What is the total sales for last quarter?" @@ -63,10 +61,9 @@ **When to Use:** - The user seeks general information about the database schema or its overall capabilities. -- The user asks about available tables, columns, relationships, schema meaning, or what questions can be asked. - The query references **missing information** (e.g., "the following items" without listing them). - The query contains **placeholder references** that cannot be resolved from context. -- The query is asking for explanation or guidance rather than retrieval, filtering, ordering, aggregation, or analysis of rows. +- The query is **incomplete for SQL generation** despite mentioning database concepts. **Requirements:** - Incorporate phrases from the user's inputs that indicate incompleteness or lack of relevance to the database schema. @@ -75,8 +72,6 @@ **Examples:** - "What is the dataset about?" - "Tell me more about the database." -- "Explain the customer table to me." -- "What tables do I have?" - "How can I analyze customer behavior with this data?" - "Show me orders for these products" (without specifying which products) - "Filter by the criteria I mentioned" (without previous context defining criteria) @@ -98,7 +93,7 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs cannot be interpreted as a database question or data retrieval request. +- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. - It appears off-topic or is simply a casual conversation starter. **Requirements:** From b7eb94f990fa3f4a778a7207a108698a2741c456 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 11:19:27 +0000 Subject: [PATCH 0677/1087] Limit table description retrieval text --- .../pipelines/indexing/table_description.py | 49 +++++++++++++++++-- .../indexing/test_table_description.py | 43 ++++++++++++++++ 2 files changed, 87 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index cc24d3331b..9ecc863802 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -18,9 +18,40 @@ logger = logging.getLogger("wren-ai-service") +MAX_TABLE_DESCRIPTION_COLUMNS = 200 +MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH = 4000 + @component class TableDescriptionChunker: + def _normalize_text(self, value: Any) -> str: + return "" if value is None else str(value) + + def _truncate_description(self, description: Any) -> str: + normalized_description = self._normalize_text(description) + if len(normalized_description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: + return normalized_description + + return ( + normalized_description[:MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH].rstrip() + + "..." + ) + + def _format_columns(self, columns: List[Any]) -> str: + normalized_columns = [self._normalize_text(column) for column in columns] + if len(normalized_columns) <= MAX_TABLE_DESCRIPTION_COLUMNS: + return ", ".join(normalized_columns) + + remaining_columns = len(normalized_columns) - MAX_TABLE_DESCRIPTION_COLUMNS + truncated_columns = normalized_columns[:MAX_TABLE_DESCRIPTION_COLUMNS] + [ + f"... (+{remaining_columns} more columns)" + ] + return ", ".join(truncated_columns) + + def _properties(self, payload: Dict[str, Any]) -> Dict[str, Any]: + properties = payload.get("properties") + return properties if isinstance(properties, dict) else {} + @component.output_types(documents=List[Document]) def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None): def _additional_meta() -> Dict[str, Any]: @@ -49,13 +80,19 @@ def _additional_meta() -> Dict[str, Any]: ] } - def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[str]: + def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[Dict[str, Any]]: def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: + properties = self._properties(payload) + return { "mdl_type": mdl_type, "name": payload.get("name"), - "columns": [column["name"] for column in payload.get("columns", [])], - "properties": payload.get("properties", {}), + "columns": [ + self._normalize_text(column.get("name", "")) + for column in payload.get("columns", []) + if isinstance(column, dict) + ], + "properties": properties, } resources = ( @@ -67,8 +104,10 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: return [ { "name": resource["name"], - "description": resource["properties"].get("description", ""), - "columns": ", ".join(resource["columns"]), + "description": self._truncate_description( + resource["properties"].get("description", "") + ), + "columns": self._format_columns(resource["columns"]), } for resource in resources if resource["name"] is not None diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index 244649e6b1..b334e9001d 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -151,6 +151,49 @@ def test_table_description_null_description(): ) +def test_table_description_excludes_generated_column_descriptions(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + { + "name": "orders", + "properties": {"description": "Customer purchase transactions."}, + "columns": [ + { + "name": "Division", + "type": "varchar", + "properties": { + "description": "Generic generated division description." + }, + }, + { + "name": "SalesAmount", + "type": "float", + "properties": { + "description": "Generic generated sales amount description." + }, + }, + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = chunker.run(mdl) + + document: Document = actual["documents"][0] + assert document.content == str( + { + "name": "orders", + "description": "Customer purchase transactions.", + "columns": "Division, SalesAmount", + } + ) + assert "Generic generated" not in document.content + + def test_table_description_truncates_long_column_lists(): chunker = TableDescriptionChunker() mdl = { From a37dce80a2a450dd40c10cb952d1796475a12409 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 17:12:55 +0530 Subject: [PATCH 0678/1087] Ground SQL generation on schema metadata --- .../generation/semantics_description.py | 7 +++- .../src/pipelines/generation/utils/sql.py | 41 ++++--------------- .../src/pipelines/indexing/db_schema.py | 10 +++-- .../src/pipelines/indexing/utils/helper.py | 5 ++- .../retrieval/db_schema_retrieval.py | 8 ++-- 5 files changed, 31 insertions(+), 40 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 25e824370d..58a63afc0e 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -38,12 +38,15 @@ Your task is to update this JSON structure by adding a `description` field inside both the `properties` attribute of each `column` and the `model` itself. Each `description` should be derived from the user-provided dataset context, the full schema, relationships, model names, column names, data types, aliases, and existing descriptions. Follow these steps: -1. **For the `model`**: Write a clear natural language business description of the model's purpose and what real-world records it represents. Insert this description in the `properties` field of the `model`. +1. **For the `model`**: Write a clear natural language business description of the model's purpose, what real-world records it represents, and the common analysis questions it can answer. Insert this description in the `properties` field of the `model`. 2. **For each `column`**: Write a clear natural language business description of the column's meaning, not just its technical name. Each column's description should be added under its respective `properties` field in the format: `'description': 'business description'`. 3. Ensure that the output is a well-formatted JSON structure, preserving the input's original format and adding the appropriate `description` fields. 4. Avoid repeating technical table or column names as the whole description. Prefer business meaning such as identifiers, dates, amounts, statuses, dimensions, ownership, and operational usage. 5. Do not use generic boilerplate such as "stores the value", "contains records for", or "field from". Explain what the data means to a business user. 6. Make every model and column description unique, human-readable, concise, factual, and useful for text-to-SQL retrieval. +7. Use the model name, display label, existing description, column names, column display labels, and data types so descriptions include searchable business terms available from the metadata. +8. Do not invent table names, column names, relationships, or business concepts that are not supported by the provided model metadata. +9. If the metadata is technical or abbreviated, describe the observable business concepts from the available names and labels instead of copying the technical names. ### Output Format: @@ -92,6 +95,8 @@ Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt and schema context. Do not omit selected models or columns. Do not copy the table or column name as the description. Use simple language that explains the business purpose, meaning, and analytical use of each field. +For each model description, include the main measures, dates, identifiers, dimensions, statuses, and business entities represented by its columns so vector retrieval can match natural-language questions to the correct model. +Keep descriptions factual and grounded in the picked model metadata only. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 0e2a2225e2..cbc07df03a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -198,20 +198,12 @@ async def _classify_generation_result( - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- Table/column aliases in schema comments are display labels only. Never use an alias as an executable table or column identifier. +- Schema comments are metadata for understanding the data. They are not SQL syntax. +- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. +- Never use a `display_label`, alias, or description as an executable table or column identifier. - Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. -- You may use aliases from schema comments only after AS in the final SELECT clause. -- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section only for final SELECT output labels. - - EXAMPLE - DATABASE SCHEMA - /* {"alias":"_orders","description":"A model representing the orders data."} */ - CREATE TABLE orders ( - -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} - ApprovedTimestamp TIMESTAMP - } - - SQL - SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; +- You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. +- Only apply numeric aggregate functions such as SUM or AVG to numeric columns or measures from the DATABASE SCHEMA. If a column is not numeric in the schema, do not aggregate it directly unless the provided SQL FUNCTIONS and database dialect support the explicit cast you use. - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -371,29 +363,11 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - For Example: - DATA SCHEMA: - `/* {"alias":"users","description":"A model representing the users data."} */ - CREATE TABLE users ( - -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} - address JSON - )` - To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. - - For Example: - DATA SCHEMA - `/* {"alias":"my_table","description":"A test my_table"} */ - CREATE TABLE my_table ( - -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} - elements JSON - )` - To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". @@ -446,7 +420,10 @@ async def _classify_generation_result( 11. Do not include ```markdown or ``` in the answer. 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +14. Use only table and column names that appear as identifiers in the DATABASE SCHEMA when writing `table:` and `column:` references. +15. Schema comments, display labels, aliases, and descriptions are context only. Do not use them as executable table or column names in the reasoning plan. +16. Do not create table or column names from words in the user's question. If a requested concept is available only through schema metadata, refer to the corresponding schema identifier. +17. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 394d087b46..4ad05b8bd0 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -4,6 +4,7 @@ import uuid from typing import Any, Dict, List, Optional +import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from hamilton.function_modifiers import extract_fields @@ -131,14 +132,17 @@ def _convert_models_and_relationships( ) -> List[Dict[str, str]]: def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) + table_name = model["name"] model_properties = { - "alias": clean_display_name(properties.get("displayName", "")), + "identifier": table_name, + "display_label": clean_display_name( + properties.get("displayName", "") + ), "description": properties.get("description", ""), } - comment = f"\n/* {str(model_properties)} */\n" + comment = f"\n/* {orjson.dumps(model_properties).decode('utf-8')} */\n" - table_name = model["name"] payload = { "type": "TABLE", "comment": comment, diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 8e0e2a6b25..61ac17084c 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -36,7 +36,10 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: display_name = props.get("displayName", "") description = props.get("description", "") column_properties = { - "alias": clean_display_name("" if display_name is None else str(display_name)), + "identifier": column.get("name", ""), + "display_label": clean_display_name( + "" if display_name is None else str(display_name) + ), "description": "" if description is None else str(description), } diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..92e043d98d 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -79,9 +79,11 @@ - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. - Be logical, concise, and ensure the output strictly follows the required JSON format. -- Use table name used in the "Create Table" statement, don't use "alias". -- Match Column names with the definition in the "Create Table" statement. -- Match Table names with the definition in the "Create Table" statement. +- Schema comments are metadata for understanding the data. They are not SQL syntax. +- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. +- Use table names from the "Create Table" statements, not display labels, aliases, or descriptions. +- Match column names exactly with the definitions in the "Create Table" statements. +- Match table names exactly with the definitions in the "Create Table" statements. Good luck! From 9cd5c133444d5afc637fda0d3c79d11e89de7c61 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 11:47:12 +0000 Subject: [PATCH 0679/1087] Guard SQL generation against unretrieved tables --- .../generation/followup_sql_generation.py | 2 + .../pipelines/generation/sql_correction.py | 2 + .../pipelines/generation/sql_generation.py | 2 + .../pipelines/generation/sql_regeneration.py | 2 + .../src/pipelines/generation/utils/sql.py | 177 ++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 66 +++++++ 6 files changed, 251 insertions(+) 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 35cfb8fccf..fb546e1b14 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,12 +147,14 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), + contexts=documents, project_id=project_id, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..46abf75bf4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -120,12 +120,14 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), + contexts=documents, project_id=project_id, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..3d2074ce9a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -139,6 +139,7 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -146,6 +147,7 @@ async def post_process( ) -> dict: return await post_processor.run( generate_sql.get("replies"), + contexts=documents, project_id=project_id, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 4b7284aa26..7cdb6d2c81 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -157,10 +157,12 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, + documents: list[str] | None = None, project_id: str | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), + contexts=documents, project_id=project_id, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index cbc07df03a..c849337c11 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, List import aiohttp @@ -16,6 +17,158 @@ logger = logging.getLogger("wren-ai-service") +_SQL_IDENTIFIER = ( + r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][\w$]*)' + r'(?:\s*\.\s*(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][\w$]*))*' +) +_DDL_TABLE_PATTERN = re.compile( + rf"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?P{_SQL_IDENTIFIER})", + re.IGNORECASE, +) +_SQL_TABLE_REFERENCE_PATTERN = re.compile( + rf"\b(?:FROM|JOIN)\s+(?P{_SQL_IDENTIFIER})", + re.IGNORECASE, +) +_SQL_CTE_PATTERN = re.compile( + rf"(?:\bWITH\s+(?:RECURSIVE\s+)?|,)\s*(?P{_SQL_IDENTIFIER})\s+AS\s*\(", + re.IGNORECASE, +) + + +def _split_identifier_parts(identifier: str) -> list[str]: + parts = [] + current = [] + quote = "" + bracket_depth = 0 + + for char in identifier.strip(): + if quote: + current.append(char) + if char == quote: + quote = "" + continue + + if bracket_depth: + current.append(char) + if char == "]": + bracket_depth = 0 + continue + + if char in {'"', "`"}: + quote = char + current.append(char) + continue + + if char == "[": + bracket_depth = 1 + current.append(char) + continue + + if char == ".": + part = "".join(current).strip() + if part: + parts.append(part) + current = [] + continue + + current.append(char) + + part = "".join(current).strip() + if part: + parts.append(part) + + return parts + + +def _normalize_identifier_part(identifier: str) -> str: + identifier = identifier.strip() + if ( + len(identifier) >= 2 + and ( + (identifier[0] == identifier[-1] and identifier[0] in {'"', "`"}) + or (identifier[0] == "[" and identifier[-1] == "]") + ) + ): + identifier = identifier[1:-1] + + return identifier.strip().lower() + + +def _normalize_identifier(identifier: str) -> str: + return ".".join( + part + for part in ( + _normalize_identifier_part(part) + for part in _split_identifier_parts(identifier) + ) + if part + ) + + +def _identifier_leaf(identifier: str) -> str: + normalized_parts = [ + _normalize_identifier_part(part) for part in _split_identifier_parts(identifier) + ] + return normalized_parts[-1] if normalized_parts else "" + + +def _allowed_table_names(contexts: list[str] | None) -> set[str]: + allowed = set() + for context in contexts or []: + for match in _DDL_TABLE_PATTERN.finditer(str(context)): + identifier = match.group("identifier") + normalized = _normalize_identifier(identifier) + leaf = _identifier_leaf(identifier) + if normalized: + allowed.add(normalized) + if leaf: + allowed.add(leaf) + + return allowed + + +def _cte_names(sql: str) -> set[str]: + ctes = set() + for match in _SQL_CTE_PATTERN.finditer(sql): + identifier = match.group("identifier") + normalized = _normalize_identifier(identifier) + leaf = _identifier_leaf(identifier) + if normalized: + ctes.add(normalized) + if leaf: + ctes.add(leaf) + + return ctes + + +def _is_table_function_reference(sql: str, end_position: int) -> bool: + remainder = sql[end_position:].lstrip() + return remainder.startswith("(") + + +def find_unavailable_sql_tables(sql: str, contexts: list[str] | None) -> list[str]: + allowed_tables = _allowed_table_names(contexts) + if not allowed_tables: + return [] + + ctes = _cte_names(sql) + unavailable = [] + for match in _SQL_TABLE_REFERENCE_PATTERN.finditer(sql): + identifier = match.group("identifier") + if _is_table_function_reference(sql, match.end("identifier")): + continue + + normalized = _normalize_identifier(identifier) + leaf = _identifier_leaf(identifier) + if normalized in allowed_tables or leaf in allowed_tables: + continue + if normalized in ctes or leaf in ctes: + continue + + unavailable.append(identifier) + + return sorted(set(unavailable)) + @component class SQLGenPostProcessor: @@ -29,6 +182,7 @@ def __init__(self, engine: Engine): async def run( self, replies: List[str] | List[List[str]], + contexts: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -44,6 +198,26 @@ async def run( "sql" ] + unavailable_tables = find_unavailable_sql_tables( + cleaned_generation_result, + contexts, + ) + if unavailable_tables: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_GROUNDING", + "error": ( + "Generated SQL references table(s) not present in the " + "retrieved schema: " + + ", ".join(unavailable_tables) + ), + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, @@ -167,6 +341,9 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. +- The DATABASE SCHEMA section is the complete and only source of executable table and column identifiers. +- MUST NOT introduce, infer, copy, or repair table or column identifiers from the user question, reasoning plan, SQL samples, failed SQL, or error messages unless the same identifiers appear in the DATABASE SCHEMA section. +- Do not copy identifiers from failed SQL into corrected SQL unless they are present in the DATABASE SCHEMA section. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 51f9d4f2e4..8d62d56c58 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,5 +1,9 @@ +import pytest + from src.pipelines.generation.utils.sql import ( + SQLGenPostProcessor, construct_instructions, + find_unavailable_sql_tables, get_json_field_instructions, get_metric_instructions, get_sql_generation_system_prompt, @@ -47,3 +51,65 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "DATABASE SCHEMA section is the complete and only source" in prompt assert "MUST NOT introduce, infer, copy, or repair" in prompt assert "Do not copy identifiers" in prompt + + +def test_find_unavailable_sql_tables_rejects_non_retrieved_table(): + contexts = [ + """ + CREATE TABLE dbo_dimCustomers ( + CustomerID VARCHAR + ); + """ + ] + + assert find_unavailable_sql_tables( + "SELECT * FROM dbo_tblPayments ORDER BY PaymentDate DESC", + contexts, + ) == ["dbo_tblPayments"] + + +def test_find_unavailable_sql_tables_allows_retrieved_table_and_cte(): + contexts = [ + """ + CREATE TABLE "dbo_tblPayments" ( + PaymentDate TIMESTAMP + ); + """ + ] + + assert ( + find_unavailable_sql_tables( + """ + WITH recent_payments AS ( + SELECT * FROM "dbo_tblPayments" + ) + SELECT * FROM recent_payments + """, + contexts, + ) + == [] + ) + + +@pytest.mark.asyncio +async def test_sql_post_processor_blocks_sql_using_non_retrieved_table(): + class Engine: + async def execute_sql(self, *_args, **_kwargs): + raise AssertionError("engine should not be called") + + post_processor = SQLGenPostProcessor(Engine()) + + result = await post_processor.run( + ['{"sql": "SELECT * FROM dbo_tblPayments"}'], + contexts=[ + """ + CREATE TABLE dbo_dimCustomers ( + CustomerID VARCHAR + ); + """ + ], + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SCHEMA_GROUNDING" + assert "dbo_tblPayments" in result["invalid_generation_result"]["error"] From c7866b9eb795b78cd1c3159aba7a8e545c8201d3 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 12:01:30 +0000 Subject: [PATCH 0680/1087] Revert "Guard SQL generation against unretrieved tables" This reverts commit 9cd5c133444d5afc637fda0d3c79d11e89de7c61. --- .../generation/followup_sql_generation.py | 2 - .../pipelines/generation/sql_correction.py | 2 - .../pipelines/generation/sql_generation.py | 2 - .../pipelines/generation/sql_regeneration.py | 2 - .../src/pipelines/generation/utils/sql.py | 177 ------------------ .../pipelines/generation/test_sql_utils.py | 66 ------- 6 files changed, 251 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 fb546e1b14..35cfb8fccf 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,14 +147,12 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), - contexts=documents, project_id=project_id, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 46abf75bf4..973b8c69a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -120,14 +120,12 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), - contexts=documents, project_id=project_id, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 3d2074ce9a..1ee4952b3e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -139,7 +139,6 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -147,7 +146,6 @@ async def post_process( ) -> dict: return await post_processor.run( generate_sql.get("replies"), - contexts=documents, project_id=project_id, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 7cdb6d2c81..4b7284aa26 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -157,12 +157,10 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, - documents: list[str] | None = None, project_id: str | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), - contexts=documents, project_id=project_id, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c849337c11..cbc07df03a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,4 @@ import logging -import re from typing import Any, Dict, List import aiohttp @@ -17,158 +16,6 @@ logger = logging.getLogger("wren-ai-service") -_SQL_IDENTIFIER = ( - r'(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][\w$]*)' - r'(?:\s*\.\s*(?:"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][\w$]*))*' -) -_DDL_TABLE_PATTERN = re.compile( - rf"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?P{_SQL_IDENTIFIER})", - re.IGNORECASE, -) -_SQL_TABLE_REFERENCE_PATTERN = re.compile( - rf"\b(?:FROM|JOIN)\s+(?P{_SQL_IDENTIFIER})", - re.IGNORECASE, -) -_SQL_CTE_PATTERN = re.compile( - rf"(?:\bWITH\s+(?:RECURSIVE\s+)?|,)\s*(?P{_SQL_IDENTIFIER})\s+AS\s*\(", - re.IGNORECASE, -) - - -def _split_identifier_parts(identifier: str) -> list[str]: - parts = [] - current = [] - quote = "" - bracket_depth = 0 - - for char in identifier.strip(): - if quote: - current.append(char) - if char == quote: - quote = "" - continue - - if bracket_depth: - current.append(char) - if char == "]": - bracket_depth = 0 - continue - - if char in {'"', "`"}: - quote = char - current.append(char) - continue - - if char == "[": - bracket_depth = 1 - current.append(char) - continue - - if char == ".": - part = "".join(current).strip() - if part: - parts.append(part) - current = [] - continue - - current.append(char) - - part = "".join(current).strip() - if part: - parts.append(part) - - return parts - - -def _normalize_identifier_part(identifier: str) -> str: - identifier = identifier.strip() - if ( - len(identifier) >= 2 - and ( - (identifier[0] == identifier[-1] and identifier[0] in {'"', "`"}) - or (identifier[0] == "[" and identifier[-1] == "]") - ) - ): - identifier = identifier[1:-1] - - return identifier.strip().lower() - - -def _normalize_identifier(identifier: str) -> str: - return ".".join( - part - for part in ( - _normalize_identifier_part(part) - for part in _split_identifier_parts(identifier) - ) - if part - ) - - -def _identifier_leaf(identifier: str) -> str: - normalized_parts = [ - _normalize_identifier_part(part) for part in _split_identifier_parts(identifier) - ] - return normalized_parts[-1] if normalized_parts else "" - - -def _allowed_table_names(contexts: list[str] | None) -> set[str]: - allowed = set() - for context in contexts or []: - for match in _DDL_TABLE_PATTERN.finditer(str(context)): - identifier = match.group("identifier") - normalized = _normalize_identifier(identifier) - leaf = _identifier_leaf(identifier) - if normalized: - allowed.add(normalized) - if leaf: - allowed.add(leaf) - - return allowed - - -def _cte_names(sql: str) -> set[str]: - ctes = set() - for match in _SQL_CTE_PATTERN.finditer(sql): - identifier = match.group("identifier") - normalized = _normalize_identifier(identifier) - leaf = _identifier_leaf(identifier) - if normalized: - ctes.add(normalized) - if leaf: - ctes.add(leaf) - - return ctes - - -def _is_table_function_reference(sql: str, end_position: int) -> bool: - remainder = sql[end_position:].lstrip() - return remainder.startswith("(") - - -def find_unavailable_sql_tables(sql: str, contexts: list[str] | None) -> list[str]: - allowed_tables = _allowed_table_names(contexts) - if not allowed_tables: - return [] - - ctes = _cte_names(sql) - unavailable = [] - for match in _SQL_TABLE_REFERENCE_PATTERN.finditer(sql): - identifier = match.group("identifier") - if _is_table_function_reference(sql, match.end("identifier")): - continue - - normalized = _normalize_identifier(identifier) - leaf = _identifier_leaf(identifier) - if normalized in allowed_tables or leaf in allowed_tables: - continue - if normalized in ctes or leaf in ctes: - continue - - unavailable.append(identifier) - - return sorted(set(unavailable)) - @component class SQLGenPostProcessor: @@ -182,7 +29,6 @@ def __init__(self, engine: Engine): async def run( self, replies: List[str] | List[List[str]], - contexts: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -198,26 +44,6 @@ async def run( "sql" ] - unavailable_tables = find_unavailable_sql_tables( - cleaned_generation_result, - contexts, - ) - if unavailable_tables: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_GROUNDING", - "error": ( - "Generated SQL references table(s) not present in the " - "retrieved schema: " - + ", ".join(unavailable_tables) - ), - "correlation_id": "", - }, - } - ( valid_generation_result, invalid_generation_result, @@ -341,9 +167,6 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. -- The DATABASE SCHEMA section is the complete and only source of executable table and column identifiers. -- MUST NOT introduce, infer, copy, or repair table or column identifiers from the user question, reasoning plan, SQL samples, failed SQL, or error messages unless the same identifiers appear in the DATABASE SCHEMA section. -- Do not copy identifiers from failed SQL into corrected SQL unless they are present in the DATABASE SCHEMA section. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 8d62d56c58..51f9d4f2e4 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,9 +1,5 @@ -import pytest - from src.pipelines.generation.utils.sql import ( - SQLGenPostProcessor, construct_instructions, - find_unavailable_sql_tables, get_json_field_instructions, get_metric_instructions, get_sql_generation_system_prompt, @@ -51,65 +47,3 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "DATABASE SCHEMA section is the complete and only source" in prompt assert "MUST NOT introduce, infer, copy, or repair" in prompt assert "Do not copy identifiers" in prompt - - -def test_find_unavailable_sql_tables_rejects_non_retrieved_table(): - contexts = [ - """ - CREATE TABLE dbo_dimCustomers ( - CustomerID VARCHAR - ); - """ - ] - - assert find_unavailable_sql_tables( - "SELECT * FROM dbo_tblPayments ORDER BY PaymentDate DESC", - contexts, - ) == ["dbo_tblPayments"] - - -def test_find_unavailable_sql_tables_allows_retrieved_table_and_cte(): - contexts = [ - """ - CREATE TABLE "dbo_tblPayments" ( - PaymentDate TIMESTAMP - ); - """ - ] - - assert ( - find_unavailable_sql_tables( - """ - WITH recent_payments AS ( - SELECT * FROM "dbo_tblPayments" - ) - SELECT * FROM recent_payments - """, - contexts, - ) - == [] - ) - - -@pytest.mark.asyncio -async def test_sql_post_processor_blocks_sql_using_non_retrieved_table(): - class Engine: - async def execute_sql(self, *_args, **_kwargs): - raise AssertionError("engine should not be called") - - post_processor = SQLGenPostProcessor(Engine()) - - result = await post_processor.run( - ['{"sql": "SELECT * FROM dbo_tblPayments"}'], - contexts=[ - """ - CREATE TABLE dbo_dimCustomers ( - CustomerID VARCHAR - ); - """ - ], - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "SCHEMA_GROUNDING" - assert "dbo_tblPayments" in result["invalid_generation_result"]["error"] From 50248cc21bc8aec682e4b155fb1643894479b57f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 18:19:08 +0530 Subject: [PATCH 0681/1087] Restore legacy schema retrieval evidence --- .../pipelines/indexing/table_description.py | 35 ++-------------- .../indexing/test_table_description.py | 16 +++++--- .../retrieval/test_db_schema_retrieval.py | 41 +++++++++---------- 3 files changed, 32 insertions(+), 60 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 9ecc863802..4d013fdfc1 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -18,36 +18,9 @@ logger = logging.getLogger("wren-ai-service") -MAX_TABLE_DESCRIPTION_COLUMNS = 200 -MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH = 4000 - @component class TableDescriptionChunker: - def _normalize_text(self, value: Any) -> str: - return "" if value is None else str(value) - - def _truncate_description(self, description: Any) -> str: - normalized_description = self._normalize_text(description) - if len(normalized_description) <= MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH: - return normalized_description - - return ( - normalized_description[:MAX_TABLE_DESCRIPTION_DESCRIPTION_LENGTH].rstrip() - + "..." - ) - - def _format_columns(self, columns: List[Any]) -> str: - normalized_columns = [self._normalize_text(column) for column in columns] - if len(normalized_columns) <= MAX_TABLE_DESCRIPTION_COLUMNS: - return ", ".join(normalized_columns) - - remaining_columns = len(normalized_columns) - MAX_TABLE_DESCRIPTION_COLUMNS - truncated_columns = normalized_columns[:MAX_TABLE_DESCRIPTION_COLUMNS] + [ - f"... (+{remaining_columns} more columns)" - ] - return ", ".join(truncated_columns) - def _properties(self, payload: Dict[str, Any]) -> Dict[str, Any]: properties = payload.get("properties") return properties if isinstance(properties, dict) else {} @@ -88,7 +61,7 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: "mdl_type": mdl_type, "name": payload.get("name"), "columns": [ - self._normalize_text(column.get("name", "")) + column.get("name", "") or "" for column in payload.get("columns", []) if isinstance(column, dict) ], @@ -104,10 +77,8 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: return [ { "name": resource["name"], - "description": self._truncate_description( - resource["properties"].get("description", "") - ), - "columns": self._format_columns(resource["columns"]), + "description": resource["properties"].get("description", "") or "", + "columns": ", ".join(resource["columns"]), } for resource in resources if resource["name"] is not None diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index b334e9001d..214ec27b82 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -194,16 +194,14 @@ def test_table_description_excludes_generated_column_descriptions(): assert "Generic generated" not in document.content -def test_table_description_truncates_long_column_lists(): +def test_table_description_keeps_complete_column_lists(): chunker = TableDescriptionChunker() + columns = [{"name": f"column_{index}"} for index in range(205)] mdl = { "models": [ { "name": "user", - "columns": [ - {"name": f"column_{index}"} - for index in range(205) - ], + "columns": columns, } ], "views": [], @@ -215,7 +213,13 @@ def test_table_description_truncates_long_column_lists(): assert len(actual["documents"]) == 1 document: Document = actual["documents"][0] - assert "... (+5 more columns)" in document.content + assert document.content == str( + { + "name": "user", + "description": "", + "columns": ", ".join(column["name"] for column in columns), + } + ) @pytest.mark.asyncio diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index fe93aa8779..1085f4136f 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -129,23 +129,21 @@ def encode(self, value): return value.split() result = check_using_db_schemas_without_pruning( - query="show top customers by invoice amount", - tables=None, construct_db_schemas=[ - { - "type": "TABLE", - "name": "orders", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "amount", - "data_type": "DOUBLE", - "comment": "", - "is_primary_key": False, - } - ], - "properties": {}, + { + "type": "TABLE", + "name": "orders", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, "primaryKey": "", } ], @@ -159,7 +157,7 @@ def encode(self, value): assert result["tokens"] > 0 -def test_check_using_db_schemas_without_pruning_selects_tables_for_question(): +def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): return value.split() @@ -183,8 +181,6 @@ def table_schema(name): } result = check_using_db_schemas_without_pruning( - query="compare recent activity by account", - tables=None, construct_db_schemas=[ table_schema("activity"), table_schema("account"), @@ -195,7 +191,10 @@ def table_schema(name): context_window_size=1000, ) - assert result["db_schemas"] == [] + assert [schema["table_name"] for schema in result["db_schemas"]] == [ + "activity", + "account", + ] assert result["tokens"] > 0 @@ -205,8 +204,6 @@ def encode(self, value): return value.split() result = check_using_db_schemas_without_pruning( - query="show records from activity", - tables=["activity"], construct_db_schemas=[ { "type": "TABLE", From 2a6d7762be8b5757b624a3148aad9fd1883140db Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 19:01:36 +0530 Subject: [PATCH 0682/1087] Replace stale SQL pair index on semantics prep --- .../src/pipelines/indexing/sql_pairs.py | 29 +++++++++++++------ .../web/v1/services/semantics_preparation.py | 2 +- .../pipelines/indexing/test_sql_pairs.py | 29 +++++++++++++++++++ 3 files changed, 50 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index eff8313ef7..9c785d3891 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -56,20 +56,25 @@ def __init__(self, sql_pairs_store: DocumentStore) -> None: @component.output_types() async def run( - self, sql_pair_ids: List[str], project_id: Optional[str] = None + self, + sql_pair_ids: List[str], + project_id: Optional[str] = None, + delete_all: bool = False, ) -> None: - filter = { - "operator": "AND", - "conditions": [ - {"field": "sql_pair_id", "operator": "in", "value": sql_pair_ids}, - ], - } + conditions = [] + + if not delete_all: + conditions.append( + {"field": "sql_pair_id", "operator": "in", "value": sql_pair_ids} + ) if project_id: - filter["conditions"].append( + conditions.append( {"field": "project_id", "operator": "==", "value": project_id} ) + filter = {"operator": "AND", "conditions": conditions} if conditions else None + return await self.store.delete_documents(filter) @@ -131,7 +136,11 @@ async def clean( ) -> Dict[str, Any]: sql_pair_ids = [sql_pair.id for sql_pair in sql_pairs] if sql_pair_ids or delete_all: - await cleaner.run(sql_pair_ids=sql_pair_ids, project_id=project_id) + await cleaner.run( + sql_pair_ids=sql_pair_ids, + project_id=project_id, + delete_all=delete_all, + ) return embedding @@ -195,6 +204,7 @@ async def run( mdl_str: str, project_id: str = "", external_pairs: Optional[Dict[str, Any]] = None, + delete_all: bool = False, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} SQL Pairs Indexing pipeline is running..." @@ -207,6 +217,7 @@ async def run( **self._external_pairs, **(external_pairs or {}), }, + "delete_all": delete_all, **self._components, } 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 2ff6215cbe..80456f5407 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -84,10 +84,10 @@ async def prepare_semantics( "db_schema", "historical_question", "table_description", - "sql_pairs", "project_meta", ] ] + tasks.append(self._pipelines["sql_pairs"].run(**input, delete_all=True)) await asyncio.gather(*tasks) diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py index 3436a64c59..e4552fc7d5 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py @@ -3,9 +3,35 @@ from src.config import settings from src.core.provider import DocumentStoreProvider from src.pipelines.indexing import SqlPairs +from src.pipelines.indexing.sql_pairs import SqlPairsCleaner from src.providers import generate_components +class _RecordingStore: + def __init__(self): + self.filters = [] + + async def delete_documents(self, filter): + self.filters.append(filter) + + +@pytest.mark.asyncio +async def test_sql_pairs_cleaner_delete_all_uses_project_scope(): + store = _RecordingStore() + cleaner = SqlPairsCleaner(store) + + await cleaner.run(sql_pair_ids=[], project_id="project-id", delete_all=True) + + assert store.filters == [ + { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-id"}, + ], + } + ] + + @pytest.mark.asyncio async def test_sql_pairs_indexing_saving_to_document_store(): pipe_components = generate_components(settings.components) @@ -94,3 +120,6 @@ async def test_sql_pairs_deletion(): await pipe.clean(sql_pairs=[], project_id="fake-id") assert await store.count_documents() == 2 + + await pipe.clean(sql_pairs=[], project_id="fake-id", delete_all=True) + assert await store.count_documents() == 0 From e576248fb52c72d46f533b5b9c5502ba67485784 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 19:06:15 +0530 Subject: [PATCH 0683/1087] Fix setup form project name assignment --- wren-ui/src/components/pages/setup/ConnectDataSource.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx index 5810fd3967..7908318a55 100644 --- a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx +++ b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx @@ -33,7 +33,7 @@ export default function ConnectDataSource(props: Props) { useEffect(() => { if (typeof router.query.projectName === 'string') { - form.setFieldValue('displayName', router.query.projectName); + form.setFieldsValue({ displayName: router.query.projectName }); } }, [form, router.query.projectName]); From 0a8a7e07f56f002d0ddc8def563b6ded957ba4b7 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 19:29:53 +0530 Subject: [PATCH 0684/1087] Align ask indexing with project-scoped semantics --- .../src/pipelines/indexing/sql_pairs.py | 16 ++++++++++++---- .../src/web/v1/services/semantics_preparation.py | 8 +++++++- .../pytest/pipelines/indexing/test_sql_pairs.py | 12 +++++++++++- .../src/apollo/server/adaptors/wrenAIAdaptor.ts | 11 ++++++++--- wren-ui/src/pages/api/v1/ask.ts | 1 + wren-ui/src/pages/api/v1/generate_sql.ts | 1 + wren-ui/src/pages/api/v1/stream/ask.ts | 1 + wren-ui/src/pages/api/v1/stream/generate_sql.ts | 1 + 8 files changed, 42 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index 9c785d3891..8027c78f36 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -96,7 +96,11 @@ def boilerplates( def sql_pairs( boilerplates: Set[str], external_pairs: Dict[str, Any], + include_default_pairs: bool = True, ) -> List[SqlPair]: + if not include_default_pairs and not external_pairs: + return [] + return [ SqlPair( id=pair.get("id"), @@ -205,19 +209,23 @@ async def run( project_id: str = "", external_pairs: Optional[Dict[str, Any]] = None, delete_all: bool = False, + include_default_pairs: bool = True, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} SQL Pairs Indexing pipeline is running..." ) + pairs = { + **(self._external_pairs if include_default_pairs else {}), + **(external_pairs or {}), + } + input = { "mdl_str": mdl_str, "project_id": project_id, - "external_pairs": { - **self._external_pairs, - **(external_pairs or {}), - }, + "external_pairs": pairs, "delete_all": delete_all, + "include_default_pairs": include_default_pairs, **self._components, } 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 80456f5407..f8dc4c1224 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -87,7 +87,13 @@ async def prepare_semantics( "project_meta", ] ] - tasks.append(self._pipelines["sql_pairs"].run(**input, delete_all=True)) + tasks.append( + self._pipelines["sql_pairs"].run( + **input, + delete_all=True, + include_default_pairs=False, + ) + ) await asyncio.gather(*tasks) diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py index e4552fc7d5..b1b0ba33c8 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py @@ -3,7 +3,7 @@ from src.config import settings from src.core.provider import DocumentStoreProvider from src.pipelines.indexing import SqlPairs -from src.pipelines.indexing.sql_pairs import SqlPairsCleaner +from src.pipelines.indexing.sql_pairs import SqlPairsCleaner, sql_pairs from src.providers import generate_components @@ -32,6 +32,16 @@ async def test_sql_pairs_cleaner_delete_all_uses_project_scope(): ] +def test_sql_pairs_can_skip_default_pairs(): + pairs = sql_pairs( + boilerplates={"default"}, + external_pairs={}, + include_default_pairs=False, + ) + + assert pairs == [] + + @pytest.mark.asyncio async def test_sql_pairs_indexing_saving_to_document_store(): pipe_components = generate_components(settings.components) diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 5acc9af772..3598f30cdf 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -247,13 +247,18 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async ask(input: AskInput): Promise { try { - const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, { + const body: Record = { query: input.query, id: input.deployId, - project_id: input.projectId, histories: this.transformHistoryInput(input.histories), configurations: input.configurations, - }); + }; + + if (input.projectId) { + body['project_id'] = String(input.projectId); + } + + const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, body); return { queryId: res.data.query_id }; } catch (err: any) { logger.debug(`Got error when asking wren AI: ${getAIServiceError(err)}`); diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index d2e98d0eb2..c7b966fde0 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -84,6 +84,7 @@ export default async function handler( const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: diff --git a/wren-ui/src/pages/api/v1/generate_sql.ts b/wren-ui/src/pages/api/v1/generate_sql.ts index fa5b3859e1..a6cef6f3bd 100644 --- a/wren-ui/src/pages/api/v1/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/generate_sql.ts @@ -79,6 +79,7 @@ export default async function handler( const task = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: diff --git a/wren-ui/src/pages/api/v1/stream/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts index 12e53eb4db..93a5a6fbfc 100644 --- a/wren-ui/src/pages/api/v1/stream/ask.ts +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -151,6 +151,7 @@ export default async function handler( const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: diff --git a/wren-ui/src/pages/api/v1/stream/generate_sql.ts b/wren-ui/src/pages/api/v1/stream/generate_sql.ts index 6102592657..9d65fc5618 100644 --- a/wren-ui/src/pages/api/v1/stream/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/stream/generate_sql.ts @@ -91,6 +91,7 @@ export default async function handler( const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: From c73cc698e92c1de35379a6cdaa61489970d3fcff Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 19:52:16 +0530 Subject: [PATCH 0685/1087] Handle empty SQL pair indexing during deploy --- .../src/pipelines/indexing/sql_pairs.py | 6 ++++ .../pipelines/indexing/test_sql_pairs.py | 28 ++++++++++++++++++- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index 8027c78f36..16ff245cd7 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -127,6 +127,9 @@ async def embedding( to_documents: Dict[str, Any], embedder: Any, ) -> Dict[str, Any]: + if not to_documents["documents"]: + return to_documents + return await embedder.run(documents=to_documents["documents"]) @@ -154,6 +157,9 @@ async def write( clean: Dict[str, Any], writer: AsyncDocumentWriter, ) -> None: + if not clean["documents"]: + return None + return await writer.run(documents=clean["documents"]) diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py index b1b0ba33c8..0ca5598ad9 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py @@ -3,7 +3,7 @@ from src.config import settings from src.core.provider import DocumentStoreProvider from src.pipelines.indexing import SqlPairs -from src.pipelines.indexing.sql_pairs import SqlPairsCleaner, sql_pairs +from src.pipelines.indexing.sql_pairs import SqlPairsCleaner, embedding, sql_pairs, write from src.providers import generate_components @@ -42,6 +42,32 @@ def test_sql_pairs_can_skip_default_pairs(): assert pairs == [] +@pytest.mark.asyncio +async def test_empty_sql_pairs_skip_embedding_and_write(): + class Embedder: + called = False + + async def run(self, documents): + self.called = True + return {"documents": documents} + + class Writer: + called = False + + async def run(self, documents): + self.called = True + + embedder = Embedder() + writer = Writer() + + result = await embedding({"documents": []}, embedder) + await write(result, writer) + + assert result == {"documents": []} + assert embedder.called is False + assert writer.called is False + + @pytest.mark.asyncio async def test_sql_pairs_indexing_saving_to_document_store(): pipe_components = generate_components(settings.components) From afaefffdc972bb2db15011ab910741b7106050a8 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 15:01:54 +0000 Subject: [PATCH 0686/1087] Ground ask generation prompts on retrieved schema --- .../pipelines/generation/data_assistance.py | 2 + .../generation/intent_classification.py | 14 ++++--- .../src/pipelines/generation/utils/sql.py | 8 +++- .../test_prompt_grounding_contracts.py | 38 +++++++++++++++++++ 4 files changed, 55 insertions(+), 7 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index 51b91197f9..7116b9d59e 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -29,6 +29,8 @@ - There should be proper line breaks, whitespace, and Markdown formatting(headers, lists, tables, etc.) in your response. - If the language is Traditional/Simplified Chinese, Korean, or Japanese, the maximum response length is 150 words; otherwise, the maximum response length is 110 words. - MUST NOT add SQL code in your response. +- Use only the provided DATABASE SCHEMA as context. Do not invent, assume, or name tables or columns that are not present in the schema. +- If the provided schema is insufficient to answer, say that the available metadata is insufficient and do not provide hypothetical schema, example table names, or example column names. - If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. ### OUTPUT FORMAT ### diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..35624513b0 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -43,14 +43,15 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes references to specific tables, columns, or data details. -- The question includes **complete information** with specific tables, columns, or data values needed for execution. -- The question provides **all necessary parameters** to generate executable SQL. +- The question (or related previous query) includes schema-resolvable references to tables, columns, or data details. +- The question asks for a data result, aggregation, ranking, listing, filtering, trend, comparison, or chart that can be answered from the provided database schema, even if the user did not type exact table or column names. +- The question includes **complete information** with schema-resolvable concepts, filters, or data values needed for execution. +- The question provides **all necessary parameters** to generate executable SQL using the provided schema. **Requirements:** - Must have complete filter criteria, specific values, or clear references to previous context. -- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. -- Reference phrases from the user's inputs that clearly relate to the schema. +- Use schema context to identify relevant tables and columns; do not require the user to write exact schema identifiers when the intent is a normal data question. +- Reference phrases from the user's inputs that clearly relate to the schema or to analytical operations that can be performed on the schema. **Examples:** - "What is the total sales for last quarter?" @@ -93,11 +94,12 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. +- The user's inputs lacks enough business meaning, values, or prior context to identify a data task from the provided database schema. - It appears off-topic or is simply a casual conversation starter. **Requirements:** - Incorporate phrases from the user's inputs that indicate lack of relevance to the database schema. +- Do not classify a data retrieval or analytics question as MISLEADING only because the user did not write exact table or column names. **Examples:** - "How are you?" diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index cbc07df03a..62f26cc4d9 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -199,10 +199,15 @@ async def _classify_generation_result( - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Schema comments are metadata for understanding the data. They are not SQL syntax. +- The DATABASE SCHEMA section is the complete and only source of executable table and column identifiers. +- MUST NOT introduce, infer, copy, or repair any table or column identifier unless the exact identifier appears in the DATABASE SCHEMA. +- Do not copy identifiers from the user question, prompt examples, SQL samples, reasoning plan, previous SQL, or error messages unless the exact identifier appears in the DATABASE SCHEMA. +- Identifiers shown in prompt examples are illustrative only and are not available for generated SQL unless they also appear in the DATABASE SCHEMA. - In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. - Never use a `display_label`, alias, or description as an executable table or column identifier. - Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. - You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. +- If the DATABASE SCHEMA does not contain the table or column needed for the user's request, do not substitute a similar, generic, or commonly known identifier. - Only apply numeric aggregate functions such as SUM or AVG to numeric columns or measures from the DATABASE SCHEMA. If a column is not numeric in the schema, do not aggregate it directly unless the provided SQL FUNCTIONS and database dialect support the explicit cast you use. - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. @@ -423,7 +428,8 @@ async def _classify_generation_result( 14. Use only table and column names that appear as identifiers in the DATABASE SCHEMA when writing `table:` and `column:` references. 15. Schema comments, display labels, aliases, and descriptions are context only. Do not use them as executable table or column names in the reasoning plan. 16. Do not create table or column names from words in the user's question. If a requested concept is available only through schema metadata, refer to the corresponding schema identifier. -17. ONLY SHOWING the reasoning plan in bullet points. +17. If the DATABASE SCHEMA does not contain an identifier needed for the request, state that the schema context is insufficient instead of naming a substitute table or column. +18. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py new file mode 100644 index 0000000000..a630ef9289 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py @@ -0,0 +1,38 @@ +from pathlib import Path + + +SERVICE_ROOT = Path(__file__).resolve().parents[4] + + +def _read_source(relative_path: str) -> str: + return (SERVICE_ROOT / relative_path).read_text(encoding="utf-8") + + +def test_intent_classification_does_not_require_exact_user_schema_names(): + source = _read_source("src/pipelines/generation/intent_classification.py") + + assert "schema-resolvable references" in source + assert "even if the user did not type exact table or column names" in source + assert "do not require the user to write exact schema identifiers" in source + assert ( + "Do not classify a data retrieval or analytics question as MISLEADING only " + "because the user did not write exact table or column names" + ) in source + + +def test_data_assistance_does_not_invent_hypothetical_schema(): + source = _read_source("src/pipelines/generation/data_assistance.py") + + assert "MUST NOT add SQL code" in source + assert "Use only the provided DATABASE SCHEMA as context" in source + assert "Do not invent, assume, or name tables or columns" in source + assert "do not provide hypothetical schema" in source + + +def test_sql_reasoning_contract_rejects_substitute_identifiers(): + source = _read_source("src/pipelines/generation/utils/sql.py") + + assert "If the DATABASE SCHEMA does not contain an identifier needed" in source + assert "instead of naming a substitute table or column" in source + assert "prompt examples" in source + assert "Identifiers shown in prompt examples are illustrative only" in source From 387a8307d135c612a27753af7d55e364709f2644 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 21:38:32 +0530 Subject: [PATCH 0687/1087] Restore SQL pair grounding during semantics prep --- .../web/v1/services/semantics_preparation.py | 1 - .../services/test_semantics_preparation.py | 45 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 wren-ai-service/tests/pytest/services/test_semantics_preparation.py 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 f8dc4c1224..f2c0f515a3 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -91,7 +91,6 @@ async def prepare_semantics( self._pipelines["sql_pairs"].run( **input, delete_all=True, - include_default_pairs=False, ) ) diff --git a/wren-ai-service/tests/pytest/services/test_semantics_preparation.py b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py new file mode 100644 index 0000000000..3c53046d9d --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py @@ -0,0 +1,45 @@ +import pytest + +from src.web.v1.services.semantics_preparation import ( + SemanticsPreparationRequest, + SemanticsPreparationService, +) + + +class _RecordingPipeline: + def __init__(self): + self.run_calls = [] + + async def run(self, **kwargs): + self.run_calls.append(kwargs) + + +@pytest.mark.asyncio +async def test_prepare_semantics_reindexes_default_sql_pairs_after_cleanup(): + pipelines = { + name: _RecordingPipeline() + for name in [ + "db_schema", + "historical_question", + "table_description", + "project_meta", + "sql_pairs", + ] + } + service = SemanticsPreparationService(pipelines) + + await service.prepare_semantics( + SemanticsPreparationRequest( + mdl='{"models": []}', + mdl_hash="mdl-hash", + project_id="project-id", + ) + ) + + assert pipelines["sql_pairs"].run_calls == [ + { + "mdl_str": '{"models": []}', + "project_id": "project-id", + "delete_all": True, + } + ] From e48659bbf21fdd884d4d5988201078b3fc2521bb Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 16:23:54 +0000 Subject: [PATCH 0688/1087] Preserve engine errors during SQL correction --- .../pipelines/generation/sql_correction.py | 1 + .../src/pipelines/generation/utils/sql.py | 2 ++ wren-ai-service/src/web/v1/services/ask.py | 24 ++++++++++++++++--- .../test_prompt_grounding_contracts.py | 18 ++++++++++++++ 4 files changed, 42 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..3e7e06fc81 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -36,6 +36,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). 2. Then, generate the syntactically correct ANSI SQL query to correct the error. +3. If the error reports an unknown table or field, replace it only with an exact executable identifier from the DATABASE SCHEMA. Do not retry the same unknown identifier. ### SQL RULES ### Make sure you follow the SQL Rules strictly. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 62f26cc4d9..58fe8013c7 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -205,6 +205,8 @@ async def _classify_generation_result( - Identifiers shown in prompt examples are illustrative only and are not available for generated SQL unless they also appear in the DATABASE SCHEMA. - In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. - Never use a `display_label`, alias, or description as an executable table or column identifier. +- Use `display_label` and `description` only to understand which executable `identifier` matches the user's business term. +- When a schema comment contains an `identifier`, generated SQL must use that exact identifier for the table or column. - Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. - You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. - If the DATABASE SCHEMA does not contain the table or column needed for the user's request, do not substitute a similar, generic, or commonly known identifier. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7a20e792c0..b1ea202d94 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -130,6 +130,21 @@ def _is_stopped(self, query_id: str, container: dict): return False + @staticmethod + def _build_sql_correction_error( + diagnosis_reasoning: Optional[str], engine_error: Optional[str] + ) -> str: + diagnosis_reasoning = (diagnosis_reasoning or "").strip() + engine_error = (engine_error or "").strip() + + if diagnosis_reasoning and engine_error: + return ( + f"{diagnosis_reasoning}\n\n" + f"Original Wren Engine validation error:\n{engine_error}" + ) + + return diagnosis_reasoning or engine_error + @observe(name="Ask Question") @trace_metadata async def ask( @@ -543,6 +558,8 @@ async def ask( sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") + else: + sql_diagnosis_reasoning = "" sql_correction_results = await self._pipelines[ "sql_correction" @@ -551,9 +568,10 @@ async def ask( instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "error": self._build_sql_correction_error( + sql_diagnosis_reasoning, + error_message, + ), }, project_id=ask_request.project_id, use_dry_plan=use_dry_plan, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py index a630ef9289..bd11a6470d 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py @@ -36,3 +36,21 @@ def test_sql_reasoning_contract_rejects_substitute_identifiers(): assert "instead of naming a substitute table or column" in source assert "prompt examples" in source assert "Identifiers shown in prompt examples are illustrative only" in source + assert "Use `display_label` and `description` only to understand" in source + assert "generated SQL must use that exact identifier" in source + + +def test_sql_correction_receives_raw_wren_engine_validation_error(): + source = _read_source("src/web/v1/services/ask.py") + + assert "_build_sql_correction_error" in source + assert "Original Wren Engine validation error" in source + assert "error_message" in source + + +def test_sql_correction_unknown_identifier_contract(): + source = _read_source("src/pipelines/generation/sql_correction.py") + + assert "If the error reports an unknown table or field" in source + assert "replace it only with an exact executable identifier" in source + assert "Do not retry the same unknown identifier" in source From 77a4bc504bd7d96f9fb1eae1a384061396d31ccb Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 16:50:38 +0000 Subject: [PATCH 0689/1087] Guard SQL generation against schema hallucinations --- .../generation/followup_sql_generation.py | 2 + .../pipelines/generation/sql_correction.py | 2 + .../pipelines/generation/sql_generation.py | 2 + .../src/pipelines/generation/utils/sql.py | 316 ++++++++++++++++++ .../pipelines/indexing/table_description.py | 33 +- .../test_prompt_grounding_contracts.py | 30 ++ 6 files changed, 381 insertions(+), 4 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 35cfb8fccf..4a4b897a93 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,6 +147,7 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -157,6 +158,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + schema_contexts=documents, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 3e7e06fc81..b6a50fe3c5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -121,6 +121,7 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -131,6 +132,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + schema_contexts=documents, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..04837adffb 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -139,6 +139,7 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -151,6 +152,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + schema_contexts=documents, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 58fe8013c7..dbda2493a4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, List import aiohttp @@ -17,6 +18,303 @@ logger = logging.getLogger("wren-ai-service") +_IDENTIFIER_ATOM_PATTERN = ( + r'"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*' +) +_QUALIFIED_IDENTIFIER_PATTERN = ( + rf"(?:{_IDENTIFIER_ATOM_PATTERN})(?:\s*\.\s*(?:{_IDENTIFIER_ATOM_PATTERN}))*" +) +_CREATE_TABLE_RE = re.compile( + rf"\bCREATE\s+TABLE\s+(?P
{_QUALIFIED_IDENTIFIER_PATTERN})\s*\((?P.*?)\n\);", + re.IGNORECASE | re.DOTALL, +) +_CREATE_VIEW_RE = re.compile( + rf"\bCREATE\s+VIEW\s+(?P
{_QUALIFIED_IDENTIFIER_PATTERN})\b", + re.IGNORECASE, +) +_TABLE_REFERENCE_RE = re.compile( + rf"\b(?:FROM|JOIN)\s+(?P
{_QUALIFIED_IDENTIFIER_PATTERN})(?:\s+(?:AS\s+)?(?P[A-Za-z_][A-Za-z0-9_$]*))?", + re.IGNORECASE, +) +_CTE_RE = re.compile( + rf"(?:\bWITH|,)\s+(?P{_IDENTIFIER_ATOM_PATTERN})\s+AS\s*\(", + re.IGNORECASE, +) +_QUALIFIED_COLUMN_RE = re.compile( + rf"(?P{_QUALIFIED_IDENTIFIER_PATTERN})\s*\.\s*(?P{_IDENTIFIER_ATOM_PATTERN})", + re.IGNORECASE, +) +_STRING_LITERAL_RE = re.compile(r"'(?:''|[^'])*'") +_COMMENT_RE = re.compile(r"--.*?$|/\*.*?\*/", re.MULTILINE | re.DOTALL) + +_SQL_STOP_WORDS = { + "AND", + "AS", + "ASC", + "BETWEEN", + "BY", + "CASE", + "CAST", + "CURRENT_DATE", + "CURRENT_TIMESTAMP", + "DATE", + "DATE_TRUNC", + "DAY", + "DESC", + "DISTINCT", + "ELSE", + "END", + "EXISTS", + "FALSE", + "FROM", + "FULL", + "GROUP", + "HAVING", + "IN", + "INNER", + "INTERVAL", + "IS", + "JOIN", + "LEFT", + "LIKE", + "LIMIT", + "LOWER", + "MONTH", + "NOT", + "NULL", + "ON", + "OR", + "ORDER", + "OUTER", + "PARTITION", + "RIGHT", + "SELECT", + "THEN", + "TRUE", + "UNION", + "WHEN", + "WHERE", + "WITH", + "YEAR", +} +_NON_COLUMN_STARTS = { + "CONSTRAINT", + "FOREIGN", + "PRIMARY", + "UNIQUE", +} + + +def _strip_identifier_quotes(identifier: str) -> str: + identifier = identifier.strip() + if ( + (identifier.startswith('"') and identifier.endswith('"')) + or (identifier.startswith("`") and identifier.endswith("`")) + or (identifier.startswith("[") and identifier.endswith("]")) + ): + return identifier[1:-1] + + return identifier + + +def _identifier_parts(identifier: str) -> list[str]: + return [ + _strip_identifier_quotes(match.group(0)) + for match in re.finditer(_IDENTIFIER_ATOM_PATTERN, identifier) + ] + + +def _normalize_identifier(identifier: str) -> str: + return ".".join(_identifier_parts(identifier)).lower() + + +def _schema_catalog_from_contexts( + schema_contexts: list[str] | None, +) -> dict[str, set[str]]: + catalog: dict[str, set[str]] = {} + + for context in schema_contexts or []: + for match in _CREATE_TABLE_RE.finditer(context): + table_name = ".".join(_identifier_parts(match.group("table"))) + columns: set[str] = set() + for line in match.group("body").splitlines(): + stripped = line.strip().rstrip(",") + if not stripped or stripped.startswith(("--", "/*", "*")): + continue + + token_match = re.match(_IDENTIFIER_ATOM_PATTERN, stripped) + if not token_match: + continue + + column_name = _strip_identifier_quotes(token_match.group(0)) + if column_name.upper() in _NON_COLUMN_STARTS: + continue + + columns.add(column_name) + + catalog[table_name] = columns + + for match in _CREATE_VIEW_RE.finditer(context): + table_name = ".".join(_identifier_parts(match.group("table"))) + catalog.setdefault(table_name, set()) + + return catalog + + +def _cte_names(sql: str) -> set[str]: + return {_normalize_identifier(match.group("name")) for match in _CTE_RE.finditer(sql)} + + +def _referenced_tables( + sql: str, catalog_by_name: dict[str, str], ctes: set[str] +) -> tuple[set[str], dict[str, str | None], list[str]]: + referenced: set[str] = set() + qualifiers: dict[str, str | None] = {} + unknown_tables: list[str] = [] + + for match in _TABLE_REFERENCE_RE.finditer(sql): + table_ref = match.group("table") + normalized_ref = _normalize_identifier(table_ref) + table_parts = _identifier_parts(table_ref) + normalized_last_part = table_parts[-1].lower() if table_parts else "" + + if normalized_ref in ctes: + qualifiers[normalized_ref] = None + continue + + if normalized_ref in catalog_by_name: + table_name = catalog_by_name[normalized_ref] + elif len(table_parts) == 1 and normalized_last_part in catalog_by_name: + table_name = catalog_by_name[normalized_last_part] + else: + unknown_tables.append(table_ref) + continue + + referenced.add(table_name) + qualifiers[table_name.lower()] = table_name + qualifiers[table_parts[-1].lower()] = table_name + + alias = match.group("alias") + if alias and alias.upper() not in _SQL_STOP_WORDS: + qualifiers[alias.lower()] = table_name + + return referenced, qualifiers, unknown_tables + + +def _clause_texts(sql: str) -> list[str]: + cleaned_sql = _COMMENT_RE.sub(" ", _STRING_LITERAL_RE.sub(" ", sql)) + clause_pattern = re.compile( + r"\b(?:WHERE|ON|GROUP\s+BY|ORDER\s+BY|HAVING|PARTITION\s+BY)\b(?P.*?)(?=\b(?:WHERE|GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|UNION|JOIN|LEFT\s+JOIN|RIGHT\s+JOIN|INNER\s+JOIN|FULL\s+JOIN)\b|$)", + re.IGNORECASE | re.DOTALL, + ) + return [match.group("body") for match in clause_pattern.finditer(cleaned_sql)] + + +def _select_aliases(sql: str) -> set[str]: + cleaned_sql = _COMMENT_RE.sub(" ", _STRING_LITERAL_RE.sub(" ", sql)) + return { + _normalize_identifier(match.group("alias")) + for match in re.finditer( + rf"\bAS\s+(?P{_IDENTIFIER_ATOM_PATTERN})\b", + cleaned_sql, + re.IGNORECASE, + ) + } + + +def _unqualified_clause_identifiers(sql: str) -> set[str]: + identifiers: set[str] = set() + aliases = _select_aliases(sql) + + for clause in _clause_texts(sql): + for match in re.finditer(_IDENTIFIER_ATOM_PATTERN, clause): + name = _strip_identifier_quotes(match.group(0)) + upper_name = name.upper() + if upper_name in _SQL_STOP_WORDS or name.lower() in aliases: + continue + + previous_char = clause[match.start() - 1] if match.start() > 0 else "" + next_char = clause[match.end()] if match.end() < len(clause) else "" + if previous_char == "." or next_char in ".(": + continue + + identifiers.add(name) + + return identifiers + + +def _format_schema_validation_error( + unknown_tables: list[str], unknown_columns: list[str] +) -> str: + parts = [ + "Generated SQL references identifiers that are not present in the retrieved DATABASE SCHEMA.", + ] + + if unknown_tables: + parts.append(f"Unknown table identifiers: {', '.join(sorted(set(unknown_tables)))}.") + + if unknown_columns: + parts.append(f"Unknown column identifiers: {', '.join(sorted(set(unknown_columns)))}.") + + parts.append( + "Use only exact table and column identifiers from the retrieved CREATE TABLE/CREATE VIEW statements; do not use physical schema prefixes, display labels, examples, or names from the user question unless they appear in the schema." + ) + return " ".join(parts) + + +def _validate_sql_against_schema_contexts( + sql: str, schema_contexts: list[str] | None +) -> str | None: + catalog = _schema_catalog_from_contexts(schema_contexts) + if not catalog: + return None + + catalog_by_name = {table.lower(): table for table in catalog} + catalog_by_name.update({table.split(".")[-1].lower(): table for table in catalog}) + + referenced_tables, qualifiers, unknown_tables = _referenced_tables( + sql, catalog_by_name, _cte_names(sql) + ) + + if unknown_tables: + return _format_schema_validation_error(unknown_tables, []) + + if not referenced_tables: + return None + + columns_by_qualifier = { + qualifier: {column.lower() for column in catalog[table_name]} + for qualifier, table_name in qualifiers.items() + if table_name in catalog + } + columns_in_referenced_tables = { + column.lower() + for table_name in referenced_tables + for column in catalog[table_name] + } + + unknown_columns: list[str] = [] + for match in _QUALIFIED_COLUMN_RE.finditer(sql): + qualifier = _normalize_identifier(match.group("qualifier")) + column = _strip_identifier_quotes(match.group("column")) + if qualifier not in columns_by_qualifier: + continue + + if column.lower() not in columns_by_qualifier[qualifier]: + unknown_columns.append(f"{match.group('qualifier')}.{column}") + + has_cte_reference = any(table_name is None for table_name in qualifiers.values()) + if not has_cte_reference: + for column in _unqualified_clause_identifiers(sql): + if column.lower() not in columns_in_referenced_tables: + unknown_columns.append(column) + + if unknown_columns: + return _format_schema_validation_error([], unknown_columns) + + return None + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -34,6 +332,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + schema_contexts: list[str] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -44,6 +343,22 @@ async def run( "sql" ] + schema_validation_error = _validate_sql_against_schema_contexts( + cleaned_generation_result, + schema_contexts, + ) + if schema_validation_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_VALIDATION", + "error": schema_validation_error, + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, @@ -98,6 +413,7 @@ async def _classify_generation_result( else: invalid_generation_result = { "sql": generation_result, + "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") else "DRY_PLAN", diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 4d013fdfc1..88b8954cf5 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -15,6 +15,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider from src.pipelines.indexing import AsyncDocumentWriter, DocumentCleaner, MDLValidator +from src.pipelines.indexing import clean_display_name logger = logging.getLogger("wren-ai-service") @@ -54,18 +55,34 @@ def _additional_meta() -> Dict[str, Any]: } def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[Dict[str, Any]]: + def _text(value: Any) -> str: + return "" if value is None else str(value).strip() + + def _column_summary(column: Dict[str, Any]) -> Dict[str, str]: + properties = self._properties(column) + return { + "identifier": _text(column.get("name", "")), + "display_label": clean_display_name( + _text(properties.get("displayName", "")) + ), + "description": _text(properties.get("description", "")), + } + def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: properties = self._properties(payload) return { "mdl_type": mdl_type, "name": payload.get("name"), + "display_label": clean_display_name( + _text(properties.get("displayName", "")) + ), + "description": _text(properties.get("description", "")), "columns": [ - column.get("name", "") or "" + _column_summary(column) for column in payload.get("columns", []) if isinstance(column, dict) ], - "properties": properties, } resources = ( @@ -77,8 +94,16 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: return [ { "name": resource["name"], - "description": resource["properties"].get("description", "") or "", - "columns": ", ".join(resource["columns"]), + "display_label": resource["display_label"], + "description": resource["description"], + "columns": ", ".join( + column["identifier"] for column in resource["columns"] + ), + "column_details": [ + column + for column in resource["columns"] + if column["display_label"] or column["description"] + ], } for resource in resources if resource["name"] is not None diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py index bd11a6470d..422473e51f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py @@ -54,3 +54,33 @@ def test_sql_correction_unknown_identifier_contract(): assert "If the error reports an unknown table or field" in source assert "replace it only with an exact executable identifier" in source assert "Do not retry the same unknown identifier" in source + + +def test_sql_post_processor_has_schema_identifier_guard(): + source = _read_source("src/pipelines/generation/utils/sql.py") + + assert "_validate_sql_against_schema_contexts" in source + assert "SCHEMA_VALIDATION" in source + assert "Unknown table identifiers" in source + assert "Unknown column identifiers" in source + assert "physical schema prefixes" in source + + +def test_generation_paths_pass_retrieved_schema_to_post_processor(): + for relative_path in [ + "src/pipelines/generation/sql_generation.py", + "src/pipelines/generation/followup_sql_generation.py", + "src/pipelines/generation/sql_correction.py", + ]: + source = _read_source(relative_path) + assert "documents: list[str] | None = None" in source + assert "schema_contexts=documents" in source + + +def test_table_description_retrieval_indexes_business_metadata(): + source = _read_source("src/pipelines/indexing/table_description.py") + + assert "display_label" in source + assert "column_details" in source + assert "clean_display_name" in source + assert "description" in source From bbfaae0dd8d3aba13606df2a5429b21cde4425ab Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 16:58:48 +0000 Subject: [PATCH 0690/1087] Fix schema guard Hamilton input type conflict --- .../src/pipelines/generation/followup_sql_generation.py | 5 +++-- .../src/pipelines/generation/sql_correction.py | 5 +++-- .../src/pipelines/generation/sql_generation.py | 5 +++-- wren-ai-service/src/pipelines/generation/utils/sql.py | 9 ++++++--- .../generation/test_prompt_grounding_contracts.py | 4 ++-- 5 files changed, 17 insertions(+), 11 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 4a4b897a93..a85c3fee5f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,7 +147,7 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str] | None = None, + schema_contexts: list[Any] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -158,7 +158,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_contexts=documents, + schema_contexts=schema_contexts, ) @@ -223,6 +223,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "schema_contexts": contexts, "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index b6a50fe3c5..f4cc4040a1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -121,7 +121,7 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str] | None = None, + schema_contexts: list[Any] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -132,7 +132,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_contexts=documents, + schema_contexts=schema_contexts, ) @@ -191,6 +191,7 @@ async def run( inputs={ "invalid_generation_result": invalid_generation_result, "documents": contexts, + "schema_contexts": contexts, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 04837adffb..a974fae5bd 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -139,7 +139,7 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str] | None = None, + schema_contexts: list[Any] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -152,7 +152,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - schema_contexts=documents, + schema_contexts=schema_contexts, ) @@ -217,6 +217,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "schema_contexts": contexts, "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index dbda2493a4..c71bffc448 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -129,11 +129,14 @@ def _normalize_identifier(identifier: str) -> str: def _schema_catalog_from_contexts( - schema_contexts: list[str] | None, + schema_contexts: list[Any] | None, ) -> dict[str, set[str]]: catalog: dict[str, set[str]] = {} for context in schema_contexts or []: + context = getattr(context, "content", context) + context = "" if context is None else str(context) + for match in _CREATE_TABLE_RE.finditer(context): table_name = ".".join(_identifier_parts(match.group("table"))) columns: set[str] = set() @@ -263,7 +266,7 @@ def _format_schema_validation_error( def _validate_sql_against_schema_contexts( - sql: str, schema_contexts: list[str] | None + sql: str, schema_contexts: list[Any] | None ) -> str | None: catalog = _schema_catalog_from_contexts(schema_contexts) if not catalog: @@ -332,7 +335,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - schema_contexts: list[str] | None = None, + schema_contexts: list[Any] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py index 422473e51f..4995ae7b64 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py @@ -73,8 +73,8 @@ def test_generation_paths_pass_retrieved_schema_to_post_processor(): "src/pipelines/generation/sql_correction.py", ]: source = _read_source(relative_path) - assert "documents: list[str] | None = None" in source - assert "schema_contexts=documents" in source + assert "schema_contexts: list[Any] | None = None" in source + assert '"schema_contexts": contexts' in source def test_table_description_retrieval_indexes_business_metadata(): From 57f03f3ef5a0aecd675eddf581c4fb46af0d38dd Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 22:54:10 +0530 Subject: [PATCH 0691/1087] Remove SQL pair default switch --- .../src/pipelines/indexing/sql_pairs.py | 8 +-- .../pipelines/indexing/test_sql_pairs.py | 65 ------------------- 2 files changed, 1 insertion(+), 72 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index 16ff245cd7..98cff3b726 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -96,11 +96,7 @@ def boilerplates( def sql_pairs( boilerplates: Set[str], external_pairs: Dict[str, Any], - include_default_pairs: bool = True, ) -> List[SqlPair]: - if not include_default_pairs and not external_pairs: - return [] - return [ SqlPair( id=pair.get("id"), @@ -215,14 +211,13 @@ async def run( project_id: str = "", external_pairs: Optional[Dict[str, Any]] = None, delete_all: bool = False, - include_default_pairs: bool = True, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} SQL Pairs Indexing pipeline is running..." ) pairs = { - **(self._external_pairs if include_default_pairs else {}), + **self._external_pairs, **(external_pairs or {}), } @@ -231,7 +226,6 @@ async def run( "project_id": project_id, "external_pairs": pairs, "delete_all": delete_all, - "include_default_pairs": include_default_pairs, **self._components, } diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py index 0ca5598ad9..3436a64c59 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py @@ -3,71 +3,9 @@ from src.config import settings from src.core.provider import DocumentStoreProvider from src.pipelines.indexing import SqlPairs -from src.pipelines.indexing.sql_pairs import SqlPairsCleaner, embedding, sql_pairs, write from src.providers import generate_components -class _RecordingStore: - def __init__(self): - self.filters = [] - - async def delete_documents(self, filter): - self.filters.append(filter) - - -@pytest.mark.asyncio -async def test_sql_pairs_cleaner_delete_all_uses_project_scope(): - store = _RecordingStore() - cleaner = SqlPairsCleaner(store) - - await cleaner.run(sql_pair_ids=[], project_id="project-id", delete_all=True) - - assert store.filters == [ - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": "project-id"}, - ], - } - ] - - -def test_sql_pairs_can_skip_default_pairs(): - pairs = sql_pairs( - boilerplates={"default"}, - external_pairs={}, - include_default_pairs=False, - ) - - assert pairs == [] - - -@pytest.mark.asyncio -async def test_empty_sql_pairs_skip_embedding_and_write(): - class Embedder: - called = False - - async def run(self, documents): - self.called = True - return {"documents": documents} - - class Writer: - called = False - - async def run(self, documents): - self.called = True - - embedder = Embedder() - writer = Writer() - - result = await embedding({"documents": []}, embedder) - await write(result, writer) - - assert result == {"documents": []} - assert embedder.called is False - assert writer.called is False - - @pytest.mark.asyncio async def test_sql_pairs_indexing_saving_to_document_store(): pipe_components = generate_components(settings.components) @@ -156,6 +94,3 @@ async def test_sql_pairs_deletion(): await pipe.clean(sql_pairs=[], project_id="fake-id") assert await store.count_documents() == 2 - - await pipe.clean(sql_pairs=[], project_id="fake-id", delete_all=True) - assert await store.count_documents() == 0 From e45fd0a1f786d0f290dac9c6befb01d9ad276064 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 22:58:00 +0530 Subject: [PATCH 0692/1087] Remove schema guard validation changes --- .../pipelines/generation/data_assistance.py | 2 - .../generation/followup_sql_generation.py | 3 - .../generation/intent_classification.py | 14 +- .../pipelines/generation/sql_correction.py | 3 - .../pipelines/generation/sql_generation.py | 3 - .../src/pipelines/generation/utils/sql.py | 327 +----------------- .../pipelines/indexing/table_description.py | 33 +- .../test_prompt_grounding_contracts.py | 86 ----- 8 files changed, 11 insertions(+), 460 deletions(-) delete mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index 7116b9d59e..51b91197f9 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -29,8 +29,6 @@ - There should be proper line breaks, whitespace, and Markdown formatting(headers, lists, tables, etc.) in your response. - If the language is Traditional/Simplified Chinese, Korean, or Japanese, the maximum response length is 150 words; otherwise, the maximum response length is 110 words. - MUST NOT add SQL code in your response. -- Use only the provided DATABASE SCHEMA as context. Do not invent, assume, or name tables or columns that are not present in the schema. -- If the provided schema is insufficient to answer, say that the available metadata is insufficient and do not provide hypothetical schema, example table names, or example column names. - If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. ### OUTPUT FORMAT ### 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 a85c3fee5f..35cfb8fccf 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,7 +147,6 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, - schema_contexts: list[Any] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -158,7 +157,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_contexts=schema_contexts, ) @@ -223,7 +221,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "schema_contexts": contexts, "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 35624513b0..4d6cd313cd 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -43,15 +43,14 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes schema-resolvable references to tables, columns, or data details. -- The question asks for a data result, aggregation, ranking, listing, filtering, trend, comparison, or chart that can be answered from the provided database schema, even if the user did not type exact table or column names. -- The question includes **complete information** with schema-resolvable concepts, filters, or data values needed for execution. -- The question provides **all necessary parameters** to generate executable SQL using the provided schema. +- The question (or related previous query) includes references to specific tables, columns, or data details. +- The question includes **complete information** with specific tables, columns, or data values needed for execution. +- The question provides **all necessary parameters** to generate executable SQL. **Requirements:** - Must have complete filter criteria, specific values, or clear references to previous context. -- Use schema context to identify relevant tables and columns; do not require the user to write exact schema identifiers when the intent is a normal data question. -- Reference phrases from the user's inputs that clearly relate to the schema or to analytical operations that can be performed on the schema. +- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. +- Reference phrases from the user's inputs that clearly relate to the schema. **Examples:** - "What is the total sales for last quarter?" @@ -94,12 +93,11 @@ **When to Use:** - The user's inputs is irrelevant to the database schema or includes SQL code. -- The user's inputs lacks enough business meaning, values, or prior context to identify a data task from the provided database schema. +- The user's inputs lacks specific details (like table names or columns) needed to generate an SQL query. - It appears off-topic or is simply a casual conversation starter. **Requirements:** - Incorporate phrases from the user's inputs that indicate lack of relevance to the database schema. -- Do not classify a data retrieval or analytics question as MISLEADING only because the user did not write exact table or column names. **Examples:** - "How are you?" diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f4cc4040a1..3e7e06fc81 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -121,7 +121,6 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, - schema_contexts: list[Any] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -132,7 +131,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_contexts=schema_contexts, ) @@ -191,7 +189,6 @@ async def run( inputs={ "invalid_generation_result": invalid_generation_result, "documents": contexts, - "schema_contexts": contexts, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index a974fae5bd..1ee4952b3e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -139,7 +139,6 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, - schema_contexts: list[Any] | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -152,7 +151,6 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - schema_contexts=schema_contexts, ) @@ -217,7 +215,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "schema_contexts": contexts, "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c71bffc448..7890524c1f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,4 @@ import logging -import re from typing import Any, Dict, List import aiohttp @@ -18,306 +17,6 @@ logger = logging.getLogger("wren-ai-service") -_IDENTIFIER_ATOM_PATTERN = ( - r'"[^"]+"|`[^`]+`|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*' -) -_QUALIFIED_IDENTIFIER_PATTERN = ( - rf"(?:{_IDENTIFIER_ATOM_PATTERN})(?:\s*\.\s*(?:{_IDENTIFIER_ATOM_PATTERN}))*" -) -_CREATE_TABLE_RE = re.compile( - rf"\bCREATE\s+TABLE\s+(?P
{_QUALIFIED_IDENTIFIER_PATTERN})\s*\((?P.*?)\n\);", - re.IGNORECASE | re.DOTALL, -) -_CREATE_VIEW_RE = re.compile( - rf"\bCREATE\s+VIEW\s+(?P
{_QUALIFIED_IDENTIFIER_PATTERN})\b", - re.IGNORECASE, -) -_TABLE_REFERENCE_RE = re.compile( - rf"\b(?:FROM|JOIN)\s+(?P
{_QUALIFIED_IDENTIFIER_PATTERN})(?:\s+(?:AS\s+)?(?P[A-Za-z_][A-Za-z0-9_$]*))?", - re.IGNORECASE, -) -_CTE_RE = re.compile( - rf"(?:\bWITH|,)\s+(?P{_IDENTIFIER_ATOM_PATTERN})\s+AS\s*\(", - re.IGNORECASE, -) -_QUALIFIED_COLUMN_RE = re.compile( - rf"(?P{_QUALIFIED_IDENTIFIER_PATTERN})\s*\.\s*(?P{_IDENTIFIER_ATOM_PATTERN})", - re.IGNORECASE, -) -_STRING_LITERAL_RE = re.compile(r"'(?:''|[^'])*'") -_COMMENT_RE = re.compile(r"--.*?$|/\*.*?\*/", re.MULTILINE | re.DOTALL) - -_SQL_STOP_WORDS = { - "AND", - "AS", - "ASC", - "BETWEEN", - "BY", - "CASE", - "CAST", - "CURRENT_DATE", - "CURRENT_TIMESTAMP", - "DATE", - "DATE_TRUNC", - "DAY", - "DESC", - "DISTINCT", - "ELSE", - "END", - "EXISTS", - "FALSE", - "FROM", - "FULL", - "GROUP", - "HAVING", - "IN", - "INNER", - "INTERVAL", - "IS", - "JOIN", - "LEFT", - "LIKE", - "LIMIT", - "LOWER", - "MONTH", - "NOT", - "NULL", - "ON", - "OR", - "ORDER", - "OUTER", - "PARTITION", - "RIGHT", - "SELECT", - "THEN", - "TRUE", - "UNION", - "WHEN", - "WHERE", - "WITH", - "YEAR", -} -_NON_COLUMN_STARTS = { - "CONSTRAINT", - "FOREIGN", - "PRIMARY", - "UNIQUE", -} - - -def _strip_identifier_quotes(identifier: str) -> str: - identifier = identifier.strip() - if ( - (identifier.startswith('"') and identifier.endswith('"')) - or (identifier.startswith("`") and identifier.endswith("`")) - or (identifier.startswith("[") and identifier.endswith("]")) - ): - return identifier[1:-1] - - return identifier - - -def _identifier_parts(identifier: str) -> list[str]: - return [ - _strip_identifier_quotes(match.group(0)) - for match in re.finditer(_IDENTIFIER_ATOM_PATTERN, identifier) - ] - - -def _normalize_identifier(identifier: str) -> str: - return ".".join(_identifier_parts(identifier)).lower() - - -def _schema_catalog_from_contexts( - schema_contexts: list[Any] | None, -) -> dict[str, set[str]]: - catalog: dict[str, set[str]] = {} - - for context in schema_contexts or []: - context = getattr(context, "content", context) - context = "" if context is None else str(context) - - for match in _CREATE_TABLE_RE.finditer(context): - table_name = ".".join(_identifier_parts(match.group("table"))) - columns: set[str] = set() - for line in match.group("body").splitlines(): - stripped = line.strip().rstrip(",") - if not stripped or stripped.startswith(("--", "/*", "*")): - continue - - token_match = re.match(_IDENTIFIER_ATOM_PATTERN, stripped) - if not token_match: - continue - - column_name = _strip_identifier_quotes(token_match.group(0)) - if column_name.upper() in _NON_COLUMN_STARTS: - continue - - columns.add(column_name) - - catalog[table_name] = columns - - for match in _CREATE_VIEW_RE.finditer(context): - table_name = ".".join(_identifier_parts(match.group("table"))) - catalog.setdefault(table_name, set()) - - return catalog - - -def _cte_names(sql: str) -> set[str]: - return {_normalize_identifier(match.group("name")) for match in _CTE_RE.finditer(sql)} - - -def _referenced_tables( - sql: str, catalog_by_name: dict[str, str], ctes: set[str] -) -> tuple[set[str], dict[str, str | None], list[str]]: - referenced: set[str] = set() - qualifiers: dict[str, str | None] = {} - unknown_tables: list[str] = [] - - for match in _TABLE_REFERENCE_RE.finditer(sql): - table_ref = match.group("table") - normalized_ref = _normalize_identifier(table_ref) - table_parts = _identifier_parts(table_ref) - normalized_last_part = table_parts[-1].lower() if table_parts else "" - - if normalized_ref in ctes: - qualifiers[normalized_ref] = None - continue - - if normalized_ref in catalog_by_name: - table_name = catalog_by_name[normalized_ref] - elif len(table_parts) == 1 and normalized_last_part in catalog_by_name: - table_name = catalog_by_name[normalized_last_part] - else: - unknown_tables.append(table_ref) - continue - - referenced.add(table_name) - qualifiers[table_name.lower()] = table_name - qualifiers[table_parts[-1].lower()] = table_name - - alias = match.group("alias") - if alias and alias.upper() not in _SQL_STOP_WORDS: - qualifiers[alias.lower()] = table_name - - return referenced, qualifiers, unknown_tables - - -def _clause_texts(sql: str) -> list[str]: - cleaned_sql = _COMMENT_RE.sub(" ", _STRING_LITERAL_RE.sub(" ", sql)) - clause_pattern = re.compile( - r"\b(?:WHERE|ON|GROUP\s+BY|ORDER\s+BY|HAVING|PARTITION\s+BY)\b(?P.*?)(?=\b(?:WHERE|GROUP\s+BY|ORDER\s+BY|HAVING|LIMIT|UNION|JOIN|LEFT\s+JOIN|RIGHT\s+JOIN|INNER\s+JOIN|FULL\s+JOIN)\b|$)", - re.IGNORECASE | re.DOTALL, - ) - return [match.group("body") for match in clause_pattern.finditer(cleaned_sql)] - - -def _select_aliases(sql: str) -> set[str]: - cleaned_sql = _COMMENT_RE.sub(" ", _STRING_LITERAL_RE.sub(" ", sql)) - return { - _normalize_identifier(match.group("alias")) - for match in re.finditer( - rf"\bAS\s+(?P{_IDENTIFIER_ATOM_PATTERN})\b", - cleaned_sql, - re.IGNORECASE, - ) - } - - -def _unqualified_clause_identifiers(sql: str) -> set[str]: - identifiers: set[str] = set() - aliases = _select_aliases(sql) - - for clause in _clause_texts(sql): - for match in re.finditer(_IDENTIFIER_ATOM_PATTERN, clause): - name = _strip_identifier_quotes(match.group(0)) - upper_name = name.upper() - if upper_name in _SQL_STOP_WORDS or name.lower() in aliases: - continue - - previous_char = clause[match.start() - 1] if match.start() > 0 else "" - next_char = clause[match.end()] if match.end() < len(clause) else "" - if previous_char == "." or next_char in ".(": - continue - - identifiers.add(name) - - return identifiers - - -def _format_schema_validation_error( - unknown_tables: list[str], unknown_columns: list[str] -) -> str: - parts = [ - "Generated SQL references identifiers that are not present in the retrieved DATABASE SCHEMA.", - ] - - if unknown_tables: - parts.append(f"Unknown table identifiers: {', '.join(sorted(set(unknown_tables)))}.") - - if unknown_columns: - parts.append(f"Unknown column identifiers: {', '.join(sorted(set(unknown_columns)))}.") - - parts.append( - "Use only exact table and column identifiers from the retrieved CREATE TABLE/CREATE VIEW statements; do not use physical schema prefixes, display labels, examples, or names from the user question unless they appear in the schema." - ) - return " ".join(parts) - - -def _validate_sql_against_schema_contexts( - sql: str, schema_contexts: list[Any] | None -) -> str | None: - catalog = _schema_catalog_from_contexts(schema_contexts) - if not catalog: - return None - - catalog_by_name = {table.lower(): table for table in catalog} - catalog_by_name.update({table.split(".")[-1].lower(): table for table in catalog}) - - referenced_tables, qualifiers, unknown_tables = _referenced_tables( - sql, catalog_by_name, _cte_names(sql) - ) - - if unknown_tables: - return _format_schema_validation_error(unknown_tables, []) - - if not referenced_tables: - return None - - columns_by_qualifier = { - qualifier: {column.lower() for column in catalog[table_name]} - for qualifier, table_name in qualifiers.items() - if table_name in catalog - } - columns_in_referenced_tables = { - column.lower() - for table_name in referenced_tables - for column in catalog[table_name] - } - - unknown_columns: list[str] = [] - for match in _QUALIFIED_COLUMN_RE.finditer(sql): - qualifier = _normalize_identifier(match.group("qualifier")) - column = _strip_identifier_quotes(match.group("column")) - if qualifier not in columns_by_qualifier: - continue - - if column.lower() not in columns_by_qualifier[qualifier]: - unknown_columns.append(f"{match.group('qualifier')}.{column}") - - has_cte_reference = any(table_name is None for table_name in qualifiers.values()) - if not has_cte_reference: - for column in _unqualified_clause_identifiers(sql): - if column.lower() not in columns_in_referenced_tables: - unknown_columns.append(column) - - if unknown_columns: - return _format_schema_validation_error([], unknown_columns) - - return None - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -335,7 +34,6 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - schema_contexts: list[Any] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -346,22 +44,6 @@ async def run( "sql" ] - schema_validation_error = _validate_sql_against_schema_contexts( - cleaned_generation_result, - schema_contexts, - ) - if schema_validation_error: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_VALIDATION", - "error": schema_validation_error, - "correlation_id": "", - }, - } - ( valid_generation_result, invalid_generation_result, @@ -416,7 +98,6 @@ async def _classify_generation_result( else: invalid_generation_result = { "sql": generation_result, - "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") else "DRY_PLAN", @@ -518,17 +199,12 @@ async def _classify_generation_result( - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Schema comments are metadata for understanding the data. They are not SQL syntax. -- The DATABASE SCHEMA section is the complete and only source of executable table and column identifiers. -- MUST NOT introduce, infer, copy, or repair any table or column identifier unless the exact identifier appears in the DATABASE SCHEMA. -- Do not copy identifiers from the user question, prompt examples, SQL samples, reasoning plan, previous SQL, or error messages unless the exact identifier appears in the DATABASE SCHEMA. -- Identifiers shown in prompt examples are illustrative only and are not available for generated SQL unless they also appear in the DATABASE SCHEMA. - In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. - Never use a `display_label`, alias, or description as an executable table or column identifier. - Use `display_label` and `description` only to understand which executable `identifier` matches the user's business term. - When a schema comment contains an `identifier`, generated SQL must use that exact identifier for the table or column. - Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. - You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. -- If the DATABASE SCHEMA does not contain the table or column needed for the user's request, do not substitute a similar, generic, or commonly known identifier. - Only apply numeric aggregate functions such as SUM or AVG to numeric columns or measures from the DATABASE SCHEMA. If a column is not numeric in the schema, do not aggregate it directly unless the provided SQL FUNCTIONS and database dialect support the explicit cast you use. - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. @@ -749,8 +425,7 @@ async def _classify_generation_result( 14. Use only table and column names that appear as identifiers in the DATABASE SCHEMA when writing `table:` and `column:` references. 15. Schema comments, display labels, aliases, and descriptions are context only. Do not use them as executable table or column names in the reasoning plan. 16. Do not create table or column names from words in the user's question. If a requested concept is available only through schema metadata, refer to the corresponding schema identifier. -17. If the DATABASE SCHEMA does not contain an identifier needed for the request, state that the schema context is insufficient instead of naming a substitute table or column. -18. ONLY SHOWING the reasoning plan in bullet points. +17. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 88b8954cf5..4d013fdfc1 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -15,7 +15,6 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider from src.pipelines.indexing import AsyncDocumentWriter, DocumentCleaner, MDLValidator -from src.pipelines.indexing import clean_display_name logger = logging.getLogger("wren-ai-service") @@ -55,34 +54,18 @@ def _additional_meta() -> Dict[str, Any]: } def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[Dict[str, Any]]: - def _text(value: Any) -> str: - return "" if value is None else str(value).strip() - - def _column_summary(column: Dict[str, Any]) -> Dict[str, str]: - properties = self._properties(column) - return { - "identifier": _text(column.get("name", "")), - "display_label": clean_display_name( - _text(properties.get("displayName", "")) - ), - "description": _text(properties.get("description", "")), - } - def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: properties = self._properties(payload) return { "mdl_type": mdl_type, "name": payload.get("name"), - "display_label": clean_display_name( - _text(properties.get("displayName", "")) - ), - "description": _text(properties.get("description", "")), "columns": [ - _column_summary(column) + column.get("name", "") or "" for column in payload.get("columns", []) if isinstance(column, dict) ], + "properties": properties, } resources = ( @@ -94,16 +77,8 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: return [ { "name": resource["name"], - "display_label": resource["display_label"], - "description": resource["description"], - "columns": ", ".join( - column["identifier"] for column in resource["columns"] - ), - "column_details": [ - column - for column in resource["columns"] - if column["display_label"] or column["description"] - ], + "description": resource["properties"].get("description", "") or "", + "columns": ", ".join(resource["columns"]), } for resource in resources if resource["name"] is not None diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py deleted file mode 100644 index 4995ae7b64..0000000000 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py +++ /dev/null @@ -1,86 +0,0 @@ -from pathlib import Path - - -SERVICE_ROOT = Path(__file__).resolve().parents[4] - - -def _read_source(relative_path: str) -> str: - return (SERVICE_ROOT / relative_path).read_text(encoding="utf-8") - - -def test_intent_classification_does_not_require_exact_user_schema_names(): - source = _read_source("src/pipelines/generation/intent_classification.py") - - assert "schema-resolvable references" in source - assert "even if the user did not type exact table or column names" in source - assert "do not require the user to write exact schema identifiers" in source - assert ( - "Do not classify a data retrieval or analytics question as MISLEADING only " - "because the user did not write exact table or column names" - ) in source - - -def test_data_assistance_does_not_invent_hypothetical_schema(): - source = _read_source("src/pipelines/generation/data_assistance.py") - - assert "MUST NOT add SQL code" in source - assert "Use only the provided DATABASE SCHEMA as context" in source - assert "Do not invent, assume, or name tables or columns" in source - assert "do not provide hypothetical schema" in source - - -def test_sql_reasoning_contract_rejects_substitute_identifiers(): - source = _read_source("src/pipelines/generation/utils/sql.py") - - assert "If the DATABASE SCHEMA does not contain an identifier needed" in source - assert "instead of naming a substitute table or column" in source - assert "prompt examples" in source - assert "Identifiers shown in prompt examples are illustrative only" in source - assert "Use `display_label` and `description` only to understand" in source - assert "generated SQL must use that exact identifier" in source - - -def test_sql_correction_receives_raw_wren_engine_validation_error(): - source = _read_source("src/web/v1/services/ask.py") - - assert "_build_sql_correction_error" in source - assert "Original Wren Engine validation error" in source - assert "error_message" in source - - -def test_sql_correction_unknown_identifier_contract(): - source = _read_source("src/pipelines/generation/sql_correction.py") - - assert "If the error reports an unknown table or field" in source - assert "replace it only with an exact executable identifier" in source - assert "Do not retry the same unknown identifier" in source - - -def test_sql_post_processor_has_schema_identifier_guard(): - source = _read_source("src/pipelines/generation/utils/sql.py") - - assert "_validate_sql_against_schema_contexts" in source - assert "SCHEMA_VALIDATION" in source - assert "Unknown table identifiers" in source - assert "Unknown column identifiers" in source - assert "physical schema prefixes" in source - - -def test_generation_paths_pass_retrieved_schema_to_post_processor(): - for relative_path in [ - "src/pipelines/generation/sql_generation.py", - "src/pipelines/generation/followup_sql_generation.py", - "src/pipelines/generation/sql_correction.py", - ]: - source = _read_source(relative_path) - assert "schema_contexts: list[Any] | None = None" in source - assert '"schema_contexts": contexts' in source - - -def test_table_description_retrieval_indexes_business_metadata(): - source = _read_source("src/pipelines/indexing/table_description.py") - - assert "display_label" in source - assert "column_details" in source - assert "clean_display_name" in source - assert "description" in source From 7f0b642c16c677051df7de839ea3e4daeba6b30f Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 17:38:28 +0000 Subject: [PATCH 0693/1087] Revert "Preserve engine errors during SQL correction" This reverts commit e48659bbf21fdd884d4d5988201078b3fc2521bb. --- .../pipelines/generation/sql_correction.py | 1 - .../src/pipelines/generation/utils/sql.py | 2 -- wren-ai-service/src/web/v1/services/ask.py | 24 +++---------------- 3 files changed, 3 insertions(+), 24 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 3e7e06fc81..973b8c69a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -36,7 +36,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). 2. Then, generate the syntactically correct ANSI SQL query to correct the error. -3. If the error reports an unknown table or field, replace it only with an exact executable identifier from the DATABASE SCHEMA. Do not retry the same unknown identifier. ### SQL RULES ### Make sure you follow the SQL Rules strictly. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 7890524c1f..cbc07df03a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -201,8 +201,6 @@ async def _classify_generation_result( - Schema comments are metadata for understanding the data. They are not SQL syntax. - In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. - Never use a `display_label`, alias, or description as an executable table or column identifier. -- Use `display_label` and `description` only to understand which executable `identifier` matches the user's business term. -- When a schema comment contains an `identifier`, generated SQL must use that exact identifier for the table or column. - Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. - You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. - Only apply numeric aggregate functions such as SUM or AVG to numeric columns or measures from the DATABASE SCHEMA. If a column is not numeric in the schema, do not aggregate it directly unless the provided SQL FUNCTIONS and database dialect support the explicit cast you use. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index b1ea202d94..7a20e792c0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -130,21 +130,6 @@ def _is_stopped(self, query_id: str, container: dict): return False - @staticmethod - def _build_sql_correction_error( - diagnosis_reasoning: Optional[str], engine_error: Optional[str] - ) -> str: - diagnosis_reasoning = (diagnosis_reasoning or "").strip() - engine_error = (engine_error or "").strip() - - if diagnosis_reasoning and engine_error: - return ( - f"{diagnosis_reasoning}\n\n" - f"Original Wren Engine validation error:\n{engine_error}" - ) - - return diagnosis_reasoning or engine_error - @observe(name="Ask Question") @trace_metadata async def ask( @@ -558,8 +543,6 @@ async def ask( sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") - else: - sql_diagnosis_reasoning = "" sql_correction_results = await self._pipelines[ "sql_correction" @@ -568,10 +551,9 @@ async def ask( instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": self._build_sql_correction_error( - sql_diagnosis_reasoning, - error_message, - ), + "error": sql_diagnosis_reasoning + if allow_sql_diagnosis + else error_message, }, project_id=ask_request.project_id, use_dry_plan=use_dry_plan, From 8bce512c552011d685d68e26e14e738e26a1f7c3 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 17:38:28 +0000 Subject: [PATCH 0694/1087] Revert "Restore SQL pair grounding during semantics prep" This reverts commit 387a8307d135c612a27753af7d55e364709f2644. --- .../web/v1/services/semantics_preparation.py | 1 + .../services/test_semantics_preparation.py | 45 ------------------- 2 files changed, 1 insertion(+), 45 deletions(-) delete mode 100644 wren-ai-service/tests/pytest/services/test_semantics_preparation.py 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 f2c0f515a3..f8dc4c1224 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -91,6 +91,7 @@ async def prepare_semantics( self._pipelines["sql_pairs"].run( **input, delete_all=True, + include_default_pairs=False, ) ) diff --git a/wren-ai-service/tests/pytest/services/test_semantics_preparation.py b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py deleted file mode 100644 index 3c53046d9d..0000000000 --- a/wren-ai-service/tests/pytest/services/test_semantics_preparation.py +++ /dev/null @@ -1,45 +0,0 @@ -import pytest - -from src.web.v1.services.semantics_preparation import ( - SemanticsPreparationRequest, - SemanticsPreparationService, -) - - -class _RecordingPipeline: - def __init__(self): - self.run_calls = [] - - async def run(self, **kwargs): - self.run_calls.append(kwargs) - - -@pytest.mark.asyncio -async def test_prepare_semantics_reindexes_default_sql_pairs_after_cleanup(): - pipelines = { - name: _RecordingPipeline() - for name in [ - "db_schema", - "historical_question", - "table_description", - "project_meta", - "sql_pairs", - ] - } - service = SemanticsPreparationService(pipelines) - - await service.prepare_semantics( - SemanticsPreparationRequest( - mdl='{"models": []}', - mdl_hash="mdl-hash", - project_id="project-id", - ) - ) - - assert pipelines["sql_pairs"].run_calls == [ - { - "mdl_str": '{"models": []}', - "project_id": "project-id", - "delete_all": True, - } - ] From 3c0012497cded449f2b85bc336ba2dac57d945a3 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 23:32:30 +0530 Subject: [PATCH 0695/1087] Align ask indexing flow with legacy --- wren-ai-service/src/pipelines/common.py | 21 +++------- .../src/pipelines/generation/utils/sql.py | 41 ++++++++++++++----- .../src/pipelines/indexing/db_schema.py | 10 ++--- .../src/pipelines/indexing/utils/helper.py | 20 +++------ .../retrieval/db_schema_retrieval.py | 8 ++-- .../web/v1/services/semantics_preparation.py | 8 +--- 6 files changed, 49 insertions(+), 59 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 940fa66eb7..f6114d63b1 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -4,21 +4,11 @@ from haystack import Document, component -def normalize_data_type(data_type: Any) -> str: - if data_type is None: - return "" - return str(data_type).strip() - - def get_engine_supported_data_type(data_type: str) -> str: """ This function makes sure downstream ai pipeline get column data types in a format that is supported by the data engine. """ - normalized_data_type = normalize_data_type(data_type) - if not normalized_data_type: - return "UNKNOWN" - - match normalized_data_type.upper(): + match data_type.upper(): case "BPCHAR" | "NAME" | "UUID" | "INET": return "VARCHAR" case "OID": @@ -34,7 +24,7 @@ def get_engine_supported_data_type(data_type: str) -> str: case "INT64": return "BIGINT" case _: - return normalized_data_type.upper() + return data_type.upper() def build_table_ddl( @@ -46,17 +36,16 @@ def build_table_ddl( for column in content["columns"]: if column["type"] == "COLUMN": - column_data_type = normalize_data_type(column.get("data_type")) if ( (not columns or (columns and column["name"] in columns)) - and column_data_type.lower() + and column["data_type"].lower() != "unknown" # quick fix: filtering out UNKNOWN column type ): if "This column is a Calculated Field" in column["comment"]: has_calculated_field = True - if column_data_type.lower() == "json": + if column["data_type"].lower() == "json": has_json_field = True - column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column_data_type)}" + column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" if column["is_primary_key"]: column_ddl += " PRIMARY KEY" columns_ddl.append(column_ddl) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index cbc07df03a..088282574e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -198,12 +198,18 @@ async def _classify_generation_result( - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- Schema comments are metadata for understanding the data. They are not SQL syntax. -- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. -- Never use a `display_label`, alias, or description as an executable table or column identifier. -- Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. -- You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. -- Only apply numeric aggregate functions such as SUM or AVG to numeric columns or measures from the DATABASE SCHEMA. If a column is not numeric in the schema, do not aggregate it directly unless the provided SQL FUNCTIONS and database dialect support the explicit cast you use. +- ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. +- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. + - EXAMPLE + DATABASE SCHEMA + /* {"alias":"_orders","description":"A model representing the orders data."} */ + CREATE TABLE orders ( + -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} + ApprovedTimestamp TIMESTAMP + } + + SQL + SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -363,11 +369,29 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields + - For Example: + DATA SCHEMA: + `/* {"alias":"users","description":"A model representing the users data."} */ + CREATE TABLE users ( + -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} + address JSON + )` + To get the city of address in user table use SQL: + `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. + - For Example: + DATA SCHEMA + `/* {"alias":"my_table","description":"A test my_table"} */ + CREATE TABLE my_table ( + -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} + elements JSON + )` + To get the number of elements in my_table table use SQL: + `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". @@ -420,10 +444,7 @@ async def _classify_generation_result( 11. Do not include ```markdown or ``` in the answer. 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. -14. Use only table and column names that appear as identifiers in the DATABASE SCHEMA when writing `table:` and `column:` references. -15. Schema comments, display labels, aliases, and descriptions are context only. Do not use them as executable table or column names in the reasoning plan. -16. Do not create table or column names from words in the user's question. If a requested concept is available only through schema metadata, refer to the corresponding schema identifier. -17. ONLY SHOWING the reasoning plan in bullet points. +14. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 4ad05b8bd0..394d087b46 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -4,7 +4,6 @@ import uuid from typing import Any, Dict, List, Optional -import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from hamilton.function_modifiers import extract_fields @@ -132,17 +131,14 @@ def _convert_models_and_relationships( ) -> List[Dict[str, str]]: def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) - table_name = model["name"] model_properties = { - "identifier": table_name, - "display_label": clean_display_name( - properties.get("displayName", "") - ), + "alias": clean_display_name(properties.get("displayName", "")), "description": properties.get("description", ""), } - comment = f"\n/* {orjson.dumps(model_properties).decode('utf-8')} */\n" + comment = f"\n/* {str(model_properties)} */\n" + table_name = model["name"] payload = { "type": "TABLE", "comment": comment, diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 61ac17084c..3829324a0a 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -29,18 +29,10 @@ def __call__(self, column: Dict[str, Any], **kwargs) -> Any: def _properties_comment(column: Dict[str, Any], **_) -> str: - props = column.get("properties") - if not isinstance(props, dict): - props = {} - - display_name = props.get("displayName", "") - description = props.get("description", "") + props = column["properties"] column_properties = { - "identifier": column.get("name", ""), - "display_label": clean_display_name( - "" if display_name is None else str(display_name) - ), - "description": "" if description is None else str(description), + "alias": clean_display_name(props.get("displayName", "")), + "description": props.get("description", ""), } # Add any nested columns if they exist @@ -64,8 +56,8 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_PREPROCESSORS = { "properties": Helper( - condition=lambda column, **_: isinstance(column.get("properties"), dict), - helper=lambda column, **_: column.get("properties", {}), + condition=lambda column, **_: "properties" in column, + helper=lambda column, **_: column.get("properties"), ), "relationship": Helper( condition=lambda column, **_: "relationship" in column, @@ -83,7 +75,7 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_COMMENT_HELPERS = { "properties": Helper( - condition=lambda column, **_: isinstance(column.get("properties"), dict), + condition=lambda column, **_: "properties" in column, helper=_properties_comment, ), "isCalculated": Helper( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 92e043d98d..6c8dd7bbe3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -79,11 +79,9 @@ - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. - Be logical, concise, and ensure the output strictly follows the required JSON format. -- Schema comments are metadata for understanding the data. They are not SQL syntax. -- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. -- Use table names from the "Create Table" statements, not display labels, aliases, or descriptions. -- Match column names exactly with the definitions in the "Create Table" statements. -- Match table names exactly with the definitions in the "Create Table" statements. +- Use table name used in the "Create Table" statement, don't use "alias". +- Match Column names with the definition in the "Create Table" statement. +- Match Table names with the definition in the "Create Table" statement. Good luck! 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 f8dc4c1224..2ff6215cbe 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -84,16 +84,10 @@ async def prepare_semantics( "db_schema", "historical_question", "table_description", + "sql_pairs", "project_meta", ] ] - tasks.append( - self._pipelines["sql_pairs"].run( - **input, - delete_all=True, - include_default_pairs=False, - ) - ) await asyncio.gather(*tasks) From 810aa289c551b4ce33732ad0848d759e6147bb14 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Mon, 27 Jul 2026 18:10:32 +0000 Subject: [PATCH 0696/1087] Revert "Remove SQL pair default switch" This reverts commit 57f03f3ef5a0aecd675eddf581c4fb46af0d38dd. --- .../src/pipelines/indexing/sql_pairs.py | 8 ++- .../pipelines/indexing/test_sql_pairs.py | 65 +++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index 98cff3b726..16ff245cd7 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -96,7 +96,11 @@ def boilerplates( def sql_pairs( boilerplates: Set[str], external_pairs: Dict[str, Any], + include_default_pairs: bool = True, ) -> List[SqlPair]: + if not include_default_pairs and not external_pairs: + return [] + return [ SqlPair( id=pair.get("id"), @@ -211,13 +215,14 @@ async def run( project_id: str = "", external_pairs: Optional[Dict[str, Any]] = None, delete_all: bool = False, + include_default_pairs: bool = True, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} SQL Pairs Indexing pipeline is running..." ) pairs = { - **self._external_pairs, + **(self._external_pairs if include_default_pairs else {}), **(external_pairs or {}), } @@ -226,6 +231,7 @@ async def run( "project_id": project_id, "external_pairs": pairs, "delete_all": delete_all, + "include_default_pairs": include_default_pairs, **self._components, } diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py index 3436a64c59..0ca5598ad9 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py @@ -3,9 +3,71 @@ from src.config import settings from src.core.provider import DocumentStoreProvider from src.pipelines.indexing import SqlPairs +from src.pipelines.indexing.sql_pairs import SqlPairsCleaner, embedding, sql_pairs, write from src.providers import generate_components +class _RecordingStore: + def __init__(self): + self.filters = [] + + async def delete_documents(self, filter): + self.filters.append(filter) + + +@pytest.mark.asyncio +async def test_sql_pairs_cleaner_delete_all_uses_project_scope(): + store = _RecordingStore() + cleaner = SqlPairsCleaner(store) + + await cleaner.run(sql_pair_ids=[], project_id="project-id", delete_all=True) + + assert store.filters == [ + { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-id"}, + ], + } + ] + + +def test_sql_pairs_can_skip_default_pairs(): + pairs = sql_pairs( + boilerplates={"default"}, + external_pairs={}, + include_default_pairs=False, + ) + + assert pairs == [] + + +@pytest.mark.asyncio +async def test_empty_sql_pairs_skip_embedding_and_write(): + class Embedder: + called = False + + async def run(self, documents): + self.called = True + return {"documents": documents} + + class Writer: + called = False + + async def run(self, documents): + self.called = True + + embedder = Embedder() + writer = Writer() + + result = await embedding({"documents": []}, embedder) + await write(result, writer) + + assert result == {"documents": []} + assert embedder.called is False + assert writer.called is False + + @pytest.mark.asyncio async def test_sql_pairs_indexing_saving_to_document_store(): pipe_components = generate_components(settings.components) @@ -94,3 +156,6 @@ async def test_sql_pairs_deletion(): await pipe.clean(sql_pairs=[], project_id="fake-id") assert await store.count_documents() == 2 + + await pipe.clean(sql_pairs=[], project_id="fake-id", delete_all=True) + assert await store.count_documents() == 0 From ee1af95c4fb05ffc645722655fdeac5821bfc534 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 27 Jul 2026 23:54:30 +0530 Subject: [PATCH 0697/1087] Revert "Align ask indexing flow with legacy" This reverts commit 3c0012497cded449f2b85bc336ba2dac57d945a3. --- wren-ai-service/src/pipelines/common.py | 21 +++++++--- .../src/pipelines/generation/utils/sql.py | 41 +++++-------------- .../src/pipelines/indexing/db_schema.py | 10 +++-- .../src/pipelines/indexing/utils/helper.py | 20 ++++++--- .../retrieval/db_schema_retrieval.py | 8 ++-- .../web/v1/services/semantics_preparation.py | 8 +++- 6 files changed, 59 insertions(+), 49 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index f6114d63b1..940fa66eb7 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -4,11 +4,21 @@ from haystack import Document, component +def normalize_data_type(data_type: Any) -> str: + if data_type is None: + return "" + return str(data_type).strip() + + def get_engine_supported_data_type(data_type: str) -> str: """ This function makes sure downstream ai pipeline get column data types in a format that is supported by the data engine. """ - match data_type.upper(): + normalized_data_type = normalize_data_type(data_type) + if not normalized_data_type: + return "UNKNOWN" + + match normalized_data_type.upper(): case "BPCHAR" | "NAME" | "UUID" | "INET": return "VARCHAR" case "OID": @@ -24,7 +34,7 @@ def get_engine_supported_data_type(data_type: str) -> str: case "INT64": return "BIGINT" case _: - return data_type.upper() + return normalized_data_type.upper() def build_table_ddl( @@ -36,16 +46,17 @@ def build_table_ddl( for column in content["columns"]: if column["type"] == "COLUMN": + column_data_type = normalize_data_type(column.get("data_type")) if ( (not columns or (columns and column["name"] in columns)) - and column["data_type"].lower() + and column_data_type.lower() != "unknown" # quick fix: filtering out UNKNOWN column type ): if "This column is a Calculated Field" in column["comment"]: has_calculated_field = True - if column["data_type"].lower() == "json": + if column_data_type.lower() == "json": has_json_field = True - column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column_data_type)}" if column["is_primary_key"]: column_ddl += " PRIMARY KEY" columns_ddl.append(column_ddl) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..cbc07df03a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -198,18 +198,12 @@ async def _classify_generation_result( - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. -- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. - - EXAMPLE - DATABASE SCHEMA - /* {"alias":"_orders","description":"A model representing the orders data."} */ - CREATE TABLE orders ( - -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} - ApprovedTimestamp TIMESTAMP - } - - SQL - SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; +- Schema comments are metadata for understanding the data. They are not SQL syntax. +- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. +- Never use a `display_label`, alias, or description as an executable table or column identifier. +- Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. +- You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. +- Only apply numeric aggregate functions such as SUM or AVG to numeric columns or measures from the DATABASE SCHEMA. If a column is not numeric in the schema, do not aggregate it directly unless the provided SQL FUNCTIONS and database dialect support the explicit cast you use. - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -369,29 +363,11 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - For Example: - DATA SCHEMA: - `/* {"alias":"users","description":"A model representing the users data."} */ - CREATE TABLE users ( - -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} - address JSON - )` - To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. - - For Example: - DATA SCHEMA - `/* {"alias":"my_table","description":"A test my_table"} */ - CREATE TABLE my_table ( - -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} - elements JSON - )` - To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". @@ -444,7 +420,10 @@ async def _classify_generation_result( 11. Do not include ```markdown or ``` in the answer. 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +14. Use only table and column names that appear as identifiers in the DATABASE SCHEMA when writing `table:` and `column:` references. +15. Schema comments, display labels, aliases, and descriptions are context only. Do not use them as executable table or column names in the reasoning plan. +16. Do not create table or column names from words in the user's question. If a requested concept is available only through schema metadata, refer to the corresponding schema identifier. +17. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 394d087b46..4ad05b8bd0 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -4,6 +4,7 @@ import uuid from typing import Any, Dict, List, Optional +import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from hamilton.function_modifiers import extract_fields @@ -131,14 +132,17 @@ def _convert_models_and_relationships( ) -> List[Dict[str, str]]: def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) + table_name = model["name"] model_properties = { - "alias": clean_display_name(properties.get("displayName", "")), + "identifier": table_name, + "display_label": clean_display_name( + properties.get("displayName", "") + ), "description": properties.get("description", ""), } - comment = f"\n/* {str(model_properties)} */\n" + comment = f"\n/* {orjson.dumps(model_properties).decode('utf-8')} */\n" - table_name = model["name"] payload = { "type": "TABLE", "comment": comment, diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 3829324a0a..61ac17084c 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -29,10 +29,18 @@ def __call__(self, column: Dict[str, Any], **kwargs) -> Any: def _properties_comment(column: Dict[str, Any], **_) -> str: - props = column["properties"] + props = column.get("properties") + if not isinstance(props, dict): + props = {} + + display_name = props.get("displayName", "") + description = props.get("description", "") column_properties = { - "alias": clean_display_name(props.get("displayName", "")), - "description": props.get("description", ""), + "identifier": column.get("name", ""), + "display_label": clean_display_name( + "" if display_name is None else str(display_name) + ), + "description": "" if description is None else str(description), } # Add any nested columns if they exist @@ -56,8 +64,8 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_PREPROCESSORS = { "properties": Helper( - condition=lambda column, **_: "properties" in column, - helper=lambda column, **_: column.get("properties"), + condition=lambda column, **_: isinstance(column.get("properties"), dict), + helper=lambda column, **_: column.get("properties", {}), ), "relationship": Helper( condition=lambda column, **_: "relationship" in column, @@ -75,7 +83,7 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_COMMENT_HELPERS = { "properties": Helper( - condition=lambda column, **_: "properties" in column, + condition=lambda column, **_: isinstance(column.get("properties"), dict), helper=_properties_comment, ), "isCalculated": Helper( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..92e043d98d 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -79,9 +79,11 @@ - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. - Be logical, concise, and ensure the output strictly follows the required JSON format. -- Use table name used in the "Create Table" statement, don't use "alias". -- Match Column names with the definition in the "Create Table" statement. -- Match Table names with the definition in the "Create Table" statement. +- Schema comments are metadata for understanding the data. They are not SQL syntax. +- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. +- Use table names from the "Create Table" statements, not display labels, aliases, or descriptions. +- Match column names exactly with the definitions in the "Create Table" statements. +- Match table names exactly with the definitions in the "Create Table" statements. Good luck! 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 2ff6215cbe..f8dc4c1224 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -84,10 +84,16 @@ async def prepare_semantics( "db_schema", "historical_question", "table_description", - "sql_pairs", "project_meta", ] ] + tasks.append( + self._pipelines["sql_pairs"].run( + **input, + delete_all=True, + include_default_pairs=False, + ) + ) await asyncio.gather(*tasks) From 329aaebb3911f78bdf9f95745e6d9b657198d2c1 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 28 Jul 2026 01:25:00 +0530 Subject: [PATCH 0698/1087] Restore legacy ask indexing flow --- wren-ai-service/src/pipelines/common.py | 21 +++------- .../src/pipelines/generation/utils/sql.py | 41 ++++++++++++++----- .../src/pipelines/indexing/db_schema.py | 10 ++--- .../src/pipelines/indexing/utils/helper.py | 20 +++------ .../retrieval/db_schema_retrieval.py | 8 ++-- .../web/v1/services/semantics_preparation.py | 8 +--- 6 files changed, 49 insertions(+), 59 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 940fa66eb7..f6114d63b1 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -4,21 +4,11 @@ from haystack import Document, component -def normalize_data_type(data_type: Any) -> str: - if data_type is None: - return "" - return str(data_type).strip() - - def get_engine_supported_data_type(data_type: str) -> str: """ This function makes sure downstream ai pipeline get column data types in a format that is supported by the data engine. """ - normalized_data_type = normalize_data_type(data_type) - if not normalized_data_type: - return "UNKNOWN" - - match normalized_data_type.upper(): + match data_type.upper(): case "BPCHAR" | "NAME" | "UUID" | "INET": return "VARCHAR" case "OID": @@ -34,7 +24,7 @@ def get_engine_supported_data_type(data_type: str) -> str: case "INT64": return "BIGINT" case _: - return normalized_data_type.upper() + return data_type.upper() def build_table_ddl( @@ -46,17 +36,16 @@ def build_table_ddl( for column in content["columns"]: if column["type"] == "COLUMN": - column_data_type = normalize_data_type(column.get("data_type")) if ( (not columns or (columns and column["name"] in columns)) - and column_data_type.lower() + and column["data_type"].lower() != "unknown" # quick fix: filtering out UNKNOWN column type ): if "This column is a Calculated Field" in column["comment"]: has_calculated_field = True - if column_data_type.lower() == "json": + if column["data_type"].lower() == "json": has_json_field = True - column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column_data_type)}" + column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" if column["is_primary_key"]: column_ddl += " PRIMARY KEY" columns_ddl.append(column_ddl) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index cbc07df03a..088282574e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -198,12 +198,18 @@ async def _classify_generation_result( - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- Schema comments are metadata for understanding the data. They are not SQL syntax. -- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. -- Never use a `display_label`, alias, or description as an executable table or column identifier. -- Use only table and column names from the CREATE TABLE statements as identifiers in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, ORDER BY, and expressions. -- You may use `display_label` or alias values from schema comments only after AS in the final SELECT clause. -- Only apply numeric aggregate functions such as SUM or AVG to numeric columns or measures from the DATABASE SCHEMA. If a column is not numeric in the schema, do not aggregate it directly unless the provided SQL FUNCTIONS and database dialect support the explicit cast you use. +- ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. +- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. + - EXAMPLE + DATABASE SCHEMA + /* {"alias":"_orders","description":"A model representing the orders data."} */ + CREATE TABLE orders ( + -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} + ApprovedTimestamp TIMESTAMP + } + + SQL + SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -363,11 +369,29 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields + - For Example: + DATA SCHEMA: + `/* {"alias":"users","description":"A model representing the users data."} */ + CREATE TABLE users ( + -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} + address JSON + )` + To get the city of address in user table use SQL: + `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. + - For Example: + DATA SCHEMA + `/* {"alias":"my_table","description":"A test my_table"} */ + CREATE TABLE my_table ( + -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} + elements JSON + )` + To get the number of elements in my_table table use SQL: + `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". @@ -420,10 +444,7 @@ async def _classify_generation_result( 11. Do not include ```markdown or ``` in the answer. 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. -14. Use only table and column names that appear as identifiers in the DATABASE SCHEMA when writing `table:` and `column:` references. -15. Schema comments, display labels, aliases, and descriptions are context only. Do not use them as executable table or column names in the reasoning plan. -16. Do not create table or column names from words in the user's question. If a requested concept is available only through schema metadata, refer to the corresponding schema identifier. -17. ONLY SHOWING the reasoning plan in bullet points. +14. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 4ad05b8bd0..394d087b46 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -4,7 +4,6 @@ import uuid from typing import Any, Dict, List, Optional -import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from hamilton.function_modifiers import extract_fields @@ -132,17 +131,14 @@ def _convert_models_and_relationships( ) -> List[Dict[str, str]]: def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) - table_name = model["name"] model_properties = { - "identifier": table_name, - "display_label": clean_display_name( - properties.get("displayName", "") - ), + "alias": clean_display_name(properties.get("displayName", "")), "description": properties.get("description", ""), } - comment = f"\n/* {orjson.dumps(model_properties).decode('utf-8')} */\n" + comment = f"\n/* {str(model_properties)} */\n" + table_name = model["name"] payload = { "type": "TABLE", "comment": comment, diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 61ac17084c..3829324a0a 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -29,18 +29,10 @@ def __call__(self, column: Dict[str, Any], **kwargs) -> Any: def _properties_comment(column: Dict[str, Any], **_) -> str: - props = column.get("properties") - if not isinstance(props, dict): - props = {} - - display_name = props.get("displayName", "") - description = props.get("description", "") + props = column["properties"] column_properties = { - "identifier": column.get("name", ""), - "display_label": clean_display_name( - "" if display_name is None else str(display_name) - ), - "description": "" if description is None else str(description), + "alias": clean_display_name(props.get("displayName", "")), + "description": props.get("description", ""), } # Add any nested columns if they exist @@ -64,8 +56,8 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_PREPROCESSORS = { "properties": Helper( - condition=lambda column, **_: isinstance(column.get("properties"), dict), - helper=lambda column, **_: column.get("properties", {}), + condition=lambda column, **_: "properties" in column, + helper=lambda column, **_: column.get("properties"), ), "relationship": Helper( condition=lambda column, **_: "relationship" in column, @@ -83,7 +75,7 @@ def _properties_comment(column: Dict[str, Any], **_) -> str: COLUMN_COMMENT_HELPERS = { "properties": Helper( - condition=lambda column, **_: isinstance(column.get("properties"), dict), + condition=lambda column, **_: "properties" in column, helper=_properties_comment, ), "isCalculated": Helper( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 92e043d98d..6c8dd7bbe3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -79,11 +79,9 @@ - Provide a reasoning list (`chain_of_thought_reasoning`) for each table, explaining why each column is necessary. - Provide the reason of selecting the table in (`table_selection_reason`) for each table. - Be logical, concise, and ensure the output strictly follows the required JSON format. -- Schema comments are metadata for understanding the data. They are not SQL syntax. -- In schema comments, `identifier` is the executable table or column name, and `display_label`/`description` are context only. -- Use table names from the "Create Table" statements, not display labels, aliases, or descriptions. -- Match column names exactly with the definitions in the "Create Table" statements. -- Match table names exactly with the definitions in the "Create Table" statements. +- Use table name used in the "Create Table" statement, don't use "alias". +- Match Column names with the definition in the "Create Table" statement. +- Match Table names with the definition in the "Create Table" statement. Good luck! 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 f8dc4c1224..2ff6215cbe 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -84,16 +84,10 @@ async def prepare_semantics( "db_schema", "historical_question", "table_description", + "sql_pairs", "project_meta", ] ] - tasks.append( - self._pipelines["sql_pairs"].run( - **input, - delete_all=True, - include_default_pairs=False, - ) - ) await asyncio.gather(*tasks) From 290fc14a941b6c7680d1dafa21587d49b5dcc65e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 28 Jul 2026 01:50:15 +0530 Subject: [PATCH 0699/1087] Restore legacy UI schema flow --- .../managers/dataSourceSchemaDetector.ts | 389 +--------------- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 423 +++++------------- 2 files changed, 125 insertions(+), 687 deletions(-) diff --git a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts index b8e41b2a3e..df45af84cf 100644 --- a/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts +++ b/wren-ui/src/apollo/server/managers/dataSourceSchemaDetector.ts @@ -1,24 +1,18 @@ -import { camelCase, differenceWith, isEmpty, uniqBy } from 'lodash'; +import { camelCase, differenceWith, isEmpty, isEqual, uniqBy } from 'lodash'; import { IContext } from '@server/types'; import { getLogger } from 'log4js'; import { SchemaChange } from '@server/repositories/schemaChangeRepository'; import { Model, ModelColumn, RelationInfo } from '../repositories'; -import type { - CompactColumn, - CompactTable, -} from '@server/services/metadataService'; -import { - handleNestedColumns, - replaceInvalidReferenceName, - transformUniqueInvalidColumnName, -} from '@server/utils/model'; const logger = getLogger('DataSourceSchemaDetector'); logger.level = 'debug'; export type DataSourceSchema = { name: string; - columns: Array & Partial>; + columns: { + name: string; + type: string; + }[]; }; export type DataSourceSchemaChange = { @@ -62,7 +56,6 @@ interface AffectedResources { export interface IDataSourceSchemaDetector { detectSchemaChange(): Promise; - syncLatestSchemaMetadata(): Promise; resolveSchemaChange(type: string): Promise; getAffectedResources( changes: DataSourceSchema[], @@ -90,10 +83,7 @@ export default class DataSourceSchemaDetector } public async detectSchemaChange() { - logger.info('Start to detect Data Source Schema changes.'); - const currentSchema = await this.getCurrentSchema(); - const latestSchema = await this.getLatestSchema(); - const diffSchema = this.getDiffSchema(currentSchema, latestSchema); + const diffSchema = await this.getDiffSchema(); if (diffSchema) { await this.addSchemaChange(diffSchema); } else { @@ -115,180 +105,7 @@ export default class DataSourceSchemaDetector } } - const hasMissingModelSync = await this.syncLatestSchemaMetadata(); - const hasSchemaMetadataSync = - await this.syncExistingModelsWithLatestSchema(latestSchema); - - return !!diffSchema || hasMissingModelSync || hasSchemaMetadataSync; - } - - public async syncLatestSchemaMetadata() { - logger.info('Start to sync latest datasource schema metadata.'); - const project = await this.ctx.projectRepository.findOneBy({ - id: this.projectId, - }); - const latestTables = - await this.ctx.projectService.getProjectDataSourceTables(project); - const models = await this.ctx.modelRepository.findAllBy({ - projectId: this.projectId, - }); - const createdModels = await this.createMissingModels(latestTables, models); - await this.createColumnsForModels(latestTables, createdModels); - logger.info('Finished syncing latest datasource schema metadata.'); - return createdModels.length > 0; - } - - private async createMissingModels( - latestTables: CompactTable[], - models: Model[], - ): Promise { - const existingSourceTableNames = new Set( - models.map((model) => model.sourceTableName), - ); - const usedReferenceNames = new Set( - models.map((model) => model.referenceName.toLowerCase()), - ); - const missingTables = latestTables.filter( - (table) => !existingSourceTableNames.has(table.name), - ); - - if (!missingTables.length) { - return []; - } - - const modelValues = missingTables.map((table) => { - const referenceName = this.getUniqueModelReferenceName( - replaceInvalidReferenceName(table.name), - usedReferenceNames, - ); - return { - projectId: this.projectId, - displayName: table.name, - referenceName, - sourceTableName: table.name, - cached: false, - refreshTime: null, - properties: table.properties ? JSON.stringify(table.properties) : null, - } as Partial; - }); - - logger.info( - `Creating ${modelValues.length} missing model(s): ${missingTables - .map((table) => table.name) - .join(', ')}`, - ); - return await this.ctx.modelRepository.createMany(modelValues); - } - - private async createColumnsForModels( - latestTables: CompactTable[], - models: Model[], - ): Promise { - if (!models.length) { - return []; - } - - const columnValues = models.flatMap((model) => { - const table = latestTables.find( - (table) => table.name === model.sourceTableName, - ); - if (!table) { - return []; - } - const usedReferenceNames = new Set(); - return this.buildColumnValues(model, table.columns, table.primaryKey, { - usedReferenceNames, - }); - }); - - if (!columnValues.length) { - return []; - } - - const columns = await this.ctx.modelColumnRepository.createMany( - columnValues, - ); - await this.createNestedColumns(latestTables, models, columns); - return columns; - } - - private buildColumnValues( - model: Model, - columns: CompactColumn[], - primaryKey: string | undefined, - { usedReferenceNames }: { usedReferenceNames: Set }, - ): Partial[] { - return columns.map( - (column) => - ({ - modelId: model.id, - isCalculated: false, - displayName: column.name, - referenceName: transformUniqueInvalidColumnName( - column.name, - usedReferenceNames, - ), - sourceColumnName: column.name, - type: column.type || 'string', - notNull: !!column.notNull, - isPk: primaryKey === column.name, - properties: column.properties - ? JSON.stringify(column.properties) - : null, - }) as Partial, - ); - } - - private async createNestedColumns( - latestTables: CompactTable[], - models: Model[], - columns: ModelColumn[], - ) { - const nestedColumnValues = models.flatMap((model) => { - const table = latestTables.find( - (table) => table.name === model.sourceTableName, - ); - if (!table) { - return []; - } - const modelColumns = columns.filter( - (column) => column.modelId === model.id, - ); - return table.columns.flatMap((compactColumn) => { - const column = modelColumns.find( - (column) => column.sourceColumnName === compactColumn.name, - ); - if (!column) { - return []; - } - return handleNestedColumns(compactColumn, { - modelId: column.modelId, - columnId: column.id, - sourceColumnName: column.sourceColumnName, - }); - }); - }); - - if (nestedColumnValues.length) { - await this.ctx.modelNestedColumnRepository.createMany(nestedColumnValues); - } - } - - private getUniqueModelReferenceName( - referenceName: string, - usedReferenceNames: Set, - ) { - const baseName = referenceName || 'model'; - let uniqueName = baseName; - let suffix = 2; - - while (usedReferenceNames.has(uniqueName.toLowerCase())) { - uniqueName = `${baseName}_${suffix}`; - suffix += 1; - } - - usedReferenceNames.add(uniqueName.toLowerCase()); - return uniqueName; + return !!diffSchema; } public async resolveSchemaChange(type: string) { @@ -296,7 +113,6 @@ export default class DataSourceSchemaDetector const supportedTypes = [ SchemaChangeType.DELETED_TABLES, SchemaChangeType.DELETED_COLUMNS, - SchemaChangeType.MODIFIED_COLUMNS, ]; if (!supportedTypes.includes(schemaChangeType)) { throw new Error('Resolved scheme change type is not supported.'); @@ -335,11 +151,10 @@ export default class DataSourceSchemaDetector }); /** - * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS / MODIFIED_COLUMNS + * Handle resolve scheme change for DELETED_TABLES / DELETED_COLUMNS * 1. Remove all affected calculated fields * 2. Remove all affected columns if DELETED_COLUMNS - * 3. Update all affected column types if MODIFIED_COLUMNS - * 4. Remove all affected tables if DELETED_TABLES + * 3. Remove all affected tables if DELETED_TABLES * * Considering that we have set up foreign keys, some data will be automatically deleted in cascade, * so there is no need to perform additional deletions. (E.g., relationships, model's column) @@ -371,27 +186,6 @@ export default class DataSourceSchemaDetector affectedColumnNames, ); } - if (schemaChangeType === SchemaChangeType.MODIFIED_COLUMNS) { - await Promise.all( - resource.columns.map(async (column) => { - const modelColumn = modelColumns.find( - (modelColumn) => - modelColumn.modelId === resource.modelId && - modelColumn.sourceColumnName === column.sourceColumnName && - !modelColumn.isCalculated, - ); - if (!modelColumn || modelColumn.type === column.type) { - return; - } - logger.debug( - `Updating column "${column.sourceColumnName}" type from "${modelColumn.type}" to "${column.type}" in model "${resource.referenceName}".`, - ); - await this.ctx.modelColumnRepository.updateOne(modelColumn.id, { - type: column.type, - }); - }), - ); - } return; }), ); @@ -542,10 +336,11 @@ export default class DataSourceSchemaDetector return affectedResources; } - private getDiffSchema( - currentSchema: DataSourceSchema[], - latestSchema: DataSourceSchema[], - ) { + private async getDiffSchema() { + logger.info('Start to detect Data Source Schema changes.'); + const currentSchema = await this.getCurrentSchema(); + const latestSchema = await this.getLatestSchema(); + const diffSchema = currentSchema.reduce((result, currentTable) => { const lastestTable = latestSchema.find( (table) => table.name === currentTable.name, @@ -563,7 +358,7 @@ export default class DataSourceSchemaDetector const diffColumns = differenceWith( currentTable.columns, lastestTable.columns, - this.isSameSchemaColumn, + isEqual, ); if (diffColumns.length > 0) { const deletedColumnChange = { name: currentTable.name, columns: [] }; @@ -674,159 +469,17 @@ export default class DataSourceSchemaDetector const result = latestDataSourceTables.map((table) => { return { name: table.name, - columns: table.columns, + columns: table.columns.map((column) => { + return { + name: column.name, + type: column.type, + }; + }), }; }); return result; } - private isSameSchemaColumn( - currentColumn: DataSourceSchema['columns'][number], - latestColumn: DataSourceSchema['columns'][number], - ) { - return ( - currentColumn.name === latestColumn.name && - currentColumn.type === latestColumn.type - ); - } - - private async syncExistingModelsWithLatestSchema( - latestSchema: DataSourceSchema[], - ): Promise { - let hasSyncedMetadata = false; - const models = await this.ctx.modelRepository.findAllBy({ - projectId: this.projectId, - }); - if (models.length === 0) { - return false; - } - - const modelIds = models.map((model) => model.id); - const modelColumns = - await this.ctx.modelColumnRepository.findColumnsByModelIds(modelIds); - - for (const model of models) { - const latestTable = latestSchema.find( - (table) => table.name === model.sourceTableName, - ); - if (!latestTable) { - continue; - } - - const existingColumns = modelColumns.filter( - (column) => column.modelId === model.id && !column.isCalculated, - ); - const latestColumnNames = new Set( - latestTable.columns.map((column) => column.name), - ); - const staleColumnNames = existingColumns - .filter((column) => !latestColumnNames.has(column.sourceColumnName)) - .map((column) => column.sourceColumnName); - if (staleColumnNames.length) { - logger.info( - `Removing stale datasource column metadata "${staleColumnNames.join( - ', ', - )}" from model "${model.referenceName}".`, - ); - await this.ctx.modelColumnRepository.deleteAllBySourceColumnNames( - model.id, - staleColumnNames, - ); - hasSyncedMetadata = true; - } - - const usedReferenceNames = new Set( - modelColumns - .filter( - (column) => - column.modelId === model.id && - !staleColumnNames.includes(column.sourceColumnName), - ) - .map((column) => column.referenceName.toLowerCase()), - ); - - for (const latestColumn of latestTable.columns) { - const existingColumn = existingColumns.find( - (column) => column.sourceColumnName === latestColumn.name, - ); - - if (!existingColumn) { - const column = await this.ctx.modelColumnRepository.createOne({ - modelId: model.id, - isCalculated: false, - displayName: latestColumn.name, - referenceName: transformUniqueInvalidColumnName( - latestColumn.name, - usedReferenceNames, - ), - sourceColumnName: latestColumn.name, - type: latestColumn.type || 'string', - notNull: latestColumn.notNull || false, - isPk: false, - properties: latestColumn.properties - ? JSON.stringify(latestColumn.properties) - : null, - }); - - await this.ctx.modelNestedColumnRepository.createMany( - handleNestedColumns(latestColumn as CompactColumn, { - modelId: column.modelId, - columnId: column.id, - sourceColumnName: column.sourceColumnName, - }), - ); - hasSyncedMetadata = true; - continue; - } - - const updateData: Partial = {}; - const latestType = latestColumn.type || 'string'; - const latestNotNull = latestColumn.notNull || false; - if (existingColumn.type !== latestType) { - updateData.type = latestType; - } - if (existingColumn.notNull !== latestNotNull) { - updateData.notNull = latestNotNull; - } - if (latestColumn.properties) { - const existingProperties = existingColumn.properties - ? JSON.parse(existingColumn.properties) - : {}; - const mergedProperties = { - ...latestColumn.properties, - ...existingProperties, - }; - const nextProperties = JSON.stringify(mergedProperties); - if ((existingColumn.properties || null) !== nextProperties) { - updateData.properties = nextProperties; - } - } - - if (!isEmpty(updateData)) { - const column = await this.ctx.modelColumnRepository.updateOne( - existingColumn.id, - updateData, - ); - hasSyncedMetadata = true; - - if (latestType.includes('STRUCT')) { - await this.ctx.modelNestedColumnRepository.deleteAllBy({ - columnId: column.id, - }); - await this.ctx.modelNestedColumnRepository.createMany( - handleNestedColumns(latestColumn as CompactColumn, { - modelId: column.modelId, - columnId: column.id, - sourceColumnName: column.sourceColumnName, - }), - ); - } - } - } - } - return hasSyncedMetadata; - } - private async updateResolveToSchemaChange( lastSchemaChange: SchemaChange, schemaChangeTypes: SchemaChangeType[], diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 08aea09c54..4f7a87671a 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -8,7 +8,6 @@ import { View, } from '../repositories'; import { - ColumnMDL, Manifest, ModelMDL, TableReference, @@ -17,30 +16,11 @@ import { import { getLogger } from '@server/utils'; import { getConfig } from '@server/config'; import { DataSourceName } from '../types'; -import { getUniqueReferenceName } from '../utils/model'; const logger = getLogger('MDLBuilder'); logger.level = 'debug'; const config = getConfig(); -const DOCUMENTED_CALCULATED_FIELD_FUNCTIONS = new Map([ - ['ABS', 'abs'], - ['AVG', 'avg'], - ['COUNT', 'count'], - ['MAX', 'max'], - ['MIN', 'min'], - ['SUM', 'sum'], - ['CBRT', 'cbrt'], - ['CEIL', 'ceil'], - ['EXP', 'exp'], - ['FLOOR', 'floor'], - ['LN', 'ln'], - ['LOG10', 'log10'], - ['ROUND', 'round'], - ['SIGN', 'sign'], - ['LENGTH', 'length'], - ['REVERSE', 'reverse'], -]); export interface MDLBuilderBuildFromOptions { project: Project; @@ -61,11 +41,6 @@ export interface IMDLBuilder { // responsible to generate a valid manifest json export class MDLBuilder implements IMDLBuilder { private manifest: Manifest; - private invalidCalculatedFields: Array<{ - modelId: number; - columnId: number; - reason: string; - }> = []; private project: Project; private readonly models: Model[]; @@ -81,12 +56,6 @@ export class MDLBuilder implements IMDLBuilder { private readonly relatedColumns: ModelColumn[]; // eslint-disable-next-line @typescript-eslint/no-unused-vars private readonly relatedRelations: RelationInfo[]; - private readonly columnNameAliases = new Map(); - private readonly manifestColumnNamesByModel = new Map>(); - private readonly manifestColumnNameBySourceByModel = new Map< - string, - Map - >(); constructor(builderOptions: MDLBuilderBuildFromOptions) { const { @@ -115,7 +84,6 @@ export class MDLBuilder implements IMDLBuilder { } public build(): Manifest { - this.invalidCalculatedFields = []; this.addProject(); this.addModel(); this.addNormalField(); @@ -123,7 +91,6 @@ export class MDLBuilder implements IMDLBuilder { this.addCalculatedField(); this.addView(); this.postProcessManifest(); - this.logInvalidCalculatedFieldSummary(); return this.getManifest(); } @@ -136,13 +103,14 @@ export class MDLBuilder implements IMDLBuilder { return; } this.manifest.models = this.models.map((model: Model) => { - const properties = this.parseProperties(model.properties); + const properties = model.properties ? JSON.parse(model.properties) : {}; // put displayName in properties if (model.displayName) { properties.displayName = model.displayName; } const tableReference = this.buildTableReference(model); - const modelMdl = { + + return { name: model.referenceName, columns: [], tableReference, @@ -160,8 +128,6 @@ export class MDLBuilder implements IMDLBuilder { }, primaryKey: '', // will be modified in addColumn } as ModelMDL; - - return modelMdl; }); } @@ -170,7 +136,7 @@ export class MDLBuilder implements IMDLBuilder { return; } this.manifest.views = this.views.map((view: View) => { - const properties = this.parseProperties(view.properties); + const properties = JSON.parse(view.properties) || {}; // filter out properties that are not null or undefined // and are in the list of properties that are allowed @@ -218,11 +184,18 @@ export class MDLBuilder implements IMDLBuilder { (model: any) => model.name === modelRefName, ); + // modify model primary key + if (column.isPk) { + model.primaryKey = column.referenceName; + } + // add column into model if (!model.columns) { model.columns = []; } - const properties = this.parseProperties(column.properties); + const properties = column.properties + ? JSON.parse(column.properties) + : {}; // put displayName in properties if (column.displayName) { properties.displayName = column.displayName; @@ -243,32 +216,9 @@ export class MDLBuilder implements IMDLBuilder { } }, {}); } - const sourceColumnName = column.sourceColumnName || column.referenceName; - const sourceColumnNames = this.getManifestSourceColumnNameMap(model); - const existingColumnName = sourceColumnNames.get( - sourceColumnName.toLowerCase(), - ); - if (existingColumnName) { - this.columnNameAliases.set(column.id, existingColumnName); - if (column.isPk) { - model.primaryKey = existingColumnName; - } - logger.debug( - `Skipping duplicate source column "${sourceColumnName}" for model "${model.name}". Reusing manifest column "${existingColumnName}".`, - ); - return; - } - - const columnName = this.getManifestColumnName(column, model); - sourceColumnNames.set(sourceColumnName.toLowerCase(), columnName); - // modify model primary key - if (column.isPk) { - model.primaryKey = columnName; - } - - const expression = this.getColumnExpression(column, model, columnName); + const expression = this.getColumnExpression(column, model); model.columns.push({ - name: columnName, + name: column.referenceName, type: column.type, isCalculated: column.isCalculated ? true : false, notNull: column.notNull ? true : false, @@ -287,56 +237,29 @@ export class MDLBuilder implements IMDLBuilder { this.columns .filter(({ isCalculated }) => isCalculated) .forEach((column: ModelColumn) => { - try { - // validate manifest.model exist - const relatedModel = this.relatedModels.find( - (model: any) => model.id === column.modelId, - ); - if (!relatedModel) { - this.recordInvalidCalculatedField( - column.modelId, - column.id, - 'can not find related model', - ); - return; - } - const model = this.manifest.models.find( - (model: any) => model.name === relatedModel.referenceName, - ); - if (!model) { - this.recordInvalidCalculatedField( - column.modelId, - column.id, - 'can not find model', - ); - return; - } - const columnName = this.getManifestColumnName(column, model); - const expression = this.getColumnExpression(column, model, columnName); - if (expression === null) { - this.recordInvalidCalculatedField( - column.modelId, - column.id, - 'invalid calculated field metadata', - ); - return; - } - const columnValue = { - name: columnName, - type: column.type, - isCalculated: true, - expression, - notNull: column.notNull ? true : false, - properties: this.parseProperties(column.properties), - }; - model.columns.push(columnValue); - } catch (error: any) { - this.recordInvalidCalculatedField( - column.modelId, - column.id, - `failed to add calculated field: ${error.message}`, + // validate manifest.model exist + const relatedModel = this.relatedModels.find( + (model: any) => model.id === column.modelId, + ); + const model = this.manifest.models.find( + (model: any) => model.name === relatedModel.referenceName, + ); + if (!model) { + logger.debug( + `Build MDL Column Error: can not find model, modelId "${column.modelId}", columnId: "${column.id}"`, ); + return; } + const expression = this.getColumnExpression(column, model); + const columnValue = { + name: column.referenceName, + type: column.type, + isCalculated: true, + expression, + notNull: column.notNull ? true : false, + properties: JSON.parse(column.properties), + }; + model.columns.push(columnValue); }); } @@ -344,44 +267,31 @@ export class MDLBuilder implements IMDLBuilder { modelName: string, calculatedField: ModelColumn, ) { - try { - const model = this.manifest.models.find( - (model: any) => model.name === modelName, - ); - if (!model) { - logger.debug(`Can not find model "${modelName}" to add calculated field`); - return; - } - const columnName = this.getManifestColumnName(calculatedField, model); - const expression = this.getColumnExpression( - calculatedField, - model, - columnName, - ); - if (expression === null) { - this.recordInvalidCalculatedField( - calculatedField.modelId, - calculatedField.id, - `insert skipped because metadata is invalid for "${calculatedField.referenceName}"`, - ); - return; - } - const columnValue = { - name: columnName, - type: calculatedField.type, - isCalculated: true, - expression, - notNull: calculatedField.notNull ? true : false, - properties: this.parseProperties(calculatedField.properties), - }; - model.columns.push(columnValue); - } catch (error: any) { - this.recordInvalidCalculatedField( - calculatedField.modelId, - calculatedField.id, - `insert failed for "${calculatedField.referenceName}": ${error.message}`, - ); + const model = this.manifest.models.find( + (model: any) => model.name === modelName, + ); + if (!model) { + logger.debug(`Can not find model "${modelName}" to add calculated field`); + return; + } + // if calculated field is already in the model, skip + if ( + model.columns.find( + (column: any) => column.name === calculatedField.referenceName, + ) + ) { + return; } + const expression = this.getColumnExpression(calculatedField, model); + const columnValue = { + name: calculatedField.referenceName, + type: calculatedField.type, + isCalculated: true, + expression, + notNull: calculatedField.notNull ? true : false, + properties: JSON.parse(calculatedField.properties), + }; + model.columns.push(columnValue); } public addRelation(): void { @@ -392,26 +302,24 @@ export class MDLBuilder implements IMDLBuilder { joinType, fromModelName, fromColumnName, - fromColumnId, toModelName, toColumnName, - toColumnId, } = relation; const condition = this.getRelationCondition(relation); this.addRelationColumn(fromModelName, { modelReferenceName: toModelName, - columnReferenceName: - this.columnNameAliases.get(toColumnId) || toColumnName, + columnReferenceName: toColumnName, relation: name, }); this.addRelationColumn(toModelName, { modelReferenceName: fromModelName, - columnReferenceName: - this.columnNameAliases.get(fromColumnId) || fromColumnName, + columnReferenceName: fromColumnName, relation: name, }); - const properties = this.parseProperties(relation.properties); + const properties = relation.properties + ? JSON.parse(relation.properties) + : {}; return { name: name, @@ -452,18 +360,13 @@ export class MDLBuilder implements IMDLBuilder { model.columns = []; } // check if the modelReferenceName is already in the model column - const modelColumnNames = this.getManifestColumnNames(model); - const modelNameDuplicated = modelColumnNames.has( - columnData.modelReferenceName.toLowerCase(), + const modelNameDuplicated = model.columns.find( + (column: any) => column.name === columnData.modelReferenceName, ); - const columnName = getUniqueReferenceName( - modelNameDuplicated + const column = { + name: modelNameDuplicated ? `${columnData.modelReferenceName}_${columnData.columnReferenceName}` : columnData.modelReferenceName, - modelColumnNames, - ); - const column = { - name: columnName, type: columnData.modelReferenceName, properties: null, relationship: columnData.relation, @@ -476,97 +379,64 @@ export class MDLBuilder implements IMDLBuilder { protected getColumnExpression( column: ModelColumn, currentModel?: Partial, - columnReferenceName = column.referenceName, - ): string | null { + ): string { if (!column.isCalculated) { // columns existed in the data source. // Provide original column name in expression to MDL if referenceName has converted. - if (column.sourceColumnName !== columnReferenceName) { + if (column.sourceColumnName !== column.referenceName) { return `"${column.sourceColumnName}"`; } return ''; } // calculated field - const lineage = this.parseLineage(column.lineage); - if (isEmpty(lineage) || !column.aggregation) { - return null; - } + const lineage = JSON.parse(column.lineage) as number[]; // lineage = [relationId1, relationId2, ..., columnId] - const fieldExpression = lineage.reduce((acc, id, index) => { - const isLast = index === lineage.length - 1; - if (isLast) { - // id is columnId - const relatedColumn = this.relatedColumns.find( - (relatedColumn) => relatedColumn.id === id, - ); - const columnReferenceName = relatedColumn - ? this.columnNameAliases.get(relatedColumn.id) || - relatedColumn.referenceName - : null; - if (!columnReferenceName) { + const fieldExpression = Object.entries(lineage).reduce( + (acc, [index, id]) => { + const isLast = parseInt(index) == lineage.length - 1; + if (isLast) { + // id is columnId + const columnReferenceName = this.relatedColumns.find( + (relatedColumn) => relatedColumn.id === id, + )?.referenceName; + acc.push(`\"${columnReferenceName}\"`); return acc; } - acc.push(`\"${columnReferenceName}\"`); - return acc; - } - // id is relationId - const usedRelation = this.relatedRelations.find( - (relatedRelation) => relatedRelation.id === id, - ); - if (!usedRelation || !currentModel?.columns) { - return acc; - } - const relationColumnName = currentModel.columns.find( - (c) => c.relationship === usedRelation.name, - )?.name; - if (!relationColumnName) { + // id is relationId + const usedRelation = this.relatedRelations.find( + (relatedRelation) => relatedRelation.id === id, + ); + const relationColumnName = currentModel!.columns.find( + (c) => c.relationship === usedRelation.name, + ).name; + // move to next model + const nextModelName = + currentModel.name === usedRelation.fromModelName + ? usedRelation.toModelName + : usedRelation.fromModelName; + const nextModel = this.manifest.models.find( + (model) => model.name === nextModelName, + ); + currentModel = nextModel; + acc.push(relationColumnName); return acc; - } - // move to next model - const nextModelName = - currentModel.name === usedRelation.fromModelName - ? usedRelation.toModelName - : usedRelation.fromModelName; - const nextModel = this.manifest.models.find( - (model) => model.name === nextModelName, - ); - currentModel = nextModel; - acc.push(relationColumnName); - return acc; - }, []); - if (fieldExpression.length !== lineage.length) { - return null; - } - const functionName = DOCUMENTED_CALCULATED_FIELD_FUNCTIONS.get( - String(column.aggregation).toUpperCase(), + }, + [], ); - if (!functionName) { - return null; - } - return `${functionName}(${fieldExpression.join('.')})`; + return `${column.aggregation}(${fieldExpression.join('.')})`; } protected getRelationCondition(relation: RelationInfo): string { //TODO phase2: implement the expression for relation condition - const { - fromColumnId, - fromColumnName, - toColumnId, - toColumnName, - fromModelName, - toModelName, - } = relation; - const fromColumnReferenceName = - this.columnNameAliases.get(fromColumnId) || fromColumnName; - const toColumnReferenceName = - this.columnNameAliases.get(toColumnId) || toColumnName; - return `"${fromModelName}".${fromColumnReferenceName} = "${toModelName}".${toColumnReferenceName}`; + const { fromColumnName, toColumnName, fromModelName, toModelName } = + relation; + return `"${fromModelName}".${fromColumnName} = "${toModelName}".${toColumnName}`; } private buildTableReference(model: Model): TableReference | null { const modelProps = model.properties && typeof model.properties === 'string' - ? this.parseProperties(model.properties) + ? JSON.parse(model.properties) : {}; if (!modelProps.table) { return null; @@ -577,91 +447,6 @@ export class MDLBuilder implements IMDLBuilder { table: modelProps.table, }; } - - private parseLineage(lineage?: string): number[] { - if (!lineage) { - return []; - } - try { - const parsedLineage = JSON.parse(lineage); - return Array.isArray(parsedLineage) ? parsedLineage : []; - } catch (error) { - logger.debug(`Can not parse calculated field lineage "${lineage}"`); - return []; - } - } - private parseProperties(properties?: string | null): Record { - if (!properties) { - return {}; - } - try { - const parsed = JSON.parse(properties); - return parsed && typeof parsed === 'object' ? parsed : {}; - } catch (error) { - logger.debug(`Can not parse properties "${properties}"`); - return {}; - } - } - private recordInvalidCalculatedField( - modelId: number, - columnId: number, - reason: string, - ) { - this.invalidCalculatedFields.push({ modelId, columnId, reason }); - } - private logInvalidCalculatedFieldSummary() { - if (this.invalidCalculatedFields.length === 0) { - return; - } - const preview = this.invalidCalculatedFields - .slice(0, 10) - .map( - ({ modelId, columnId, reason }) => - `modelId="${modelId}", columnId="${columnId}", reason="${reason}"`, - ) - .join('; '); - logger.warn( - `Skipped ${this.invalidCalculatedFields.length} invalid calculated field(s) while building MDL. ${preview}${this.invalidCalculatedFields.length > 10 ? '; ...' : ''}`, - ); - } - - private getManifestColumnName( - column: ModelColumn, - model: Partial, - ): string { - if (this.columnNameAliases.has(column.id)) { - return this.columnNameAliases.get(column.id)!; - } - const columnName = getUniqueReferenceName( - column.referenceName, - this.getManifestColumnNames(model), - ); - this.columnNameAliases.set(column.id, columnName); - return columnName; - } - - private getManifestColumnNames(model: Partial): Set { - const modelName = model.name || ''; - if (!this.manifestColumnNamesByModel.has(modelName)) { - const existingColumns = (model.columns || []) as ColumnMDL[]; - this.manifestColumnNamesByModel.set( - modelName, - new Set(existingColumns.map((column) => column.name.toLowerCase())), - ); - } - return this.manifestColumnNamesByModel.get(modelName)!; - } - - private getManifestSourceColumnNameMap( - model: Partial, - ): Map { - const modelName = model.name || ''; - if (!this.manifestColumnNameBySourceByModel.has(modelName)) { - this.manifestColumnNameBySourceByModel.set(modelName, new Map()); - } - return this.manifestColumnNameBySourceByModel.get(modelName)!; - } - private postProcessManifest() { if (this.useRustWrenEngine()) { // 1. remove all the key that the value is null From 033ea3997b42f15aadc0288ad99404a8058b1d4c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 28 Jul 2026 02:37:29 +0530 Subject: [PATCH 0700/1087] Restore legacy AI ask deploy contract --- wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts | 13 +++---------- wren-ui/src/apollo/server/models/adaptor.ts | 2 -- wren-ui/src/apollo/server/services/askingService.ts | 1 - wren-ui/src/apollo/server/services/deployService.ts | 1 - 4 files changed, 3 insertions(+), 14 deletions(-) diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 3598f30cdf..4c691367c7 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -247,18 +247,12 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async ask(input: AskInput): Promise { try { - const body: Record = { + const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, { query: input.query, id: input.deployId, histories: this.transformHistoryInput(input.histories), configurations: input.configurations, - }; - - if (input.projectId) { - body['project_id'] = String(input.projectId); - } - - const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, body); + }); return { queryId: res.data.query_id }; } catch (err: any) { logger.debug(`Got error when asking wren AI: ${getAIServiceError(err)}`); @@ -352,14 +346,13 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } public async deploy(deployData: DeployData): Promise { - const { manifest, hash, projectId } = deployData; + const { manifest, hash } = deployData; try { const res = await axios.post( `${this.wrenAIBaseEndpoint}/v1/semantics-preparations`, { mdl: JSON.stringify(manifest), id: hash, - project_id: projectId.toString(), }, ); const deployId = res.data.id; diff --git a/wren-ui/src/apollo/server/models/adaptor.ts b/wren-ui/src/apollo/server/models/adaptor.ts index 0a60e14ebe..102ed91a5a 100644 --- a/wren-ui/src/apollo/server/models/adaptor.ts +++ b/wren-ui/src/apollo/server/models/adaptor.ts @@ -51,7 +51,6 @@ export enum WrenAILanguage { export interface DeployData { manifest: Manifest; hash: string; - projectId: number; } // ask @@ -74,7 +73,6 @@ export interface ProjectConfigurations { export interface AskInput { query: string; deployId: string; - projectId?: string; histories?: ThreadResponse[]; configurations?: ProjectConfigurations; } diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 4993320ec0..33736f87b7 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -698,7 +698,6 @@ export class AskingService implements IAskingService { query: input.question, histories, deployId, - projectId: projectId.toString(), configurations: { language }, rerunFromCancelled, previousTaskId, diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 746e5fa12e..45f465bdd0 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -138,7 +138,6 @@ export class DeployService implements IDeployService { await this.wrenAIAdaptor.deploy({ manifest, hash, - projectId, }); // update deploy status From 4b579203ab50bdb26a7fe30101a3f0335df39c34 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 14:28:55 +0530 Subject: [PATCH 0701/1087] Restore legacy AI retrieval flow --- .../pipelines/generation/test_sql_utils.py | 6 +- .../test_ask_heuristic_text_to_sql.py | 596 ---- .../pytest/services/test_ask_sales_sql.py | 2561 ----------------- .../services/test_ask_unqueryable_metrics.py | 102 - 4 files changed, 3 insertions(+), 3262 deletions(-) delete mode 100644 wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py delete mode 100644 wren-ai-service/tests/pytest/services/test_ask_sales_sql.py delete mode 100644 wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 51f9d4f2e4..5d51428c4e 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -44,6 +44,6 @@ def test_get_json_field_instructions_uses_sql_knowledge_override(): def test_sql_generation_system_prompt_grounding_contract(): prompt = get_sql_generation_system_prompt() - assert "DATABASE SCHEMA section is the complete and only source" in prompt - assert "MUST NOT introduce, infer, copy, or repair" in prompt - assert "Do not copy identifiers" in prompt + assert "ONLY USE table/column alias in the final SELECT clause" in prompt + assert "Refer to the value of alias from the comment section" in prompt + assert 'SELECT "_orders"."ApprovedTimestamp" AS "_timestamp"' in prompt diff --git a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py b/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py deleted file mode 100644 index c9ed99629f..0000000000 --- a/wren-ai-service/tests/pytest/services/test_ask_heuristic_text_to_sql.py +++ /dev/null @@ -1,596 +0,0 @@ -import asyncio - -from src.web.v1.services.ask import AskHistory, AskService - - -class _FakeSchemaRetrievalPipeline: - def __init__(self): - self.calls = [] - - async def run(self, **kwargs): - self.calls.append(kwargs) - return { - "construct_retrieval_results": { - "retrieval_results": [ - { - "table_name": "dbo_orders", - "table_ddl": ( - "CREATE TABLE dbo_orders (" - "OrderId INT, CustomerId INT, OrderDate TIMESTAMP" - ")" - ), - } - ], - "has_calculated_field": False, - "has_metric": False, - "has_json_field": False, - } - } - - -def test_independent_question_does_not_reuse_historical_sql(): - service = AskService(pipelines={}) - - assert not service._should_reuse_historical_question_sql( - "Show monthly order count by market.", - [], - ) - assert not service._should_reuse_historical_question_sql( - "Show monthly order count by market.", - [AskHistory(question="previous", sql="SELECT 1")], - ) - - -def test_contextual_followup_does_not_reuse_historical_sql(): - service = AskService(pipelines={}) - - assert not service._should_reuse_historical_question_sql( - "Use the same table and show it by month.", - [AskHistory(question="previous", sql="SELECT 1")], - ) - - -def test_metadata_table_question_is_not_sql_or_chart_intent(): - service = AskService(pipelines={}) - - assert service._get_metadata_question_kind( - "What tables are there in this datasource?" - ) == "tables" - assert service._get_metadata_question_kind( - "List the available models in the semantic layer" - ) == "tables" - assert service._get_metadata_question_kind( - "Create a bar chart of orders by table category" - ) is None - - -def test_metadata_column_question_is_not_sql_or_chart_intent(): - service = AskService(pipelines={}) - - assert service._get_metadata_question_kind( - "What columns are available in dbo_orders?" - ) == "columns" - assert service._get_metadata_question_kind( - "Show fields in the CustomerMaster table" - ) == "columns" - assert service._get_metadata_question_kind( - "Show schema for CustomerMaster" - ) == "schema" - assert service._get_metadata_question_kind( - "Show a line chart of monthly order count by customer field" - ) is None - assert service._get_metadata_question_kind( - "What is the row count for dbo_orders?" - ) is None - - -def test_metadata_relationship_and_count_questions_have_specific_intents(): - service = AskService(pipelines={}) - - assert service._get_metadata_question_kind( - "What relationships exist between tables?" - ) == "relationships" - assert service._get_metadata_question_kind( - "How many tables are in this datasource?" - ) == "table_count" - assert service._get_metadata_question_kind( - "How many columns are in dbo_orders?" - ) == "column_count" - - -def test_metadata_table_answer_lists_deployed_tables(): - service = AskService(pipelines={}) - answer = service._build_metadata_response( - "What tables are there in this datasource?", - [ - """ - CREATE TABLE dbo_orders ( - OrderId INT, - CustomerName VARCHAR - ); - """, - """ - CREATE TABLE dbo_customers ( - CustomerId INT, - Region VARCHAR - ); - """, - ], - [], - ) - - assert "active datasource has 2 deployed tables" in answer - assert "- dbo_orders" in answer - assert "- dbo_customers" in answer - - -def test_metadata_column_answer_lists_matching_table_columns(): - service = AskService(pipelines={}) - answer = service._build_metadata_response( - "What columns are available in dbo_orders?", - [ - """ - CREATE TABLE dbo_orders ( - OrderId INT, - CustomerName VARCHAR, - OrderDate TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_customers ( - CustomerId INT, - Region VARCHAR - ); - """, - ], - [], - ) - - assert "dbo_orders" in answer - assert "OrderId (INT)" in answer - assert "CustomerName (VARCHAR)" in answer - assert "OrderDate (TIMESTAMP)" in answer - assert "dbo_customers" not in answer - - -def test_metadata_schema_answer_includes_columns_and_relationships(): - service = AskService(pipelines={}) - answer = service._build_metadata_response( - "Show schema for dbo_orders", - [ - """ - CREATE TABLE dbo_orders ( - OrderId INT, - CustomerId INT, - CONSTRAINT fk_customer FOREIGN KEY (CustomerId) - REFERENCES dbo_customers(CustomerId) - ); - """, - """ - CREATE TABLE dbo_customers ( - CustomerId INT, - Region VARCHAR - ); - """, - ], - [], - ) - - assert "Schema details from the active datasource metadata" in answer - assert "- dbo_orders" in answer - assert "OrderId (INT)" in answer - assert "Relationships:" in answer - assert "dbo_orders(CustomerId) -> dbo_customers(CustomerId)" in answer - - -def test_metadata_relationship_answer_lists_foreign_keys(): - service = AskService(pipelines={}) - answer = service._build_metadata_response( - "What relationships exist between tables?", - [ - """ - CREATE TABLE dbo_orders ( - OrderId INT, - CustomerId INT, - FOREIGN KEY (CustomerId) REFERENCES dbo_customers(CustomerId) - ); - """, - ], - [], - ) - - assert "active datasource metadata has 1 relationship" in answer - assert "dbo_orders(CustomerId) -> dbo_customers(CustomerId)" in answer - - -def test_metadata_count_answers_are_intent_specific(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_orders ( - OrderId INT, - CustomerId INT - ); - """, - """ - CREATE TABLE dbo_customers ( - CustomerId INT, - Region VARCHAR, - Segment VARCHAR - ); - """, - ] - - table_count = service._build_metadata_response( - "How many tables are in this datasource?", - table_ddls, - [], - ) - column_count = service._build_metadata_response( - "How many columns are in dbo_customers?", - table_ddls, - [], - ) - - assert table_count == "The active datasource has 2 deployed tables." - assert column_count == "dbo_customers has 3 deployed columns." - - -def test_manufacturing_throughput_trend_uses_debug_entry_business_unit(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_DebugEntries ( - DebugEntryId VARCHAR, - BusinessUnit VARCHAR, - DateIn TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_batch_records ( - id VARCHAR, - board_model VARCHAR, - production_date TIMESTAMP - ); - """, - ] - - sql = service._build_manufacturing_throughput_sql( - "Show throughput trends across different manufacturing units.", - table_ddls, - table_names=["dbo_DebugEntries", "dbo_batch_records"], - ) - - assert sql - assert '"dbo_DebugEntries"."BusinessUnit"' in sql - assert '"dbo_DebugEntries"."DateIn"' in sql - assert "dbo_batch_records" not in sql - assert 'COUNT(*) AS "throughput"' in sql - assert "DATEPART(MONTH" in sql - - -def test_manufacturing_throughput_fallback_requires_business_unit_column(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_batch_records ( - id VARCHAR, - board_model VARCHAR, - production_date TIMESTAMP - ); - """ - ] - - assert ( - service._build_manufacturing_throughput_sql( - "Show throughput trends across different manufacturing units.", - table_ddls, - table_names=["dbo_batch_records"], - ) - is None - ) - - -def test_repair_failure_count_uses_repair_log_failure_code(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - created_at TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_reports ( - id VARCHAR, - name VARCHAR - ); - """, - ] - - sql = service._build_repair_failure_count_sql( - "Create a bar chart of repair counts grouped by failure category.", - table_ddls, - table_names=["dbo_repair_logs", "dbo_reports"], - ) - - assert sql - assert '"dbo_repair_logs"."failure_code" AS "failure_category"' in sql - assert 'COUNT(*) AS "repair_count"' in sql - assert "Failure Category" not in sql - assert "dbo_reports" not in sql - - -def test_repair_failure_count_prefers_debug_fix_description_when_available(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_DebugEntries ( - DebugEntryId VARCHAR, - FailureSys VARCHAR - ); - """, - """ - CREATE TABLE dbo_DebugFixLogs ( - DebugEntryId VARCHAR, - FixId VARCHAR - ); - """, - """ - CREATE TABLE dbo_DebugFixes ( - Id VARCHAR, - Description VARCHAR - ); - """, - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - failure_code VARCHAR - ); - """, - ] - - sql = service._build_repair_failure_count_sql( - "Create a bar chart of repair counts grouped by failure category.", - table_ddls, - table_names=[ - "dbo_DebugEntries", - "dbo_DebugFixLogs", - "dbo_DebugFixes", - "dbo_repair_logs", - ], - ) - - assert sql - assert '"dbo_DebugFixes"."Description" AS "failure_category"' in sql - assert '"dbo_DebugFixLogs"."FixId" = "dbo_DebugFixes"."Id"' in sql - assert '"dbo_repair_logs"."failure_code"' not in sql - assert "FailurePatternID" not in sql - - -def test_common_pcb_failures_uses_repair_log_failure_code(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - created_at TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_DebugEntries ( - DebugEntryId VARCHAR, - FailureSys VARCHAR - ); - """, - ] - - sql = service._build_repair_failure_count_sql( - "Show top 10 most common PCB failures in a bar chart.", - table_ddls, - table_names=["dbo_repair_logs", "dbo_DebugEntries"], - ) - - assert sql - assert sql.startswith('SELECT "dbo_repair_logs"."failure_code"') - assert '"dbo_repair_logs"."failure_code" AS "failure_category"' in sql - assert 'COUNT(*) AS "repair_count"' in sql - assert "FailureSys" not in sql - assert "TOP 10" not in sql - assert sql.endswith("LIMIT 10") - - -def test_common_pcb_failures_uses_direct_heuristic_route(): - service = AskService(pipelines={}) - - assert service._is_direct_heuristic_sql_query( - "Show top 10 most common PCB failures in a bar chart." - ) - - -def test_repair_sla_compliance_uses_status_when_no_sla_duration_field(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP, - data VARCHAR - ); - """ - ] - - sql = service._build_repair_sla_compliance_sql( - "Generate a dashboard chart for repair SLA compliance.", - table_ddls, - table_names=["dbo_repair_logs"], - ) - - assert sql - assert '"dbo_repair_logs"."status" AS "sla_status"' in sql - assert 'COUNT(*) AS "repair_count"' in sql - assert '"dbo_repair_logs"."turnaround_time"' not in sql - assert '"DAY"' not in sql - assert '"MONTH"' not in sql - assert "DATEDIFF" not in sql.upper() - - -def test_repair_sla_compliance_uses_direct_heuristic_route(): - service = AskService(pipelines={}) - - assert service._is_direct_heuristic_sql_query( - "Generate a dashboard chart for repair SLA compliance." - ) - - -def test_monthly_repair_volume_uses_created_at_bucket(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP, - data VARCHAR - ); - """ - ] - - sql = service._build_monthly_repair_volume_sql( - "Generate a line chart showing monthly repair volume for the last 12 months.", - table_ddls, - table_names=["dbo_repair_logs"], - ) - - assert sql - assert 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year"' in sql - assert 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month"' in sql - assert 'COUNT(*) AS "repair_count"' in sql - assert '"dbo_repair_logs"."MONTH"' not in sql - assert '"MONTH"' not in sql - assert "DATEADD" not in sql.upper() - assert "GETDATE" not in sql.upper() - - -def test_monthly_repair_volume_uses_direct_heuristic_route(): - service = AskService(pipelines={}) - - assert service._is_direct_heuristic_sql_query( - "Generate a line chart showing monthly repair volume for the last 12 months." - ) - - -def test_repair_failure_count_requires_schema_backed_failure_dimension(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - status VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - assert ( - service._build_repair_failure_count_sql( - "Create a bar chart of repair counts grouped by failure category.", - table_ddls, - table_names=["dbo_repair_logs"], - ) - is None - ) - - -def test_ask_result_validation_requires_select_sql(): - service = AskService(pipelines={}) - - assert service._build_ask_result_from_sql("SELECT 1") - assert service._build_ask_result_from_sql("WITH rows AS (SELECT 1) SELECT * FROM rows") - assert service._build_ask_result_from_sql("") is None - assert service._build_ask_result_from_sql("DELETE FROM dbo_repair_logs") is None - assert service._build_ask_result_from_sql(None) is None - - -def test_retrieval_metadata_ignores_malformed_documents(): - service = AskService(pipelines={}) - - documents, table_names, table_ddls = service._extract_retrieval_metadata( - { - "construct_retrieval_results": { - "retrieval_results": [ - {"table_name": "dbo_repair_logs", "table_ddl": "CREATE TABLE dbo_repair_logs (id varchar)"}, - {}, - "bad-document", - ] - } - } - ) - - assert len(documents) == 1 - assert table_names == ["dbo_repair_logs"] - assert table_ddls == ["CREATE TABLE dbo_repair_logs (id varchar)"] - - -def test_complete_sql_generation_context_refetches_full_selected_schema(): - pipeline = _FakeSchemaRetrievalPipeline() - service = AskService( - pipelines={"db_schema_retrieval": pipeline}, - schema_retrieval_timeout_seconds=180, - ) - - documents, table_names, table_ddls, retrieval_result = asyncio.run( - service._complete_sql_generation_context( - query="Show monthly orders by customer.", - project_id="project-1", - documents=[ - { - "table_name": "dbo_orders", - "table_ddl": "CREATE TABLE dbo_orders (OrderId INT)", - } - ], - table_names=["dbo_orders"], - table_ddls=["CREATE TABLE dbo_orders (OrderId INT)"], - ) - ) - - assert table_names == ["dbo_orders"] - assert table_ddls == [ - "CREATE TABLE dbo_orders (OrderId INT, CustomerId INT, OrderDate TIMESTAMP)" - ] - assert documents == [ - { - "table_name": "dbo_orders", - "table_ddl": ( - "CREATE TABLE dbo_orders (" - "OrderId INT, CustomerId INT, OrderDate TIMESTAMP" - ")" - ), - } - ] - assert retrieval_result["retrieval_results"] == documents - assert pipeline.calls == [ - { - "query": "Show monthly orders by customer.", - "tables": ["dbo_orders"], - "project_id": "project-1", - "histories": [], - "enable_column_pruning": False, - } - ] diff --git a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py b/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py deleted file mode 100644 index 9593fca054..0000000000 --- a/wren-ai-service/tests/pytest/services/test_ask_sales_sql.py +++ /dev/null @@ -1,2561 +0,0 @@ -from src.web.v1.services.ask import AskService - - -def test_build_schema_grounded_sales_sql_for_salesperson_performance(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Create a SalesPerson performance ranking chart", - [ - """ - CREATE TABLE dbo_tblSales ( - SalesPerson VARCHAR, - SalesValue DOUBLE, - CustNo VARCHAR, - Country VARCHAR, - "MRO%" DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_tblSales"."SalesPerson" AS "SalesPerson", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."SalesPerson" IS NOT NULL ' - 'GROUP BY "dbo_tblSales"."SalesPerson" ' - 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' - ) - assert "MRO" not in sql - assert "CustID" not in sql - - -def test_build_schema_grounded_sales_sql_for_sales_by_product_category(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "What is the distribution of sales across product categories?", - [ - """ - CREATE TABLE dbo_qSalesMargin ( - ProductCategory VARCHAR, - ProdName VARCHAR, - SalesValue DOUBLE, - OrdNo VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_qSalesMargin"."ProductCategory" AS "ProductCategory", ' - 'SUM("dbo_qSalesMargin"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_qSalesMargin" ' - 'WHERE "dbo_qSalesMargin"."ProductCategory" IS NOT NULL ' - 'GROUP BY "dbo_qSalesMargin"."ProductCategory" ' - 'ORDER BY SUM("dbo_qSalesMargin"."SalesValue") DESC' - ) - assert "COUNT" not in sql - - -def test_build_schema_grounded_sales_sql_for_total_quantity_sold_by_product(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show total quantity sold by product.", - [ - """ - CREATE TABLE dbo_tblOrderLines ( - ProductName VARCHAR, - Quantity DOUBLE, - OrderNo VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblOrderLines"."ProductName" AS "ProductName", ' - 'SUM("dbo_tblOrderLines"."Quantity") AS "TotalQuantity" ' - 'FROM "dbo_tblOrderLines" ' - 'WHERE "dbo_tblOrderLines"."ProductName" IS NOT NULL ' - 'GROUP BY "dbo_tblOrderLines"."ProductName" ' - 'ORDER BY SUM("dbo_tblOrderLines"."Quantity") DESC' - ) - assert "COUNT" not in sql - - -def test_broad_request_explicit_tables_are_not_forced_schema_scope(): - service = AskService.__new__(AskService) - table_names = [f"table_{index}" for index in range(6)] - - assert service._forced_explicit_table_names(table_names) == [] - assert service._forced_explicit_table_names(table_names[:5]) == table_names[:5] - - -def test_full_schema_loading_gate_allows_only_metadata_questions(): - service = AskService.__new__(AskService) - - assert service._should_load_full_schema_for_question( - "List all tables in this datasource." - ) - assert service._should_load_full_schema_for_question( - "How many deployed models are available?" - ) - assert service._should_load_full_schema_for_question( - "Show the schema metadata." - ) - - assert not service._should_load_full_schema_for_question( - "Show top 5 customers by order count." - ) - assert not service._should_load_full_schema_for_question( - "From dbo_tblNewOrders, show the top 5 customers by order count using CustName." - ) - assert not service._should_load_full_schema_for_question( - "What is the distribution of sales across product categories?" - ) - - -def test_schema_grounded_sql_fallback_requires_retrieved_documents(): - service = AskService.__new__(AskService) - - assert not service._can_use_schema_grounded_sql_fallback( - [], - [ - """ - CREATE TABLE dbo_orders ( - CustomerName VARCHAR, - OrderId VARCHAR - ); - """ - ], - "Show top customers by order count.", - ) - assert not service._can_use_schema_grounded_sql_fallback( - [{"table_name": "dbo_orders"}], - [], - "Show top customers by order count.", - ) - assert service._can_use_schema_grounded_sql_fallback( - [{"table_name": "dbo_orders"}], - [ - """ - CREATE TABLE dbo_orders ( - CustomerName VARCHAR, - OrderId VARCHAR - ); - """ - ], - "Show top customers by order count.", - ) - - -def test_build_direct_orders_sales_sql_for_salesperson_order_count(): - service = AskService.__new__(AskService) - sql = service._build_direct_orders_sales_sql( - "Create a bar chart of top 10 SalesPerson by order count" - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_tblSales"."SalesPerson" AS "SalesPerson", ' - 'COUNT(*) AS "OrderCount" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."SalesPerson" IS NOT NULL ' - 'GROUP BY "dbo_tblSales"."SalesPerson" ' - 'ORDER BY COUNT(*) DESC' - ) - - -def test_direct_heuristic_gate_is_disabled_for_generic_schema_selection(): - service = AskService.__new__(AskService) - - assert not service._is_direct_heuristic_sql_query( - "Show top 10 customers by invoice amount." - ) - assert not service._is_direct_heuristic_sql_query( - "Show throughput trends across different manufacturing units." - ) - - -def test_data_query_timeout_retry_does_not_allow_full_project_schema(): - service = AskService.__new__(AskService) - - assert not service._should_retry_selected_schema_after_retrieval_timeout(None) - assert not service._should_retry_selected_schema_after_retrieval_timeout([]) - assert service._should_retry_selected_schema_after_retrieval_timeout(["orders"]) - - -def test_rewrite_query_for_text_to_sql_preserves_user_question(): - service = AskService.__new__(AskService) - - query = "Show total invoice amount by currency." - rewritten = service._rewrite_query_for_text_to_sql(query) - - assert rewritten == query - - -def test_validated_sql_rejects_count_for_total_amount_question(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_Invoices"."Currency" AS "Currency", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_Invoices" ' - 'GROUP BY "dbo_Invoices"."Currency"' - ), - [ - """ - CREATE TABLE dbo_Invoices ( - Currency VARCHAR, - InvoiceAmount DOUBLE - ); - """ - ], - "Show total invoice amount by currency.", - ) - - assert result is None - - -def test_validated_sql_rejects_detail_rows_for_customer_order_count_question(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT TOP 5 "dbo_tblNewOrders"."CustName" AS "CustName", ' - '"dbo_tblNewOrders"."OrdNo" AS "OrdNo" ' - 'FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" IS NOT NULL' - ), - [ - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - "From dbo_tblNewOrders, show the top 5 customers by order count using CustName.", - ) - - assert result is None - - -def test_schema_grounded_table_question_groups_top_customers_by_order_count(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "From dbo_tblNewOrders, show the top 5 customers by order count using CustName.", - [ - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 5 "dbo_tblNewOrders"."CustName" AS "CustName", ' - 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "RecordCount" ' - 'FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" IS NOT NULL ' - 'AND LTRIM(RTRIM("dbo_tblNewOrders"."CustName")) <> \'\' ' - 'GROUP BY "dbo_tblNewOrders"."CustName" ' - 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' - ) - - -def test_validated_sql_rejects_country_question_without_country_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" AS ' - '"Commodity_Line_Value", COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Commodity_Line_Value" ' - 'ORDER BY COUNT(*) DESC' - ), - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - "Show the total commodity line value by country.", - ) - - assert result is None - - -def test_validated_sql_rejects_unrequested_test_table_for_market_distribution(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT TOP 10 "dbo_xStageLoad8_Test"."Market" AS "Market", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_xStageLoad8_Test" ' - 'WHERE "dbo_xStageLoad8_Test"."Market" IS NOT NULL ' - 'GROUP BY "dbo_xStageLoad8_Test"."Market" ' - 'ORDER BY COUNT(*) DESC' - ), - [ - """ - CREATE TABLE dbo_xStageLoad8_Test ( - Market VARCHAR, - OrdNo VARCHAR - ); - """ - ], - "Show order distribution across markets.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_groups_commodity_value_by_country(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Show the total commodity line value by country.", - [ - """ - CREATE TABLE dbo_ytblTarrifsExportsA ( - Country_of_Ultimate_Destination_Code VARCHAR, - Commodity_Line_Value DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'AS "Country_of_Ultimate_Destination_Code", ' - 'SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") ' - 'AS "TotalCommodity_Line_Value" ' - 'FROM "dbo_ytblTarrifsExportsA" ' - 'WHERE "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsExportsA"."Country_of_Ultimate_Destination_Code" ' - 'ORDER BY SUM("dbo_ytblTarrifsExportsA"."Commodity_Line_Value") DESC' - ) - - -def test_validated_sql_rejects_entity_lookup_using_non_customer_column(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT * FROM "dbo_tnoStageNewOrders" ' - 'WHERE "dbo_tnoStageNewOrders"."Division" = ' - "'Daimler Trucks North America'" - ), - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """ - ], - "List orders for Daimler Trucks North America.", - ) - - assert result is None - - -def test_schema_grounded_sales_sql_filters_entity_lookup_by_customer_name(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "List orders for Daimler Trucks North America.", - [ - """ - CREATE TABLE dbo_tnoStageNewOrders ( - Division VARCHAR, - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - """ - CREATE TABLE dbo_tblNewOrders ( - CustName VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """, - ], - ) - - assert sql == ( - 'SELECT TOP 500 * FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."CustName" = ' - "'Daimler Trucks North America' " - 'ORDER BY "dbo_tblNewOrders"."OrdDate" DESC' - ) - - -def test_validated_sql_rejects_distinct_with_extra_entity_columns_for_no_duplicates(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT DISTINCT "dbo_Customers"."CustomerName" AS "CustomerName", ' - '"dbo_Customers"."CustomerId" AS "CustomerId" ' - 'FROM "dbo_Customers"' - ), - [ - """ - CREATE TABLE dbo_Customers ( - CustomerName VARCHAR, - CustomerId VARCHAR - ); - """ - ], - "List customer names with no duplicates.", - ) - - assert result is None - - -def test_validated_sql_rejects_grouped_extra_columns_for_no_duplicates(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_Customers"."CustomerName" AS "CustomerName", ' - '"dbo_Customers"."CustomerId" AS "CustomerId" ' - 'FROM "dbo_Customers" ' - 'GROUP BY "dbo_Customers"."CustomerName", "dbo_Customers"."CustomerId"' - ), - [ - """ - CREATE TABLE dbo_Customers ( - CustomerName VARCHAR, - CustomerId VARCHAR - ); - """ - ], - "List customer names with no duplicates.", - ) - - assert result is None - - -def test_build_direct_orders_sales_sql_for_top_new_orders_q1(): - service = AskService.__new__(AskService) - sql = service._build_direct_orders_sales_sql( - "Show the top 20 new orders for period 2026-Q1" - ) - - assert sql == ( - 'SELECT TOP 20 "dbo_tblSales"."BU" AS "BU", ' - '"dbo_tblSales"."Market" AS "Market", ' - '"dbo_tblSales"."Customer" AS "Customer", ' - '"dbo_tblSales"."ProdName" AS "ProdName", ' - '"dbo_tblSales"."SalesValue" AS "SalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."OrdDate" >= \'2026-01-01 00:00:00\' ' - 'AND "dbo_tblSales"."OrdDate" < \'2026-04-01 00:00:00\' ' - 'ORDER BY "dbo_tblSales"."SalesValue" DESC' - ) - - -def test_build_direct_orders_sales_sql_for_market_growth_comparison(): - service = AskService.__new__(AskService) - sql = service._build_direct_orders_sales_sql( - "Which markets had the highest growth in the first 6 months of this year compared with the same period last year?" - ) - - assert sql is not None - assert '"dbo_tblSales"."Market" AS "Market"' in sql - assert '"CurrentPeriodSales"' in sql - assert '"PreviousPeriodSales"' in sql - assert '"SalesGrowth"' in sql - assert "2026-01-01" in sql - assert "2025-01-01" in sql - - -def test_build_schema_grounded_sales_sql_requires_sales_schema(): - service = AskService.__new__(AskService) - - assert ( - service._build_schema_grounded_sales_sql( - "Create a SalesPerson performance ranking chart", - ["CREATE TABLE dbo_other (SalesPerson VARCHAR);"], - ) - is None - ) - - -def test_build_explicit_table_preview_sql_for_named_table(): - service = AskService.__new__(AskService) - result = service._build_explicit_table_preview_sql( - "Show the first 5 rows from tblNewOrders", - [ - """ - CREATE TABLE tblNewOrders ( - OrdNo VARCHAR, - Customer VARCHAR, - InvDate TIMESTAMP - ); - """ - ], - ) - - assert result == ('SELECT TOP 5 * FROM "tblNewOrders"', "tblNewOrders") - - -def test_build_explicit_table_preview_sql_for_show_data_prompt(): - service = AskService.__new__(AskService) - result = service._build_explicit_table_preview_sql( - "Show data from CustomerMaster", - [ - """ - CREATE TABLE CustomerMaster ( - CustomerId VARCHAR, - CustomerName VARCHAR - ); - """ - ], - ) - - assert result == ('SELECT TOP 10 * FROM "CustomerMaster"', "CustomerMaster") - - -def test_extract_explicit_table_names_from_query(): - service = AskService.__new__(AskService) - - assert service._extract_explicit_table_names_from_query( - "Show the first 10 rows from tblNewOrders" - ) == ["tblNewOrders"] - assert service._extract_explicit_table_names_from_query( - "Show the latest records from last month" - ) == [] - assert service._extract_explicit_table_names_from_query( - "Show all customers names" - ) == [] - - -def test_extract_explicit_table_names_from_using_clause(): - service = AskService.__new__(AskService) - - assert service._extract_explicit_table_names_from_query( - "Show new orders by CustName using dbo.XStageNewOrders OrdDate and CustName" - ) == ["dbo.XStageNewOrders"] - assert service._extract_explicit_table_names_from_query( - "Show new orders using OrdDate and CustName" - ) == [] - - -def test_extract_explicit_table_names_from_in_clause(): - service = AskService.__new__(AskService) - - assert service._extract_explicit_table_names_from_query( - "Which customers have the highest number of orders in dbo.tblNewOrders?" - ) == ["dbo.tblNewOrders"] - assert service._extract_explicit_table_names_from_query( - "Which customers have the highest number of orders in market?" - ) == [] - assert service._extract_explicit_table_names_from_query( - "Show top 10 customers by invoice amount in the current year" - ) == [] - - -def test_build_schema_ranked_measure_sql_uses_matching_dimension_and_measure(): - service = AskService.__new__(AskService) - - sql = service._build_schema_ranked_measure_sql( - "Show top 10 customers by invoice amount", - [ - """ - CREATE TABLE sales_fact ( - Customer_Name VARCHAR, - Product_Name VARCHAR, - Transaction_Amount FLOAT, - Invoice_ID VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "sales_fact"."Customer_Name" AS "Customer_Name", ' - 'SUM("sales_fact"."Transaction_Amount") AS "TotalTransaction_Amount" ' - 'FROM "sales_fact" ' - 'WHERE "sales_fact"."Customer_Name" IS NOT NULL ' - 'AND "sales_fact"."Transaction_Amount" IS NOT NULL ' - 'GROUP BY "sales_fact"."Customer_Name" ' - 'ORDER BY SUM("sales_fact"."Transaction_Amount") DESC' - ) - - -def test_extract_explicit_table_names_from_repair_logs_phrase(): - service = AskService.__new__(AskService) - - assert service._extract_explicit_table_names_from_query( - "Count repair logs by failure_code in repair logs." - ) == [] - - -def test_extract_explicit_table_names_from_pcb_repair_phrases(): - service = AskService.__new__(AskService) - - assert service._extract_explicit_table_names_from_query( - "How many different board models are present in the dbo.repair_logs table?" - ) == ["dbo.repair_logs"] - assert service._extract_explicit_table_names_from_query( - "Display top 10 ticket labels." - ) == [] - - -def test_explicit_table_name_candidates_include_dotted_and_short_forms(): - service = AskService.__new__(AskService) - - assert service._explicit_table_name_candidates("dbo_tblNewOrders") == [ - "dbo_tblNewOrders", - "dbo.tblNewOrders", - "tblNewOrders", - ] - assert service._explicit_table_name_candidates("dbo.tblNewOrders") == [ - "dbo.tblNewOrders", - "dbo_tblNewOrders", - "tblNewOrders", - ] - - -def test_filter_retrieval_metadata_for_explicit_query_keeps_only_named_table(): - service = AskService.__new__(AskService) - documents = [ - { - "table_name": "dbo_knowledge_articles", - "table_ddl": """ - CREATE TABLE dbo_knowledge_articles ( - id INTEGER, - last_run_date TIMESTAMP, - category VARCHAR - ); - """, - }, - { - "table_name": "dbo_failure_patterns", - "table_ddl": """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - created_at TIMESTAMP - ); - """, - }, - ] - - filtered_documents, table_names, table_ddls = ( - service._filter_retrieval_metadata_for_explicit_query( - "show monthly record count by created_at in dbo.failure_patterns.", - documents, - ) - ) - - assert filtered_documents == [documents[1]] - assert table_names == ["dbo_failure_patterns"] - assert table_ddls == [documents[1]["table_ddl"]] - - -def test_filter_retrieval_metadata_for_explicit_query_matches_dotted_table_name(): - service = AskService.__new__(AskService) - documents = [ - { - "table_name": "dbo.tblNewOrders", - "table_ddl": """ - CREATE TABLE "dbo.tblNewOrders" ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - }, - { - "table_name": "dbo_other", - "table_ddl": """ - CREATE TABLE dbo_other ( - CustName VARCHAR, - OrdNo VARCHAR - ); - """, - }, - ] - - filtered_documents, table_names, table_ddls = ( - service._filter_retrieval_metadata_for_explicit_query( - "Show the top 5 CustName values from dbo_tblNewOrders by number of orders.", - documents, - ["dbo_tblNewOrders"], - ) - ) - - assert filtered_documents == [documents[0]] - assert table_names == ["dbo.tblNewOrders"] - assert table_ddls == [documents[0]["table_ddl"]] - - -def test_build_validated_ask_result_rejects_sql_for_different_explicit_table(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date") AS "year", ' - '"dbo_knowledge_articles"."category" AS "category", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_knowledge_articles" ' - 'GROUP BY DATEPART(YEAR, "dbo_knowledge_articles"."last_run_date"), ' - '"dbo_knowledge_articles"."category"' - ), - [ - """ - CREATE TABLE dbo_knowledge_articles ( - id INTEGER, - last_run_date TIMESTAMP, - category VARCHAR - ); - """, - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - created_at TIMESTAMP - ); - """, - ], - "show monthly record count by created_at in dbo.failure_patterns.", - ) - - assert result is None - - -def test_build_validated_ask_result_accepts_sql_for_explicit_table_alias(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_failure_patterns" ' - 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at")' - ), - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - created_at TIMESTAMP - ); - """ - ], - "show monthly record count by created_at in dbo.failure_patterns.", - ) - - assert result is not None - - -def test_needs_conversation_context_only_for_true_followups(): - service = AskService.__new__(AskService) - - assert not service._needs_conversation_context( - "Show refund status distribution by Refund_Status" - ) - assert service._needs_conversation_context( - "What about the same period from the previous result?" - ) - - -def test_prune_sql_generation_context_prefers_referenced_table_and_columns(): - service = AskService.__new__(AskService) - table_ddls = [ - """ - CREATE TABLE dbo_Customers ( - CustomerId VARCHAR, - CustomerName VARCHAR - ); - """, - """ - CREATE TABLE dbo_Products ( - ProductId VARCHAR, - ProductName VARCHAR - ); - """, - """ - CREATE TABLE dbo_XStageNewOrders ( - OrdNo VARCHAR, - OrdDate TIMESTAMP, - CustName VARCHAR - ); - """, - ] - documents = [ - {"table_name": "dbo_Customers", "table_ddl": table_ddls[0]}, - {"table_name": "dbo_Products", "table_ddl": table_ddls[1]}, - {"table_name": "dbo_XStageNewOrders", "table_ddl": table_ddls[2]}, - ] - - _, table_names, pruned_ddls = service._prune_sql_generation_context( - "Show new orders by CustName using dbo.XStageNewOrders OrdDate and CustName", - documents, - [document["table_name"] for document in documents], - table_ddls, - max_tables=1, - ) - - assert table_names == ["dbo_XStageNewOrders"] - assert pruned_ddls == [table_ddls[2]] - - -def test_prune_sql_generation_context_keeps_related_join_table(): - service = AskService.__new__(AskService) - table_ddls = [ - """ - CREATE TABLE dbo_Customers ( - CustomerId VARCHAR, - CustomerName VARCHAR - ); - """, - """ - CREATE TABLE dbo_Products ( - ProductId VARCHAR, - ProductName VARCHAR - ); - """, - """ - CREATE TABLE dbo_Orders ( - OrderId VARCHAR, - CustomerId VARCHAR, - ProductId VARCHAR, - OrderDate TIMESTAMP, - FOREIGN KEY (CustomerId) REFERENCES dbo_Customers(CustomerId) - ); - """, - ] - documents = [ - {"table_name": "dbo_Customers", "table_ddl": table_ddls[0]}, - {"table_name": "dbo_Products", "table_ddl": table_ddls[1]}, - {"table_name": "dbo_Orders", "table_ddl": table_ddls[2]}, - ] - - _, table_names, _ = service._prune_sql_generation_context( - "Which customer names have the highest number of orders?", - documents, - [document["table_name"] for document in documents], - table_ddls, - max_tables=2, - ) - - assert table_names == ["dbo_Customers", "dbo_Orders"] - - -def test_build_schema_grounded_sales_sql_for_top_markets(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "What are the Top 10 Markets by New Order Value this year?", - [ - """ - CREATE TABLE dbo_tblSales ( - Market VARCHAR, - SalesValue DOUBLE, - OrdDate TIMESTAMP, - Division VARCHAR, - ProdType VARCHAR - ); - """, - """ - CREATE TABLE dbo_tblStageNewOrders ( - Market VARCHAR, - NewOrderValue DOUBLE, - OrdDate TIMESTAMP - ); - """, - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_tblSales"."Market" AS "Market", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."OrdDate" >= \'2026-01-01 00:00:00\' ' - 'AND "dbo_tblSales"."OrdDate" < \'2027-01-01 00:00:00\' ' - 'AND "dbo_tblSales"."Market" IS NOT NULL ' - 'GROUP BY "dbo_tblSales"."Market" ' - 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' - ) - assert "dbo_tblStageNewOrders" not in sql - - -def test_build_schema_grounded_sales_sql_for_market_performance_over_time_with_count_fallback(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Compare market performance over time.", - [ - """ - CREATE TABLE dbo_OrderEvents ( - Market VARCHAR, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_OrderEvents"."OrdDate") AS "year", ' - 'DATEPART(MONTH, "dbo_OrderEvents"."OrdDate") AS "month", ' - '"dbo_OrderEvents"."Market" AS "Market", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_OrderEvents" ' - 'WHERE "dbo_OrderEvents"."Market" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_OrderEvents"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_OrderEvents"."OrdDate"), "dbo_OrderEvents"."Market" ' - 'ORDER BY DATEPART(YEAR, "dbo_OrderEvents"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_OrderEvents"."OrdDate")' - ) - - -def test_build_schema_grounded_sales_sql_for_invoice_distribution_by_currency(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show invoice distribution by currency.", - [ - """ - CREATE TABLE dbo_Invoices ( - InvoiceNo VARCHAR, - Currency VARCHAR, - InvoiceAmount DOUBLE, - InvoiceDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_Invoices"."Currency" AS "Currency", ' - 'COUNT(DISTINCT "dbo_Invoices"."InvoiceNo") AS "OrderCount" ' - 'FROM "dbo_Invoices" ' - 'WHERE "dbo_Invoices"."Currency" IS NOT NULL ' - 'GROUP BY "dbo_Invoices"."Currency" ' - 'ORDER BY COUNT(DISTINCT "dbo_Invoices"."InvoiceNo") DESC' - ) - - -def test_validated_sql_normalizes_direction_keywords_before_parser_execution(): - service = AskService.__new__(AskService) - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_tblSales"."ProdName" AS "ProdName", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'GROUP BY "dbo_tblSales"."ProdName" ' - 'ORDER BY SUM("dbo_tblSales"."SalesValue") Desc' - ), - [ - """ - CREATE TABLE dbo_tblSales ( - ProdName VARCHAR, - SalesValue DOUBLE - ); - """ - ], - "Show top-selling products.", - ) - - assert result is not None - assert result.sql.endswith('ORDER BY SUM("dbo_tblSales"."SalesValue") DESC') - - -def test_build_schema_grounded_sales_sql_for_division_revenue_trend(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Create a Division-wise revenue trend line chart.", - [ - """ - CREATE TABLE dbo_tblSales ( - Division VARCHAR, - SalesValue DOUBLE, - OrdDate TIMESTAMP, - Market VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_tblSales"."OrdDate") AS "year", ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") AS "month", ' - '"dbo_tblSales"."Division" AS "Division", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Division" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), "dbo_tblSales"."Division" ' - 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate")' - ) - - -def test_build_schema_grounded_sales_sql_for_orders_by_dimensions(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show New Orders by Division, Market, and Product Type.", - [ - """ - CREATE TABLE dbo_tblSales ( - Division VARCHAR, - Market VARCHAR, - ProdType VARCHAR, - SalesValue DOUBLE, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblSales"."Market" AS "Market", ' - '"dbo_tblSales"."Division" AS "Division", ' - '"dbo_tblSales"."ProdType" AS "ProdType", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."Division" IS NOT NULL ' - 'AND "dbo_tblSales"."ProdType" IS NOT NULL ' - 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Division", ' - '"dbo_tblSales"."ProdType" ' - 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' - ) - - -def test_build_schema_grounded_sales_sql_for_order_date_distribution_by_dimensions(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - 'In the "Backlog" category, what is the distribution of order dates ' - "(OrdDate) for each product type (ProdType) sold in each market " - "segment (Market), considering the salesperson responsible " - "(SalesPerson)?", - [ - """ - CREATE TABLE dbo_tblSales ( - Category VARCHAR, - Market VARCHAR, - ProdType VARCHAR, - SalesPerson VARCHAR, - SalesValue DOUBLE, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_tblSales"."OrdDate") AS "year", ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") AS "month", ' - '"dbo_tblSales"."SalesPerson" AS "SalesPerson", ' - '"dbo_tblSales"."Market" AS "Market", ' - '"dbo_tblSales"."ProdType" AS "ProdType", ' - 'COUNT(*) AS "OrderCount" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."OrdDate" IS NOT NULL ' - 'AND "dbo_tblSales"."SalesPerson" IS NOT NULL ' - 'AND "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."ProdType" IS NOT NULL ' - 'AND "dbo_tblSales"."Category" = \'Backlog\' ' - 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), ' - '"dbo_tblSales"."SalesPerson", "dbo_tblSales"."Market", ' - '"dbo_tblSales"."ProdType" ' - 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate"), COUNT(*) DESC' - ) - - -def test_build_schema_grounded_sales_sql_for_top_new_order_detail_rows(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show the Top 20 New Orders for Period X, including Business Unit, " - "Market, Customer, Product, and Order Value.", - [ - """ - CREATE TABLE dbo_tblSales ( - BU VARCHAR, - Market VARCHAR, - Customer VARCHAR, - ProdName VARCHAR, - SalesValue DOUBLE, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 20 "dbo_tblSales"."BU" AS "BU", ' - '"dbo_tblSales"."Market" AS "Market", ' - '"dbo_tblSales"."ProdName" AS "ProdName", ' - '"dbo_tblSales"."Customer" AS "Customer", ' - '"dbo_tblSales"."SalesValue" AS "SalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."BU" IS NOT NULL ' - 'AND "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."ProdName" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' - 'ORDER BY "dbo_tblSales"."SalesValue" DESC' - ) - - -def test_build_schema_grounded_sales_sql_ignores_missing_metadata_entries(): - service = AskService.__new__(AskService) - - assert ( - service._build_schema_grounded_sales_sql( - "Show the Top 20 New Orders including Market and Customer.", - [ - None, - """ - CREATE TABLE dbo_tblSales ( - Market VARCHAR, - Customer VARCHAR, - SalesValue DOUBLE - ); - """, - ], - ) - == 'SELECT TOP 20 "dbo_tblSales"."Market" AS "Market", ' - '"dbo_tblSales"."Customer" AS "Customer", ' - '"dbo_tblSales"."SalesValue" AS "SalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' - 'ORDER BY "dbo_tblSales"."SalesValue" DESC' - ) - - -def test_build_schema_grounded_sales_sql_for_order_invoice_conversion_rate(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show Order-to-Invoice conversion rate by Month.", - [ - """ - CREATE TABLE dbo_tblSales ( - OrdNo VARCHAR, - InvoiceNo VARCHAR, - OrdDate TIMESTAMP, - SalesValue DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_tblSales"."OrdDate") AS "year", ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") AS "month", ' - 'COUNT(DISTINCT "dbo_tblSales"."OrdNo") AS "OrderCount", ' - 'COUNT(DISTINCT "dbo_tblSales"."InvoiceNo") AS "InvoiceCount", ' - '(COUNT(DISTINCT "dbo_tblSales"."InvoiceNo") * 100.0 / ' - 'NULLIF(COUNT(DISTINCT "dbo_tblSales"."OrdNo"), 0)) AS "ConversionRate" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."OrdNo" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate") ' - 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."OrdDate")' - ) - assert "P-M" not in sql - - -def test_build_schema_grounded_sales_sql_for_monthly_order_count_by_invdate(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show monthly order count by InvDate.", - [ - """ - CREATE TABLE dbo_tblSales ( - OrdNo VARCHAR, - InvDate TIMESTAMP, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_tblSales"."InvDate") AS "year", ' - 'DATEPART(MONTH, "dbo_tblSales"."InvDate") AS "month", ' - 'COUNT(*) AS "OrderCount" ' - 'FROM "dbo_tblSales" ' - 'GROUP BY DATEPART(YEAR, "dbo_tblSales"."InvDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."InvDate") ' - 'ORDER BY DATEPART(YEAR, "dbo_tblSales"."InvDate"), ' - 'DATEPART(MONTH, "dbo_tblSales"."InvDate")' - ) - - -def test_build_schema_grounded_sales_sql_counts_new_orders_by_customer_over_time(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show new orders by CustName over the last 12 months using dbo.XStageNewOrders OrdDate and CustName.", - [ - """ - CREATE TABLE dbo_XStageNewOrders ( - OrdNo VARCHAR, - OrdDate TIMESTAMP, - CustName VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate") AS "year", ' - 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate") AS "month", ' - '"dbo_XStageNewOrders"."CustName" AS "CustName", ' - 'COUNT(DISTINCT "dbo_XStageNewOrders"."OrdNo") AS "OrderCount" ' - 'FROM "dbo_XStageNewOrders" ' - 'WHERE "dbo_XStageNewOrders"."CustName" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate"), ' - '"dbo_XStageNewOrders"."CustName" ' - 'ORDER BY DATEPART(YEAR, "dbo_XStageNewOrders"."OrdDate"), ' - 'DATEPART(MONTH, "dbo_XStageNewOrders"."OrdDate")' - ) - - -def test_build_schema_grounded_sales_sql_for_highest_invoice_value(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Which Orders have the highest invoice value for by product and by customer", - [ - """ - CREATE TABLE dbo_tblSales ( - Customer VARCHAR, - ProdName VARCHAR, - SalesValue DOUBLE, - InvoiceNo VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblSales"."ProdName" AS "ProdName", ' - '"dbo_tblSales"."Customer" AS "Customer", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."ProdName" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' - 'GROUP BY "dbo_tblSales"."ProdName", "dbo_tblSales"."Customer" ' - 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' - ) - - -def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_country(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Which countries have the highest order revenue?", - [ - """ - CREATE TABLE dbo_tblSales ( - Country VARCHAR, - OrderValue DOUBLE, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblSales"."Country" AS "Country", ' - 'SUM("dbo_tblSales"."OrderValue") AS "TotalOrderValue" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Country" IS NOT NULL ' - 'GROUP BY "dbo_tblSales"."Country" ' - 'ORDER BY SUM("dbo_tblSales"."OrderValue") DESC' - ) - - -def test_build_schema_grounded_sales_sql_for_highest_order_revenue_by_prefixed_country(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Which countries have the highest order revenue?", - [ - """ - CREATE TABLE dbo_xStageLoad8 ( - col_07_Country VARCHAR, - TotalOrderValue DOUBLE, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_xStageLoad8"."col_07_Country" AS "col_07_Country", ' - 'SUM("dbo_xStageLoad8"."TotalOrderValue") AS "TotalTotalOrderValue" ' - 'FROM "dbo_xStageLoad8" ' - 'WHERE "dbo_xStageLoad8"."col_07_Country" IS NOT NULL ' - 'GROUP BY "dbo_xStageLoad8"."col_07_Country" ' - 'ORDER BY SUM("dbo_xStageLoad8"."TotalOrderValue") DESC' - ) - - -def test_build_schema_grounded_sales_sql_for_losing_order_value_by_market(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Which markets are losing order value?", - [ - """ - CREATE TABLE dbo_tblStageNewOrders ( - Market VARCHAR, - TotalOrderValue DOUBLE, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblStageNewOrders"."Market" AS "Market", ' - 'SUM("dbo_tblStageNewOrders"."TotalOrderValue") AS "TotalTotalOrderValue" ' - 'FROM "dbo_tblStageNewOrders" ' - 'WHERE "dbo_tblStageNewOrders"."Market" IS NOT NULL ' - 'GROUP BY "dbo_tblStageNewOrders"."Market" ' - 'ORDER BY SUM("dbo_tblStageNewOrders"."TotalOrderValue") ASC' - ) - - -def test_build_schema_grounded_sales_sql_for_highest_customers_each_market(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Which Customers have the highest New Orders in each Market?", - [ - """ - CREATE TABLE dbo_tblSales ( - Market VARCHAR, - Customer VARCHAR, - OrdNo VARCHAR, - SalesValue DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'WITH grouped_results AS (SELECT "dbo_tblSales"."Market" AS "Market", ' - '"dbo_tblSales"."Customer" AS "Customer", ' - 'COUNT(DISTINCT "dbo_tblSales"."OrdNo") AS "OrderCount" ' - 'FROM "dbo_tblSales" ' - 'WHERE "dbo_tblSales"."Market" IS NOT NULL ' - 'AND "dbo_tblSales"."Customer" IS NOT NULL ' - 'GROUP BY "dbo_tblSales"."Market", "dbo_tblSales"."Customer"), ' - 'ranked_results AS (SELECT "Market", "Customer", "OrderCount", ' - 'ROW_NUMBER() OVER (PARTITION BY "Market" ' - 'ORDER BY "OrderCount" DESC) AS "rank" ' - 'FROM grouped_results) ' - 'SELECT "Market", "Customer", "OrderCount" ' - 'FROM ranked_results ' - 'WHERE "rank" = 1 ' - 'ORDER BY "OrderCount" DESC' - ) - - -def test_build_schema_grounded_sales_sql_for_product_type_contribution(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Create a Product Type contribution pie chart.", - [ - """ - CREATE TABLE dbo_tblSales ( - ProdType VARCHAR, - SalesValue DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblSales"."ProdType" AS "ProdType", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'GROUP BY "dbo_tblSales"."ProdType" ' - 'ORDER BY SUM("dbo_tblSales"."SalesValue") DESC' - ) - - -def test_build_schema_grounded_sql_counts_categorical_status_values(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show refund status distribution.", - [ - """ - CREATE TABLE dbo_ytblRefund ( - Refund_Id VARCHAR, - Refund_Status VARCHAR, - CustomerName VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_ytblRefund"."Refund_Status" AS "Refund_Status", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ytblRefund" ' - 'WHERE "dbo_ytblRefund"."Refund_Status" IS NOT NULL ' - 'GROUP BY "dbo_ytblRefund"."Refund_Status" ' - 'ORDER BY COUNT(*) DESC' - ) - - -def test_build_schema_grounded_sql_counts_destination_databases_restored_most_often(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Which destination databases were restored most often?", - [ - """ - CREATE TABLE dbo_db_policies ( - id VARCHAR, - policy_name VARCHAR, - destination_phys_name VARCHAR, - restore_date TIMESTAMP, - restore_type VARCHAR, - status VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_db_policies"."destination_phys_name" AS ' - '"destination_phys_name", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_db_policies" ' - 'WHERE "dbo_db_policies"."destination_phys_name" IS NOT NULL ' - 'GROUP BY "dbo_db_policies"."destination_phys_name" ' - 'ORDER BY COUNT(*) DESC' - ) - assert "destination_database_name" not in sql - - -def test_build_schema_grounded_sales_sql_for_yoy_waterfall_dimensions(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show a waterfall of YOY changes by Customer, Product, and Market.", - [ - """ - CREATE TABLE dbo_tblSales ( - YearInd INTEGER, - Customer VARCHAR, - ProdName VARCHAR, - Market VARCHAR, - SalesValue DOUBLE - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tblSales"."YearInd" AS "year", ' - '"dbo_tblSales"."Customer" AS "Customer", ' - '"dbo_tblSales"."ProdName" AS "ProdName", ' - '"dbo_tblSales"."Market" AS "Market", ' - 'SUM("dbo_tblSales"."SalesValue") AS "TotalSalesValue" ' - 'FROM "dbo_tblSales" ' - 'GROUP BY "dbo_tblSales"."YearInd", "dbo_tblSales"."Customer", ' - '"dbo_tblSales"."ProdName", "dbo_tblSales"."Market" ' - 'ORDER BY "dbo_tblSales"."YearInd", SUM("dbo_tblSales"."SalesValue") DESC' - ) - - -def test_build_schema_grounded_sql_for_ticket_category_request_uses_existing_columns(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Create a bar chart of tickets by category.", - [ - """ - CREATE TABLE dbo_tickets ( - id VARCHAR, - org_id VARCHAR, - title VARCHAR, - description VARCHAR, - status VARCHAR, - priority VARCHAR, - assignee_user_id VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tickets"."status" AS "status", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_tickets" ' - 'GROUP BY "dbo_tickets"."status" ' - 'ORDER BY COUNT(*) DESC' - ) - assert "category" not in sql - - -def test_build_explicit_group_count_sql_for_schema_table_column_reference(): - service = AskService.__new__(AskService) - sql = service._build_explicit_group_count_sql( - "Show a pie chart grouped by dbo.tickets.status." - ) - - assert sql == ( - 'SELECT "dbo_tickets"."status" AS "status", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_tickets" ' - 'GROUP BY "dbo_tickets"."status" ' - 'ORDER BY COUNT(*) DESC' - ) - - -def test_build_schema_grounded_sql_for_knowledge_source_request_uses_existing_columns(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show knowledge article count by source.", - [ - """ - CREATE TABLE dbo_knowledge_articles ( - id VARCHAR, - org_id VARCHAR, - title VARCHAR, - category VARCHAR, - subcategory VARCHAR, - content VARCHAR, - author VARCHAR, - tags VARCHAR, - views INTEGER, - helpful INTEGER, - data VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_knowledge_articles"."author" AS "author", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_knowledge_articles" ' - 'GROUP BY "dbo_knowledge_articles"."author" ' - 'ORDER BY COUNT(*) DESC' - ) - assert "source" not in sql - - -def test_build_schema_grounded_sql_for_ticket_throughput_trend(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "Show throughput trends across different manufacturing units.", - [ - """ - CREATE TABLE dbo_tickets ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - assignee_user_id VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_tickets"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_tickets"."created_at") AS "month", ' - '"dbo_tickets"."assignee_user_id" AS "assignee_user_id", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_tickets" ' - 'GROUP BY DATEPART(YEAR, "dbo_tickets"."created_at"), ' - 'DATEPART(MONTH, "dbo_tickets"."created_at"), ' - '"dbo_tickets"."assignee_user_id" ' - 'ORDER BY DATEPART(YEAR, "dbo_tickets"."created_at"), ' - 'DATEPART(MONTH, "dbo_tickets"."created_at")' - ) - - -def test_build_manufacturing_throughput_sql_uses_active_unit_and_date_columns(): - service = AskService.__new__(AskService) - - sql = service._build_manufacturing_throughput_sql( - "Show throughput trends across different manufacturing units.", - [ - """ - CREATE TABLE dbo_production_events ( - id INTEGER, - manufacturing_unit VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_production_events"."manufacturing_unit" AS "manufacturing_unit", ' - 'DATEPART(YEAR, "dbo_production_events"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_production_events"."created_at") AS "month", ' - 'COUNT(*) AS "throughput" ' - 'FROM "dbo_production_events" ' - 'WHERE "dbo_production_events"."manufacturing_unit" IS NOT NULL ' - 'AND "dbo_production_events"."created_at" IS NOT NULL ' - 'GROUP BY "dbo_production_events"."manufacturing_unit", ' - 'DATEPART(YEAR, "dbo_production_events"."created_at"), ' - 'DATEPART(MONTH, "dbo_production_events"."created_at") ' - 'ORDER BY "dbo_production_events"."manufacturing_unit" ASC, ' - 'DATEPART(YEAR, "dbo_production_events"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_production_events"."created_at") ASC' - ) - - -def test_build_monthly_repair_volume_sql_uses_repair_log_date_column(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Generate a line chart showing monthly repair volume for the last 12 months.", - [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."created_at" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC' - ) - assert '"status"' not in sql - - -def test_build_monthly_repair_volume_sql_uses_debug_entry_date_column(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_sales_sql( - "Generate a line chart showing monthly repair volume for the last 12 months.", - [ - """ - CREATE TABLE dbo_DebugEntries ( - DebugEntryId VARCHAR, - Status VARCHAR, - DateIn TIMESTAMP - ); - """, - """ - CREATE TABLE dbo_DebugFixLogs ( - DebugEntryId VARCHAR, - FixId VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_DebugEntries"."DateIn") AS "year", ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") AS "month", ' - 'COUNT(*) AS "repair_count" ' - 'FROM "dbo_DebugEntries" ' - 'WHERE "dbo_DebugEntries"."DateIn" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn"), ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ' - 'ORDER BY DATEPART(YEAR, "dbo_DebugEntries"."DateIn") ASC, ' - 'DATEPART(MONTH, "dbo_DebugEntries"."DateIn") ASC' - ) - assert '"Status"' not in sql - - -def test_get_unqueryable_metric_message_for_throughput_without_unit_column(): - service = AskService.__new__(AskService) - - message = service._get_unqueryable_metric_message( - "Show throughput trends across different manufacturing units.", - [ - """ - CREATE TABLE dbo_failure_patterns ( - name VARCHAR, - occurrences INTEGER, - created_at TIMESTAMP - ); - """ - ], - ) - - assert message is not None - assert "unit" in message - - -def test_build_schema_grounded_sql_for_ticket_workflow_total_time_uses_first_class_columns(): - service = AskService.__new__(AskService) - sql = service._build_schema_grounded_sales_sql( - "What is the estimated total time for each workflow?", - [ - """ - CREATE TABLE dbo_tickets ( - id VARCHAR, - org_id VARCHAR, - title VARCHAR, - description VARCHAR, - status VARCHAR, - priority VARCHAR, - assignee_user_id VARCHAR, - data VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_tickets"."status" AS "workflow", ' - 'SUM(DATEDIFF(\'second\', "dbo_tickets"."created_at", ' - '"dbo_tickets"."updated_at")) AS "total_time_seconds" ' - 'FROM "dbo_tickets" ' - 'WHERE "dbo_tickets"."created_at" IS NOT NULL ' - 'AND "dbo_tickets"."updated_at" IS NOT NULL ' - 'AND "dbo_tickets"."status" IS NOT NULL ' - 'GROUP BY "dbo_tickets"."status" ' - 'ORDER BY "total_time_seconds" DESC' - ) - assert "data" not in sql - assert "JSON" not in sql - - -def test_build_audit_log_activity_sql_uses_existing_condition_columns(): - service = AskService.__new__(AskService) - sql = service._build_audit_log_activity_sql( - "Show audit log activity by condition name over time.", - [ - """ - CREATE TABLE dbo_audit_log ( - id VARCHAR, - action VARCHAR, - actor_name VARCHAR, - actor_user_id VARCHAR, - after_state VARCHAR, - before_state VARCHAR, - created_at TIMESTAMP, - entity_type VARCHAR, - is_name_condition BOOLEAN, - name VARCHAR - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_audit_log"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_audit_log"."created_at") AS "month", ' - '"dbo_audit_log"."is_name_condition" AS "is_name_condition", ' - 'COUNT(*) AS "activity_count" ' - 'FROM "dbo_audit_log" ' - 'WHERE "dbo_audit_log"."created_at" IS NOT NULL ' - 'AND "dbo_audit_log"."is_name_condition" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_audit_log"."created_at"), ' - 'DATEPART(MONTH, "dbo_audit_log"."created_at"), ' - '"dbo_audit_log"."is_name_condition" ' - 'ORDER BY DATEPART(YEAR, "dbo_audit_log"."created_at"), ' - 'DATEPART(MONTH, "dbo_audit_log"."created_at"), ' - '"activity_count" DESC' - ) - assert "condition_name" not in sql - assert "timestamp" not in sql - - -def test_build_validated_ask_result_from_sql_uses_local_schema_validation(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - 'SELECT "dbo_tblSales"."SalesPerson" FROM "dbo_tblSales"', - [ - """ - CREATE TABLE dbo_tblSales ( - SalesPerson VARCHAR, - SalesValue INTEGER - ); - """ - ], - ) - - assert result is not None - assert result.sql == 'SELECT "dbo_tblSales"."SalesPerson" FROM "dbo_tblSales"' - - -def test_build_validated_ask_result_from_sql_normalizes_column_case_to_schema(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - 'SELECT "dbo_tblFactSales"."TimeID" FROM "dbo_tblFactSales"', - [ - """ - CREATE TABLE dbo_tblFactSales ( - account VARCHAR, - timeid VARCHAR, - amount DOUBLE - ); - """ - ], - ) - - assert result is not None - assert result.sql == 'SELECT "dbo_tblFactSales"."timeid" FROM "dbo_tblFactSales"' - - -def test_build_validated_ask_result_from_sql_normalizes_customer_to_account(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - 'SELECT "dbo_tblFactSales"."customer" FROM "dbo_tblFactSales"', - [ - """ - CREATE TABLE dbo_tblFactSales ( - account VARCHAR, - customerpo VARCHAR, - timeid VARCHAR, - amount DOUBLE - ); - """ - ], - ) - - assert result is not None - assert result.sql == 'SELECT "dbo_tblFactSales"."account" FROM "dbo_tblFactSales"' - - -def test_build_validated_ask_result_from_sql_normalizes_active_table_reference(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "public"."dbo_failure"."created_at" ' - 'FROM "public"."dbo_failure"' - ), - [ - """ - CREATE TABLE dbo_failure_patterns ( - id VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP - ); - """ - ], - "Show monthly record count by created_at in dbo.failure_patterns", - ) - - assert result is not None - assert result.sql == ( - 'SELECT "dbo_failure_patterns"."created_at" ' - 'FROM "dbo_failure_patterns"' - ) - - -def test_build_schema_grounded_operational_sql_prefers_failure_column_for_error_rate(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_operational_sql( - "What is the error rate in the repair data collection system for different types of failures?", - [ - { - "name": "dbo_repair_logs", - "columns": [ - {"name": "status", "type": "varchar"}, - {"name": "failure_code", "type": "varchar"}, - {"name": "created_at", "type": "timestamp"}, - ], - } - ], - ) - - assert sql is not None - assert '"dbo_repair_logs"."failure_code" AS "failure_code"' in sql - assert '"dbo_repair_logs"."status" AS "status"' not in sql - - -def test_build_validated_ask_result_rejects_status_when_failure_field_matches_question(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_repair_logs"."status" AS "status", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_repair_logs" ' - 'GROUP BY "dbo_repair_logs"."status"' - ), - [ - """ - CREATE TABLE dbo_repair_logs ( - status VARCHAR, - failure_code VARCHAR, - created_at TIMESTAMP - ); - """ - ], - "What is the error rate for different types of failures?", - ) - - assert result is None - - -def test_build_validated_ask_result_rejects_sql_for_wrong_explicit_table(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_qMarginSales"."Customer" AS "Customer", ' - 'COUNT(DISTINCT "dbo_qMarginSales"."OrdNo") AS "OrderCount" ' - 'FROM "dbo_qMarginSales" ' - 'WHERE "dbo_qMarginSales"."Customer" IS NOT NULL ' - 'GROUP BY "dbo_qMarginSales"."Customer"' - ), - [ - """ - CREATE TABLE dbo_tblNewOrders ( - Customer VARCHAR, - OrdNo VARCHAR - ); - """, - """ - CREATE TABLE dbo_qMarginSales ( - Customer VARCHAR, - OrdNo VARCHAR - ); - """, - ], - "Which customers have the highest number of orders in dbo.tblNewOrders?", - ) - - assert result is None - - -def test_reusable_historical_question_allows_exact_recommended_question(): - assert AskService._is_reusable_historical_question( - "How many tickets are currently open?", - "How many tickets are currently open?", - ) - - -def test_reusable_historical_question_rejects_materially_different_agent_question(): - assert not AskService._is_reusable_historical_question( - "What is the estimated total time for each workflow?", - "How many tickets are currently open?", - ) - - -def test_reusable_historical_question_rejects_similar_but_different_failure_question(): - assert not AskService._is_reusable_historical_question( - "Which name values have the highest occurrences in dbo.failure_patterns?", - "What is the distribution of name in dbo.failure_patterns?", - ) - - -def test_should_not_use_histories_for_independent_same_thread_question(): - assert not AskService._should_use_histories_for_query( - "Which name values have the highest occurrences in dbo.failure_patterns?" - ) - assert not AskService._should_use_histories_for_query( - "Show monthly record count by created_at in dbo.failure_patterns" - ) - - -def test_should_use_histories_for_contextual_followup_question(): - assert AskService._should_use_histories_for_query("What about by month?") - assert AskService._should_use_histories_for_query("Show the same for last year") - - -def test_build_schema_grounded_table_question_sql_for_record_count(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "How many records are in dbo.failure_patterns?", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == 'SELECT COUNT(*) AS "RecordCount" FROM "dbo_failure_patterns"' - - -def test_build_schema_grounded_table_question_sql_matches_dot_table_to_underscore_table(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "How many records are in dbo.ytblTarrifsRec?", - [ - """ - CREATE TABLE dbo_ytblTarrifsRec ( - EntrySummaryNumber2 VARCHAR, - LiquidationStatus VARCHAR, - LiquidationDate TIMESTAMP - ); - """ - ], - ) - - assert sql == 'SELECT COUNT(*) AS "RecordCount" FROM "dbo_ytblTarrifsRec"' - - -def test_build_schema_grounded_table_question_sql_for_name_distribution(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "What is the distribution of name in dbo.failure_patterns?", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_failure_patterns"."name" AS "name", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."name" IS NOT NULL ' - 'GROUP BY "dbo_failure_patterns"."name" ' - 'ORDER BY COUNT(*) DESC' - ) - - -def test_build_schema_grounded_table_question_sql_for_repair_log_failure_code_count(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Count repair logs by failure_code in repair logs.", - [ - """ - CREATE TABLE dbo_repair_logs ( - org_id VARCHAR, - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP, - data JSON - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_repair_logs"."failure_code" AS "failure_code", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."failure_code" IS NOT NULL ' - 'GROUP BY "dbo_repair_logs"."failure_code" ' - 'ORDER BY COUNT(*) DESC' - ) - - -def test_build_schema_grounded_table_question_sql_for_highest_orders_by_customer(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Which customers have the highest number of orders in dbo.tblNewOrders?", - [ - """ - CREATE TABLE dbo_tblNewOrders ( - Customer VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_tblNewOrders"."Customer" AS "Customer", ' - 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") AS "RecordCount" ' - 'FROM "dbo_tblNewOrders" ' - 'WHERE "dbo_tblNewOrders"."Customer" IS NOT NULL ' - 'AND LTRIM(RTRIM("dbo_tblNewOrders"."Customer")) <> \'\' ' - 'GROUP BY "dbo_tblNewOrders"."Customer" ' - 'ORDER BY COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo") DESC' - ) - - -def test_schema_grounded_analytics_prefers_order_table_for_order_count_question(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Which customers have the highest number of orders?", - [ - """ - CREATE TABLE dbo_qMarginSales ( - Customer VARCHAR, - OrdNo VARCHAR, - SalesValue DOUBLE - ); - """, - """ - CREATE TABLE dbo_tblNewOrders ( - Customer VARCHAR, - OrdNo VARCHAR, - OrdDate TIMESTAMP - ); - """, - ], - ) - - assert sql is not None - assert 'FROM "dbo_tblNewOrders"' in sql - assert 'FROM "dbo_qMarginSales"' not in sql - assert 'COUNT(DISTINCT "dbo_tblNewOrders"."OrdNo")' in sql - - -def test_build_pcb_direct_question_sql_for_board_model_distribution_over_time(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "How many different board models are present in the dbo.repair_logs table, and what is their distribution over time?", - [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_repair_logs"."board_model" AS "board_model", ' - 'DATEPART(YEAR, "dbo_repair_logs"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") AS "month", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."board_model" IS NOT NULL ' - 'AND "dbo_repair_logs"."created_at" IS NOT NULL ' - 'GROUP BY "dbo_repair_logs"."board_model", ' - 'DATEPART(YEAR, "dbo_repair_logs"."created_at"), ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_repair_logs"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_repair_logs"."created_at") ASC, ' - '"dbo_repair_logs"."board_model" ASC' - ) - - -def test_build_pcb_direct_question_sql_for_recurring_pcb_failures_by_product(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Show recurring PCB failures by product.", - [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_repair_logs"."board_model" AS "board_model", ' - '"dbo_repair_logs"."failure_code" AS "failure_code", ' - 'COUNT(*) AS "failure_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."board_model" IS NOT NULL ' - 'AND "dbo_repair_logs"."failure_code" IS NOT NULL ' - 'GROUP BY "dbo_repair_logs"."board_model", ' - '"dbo_repair_logs"."failure_code" ' - 'ORDER BY "failure_count" DESC' - ) - - -def test_build_pcb_direct_question_sql_for_highest_priority_repairs(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Which repair logs have the highest priority?", - [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql is not None - assert sql.startswith('SELECT TOP 10 "dbo_repair_logs"."id" AS "id"') - assert 'FROM "dbo_repair_logs"' in sql - assert 'CASE LOWER("dbo_repair_logs"."priority")' in sql - assert "WHEN 'critical' THEN 1" in sql - assert "WHEN 'high' THEN 2" in sql - - -def test_build_pcb_direct_question_sql_for_repair_ticket_distribution(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Show repair ticket again distribution.", - [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT "dbo_repair_logs"."status" AS "status", ' - 'COUNT(*) AS "ticket_count" ' - 'FROM "dbo_repair_logs" ' - 'WHERE "dbo_repair_logs"."status" IS NOT NULL ' - 'GROUP BY "dbo_repair_logs"."status" ' - 'ORDER BY "ticket_count" DESC' - ) - - -def test_build_pcb_direct_question_sql_for_top_ticket_labels(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_analytics_sql( - "Display top 10 ticket labels.", - [ - """ - CREATE TABLE dbo_ticket_labels ( - id VARCHAR, - name VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_ticket_labels"."name" AS "name", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ticket_labels" ' - 'WHERE "dbo_ticket_labels"."name" IS NOT NULL ' - 'GROUP BY "dbo_ticket_labels"."name" ' - 'ORDER BY COUNT(*) DESC' - ) - - -def test_build_schema_grounded_table_question_sql_for_numeric_column_distribution(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "What is the distribution of EntrySummaryNumber2 in dbo.ytblTarrifsRec?", - [ - """ - CREATE TABLE dbo_ytblTarrifsRec ( - EntrySummaryNumber2 BIGINT, - LiquidationStatus VARCHAR, - LiquidationDate TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_ytblTarrifsRec"."EntrySummaryNumber2" AS ' - '"EntrySummaryNumber2", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_ytblTarrifsRec" ' - 'WHERE "dbo_ytblTarrifsRec"."EntrySummaryNumber2" IS NOT NULL ' - 'GROUP BY "dbo_ytblTarrifsRec"."EntrySummaryNumber2" ' - 'ORDER BY COUNT(*) DESC' - ) - - -def test_build_schema_grounded_table_question_sql_for_highest_occurrences(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Which name values have the highest occurrences in dbo.failure_patterns?", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR, - occurrences INTEGER, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_failure_patterns"."name" AS "name", ' - '"dbo_failure_patterns"."occurrences" AS "occurrences" ' - 'FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."name" IS NOT NULL ' - 'ORDER BY "dbo_failure_patterns"."occurrences" DESC' - ) - - -def test_build_schema_grounded_table_question_sql_for_plural_names_by_occurrences(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Show top 10 names by occurrences in dbo.failure_patterns.", - [ - """ - CREATE TABLE dbo_failure_patterns ( - name VARCHAR, - occurrences INTEGER, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 "dbo_failure_patterns"."name" AS "name", ' - '"dbo_failure_patterns"."occurrences" AS "occurrences" ' - 'FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."name" IS NOT NULL ' - 'ORDER BY "dbo_failure_patterns"."occurrences" DESC' - ) - - -def test_build_schema_grounded_table_question_sql_for_latest_records_by_created_at(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Show the latest records from dbo.failure_patterns by created_at.", - [ - """ - CREATE TABLE dbo_failure_patterns ( - name VARCHAR, - occurrences INTEGER, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT TOP 10 * FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' - 'ORDER BY "dbo_failure_patterns"."created_at" DESC' - ) - - -def test_build_schema_grounded_table_question_sql_for_monthly_created_at_count(): - service = AskService.__new__(AskService) - - sql = service._build_schema_grounded_table_question_sql( - "Show monthly record count by created_at in dbo.failure_patterns", - [ - """ - CREATE TABLE dbo_failure_patterns ( - id INTEGER, - name VARCHAR, - created_at TIMESTAMP - ); - """ - ], - ) - - assert sql == ( - 'SELECT DATEPART(YEAR, "dbo_failure_patterns"."created_at") AS "year", ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") AS "month", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_failure_patterns" ' - 'WHERE "dbo_failure_patterns"."created_at" IS NOT NULL ' - 'GROUP BY DATEPART(YEAR, "dbo_failure_patterns"."created_at"), ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ' - 'ORDER BY DATEPART(YEAR, "dbo_failure_patterns"."created_at") ASC, ' - 'DATEPART(MONTH, "dbo_failure_patterns"."created_at") ASC' - ) - - -def test_build_validated_ask_result_rejects_status_for_product_line_pcb_question(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_repair_logs"."status" AS "status" ' - 'FROM "dbo_repair_logs"' - ), - [ - """ - CREATE TABLE dbo_repair_logs ( - status VARCHAR, - product_line VARCHAR, - pcb_issue VARCHAR, - created_at TIMESTAMP - ); - """ - ], - "Create a visualization of recurring PCB issues by product line.", - ) - - assert result is None - - -def test_build_validated_ask_result_rejects_status_for_repair_cost_question(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_repair_logs"."status" AS "status" ' - 'FROM "dbo_repair_logs"' - ), - [ - """ - CREATE TABLE dbo_repair_logs ( - status VARCHAR, - repair_cost DOUBLE, - created_at TIMESTAMP - ); - """ - ], - "Generate a quarterly repair cost analysis chart.", - ) - - assert result is None - - -def test_build_validated_ask_result_rejects_status_for_critical_repairs_question(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_repair_logs"."status" AS "status" ' - 'FROM "dbo_repair_logs"' - ), - [ - """ - CREATE TABLE dbo_repair_logs ( - status VARCHAR, - severity VARCHAR, - repair_id INTEGER - ); - """ - ], - "Create a stacked bar chart comparing critical vs non-critical repairs.", - ) - - assert result is None - - -def test_build_validated_ask_result_accepts_product_line_pcb_question_sql(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT "dbo_repair_logs"."product_line" AS "product_line", ' - '"dbo_repair_logs"."pcb_issue" AS "pcb_issue", ' - 'COUNT(*) AS "RecordCount" ' - 'FROM "dbo_repair_logs" ' - 'GROUP BY "dbo_repair_logs"."product_line", ' - '"dbo_repair_logs"."pcb_issue"' - ), - [ - """ - CREATE TABLE dbo_repair_logs ( - product_line VARCHAR, - pcb_issue VARCHAR - ); - """ - ], - "Create a visualization of recurring PCB issues by product line.", - ) - - assert result is not None - - -def test_build_validated_ask_result_rejects_unqualified_invalid_columns(): - service = AskService.__new__(AskService) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - created_at TIMESTAMP, - updated_at TIMESTAMP, - status VARCHAR, - name VARCHAR, - occurrences INTEGER - ); - """ - ] - - invalid_sqls = [ - 'SELECT DATEPART(MONTH, execution_date) AS "month", COUNT(*) AS "RecordCount" FROM "dbo_repair_logs" GROUP BY DATEPART(MONTH, execution_date)', - 'SELECT created_by AS "created_by", COUNT(*) AS "RecordCount" FROM "dbo_repair_logs" GROUP BY created_by', - 'SELECT physical_name AS "physical_name", occurrences FROM "dbo_repair_logs" ORDER BY occurrences DESC', - 'SELECT name AS "created_by" FROM "dbo_repair_logs" ORDER BY created_by', - ] - - for sql in invalid_sqls: - assert ( - service._build_validated_ask_result_from_sql( - sql, - table_ddls, - "Show top 10 failure pattern names by occurrences.", - ) - is None - ) - - -def test_build_validated_ask_result_accepts_unqualified_valid_columns(): - service = AskService.__new__(AskService) - - result = service._build_validated_ask_result_from_sql( - ( - 'SELECT name AS "name", occurrences AS "occurrences" ' - 'FROM "dbo_repair_logs" ' - 'ORDER BY occurrences DESC' - ), - [ - """ - CREATE TABLE dbo_repair_logs ( - name VARCHAR, - occurrences INTEGER - ); - """ - ], - "Show top 10 failure pattern names by occurrences.", - ) - - assert result is not None diff --git a/wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py b/wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py deleted file mode 100644 index e13af38fc6..0000000000 --- a/wren-ai-service/tests/pytest/services/test_ask_unqueryable_metrics.py +++ /dev/null @@ -1,102 +0,0 @@ -from src.web.v1.services.ask import AskService - - -def test_first_pass_yield_requires_queryable_attempt_fields(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - org_id VARCHAR, - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP, - data JSON - ); - """ - ] - - message = service._get_unqueryable_metric_message( - "Show First Pass Yield percentage trend over time.", - table_ddls, - ) - - assert message - assert "first-pass yield" in message.lower() - assert "first-class columns" in message - - -def test_first_pass_yield_guard_allows_queryable_attempt_fields(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - created_at TIMESTAMP, - attempt_number INTEGER, - pass_fail VARCHAR - ); - """ - ] - - assert ( - service._get_unqueryable_metric_message( - "Show FPY trend over time.", - table_ddls, - ) - is None - ) - - -def test_repair_cost_requires_queryable_cost_field(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - org_id VARCHAR, - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP, - updated_at TIMESTAMP, - data JSON - ); - """ - ] - - message = service._get_unqueryable_metric_message( - "Create a line chart comparing repair cost and turnaround time.", - table_ddls, - ) - - assert message - assert "repair cost" in message.lower() - assert "first-class column" in message - assert "JSON/text" in message - - -def test_repair_cost_guard_allows_queryable_cost_field(): - service = AskService(pipelines={}) - table_ddls = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - repair_cost DOUBLE, - created_at TIMESTAMP, - updated_at TIMESTAMP - ); - """ - ] - - assert ( - service._get_unqueryable_metric_message( - "Create a line chart comparing repair cost and turnaround time.", - table_ddls, - ) - is None - ) From 2621dd9691c85a30d30a30e962c6d496c754725b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 15:05:35 +0530 Subject: [PATCH 0702/1087] Scope ask retrieval by project --- .../src/web/v1/services/__init__.py | 5 ++- .../adaptors/tests/wrenAIAdaptor.test.ts | 31 +++++++++++++++++++ .../apollo/server/adaptors/wrenAIAdaptor.ts | 1 + wren-ui/src/apollo/server/models/adaptor.ts | 1 + .../apollo/server/services/askingService.ts | 1 + 5 files changed, 38 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/web/v1/services/__init__.py b/wren-ai-service/src/web/v1/services/__init__.py index 250cffb2dc..5296b274b4 100644 --- a/wren-ai-service/src/web/v1/services/__init__.py +++ b/wren-ai-service/src/web/v1/services/__init__.py @@ -57,7 +57,10 @@ def serialize(self): # for POST, PATCH, UPDATE, DELETE requests class BaseRequest(BaseModel): query_id: Optional[str] = Field(default=None, exclude=True) - project_id: Optional[str] = None + project_id: Optional[str] = Field( + default=None, + validation_alias=AliasChoices("project_id", "projectId"), + ) thread_id: Optional[str] = None configurations: Configuration = Field( default_factory=Configuration, diff --git a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts index 69252623ff..92bcbb0387 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts @@ -1,6 +1,7 @@ import axios from 'axios'; import { WrenAIAdaptor } from '../wrenAIAdaptor'; import { + AskInput, RecommendationQuestionsInput, RecommendationQuestionStatus, } from '@server/models/adaptor'; @@ -33,6 +34,36 @@ describe('WrenAIAdaptor', () => { jest.clearAllMocks(); }); + describe('ask', () => { + const mockInput: AskInput = { + query: 'Show active records', + deployId: 'deploy-hash', + projectId: 'project-123', + histories: [], + configurations: { + language: 'English', + }, + }; + + it('should send the project id to scope AI retrieval', async () => { + const mockQueryId = 'query-123'; + mockedAxios.post.mockResolvedValueOnce({ + data: { query_id: mockQueryId }, + }); + + const result = await adaptor.ask(mockInput); + + expect(result).toEqual({ queryId: mockQueryId }); + expect(mockedAxios.post).toHaveBeenCalledWith(`${baseEndpoint}/v1/asks`, { + query: mockInput.query, + id: mockInput.deployId, + project_id: mockInput.projectId, + histories: [], + configurations: mockInput.configurations, + }); + }); + }); + describe('generateRecommendationQuestions', () => { const mockInput: RecommendationQuestionsInput = { manifest: sampleManifest, diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 4c691367c7..76407ef301 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -250,6 +250,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, { query: input.query, id: input.deployId, + project_id: input.projectId, histories: this.transformHistoryInput(input.histories), configurations: input.configurations, }); diff --git a/wren-ui/src/apollo/server/models/adaptor.ts b/wren-ui/src/apollo/server/models/adaptor.ts index 102ed91a5a..75b619fe05 100644 --- a/wren-ui/src/apollo/server/models/adaptor.ts +++ b/wren-ui/src/apollo/server/models/adaptor.ts @@ -73,6 +73,7 @@ export interface ProjectConfigurations { export interface AskInput { query: string; deployId: string; + projectId?: string; histories?: ThreadResponse[]; configurations?: ProjectConfigurations; } diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 33736f87b7..4993320ec0 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -698,6 +698,7 @@ export class AskingService implements IAskingService { query: input.question, histories, deployId, + projectId: projectId.toString(), configurations: { language }, rerunFromCancelled, previousTaskId, From 7f4c5373a4524b2c58f8b6ca5b89522afca7f183 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 15:20:43 +0530 Subject: [PATCH 0703/1087] Scope semantic indexing by project --- .../adaptors/tests/wrenAIAdaptor.test.ts | 27 +++++++++++++++++++ .../apollo/server/adaptors/wrenAIAdaptor.ts | 3 ++- wren-ui/src/apollo/server/models/adaptor.ts | 1 + .../apollo/server/services/deployService.ts | 1 + .../services/tests/deployService.test.ts | 5 ++++ 5 files changed, 36 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts index 92bcbb0387..36bdf9be4a 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts @@ -34,6 +34,33 @@ describe('WrenAIAdaptor', () => { jest.clearAllMocks(); }); + describe('deploy', () => { + it('should send the project id to scope indexed semantics', async () => { + const mockInput = { + manifest: sampleManifest, + hash: 'deploy-hash', + projectId: 123, + }; + mockedAxios.post.mockResolvedValueOnce({ + data: { id: mockInput.hash }, + }); + mockedAxios.get.mockResolvedValueOnce({ + data: { status: 'finished' }, + }); + + await adaptor.deploy(mockInput); + + expect(mockedAxios.post).toHaveBeenCalledWith( + `${baseEndpoint}/v1/semantics-preparations`, + { + mdl: JSON.stringify(mockInput.manifest), + id: mockInput.hash, + project_id: mockInput.projectId.toString(), + }, + ); + }); + }); + describe('ask', () => { const mockInput: AskInput = { query: 'Show active records', diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 76407ef301..5acc9af772 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -347,13 +347,14 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } public async deploy(deployData: DeployData): Promise { - const { manifest, hash } = deployData; + const { manifest, hash, projectId } = deployData; try { const res = await axios.post( `${this.wrenAIBaseEndpoint}/v1/semantics-preparations`, { mdl: JSON.stringify(manifest), id: hash, + project_id: projectId.toString(), }, ); const deployId = res.data.id; diff --git a/wren-ui/src/apollo/server/models/adaptor.ts b/wren-ui/src/apollo/server/models/adaptor.ts index 75b619fe05..0a60e14ebe 100644 --- a/wren-ui/src/apollo/server/models/adaptor.ts +++ b/wren-ui/src/apollo/server/models/adaptor.ts @@ -51,6 +51,7 @@ export enum WrenAILanguage { export interface DeployData { manifest: Manifest; hash: string; + projectId: number; } // ask diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 45f465bdd0..746e5fa12e 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -138,6 +138,7 @@ export class DeployService implements IDeployService { await this.wrenAIAdaptor.deploy({ manifest, hash, + projectId, }); // update deploy status diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 8788ef3f57..ce60461bf7 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -36,6 +36,11 @@ describe('DeployService', () => { const response = await deployService.deploy(manifest, projectId); expect(response.status).toEqual(DeployStatusEnum.SUCCESS); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash: deployService.createMDLHash(manifest, projectId), + projectId, + }); expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { status: DeployStatusEnum.SUCCESS, error: undefined, From d9fadec6df46392d17cee47e90039425dfc6a265 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 15:49:23 +0530 Subject: [PATCH 0704/1087] Enrich semantic context for SQL grounding --- .../src/pipelines/generation/utils/sql.py | 10 ++- .../pipelines/indexing/table_description.py | 79 ++++++++++++++++++- .../pipelines/generation/test_sql_utils.py | 4 + .../indexing/test_table_description.py | 37 ++++++++- 4 files changed, 123 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..2dd952690c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -169,8 +169,13 @@ async def _classify_generation_result( - ONLY USE the tables and columns mentioned in the database schema. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. +- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. +- Comments, aliases, display labels, and descriptions are only semantic context. Do not use them as executable table or column identifiers, except as aliases in the final SELECT clause. +- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. +- If a requested concept, filter, sort, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! +- When using multiple tables, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. @@ -444,7 +449,10 @@ async def _classify_generation_result( 11. Do not include ```markdown or ``` in the answer. 12. A table name in the reasoning plan must be in this format: `table: `. 13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +14. Use only exact table and column names that appear in the DATABASE SCHEMA section. +15. Comments, aliases, display labels, and descriptions are semantic hints only; do not turn them into table or column names in the reasoning plan. +16. If the question asks for a concept such as recent, latest, status, amount, customer, supplier, order, payment, or market, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. +17. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 4d013fdfc1..a7f5652d1f 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -54,6 +54,9 @@ def _additional_meta() -> Dict[str, Any]: } def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[Dict[str, Any]]: + def _text(value: Any) -> str: + return "" if value is None else str(value) + def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: properties = self._properties(payload) @@ -61,25 +64,93 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: "mdl_type": mdl_type, "name": payload.get("name"), "columns": [ - column.get("name", "") or "" + { + "name": _text(column.get("name", "")), + "type": _text(column.get("type", "")), + "description": _text( + self._properties(column).get("description", "") + ), + "displayName": _text( + self._properties(column).get("displayName", "") + ), + } for column in payload.get("columns", []) if isinstance(column, dict) ], "properties": properties, } + def _relationship_context_by_model() -> Dict[str, List[str]]: + relationships = {model.get("name"): [] for model in mdl.get("models", [])} + + for relationship in mdl.get("relationships", []) or []: + models = relationship.get("models", []) + if len(models) != 2: + continue + + summary = " ".join( + part + for part in [ + _text(relationship.get("name", "")), + _text(relationship.get("joinType", "")), + _text(relationship.get("condition", "")), + ] + if part + ) + if not summary: + continue + + for model_name in models: + relationships.setdefault(model_name, []).append(summary) + + return relationships + + def _column_context(columns: List[Dict[str, Any]]) -> str: + details = [] + + for column in columns: + semantic_parts = [ + column["type"], + column["displayName"], + column["description"], + ] + if not any(semantic_parts): + continue + + details.append( + " ".join( + part for part in [column["name"], *semantic_parts] if part + ) + ) + + return "; ".join(detail for detail in details if detail) + + relationship_context = _relationship_context_by_model() resources = ( [_structure_data("MODEL", model) for model in mdl["models"]] + [_structure_data("METRIC", metric) for metric in mdl["metrics"]] + [_structure_data("VIEW", view) for view in mdl["views"]] ) - return [ - { + def _resource_description(resource: Dict[str, Any]) -> Dict[str, str]: + description = { "name": resource["name"], "description": resource["properties"].get("description", "") or "", - "columns": ", ".join(resource["columns"]), + "columns": ", ".join( + column["name"] for column in resource["columns"] + ), } + + if column_context := _column_context(resource["columns"]): + description["column_context"] = column_context + + if relationships := "; ".join(relationship_context.get(resource["name"], [])): + description["relationships"] = relationships + + return description + + return [ + _resource_description(resource) for resource in resources if resource["name"] is not None ] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 5d51428c4e..db197a7833 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -24,6 +24,9 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "ONLY USE the tables and columns mentioned in the database schema" in rules assert 'ONLY USE "*" if the user query asks for all the columns' in rules + assert "Do not use them as executable table or column identifiers" in rules + assert "do not invent a field" in rules + assert "join only through the FOREIGN KEY relationships shown" in rules def test_get_text_to_sql_rules_uses_sql_knowledge_override(): @@ -46,4 +49,5 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "ONLY USE table/column alias in the final SELECT clause" in prompt assert "Refer to the value of alias from the comment section" in prompt + assert "source of executable table and column identifiers" in prompt assert 'SELECT "_orders"."ApprovedTimestamp" AS "_timestamp"' in prompt diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index 214ec27b82..8898ced29d 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -151,7 +151,7 @@ def test_table_description_null_description(): ) -def test_table_description_excludes_generated_column_descriptions(): +def test_table_description_includes_column_semantic_context(): chunker = TableDescriptionChunker() mdl = { "models": [ @@ -189,9 +189,42 @@ def test_table_description_excludes_generated_column_descriptions(): "name": "orders", "description": "Customer purchase transactions.", "columns": "Division, SalesAmount", + "column_context": ( + "Division varchar Generic generated division description.; " + "SalesAmount float Generic generated sales amount description." + ), } ) - assert "Generic generated" not in document.content + + +def test_table_description_includes_relationship_context(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + {"name": "source", "columns": [{"name": "source_id"}]}, + {"name": "target", "columns": [{"name": "source_id"}]}, + ], + "views": [], + "relationships": [ + { + "name": "source_to_target", + "models": ["source", "target"], + "joinType": "ONE_TO_MANY", + "condition": "source.source_id = target.source_id", + } + ], + "metrics": [], + } + + actual = chunker.run(mdl) + + assert len(actual["documents"]) == 2 + for document in actual["documents"]: + assert document.meta["type"] == "TABLE_DESCRIPTION" + assert ( + "source_to_target ONE_TO_MANY source.source_id = target.source_id" + in document.content + ) def test_table_description_keeps_complete_column_lists(): From 317b1d4e6bd41986b9d1c2695b7e39a57f470377 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 16:17:31 +0530 Subject: [PATCH 0705/1087] Strengthen SQL grounding with metadata semantics --- .../src/pipelines/generation/utils/sql.py | 28 ++++-- .../pipelines/indexing/table_description.py | 47 +++++++++- .../pipelines/generation/test_sql_utils.py | 10 +- .../indexing/test_table_description.py | 92 +++++++++++-------- 4 files changed, 126 insertions(+), 51 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2dd952690c..a26dedf0f1 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -163,19 +163,26 @@ async def _classify_generation_result( return valid_generation_result, invalid_generation_result +_MANDATORY_SQL_GROUNDING_RULES = """ +### MANDATORY SQL GROUNDING RULES ### +- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. +- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. +- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. Do not use them as executable table or column identifiers, except as aliases in the final SELECT clause. +- Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. +- If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. +- When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. +- When using multiple tables, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +""" + + _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. - ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. -- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. -- Comments, aliases, display labels, and descriptions are only semantic context. Do not use them as executable table or column identifiers, except as aliases in the final SELECT clause. -- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. -- If a requested concept, filter, sort, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! -- When using multiple tables, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. @@ -451,8 +458,10 @@ async def _classify_generation_result( 13. A column name in the reasoning plan must be in this format: `column: .`. 14. Use only exact table and column names that appear in the DATABASE SCHEMA section. 15. Comments, aliases, display labels, and descriptions are semantic hints only; do not turn them into table or column names in the reasoning plan. -16. If the question asks for a concept such as recent, latest, status, amount, customer, supplier, order, payment, or market, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. -17. ONLY SHOWING the reasoning plan in bullet points. +16. Do not write SQL, possible SQL, sample SQL, or assumed SQL in the reasoning plan. +17. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the schema does not show the exact table or column needed, state that the available schema does not include that part. +18. If the question asks for a concept, filter, sort, or timeframe, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. +19. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -470,12 +479,13 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: + rules = _DEFAULT_TEXT_TO_SQL_RULES if sql_knowledge is not None: - return _extract_from_sql_knowledge( + rules = _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) - return _DEFAULT_TEXT_TO_SQL_RULES + return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index a7f5652d1f..412fc5c6ef 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -57,16 +57,51 @@ def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[Dict[str, Any]]: def _text(value: Any) -> str: return "" if value is None else str(value) + def _source_context(payload: Dict[str, Any]) -> str: + table_reference = payload.get("tableReference") + if isinstance(table_reference, dict): + reference_parts = [ + _text(table_reference.get("catalog", "")), + _text(table_reference.get("schema", "")), + _text(table_reference.get("table", "")), + ] + return ".".join(part for part in reference_parts if part) + + return _text(payload.get("baseObject", "")) + + def _columns(payload: Dict[str, Any]) -> List[Dict[str, Any]]: + columns = payload.get("columns", []) + if columns: + return [ + {**column, "role": _text(column.get("role", ""))} + for column in columns + if isinstance(column, dict) + ] + + metric_columns = [] + for role, key in [("dimension", "dimension"), ("measure", "measure")]: + metric_columns += [ + {**column, "role": role} + for column in payload.get(key, []) or [] + if isinstance(column, dict) + ] + + return metric_columns + def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: properties = self._properties(payload) return { "mdl_type": mdl_type, "name": payload.get("name"), + "displayName": _text(properties.get("displayName", "")), + "source": _source_context(payload), "columns": [ { "name": _text(column.get("name", "")), "type": _text(column.get("type", "")), + "role": _text(column.get("role", "")), + "expression": _text(column.get("expression", "")), "description": _text( self._properties(column).get("description", "") ), @@ -74,8 +109,7 @@ def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: self._properties(column).get("displayName", "") ), } - for column in payload.get("columns", []) - if isinstance(column, dict) + for column in _columns(payload) ], "properties": properties, } @@ -111,8 +145,10 @@ def _column_context(columns: List[Dict[str, Any]]) -> str: for column in columns: semantic_parts = [ column["type"], + column["role"], column["displayName"], column["description"], + column["expression"], ] if not any(semantic_parts): continue @@ -135,12 +171,19 @@ def _column_context(columns: List[Dict[str, Any]]) -> str: def _resource_description(resource: Dict[str, Any]) -> Dict[str, str]: description = { "name": resource["name"], + "resource_type": resource["mdl_type"], "description": resource["properties"].get("description", "") or "", "columns": ", ".join( column["name"] for column in resource["columns"] ), } + if resource["displayName"]: + description["displayName"] = resource["displayName"] + + if resource["source"]: + description["source"] = resource["source"] + if column_context := _column_context(resource["columns"]): description["column_context"] = column_context diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index db197a7833..bb738433b3 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -27,10 +27,15 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "Do not use them as executable table or column identifiers" in rules assert "do not invent a field" in rules assert "join only through the FOREIGN KEY relationships shown" in rules + assert "Never generate SQL from assumptions" in rules -def test_get_text_to_sql_rules_uses_sql_knowledge_override(): - assert get_text_to_sql_rules(_SqlKnowledge()) == _SqlKnowledge.text_to_sql_rule +def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): + rules = get_text_to_sql_rules(_SqlKnowledge()) + + assert _SqlKnowledge.text_to_sql_rule in rules + assert "MANDATORY SQL GROUNDING RULES" in rules + assert "Every table and column referenced" in rules def test_get_metric_instructions_uses_sql_knowledge_override(): @@ -50,4 +55,5 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "ONLY USE table/column alias in the final SELECT clause" in prompt assert "Refer to the value of alias from the comment section" in prompt assert "source of executable table and column identifiers" in prompt + assert "Never generate SQL from assumptions" in prompt assert 'SELECT "_orders"."ApprovedTimestamp" AS "_timestamp"' in prompt diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index 8898ced29d..079b52f55a 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -24,8 +24,8 @@ def test_single_table_description(): mdl = { "models": [ { - "name": "user", - "properties": {"description": "A table containing user information."}, + "name": "entity", + "properties": {"description": "A generic entity resource."}, } ], "views": [], @@ -37,11 +37,12 @@ def test_single_table_description(): assert len(actual["documents"]) == 1 document: Document = actual["documents"][0] - assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "user"} + assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "entity"} assert document.content == str( { - "name": "user", - "description": "A table containing user information.", + "name": "entity", + "resource_type": "MODEL", + "description": "A generic entity resource.", "columns": "", } ) @@ -52,12 +53,12 @@ def test_multiple_table_descriptions(): mdl = { "models": [ { - "name": "user", - "properties": {"description": "A table containing user information."}, + "name": "entity", + "properties": {"description": "A generic entity resource."}, }, { - "name": "order", - "properties": {"description": "A table containing order details."}, + "name": "activity", + "properties": {"description": "A generic activity resource."}, }, ], "views": [], @@ -71,22 +72,24 @@ def test_multiple_table_descriptions(): document_1: Document = actual["documents"][0] assert document_1.meta == { "type": "TABLE_DESCRIPTION", - "name": "user", + "name": "entity", } assert document_1.content == str( { - "name": "user", - "description": "A table containing user information.", + "name": "entity", + "resource_type": "MODEL", + "description": "A generic entity resource.", "columns": "", } ) document_2: Document = actual["documents"][1] - assert document_2.meta == {"type": "TABLE_DESCRIPTION", "name": "order"} + assert document_2.meta == {"type": "TABLE_DESCRIPTION", "name": "activity"} assert document_2.content == str( { - "name": "order", - "description": "A table containing order details.", + "name": "activity", + "resource_type": "MODEL", + "description": "A generic activity resource.", "columns": "", } ) @@ -112,7 +115,7 @@ def test_table_description_missing_name(): def test_table_description_missing_description(): chunker = TableDescriptionChunker() mdl = { - "models": [{"name": "user"}], + "models": [{"name": "entity"}], "views": [], "relationships": [], "metrics": [], @@ -122,8 +125,10 @@ def test_table_description_missing_description(): assert len(actual["documents"]) == 1 document: Document = actual["documents"][0] - assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "user"} - assert document.content == str({"name": "user", "description": "", "columns": ""}) + assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "entity"} + assert document.content == str( + {"name": "entity", "resource_type": "MODEL", "description": "", "columns": ""} + ) def test_table_description_null_description(): @@ -131,7 +136,7 @@ def test_table_description_null_description(): mdl = { "models": [ { - "name": "user", + "name": "entity", "properties": {"description": None, "displayName": None}, "columns": [{"name": "id"}, {"name": None}], } @@ -145,9 +150,14 @@ def test_table_description_null_description(): assert len(actual["documents"]) == 1 document: Document = actual["documents"][0] - assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "user"} + assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "entity"} assert document.content == str( - {"name": "user", "description": "", "columns": "id, "} + { + "name": "entity", + "resource_type": "MODEL", + "description": "", + "columns": "id, ", + } ) @@ -156,21 +166,24 @@ def test_table_description_includes_column_semantic_context(): mdl = { "models": [ { - "name": "orders", - "properties": {"description": "Customer purchase transactions."}, + "name": "resource", + "properties": { + "description": "A generic described resource.", + "displayName": "Resource", + }, "columns": [ { - "name": "Division", + "name": "AttributeOne", "type": "varchar", "properties": { - "description": "Generic generated division description." + "description": "Generic generated attribute description." }, }, { - "name": "SalesAmount", + "name": "MeasureOne", "type": "float", "properties": { - "description": "Generic generated sales amount description." + "description": "Generic generated measure description." }, }, ], @@ -186,12 +199,14 @@ def test_table_description_includes_column_semantic_context(): document: Document = actual["documents"][0] assert document.content == str( { - "name": "orders", - "description": "Customer purchase transactions.", - "columns": "Division, SalesAmount", + "name": "resource", + "resource_type": "MODEL", + "description": "A generic described resource.", + "columns": "AttributeOne, MeasureOne", + "displayName": "Resource", "column_context": ( - "Division varchar Generic generated division description.; " - "SalesAmount float Generic generated sales amount description." + "AttributeOne varchar Generic generated attribute description.; " + "MeasureOne float Generic generated measure description." ), } ) @@ -233,7 +248,7 @@ def test_table_description_keeps_complete_column_lists(): mdl = { "models": [ { - "name": "user", + "name": "entity", "columns": columns, } ], @@ -248,7 +263,8 @@ def test_table_description_keeps_complete_column_lists(): document: Document = actual["documents"][0] assert document.content == str( { - "name": "user", + "name": "entity", + "resource_type": "MODEL", "description": "", "columns": ", ".join(column["name"] for column in columns), } @@ -260,12 +276,12 @@ async def test_pipeline_run(mocker: MockFixture): test_mdl = { "models": [ { - "name": "user", - "properties": {"description": "A table containing user information."}, + "name": "entity", + "properties": {"description": "A generic entity resource."}, }, { - "name": "order", - "properties": {"description": "A table containing order details."}, + "name": "activity", + "properties": {"description": "A generic activity resource."}, }, ], "views": [], From 3beca535ddbbac983ebfce4e4c6964e75f999ce1 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 16:44:45 +0530 Subject: [PATCH 0706/1087] Prevent assumed identifiers in SQL prompts --- .../generation/followup_sql_generation.py | 2 ++ .../followup_sql_generation_reasoning.py | 2 ++ .../src/pipelines/generation/sql_correction.py | 2 ++ .../src/pipelines/generation/sql_generation.py | 1 + .../generation/sql_generation_reasoning.py | 1 + .../src/pipelines/generation/sql_regeneration.py | 2 ++ .../src/pipelines/generation/utils/sql.py | 16 ++++++++++++---- .../pipelines/generation/test_sql_utils.py | 5 +++++ 8 files changed, 27 insertions(+), 4 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 35cfb8fccf..deeaf14a04 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -60,6 +60,7 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. {% for sample in sql_samples %} Summary: {{sample.summary}} @@ -79,6 +80,7 @@ User's Follow-up Question: {{ query }} ### REASONING PLAN ### +Use this reasoning plan only where it is consistent with the current DATABASE SCHEMA and SQL RULES. {{ sql_generation_reasoning }} Let's think step by step. diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 42b28c5b8f..f68c363b20 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -30,6 +30,7 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Use table names, column names, values, and functions only if they are present in the current DATABASE SCHEMA or SQL FUNCTIONS. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} @@ -46,6 +47,7 @@ {% endif %} ### User's QUERY HISTORY ### +Query history is context only. Do not reuse prior table names, column names, values, or functions unless they are present in the current DATABASE SCHEMA or SQL FUNCTIONS. {% for history in histories %} Question: {{ history.question }} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..8143090eb4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -36,6 +36,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). 2. Then, generate the syntactically correct ANSI SQL query to correct the error. +3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. +4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. ### SQL RULES ### Make sure you follow the SQL Rules strictly. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..86ec514164 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -54,6 +54,7 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. {% for sample in sql_samples %} Question: {{sample.question}} 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 00b731cb2c..d581e1a221 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -29,6 +29,7 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Use table names, column names, values, and functions only if they are present in the current DATABASE SCHEMA or SQL FUNCTIONS. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 4b7284aa26..886b867096 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -77,6 +77,7 @@ def get_sql_regeneration_system_prompt( {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. {% for sample in sql_samples %} Question: {{sample.question}} @@ -93,6 +94,7 @@ def get_sql_regeneration_system_prompt( {% endif %} ### QUESTION ### +Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a26dedf0f1..240fb89769 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -172,6 +172,11 @@ async def _classify_generation_result( - If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - When using multiple tables, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. +- SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. +- Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. +- Apply relative date or time filters only to schema fields whose type or metadata clearly supports date/time semantics. Do not compare text fields to date functions. +- For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. """ @@ -421,8 +426,9 @@ async def _classify_generation_result( - Table structures and relationships used - Specific functions and operators employed - Query patterns and techniques demonstrated -3. Use these samples as reference patterns when generating similar queries +3. Use these samples as reference patterns when generating similar queries, but treat the DATABASE SCHEMA as the only valid source of executable table and column names 4. Adapt the techniques shown in the samples to match new query requirements while maintaining consistent style and approach +5. Never copy table names, column names, aliases, literal values, or functions from samples unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS The samples will help you understand: - Preferred table join patterns @@ -461,7 +467,9 @@ async def _classify_generation_result( 16. Do not write SQL, possible SQL, sample SQL, or assumed SQL in the reasoning plan. 17. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the schema does not show the exact table or column needed, state that the available schema does not include that part. 18. If the question asks for a concept, filter, sort, or timeframe, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. -19. ONLY SHOWING the reasoning plan in bullet points. +19. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, or functions from them unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS. +20. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +21. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -529,8 +537,8 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. +4. YOU MUST FOLLOW the reasoning plan step by step only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains assumed SQL, placeholder identifiers, or identifiers missing from DATABASE SCHEMA, ignore those parts. 5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index bb738433b3..3cf5dbdd24 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -28,6 +28,9 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "do not invent a field" in rules assert "join only through the FOREIGN KEY relationships shown" in rules assert "Never generate SQL from assumptions" in rules + assert "Do not query INFORMATION_SCHEMA" in rules + assert "SQL samples and query history are examples of intent and style only" in rules + assert "order by that alias" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -36,6 +39,7 @@ def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): assert _SqlKnowledge.text_to_sql_rule in rules assert "MANDATORY SQL GROUNDING RULES" in rules assert "Every table and column referenced" in rules + assert "Do not query INFORMATION_SCHEMA" in rules def test_get_metric_instructions_uses_sql_knowledge_override(): @@ -56,4 +60,5 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "Refer to the value of alias from the comment section" in prompt assert "source of executable table and column identifiers" in prompt assert "Never generate SQL from assumptions" in prompt + assert "ignore those parts" in prompt assert 'SELECT "_orders"."ApprovedTimestamp" AS "_timestamp"' in prompt From 9d072f1d1ee8f80a00a697420c7815587b95ec22 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 17:23:55 +0530 Subject: [PATCH 0707/1087] Regenerate failed SQL from user intent --- .../generation/followup_sql_generation.py | 1 + .../pipelines/generation/sql_correction.py | 17 ++++++++++++++++ .../pipelines/generation/sql_generation.py | 1 + .../pipelines/generation/sql_regeneration.py | 7 +++++++ .../src/pipelines/generation/utils/sql.py | 14 +++++++++---- wren-ai-service/src/web/v1/services/ask.py | 2 ++ .../src/web/v1/services/ask_feedback.py | 3 +++ .../pipelines/generation/test_sql_utils.py | 20 +++++++++++++++++++ 8 files changed, 61 insertions(+), 4 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 deeaf14a04..8e315c17d3 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -78,6 +78,7 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. ### REASONING PLAN ### Use this reasoning plan only where it is consistent with the current DATABASE SCHEMA and SQL RULES. diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 8143090eb4..c371309a52 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -38,6 +38,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 2. Then, generate the syntactically correct ANSI SQL query to correct the error. 3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. 4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. +5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. +6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -76,6 +78,13 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### +{% if query %} +User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. +{% endif %} +{% if sql_generation_reasoning %} +SQL generation reasoning: {{ sql_generation_reasoning }} +{% endif %} SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} @@ -89,12 +98,16 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, + query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( + query=query, documents=documents, invalid_generation_result=invalid_generation_result, + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -171,6 +184,8 @@ async def run( self, contexts: List[Document], invalid_generation_result: Dict[str, str], + query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, @@ -189,7 +204,9 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, + "query": query, "documents": contexts, + "sql_generation_reasoning": sql_generation_reasoning, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 86ec514164..2b94885786 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -72,6 +72,7 @@ ### QUESTION ### User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. {% if sql_generation_reasoning %} ### REASONING PLAN ### diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 886b867096..9bb3255b94 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -38,6 +38,7 @@ def get_sql_regeneration_system_prompt( please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. While generating the new SQL query, you should use the original SQL query as a reference. While generating the new SQL query, make sure to use the database schema to generate the SQL query. +If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. {text_to_sql_rules} @@ -94,6 +95,8 @@ def get_sql_regeneration_system_prompt( {% endif %} ### QUESTION ### +User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} @@ -105,6 +108,7 @@ def get_sql_regeneration_system_prompt( ## Start of Pipeline @observe(capture_input=False) def prompt( + query: str, documents: list[str], sql_generation_reasoning: str, sql: str, @@ -118,6 +122,7 @@ def prompt( sql_knowledge: SqlKnowledge | None = None, ) -> dict: _prompt = prompt_builder.run( + query=query, sql=sql, documents=documents, sql_generation_reasoning=sql_generation_reasoning, @@ -197,6 +202,7 @@ def __init__( async def run( self, contexts: list[str], + query: str, sql_generation_reasoning: str, sql: str, sql_samples: list[dict] | None = None, @@ -214,6 +220,7 @@ async def run( ["post_process"], inputs={ "documents": contexts, + "query": query, "sql_generation_reasoning": sql_generation_reasoning, "sql": sql, "sql_samples": sql_samples, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 240fb89769..526b1a1f83 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -168,10 +168,13 @@ async def _classify_generation_result( - Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. - Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. Do not use them as executable table or column identifiers, except as aliases in the final SELECT clause. +- Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. +- When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. - If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - When using multiple tables, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. @@ -467,9 +470,11 @@ async def _classify_generation_result( 16. Do not write SQL, possible SQL, sample SQL, or assumed SQL in the reasoning plan. 17. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the schema does not show the exact table or column needed, state that the available schema does not include that part. 18. If the question asks for a concept, filter, sort, or timeframe, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. -19. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, or functions from them unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS. -20. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -21. ONLY SHOWING the reasoning plan in bullet points. +19. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but name only exact tables and columns from DATABASE SCHEMA in the reasoning plan. +20. If multiple schema objects are required to answer the intent, include each required object only when DATABASE SCHEMA provides both the needed fields and the relationship path between them. +21. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, or functions from them unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS. +22. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +23. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -539,7 +544,8 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. 4. YOU MUST FOLLOW the reasoning plan step by step only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains assumed SQL, placeholder identifiers, or identifiers missing from DATABASE SCHEMA, ignore those parts. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. +6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7a20e792c0..e5554a4e1a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -548,6 +548,8 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, + query=user_query, + sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 25044de18c..f1971afa15 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -176,6 +176,7 @@ async def ask_feedback( "sql_regeneration" ].run( contexts=table_ddls, + query=ask_feedback_request.question, sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, @@ -239,6 +240,8 @@ async def ask_feedback( "sql_correction" ].run( contexts=table_ddls, + query=ask_feedback_request.question, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "original_sql": original_sql, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 3cf5dbdd24..d3d36c6931 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -5,6 +5,8 @@ get_sql_generation_system_prompt, get_text_to_sql_rules, ) +from src.pipelines.generation.sql_correction import get_sql_correction_system_prompt +from src.pipelines.generation.sql_regeneration import get_sql_regeneration_system_prompt class _SqlKnowledge: @@ -31,6 +33,9 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "Do not query INFORMATION_SCHEMA" in rules assert "SQL samples and query history are examples of intent and style only" in rules assert "order by that alias" in rules + assert "Interpret the user's intent" in rules + assert "schema descriptions, aliases, display labels" in rules + assert "use all required related tables" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -61,4 +66,19 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "source of executable table and column identifiers" in prompt assert "Never generate SQL from assumptions" in prompt assert "ignore those parts" in prompt + assert "answer the user's intent" in prompt assert 'SELECT "_orders"."ApprovedTimestamp" AS "_timestamp"' in prompt + + +def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): + prompt = get_sql_regeneration_system_prompt() + + assert "regenerate from the user's question" in prompt + assert "unsupported identifiers" in prompt + + +def test_sql_correction_system_prompt_discards_invalid_identifier_context(): + prompt = get_sql_correction_system_prompt() + + assert "treat it as the source of intent" in prompt + assert "Do not copy placeholders" in prompt From a866cb878d8b82b1b5a0aeaa09289a0a6646124a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 28 Jul 2026 19:52:54 +0530 Subject: [PATCH 0708/1087] Fix schema retrieval and SQL correction flow --- .../generation/followup_sql_generation.py | 9 +----- .../generation/intent_classification.py | 12 +++---- .../pipelines/generation/sql_correction.py | 6 ---- .../pipelines/generation/sql_generation.py | 8 ----- .../pipelines/generation/sql_regeneration.py | 14 +++------ .../src/pipelines/generation/utils/sql.py | 7 ++--- .../retrieval/db_schema_retrieval.py | 17 +++++----- wren-ai-service/src/web/v1/services/ask.py | 6 ++-- .../src/web/v1/services/ask_feedback.py | 7 ++--- .../src/web/v1/services/sql_corrections.py | 9 ------ .../retrieval/test_db_schema_retrieval.py | 31 +++++++++++++++++++ 11 files changed, 61 insertions(+), 65 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 8e315c17d3..e71894541e 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -80,10 +80,6 @@ User's Follow-up Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. -### REASONING PLAN ### -Use this reasoning plan only where it is consistent with the current DATABASE SCHEMA and SQL RULES. -{{ sql_generation_reasoning }} - Let's think step by step. """ @@ -93,7 +89,6 @@ def prompt( query: str, documents: list[str], - sql_generation_reasoning: str, prompt_builder: PromptBuilder, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -106,7 +101,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -199,7 +193,7 @@ async def run( self, query: str, contexts: list[str], - sql_generation_reasoning: str, + sql_generation_reasoning: str | None, histories: list[AskHistory], sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -224,7 +218,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, "sql_samples": sql_samples, diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..2d68d747b6 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -189,10 +189,7 @@ async def table_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - return await table_retriever.run( - query_embedding=embedding.get("embedding"), - filters=filters, - ) + return await table_retriever.run(query_embedding=[], filters=filters) @observe(capture_input=False) @@ -212,6 +209,9 @@ async def dbschema_retrieval( for table_name in table_names ] + if not table_name_conditions: + return [] + filters = { "operator": "AND", "conditions": [ @@ -225,9 +225,7 @@ async def dbschema_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - results = await dbschema_retriever.run( - query_embedding=embedding.get("embedding"), filters=filters - ) + results = await dbschema_retriever.run(query_embedding=[], filters=filters) return results["documents"] diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index c371309a52..99b6d081fa 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -82,9 +82,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. {% endif %} -{% if sql_generation_reasoning %} -SQL generation reasoning: {{ sql_generation_reasoning }} -{% endif %} SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} @@ -99,7 +96,6 @@ def prompt( invalid_generation_result: Dict, prompt_builder: PromptBuilder, query: str | None = None, - sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: @@ -107,7 +103,6 @@ def prompt( query=query, documents=documents, invalid_generation_result=invalid_generation_result, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -206,7 +201,6 @@ async def run( "invalid_generation_result": invalid_generation_result, "query": query, "documents": contexts, - "sql_generation_reasoning": sql_generation_reasoning, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 2b94885786..6761b37b40 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -74,11 +74,6 @@ User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} -{% endif %} - Let's think step by step. """ @@ -89,7 +84,6 @@ def prompt( query: str, documents: list[str], prompt_builder: PromptBuilder, - sql_generation_reasoning: str | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -101,7 +95,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -217,7 +210,6 @@ async def run( inputs={ "query": query, "documents": contexts, - "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 9bb3255b94..616da83747 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,11 +34,11 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query, -please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. -While generating the new SQL query, you should use the original SQL query as a reference. +You are a great ANSI SQL expert. Now you are given database schema, the user's question, and an original SQL query. +Generate a new SQL query that answers the user's question. +While generating the new SQL query, use the original SQL query as intent context only. While generating the new SQL query, make sure to use the database schema to generate the SQL query. -If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. +If the original SQL query contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. {text_to_sql_rules} @@ -98,7 +98,6 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. -SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} Let's think step by step. @@ -110,7 +109,6 @@ def get_sql_regeneration_system_prompt( def prompt( query: str, documents: list[str], - sql_generation_reasoning: str, sql: str, prompt_builder: PromptBuilder, sql_samples: list[dict] | None = None, @@ -125,7 +123,6 @@ def prompt( query=query, sql=sql, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -203,7 +200,7 @@ async def run( self, contexts: list[str], query: str, - sql_generation_reasoning: str, + sql_generation_reasoning: str | None, sql: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -221,7 +218,6 @@ async def run( inputs={ "documents": contexts, "query": query, - "sql_generation_reasoning": sql_generation_reasoning, "sql": sql, "sql_samples": sql_samples, "instructions": instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 526b1a1f83..d50907053f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -536,16 +536,15 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" You are a helpful assistant that converts natural language queries into ANSI SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. +Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -4. YOU MUST FOLLOW the reasoning plan step by step only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains assumed SQL, placeholder identifiers, or identifiers missing from DATABASE SCHEMA, ignore those parts. -5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +4. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. +5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..4c21ce5cca 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -138,7 +138,10 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> d @observe(capture_input=False) async def table_retrieval( - embedding: dict, project_id: str, tables: list[str], table_retriever: Any + embedding: dict, + project_id: str, + tables: Optional[list[str]], + table_retriever: Any, ) -> dict: filters = { "operator": "AND", @@ -152,12 +155,7 @@ async def table_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - if embedding: - return await table_retriever.run( - query_embedding=embedding.get("embedding"), - filters=filters, - ) - else: + if tables: filters["conditions"].append( {"field": "name", "operator": "in", "value": tables} ) @@ -167,6 +165,11 @@ async def table_retrieval( filters=filters, ) + return await table_retriever.run( + query_embedding=[], + filters=filters, + ) + @observe(capture_input=False) async def dbschema_retrieval( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e5554a4e1a..8b4fb04e04 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -464,7 +464,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, histories=histories, project_id=ask_request.project_id, sql_samples=sql_samples, @@ -483,7 +483,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, project_id=ask_request.project_id, sql_samples=sql_samples, instructions=instructions, @@ -549,7 +549,7 @@ async def ask( ].run( contexts=table_ddls, query=user_query, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=None, instructions=instructions, invalid_generation_result={ "sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index f1971afa15..e8297e7e35 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -120,7 +120,6 @@ async def ask_feedback( instructions_task, ) = await asyncio.gather( self._pipelines["db_schema_retrieval"].run( - tables=ask_feedback_request.tables, project_id=ask_feedback_request.project_id, ), self._pipelines["sql_pairs_retrieval"].run( @@ -177,7 +176,7 @@ async def ask_feedback( ].run( contexts=table_ddls, query=ask_feedback_request.question, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + sql_generation_reasoning=None, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, sql_samples=sql_samples, @@ -241,11 +240,11 @@ async def ask_feedback( ].run( contexts=table_ddls, query=ask_feedback_request.question, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + sql_generation_reasoning=None, instructions=instructions, invalid_generation_result={ "original_sql": original_sql, - "sql": invalid_sql, + "sql": original_sql, "error": correction_error_message, }, project_id=ask_feedback_request.project_id, diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 86d0f55301..0d6c3f2785 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -77,7 +77,6 @@ async def correct( sql = request.sql error = request.error project_id = request.project_id - retrieved_tables = request.retrieved_tables use_dry_plan = request.use_dry_plan allow_dry_plan_fallback = request.allow_dry_plan_fallback sql_knowledge = None @@ -88,13 +87,6 @@ async def correct( "error": error, } - if not retrieved_tables: - retrieved_tables = ( - await self._pipelines["sql_tables_extraction"].run( - sql=sql, - ) - )["post_process"] - if self._allow_sql_knowledge_retrieval: sql_knowledge = await self._pipelines["sql_knowledge_retrieval"].run( project_id=project_id, @@ -104,7 +96,6 @@ async def correct( ( await self._pipelines["db_schema_retrieval"].run( project_id=project_id, - tables=retrieved_tables, ) ) .get("construct_retrieval_results", {}) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 1085f4136f..fa4e93d4e0 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -37,6 +37,37 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_table_retrieval_fetches_current_project_table_descriptions(): + class Retriever: + def __init__(self): + self.query_embedding = None + self.filters = None + + async def run(self, query_embedding, filters): + self.query_embedding = query_embedding + self.filters = filters + return {"documents": []} + + retriever = Retriever() + + await table_retrieval( + embedding={"embedding": [0.1, 0.2]}, + project_id="project-1", + tables=None, + table_retriever=retriever, + ) + + assert retriever.query_embedding == [] + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + + @pytest.mark.asyncio async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: From 9abb31f2948eee07a597927d80166d587c0e0243 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 28 Jul 2026 20:03:35 +0530 Subject: [PATCH 0709/1087] Revert "Fix schema retrieval and SQL correction flow" This reverts commit a866cb878d8b82b1b5a0aeaa09289a0a6646124a. --- .../generation/followup_sql_generation.py | 9 +++++- .../generation/intent_classification.py | 12 ++++--- .../pipelines/generation/sql_correction.py | 6 ++++ .../pipelines/generation/sql_generation.py | 8 +++++ .../pipelines/generation/sql_regeneration.py | 14 ++++++--- .../src/pipelines/generation/utils/sql.py | 7 +++-- .../retrieval/db_schema_retrieval.py | 17 +++++----- wren-ai-service/src/web/v1/services/ask.py | 6 ++-- .../src/web/v1/services/ask_feedback.py | 7 +++-- .../src/web/v1/services/sql_corrections.py | 9 ++++++ .../retrieval/test_db_schema_retrieval.py | 31 ------------------- 11 files changed, 65 insertions(+), 61 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 e71894541e..8e315c17d3 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -80,6 +80,10 @@ User's Follow-up Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. +### REASONING PLAN ### +Use this reasoning plan only where it is consistent with the current DATABASE SCHEMA and SQL RULES. +{{ sql_generation_reasoning }} + Let's think step by step. """ @@ -89,6 +93,7 @@ def prompt( query: str, documents: list[str], + sql_generation_reasoning: str, prompt_builder: PromptBuilder, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -101,6 +106,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -193,7 +199,7 @@ async def run( self, query: str, contexts: list[str], - sql_generation_reasoning: str | None, + sql_generation_reasoning: str, histories: list[AskHistory], sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -218,6 +224,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, "sql_samples": sql_samples, diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 2d68d747b6..4d6cd313cd 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -189,7 +189,10 @@ async def table_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - return await table_retriever.run(query_embedding=[], filters=filters) + return await table_retriever.run( + query_embedding=embedding.get("embedding"), + filters=filters, + ) @observe(capture_input=False) @@ -209,9 +212,6 @@ async def dbschema_retrieval( for table_name in table_names ] - if not table_name_conditions: - return [] - filters = { "operator": "AND", "conditions": [ @@ -225,7 +225,9 @@ async def dbschema_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - results = await dbschema_retriever.run(query_embedding=[], filters=filters) + results = await dbschema_retriever.run( + query_embedding=embedding.get("embedding"), filters=filters + ) return results["documents"] diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 99b6d081fa..c371309a52 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -82,6 +82,9 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. {% endif %} +{% if sql_generation_reasoning %} +SQL generation reasoning: {{ sql_generation_reasoning }} +{% endif %} SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} @@ -96,6 +99,7 @@ def prompt( invalid_generation_result: Dict, prompt_builder: PromptBuilder, query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: @@ -103,6 +107,7 @@ def prompt( query=query, documents=documents, invalid_generation_result=invalid_generation_result, + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -201,6 +206,7 @@ async def run( "invalid_generation_result": invalid_generation_result, "query": query, "documents": contexts, + "sql_generation_reasoning": sql_generation_reasoning, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 6761b37b40..2b94885786 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -74,6 +74,11 @@ User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. +{% if sql_generation_reasoning %} +### REASONING PLAN ### +{{ sql_generation_reasoning }} +{% endif %} + Let's think step by step. """ @@ -84,6 +89,7 @@ def prompt( query: str, documents: list[str], prompt_builder: PromptBuilder, + sql_generation_reasoning: str | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -95,6 +101,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -210,6 +217,7 @@ async def run( inputs={ "query": query, "documents": contexts, + "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 616da83747..9bb3255b94 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,11 +34,11 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema, the user's question, and an original SQL query. -Generate a new SQL query that answers the user's question. -While generating the new SQL query, use the original SQL query as intent context only. +You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query, +please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. +While generating the new SQL query, you should use the original SQL query as a reference. While generating the new SQL query, make sure to use the database schema to generate the SQL query. -If the original SQL query contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. +If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. {text_to_sql_rules} @@ -98,6 +98,7 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. +SQL generation reasoning: {{ sql_generation_reasoning }} Original SQL query: {{ sql }} Let's think step by step. @@ -109,6 +110,7 @@ def get_sql_regeneration_system_prompt( def prompt( query: str, documents: list[str], + sql_generation_reasoning: str, sql: str, prompt_builder: PromptBuilder, sql_samples: list[dict] | None = None, @@ -123,6 +125,7 @@ def prompt( query=query, sql=sql, documents=documents, + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -200,7 +203,7 @@ async def run( self, contexts: list[str], query: str, - sql_generation_reasoning: str | None, + sql_generation_reasoning: str, sql: str, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, @@ -218,6 +221,7 @@ async def run( inputs={ "documents": contexts, "query": query, + "sql_generation_reasoning": sql_generation_reasoning, "sql": sql, "sql_samples": sql_samples, "instructions": instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index d50907053f..526b1a1f83 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -536,15 +536,16 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" You are a helpful assistant that converts natural language queries into ANSI SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query. +Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -4. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +4. YOU MUST FOLLOW the reasoning plan step by step only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains assumed SQL, placeholder identifiers, or identifiers missing from DATABASE SCHEMA, ignore those parts. +5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. +6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 4c21ce5cca..6c8dd7bbe3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -138,10 +138,7 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> d @observe(capture_input=False) async def table_retrieval( - embedding: dict, - project_id: str, - tables: Optional[list[str]], - table_retriever: Any, + embedding: dict, project_id: str, tables: list[str], table_retriever: Any ) -> dict: filters = { "operator": "AND", @@ -155,7 +152,12 @@ async def table_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - if tables: + if embedding: + return await table_retriever.run( + query_embedding=embedding.get("embedding"), + filters=filters, + ) + else: filters["conditions"].append( {"field": "name", "operator": "in", "value": tables} ) @@ -165,11 +167,6 @@ async def table_retrieval( filters=filters, ) - return await table_retriever.run( - query_embedding=[], - filters=filters, - ) - @observe(capture_input=False) async def dbschema_retrieval( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8b4fb04e04..e5554a4e1a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -464,7 +464,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, sql_samples=sql_samples, @@ -483,7 +483,7 @@ async def ask( ].run( query=user_query, contexts=table_ddls, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, sql_samples=sql_samples, instructions=instructions, @@ -549,7 +549,7 @@ async def ask( ].run( contexts=table_ddls, query=user_query, - sql_generation_reasoning=None, + sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index e8297e7e35..f1971afa15 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -120,6 +120,7 @@ async def ask_feedback( instructions_task, ) = await asyncio.gather( self._pipelines["db_schema_retrieval"].run( + tables=ask_feedback_request.tables, project_id=ask_feedback_request.project_id, ), self._pipelines["sql_pairs_retrieval"].run( @@ -176,7 +177,7 @@ async def ask_feedback( ].run( contexts=table_ddls, query=ask_feedback_request.question, - sql_generation_reasoning=None, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, sql_samples=sql_samples, @@ -240,11 +241,11 @@ async def ask_feedback( ].run( contexts=table_ddls, query=ask_feedback_request.question, - sql_generation_reasoning=None, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "original_sql": original_sql, - "sql": original_sql, + "sql": invalid_sql, "error": correction_error_message, }, project_id=ask_feedback_request.project_id, diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 0d6c3f2785..86d0f55301 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -77,6 +77,7 @@ async def correct( sql = request.sql error = request.error project_id = request.project_id + retrieved_tables = request.retrieved_tables use_dry_plan = request.use_dry_plan allow_dry_plan_fallback = request.allow_dry_plan_fallback sql_knowledge = None @@ -87,6 +88,13 @@ async def correct( "error": error, } + if not retrieved_tables: + retrieved_tables = ( + await self._pipelines["sql_tables_extraction"].run( + sql=sql, + ) + )["post_process"] + if self._allow_sql_knowledge_retrieval: sql_knowledge = await self._pipelines["sql_knowledge_retrieval"].run( project_id=project_id, @@ -96,6 +104,7 @@ async def correct( ( await self._pipelines["db_schema_retrieval"].run( project_id=project_id, + tables=retrieved_tables, ) ) .get("construct_retrieval_results", {}) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index fa4e93d4e0..1085f4136f 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -37,37 +37,6 @@ async def run(self, query_embedding, filters): } -@pytest.mark.asyncio -async def test_table_retrieval_fetches_current_project_table_descriptions(): - class Retriever: - def __init__(self): - self.query_embedding = None - self.filters = None - - async def run(self, query_embedding, filters): - self.query_embedding = query_embedding - self.filters = filters - return {"documents": []} - - retriever = Retriever() - - await table_retrieval( - embedding={"embedding": [0.1, 0.2]}, - project_id="project-1", - tables=None, - table_retriever=retriever, - ) - - assert retriever.query_embedding == [] - assert retriever.filters == { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, - {"field": "project_id", "operator": "==", "value": "project-1"}, - ], - } - - @pytest.mark.asyncio async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: From 14988ace5a8a39ba3e95b707612fa1de9fe8b506 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 20:52:08 +0530 Subject: [PATCH 0710/1087] Improve metadata-grounded SQL generation --- wren-ai-service/src/pipelines/common.py | 16 +- .../src/pipelines/generation/utils/sql.py | 152 ++---------------- .../src/pipelines/indexing/db_schema.py | 86 ++++++++-- .../pipelines/indexing/table_description.py | 3 + .../pipelines/generation/test_sql_utils.py | 4 +- .../pipelines/indexing/test_db_schema.py | 29 +--- .../retrieval/test_db_schema_retrieval.py | 50 ++++++ 7 files changed, 166 insertions(+), 174 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index f6114d63b1..f0ca079a72 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -33,11 +33,23 @@ def build_table_ddl( columns_ddl = [] has_calculated_field = False has_json_field = False + relationship_columns = { + column.get("column") + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + } + relationship_columns.discard(None) for column in content["columns"]: if column["type"] == "COLUMN": if ( - (not columns or (columns and column["name"] in columns)) + ( + not columns + or column["name"] in columns + or column["name"] in relationship_columns + or column["is_primary_key"] + ) and column["data_type"].lower() != "unknown" # quick fix: filtering out UNKNOWN column type ): @@ -50,7 +62,7 @@ def build_table_ddl( column_ddl += " PRIMARY KEY" columns_ddl.append(column_ddl) elif column["type"] == "FOREIGN_KEY": - if not tables or (tables and set(column["tables"]).issubset(tables)): + if not tables or (tables and set(column.get("tables", [])).issubset(tables)): columns_ddl.append(f"{column['comment']}{column['constraint']}") return ( diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 526b1a1f83..33e42f2586 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -196,50 +196,23 @@ async def _classify_generation_result( - Put double quotes around column and table names. - Put single quotes around string literals. - Never quote numeric literals. - For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992; -- YOU MUST USE "lower(.) like lower()" function or "lower(.) = lower()" function for case-insensitive comparison! - - Use "lower(.) LIKE lower()" when: - - The user requests a pattern or partial match. - - The value is not specific enough to be a single, exact value. - - Wildcards (%) are needed to capture the pattern. - - Use "lower(.) = lower()" when: - - The user requests an exact, specific value. - - There is no ambiguity or pattern in the value. -- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query - - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds - - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds - - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds -- ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query - - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) -- If the user asks for a specific date, please give the date range in SQL query - - example: "What is the total revenue for the month of 2024-11-01?" - - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" +- For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. +- For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only SQL FUNCTIONS-supported date/time functions and casts for the active datasource. +- If the question asks for a specific date, generate a bounded date/time filter on an exact date/time schema column when supported by SQL FUNCTIONS. If no exact date/time schema column is available, do not invent one. - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. - Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. - - EXAMPLE - DATABASE SCHEMA - /* {"alias":"_orders","description":"A model representing the orders data."} */ - CREATE TABLE orders ( - -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} - ApprovedTimestamp TIMESTAMP - } - - SQL - SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. - DON'T USE "EXTRACT()" function with INTERVAL data types as arguments - DON'T USE INTERVAL or generate INTERVAL-like expression in the generated SQL query. - DON'T USE "TO_CHAR" function in the generated SQL query. +- DON'T USE unsupported statistical, date/time, or formatting functions. If SQL FUNCTIONS does not list a function needed by the intent, answer with the closest supported aggregation/filter over exact schema fields. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -- For the ranking problem, you must add the ranking column to the final SELECT clause. +- For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. """ @@ -247,47 +220,9 @@ async def _classify_generation_result( #### Instructions for Calculated Field #### The first structure is the special column marked as "Calculated Field". You need to interpret the purpose and calculation basis for these columns, then utilize them in the following text-to-sql generation tasks. -First, provide a brief explanation of what each field represents in the context of the schema, including how each field is computed using the relationships between models. -Then, during the following tasks, if the user queries pertain to any calculated fields defined in the database schema, ensure to utilize those calculated fields appropriately in the output SQL queries. -The goal is to accurately reflect the intent of the question in the SQL syntax, leveraging the pre-computed logic embedded within the calculated fields. - -EXAMPLES: -The given schema is created by the SQL command: - -CREATE TABLE orders ( - OrderId VARCHAR PRIMARY KEY, - CustomerId VARCHAR, - -- This column is a Calculated Field - -- column expression: avg(reviews.Score) - Rating DOUBLE, - -- This column is a Calculated Field - -- column expression: count(reviews.Id) - ReviewCount BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) - Size BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) > 1 - Large BOOLEAN, - FOREIGN KEY (CustomerId) REFERENCES customers(Id) -); - -Interpret the columns that are marked as Calculated Fields in the schema: -Rating (DOUBLE) - Calculated as the average score (avg) of the Score field from the reviews table where the reviews are associated with the order. This field represents the overall customer satisfaction rating for the order based on review scores. -ReviewCount (BIGINT) - Calculated by counting (count) the number of entries in the reviews table associated with this order. It measures the volume of customer feedback received for the order. -Size (BIGINT) - Represents the total number of items in the order, calculated by counting the number of item entries (ItemNumber) in the order_items table linked to this order. This field is useful for understanding the scale or size of an order. -Large (BOOLEAN) - A boolean value calculated to check if the number of items in the order exceeds one (count(order_items.ItemNumber) > 1). It indicates whether the order is considered large in terms of item quantity. - -And if the user input queries like these: -1. "How many large orders have been placed by customer with ID 'C1234'?" -2. "What is the average customer rating for orders that were rated by more than 10 reviewers?" - -For the first query: -First try to intepret the user query, the user wants to know the average rating for orders which have attracted significant review activity, specifically those with more than 10 reviews. -Then, according to the above intepretation about the given schema, the term 'Rating' is predefined in the Calculated Field of the 'orders' model. And, the number of reviews is also predefined in the 'ReviewCount' Calculated Field. -So utilize those Calculated Fields in the SQL generation process to give an answer like this: - -SQL Query: SELECT AVG(Rating) FROM orders WHERE ReviewCount > 10 +First, interpret each calculated field from its expression, data type, comments, aliases, descriptions, and relationship context in the provided DATABASE SCHEMA. +Then, if the user query matches a concept already represented by a calculated field, use that exact calculated field name from DATABASE SCHEMA instead of recreating or inventing the calculation. +Calculated field expressions are semantic definitions; do not copy identifiers from an expression unless they also appear as executable identifiers in the current DATABASE SCHEMA. """ _DEFAULT_METRIC_INSTRUCTIONS = """ @@ -317,68 +252,7 @@ async def _classify_generation_result( If the given schema contains the structures marked as 'metric', you should first interpret the metric schema based on the above definition. Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. - -EXAMPLES: -The given schema is created by the SQL command: - -/* This table is a metric */ -/* Metric Base Object: orders */ -CREATE TABLE Revenue ( - -- This column is a dimension - PurchaseTimestamp TIMESTAMP, - -- This column is a dimension - CustomerId VARCHAR, - -- This column is a dimension - Status VARCHAR, - -- This column is a measure - -- expression: sum(order_items.Price) - PriceSum DOUBLE, - -- This column is a measure - -- expression: count(OrderId) - NumberOfOrders BIGINT -); - -Interpret the metric with the understanding of the metric structure: -1. Base Object: orders -This is the primary data source for the metric. -The orders table provides the underlying data from which dimensions and measures are derived. -It is the foundation upon which the metric is built, though it itself is not directly used in queries against the Revenue table. -It shows the reference between the 'Revenue' metric and the 'orders' model. For the user queries pretain to the 'Revenue' of 'orders', the metric should be utilize in the sql generation process. -2. Dimensions -The metric contains the columns marked as 'dimension'. They can be interpreted as below: -- PurchaseTimestamp (TIMESTAMP) - Acts as a temporal dimension, allowing analysis of revenue over time. This can be used to observe trends, seasonal variations, or performance over specific periods. -- CustomerId (VARCHAR) - A key dimension for customer segmentation, it enables the analysis of revenue generated from individual customers or customer groups. -- Status (VARCHAR) - Reflects the current state of an order (e.g., pending, completed, cancelled). This dimension is crucial for analyses that differentiate performance based on order status. -3. Measures -The metric contains the columns marked as 'measure'. They can be interpreted as below: -- PriceSum (DOUBLE) - A financial measure calculated as sum(order_items.Price), representing the total revenue generated from orders. This measure is vital for tracking overall sales performance and is the primary output of interest in many financial and business analyses. -- NumberOfOrders (BIGINT) - A count measure that provides the total number of orders. This is essential for operational metrics, such as assessing the volume of business activity and evaluating the efficiency of sales processes. - -Now, if the user input queries like this: -Question: "What was the total revenue from each customer last month?" - -First try to intepret the user query, the user asks for a breakdown of the total revenue generated by each customer in the previous calendar month. -The user is specifically interested in understanding how much each customer contributed to the total sales during this period. -To answer this question, it is suitable to use the following components from the metric: -1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. -2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. -3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. -So utilize those metric components in the SQL generation process to give an answer like this: - -SQL Query: -SELECT - CustomerId, - PriceSum AS TotalRevenue -FROM - Revenue -WHERE - PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND - PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) +Use metric columns exactly as declared in DATABASE SCHEMA. Treat dimensions as grouping/filtering fields and measures as pre-defined numeric outputs. Metric base objects and measure expressions are semantic context only; do not copy identifiers from them unless those identifiers also appear in the current DATABASE SCHEMA. """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -454,8 +328,8 @@ async def _classify_generation_result( 2. Explicitly state the following information in the reasoning plan: if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; otherwise, you will put the relative timeframe in the SQL query. -3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. +3. For top, bottom, first, last, highest, or lowest requests, plan to sort by an exact selected column or aggregate alias and limit the result. Include a ranking column only when the user explicitly asks for rank values. +4. Do not plan to use a SQL function unless it appears in SQL FUNCTIONS for this request or is already part of a valid metric/calculated-field definition in DATABASE SCHEMA. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. @@ -534,7 +408,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" -You are a helpful assistant that converts natural language queries into ANSI SQL queries. +You are a helpful assistant that converts natural language queries into Wren SQL queries. Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. @@ -550,7 +424,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be a Wren SQL query in JSON format: {{ "sql": diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 394d087b46..f88cefe4b8 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -25,6 +25,21 @@ logger = logging.getLogger("wren-ai-service") +MAX_DB_SCHEMA_DOCUMENT_LENGTH = 4000 +MAX_DB_SCHEMA_METADATA_TEXT_LENGTH = 1000 + + +def _truncate_metadata_text(value: Any) -> Any: + if value is None: + return "" + if isinstance(value, str) and len(value) > MAX_DB_SCHEMA_METADATA_TEXT_LENGTH: + return value[: MAX_DB_SCHEMA_METADATA_TEXT_LENGTH - 3] + "..." + if isinstance(value, dict): + return {key: _truncate_metadata_text(item) for key, item in value.items()} + if isinstance(value, list): + return [_truncate_metadata_text(item) for item in value] + return value + @component class DDLChunker: @@ -70,7 +85,7 @@ def _column_preprocessor( column: Dict[str, Any], addition: Dict[str, Any] ) -> Dict[str, Any]: addition = { - key: helper(column, **addition) + key: _truncate_metadata_text(helper(column, **addition)) for key, helper in helper.COLUMN_PREPROCESSORS.items() if helper.condition(column, **addition) } @@ -83,7 +98,7 @@ def _column_preprocessor( async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: addition = { - key: await helper(model, **kwargs) + key: _truncate_metadata_text(await helper(model, **kwargs)) for key, helper in helper.MODEL_PREPROCESSORS.items() if helper.condition(model, **kwargs) } @@ -95,7 +110,7 @@ async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: ] return { "name": model.get("name", ""), - "properties": model.get("properties", {}), + "properties": _truncate_metadata_text(model.get("properties", {})), "columns": columns, "primaryKey": model.get("primaryKey", ""), } @@ -182,20 +197,69 @@ def _relationship_command( if join_type not in ["MANY_TO_ONE", "ONE_TO_MANY", "ONE_TO_ONE"]: return None - # Get related table and foreign key column - is_source = table_name == models[0] - related_table = models[1] if is_source else models[0] - condition_parts = condition.split(" = ") - fk_column = condition_parts[0 if is_source else 1].split(".")[1] + condition_parts = [ + part.strip() for part in condition.split("=", maxsplit=1) + ] + if len(condition_parts) != 2: + return None + + model_columns = [] + for condition_part in condition_parts: + name_parts = [ + part.strip() for part in condition_part.split(".", maxsplit=1) + ] + if len(name_parts) != 2: + return None + model_columns.append( + {"table": name_parts[0], "column": name_parts[1]} + ) + + left, right = model_columns + + if join_type == "MANY_TO_ONE": + foreign_side, referenced_side = left, right + elif join_type == "ONE_TO_MANY": + foreign_side, referenced_side = right, left + elif table_name == left["table"]: + foreign_side, referenced_side = left, right + else: + foreign_side, referenced_side = right, left - # Build foreign key constraint - fk_constraint = f"FOREIGN KEY ({fk_column}) REFERENCES {related_table}({primary_keys_map[related_table]})" + if table_name != foreign_side["table"]: + return None + + related_table = referenced_side["table"] + referenced_column = referenced_side["column"] or primary_keys_map.get( + related_table, "" + ) + fk_column = foreign_side["column"] + fk_constraint = ( + f"FOREIGN KEY ({fk_column}) " + f"REFERENCES {related_table}({referenced_column})" + ) + + properties = relationship.get("properties", {}) + relationship_properties = { + "name": relationship.get("name", ""), + "condition": condition, + "joinType": join_type, + "description": _truncate_metadata_text( + properties.get("description", "") + ) + if isinstance(properties, dict) + else "", + "from": f"{foreign_side['table']}.{foreign_side['column']}", + "to": f"{referenced_side['table']}.{referenced_side['column']}", + } return { "type": "FOREIGN_KEY", - "comment": f'-- {{"condition": {condition}, "joinType": {join_type}}}\n ', + "comment": f"-- {relationship_properties}\n ", "constraint": fk_constraint, "tables": models, + "column": fk_column, + "referenced_table": related_table, + "referenced_column": referenced_column, } def _column_batch( diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 412fc5c6ef..a8f8fd9d1f 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -122,12 +122,15 @@ def _relationship_context_by_model() -> Dict[str, List[str]]: if len(models) != 2: continue + properties = self._properties(relationship) summary = " ".join( part for part in [ _text(relationship.get("name", "")), _text(relationship.get("joinType", "")), _text(relationship.get("condition", "")), + _text(properties.get("description", "")), + f"models {' <-> '.join(_text(model) for model in models)}", ] if part ) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index d3d36c6931..c9805c3e8f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -67,7 +67,9 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "Never generate SQL from assumptions" in prompt assert "ignore those parts" in prompt assert "answer the user's intent" in prompt - assert 'SELECT "_orders"."ApprovedTimestamp" AS "_timestamp"' in prompt + assert "Wren SQL query" in prompt + assert "use a normal equality or LIKE comparison" in prompt + assert "unless the user explicitly asks for rank values" in prompt def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py index 3a1d3b05e1..f2b15cef65 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py @@ -384,7 +384,7 @@ async def test_column_with_relationship(): } actual = await chunker.run(mdl, column_batch_size=1) - assert len(actual["documents"]) == 6 + assert len(actual["documents"]) == 5 document_0: Document = actual["documents"][0] assert document_0.meta == {"type": "TABLE_SCHEMA", "name": "user"} @@ -403,33 +403,20 @@ async def test_column_with_relationship(): } ) - document_1: Document = actual["documents"][1] - assert document_1.meta == {"type": "TABLE_SCHEMA", "name": "user"} - assert document_1.content == str( - { - "type": "TABLE_COLUMNS", - "columns": [ - { - "type": "FOREIGN_KEY", - "comment": '-- {"condition": user.id = order.user_id, "joinType": ONE_TO_MANY}\n ', - "constraint": "FOREIGN KEY (id) REFERENCES order(user_id)", - "tables": ["user", "order"], - } - ], - } - ) - - document_4: Document = actual["documents"][4] - assert document_4.meta == {"type": "TABLE_SCHEMA", "name": "order"} - assert document_4.content == str( + document_3: Document = actual["documents"][3] + assert document_3.meta == {"type": "TABLE_SCHEMA", "name": "order"} + assert document_3.content == str( { "type": "TABLE_COLUMNS", "columns": [ { "type": "FOREIGN_KEY", - "comment": '-- {"condition": user.id = order.user_id, "joinType": ONE_TO_MANY}\n ', + "comment": "-- {'name': 'relationship_1', 'condition': 'user.id = order.user_id', 'joinType': 'ONE_TO_MANY', 'description': '', 'from': 'order.user_id', 'to': 'user.id'}\n ", "constraint": "FOREIGN KEY (user_id) REFERENCES user(id)", "tables": ["user", "order"], + "column": "user_id", + "referenced_table": "user", + "referenced_column": "id", } ], } diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 1085f4136f..bb23a90768 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1,6 +1,7 @@ import pytest from haystack import Document +from src.pipelines.common import build_table_ddl from src.pipelines.retrieval.db_schema_retrieval import ( check_using_db_schemas_without_pruning, dbschema_retrieval, @@ -229,3 +230,52 @@ def encode(self, value): ) assert [schema["table_name"] for schema in result["db_schemas"]] == ["activity"] + + +def test_build_table_ddl_preserves_join_columns_when_pruned(): + ddl, _, _ = build_table_ddl( + { + "type": "TABLE", + "name": "detail", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "detail_id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": True, + }, + { + "type": "COLUMN", + "name": "parent_id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + }, + { + "type": "FOREIGN_KEY", + "comment": "", + "constraint": "FOREIGN KEY (parent_id) REFERENCES parent(parent_id)", + "tables": ["parent", "detail"], + "column": "parent_id", + "referenced_table": "parent", + "referenced_column": "parent_id", + }, + ], + }, + columns={"amount"}, + tables={"parent", "detail"}, + ) + + assert "detail_id INTEGER PRIMARY KEY" in ddl + assert "parent_id INTEGER" in ddl + assert "amount DOUBLE" in ddl + assert "FOREIGN KEY (parent_id) REFERENCES parent(parent_id)" in ddl From 7e1159ff2b61ecbd7b7bd125e93567255faf3eae Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 21:25:29 +0530 Subject: [PATCH 0711/1087] Tighten SQL generation grounding prompt --- .../src/pipelines/generation/utils/sql.py | 27 +++++-------------- .../pipelines/generation/test_sql_utils.py | 14 ++++++++++ 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 33e42f2586..06d1219a3a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -180,6 +180,8 @@ async def _classify_generation_result( - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. - Apply relative date or time filters only to schema fields whose type or metadata clearly supports date/time semantics. Do not compare text fields to date functions. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. +- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part or answer with the closest valid SQL over grounded fields only. +- If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. """ @@ -263,31 +265,13 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - For Example: - DATA SCHEMA: - `/* {"alias":"users","description":"A model representing the users data."} */ - CREATE TABLE users ( - -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} - address JSON - )` - To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` + - JSON paths and nested field names must come from the json_fields metadata attached to the exact JSON column in DATABASE SCHEMA. - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. - - For Example: - DATA SCHEMA - `/* {"alias":"my_table","description":"A test my_table"} */ - CREATE TABLE my_table ( - -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} - elements JSON - )` - To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` + - Do not copy JSON examples, placeholder aliases, or nested paths from prior context. Use only the current table name, JSON column name, and json_fields metadata in DATABASE SCHEMA. - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". - DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". """ @@ -419,7 +403,8 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. 4. YOU MUST FOLLOW the reasoning plan step by step only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains assumed SQL, placeholder identifiers, or identifiers missing from DATABASE SCHEMA, ignore those parts. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -6. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +6. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. +7. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index c9805c3e8f..9c8d9b3e2e 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -36,6 +36,8 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "Interpret the user's intent" in rules assert "schema descriptions, aliases, display labels" in rules assert "use all required related tables" in rules + assert "silently check that each identifier and function" in rules + assert "instead of inventing a replacement" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -70,6 +72,18 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "Wren SQL query" in prompt assert "use a normal equality or LIKE comparison" in prompt assert "unless the user explicitly asks for rank values" in prompt + assert "perform a silent grounding check" in prompt + assert "closest grounded expression" in prompt + + +def test_json_field_instructions_do_not_include_placeholder_identifiers(): + prompt = get_json_field_instructions() + + assert "json_fields metadata" in prompt + assert "Do not copy JSON examples" in prompt + assert "CREATE TABLE users" not in prompt + assert "my_table" not in prompt + assert "parent_table" not in prompt def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): From d0fdc1fae2629a88da054946adb7528782b6a36f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 22:08:51 +0530 Subject: [PATCH 0712/1087] Ground SQL generation in metadata semantics --- .../generation/followup_sql_generation.py | 5 ++-- .../followup_sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_correction.py | 17 +++++++----- .../pipelines/generation/sql_generation.py | 3 ++- .../generation/sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_regeneration.py | 20 ++++++++------ .../src/pipelines/generation/utils/sql.py | 19 ++++++++----- .../pipelines/indexing/table_description.py | 15 +++++++++-- .../pipelines/generation/test_sql_utils.py | 27 +++++++++++++++++++ .../indexing/test_table_description.py | 5 ++++ 10 files changed, 87 insertions(+), 28 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 8e315c17d3..4023c6f770 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -81,10 +81,11 @@ Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. ### REASONING PLAN ### -Use this reasoning plan only where it is consistent with the current DATABASE SCHEMA and SQL RULES. +Use this reasoning plan only as non-executable context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. +Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} -Let's think step by step. +Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index f68c363b20..2a2e35c5ce 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -60,7 +60,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Let's think step by step. +Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index c371309a52..15255f8f39 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,16 +30,17 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query. +You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. ### SQL CORRECTION INSTRUCTIONS ### -1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). -2. Then, generate the syntactically correct ANSI SQL query to correct the error. +1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. +2. Then, generate a syntactically correct ANSI SQL query from the user's intent and the current DATABASE SCHEMA. 3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. 4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. 5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. 6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. +7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -83,12 +84,16 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. {% endif %} {% if sql_generation_reasoning %} -SQL generation reasoning: {{ sql_generation_reasoning }} +### REASONING PLAN ### +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. +{{ sql_generation_reasoning }} {% endif %} -SQL: {{ invalid_generation_result.sql }} +### FAILED SQL ### +This SQL failed dry run. Do not preserve any identifier or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +{{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} -Let's think step by step. +Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 2b94885786..758eae2c8c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -76,10 +76,11 @@ {% if sql_generation_reasoning %} ### REASONING PLAN ### +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} {% endif %} -Let's think step by step. +Return only the final JSON SQL response. """ 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 d581e1a221..e6d5c2c625 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -50,7 +50,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Let's think step by step. +Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 9bb3255b94..8e636bf317 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,10 +34,10 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query, -please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. -While generating the new SQL query, you should use the original SQL query as a reference. -While generating the new SQL query, make sure to use the database schema to generate the SQL query. +You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query. +Carefully review the user's question and current DATABASE SCHEMA, then generate a new SQL query that answers the user's intent. +Use the original SQL query only as non-executable intent context. +While generating the new SQL query, make sure to use the database schema as the only source of executable table and column identifiers. If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. {text_to_sql_rules} @@ -98,10 +98,14 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. -SQL generation reasoning: {{ sql_generation_reasoning }} -Original SQL query: {{ sql }} - -Let's think step by step. +### REASONING PLAN ### +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. +{{ sql_generation_reasoning }} +### ORIGINAL SQL QUERY ### +Use this SQL only as non-executable intent context. Do not preserve any identifier or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +{{ sql }} + +Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 06d1219a3a..62ca4aa38c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -179,6 +179,8 @@ async def _classify_generation_result( - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. - Apply relative date or time filters only to schema fields whose type or metadata clearly supports date/time semantics. Do not compare text fields to date functions. +- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, or unsupported functions from them. +- If a column comment, alias, display label, or description names a business concept, first locate the exact declared column for that concept in DATABASE SCHEMA. If no exact declared column exists, omit that concept. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part or answer with the closest valid SQL over grounded fields only. - If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. @@ -305,7 +307,7 @@ async def _classify_generation_result( sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who is great at thinking deeply and reasoning about the user's question and the database schema, and you provide a step-by-step reasoning plan in order to answer the user's question. +You are a helpful data analyst who maps a user's intent to the provided database schema and provides a concise reasoning plan for answering the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. @@ -325,14 +327,16 @@ async def _classify_generation_result( 13. A column name in the reasoning plan must be in this format: `column: .`. 14. Use only exact table and column names that appear in the DATABASE SCHEMA section. 15. Comments, aliases, display labels, and descriptions are semantic hints only; do not turn them into table or column names in the reasoning plan. -16. Do not write SQL, possible SQL, sample SQL, or assumed SQL in the reasoning plan. +16. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. 17. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the schema does not show the exact table or column needed, state that the available schema does not include that part. 18. If the question asks for a concept, filter, sort, or timeframe, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. 19. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but name only exact tables and columns from DATABASE SCHEMA in the reasoning plan. 20. If multiple schema objects are required to answer the intent, include each required object only when DATABASE SCHEMA provides both the needed fields and the relationship path between them. 21. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, or functions from them unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS. 22. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -23. ONLY SHOWING the reasoning plan in bullet points. +23. Use comments, aliases, display labels, and descriptions to explain why an exact schema column is relevant, not as replacement names. +24. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. +25. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -394,17 +398,18 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" You are a helpful assistant that converts natural language queries into Wren SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. +Given the user's question and database schema, generate one grounded Wren SQL query. The DATABASE SCHEMA is the only source of executable identifiers. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -4. YOU MUST FOLLOW the reasoning plan step by step only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains assumed SQL, placeholder identifiers, or identifiers missing from DATABASE SCHEMA, ignore those parts. +4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains SQL fragments, assumed SQL, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers missing from DATABASE SCHEMA, ignore those parts. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -6. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. -7. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +6. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +7. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. +8. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index a8f8fd9d1f..e0914ccc81 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -172,6 +172,8 @@ def _column_context(columns: List[Dict[str, Any]]) -> str: ) def _resource_description(resource: Dict[str, Any]) -> Dict[str, str]: + column_context = _column_context(resource["columns"]) + relationships = "; ".join(relationship_context.get(resource["name"], [])) description = { "name": resource["name"], "resource_type": resource["mdl_type"], @@ -187,12 +189,21 @@ def _resource_description(resource: Dict[str, Any]) -> Dict[str, str]: if resource["source"]: description["source"] = resource["source"] - if column_context := _column_context(resource["columns"]): + if column_context: description["column_context"] = column_context - if relationships := "; ".join(relationship_context.get(resource["name"], [])): + if relationships: description["relationships"] = relationships + semantic_parts = [ + description.get("displayName", ""), + description.get("source", ""), + column_context, + relationships, + ] + if semantic_context := "; ".join(part for part in semantic_parts if part): + description["semantic_context"] = semantic_context + return description return [ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 9c8d9b3e2e..ea4b867025 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -4,6 +4,7 @@ get_metric_instructions, get_sql_generation_system_prompt, get_text_to_sql_rules, + sql_generation_reasoning_system_prompt, ) from src.pipelines.generation.sql_correction import get_sql_correction_system_prompt from src.pipelines.generation.sql_regeneration import get_sql_regeneration_system_prompt @@ -38,6 +39,11 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "use all required related tables" in rules assert "silently check that each identifier and function" in rules assert "instead of inventing a replacement" in rules + assert ( + "Treat reasoning plans, correction notes, and error messages as non-executable context" + in rules + ) + assert "first locate the exact declared column" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -74,6 +80,9 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "unless the user explicitly asks for rank values" in prompt assert "perform a silent grounding check" in prompt assert "closest grounded expression" in prompt + assert "DATABASE SCHEMA is the only source of executable identifiers" in prompt + assert "reasoning plan only as non-executable context" in prompt + assert "include those objects only when DATABASE SCHEMA shows" in prompt def test_json_field_instructions_do_not_include_placeholder_identifiers(): @@ -91,6 +100,11 @@ def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): assert "regenerate from the user's question" in prompt assert "unsupported identifiers" in prompt + assert "Use the original SQL query only as non-executable intent context" in prompt + assert ( + "database schema as the only source of executable table and column identifiers" + in prompt + ) def test_sql_correction_system_prompt_discards_invalid_identifier_context(): @@ -98,3 +112,16 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "treat it as the source of intent" in prompt assert "Do not copy placeholders" in prompt + assert "Regenerate a grounded Wren SQL query" in prompt + assert ( + "Do not preserve a table, column, join, filter, grouping, ordering, or function" + in prompt + ) + + +def test_sql_reasoning_prompt_forbids_executable_sql_context(): + prompt = sql_generation_reasoning_system_prompt + + assert "Do not write SQL, possible SQL, sample SQL, assumed SQL" in prompt + assert "SQL clauses, SQL functions, code blocks, or executable expressions" in prompt + assert "The reasoning plan is non-executable context" in prompt diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index 079b52f55a..db5051887b 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -208,6 +208,11 @@ def test_table_description_includes_column_semantic_context(): "AttributeOne varchar Generic generated attribute description.; " "MeasureOne float Generic generated measure description." ), + "semantic_context": ( + "Resource; " + "AttributeOne varchar Generic generated attribute description.; " + "MeasureOne float Generic generated measure description." + ), } ) From 8ab4080dfc31d0a9f47eda18f43cf6f5fac0908c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 22:59:53 +0530 Subject: [PATCH 0713/1087] Stop semantic labels from becoming SQL identifiers --- .../generation/followup_sql_generation.py | 2 +- .../pipelines/generation/sql_correction.py | 2 +- .../pipelines/generation/sql_generation.py | 2 +- .../pipelines/generation/sql_regeneration.py | 2 +- .../src/pipelines/generation/utils/sql.py | 28 ++++++++++--------- .../pipelines/generation/test_sql_utils.py | 12 +++++--- 6 files changed, 27 insertions(+), 21 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 4023c6f770..3de5b8611a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -78,7 +78,7 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. ### REASONING PLAN ### Use this reasoning plan only as non-executable context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 15255f8f39..1b5a4c7b63 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -81,7 +81,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### QUESTION ### {% if query %} User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 758eae2c8c..971e3dd940 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -72,7 +72,7 @@ ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. {% if sql_generation_reasoning %} ### REASONING PLAN ### diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 8e636bf317..e157739d07 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -96,7 +96,7 @@ def get_sql_regeneration_system_prompt( ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact table and column names from DATABASE SCHEMA. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 62ca4aa38c..4ae103031c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -167,9 +167,10 @@ async def _classify_generation_result( ### MANDATORY SQL GROUNDING RULES ### - Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. - Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. -- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. Do not use them as executable table or column identifiers, except as aliases in the final SELECT clause. +- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. - Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. +- The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. - If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. @@ -178,9 +179,9 @@ async def _classify_generation_result( - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. -- Apply relative date or time filters only to schema fields whose type or metadata clearly supports date/time semantics. Do not compare text fields to date functions. +- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. - Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, or unsupported functions from them. -- If a column comment, alias, display label, or description names a business concept, first locate the exact declared column for that concept in DATABASE SCHEMA. If no exact declared column exists, omit that concept. +- If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part or answer with the closest valid SQL over grounded fields only. - If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. @@ -201,13 +202,13 @@ async def _classify_generation_result( - Put single quotes around string literals. - Never quote numeric literals. - For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. -- For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only SQL FUNCTIONS-supported date/time functions and casts for the active datasource. -- If the question asks for a specific date, generate a bounded date/time filter on an exact date/time schema column when supported by SQL FUNCTIONS. If no exact date/time schema column is available, do not invent one. +- For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only date/time functions and casts whose exact syntax is provided in SQL FUNCTIONS for this request. +- If the question asks for a specific or relative date, generate a bounded date/time filter only when both the exact date/time schema column and required SQL FUNCTIONS-supported operation are available. If either is missing, do not invent a field or function. - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. -- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. -- DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. +- Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. +- Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. +- DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. - DON'T USE "EXTRACT()" function with INTERVAL data types as arguments @@ -311,9 +312,9 @@ async def _classify_generation_result( ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; -otherwise, you will put the relative timeframe in the SQL query. +2. Explicitly state the following information in the reasoning plan: +if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will use that absolute time frame only with an exact date/time schema column and supported SQL FUNCTIONS; +otherwise, you will use a relative timeframe only when an exact date/time schema column and supported SQL FUNCTIONS are available. If they are not available, state that the available schema/functions do not support that time filter. 3. For top, bottom, first, last, highest, or lowest requests, plan to sort by an exact selected column or aggregate alias and limit the result. Include a ranking column only when the user explicitly asks for rank values. 4. Do not plan to use a SQL function unless it appears in SQL FUNCTIONS for this request or is already part of a valid metric/calculated-field definition in DATABASE SCHEMA. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. @@ -336,7 +337,8 @@ async def _classify_generation_result( 22. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. 23. Use comments, aliases, display labels, and descriptions to explain why an exact schema column is relevant, not as replacement names. 24. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. -25. ONLY SHOWING the reasoning plan in bullet points. +25. Do not derive executable table or column names from natural language, comments, aliases, display labels, or descriptions. Only cite exact declared names from DATABASE SCHEMA. +26. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -403,7 +405,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. -2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. +2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. 4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains SQL fragments, assumed SQL, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers missing from DATABASE SCHEMA, ignore those parts. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index ea4b867025..525ea01ed3 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -27,10 +27,11 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "ONLY USE the tables and columns mentioned in the database schema" in rules assert 'ONLY USE "*" if the user query asks for all the columns' in rules - assert "Do not use them as executable table or column identifiers" in rules + assert "They are never source table or source column identifiers" in rules assert "do not invent a field" in rules assert "join only through the FOREIGN KEY relationships shown" in rules assert "Never generate SQL from assumptions" in rules + assert "Do not derive executable identifiers" in rules assert "Do not query INFORMATION_SCHEMA" in rules assert "SQL samples and query history are examples of intent and style only" in rules assert "order by that alias" in rules @@ -39,11 +40,12 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "use all required related tables" in rules assert "silently check that each identifier and function" in rules assert "instead of inventing a replacement" in rules + assert "exact date/time schema column and required SQL FUNCTIONS-supported operation" in rules assert ( "Treat reasoning plans, correction notes, and error messages as non-executable context" in rules ) - assert "first locate the exact declared column" in rules + assert "first locate the exact declared source column" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -69,8 +71,8 @@ def test_get_json_field_instructions_uses_sql_knowledge_override(): def test_sql_generation_system_prompt_grounding_contract(): prompt = get_sql_generation_system_prompt() - assert "ONLY USE table/column alias in the final SELECT clause" in prompt - assert "Refer to the value of alias from the comment section" in prompt + assert "Output aliases are labels for result columns only" in prompt + assert "must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY" in prompt assert "source of executable table and column identifiers" in prompt assert "Never generate SQL from assumptions" in prompt assert "ignore those parts" in prompt @@ -83,6 +85,7 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "DATABASE SCHEMA is the only source of executable identifiers" in prompt assert "reasoning plan only as non-executable context" in prompt assert "include those objects only when DATABASE SCHEMA shows" in prompt + assert "Use the exact supported syntax shown there" in prompt def test_json_field_instructions_do_not_include_placeholder_identifiers(): @@ -125,3 +128,4 @@ def test_sql_reasoning_prompt_forbids_executable_sql_context(): assert "Do not write SQL, possible SQL, sample SQL, assumed SQL" in prompt assert "SQL clauses, SQL functions, code blocks, or executable expressions" in prompt assert "The reasoning plan is non-executable context" in prompt + assert "Only cite exact declared names from DATABASE SCHEMA" in prompt From 356bbd618a5a2d9ae022d644c7bccbcf4d28655d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 28 Jul 2026 23:48:06 +0530 Subject: [PATCH 0714/1087] Keep source metadata out of generated SQL --- .../generation/followup_sql_generation.py | 6 ++-- .../pipelines/generation/sql_correction.py | 8 +++-- .../pipelines/generation/sql_generation.py | 6 ++-- .../pipelines/generation/sql_regeneration.py | 9 +++--- .../src/pipelines/generation/utils/sql.py | 12 +++++-- .../pipelines/generation/test_sql_utils.py | 31 ++++++++++++++++++- 6 files changed, 55 insertions(+), 17 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 3de5b8611a..e7efdec613 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -60,7 +60,7 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. +These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. Treat physical/source/lineage names in samples as non-executable unless the exact same identifier is declared in DATABASE SCHEMA. {% for sample in sql_samples %} Summary: {{sample.summary}} @@ -78,11 +78,11 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. ### REASONING PLAN ### Use this reasoning plan only as non-executable context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. -Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 1b5a4c7b63..30f5809890 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -41,6 +41,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. 6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. 7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -81,15 +83,15 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### QUESTION ### {% if query %} User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} {% endif %} ### FAILED SQL ### -This SQL failed dry run. Do not preserve any identifier or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +This SQL failed dry run. Do not preserve any identifier, source/physical/lineage name, or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. Do not use the failed SQL or error message as a source for alternate similar names; regenerate from the question and current schema. {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 971e3dd940..03a1a1ae90 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -54,7 +54,7 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. +These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. Treat physical/source/lineage names in samples as non-executable unless the exact same identifier is declared in DATABASE SCHEMA. {% for sample in sql_samples %} Question: {{sample.question}} @@ -72,11 +72,11 @@ ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index e157739d07..b508d7135c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -39,6 +39,7 @@ def get_sql_regeneration_system_prompt( Use the original SQL query only as non-executable intent context. While generating the new SQL query, make sure to use the database schema as the only source of executable table and column identifiers. If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. +Treat physical/source/lineage names from the original SQL, reasoning, samples, comments, or descriptions as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. {text_to_sql_rules} @@ -78,7 +79,7 @@ def get_sql_regeneration_system_prompt( {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. +These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. Treat physical/source/lineage names in samples as non-executable unless the exact same identifier is declared in DATABASE SCHEMA. {% for sample in sql_samples %} Question: {{sample.question}} @@ -96,13 +97,13 @@ def get_sql_regeneration_system_prompt( ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels or inferred names into executable SQL. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} ### ORIGINAL SQL QUERY ### -Use this SQL only as non-executable intent context. Do not preserve any identifier or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +Use this SQL only as non-executable intent context. Do not preserve any identifier, source/physical/lineage name, or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. {{ sql }} Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4ae103031c..b758c36504 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -168,12 +168,14 @@ async def _classify_generation_result( - Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. - Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. +- Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. - Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. - If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. +- Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. - When using multiple tables, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. - If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. @@ -208,6 +210,7 @@ async def _classify_generation_result( - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. - Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. +- Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. - DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -337,8 +340,9 @@ async def _classify_generation_result( 22. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. 23. Use comments, aliases, display labels, and descriptions to explain why an exact schema column is relevant, not as replacement names. 24. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. -25. Do not derive executable table or column names from natural language, comments, aliases, display labels, or descriptions. Only cite exact declared names from DATABASE SCHEMA. -26. ONLY SHOWING the reasoning plan in bullet points. +25. Do not derive executable table or column names from natural language, comments, aliases, display labels, descriptions, source metadata, physical datasource names, or lineage names. Only cite exact declared names from DATABASE SCHEMA. +26. If source metadata or lineage names help identify a concept, map that concept only to exact declared DATABASE SCHEMA objects; if no exact object exists, state that the available schema does not include that part. +27. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -411,7 +415,9 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 7. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. -8. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +8. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +9. If an identifier appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request. +10. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 525ea01ed3..216ecde889 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -6,8 +6,16 @@ get_text_to_sql_rules, sql_generation_reasoning_system_prompt, ) -from src.pipelines.generation.sql_correction import get_sql_correction_system_prompt +from src.pipelines.generation.followup_sql_generation import ( + text_to_sql_with_followup_user_prompt_template, +) +from src.pipelines.generation.sql_correction import ( + get_sql_correction_system_prompt, + sql_correction_user_prompt_template, +) +from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template from src.pipelines.generation.sql_regeneration import get_sql_regeneration_system_prompt +from src.pipelines.generation.sql_regeneration import sql_regeneration_user_prompt_template class _SqlKnowledge: @@ -46,6 +54,9 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): in rules ) assert "first locate the exact declared source column" in rules + assert "Physical datasource names, source database names" in rules + assert "Do not replace an invalid identifier with a similar-looking physical" in rules + assert "source/lineage names from metadata may guide meaning" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -86,6 +97,8 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "reasoning plan only as non-executable context" in prompt assert "include those objects only when DATABASE SCHEMA shows" in prompt assert "Use the exact supported syntax shown there" in prompt + assert "source database/schema/table names" in prompt + assert "appears only in SQL samples, reasoning, failed SQL" in prompt def test_json_field_instructions_do_not_include_placeholder_identifiers(): @@ -108,6 +121,7 @@ def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): "database schema as the only source of executable table and column identifiers" in prompt ) + assert "Treat physical/source/lineage names from the original SQL" in prompt def test_sql_correction_system_prompt_discards_invalid_identifier_context(): @@ -120,6 +134,8 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): "Do not preserve a table, column, join, filter, grouping, ordering, or function" in prompt ) + assert "Treat physical/source/lineage names from the failed SQL" in prompt + assert "do not try a similar replacement from source metadata" in prompt def test_sql_reasoning_prompt_forbids_executable_sql_context(): @@ -129,3 +145,16 @@ def test_sql_reasoning_prompt_forbids_executable_sql_context(): assert "SQL clauses, SQL functions, code blocks, or executable expressions" in prompt assert "The reasoning plan is non-executable context" in prompt assert "Only cite exact declared names from DATABASE SCHEMA" in prompt + assert "source metadata, physical datasource names, or lineage names" in prompt + + +def test_user_prompt_templates_keep_source_metadata_non_executable(): + for prompt in ( + sql_generation_user_prompt_template, + text_to_sql_with_followup_user_prompt_template, + sql_regeneration_user_prompt_template, + sql_correction_user_prompt_template, + ): + assert "source/physical/lineage names" in prompt + assert "omit that unsupported part instead of inventing" in prompt + assert "exact declared table and column names from DATABASE SCHEMA" in prompt From 10417fe14ebaf77e3dcc71a6c4e04a5985fdcea5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 01:34:39 +0530 Subject: [PATCH 0715/1087] Improve SQL correction dry-run context --- wren-ai-service/src/web/v1/services/ask.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index e5554a4e1a..6ddd3b14e2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -553,9 +553,12 @@ async def ask( instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "error": ( + f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" + if allow_sql_diagnosis + and sql_diagnosis_reasoning + else error_message + ), }, project_id=ask_request.project_id, use_dry_plan=use_dry_plan, From 400d29e18284ae1d5b2f7f0dd7970d68039ff63f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 01:40:49 +0530 Subject: [PATCH 0716/1087] Add wren-engine CTE rewrite patch --- ...002-Fix-CTE-rewrite-for-set-operations.patch | Bin 0 -> 19692 bytes 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 patches/0002-Fix-CTE-rewrite-for-set-operations.patch diff --git a/patches/0002-Fix-CTE-rewrite-for-set-operations.patch b/patches/0002-Fix-CTE-rewrite-for-set-operations.patch new file mode 100644 index 0000000000000000000000000000000000000000..1fb74c171f49ce88895d7cb78494ce27de566f4d GIT binary patch literal 19692 zcmeI4TXP)6702f}RrwAZ`2k_sYwLoAsW=4~yFv;uAQ2u^t}W|k6Le!&3OIx>kq?~Y z|Le2Ir+cPncV!8i2Wo3~W@l&4{nn?u_kaI=+>D#k=C~O&_w?Fq_M2zTdwQ?xe6!ip zwY6p=U0c=lgXXZguQ%8Cb^ zkbMq>forSHADcfkPR|Fr^RI$4ZYIr(=DX%xpU-sd&pPhv>`3oPbD{IG&W`myXp0Q3a~)45pM%87x$Z$~lY|M0+)gjf1_{Q2e!ZjL$HIOmyK*e9j&uc@4w{jE&vA|= zk26VVFn#@7$NNbiE8=dEm(dZQZHmA;5$=g-xDbu-dMDAs)zMtKpvW{*QUgZAwWd1` z1@|an8VDjbaUxy2M(;~Y1KolAkY`A9Svou@Cq+j8@{@$^G#Xqo{mGk@J~DzuRKWhzA7uet9v)}0)lDu zx_ASk>2Dxfu0vWt{7^46u`0T6B_Caxo)b~**F)*sBH5(*LArge_o*xZdpMTu<8gdL zR%gi)&ZJM`3*Lgy9Zc=457JO~oh6$(7R-O?e!S`=K_&tnO|Oo0L?lHs%FpHq#X`M7r;m)!Y6PonEK+|PCGOfUTLxiVhdX;AI+hMs$rJSn zRM*X_@j`=4_I;rD7s47I4V^2xTE09B*C?bu=z0wLfubXEHhkq%uB+UZMg_$sq~#>H zEjxXoGq|{{TS{NqY&yzL} zC2#7cb79Dzcj}=1=jNAbM(xM#I>_;2^GG_sHkr%Cd~V!YEC$vrWu?cW@wapFJ5F@{ zERL!A+{U2SEZr%`OUW&{72ro$4)!z9(R_J%%H>3oin{PP>E8I~Ig>}|q15j$Q>sgE zyrJi!T6=qEP~}l)O%Y~8ypgq66n}W#QjK!;m}43sqmtuui?n*!?Z>U8d(3rglH~KY zpp`G`2h*L!`cW2zi9E#2bNL|pf3EYPKCzzVX6e7(X*Jx9e3)bd%k zxb9W_rlepb+%FYREboQ|#?rGCD;y}C{HbM4UD41 zTiwoG){eoCuK9V=4Bc4;U{W$GG&+W=R;=z=?xE~K)dP*Pu17hAZns3puD_5`>a_KF zQXkxaVjeaRa-ZFX-9hmkkVdrFFY>L&AI1~*+;vBu*MZ!YWuNuU^KqhT}- zzT18rmn0X9^l{D4>h%gWj``wF%4y@&+aepehxU0izk5||k;8r42)8vDwM8z6DOKxvm3~X+PzTk^0nn94E;qjK0}aT1NsIQJ0(+6|z~1(P-nV zkjN-SN=C-DC*cfg$p7%=nccHuj6A8cXHyt@M^>Y$jOf-{6`sqmt|zZG=S$FZHcMB5 zztoqFFXmgB>&lOH4;V^fXR4d5OO6dQKB>!ks2I!dL_r&0&XcI)MxL+B@k43XyH|{q zJX_RdH(SF_j~`8=A?cwklR1Y;vZe45llfw>FRz*Cd0FqC5v9l4+zz46J86s2;(VIy zF{jJ8f09zw?~zle_m{>W$cm{eCM`i#jMZBn89;9u<1RvuX&f_sj`2croA%( z{u+$hwH&TlI{hM}wxrqZB_gxOdXLH1laXiHTkeZ*KQ~3S7n#jvF^h|C?C#U3%iNEy zVXu+RCob0mtZUUVl=aIjE*7Qvb?OthF`hk=HG9r?iyNjo!!yz1nZ6cTiSCBy@-e+` z0YIE3x#&f8=^--oX9#jP@hZBWv$g43wV7aMmzZk_s_Di+ZS+AsMhq2eZJx_U=hPUp zdFs|9i&iAYZ%|cw6Eu31ii(5ieIyL$ zW~RgawV$WU>cjEp5FNW|V>Z90{l#pc-x@DV_0^)a&yy@NtJWQP-;fqQmUpmDgHBfX z?x6Y4pUTcGGpR9an$K`Ye2V#x$Lgs)(W;_5TC?*+9kVZbEWNDrToyG=L}eG_$9W^o zx-RotE@1uI^%@m4T^X|vx79d??7O4wzFN&;nss|k*449T%FvX1pVhL&akA&W;y!d7 z$wR%i-PRRb3*h6VE_~;r^tQkj)`>>?pEs`xC_YEznVrl=bjv;5a}g`E82QK3_LSLF zeOA_ewQeUawa^SUvMd>KoCqJYjCA!PnuR@C_2uIaI)98o;u&U}%DF+O#--Ulx^#S{ zOucIq7@)FWi*4xL??mu+&!}H_AtkfMT>JTNrIu#fLp;~@iJdH3Jr$?qKiBAHr`zGTWRz?mW`Ty^`e4_%j&)P3Dk{2}h| z6hO$dXSqas#BfZ^gF$3rq-Lc()+4w+sl>ESA&M6t*)jHW6k{gnqS_Pg!j|BZdQHsZBJ{h z*uk-p)=Mw8>ig=H)z^Ms9j>n+gAZ1!+6uD#o81w%#MqJ!lI4q*LvE{TI5I zc|}`8!)zj}X|0>i2#easev2qr{)!dhAF(F{+h^=R4Y?#-7lQD;j=O1ng4^V-{DDqm zdjw~*^v5af=3_C`*b&_bdrHRQawwy!I+sf%D(@^>rckxH#?l%l*QCm2iQ2y zqkX=DzM;E6m&Wk=Bb|H2$F6BC-R5A3L5#oVt$?rNxNKf) zb5my72jZq(4THBkqOQkh%gu{er|q)1nmxOm+tJbIuXj=n(CxgG&kd6v*$uKI?qeLi zq`8~u+C>*7b^?`^U7Ks5g_I1Mb>~e7u~(oj;a`W!peV1s`|8iS&Dkur(-+6mle${5 zUH^E)PpbKrQY_;eV-+fN{z0HSma};ivVuQ{H`e?QYP|+vxBC58H30t(aCaI<-e%-r z+w|jed3U!!wcIbtFgtj|HX^~BLC#IIz9Kx0oFF!NmyjL4qrI#iVNNyJjaNX&Mu zP|ls?Cet(4O~-~Afr*}>3OasOGv^%=_kTPd|K~qLR#Zi->UB%~xOIKMz6jv@!|PtY zlJ_uC%FLau!hWQ9$E>maKLK;B{nB)~j+Hxrlke&#W^QTE(r0NG)3!LFyW}e`SXqA} zogw8P^~nsX&%{Lkqx6ujrP%rjwVEhadYir2YB1JkV=*uEl>O@qyBc0p?5FDy zW;`}3xtUi^F9!YcC(t9cY(00_AhOw^y`qK{OotAar^LUEo zSj!u^y|gp%pn0c-Zd!dq~#`+H?nOw+*j!oR(QQ+t31Uo!c^BxLFK#Y0;~)(U7W9J`V$U zJRjfC+1uOOKMS)_34K0lW_)bXah$Ra|1U!HWL5&N+10E|WcX$MX8(=ip3b_0pQiB* e%qD|1 Date: Wed, 29 Jul 2026 01:41:30 +0530 Subject: [PATCH 0717/1087] Store wren-engine CTE patch as text --- ...2-Fix-CTE-rewrite-for-set-operations.patch | Bin 19692 -> 9620 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/patches/0002-Fix-CTE-rewrite-for-set-operations.patch b/patches/0002-Fix-CTE-rewrite-for-set-operations.patch index 1fb74c171f49ce88895d7cb78494ce27de566f4d..80a84bd81f3be6d1731c2463888289c92a7faa7a 100644 GIT binary patch literal 9620 zcmds7-BQ~~7QXjWbU7De2P6CgICc$FnPs5DR7hq5*<4VRN|xGK8(VVJQh?3mb{}9L zai3(rbGoH&85@(#-0Z3-Y)Plj|6iZ>X_hX;;o&d}4ualr7#{ZbgM;3IoQ|g9Vbq(x zm=2Mvn>ZJNX8KX?C);5faT3WpaoC^8;eIGR z?**(m8px68b-RcA`=qF+&}u+3@f3P5Uy9w{Vdqfn;QvtIc1mugKFx3DX)fZ$GR<-! z@0K00go{W&O=P16DnU|#5BADZ=A4f9D<7q52RpMCtawDow1rSZ|{V<&u zi=;xb@_6g=JeZJCjTq%>{j{;;it31=LbztRCeMP;kV5+WbcLV_XL1pgpy$DI87DUv z{IDS!;+Z%}RFO%M&!qSO?-o%U%n=+SRtSfiSmiR4k+?kJ2- z-hk{AA~+QV)H@wd*|NXp9BlJ@{WKCC8H|~aeW09djxwej>H=_Z)2y# zMzLyPJEAg2f)JihJ+co8q_yK~Py^JmXb2sYdfy@H8OI_~R2HJ`AvG*it39^QXjKg( z7$ojJWRFNrgJPbyT;{%fZPynQ_nH}ma>@G&fd@oLA8@%!8#`9LnaqnU5nqrmNl}7S zO_)_4WLkD5GKJZSEEZ#+G-+kW>P6Ca?Oqi=5z9W1%}Iuq+Vf>#X^+B)uB3I z@IFAbEeF++HCo%6hftD9m@p#&Ne7u@8W)ryV5SvVXj= zF{ZBQO63`)9x|e1TO?a58#?B|bwuvqUGP%}3~Hg6*3e^X2NKvMFNdHFGj7Rr)k{RZ z;%5x`QTP!j)~##x3pVc)1Y?bzoIjdX;c% zgjXyhw9>6Q%iiZ;UPx5H6)p*ZL@%_N>@gRtHphU&+e6+76->Y!d-*I@mhFeYn*`zJ z30?VA8^9o=KI^@1o`QQlRB^cw23jI48d{T2;` zcw+Ygl}(Yn74|d~bthF$U-d5=!J0-}?EZw7FXX81INzvyxW0CzrM7150N6zu6uSQM z)fMa;Bet&Qzb15S>qWs+r%*YwBcryXZa^Px;jH;+iOC{${Ar-7%88&LKzV#$=Y}|P1V=)yB_(c!Dm3l7^K*y89rlB)f_>_zLY6(`k!~Q#2|U`P zpgOA+hp!a8E?q!n4v*MtaJNdaHKE5XT5k^Qv2BA%ZH&1V6)`D&mBc&m7MFt@1`!(z#v^!;!+5*G@K~VqEZKgiaqBp z=8|d$s!at>7BIHrRc!Y9Yf1`aP3m^w2PBrKU7mslU8#PJSJ05N;e;)%GKejfqa#{= zN2s*OLdkx`3PO&erb|Y?exo*QQHd8~!)bFElP<2MbR<;ms`6D$ZsrG0ss-0~Mw+?? z>IS&tTWTChEssJHfCaRmTv4`kLp4?*w=1(*YFmm`lalms`~W+x=j`)v0PXk6a#5y+BH*J6LL7NUX*TB02nI`m zVAIg-34m1D9%8`Mjl%)rP+U29oQiF{^)$f5NrQ7|i~|{jnX{z6|7A1v4%kfmQJU$m zu!u2km=*~i&?v{gj~O&hawFLPyg9ns6vjBzo^LQkTb~B>dBEG_kJ1$1sXv{L5Jre?oL28gVE`JB$Z0MuV`6-(L$NAkry%A?d!|BdQ+McxqX>tVPxi>8))Xg@3arRn4a%?%;9?u$}nR zte=fYd6BvoKQWVL>d*BsU-?sPS|&mGkw)25^#8nq{-Hzoy^z^`HTlOu8LF|_;0RiW zgC?vLd5T5>baR@tCkp9dLD`pcAC=odcSI+?gI=%GcTRjCF*}JSmiNT8mjyAhMC;@( zlsK196PHKgGrRT7FDhEii<8$U$CucR(SSh|$;TwUO?=*DQ0~xnsMLYkPnwh?Q?z6H zWZtI%l`n&E=HnM5NjNJtQ|TX?xsds8^>YuQ;c(<<5)m!t4lIi&1!9-GZJPJc$v>;5 z!j;{^5Bz$I=%2rTeR<}ezy0OpwSWHeo3qoCi%VyYq%#)QIigz^Bc1$U6p>;+h3KR+ zOV&sktm20i`1fzl-o6oLRp^HZ8oP+@kXo~*v#K_f)+vC?E=fZvqAI~xgd3mT&OAq8@LjQru7(P(<-e=S=XWPfO@l16@%7@JOnoC?KSmy z1qb*>vCI^N<6l_(*>iJeJ^f&VqD+n9d|4*vn0!V8zl6eQ7J=QOc81J6jqlLf?}|p> SI~aIf?-r>&y;WYn(fBXE(`Ft3 literal 19692 zcmeI4TXP)6702f}RrwAZ`2k_sYwLoAsW=4~yFv;uAQ2u^t}W|k6Le!&3OIx>kq?~Y z|Le2Ir+cPncV!8i2Wo3~W@l&4{nn?u_kaI=+>D#k=C~O&_w?Fq_M2zTdwQ?xe6!ip zwY6p=U0c=lgXXZguQ%8Cb^ zkbMq>forSHADcfkPR|Fr^RI$4ZYIr(=DX%xpU-sd&pPhv>`3oPbD{IG&W`myXp0Q3a~)45pM%87x$Z$~lY|M0+)gjf1_{Q2e!ZjL$HIOmyK*e9j&uc@4w{jE&vA|= zk26VVFn#@7$NNbiE8=dEm(dZQZHmA;5$=g-xDbu-dMDAs)zMtKpvW{*QUgZAwWd1` z1@|an8VDjbaUxy2M(;~Y1KolAkY`A9Svou@Cq+j8@{@$^G#Xqo{mGk@J~DzuRKWhzA7uet9v)}0)lDu zx_ASk>2Dxfu0vWt{7^46u`0T6B_Caxo)b~**F)*sBH5(*LArge_o*xZdpMTu<8gdL zR%gi)&ZJM`3*Lgy9Zc=457JO~oh6$(7R-O?e!S`=K_&tnO|Oo0L?lHs%FpHq#X`M7r;m)!Y6PonEK+|PCGOfUTLxiVhdX;AI+hMs$rJSn zRM*X_@j`=4_I;rD7s47I4V^2xTE09B*C?bu=z0wLfubXEHhkq%uB+UZMg_$sq~#>H zEjxXoGq|{{TS{NqY&yzL} zC2#7cb79Dzcj}=1=jNAbM(xM#I>_;2^GG_sHkr%Cd~V!YEC$vrWu?cW@wapFJ5F@{ zERL!A+{U2SEZr%`OUW&{72ro$4)!z9(R_J%%H>3oin{PP>E8I~Ig>}|q15j$Q>sgE zyrJi!T6=qEP~}l)O%Y~8ypgq66n}W#QjK!;m}43sqmtuui?n*!?Z>U8d(3rglH~KY zpp`G`2h*L!`cW2zi9E#2bNL|pf3EYPKCzzVX6e7(X*Jx9e3)bd%k zxb9W_rlepb+%FYREboQ|#?rGCD;y}C{HbM4UD41 zTiwoG){eoCuK9V=4Bc4;U{W$GG&+W=R;=z=?xE~K)dP*Pu17hAZns3puD_5`>a_KF zQXkxaVjeaRa-ZFX-9hmkkVdrFFY>L&AI1~*+;vBu*MZ!YWuNuU^KqhT}- zzT18rmn0X9^l{D4>h%gWj``wF%4y@&+aepehxU0izk5||k;8r42)8vDwM8z6DOKxvm3~X+PzTk^0nn94E;qjK0}aT1NsIQJ0(+6|z~1(P-nV zkjN-SN=C-DC*cfg$p7%=nccHuj6A8cXHyt@M^>Y$jOf-{6`sqmt|zZG=S$FZHcMB5 zztoqFFXmgB>&lOH4;V^fXR4d5OO6dQKB>!ks2I!dL_r&0&XcI)MxL+B@k43XyH|{q zJX_RdH(SF_j~`8=A?cwklR1Y;vZe45llfw>FRz*Cd0FqC5v9l4+zz46J86s2;(VIy zF{jJ8f09zw?~zle_m{>W$cm{eCM`i#jMZBn89;9u<1RvuX&f_sj`2croA%( z{u+$hwH&TlI{hM}wxrqZB_gxOdXLH1laXiHTkeZ*KQ~3S7n#jvF^h|C?C#U3%iNEy zVXu+RCob0mtZUUVl=aIjE*7Qvb?OthF`hk=HG9r?iyNjo!!yz1nZ6cTiSCBy@-e+` z0YIE3x#&f8=^--oX9#jP@hZBWv$g43wV7aMmzZk_s_Di+ZS+AsMhq2eZJx_U=hPUp zdFs|9i&iAYZ%|cw6Eu31ii(5ieIyL$ zW~RgawV$WU>cjEp5FNW|V>Z90{l#pc-x@DV_0^)a&yy@NtJWQP-;fqQmUpmDgHBfX z?x6Y4pUTcGGpR9an$K`Ye2V#x$Lgs)(W;_5TC?*+9kVZbEWNDrToyG=L}eG_$9W^o zx-RotE@1uI^%@m4T^X|vx79d??7O4wzFN&;nss|k*449T%FvX1pVhL&akA&W;y!d7 z$wR%i-PRRb3*h6VE_~;r^tQkj)`>>?pEs`xC_YEznVrl=bjv;5a}g`E82QK3_LSLF zeOA_ewQeUawa^SUvMd>KoCqJYjCA!PnuR@C_2uIaI)98o;u&U}%DF+O#--Ulx^#S{ zOucIq7@)FWi*4xL??mu+&!}H_AtkfMT>JTNrIu#fLp;~@iJdH3Jr$?qKiBAHr`zGTWRz?mW`Ty^`e4_%j&)P3Dk{2}h| z6hO$dXSqas#BfZ^gF$3rq-Lc()+4w+sl>ESA&M6t*)jHW6k{gnqS_Pg!j|BZdQHsZBJ{h z*uk-p)=Mw8>ig=H)z^Ms9j>n+gAZ1!+6uD#o81w%#MqJ!lI4q*LvE{TI5I zc|}`8!)zj}X|0>i2#easev2qr{)!dhAF(F{+h^=R4Y?#-7lQD;j=O1ng4^V-{DDqm zdjw~*^v5af=3_C`*b&_bdrHRQawwy!I+sf%D(@^>rckxH#?l%l*QCm2iQ2y zqkX=DzM;E6m&Wk=Bb|H2$F6BC-R5A3L5#oVt$?rNxNKf) zb5my72jZq(4THBkqOQkh%gu{er|q)1nmxOm+tJbIuXj=n(CxgG&kd6v*$uKI?qeLi zq`8~u+C>*7b^?`^U7Ks5g_I1Mb>~e7u~(oj;a`W!peV1s`|8iS&Dkur(-+6mle${5 zUH^E)PpbKrQY_;eV-+fN{z0HSma};ivVuQ{H`e?QYP|+vxBC58H30t(aCaI<-e%-r z+w|jed3U!!wcIbtFgtj|HX^~BLC#IIz9Kx0oFF!NmyjL4qrI#iVNNyJjaNX&Mu zP|ls?Cet(4O~-~Afr*}>3OasOGv^%=_kTPd|K~qLR#Zi->UB%~xOIKMz6jv@!|PtY zlJ_uC%FLau!hWQ9$E>maKLK;B{nB)~j+Hxrlke&#W^QTE(r0NG)3!LFyW}e`SXqA} zogw8P^~nsX&%{Lkqx6ujrP%rjwVEhadYir2YB1JkV=*uEl>O@qyBc0p?5FDy zW;`}3xtUi^F9!YcC(t9cY(00_AhOw^y`qK{OotAar^LUEo zSj!u^y|gp%pn0c-Zd!dq~#`+H?nOw+*j!oR(QQ+t31Uo!c^BxLFK#Y0;~)(U7W9J`V$U zJRjfC+1uOOKMS)_34K0lW_)bXah$Ra|1U!HWL5&N+10E|WcX$MX8(=ip3b_0pQiB* e%qD|1 Date: Wed, 29 Jul 2026 02:06:41 +0530 Subject: [PATCH 0718/1087] Keep SQL reasoning non-executable --- .../generation/followup_sql_generation.py | 4 +- .../followup_sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_correction.py | 2 +- .../pipelines/generation/sql_generation.py | 2 +- .../generation/sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_regeneration.py | 2 +- .../src/pipelines/generation/utils/sql.py | 40 ++++++++----------- 7 files changed, 23 insertions(+), 31 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 e7efdec613..b7770d1395 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -81,8 +81,8 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. ### REASONING PLAN ### -Use this reasoning plan only as non-executable context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. -Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. +Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 2a2e35c5ce..45778b82e1 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -60,7 +60,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. +Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 30f5809890..977a688807 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -87,7 +87,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} {% endif %} ### FAILED SQL ### diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 03a1a1ae90..8d3275431b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -76,7 +76,7 @@ {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} {% endif %} 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 e6d5c2c625..3f24a526e1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -50,7 +50,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. +Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index b508d7135c..d0e9d403f4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -100,7 +100,7 @@ def get_sql_regeneration_system_prompt( Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} ### ORIGINAL SQL QUERY ### Use this SQL only as non-executable intent context. Do not preserve any identifier, source/physical/lineage name, or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b758c36504..b3fd269a4a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -311,38 +311,30 @@ async def _classify_generation_result( sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who maps a user's intent to the provided database schema and provides a concise reasoning plan for answering the user's question. +You are a helpful data analyst who explains the user's analytical intent and provides a concise, non-executable reasoning plan for answering the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will use that absolute time frame only with an exact date/time schema column and supported SQL FUNCTIONS; -otherwise, you will use a relative timeframe only when an exact date/time schema column and supported SQL FUNCTIONS are available. If they are not available, state that the available schema/functions do not support that time filter. -3. For top, bottom, first, last, highest, or lowest requests, plan to sort by an exact selected column or aggregate alias and limit the result. Include a ranking column only when the user explicitly asks for rank values. -4. Do not plan to use a SQL function unless it appears in SQL FUNCTIONS for this request or is already part of a valid metric/calculated-field definition in DATABASE SCHEMA. +2. Explicitly state requested timeframes in natural language only. Do not name date columns, functions, expressions, or SQL clauses. +3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language only. Do not name columns, aliases, aggregate expressions, or SQL clauses. +4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. 8. The reasoning plan should be in the language same as the language user provided in the input. -9. Don't include SQL in the reasoning plan. +9. Do not include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. A table name in the reasoning plan must be in this format: `table: `. -13. A column name in the reasoning plan must be in this format: `column: .`. -14. Use only exact table and column names that appear in the DATABASE SCHEMA section. -15. Comments, aliases, display labels, and descriptions are semantic hints only; do not turn them into table or column names in the reasoning plan. -16. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. -17. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the schema does not show the exact table or column needed, state that the available schema does not include that part. -18. If the question asks for a concept, filter, sort, or timeframe, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. -19. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but name only exact tables and columns from DATABASE SCHEMA in the reasoning plan. -20. If multiple schema objects are required to answer the intent, include each required object only when DATABASE SCHEMA provides both the needed fields and the relationship path between them. -21. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, or functions from them unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS. -22. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -23. Use comments, aliases, display labels, and descriptions to explain why an exact schema column is relevant, not as replacement names. -24. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. -25. Do not derive executable table or column names from natural language, comments, aliases, display labels, descriptions, source metadata, physical datasource names, or lineage names. Only cite exact declared names from DATABASE SCHEMA. -26. If source metadata or lineage names help identify a concept, map that concept only to exact declared DATABASE SCHEMA objects; if no exact object exists, state that the available schema does not include that part. -27. ONLY SHOWING the reasoning plan in bullet points. +12. Do not mention table names, view names, metric names, column names, aliases, source names, physical names, lineage names, schema names, database names, or identifier-like labels in the reasoning plan. +13. Do not write possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, executable expressions, or date/time expressions in the reasoning plan. +14. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. +15. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language only. +16. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but do not name the underlying schema objects in the reasoning plan. +17. If multiple schema objects may be required to answer the intent, describe the need to combine related data in natural language only. +18. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, functions, or SQL patterns from them into the reasoning plan. +19. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +20. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. +21. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -411,7 +403,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains SQL fragments, assumed SQL, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers missing from DATABASE SCHEMA, ignore those parts. +4. YOU MUST use the reasoning plan only as non-executable intent context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 7. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. From 376043c8a707c4927c0d7f47bcaa4e4eb132462a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 02:23:27 +0530 Subject: [PATCH 0719/1087] Add MSSQL date transpilation engine patch --- .../0003-Fix-MSSQL-date-transpilation.patch | 175 ++++++++++++++++++ 1 file changed, 175 insertions(+) create mode 100644 patches/0003-Fix-MSSQL-date-transpilation.patch diff --git a/patches/0003-Fix-MSSQL-date-transpilation.patch b/patches/0003-Fix-MSSQL-date-transpilation.patch new file mode 100644 index 0000000000..7667fb9e94 --- /dev/null +++ b/patches/0003-Fix-MSSQL-date-transpilation.patch @@ -0,0 +1,175 @@ +From 20c40106e6de7a5696fe25d6d2d8ff88eb9221ed Mon Sep 17 00:00:00 2001 +From: Harshitha +Date: Wed, 29 Jul 2026 02:23:00 +0530 +Subject: [PATCH] Fix MSSQL date transpilation + +--- + ibis-server/app/mdl/rewriter.py | 99 ++++++++++++++++++- + .../tests/model/test_mssql_connector.py | 15 +++ + 2 files changed, 113 insertions(+), 1 deletion(-) + +diff --git a/ibis-server/app/mdl/rewriter.py b/ibis-server/app/mdl/rewriter.py +index 3fe7e8af..6515c9de 100644 +--- a/ibis-server/app/mdl/rewriter.py ++++ b/ibis-server/app/mdl/rewriter.py +@@ -2,6 +2,7 @@ import importlib + + import httpx + import sqlglot ++from sqlglot import exp + from anyio import to_thread + from loguru import logger + from opentelemetry import trace +@@ -27,6 +28,97 @@ importlib.import_module("app.custom_sqlglot.dialects") + tracer = trace.get_tracer(__name__) + + ++def rewrite_mssql_date_operations(sql: str) -> str: ++ ast = sqlglot.parse_one(sql, dialect="tsql") ++ ast = ast.transform(_rewrite_mssql_interval_arithmetic) ++ ast = ast.transform(_rewrite_mssql_date_trunc) ++ return ast.sql(dialect="tsql") ++ ++ ++def _rewrite_mssql_interval_arithmetic(node: exp.Expression) -> exp.Expression: ++ if isinstance(node, exp.Sub) and isinstance(node.expression, exp.Interval): ++ return _mssql_dateadd( ++ node.expression, ++ node.this.copy(), ++ negative=True, ++ ) or node ++ ++ if isinstance(node, exp.Add): ++ if isinstance(node.expression, exp.Interval): ++ return _mssql_dateadd(node.expression, node.this.copy()) or node ++ if isinstance(node.this, exp.Interval): ++ return _mssql_dateadd(node.this, node.expression.copy()) or node ++ ++ return node ++ ++ ++def _rewrite_mssql_date_trunc(node: exp.Expression) -> exp.Expression: ++ if not isinstance(node, (exp.DateTrunc, exp.TimestampTrunc)): ++ return node ++ ++ unit = _mssql_interval_unit(node.args.get("unit")) ++ if unit is None: ++ return node ++ ++ value = node.this.this.copy() if isinstance(node.this, exp.Paren) else node.this.copy() ++ return exp.Anonymous( ++ this="DATEADD", ++ expressions=[ ++ unit, ++ exp.Anonymous( ++ this="DATEDIFF", ++ expressions=[ ++ unit.copy(), ++ exp.Literal.number(0), ++ value, ++ ], ++ ), ++ exp.Literal.number(0), ++ ], ++ ) ++ ++ ++def _mssql_dateadd( ++ interval: exp.Interval, date_expression: exp.Expression, *, negative: bool = False ++) -> exp.Expression | None: ++ amount = _mssql_interval_amount(interval) ++ if negative: ++ amount = exp.Neg(this=amount) ++ ++ unit = _mssql_interval_unit(interval.args.get("unit")) ++ if unit is None: ++ return None ++ ++ return exp.Anonymous( ++ this="DATEADD", ++ expressions=[ ++ unit, ++ amount, ++ date_expression, ++ ], ++ ) ++ ++ ++def _mssql_interval_amount(interval: exp.Interval) -> exp.Expression: ++ amount = interval.this.copy() ++ if isinstance(amount, exp.Literal) and amount.is_string: ++ text = str(amount.this).strip() ++ try: ++ return exp.Literal.number(int(text)) ++ except ValueError: ++ try: ++ return exp.Literal.number(float(text)) ++ except ValueError: ++ return amount ++ return amount ++ ++ ++def _mssql_interval_unit(unit: exp.Expression | None) -> exp.Var | None: ++ if unit is None: ++ return None ++ return exp.Var(this=unit.name.upper()) ++ ++ + class Rewriter: + def __init__( + self, +@@ -53,7 +145,10 @@ class Rewriter: + try: + read = self._get_read_dialect(self.experiment) + write = self._get_write_dialect(self.data_source) +- return sqlglot.transpile(planned_sql, read=read, write=write)[0] ++ dialect_sql = sqlglot.transpile(planned_sql, read=read, write=write)[0] ++ if self.data_source == DataSource.mssql: ++ return rewrite_mssql_date_operations(dialect_sql) ++ return dialect_sql + except Exception as e: + raise WrenError( + ErrorCode.SQLGLOT_ERROR, +@@ -99,6 +194,8 @@ class Rewriter: + DataSource.gcs_file, + }: + return "duckdb" ++ elif data_source == DataSource.mssql: ++ return "tsql" + return data_source.name + + +diff --git a/ibis-server/tests/model/test_mssql_connector.py b/ibis-server/tests/model/test_mssql_connector.py +index e1d5d204..f26332e7 100644 +--- a/ibis-server/tests/model/test_mssql_connector.py ++++ b/ibis-server/tests/model/test_mssql_connector.py +@@ -1,4 +1,6 @@ ++from app.mdl.rewriter import Rewriter, rewrite_mssql_date_operations + from app.model.connector import MSSqlConnector ++from app.model.data_source import DataSource + from app.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError + + +@@ -44,6 +46,19 @@ def _connector(connection): + return connector + + ++def test_mssql_rewriter_uses_tsql_dialect(): ++ assert Rewriter._get_write_dialect(DataSource.mssql) == "tsql" ++ ++ ++def test_rewrite_mssql_date_operations_uses_tsql_dateadd(): ++ assert rewrite_mssql_date_operations( ++ "SELECT TIMESTAMP_TRUNC((GETDATE() - INTERVAL '1' MONTH), MONTH)" ++ ) == ( ++ "SELECT DATEADD(MONTH, DATEDIFF(MONTH, 0, " ++ "DATEADD(MONTH, -1, GETDATE())), 0)" ++ ) ++ ++ + def test_query_uses_raw_sql_for_grouped_aggregate_results(): + connection = FakeConnection( + [ +-- +2.53.0.windows.2 + From a6fa8d913dde94eae847a13804096bb202836e91 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 02:39:22 +0530 Subject: [PATCH 0720/1087] Revert "Add MSSQL date transpilation engine patch" This reverts commit 376043c8a707c4927c0d7f47bcaa4e4eb132462a. --- .../0003-Fix-MSSQL-date-transpilation.patch | 175 ------------------ 1 file changed, 175 deletions(-) delete mode 100644 patches/0003-Fix-MSSQL-date-transpilation.patch diff --git a/patches/0003-Fix-MSSQL-date-transpilation.patch b/patches/0003-Fix-MSSQL-date-transpilation.patch deleted file mode 100644 index 7667fb9e94..0000000000 --- a/patches/0003-Fix-MSSQL-date-transpilation.patch +++ /dev/null @@ -1,175 +0,0 @@ -From 20c40106e6de7a5696fe25d6d2d8ff88eb9221ed Mon Sep 17 00:00:00 2001 -From: Harshitha -Date: Wed, 29 Jul 2026 02:23:00 +0530 -Subject: [PATCH] Fix MSSQL date transpilation - ---- - ibis-server/app/mdl/rewriter.py | 99 ++++++++++++++++++- - .../tests/model/test_mssql_connector.py | 15 +++ - 2 files changed, 113 insertions(+), 1 deletion(-) - -diff --git a/ibis-server/app/mdl/rewriter.py b/ibis-server/app/mdl/rewriter.py -index 3fe7e8af..6515c9de 100644 ---- a/ibis-server/app/mdl/rewriter.py -+++ b/ibis-server/app/mdl/rewriter.py -@@ -2,6 +2,7 @@ import importlib - - import httpx - import sqlglot -+from sqlglot import exp - from anyio import to_thread - from loguru import logger - from opentelemetry import trace -@@ -27,6 +28,97 @@ importlib.import_module("app.custom_sqlglot.dialects") - tracer = trace.get_tracer(__name__) - - -+def rewrite_mssql_date_operations(sql: str) -> str: -+ ast = sqlglot.parse_one(sql, dialect="tsql") -+ ast = ast.transform(_rewrite_mssql_interval_arithmetic) -+ ast = ast.transform(_rewrite_mssql_date_trunc) -+ return ast.sql(dialect="tsql") -+ -+ -+def _rewrite_mssql_interval_arithmetic(node: exp.Expression) -> exp.Expression: -+ if isinstance(node, exp.Sub) and isinstance(node.expression, exp.Interval): -+ return _mssql_dateadd( -+ node.expression, -+ node.this.copy(), -+ negative=True, -+ ) or node -+ -+ if isinstance(node, exp.Add): -+ if isinstance(node.expression, exp.Interval): -+ return _mssql_dateadd(node.expression, node.this.copy()) or node -+ if isinstance(node.this, exp.Interval): -+ return _mssql_dateadd(node.this, node.expression.copy()) or node -+ -+ return node -+ -+ -+def _rewrite_mssql_date_trunc(node: exp.Expression) -> exp.Expression: -+ if not isinstance(node, (exp.DateTrunc, exp.TimestampTrunc)): -+ return node -+ -+ unit = _mssql_interval_unit(node.args.get("unit")) -+ if unit is None: -+ return node -+ -+ value = node.this.this.copy() if isinstance(node.this, exp.Paren) else node.this.copy() -+ return exp.Anonymous( -+ this="DATEADD", -+ expressions=[ -+ unit, -+ exp.Anonymous( -+ this="DATEDIFF", -+ expressions=[ -+ unit.copy(), -+ exp.Literal.number(0), -+ value, -+ ], -+ ), -+ exp.Literal.number(0), -+ ], -+ ) -+ -+ -+def _mssql_dateadd( -+ interval: exp.Interval, date_expression: exp.Expression, *, negative: bool = False -+) -> exp.Expression | None: -+ amount = _mssql_interval_amount(interval) -+ if negative: -+ amount = exp.Neg(this=amount) -+ -+ unit = _mssql_interval_unit(interval.args.get("unit")) -+ if unit is None: -+ return None -+ -+ return exp.Anonymous( -+ this="DATEADD", -+ expressions=[ -+ unit, -+ amount, -+ date_expression, -+ ], -+ ) -+ -+ -+def _mssql_interval_amount(interval: exp.Interval) -> exp.Expression: -+ amount = interval.this.copy() -+ if isinstance(amount, exp.Literal) and amount.is_string: -+ text = str(amount.this).strip() -+ try: -+ return exp.Literal.number(int(text)) -+ except ValueError: -+ try: -+ return exp.Literal.number(float(text)) -+ except ValueError: -+ return amount -+ return amount -+ -+ -+def _mssql_interval_unit(unit: exp.Expression | None) -> exp.Var | None: -+ if unit is None: -+ return None -+ return exp.Var(this=unit.name.upper()) -+ -+ - class Rewriter: - def __init__( - self, -@@ -53,7 +145,10 @@ class Rewriter: - try: - read = self._get_read_dialect(self.experiment) - write = self._get_write_dialect(self.data_source) -- return sqlglot.transpile(planned_sql, read=read, write=write)[0] -+ dialect_sql = sqlglot.transpile(planned_sql, read=read, write=write)[0] -+ if self.data_source == DataSource.mssql: -+ return rewrite_mssql_date_operations(dialect_sql) -+ return dialect_sql - except Exception as e: - raise WrenError( - ErrorCode.SQLGLOT_ERROR, -@@ -99,6 +194,8 @@ class Rewriter: - DataSource.gcs_file, - }: - return "duckdb" -+ elif data_source == DataSource.mssql: -+ return "tsql" - return data_source.name - - -diff --git a/ibis-server/tests/model/test_mssql_connector.py b/ibis-server/tests/model/test_mssql_connector.py -index e1d5d204..f26332e7 100644 ---- a/ibis-server/tests/model/test_mssql_connector.py -+++ b/ibis-server/tests/model/test_mssql_connector.py -@@ -1,4 +1,6 @@ -+from app.mdl.rewriter import Rewriter, rewrite_mssql_date_operations - from app.model.connector import MSSqlConnector -+from app.model.data_source import DataSource - from app.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError - - -@@ -44,6 +46,19 @@ def _connector(connection): - return connector - - -+def test_mssql_rewriter_uses_tsql_dialect(): -+ assert Rewriter._get_write_dialect(DataSource.mssql) == "tsql" -+ -+ -+def test_rewrite_mssql_date_operations_uses_tsql_dateadd(): -+ assert rewrite_mssql_date_operations( -+ "SELECT TIMESTAMP_TRUNC((GETDATE() - INTERVAL '1' MONTH), MONTH)" -+ ) == ( -+ "SELECT DATEADD(MONTH, DATEDIFF(MONTH, 0, " -+ "DATEADD(MONTH, -1, GETDATE())), 0)" -+ ) -+ -+ - def test_query_uses_raw_sql_for_grouped_aggregate_results(): - connection = FakeConnection( - [ --- -2.53.0.windows.2 - From 3ebf125e62ad0baf46602716d07030fcc3a9e5da Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 02:39:23 +0530 Subject: [PATCH 0721/1087] Revert "Keep SQL reasoning non-executable" This reverts commit 4ffb6d43a8cd2a214957935b78d99b48b68a9a85. --- .../generation/followup_sql_generation.py | 4 +- .../followup_sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_correction.py | 2 +- .../pipelines/generation/sql_generation.py | 2 +- .../generation/sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_regeneration.py | 2 +- .../src/pipelines/generation/utils/sql.py | 40 +++++++++++-------- 7 files changed, 31 insertions(+), 23 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 b7770d1395..e7efdec613 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -81,8 +81,8 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. ### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. -Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +Use this reasoning plan only as non-executable context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. +Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 45778b82e1..2a2e35c5ce 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -60,7 +60,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, functions, or identifier-like labels. +Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 977a688807..30f5809890 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -87,7 +87,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} {% endif %} ### FAILED SQL ### diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 8d3275431b..03a1a1ae90 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -76,7 +76,7 @@ {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} {% endif %} 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 3f24a526e1..e6d5c2c625 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -50,7 +50,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, functions, or identifier-like labels. +Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index d0e9d403f4..b508d7135c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -100,7 +100,7 @@ def get_sql_regeneration_system_prompt( Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. {{ sql_generation_reasoning }} ### ORIGINAL SQL QUERY ### Use this SQL only as non-executable intent context. Do not preserve any identifier, source/physical/lineage name, or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b3fd269a4a..b758c36504 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -311,30 +311,38 @@ async def _classify_generation_result( sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who explains the user's analytical intent and provides a concise, non-executable reasoning plan for answering the user's question. +You are a helpful data analyst who maps a user's intent to the provided database schema and provides a concise reasoning plan for answering the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state requested timeframes in natural language only. Do not name date columns, functions, expressions, or SQL clauses. -3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language only. Do not name columns, aliases, aggregate expressions, or SQL clauses. -4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. +2. Explicitly state the following information in the reasoning plan: +if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will use that absolute time frame only with an exact date/time schema column and supported SQL FUNCTIONS; +otherwise, you will use a relative timeframe only when an exact date/time schema column and supported SQL FUNCTIONS are available. If they are not available, state that the available schema/functions do not support that time filter. +3. For top, bottom, first, last, highest, or lowest requests, plan to sort by an exact selected column or aggregate alias and limit the result. Include a ranking column only when the user explicitly asks for rank values. +4. Do not plan to use a SQL function unless it appears in SQL FUNCTIONS for this request or is already part of a valid metric/calculated-field definition in DATABASE SCHEMA. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. 8. The reasoning plan should be in the language same as the language user provided in the input. -9. Do not include SQL in the reasoning plan. +9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. Do not mention table names, view names, metric names, column names, aliases, source names, physical names, lineage names, schema names, database names, or identifier-like labels in the reasoning plan. -13. Do not write possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, executable expressions, or date/time expressions in the reasoning plan. -14. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. -15. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language only. -16. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but do not name the underlying schema objects in the reasoning plan. -17. If multiple schema objects may be required to answer the intent, describe the need to combine related data in natural language only. -18. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, functions, or SQL patterns from them into the reasoning plan. -19. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -20. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. -21. ONLY SHOWING the reasoning plan in bullet points. +12. A table name in the reasoning plan must be in this format: `table: `. +13. A column name in the reasoning plan must be in this format: `column: .`. +14. Use only exact table and column names that appear in the DATABASE SCHEMA section. +15. Comments, aliases, display labels, and descriptions are semantic hints only; do not turn them into table or column names in the reasoning plan. +16. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. +17. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the schema does not show the exact table or column needed, state that the available schema does not include that part. +18. If the question asks for a concept, filter, sort, or timeframe, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. +19. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but name only exact tables and columns from DATABASE SCHEMA in the reasoning plan. +20. If multiple schema objects are required to answer the intent, include each required object only when DATABASE SCHEMA provides both the needed fields and the relationship path between them. +21. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, or functions from them unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS. +22. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +23. Use comments, aliases, display labels, and descriptions to explain why an exact schema column is relevant, not as replacement names. +24. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. +25. Do not derive executable table or column names from natural language, comments, aliases, display labels, descriptions, source metadata, physical datasource names, or lineage names. Only cite exact declared names from DATABASE SCHEMA. +26. If source metadata or lineage names help identify a concept, map that concept only to exact declared DATABASE SCHEMA objects; if no exact object exists, state that the available schema does not include that part. +27. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -403,7 +411,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -4. YOU MUST use the reasoning plan only as non-executable intent context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains SQL fragments, assumed SQL, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers missing from DATABASE SCHEMA, ignore those parts. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 7. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. From 937df1b0a57c7560a00af4cc4639ccef72b51fc2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 02:54:34 +0530 Subject: [PATCH 0722/1087] Keep SQL generation grounded to current schema --- .../generation/followup_sql_generation.py | 8 +-- .../followup_sql_generation_reasoning.py | 10 +-- .../generation/intent_classification.py | 5 +- .../pipelines/generation/sql_correction.py | 5 +- .../pipelines/generation/sql_generation.py | 6 +- .../generation/sql_generation_reasoning.py | 6 +- .../pipelines/generation/sql_regeneration.py | 9 +-- .../src/pipelines/generation/utils/sql.py | 67 ++++++++----------- 8 files changed, 43 insertions(+), 73 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 e7efdec613..5beb9cae14 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -60,12 +60,10 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. Treat physical/source/lineage names in samples as non-executable unless the exact same identifier is declared in DATABASE SCHEMA. +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Summary: {{sample.summary}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -81,8 +79,8 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. ### REASONING PLAN ### -Use this reasoning plan only as non-executable context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. -Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. +Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 2a2e35c5ce..4e285a663b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -30,12 +30,10 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use table names, column names, values, and functions only if they are present in the current DATABASE SCHEMA or SQL FUNCTIONS. +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} @@ -47,12 +45,10 @@ {% endif %} ### User's QUERY HISTORY ### -Query history is context only. Do not reuse prior table names, column names, values, or functions unless they are present in the current DATABASE SCHEMA or SQL FUNCTIONS. +Query history is intent context only. Prior SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for history in histories %} Question: {{ history.question }} -SQL: -{{ history.sql }} {% endfor %} ### QUESTION ### @@ -60,7 +56,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. +Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..7891c0a4ae 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -123,11 +123,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are intent examples only. SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} @@ -149,8 +148,6 @@ {% for history in histories %} Question: {{ history.question }} -SQL: -{{ history.sql }} {% endfor %} {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 30f5809890..95fb5fbe02 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -87,12 +87,11 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} {% endif %} ### FAILED SQL ### -This SQL failed dry run. Do not preserve any identifier, source/physical/lineage name, or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. Do not use the failed SQL or error message as a source for alternate similar names; regenerate from the question and current schema. -{{ invalid_generation_result.sql }} +The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. Use the error message only to understand why the previous attempt failed. Error Message: {{ invalid_generation_result.error }} Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 03a1a1ae90..13e450c556 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -54,12 +54,10 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. Treat physical/source/lineage names in samples as non-executable unless the exact same identifier is declared in DATABASE SCHEMA. +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -76,7 +74,7 @@ {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} {% endif %} 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 e6d5c2c625..9867af1b2f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -29,12 +29,10 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use table names, column names, values, and functions only if they are present in the current DATABASE SCHEMA or SQL FUNCTIONS. +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} @@ -50,7 +48,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL or SQL-like expressions. +Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index b508d7135c..11fc4eaa98 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -79,12 +79,10 @@ def get_sql_regeneration_system_prompt( {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Use executable table names, column names, literals, and functions from the current DATABASE SCHEMA and SQL FUNCTIONS only. Treat physical/source/lineage names in samples as non-executable unless the exact same identifier is declared in DATABASE SCHEMA. +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -100,11 +98,10 @@ def get_sql_regeneration_system_prompt( Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### -Use this reasoning plan only as non-executable context. Ignore any SQL fragments, placeholder identifiers, inferred identifiers, source/physical/lineage names, unsupported functions, or identifiers not present in DATABASE SCHEMA. +Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. {{ sql_generation_reasoning }} ### ORIGINAL SQL QUERY ### -Use this SQL only as non-executable intent context. Do not preserve any identifier, source/physical/lineage name, or function from it unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. -{{ sql }} +The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b758c36504..dc5e9483de 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -285,41 +285,34 @@ async def _classify_generation_result( sql_samples_instructions = """ #### Instructions for SQL Samples #### -Finally, you will learn from the sample SQL queries provided in the input. These samples demonstrate best practices and common patterns for querying this specific database. +Finally, you will learn from the sample questions provided in the input. These samples demonstrate intent and response style for this specific database. For each sample, you should: 1. Study the question that explains what the query aims to accomplish -2. Analyze the SQL implementation to understand: - - Table structures and relationships used - - Specific functions and operators employed - - Query patterns and techniques demonstrated -3. Use these samples as reference patterns when generating similar queries, but treat the DATABASE SCHEMA as the only valid source of executable table and column names -4. Adapt the techniques shown in the samples to match new query requirements while maintaining consistent style and approach -5. Never copy table names, column names, aliases, literal values, or functions from samples unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS +2. Use these samples as intent and style context only, but treat the DATABASE SCHEMA as the only valid source of executable table and column names +3. Adapt the intent patterns to match new query requirements while maintaining consistent style and approach +4. Never copy table names, column names, aliases, literal values, placeholders, or functions from samples The samples will help you understand: -- Preferred table join patterns -- Common aggregation methods -- Specific function usage -- Query structure and formatting conventions +- Common analytical intents +- Common aggregation requests +- Preferred answer style -When generating new queries, try to follow similar patterns when applicable, while adapting them to the specific requirements of each new query. +When generating new queries, follow similar intent patterns when applicable, while adapting them to the specific requirements of each new query. -Learn about the usage of the schema structures and generate SQL based on them. +Learn about the user's intent from the samples and generate SQL from the current DATABASE SCHEMA and SQL FUNCTIONS only. """ sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who maps a user's intent to the provided database schema and provides a concise reasoning plan for answering the user's question. +You are a helpful data analyst who explains the user's analytical intent and provides a concise, non-executable reasoning plan for answering the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will use that absolute time frame only with an exact date/time schema column and supported SQL FUNCTIONS; -otherwise, you will use a relative timeframe only when an exact date/time schema column and supported SQL FUNCTIONS are available. If they are not available, state that the available schema/functions do not support that time filter. -3. For top, bottom, first, last, highest, or lowest requests, plan to sort by an exact selected column or aggregate alias and limit the result. Include a ranking column only when the user explicitly asks for rank values. -4. Do not plan to use a SQL function unless it appears in SQL FUNCTIONS for this request or is already part of a valid metric/calculated-field definition in DATABASE SCHEMA. +2. Explicitly state requested timeframes in natural language only. Do not name date columns, functions, expressions, or SQL clauses. +3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language only. Do not name columns, aliases, aggregate expressions, or SQL clauses. +4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. @@ -327,22 +320,16 @@ async def _classify_generation_result( 9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. A table name in the reasoning plan must be in this format: `table: `. -13. A column name in the reasoning plan must be in this format: `column: .`. -14. Use only exact table and column names that appear in the DATABASE SCHEMA section. -15. Comments, aliases, display labels, and descriptions are semantic hints only; do not turn them into table or column names in the reasoning plan. -16. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. -17. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the schema does not show the exact table or column needed, state that the available schema does not include that part. -18. If the question asks for a concept, filter, sort, or timeframe, map it only to exact available schema columns. If no exact schema column supports part of the request, state that the available schema does not include that part instead of inventing a column. -19. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but name only exact tables and columns from DATABASE SCHEMA in the reasoning plan. -20. If multiple schema objects are required to answer the intent, include each required object only when DATABASE SCHEMA provides both the needed fields and the relationship path between them. -21. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, or functions from them unless they also appear in the current DATABASE SCHEMA or SQL FUNCTIONS. -22. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -23. Use comments, aliases, display labels, and descriptions to explain why an exact schema column is relevant, not as replacement names. -24. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. -25. Do not derive executable table or column names from natural language, comments, aliases, display labels, descriptions, source metadata, physical datasource names, or lineage names. Only cite exact declared names from DATABASE SCHEMA. -26. If source metadata or lineage names help identify a concept, map that concept only to exact declared DATABASE SCHEMA objects; if no exact object exists, state that the available schema does not include that part. -27. ONLY SHOWING the reasoning plan in bullet points. +12. Do not mention table names, view names, metric names, column names, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels in the reasoning plan. +13. Do not write possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, executable expressions, or date/time expressions in the reasoning plan. +14. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. +15. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language only. +16. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but do not name the underlying schema objects in the reasoning plan. +17. If multiple schema objects may be required to answer the intent, describe the need to combine related data in natural language only. +18. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan. +19. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +20. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. +21. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -410,13 +397,13 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. -3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, or functions from samples unless they are valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. If the reasoning plan contains SQL fragments, assumed SQL, placeholder identifiers, inferred identifiers, unsupported functions, or identifiers missing from DATABASE SCHEMA, ignore those parts. +3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. +4. YOU MUST use the reasoning plan only as non-executable intent context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 7. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. 8. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If an identifier appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request. +9. If an identifier, literal value, placeholder, or function appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request. 10. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} @@ -471,7 +458,7 @@ def construct_ask_history_messages( ) messages.append( ChatMessage.from_assistant( - history.sql if hasattr(history, "sql") else history["sql"] + "Previous SQL omitted. Use only the current DATABASE SCHEMA and SQL FUNCTIONS for executable SQL." ) ) return messages From d846595b6d2f5e76d82362a8e29064b611798b8e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 14:16:13 +0530 Subject: [PATCH 0723/1087] Expand schema retrieval through relationships --- .../src/pipelines/generation/utils/sql.py | 17 +-- .../retrieval/db_schema_retrieval.py | 128 ++++++++++++++--- .../retrieval/test_db_schema_retrieval.py | 129 ++++++++++++++++++ 3 files changed, 246 insertions(+), 28 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index dc5e9483de..490eda1458 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -321,15 +321,16 @@ async def _classify_generation_result( 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. 12. Do not mention table names, view names, metric names, column names, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels in the reasoning plan. -13. Do not write possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, executable expressions, or date/time expressions in the reasoning plan. +13. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. 14. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. 15. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language only. 16. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but do not name the underlying schema objects in the reasoning plan. -17. If multiple schema objects may be required to answer the intent, describe the need to combine related data in natural language only. -18. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan. -19. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -20. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. -21. ONLY SHOWING the reasoning plan in bullet points. +17. Only cite exact declared names from DATABASE SCHEMA if an internal grounding note requires it; do not expose table names, column names, source metadata, physical datasource names, or lineage names in the reasoning plan. +18. If multiple schema objects may be required to answer the intent, describe the need to combine related data in natural language only. +19. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan. +20. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +21. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. +22. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -398,12 +399,12 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. -4. YOU MUST use the reasoning plan only as non-executable intent context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 7. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. 8. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If an identifier, literal value, placeholder, or function appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request. +9. If an identifier, literal value, placeholder, or function appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. 10. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..62a1a659d3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -29,7 +29,13 @@ ### TASK ### You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. -The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. +The database schema includes structural, semantic, and business modeling metadata: +- Models are logical datasets backed by physical tables or SQL definitions. +- Columns are exposed fields, including renamed fields, expressions, primary keys, and calculated fields. +- Relationships are reusable join logic between models. +- Calculated fields are business logic defined once and reused across queries. +- Views are named SQL statements that behave like stable virtual tables. +- Metrics are structured aggregation objects with measures and dimensions. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -39,6 +45,12 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +8. Map the business question to the modeled datasets whose descriptions, aliases, columns, calculated fields, views, metrics, and relationships support the intent. +9. Prefer modeled analytical interfaces such as views and metrics when they expose the fields needed to answer the question. +10. If the answer needs fields, filters, time dimensions, ordering, aggregations, or relationship keys from multiple related datasets, include every required related dataset and the columns needed from each one. +11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. +12. Follow only the relationships shown in the provided schema when selecting columns across datasets. +13. Do not stop at a single top candidate when the question needs multiple related datasets. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -172,35 +184,111 @@ async def table_retrieval( async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any ) -> list[Document]: - tables = table_retrieval.get("documents", []) + table_names = _table_names_from_description_documents( + table_retrieval.get("documents", []) + ) + + if table_names: + documents = [] + retrieved_table_names = set() + pending_table_names = table_names + + while pending_table_names: + retrieved_table_names.update(pending_table_names) + retrieved_documents = await _retrieve_schema_documents( + pending_table_names, project_id, dbschema_retriever + ) + documents = _dedupe_documents(documents + retrieved_documents) + pending_table_names = [ + table_name + for table_name in _related_table_names(documents) + if table_name not in retrieved_table_names + ] + + return documents + + return [] + + +def _table_names_from_description_documents(documents: list[Document]) -> list[str]: table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + table_name = content["name"] + if table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + + return table_names + +async def _retrieve_schema_documents( + table_names: list[str], project_id: str, dbschema_retriever: Any +) -> list[Document]: table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names ] - if table_name_conditions: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } + if not table_name_conditions: + return [] - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } + + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] - return [] + +def _related_table_names(documents: list[Document]) -> list[str]: + related_table_names = [] + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + if content.get("type") != "TABLE_COLUMNS": + continue + + for column in content.get("columns", []): + if column.get("type") != "FOREIGN_KEY": + continue + + for table_name in column.get("tables", []): + if table_name not in seen: + related_table_names.append(table_name) + seen.add(table_name) + + return related_table_names + + +def _dedupe_documents(documents: list[Document]) -> list[Document]: + deduped = [] + seen = set() + + for document in documents: + identity = ( + document.meta.get("type"), + document.meta.get("name"), + document.content, + ) + if identity in seen: + continue + deduped.append(document) + seen.add(identity) + + return deduped @observe() diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index bb23a90768..befe2497cd 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -102,6 +102,135 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_dbschema_retrieval_expands_declared_relationships(): + selected_model = "model_a" + related_model = "model_b" + downstream_model = "model_c" + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + names = [ + condition["value"] + for condition in filters["conditions"][1]["conditions"] + ] + self.calls.append(names) + + if names == [selected_model]: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": selected_model, + } + ), + meta={"type": "TABLE_SCHEMA", "name": selected_model}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "FOREIGN_KEY", + "tables": [ + selected_model, + related_model, + ], + "column": "model_b_id", + "referenced_table": related_model, + "referenced_column": "id", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": selected_model}, + ), + ] + } + + if names == [related_model]: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": related_model, + } + ), + meta={"type": "TABLE_SCHEMA", "name": related_model}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "FOREIGN_KEY", + "tables": [ + related_model, + downstream_model, + ], + "column": "model_c_id", + "referenced_table": downstream_model, + "referenced_column": "id", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": related_model}, + ), + ] + } + + if names == [downstream_model]: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": downstream_model, + } + ), + meta={"type": "TABLE_SCHEMA", "name": downstream_model}, + ) + ] + } + + return {"documents": []} + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={ + "documents": [ + Document( + content=str({"name": selected_model}), + meta={"type": "TABLE_DESCRIPTION", "name": selected_model}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + ) + + assert retriever.calls == [[selected_model], [related_model], [downstream_model]] + assert [document.meta["name"] for document in documents] == [ + selected_model, + selected_model, + related_model, + related_model, + downstream_model, + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_does_not_load_full_schema_for_unmatched_question(): class Retriever: From f573ba546a7ebbefe835f8600f23ec0bbd7af77b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 14:57:06 +0530 Subject: [PATCH 0724/1087] Keep untrusted planning out of SQL prompts --- .../generation/followup_sql_generation.py | 6 +- .../pipelines/generation/sql_correction.py | 8 +-- .../pipelines/generation/sql_generation.py | 6 +- .../pipelines/generation/sql_regeneration.py | 11 ++- .../pipelines/generation/test_sql_utils.py | 67 ++++++++++++++++++- 5 files changed, 78 insertions(+), 20 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 5beb9cae14..8d323d8376 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -79,9 +79,7 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. ### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context and only where it is consistent with the current DATABASE SCHEMA and SQL RULES. -Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. -{{ sql_generation_reasoning }} +The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. Return only the final JSON SQL response. """ @@ -105,7 +103,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=bool(sql_generation_reasoning), instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 95fb5fbe02..26c2888b95 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -87,12 +87,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. -{{ sql_generation_reasoning }} +The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. {% endif %} ### FAILED SQL ### -The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. Use the error message only to understand why the previous attempt failed. -Error Message: {{ invalid_generation_result.error }} +The failed SQL and raw dry-run error are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. Regenerate from the user question and current DATABASE SCHEMA. Return only the final JSON SQL response. """ @@ -113,7 +111,7 @@ def prompt( query=query, documents=documents, invalid_generation_result=invalid_generation_result, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=bool(sql_generation_reasoning), instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 13e450c556..4f71cfc091 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -73,9 +73,7 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. {% if sql_generation_reasoning %} -### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. -{{ sql_generation_reasoning }} +The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. {% endif %} Return only the final JSON SQL response. @@ -100,7 +98,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=bool(sql_generation_reasoning), instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 11fc4eaa98..07074d1433 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,9 +34,9 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query. +You are a great ANSI SQL expert. Now you are given database schema and a user's question. Carefully review the user's question and current DATABASE SCHEMA, then generate a new SQL query that answers the user's intent. -Use the original SQL query only as non-executable intent context. +The original SQL query and UI planning text are intentionally omitted from the prompt and must not be used as executable context. While generating the new SQL query, make sure to use the database schema as the only source of executable table and column identifiers. If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. Treat physical/source/lineage names from the original SQL, reasoning, samples, comments, or descriptions as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. @@ -96,10 +96,9 @@ def get_sql_regeneration_system_prompt( ### QUESTION ### User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. -Use the original SQL query only as intent context. Regenerate with executable identifiers from the current DATABASE SCHEMA only. +Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### -Use this reasoning plan only as non-executable intent context. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from it. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. -{{ sql_generation_reasoning }} +The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. ### ORIGINAL SQL QUERY ### The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. @@ -127,7 +126,7 @@ def prompt( query=query, sql=sql, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, + sql_generation_reasoning=bool(sql_generation_reasoning), instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 216ecde889..ef6346973f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,3 +1,5 @@ +from haystack.components.builders.prompt_builder import PromptBuilder + from src.pipelines.generation.utils.sql import ( construct_instructions, get_json_field_instructions, @@ -116,7 +118,7 @@ def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): assert "regenerate from the user's question" in prompt assert "unsupported identifiers" in prompt - assert "Use the original SQL query only as non-executable intent context" in prompt + assert "original SQL query and UI planning text are intentionally omitted" in prompt assert ( "database schema as the only source of executable table and column identifiers" in prompt @@ -158,3 +160,66 @@ def test_user_prompt_templates_keep_source_metadata_non_executable(): assert "source/physical/lineage names" in prompt assert "omit that unsupported part instead of inventing" in prompt assert "exact declared table and column names from DATABASE SCHEMA" in prompt + + +def test_executable_prompt_templates_omit_planning_error_and_original_sql_context(): + marker = "UNTRUSTED_CONTEXT_MARKER" + + generation_prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( + query="Question", + documents=["SCHEMA_CONTEXT"], + sql_generation_reasoning=marker, + instructions=[], + calculated_field_instructions="", + metric_instructions="", + json_field_instructions="", + sql_samples=[], + sql_functions=[], + )["prompt"] + + followup_prompt = PromptBuilder( + template=text_to_sql_with_followup_user_prompt_template + ).run( + query="Question", + documents=["SCHEMA_CONTEXT"], + sql_generation_reasoning=marker, + instructions=[], + calculated_field_instructions="", + metric_instructions="", + json_field_instructions="", + sql_samples=[], + sql_functions=[], + )["prompt"] + + correction_prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( + query="Question", + documents=["SCHEMA_CONTEXT"], + invalid_generation_result={"error": marker}, + sql_generation_reasoning=marker, + instructions=[], + sql_functions=[], + )["prompt"] + + regeneration_prompt = PromptBuilder( + template=sql_regeneration_user_prompt_template + ).run( + query="Question", + sql=marker, + documents=["SCHEMA_CONTEXT"], + sql_generation_reasoning=marker, + instructions=[], + calculated_field_instructions="", + metric_instructions="", + json_field_instructions="", + sql_samples=[], + sql_functions=[], + )["prompt"] + + for prompt in ( + generation_prompt, + followup_prompt, + correction_prompt, + regeneration_prompt, + ): + assert marker not in prompt + assert "intentionally omitted" in prompt From 1af6fdab4796cf333c5df169cc3f76168ecfc1c6 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 15:26:35 +0530 Subject: [PATCH 0725/1087] Fix Ask dry-run manifest consistency --- .../generation/followup_sql_generation.py | 4 ++ .../pipelines/generation/sql_correction.py | 33 ++++++--------- .../src/pipelines/generation/sql_diagnosis.py | 3 +- .../pipelines/generation/sql_generation.py | 4 ++ .../pipelines/generation/sql_regeneration.py | 8 +++- .../src/pipelines/generation/utils/sql.py | 17 +++++++- wren-ai-service/src/providers/engine/wren.py | 2 + wren-ai-service/src/web/v1/services/ask.py | 14 +++---- .../pipelines/generation/test_sql_utils.py | 42 +++++++++---------- wren-ui/src/apollo/server/models/model.ts | 1 + .../apollo/server/resolvers/modelResolver.ts | 29 +++++++------ wren-ui/src/apollo/server/schema.ts | 1 + 12 files changed, 91 insertions(+), 67 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 8d323d8376..f8677de2a1 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -148,12 +148,14 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -201,6 +203,7 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -224,6 +227,7 @@ async def run( "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, + "mdl_hash": mdl_hash, "sql_samples": sql_samples, "instructions": instructions, "has_calculated_field": has_calculated_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 26c2888b95..4380b28540 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,19 +30,14 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. +You are a Wren SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect Wren SQL query. ### SQL CORRECTION INSTRUCTIONS ### -1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. -2. Then, generate a syntactically correct ANSI SQL query from the user's intent and the current DATABASE SCHEMA. -3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. -4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. -5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. -6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. -7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. -8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. +1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). +2. Then, generate the syntactically correct Wren SQL query to correct the error. +3. Keep executable table and column identifiers grounded in DATABASE SCHEMA. +4. Use SQL FUNCTIONS only when their exact syntax is provided. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -81,18 +76,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -{% if query %} -User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. -{% endif %} -{% if sql_generation_reasoning %} -### REASONING PLAN ### -The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. -{% endif %} -### FAILED SQL ### -The failed SQL and raw dry-run error are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. Regenerate from the user question and current DATABASE SCHEMA. +SQL: {{ invalid_generation_result.sql }} +Error Message: {{ invalid_generation_result.error }} -Return only the final JSON SQL response. +Let's think step by step. """ @@ -140,12 +127,14 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -193,6 +182,7 @@ async def run( instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, @@ -214,6 +204,7 @@ async def run( "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, + "mdl_hash": mdl_hash, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), diff --git a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py index 3f22b9d512..e65b5e6719 100644 --- a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py +++ b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py @@ -20,7 +20,7 @@ sql_diagnosis_system_prompt = """ ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to diagnose the issue with the given SQL query, error message and database schema. +You are a Wren SQL expert with exceptional logical thinking skills and debugging skills, you need to diagnose the issue with the given SQL query, error message and database schema. ### SQL DIAGNOSIS INSTRUCTIONS ### @@ -29,6 +29,7 @@ 3. Then, return the reasoning behind the diagnosis.(You should give me the part of the original SQL query that is incorrect and the reason why it is incorrect) 4. Reasoning should be in the language same as the language user provided in the INPUTS section. 5. Reasoning should be concise and to the point and within 50 words. +6. Diagnose against Wren SQL syntax and the provided DATABASE SCHEMA. Do not suggest datasource-specific SQL syntax. ### FINAL ANSWER FORMAT ### The final answer must be in JSON format: diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 4f71cfc091..752c7caf83 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -139,6 +139,7 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, @@ -146,6 +147,7 @@ async def post_process( return await post_processor.run( generate_sql.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -193,6 +195,7 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -218,6 +221,7 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, + "mdl_hash": mdl_hash, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 07074d1433..5c01c84b11 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,7 +34,7 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema and a user's question. +You are a great Wren SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query. Carefully review the user's question and current DATABASE SCHEMA, then generate a new SQL query that answers the user's intent. The original SQL query and UI planning text are intentionally omitted from the prompt and must not be used as executable context. While generating the new SQL query, make sure to use the database schema as the only source of executable table and column identifiers. @@ -44,7 +44,7 @@ def get_sql_regeneration_system_prompt( {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be a Wren SQL query in JSON format: {{ "sql": @@ -166,10 +166,12 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, + mdl_hash: str | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, ) @@ -209,6 +211,7 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -227,6 +230,7 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, + "mdl_hash": mdl_hash, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 490eda1458..1049c2a7fe 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -30,6 +30,7 @@ async def run( self, replies: List[str] | List[List[str]], project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, data_source: str = "", @@ -50,6 +51,7 @@ async def run( ) = await self._classify_generation_result( cleaned_generation_result, project_id=project_id, + mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, @@ -72,6 +74,7 @@ async def _classify_generation_result( self, generation_result: str, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, data_source: str = "", @@ -109,6 +112,7 @@ async def _classify_generation_result( generation_result, session, project_id=project_id, + mdl_hash=mdl_hash, limit=1, dry_run=True, ) @@ -134,6 +138,7 @@ async def _classify_generation_result( generation_result, session, project_id=project_id, + mdl_hash=mdl_hash, limit=1, dry_run=False, ) @@ -165,6 +170,9 @@ async def _classify_generation_result( _MANDATORY_SQL_GROUNDING_RULES = """ ### MANDATORY SQL GROUNDING RULES ### +- Generate Wren SQL that can be parsed by the Wren engine before any datasource dialect rewrite happens. +- Use the Wren SQL clause order for result limiting: ORDER BY comes before LIMIT, and LIMIT belongs at the end of the SELECT or final UNION result. +- Datasource-specific SQL syntax is execution-target context only. Do not emit syntax that requires a datasource parser before the Wren engine rewrite step. - Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. - Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. @@ -350,9 +358,16 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: rules = _DEFAULT_TEXT_TO_SQL_RULES if sql_knowledge is not None: - rules = _extract_from_sql_knowledge( + additional_rules = _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) + if additional_rules != _DEFAULT_TEXT_TO_SQL_RULES: + rules = ( + f"{rules}\n\n" + "### ADDITIONAL SQL KNOWLEDGE ###\n" + "Use this section only when it is compatible with Wren SQL and the current SQL FUNCTIONS.\n" + f"{additional_rules}" + ) return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 3a92853e04..3e812f0524 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -28,6 +28,7 @@ async def execute_sql( sql: str, session: aiohttp.ClientSession, project_id: str | None = None, + mdl_hash: str | None = None, dry_run: bool = True, timeout: float = settings.engine_timeout, limit: int = 500, @@ -36,6 +37,7 @@ async def execute_sql( data = { "sql": remove_limit_statement(sql), "projectId": project_id, + "hash": mdl_hash, } if dry_run: data["dryRun"] = True diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6ddd3b14e2..2c4a55d457 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -467,6 +467,7 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -485,6 +486,7 @@ async def ask( contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -548,19 +550,15 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - query=user_query, - sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": ( - f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" - if allow_sql_diagnosis - and sql_diagnosis_reasoning - else error_message - ), + "error": sql_diagnosis_reasoning + if allow_sql_diagnosis + else error_message, }, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index ef6346973f..ba45b460fc 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -35,6 +35,9 @@ def test_construct_instructions_uses_instruction_text(): def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): rules = get_text_to_sql_rules() + assert "Generate Wren SQL that can be parsed by the Wren engine" in rules + assert "ORDER BY comes before LIMIT" in rules + assert "Datasource-specific SQL syntax is execution-target context only" in rules assert "ONLY USE the tables and columns mentioned in the database schema" in rules assert 'ONLY USE "*" if the user query asks for all the columns' in rules assert "They are never source table or source column identifiers" in rules @@ -64,8 +67,12 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): rules = get_text_to_sql_rules(_SqlKnowledge()) + assert "ONLY USE the tables and columns mentioned in the database schema" in rules + assert "For top, bottom, highest, lowest, first, or last requests" in rules assert _SqlKnowledge.text_to_sql_rule in rules + assert "Use this section only when it is compatible with Wren SQL" in rules assert "MANDATORY SQL GROUNDING RULES" in rules + assert "Generate Wren SQL that can be parsed by the Wren engine" in rules assert "Every table and column referenced" in rules assert "Do not query INFORMATION_SCHEMA" in rules @@ -116,6 +123,8 @@ def test_json_field_instructions_do_not_include_placeholder_identifiers(): def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): prompt = get_sql_regeneration_system_prompt() + assert "Wren SQL query" in prompt + assert "ANSI SQL" not in prompt assert "regenerate from the user's question" in prompt assert "unsupported identifiers" in prompt assert "original SQL query and UI planning text are intentionally omitted" in prompt @@ -129,15 +138,11 @@ def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): def test_sql_correction_system_prompt_discards_invalid_identifier_context(): prompt = get_sql_correction_system_prompt() - assert "treat it as the source of intent" in prompt - assert "Do not copy placeholders" in prompt - assert "Regenerate a grounded Wren SQL query" in prompt - assert ( - "Do not preserve a table, column, join, filter, grouping, ordering, or function" - in prompt - ) - assert "Treat physical/source/lineage names from the failed SQL" in prompt - assert "do not try a similar replacement from source metadata" in prompt + assert "Wren SQL query" in prompt + assert "ANSI SQL" not in prompt + assert "fix the syntactically incorrect Wren SQL query" in prompt + assert "generate the syntactically correct Wren SQL query" in prompt + assert "Keep executable table and column identifiers grounded" in prompt def test_sql_reasoning_prompt_forbids_executable_sql_context(): @@ -155,14 +160,13 @@ def test_user_prompt_templates_keep_source_metadata_non_executable(): sql_generation_user_prompt_template, text_to_sql_with_followup_user_prompt_template, sql_regeneration_user_prompt_template, - sql_correction_user_prompt_template, ): assert "source/physical/lineage names" in prompt assert "omit that unsupported part instead of inventing" in prompt assert "exact declared table and column names from DATABASE SCHEMA" in prompt -def test_executable_prompt_templates_omit_planning_error_and_original_sql_context(): +def test_executable_prompt_templates_omit_planning_and_original_sql_context(): marker = "UNTRUSTED_CONTEXT_MARKER" generation_prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( @@ -191,15 +195,6 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex sql_functions=[], )["prompt"] - correction_prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( - query="Question", - documents=["SCHEMA_CONTEXT"], - invalid_generation_result={"error": marker}, - sql_generation_reasoning=marker, - instructions=[], - sql_functions=[], - )["prompt"] - regeneration_prompt = PromptBuilder( template=sql_regeneration_user_prompt_template ).run( @@ -218,8 +213,13 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex for prompt in ( generation_prompt, followup_prompt, - correction_prompt, regeneration_prompt, ): assert marker not in prompt assert "intentionally omitted" in prompt + + +def test_sql_correction_user_prompt_follows_legacy_failed_sql_flow(): + assert "SQL: {{ invalid_generation_result.sql }}" in sql_correction_user_prompt_template + assert "Error Message: {{ invalid_generation_result.error }}" in sql_correction_user_prompt_template + assert "Let's think step by step." in sql_correction_user_prompt_template diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index c6403db9b9..4f3ff216eb 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -98,6 +98,7 @@ export interface CheckCalculatedFieldCanQueryData { export interface PreviewSQLData { sql: string; projectId?: string; + hash?: string; limit?: number; dryRun?: boolean; } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 3cd42c4505..5cabdea849 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -236,19 +236,16 @@ export class ModelResolver { public async checkModelSync(_root: any, _args: any, ctx: IContext) { try { const { id } = await ctx.projectService.getCurrentProject(); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const currentHash = ctx.deployService.createMDLHash(manifest, id); + const lastDeploy = await ctx.deployService.getLastDeployment(id); + const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - - const project = await ctx.projectService.getCurrentProject(); - if (dirtyProjectIds.has(project.id)) { - return { status: SyncStatusEnum.UNSYNCRONIZED }; - } - - const lastDeploy = await ctx.deployService.getLastDeployment(project.id); - return lastDeploy + return currentHash == lastDeployHash ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { @@ -1207,7 +1204,7 @@ export class ModelResolver { // create view const project = await ctx.projectService.getCurrentProject(); - const manifest = await this.getLastDeployedManifest(ctx, project.id); + const manifest = await this.getDeployedManifest(ctx, project.id); // get sql statement of a response const response = await ctx.askingService.getResponse(responseId); @@ -1334,11 +1331,11 @@ export class ModelResolver { args: { data: PreviewSQLData }, ctx: IContext, ) { - const { sql, projectId, limit, dryRun } = args.data; + const { sql, projectId, hash, limit, dryRun } = args.data; const project = projectId ? await ctx.projectService.getProjectById(parseInt(projectId)) : await ctx.projectService.getCurrentProject(); - const manifest = await this.getLastDeployedManifest(ctx, project.id); + const manifest = await this.getDeployedManifest(ctx, project.id, hash); return await ctx.queryService.preview(sql, { project, limit: limit, @@ -1487,8 +1484,14 @@ export class ModelResolver { }; } - private async getLastDeployedManifest(ctx: IContext, projectId: number) { - const deployment = await ctx.deployService.getLastDeployment(projectId); + private async getDeployedManifest( + ctx: IContext, + projectId: number, + hash?: string, + ) { + const deployment = hash + ? await ctx.deployLogRepository.findOneBy({ projectId, hash }) + : await ctx.deployService.getLastDeployment(projectId); if (!deployment?.manifest) { throw new Error( 'Project has not been deployed successfully yet. Deploy the model before previewing or validating SQL.', diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 37e2549949..3c656de8c2 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -923,6 +923,7 @@ export const typeDefs = gql` input PreviewSQLDataInput { sql: String! projectId: String + hash: String limit: Int dryRun: Boolean } From 749ad225498365379e721e7f4a887970f354307c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 15:51:07 +0530 Subject: [PATCH 0726/1087] Fix deploy sync status for equivalent manifests --- .../apollo/server/resolvers/modelResolver.ts | 4 +-- .../apollo/server/services/deployService.ts | 2 +- .../services/tests/deployService.test.ts | 27 +++++++++++++++++++ 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 5cabdea849..754394474b 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -237,15 +237,13 @@ export class ModelResolver { try { const { id } = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest, id); const lastDeploy = await ctx.deployService.getLastDeployment(id); - const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - return currentHash == lastDeployHash + return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 746e5fa12e..3f241a7455 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -107,7 +107,7 @@ export class DeployService implements IDeployService { // check if the model current deployment const lastDeploy = await this.deployLogRepository.findLastProjectDeployLog(projectId); - if (lastDeploy && lastDeploy.hash === hash) { + if (this.isSameDeployment(manifest, projectId, lastDeploy)) { logger.log(`Model has been deployed, hash: ${hash}`); await this.deployLogRepository.updateOne(lastDeploy.id, { status: DeployStatusEnum.SUCCESS, diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index ce60461bf7..632cd39da0 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -101,6 +101,33 @@ describe('DeployService', () => { }); }); + it('should skip deployment if an existing deployment has the same manifest with a different hash', async () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + ], + }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ + id: 123, + hash: 'legacy-hash', + manifest, + }); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.SUCCESS); + expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.SUCCESS, + error: null, + }); + }); + it('should create the same deployment hash for equivalent manifests', () => { const manifest = { models: [ From 5e1c1a60d22e0058e00e53a08ff3d4c47433be4f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:08:30 +0530 Subject: [PATCH 0727/1087] Keep deploy sync stable after persistence --- .../apollo/server/resolvers/modelResolver.ts | 114 +----------------- .../apollo/server/services/deployService.ts | 22 +++- .../services/tests/deployService.test.ts | 39 +++++- wren-ui/src/apollo/server/types/context.ts | 2 +- 4 files changed, 57 insertions(+), 120 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 754394474b..1bcd27fcfd 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -258,22 +258,11 @@ export class ModelResolver { ctx: IContext, ): Promise { const project = await this.prepareProjectForDeploy(ctx); - const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const lastDeploy = await ctx.deployService.getLastDeployment(project.id); - const hasModelingChangesAfterDeploy = - !(await this.isLastDeployNewerThanModelingChanges( - ctx, - project.id, - lastDeploy, - )); - const shouldForceDeploy = - args.force || - hasModelingChangesAfterDeploy || - !ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy); + const { manifest } = await ctx.mdlService.makeModelMDL(project); const deployRes = await ctx.deployService.deploy( manifest, project.id, - shouldForceDeploy, + args.force, ); if (deployRes.status === 'SUCCESS') { dirtyProjectIds.delete(project.id); @@ -301,105 +290,6 @@ export class ModelResolver { return project; } - private async isLastDeployNewerThanModelingChanges( - ctx: IContext, - projectId: number, - lastDeploy?: { createdAt?: Date; updatedAt?: Date } | null, - ): Promise { - if (!lastDeploy) { - return false; - } - - const deployedAt = this.toTime(lastDeploy.updatedAt || lastDeploy.createdAt); - if (!deployedAt) { - return false; - } - - const models = await ctx.modelRepository.findAllBy({ projectId }); - const modelIds = models.map((model) => model.id); - const [columns, nestedColumns, relations, views] = await Promise.all([ - modelIds.length - ? ctx.modelColumnRepository.findColumnsByModelIds(modelIds) - : Promise.resolve([]), - modelIds.length - ? ctx.modelNestedColumnRepository.findNestedColumnsByModelIds(modelIds) - : Promise.resolve([]), - ctx.relationRepository.findRelationInfoBy({ projectId }), - ctx.viewRepository.findAllBy({ projectId }), - ]); - - const latestModelingChangeAt = [ - ...models, - ...columns, - ...nestedColumns, - ...relations, - ...views, - ].reduce((latest, item: any) => { - return Math.max(latest, this.toTime(item.updatedAt || item.createdAt)); - }, 0); - - return deployedAt >= latestModelingChangeAt; - } - - private toTime(value?: Date | string | null): number { - if (!value) { - return 0; - } - const time = new Date(value).getTime(); - return Number.isFinite(time) ? time : 0; - } - - private isSameDeploymentIgnoringColumnNullability( - manifest: any, - lastDeploy?: { manifest?: any } | null, - ): boolean { - if (!lastDeploy?.manifest) { - return false; - } - - return ( - this.stableStringify(this.omitColumnNullability(lastDeploy.manifest)) === - this.stableStringify(this.omitColumnNullability(manifest)) - ); - } - - private omitColumnNullability(value: any): any { - if (Array.isArray(value)) { - return value.map((item) => this.omitColumnNullability(item)); - } - if (!value || typeof value !== 'object') { - return value; - } - - const result: Record = {}; - for (const key of Object.keys(value)) { - if (key === 'notNull') { - continue; - } - result[key] = this.omitColumnNullability(value[key]); - } - return result; - } - - private stableStringify(value: any): string { - if (Array.isArray(value)) { - const serializedItems = value.map((item) => this.stableStringify(item)); - if (value.every((item) => item && typeof item === 'object')) { - serializedItems.sort(); - } - return `[${serializedItems.join(',')}]`; - } - - if (value && typeof value === 'object') { - return `{${Object.keys(value) - .sort() - .map((key) => `${JSON.stringify(key)}:${this.stableStringify(value[key])}`) - .join(',')}}`; - } - - return JSON.stringify(value); - } - private markProjectDirty(projectId: number) { dirtyProjectIds.add(projectId); } diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 3f241a7455..9b54a46d1f 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -215,9 +215,15 @@ export class DeployService implements IDeployService { ); } - private canonicalStringify(value: any): string { + private canonicalStringify(value: any): string | undefined { + if (value === undefined) { + return undefined; + } + if (Array.isArray(value)) { - const serializedItems = value.map((item) => this.canonicalStringify(item)); + const serializedItems = value.map( + (item) => this.canonicalStringify(item) ?? 'null', + ); if (value.every((item) => item && typeof item === 'object')) { serializedItems.sort(); } @@ -225,10 +231,16 @@ export class DeployService implements IDeployService { } if (value && typeof value === 'object') { - return `{${Object.keys(value) + const serializedProperties = Object.keys(value) .sort() - .map((key) => `${JSON.stringify(key)}:${this.canonicalStringify(value[key])}`) - .join(',')}}`; + .map((key) => { + const serializedValue = this.canonicalStringify(value[key]); + return serializedValue === undefined + ? undefined + : `${JSON.stringify(key)}:${serializedValue}`; + }) + .filter((property): property is string => property !== undefined); + return `{${serializedProperties.join(',')}}`; } return JSON.stringify(value); diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 632cd39da0..0c1b3bd1e5 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -105,8 +105,8 @@ describe('DeployService', () => { const manifest = { models: [ { - name: 'orders', - columns: [{ name: 'id' }, { name: 'amount' }], + name: 'model', + columns: [{ name: 'field' }, { name: 'measure' }], }, ], }; @@ -128,6 +128,41 @@ describe('DeployService', () => { }); }); + it('should treat JSON-persisted manifests with omitted undefined properties as the same deployment', () => { + const manifest = { + models: [ + { + name: 'model', + properties: { + displayName: 'Model', + description: undefined, + }, + columns: [ + { + name: 'field', + expression: undefined, + properties: { + displayName: 'Field', + description: undefined, + }, + }, + ], + }, + ], + }; + const persistedManifest = JSON.parse(JSON.stringify(manifest)); + + expect(deployService.createMDLHash(manifest, 1)).toEqual( + deployService.createMDLHash(persistedManifest, 1), + ); + expect( + deployService.isSameDeployment(manifest, 1, { + hash: 'previous-hash', + manifest: persistedManifest, + }), + ).toBe(true); + }); + it('should create the same deployment hash for equivalent manifests', () => { const manifest = { models: [ diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index 4aed639176..9202777672 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -75,7 +75,7 @@ export interface IContext { modelNestedColumnRepository: IModelNestedColumnRepository; relationRepository: IRelationRepository; viewRepository: IViewRepository; - deployRepository: IDeployLogRepository; + deployLogRepository: IDeployLogRepository; schemaChangeRepository: ISchemaChangeRepository; learningRepository: ILearningRepository; dashboardRepository: IDashboardRepository; From d7aae0536c507bc3ca8f66fbea79dfece23fd6b2 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:31:08 +0530 Subject: [PATCH 0728/1087] Keep model sync stable across navigation --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 1bcd27fcfd..107261dd18 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -236,14 +236,18 @@ export class ModelResolver { public async checkModelSync(_root: any, _args: any, ctx: IContext) { try { const { id } = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const lastDeploy = await ctx.deployService.getLastDeployment(id); const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) + + if (dirtyProjectIds.has(id)) { + return { status: SyncStatusEnum.UNSYNCRONIZED }; + } + + const lastDeploy = await ctx.deployService.getLastDeployment(id); + return lastDeploy ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { From 0e0fdd535367313afbdd9d5b377ceee728816670 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:45:45 +0530 Subject: [PATCH 0729/1087] Base model sync on persisted modeling changes --- .../apollo/server/resolvers/modelResolver.ts | 82 ++++++++++++------- 1 file changed, 51 insertions(+), 31 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 107261dd18..2ed40e2b68 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -37,8 +37,6 @@ import DataSourceSchemaDetector, { const logger = getLogger('ModelResolver'); logger.level = 'debug'; -const dirtyProjectIds = new Set(); - const isSameId = (left: string | number, right: string | number) => String(left) === String(right); @@ -111,7 +109,6 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_RELATION; try { const relation = await ctx.modelService.createRelation(data); - this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -134,7 +131,6 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_UPDATE_RELATION; try { const relation = await ctx.modelService.updateRelation(data, where.id); - this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -153,10 +149,8 @@ export class ModelResolver { args: { where: { id: number } }, ctx: IContext, ) { - const project = await ctx.projectService.getCurrentProject(); const relationId = args.where.id; await ctx.modelService.deleteRelation(relationId); - this.markProjectDirty(project.id); return true; } @@ -168,8 +162,6 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_CF; try { const column = await ctx.modelService.createCalculatedField(_args.data); - const project = await ctx.projectService.getCurrentProject(); - this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data: _args.data }); return column; } catch (err: any) { @@ -205,8 +197,6 @@ export class ModelResolver { data, where.id, ); - const project = await ctx.projectService.getCurrentProject(); - this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data }); return column; } catch (err: any) { @@ -227,9 +217,7 @@ export class ModelResolver { if (!column || !column.isCalculated) { throw new Error('Calculated field not found'); } - const project = await ctx.projectService.getCurrentProject(); await ctx.modelColumnRepository.deleteOne(columnId); - this.markProjectDirty(project.id); return true; } @@ -242,12 +230,14 @@ export class ModelResolver { return { status: SyncStatusEnum.IN_PROGRESS }; } - if (dirtyProjectIds.has(id)) { + const lastDeploy = await ctx.deployService.getLastDeployment(id); + if (!lastDeploy) { return { status: SyncStatusEnum.UNSYNCRONIZED }; } - const lastDeploy = await ctx.deployService.getLastDeployment(id); - return lastDeploy + const deployedAfterModelingChanges = + await this.isLastDeployNewerThanModelingChanges(ctx, id, lastDeploy); + return deployedAfterModelingChanges ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { @@ -268,9 +258,6 @@ export class ModelResolver { project.id, args.force, ); - if (deployRes.status === 'SUCCESS') { - dirtyProjectIds.delete(project.id); - } if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { ctx.projectService.generateProjectRecommendationQuestions().catch((err) => logger.warn( @@ -294,8 +281,52 @@ export class ModelResolver { return project; } - private markProjectDirty(projectId: number) { - dirtyProjectIds.add(projectId); + private async isLastDeployNewerThanModelingChanges( + ctx: IContext, + projectId: number, + lastDeploy?: { createdAt?: Date; updatedAt?: Date } | null, + ): Promise { + if (!lastDeploy) { + return false; + } + + const deployedAt = this.toTime(lastDeploy.updatedAt || lastDeploy.createdAt); + if (!deployedAt) { + return false; + } + + const models = await ctx.modelRepository.findAllBy({ projectId }); + const modelIds = models.map((model) => model.id); + const [columns, nestedColumns, relations, views] = await Promise.all([ + modelIds.length + ? ctx.modelColumnRepository.findColumnsByModelIds(modelIds) + : Promise.resolve([]), + modelIds.length + ? ctx.modelNestedColumnRepository.findNestedColumnsByModelIds(modelIds) + : Promise.resolve([]), + ctx.relationRepository.findRelationInfoBy({ projectId }), + ctx.viewRepository.findAllBy({ projectId }), + ]); + + const latestModelingChangeAt = [ + ...models, + ...columns, + ...nestedColumns, + ...relations, + ...views, + ].reduce((latest, item: any) => { + return Math.max(latest, this.toTime(item.updatedAt || item.createdAt)); + }, 0); + + return deployedAt >= latestModelingChangeAt; + } + + private toTime(value?: Date | string | null): number { + if (!value) { + return 0; + } + const time = new Date(value).getTime(); + return Number.isFinite(time) ? time : 0; } private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { @@ -418,7 +449,6 @@ export class ModelResolver { }), ); - this.markProjectDirty(project.id); return { savedCount: args.data?.length || 0 }; } @@ -495,7 +525,6 @@ export class ModelResolver { toColumnId: toColumn.id, type: relationship.type, }); - this.markProjectDirty(savedRelation.projectId); createdCount += 1; } catch (err: any) { logger.warn( @@ -601,7 +630,6 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_CREATE_MODEL, { data: args.data, }); - this.markProjectDirty(model.projectId); return model; } catch (error: any) { ctx.telemetry.sendEvent( @@ -700,7 +728,6 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_UPDATE_MODEL, { data: args.data, }); - this.markProjectDirty(model.projectId); return model; } catch (err: any) { ctx.telemetry.sendEvent( @@ -830,7 +857,6 @@ export class ModelResolver { // related columns and relationships will be deleted in cascade await ctx.modelRepository.deleteOne(modelId); - this.markProjectDirty(model.projectId); return true; } @@ -876,7 +902,6 @@ export class ModelResolver { } ctx.telemetry.sendEvent(eventName, { data }); - this.markProjectDirty(model.projectId); return true; } catch (err: any) { ctx.telemetry.sendEvent( @@ -1146,8 +1171,6 @@ export class ModelResolver { // telemetry ctx.telemetry.sendEvent(eventName, eventProperties); - this.markProjectDirty(project.id); - return { ...view, displayName }; } catch (err: any) { ctx.telemetry.sendEvent( @@ -1172,7 +1195,6 @@ export class ModelResolver { throw new Error('View not found'); } await ctx.viewRepository.deleteOne(viewId); - this.markProjectDirty(view.projectId); return true; } @@ -1329,8 +1351,6 @@ export class ModelResolver { name: newName, properties: JSON.stringify(properties), }); - this.markProjectDirty(view.projectId); - return true; } From ce1d1a0454f55545a0ff4856a7e9d6aa89c0af91 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:56:34 +0530 Subject: [PATCH 0730/1087] Revert "Base model sync on persisted modeling changes" This reverts commit 0e0fdd535367313afbdd9d5b377ceee728816670. --- .../apollo/server/resolvers/modelResolver.ts | 82 +++++++------------ 1 file changed, 31 insertions(+), 51 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 2ed40e2b68..107261dd18 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -37,6 +37,8 @@ import DataSourceSchemaDetector, { const logger = getLogger('ModelResolver'); logger.level = 'debug'; +const dirtyProjectIds = new Set(); + const isSameId = (left: string | number, right: string | number) => String(left) === String(right); @@ -109,6 +111,7 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_RELATION; try { const relation = await ctx.modelService.createRelation(data); + this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -131,6 +134,7 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_UPDATE_RELATION; try { const relation = await ctx.modelService.updateRelation(data, where.id); + this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -149,8 +153,10 @@ export class ModelResolver { args: { where: { id: number } }, ctx: IContext, ) { + const project = await ctx.projectService.getCurrentProject(); const relationId = args.where.id; await ctx.modelService.deleteRelation(relationId); + this.markProjectDirty(project.id); return true; } @@ -162,6 +168,8 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_CF; try { const column = await ctx.modelService.createCalculatedField(_args.data); + const project = await ctx.projectService.getCurrentProject(); + this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data: _args.data }); return column; } catch (err: any) { @@ -197,6 +205,8 @@ export class ModelResolver { data, where.id, ); + const project = await ctx.projectService.getCurrentProject(); + this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data }); return column; } catch (err: any) { @@ -217,7 +227,9 @@ export class ModelResolver { if (!column || !column.isCalculated) { throw new Error('Calculated field not found'); } + const project = await ctx.projectService.getCurrentProject(); await ctx.modelColumnRepository.deleteOne(columnId); + this.markProjectDirty(project.id); return true; } @@ -230,14 +242,12 @@ export class ModelResolver { return { status: SyncStatusEnum.IN_PROGRESS }; } - const lastDeploy = await ctx.deployService.getLastDeployment(id); - if (!lastDeploy) { + if (dirtyProjectIds.has(id)) { return { status: SyncStatusEnum.UNSYNCRONIZED }; } - const deployedAfterModelingChanges = - await this.isLastDeployNewerThanModelingChanges(ctx, id, lastDeploy); - return deployedAfterModelingChanges + const lastDeploy = await ctx.deployService.getLastDeployment(id); + return lastDeploy ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { @@ -258,6 +268,9 @@ export class ModelResolver { project.id, args.force, ); + if (deployRes.status === 'SUCCESS') { + dirtyProjectIds.delete(project.id); + } if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { ctx.projectService.generateProjectRecommendationQuestions().catch((err) => logger.warn( @@ -281,52 +294,8 @@ export class ModelResolver { return project; } - private async isLastDeployNewerThanModelingChanges( - ctx: IContext, - projectId: number, - lastDeploy?: { createdAt?: Date; updatedAt?: Date } | null, - ): Promise { - if (!lastDeploy) { - return false; - } - - const deployedAt = this.toTime(lastDeploy.updatedAt || lastDeploy.createdAt); - if (!deployedAt) { - return false; - } - - const models = await ctx.modelRepository.findAllBy({ projectId }); - const modelIds = models.map((model) => model.id); - const [columns, nestedColumns, relations, views] = await Promise.all([ - modelIds.length - ? ctx.modelColumnRepository.findColumnsByModelIds(modelIds) - : Promise.resolve([]), - modelIds.length - ? ctx.modelNestedColumnRepository.findNestedColumnsByModelIds(modelIds) - : Promise.resolve([]), - ctx.relationRepository.findRelationInfoBy({ projectId }), - ctx.viewRepository.findAllBy({ projectId }), - ]); - - const latestModelingChangeAt = [ - ...models, - ...columns, - ...nestedColumns, - ...relations, - ...views, - ].reduce((latest, item: any) => { - return Math.max(latest, this.toTime(item.updatedAt || item.createdAt)); - }, 0); - - return deployedAt >= latestModelingChangeAt; - } - - private toTime(value?: Date | string | null): number { - if (!value) { - return 0; - } - const time = new Date(value).getTime(); - return Number.isFinite(time) ? time : 0; + private markProjectDirty(projectId: number) { + dirtyProjectIds.add(projectId); } private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { @@ -449,6 +418,7 @@ export class ModelResolver { }), ); + this.markProjectDirty(project.id); return { savedCount: args.data?.length || 0 }; } @@ -525,6 +495,7 @@ export class ModelResolver { toColumnId: toColumn.id, type: relationship.type, }); + this.markProjectDirty(savedRelation.projectId); createdCount += 1; } catch (err: any) { logger.warn( @@ -630,6 +601,7 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_CREATE_MODEL, { data: args.data, }); + this.markProjectDirty(model.projectId); return model; } catch (error: any) { ctx.telemetry.sendEvent( @@ -728,6 +700,7 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_UPDATE_MODEL, { data: args.data, }); + this.markProjectDirty(model.projectId); return model; } catch (err: any) { ctx.telemetry.sendEvent( @@ -857,6 +830,7 @@ export class ModelResolver { // related columns and relationships will be deleted in cascade await ctx.modelRepository.deleteOne(modelId); + this.markProjectDirty(model.projectId); return true; } @@ -902,6 +876,7 @@ export class ModelResolver { } ctx.telemetry.sendEvent(eventName, { data }); + this.markProjectDirty(model.projectId); return true; } catch (err: any) { ctx.telemetry.sendEvent( @@ -1171,6 +1146,8 @@ export class ModelResolver { // telemetry ctx.telemetry.sendEvent(eventName, eventProperties); + this.markProjectDirty(project.id); + return { ...view, displayName }; } catch (err: any) { ctx.telemetry.sendEvent( @@ -1195,6 +1172,7 @@ export class ModelResolver { throw new Error('View not found'); } await ctx.viewRepository.deleteOne(viewId); + this.markProjectDirty(view.projectId); return true; } @@ -1351,6 +1329,8 @@ export class ModelResolver { name: newName, properties: JSON.stringify(properties), }); + this.markProjectDirty(view.projectId); + return true; } From d24980c00df4294e22478807637babb85c7578ba Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:56:34 +0530 Subject: [PATCH 0731/1087] Revert "Keep model sync stable across navigation" This reverts commit d7aae0536c507bc3ca8f66fbea79dfece23fd6b2. --- wren-ui/src/apollo/server/resolvers/modelResolver.ts | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 107261dd18..1bcd27fcfd 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -236,18 +236,14 @@ export class ModelResolver { public async checkModelSync(_root: any, _args: any, ctx: IContext) { try { const { id } = await ctx.projectService.getCurrentProject(); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const lastDeploy = await ctx.deployService.getLastDeployment(id); const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - - if (dirtyProjectIds.has(id)) { - return { status: SyncStatusEnum.UNSYNCRONIZED }; - } - - const lastDeploy = await ctx.deployService.getLastDeployment(id); - return lastDeploy + return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { From 05aad23512d6760dbd12899fa02ca0c6f711d701 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:56:34 +0530 Subject: [PATCH 0732/1087] Revert "Keep deploy sync stable after persistence" This reverts commit 5e1c1a60d22e0058e00e53a08ff3d4c47433be4f. --- .../apollo/server/resolvers/modelResolver.ts | 114 +++++++++++++++++- .../apollo/server/services/deployService.ts | 22 +--- .../services/tests/deployService.test.ts | 39 +----- wren-ui/src/apollo/server/types/context.ts | 2 +- 4 files changed, 120 insertions(+), 57 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 1bcd27fcfd..754394474b 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -258,11 +258,22 @@ export class ModelResolver { ctx: IContext, ): Promise { const project = await this.prepareProjectForDeploy(ctx); - const { manifest } = await ctx.mdlService.makeModelMDL(project); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + const hasModelingChangesAfterDeploy = + !(await this.isLastDeployNewerThanModelingChanges( + ctx, + project.id, + lastDeploy, + )); + const shouldForceDeploy = + args.force || + hasModelingChangesAfterDeploy || + !ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy); const deployRes = await ctx.deployService.deploy( manifest, project.id, - args.force, + shouldForceDeploy, ); if (deployRes.status === 'SUCCESS') { dirtyProjectIds.delete(project.id); @@ -290,6 +301,105 @@ export class ModelResolver { return project; } + private async isLastDeployNewerThanModelingChanges( + ctx: IContext, + projectId: number, + lastDeploy?: { createdAt?: Date; updatedAt?: Date } | null, + ): Promise { + if (!lastDeploy) { + return false; + } + + const deployedAt = this.toTime(lastDeploy.updatedAt || lastDeploy.createdAt); + if (!deployedAt) { + return false; + } + + const models = await ctx.modelRepository.findAllBy({ projectId }); + const modelIds = models.map((model) => model.id); + const [columns, nestedColumns, relations, views] = await Promise.all([ + modelIds.length + ? ctx.modelColumnRepository.findColumnsByModelIds(modelIds) + : Promise.resolve([]), + modelIds.length + ? ctx.modelNestedColumnRepository.findNestedColumnsByModelIds(modelIds) + : Promise.resolve([]), + ctx.relationRepository.findRelationInfoBy({ projectId }), + ctx.viewRepository.findAllBy({ projectId }), + ]); + + const latestModelingChangeAt = [ + ...models, + ...columns, + ...nestedColumns, + ...relations, + ...views, + ].reduce((latest, item: any) => { + return Math.max(latest, this.toTime(item.updatedAt || item.createdAt)); + }, 0); + + return deployedAt >= latestModelingChangeAt; + } + + private toTime(value?: Date | string | null): number { + if (!value) { + return 0; + } + const time = new Date(value).getTime(); + return Number.isFinite(time) ? time : 0; + } + + private isSameDeploymentIgnoringColumnNullability( + manifest: any, + lastDeploy?: { manifest?: any } | null, + ): boolean { + if (!lastDeploy?.manifest) { + return false; + } + + return ( + this.stableStringify(this.omitColumnNullability(lastDeploy.manifest)) === + this.stableStringify(this.omitColumnNullability(manifest)) + ); + } + + private omitColumnNullability(value: any): any { + if (Array.isArray(value)) { + return value.map((item) => this.omitColumnNullability(item)); + } + if (!value || typeof value !== 'object') { + return value; + } + + const result: Record = {}; + for (const key of Object.keys(value)) { + if (key === 'notNull') { + continue; + } + result[key] = this.omitColumnNullability(value[key]); + } + return result; + } + + private stableStringify(value: any): string { + if (Array.isArray(value)) { + const serializedItems = value.map((item) => this.stableStringify(item)); + if (value.every((item) => item && typeof item === 'object')) { + serializedItems.sort(); + } + return `[${serializedItems.join(',')}]`; + } + + if (value && typeof value === 'object') { + return `{${Object.keys(value) + .sort() + .map((key) => `${JSON.stringify(key)}:${this.stableStringify(value[key])}`) + .join(',')}}`; + } + + return JSON.stringify(value); + } + private markProjectDirty(projectId: number) { dirtyProjectIds.add(projectId); } diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 9b54a46d1f..3f241a7455 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -215,15 +215,9 @@ export class DeployService implements IDeployService { ); } - private canonicalStringify(value: any): string | undefined { - if (value === undefined) { - return undefined; - } - + private canonicalStringify(value: any): string { if (Array.isArray(value)) { - const serializedItems = value.map( - (item) => this.canonicalStringify(item) ?? 'null', - ); + const serializedItems = value.map((item) => this.canonicalStringify(item)); if (value.every((item) => item && typeof item === 'object')) { serializedItems.sort(); } @@ -231,16 +225,10 @@ export class DeployService implements IDeployService { } if (value && typeof value === 'object') { - const serializedProperties = Object.keys(value) + return `{${Object.keys(value) .sort() - .map((key) => { - const serializedValue = this.canonicalStringify(value[key]); - return serializedValue === undefined - ? undefined - : `${JSON.stringify(key)}:${serializedValue}`; - }) - .filter((property): property is string => property !== undefined); - return `{${serializedProperties.join(',')}}`; + .map((key) => `${JSON.stringify(key)}:${this.canonicalStringify(value[key])}`) + .join(',')}}`; } return JSON.stringify(value); diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 0c1b3bd1e5..632cd39da0 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -105,8 +105,8 @@ describe('DeployService', () => { const manifest = { models: [ { - name: 'model', - columns: [{ name: 'field' }, { name: 'measure' }], + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], }, ], }; @@ -128,41 +128,6 @@ describe('DeployService', () => { }); }); - it('should treat JSON-persisted manifests with omitted undefined properties as the same deployment', () => { - const manifest = { - models: [ - { - name: 'model', - properties: { - displayName: 'Model', - description: undefined, - }, - columns: [ - { - name: 'field', - expression: undefined, - properties: { - displayName: 'Field', - description: undefined, - }, - }, - ], - }, - ], - }; - const persistedManifest = JSON.parse(JSON.stringify(manifest)); - - expect(deployService.createMDLHash(manifest, 1)).toEqual( - deployService.createMDLHash(persistedManifest, 1), - ); - expect( - deployService.isSameDeployment(manifest, 1, { - hash: 'previous-hash', - manifest: persistedManifest, - }), - ).toBe(true); - }); - it('should create the same deployment hash for equivalent manifests', () => { const manifest = { models: [ diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index 9202777672..4aed639176 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -75,7 +75,7 @@ export interface IContext { modelNestedColumnRepository: IModelNestedColumnRepository; relationRepository: IRelationRepository; viewRepository: IViewRepository; - deployLogRepository: IDeployLogRepository; + deployRepository: IDeployLogRepository; schemaChangeRepository: ISchemaChangeRepository; learningRepository: ILearningRepository; dashboardRepository: IDashboardRepository; From 9e2870ab36e5248734a10fe3aed07f77aafc26da Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:56:34 +0530 Subject: [PATCH 0733/1087] Revert "Fix deploy sync status for equivalent manifests" This reverts commit 749ad225498365379e721e7f4a887970f354307c. --- .../apollo/server/resolvers/modelResolver.ts | 4 ++- .../apollo/server/services/deployService.ts | 2 +- .../services/tests/deployService.test.ts | 27 ------------------- 3 files changed, 4 insertions(+), 29 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 754394474b..5cabdea849 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -237,13 +237,15 @@ export class ModelResolver { try { const { id } = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const currentHash = ctx.deployService.createMDLHash(manifest, id); const lastDeploy = await ctx.deployService.getLastDeployment(id); + const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - return ctx.deployService.isSameDeployment(manifest, id, lastDeploy) + return currentHash == lastDeployHash ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 3f241a7455..746e5fa12e 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -107,7 +107,7 @@ export class DeployService implements IDeployService { // check if the model current deployment const lastDeploy = await this.deployLogRepository.findLastProjectDeployLog(projectId); - if (this.isSameDeployment(manifest, projectId, lastDeploy)) { + if (lastDeploy && lastDeploy.hash === hash) { logger.log(`Model has been deployed, hash: ${hash}`); await this.deployLogRepository.updateOne(lastDeploy.id, { status: DeployStatusEnum.SUCCESS, diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 632cd39da0..ce60461bf7 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -101,33 +101,6 @@ describe('DeployService', () => { }); }); - it('should skip deployment if an existing deployment has the same manifest with a different hash', async () => { - const manifest = { - models: [ - { - name: 'orders', - columns: [{ name: 'id' }, { name: 'amount' }], - }, - ], - }; - const projectId = 1; - - mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ - id: 123, - hash: 'legacy-hash', - manifest, - }); - - const response = await deployService.deploy(manifest, projectId); - - expect(response.status).toEqual(DeployStatusEnum.SUCCESS); - expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); - expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { - status: DeployStatusEnum.SUCCESS, - error: null, - }); - }); - it('should create the same deployment hash for equivalent manifests', () => { const manifest = { models: [ From 38d66c023d48990324feaa965144e9fe57c35e0c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 16:56:35 +0530 Subject: [PATCH 0734/1087] Revert "Fix Ask dry-run manifest consistency" This reverts commit 1af6fdab4796cf333c5df169cc3f76168ecfc1c6. --- .../generation/followup_sql_generation.py | 4 -- .../pipelines/generation/sql_correction.py | 33 +++++++++------ .../src/pipelines/generation/sql_diagnosis.py | 3 +- .../pipelines/generation/sql_generation.py | 4 -- .../pipelines/generation/sql_regeneration.py | 8 +--- .../src/pipelines/generation/utils/sql.py | 17 +------- wren-ai-service/src/providers/engine/wren.py | 2 - wren-ai-service/src/web/v1/services/ask.py | 14 ++++--- .../pipelines/generation/test_sql_utils.py | 42 +++++++++---------- wren-ui/src/apollo/server/models/model.ts | 1 - .../apollo/server/resolvers/modelResolver.ts | 29 ++++++------- wren-ui/src/apollo/server/schema.ts | 1 - 12 files changed, 67 insertions(+), 91 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 f8677de2a1..8d323d8376 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -148,14 +148,12 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, project_id: str | None = None, - mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), project_id=project_id, - mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -203,7 +201,6 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, - mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -227,7 +224,6 @@ async def run( "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, - "mdl_hash": mdl_hash, "sql_samples": sql_samples, "instructions": instructions, "has_calculated_field": has_calculated_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 4380b28540..26c2888b95 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,14 +30,19 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are a Wren SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect Wren SQL query. +You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. ### SQL CORRECTION INSTRUCTIONS ### -1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). -2. Then, generate the syntactically correct Wren SQL query to correct the error. -3. Keep executable table and column identifiers grounded in DATABASE SCHEMA. -4. Use SQL FUNCTIONS only when their exact syntax is provided. +1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. +2. Then, generate a syntactically correct ANSI SQL query from the user's intent and the current DATABASE SCHEMA. +3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. +4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. +5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. +6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. +7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -76,10 +81,18 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -SQL: {{ invalid_generation_result.sql }} -Error Message: {{ invalid_generation_result.error }} +{% if query %} +User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. +{% endif %} +{% if sql_generation_reasoning %} +### REASONING PLAN ### +The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. +{% endif %} +### FAILED SQL ### +The failed SQL and raw dry-run error are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. Regenerate from the user question and current DATABASE SCHEMA. -Let's think step by step. +Return only the final JSON SQL response. """ @@ -127,14 +140,12 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, project_id: str | None = None, - mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), project_id=project_id, - mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -182,7 +193,6 @@ async def run( instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, - mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, @@ -204,7 +214,6 @@ async def run( "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, - "mdl_hash": mdl_hash, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), diff --git a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py index e65b5e6719..3f22b9d512 100644 --- a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py +++ b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py @@ -20,7 +20,7 @@ sql_diagnosis_system_prompt = """ ### TASK ### -You are a Wren SQL expert with exceptional logical thinking skills and debugging skills, you need to diagnose the issue with the given SQL query, error message and database schema. +You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to diagnose the issue with the given SQL query, error message and database schema. ### SQL DIAGNOSIS INSTRUCTIONS ### @@ -29,7 +29,6 @@ 3. Then, return the reasoning behind the diagnosis.(You should give me the part of the original SQL query that is incorrect and the reason why it is incorrect) 4. Reasoning should be in the language same as the language user provided in the INPUTS section. 5. Reasoning should be concise and to the point and within 50 words. -6. Diagnose against Wren SQL syntax and the provided DATABASE SCHEMA. Do not suggest datasource-specific SQL syntax. ### FINAL ANSWER FORMAT ### The final answer must be in JSON format: diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 752c7caf83..4f71cfc091 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -139,7 +139,6 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, project_id: str | None = None, - mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, @@ -147,7 +146,6 @@ async def post_process( return await post_processor.run( generate_sql.get("replies"), project_id=project_id, - mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -195,7 +193,6 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, - mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -221,7 +218,6 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, - "mdl_hash": mdl_hash, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 5c01c84b11..07074d1433 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,7 +34,7 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great Wren SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query. +You are a great ANSI SQL expert. Now you are given database schema and a user's question. Carefully review the user's question and current DATABASE SCHEMA, then generate a new SQL query that answers the user's intent. The original SQL query and UI planning text are intentionally omitted from the prompt and must not be used as executable context. While generating the new SQL query, make sure to use the database schema as the only source of executable table and column identifiers. @@ -44,7 +44,7 @@ def get_sql_regeneration_system_prompt( {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a Wren SQL query in JSON format: +The final answer must be a ANSI SQL query in JSON format: {{ "sql": @@ -166,12 +166,10 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, - mdl_hash: str | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, - mdl_hash=mdl_hash, ) @@ -211,7 +209,6 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, - mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -230,7 +227,6 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, - "mdl_hash": mdl_hash, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1049c2a7fe..490eda1458 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -30,7 +30,6 @@ async def run( self, replies: List[str] | List[List[str]], project_id: str | None = None, - mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, data_source: str = "", @@ -51,7 +50,6 @@ async def run( ) = await self._classify_generation_result( cleaned_generation_result, project_id=project_id, - mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, @@ -74,7 +72,6 @@ async def _classify_generation_result( self, generation_result: str, project_id: str | None = None, - mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, data_source: str = "", @@ -112,7 +109,6 @@ async def _classify_generation_result( generation_result, session, project_id=project_id, - mdl_hash=mdl_hash, limit=1, dry_run=True, ) @@ -138,7 +134,6 @@ async def _classify_generation_result( generation_result, session, project_id=project_id, - mdl_hash=mdl_hash, limit=1, dry_run=False, ) @@ -170,9 +165,6 @@ async def _classify_generation_result( _MANDATORY_SQL_GROUNDING_RULES = """ ### MANDATORY SQL GROUNDING RULES ### -- Generate Wren SQL that can be parsed by the Wren engine before any datasource dialect rewrite happens. -- Use the Wren SQL clause order for result limiting: ORDER BY comes before LIMIT, and LIMIT belongs at the end of the SELECT or final UNION result. -- Datasource-specific SQL syntax is execution-target context only. Do not emit syntax that requires a datasource parser before the Wren engine rewrite step. - Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. - Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. @@ -358,16 +350,9 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: rules = _DEFAULT_TEXT_TO_SQL_RULES if sql_knowledge is not None: - additional_rules = _extract_from_sql_knowledge( + rules = _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) - if additional_rules != _DEFAULT_TEXT_TO_SQL_RULES: - rules = ( - f"{rules}\n\n" - "### ADDITIONAL SQL KNOWLEDGE ###\n" - "Use this section only when it is compatible with Wren SQL and the current SQL FUNCTIONS.\n" - f"{additional_rules}" - ) return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 3e812f0524..3a92853e04 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -28,7 +28,6 @@ async def execute_sql( sql: str, session: aiohttp.ClientSession, project_id: str | None = None, - mdl_hash: str | None = None, dry_run: bool = True, timeout: float = settings.engine_timeout, limit: int = 500, @@ -37,7 +36,6 @@ async def execute_sql( data = { "sql": remove_limit_statement(sql), "projectId": project_id, - "hash": mdl_hash, } if dry_run: data["dryRun"] = True diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2c4a55d457..6ddd3b14e2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -467,7 +467,6 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -486,7 +485,6 @@ async def ask( contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -550,15 +548,19 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, + query=user_query, + sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "error": ( + f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" + if allow_sql_diagnosis + and sql_diagnosis_reasoning + else error_message + ), }, project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index ba45b460fc..ef6346973f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -35,9 +35,6 @@ def test_construct_instructions_uses_instruction_text(): def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): rules = get_text_to_sql_rules() - assert "Generate Wren SQL that can be parsed by the Wren engine" in rules - assert "ORDER BY comes before LIMIT" in rules - assert "Datasource-specific SQL syntax is execution-target context only" in rules assert "ONLY USE the tables and columns mentioned in the database schema" in rules assert 'ONLY USE "*" if the user query asks for all the columns' in rules assert "They are never source table or source column identifiers" in rules @@ -67,12 +64,8 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): rules = get_text_to_sql_rules(_SqlKnowledge()) - assert "ONLY USE the tables and columns mentioned in the database schema" in rules - assert "For top, bottom, highest, lowest, first, or last requests" in rules assert _SqlKnowledge.text_to_sql_rule in rules - assert "Use this section only when it is compatible with Wren SQL" in rules assert "MANDATORY SQL GROUNDING RULES" in rules - assert "Generate Wren SQL that can be parsed by the Wren engine" in rules assert "Every table and column referenced" in rules assert "Do not query INFORMATION_SCHEMA" in rules @@ -123,8 +116,6 @@ def test_json_field_instructions_do_not_include_placeholder_identifiers(): def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): prompt = get_sql_regeneration_system_prompt() - assert "Wren SQL query" in prompt - assert "ANSI SQL" not in prompt assert "regenerate from the user's question" in prompt assert "unsupported identifiers" in prompt assert "original SQL query and UI planning text are intentionally omitted" in prompt @@ -138,11 +129,15 @@ def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): def test_sql_correction_system_prompt_discards_invalid_identifier_context(): prompt = get_sql_correction_system_prompt() - assert "Wren SQL query" in prompt - assert "ANSI SQL" not in prompt - assert "fix the syntactically incorrect Wren SQL query" in prompt - assert "generate the syntactically correct Wren SQL query" in prompt - assert "Keep executable table and column identifiers grounded" in prompt + assert "treat it as the source of intent" in prompt + assert "Do not copy placeholders" in prompt + assert "Regenerate a grounded Wren SQL query" in prompt + assert ( + "Do not preserve a table, column, join, filter, grouping, ordering, or function" + in prompt + ) + assert "Treat physical/source/lineage names from the failed SQL" in prompt + assert "do not try a similar replacement from source metadata" in prompt def test_sql_reasoning_prompt_forbids_executable_sql_context(): @@ -160,13 +155,14 @@ def test_user_prompt_templates_keep_source_metadata_non_executable(): sql_generation_user_prompt_template, text_to_sql_with_followup_user_prompt_template, sql_regeneration_user_prompt_template, + sql_correction_user_prompt_template, ): assert "source/physical/lineage names" in prompt assert "omit that unsupported part instead of inventing" in prompt assert "exact declared table and column names from DATABASE SCHEMA" in prompt -def test_executable_prompt_templates_omit_planning_and_original_sql_context(): +def test_executable_prompt_templates_omit_planning_error_and_original_sql_context(): marker = "UNTRUSTED_CONTEXT_MARKER" generation_prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( @@ -195,6 +191,15 @@ def test_executable_prompt_templates_omit_planning_and_original_sql_context(): sql_functions=[], )["prompt"] + correction_prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( + query="Question", + documents=["SCHEMA_CONTEXT"], + invalid_generation_result={"error": marker}, + sql_generation_reasoning=marker, + instructions=[], + sql_functions=[], + )["prompt"] + regeneration_prompt = PromptBuilder( template=sql_regeneration_user_prompt_template ).run( @@ -213,13 +218,8 @@ def test_executable_prompt_templates_omit_planning_and_original_sql_context(): for prompt in ( generation_prompt, followup_prompt, + correction_prompt, regeneration_prompt, ): assert marker not in prompt assert "intentionally omitted" in prompt - - -def test_sql_correction_user_prompt_follows_legacy_failed_sql_flow(): - assert "SQL: {{ invalid_generation_result.sql }}" in sql_correction_user_prompt_template - assert "Error Message: {{ invalid_generation_result.error }}" in sql_correction_user_prompt_template - assert "Let's think step by step." in sql_correction_user_prompt_template diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index 4f3ff216eb..c6403db9b9 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -98,7 +98,6 @@ export interface CheckCalculatedFieldCanQueryData { export interface PreviewSQLData { sql: string; projectId?: string; - hash?: string; limit?: number; dryRun?: boolean; } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 5cabdea849..3cd42c4505 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -236,16 +236,19 @@ export class ModelResolver { public async checkModelSync(_root: any, _args: any, ctx: IContext) { try { const { id } = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest, id); - const lastDeploy = await ctx.deployService.getLastDeployment(id); - const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } - return currentHash == lastDeployHash + + const project = await ctx.projectService.getCurrentProject(); + if (dirtyProjectIds.has(project.id)) { + return { status: SyncStatusEnum.UNSYNCRONIZED }; + } + + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + return lastDeploy ? { status: SyncStatusEnum.SYNCRONIZED } : { status: SyncStatusEnum.UNSYNCRONIZED }; } catch (err: any) { @@ -1204,7 +1207,7 @@ export class ModelResolver { // create view const project = await ctx.projectService.getCurrentProject(); - const manifest = await this.getDeployedManifest(ctx, project.id); + const manifest = await this.getLastDeployedManifest(ctx, project.id); // get sql statement of a response const response = await ctx.askingService.getResponse(responseId); @@ -1331,11 +1334,11 @@ export class ModelResolver { args: { data: PreviewSQLData }, ctx: IContext, ) { - const { sql, projectId, hash, limit, dryRun } = args.data; + const { sql, projectId, limit, dryRun } = args.data; const project = projectId ? await ctx.projectService.getProjectById(parseInt(projectId)) : await ctx.projectService.getCurrentProject(); - const manifest = await this.getDeployedManifest(ctx, project.id, hash); + const manifest = await this.getLastDeployedManifest(ctx, project.id); return await ctx.queryService.preview(sql, { project, limit: limit, @@ -1484,14 +1487,8 @@ export class ModelResolver { }; } - private async getDeployedManifest( - ctx: IContext, - projectId: number, - hash?: string, - ) { - const deployment = hash - ? await ctx.deployLogRepository.findOneBy({ projectId, hash }) - : await ctx.deployService.getLastDeployment(projectId); + private async getLastDeployedManifest(ctx: IContext, projectId: number) { + const deployment = await ctx.deployService.getLastDeployment(projectId); if (!deployment?.manifest) { throw new Error( 'Project has not been deployed successfully yet. Deploy the model before previewing or validating SQL.', diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 3c656de8c2..37e2549949 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -923,7 +923,6 @@ export const typeDefs = gql` input PreviewSQLDataInput { sql: String! projectId: String - hash: String limit: Int dryRun: Boolean } From 0d017ed163e95a01d4d3fb323a33ec7830c7b78f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 17:27:54 +0530 Subject: [PATCH 0735/1087] Add structured retrieved schema context --- .../src/pipelines/generation/utils/sql.py | 13 +- .../retrieval/db_schema_retrieval.py | 129 ++++++++++++++++-- .../pipelines/generation/test_sql_utils.py | 3 + .../retrieval/test_db_schema_retrieval.py | 4 + 4 files changed, 134 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 490eda1458..bf7180bdb9 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -170,9 +170,11 @@ async def _classify_generation_result( - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. - Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. - Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. +- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's executable name, executable columns, semantic metadata, relationships, views, metrics, and calculated fields. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. +- Never generate placeholder identifiers or placeholder table names. If the retrieved metadata does not contain an executable object or column for a requested concept, use the closest executable object and column whose semantic metadata supports the intent, or omit that unsupported concept. - If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. @@ -401,11 +403,12 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. 4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -6. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -7. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. -8. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If an identifier, literal value, placeholder, or function appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. -10. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +6. YOU MUST first read any WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use executable_name values and the following DDL declarations as the executable grounding, and use semantic_metadata only to understand business meaning. +7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +8. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. +9. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +10. If an identifier, literal value, placeholder, or function appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +11. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 62a1a659d3..36e000cb06 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -112,6 +112,22 @@ def _build_metric_ddl(content: dict) -> str: + context = _format_semantic_context( + { + "object_type": "metric", + "executable_name": content["name"], + "semantic_role": "stable analytical aggregation interface", + "columns": [ + { + "executable_name": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "semantic_metadata": column["comment"], + } + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + } + ) columns_ddl = [ f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" for column in content["columns"] @@ -120,16 +136,105 @@ def _build_metric_ddl(content: dict) -> str: ] return ( - f"{content['comment']}CREATE TABLE {content['name']} (\n " + f"{context}{content['comment']}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) def _build_view_ddl(content: dict) -> str: + context = _format_semantic_context( + { + "object_type": "view", + "executable_name": content["name"], + "semantic_role": "stable virtual table interface", + "definition_is_semantic_context": True, + } + ) + return ( + f"{context}{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" + ) + + +def _format_semantic_context(context: dict) -> str: return ( - f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" + "/*\n" + "WREN RETRIEVED SEMANTIC CONTEXT\n" + f"{orjson.dumps(context).decode('utf-8')}\n" + "Only executable_name values in this retrieved context and identifiers declared in the following DDL are executable in Wren SQL. Semantic metadata explains meaning but is not an executable identifier.\n" + "*/\n" + ) + + +def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: + relationship_columns = { + column.get("column") + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + } + relationship_columns.discard(None) + return relationship_columns + + +def _included_columns( + content: dict, columns: Optional[set[str]], tables: Optional[set[str]] +) -> list[dict]: + relationship_columns = _included_relationship_columns(content, tables) + return [ + column + for column in content["columns"] + if column["type"] == "COLUMN" + and ( + not columns + or column["name"] in columns + or column["name"] in relationship_columns + or column["is_primary_key"] + ) + and column["data_type"].lower() != "unknown" + ] + + +def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: + return [ + column + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + ] + + +def _build_table_retrieval_context( + content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None +) -> tuple[str, bool, bool]: + ddl, has_calculated_field, has_json_field = build_table_ddl( + content, columns=columns, tables=tables + ) + context = _format_semantic_context( + { + "object_type": "model", + "executable_name": content["name"], + "semantic_metadata": content["comment"], + "columns": [ + { + "executable_name": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "is_primary_key": column["is_primary_key"], + "semantic_metadata": column["comment"], + } + for column in _included_columns(content, columns, tables) + ], + "relationships": [ + { + "semantic_metadata": relationship["comment"], + "constraint": relationship["constraint"], + "related_models": relationship.get("tables", []), + } + for relationship in _included_relationships(content, tables) + ], + } ) + return f"{context}{ddl}", has_calculated_field, has_json_field ## Start of Pipeline @@ -334,7 +439,9 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) retrieval_results.append( { "table_name": table_schema["name"], @@ -397,7 +504,7 @@ def prompt( ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ - build_table_ddl(construct_db_schema)[0] + _build_table_retrieval_context(construct_db_schema)[0] for construct_db_schema in construct_db_schemas ] @@ -452,12 +559,14 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - ddl, _has_calculated_field, _has_json_field = build_table_ddl( - table_schema, - columns=set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ), - tables=tables, + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context( + table_schema, + columns=set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ), + tables=tables, + ) ) if _has_calculated_field: has_calculated_field = True diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index ef6346973f..c85f25a487 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -47,6 +47,8 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "order by that alias" in rules assert "Interpret the user's intent" in rules assert "schema descriptions, aliases, display labels" in rules + assert "WREN RETRIEVED SEMANTIC CONTEXT" in rules + assert "Never generate placeholder identifiers" in rules assert "use all required related tables" in rules assert "silently check that each identifier and function" in rules assert "instead of inventing a replacement" in rules @@ -99,6 +101,7 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "reasoning plan only as non-executable context" in prompt assert "include those objects only when DATABASE SCHEMA shows" in prompt assert "Use the exact supported syntax shown there" in prompt + assert "Use executable_name values and the following DDL declarations" in prompt assert "source database/schema/table names" in prompt assert "appears only in SQL samples, reasoning, failed SQL" in prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index befe2497cd..159d0a15e9 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -325,6 +325,10 @@ def table_schema(name): "activity", "account", ] + assert all( + "WREN RETRIEVED SEMANTIC CONTEXT" in schema["table_ddl"] + for schema in result["db_schemas"] + ) assert result["tokens"] > 0 From 6744d8d93be481cb3f1e06b2bb3155eb74f38e9b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 17:52:27 +0530 Subject: [PATCH 0736/1087] Clarify retrieved SQL identifier contract --- .../src/pipelines/generation/utils/sql.py | 7 +- .../retrieval/db_schema_retrieval.py | 73 ++++++++++++++----- .../pipelines/generation/test_sql_utils.py | 9 ++- .../retrieval/test_db_schema_retrieval.py | 52 +++++++++++++ 4 files changed, 119 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index bf7180bdb9..354b4376ea 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -170,7 +170,9 @@ async def _classify_generation_result( - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. - Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. - Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. -- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's executable name, executable columns, semantic metadata, relationships, views, metrics, and calculated fields. +- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. +- In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. +- Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. @@ -203,6 +205,7 @@ async def _classify_generation_result( - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. + - Use Wren SQL identifier quoting with double quotes only; the engine rewrite step converts grounded Wren SQL to the active connector dialect. - Put single quotes around string literals. - Never quote numeric literals. - For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. @@ -403,7 +406,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. 4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -6. YOU MUST first read any WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use executable_name values and the following DDL declarations as the executable grounding, and use semantic_metadata only to understand business meaning. +6. YOU MUST first read any WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. 7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 8. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. 9. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 36e000cb06..7c1fb5f2aa 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -51,6 +51,8 @@ 11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. 12. Follow only the relationships shown in the provided schema when selecting columns across datasets. 13. Do not stop at a single top candidate when the question needs multiple related datasets. +14. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. +15. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -112,19 +114,31 @@ def _build_metric_ddl(content: dict) -> str: + columns = [ + column + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ] context = _format_semantic_context( { "object_type": "metric", - "executable_name": content["name"], - "semantic_role": "stable analytical aggregation interface", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in columns + ], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable analytical aggregation interface", + "description": content["comment"], + }, "columns": [ { - "executable_name": column["name"], + "sql_column_name_use_exactly": column["name"], "data_type": get_engine_supported_data_type(column["data_type"]), - "semantic_metadata": column["comment"], + "semantic_context_not_sql_identifier": column["comment"], } - for column in content["columns"] - if column["data_type"].lower() != "unknown" + for column in columns ], } ) @@ -146,9 +160,14 @@ def _build_view_ddl(content: dict) -> str: context = _format_semantic_context( { "object_type": "view", - "executable_name": content["name"], - "semantic_role": "stable virtual table interface", - "definition_is_semantic_context": True, + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable virtual table interface", + "description": content["comment"], + "definition_is_semantic_context": True, + }, } ) return ( @@ -161,7 +180,8 @@ def _format_semantic_context(context: dict) -> str: "/*\n" "WREN RETRIEVED SEMANTIC CONTEXT\n" f"{orjson.dumps(context).decode('utf-8')}\n" - "Only executable_name values in this retrieved context and identifiers declared in the following DDL are executable in Wren SQL. Semantic metadata explains meaning but is not an executable identifier.\n" + "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" + "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" "*/\n" ) @@ -210,27 +230,42 @@ def _build_table_retrieval_context( ddl, has_calculated_field, has_json_field = build_table_ddl( content, columns=columns, tables=tables ) + included_columns = _included_columns(content, columns, tables) + included_relationships = _included_relationships(content, tables) context = _format_semantic_context( { "object_type": "model", - "executable_name": content["name"], - "semantic_metadata": content["comment"], + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in included_columns + ], + "relationship_constraints_use_exactly": [ + relationship["constraint"] + for relationship in included_relationships + ], + }, + "semantic_context_not_sql_identifiers": { + "description": content["comment"], + }, "columns": [ { - "executable_name": column["name"], + "sql_column_name_use_exactly": column["name"], "data_type": get_engine_supported_data_type(column["data_type"]), "is_primary_key": column["is_primary_key"], - "semantic_metadata": column["comment"], + "semantic_context_not_sql_identifier": column["comment"], } - for column in _included_columns(content, columns, tables) + for column in included_columns ], "relationships": [ { - "semantic_metadata": relationship["comment"], - "constraint": relationship["constraint"], - "related_models": relationship.get("tables", []), + "semantic_context_not_sql_identifier": relationship["comment"], + "sql_relationship_constraint_use_exactly": relationship[ + "constraint" + ], + "related_models_use_exactly": relationship.get("tables", []), } - for relationship in _included_relationships(content, tables) + for relationship in included_relationships ], } ) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index c85f25a487..118970f458 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -48,6 +48,10 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "Interpret the user's intent" in rules assert "schema descriptions, aliases, display labels" in rules assert "WREN RETRIEVED SEMANTIC CONTEXT" in rules + assert "sql_table_name_use_exactly" in rules + assert "sql_column_name_use_exactly" in rules + assert "semantic_context_not_sql_identifier" in rules + assert "Do not combine words, labels, ordinals" in rules assert "Never generate placeholder identifiers" in rules assert "use all required related tables" in rules assert "silently check that each identifier and function" in rules @@ -101,7 +105,10 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "reasoning plan only as non-executable context" in prompt assert "include those objects only when DATABASE SCHEMA shows" in prompt assert "Use the exact supported syntax shown there" in prompt - assert "Use executable_name values and the following DDL declarations" in prompt + assert "Use sql_table_name_use_exactly" in prompt + assert "sql_column_names_use_exactly" in prompt + assert "semantic_context_not_sql_identifiers" in prompt + assert "Use Wren SQL identifier quoting with double quotes only" in prompt assert "source database/schema/table names" in prompt assert "appears only in SQL samples, reasoning, failed SQL" in prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 159d0a15e9..7b2c5f7804 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -329,9 +329,61 @@ def table_schema(name): "WREN RETRIEVED SEMANTIC CONTEXT" in schema["table_ddl"] for schema in result["db_schemas"] ) + assert all( + "sql_table_name_use_exactly" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert all( + "sql_column_name_use_exactly" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert all( + "semantic_context_not_sql_identifier" in schema["table_ddl"] + for schema in result["db_schemas"] + ) assert result["tokens"] > 0 +def test_retrieved_schema_separates_exact_sql_names_from_semantic_context(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "Business-facing dataset description.", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "Business-facing attribute label.", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + table_ddl = result["db_schemas"][0]["table_ddl"] + + assert '"sql_table_name_use_exactly":"modeled_dataset"' in table_ddl + assert '"sql_column_name_use_exactly":"stored_attribute"' in table_ddl + assert ( + '"semantic_context_not_sql_identifier":"Business-facing attribute label."' + in table_ddl + ) + + def test_check_using_db_schemas_without_pruning_keeps_explicit_table_fast_path(): class Encoding: def encode(self, value): From d5d42bbbc42525d0301b91d6961739b6f5227fbf Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 18:30:20 +0530 Subject: [PATCH 0737/1087] Keep executable DDL free of semantic labels --- wren-ai-service/src/pipelines/common.py | 32 ++++++++++++++----- .../pipelines/generation/sql_correction.py | 12 ++++++- .../retrieval/db_schema_retrieval.py | 13 ++++---- .../pipelines/generation/test_sql_utils.py | 23 ++++++++----- .../retrieval/test_db_schema_retrieval.py | 28 ++++++++++++++++ 5 files changed, 85 insertions(+), 23 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index f0ca079a72..605238bb64 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -4,10 +4,13 @@ from haystack import Document, component -def get_engine_supported_data_type(data_type: str) -> str: +def get_engine_supported_data_type(data_type: str | None) -> str: """ This function makes sure downstream ai pipeline get column data types in a format that is supported by the data engine. """ + if not data_type: + return "UNKNOWN" + match data_type.upper(): case "BPCHAR" | "NAME" | "UUID" | "INET": return "VARCHAR" @@ -28,7 +31,10 @@ def get_engine_supported_data_type(data_type: str) -> str: def build_table_ddl( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None + content: dict, + columns: Optional[set[str]] = None, + tables: Optional[set[str]] = None, + include_semantic_comments: bool = True, ) -> Tuple[str, bool, bool]: columns_ddl = [] has_calculated_field = False @@ -43,6 +49,8 @@ def build_table_ddl( for column in content["columns"]: if column["type"] == "COLUMN": + raw_data_type = column["data_type"] + supported_data_type = get_engine_supported_data_type(raw_data_type) if ( ( not columns @@ -50,24 +58,32 @@ def build_table_ddl( or column["name"] in relationship_columns or column["is_primary_key"] ) - and column["data_type"].lower() - != "unknown" # quick fix: filtering out UNKNOWN column type + and ( + raw_data_type is None + or supported_data_type.lower() + != "unknown" # quick fix: filtering out UNKNOWN column type + ) ): if "This column is a Calculated Field" in column["comment"]: has_calculated_field = True - if column["data_type"].lower() == "json": + if supported_data_type.lower() == "json": has_json_field = True - column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + column_comment = column["comment"] if include_semantic_comments else "" + column_ddl = f"{column_comment}{column['name']} {supported_data_type}" if column["is_primary_key"]: column_ddl += " PRIMARY KEY" columns_ddl.append(column_ddl) elif column["type"] == "FOREIGN_KEY": if not tables or (tables and set(column.get("tables", [])).issubset(tables)): - columns_ddl.append(f"{column['comment']}{column['constraint']}") + relationship_comment = ( + column["comment"] if include_semantic_comments else "" + ) + columns_ddl.append(f"{relationship_comment}{column['constraint']}") + table_comment = content["comment"] if include_semantic_comments else "" return ( ( - f"{content['comment']}CREATE TABLE {content['name']} (\n " + f"{table_comment}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ), diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 26c2888b95..afaa4e54aa 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -90,7 +90,17 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. {% endif %} ### FAILED SQL ### -The failed SQL and raw dry-run error are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. Regenerate from the user question and current DATABASE SCHEMA. +The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. + +### DRY-RUN DIAGNOSTIC ### +{% if invalid_generation_result and invalid_generation_result.error %} +Diagnostic text: +{{ invalid_generation_result.error }} +{% else %} +No diagnostic text was provided. +{% endif %} + +Use the diagnostic text only to understand the failure category. Do not copy identifiers, literal values, functions, SQL snippets, physical names, source names, or replacement candidates from the diagnostic text. Regenerate from the user question and current DATABASE SCHEMA. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 7c1fb5f2aa..89a2e50486 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -143,14 +143,12 @@ def _build_metric_ddl(content: dict) -> str: } ) columns_ddl = [ - f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" - for column in content["columns"] - if column["data_type"].lower() - != "unknown" # quick fix: filtering out UNKNOWN column type + f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in columns ] return ( - f"{context}{content['comment']}CREATE TABLE {content['name']} (\n " + f"{context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) @@ -228,7 +226,10 @@ def _build_table_retrieval_context( content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None ) -> tuple[str, bool, bool]: ddl, has_calculated_field, has_json_field = build_table_ddl( - content, columns=columns, tables=tables + content, + columns=columns, + tables=tables, + include_semantic_comments=False, ) included_columns = _included_columns(content, columns, tables) included_relationships = _included_relationships(content, tables) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 118970f458..55124870ca 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -173,12 +173,13 @@ def test_user_prompt_templates_keep_source_metadata_non_executable(): def test_executable_prompt_templates_omit_planning_error_and_original_sql_context(): - marker = "UNTRUSTED_CONTEXT_MARKER" + reasoning_marker = "UNTRUSTED_REASONING_CONTEXT_MARKER" + diagnostic_marker = "UNTRUSTED_DIAGNOSTIC_CONTEXT_MARKER" generation_prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( query="Question", documents=["SCHEMA_CONTEXT"], - sql_generation_reasoning=marker, + sql_generation_reasoning=reasoning_marker, instructions=[], calculated_field_instructions="", metric_instructions="", @@ -192,7 +193,7 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex ).run( query="Question", documents=["SCHEMA_CONTEXT"], - sql_generation_reasoning=marker, + sql_generation_reasoning=reasoning_marker, instructions=[], calculated_field_instructions="", metric_instructions="", @@ -204,8 +205,8 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex correction_prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( query="Question", documents=["SCHEMA_CONTEXT"], - invalid_generation_result={"error": marker}, - sql_generation_reasoning=marker, + invalid_generation_result={"error": diagnostic_marker}, + sql_generation_reasoning=reasoning_marker, instructions=[], sql_functions=[], )["prompt"] @@ -214,9 +215,9 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex template=sql_regeneration_user_prompt_template ).run( query="Question", - sql=marker, + sql=reasoning_marker, documents=["SCHEMA_CONTEXT"], - sql_generation_reasoning=marker, + sql_generation_reasoning=reasoning_marker, instructions=[], calculated_field_instructions="", metric_instructions="", @@ -231,5 +232,11 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex correction_prompt, regeneration_prompt, ): - assert marker not in prompt + assert reasoning_marker not in prompt assert "intentionally omitted" in prompt + + assert diagnostic_marker in correction_prompt + assert "Use the diagnostic text only to understand the failure category" in ( + correction_prompt + ) + assert "Do not copy identifiers" in correction_prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7b2c5f7804..bcb803b57c 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -375,6 +375,7 @@ def encode(self, value): ) table_ddl = result["db_schemas"][0]["table_ddl"] + executable_ddl = table_ddl.split("*/", maxsplit=1)[1] assert '"sql_table_name_use_exactly":"modeled_dataset"' in table_ddl assert '"sql_column_name_use_exactly":"stored_attribute"' in table_ddl @@ -382,6 +383,33 @@ def encode(self, value): '"semantic_context_not_sql_identifier":"Business-facing attribute label."' in table_ddl ) + assert "Business-facing attribute label." not in executable_ddl + assert "Business-facing dataset description." not in executable_ddl + assert "CREATE TABLE modeled_dataset" in executable_ddl + assert "stored_attribute VARCHAR" in executable_ddl + + +def test_build_table_ddl_can_render_executable_schema_without_semantic_comments(): + ddl, has_calculated_field, has_json_field = build_table_ddl( + { + "comment": "/* semantic table context */\n", + "name": "modeled_dataset", + "columns": [ + { + "type": "COLUMN", + "comment": "-- semantic field context\n ", + "name": "stored_attribute", + "data_type": "VARCHAR", + "is_primary_key": False, + } + ], + }, + include_semantic_comments=False, + ) + + assert ddl == "CREATE TABLE modeled_dataset (\n stored_attribute VARCHAR\n);" + assert not has_calculated_field + assert not has_json_field def test_check_using_db_schemas_without_pruning_keeps_explicit_table_fast_path(): From 2d6c1c7e9ed5ebb6661080aa1d5baeaede844927 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 19:00:24 +0530 Subject: [PATCH 0738/1087] Ground retrieval on current ask context --- wren-ai-service/src/config.py | 2 +- .../src/pipelines/generation/utils/sql.py | 16 +---- .../retrieval/db_schema_retrieval.py | 15 +---- wren-ai-service/tests/data/config.test.yaml | 2 +- .../pipelines/generation/test_sql_utils.py | 12 ++++ .../retrieval/test_db_schema_retrieval.py | 58 +++++++++++++++++++ wren-ai-service/tests/pytest/test_config.py | 2 +- .../tools/config/config.example.yaml | 2 +- wren-ai-service/tools/config/config.full.yaml | 2 +- 9 files changed, 77 insertions(+), 34 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index c5acf4ae47..04c6c05762 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -27,7 +27,7 @@ class Settings(BaseSettings): # indexing and retrieval config column_indexing_batch_size: int = Field(default=50) - table_retrieval_size: int = Field(default=10) + table_retrieval_size: int = Field(default=50) table_column_retrieval_size: int = Field(default=100) enable_column_pruning: bool = Field(default=False) historical_question_retrieval_similarity_threshold: float = Field(default=0.9) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 354b4376ea..01a00a97d3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -454,18 +454,4 @@ def construct_instructions( def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: - messages = [] - for history in histories: - messages.append( - ChatMessage.from_user( - history.question - if hasattr(history, "question") - else history["question"] - ) - ) - messages.append( - ChatMessage.from_assistant( - "Previous SQL omitted. Use only the current DATABASE SCHEMA and SQL FUNCTIONS for executable SQL." - ) - ) - return messages + return [] diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 89a2e50486..7934b1d51c 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -277,13 +277,6 @@ def _build_table_retrieval_context( @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: - if histories: - previous_query_summaries = [history.question for history in histories] - else: - previous_query_summaries = [] - - query = "\n".join(previous_query_summaries) + "\n" + query - return await embedder.run(query) else: return {} @@ -544,12 +537,6 @@ def prompt( for construct_db_schema in construct_db_schemas ] - previous_query_summaries = ( - [history.question for history in histories] if histories else [] - ) - - query = "\n".join(previous_query_summaries) + "\n" + query - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: @@ -688,7 +675,7 @@ def __init__( llm_provider: LLMProvider, embedder_provider: EmbedderProvider, document_store_provider: DocumentStoreProvider, - table_retrieval_size: int = 10, + table_retrieval_size: int = 50, table_column_retrieval_size: int = 100, **kwargs, ): diff --git a/wren-ai-service/tests/data/config.test.yaml b/wren-ai-service/tests/data/config.test.yaml index 65613d2c83..2c57e4d28a 100644 --- a/wren-ai-service/tests/data/config.test.yaml +++ b/wren-ai-service/tests/data/config.test.yaml @@ -80,7 +80,7 @@ settings: column_indexing_batch_size: 50 doc_endpoint: https://docs.getwren.ai is_oss: true - table_retrieval_size: 10 + table_retrieval_size: 50 table_column_retrieval_size: 1000 query_cache_maxsize: 1000 query_cache_ttl: 3600 diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 55124870ca..5f6ea69ac8 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,6 +1,7 @@ from haystack.components.builders.prompt_builder import PromptBuilder from src.pipelines.generation.utils.sql import ( + construct_ask_history_messages, construct_instructions, get_json_field_instructions, get_metric_instructions, @@ -32,6 +33,17 @@ def test_construct_instructions_uses_instruction_text(): ) == ["First rule.", "Second rule."] +def test_construct_ask_history_messages_omits_executable_history_context(): + histories = [ + { + "question": "previous natural language request", + "sql": "SELECT * FROM previous_model", + } + ] + + assert construct_ask_history_messages(histories) == [] + + def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): rules = get_text_to_sql_rules() diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index bcb803b57c..1e6a02d0d6 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1,14 +1,72 @@ import pytest from haystack import Document +from haystack.components.builders.prompt_builder import PromptBuilder from src.pipelines.common import build_table_ddl from src.pipelines.retrieval.db_schema_retrieval import ( check_using_db_schemas_without_pruning, dbschema_retrieval, + embedding, + prompt as build_column_selection_prompt, table_retrieval, + table_columns_selection_user_prompt_template, ) +@pytest.mark.asyncio +async def test_embedding_uses_current_query_without_history_text(): + class Embedder: + def __init__(self): + self.query = None + + async def run(self, query): + self.query = query + return {"embedding": [1.0]} + + embedder = Embedder() + + result = await embedding( + query="current request", + embedder=embedder, + histories=[{"question": "previous request"}], + ) + + assert result == {"embedding": [1.0]} + assert embedder.query == "current request" + + +def test_column_pruning_prompt_uses_current_query_without_history_text(): + result = build_column_selection_prompt( + query="current request", + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + prompt_builder=PromptBuilder( + template=table_columns_selection_user_prompt_template + ), + check_using_db_schemas_without_pruning={"db_schemas": []}, + histories=[{"question": "previous request"}], + ) + + assert "current request" in result["prompt"] + assert "previous request" not in result["prompt"] + + @pytest.mark.asyncio async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: diff --git a/wren-ai-service/tests/pytest/test_config.py b/wren-ai-service/tests/pytest/test_config.py index 70e3ddace3..4a745d49e9 100644 --- a/wren-ai-service/tests/pytest/test_config.py +++ b/wren-ai-service/tests/pytest/test_config.py @@ -12,7 +12,7 @@ def test_settings_default_values(): assert settings.port == 5555 assert settings.column_indexing_batch_size == 50 - assert settings.table_retrieval_size == 10 + assert settings.table_retrieval_size == 50 assert settings.table_column_retrieval_size == 100 assert settings.query_cache_ttl == 3600 diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index b7675000e0..b1b98a0962 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -186,7 +186,7 @@ settings: is_oss: true engine_timeout: 30 column_indexing_batch_size: 50 - table_retrieval_size: 10 + table_retrieval_size: 50 table_column_retrieval_size: 100 allow_intent_classification: false allow_sql_generation_reasoning: true diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 40657375ea..265292b9f1 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -183,7 +183,7 @@ settings: is_oss: true engine_timeout: 30 column_indexing_batch_size: 50 - table_retrieval_size: 10 + table_retrieval_size: 50 table_column_retrieval_size: 100 query_cache_maxsize: 1000 allow_intent_classification: false From 006054fa760c5fb9bc8bfff722b589ac6cdf3ac8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 19:14:46 +0530 Subject: [PATCH 0739/1087] Preserve semantic interfaces during retrieval pruning --- .../generation/followup_sql_generation.py | 4 +- .../retrieval/db_schema_retrieval.py | 35 +++++----- .../pipelines/generation/test_sql_utils.py | 7 ++ .../retrieval/test_db_schema_retrieval.py | 69 +++++++++++++++++++ 4 files changed, 95 insertions(+), 20 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 8d323d8376..762434b248 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -31,8 +31,8 @@ text_to_sql_with_followup_user_prompt_template = """ ### TASK ### -Given the following user's follow-up question and previous SQL query and summary, -generate one SQL query to best answer user's question. +Given the user's current follow-up question and the current retrieved DATABASE SCHEMA, +generate one SQL query to best answer the user's question. ### DATABASE SCHEMA ### {% for document in documents %} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 7934b1d51c..94210428ab 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -604,24 +604,23 @@ def construct_retrieval_results( ) for document in dbschema_retrieval: - if document.meta["name"] in columns_and_tables_needed: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - } - ) + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) return { "retrieval_results": retrieval_results, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 5f6ea69ac8..8fea0e8946 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -184,6 +184,13 @@ def test_user_prompt_templates_keep_source_metadata_non_executable(): assert "exact declared table and column names from DATABASE SCHEMA" in prompt +def test_followup_sql_prompt_does_not_expect_previous_sql_context(): + assert "previous SQL query" not in text_to_sql_with_followup_user_prompt_template + assert "current retrieved DATABASE SCHEMA" in ( + text_to_sql_with_followup_user_prompt_template + ) + + def test_executable_prompt_templates_omit_planning_error_and_original_sql_context(): reasoning_marker = "UNTRUSTED_REASONING_CONTEXT_MARKER" diagnostic_marker = "UNTRUSTED_DIAGNOSTIC_CONTEXT_MARKER" diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 1e6a02d0d6..de9c136e89 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -5,6 +5,7 @@ from src.pipelines.common import build_table_ddl from src.pipelines.retrieval.db_schema_retrieval import ( check_using_db_schemas_without_pruning, + construct_retrieval_results, dbschema_retrieval, embedding, prompt as build_column_selection_prompt, @@ -345,6 +346,74 @@ def encode(self, value): assert result["tokens"] > 0 +def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_attribute"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "name": "semantic_metric", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "semantic_metric"}, + ) + ], + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "modeled_dataset", + "semantic_metric", + ] + assert result["has_metric"] is True + + def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): From 26bd793a33079f62fb0d8f134bb88c0b723ed2a8 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 20:14:36 +0530 Subject: [PATCH 0740/1087] Restore schema-grounded reasoning handoff --- .../generation/followup_sql_generation.py | 5 ++- .../pipelines/generation/sql_correction.py | 5 ++- .../pipelines/generation/sql_generation.py | 6 ++- .../generation/sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_regeneration.py | 5 ++- .../src/pipelines/generation/utils/sql.py | 34 +++++++-------- .../retrieval/db_schema_retrieval.py | 41 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 31 +++++++++----- .../retrieval/test_db_schema_retrieval.py | 8 ++++ 9 files changed, 101 insertions(+), 36 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 762434b248..ff7cf38092 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -79,7 +79,8 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. ### REASONING PLAN ### -The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. +Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +{{ sql_generation_reasoning }} Return only the final JSON SQL response. """ @@ -103,7 +104,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=bool(sql_generation_reasoning), + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index afaa4e54aa..abb690d0a0 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -87,7 +87,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### -The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. +Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +{{ sql_generation_reasoning }} {% endif %} ### FAILED SQL ### The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. @@ -121,7 +122,7 @@ def prompt( query=query, documents=documents, invalid_generation_result=invalid_generation_result, - sql_generation_reasoning=bool(sql_generation_reasoning), + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 4f71cfc091..a1ed60a4d6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -73,7 +73,9 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. {% if sql_generation_reasoning %} -The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. +### REASONING PLAN ### +Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +{{ sql_generation_reasoning }} {% endif %} Return only the final JSON SQL response. @@ -98,7 +100,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=bool(sql_generation_reasoning), + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), 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 9867af1b2f..5311536d03 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -48,7 +48,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. +Return only the reasoning plan described by the system instructions. Ground the plan with exact `table: ` and `column: .` references from DATABASE SCHEMA when they are relevant. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 07074d1433..1a96d18c84 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -98,7 +98,8 @@ def get_sql_regeneration_system_prompt( Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### -The UI planning text is intentionally omitted from this executable SQL prompt so it cannot provide table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions. +Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +{{ sql_generation_reasoning }} ### ORIGINAL SQL QUERY ### The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. @@ -126,7 +127,7 @@ def prompt( query=query, sql=sql, documents=documents, - sql_generation_reasoning=bool(sql_generation_reasoning), + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 01a00a97d3..a38f481751 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -171,6 +171,7 @@ async def _classify_generation_result( - Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. - Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. - When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. +- When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. - In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. - Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. @@ -315,8 +316,8 @@ async def _classify_generation_result( ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state requested timeframes in natural language only. Do not name date columns, functions, expressions, or SQL clauses. -3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language only. Do not name columns, aliases, aggregate expressions, or SQL clauses. +2. Explicitly state requested timeframes in natural language only. Mention exact date/time columns only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language. Mention exact ordering columns or measures only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. 4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. @@ -325,17 +326,18 @@ async def _classify_generation_result( 9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. Do not mention table names, view names, metric names, column names, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels in the reasoning plan. -13. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. -14. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. -15. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language only. -16. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, but do not name the underlying schema objects in the reasoning plan. -17. Only cite exact declared names from DATABASE SCHEMA if an internal grounding note requires it; do not expose table names, column names, source metadata, physical datasource names, or lineage names in the reasoning plan. -18. If multiple schema objects may be required to answer the intent, describe the need to combine related data in natural language only. -19. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan. -20. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -21. The reasoning plan is non-executable context. Do not include anything that could be copied as SQL. -22. ONLY SHOWING the reasoning plan in bullet points. +12. Mention table names only in this exact format: `table: `, and only when `` is declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +13. Mention column names only in this exact format: `column: .`, and only when both `` and `` are declared together in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +14. Do not mention aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels from comments, SQL samples, failed SQL, or user wording as executable identifiers. +15. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. +16. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. +17. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language and cite exact declared tables or columns only when they are grounded by DATABASE SCHEMA. +18. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, then ground the plan in exact declared schema identifiers. +19. If multiple schema objects are required, identify the exact declared relationship path from DATABASE SCHEMA. If no relationship path is declared, say that the retrieved metadata does not provide a join path. +20. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan unless they also appear exactly in DATABASE SCHEMA. +21. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +22. The reasoning plan is a grounding handoff for SQL generation, not SQL. Do not include executable SQL. +23. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -404,13 +406,13 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. -4. YOU MUST use the reasoning plan only as non-executable context, and only when it is consistent with DATABASE SCHEMA and SQL Rules. Do not copy table names, column names, aliases, source names, physical names, lineage names, SQL fragments, date expressions, literal values, placeholders, or functions from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +4. YOU MUST use the reasoning plan as the legacy grounding handoff only when it is consistent with DATABASE SCHEMA and SQL Rules. Table and column references in the reasoning plan are executable only when they also appear exactly in DATABASE SCHEMA. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -6. YOU MUST first read any WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. +6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. 7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 8. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. 9. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -10. If an identifier, literal value, placeholder, or function appears only in SQL samples, reasoning, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +10. If an identifier, literal value, placeholder, or function appears only in SQL samples, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. Reasoning-plan identifiers are usable only when they exactly match DATABASE SCHEMA. 11. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 94210428ab..3c4e894831 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -178,12 +178,53 @@ def _format_semantic_context(context: dict) -> str: "/*\n" "WREN RETRIEVED SEMANTIC CONTEXT\n" f"{orjson.dumps(context).decode('utf-8')}\n" + f"{_format_identifier_contract(context)}" "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" "*/\n" ) +def _format_identifier_contract(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "WREN SQL IDENTIFIER CONTRACT", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"sql_table_name_use_exactly: {table_name}") + if column_names: + lines.append("sql_column_names_use_exactly:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationship_constraints_use_exactly:") + lines.extend( + f"- {relationship_constraint}" + for relationship_constraint in relationship_constraints + ) + lines.extend( + [ + "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", + "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", + "END WREN SQL IDENTIFIER CONTRACT", + "", + ] + ) + return "\n".join(lines) + + def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: relationship_columns = { column.get("column") diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 8fea0e8946..4e020471a3 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -60,6 +60,8 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "Interpret the user's intent" in rules assert "schema descriptions, aliases, display labels" in rules assert "WREN RETRIEVED SEMANTIC CONTEXT" in rules + assert "WREN SQL IDENTIFIER CONTRACT" in rules + assert "compact authoritative list of executable identifiers" in rules assert "sql_table_name_use_exactly" in rules assert "sql_column_name_use_exactly" in rules assert "semantic_context_not_sql_identifier" in rules @@ -114,15 +116,16 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "perform a silent grounding check" in prompt assert "closest grounded expression" in prompt assert "DATABASE SCHEMA is the only source of executable identifiers" in prompt - assert "reasoning plan only as non-executable context" in prompt + assert "reasoning plan as the legacy grounding handoff" in prompt assert "include those objects only when DATABASE SCHEMA shows" in prompt assert "Use the exact supported syntax shown there" in prompt + assert "WREN SQL IDENTIFIER CONTRACT" in prompt assert "Use sql_table_name_use_exactly" in prompt assert "sql_column_names_use_exactly" in prompt assert "semantic_context_not_sql_identifiers" in prompt assert "Use Wren SQL identifier quoting with double quotes only" in prompt assert "source database/schema/table names" in prompt - assert "appears only in SQL samples, reasoning, failed SQL" in prompt + assert "appears only in SQL samples, failed SQL" in prompt def test_json_field_instructions_do_not_include_placeholder_identifiers(): @@ -140,7 +143,8 @@ def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): assert "regenerate from the user's question" in prompt assert "unsupported identifiers" in prompt - assert "original SQL query and UI planning text are intentionally omitted" in prompt + assert "original SQL query" in prompt + assert "intentionally omitted" in prompt assert ( "database schema as the only source of executable table and column identifiers" in prompt @@ -162,14 +166,16 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "do not try a similar replacement from source metadata" in prompt -def test_sql_reasoning_prompt_forbids_executable_sql_context(): +def test_sql_reasoning_prompt_uses_legacy_schema_grounding_handoff(): prompt = sql_generation_reasoning_system_prompt assert "Do not write SQL, possible SQL, sample SQL, assumed SQL" in prompt assert "SQL clauses, SQL functions, code blocks, or executable expressions" in prompt - assert "The reasoning plan is non-executable context" in prompt - assert "Only cite exact declared names from DATABASE SCHEMA" in prompt - assert "source metadata, physical datasource names, or lineage names" in prompt + assert "table: " in prompt + assert "column: ." in prompt + assert "The reasoning plan is a grounding handoff" in prompt + assert "declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT" in prompt + assert "source names, physical names, lineage names" in prompt def test_user_prompt_templates_keep_source_metadata_non_executable(): @@ -191,9 +197,10 @@ def test_followup_sql_prompt_does_not_expect_previous_sql_context(): ) -def test_executable_prompt_templates_omit_planning_error_and_original_sql_context(): +def test_executable_prompt_templates_include_grounded_reasoning_but_omit_sql_context(): reasoning_marker = "UNTRUSTED_REASONING_CONTEXT_MARKER" diagnostic_marker = "UNTRUSTED_DIAGNOSTIC_CONTEXT_MARKER" + original_sql_marker = "UNTRUSTED_ORIGINAL_SQL_MARKER" generation_prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( query="Question", @@ -234,7 +241,7 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex template=sql_regeneration_user_prompt_template ).run( query="Question", - sql=reasoning_marker, + sql=original_sql_marker, documents=["SCHEMA_CONTEXT"], sql_generation_reasoning=reasoning_marker, instructions=[], @@ -251,9 +258,11 @@ def test_executable_prompt_templates_omit_planning_error_and_original_sql_contex correction_prompt, regeneration_prompt, ): - assert reasoning_marker not in prompt - assert "intentionally omitted" in prompt + assert reasoning_marker in prompt + assert "legacy grounding handoff" in prompt + assert original_sql_marker not in regeneration_prompt + assert "original SQL is intentionally omitted" in regeneration_prompt assert diagnostic_marker in correction_prompt assert "Use the diagnostic text only to understand the failure category" in ( correction_prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index de9c136e89..80db4bc187 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -456,6 +456,10 @@ def table_schema(name): "WREN RETRIEVED SEMANTIC CONTEXT" in schema["table_ddl"] for schema in result["db_schemas"] ) + assert all( + "WREN SQL IDENTIFIER CONTRACT" in schema["table_ddl"] + for schema in result["db_schemas"] + ) assert all( "sql_table_name_use_exactly" in schema["table_ddl"] for schema in result["db_schemas"] @@ -506,6 +510,10 @@ def encode(self, value): assert '"sql_table_name_use_exactly":"modeled_dataset"' in table_ddl assert '"sql_column_name_use_exactly":"stored_attribute"' in table_ddl + assert "WREN SQL IDENTIFIER CONTRACT" in table_ddl + assert "sql_table_name_use_exactly: modeled_dataset" in table_ddl + assert "sql_column_names_use_exactly:\n- stored_attribute" in table_ddl + assert "END WREN SQL IDENTIFIER CONTRACT" in table_ddl assert ( '"semantic_context_not_sql_identifier":"Business-facing attribute label."' in table_ddl From f991632cb746d16a5a5df00f1908b9e908da18e2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 20:33:41 +0530 Subject: [PATCH 0741/1087] Default missing project language in settings --- wren-ui/src/apollo/server/resolvers/projectResolver.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 5253d5d197..a827527901 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -40,6 +40,7 @@ import { TelemetryEvent } from '../telemetry/telemetry'; const logger = getLogger('DataSourceResolver'); logger.level = 'debug'; +const DEFAULT_PROJECT_LANGUAGE = 'EN'; export enum OnboardingStatusEnum { NOT_STARTED = 'NOT_STARTED', @@ -84,7 +85,7 @@ export class ProjectResolver { } as DataSourceProperties, sampleDataset: project.sampleDataset, }, - language: project.language, + language: project.language || DEFAULT_PROJECT_LANGUAGE, }; } From b591050b30a6859276bdc3cd1dfb2a6131fd848d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 20:56:43 +0530 Subject: [PATCH 0742/1087] Ground SQL generation on retrieved schema --- .../followup_sql_generation_reasoning.py | 2 +- .../generation/question_recommendation.py | 53 +++++-------------- .../src/pipelines/generation/utils/sql.py | 14 +++-- .../retrieval/db_schema_retrieval.py | 6 ++- .../pipelines/generation/test_sql_utils.py | 3 ++ .../retrieval/test_db_schema_retrieval.py | 10 ++++ 6 files changed, 41 insertions(+), 47 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 4e285a663b..f9fb31847d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -56,7 +56,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Do not include SQL, SQL-like expressions, table names, column names, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. +Return only the reasoning plan described by the system instructions. Ground the plan with exact `table: ` and `column: .` references from DATABASE SCHEMA when they are relevant. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index a6e7c17b02..83dfb5b938 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -18,7 +18,16 @@ system_prompt = """ -You are an expert in data analysis and SQL query generation. Given a data model specification, optionally a user's question, and a list of categories, your task is to generate insightful, specific questions that can be answered using the provided data model. Each question should be accompanied by a brief explanation of its relevance or importance. +You are an expert in data analysis and SQL query generation. Given a data model specification, optionally a user's question, and a list of categories, your task is to generate insightful, specific questions that can be answered using the provided data model. + +### Grounding Rules + +- The DATABASE SCHEMA is the only source for answerable business concepts. +- Generate questions only from tables, views, metrics, columns, measures, dimensions, calculated fields, and relationships that are present in DATABASE SCHEMA. +- Use aliases, descriptions, and comments only to understand meaning. Do not introduce nouns, measures, dimensions, periods, or entities that are not supported by the schema. +- If the same business concept appears in multiple modeled datasets, generate questions that can use each relevant modeled dataset, provided the schema exposes the needed fields. +- If a question would require fields from multiple datasets, generate it only when DATABASE SCHEMA provides either a relationship path, a view, a metric, or compatible fields that can be combined as separate rows. +- Do not use generic analytics examples, common business templates, or prior wording as a source of answerable concepts unless the concept is represented in DATABASE SCHEMA. ### JSON Output Structure @@ -58,6 +67,7 @@ - Generate questions that are closely related to the user's previous question, ensuring that the new questions build upon or provide deeper insights into the original query. - Use **random category selection** to introduce diverse perspectives while maintaining a focus on the context of the previous question. - Apply the analysis techniques above to enhance the relevance and depth of the generated questions. + - Keep only the parts of the previous question that are supported by DATABASE SCHEMA. 4. **If No User Question is Provided:** @@ -68,58 +78,23 @@ - Ensure questions can be answered using the data model. - Mix simple and complex questions. - Avoid open-ended questions - each should have a definite answer. - - Incorporate time-based analysis where relevant. - - Combine multiple analysis techniques when appropriate for deeper insights. + - Incorporate time-based analysis only when DATABASE SCHEMA exposes relevant time fields. + - Combine multiple analysis techniques when DATABASE SCHEMA supports the required fields and relationships. ### Categories of Questions 1. **Descriptive Questions** Summarize historical data. - - Example: _"What was the total sales volume for each product last quarter?"_ - 2. **Segmentation Questions** Identify meaningful data segments. - - Example: _"Which customer segments contributed most to revenue growth?"_ - 3. **Comparative Questions** Compare data across segments or periods. - - Example: _"How did Product A perform compared to Product B last year?"_ - 4. **Data Quality/Accuracy Questions** Assess data reliability and completeness. - - Example: _"Are there inconsistencies in the sales records for Q1?"_ - ---- - -### Example JSON Output - -```json -{ - "questions": [ - { - "question": "What was the total revenue generated by each region in the last year?", - "category": "Descriptive Questions" - }, - { - "question": "How do customer preferences differ between age groups?", - "category": "Segmentation Questions" - }, - { - "question": "How does the conversion rate vary across different lead sources?", - "category": "Comparative Questions" - }, - { - "question": "What percentage of contacts have incomplete or missing key properties (e.g., email, lifecycle stage, or deal association)", - "category": "Data Quality/Accuracy Questions" - } - ] -} -``` - --- ### Additional Instructions for Randomization @@ -156,7 +131,7 @@ {% endfor %} {% endif %} -Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based on the provided data model. Both the questions and category names should be translated into {{language}}{% if user_question %} and be related to the user's question{% endif %}. The output format should maintain the structure but with localized text. +Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based only on the provided data model. Both the questions and category names should be translated into {{language}}{% if user_question %} and be related to the user's question{% endif %}. The output format should maintain the structure but with localized text. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a38f481751..42bf897b93 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -181,7 +181,9 @@ async def _classify_generation_result( - If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. -- When using multiple tables, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- When the same requested result can be answered from multiple schema objects with compatible columns or metrics, include all relevant schema objects by combining separate result rows with UNION ALL instead of choosing only one object. +- Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. - If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. @@ -192,6 +194,7 @@ async def _classify_generation_result( - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part or answer with the closest valid SQL over grounded fields only. - If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. +- If a requested noun, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. """ @@ -410,10 +413,11 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. 7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -8. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. -9. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -10. If an identifier, literal value, placeholder, or function appears only in SQL samples, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. Reasoning-plan identifiers are usable only when they exactly match DATABASE SCHEMA. -11. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +8. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. +10. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +11. If an identifier, literal value, placeholder, or function appears only in SQL samples, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. Reasoning-plan identifiers are usable only when they exactly match DATABASE SCHEMA. +12. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 3c4e894831..5909ea74b9 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -51,8 +51,10 @@ 11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. 12. Follow only the relationships shown in the provided schema when selecting columns across datasets. 13. Do not stop at a single top candidate when the question needs multiple related datasets. -14. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. -15. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. +14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. +15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. +16. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. +17. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 4e020471a3..1416b3cf1a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -79,6 +79,9 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "Physical datasource names, source database names" in rules assert "Do not replace an invalid identifier with a similar-looking physical" in rules assert "source/lineage names from metadata may guide meaning" in rules + assert "combining separate result rows with UNION ALL" in rules + assert "independently valid from DATABASE SCHEMA" in rules + assert "do not translate it into a generic object name" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 80db4bc187..d7e151baf4 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -10,6 +10,7 @@ embedding, prompt as build_column_selection_prompt, table_retrieval, + table_columns_selection_system_prompt, table_columns_selection_user_prompt_template, ) @@ -68,6 +69,15 @@ def test_column_pruning_prompt_uses_current_query_without_history_text(): assert "previous request" not in result["prompt"] +def test_table_selection_prompt_keeps_multiple_relevant_datasets(): + assert "same business concept is represented by multiple modeled datasets" in ( + table_columns_selection_system_prompt + ) + assert "compatible fields for the same requested result shape" in ( + table_columns_selection_system_prompt + ) + + @pytest.mark.asyncio async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: From 23fa4101612287938ddb0a2d8867fe6defaa94fe Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 29 Jul 2026 21:31:33 +0530 Subject: [PATCH 0743/1087] Ground SQL prompts on schema identifiers only --- .../generation/followup_sql_generation.py | 2 +- .../followup_sql_generation_reasoning.py | 2 +- .../generation/question_recommendation.py | 4 ++-- .../pipelines/generation/sql_correction.py | 4 ++-- .../pipelines/generation/sql_generation.py | 2 +- .../generation/sql_generation_reasoning.py | 2 +- .../pipelines/generation/sql_regeneration.py | 4 ++-- .../src/pipelines/generation/utils/sql.py | 16 +++++++-------- .../pipelines/generation/test_sql_utils.py | 20 +++++++++++++------ 9 files changed, 32 insertions(+), 24 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 ff7cf38092..9294310fd2 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -79,7 +79,7 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. ### REASONING PLAN ### -Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. {{ sql_generation_reasoning }} Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index f9fb31847d..cc4c0cd26f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -56,7 +56,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Ground the plan with exact `table: ` and `column: .` references from DATABASE SCHEMA when they are relevant. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. +Return only the reasoning plan described by the system instructions. When relevant, ground the plan by using the literal prefix `table:` followed by an exact declared table name from DATABASE SCHEMA, or the literal prefix `column:` followed by an exact declared table name, a dot, and an exact declared column name. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, template markers, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index 83dfb5b938..cccb1fc2e8 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -37,8 +37,8 @@ { "questions": [ { - "question": "", - "category": "" + "question": "schema-grounded question text", + "category": "question category" }, ... ] diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index abb690d0a0..7a0d4dfdec 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -53,7 +53,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) The final answer must be in JSON format: {{ - "sql": + "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA" }} """ @@ -87,7 +87,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. {{ sql_generation_reasoning }} {% endif %} ### FAILED SQL ### diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index a1ed60a4d6..82b1e2d889 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -74,7 +74,7 @@ {% if sql_generation_reasoning %} ### REASONING PLAN ### -Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. {{ sql_generation_reasoning }} {% endif %} 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 5311536d03..a26611ffd9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -48,7 +48,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. Ground the plan with exact `table: ` and `column: .` references from DATABASE SCHEMA when they are relevant. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, functions, or identifier-like labels. +Return only the reasoning plan described by the system instructions. When relevant, ground the plan by using the literal prefix `table:` followed by an exact declared table name from DATABASE SCHEMA, or the literal prefix `column:` followed by an exact declared table name, a dot, and an exact declared column name. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, template markers, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 1a96d18c84..b29b4344d4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -47,7 +47,7 @@ def get_sql_regeneration_system_prompt( The final answer must be a ANSI SQL query in JSON format: {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA" }} """ @@ -98,7 +98,7 @@ def get_sql_regeneration_system_prompt( Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### REASONING PLAN ### -Use this plan as the legacy grounding handoff only when each referenced table or column appears exactly in DATABASE SCHEMA. +Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. {{ sql_generation_reasoning }} ### ORIGINAL SQL QUERY ### The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 42bf897b93..7e1321278f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -177,7 +177,7 @@ async def _classify_generation_result( - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. -- Never generate placeholder identifiers or placeholder table names. If the retrieved metadata does not contain an executable object or column for a requested concept, use the closest executable object and column whose semantic metadata supports the intent, or omit that unsupported concept. +- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, use the closest executable object and column whose semantic metadata supports the intent, or omit that unsupported concept. - If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. @@ -189,7 +189,7 @@ async def _classify_generation_result( - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. - Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. -- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, or unsupported functions from them. +- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. - If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part or answer with the closest valid SQL over grounded fields only. @@ -329,8 +329,8 @@ async def _classify_generation_result( 9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. Mention table names only in this exact format: `table: `, and only when `` is declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. -13. Mention column names only in this exact format: `column: .`, and only when both `` and `` are declared together in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +12. Mention table names only by writing the literal prefix `table:` followed by an exact table name declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +13. Mention column names only by writing the literal prefix `column:` followed by an exact declared table name, a dot, and an exact column name declared for that table in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. 14. Do not mention aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels from comments, SQL samples, failed SQL, or user wording as executable identifiers. 15. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. 16. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. @@ -339,7 +339,7 @@ async def _classify_generation_result( 19. If multiple schema objects are required, identify the exact declared relationship path from DATABASE SCHEMA. If no relationship path is declared, say that the retrieved metadata does not provide a join path. 20. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan unless they also appear exactly in DATABASE SCHEMA. 21. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -22. The reasoning plan is a grounding handoff for SQL generation, not SQL. Do not include executable SQL. +22. The reasoning plan is semantic context for intent only, not a source of executable identifiers. SQL generation must re-read DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT before using any identifier. 23. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### @@ -409,14 +409,14 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. 2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. 3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. -4. YOU MUST use the reasoning plan as the legacy grounding handoff only when it is consistent with DATABASE SCHEMA and SQL Rules. Table and column references in the reasoning plan are executable only when they also appear exactly in DATABASE SCHEMA. Choose every executable identifier only from DATABASE SCHEMA and every function only from SQL FUNCTIONS. +4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. 7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 8. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. 9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. 10. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -11. If an identifier, literal value, placeholder, or function appears only in SQL samples, failed SQL, descriptions, lineage, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. Reasoning-plan identifiers are usable only when they exactly match DATABASE SCHEMA. +11. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. 12. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} @@ -425,7 +425,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) The final answer must be a Wren SQL query in JSON format: {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA" }} """ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 1416b3cf1a..13feb0228a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -119,7 +119,8 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "perform a silent grounding check" in prompt assert "closest grounded expression" in prompt assert "DATABASE SCHEMA is the only source of executable identifiers" in prompt - assert "reasoning plan as the legacy grounding handoff" in prompt + assert "reasoning plan as semantic context for intent only" in prompt + assert "Do not copy identifiers, functions, literal values" in prompt assert "include those objects only when DATABASE SCHEMA shows" in prompt assert "Use the exact supported syntax shown there" in prompt assert "WREN SQL IDENTIFIER CONTRACT" in prompt @@ -129,6 +130,7 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "Use Wren SQL identifier quoting with double quotes only" in prompt assert "source database/schema/table names" in prompt assert "appears only in SQL samples, failed SQL" in prompt + assert "" not in prompt def test_json_field_instructions_do_not_include_placeholder_identifiers(): @@ -169,16 +171,18 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "do not try a similar replacement from source metadata" in prompt -def test_sql_reasoning_prompt_uses_legacy_schema_grounding_handoff(): +def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): prompt = sql_generation_reasoning_system_prompt assert "Do not write SQL, possible SQL, sample SQL, assumed SQL" in prompt assert "SQL clauses, SQL functions, code blocks, or executable expressions" in prompt - assert "table: " in prompt - assert "column: ." in prompt - assert "The reasoning plan is a grounding handoff" in prompt + assert "literal prefix `table:`" in prompt + assert "literal prefix `column:`" in prompt + assert "reasoning plan is semantic context for intent only" in prompt assert "declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT" in prompt assert "source names, physical names, lineage names" in prompt + assert "" not in prompt + assert "" not in prompt def test_user_prompt_templates_keep_source_metadata_non_executable(): @@ -262,7 +266,11 @@ def test_executable_prompt_templates_include_grounded_reasoning_but_omit_sql_con regeneration_prompt, ): assert reasoning_marker in prompt - assert "legacy grounding handoff" in prompt + assert "semantic context for the user's intent only" in prompt + assert "Do not copy identifiers, literal values, functions" in prompt + assert "copy only exact declared identifiers" in prompt + assert "" not in prompt + assert "" not in prompt assert original_sql_marker not in regeneration_prompt assert "original SQL is intentionally omitted" in regeneration_prompt From 3076621a25ee2cae4d156ade70a5f9fbc9806683 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 22:01:06 +0530 Subject: [PATCH 0744/1087] Validate generated SQL with dry plan by default --- .../src/web/v1/routers/sql_corrections.py | 4 +- wren-ai-service/src/web/v1/services/ask.py | 4 +- .../src/web/v1/services/sql_corrections.py | 4 +- .../pytest/services/test_dry_plan_defaults.py | 40 +++++++++++++++++++ 4 files changed, 46 insertions(+), 6 deletions(-) create mode 100644 wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py diff --git a/wren-ai-service/src/web/v1/routers/sql_corrections.py b/wren-ai-service/src/web/v1/routers/sql_corrections.py index b74be58bcf..60dbd9246b 100644 --- a/wren-ai-service/src/web/v1/routers/sql_corrections.py +++ b/wren-ai-service/src/web/v1/routers/sql_corrections.py @@ -20,8 +20,8 @@ class PostRequest(BaseRequest): sql: str error: str retrieved_tables: Optional[List[str]] = None - use_dry_plan: bool = False - allow_dry_plan_fallback: bool = True + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False class PostResponse(BaseModel): diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 6ddd3b14e2..71a7639a9c 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -27,8 +27,8 @@ class AskRequest(BaseRequest): histories: Optional[list[AskHistory]] = Field(default_factory=list) ignore_sql_generation_reasoning: bool = False enable_column_pruning: bool = False - use_dry_plan: bool = False - allow_dry_plan_fallback: bool = True + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False custom_instruction: Optional[str] = None diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 86d0f55301..4336186b6f 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -61,8 +61,8 @@ class CorrectionRequest(BaseRequest): sql: str error: str retrieved_tables: Optional[List[str]] = None - use_dry_plan: bool = False - allow_dry_plan_fallback: bool = True + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False @observe(name="SQL Correction") @trace_metadata diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py new file mode 100644 index 0000000000..14460df5b7 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py @@ -0,0 +1,40 @@ +from src.web.v1.routers.sql_corrections import PostRequest +from src.web.v1.services.ask import AskRequest +from src.web.v1.services.sql_corrections import SqlCorrectionService + + +def test_ask_request_defaults_to_planner_validation_without_fallback(): + request = AskRequest(query="How many records are available?", id="deploy-id") + + assert request.use_dry_plan is True + assert request.allow_dry_plan_fallback is False + + +def test_ask_request_allows_explicit_planner_override(): + request = AskRequest( + query="How many records are available?", + id="deploy-id", + use_dry_plan=False, + allow_dry_plan_fallback=True, + ) + + assert request.use_dry_plan is False + assert request.allow_dry_plan_fallback is True + + +def test_sql_correction_router_defaults_to_planner_validation_without_fallback(): + request = PostRequest(sql="SELECT 1", error="dry run failed") + + assert request.use_dry_plan is True + assert request.allow_dry_plan_fallback is False + + +def test_sql_correction_service_defaults_to_planner_validation_without_fallback(): + request = SqlCorrectionService.CorrectionRequest( + event_id="event-id", + sql="SELECT 1", + error="dry run failed", + ) + + assert request.use_dry_plan is True + assert request.allow_dry_plan_fallback is False From 350e20d046f30b90441bb72312706cbb2d7f2aeb Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 22:33:00 +0530 Subject: [PATCH 0745/1087] Route ask validation through dry plan --- .../src/pipelines/generation/utils/sql.py | 1 + wren-ai-service/src/providers/engine/wren.py | 68 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 28 ++++++++ .../providers/test_wren_engine_provider.py | 61 +++++++++++++++++ .../src/apollo/server/adaptors/ibisAdaptor.ts | 9 ++- wren-ui/src/apollo/server/models/model.ts | 6 ++ wren-ui/src/apollo/server/resolvers.ts | 1 + .../apollo/server/resolvers/modelResolver.ts | 30 ++++++++ wren-ui/src/apollo/server/schema.ts | 7 ++ 9 files changed, 210 insertions(+), 1 deletion(-) create mode 100644 wren-ai-service/tests/pytest/providers/test_wren_engine_provider.py diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 7e1321278f..32343f8d4d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -87,6 +87,7 @@ async def _classify_generation_result( session, generation_result, data_source, + project_id=project_id, allow_fallback=allow_dry_plan_fallback, ) diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 3a92853e04..8a4eab6479 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -138,6 +138,44 @@ async def execute_sql( {"error_message": f"Request timed out: {timeout} seconds"}, ) + async def dry_plan( + self, + session: aiohttp.ClientSession, + sql: str, + data_source: str, + project_id: str | None = None, + timeout: float = settings.engine_timeout, + allow_fallback: bool = True, + **kwargs, + ) -> Tuple[bool, str]: + data = { + "sql": sql, + "projectId": project_id, + "allowFallback": allow_fallback, + } + + try: + async with session.post( + f"{self._endpoint}/api/graphql", + json={ + "query": "mutation DryPlanSql($data: DryPlanSQLDataInput) { dryPlanSql(data: $data) }", + "variables": {"data": data}, + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + res_json = await response.json() + if res_data := res_json.get("data"): + return bool(res_data.get("dryPlanSql")), "" + + error_message = res_json.get("errors", [{}])[0].get( + "message", "Unknown error" + ) + logger.error(f"Error dry planning SQL: {error_message}") + return False, error_message + except asyncio.TimeoutError: + logger.error(f"Request timed out: {timeout} seconds") + return False, f"Request timed out: {timeout} seconds" + @provider("wren_ibis") class WrenIbis(Engine): @@ -348,3 +386,33 @@ async def execute_sql( ) except asyncio.TimeoutError: return False, None, f"Request timed out: {timeout} seconds" + + async def dry_plan( + self, + session: aiohttp.ClientSession, + sql: str, + timeout: float = settings.engine_timeout, + **kwargs, + ) -> Tuple[bool, str]: + api_endpoint = f"{self._endpoint}/v1/mdl/dry-plan" + + try: + async with session.get( + api_endpoint, + json={ + "manifest": orjson.loads(base64.b64decode(self._manifest)) + if self._manifest + else {}, + "sql": sql, + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + res = await response.text() + + if response.status == 200: + return True, "" + + return False, res + except asyncio.TimeoutError: + logger.error(f"Request timed out: {timeout} seconds") + return False, f"Request timed out: {timeout} seconds" diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 13feb0228a..078ff4191b 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -1,6 +1,8 @@ +import pytest from haystack.components.builders.prompt_builder import PromptBuilder from src.pipelines.generation.utils.sql import ( + SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, get_json_field_instructions, @@ -27,6 +29,32 @@ class _SqlKnowledge: json_field_instructions = "Use the supplied JSON field definitions only." +class _DryPlanEngine: + def __init__(self): + self.calls = [] + + async def dry_plan(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return True, "" + + +@pytest.mark.asyncio +async def test_sql_postprocessor_passes_project_context_to_dry_plan(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=["SELECT 1"], + project_id="project-id", + use_dry_plan=True, + allow_dry_plan_fallback=False, + data_source="source", + ) + + assert result["valid_generation_result"]["sql"] == "SELECT 1" + assert engine.calls[0][1]["project_id"] == "project-id" + assert engine.calls[0][1]["allow_fallback"] is False + + def test_construct_instructions_uses_instruction_text(): assert construct_instructions( [{"instruction": "First rule."}, {"instruction": "Second rule."}] diff --git a/wren-ai-service/tests/pytest/providers/test_wren_engine_provider.py b/wren-ai-service/tests/pytest/providers/test_wren_engine_provider.py new file mode 100644 index 0000000000..f05e86b6ba --- /dev/null +++ b/wren-ai-service/tests/pytest/providers/test_wren_engine_provider.py @@ -0,0 +1,61 @@ +import aiohttp +import pytest +from aioresponses import CallbackResult, aioresponses + +from src.providers.engine.wren import WrenUI + + +@pytest.mark.asyncio +async def test_wren_ui_dry_plan_calls_graphql_planner(): + endpoint = "http://engine-host" + captured_request = {} + + def callback(_url, **kwargs): + captured_request.update(kwargs) + return CallbackResult(payload={"data": {"dryPlanSql": True}}) + + with aioresponses() as mocked: + mocked.post(f"{endpoint}/api/graphql", callback=callback) + + async with aiohttp.ClientSession() as session: + success, error_message = await WrenUI(endpoint=endpoint).dry_plan( + session, + sql="SELECT 1", + data_source="source", + project_id="project-id", + allow_fallback=False, + ) + + assert success is True + assert error_message == "" + assert captured_request["json"] == { + "query": "mutation DryPlanSql($data: DryPlanSQLDataInput) { dryPlanSql(data: $data) }", + "variables": { + "data": { + "sql": "SELECT 1", + "projectId": "project-id", + "allowFallback": False, + } + }, + } + + +@pytest.mark.asyncio +async def test_wren_ui_dry_plan_returns_graphql_error_message(): + endpoint = "http://engine-host" + + with aioresponses() as mocked: + mocked.post( + f"{endpoint}/api/graphql", + payload={"errors": [{"message": "planner failed"}]}, + ) + + async with aiohttp.ClientSession() as session: + success, error_message = await WrenUI(endpoint=endpoint).dry_plan( + session, + sql="SELECT 1", + data_source="source", + ) + + assert success is False + assert error_message == "planner failed" diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 3b16be0c34..5d421cd975 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -192,6 +192,7 @@ export interface IbisDryPlanOptions { mdl: Manifest; // TODO: replace sql type with WrenSQL sql: string; + allowFallback?: boolean; } export interface IIbisAdaptor { @@ -268,7 +269,7 @@ export class IbisAdaptor implements IIbisAdaptor { this.ibisServerEndpoint = ibisServerEndpoint; } public async getNativeSql(options: IbisDryPlanOptions): Promise { - const { dataSource, mdl, sql } = options; + const { dataSource, mdl, sql, allowFallback } = options; const body = { sql, manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), @@ -277,6 +278,12 @@ export class IbisAdaptor implements IIbisAdaptor { const res = await axios.post( `${this.ibisServerEndpoint}/${this.getIbisApiVersion(IBIS_API_TYPE.DRY_PLAN)}/connector/${dataSourceUrlMap[dataSource]}/dry-plan`, body, + { + headers: { + 'x-wren-fallback_disable': + allowFallback === false ? 'true' : 'false', + }, + }, ); return res.data; } catch (e) { diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index c6403db9b9..42262b3626 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -101,3 +101,9 @@ export interface PreviewSQLData { limit?: number; dryRun?: boolean; } + +export interface DryPlanSQLData { + sql: string; + projectId?: string; + allowFallback?: boolean; +} diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 95c9afd840..11e3487602 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -163,6 +163,7 @@ const resolvers = { // preview previewSql: modelResolver.previewSql, + dryPlanSql: modelResolver.dryPlanSql, // Learning saveLearningRecord: learningResolver.saveLearningRecord, diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 3cd42c4505..ce98d0f4db 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -6,6 +6,7 @@ import { CreateCalculatedFieldData, UpdateCalculatedFieldData, UpdateViewMetadataInput, + DryPlanSQLData, PreviewSQLData, } from '../models'; import { @@ -82,6 +83,7 @@ export class ModelResolver { this.previewModelData = this.previewModelData.bind(this); this.previewViewData = this.previewViewData.bind(this); this.previewSql = this.previewSql.bind(this); + this.dryPlanSql = this.dryPlanSql.bind(this); this.getNativeSql = this.getNativeSql.bind(this); // calculated field @@ -1348,6 +1350,34 @@ export class ModelResolver { }); } + public async dryPlanSql( + _root: any, + args: { data: DryPlanSQLData }, + ctx: IContext, + ): Promise { + const { sql, projectId, allowFallback } = args.data; + const project = projectId + ? await ctx.projectService.getProjectById(parseInt(projectId)) + : await ctx.projectService.getCurrentProject(); + const manifest = await this.getLastDeployedManifest(ctx, project.id); + + if (project.type === DataSourceName.DUCKDB) { + await ctx.wrenEngineAdaptor.getNativeSQL(sql, { + manifest, + modelingOnly: false, + }); + } else { + await ctx.ibisServerAdaptor.getNativeSql({ + dataSource: project.type, + sql, + mdl: manifest, + allowFallback, + }); + } + + return true; + } + public async getNativeSql( _root: any, args: { responseId: number }, diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 37e2549949..90fd686c77 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -927,6 +927,12 @@ export const typeDefs = gql` dryRun: Boolean } + input DryPlanSQLDataInput { + sql: String! + projectId: String + allowFallback: Boolean + } + # Schema Change type SchemaChange { deletedTables: [DetailedChangeTable!] @@ -1425,6 +1431,7 @@ export const typeDefs = gql` # preview previewSql(data: PreviewSQLDataInput): JSON! + dryPlanSql(data: DryPlanSQLDataInput): Boolean! # Learning saveLearningRecord(data: SaveLearningRecordInput!): LearningRecord! From 88d7bfa8a7c824c5f3cb46578b7939efeaa5e408 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 23:12:55 +0530 Subject: [PATCH 0746/1087] Validate generated SQL with dry run after planning --- .../src/pipelines/generation/utils/sql.py | 28 +++++++++++++++ .../web/v1/routers/question_recommendation.py | 4 ++- .../v1/services/question_recommendation.py | 14 ++++++-- .../pipelines/generation/test_sql_utils.py | 36 ++++++++++++++++--- .../pytest/services/test_dry_plan_defaults.py | 24 +++++++++++-- 5 files changed, 96 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 32343f8d4d..faf4d6df64 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -99,12 +99,40 @@ async def _classify_generation_result( else: invalid_generation_result = { "sql": generation_result, + "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") else "DRY_PLAN", "error": error_message, "correlation_id": "", } + return valid_generation_result, invalid_generation_result + + success, _, addition = await self._engine.execute_sql( + generation_result, + session, + project_id=project_id, + limit=1, + dry_run=True, + ) + addition = addition if isinstance(addition, dict) else {} + + if success: + valid_generation_result = { + "sql": generation_result, + "correlation_id": addition.get("correlation_id", ""), + } + else: + error_message = addition.get("error_message", "") + invalid_generation_result = { + "sql": addition.get("error_sql", generation_result), + "original_sql": generation_result, + "type": "TIME_OUT" + if error_message.startswith("Request timed out") + else "DRY_RUN", + "error": error_message, + "correlation_id": addition.get("correlation_id", ""), + } elif use_dry_run: success, _, addition = await self._engine.execute_sql( generation_result, diff --git a/wren-ai-service/src/web/v1/routers/question_recommendation.py b/wren-ai-service/src/web/v1/routers/question_recommendation.py index f029804996..8aa89e3f99 100644 --- a/wren-ai-service/src/web/v1/routers/question_recommendation.py +++ b/wren-ai-service/src/web/v1/routers/question_recommendation.py @@ -22,7 +22,9 @@ class PostRequest(BaseRequest): max_questions: int = 5 max_categories: int = 3 regenerate: bool = False - allow_data_preview: bool = True + allow_data_preview: bool = False + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False class PostResponse(BaseModel): diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 6033237a45..694d044bfa 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -67,7 +67,9 @@ async def _validate_question( max_questions: int, max_categories: int, project_id: Optional[str] = None, - allow_data_preview: bool = True, + allow_data_preview: bool = False, + use_dry_plan: bool = True, + allow_dry_plan_fallback: bool = False, ): async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( @@ -131,6 +133,8 @@ async def _instructions_retrieval() -> list[dict]: has_metric=has_metric, has_json_field=has_json_field, sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, ) @@ -172,7 +176,9 @@ class Request(BaseRequest): max_questions: int = 5 max_categories: int = 3 regenerate: bool = False - allow_data_preview: bool = True + allow_data_preview: bool = False + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False async def _recommend(self, request: dict): resp = await self._pipelines["question_recommendation"].run(**request) @@ -185,6 +191,8 @@ async def _recommend(self, request: dict): request["max_categories"], project_id=request["project_id"], allow_data_preview=request["allow_data_preview"], + use_dry_plan=request["use_dry_plan"], + allow_dry_plan_fallback=request["allow_dry_plan_fallback"], ) for question in questions ] @@ -218,6 +226,8 @@ async def recommend(self, input: Request, **kwargs) -> Event: "project_id": input.project_id, "event_id": input.event_id, "allow_data_preview": input.allow_data_preview, + "use_dry_plan": input.use_dry_plan, + "allow_dry_plan_fallback": input.allow_dry_plan_fallback, } await self._recommend(request) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 078ff4191b..90034eac1f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -31,15 +31,20 @@ class _SqlKnowledge: class _DryPlanEngine: def __init__(self): - self.calls = [] + self.dry_plan_calls = [] + self.execute_sql_calls = [] async def dry_plan(self, *args, **kwargs): - self.calls.append((args, kwargs)) + self.dry_plan_calls.append((args, kwargs)) return True, "" + async def execute_sql(self, *args, **kwargs): + self.execute_sql_calls.append((args, kwargs)) + return True, {}, {"correlation_id": "correlation-id"} + @pytest.mark.asyncio -async def test_sql_postprocessor_passes_project_context_to_dry_plan(): +async def test_sql_postprocessor_validates_with_dry_plan_then_dry_run(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -51,8 +56,29 @@ async def test_sql_postprocessor_passes_project_context_to_dry_plan(): ) assert result["valid_generation_result"]["sql"] == "SELECT 1" - assert engine.calls[0][1]["project_id"] == "project-id" - assert engine.calls[0][1]["allow_fallback"] is False + assert result["valid_generation_result"]["correlation_id"] == "correlation-id" + assert engine.dry_plan_calls[0][1]["project_id"] == "project-id" + assert engine.dry_plan_calls[0][1]["allow_fallback"] is False + assert engine.execute_sql_calls[0][1]["project_id"] == "project-id" + assert engine.execute_sql_calls[0][1]["dry_run"] is True + + +class _FailingDryPlanEngine: + async def dry_plan(self, *args, **kwargs): + return False, "planner failed" + + +@pytest.mark.asyncio +async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): + result = await SQLGenPostProcessor(_FailingDryPlanEngine()).run( + replies=["SELECT 1"], + use_dry_plan=True, + data_source="source", + ) + + assert result["invalid_generation_result"]["sql"] == "SELECT 1" + assert result["invalid_generation_result"]["original_sql"] == "SELECT 1" + assert result["invalid_generation_result"]["type"] == "DRY_PLAN" def test_construct_instructions_uses_instruction_text(): diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py index 14460df5b7..8c2064c466 100644 --- a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py +++ b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py @@ -1,5 +1,9 @@ -from src.web.v1.routers.sql_corrections import PostRequest +from src.web.v1.routers.question_recommendation import ( + PostRequest as QuestionRecommendationPostRequest, +) +from src.web.v1.routers.sql_corrections import PostRequest as SqlCorrectionPostRequest from src.web.v1.services.ask import AskRequest +from src.web.v1.services.question_recommendation import QuestionRecommendation from src.web.v1.services.sql_corrections import SqlCorrectionService @@ -23,7 +27,7 @@ def test_ask_request_allows_explicit_planner_override(): def test_sql_correction_router_defaults_to_planner_validation_without_fallback(): - request = PostRequest(sql="SELECT 1", error="dry run failed") + request = SqlCorrectionPostRequest(sql="SELECT 1", error="dry run failed") assert request.use_dry_plan is True assert request.allow_dry_plan_fallback is False @@ -38,3 +42,19 @@ def test_sql_correction_service_defaults_to_planner_validation_without_fallback( assert request.use_dry_plan is True assert request.allow_dry_plan_fallback is False + + +def test_question_recommendation_router_defaults_to_planner_validation_without_fallback(): + request = QuestionRecommendationPostRequest(mdl='{"models":[]}') + + assert request.allow_data_preview is False + assert request.use_dry_plan is True + assert request.allow_dry_plan_fallback is False + + +def test_question_recommendation_service_defaults_to_planner_validation_without_fallback(): + request = QuestionRecommendation.Request(event_id="event-id", mdl='{"models":[]}') + + assert request.allow_data_preview is False + assert request.use_dry_plan is True + assert request.allow_dry_plan_fallback is False From 4707e743d134df48aea97222751e449f74cc18f9 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Wed, 29 Jul 2026 23:40:25 +0530 Subject: [PATCH 0747/1087] Prevent failed SQL from being accepted --- .../src/pipelines/generation/utils/sql.py | 7 +--- wren-ai-service/src/web/v1/services/ask.py | 24 +----------- .../pipelines/generation/test_sql_utils.py | 22 +++++++++++ .../pytest/services/test_dry_plan_defaults.py | 39 ++++++++++++++++++- 4 files changed, 62 insertions(+), 30 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index faf4d6df64..b92d45c4f6 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -91,12 +91,7 @@ async def _classify_generation_result( allow_fallback=allow_dry_plan_fallback, ) - if dry_plan_result: - valid_generation_result = { - "sql": generation_result, - "correlation_id": "", - } - else: + if not dry_plan_result: invalid_generation_result = { "sql": generation_result, "original_sql": generation_result, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 71a7639a9c..fe3bd0330f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -189,29 +189,7 @@ async def ask( is_followup=True if histories else False, ) - historical_question = await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ) - - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - - if historical_question_result: - api_results = [ - AskResult( - **{ - "sql": result.get("statement"), - "type": "view" if result.get("viewId") else "llm", - "viewId": result.get("viewId"), - } - ) - for result in historical_question_result - ] - sql_generation_reasoning = "" - else: + if not api_results: # Run both pipeline operations concurrently sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 90034eac1f..a13fc02a34 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -68,6 +68,14 @@ async def dry_plan(self, *args, **kwargs): return False, "planner failed" +class _FailingDryRunAfterDryPlanEngine: + async def dry_plan(self, *args, **kwargs): + return True, "" + + async def execute_sql(self, *args, **kwargs): + return False, {}, {"error_message": "dry run failed"} + + @pytest.mark.asyncio async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): result = await SQLGenPostProcessor(_FailingDryPlanEngine()).run( @@ -81,6 +89,20 @@ async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): assert result["invalid_generation_result"]["type"] == "DRY_PLAN" +@pytest.mark.asyncio +async def test_sql_postprocessor_does_not_keep_valid_result_when_dry_run_fails(): + result = await SQLGenPostProcessor(_FailingDryRunAfterDryPlanEngine()).run( + replies=["SELECT 1"], + use_dry_plan=True, + data_source="source", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["sql"] == "SELECT 1" + assert result["invalid_generation_result"]["original_sql"] == "SELECT 1" + assert result["invalid_generation_result"]["type"] == "DRY_RUN" + + def test_construct_instructions_uses_instruction_text(): assert construct_instructions( [{"instruction": "First rule."}, {"instruction": "Second rule."}] diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py index 8c2064c466..daca3770de 100644 --- a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py +++ b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py @@ -1,8 +1,12 @@ +from unittest.mock import AsyncMock + +import pytest + from src.web.v1.routers.question_recommendation import ( PostRequest as QuestionRecommendationPostRequest, ) from src.web.v1.routers.sql_corrections import PostRequest as SqlCorrectionPostRequest -from src.web.v1.services.ask import AskRequest +from src.web.v1.services.ask import AskRequest, AskResultRequest, AskService from src.web.v1.services.question_recommendation import QuestionRecommendation from src.web.v1.services.sql_corrections import SqlCorrectionService @@ -58,3 +62,36 @@ def test_question_recommendation_service_defaults_to_planner_validation_without_ assert request.allow_data_preview is False assert request.use_dry_plan is True assert request.allow_dry_plan_fallback is False + + +@pytest.mark.asyncio +async def test_ask_service_does_not_require_historical_sql_shortcut(): + service = AskService( + { + "sql_pairs_retrieval": AsyncMock( + run=AsyncMock(return_value={"formatted_output": {"documents": []}}) + ), + "instructions_retrieval": AsyncMock( + run=AsyncMock(return_value={"formatted_output": {"documents": []}}) + ), + "intent_classification": AsyncMock( + run=AsyncMock( + return_value={ + "post_process": { + "intent": "MISLEADING_QUERY", + "reasoning": "No matching schema context.", + } + } + ) + ), + "misleading_assistance": AsyncMock(run=AsyncMock(return_value={})), + } + ) + request = AskRequest(query="Can this be answered?", id="deploy-id") + request.query_id = "query-id" + + await service.ask(request) + result = service.get_ask_result(AskResultRequest(query_id="query-id")) + + assert result.status == "finished" + assert result.type == "GENERAL" From 03c4650ab6c98085a0e77b3ca2cb9bd781a5dc43 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 00:14:10 +0530 Subject: [PATCH 0748/1087] Ground SQL generation with schema retrieval --- .../pipelines/generation/sql_correction.py | 4 +- .../pipelines/generation/sql_regeneration.py | 4 +- .../src/pipelines/generation/utils/sql.py | 21 +++-- .../retrieval/db_schema_retrieval.py | 65 +++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 7 +- .../pipelines/generation/test_sql_utils.py | 20 +++++ .../retrieval/test_db_schema_retrieval.py | 88 +++++++++++++++++-- 7 files changed, 190 insertions(+), 19 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 7a0d4dfdec..9ccf38d417 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -50,10 +50,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA" + "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index b29b4344d4..27d076f2d5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -44,10 +44,10 @@ def get_sql_regeneration_system_prompt( {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA" + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b92d45c4f6..a7749ae887 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -40,9 +40,9 @@ async def run( # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' if cleaned_generation_result.startswith("{"): - cleaned_generation_result = orjson.loads(cleaned_generation_result)[ + cleaned_generation_result = orjson.loads(cleaned_generation_result).get( "sql" - ] + ) ( valid_generation_result, @@ -70,7 +70,7 @@ async def run( async def _classify_generation_result( self, - generation_result: str, + generation_result: str | None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, @@ -81,6 +81,15 @@ async def _classify_generation_result( invalid_generation_result = {} use_dry_run = not allow_data_preview + if not generation_result: + return valid_generation_result, { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": "No grounded SQL was generated from the current schema.", + "correlation_id": "", + } + async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( @@ -446,16 +455,16 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a Wren SQL query in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA" + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ class SqlGenerationResult(BaseModel): - sql: str + sql: str | None SQL_GENERATION_MODEL_KWARGS = { diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 5909ea74b9..738d139565 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -359,14 +359,22 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, project_id: str, dbschema_retriever: Any + table_retrieval: dict, project_id: str, dbschema_retriever: Any, embedding: dict ) -> list[Document]: table_names = _table_names_from_description_documents( table_retrieval.get("documents", []) ) + documents = [] + if embedding: + documents = await _retrieve_semantic_schema_documents( + embedding, project_id, dbschema_retriever + ) + table_names = _merge_names( + table_names, + _table_names_from_schema_documents(documents), + ) if table_names: - documents = [] retrieved_table_names = set() pending_table_names = table_names @@ -387,6 +395,59 @@ async def dbschema_retrieval( return [] +async def _retrieve_semantic_schema_documents( + embedding: dict, project_id: str, dbschema_retriever: Any +) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + results = await dbschema_retriever.run( + query_embedding=embedding.get("embedding"), + filters=filters, + ) + return results["documents"] + + +def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: + table_names = [] + seen = set() + + for document in documents: + table_name = document.meta.get("name") + if not table_name: + content = ast.literal_eval(document.content) + table_name = content.get("name") + + if table_name and table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + + return table_names + + +def _merge_names(*name_groups: list[str]) -> list[str]: + merged = [] + seen = set() + + for names in name_groups: + for name in names: + if name in seen: + continue + merged.append(name) + seen.add(name) + + return merged + + def _table_names_from_description_documents(documents: list[Document]) -> list[str]: table_names = [] seen = set() diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index fe3bd0330f..49beeffbc3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -489,7 +489,12 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + if failed_dry_run_result["type"] in ( + "TIME_OUT", + "NO_RELEVANT_SQL", + ): + error_message = failed_dry_run_result["error"] + invalid_sql = failed_dry_run_result["sql"] break original_sql = failed_dry_run_result["original_sql"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index a13fc02a34..4a38ddf29f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -2,6 +2,7 @@ from haystack.components.builders.prompt_builder import PromptBuilder from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, @@ -103,6 +104,19 @@ async def test_sql_postprocessor_does_not_keep_valid_result_when_dry_run_fails() assert result["invalid_generation_result"]["type"] == "DRY_RUN" +@pytest.mark.asyncio +async def test_sql_postprocessor_rejects_null_sql_generation_result(): + result = await SQLGenPostProcessor(_DryPlanEngine()).run( + replies=['{"sql": null}'], + use_dry_plan=True, + data_source="source", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert result["invalid_generation_result"]["sql"] == "" + + def test_construct_instructions_uses_instruction_text(): assert construct_instructions( [{"instruction": "First rule."}, {"instruction": "Second rule."}] @@ -169,6 +183,12 @@ def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): assert "Do not query INFORMATION_SCHEMA" in rules +def test_sql_generation_schema_allows_null_when_sql_cannot_be_grounded(): + schema = SQL_GENERATION_MODEL_KWARGS["response_format"]["json_schema"]["schema"] + + assert {"type": "null"} in schema["properties"]["sql"]["anyOf"] + + def test_get_metric_instructions_uses_sql_knowledge_override(): assert get_metric_instructions(_SqlKnowledge()) == _SqlKnowledge.metric_instructions diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index d7e151baf4..ba5fe77b99 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -153,6 +153,7 @@ async def run(self, query_embedding, filters): }, project_id="project-1", dbschema_retriever=retriever, + embedding={}, ) assert [document.meta["name"] for document in documents] == ["orders", "customers"] @@ -288,6 +289,7 @@ async def run(self, query_embedding, filters): }, project_id="project-1", dbschema_retriever=retriever, + embedding={}, ) assert retriever.calls == [[selected_model], [related_model], [downstream_model]] @@ -301,14 +303,61 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_dbschema_retrieval_does_not_load_full_schema_for_unmatched_question(): +async def test_dbschema_retrieval_uses_semantic_schema_hits_when_table_retrieval_misses(): + semantic_model = "semantic_dataset" + class Retriever: def __init__(self): - self.called = False + self.calls = [] async def run(self, query_embedding, filters): - self.called = True - return {"documents": []} + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + + if query_embedding: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "semantic_measure", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": semantic_model}, + ) + ] + } + + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": semantic_model, + "comment": "", + "columns": [], + "properties": {}, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": semantic_model}, + ) + ] + } retriever = Retriever() @@ -316,10 +365,37 @@ async def run(self, query_embedding, filters): table_retrieval={"documents": []}, project_id="project-1", dbschema_retriever=retriever, + embedding={"embedding": [0.25]}, ) - assert documents == [] - assert not retriever.called + assert retriever.calls[0] == { + "query_embedding": [0.25], + "filters": { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + }, + } + assert retriever.calls[1]["query_embedding"] == [] + assert retriever.calls[1]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": semantic_model}, + ], + }, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert [document.meta["name"] for document in documents] == [ + semantic_model, + semantic_model, + ] def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): From 8679d2ec0c5aa79618af68ba60fac4946cfd0590 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 00:34:51 +0530 Subject: [PATCH 0749/1087] Keep SQL generation grounded to schema prompts --- wren-ai-service/src/config.py | 2 +- .../src/pipelines/generation/followup_sql_generation.py | 4 ---- .../src/pipelines/generation/sql_correction.py | 5 ----- .../src/pipelines/generation/sql_generation.py | 6 ------ .../src/pipelines/generation/sql_regeneration.py | 3 --- wren-ai-service/src/web/v1/services/ask.py | 2 +- .../tests/pytest/pipelines/generation/test_sql_utils.py | 8 +++----- .../tests/pytest/services/test_dry_plan_defaults.py | 6 ++++++ 8 files changed, 11 insertions(+), 25 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 04c6c05762..d6d2641224 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -29,7 +29,7 @@ class Settings(BaseSettings): column_indexing_batch_size: int = Field(default=50) table_retrieval_size: int = Field(default=50) table_column_retrieval_size: int = Field(default=100) - enable_column_pruning: bool = Field(default=False) + enable_column_pruning: bool = Field(default=True) 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) 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 9294310fd2..ded5ac3b48 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -78,10 +78,6 @@ User's Follow-up Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. -### REASONING PLAN ### -Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. -{{ sql_generation_reasoning }} - Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 9ccf38d417..c7b8d4c20e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -85,11 +85,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. {% endif %} -{% if sql_generation_reasoning %} -### REASONING PLAN ### -Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. -{{ sql_generation_reasoning }} -{% endif %} ### FAILED SQL ### The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 82b1e2d889..5856205a37 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -72,12 +72,6 @@ User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. -{% if sql_generation_reasoning %} -### REASONING PLAN ### -Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. -{{ sql_generation_reasoning }} -{% endif %} - Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 27d076f2d5..c66875dc0c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -97,9 +97,6 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. Regenerate with executable identifiers from the current DATABASE SCHEMA only. -### REASONING PLAN ### -Use this plan as semantic context for the user's intent only. Do not copy identifiers, literal values, functions, SQL fragments, template markers, or placeholders from it. Before using any table, column, relationship, metric, view, or function mentioned by the plan, re-read DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS, then copy only exact declared identifiers from those sections. -{{ sql_generation_reasoning }} ### ORIGINAL SQL QUERY ### The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 49beeffbc3..ca55061902 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -103,7 +103,7 @@ def __init__( allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = False, + enable_column_pruning: bool = True, max_sql_correction_retries: int = 3, max_histories: int = 5, maxsize: int = 1_000_000, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 4a38ddf29f..a6c6eef677 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -300,7 +300,7 @@ def test_followup_sql_prompt_does_not_expect_previous_sql_context(): ) -def test_executable_prompt_templates_include_grounded_reasoning_but_omit_sql_context(): +def test_executable_prompt_templates_omit_untrusted_reasoning_and_sql_context(): reasoning_marker = "UNTRUSTED_REASONING_CONTEXT_MARKER" diagnostic_marker = "UNTRUSTED_DIAGNOSTIC_CONTEXT_MARKER" original_sql_marker = "UNTRUSTED_ORIGINAL_SQL_MARKER" @@ -361,10 +361,8 @@ def test_executable_prompt_templates_include_grounded_reasoning_but_omit_sql_con correction_prompt, regeneration_prompt, ): - assert reasoning_marker in prompt - assert "semantic context for the user's intent only" in prompt - assert "Do not copy identifiers, literal values, functions" in prompt - assert "copy only exact declared identifiers" in prompt + assert reasoning_marker not in prompt + assert "REASONING PLAN" not in prompt assert "" not in prompt assert "" not in prompt diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py index daca3770de..30d9957c0d 100644 --- a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py +++ b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py @@ -30,6 +30,12 @@ def test_ask_request_allows_explicit_planner_override(): assert request.allow_dry_plan_fallback is True +def test_ask_service_defaults_to_column_pruning(): + service = AskService({}) + + assert service._enable_column_pruning is True + + def test_sql_correction_router_defaults_to_planner_validation_without_fallback(): request = SqlCorrectionPostRequest(sql="SELECT 1", error="dry run failed") From 97386292d102bb9d3cb89800fd844aecd4a169b4 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 01:26:05 +0530 Subject: [PATCH 0750/1087] Preserve schema when column pruning misses --- .../retrieval/db_schema_retrieval.py | 21 ++++++- .../retrieval/test_db_schema_retrieval.py | 59 +++++++++++++++++++ 2 files changed, 77 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 738d139565..021272d73f 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -265,6 +265,15 @@ def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[d ] +def _has_selected_executable_column(content: dict, columns: set[str]) -> bool: + return any( + column["type"] == "COLUMN" + and column["data_type"].lower() != "unknown" + and column["name"] in columns + for column in content["columns"] + ) + + def _build_table_retrieval_context( content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None ) -> tuple[str, bool, bool]: @@ -686,12 +695,18 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: + selected_columns = set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ) + columns = ( + selected_columns + if _has_selected_executable_column(table_schema, selected_columns) + else None + ) ddl, _has_calculated_field, _has_json_field = ( _build_table_retrieval_context( table_schema, - columns=set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ), + columns=columns, tables=tables, ) ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index ba5fe77b99..f20c0b358d 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -500,6 +500,65 @@ def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): assert result["has_metric"] is True +def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_columns(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["semantic_label"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_dimension", + "data_type": "VARCHAR", + "comment": "Semantic dimension label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "stored_measure", + "data_type": "DOUBLE", + "comment": "Semantic measure label.", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + table_ddl = result["retrieval_results"][0]["table_ddl"] + + assert "semantic_label" not in table_ddl + assert "stored_dimension VARCHAR" in table_ddl + assert "stored_measure DOUBLE" in table_ddl + assert "sql_column_names_use_exactly:\n- stored_dimension\n- stored_measure" in ( + table_ddl + ) + + def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): From c769703c685f05db19643f977a8916290e97bfda Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 01:49:35 +0530 Subject: [PATCH 0751/1087] Require pruned columns to match schema --- .../retrieval/db_schema_retrieval.py | 16 ++--- .../retrieval/test_db_schema_retrieval.py | 59 +++++++++++++++++++ 2 files changed, 68 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 021272d73f..64f52c5287 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -265,13 +265,13 @@ def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[d ] -def _has_selected_executable_column(content: dict, columns: set[str]) -> bool: - return any( - column["type"] == "COLUMN" - and column["data_type"].lower() != "unknown" - and column["name"] in columns +def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: + executable_columns = { + column["name"] for column in content["columns"] - ) + if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" + } + return bool(columns) and columns.issubset(executable_columns) def _build_table_retrieval_context( @@ -700,7 +700,9 @@ def construct_retrieval_results( ) columns = ( selected_columns - if _has_selected_executable_column(table_schema, selected_columns) + if _selected_columns_are_executable( + table_schema, selected_columns + ) else None ) ddl, _has_calculated_field, _has_json_field = ( diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index f20c0b358d..2141eb0a02 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -559,6 +559,65 @@ def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_co ) +def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_unknown_columns(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed fields."], + "columns": ["stored_measure", "semantic_label"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_dimension", + "data_type": "VARCHAR", + "comment": "Semantic dimension label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "stored_measure", + "data_type": "DOUBLE", + "comment": "Semantic measure label.", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + table_ddl = result["retrieval_results"][0]["table_ddl"] + + assert "semantic_label" not in table_ddl + assert "stored_dimension VARCHAR" in table_ddl + assert "stored_measure DOUBLE" in table_ddl + assert "sql_column_names_use_exactly:\n- stored_dimension\n- stored_measure" in ( + table_ddl + ) + + def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): From 7a541cd6b65c882a4933fff4e2f98d5cdea3c4a3 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 02:22:11 +0530 Subject: [PATCH 0752/1087] Keep full schema for grounded SQL generation --- .../src/pipelines/generation/utils/sql.py | 14 +++-- .../retrieval/db_schema_retrieval.py | 35 +++++------ .../retrieval/test_db_schema_retrieval.py | 63 +++++++++++++++++++ 3 files changed, 86 insertions(+), 26 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a7749ae887..6de2c8794f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -206,6 +206,7 @@ async def _classify_generation_result( - When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. - When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. - In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. +- In sql_column_semantic_context, each semantic_context_not_sql_identifier describes the meaning of the paired sql_column_name_use_exactly. When answering a requested business concept, use the paired sql_column_name_use_exactly for that concept. - Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. @@ -445,12 +446,13 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. -7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -8. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. -9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. -10. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -11. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. -12. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +7. YOU MUST treat sql_column_semantic_context entries as column meaning mappings: semantic_context_not_sql_identifier explains the paired sql_column_name_use_exactly. Use that paired exact column when it represents the requested business concept. +8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. +11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +13. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 64f52c5287..8116e2bdaf 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -188,6 +188,9 @@ def _format_semantic_context(context: dict) -> str: def _format_identifier_contract(context: dict) -> str: + def _one_line(value: Any) -> str: + return " ".join(str(value or "").split()) + contract = context.get("sql_identifier_contract", {}) table_name = contract.get("sql_table_name_use_exactly") column_names = contract.get("sql_column_names_use_exactly") or [ @@ -210,6 +213,17 @@ def _format_identifier_contract(context: dict) -> str: if column_names: lines.append("sql_column_names_use_exactly:") lines.extend(f"- {column_name}" for column_name in column_names) + if context.get("columns"): + lines.append("sql_column_semantic_context:") + for column in context["columns"]: + lines.append( + f"- sql_column_name_use_exactly: {column['sql_column_name_use_exactly']}" + ) + lines.append(f" data_type: {column.get('data_type', '')}") + lines.append( + " semantic_context_not_sql_identifier: " + f"{_one_line(column.get('semantic_context_not_sql_identifier'))}" + ) if relationship_constraints: lines.append("relationship_constraints_use_exactly:") lines.extend( @@ -265,15 +279,6 @@ def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[d ] -def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: - executable_columns = { - column["name"] - for column in content["columns"] - if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" - } - return bool(columns) and columns.issubset(executable_columns) - - def _build_table_retrieval_context( content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None ) -> tuple[str, bool, bool]: @@ -695,20 +700,10 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - selected_columns = set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ) - columns = ( - selected_columns - if _selected_columns_are_executable( - table_schema, selected_columns - ) - else None - ) ddl, _has_calculated_field, _has_json_field = ( _build_table_retrieval_context( table_schema, - columns=columns, + columns=None, tables=tables, ) ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 2141eb0a02..0a62246614 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -500,6 +500,64 @@ def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): assert result["has_metric"] is True +def test_construct_retrieval_results_keeps_full_table_schema_after_table_selection(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_measure"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_role", + "data_type": "VARCHAR", + "comment": "Semantic role label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "stored_measure", + "data_type": "DOUBLE", + "comment": "Semantic measure label.", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + table_ddl = result["retrieval_results"][0]["table_ddl"] + + assert "stored_role VARCHAR" in table_ddl + assert "stored_measure DOUBLE" in table_ddl + assert "sql_column_names_use_exactly:\n- stored_role\n- stored_measure" in ( + table_ddl + ) + + def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_columns(): result = construct_retrieval_results( check_using_db_schemas_without_pruning={}, @@ -717,6 +775,11 @@ def encode(self, value): assert "WREN SQL IDENTIFIER CONTRACT" in table_ddl assert "sql_table_name_use_exactly: modeled_dataset" in table_ddl assert "sql_column_names_use_exactly:\n- stored_attribute" in table_ddl + assert "sql_column_semantic_context:" in table_ddl + assert "- sql_column_name_use_exactly: stored_attribute" in table_ddl + assert "semantic_context_not_sql_identifier: Business-facing attribute label." in ( + table_ddl + ) assert "END WREN SQL IDENTIFIER CONTRACT" in table_ddl assert ( '"semantic_context_not_sql_identifier":"Business-facing attribute label."' From 593ce5fad7a3e27bfe66dd903752445b0fa00524 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 02:41:51 +0530 Subject: [PATCH 0753/1087] Revert "Keep full schema for grounded SQL generation" This reverts commit 7a541cd6b65c882a4933fff4e2f98d5cdea3c4a3. --- .../src/pipelines/generation/utils/sql.py | 14 ++--- .../retrieval/db_schema_retrieval.py | 35 ++++++----- .../retrieval/test_db_schema_retrieval.py | 63 ------------------- 3 files changed, 26 insertions(+), 86 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 6de2c8794f..a7749ae887 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -206,7 +206,6 @@ async def _classify_generation_result( - When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. - When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. - In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. -- In sql_column_semantic_context, each semantic_context_not_sql_identifier describes the meaning of the paired sql_column_name_use_exactly. When answering a requested business concept, use the paired sql_column_name_use_exactly for that concept. - Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. @@ -446,13 +445,12 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. -7. YOU MUST treat sql_column_semantic_context entries as column meaning mappings: semantic_context_not_sql_identifier explains the paired sql_column_name_use_exactly. Use that paired exact column when it represents the requested business concept. -8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. -10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. -11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. -13. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +8. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. +10. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +11. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +12. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 8116e2bdaf..64f52c5287 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -188,9 +188,6 @@ def _format_semantic_context(context: dict) -> str: def _format_identifier_contract(context: dict) -> str: - def _one_line(value: Any) -> str: - return " ".join(str(value or "").split()) - contract = context.get("sql_identifier_contract", {}) table_name = contract.get("sql_table_name_use_exactly") column_names = contract.get("sql_column_names_use_exactly") or [ @@ -213,17 +210,6 @@ def _one_line(value: Any) -> str: if column_names: lines.append("sql_column_names_use_exactly:") lines.extend(f"- {column_name}" for column_name in column_names) - if context.get("columns"): - lines.append("sql_column_semantic_context:") - for column in context["columns"]: - lines.append( - f"- sql_column_name_use_exactly: {column['sql_column_name_use_exactly']}" - ) - lines.append(f" data_type: {column.get('data_type', '')}") - lines.append( - " semantic_context_not_sql_identifier: " - f"{_one_line(column.get('semantic_context_not_sql_identifier'))}" - ) if relationship_constraints: lines.append("relationship_constraints_use_exactly:") lines.extend( @@ -279,6 +265,15 @@ def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[d ] +def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: + executable_columns = { + column["name"] + for column in content["columns"] + if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" + } + return bool(columns) and columns.issubset(executable_columns) + + def _build_table_retrieval_context( content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None ) -> tuple[str, bool, bool]: @@ -700,10 +695,20 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: + selected_columns = set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ) + columns = ( + selected_columns + if _selected_columns_are_executable( + table_schema, selected_columns + ) + else None + ) ddl, _has_calculated_field, _has_json_field = ( _build_table_retrieval_context( table_schema, - columns=None, + columns=columns, tables=tables, ) ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 0a62246614..2141eb0a02 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -500,64 +500,6 @@ def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): assert result["has_metric"] is True -def test_construct_retrieval_results_keeps_full_table_schema_after_table_selection(): - result = construct_retrieval_results( - check_using_db_schemas_without_pruning={}, - filter_columns_in_tables={ - "replies": [ - """ - { - "results": [ - { - "table_name": "modeled_dataset", - "table_selection_reason": "Selected for the current request.", - "table_contents": { - "chain_of_thought_reasoning": ["Needed field."], - "columns": ["stored_measure"] - } - } - ] - } - """ - ] - }, - construct_db_schemas=[ - { - "type": "TABLE", - "name": "modeled_dataset", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "stored_role", - "data_type": "VARCHAR", - "comment": "Semantic role label.", - "is_primary_key": False, - }, - { - "type": "COLUMN", - "name": "stored_measure", - "data_type": "DOUBLE", - "comment": "Semantic measure label.", - "is_primary_key": False, - }, - ], - "properties": {}, - "primaryKey": "", - } - ], - dbschema_retrieval=[], - ) - - table_ddl = result["retrieval_results"][0]["table_ddl"] - - assert "stored_role VARCHAR" in table_ddl - assert "stored_measure DOUBLE" in table_ddl - assert "sql_column_names_use_exactly:\n- stored_role\n- stored_measure" in ( - table_ddl - ) - - def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_columns(): result = construct_retrieval_results( check_using_db_schemas_without_pruning={}, @@ -775,11 +717,6 @@ def encode(self, value): assert "WREN SQL IDENTIFIER CONTRACT" in table_ddl assert "sql_table_name_use_exactly: modeled_dataset" in table_ddl assert "sql_column_names_use_exactly:\n- stored_attribute" in table_ddl - assert "sql_column_semantic_context:" in table_ddl - assert "- sql_column_name_use_exactly: stored_attribute" in table_ddl - assert "semantic_context_not_sql_identifier: Business-facing attribute label." in ( - table_ddl - ) assert "END WREN SQL IDENTIFIER CONTRACT" in table_ddl assert ( '"semantic_context_not_sql_identifier":"Business-facing attribute label."' From f6f07b4d96b51f93737eb37b6656bbd886a57433 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 03:18:33 +0530 Subject: [PATCH 0754/1087] Update wren-engine submodule --- wren-engine | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-engine b/wren-engine index 06fb43c3db..44d0811961 160000 --- a/wren-engine +++ b/wren-engine @@ -1 +1 @@ -Subproject commit 06fb43c3dbb6486c05b93e722d03c5aca2c0c5f4 +Subproject commit 44d08119612dca9c8ff007fa12b305fd1aad7593 From 9a2e4aff4716975f08d92324c97fe61f063ffb6c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 15:17:08 +0530 Subject: [PATCH 0755/1087] Prevent ungrounded SQL generation --- .../generation/followup_sql_generation.py | 2 +- .../pipelines/generation/sql_correction.py | 2 +- .../pipelines/generation/sql_generation.py | 2 +- .../pipelines/generation/sql_regeneration.py | 2 +- .../src/pipelines/generation/utils/sql.py | 4 +- .../pipelines/generation/test_sql_utils.py | 4 + .../services/tests/queryService.test.ts | 795 +++--------------- 7 files changed, 125 insertions(+), 686 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 ded5ac3b48..5d4edb01b6 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -76,7 +76,7 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index c7b8d4c20e..f28d612aaa 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -83,7 +83,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### QUESTION ### {% if query %} User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. {% endif %} ### FAILED SQL ### The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 5856205a37..2d5f1a2b96 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -70,7 +70,7 @@ ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index c66875dc0c..b010d19c1b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -95,7 +95,7 @@ def get_sql_regeneration_system_prompt( ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### ORIGINAL SQL QUERY ### The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a7749ae887..4290b7847b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -228,6 +228,8 @@ async def _classify_generation_result( - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part or answer with the closest valid SQL over grounded fields only. - If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. - If a requested noun, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. +- If the user's primary requested subject, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query over unrelated schema objects. +- Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. """ @@ -455,7 +457,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS. If no fully grounded SQL can be generated, return null for sql. +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. If the retrieved schema does not ground the requested intent, return null for sql. {{ "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index a6c6eef677..87bf9b09de 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -172,6 +172,8 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "combining separate result rows with UNION ALL" in rules assert "independently valid from DATABASE SCHEMA" in rules assert "do not translate it into a generic object name" in rules + assert "return null for sql instead of producing an approximate query" in rules + assert "A retrieved object is usable only when" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -226,6 +228,7 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "Use Wren SQL identifier quoting with double quotes only" in prompt assert "source database/schema/table names" in prompt assert "appears only in SQL samples, failed SQL" in prompt + assert "retrieved schema does not ground the requested intent" in prompt assert "" not in prompt @@ -291,6 +294,7 @@ def test_user_prompt_templates_keep_source_metadata_non_executable(): assert "source/physical/lineage names" in prompt assert "omit that unsupported part instead of inventing" in prompt assert "exact declared table and column names from DATABASE SCHEMA" in prompt + assert "return null for sql instead of querying an unrelated object" in prompt def test_followup_sql_prompt_does_not_expect_previous_sql_context(): diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 38c5d2fdef..fadbbd2801 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -3,6 +3,11 @@ import { DataSourceName } from '../../types'; import { QueryService } from '../queryService'; describe('QueryService', () => { + const sql = 'selectable statement'; + const dataSource = DataSourceName.POSTGRES; + const project = { type: dataSource, connectionInfo: {} }; + const manifest = {}; + let mockIbisAdaptor; let mockWrenEngineAdaptor; let mockTelemetry; @@ -28,729 +33,157 @@ describe('QueryService', () => { jest.clearAllMocks(); }); - it('should return true and send event when previewing via ibis dry run succeeds', async () => { + it('passes dry-run requests to ibis and records success telemetry', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', + correlationId: 'correlation-id', + processTime: 'process-time', }); - const res = await queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, - dryRun: true, - }); - - expect(res).toEqual({ correlationId: '123' }); - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_DRY_RUN, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, - }, - actionSuccess: true, - }); - }); - - it('should normalize deployed dbo-prefixed table references for non-mssql previews', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview('SELECT * FROM "dbo_search_queries"', { - project: { - type: DataSourceName.POSTGRES, - connectionInfo: {}, - schema: 'public', - }, - manifest: { - schema: 'public', - models: [ - { - name: 'dbo_search_queries', - tableReference: { - catalog: 'wrenai', - schema: 'public', - table: 'dbo_search_queries', - }, - }, - { - name: 'dbo_tickets', - tableReference: { - catalog: 'wrenai', - schema: 'dbo', - table: 'tickets', - }, - }, - ], - }, + const res: any = await queryService.preview(sql, { + project, + manifest, dryRun: true, }); + expect(res).toEqual({ correlationId: 'correlation-id' }); expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT * FROM "dbo_search_queries"', + sql, expect.objectContaining({ - mdl: expect.objectContaining({ - models: [ - expect.objectContaining({ - name: 'dbo_search_queries', - tableReference: { - catalog: 'wrenai', - schema: 'public', - table: 'search_queries', - }, - }), - expect.objectContaining({ - name: 'dbo_tickets', - tableReference: { - catalog: 'wrenai', - schema: 'public', - table: 'tickets', - }, - }), - ], - }), + dataSource, + connectionInfo: project.connectionInfo, + mdl: manifest, }), ); - }); - - it('should preserve deployed dbo-prefixed table references for mssql previews', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview('SELECT * FROM "dbo_search_queries"', { - project: { - type: DataSourceName.MSSQL, - connectionInfo: {}, - schema: 'public', - }, - manifest: { - schema: 'public', - models: [ - { - name: 'dbo_search_queries', - tableReference: { - catalog: null, - schema: 'dbo', - table: 'search_queries', - }, - }, - ], - }, - dryRun: true, - }); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT * FROM "dbo_search_queries"', - expect.objectContaining({ - mdl: expect.objectContaining({ - models: [ - expect.objectContaining({ - tableReference: { - catalog: null, - schema: 'dbo', - table: 'search_queries', - }, - }), - ], - }), - }), - ); - }); - - it('should repair old non-mssql deployments that still use dbo refSql', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview('SELECT * FROM "dbo_search_queries"', { - project: { - type: DataSourceName.POSTGRES, - connectionInfo: {}, - catalog: 'wrenai', - schema: 'public', - }, - manifest: { - catalog: 'wrenai', - schema: 'public', - models: [ - { - name: 'dbo_search_queries', - refSql: 'SELECT * FROM wrenai.public.dbo_search_queries', - }, - ], - }, - dryRun: true, - }); - - const dryRunOptions = mockIbisAdaptor.dryRun.mock.calls[0][1]; - expect(dryRunOptions.mdl.models[0]).toEqual({ - name: 'dbo_search_queries', - tableReference: { - catalog: 'wrenai', - schema: 'public', - table: 'search_queries', - }, - }); - }); - - it('should normalize non-mssql dbo model names before previewing with ibis', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - 'SELECT created_at FROM dbo_search_queries ORDER BY created_at', - { - project: { - type: DataSourceName.POSTGRES, - connectionInfo: {}, - schema: 'public', - }, - manifest: {}, - dryRun: true, - }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT created_at FROM "dbo_search_queries" ORDER BY created_at', - expect.any(Object), - ); - }); - - it('should normalize generated datediff calls for non-mssql previews', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - `SELECT DATEDIFF('day', "dbo_tickets"."created_at", CURRENT_DATE) AS ticket_age FROM "dbo_tickets"`, + expect(mockTelemetry.records).toEqual([ { - project: { - type: DataSourceName.POSTGRES, - connectionInfo: {}, - schema: 'public', + event: TelemetryEvent.IBIS_DRY_RUN, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, }, - manifest: {}, - dryRun: true, + actionSuccess: true, }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT EXTRACT(DAY FROM (CURRENT_DATE - "dbo_tickets"."created_at")) AS ticket_age FROM "dbo_tickets"', - expect.any(Object), - ); + ]); }); - it('should send event when previewing via ibis dry run fails', async () => { - mockIbisAdaptor.dryRun.mockRejectedValue({ - message: 'Error message', + it('records dry-run failure telemetry and rethrows the adaptor error', async () => { + const error = { + message: 'adaptor failure', extensions: { other: { - correlationId: '123', - processTime: '1s', + correlationId: 'correlation-id', + processTime: 'process-time', }, }, - }); - - try { - await queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, - dryRun: true, - }); - } catch (e) { - expect(e.message).toEqual('Error message'); - expect(e.extensions.other.correlationId).toEqual('123'); - expect(e.extensions.other.processTime).toEqual('1s'); - } - - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_DRY_RUN, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, - error: 'Error message', - }, - actionSuccess: false, - service: undefined, - }); - }); - - it('should return data and send event when previewing via ibis query succeeds', async () => { - mockIbisAdaptor.query.mockResolvedValue({ - data: [], - columns: [], - dtypes: [], - correlationId: '123', - processTime: '1s', - }); - - const res = await queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, - limit: 10, - }); - - expect(res.data).toEqual([]); - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_QUERY, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, - }, - actionSuccess: true, - }); - }); - - it('should send event when previewing via ibis query fails', async () => { - mockIbisAdaptor.query.mockRejectedValue({ - message: 'Error message', - extensions: { - other: { - correlationId: '123', - processTime: '1s', - }, - }, - }); - - await expect( - queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, - }), - ).rejects.toMatchObject({ - message: 'Error message', - extensions: { - other: { - correlationId: '123', - processTime: '1s', - }, - }, - }); - - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_QUERY, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, - error: 'Error message', - }, - actionSuccess: false, - service: undefined, - }); - }); - - it('should reject sql that references tables outside the active manifest before ibis dry run', async () => { - await expect( - queryService.preview('SELECT * FROM dbo_failure_patterns', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_tblSales', - tableReference: { table: 'dbo_tblSales' }, - }, - ], - }, - dryRun: true, - }), - ).rejects.toThrow( - 'Generated SQL references table(s) not present in the active datasource metadata: dbo_failure_patterns', - ); - - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); - }); + }; + mockIbisAdaptor.dryRun.mockRejectedValue(error); - it('should reject sql that references columns outside the active manifest before ibis dry run', async () => { await expect( - queryService.preview('SELECT "orders"."OTD_Date" FROM "orders"', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'orders', - tableReference: { table: 'orders' }, - columns: [ - { name: 'order_date', type: 'timestamp', isCalculated: false }, - { name: 'quantity', type: 'integer', isCalculated: false }, - ], - }, - ], - }, + queryService.preview(sql, { + project, + manifest, dryRun: true, }), - ).rejects.toThrow( - 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: orders.OTD_Date', - ); + ).rejects.toMatchObject(error); - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); - }); - - it('should not treat model column expressions as missing table references', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - 'SELECT DATEPART(WEEK, "dbo_DebugEntries"."DateIn") AS "week" FROM "dbo_DebugEntries"', + expect(mockTelemetry.records).toEqual([ { - project: { type: DataSourceName.MSSQL, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_DebugEntries', - tableReference: { table: 'dbo_DebugEntries' }, - columns: [ - { name: 'DateIn', type: 'timestamp', isCalculated: false }, - ], - }, - ], + event: TelemetryEvent.IBIS_DRY_RUN, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, + error: 'adaptor failure', }, - dryRun: true, + actionSuccess: false, + service: undefined, }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalled(); - }); - - it('should reject unqualified projected columns outside a single active manifest table', async () => { - await expect( - queryService.preview('SELECT tools_required FROM "knowledge_articles"', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'knowledge_articles', - tableReference: { table: 'knowledge_articles' }, - columns: [ - { name: 'id', type: 'integer', isCalculated: false }, - { name: 'content', type: 'string', isCalculated: false }, - ], - }, - ], - }, - dryRun: true, - }), - ).rejects.toThrow( - 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: tools_required', - ); - - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); - }); - - it('should reject unqualified filter columns outside a single active manifest table', async () => { - await expect( - queryService.preview( - 'SELECT id FROM "policies" WHERE policy_category_id = 1', - { - project: { type: DataSourceName.MSSQL, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'policies', - tableReference: { table: 'policies' }, - columns: [ - { name: 'id', type: 'integer', isCalculated: false }, - { name: 'policy_name', type: 'string', isCalculated: false }, - ], - }, - ], - }, - dryRun: true, - }, - ), - ).rejects.toThrow( - 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: policy_category_id', - ); - - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); + ]); }); - it('should reject unknown unqualified function argument columns before ibis planning', async () => { - await expect( - queryService.preview('SELECT COUNT(policy_category_id) FROM "policies"', { - project: { type: DataSourceName.MSSQL, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'policies', - tableReference: { table: 'policies' }, - columns: [ - { name: 'id', type: 'integer', isCalculated: false }, - { name: 'policy_name', type: 'string', isCalculated: false }, - ], - }, - ], - }, - dryRun: true, - }), - ).rejects.toThrow( - 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: policy_category_id', - ); - - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); - }); - - it('should reject numeric aggregates on non-numeric manifest columns before ibis planning', async () => { - await expect( - queryService.preview('SELECT AVG("orders"."quantity") FROM "orders"', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'orders', - tableReference: { table: 'orders' }, - columns: [ - { name: 'quantity', type: 'string', isCalculated: false }, - ], - }, - ], - }, - dryRun: true, + it('passes query requests to ibis, transforms column metadata, and records success telemetry', async () => { + mockIbisAdaptor.query.mockResolvedValue({ + data: [['value']], + columns: ['field'], + dtypes: { field: 'object' }, + correlationId: 'correlation-id', + processTime: 'process-time', + }); + + const res = await queryService.preview(sql, { + project, + manifest, + limit: 1, + }); + + expect(res).toEqual({ + columns: [{ name: 'field', type: 'string' }], + data: [['value']], + correlationId: 'correlation-id', + cacheHit: false, + cacheCreatedAt: undefined, + cacheOverrodeAt: undefined, + override: false, + }); + expect(mockIbisAdaptor.query).toHaveBeenCalledWith( + sql, + expect.objectContaining({ + dataSource, + connectionInfo: project.connectionInfo, + mdl: manifest, + limit: 1, }), - ).rejects.toThrow( - 'Generated SQL references column(s) or expressions not valid for the active datasource metadata: AVG(orders.quantity) uses a non-numeric column', - ); - - expect(mockIbisAdaptor.dryRun).not.toHaveBeenCalled(); - }); - - it('should allow numeric aggregates on numeric manifest columns before ibis planning', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview('SELECT AVG("orders"."quantity") FROM "orders"', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'orders', - tableReference: { table: 'orders' }, - columns: [ - { name: 'quantity', type: 'integer', isCalculated: false }, - ], - }, - ], - }, - dryRun: true, - }); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); - }); - - it('should allow active manifest table references before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview('SELECT * FROM wrenai.public.dbo_tblSales', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_tblSales', - tableReference: { table: 'dbo_tblSales' }, - }, - ], - }, - dryRun: true, - }); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); - }); - - it('should rewrite physical table references to active model names before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - 'SELECT "wrenai.public.dbo_failure"."created_at" FROM "wrenai.public.dbo_failure"', - { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_failure_patterns', - tableReference: { - catalog: 'wrenai', - schema: 'public', - table: 'dbo_failure', - }, - columns: [{ name: 'created_at' }], - }, - ], - }, - dryRun: true, - }, ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', - expect.any(Object), - ); - }); - - it('should rewrite multipart physical table references to active model names before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - 'SELECT "wrenai"."public"."dbo_failure"."created_at" FROM "wrenai"."public"."dbo_failure"', + expect(mockTelemetry.records).toEqual([ { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_failure_patterns', - tableReference: { - catalog: 'wrenai', - schema: 'public', - table: 'dbo_failure', - }, - columns: [{ name: 'created_at' }], - }, - ], + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, }, - dryRun: true, + actionSuccess: true, }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', - expect.any(Object), - ); + ]); }); - it('should rewrite generated base pattern table references to active pattern model names before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - 'SELECT "dbo_failure"."created_at" FROM "dbo_failure"', - { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_failure_patterns', - columns: [{ name: 'created_at' }], - }, - ], - }, - dryRun: true, - }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', - expect.any(Object), - ); - }); - - it('should rewrite dotted dbo pattern model references before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview( - 'SELECT "dbo"."failure_patterns"."created_at" FROM "dbo"."failure_patterns"', - { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_failure_patterns', - columns: [{ name: 'created_at' }], - }, - ], + it('records query failure telemetry and rethrows the adaptor error', async () => { + const error = { + message: 'adaptor failure', + extensions: { + other: { + correlationId: 'correlation-id', + processTime: 'process-time', }, - dryRun: true, }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', - expect.any(Object), - ); - }); + }; + mockIbisAdaptor.query.mockRejectedValue(error); - it('should rewrite dotted dbo base pattern references before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); + await expect( + queryService.preview(sql, { + project, + manifest, + }), + ).rejects.toMatchObject(error); - await queryService.preview( - 'SELECT "dbo"."failure"."created_at" FROM "dbo"."failure"', + expect(mockTelemetry.records).toEqual([ { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'dbo_failure_patterns', - columns: [{ name: 'created_at' }], - }, - ], + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, + error: 'adaptor failure', }, - dryRun: true, - }, - ); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( - 'SELECT "dbo_failure_patterns"."created_at" FROM "dbo_failure_patterns"', - expect.any(Object), - ); - }); - - it('should allow source tables referenced by active manifest refSql before ibis dry run', async () => { - mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', - }); - - await queryService.preview('SELECT * FROM dbo_repair_logs', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: { - models: [ - { - name: 'repair_logs', - refSql: 'SELECT created_at, ticket_id FROM dbo_repair_logs', - }, - ], + actionSuccess: false, + service: undefined, }, - dryRun: true, - }); - - expect(mockIbisAdaptor.dryRun).toHaveBeenCalledTimes(1); + ]); }); }); @@ -759,7 +192,7 @@ class MockTelemetry { sendEvent( event: TelemetryEvent, properties: Record = {}, - service: any, + service: any = undefined, actionSuccess: boolean = true, ) { this.records.push({ event, properties, service, actionSuccess }); From 89a40253c8e618846a24322628668bde7db76d39 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 16:35:18 +0530 Subject: [PATCH 0756/1087] Enforce retrieved schema identifiers before SQL validation --- .../generation/followup_sql_generation.py | 4 + .../pipelines/generation/sql_correction.py | 4 + .../pipelines/generation/sql_generation.py | 4 + .../pipelines/generation/sql_regeneration.py | 4 + .../src/pipelines/generation/utils/sql.py | 489 ++++++++++++++++++ .../retrieval/db_schema_retrieval.py | 44 ++ wren-ai-service/src/web/v1/services/ask.py | 14 +- .../src/web/v1/services/ask_feedback.py | 12 +- .../v1/services/question_recommendation.py | 24 +- .../src/web/v1/services/sql_corrections.py | 6 + .../pipelines/generation/test_sql_utils.py | 94 ++++ .../retrieval/test_db_schema_retrieval.py | 39 ++ 12 files changed, 733 insertions(+), 5 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 5d4edb01b6..19dfd24c3a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,6 +147,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -154,6 +155,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + identifier_contracts=identifier_contracts, ) @@ -205,6 +207,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -231,6 +234,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f28d612aaa..2b90870194 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -148,6 +148,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -155,6 +156,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + identifier_contracts=identifier_contracts, ) @@ -202,6 +204,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -224,6 +227,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **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 2d5f1a2b96..97ffd0c3b3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -138,6 +138,7 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -146,6 +147,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + identifier_contracts=identifier_contracts, ) @@ -197,6 +199,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -223,6 +226,7 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index b010d19c1b..59b7f7f7a9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -164,10 +164,12 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, + identifier_contracts=identifier_contracts, ) @@ -212,6 +214,7 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -230,6 +233,7 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4290b7847b..6b88fb5cef 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3,9 +3,19 @@ import aiohttp import orjson +import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel +from sqlparse.sql import ( + Function, + Identifier, + IdentifierList, + Parenthesis, + Statement, + TokenList, +) +from sqlparse.tokens import DML, Keyword, Name, Punctuation, Whitespace, Wildcard from src.core.engine import ( Engine, @@ -17,6 +27,470 @@ logger = logging.getLogger("wren-ai-service") +class SQLGroundingValidator: + def __init__(self, identifier_contracts: list[dict] | None = None): + self._tables = { + contract["table_name"]: ( + set(contract["columns"]) + if contract.get("columns") is not None + else None + ) + for contract in identifier_contracts or [] + if contract and contract.get("table_name") + } + + def validate(self, sql: str) -> tuple[bool, str]: + if not self._tables: + return True, "" + + try: + statements = sqlparse.parse(sql) + except Exception as exc: + return False, f"Unable to parse generated SQL for schema grounding: {exc}" + + if not statements: + return False, "No SQL statement was generated." + + for statement in statements: + context = _ValidationContext(schema_tables=self._tables) + valid, error = context.validate_statement(statement) + if not valid: + return False, error + + return True, "" + + +class _ValidationContext: + def __init__( + self, + schema_tables: dict[str, set[str] | None], + ctes: dict[str, set[str] | None] | None = None, + ): + self.schema_tables = schema_tables + self.ctes = ctes or {} + + def validate_statement(self, statement: Statement | TokenList) -> tuple[bool, str]: + ctes = dict(self.ctes) + valid, error = self._collect_ctes(statement, ctes) + if not valid: + return False, error + + scope: dict[str, set[str] | None] = {} + source_tokens: set[int] = set() + valid, error = self._collect_sources(statement, ctes, scope, source_tokens) + if not valid: + return False, error + + output_columns = self._infer_select_output_columns(statement, scope) + return self._validate_expressions( + statement, + scope=scope, + source_tokens=source_tokens, + output_columns=output_columns, + ) + + def _collect_ctes( + self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] + ) -> tuple[bool, str]: + tokens = _meaningful_tokens(statement) + if not tokens or not tokens[0].match(Keyword.CTE, "WITH"): + return True, "" + + for token in tokens[1:]: + if token.ttype is DML and token.normalized == "SELECT": + break + + identifiers = ( + list(token.get_identifiers()) + if isinstance(token, IdentifierList) + else [token] + if isinstance(token, Identifier) + else [] + ) + for identifier in identifiers: + cte_name = _identifier_name(identifier) + if not cte_name: + continue + subquery = _subquery_from_identifier(identifier) + if subquery is None: + continue + child_context = _ValidationContext(self.schema_tables, ctes=ctes) + valid, error = child_context.validate_statement(subquery) + if not valid: + return False, error + ctes[cte_name] = child_context._infer_select_output_columns( + subquery, child_context._collect_scope_for_inference(subquery, ctes) + ) + + return True, "" + + def _collect_scope_for_inference( + self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] + ) -> dict[str, set[str] | None]: + scope: dict[str, set[str] | None] = {} + self._collect_sources(statement, ctes, scope, set()) + return scope + + def _collect_sources( + self, + statement: Statement | TokenList, + ctes: dict[str, set[str] | None], + scope: dict[str, set[str] | None], + source_tokens: set[int], + ) -> tuple[bool, str]: + expect_source = False + for token in _meaningful_tokens(statement): + if token.ttype is DML and token.normalized == "SELECT": + expect_source = False + continue + + if _starts_new_clause(token): + expect_source = False + + if token.match( + Keyword, + ( + "FROM", + "JOIN", + "INNER JOIN", + "LEFT JOIN", + "LEFT OUTER JOIN", + "RIGHT JOIN", + "RIGHT OUTER JOIN", + "FULL JOIN", + "FULL OUTER JOIN", + "CROSS JOIN", + ), + ): + expect_source = True + continue + + if not expect_source: + continue + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + valid, error = self._add_source(identifier, ctes, scope) + if not valid: + return False, error + source_tokens.add(id(identifier)) + continue + + if isinstance(token, Identifier): + valid, error = self._add_source(token, ctes, scope) + if not valid: + return False, error + source_tokens.add(id(token)) + continue + + if isinstance(token, Parenthesis): + subquery = _statement_from_parenthesis(token) + if subquery is not None: + child_context = _ValidationContext(self.schema_tables, ctes=ctes) + valid, error = child_context.validate_statement(subquery) + if not valid: + return False, error + continue + + return True, "" + + def _add_source( + self, + identifier: Identifier, + ctes: dict[str, set[str] | None], + scope: dict[str, set[str] | None], + ) -> tuple[bool, str]: + alias = identifier.get_alias() + subquery = _subquery_from_identifier(identifier) + if subquery is not None: + child_context = _ValidationContext(self.schema_tables, ctes=ctes) + valid, error = child_context.validate_statement(subquery) + if not valid: + return False, error + if alias: + scope[alias] = child_context._infer_select_output_columns( + subquery, child_context._collect_scope_for_inference(subquery, ctes) + ) + return True, "" + + table_name = identifier.get_real_name() + if not table_name: + return True, "" + + if table_name in ctes: + columns = ctes[table_name] + elif table_name in self.schema_tables: + columns = self.schema_tables[table_name] + else: + return ( + False, + "Generated SQL references a table that is not declared in the retrieved schema.", + ) + + scope[alias or table_name] = columns + scope[table_name] = columns + return True, "" + + def _infer_select_output_columns( + self, statement: Statement | TokenList, scope: dict[str, set[str] | None] + ) -> set[str] | None: + columns: set[str] = set() + in_select = False + + for token in _meaningful_tokens(statement): + if token.ttype is DML and token.normalized == "SELECT": + in_select = True + continue + + if in_select and token.match(Keyword, "FROM"): + return columns + + if not in_select: + continue + + if token.ttype is Wildcard: + return None + + identifiers = ( + list(token.get_identifiers()) + if isinstance(token, IdentifierList) + else [token] + if isinstance(token, Identifier) + else [] + ) + for identifier in identifiers: + if any(child.ttype is Wildcard for child in identifier.flatten()): + return None + name = identifier.get_alias() or identifier.get_real_name() + if name: + columns.add(name) + + return columns + + def _validate_expressions( + self, + token: TokenList, + scope: dict[str, set[str] | None], + source_tokens: set[int], + output_columns: set[str] | None, + in_order_by: bool = False, + ) -> tuple[bool, str]: + children = getattr(token, "tokens", []) + next_in_order_by = in_order_by + + for child in children: + if child.is_whitespace or child.ttype in Whitespace: + continue + + if id(child) in source_tokens: + continue + + if child.match(Keyword, "ORDER BY"): + next_in_order_by = True + continue + if _starts_new_clause(child) and not child.match(Keyword, "ORDER BY"): + next_in_order_by = False + + if isinstance(child, IdentifierList): + valid, error = self._validate_expressions( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if isinstance(child, Function): + valid, error = self._validate_function( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if isinstance(child, Identifier): + if any(isinstance(grandchild, Function) for grandchild in child.tokens): + valid, error = self._validate_identifier_children( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + valid, error = self._validate_identifier( + child, scope, output_columns, next_in_order_by + ) + if not valid: + return False, error + valid, error = self._validate_identifier_children( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if isinstance(child, Parenthesis): + subquery = _statement_from_parenthesis(child) + if subquery is not None: + valid, error = _ValidationContext( + self.schema_tables, self.ctes + ).validate_statement(subquery) + if not valid: + return False, error + continue + + if isinstance(child, TokenList): + valid, error = self._validate_expressions( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if child.ttype is Name: + valid, error = self._validate_unqualified_column( + str(child), scope, output_columns, next_in_order_by + ) + if not valid: + return False, error + + return True, "" + + def _validate_function( + self, + function: Function, + scope: dict[str, set[str] | None], + source_tokens: set[int], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + for child in function.tokens: + if isinstance(child, Parenthesis): + return self._validate_expressions( + child, scope, source_tokens, output_columns, in_order_by + ) + return True, "" + + def _validate_identifier_children( + self, + identifier: Identifier, + scope: dict[str, set[str] | None], + source_tokens: set[int], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + for child in identifier.tokens: + if isinstance(child, Function): + return self._validate_function( + child, scope, source_tokens, output_columns, in_order_by + ) + return True, "" + + def _validate_identifier( + self, + identifier: Identifier, + scope: dict[str, set[str] | None], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + parent_name = identifier.get_parent_name() + column_name = identifier.get_real_name() + + if not column_name: + return True, "" + + if parent_name: + if parent_name not in scope: + return ( + False, + "Generated SQL references a table alias that is not declared in the query scope.", + ) + columns = scope[parent_name] + if columns is None or column_name in columns: + return True, "" + return ( + False, + "Generated SQL references a column that is not declared in the retrieved schema.", + ) + + return self._validate_unqualified_column( + column_name, scope, output_columns, in_order_by + ) + + def _validate_unqualified_column( + self, + column_name: str, + scope: dict[str, set[str] | None], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + if in_order_by and output_columns is not None and column_name in output_columns: + return True, "" + + if column_name in scope: + return True, "" + + for columns in scope.values(): + if columns is None or column_name in columns: + return True, "" + + return ( + False, + "Generated SQL references a column that is not declared in the retrieved schema.", + ) + + +def _meaningful_tokens(token_list: TokenList) -> list: + return [ + token + for token in token_list.tokens + if not token.is_whitespace + and token.ttype not in Whitespace + and token.ttype not in Punctuation + ] + + +def _identifier_name(identifier: Identifier) -> str | None: + return identifier.get_real_name() or identifier.get_name() + + +def _subquery_from_identifier(identifier: Identifier) -> Statement | None: + for token in identifier.tokens: + if isinstance(token, Parenthesis): + return _statement_from_parenthesis(token) + return None + + +def _statement_from_parenthesis(parenthesis: Parenthesis) -> Statement | None: + inner_sql = str(parenthesis)[1:-1].strip() + if not inner_sql: + return None + + parsed = sqlparse.parse(inner_sql) + if parsed and any( + token.ttype is DML and token.normalized == "SELECT" + for token in _meaningful_tokens(parsed[0]) + ): + return parsed[0] + return None + + +def _starts_new_clause(token) -> bool: + return token.match( + Keyword, + ( + "WHERE", + "GROUP BY", + "HAVING", + "ORDER BY", + "LIMIT", + "UNION", + "UNION ALL", + "EXCEPT", + "INTERSECT", + ), + ) + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -34,6 +508,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + identifier_contracts: list[dict] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -54,6 +529,7 @@ async def run( allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, allow_data_preview=allow_data_preview, + identifier_contracts=identifier_contracts, ) return { @@ -76,6 +552,7 @@ async def _classify_generation_result( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + identifier_contracts: list[dict] | None = None, ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} @@ -90,6 +567,18 @@ async def _classify_generation_result( "correlation_id": "", } + is_grounded, grounding_error = SQLGroundingValidator( + identifier_contracts + ).validate(generation_result) + if not is_grounded: + return valid_generation_result, { + "sql": generation_result, + "original_sql": generation_result, + "type": "SCHEMA_GROUNDING", + "error": grounding_error, + "correlation_id": "", + } + async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 64f52c5287..472cd95fea 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -156,6 +156,17 @@ def _build_metric_ddl(content: dict) -> str: ) +def _build_metric_identifier_contract(content: dict) -> dict: + return { + "table_name": content["name"], + "columns": [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + } + + def _build_view_ddl(content: dict) -> str: context = _format_semantic_context( { @@ -175,6 +186,13 @@ def _build_view_ddl(content: dict) -> str: ) +def _build_view_identifier_contract(content: dict) -> dict: + return { + "table_name": content["name"], + "columns": None, + } + + def _format_semantic_context(context: dict) -> str: return ( "/*\n" @@ -325,6 +343,16 @@ def _build_table_retrieval_context( return f"{context}{ddl}", has_calculated_field, has_json_field +def _build_table_identifier_contract( + content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None +) -> dict: + included_columns = _included_columns(content, columns, tables) + return { + "table_name": content["name"], + "columns": [column["name"] for column in included_columns], + } + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -588,6 +616,9 @@ def check_using_db_schemas_without_pruning( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_contract": _build_table_identifier_contract( + table_schema + ), } ) if _has_calculated_field: @@ -603,6 +634,7 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_contract": _build_metric_identifier_contract(content), } ) has_metric = True @@ -611,6 +643,7 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_contract": _build_view_identifier_contract(content), } ) @@ -721,6 +754,11 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_contract": _build_table_identifier_contract( + table_schema, + columns=columns, + tables=tables, + ), } ) @@ -732,6 +770,9 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_contract": _build_metric_identifier_contract( + content + ), } ) has_metric = True @@ -740,6 +781,9 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_contract": _build_view_identifier_contract( + content + ), } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ca55061902..106a3d4e41 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -332,6 +332,11 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -454,6 +459,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -472,6 +478,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -501,6 +508,7 @@ async def ask( invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] current_sql_correction_retries += 1 + sql_diagnosis_reasoning = None self._ask_results[query_id] = AskResultResponse( status="correcting", @@ -513,7 +521,10 @@ async def ask( is_followup=True if histories else False, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -548,6 +559,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index f1971afa15..024dcfc3dd 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -161,6 +161,11 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] sql_samples = sql_samples_task["formatted_output"].get("documents", []) instructions = instructions_task["formatted_output"].get( "documents", [] @@ -187,6 +192,7 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -216,7 +222,10 @@ async def ask_feedback( trace_id=trace_id, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -251,6 +260,7 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 694d044bfa..1e7383bbf8 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -71,7 +71,7 @@ async def _validate_question( use_dry_plan: bool = True, allow_dry_plan_fallback: bool = False, ): - async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: + async def _document_retrieval() -> tuple[list[str], list[dict], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -79,10 +79,21 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return table_ddls, has_calculated_field, has_metric, has_json_field + return ( + table_ddls, + identifier_contracts, + has_calculated_field, + has_metric, + has_json_field, + ) async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -107,7 +118,13 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - table_ddls, has_calculated_field, has_metric, has_json_field = _document + ( + table_ddls, + identifier_contracts, + has_calculated_field, + has_metric, + has_json_field, + ) = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -137,6 +154,7 @@ async def _instructions_retrieval() -> list[dict]: allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) post_process = generated_sql["post_process"] diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 4336186b6f..bec96f50fc 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -111,6 +111,11 @@ async def correct( .get("retrieval_results", []) ) table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] res = await self._pipelines["sql_correction"].run( contexts=table_ddls, @@ -119,6 +124,7 @@ async def correct( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) post_process = res["post_process"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 87bf9b09de..86369e2ab1 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -4,6 +4,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + SQLGroundingValidator, construct_ask_history_messages, construct_instructions, get_json_field_instructions, @@ -77,6 +78,99 @@ async def execute_sql(self, *args, **kwargs): return False, {}, {"error_message": "dry run failed"} +def test_sql_grounding_validator_accepts_retrieved_schema_identifiers(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_dimension", "retrieved_measure"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'SELECT t."retrieved_dimension", SUM(t."retrieved_measure") AS result_label ' + 'FROM "retrieved_model" t ' + 'WHERE t."retrieved_dimension" = \'filter_value\' ' + 'GROUP BY t."retrieved_dimension" ' + "ORDER BY result_label" + ) + + assert valid is True + assert error == "" + + +def test_sql_grounding_validator_rejects_unretrieved_table(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'SELECT * FROM "unretrieved_model" WHERE "unretrieved_column" = \'filter_value\'' + ) + + assert valid is False + assert "table that is not declared" in error + + +def test_sql_grounding_validator_rejects_unretrieved_column(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'SELECT * FROM "retrieved_model" WHERE "unretrieved_column" = \'filter_value\'' + ) + + assert valid is False + assert "column that is not declared" in error + + +def test_sql_grounding_validator_handles_cte_output_columns(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'WITH scoped_result AS (SELECT "retrieved_column" AS result_column FROM "retrieved_model") ' + "SELECT * FROM scoped_result WHERE result_column = 'filter_value'" + ) + + assert valid is True + assert error == "" + + +@pytest.mark.asyncio +async def test_sql_postprocessor_rejects_ungrounded_sql_before_engine_validation(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=[ + '{"sql": "SELECT * FROM \\"unretrieved_model\\" WHERE \\"unretrieved_column\\" = \'filter_value\'"}' + ], + use_dry_plan=True, + data_source="source", + identifier_contracts=[ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ], + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SCHEMA_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + @pytest.mark.asyncio async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): result = await SQLGenPostProcessor(_FailingDryPlanEngine()).run( diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 2141eb0a02..0e0f5673be 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -69,6 +69,45 @@ def test_column_pruning_prompt_uses_current_query_without_history_text(): assert "previous request" not in result["prompt"] +def test_retrieval_results_include_exact_identifier_contract(): + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "retrieved_model", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "retrieved_column_a", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "retrieved_column_b", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=type("Encoding", (), {"encode": lambda self, text: [1]})(), + enable_column_pruning=False, + context_window_size=100, + ) + + assert result["db_schemas"][0]["identifier_contract"] == { + "table_name": "retrieved_model", + "columns": ["retrieved_column_a", "retrieved_column_b"], + } + + def test_table_selection_prompt_keeps_multiple_relevant_datasets(): assert "same business concept is represented by multiple modeled datasets" in ( table_columns_selection_system_prompt From fffd3054b012ef0812d7fa3b0ff3ec25cf96cbdb Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 16:52:57 +0530 Subject: [PATCH 0757/1087] Revert "Enforce retrieved schema identifiers before SQL validation" This reverts commit 89a40253c8e618846a24322628668bde7db76d39. --- .../generation/followup_sql_generation.py | 4 - .../pipelines/generation/sql_correction.py | 4 - .../pipelines/generation/sql_generation.py | 4 - .../pipelines/generation/sql_regeneration.py | 4 - .../src/pipelines/generation/utils/sql.py | 489 ------------------ .../retrieval/db_schema_retrieval.py | 44 -- wren-ai-service/src/web/v1/services/ask.py | 14 +- .../src/web/v1/services/ask_feedback.py | 12 +- .../v1/services/question_recommendation.py | 24 +- .../src/web/v1/services/sql_corrections.py | 6 - .../pipelines/generation/test_sql_utils.py | 94 ---- .../retrieval/test_db_schema_retrieval.py | 39 -- 12 files changed, 5 insertions(+), 733 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 19dfd24c3a..5d4edb01b6 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,7 +147,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -155,7 +154,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - identifier_contracts=identifier_contracts, ) @@ -207,7 +205,6 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -234,7 +231,6 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 2b90870194..f28d612aaa 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -148,7 +148,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -156,7 +155,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - identifier_contracts=identifier_contracts, ) @@ -204,7 +202,6 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -227,7 +224,6 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **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 97ffd0c3b3..2d5f1a2b96 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -138,7 +138,6 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -147,7 +146,6 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - identifier_contracts=identifier_contracts, ) @@ -199,7 +197,6 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -226,7 +223,6 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 59b7f7f7a9..b010d19c1b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -164,12 +164,10 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, - identifier_contracts=identifier_contracts, ) @@ -214,7 +212,6 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -233,7 +230,6 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 6b88fb5cef..4290b7847b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3,19 +3,9 @@ import aiohttp import orjson -import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel -from sqlparse.sql import ( - Function, - Identifier, - IdentifierList, - Parenthesis, - Statement, - TokenList, -) -from sqlparse.tokens import DML, Keyword, Name, Punctuation, Whitespace, Wildcard from src.core.engine import ( Engine, @@ -27,470 +17,6 @@ logger = logging.getLogger("wren-ai-service") -class SQLGroundingValidator: - def __init__(self, identifier_contracts: list[dict] | None = None): - self._tables = { - contract["table_name"]: ( - set(contract["columns"]) - if contract.get("columns") is not None - else None - ) - for contract in identifier_contracts or [] - if contract and contract.get("table_name") - } - - def validate(self, sql: str) -> tuple[bool, str]: - if not self._tables: - return True, "" - - try: - statements = sqlparse.parse(sql) - except Exception as exc: - return False, f"Unable to parse generated SQL for schema grounding: {exc}" - - if not statements: - return False, "No SQL statement was generated." - - for statement in statements: - context = _ValidationContext(schema_tables=self._tables) - valid, error = context.validate_statement(statement) - if not valid: - return False, error - - return True, "" - - -class _ValidationContext: - def __init__( - self, - schema_tables: dict[str, set[str] | None], - ctes: dict[str, set[str] | None] | None = None, - ): - self.schema_tables = schema_tables - self.ctes = ctes or {} - - def validate_statement(self, statement: Statement | TokenList) -> tuple[bool, str]: - ctes = dict(self.ctes) - valid, error = self._collect_ctes(statement, ctes) - if not valid: - return False, error - - scope: dict[str, set[str] | None] = {} - source_tokens: set[int] = set() - valid, error = self._collect_sources(statement, ctes, scope, source_tokens) - if not valid: - return False, error - - output_columns = self._infer_select_output_columns(statement, scope) - return self._validate_expressions( - statement, - scope=scope, - source_tokens=source_tokens, - output_columns=output_columns, - ) - - def _collect_ctes( - self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] - ) -> tuple[bool, str]: - tokens = _meaningful_tokens(statement) - if not tokens or not tokens[0].match(Keyword.CTE, "WITH"): - return True, "" - - for token in tokens[1:]: - if token.ttype is DML and token.normalized == "SELECT": - break - - identifiers = ( - list(token.get_identifiers()) - if isinstance(token, IdentifierList) - else [token] - if isinstance(token, Identifier) - else [] - ) - for identifier in identifiers: - cte_name = _identifier_name(identifier) - if not cte_name: - continue - subquery = _subquery_from_identifier(identifier) - if subquery is None: - continue - child_context = _ValidationContext(self.schema_tables, ctes=ctes) - valid, error = child_context.validate_statement(subquery) - if not valid: - return False, error - ctes[cte_name] = child_context._infer_select_output_columns( - subquery, child_context._collect_scope_for_inference(subquery, ctes) - ) - - return True, "" - - def _collect_scope_for_inference( - self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] - ) -> dict[str, set[str] | None]: - scope: dict[str, set[str] | None] = {} - self._collect_sources(statement, ctes, scope, set()) - return scope - - def _collect_sources( - self, - statement: Statement | TokenList, - ctes: dict[str, set[str] | None], - scope: dict[str, set[str] | None], - source_tokens: set[int], - ) -> tuple[bool, str]: - expect_source = False - for token in _meaningful_tokens(statement): - if token.ttype is DML and token.normalized == "SELECT": - expect_source = False - continue - - if _starts_new_clause(token): - expect_source = False - - if token.match( - Keyword, - ( - "FROM", - "JOIN", - "INNER JOIN", - "LEFT JOIN", - "LEFT OUTER JOIN", - "RIGHT JOIN", - "RIGHT OUTER JOIN", - "FULL JOIN", - "FULL OUTER JOIN", - "CROSS JOIN", - ), - ): - expect_source = True - continue - - if not expect_source: - continue - - if isinstance(token, IdentifierList): - for identifier in token.get_identifiers(): - valid, error = self._add_source(identifier, ctes, scope) - if not valid: - return False, error - source_tokens.add(id(identifier)) - continue - - if isinstance(token, Identifier): - valid, error = self._add_source(token, ctes, scope) - if not valid: - return False, error - source_tokens.add(id(token)) - continue - - if isinstance(token, Parenthesis): - subquery = _statement_from_parenthesis(token) - if subquery is not None: - child_context = _ValidationContext(self.schema_tables, ctes=ctes) - valid, error = child_context.validate_statement(subquery) - if not valid: - return False, error - continue - - return True, "" - - def _add_source( - self, - identifier: Identifier, - ctes: dict[str, set[str] | None], - scope: dict[str, set[str] | None], - ) -> tuple[bool, str]: - alias = identifier.get_alias() - subquery = _subquery_from_identifier(identifier) - if subquery is not None: - child_context = _ValidationContext(self.schema_tables, ctes=ctes) - valid, error = child_context.validate_statement(subquery) - if not valid: - return False, error - if alias: - scope[alias] = child_context._infer_select_output_columns( - subquery, child_context._collect_scope_for_inference(subquery, ctes) - ) - return True, "" - - table_name = identifier.get_real_name() - if not table_name: - return True, "" - - if table_name in ctes: - columns = ctes[table_name] - elif table_name in self.schema_tables: - columns = self.schema_tables[table_name] - else: - return ( - False, - "Generated SQL references a table that is not declared in the retrieved schema.", - ) - - scope[alias or table_name] = columns - scope[table_name] = columns - return True, "" - - def _infer_select_output_columns( - self, statement: Statement | TokenList, scope: dict[str, set[str] | None] - ) -> set[str] | None: - columns: set[str] = set() - in_select = False - - for token in _meaningful_tokens(statement): - if token.ttype is DML and token.normalized == "SELECT": - in_select = True - continue - - if in_select and token.match(Keyword, "FROM"): - return columns - - if not in_select: - continue - - if token.ttype is Wildcard: - return None - - identifiers = ( - list(token.get_identifiers()) - if isinstance(token, IdentifierList) - else [token] - if isinstance(token, Identifier) - else [] - ) - for identifier in identifiers: - if any(child.ttype is Wildcard for child in identifier.flatten()): - return None - name = identifier.get_alias() or identifier.get_real_name() - if name: - columns.add(name) - - return columns - - def _validate_expressions( - self, - token: TokenList, - scope: dict[str, set[str] | None], - source_tokens: set[int], - output_columns: set[str] | None, - in_order_by: bool = False, - ) -> tuple[bool, str]: - children = getattr(token, "tokens", []) - next_in_order_by = in_order_by - - for child in children: - if child.is_whitespace or child.ttype in Whitespace: - continue - - if id(child) in source_tokens: - continue - - if child.match(Keyword, "ORDER BY"): - next_in_order_by = True - continue - if _starts_new_clause(child) and not child.match(Keyword, "ORDER BY"): - next_in_order_by = False - - if isinstance(child, IdentifierList): - valid, error = self._validate_expressions( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if isinstance(child, Function): - valid, error = self._validate_function( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if isinstance(child, Identifier): - if any(isinstance(grandchild, Function) for grandchild in child.tokens): - valid, error = self._validate_identifier_children( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - valid, error = self._validate_identifier( - child, scope, output_columns, next_in_order_by - ) - if not valid: - return False, error - valid, error = self._validate_identifier_children( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if isinstance(child, Parenthesis): - subquery = _statement_from_parenthesis(child) - if subquery is not None: - valid, error = _ValidationContext( - self.schema_tables, self.ctes - ).validate_statement(subquery) - if not valid: - return False, error - continue - - if isinstance(child, TokenList): - valid, error = self._validate_expressions( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if child.ttype is Name: - valid, error = self._validate_unqualified_column( - str(child), scope, output_columns, next_in_order_by - ) - if not valid: - return False, error - - return True, "" - - def _validate_function( - self, - function: Function, - scope: dict[str, set[str] | None], - source_tokens: set[int], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - for child in function.tokens: - if isinstance(child, Parenthesis): - return self._validate_expressions( - child, scope, source_tokens, output_columns, in_order_by - ) - return True, "" - - def _validate_identifier_children( - self, - identifier: Identifier, - scope: dict[str, set[str] | None], - source_tokens: set[int], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - for child in identifier.tokens: - if isinstance(child, Function): - return self._validate_function( - child, scope, source_tokens, output_columns, in_order_by - ) - return True, "" - - def _validate_identifier( - self, - identifier: Identifier, - scope: dict[str, set[str] | None], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - parent_name = identifier.get_parent_name() - column_name = identifier.get_real_name() - - if not column_name: - return True, "" - - if parent_name: - if parent_name not in scope: - return ( - False, - "Generated SQL references a table alias that is not declared in the query scope.", - ) - columns = scope[parent_name] - if columns is None or column_name in columns: - return True, "" - return ( - False, - "Generated SQL references a column that is not declared in the retrieved schema.", - ) - - return self._validate_unqualified_column( - column_name, scope, output_columns, in_order_by - ) - - def _validate_unqualified_column( - self, - column_name: str, - scope: dict[str, set[str] | None], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - if in_order_by and output_columns is not None and column_name in output_columns: - return True, "" - - if column_name in scope: - return True, "" - - for columns in scope.values(): - if columns is None or column_name in columns: - return True, "" - - return ( - False, - "Generated SQL references a column that is not declared in the retrieved schema.", - ) - - -def _meaningful_tokens(token_list: TokenList) -> list: - return [ - token - for token in token_list.tokens - if not token.is_whitespace - and token.ttype not in Whitespace - and token.ttype not in Punctuation - ] - - -def _identifier_name(identifier: Identifier) -> str | None: - return identifier.get_real_name() or identifier.get_name() - - -def _subquery_from_identifier(identifier: Identifier) -> Statement | None: - for token in identifier.tokens: - if isinstance(token, Parenthesis): - return _statement_from_parenthesis(token) - return None - - -def _statement_from_parenthesis(parenthesis: Parenthesis) -> Statement | None: - inner_sql = str(parenthesis)[1:-1].strip() - if not inner_sql: - return None - - parsed = sqlparse.parse(inner_sql) - if parsed and any( - token.ttype is DML and token.normalized == "SELECT" - for token in _meaningful_tokens(parsed[0]) - ): - return parsed[0] - return None - - -def _starts_new_clause(token) -> bool: - return token.match( - Keyword, - ( - "WHERE", - "GROUP BY", - "HAVING", - "ORDER BY", - "LIMIT", - "UNION", - "UNION ALL", - "EXCEPT", - "INTERSECT", - ), - ) - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -508,7 +34,6 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - identifier_contracts: list[dict] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -529,7 +54,6 @@ async def run( allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, allow_data_preview=allow_data_preview, - identifier_contracts=identifier_contracts, ) return { @@ -552,7 +76,6 @@ async def _classify_generation_result( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - identifier_contracts: list[dict] | None = None, ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} @@ -567,18 +90,6 @@ async def _classify_generation_result( "correlation_id": "", } - is_grounded, grounding_error = SQLGroundingValidator( - identifier_contracts - ).validate(generation_result) - if not is_grounded: - return valid_generation_result, { - "sql": generation_result, - "original_sql": generation_result, - "type": "SCHEMA_GROUNDING", - "error": grounding_error, - "correlation_id": "", - } - async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 472cd95fea..64f52c5287 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -156,17 +156,6 @@ def _build_metric_ddl(content: dict) -> str: ) -def _build_metric_identifier_contract(content: dict) -> dict: - return { - "table_name": content["name"], - "columns": [ - column["name"] - for column in content["columns"] - if column["data_type"].lower() != "unknown" - ], - } - - def _build_view_ddl(content: dict) -> str: context = _format_semantic_context( { @@ -186,13 +175,6 @@ def _build_view_ddl(content: dict) -> str: ) -def _build_view_identifier_contract(content: dict) -> dict: - return { - "table_name": content["name"], - "columns": None, - } - - def _format_semantic_context(context: dict) -> str: return ( "/*\n" @@ -343,16 +325,6 @@ def _build_table_retrieval_context( return f"{context}{ddl}", has_calculated_field, has_json_field -def _build_table_identifier_contract( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None -) -> dict: - included_columns = _included_columns(content, columns, tables) - return { - "table_name": content["name"], - "columns": [column["name"] for column in included_columns], - } - - ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -616,9 +588,6 @@ def check_using_db_schemas_without_pruning( { "table_name": table_schema["name"], "table_ddl": ddl, - "identifier_contract": _build_table_identifier_contract( - table_schema - ), } ) if _has_calculated_field: @@ -634,7 +603,6 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), - "identifier_contract": _build_metric_identifier_contract(content), } ) has_metric = True @@ -643,7 +611,6 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), - "identifier_contract": _build_view_identifier_contract(content), } ) @@ -754,11 +721,6 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, - "identifier_contract": _build_table_identifier_contract( - table_schema, - columns=columns, - tables=tables, - ), } ) @@ -770,9 +732,6 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), - "identifier_contract": _build_metric_identifier_contract( - content - ), } ) has_metric = True @@ -781,9 +740,6 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), - "identifier_contract": _build_view_identifier_contract( - content - ), } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 106a3d4e41..ca55061902 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -332,11 +332,6 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -459,7 +454,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -478,7 +472,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -508,7 +501,6 @@ async def ask( invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] current_sql_correction_retries += 1 - sql_diagnosis_reasoning = None self._ask_results[query_id] = AskResultResponse( status="correcting", @@ -521,10 +513,7 @@ async def ask( is_followup=True if histories else False, ) - if ( - allow_sql_diagnosis - and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -559,7 +548,6 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 024dcfc3dd..f1971afa15 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -161,11 +161,6 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] sql_samples = sql_samples_task["formatted_output"].get("documents", []) instructions = instructions_task["formatted_output"].get( "documents", [] @@ -192,7 +187,6 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -222,10 +216,7 @@ async def ask_feedback( trace_id=trace_id, ) - if ( - allow_sql_diagnosis - and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -260,7 +251,6 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 1e7383bbf8..694d044bfa 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -71,7 +71,7 @@ async def _validate_question( use_dry_plan: bool = True, allow_dry_plan_fallback: bool = False, ): - async def _document_retrieval() -> tuple[list[str], list[dict], bool, bool, bool]: + async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -79,21 +79,10 @@ async def _document_retrieval() -> tuple[list[str], list[dict], bool, bool, bool _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return ( - table_ddls, - identifier_contracts, - has_calculated_field, - has_metric, - has_json_field, - ) + return table_ddls, has_calculated_field, has_metric, has_json_field async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -118,13 +107,7 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - ( - table_ddls, - identifier_contracts, - has_calculated_field, - has_metric, - has_json_field, - ) = _document + table_ddls, has_calculated_field, has_metric, has_json_field = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -154,7 +137,6 @@ async def _instructions_retrieval() -> list[dict]: allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) post_process = generated_sql["post_process"] diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index bec96f50fc..4336186b6f 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -111,11 +111,6 @@ async def correct( .get("retrieval_results", []) ) table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] res = await self._pipelines["sql_correction"].run( contexts=table_ddls, @@ -124,7 +119,6 @@ async def correct( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) post_process = res["post_process"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 86369e2ab1..87bf9b09de 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -4,7 +4,6 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, - SQLGroundingValidator, construct_ask_history_messages, construct_instructions, get_json_field_instructions, @@ -78,99 +77,6 @@ async def execute_sql(self, *args, **kwargs): return False, {}, {"error_message": "dry run failed"} -def test_sql_grounding_validator_accepts_retrieved_schema_identifiers(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_dimension", "retrieved_measure"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'SELECT t."retrieved_dimension", SUM(t."retrieved_measure") AS result_label ' - 'FROM "retrieved_model" t ' - 'WHERE t."retrieved_dimension" = \'filter_value\' ' - 'GROUP BY t."retrieved_dimension" ' - "ORDER BY result_label" - ) - - assert valid is True - assert error == "" - - -def test_sql_grounding_validator_rejects_unretrieved_table(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'SELECT * FROM "unretrieved_model" WHERE "unretrieved_column" = \'filter_value\'' - ) - - assert valid is False - assert "table that is not declared" in error - - -def test_sql_grounding_validator_rejects_unretrieved_column(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'SELECT * FROM "retrieved_model" WHERE "unretrieved_column" = \'filter_value\'' - ) - - assert valid is False - assert "column that is not declared" in error - - -def test_sql_grounding_validator_handles_cte_output_columns(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'WITH scoped_result AS (SELECT "retrieved_column" AS result_column FROM "retrieved_model") ' - "SELECT * FROM scoped_result WHERE result_column = 'filter_value'" - ) - - assert valid is True - assert error == "" - - -@pytest.mark.asyncio -async def test_sql_postprocessor_rejects_ungrounded_sql_before_engine_validation(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=[ - '{"sql": "SELECT * FROM \\"unretrieved_model\\" WHERE \\"unretrieved_column\\" = \'filter_value\'"}' - ], - use_dry_plan=True, - data_source="source", - identifier_contracts=[ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ], - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "SCHEMA_GROUNDING" - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] - - @pytest.mark.asyncio async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): result = await SQLGenPostProcessor(_FailingDryPlanEngine()).run( diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 0e0f5673be..2141eb0a02 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -69,45 +69,6 @@ def test_column_pruning_prompt_uses_current_query_without_history_text(): assert "previous request" not in result["prompt"] -def test_retrieval_results_include_exact_identifier_contract(): - result = check_using_db_schemas_without_pruning( - construct_db_schemas=[ - { - "type": "TABLE", - "name": "retrieved_model", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "retrieved_column_a", - "data_type": "VARCHAR", - "comment": "", - "is_primary_key": False, - }, - { - "type": "COLUMN", - "name": "retrieved_column_b", - "data_type": "VARCHAR", - "comment": "", - "is_primary_key": False, - }, - ], - "properties": {}, - "primaryKey": "", - } - ], - dbschema_retrieval=[], - encoding=type("Encoding", (), {"encode": lambda self, text: [1]})(), - enable_column_pruning=False, - context_window_size=100, - ) - - assert result["db_schemas"][0]["identifier_contract"] == { - "table_name": "retrieved_model", - "columns": ["retrieved_column_a", "retrieved_column_b"], - } - - def test_table_selection_prompt_keeps_multiple_relevant_datasets(): assert "same business concept is represented by multiple modeled datasets" in ( table_columns_selection_system_prompt From 3df3c72a4181ee737d833160812b0cc5d3532d7e Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 17:27:28 +0530 Subject: [PATCH 0758/1087] Prevent ungrounded SQL identifiers --- .../generation/followup_sql_generation.py | 4 + .../pipelines/generation/sql_correction.py | 4 + .../pipelines/generation/sql_generation.py | 4 + .../pipelines/generation/sql_regeneration.py | 4 + .../src/pipelines/generation/utils/sql.py | 490 ++++++++++++++++++ .../retrieval/db_schema_retrieval.py | 44 ++ wren-ai-service/src/providers/engine/wren.py | 2 + .../src/web/v1/routers/sql_corrections.py | 1 + wren-ai-service/src/web/v1/services/ask.py | 14 +- .../src/web/v1/services/ask_feedback.py | 12 +- .../v1/services/question_recommendation.py | 24 +- .../src/web/v1/services/sql_corrections.py | 9 + .../pipelines/generation/test_sql_utils.py | 94 ++++ .../retrieval/test_db_schema_retrieval.py | 39 ++ wren-engine | 2 +- .../src/apollo/client/graphql/__types__.ts | 1 + .../src/apollo/server/adaptors/ibisAdaptor.ts | 11 + .../server/adaptors/tests/ibisAdaptor.test.ts | 35 ++ wren-ui/src/apollo/server/models/model.ts | 1 + .../apollo/server/resolvers/modelResolver.ts | 3 +- wren-ui/src/apollo/server/schema.ts | 1 + .../apollo/server/services/queryService.ts | 8 + 22 files changed, 800 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 5d4edb01b6..19dfd24c3a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,6 +147,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -154,6 +155,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + identifier_contracts=identifier_contracts, ) @@ -205,6 +207,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -231,6 +234,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f28d612aaa..2b90870194 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -148,6 +148,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -155,6 +156,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + identifier_contracts=identifier_contracts, ) @@ -202,6 +204,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -224,6 +227,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **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 2d5f1a2b96..97ffd0c3b3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -138,6 +138,7 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -146,6 +147,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + identifier_contracts=identifier_contracts, ) @@ -197,6 +199,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -223,6 +226,7 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index b010d19c1b..59b7f7f7a9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -164,10 +164,12 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, + identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, + identifier_contracts=identifier_contracts, ) @@ -212,6 +214,7 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + identifier_contracts: list[dict] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -230,6 +233,7 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, + "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4290b7847b..cef2a01047 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3,9 +3,19 @@ import aiohttp import orjson +import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel +from sqlparse.sql import ( + Function, + Identifier, + IdentifierList, + Parenthesis, + Statement, + TokenList, +) +from sqlparse.tokens import DML, Keyword, Name, Punctuation, Whitespace, Wildcard from src.core.engine import ( Engine, @@ -17,6 +27,470 @@ logger = logging.getLogger("wren-ai-service") +class SQLGroundingValidator: + def __init__(self, identifier_contracts: list[dict] | None = None): + self._tables = { + contract["table_name"]: ( + set(contract["columns"]) + if contract.get("columns") is not None + else None + ) + for contract in identifier_contracts or [] + if contract and contract.get("table_name") + } + + def validate(self, sql: str) -> tuple[bool, str]: + if not self._tables: + return True, "" + + try: + statements = sqlparse.parse(sql) + except Exception as exc: + return False, f"Unable to parse generated SQL for schema grounding: {exc}" + + if not statements: + return False, "No SQL statement was generated." + + for statement in statements: + context = _ValidationContext(schema_tables=self._tables) + valid, error = context.validate_statement(statement) + if not valid: + return False, error + + return True, "" + + +class _ValidationContext: + def __init__( + self, + schema_tables: dict[str, set[str] | None], + ctes: dict[str, set[str] | None] | None = None, + ): + self.schema_tables = schema_tables + self.ctes = ctes or {} + + def validate_statement(self, statement: Statement | TokenList) -> tuple[bool, str]: + ctes = dict(self.ctes) + valid, error = self._collect_ctes(statement, ctes) + if not valid: + return False, error + + scope: dict[str, set[str] | None] = {} + source_tokens: set[int] = set() + valid, error = self._collect_sources(statement, ctes, scope, source_tokens) + if not valid: + return False, error + + output_columns = self._infer_select_output_columns(statement, scope) + return self._validate_expressions( + statement, + scope=scope, + source_tokens=source_tokens, + output_columns=output_columns, + ) + + def _collect_ctes( + self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] + ) -> tuple[bool, str]: + tokens = _meaningful_tokens(statement) + if not tokens or not tokens[0].match(Keyword.CTE, "WITH"): + return True, "" + + for token in tokens[1:]: + if token.ttype is DML and token.normalized == "SELECT": + break + + identifiers = ( + list(token.get_identifiers()) + if isinstance(token, IdentifierList) + else [token] + if isinstance(token, Identifier) + else [] + ) + for identifier in identifiers: + cte_name = _identifier_name(identifier) + if not cte_name: + continue + subquery = _subquery_from_identifier(identifier) + if subquery is None: + continue + child_context = _ValidationContext(self.schema_tables, ctes=ctes) + valid, error = child_context.validate_statement(subquery) + if not valid: + return False, error + ctes[cte_name] = child_context._infer_select_output_columns( + subquery, child_context._collect_scope_for_inference(subquery, ctes) + ) + + return True, "" + + def _collect_scope_for_inference( + self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] + ) -> dict[str, set[str] | None]: + scope: dict[str, set[str] | None] = {} + self._collect_sources(statement, ctes, scope, set()) + return scope + + def _collect_sources( + self, + statement: Statement | TokenList, + ctes: dict[str, set[str] | None], + scope: dict[str, set[str] | None], + source_tokens: set[int], + ) -> tuple[bool, str]: + expect_source = False + for token in _meaningful_tokens(statement): + if token.ttype is DML and token.normalized == "SELECT": + expect_source = False + continue + + if _starts_new_clause(token): + expect_source = False + + if token.match( + Keyword, + ( + "FROM", + "JOIN", + "INNER JOIN", + "LEFT JOIN", + "LEFT OUTER JOIN", + "RIGHT JOIN", + "RIGHT OUTER JOIN", + "FULL JOIN", + "FULL OUTER JOIN", + "CROSS JOIN", + ), + ): + expect_source = True + continue + + if not expect_source: + continue + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + valid, error = self._add_source(identifier, ctes, scope) + if not valid: + return False, error + source_tokens.add(id(identifier)) + continue + + if isinstance(token, Identifier): + valid, error = self._add_source(token, ctes, scope) + if not valid: + return False, error + source_tokens.add(id(token)) + continue + + if isinstance(token, Parenthesis): + subquery = _statement_from_parenthesis(token) + if subquery is not None: + child_context = _ValidationContext(self.schema_tables, ctes=ctes) + valid, error = child_context.validate_statement(subquery) + if not valid: + return False, error + continue + + return True, "" + + def _add_source( + self, + identifier: Identifier, + ctes: dict[str, set[str] | None], + scope: dict[str, set[str] | None], + ) -> tuple[bool, str]: + alias = identifier.get_alias() + subquery = _subquery_from_identifier(identifier) + if subquery is not None: + child_context = _ValidationContext(self.schema_tables, ctes=ctes) + valid, error = child_context.validate_statement(subquery) + if not valid: + return False, error + if alias: + scope[alias] = child_context._infer_select_output_columns( + subquery, child_context._collect_scope_for_inference(subquery, ctes) + ) + return True, "" + + table_name = identifier.get_real_name() + if not table_name: + return True, "" + + if table_name in ctes: + columns = ctes[table_name] + elif table_name in self.schema_tables: + columns = self.schema_tables[table_name] + else: + return ( + False, + "Generated SQL references a table that is not declared in the retrieved schema.", + ) + + scope[alias or table_name] = columns + scope[table_name] = columns + return True, "" + + def _infer_select_output_columns( + self, statement: Statement | TokenList, scope: dict[str, set[str] | None] + ) -> set[str] | None: + columns: set[str] = set() + in_select = False + + for token in _meaningful_tokens(statement): + if token.ttype is DML and token.normalized == "SELECT": + in_select = True + continue + + if in_select and token.match(Keyword, "FROM"): + return columns + + if not in_select: + continue + + if token.ttype is Wildcard: + return None + + identifiers = ( + list(token.get_identifiers()) + if isinstance(token, IdentifierList) + else [token] + if isinstance(token, Identifier) + else [] + ) + for identifier in identifiers: + if any(child.ttype is Wildcard for child in identifier.flatten()): + return None + name = identifier.get_alias() or identifier.get_real_name() + if name: + columns.add(name) + + return columns + + def _validate_expressions( + self, + token: TokenList, + scope: dict[str, set[str] | None], + source_tokens: set[int], + output_columns: set[str] | None, + in_order_by: bool = False, + ) -> tuple[bool, str]: + children = getattr(token, "tokens", []) + next_in_order_by = in_order_by + + for child in children: + if child.is_whitespace or child.ttype in Whitespace: + continue + + if id(child) in source_tokens: + continue + + if child.match(Keyword, "ORDER BY"): + next_in_order_by = True + continue + if _starts_new_clause(child) and not child.match(Keyword, "ORDER BY"): + next_in_order_by = False + + if isinstance(child, IdentifierList): + valid, error = self._validate_expressions( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if isinstance(child, Function): + valid, error = self._validate_function( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if isinstance(child, Identifier): + if any(isinstance(grandchild, Function) for grandchild in child.tokens): + valid, error = self._validate_identifier_children( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + valid, error = self._validate_identifier( + child, scope, output_columns, next_in_order_by + ) + if not valid: + return False, error + valid, error = self._validate_identifier_children( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if isinstance(child, Parenthesis): + subquery = _statement_from_parenthesis(child) + if subquery is not None: + valid, error = _ValidationContext( + self.schema_tables, self.ctes + ).validate_statement(subquery) + if not valid: + return False, error + continue + + if isinstance(child, TokenList): + valid, error = self._validate_expressions( + child, scope, source_tokens, output_columns, next_in_order_by + ) + if not valid: + return False, error + continue + + if child.ttype is Name: + valid, error = self._validate_unqualified_column( + str(child), scope, output_columns, next_in_order_by + ) + if not valid: + return False, error + + return True, "" + + def _validate_function( + self, + function: Function, + scope: dict[str, set[str] | None], + source_tokens: set[int], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + for child in function.tokens: + if isinstance(child, Parenthesis): + return self._validate_expressions( + child, scope, source_tokens, output_columns, in_order_by + ) + return True, "" + + def _validate_identifier_children( + self, + identifier: Identifier, + scope: dict[str, set[str] | None], + source_tokens: set[int], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + for child in identifier.tokens: + if isinstance(child, Function): + return self._validate_function( + child, scope, source_tokens, output_columns, in_order_by + ) + return True, "" + + def _validate_identifier( + self, + identifier: Identifier, + scope: dict[str, set[str] | None], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + parent_name = identifier.get_parent_name() + column_name = identifier.get_real_name() + + if not column_name: + return True, "" + + if parent_name: + if parent_name not in scope: + return ( + False, + "Generated SQL references a table alias that is not declared in the query scope.", + ) + columns = scope[parent_name] + if columns is None or column_name in columns: + return True, "" + return ( + False, + "Generated SQL references a column that is not declared in the retrieved schema.", + ) + + return self._validate_unqualified_column( + column_name, scope, output_columns, in_order_by + ) + + def _validate_unqualified_column( + self, + column_name: str, + scope: dict[str, set[str] | None], + output_columns: set[str] | None, + in_order_by: bool, + ) -> tuple[bool, str]: + if in_order_by and output_columns is not None and column_name in output_columns: + return True, "" + + if column_name in scope: + return True, "" + + for columns in scope.values(): + if columns is None or column_name in columns: + return True, "" + + return ( + False, + "Generated SQL references a column that is not declared in the retrieved schema.", + ) + + +def _meaningful_tokens(token_list: TokenList) -> list: + return [ + token + for token in token_list.tokens + if not token.is_whitespace + and token.ttype not in Whitespace + and token.ttype not in Punctuation + ] + + +def _identifier_name(identifier: Identifier) -> str | None: + return identifier.get_real_name() or identifier.get_name() + + +def _subquery_from_identifier(identifier: Identifier) -> Statement | None: + for token in identifier.tokens: + if isinstance(token, Parenthesis): + return _statement_from_parenthesis(token) + return None + + +def _statement_from_parenthesis(parenthesis: Parenthesis) -> Statement | None: + inner_sql = str(parenthesis)[1:-1].strip() + if not inner_sql: + return None + + parsed = sqlparse.parse(inner_sql) + if parsed and any( + token.ttype is DML and token.normalized == "SELECT" + for token in _meaningful_tokens(parsed[0]) + ): + return parsed[0] + return None + + +def _starts_new_clause(token) -> bool: + return token.match( + Keyword, + ( + "WHERE", + "GROUP BY", + "HAVING", + "ORDER BY", + "LIMIT", + "UNION", + "UNION ALL", + "EXCEPT", + "INTERSECT", + ), + ) + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -34,6 +508,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + identifier_contracts: list[dict] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -54,6 +529,7 @@ async def run( allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, allow_data_preview=allow_data_preview, + identifier_contracts=identifier_contracts, ) return { @@ -76,6 +552,7 @@ async def _classify_generation_result( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + identifier_contracts: list[dict] | None = None, ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} @@ -90,6 +567,18 @@ async def _classify_generation_result( "correlation_id": "", } + is_grounded, grounding_error = SQLGroundingValidator( + identifier_contracts + ).validate(generation_result) + if not is_grounded: + return valid_generation_result, { + "sql": generation_result, + "original_sql": generation_result, + "type": "SCHEMA_GROUNDING", + "error": grounding_error, + "correlation_id": "", + } + async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( @@ -118,6 +607,7 @@ async def _classify_generation_result( project_id=project_id, limit=1, dry_run=True, + allow_fallback=allow_dry_plan_fallback, ) addition = addition if isinstance(addition, dict) else {} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 64f52c5287..472cd95fea 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -156,6 +156,17 @@ def _build_metric_ddl(content: dict) -> str: ) +def _build_metric_identifier_contract(content: dict) -> dict: + return { + "table_name": content["name"], + "columns": [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + } + + def _build_view_ddl(content: dict) -> str: context = _format_semantic_context( { @@ -175,6 +186,13 @@ def _build_view_ddl(content: dict) -> str: ) +def _build_view_identifier_contract(content: dict) -> dict: + return { + "table_name": content["name"], + "columns": None, + } + + def _format_semantic_context(context: dict) -> str: return ( "/*\n" @@ -325,6 +343,16 @@ def _build_table_retrieval_context( return f"{context}{ddl}", has_calculated_field, has_json_field +def _build_table_identifier_contract( + content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None +) -> dict: + included_columns = _included_columns(content, columns, tables) + return { + "table_name": content["name"], + "columns": [column["name"] for column in included_columns], + } + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -588,6 +616,9 @@ def check_using_db_schemas_without_pruning( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_contract": _build_table_identifier_contract( + table_schema + ), } ) if _has_calculated_field: @@ -603,6 +634,7 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_contract": _build_metric_identifier_contract(content), } ) has_metric = True @@ -611,6 +643,7 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_contract": _build_view_identifier_contract(content), } ) @@ -721,6 +754,11 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_contract": _build_table_identifier_contract( + table_schema, + columns=columns, + tables=tables, + ), } ) @@ -732,6 +770,9 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_contract": _build_metric_identifier_contract( + content + ), } ) has_metric = True @@ -740,6 +781,9 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_contract": _build_view_identifier_contract( + content + ), } ) diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 8a4eab6479..941a8776c8 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -31,11 +31,13 @@ async def execute_sql( dry_run: bool = True, timeout: float = settings.engine_timeout, limit: int = 500, + allow_fallback: bool = True, **kwargs, ) -> Tuple[bool, Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: data = { "sql": remove_limit_statement(sql), "projectId": project_id, + "allowFallback": allow_fallback, } if dry_run: data["dryRun"] = True diff --git a/wren-ai-service/src/web/v1/routers/sql_corrections.py b/wren-ai-service/src/web/v1/routers/sql_corrections.py index 60dbd9246b..52592f950e 100644 --- a/wren-ai-service/src/web/v1/routers/sql_corrections.py +++ b/wren-ai-service/src/web/v1/routers/sql_corrections.py @@ -19,6 +19,7 @@ class PostRequest(BaseRequest): sql: str error: str + query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = True allow_dry_plan_fallback: bool = False diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ca55061902..106a3d4e41 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -332,6 +332,11 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -454,6 +459,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -472,6 +478,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -501,6 +508,7 @@ async def ask( invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] current_sql_correction_retries += 1 + sql_diagnosis_reasoning = None self._ask_results[query_id] = AskResultResponse( status="correcting", @@ -513,7 +521,10 @@ async def ask( is_followup=True if histories else False, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -548,6 +559,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index f1971afa15..024dcfc3dd 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -161,6 +161,11 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] sql_samples = sql_samples_task["formatted_output"].get("documents", []) instructions = instructions_task["formatted_output"].get( "documents", [] @@ -187,6 +192,7 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -216,7 +222,10 @@ async def ask_feedback( trace_id=trace_id, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -251,6 +260,7 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 694d044bfa..1e7383bbf8 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -71,7 +71,7 @@ async def _validate_question( use_dry_plan: bool = True, allow_dry_plan_fallback: bool = False, ): - async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: + async def _document_retrieval() -> tuple[list[str], list[dict], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -79,10 +79,21 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return table_ddls, has_calculated_field, has_metric, has_json_field + return ( + table_ddls, + identifier_contracts, + has_calculated_field, + has_metric, + has_json_field, + ) async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -107,7 +118,13 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - table_ddls, has_calculated_field, has_metric, has_json_field = _document + ( + table_ddls, + identifier_contracts, + has_calculated_field, + has_metric, + has_json_field, + ) = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -137,6 +154,7 @@ async def _instructions_retrieval() -> list[dict]: allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) post_process = generated_sql["post_process"] diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 4336186b6f..aab7dd2665 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -60,6 +60,7 @@ class CorrectionRequest(BaseRequest): event_id: str sql: str error: str + query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = True allow_dry_plan_fallback: bool = False @@ -76,6 +77,7 @@ async def correct( event_id = request.event_id sql = request.sql error = request.error + query = request.query project_id = request.project_id retrieved_tables = request.retrieved_tables use_dry_plan = request.use_dry_plan @@ -111,14 +113,21 @@ async def correct( .get("retrieval_results", []) ) table_ddls = [document.get("table_ddl") for document in documents] + identifier_contracts = [ + document.get("identifier_contract") + for document in documents + if document.get("identifier_contract") + ] res = await self._pipelines["sql_correction"].run( contexts=table_ddls, + query=query, invalid_generation_result=_invalid, project_id=project_id, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + identifier_contracts=identifier_contracts, ) post_process = res["post_process"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 87bf9b09de..86369e2ab1 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -4,6 +4,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + SQLGroundingValidator, construct_ask_history_messages, construct_instructions, get_json_field_instructions, @@ -77,6 +78,99 @@ async def execute_sql(self, *args, **kwargs): return False, {}, {"error_message": "dry run failed"} +def test_sql_grounding_validator_accepts_retrieved_schema_identifiers(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_dimension", "retrieved_measure"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'SELECT t."retrieved_dimension", SUM(t."retrieved_measure") AS result_label ' + 'FROM "retrieved_model" t ' + 'WHERE t."retrieved_dimension" = \'filter_value\' ' + 'GROUP BY t."retrieved_dimension" ' + "ORDER BY result_label" + ) + + assert valid is True + assert error == "" + + +def test_sql_grounding_validator_rejects_unretrieved_table(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'SELECT * FROM "unretrieved_model" WHERE "unretrieved_column" = \'filter_value\'' + ) + + assert valid is False + assert "table that is not declared" in error + + +def test_sql_grounding_validator_rejects_unretrieved_column(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'SELECT * FROM "retrieved_model" WHERE "unretrieved_column" = \'filter_value\'' + ) + + assert valid is False + assert "column that is not declared" in error + + +def test_sql_grounding_validator_handles_cte_output_columns(): + contracts = [ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ] + + valid, error = SQLGroundingValidator(contracts).validate( + 'WITH scoped_result AS (SELECT "retrieved_column" AS result_column FROM "retrieved_model") ' + "SELECT * FROM scoped_result WHERE result_column = 'filter_value'" + ) + + assert valid is True + assert error == "" + + +@pytest.mark.asyncio +async def test_sql_postprocessor_rejects_ungrounded_sql_before_engine_validation(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=[ + '{"sql": "SELECT * FROM \\"unretrieved_model\\" WHERE \\"unretrieved_column\\" = \'filter_value\'"}' + ], + use_dry_plan=True, + data_source="source", + identifier_contracts=[ + { + "table_name": "retrieved_model", + "columns": ["retrieved_column"], + } + ], + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SCHEMA_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + @pytest.mark.asyncio async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): result = await SQLGenPostProcessor(_FailingDryPlanEngine()).run( diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 2141eb0a02..0e0f5673be 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -69,6 +69,45 @@ def test_column_pruning_prompt_uses_current_query_without_history_text(): assert "previous request" not in result["prompt"] +def test_retrieval_results_include_exact_identifier_contract(): + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "retrieved_model", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "retrieved_column_a", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "retrieved_column_b", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=type("Encoding", (), {"encode": lambda self, text: [1]})(), + enable_column_pruning=False, + context_window_size=100, + ) + + assert result["db_schemas"][0]["identifier_contract"] == { + "table_name": "retrieved_model", + "columns": ["retrieved_column_a", "retrieved_column_b"], + } + + def test_table_selection_prompt_keeps_multiple_relevant_datasets(): assert "same business concept is represented by multiple modeled datasets" in ( table_columns_selection_system_prompt diff --git a/wren-engine b/wren-engine index 44d0811961..abc150f8f7 160000 --- a/wren-engine +++ b/wren-engine @@ -1 +1 @@ -Subproject commit 44d08119612dca9c8ff007fa12b305fd1aad7593 +Subproject commit abc150f8f703b04bfc1a63df70e50e6e81218dc2 diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index f3da8f04fd..d04ed87cfa 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -1142,6 +1142,7 @@ export type PreviewItemSqlInput = { }; export type PreviewSqlDataInput = { + allowFallback?: InputMaybe; dryRun?: InputMaybe; limit?: InputMaybe; projectId?: InputMaybe; diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 5d421cd975..07b8b9b098 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -186,6 +186,7 @@ export interface IbisQueryOptions extends IbisBaseOptions { limit?: number; refresh?: boolean; cacheEnabled?: boolean; + allowFallback?: boolean; } export interface IbisDryPlanOptions { dataSource: DataSourceName; @@ -317,6 +318,10 @@ export class IbisAdaptor implements IIbisAdaptor { params: { limit: options.limit || DEFAULT_PREVIEW_LIMIT, }, + headers: { + 'x-wren-fallback_disable': + options.allowFallback === false ? 'true' : 'false', + }, }, ); return { @@ -357,6 +362,12 @@ export class IbisAdaptor implements IIbisAdaptor { const response = await axios.post( `${this.ibisServerEndpoint}/${this.getIbisApiVersion(IBIS_API_TYPE.DRY_RUN)}/connector/${dataSourceUrlMap[dataSource]}/query?dryRun=true`, body, + { + headers: { + 'x-wren-fallback_disable': + options.allowFallback === false ? 'true' : 'false', + }, + }, ); logger.debug(`Ibis server Dry run success`); return { diff --git a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts index 44bf4dc5ee..c51fdeca2e 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts @@ -582,6 +582,9 @@ describe('IbisAdaptor', () => { params: { limit: DEFAULT_PREVIEW_LIMIT, }, + headers: { + 'x-wren-fallback_disable': 'false', + }, }, ); }); @@ -617,6 +620,9 @@ describe('IbisAdaptor', () => { params: { limit: customLimit, }, + headers: { + 'x-wren-fallback_disable': 'false', + }, }, ); }); @@ -678,6 +684,35 @@ describe('IbisAdaptor', () => { expect(res.processTime).toEqual('1s'); }); + it('should disable v3 fallback during dry run when requested', async () => { + mockedAxios.post.mockResolvedValue({ + headers: { + 'x-correlation-id': '123', + 'x-process-time': '1s', + }, + }); + mockedEncryptor.prototype.decrypt.mockReturnValue( + JSON.stringify({ password: mockPostgresConnectionInfo.password }), + ); + + await ibisAdaptor.dryRun('SELECT * FROM test_table', { + dataSource: DataSourceName.POSTGRES, + connectionInfo: mockPostgresConnectionInfo, + mdl: mockManifest, + allowFallback: false, + }); + + expect(mockedAxios.post).toHaveBeenCalledWith( + expect.any(String), + expect.any(Object), + { + headers: { + 'x-wren-fallback_disable': 'true', + }, + }, + ); + }); + it('should throw an exception with correlationId and processTime when dry run fails', async () => { const mockError = { response: { diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index 42262b3626..7b6aa3ba22 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -100,6 +100,7 @@ export interface PreviewSQLData { projectId?: string; limit?: number; dryRun?: boolean; + allowFallback?: boolean; } export interface DryPlanSQLData { diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index ce98d0f4db..6eee24d818 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -1336,7 +1336,7 @@ export class ModelResolver { args: { data: PreviewSQLData }, ctx: IContext, ) { - const { sql, projectId, limit, dryRun } = args.data; + const { sql, projectId, limit, dryRun, allowFallback } = args.data; const project = projectId ? await ctx.projectService.getProjectById(parseInt(projectId)) : await ctx.projectService.getCurrentProject(); @@ -1347,6 +1347,7 @@ export class ModelResolver { modelingOnly: false, manifest, dryRun, + allowFallback, }); } diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 90fd686c77..23c7d0c7f3 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -925,6 +925,7 @@ export const typeDefs = gql` projectId: String limit: Int dryRun: Boolean + allowFallback: Boolean } input DryPlanSQLDataInput { diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index ee7b341fdd..24808037d9 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -42,6 +42,7 @@ export interface PreviewOptions { manifest: Manifest; limit?: number; dryRun?: boolean; + allowFallback?: boolean; refresh?: boolean; cacheEnabled?: boolean; } @@ -104,6 +105,7 @@ export class QueryService implements IQueryService { manifest: mdl, limit, dryRun, + allowFallback, refresh, cacheEnabled, } = options; @@ -134,6 +136,7 @@ export class QueryService implements IQueryService { dataSource, connectionInfo, mdl, + allowFallback, ); } else { return await this.ibisQuery( @@ -144,6 +147,7 @@ export class QueryService implements IQueryService { limit, refresh, cacheEnabled, + allowFallback, ); } } @@ -202,6 +206,7 @@ export class QueryService implements IQueryService { dataSource: DataSourceName, connectionInfo: any, mdl: Manifest, + allowFallback?: boolean, ): Promise { const event = TelemetryEvent.IBIS_DRY_RUN; try { @@ -209,6 +214,7 @@ export class QueryService implements IQueryService { dataSource, connectionInfo, mdl, + allowFallback, }); this.sendIbisEvent(event, res, { dataSource, sql }); return { @@ -231,6 +237,7 @@ export class QueryService implements IQueryService { limit: number, refresh?: boolean, cacheEnabled?: boolean, + allowFallback?: boolean, ): Promise { const event = TelemetryEvent.IBIS_QUERY; try { @@ -241,6 +248,7 @@ export class QueryService implements IQueryService { limit, refresh, cacheEnabled, + allowFallback, }); this.sendIbisEvent(event, res, { dataSource, From 6d2348d24294b8a8a6fade6e9592acad86c5a88a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 17:39:51 +0530 Subject: [PATCH 0759/1087] Revert "Prevent ungrounded SQL identifiers" This reverts commit 3df3c72a4181ee737d833160812b0cc5d3532d7e. --- .../generation/followup_sql_generation.py | 4 - .../pipelines/generation/sql_correction.py | 4 - .../pipelines/generation/sql_generation.py | 4 - .../pipelines/generation/sql_regeneration.py | 4 - .../src/pipelines/generation/utils/sql.py | 490 ------------------ .../retrieval/db_schema_retrieval.py | 44 -- wren-ai-service/src/providers/engine/wren.py | 2 - .../src/web/v1/routers/sql_corrections.py | 1 - wren-ai-service/src/web/v1/services/ask.py | 14 +- .../src/web/v1/services/ask_feedback.py | 12 +- .../v1/services/question_recommendation.py | 24 +- .../src/web/v1/services/sql_corrections.py | 9 - .../pipelines/generation/test_sql_utils.py | 94 ---- .../retrieval/test_db_schema_retrieval.py | 39 -- wren-engine | 2 +- .../src/apollo/client/graphql/__types__.ts | 1 - .../src/apollo/server/adaptors/ibisAdaptor.ts | 11 - .../server/adaptors/tests/ibisAdaptor.test.ts | 35 -- wren-ui/src/apollo/server/models/model.ts | 1 - .../apollo/server/resolvers/modelResolver.ts | 3 +- wren-ui/src/apollo/server/schema.ts | 1 - .../apollo/server/services/queryService.ts | 8 - 22 files changed, 7 insertions(+), 800 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 19dfd24c3a..5d4edb01b6 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -147,7 +147,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -155,7 +154,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - identifier_contracts=identifier_contracts, ) @@ -207,7 +205,6 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -234,7 +231,6 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 2b90870194..f28d612aaa 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -148,7 +148,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -156,7 +155,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - identifier_contracts=identifier_contracts, ) @@ -204,7 +202,6 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -227,7 +224,6 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **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 97ffd0c3b3..2d5f1a2b96 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -138,7 +138,6 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -147,7 +146,6 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - identifier_contracts=identifier_contracts, ) @@ -199,7 +197,6 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -226,7 +223,6 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 59b7f7f7a9..b010d19c1b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -164,12 +164,10 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, - identifier_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, - identifier_contracts=identifier_contracts, ) @@ -214,7 +212,6 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - identifier_contracts: list[dict] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -233,7 +230,6 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, - "identifier_contracts": identifier_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index cef2a01047..4290b7847b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3,19 +3,9 @@ import aiohttp import orjson -import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel -from sqlparse.sql import ( - Function, - Identifier, - IdentifierList, - Parenthesis, - Statement, - TokenList, -) -from sqlparse.tokens import DML, Keyword, Name, Punctuation, Whitespace, Wildcard from src.core.engine import ( Engine, @@ -27,470 +17,6 @@ logger = logging.getLogger("wren-ai-service") -class SQLGroundingValidator: - def __init__(self, identifier_contracts: list[dict] | None = None): - self._tables = { - contract["table_name"]: ( - set(contract["columns"]) - if contract.get("columns") is not None - else None - ) - for contract in identifier_contracts or [] - if contract and contract.get("table_name") - } - - def validate(self, sql: str) -> tuple[bool, str]: - if not self._tables: - return True, "" - - try: - statements = sqlparse.parse(sql) - except Exception as exc: - return False, f"Unable to parse generated SQL for schema grounding: {exc}" - - if not statements: - return False, "No SQL statement was generated." - - for statement in statements: - context = _ValidationContext(schema_tables=self._tables) - valid, error = context.validate_statement(statement) - if not valid: - return False, error - - return True, "" - - -class _ValidationContext: - def __init__( - self, - schema_tables: dict[str, set[str] | None], - ctes: dict[str, set[str] | None] | None = None, - ): - self.schema_tables = schema_tables - self.ctes = ctes or {} - - def validate_statement(self, statement: Statement | TokenList) -> tuple[bool, str]: - ctes = dict(self.ctes) - valid, error = self._collect_ctes(statement, ctes) - if not valid: - return False, error - - scope: dict[str, set[str] | None] = {} - source_tokens: set[int] = set() - valid, error = self._collect_sources(statement, ctes, scope, source_tokens) - if not valid: - return False, error - - output_columns = self._infer_select_output_columns(statement, scope) - return self._validate_expressions( - statement, - scope=scope, - source_tokens=source_tokens, - output_columns=output_columns, - ) - - def _collect_ctes( - self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] - ) -> tuple[bool, str]: - tokens = _meaningful_tokens(statement) - if not tokens or not tokens[0].match(Keyword.CTE, "WITH"): - return True, "" - - for token in tokens[1:]: - if token.ttype is DML and token.normalized == "SELECT": - break - - identifiers = ( - list(token.get_identifiers()) - if isinstance(token, IdentifierList) - else [token] - if isinstance(token, Identifier) - else [] - ) - for identifier in identifiers: - cte_name = _identifier_name(identifier) - if not cte_name: - continue - subquery = _subquery_from_identifier(identifier) - if subquery is None: - continue - child_context = _ValidationContext(self.schema_tables, ctes=ctes) - valid, error = child_context.validate_statement(subquery) - if not valid: - return False, error - ctes[cte_name] = child_context._infer_select_output_columns( - subquery, child_context._collect_scope_for_inference(subquery, ctes) - ) - - return True, "" - - def _collect_scope_for_inference( - self, statement: Statement | TokenList, ctes: dict[str, set[str] | None] - ) -> dict[str, set[str] | None]: - scope: dict[str, set[str] | None] = {} - self._collect_sources(statement, ctes, scope, set()) - return scope - - def _collect_sources( - self, - statement: Statement | TokenList, - ctes: dict[str, set[str] | None], - scope: dict[str, set[str] | None], - source_tokens: set[int], - ) -> tuple[bool, str]: - expect_source = False - for token in _meaningful_tokens(statement): - if token.ttype is DML and token.normalized == "SELECT": - expect_source = False - continue - - if _starts_new_clause(token): - expect_source = False - - if token.match( - Keyword, - ( - "FROM", - "JOIN", - "INNER JOIN", - "LEFT JOIN", - "LEFT OUTER JOIN", - "RIGHT JOIN", - "RIGHT OUTER JOIN", - "FULL JOIN", - "FULL OUTER JOIN", - "CROSS JOIN", - ), - ): - expect_source = True - continue - - if not expect_source: - continue - - if isinstance(token, IdentifierList): - for identifier in token.get_identifiers(): - valid, error = self._add_source(identifier, ctes, scope) - if not valid: - return False, error - source_tokens.add(id(identifier)) - continue - - if isinstance(token, Identifier): - valid, error = self._add_source(token, ctes, scope) - if not valid: - return False, error - source_tokens.add(id(token)) - continue - - if isinstance(token, Parenthesis): - subquery = _statement_from_parenthesis(token) - if subquery is not None: - child_context = _ValidationContext(self.schema_tables, ctes=ctes) - valid, error = child_context.validate_statement(subquery) - if not valid: - return False, error - continue - - return True, "" - - def _add_source( - self, - identifier: Identifier, - ctes: dict[str, set[str] | None], - scope: dict[str, set[str] | None], - ) -> tuple[bool, str]: - alias = identifier.get_alias() - subquery = _subquery_from_identifier(identifier) - if subquery is not None: - child_context = _ValidationContext(self.schema_tables, ctes=ctes) - valid, error = child_context.validate_statement(subquery) - if not valid: - return False, error - if alias: - scope[alias] = child_context._infer_select_output_columns( - subquery, child_context._collect_scope_for_inference(subquery, ctes) - ) - return True, "" - - table_name = identifier.get_real_name() - if not table_name: - return True, "" - - if table_name in ctes: - columns = ctes[table_name] - elif table_name in self.schema_tables: - columns = self.schema_tables[table_name] - else: - return ( - False, - "Generated SQL references a table that is not declared in the retrieved schema.", - ) - - scope[alias or table_name] = columns - scope[table_name] = columns - return True, "" - - def _infer_select_output_columns( - self, statement: Statement | TokenList, scope: dict[str, set[str] | None] - ) -> set[str] | None: - columns: set[str] = set() - in_select = False - - for token in _meaningful_tokens(statement): - if token.ttype is DML and token.normalized == "SELECT": - in_select = True - continue - - if in_select and token.match(Keyword, "FROM"): - return columns - - if not in_select: - continue - - if token.ttype is Wildcard: - return None - - identifiers = ( - list(token.get_identifiers()) - if isinstance(token, IdentifierList) - else [token] - if isinstance(token, Identifier) - else [] - ) - for identifier in identifiers: - if any(child.ttype is Wildcard for child in identifier.flatten()): - return None - name = identifier.get_alias() or identifier.get_real_name() - if name: - columns.add(name) - - return columns - - def _validate_expressions( - self, - token: TokenList, - scope: dict[str, set[str] | None], - source_tokens: set[int], - output_columns: set[str] | None, - in_order_by: bool = False, - ) -> tuple[bool, str]: - children = getattr(token, "tokens", []) - next_in_order_by = in_order_by - - for child in children: - if child.is_whitespace or child.ttype in Whitespace: - continue - - if id(child) in source_tokens: - continue - - if child.match(Keyword, "ORDER BY"): - next_in_order_by = True - continue - if _starts_new_clause(child) and not child.match(Keyword, "ORDER BY"): - next_in_order_by = False - - if isinstance(child, IdentifierList): - valid, error = self._validate_expressions( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if isinstance(child, Function): - valid, error = self._validate_function( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if isinstance(child, Identifier): - if any(isinstance(grandchild, Function) for grandchild in child.tokens): - valid, error = self._validate_identifier_children( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - valid, error = self._validate_identifier( - child, scope, output_columns, next_in_order_by - ) - if not valid: - return False, error - valid, error = self._validate_identifier_children( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if isinstance(child, Parenthesis): - subquery = _statement_from_parenthesis(child) - if subquery is not None: - valid, error = _ValidationContext( - self.schema_tables, self.ctes - ).validate_statement(subquery) - if not valid: - return False, error - continue - - if isinstance(child, TokenList): - valid, error = self._validate_expressions( - child, scope, source_tokens, output_columns, next_in_order_by - ) - if not valid: - return False, error - continue - - if child.ttype is Name: - valid, error = self._validate_unqualified_column( - str(child), scope, output_columns, next_in_order_by - ) - if not valid: - return False, error - - return True, "" - - def _validate_function( - self, - function: Function, - scope: dict[str, set[str] | None], - source_tokens: set[int], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - for child in function.tokens: - if isinstance(child, Parenthesis): - return self._validate_expressions( - child, scope, source_tokens, output_columns, in_order_by - ) - return True, "" - - def _validate_identifier_children( - self, - identifier: Identifier, - scope: dict[str, set[str] | None], - source_tokens: set[int], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - for child in identifier.tokens: - if isinstance(child, Function): - return self._validate_function( - child, scope, source_tokens, output_columns, in_order_by - ) - return True, "" - - def _validate_identifier( - self, - identifier: Identifier, - scope: dict[str, set[str] | None], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - parent_name = identifier.get_parent_name() - column_name = identifier.get_real_name() - - if not column_name: - return True, "" - - if parent_name: - if parent_name not in scope: - return ( - False, - "Generated SQL references a table alias that is not declared in the query scope.", - ) - columns = scope[parent_name] - if columns is None or column_name in columns: - return True, "" - return ( - False, - "Generated SQL references a column that is not declared in the retrieved schema.", - ) - - return self._validate_unqualified_column( - column_name, scope, output_columns, in_order_by - ) - - def _validate_unqualified_column( - self, - column_name: str, - scope: dict[str, set[str] | None], - output_columns: set[str] | None, - in_order_by: bool, - ) -> tuple[bool, str]: - if in_order_by and output_columns is not None and column_name in output_columns: - return True, "" - - if column_name in scope: - return True, "" - - for columns in scope.values(): - if columns is None or column_name in columns: - return True, "" - - return ( - False, - "Generated SQL references a column that is not declared in the retrieved schema.", - ) - - -def _meaningful_tokens(token_list: TokenList) -> list: - return [ - token - for token in token_list.tokens - if not token.is_whitespace - and token.ttype not in Whitespace - and token.ttype not in Punctuation - ] - - -def _identifier_name(identifier: Identifier) -> str | None: - return identifier.get_real_name() or identifier.get_name() - - -def _subquery_from_identifier(identifier: Identifier) -> Statement | None: - for token in identifier.tokens: - if isinstance(token, Parenthesis): - return _statement_from_parenthesis(token) - return None - - -def _statement_from_parenthesis(parenthesis: Parenthesis) -> Statement | None: - inner_sql = str(parenthesis)[1:-1].strip() - if not inner_sql: - return None - - parsed = sqlparse.parse(inner_sql) - if parsed and any( - token.ttype is DML and token.normalized == "SELECT" - for token in _meaningful_tokens(parsed[0]) - ): - return parsed[0] - return None - - -def _starts_new_clause(token) -> bool: - return token.match( - Keyword, - ( - "WHERE", - "GROUP BY", - "HAVING", - "ORDER BY", - "LIMIT", - "UNION", - "UNION ALL", - "EXCEPT", - "INTERSECT", - ), - ) - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -508,7 +34,6 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - identifier_contracts: list[dict] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -529,7 +54,6 @@ async def run( allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, allow_data_preview=allow_data_preview, - identifier_contracts=identifier_contracts, ) return { @@ -552,7 +76,6 @@ async def _classify_generation_result( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - identifier_contracts: list[dict] | None = None, ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} @@ -567,18 +90,6 @@ async def _classify_generation_result( "correlation_id": "", } - is_grounded, grounding_error = SQLGroundingValidator( - identifier_contracts - ).validate(generation_result) - if not is_grounded: - return valid_generation_result, { - "sql": generation_result, - "original_sql": generation_result, - "type": "SCHEMA_GROUNDING", - "error": grounding_error, - "correlation_id": "", - } - async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( @@ -607,7 +118,6 @@ async def _classify_generation_result( project_id=project_id, limit=1, dry_run=True, - allow_fallback=allow_dry_plan_fallback, ) addition = addition if isinstance(addition, dict) else {} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 472cd95fea..64f52c5287 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -156,17 +156,6 @@ def _build_metric_ddl(content: dict) -> str: ) -def _build_metric_identifier_contract(content: dict) -> dict: - return { - "table_name": content["name"], - "columns": [ - column["name"] - for column in content["columns"] - if column["data_type"].lower() != "unknown" - ], - } - - def _build_view_ddl(content: dict) -> str: context = _format_semantic_context( { @@ -186,13 +175,6 @@ def _build_view_ddl(content: dict) -> str: ) -def _build_view_identifier_contract(content: dict) -> dict: - return { - "table_name": content["name"], - "columns": None, - } - - def _format_semantic_context(context: dict) -> str: return ( "/*\n" @@ -343,16 +325,6 @@ def _build_table_retrieval_context( return f"{context}{ddl}", has_calculated_field, has_json_field -def _build_table_identifier_contract( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None -) -> dict: - included_columns = _included_columns(content, columns, tables) - return { - "table_name": content["name"], - "columns": [column["name"] for column in included_columns], - } - - ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -616,9 +588,6 @@ def check_using_db_schemas_without_pruning( { "table_name": table_schema["name"], "table_ddl": ddl, - "identifier_contract": _build_table_identifier_contract( - table_schema - ), } ) if _has_calculated_field: @@ -634,7 +603,6 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), - "identifier_contract": _build_metric_identifier_contract(content), } ) has_metric = True @@ -643,7 +611,6 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), - "identifier_contract": _build_view_identifier_contract(content), } ) @@ -754,11 +721,6 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, - "identifier_contract": _build_table_identifier_contract( - table_schema, - columns=columns, - tables=tables, - ), } ) @@ -770,9 +732,6 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), - "identifier_contract": _build_metric_identifier_contract( - content - ), } ) has_metric = True @@ -781,9 +740,6 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), - "identifier_contract": _build_view_identifier_contract( - content - ), } ) diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 941a8776c8..8a4eab6479 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -31,13 +31,11 @@ async def execute_sql( dry_run: bool = True, timeout: float = settings.engine_timeout, limit: int = 500, - allow_fallback: bool = True, **kwargs, ) -> Tuple[bool, Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: data = { "sql": remove_limit_statement(sql), "projectId": project_id, - "allowFallback": allow_fallback, } if dry_run: data["dryRun"] = True diff --git a/wren-ai-service/src/web/v1/routers/sql_corrections.py b/wren-ai-service/src/web/v1/routers/sql_corrections.py index 52592f950e..60dbd9246b 100644 --- a/wren-ai-service/src/web/v1/routers/sql_corrections.py +++ b/wren-ai-service/src/web/v1/routers/sql_corrections.py @@ -19,7 +19,6 @@ class PostRequest(BaseRequest): sql: str error: str - query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = True allow_dry_plan_fallback: bool = False diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 106a3d4e41..ca55061902 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -332,11 +332,6 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -459,7 +454,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -478,7 +472,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -508,7 +501,6 @@ async def ask( invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] current_sql_correction_retries += 1 - sql_diagnosis_reasoning = None self._ask_results[query_id] = AskResultResponse( status="correcting", @@ -521,10 +513,7 @@ async def ask( is_followup=True if histories else False, ) - if ( - allow_sql_diagnosis - and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -559,7 +548,6 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 024dcfc3dd..f1971afa15 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -161,11 +161,6 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] sql_samples = sql_samples_task["formatted_output"].get("documents", []) instructions = instructions_task["formatted_output"].get( "documents", [] @@ -192,7 +187,6 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -222,10 +216,7 @@ async def ask_feedback( trace_id=trace_id, ) - if ( - allow_sql_diagnosis - and failed_dry_run_result["type"] != "SCHEMA_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -260,7 +251,6 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 1e7383bbf8..694d044bfa 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -71,7 +71,7 @@ async def _validate_question( use_dry_plan: bool = True, allow_dry_plan_fallback: bool = False, ): - async def _document_retrieval() -> tuple[list[str], list[dict], bool, bool, bool]: + async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -79,21 +79,10 @@ async def _document_retrieval() -> tuple[list[str], list[dict], bool, bool, bool _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return ( - table_ddls, - identifier_contracts, - has_calculated_field, - has_metric, - has_json_field, - ) + return table_ddls, has_calculated_field, has_metric, has_json_field async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -118,13 +107,7 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - ( - table_ddls, - identifier_contracts, - has_calculated_field, - has_metric, - has_json_field, - ) = _document + table_ddls, has_calculated_field, has_metric, has_json_field = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -154,7 +137,6 @@ async def _instructions_retrieval() -> list[dict]: allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) post_process = generated_sql["post_process"] diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index aab7dd2665..4336186b6f 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -60,7 +60,6 @@ class CorrectionRequest(BaseRequest): event_id: str sql: str error: str - query: Optional[str] = None retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = True allow_dry_plan_fallback: bool = False @@ -77,7 +76,6 @@ async def correct( event_id = request.event_id sql = request.sql error = request.error - query = request.query project_id = request.project_id retrieved_tables = request.retrieved_tables use_dry_plan = request.use_dry_plan @@ -113,21 +111,14 @@ async def correct( .get("retrieval_results", []) ) table_ddls = [document.get("table_ddl") for document in documents] - identifier_contracts = [ - document.get("identifier_contract") - for document in documents - if document.get("identifier_contract") - ] res = await self._pipelines["sql_correction"].run( contexts=table_ddls, - query=query, invalid_generation_result=_invalid, project_id=project_id, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - identifier_contracts=identifier_contracts, ) post_process = res["post_process"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 86369e2ab1..87bf9b09de 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -4,7 +4,6 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, - SQLGroundingValidator, construct_ask_history_messages, construct_instructions, get_json_field_instructions, @@ -78,99 +77,6 @@ async def execute_sql(self, *args, **kwargs): return False, {}, {"error_message": "dry run failed"} -def test_sql_grounding_validator_accepts_retrieved_schema_identifiers(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_dimension", "retrieved_measure"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'SELECT t."retrieved_dimension", SUM(t."retrieved_measure") AS result_label ' - 'FROM "retrieved_model" t ' - 'WHERE t."retrieved_dimension" = \'filter_value\' ' - 'GROUP BY t."retrieved_dimension" ' - "ORDER BY result_label" - ) - - assert valid is True - assert error == "" - - -def test_sql_grounding_validator_rejects_unretrieved_table(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'SELECT * FROM "unretrieved_model" WHERE "unretrieved_column" = \'filter_value\'' - ) - - assert valid is False - assert "table that is not declared" in error - - -def test_sql_grounding_validator_rejects_unretrieved_column(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'SELECT * FROM "retrieved_model" WHERE "unretrieved_column" = \'filter_value\'' - ) - - assert valid is False - assert "column that is not declared" in error - - -def test_sql_grounding_validator_handles_cte_output_columns(): - contracts = [ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ] - - valid, error = SQLGroundingValidator(contracts).validate( - 'WITH scoped_result AS (SELECT "retrieved_column" AS result_column FROM "retrieved_model") ' - "SELECT * FROM scoped_result WHERE result_column = 'filter_value'" - ) - - assert valid is True - assert error == "" - - -@pytest.mark.asyncio -async def test_sql_postprocessor_rejects_ungrounded_sql_before_engine_validation(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=[ - '{"sql": "SELECT * FROM \\"unretrieved_model\\" WHERE \\"unretrieved_column\\" = \'filter_value\'"}' - ], - use_dry_plan=True, - data_source="source", - identifier_contracts=[ - { - "table_name": "retrieved_model", - "columns": ["retrieved_column"], - } - ], - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "SCHEMA_GROUNDING" - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] - - @pytest.mark.asyncio async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): result = await SQLGenPostProcessor(_FailingDryPlanEngine()).run( diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 0e0f5673be..2141eb0a02 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -69,45 +69,6 @@ def test_column_pruning_prompt_uses_current_query_without_history_text(): assert "previous request" not in result["prompt"] -def test_retrieval_results_include_exact_identifier_contract(): - result = check_using_db_schemas_without_pruning( - construct_db_schemas=[ - { - "type": "TABLE", - "name": "retrieved_model", - "comment": "", - "columns": [ - { - "type": "COLUMN", - "name": "retrieved_column_a", - "data_type": "VARCHAR", - "comment": "", - "is_primary_key": False, - }, - { - "type": "COLUMN", - "name": "retrieved_column_b", - "data_type": "VARCHAR", - "comment": "", - "is_primary_key": False, - }, - ], - "properties": {}, - "primaryKey": "", - } - ], - dbschema_retrieval=[], - encoding=type("Encoding", (), {"encode": lambda self, text: [1]})(), - enable_column_pruning=False, - context_window_size=100, - ) - - assert result["db_schemas"][0]["identifier_contract"] == { - "table_name": "retrieved_model", - "columns": ["retrieved_column_a", "retrieved_column_b"], - } - - def test_table_selection_prompt_keeps_multiple_relevant_datasets(): assert "same business concept is represented by multiple modeled datasets" in ( table_columns_selection_system_prompt diff --git a/wren-engine b/wren-engine index abc150f8f7..44d0811961 160000 --- a/wren-engine +++ b/wren-engine @@ -1 +1 @@ -Subproject commit abc150f8f703b04bfc1a63df70e50e6e81218dc2 +Subproject commit 44d08119612dca9c8ff007fa12b305fd1aad7593 diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index d04ed87cfa..f3da8f04fd 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -1142,7 +1142,6 @@ export type PreviewItemSqlInput = { }; export type PreviewSqlDataInput = { - allowFallback?: InputMaybe; dryRun?: InputMaybe; limit?: InputMaybe; projectId?: InputMaybe; diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 07b8b9b098..5d421cd975 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -186,7 +186,6 @@ export interface IbisQueryOptions extends IbisBaseOptions { limit?: number; refresh?: boolean; cacheEnabled?: boolean; - allowFallback?: boolean; } export interface IbisDryPlanOptions { dataSource: DataSourceName; @@ -318,10 +317,6 @@ export class IbisAdaptor implements IIbisAdaptor { params: { limit: options.limit || DEFAULT_PREVIEW_LIMIT, }, - headers: { - 'x-wren-fallback_disable': - options.allowFallback === false ? 'true' : 'false', - }, }, ); return { @@ -362,12 +357,6 @@ export class IbisAdaptor implements IIbisAdaptor { const response = await axios.post( `${this.ibisServerEndpoint}/${this.getIbisApiVersion(IBIS_API_TYPE.DRY_RUN)}/connector/${dataSourceUrlMap[dataSource]}/query?dryRun=true`, body, - { - headers: { - 'x-wren-fallback_disable': - options.allowFallback === false ? 'true' : 'false', - }, - }, ); logger.debug(`Ibis server Dry run success`); return { diff --git a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts index c51fdeca2e..44bf4dc5ee 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts @@ -582,9 +582,6 @@ describe('IbisAdaptor', () => { params: { limit: DEFAULT_PREVIEW_LIMIT, }, - headers: { - 'x-wren-fallback_disable': 'false', - }, }, ); }); @@ -620,9 +617,6 @@ describe('IbisAdaptor', () => { params: { limit: customLimit, }, - headers: { - 'x-wren-fallback_disable': 'false', - }, }, ); }); @@ -684,35 +678,6 @@ describe('IbisAdaptor', () => { expect(res.processTime).toEqual('1s'); }); - it('should disable v3 fallback during dry run when requested', async () => { - mockedAxios.post.mockResolvedValue({ - headers: { - 'x-correlation-id': '123', - 'x-process-time': '1s', - }, - }); - mockedEncryptor.prototype.decrypt.mockReturnValue( - JSON.stringify({ password: mockPostgresConnectionInfo.password }), - ); - - await ibisAdaptor.dryRun('SELECT * FROM test_table', { - dataSource: DataSourceName.POSTGRES, - connectionInfo: mockPostgresConnectionInfo, - mdl: mockManifest, - allowFallback: false, - }); - - expect(mockedAxios.post).toHaveBeenCalledWith( - expect.any(String), - expect.any(Object), - { - headers: { - 'x-wren-fallback_disable': 'true', - }, - }, - ); - }); - it('should throw an exception with correlationId and processTime when dry run fails', async () => { const mockError = { response: { diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index 7b6aa3ba22..42262b3626 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -100,7 +100,6 @@ export interface PreviewSQLData { projectId?: string; limit?: number; dryRun?: boolean; - allowFallback?: boolean; } export interface DryPlanSQLData { diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 6eee24d818..ce98d0f4db 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -1336,7 +1336,7 @@ export class ModelResolver { args: { data: PreviewSQLData }, ctx: IContext, ) { - const { sql, projectId, limit, dryRun, allowFallback } = args.data; + const { sql, projectId, limit, dryRun } = args.data; const project = projectId ? await ctx.projectService.getProjectById(parseInt(projectId)) : await ctx.projectService.getCurrentProject(); @@ -1347,7 +1347,6 @@ export class ModelResolver { modelingOnly: false, manifest, dryRun, - allowFallback, }); } diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 23c7d0c7f3..90fd686c77 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -925,7 +925,6 @@ export const typeDefs = gql` projectId: String limit: Int dryRun: Boolean - allowFallback: Boolean } input DryPlanSQLDataInput { diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 24808037d9..ee7b341fdd 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -42,7 +42,6 @@ export interface PreviewOptions { manifest: Manifest; limit?: number; dryRun?: boolean; - allowFallback?: boolean; refresh?: boolean; cacheEnabled?: boolean; } @@ -105,7 +104,6 @@ export class QueryService implements IQueryService { manifest: mdl, limit, dryRun, - allowFallback, refresh, cacheEnabled, } = options; @@ -136,7 +134,6 @@ export class QueryService implements IQueryService { dataSource, connectionInfo, mdl, - allowFallback, ); } else { return await this.ibisQuery( @@ -147,7 +144,6 @@ export class QueryService implements IQueryService { limit, refresh, cacheEnabled, - allowFallback, ); } } @@ -206,7 +202,6 @@ export class QueryService implements IQueryService { dataSource: DataSourceName, connectionInfo: any, mdl: Manifest, - allowFallback?: boolean, ): Promise { const event = TelemetryEvent.IBIS_DRY_RUN; try { @@ -214,7 +209,6 @@ export class QueryService implements IQueryService { dataSource, connectionInfo, mdl, - allowFallback, }); this.sendIbisEvent(event, res, { dataSource, sql }); return { @@ -237,7 +231,6 @@ export class QueryService implements IQueryService { limit: number, refresh?: boolean, cacheEnabled?: boolean, - allowFallback?: boolean, ): Promise { const event = TelemetryEvent.IBIS_QUERY; try { @@ -248,7 +241,6 @@ export class QueryService implements IQueryService { limit, refresh, cacheEnabled, - allowFallback, }); this.sendIbisEvent(event, res, { dataSource, From 611a6225572db22cf0466cfe7410dd0885c41a0a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 18:47:41 +0530 Subject: [PATCH 0760/1087] Clear project semantics before rebuilding indexes --- .../src/pipelines/indexing/instructions.py | 27 +++++++++++------ .../web/v1/services/semantics_preparation.py | 29 ++++++++++++++----- 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 b23f3cf2ab..006de9237f 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -61,20 +61,25 @@ def __init__(self, instructions_store: DocumentStore) -> None: @component.output_types() async def run( - self, instruction_ids: List[str], project_id: Optional[str] = None + self, + instruction_ids: List[str], + project_id: Optional[str] = None, + delete_all: bool = False, ) -> None: - filter = { - "operator": "AND", - "conditions": [ - {"field": "instruction_id", "operator": "in", "value": instruction_ids}, - ], - } + conditions = [] + + if not delete_all: + conditions.append( + {"field": "instruction_id", "operator": "in", "value": instruction_ids} + ) if project_id: - filter["conditions"].append( + conditions.append( {"field": "project_id", "operator": "==", "value": project_id} ) + filter = {"operator": "AND", "conditions": conditions} if conditions else None + return await self.store.delete_documents(filter) @@ -108,7 +113,11 @@ async def clean( ) -> 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) + await cleaner.run( + instruction_ids=instruction_ids, + project_id=project_id, + delete_all=delete_all, + ) return embedding 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 2ff6215cbe..327630203e 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -77,16 +77,31 @@ async def prepare_semantics( "mdl_str": prepare_semantics_request.mdl, "project_id": prepare_semantics_request.project_id, } + project_scoped_index_names = [ + "db_schema", + "historical_question", + "table_description", + "sql_pairs", + "project_meta", + ] + + await asyncio.gather( + *[ + self._pipelines[name].clean( + project_id=prepare_semantics_request.project_id, + delete_all=True, + ) + if name == "sql_pairs" + else self._pipelines[name].clean( + project_id=prepare_semantics_request.project_id + ) + for name in project_scoped_index_names + ] + ) tasks = [ self._pipelines[name].run(**input) - for name in [ - "db_schema", - "historical_question", - "table_description", - "sql_pairs", - "project_meta", - ] + for name in project_scoped_index_names ] await asyncio.gather(*tasks) From ba865f064c86952c551aebc9750c790663d0f461 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 20:01:26 +0530 Subject: [PATCH 0761/1087] Export physical column source metadata --- wren-ui/src/apollo/server/mdl/mdlBuilder.ts | 3 +++ wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts | 6 ++++++ wren-ui/src/apollo/server/mdl/type.ts | 1 + wren-ui/src/apollo/server/types/manifest.ts | 1 + 4 files changed, 11 insertions(+) diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 4f7a87671a..95865173e4 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -200,6 +200,9 @@ export class MDLBuilder implements IMDLBuilder { if (column.displayName) { properties.displayName = column.displayName; } + if (column.sourceColumnName) { + properties.sourceColumnName = column.sourceColumnName; + } // put nested columns in properties if (column.type.includes('STRUCT')) { const nestedColumns = this.nestedColumns.filter( diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 8a950f1931..d171ea56da 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -707,6 +707,12 @@ describe('MDLBuilder', () => { expect(manifest.models[0].refSql).toEqual( 'SELECT "Source Name" AS "SourceName", "Source Type" AS "SourceType" FROM "physical_catalog"."physical_schema"."physical_table"', ); + expect(manifest.models[0].columns[0].properties.sourceColumnName).toEqual( + 'Source Name', + ); + expect(manifest.models[0].columns[1].properties.sourceColumnName).toEqual( + 'Source Type', + ); }); it('should preserve refSql when a model has no tableReference.', () => { diff --git a/wren-ui/src/apollo/server/mdl/type.ts b/wren-ui/src/apollo/server/mdl/type.ts index bb7c8be10f..fc2bea59fe 100644 --- a/wren-ui/src/apollo/server/mdl/type.ts +++ b/wren-ui/src/apollo/server/mdl/type.ts @@ -7,6 +7,7 @@ export interface ColumnMDL { properties?: { description?: string; // eg: "the key of each order" displayName?: string; // eg: "Order Key" + sourceColumnName?: string; }; expression?: string; // eg: "SUM(orders.totalprice)" } diff --git a/wren-ui/src/apollo/server/types/manifest.ts b/wren-ui/src/apollo/server/types/manifest.ts index e96a55d600..50320b644b 100644 --- a/wren-ui/src/apollo/server/types/manifest.ts +++ b/wren-ui/src/apollo/server/types/manifest.ts @@ -31,6 +31,7 @@ export interface Measure { export interface CumulativeMetricProperties { description?: string; + sourceColumnName?: string; } export interface Window { From a4e58637fe3e5cbbc43a23acd7e2217699feae54 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 21:07:54 +0530 Subject: [PATCH 0762/1087] Improve schema retrieval grounding for asks --- .../src/pipelines/generation/sql_answer.py | 3 + .../retrieval/db_schema_retrieval.py | 7 +- .../pipelines/generation/test_sql_utils.py | 10 +++ .../retrieval/test_db_schema_retrieval.py | 67 +++++++++++++++++++ 4 files changed, 82 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 81289081b5..a211c925d8 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -32,6 +32,9 @@ 6. Answer must be in the same language user specified. 7. Do not include ```markdown or ``` in the answer. 8. If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. +9. Use only the columns and rows provided in Data. Do not invent, duplicate, reorder, aggregate, rank, or label rows unless that operation is directly represented by the provided SQL result. +10. If the Data has aggregate rows, summarize those exact aggregate rows instead of describing them as separate top examples. +11. If the Data is empty, state that no matching records were returned. ### OUTPUT FORMAT diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 64f52c5287..e22ed2b479 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -374,14 +374,11 @@ async def dbschema_retrieval( table_retrieval.get("documents", []) ) documents = [] - if embedding: + if embedding and not table_names: documents = await _retrieve_semantic_schema_documents( embedding, project_id, dbschema_retriever ) - table_names = _merge_names( - table_names, - _table_names_from_schema_documents(documents), - ) + table_names = _table_names_from_schema_documents(documents) if table_names: retrieved_table_names = set() diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 87bf9b09de..d2634bd878 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -19,6 +19,7 @@ get_sql_correction_system_prompt, sql_correction_user_prompt_template, ) +from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template from src.pipelines.generation.sql_regeneration import get_sql_regeneration_system_prompt from src.pipelines.generation.sql_regeneration import sql_regeneration_user_prompt_template @@ -304,6 +305,15 @@ def test_followup_sql_prompt_does_not_expect_previous_sql_context(): ) +def test_sql_answer_prompt_uses_only_returned_data_rows(): + assert "Use only the columns and rows provided in Data" in sql_to_answer_system_prompt + assert "Do not invent, duplicate, reorder, aggregate, rank, or label rows" in ( + sql_to_answer_system_prompt + ) + assert "summarize those exact aggregate rows" in sql_to_answer_system_prompt + assert "If the Data is empty" in sql_to_answer_system_prompt + + def test_executable_prompt_templates_omit_untrusted_reasoning_and_sql_context(): reasoning_marker = "UNTRUSTED_REASONING_CONTEXT_MARKER" diagnostic_marker = "UNTRUSTED_DIAGNOSTIC_CONTEXT_MARKER" diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 2141eb0a02..65ccfb396b 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -398,6 +398,73 @@ async def run(self, query_embedding, filters): ] +@pytest.mark.asyncio +async def test_dbschema_retrieval_prefers_table_description_hits_over_schema_chunk_hits(): + described_model = "described_dataset" + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": described_model, + "comment": "", + "columns": [], + "properties": {}, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": described_model}, + ) + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={ + "documents": [ + Document( + content=str({"name": described_model}), + meta={"type": "TABLE_DESCRIPTION", "name": described_model}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + embedding={"embedding": [0.25]}, + ) + + assert [call["query_embedding"] for call in retriever.calls] == [[]] + assert retriever.calls[0]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": described_model}, + ], + }, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert [document.meta["name"] for document in documents] == [described_model] + + def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): class Encoding: def encode(self, value): From ba8cf833a79daaf8fd53cff0c9a3eeaffde1e13e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 21:31:50 +0530 Subject: [PATCH 0763/1087] Strengthen Wren SQL grounding rules --- .../src/pipelines/generation/sql_correction.py | 5 +++-- wren-ai-service/src/pipelines/generation/utils/sql.py | 4 ++++ .../tests/pytest/pipelines/generation/test_sql_utils.py | 8 ++++++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f28d612aaa..d602494f75 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,12 +30,12 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. +You are a Wren SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. ### SQL CORRECTION INSTRUCTIONS ### 1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. -2. Then, generate a syntactically correct ANSI SQL query from the user's intent and the current DATABASE SCHEMA. +2. Then, generate a syntactically correct Wren SQL query from the user's intent and the current DATABASE SCHEMA. 3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. 4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. 5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. @@ -43,6 +43,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. 8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. 9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. +10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. ### SQL RULES ### Make sure you follow the SQL Rules strictly. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4290b7847b..ddc4542ca3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -247,6 +247,10 @@ async def _classify_generation_result( - Use Wren SQL identifier quoting with double quotes only; the engine rewrite step converts grounded Wren SQL to the active connector dialect. - Put single quotes around string literals. - Never quote numeric literals. +- Generate Wren SQL syntax only, not connector-specific SQL syntax. +- Never use SELECT TOP, TOP(...), FETCH FIRST, square-bracket identifiers, or backtick identifiers. For top or limit requests, sort with ORDER BY and put LIMIT at the end of the query. +- Preserve every deployed table and column identifier exactly as it appears in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, including spaces, digits, underscores, case, and punctuation, then wrap that exact identifier in double quotes in SQL. +- Do not convert deployed identifiers into display-friendly variants by replacing spaces with underscores, removing prefixes, changing case, shortening names, or expanding abbreviations. - For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. - For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only date/time functions and casts whose exact syntax is provided in SQL FUNCTIONS for this request. - If the question asks for a specific or relative date, generate a bounded date/time filter only when both the exact date/time schema column and required SQL FUNCTIONS-supported operation are available. If either is missing, do not invent a field or function. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index d2634bd878..6458aba8df 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -175,6 +175,11 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "do not translate it into a generic object name" in rules assert "return null for sql instead of producing an approximate query" in rules assert "A retrieved object is usable only when" in rules + assert "Generate Wren SQL syntax only" in rules + assert "Never use SELECT TOP" in rules + assert "square-bracket identifiers" in rules + assert "Preserve every deployed table and column identifier exactly" in rules + assert "Do not convert deployed identifiers into display-friendly variants" in rules def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): @@ -263,12 +268,15 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "treat it as the source of intent" in prompt assert "Do not copy placeholders" in prompt assert "Regenerate a grounded Wren SQL query" in prompt + assert "Wren SQL expert" in prompt + assert "syntactically correct Wren SQL query" in prompt assert ( "Do not preserve a table, column, join, filter, grouping, ordering, or function" in prompt ) assert "Treat physical/source/lineage names from the failed SQL" in prompt assert "do not try a similar replacement from source metadata" in prompt + assert "connector-specific syntax" in prompt def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): From fc7123c194a34702d7a3252d780e8c397e2f99e1 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 21:38:25 +0530 Subject: [PATCH 0764/1087] Prevent LLM from assuming SQL identifiers --- .../generation/followup_sql_generation.py | 2 +- .../pipelines/generation/sql_correction.py | 6 ++-- .../pipelines/generation/sql_generation.py | 2 +- .../pipelines/generation/sql_regeneration.py | 4 +-- .../src/pipelines/generation/utils/sql.py | 17 ++++++----- .../src/pipelines/indexing/db_schema.py | 13 +++++++++ .../retrieval/db_schema_retrieval.py | 29 +++++++++++++++++-- .../pipelines/generation/test_sql_utils.py | 11 +++++-- .../pipelines/indexing/test_db_schema.py | 2 ++ .../retrieval/test_db_schema_retrieval.py | 25 ++++++++++++++++ 10 files changed, 91 insertions(+), 20 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 5d4edb01b6..aa419bbcd0 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -76,7 +76,7 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index d602494f75..50fe59ba51 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -42,7 +42,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. 7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. 8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. +9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA. If the unsupported part is needed to answer the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql instead of substituting non-schema identifiers. 10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. ### SQL RULES ### @@ -51,7 +51,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS. If no fully grounded SQL can be generated, return null for sql. +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. {{ "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" @@ -84,7 +84,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### QUESTION ### {% if query %} User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. {% endif %} ### FAILED SQL ### The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 2d5f1a2b96..7818197640 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -70,7 +70,7 @@ ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index b010d19c1b..70a7b00e21 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -44,7 +44,7 @@ def get_sql_regeneration_system_prompt( {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS. If no fully grounded SQL can be generated, return null for sql. +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. {{ "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" @@ -95,7 +95,7 @@ def get_sql_regeneration_system_prompt( ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, or inferred names into executable SQL. If a needed table, column, relation, date field, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, omit that unsupported part instead of inventing or substituting a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### ORIGINAL SQL QUERY ### The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ddc4542ca3..28348363d7 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -210,8 +210,9 @@ async def _classify_generation_result( - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. -- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, use the closest executable object and column whose semantic metadata supports the intent, or omit that unsupported concept. -- If a requested concept, filter, sort, join, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. Generate the closest valid SQL using only available schema fields. +- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. +- Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. +- If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. - When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. @@ -225,10 +226,10 @@ async def _classify_generation_result( - Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. - If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. -- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part or answer with the closest valid SQL over grounded fields only. +- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. - If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. -- If a requested noun, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. -- If the user's primary requested subject, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query over unrelated schema objects. +- If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. +- If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. """ @@ -265,7 +266,7 @@ async def _classify_generation_result( - DON'T USE "EXTRACT()" function with INTERVAL data types as arguments - DON'T USE INTERVAL or generate INTERVAL-like expression in the generated SQL query. - DON'T USE "TO_CHAR" function in the generated SQL query. -- DON'T USE unsupported statistical, date/time, or formatting functions. If SQL FUNCTIONS does not list a function needed by the intent, answer with the closest supported aggregation/filter over exact schema fields. +- DON'T USE unsupported statistical, date/time, or formatting functions. If SQL FUNCTIONS does not list a function needed by the requested intent, omit the function-dependent part. If that function is required to answer the request, return null for sql. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. - For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. @@ -453,7 +454,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. 7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 8. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. -9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element or use the closest grounded expression. +9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. 10. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. 11. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. 12. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. @@ -461,7 +462,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. If the retrieved schema does not ground the requested intent, return null for sql. +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. Do not create table or column identifiers from the user's wording. If the retrieved schema does not ground the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. {{ "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index f88cefe4b8..97cf6110f0 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -298,6 +298,18 @@ def _column_batch( ] def _convert_views(self, views: List[Dict[str, Any]]) -> List[Dict[str, str]]: + def _columns(view: Dict[str, Any]) -> List[dict]: + properties = view.get("properties", {}) or {} + return [ + { + "name": column.get("name", ""), + "data_type": column.get("type", ""), + "comment": column.get("description", ""), + } + for column in properties.get("columns", []) + if column.get("name") + ] + def _payload(view: Dict[str, Any]) -> dict: return { "type": "VIEW", @@ -306,6 +318,7 @@ def _payload(view: Dict[str, Any]) -> dict: else "", "name": view["name"], "statement": view["statement"], + "columns": _columns(view), } return [ diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e22ed2b479..348dac8adf 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -157,21 +157,46 @@ def _build_metric_ddl(content: dict) -> str: def _build_view_ddl(content: dict) -> str: + columns = [ + column + for column in content.get("columns", []) + if column.get("name") and column.get("data_type", "").lower() != "unknown" + ] context = _format_semantic_context( { "object_type": "view", "sql_identifier_contract": { "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in columns + ], }, "semantic_context_not_sql_identifiers": { "role": "stable virtual table interface", "description": content["comment"], - "definition_is_semantic_context": True, + "definition_omitted_from_executable_schema": True, }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type( + column.get("data_type") + ), + "semantic_context_not_sql_identifier": column.get("comment", ""), + } + for column in columns + ], } ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" + for column in columns + ] + return ( - f"{context}{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" + f"{context}CREATE TABLE {content['name']} (\n " + + ",\n ".join(columns_ddl) + + "\n);" ) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6458aba8df..106a09d3ba 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -158,6 +158,7 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "semantic_context_not_sql_identifier" in rules assert "Do not combine words, labels, ordinals" in rules assert "Never generate placeholder identifiers" in rules + assert "Never create an identifier from user question wording" in rules assert "use all required related tables" in rules assert "silently check that each identifier and function" in rules assert "instead of inventing a replacement" in rules @@ -174,6 +175,7 @@ def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): assert "independently valid from DATABASE SCHEMA" in rules assert "do not translate it into a generic object name" in rules assert "return null for sql instead of producing an approximate query" in rules + assert "If that field is required to answer the request, return null for sql" in rules assert "A retrieved object is usable only when" in rules assert "Generate Wren SQL syntax only" in rules assert "Never use SELECT TOP" in rules @@ -221,7 +223,8 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "use a normal equality or LIKE comparison" in prompt assert "unless the user explicitly asks for rank values" in prompt assert "perform a silent grounding check" in prompt - assert "closest grounded expression" in prompt + assert "return null for sql" in prompt + assert "Do not create table or column identifiers from the user's wording" in prompt assert "DATABASE SCHEMA is the only source of executable identifiers" in prompt assert "reasoning plan as semantic context for intent only" in prompt assert "Do not copy identifiers, functions, literal values" in prompt @@ -234,7 +237,7 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "Use Wren SQL identifier quoting with double quotes only" in prompt assert "source database/schema/table names" in prompt assert "appears only in SQL samples, failed SQL" in prompt - assert "retrieved schema does not ground the requested intent" in prompt + assert "retrieved schema does not ground the requested subject" in prompt assert "" not in prompt @@ -277,6 +280,7 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "Treat physical/source/lineage names from the failed SQL" in prompt assert "do not try a similar replacement from source metadata" in prompt assert "connector-specific syntax" in prompt + assert "return null for sql instead of substituting non-schema identifiers" in prompt def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): @@ -301,8 +305,9 @@ def test_user_prompt_templates_keep_source_metadata_non_executable(): sql_correction_user_prompt_template, ): assert "source/physical/lineage names" in prompt - assert "omit that unsupported part instead of inventing" in prompt + assert "return null for sql instead of inventing" in prompt assert "exact declared table and column names from DATABASE SCHEMA" in prompt + assert "user question words" in prompt assert "return null for sql instead of querying an unrelated object" in prompt diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py index f2b15cef65..da7add8e8f 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py @@ -571,6 +571,7 @@ async def test_view(): "comment": "", "name": "view_1", "statement": "SELECT * FROM user", + "columns": [], } ) @@ -601,6 +602,7 @@ async def test_view_with_properties(): "comment": "/* {'description': 'A view containing user information.'} */\n", "name": "view_1", "statement": "SELECT * FROM user", + "columns": [], } ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 65ccfb396b..cca9daa46f 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -4,6 +4,7 @@ from src.pipelines.common import build_table_ddl from src.pipelines.retrieval.db_schema_retrieval import ( + _build_view_ddl, check_using_db_schemas_without_pruning, construct_retrieval_results, dbschema_retrieval, @@ -78,6 +79,30 @@ def test_table_selection_prompt_keeps_multiple_relevant_datasets(): ) +def test_view_schema_context_uses_declared_view_columns_not_view_definition(): + result = _build_view_ddl( + { + "type": "VIEW", + "comment": "Semantic description.", + "name": "retrieved_view", + "statement": "NON_EXECUTABLE_DEFINITION_TOKEN", + "columns": [ + { + "name": "visible_attribute", + "data_type": "VARCHAR", + "comment": "Semantic field.", + } + ], + } + ) + + assert "CREATE TABLE retrieved_view" in result + assert "visible_attribute VARCHAR" in result + assert "sql_column_names_use_exactly" in result + assert "definition_omitted_from_executable_schema" in result + assert "NON_EXECUTABLE_DEFINITION_TOKEN" not in result + + @pytest.mark.asyncio async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: From a86b1d27faa07021a223dd19ee9a9b08a9136c01 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 21:55:52 +0530 Subject: [PATCH 0765/1087] Omit failed SQL diagnostics from correction prompt --- .../src/pipelines/generation/sql_correction.py | 9 ++------- wren-ai-service/src/pipelines/generation/utils/sql.py | 2 ++ .../pytest/pipelines/generation/test_sql_utils.py | 11 ++++++++--- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 50fe59ba51..dee73259cc 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -90,14 +90,9 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. ### DRY-RUN DIAGNOSTIC ### -{% if invalid_generation_result and invalid_generation_result.error %} -Diagnostic text: -{{ invalid_generation_result.error }} -{% else %} -No diagnostic text was provided. -{% endif %} +The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. -Use the diagnostic text only to understand the failure category. Do not copy identifiers, literal values, functions, SQL snippets, physical names, source names, or replacement candidates from the diagnostic text. Regenerate from the user question and current DATABASE SCHEMA. +Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 28348363d7..c7f35c30bf 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -381,6 +381,8 @@ async def _classify_generation_result( 21. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. 22. The reasoning plan is semantic context for intent only, not a source of executable identifiers. SQL generation must re-read DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT before using any identifier. 23. ONLY SHOWING the reasoning plan in bullet points. +24. Do not use the words "assume", "assuming", "likely", "possible", "might", or "example" when describing tables, columns, filters, or SQL. +25. If exact deployed table and column identifiers are not available for a requested part, say only that the retrieved metadata does not support that part. Do not propose a replacement name. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 106a09d3ba..f8a359f9f6 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -293,6 +293,9 @@ def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): assert "reasoning plan is semantic context for intent only" in prompt assert "declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT" in prompt assert "source names, physical names, lineage names" in prompt + assert 'Do not use the words "assume", "assuming", "likely"' in prompt + assert "retrieved metadata does not support that part" in prompt + assert "Do not propose a replacement name" in prompt assert "" not in prompt assert "" not in prompt @@ -395,8 +398,10 @@ def test_executable_prompt_templates_omit_untrusted_reasoning_and_sql_context(): assert original_sql_marker not in regeneration_prompt assert "original SQL is intentionally omitted" in regeneration_prompt - assert diagnostic_marker in correction_prompt - assert "Use the diagnostic text only to understand the failure category" in ( + assert diagnostic_marker not in correction_prompt + assert "dry-run diagnostic text is intentionally omitted" in ( + correction_prompt + ) + assert "Regenerate from the user question and current DATABASE SCHEMA only" in ( correction_prompt ) - assert "Do not copy identifiers" in correction_prompt From f6189b76d385cbcc2ff228518d255bbb2c8f898c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 22:03:40 +0530 Subject: [PATCH 0766/1087] Require exact schema identifiers for SQL generation --- .../src/pipelines/generation/followup_sql_generation.py | 2 ++ wren-ai-service/src/pipelines/generation/sql_correction.py | 1 + wren-ai-service/src/pipelines/generation/sql_generation.py | 2 ++ wren-ai-service/src/pipelines/generation/utils/sql.py | 5 ++++- .../tests/pytest/pipelines/generation/test_sql_utils.py | 4 ++++ 5 files changed, 13 insertions(+), 1 deletion(-) 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 aa419bbcd0..1bcb3f853d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -77,6 +77,8 @@ ### QUESTION ### User's Follow-up Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index dee73259cc..e92b8fc648 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -85,6 +85,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% if query %} User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. {% endif %} ### FAILED SQL ### The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 7818197640..2640082430 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -71,6 +71,8 @@ ### QUESTION ### User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c7f35c30bf..5473f0d9e0 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -383,6 +383,8 @@ async def _classify_generation_result( 23. ONLY SHOWING the reasoning plan in bullet points. 24. Do not use the words "assume", "assuming", "likely", "possible", "might", or "example" when describing tables, columns, filters, or SQL. 25. If exact deployed table and column identifiers are not available for a requested part, say only that the retrieved metadata does not support that part. Do not propose a replacement name. +26. Do not write table names or column names from the user's wording unless the same identifier appears exactly in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +27. Do not include code blocks, inline SQL fragments, SELECT statements, WHERE clauses, join clauses, or any query-shaped text in the reasoning plan. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -459,7 +461,8 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. 10. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. 11. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. -12. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +12. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. +13. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index f8a359f9f6..e86b2f5933 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -238,6 +238,8 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "source database/schema/table names" in prompt assert "appears only in SQL samples, failed SQL" in prompt assert "retrieved schema does not ground the requested subject" in prompt + assert "If any planned SQL identifier cannot be copied exactly" in prompt + assert "Never create a table or column from the user's wording" in prompt assert "" not in prompt @@ -296,6 +298,8 @@ def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): assert 'Do not use the words "assume", "assuming", "likely"' in prompt assert "retrieved metadata does not support that part" in prompt assert "Do not propose a replacement name" in prompt + assert "Do not write table names or column names from the user's wording" in prompt + assert "Do not include code blocks, inline SQL fragments" in prompt assert "" not in prompt assert "" not in prompt From 2832802d326565ef84b9f8084c8351f17d5c358b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Thu, 30 Jul 2026 22:18:46 +0530 Subject: [PATCH 0767/1087] Expose executable identifier catalog in schema context --- .../src/pipelines/generation/utils/sql.py | 15 +++---- .../retrieval/db_schema_retrieval.py | 39 +++++++++++++++++++ .../pipelines/generation/test_sql_utils.py | 2 + .../retrieval/test_db_schema_retrieval.py | 8 ++++ 4 files changed, 57 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5473f0d9e0..1633c399e4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -456,13 +456,14 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. 5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. -7. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -8. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. -9. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. -10. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -11. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. -12. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. -13. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. +8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. +11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +13. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. +14. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 348dac8adf..759c99e84c 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -209,9 +209,48 @@ def _format_semantic_context(context: dict) -> str: "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" "*/\n" + f"{_format_executable_identifier_catalog(context)}" ) +def _format_executable_identifier_catalog(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "### EXECUTABLE WREN IDENTIFIER CATALOG ###", + "Copy SQL identifiers only from this catalog or the following DDL.", + "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"table: {table_name}") + if column_names: + lines.append("columns:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationships:") + lines.extend(f"- {constraint}" for constraint in relationship_constraints) + lines.extend( + [ + "If a needed table, column, or relationship is not listed here or declared in the following DDL, return null for sql.", + "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", + "", + ] + ) + return "\n".join(lines) + + def _format_identifier_contract(context: dict) -> str: contract = context.get("sql_identifier_contract", {}) table_name = contract.get("sql_table_name_use_exactly") diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index e86b2f5933..0f4ed59d26 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -231,6 +231,8 @@ def test_sql_generation_system_prompt_grounding_contract(): assert "include those objects only when DATABASE SCHEMA shows" in prompt assert "Use the exact supported syntax shown there" in prompt assert "WREN SQL IDENTIFIER CONTRACT" in prompt + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in prompt + assert "first and clearest list of allowed executable identifiers" in prompt assert "Use sql_table_name_use_exactly" in prompt assert "sql_column_names_use_exactly" in prompt assert "semantic_context_not_sql_identifiers" in prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index cca9daa46f..73058b9a02 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -756,6 +756,10 @@ def table_schema(name): "WREN SQL IDENTIFIER CONTRACT" in schema["table_ddl"] for schema in result["db_schemas"] ) + assert all( + "EXECUTABLE WREN IDENTIFIER CATALOG" in schema["table_ddl"] + for schema in result["db_schemas"] + ) assert all( "sql_table_name_use_exactly" in schema["table_ddl"] for schema in result["db_schemas"] @@ -810,6 +814,10 @@ def encode(self, value): assert "sql_table_name_use_exactly: modeled_dataset" in table_ddl assert "sql_column_names_use_exactly:\n- stored_attribute" in table_ddl assert "END WREN SQL IDENTIFIER CONTRACT" in table_ddl + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in table_ddl + assert "table: modeled_dataset" in table_ddl + assert "columns:\n- stored_attribute" in table_ddl + assert "Do not create identifiers from user wording" in table_ddl assert ( '"semantic_context_not_sql_identifier":"Business-facing attribute label."' in table_ddl From 53a08a90978b5fdbe6a48780ab35dc49b31f2c15 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 22:47:05 +0530 Subject: [PATCH 0768/1087] Stop speculative SQL planning retries --- wren-ai-service/src/config.py | 4 +- wren-ai-service/src/web/v1/services/ask.py | 13 ++- .../src/web/v1/services/sql_corrections.py | 13 +-- .../pytest/services/test_dry_plan_defaults.py | 86 +++++++++++++++++++ 4 files changed, 100 insertions(+), 16 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index d6d2641224..bf35097a31 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -38,12 +38,12 @@ class Settings(BaseSettings): # generation config allow_intent_classification: bool = Field(default=True) - allow_sql_generation_reasoning: bool = Field(default=True) + allow_sql_generation_reasoning: bool = Field(default=False) allow_sql_functions_retrieval: bool = Field(default=True) allow_sql_diagnosis: bool = Field(default=True) allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) - max_sql_correction_retries: int = Field(default=3) + max_sql_correction_retries: int = Field(default=0) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ca55061902..67986120b7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -25,7 +25,7 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = False + ignore_sql_generation_reasoning: bool = True enable_column_pruning: bool = False use_dry_plan: bool = True allow_dry_plan_fallback: bool = False @@ -99,12 +99,12 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = True, + allow_sql_generation_reasoning: bool = False, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = True, - max_sql_correction_retries: int = 3, + max_sql_correction_retries: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -161,17 +161,14 @@ async def ask( table_names = [] error_message = None invalid_sql = None - allow_sql_generation_reasoning = ( - self._allow_sql_generation_reasoning - and not ask_request.ignore_sql_generation_reasoning - ) + allow_sql_generation_reasoning = False enable_column_pruning = ( self._enable_column_pruning or ask_request.enable_column_pruning ) allow_sql_functions_retrieval = self._allow_sql_functions_retrieval allow_sql_diagnosis = self._allow_sql_diagnosis allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - max_sql_correction_retries = self._max_sql_correction_retries + max_sql_correction_retries = 0 current_sql_correction_retries = 0 use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 4336186b6f..f805024ad2 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -89,11 +89,13 @@ async def correct( } if not retrieved_tables: - retrieved_tables = ( - await self._pipelines["sql_tables_extraction"].run( - sql=sql, - ) - )["post_process"] + self._handle_exception( + event_id, + "SQL correction requires retrieved table context from the original ask result.", + trace_id=trace_id, + request_from=request.request_from, + ) + return self._cache[event_id].with_metadata() if self._allow_sql_knowledge_retrieval: sql_knowledge = await self._pipelines["sql_knowledge_retrieval"].run( @@ -131,7 +133,6 @@ async def correct( event_id, f"An error occurred during SQL correction: {error_message}", trace_id=trace_id, - invalid_sql=invalid["sql"], request_from=request.request_from, ) else: diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py index 30d9957c0d..f5f15b14bf 100644 --- a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py +++ b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py @@ -16,6 +16,7 @@ def test_ask_request_defaults_to_planner_validation_without_fallback(): assert request.use_dry_plan is True assert request.allow_dry_plan_fallback is False + assert request.ignore_sql_generation_reasoning is True def test_ask_request_allows_explicit_planner_override(): @@ -34,6 +35,91 @@ def test_ask_service_defaults_to_column_pruning(): service = AskService({}) assert service._enable_column_pruning is True + assert service._allow_sql_generation_reasoning is False + assert service._max_sql_correction_retries == 0 + + +@pytest.mark.asyncio +async def test_sql_correction_service_requires_retrieved_tables(): + sql_tables_extraction = AsyncMock(run=AsyncMock(return_value={"post_process": []})) + service = SqlCorrectionService({"sql_tables_extraction": sql_tables_extraction}) + request = SqlCorrectionService.CorrectionRequest( + event_id="event-id", + sql="SELECT 1", + error="dry run failed", + ) + + await service.correct(request) + result = service["event-id"] + + assert result.status == "failed" + assert "retrieved table context" in result.error.message + assert result.invalid_sql is None + sql_tables_extraction.run.assert_not_called() + + +@pytest.mark.asyncio +async def test_ask_service_does_not_run_speculative_reasoning_or_correction_after_failed_generation(): + sql_generation_reasoning = AsyncMock(run=AsyncMock(return_value={})) + sql_correction = AsyncMock(run=AsyncMock(return_value={})) + service = AskService( + { + "sql_pairs_retrieval": AsyncMock( + run=AsyncMock(return_value={"formatted_output": {"documents": []}}) + ), + "instructions_retrieval": AsyncMock( + run=AsyncMock(return_value={"formatted_output": {"documents": []}}) + ), + "db_schema_retrieval": AsyncMock( + run=AsyncMock( + return_value={ + "construct_retrieval_results": { + "retrieval_results": [ + { + "table_name": "retrieved_model", + "table_ddl": "CREATE TABLE retrieved_model (retrieved_field VARCHAR);", + } + ], + "has_calculated_field": False, + "has_metric": False, + "has_json_field": False, + } + } + ) + ), + "sql_functions_retrieval": AsyncMock(run=AsyncMock(return_value=[])), + "sql_knowledge_retrieval": AsyncMock(run=AsyncMock(return_value=None)), + "sql_generation_reasoning": sql_generation_reasoning, + "sql_generation": AsyncMock( + run=AsyncMock( + return_value={ + "post_process": { + "valid_generation_result": {}, + "invalid_generation_result": { + "type": "DRY_RUN", + "sql": "SELECT 1", + "original_sql": "SELECT 1", + "error": "dry run failed", + }, + } + } + ) + ), + "sql_correction": sql_correction, + }, + allow_intent_classification=False, + ) + request = AskRequest(query="Can this be answered?", id="deploy-id") + request.query_id = "query-id" + + await service.ask(request) + result = service.get_ask_result(AskResultRequest(query_id="query-id")) + + assert result.status == "failed" + assert result.error.code == "NO_RELEVANT_SQL" + assert result.invalid_sql is None + sql_generation_reasoning.run.assert_not_called() + sql_correction.run.assert_not_called() def test_sql_correction_router_defaults_to_planner_validation_without_fallback(): From 515314c71c5ecc77efeff7562a7f8718d48a8f7d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Thu, 30 Jul 2026 23:36:44 +0530 Subject: [PATCH 0769/1087] Use executable schema context for SQL generation --- wren-ai-service/src/config.py | 2 +- .../retrieval/db_schema_retrieval.py | 203 ++++++++++-------- wren-ai-service/src/web/v1/services/ask.py | 4 +- .../retrieval/test_db_schema_retrieval.py | 48 ++--- .../pytest/services/test_dry_plan_defaults.py | 27 ++- 5 files changed, 147 insertions(+), 137 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index bf35097a31..ecb7b1cc0d 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -43,7 +43,7 @@ class Settings(BaseSettings): allow_sql_diagnosis: bool = Field(default=True) allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) - max_sql_correction_retries: int = Field(default=0) + max_sql_correction_retries: int = Field(default=3) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 759c99e84c..ff1931cbec 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -115,34 +115,35 @@ """ -def _build_metric_ddl(content: dict) -> str: +def _build_metric_ddl(content: dict, include_semantic_context: bool = True) -> str: columns = [ column for column in content["columns"] if column["data_type"].lower() != "unknown" ] - context = _format_semantic_context( - { - "object_type": "metric", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in columns - ], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable analytical aggregation interface", - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "semantic_context_not_sql_identifier": column["comment"], - } - for column in columns - ], - } + context = { + "object_type": "metric", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [column["name"] for column in columns], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable analytical aggregation interface", + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "semantic_context_not_sql_identifier": column["comment"], + } + for column in columns + ], + } + schema_context = ( + _format_semantic_context(context) + if include_semantic_context + else _format_executable_identifier_catalog(context) ) columns_ddl = [ f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" @@ -150,43 +151,42 @@ def _build_metric_ddl(content: dict) -> str: ] return ( - f"{context}CREATE TABLE {content['name']} (\n " + f"{schema_context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) -def _build_view_ddl(content: dict) -> str: +def _build_view_ddl(content: dict, include_semantic_context: bool = True) -> str: columns = [ column for column in content.get("columns", []) if column.get("name") and column.get("data_type", "").lower() != "unknown" ] - context = _format_semantic_context( - { - "object_type": "view", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in columns - ], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable virtual table interface", - "description": content["comment"], - "definition_omitted_from_executable_schema": True, - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type( - column.get("data_type") - ), - "semantic_context_not_sql_identifier": column.get("comment", ""), - } - for column in columns - ], - } + context = { + "object_type": "view", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [column["name"] for column in columns], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable virtual table interface", + "description": content["comment"], + "definition_omitted_from_executable_schema": True, + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column.get("data_type")), + "semantic_context_not_sql_identifier": column.get("comment", ""), + } + for column in columns + ], + } + schema_context = ( + _format_semantic_context(context) + if include_semantic_context + else _format_executable_identifier_catalog(context) ) columns_ddl = [ f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" @@ -194,7 +194,7 @@ def _build_view_ddl(content: dict) -> str: ] return ( - f"{context}CREATE TABLE {content['name']} (\n " + f"{schema_context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) @@ -339,7 +339,10 @@ def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: def _build_table_retrieval_context( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None + content: dict, + columns: Optional[set[str]] = None, + tables: Optional[set[str]] = None, + include_semantic_context: bool = True, ) -> tuple[str, bool, bool]: ddl, has_calculated_field, has_json_field = build_table_ddl( content, @@ -349,44 +352,44 @@ def _build_table_retrieval_context( ) included_columns = _included_columns(content, columns, tables) included_relationships = _included_relationships(content, tables) - context = _format_semantic_context( - { - "object_type": "model", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in included_columns - ], - "relationship_constraints_use_exactly": [ - relationship["constraint"] - for relationship in included_relationships - ], - }, - "semantic_context_not_sql_identifiers": { - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "is_primary_key": column["is_primary_key"], - "semantic_context_not_sql_identifier": column["comment"], - } - for column in included_columns + context = { + "object_type": "model", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in included_columns ], - "relationships": [ - { - "semantic_context_not_sql_identifier": relationship["comment"], - "sql_relationship_constraint_use_exactly": relationship[ - "constraint" - ], - "related_models_use_exactly": relationship.get("tables", []), - } - for relationship in included_relationships + "relationship_constraints_use_exactly": [ + relationship["constraint"] for relationship in included_relationships ], - } + }, + "semantic_context_not_sql_identifiers": { + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "is_primary_key": column["is_primary_key"], + "semantic_context_not_sql_identifier": column["comment"], + } + for column in included_columns + ], + "relationships": [ + { + "semantic_context_not_sql_identifier": relationship["comment"], + "sql_relationship_constraint_use_exactly": relationship["constraint"], + "related_models_use_exactly": relationship.get("tables", []), + } + for relationship in included_relationships + ], + } + schema_context = ( + _format_semantic_context(context) + if include_semantic_context + else _format_executable_identifier_catalog(context) ) - return f"{context}{ddl}", has_calculated_field, has_json_field + return f"{schema_context}{ddl}", has_calculated_field, has_json_field ## Start of Pipeline @@ -643,7 +646,10 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context(table_schema) + _build_table_retrieval_context( + table_schema, + include_semantic_context=False, + ) ) retrieval_results.append( { @@ -663,7 +669,10 @@ def check_using_db_schemas_without_pruning( retrieval_results.append( { "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), + "table_ddl": _build_metric_ddl( + content, + include_semantic_context=False, + ), } ) has_metric = True @@ -671,7 +680,10 @@ def check_using_db_schemas_without_pruning( retrieval_results.append( { "table_name": content["name"], - "table_ddl": _build_view_ddl(content), + "table_ddl": _build_view_ddl( + content, + include_semantic_context=False, + ), } ) @@ -771,6 +783,7 @@ def construct_retrieval_results( table_schema, columns=columns, tables=tables, + include_semantic_context=False, ) ) if _has_calculated_field: @@ -792,7 +805,10 @@ def construct_retrieval_results( retrieval_results.append( { "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), + "table_ddl": _build_metric_ddl( + content, + include_semantic_context=False, + ), } ) has_metric = True @@ -800,7 +816,10 @@ def construct_retrieval_results( retrieval_results.append( { "table_name": content["name"], - "table_ddl": _build_view_ddl(content), + "table_ddl": _build_view_ddl( + content, + include_semantic_context=False, + ), } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 67986120b7..c0b6f63442 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -104,7 +104,7 @@ def __init__( allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = True, - max_sql_correction_retries: int = 0, + max_sql_correction_retries: int = 3, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -168,7 +168,7 @@ async def ask( allow_sql_functions_retrieval = self._allow_sql_functions_retrieval allow_sql_diagnosis = self._allow_sql_diagnosis allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - max_sql_correction_retries = 0 + max_sql_correction_retries = self._max_sql_correction_retries current_sql_correction_retries = 0 use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 73058b9a02..59cddee4d7 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -646,9 +646,7 @@ def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_co assert "semantic_label" not in table_ddl assert "stored_dimension VARCHAR" in table_ddl assert "stored_measure DOUBLE" in table_ddl - assert "sql_column_names_use_exactly:\n- stored_dimension\n- stored_measure" in ( - table_ddl - ) + assert "columns:\n- stored_dimension\n- stored_measure" in table_ddl def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_unknown_columns(): @@ -705,9 +703,7 @@ def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_un assert "semantic_label" not in table_ddl assert "stored_dimension VARCHAR" in table_ddl assert "stored_measure DOUBLE" in table_ddl - assert "sql_column_names_use_exactly:\n- stored_dimension\n- stored_measure" in ( - table_ddl - ) + assert "columns:\n- stored_dimension\n- stored_measure" in table_ddl def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): @@ -748,28 +744,18 @@ def table_schema(name): "activity", "account", ] - assert all( - "WREN RETRIEVED SEMANTIC CONTEXT" in schema["table_ddl"] - for schema in result["db_schemas"] - ) - assert all( - "WREN SQL IDENTIFIER CONTRACT" in schema["table_ddl"] - for schema in result["db_schemas"] - ) assert all( "EXECUTABLE WREN IDENTIFIER CATALOG" in schema["table_ddl"] for schema in result["db_schemas"] ) + assert all("table: " in schema["table_ddl"] for schema in result["db_schemas"]) + assert all("columns:" in schema["table_ddl"] for schema in result["db_schemas"]) assert all( - "sql_table_name_use_exactly" in schema["table_ddl"] - for schema in result["db_schemas"] - ) - assert all( - "sql_column_name_use_exactly" in schema["table_ddl"] + "WREN RETRIEVED SEMANTIC CONTEXT" not in schema["table_ddl"] for schema in result["db_schemas"] ) assert all( - "semantic_context_not_sql_identifier" in schema["table_ddl"] + "semantic_context_not_sql_identifier" not in schema["table_ddl"] for schema in result["db_schemas"] ) assert result["tokens"] > 0 @@ -806,26 +792,16 @@ def encode(self, value): ) table_ddl = result["db_schemas"][0]["table_ddl"] - executable_ddl = table_ddl.split("*/", maxsplit=1)[1] - - assert '"sql_table_name_use_exactly":"modeled_dataset"' in table_ddl - assert '"sql_column_name_use_exactly":"stored_attribute"' in table_ddl - assert "WREN SQL IDENTIFIER CONTRACT" in table_ddl - assert "sql_table_name_use_exactly: modeled_dataset" in table_ddl - assert "sql_column_names_use_exactly:\n- stored_attribute" in table_ddl - assert "END WREN SQL IDENTIFIER CONTRACT" in table_ddl assert "EXECUTABLE WREN IDENTIFIER CATALOG" in table_ddl assert "table: modeled_dataset" in table_ddl assert "columns:\n- stored_attribute" in table_ddl assert "Do not create identifiers from user wording" in table_ddl - assert ( - '"semantic_context_not_sql_identifier":"Business-facing attribute label."' - in table_ddl - ) - assert "Business-facing attribute label." not in executable_ddl - assert "Business-facing dataset description." not in executable_ddl - assert "CREATE TABLE modeled_dataset" in executable_ddl - assert "stored_attribute VARCHAR" in executable_ddl + assert "Business-facing attribute label." not in table_ddl + assert "Business-facing dataset description." not in table_ddl + assert "WREN RETRIEVED SEMANTIC CONTEXT" not in table_ddl + assert "semantic_context_not_sql_identifier" not in table_ddl + assert "CREATE TABLE modeled_dataset" in table_ddl + assert "stored_attribute VARCHAR" in table_ddl def test_build_table_ddl_can_render_executable_schema_without_semantic_comments(): diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py index f5f15b14bf..244bfe6a5d 100644 --- a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py +++ b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py @@ -36,7 +36,7 @@ def test_ask_service_defaults_to_column_pruning(): assert service._enable_column_pruning is True assert service._allow_sql_generation_reasoning is False - assert service._max_sql_correction_retries == 0 + assert service._max_sql_correction_retries == 3 @pytest.mark.asyncio @@ -59,9 +59,20 @@ async def test_sql_correction_service_requires_retrieved_tables(): @pytest.mark.asyncio -async def test_ask_service_does_not_run_speculative_reasoning_or_correction_after_failed_generation(): +async def test_ask_service_skips_speculative_reasoning_and_corrects_with_retrieved_schema(): sql_generation_reasoning = AsyncMock(run=AsyncMock(return_value={})) - sql_correction = AsyncMock(run=AsyncMock(return_value={})) + sql_correction = AsyncMock( + run=AsyncMock( + return_value={ + "post_process": { + "valid_generation_result": { + "sql": "SELECT retrieved_field FROM retrieved_model" + }, + "invalid_generation_result": {}, + } + } + ) + ) service = AskService( { "sql_pairs_retrieval": AsyncMock( @@ -108,6 +119,7 @@ async def test_ask_service_does_not_run_speculative_reasoning_or_correction_afte "sql_correction": sql_correction, }, allow_intent_classification=False, + allow_sql_diagnosis=False, ) request = AskRequest(query="Can this be answered?", id="deploy-id") request.query_id = "query-id" @@ -115,11 +127,14 @@ async def test_ask_service_does_not_run_speculative_reasoning_or_correction_afte await service.ask(request) result = service.get_ask_result(AskResultRequest(query_id="query-id")) - assert result.status == "failed" - assert result.error.code == "NO_RELEVANT_SQL" + assert result.status == "finished" + assert result.response[0].sql == "SELECT retrieved_field FROM retrieved_model" assert result.invalid_sql is None sql_generation_reasoning.run.assert_not_called() - sql_correction.run.assert_not_called() + sql_correction.run.assert_called_once() + assert sql_correction.run.call_args.kwargs["contexts"] == [ + "CREATE TABLE retrieved_model (retrieved_field VARCHAR);" + ] def test_sql_correction_router_defaults_to_planner_validation_without_fallback(): From 072e3b8bbd6b3398591709d1ef576eb8ff2b360f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 00:05:09 +0530 Subject: [PATCH 0770/1087] Enforce manifest-grounded SQL generation --- .../generation/followup_sql_generation.py | 4 + .../pipelines/generation/sql_correction.py | 4 + .../pipelines/generation/sql_generation.py | 4 + .../pipelines/generation/sql_regeneration.py | 4 + .../src/pipelines/generation/utils/sql.py | 393 ++++++++++++++++++ .../retrieval/db_schema_retrieval.py | 32 ++ wren-ai-service/src/web/v1/services/ask.py | 9 + .../src/web/v1/services/ask_feedback.py | 8 + .../v1/services/question_recommendation.py | 26 +- .../src/web/v1/services/sql_corrections.py | 6 + .../pipelines/generation/test_sql_utils.py | 84 ++++ 11 files changed, 571 insertions(+), 3 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 1bcb3f853d..19f3c69ed0 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -149,6 +149,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -156,6 +157,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + schema_manifest=schema_manifest, ) @@ -207,6 +209,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + schema_manifest: dict[str, list[str]] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -233,6 +236,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "schema_manifest": schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index e92b8fc648..5caa8ea1a4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -145,6 +145,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -152,6 +153,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + schema_manifest=schema_manifest, ) @@ -199,6 +201,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + schema_manifest: dict[str, list[str]] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -221,6 +224,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "schema_manifest": schema_manifest, **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 2640082430..4dcb670fa6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -140,6 +140,7 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -148,6 +149,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + schema_manifest=schema_manifest, ) @@ -199,6 +201,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + schema_manifest: dict[str, list[str]] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -225,6 +228,7 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, + "schema_manifest": schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 70a7b00e21..2eacb3e981 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -164,10 +164,12 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, + schema_manifest=schema_manifest, ) @@ -212,6 +214,7 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_manifest: dict[str, list[str]] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -230,6 +233,7 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, + "schema_manifest": schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1633c399e4..1f4bee36eb 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,11 +1,15 @@ import logging +from dataclasses import dataclass, field from typing import Any, Dict, List import aiohttp import orjson +import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel +from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList +from sqlparse.tokens import DML, Keyword, Whitespace from src.core.engine import ( Engine, @@ -17,6 +21,380 @@ logger = logging.getLogger("wren-ai-service") +@dataclass +class ManifestGroundingResult: + is_grounded: bool + error: str = "" + + +@dataclass +class _SqlScope: + table_aliases: dict[str, str] = field(default_factory=dict) + cte_names: set[str] = field(default_factory=set) + derived_aliases: set[str] = field(default_factory=set) + output_aliases: set[str] = field(default_factory=set) + + +_CLAUSE_BOUNDARY_KEYWORDS = { + "ON", + "JOIN", + "WHERE", + "GROUP BY", + "HAVING", + "ORDER BY", + "LIMIT", + "UNION", + "UNION ALL", + "EXCEPT", + "INTERSECT", + "QUALIFY", + "WINDOW", +} + + +def _manifest_table_columns( + schema_manifest: dict[str, list[str]] | None, +) -> dict[str, set[str]]: + if not schema_manifest: + return {} + + return { + table_name: set(column_names or []) + for table_name, column_names in schema_manifest.items() + if table_name + } + + +def _non_whitespace_tokens(token_list: TokenList) -> list: + return [ + token + for token in token_list.tokens + if not token.is_whitespace and token.ttype is not Whitespace + ] + + +def _keyword_value(token) -> str: + return token.normalized if token.ttype in Keyword else "" + + +def _is_from_or_join_keyword(token) -> bool: + keyword = _keyword_value(token) + return keyword in {"FROM", "JOIN"} or keyword.endswith(" JOIN") + + +def _is_clause_boundary(token) -> bool: + keyword = _keyword_value(token) + return keyword in _CLAUSE_BOUNDARY_KEYWORDS or keyword.endswith(" JOIN") + + +def _identifier_name(identifier: Identifier | Function | None) -> str | None: + if identifier is None: + return None + return identifier.get_real_name() or identifier.get_name() + + +def _select_parenthesis(parenthesis: Parenthesis): + for token in parenthesis.tokens: + if isinstance(token, TokenList): + for child in token.flatten(): + if child.ttype is DML and child.normalized == "SELECT": + return token + return None + + +def _identifier_table_name(identifier: Identifier, manifest_tables: set[str]) -> str | None: + name = _identifier_name(identifier) + if name in manifest_tables: + return name + + value_tokens = [] + for token in identifier.tokens: + if token.is_whitespace or token.ttype is Whitespace: + break + if token.ttype in Keyword: + break + value_tokens.append(token.value) + full_name = "".join(value_tokens).strip('"') + return full_name if full_name in manifest_tables else name + + +def _identifier_is_subquery(identifier: Identifier) -> bool: + return any( + isinstance(token, Parenthesis) and _select_parenthesis(token) is not None + for token in identifier.tokens + ) + + +def _subquery_from_identifier(identifier: Identifier): + for token in identifier.tokens: + if isinstance(token, Parenthesis): + statement = _select_parenthesis(token) + if statement is not None: + return statement + return None + + +def _collect_ctes( + statement: TokenList, manifest: dict[str, set[str]], issues: list[str] +) -> set[str]: + cte_names: set[str] = set() + tokens = _non_whitespace_tokens(statement) + if not tokens or _keyword_value(tokens[0]) != "WITH": + return cte_names + + cte_token = tokens[1] if len(tokens) > 1 else None + cte_identifiers = ( + list(cte_token.get_identifiers()) + if isinstance(cte_token, IdentifierList) + else [cte_token] + if isinstance(cte_token, Identifier) + else [] + ) + + for cte_identifier in cte_identifiers: + cte_name = _identifier_name(cte_identifier) + if cte_name: + cte_names.add(cte_name) + subquery = _subquery_from_identifier(cte_identifier) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=cte_names) + + return cte_names + + +def _register_table_identifier( + identifier: Identifier, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if _identifier_is_subquery(identifier): + subquery = _subquery_from_identifier(identifier) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + if alias := identifier.get_alias(): + scope.derived_aliases.add(alias) + return + + table_name = _identifier_table_name(identifier, set(manifest)) + alias = identifier.get_alias() + if table_name in scope.cte_names: + if alias: + scope.derived_aliases.add(alias) + return + + if table_name not in manifest: + issues.append("Generated SQL references a table outside the retrieved Wren schema.") + return + + scope.table_aliases[alias or table_name] = table_name + + +def _register_table_token( + token, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + _register_table_identifier(identifier, scope, manifest, issues) + elif isinstance(token, Identifier): + _register_table_identifier(token, scope, manifest, issues) + elif isinstance(token, Parenthesis): + subquery = _select_parenthesis(token) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + + +def _collect_tables( + statement: TokenList, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + tokens = _non_whitespace_tokens(statement) + index = 0 + + while index < len(tokens): + token = tokens[index] + if _is_from_or_join_keyword(token): + index += 1 + while index < len(tokens) and not _is_clause_boundary(tokens[index]): + _register_table_token(tokens[index], scope, manifest, issues) + index += 1 + continue + + if isinstance(token, Parenthesis): + subquery = _select_parenthesis(token) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + + index += 1 + + +def _collect_select_aliases(statement: TokenList, scope: _SqlScope) -> None: + tokens = _non_whitespace_tokens(statement) + in_select = False + + for token in tokens: + if token.ttype is DML and token.normalized == "SELECT": + in_select = True + continue + if in_select and _keyword_value(token) == "FROM": + return + if not in_select: + continue + + identifiers = ( + token.get_identifiers() + if isinstance(token, IdentifierList) + else [token] + if isinstance(token, Identifier) + else [] + ) + for identifier in identifiers: + if alias := identifier.get_alias(): + scope.output_aliases.add(alias) + + +def _column_is_grounded(column_name: str, scope: _SqlScope, manifest: dict[str, set[str]]) -> bool: + manifest_tables = set(scope.table_aliases.values()) + if not manifest_tables: + return True + + return any(column_name in manifest[table_name] for table_name in manifest_tables) + + +def _validate_identifier_columns( + identifier: Identifier, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if isinstance(identifier, Function): + _validate_token_columns(identifier, scope, manifest, issues) + return + + if any(isinstance(token, Function) for token in identifier.tokens): + for token in identifier.tokens: + if isinstance(token, Function): + _validate_token_columns(token, scope, manifest, issues) + return + + column_name = _identifier_name(identifier) + if not column_name or column_name == "*": + return + if column_name in scope.output_aliases: + return + + parent_name = identifier.get_parent_name() + if parent_name: + if parent_name in scope.derived_aliases or parent_name in scope.cte_names: + return + table_name = scope.table_aliases.get(parent_name, parent_name) + if table_name not in manifest or column_name not in manifest[table_name]: + issues.append( + "Generated SQL references a column outside the retrieved Wren schema." + ) + return + + if not _column_is_grounded(column_name, scope, manifest): + issues.append("Generated SQL references a column outside the retrieved Wren schema.") + + +def _validate_token_columns( + token, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if isinstance(token, Function): + for child in token.tokens: + if isinstance(child, Parenthesis): + _validate_token_columns(child, scope, manifest, issues) + return + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + _validate_identifier_columns(identifier, scope, manifest, issues) + return + + if isinstance(token, Identifier): + _validate_identifier_columns(token, scope, manifest, issues) + return + + if isinstance(token, Parenthesis): + subquery = _select_parenthesis(token) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + return + + if isinstance(token, TokenList): + for child in token.tokens: + _validate_token_columns(child, scope, manifest, issues) + + +def _validate_columns( + statement: TokenList, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + tokens = _non_whitespace_tokens(statement) + index = 0 + + while index < len(tokens): + token = tokens[index] + if _keyword_value(token) == "WITH": + index += 2 + continue + + if _is_from_or_join_keyword(token): + index += 1 + while index < len(tokens) and not _is_clause_boundary(tokens[index]): + index += 1 + continue + + _validate_token_columns(token, scope, manifest, issues) + index += 1 + + +def _validate_statement( + statement: TokenList, + manifest: dict[str, set[str]], + issues: list[str], + parent_ctes: set[str] | None = None, +) -> None: + scope = _SqlScope(cte_names=set(parent_ctes or [])) + scope.cte_names.update(_collect_ctes(statement, manifest, issues)) + _collect_tables(statement, scope, manifest, issues) + _collect_select_aliases(statement, scope) + _validate_columns(statement, scope, manifest, issues) + + +def validate_sql_grounded_in_manifest( + sql: str, schema_manifest: dict[str, list[str]] | None +) -> ManifestGroundingResult: + manifest = _manifest_table_columns(schema_manifest) + if not manifest: + return ManifestGroundingResult(is_grounded=True) + + statements = [statement for statement in sqlparse.parse(sql) if statement.tokens] + if len(statements) != 1: + return ManifestGroundingResult( + is_grounded=False, + error="Generated SQL must contain one grounded SELECT statement.", + ) + + issues: list[str] = [] + _validate_statement(statements[0], manifest, issues) + if issues: + return ManifestGroundingResult(is_grounded=False, error=issues[0]) + + return ManifestGroundingResult(is_grounded=True) + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -34,6 +412,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -54,6 +433,7 @@ async def run( allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, allow_data_preview=allow_data_preview, + schema_manifest=schema_manifest, ) return { @@ -76,6 +456,7 @@ async def _classify_generation_result( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + schema_manifest: dict[str, list[str]] | None = None, ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} @@ -90,6 +471,18 @@ async def _classify_generation_result( "correlation_id": "", } + grounding_result = validate_sql_grounded_in_manifest( + generation_result, schema_manifest + ) + if not grounding_result.is_grounded: + return valid_generation_result, { + "sql": generation_result, + "original_sql": generation_result, + "type": "MANIFEST_GROUNDING", + "error": grounding_result.error, + "correlation_id": "", + } + async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ff1931cbec..c764fe07f8 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -338,6 +338,12 @@ def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: return bool(columns) and columns.issubset(executable_columns) +def _included_column_names( + content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None +) -> list[str]: + return [column["name"] for column in _included_columns(content, columns, tables)] + + def _build_table_retrieval_context( content: dict, columns: Optional[set[str]] = None, @@ -655,6 +661,7 @@ def check_using_db_schemas_without_pruning( { "table_name": table_schema["name"], "table_ddl": ddl, + "column_names": _included_column_names(table_schema), } ) if _has_calculated_field: @@ -673,6 +680,11 @@ def check_using_db_schemas_without_pruning( content, include_semantic_context=False, ), + "column_names": [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], } ) has_metric = True @@ -684,6 +696,12 @@ def check_using_db_schemas_without_pruning( content, include_semantic_context=False, ), + "column_names": [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], } ) @@ -795,6 +813,9 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, + "column_names": _included_column_names( + table_schema, columns=columns, tables=tables + ), } ) @@ -809,6 +830,11 @@ def construct_retrieval_results( content, include_semantic_context=False, ), + "column_names": [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], } ) has_metric = True @@ -820,6 +846,12 @@ def construct_retrieval_results( content, include_semantic_context=False, ), + "column_names": [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c0b6f63442..576330d73e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -159,6 +159,7 @@ async def ask( instructions = [] api_results = [] table_names = [] + schema_manifest = {} error_message = None invalid_sql = None allow_sql_generation_reasoning = False @@ -329,6 +330,11 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] + schema_manifest = { + document.get("table_name"): document.get("column_names", []) + for document in documents + if document.get("table_name") + } if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -451,6 +457,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_manifest=schema_manifest, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -469,6 +476,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_manifest=schema_manifest, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -545,6 +553,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + schema_manifest=schema_manifest, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index f1971afa15..8718547e2f 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -105,6 +105,7 @@ async def ask_feedback( error_message = None invalid_sql = None sql_knowledge = None + schema_manifest = {} allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval try: @@ -161,6 +162,11 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + schema_manifest = { + document.get("table_name"): document.get("column_names", []) + for document in documents + if document.get("table_name") + } sql_samples = sql_samples_task["formatted_output"].get("documents", []) instructions = instructions_task["formatted_output"].get( "documents", [] @@ -187,6 +193,7 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + schema_manifest=schema_manifest, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -251,6 +258,7 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + schema_manifest=schema_manifest, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 694d044bfa..a582be7c75 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -71,7 +71,9 @@ async def _validate_question( use_dry_plan: bool = True, allow_dry_plan_fallback: bool = False, ): - async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: + async def _document_retrieval() -> tuple[ + list[str], dict[str, list[str]], bool, bool, bool + ]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -79,10 +81,21 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + schema_manifest = { + document.get("table_name"): document.get("column_names", []) + for document in documents + if document.get("table_name") + } has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return table_ddls, has_calculated_field, has_metric, has_json_field + return ( + table_ddls, + schema_manifest, + has_calculated_field, + has_metric, + has_json_field, + ) async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -107,7 +120,13 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - table_ddls, has_calculated_field, has_metric, has_json_field = _document + ( + table_ddls, + schema_manifest, + has_calculated_field, + has_metric, + has_json_field, + ) = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -137,6 +156,7 @@ async def _instructions_retrieval() -> list[dict]: allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, + schema_manifest=schema_manifest, ) post_process = generated_sql["post_process"] diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index f805024ad2..d8d65a368e 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -113,6 +113,11 @@ async def correct( .get("retrieval_results", []) ) table_ddls = [document.get("table_ddl") for document in documents] + schema_manifest = { + document.get("table_name"): document.get("column_names", []) + for document in documents + if document.get("table_name") + } res = await self._pipelines["sql_correction"].run( contexts=table_ddls, @@ -121,6 +126,7 @@ async def correct( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_manifest=schema_manifest, ) post_process = res["post_process"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 0f4ed59d26..f56bdb81b2 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -118,6 +118,90 @@ async def test_sql_postprocessor_rejects_null_sql_generation_result(): assert result["invalid_generation_result"]["sql"] == "" +@pytest.mark.asyncio +async def test_sql_postprocessor_rejects_table_outside_retrieved_manifest(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=['{"sql": "SELECT \\"AvailableField\\" FROM \\"UnretrievedObject\\""}'], + use_dry_plan=True, + data_source="source", + schema_manifest={"RetrievedObject": ["AvailableField"]}, + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + +@pytest.mark.asyncio +async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=['{"sql": "SELECT \\"UnretrievedField\\" FROM \\"RetrievedObject\\""}'], + use_dry_plan=True, + data_source="source", + schema_manifest={"RetrievedObject": ["AvailableField"]}, + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + +@pytest.mark.asyncio +async def test_sql_postprocessor_allows_exact_retrieved_manifest_identifiers(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=[ + ( + '{"sql": "SELECT SUM(\\"AvailableField\\") AS \\"TotalField\\" ' + 'FROM \\"RetrievedObject\\""}' + ) + ], + use_dry_plan=True, + data_source="source", + schema_manifest={"RetrievedObject": ["AvailableField"]}, + ) + + assert result["valid_generation_result"]["sql"] == ( + 'SELECT SUM("AvailableField") AS "TotalField" FROM "RetrievedObject"' + ) + assert len(engine.dry_plan_calls) == 1 + assert len(engine.execute_sql_calls) == 1 + + +@pytest.mark.asyncio +async def test_sql_postprocessor_validates_join_predicate_columns(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=[ + ( + '{"sql": "SELECT a.\\"AvailableField\\" FROM \\"RetrievedObject\\" a ' + 'JOIN \\"RelatedObject\\" b ON a.\\"JoinField\\" = b.\\"JoinField\\""}' + ) + ], + use_dry_plan=True, + data_source="source", + schema_manifest={ + "RetrievedObject": ["AvailableField", "JoinField"], + "RelatedObject": ["JoinField"], + }, + ) + + assert result["valid_generation_result"]["sql"] == ( + 'SELECT a."AvailableField" FROM "RetrievedObject" a JOIN ' + '"RelatedObject" b ON a."JoinField" = b."JoinField"' + ) + assert len(engine.dry_plan_calls) == 1 + assert len(engine.execute_sql_calls) == 1 + + def test_construct_instructions_uses_instruction_text(): assert construct_instructions( [{"instruction": "First rule."}, {"instruction": "Second rule."}] From 970bbce4e8f54e4cf41dbf91fee2f921c96f7cc8 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 00:19:02 +0530 Subject: [PATCH 0771/1087] Validate manifest columns in SQL filters --- .../src/pipelines/generation/utils/sql.py | 8 +++- .../pipelines/generation/test_sql_utils.py | 44 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1f4bee36eb..1f122dd79d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -8,7 +8,7 @@ from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel -from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList +from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList, Where from sqlparse.tokens import DML, Keyword, Whitespace from src.core.engine import ( @@ -83,6 +83,9 @@ def _is_from_or_join_keyword(token) -> bool: def _is_clause_boundary(token) -> bool: + if isinstance(token, Where): + return True + keyword = _keyword_value(token) return keyword in _CLAUSE_BOUNDARY_KEYWORDS or keyword.endswith(" JOIN") @@ -317,7 +320,8 @@ def _validate_token_columns( if isinstance(token, IdentifierList): for identifier in token.get_identifiers(): - _validate_identifier_columns(identifier, scope, manifest, issues) + if isinstance(identifier, Identifier): + _validate_identifier_columns(identifier, scope, manifest, issues) return if isinstance(token, Identifier): diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index f56bdb81b2..746fd7668f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -152,6 +152,50 @@ async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): assert engine.execute_sql_calls == [] +@pytest.mark.asyncio +async def test_sql_postprocessor_rejects_filter_column_outside_retrieved_manifest(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=[ + ( + '{"sql": "SELECT * FROM \\"RetrievedObject\\" ' + 'WHERE \\"UnretrievedDate\\" >= CURRENT_DATE"}' + ) + ], + use_dry_plan=True, + data_source="source", + schema_manifest={"RetrievedObject": ["AvailableField"]}, + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + +@pytest.mark.asyncio +async def test_sql_postprocessor_rejects_function_argument_column_outside_manifest(): + engine = _DryPlanEngine() + + result = await SQLGenPostProcessor(engine).run( + replies=[ + ( + '{"sql": "SELECT * FROM \\"RetrievedObject\\" ' + "WHERE DATE_TRUNC('month', \\\"UnretrievedDate\\\") = CURRENT_DATE\"}" + ) + ], + use_dry_plan=True, + data_source="source", + schema_manifest={"RetrievedObject": ["AvailableField"]}, + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + @pytest.mark.asyncio async def test_sql_postprocessor_allows_exact_retrieved_manifest_identifiers(): engine = _DryPlanEngine() From fbb6f5d137fc942372fa7390de419b8d1fd4a2a0 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 00:29:51 +0530 Subject: [PATCH 0772/1087] Use full retrieved columns for SQL grounding --- .../retrieval/db_schema_retrieval.py | 11 ++++ wren-ai-service/src/web/v1/services/ask.py | 5 +- .../src/web/v1/services/ask_feedback.py | 5 +- .../v1/services/question_recommendation.py | 5 +- .../src/web/v1/services/sql_corrections.py | 5 +- .../retrieval/test_db_schema_retrieval.py | 58 +++++++++++++++++++ 6 files changed, 85 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index c764fe07f8..59b96ff83c 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -344,6 +344,14 @@ def _included_column_names( return [column["name"] for column in _included_columns(content, columns, tables)] +def _executable_column_names(content: dict) -> list[str]: + return [ + column["name"] + for column in content["columns"] + if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" + ] + + def _build_table_retrieval_context( content: dict, columns: Optional[set[str]] = None, @@ -816,6 +824,9 @@ def construct_retrieval_results( "column_names": _included_column_names( table_schema, columns=columns, tables=tables ), + "manifest_column_names": _executable_column_names( + table_schema + ), } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 576330d73e..a6990e773f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -331,7 +331,10 @@ async def ask( table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] schema_manifest = { - document.get("table_name"): document.get("column_names", []) + document.get("table_name"): document.get( + "manifest_column_names", + document.get("column_names", []), + ) for document in documents if document.get("table_name") } diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 8718547e2f..9fda32e7f4 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -163,7 +163,10 @@ async def ask_feedback( documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] schema_manifest = { - document.get("table_name"): document.get("column_names", []) + document.get("table_name"): document.get( + "manifest_column_names", + document.get("column_names", []), + ) for document in documents if document.get("table_name") } diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index a582be7c75..3203b3cca2 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -82,7 +82,10 @@ async def _document_retrieval() -> tuple[ documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] schema_manifest = { - document.get("table_name"): document.get("column_names", []) + document.get("table_name"): document.get( + "manifest_column_names", + document.get("column_names", []), + ) for document in documents if document.get("table_name") } diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index d8d65a368e..6089c6e62d 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -114,7 +114,10 @@ async def correct( ) table_ddls = [document.get("table_ddl") for document in documents] schema_manifest = { - document.get("table_name"): document.get("column_names", []) + document.get("table_name"): document.get( + "manifest_column_names", + document.get("column_names", []), + ) for document in documents if document.get("table_name") } diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 59cddee4d7..06c1f72d45 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -706,6 +706,64 @@ def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_un assert "columns:\n- stored_dimension\n- stored_measure" in table_ddl +def test_construct_retrieval_results_uses_full_columns_for_grounding_manifest(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_measure"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_dimension", + "data_type": "VARCHAR", + "comment": "Semantic dimension label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "stored_measure", + "data_type": "DOUBLE", + "comment": "Semantic measure label.", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + retrieved = result["retrieval_results"][0] + + assert retrieved["column_names"] == ["stored_measure"] + assert retrieved["manifest_column_names"] == [ + "stored_dimension", + "stored_measure", + ] + + def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): From 6663877ff4d6a21aa57953cd344eb942bab31fc9 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 00:51:11 +0530 Subject: [PATCH 0773/1087] Keep full columns in SQL generation context --- .../retrieval/db_schema_retrieval.py | 22 +------------------ .../retrieval/test_db_schema_retrieval.py | 9 ++++++-- 2 files changed, 8 insertions(+), 23 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 59b96ff83c..014e1ec616 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -329,15 +329,6 @@ def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[d ] -def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: - executable_columns = { - column["name"] - for column in content["columns"] - if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" - } - return bool(columns) and columns.issubset(executable_columns) - - def _included_column_names( content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None ) -> list[str]: @@ -794,20 +785,9 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - selected_columns = set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ) - columns = ( - selected_columns - if _selected_columns_are_executable( - table_schema, selected_columns - ) - else None - ) ddl, _has_calculated_field, _has_json_field = ( _build_table_retrieval_context( table_schema, - columns=columns, tables=tables, include_semantic_context=False, ) @@ -822,7 +802,7 @@ def construct_retrieval_results( "table_name": table_schema["name"], "table_ddl": ddl, "column_names": _included_column_names( - table_schema, columns=columns, tables=tables + table_schema, tables=tables ), "manifest_column_names": _executable_column_names( table_schema diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 06c1f72d45..6b8be0957a 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -706,7 +706,7 @@ def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_un assert "columns:\n- stored_dimension\n- stored_measure" in table_ddl -def test_construct_retrieval_results_uses_full_columns_for_grounding_manifest(): +def test_construct_retrieval_results_uses_full_columns_for_sql_generation(): result = construct_retrieval_results( check_using_db_schemas_without_pruning={}, filter_columns_in_tables={ @@ -757,7 +757,12 @@ def test_construct_retrieval_results_uses_full_columns_for_grounding_manifest(): retrieved = result["retrieval_results"][0] - assert retrieved["column_names"] == ["stored_measure"] + assert "stored_dimension VARCHAR" in retrieved["table_ddl"] + assert "stored_measure DOUBLE" in retrieved["table_ddl"] + assert retrieved["column_names"] == [ + "stored_dimension", + "stored_measure", + ] assert retrieved["manifest_column_names"] == [ "stored_dimension", "stored_measure", From 90d351dd27497eb6f6ae9a9020cfaec449855551 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 01:04:46 +0530 Subject: [PATCH 0774/1087] Regenerate SQL from manifest grounding failures --- .../pipelines/generation/sql_correction.py | 7 ++++++ .../src/pipelines/generation/utils/sql.py | 19 +++++++++++--- .../retrieval/db_schema_retrieval.py | 2 +- wren-ai-service/src/web/v1/services/ask.py | 8 +++++- .../src/web/v1/services/ask_feedback.py | 7 +++++- .../pipelines/generation/test_sql_utils.py | 25 +++++++++++++++++++ .../retrieval/test_db_schema_retrieval.py | 2 ++ 7 files changed, 63 insertions(+), 7 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 5caa8ea1a4..04830e2d17 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -93,6 +93,13 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### DRY-RUN DIAGNOSTIC ### The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. +{% if invalid_generation_result.get("type") == "MANIFEST_GROUNDING" %} +### MANIFEST GROUNDING FAILURE ### +The previous generated SQL was rejected before dry-run because it was not fully grounded in the retrieved Wren schema. +{{ invalid_generation_result.get("error", "") }} +Do not reuse rejected identifiers. Regenerate from the user question and the exact table and column identifiers in DATABASE SCHEMA only. +{% endif %} + Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1f122dd79d..9c2dfd4c18 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -8,7 +8,8 @@ from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel -from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList, Where +from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList +from sqlparse.sql import Where from sqlparse.tokens import DML, Keyword, Whitespace from src.core.engine import ( @@ -187,7 +188,9 @@ def _register_table_identifier( return if table_name not in manifest: - issues.append("Generated SQL references a table outside the retrieved Wren schema.") + issues.append( + f"Generated SQL references table `{table_name}` outside the retrieved Wren schema." + ) return scope.table_aliases[alias or table_name] = table_name @@ -298,12 +301,20 @@ def _validate_identifier_columns( table_name = scope.table_aliases.get(parent_name, parent_name) if table_name not in manifest or column_name not in manifest[table_name]: issues.append( - "Generated SQL references a column outside the retrieved Wren schema." + "Generated SQL references column " + f"`{column_name}` outside the retrieved Wren schema " + f"for table `{table_name}`." ) return if not _column_is_grounded(column_name, scope, manifest): - issues.append("Generated SQL references a column outside the retrieved Wren schema.") + table_names = sorted(set(scope.table_aliases.values())) + table_context = ( + f" for table `{table_names[0]}`" if len(table_names) == 1 else "" + ) + issues.append( + f"Generated SQL references column `{column_name}` outside the retrieved Wren schema{table_context}." + ) def _validate_token_columns( diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 014e1ec616..75941d5ea5 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -789,7 +789,7 @@ def construct_retrieval_results( _build_table_retrieval_context( table_schema, tables=tables, - include_semantic_context=False, + include_semantic_context=True, ) ) if _has_calculated_field: diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a6990e773f..f2c09fdc2e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -508,6 +508,7 @@ async def ask( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + error_type = failed_dry_run_result["type"] current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -521,7 +522,10 @@ async def ask( is_followup=True if histories else False, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and error_type != "MANIFEST_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -543,10 +547,12 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ + "type": error_type, "sql": original_sql, "error": ( f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" if allow_sql_diagnosis + and error_type != "MANIFEST_GROUNDING" and sql_diagnosis_reasoning else error_message ), diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 9fda32e7f4..7c5c7ebcfd 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -217,6 +217,7 @@ async def ask_feedback( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + error_type = failed_dry_run_result["type"] sql_diagnosis_reasoning = None self._ask_feedback_results[ @@ -226,7 +227,10 @@ async def ask_feedback( trace_id=trace_id, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and error_type != "MANIFEST_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -254,6 +258,7 @@ async def ask_feedback( sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, instructions=instructions, invalid_generation_result={ + "type": error_type, "original_sql": original_sql, "sql": invalid_sql, "error": correction_error_message, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 746fd7668f..13e481932c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -131,6 +131,7 @@ async def test_sql_postprocessor_rejects_table_outside_retrieved_manifest(): assert result["valid_generation_result"] == {} assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert "UnretrievedObject" in result["invalid_generation_result"]["error"] assert engine.dry_plan_calls == [] assert engine.execute_sql_calls == [] @@ -148,6 +149,8 @@ async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): assert result["valid_generation_result"] == {} assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert "UnretrievedField" in result["invalid_generation_result"]["error"] + assert "RetrievedObject" in result["invalid_generation_result"]["error"] assert engine.dry_plan_calls == [] assert engine.execute_sql_calls == [] @@ -415,6 +418,28 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "return null for sql instead of substituting non-schema identifiers" in prompt +def test_sql_correction_prompt_includes_manifest_grounding_failure(): + prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( + query="Question", + documents=["SCHEMA_CONTEXT"], + invalid_generation_result={ + "type": "MANIFEST_GROUNDING", + "error": ( + "Generated SQL references column `RejectedField` outside the " + "retrieved Wren schema for table `RetrievedObject`." + ), + }, + sql_generation_reasoning=None, + instructions=[], + sql_functions=[], + )["prompt"] + + assert "MANIFEST GROUNDING FAILURE" in prompt + assert "RejectedField" in prompt + assert "RetrievedObject" in prompt + assert "Do not reuse rejected identifiers" in prompt + + def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): prompt = sql_generation_reasoning_system_prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 6b8be0957a..7bbc25248d 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -759,6 +759,8 @@ def test_construct_retrieval_results_uses_full_columns_for_sql_generation(): assert "stored_dimension VARCHAR" in retrieved["table_ddl"] assert "stored_measure DOUBLE" in retrieved["table_ddl"] + assert "WREN RETRIEVED SEMANTIC CONTEXT" in retrieved["table_ddl"] + assert "sql_column_name_use_exactly" in retrieved["table_ddl"] assert retrieved["column_names"] == [ "stored_dimension", "stored_measure", From 893c1f1dfffb2739f94042067ae1d0d324572b8f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 01:48:14 +0530 Subject: [PATCH 0775/1087] Constrain SQL generation to schema identifiers --- .../generation/followup_sql_generation.py | 1 + .../src/pipelines/generation/sql_correction.py | 1 + .../src/pipelines/generation/sql_generation.py | 1 + .../src/pipelines/generation/utils/sql.py | 5 +++++ .../pipelines/generation/test_sql_utils.py | 18 ++++++++++++++++++ 5 files changed, 26 insertions(+) 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 19f3c69ed0..f06f19feaa 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -76,6 +76,7 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 04830e2d17..cda7988af3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -84,6 +84,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### QUESTION ### {% if query %} User's Question: {{ query }} +Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 4dcb670fa6..50d3f2b6eb 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -70,6 +70,7 @@ ### QUESTION ### User's Question: {{ query }} +Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 9c2dfd4c18..c10fcb10e8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -611,12 +611,17 @@ async def _classify_generation_result( - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. - Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. - Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. +- Read every retrieved DATABASE SCHEMA object before choosing the SQL source. Retrieval rank only lists candidates; it does not decide the table, view, or metric to query. +- Choose table, view, metric, and column identifiers only from retrieved schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. +- Do not default to the first retrieved object or to a broad object merely because it contains one requested word. If no retrieved object supports the requested intent with declared identifiers and declared relationships, return null for sql. - When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. - When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. - In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. - Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. +- Treat every declared table and column identifier as an indivisible string. Never splice, recombine, or transfer prefixes, numeric ordinals, suffixes, underscores, casing, punctuation, or business words between different declared identifiers. +- If a declared identifier contains generated prefixes, numeric ordinals, abbreviations, spaces, punctuation, or suffixes, copy the entire identifier exactly as declared. Do not rebuild it from the business meaning. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. - Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. - Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 13e481932c..6399366c7b 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -418,6 +418,24 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "return null for sql instead of substituting non-schema identifiers" in prompt +def test_sql_generation_prompt_requires_schema_object_selection_before_sql(): + prompt = get_sql_generation_system_prompt() + + assert "Read every retrieved DATABASE SCHEMA object before choosing" in prompt + assert "Retrieval rank only lists candidates" in prompt + assert "Do not default to the first retrieved object" in prompt + assert "declared fields support the user's requested subject" in prompt + + +def test_sql_generation_prompt_treats_identifiers_as_indivisible_strings(): + prompt = get_sql_generation_system_prompt() + + assert "Treat every declared table and column identifier as an indivisible string" in prompt + assert "Never splice, recombine, or transfer prefixes" in prompt + assert "copy the entire identifier exactly as declared" in prompt + assert "Do not rebuild it from the business meaning" in prompt + + def test_sql_correction_prompt_includes_manifest_grounding_failure(): prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( query="Question", From f03b90f8ddf8b0a5f26eb32a6f37e8621208e86c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 02:13:24 +0530 Subject: [PATCH 0776/1087] Add manifest identifier catalog to SQL prompts --- .../generation/followup_sql_generation.py | 9 +++++ .../pipelines/generation/sql_correction.py | 9 +++++ .../pipelines/generation/sql_generation.py | 9 +++++ .../pipelines/generation/sql_regeneration.py | 9 +++++ .../src/pipelines/generation/utils/sql.py | 22 ++++++++++++ .../pipelines/generation/test_sql_utils.py | 34 +++++++++++++++++++ 6 files changed, 92 insertions(+) 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 f06f19feaa..b3551608b9 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -15,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, + construct_executable_identifier_catalog, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -39,6 +40,10 @@ {{ document }} {% endfor %} +{% if executable_identifier_catalog %} +{{ executable_identifier_catalog }} +{% endif %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -99,10 +104,14 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + executable_identifier_catalog=construct_executable_identifier_catalog( + schema_manifest + ), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index cda7988af3..4b127c5cb6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + construct_executable_identifier_catalog, construct_instructions, get_text_to_sql_rules, ) @@ -65,6 +66,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% for document in documents %} {{ document }} {% endfor %} + +{% if executable_identifier_catalog %} +{{ executable_identifier_catalog }} +{% endif %} {% endif %} {% if sql_functions %} @@ -117,10 +122,14 @@ def prompt( sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + executable_identifier_catalog=construct_executable_identifier_catalog( + schema_manifest + ), invalid_generation_result=invalid_generation_result, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 50d3f2b6eb..3b78418728 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + construct_executable_identifier_catalog, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -33,6 +34,10 @@ {{ document }} {% endfor %} +{% if executable_identifier_catalog %} +{{ executable_identifier_catalog }} +{% endif %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -93,10 +98,14 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + executable_identifier_catalog=construct_executable_identifier_catalog( + schema_manifest + ), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 2eacb3e981..27e10c3d87 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + construct_executable_identifier_catalog, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -58,6 +59,10 @@ def get_sql_regeneration_system_prompt( {{ document }} {% endfor %} +{% if executable_identifier_catalog %} +{{ executable_identifier_catalog }} +{% endif %} + {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -119,11 +124,15 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, sql=sql, documents=documents, + executable_identifier_catalog=construct_executable_identifier_catalog( + schema_manifest + ), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c10fcb10e8..943c06b3c4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -916,6 +916,28 @@ def construct_instructions( return _instructions +def construct_executable_identifier_catalog( + schema_manifest: dict[str, list[str]] | None, +) -> str: + if not schema_manifest: + return "" + + lines = [ + "### EXECUTABLE WREN IDENTIFIER CATALOG ###", + "Use this catalog as the compact authoritative list of executable table and column identifiers for the generated SQL.", + "Copy identifiers exactly as written here. Do not derive, rebuild, or recombine table or column names from the user question or semantic descriptions.", + ] + for table_name, column_names in schema_manifest.items(): + if not table_name: + continue + lines.append(f'Table: "{table_name}"') + if column_names: + lines.append("Columns:") + lines.extend(f'- "{column_name}"' for column_name in column_names) + + return "\n".join(lines) + + def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 6399366c7b..9a41161f48 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -5,6 +5,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, + construct_executable_identifier_catalog, construct_instructions, get_json_field_instructions, get_metric_instructions, @@ -436,6 +437,39 @@ def test_sql_generation_prompt_treats_identifiers_as_indivisible_strings(): assert "Do not rebuild it from the business meaning" in prompt +def test_construct_executable_identifier_catalog_lists_manifest_identifiers(): + catalog = construct_executable_identifier_catalog( + {"ObjectA": ["FieldA", "FieldB"], "ObjectB": ["FieldC"]} + ) + + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in catalog + assert 'Table: "ObjectA"' in catalog + assert '- "FieldA"' in catalog + assert 'Table: "ObjectB"' in catalog + assert "Copy identifiers exactly as written here" in catalog + + +def test_sql_generation_prompt_can_include_executable_identifier_catalog(): + catalog = construct_executable_identifier_catalog({"ObjectA": ["FieldA"]}) + prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( + query="Question", + documents=["SCHEMA_CONTEXT"], + executable_identifier_catalog=catalog, + sql_generation_reasoning=None, + instructions=[], + calculated_field_instructions="", + metric_instructions="", + json_field_instructions="", + sql_samples=[], + sql_functions=[], + )["prompt"] + + assert "SCHEMA_CONTEXT" in prompt + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in prompt + assert 'Table: "ObjectA"' in prompt + assert '- "FieldA"' in prompt + + def test_sql_correction_prompt_includes_manifest_grounding_failure(): prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( query="Question", From af343a109b322f3e5e517dc99da5424c16e3ceaa Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 02:53:51 +0530 Subject: [PATCH 0777/1087] Let engine handle generated SQL execution --- .../pipelines/generation/sql_correction.py | 7 - .../src/pipelines/generation/utils/sql.py | 405 ------------------ wren-ai-service/src/web/v1/services/ask.py | 9 +- .../src/web/v1/services/ask_feedback.py | 5 +- .../pipelines/generation/test_sql_utils.py | 70 ++- 5 files changed, 28 insertions(+), 468 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 4b127c5cb6..107bbddc0c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -99,13 +99,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### DRY-RUN DIAGNOSTIC ### The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. -{% if invalid_generation_result.get("type") == "MANIFEST_GROUNDING" %} -### MANIFEST GROUNDING FAILURE ### -The previous generated SQL was rejected before dry-run because it was not fully grounded in the retrieved Wren schema. -{{ invalid_generation_result.get("error", "") }} -Do not reuse rejected identifiers. Regenerate from the user question and the exact table and column identifiers in DATABASE SCHEMA only. -{% endif %} - Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 943c06b3c4..47910d94d4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,16 +1,11 @@ import logging -from dataclasses import dataclass, field from typing import Any, Dict, List import aiohttp import orjson -import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel -from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList -from sqlparse.sql import Where -from sqlparse.tokens import DML, Keyword, Whitespace from src.core.engine import ( Engine, @@ -22,394 +17,6 @@ logger = logging.getLogger("wren-ai-service") -@dataclass -class ManifestGroundingResult: - is_grounded: bool - error: str = "" - - -@dataclass -class _SqlScope: - table_aliases: dict[str, str] = field(default_factory=dict) - cte_names: set[str] = field(default_factory=set) - derived_aliases: set[str] = field(default_factory=set) - output_aliases: set[str] = field(default_factory=set) - - -_CLAUSE_BOUNDARY_KEYWORDS = { - "ON", - "JOIN", - "WHERE", - "GROUP BY", - "HAVING", - "ORDER BY", - "LIMIT", - "UNION", - "UNION ALL", - "EXCEPT", - "INTERSECT", - "QUALIFY", - "WINDOW", -} - - -def _manifest_table_columns( - schema_manifest: dict[str, list[str]] | None, -) -> dict[str, set[str]]: - if not schema_manifest: - return {} - - return { - table_name: set(column_names or []) - for table_name, column_names in schema_manifest.items() - if table_name - } - - -def _non_whitespace_tokens(token_list: TokenList) -> list: - return [ - token - for token in token_list.tokens - if not token.is_whitespace and token.ttype is not Whitespace - ] - - -def _keyword_value(token) -> str: - return token.normalized if token.ttype in Keyword else "" - - -def _is_from_or_join_keyword(token) -> bool: - keyword = _keyword_value(token) - return keyword in {"FROM", "JOIN"} or keyword.endswith(" JOIN") - - -def _is_clause_boundary(token) -> bool: - if isinstance(token, Where): - return True - - keyword = _keyword_value(token) - return keyword in _CLAUSE_BOUNDARY_KEYWORDS or keyword.endswith(" JOIN") - - -def _identifier_name(identifier: Identifier | Function | None) -> str | None: - if identifier is None: - return None - return identifier.get_real_name() or identifier.get_name() - - -def _select_parenthesis(parenthesis: Parenthesis): - for token in parenthesis.tokens: - if isinstance(token, TokenList): - for child in token.flatten(): - if child.ttype is DML and child.normalized == "SELECT": - return token - return None - - -def _identifier_table_name(identifier: Identifier, manifest_tables: set[str]) -> str | None: - name = _identifier_name(identifier) - if name in manifest_tables: - return name - - value_tokens = [] - for token in identifier.tokens: - if token.is_whitespace or token.ttype is Whitespace: - break - if token.ttype in Keyword: - break - value_tokens.append(token.value) - full_name = "".join(value_tokens).strip('"') - return full_name if full_name in manifest_tables else name - - -def _identifier_is_subquery(identifier: Identifier) -> bool: - return any( - isinstance(token, Parenthesis) and _select_parenthesis(token) is not None - for token in identifier.tokens - ) - - -def _subquery_from_identifier(identifier: Identifier): - for token in identifier.tokens: - if isinstance(token, Parenthesis): - statement = _select_parenthesis(token) - if statement is not None: - return statement - return None - - -def _collect_ctes( - statement: TokenList, manifest: dict[str, set[str]], issues: list[str] -) -> set[str]: - cte_names: set[str] = set() - tokens = _non_whitespace_tokens(statement) - if not tokens or _keyword_value(tokens[0]) != "WITH": - return cte_names - - cte_token = tokens[1] if len(tokens) > 1 else None - cte_identifiers = ( - list(cte_token.get_identifiers()) - if isinstance(cte_token, IdentifierList) - else [cte_token] - if isinstance(cte_token, Identifier) - else [] - ) - - for cte_identifier in cte_identifiers: - cte_name = _identifier_name(cte_identifier) - if cte_name: - cte_names.add(cte_name) - subquery = _subquery_from_identifier(cte_identifier) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=cte_names) - - return cte_names - - -def _register_table_identifier( - identifier: Identifier, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if _identifier_is_subquery(identifier): - subquery = _subquery_from_identifier(identifier) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - if alias := identifier.get_alias(): - scope.derived_aliases.add(alias) - return - - table_name = _identifier_table_name(identifier, set(manifest)) - alias = identifier.get_alias() - if table_name in scope.cte_names: - if alias: - scope.derived_aliases.add(alias) - return - - if table_name not in manifest: - issues.append( - f"Generated SQL references table `{table_name}` outside the retrieved Wren schema." - ) - return - - scope.table_aliases[alias or table_name] = table_name - - -def _register_table_token( - token, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if isinstance(token, IdentifierList): - for identifier in token.get_identifiers(): - _register_table_identifier(identifier, scope, manifest, issues) - elif isinstance(token, Identifier): - _register_table_identifier(token, scope, manifest, issues) - elif isinstance(token, Parenthesis): - subquery = _select_parenthesis(token) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - - -def _collect_tables( - statement: TokenList, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - tokens = _non_whitespace_tokens(statement) - index = 0 - - while index < len(tokens): - token = tokens[index] - if _is_from_or_join_keyword(token): - index += 1 - while index < len(tokens) and not _is_clause_boundary(tokens[index]): - _register_table_token(tokens[index], scope, manifest, issues) - index += 1 - continue - - if isinstance(token, Parenthesis): - subquery = _select_parenthesis(token) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - - index += 1 - - -def _collect_select_aliases(statement: TokenList, scope: _SqlScope) -> None: - tokens = _non_whitespace_tokens(statement) - in_select = False - - for token in tokens: - if token.ttype is DML and token.normalized == "SELECT": - in_select = True - continue - if in_select and _keyword_value(token) == "FROM": - return - if not in_select: - continue - - identifiers = ( - token.get_identifiers() - if isinstance(token, IdentifierList) - else [token] - if isinstance(token, Identifier) - else [] - ) - for identifier in identifiers: - if alias := identifier.get_alias(): - scope.output_aliases.add(alias) - - -def _column_is_grounded(column_name: str, scope: _SqlScope, manifest: dict[str, set[str]]) -> bool: - manifest_tables = set(scope.table_aliases.values()) - if not manifest_tables: - return True - - return any(column_name in manifest[table_name] for table_name in manifest_tables) - - -def _validate_identifier_columns( - identifier: Identifier, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if isinstance(identifier, Function): - _validate_token_columns(identifier, scope, manifest, issues) - return - - if any(isinstance(token, Function) for token in identifier.tokens): - for token in identifier.tokens: - if isinstance(token, Function): - _validate_token_columns(token, scope, manifest, issues) - return - - column_name = _identifier_name(identifier) - if not column_name or column_name == "*": - return - if column_name in scope.output_aliases: - return - - parent_name = identifier.get_parent_name() - if parent_name: - if parent_name in scope.derived_aliases or parent_name in scope.cte_names: - return - table_name = scope.table_aliases.get(parent_name, parent_name) - if table_name not in manifest or column_name not in manifest[table_name]: - issues.append( - "Generated SQL references column " - f"`{column_name}` outside the retrieved Wren schema " - f"for table `{table_name}`." - ) - return - - if not _column_is_grounded(column_name, scope, manifest): - table_names = sorted(set(scope.table_aliases.values())) - table_context = ( - f" for table `{table_names[0]}`" if len(table_names) == 1 else "" - ) - issues.append( - f"Generated SQL references column `{column_name}` outside the retrieved Wren schema{table_context}." - ) - - -def _validate_token_columns( - token, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if isinstance(token, Function): - for child in token.tokens: - if isinstance(child, Parenthesis): - _validate_token_columns(child, scope, manifest, issues) - return - - if isinstance(token, IdentifierList): - for identifier in token.get_identifiers(): - if isinstance(identifier, Identifier): - _validate_identifier_columns(identifier, scope, manifest, issues) - return - - if isinstance(token, Identifier): - _validate_identifier_columns(token, scope, manifest, issues) - return - - if isinstance(token, Parenthesis): - subquery = _select_parenthesis(token) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - return - - if isinstance(token, TokenList): - for child in token.tokens: - _validate_token_columns(child, scope, manifest, issues) - - -def _validate_columns( - statement: TokenList, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - tokens = _non_whitespace_tokens(statement) - index = 0 - - while index < len(tokens): - token = tokens[index] - if _keyword_value(token) == "WITH": - index += 2 - continue - - if _is_from_or_join_keyword(token): - index += 1 - while index < len(tokens) and not _is_clause_boundary(tokens[index]): - index += 1 - continue - - _validate_token_columns(token, scope, manifest, issues) - index += 1 - - -def _validate_statement( - statement: TokenList, - manifest: dict[str, set[str]], - issues: list[str], - parent_ctes: set[str] | None = None, -) -> None: - scope = _SqlScope(cte_names=set(parent_ctes or [])) - scope.cte_names.update(_collect_ctes(statement, manifest, issues)) - _collect_tables(statement, scope, manifest, issues) - _collect_select_aliases(statement, scope) - _validate_columns(statement, scope, manifest, issues) - - -def validate_sql_grounded_in_manifest( - sql: str, schema_manifest: dict[str, list[str]] | None -) -> ManifestGroundingResult: - manifest = _manifest_table_columns(schema_manifest) - if not manifest: - return ManifestGroundingResult(is_grounded=True) - - statements = [statement for statement in sqlparse.parse(sql) if statement.tokens] - if len(statements) != 1: - return ManifestGroundingResult( - is_grounded=False, - error="Generated SQL must contain one grounded SELECT statement.", - ) - - issues: list[str] = [] - _validate_statement(statements[0], manifest, issues) - if issues: - return ManifestGroundingResult(is_grounded=False, error=issues[0]) - - return ManifestGroundingResult(is_grounded=True) - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -486,18 +93,6 @@ async def _classify_generation_result( "correlation_id": "", } - grounding_result = validate_sql_grounded_in_manifest( - generation_result, schema_manifest - ) - if not grounding_result.is_grounded: - return valid_generation_result, { - "sql": generation_result, - "original_sql": generation_result, - "type": "MANIFEST_GROUNDING", - "error": grounding_result.error, - "correlation_id": "", - } - async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f2c09fdc2e..4c2786f3c7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -522,10 +522,7 @@ async def ask( is_followup=True if histories else False, ) - if ( - allow_sql_diagnosis - and error_type != "MANIFEST_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -551,9 +548,7 @@ async def ask( "sql": original_sql, "error": ( f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" - if allow_sql_diagnosis - and error_type != "MANIFEST_GROUNDING" - and sql_diagnosis_reasoning + if allow_sql_diagnosis and sql_diagnosis_reasoning else error_message ), }, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 7c5c7ebcfd..e745a90da2 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -227,10 +227,7 @@ async def ask_feedback( trace_id=trace_id, ) - if ( - allow_sql_diagnosis - and error_type != "MANIFEST_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 9a41161f48..c6754da03a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -120,7 +120,7 @@ async def test_sql_postprocessor_rejects_null_sql_generation_result(): @pytest.mark.asyncio -async def test_sql_postprocessor_rejects_table_outside_retrieved_manifest(): +async def test_sql_postprocessor_sends_generated_sql_to_engine_with_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -130,15 +130,15 @@ async def test_sql_postprocessor_rejects_table_outside_retrieved_manifest(): schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert "UnretrievedObject" in result["invalid_generation_result"]["error"] - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] + assert result["valid_generation_result"]["sql"] == ( + 'SELECT "AvailableField" FROM "UnretrievedObject"' + ) + assert len(engine.dry_plan_calls) == 1 + assert len(engine.execute_sql_calls) == 1 @pytest.mark.asyncio -async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): +async def test_sql_postprocessor_uses_engine_for_column_errors_with_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -148,16 +148,15 @@ async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert "UnretrievedField" in result["invalid_generation_result"]["error"] - assert "RetrievedObject" in result["invalid_generation_result"]["error"] - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] + assert result["valid_generation_result"]["sql"] == ( + 'SELECT "UnretrievedField" FROM "RetrievedObject"' + ) + assert len(engine.dry_plan_calls) == 1 + assert len(engine.execute_sql_calls) == 1 @pytest.mark.asyncio -async def test_sql_postprocessor_rejects_filter_column_outside_retrieved_manifest(): +async def test_sql_postprocessor_sends_filter_sql_to_engine_with_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -172,14 +171,15 @@ async def test_sql_postprocessor_rejects_filter_column_outside_retrieved_manifes schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] + assert result["valid_generation_result"]["sql"] == ( + 'SELECT * FROM "RetrievedObject" WHERE "UnretrievedDate" >= CURRENT_DATE' + ) + assert len(engine.dry_plan_calls) == 1 + assert len(engine.execute_sql_calls) == 1 @pytest.mark.asyncio -async def test_sql_postprocessor_rejects_function_argument_column_outside_manifest(): +async def test_sql_postprocessor_sends_function_sql_to_engine_with_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -194,10 +194,12 @@ async def test_sql_postprocessor_rejects_function_argument_column_outside_manife schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] + assert result["valid_generation_result"]["sql"] == ( + 'SELECT * FROM "RetrievedObject" ' + 'WHERE DATE_TRUNC(\'month\', "UnretrievedDate") = CURRENT_DATE' + ) + assert len(engine.dry_plan_calls) == 1 + assert len(engine.execute_sql_calls) == 1 @pytest.mark.asyncio @@ -470,28 +472,6 @@ def test_sql_generation_prompt_can_include_executable_identifier_catalog(): assert '- "FieldA"' in prompt -def test_sql_correction_prompt_includes_manifest_grounding_failure(): - prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( - query="Question", - documents=["SCHEMA_CONTEXT"], - invalid_generation_result={ - "type": "MANIFEST_GROUNDING", - "error": ( - "Generated SQL references column `RejectedField` outside the " - "retrieved Wren schema for table `RetrievedObject`." - ), - }, - sql_generation_reasoning=None, - instructions=[], - sql_functions=[], - )["prompt"] - - assert "MANIFEST GROUNDING FAILURE" in prompt - assert "RejectedField" in prompt - assert "RetrievedObject" in prompt - assert "Do not reuse rejected identifiers" in prompt - - def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): prompt = sql_generation_reasoning_system_prompt From fdffe1a6e38d351064b5373076cfff23f96a9420 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 03:00:43 +0530 Subject: [PATCH 0778/1087] Revert "Let engine handle generated SQL execution" This reverts commit af343a109b322f3e5e517dc99da5424c16e3ceaa. --- .../pipelines/generation/sql_correction.py | 7 + .../src/pipelines/generation/utils/sql.py | 405 ++++++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 9 +- .../src/web/v1/services/ask_feedback.py | 5 +- .../pipelines/generation/test_sql_utils.py | 70 +-- 5 files changed, 468 insertions(+), 28 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 107bbddc0c..4b127c5cb6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -99,6 +99,13 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### DRY-RUN DIAGNOSTIC ### The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. +{% if invalid_generation_result.get("type") == "MANIFEST_GROUNDING" %} +### MANIFEST GROUNDING FAILURE ### +The previous generated SQL was rejected before dry-run because it was not fully grounded in the retrieved Wren schema. +{{ invalid_generation_result.get("error", "") }} +Do not reuse rejected identifiers. Regenerate from the user question and the exact table and column identifiers in DATABASE SCHEMA only. +{% endif %} + Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. Return only the final JSON SQL response. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 47910d94d4..943c06b3c4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,11 +1,16 @@ import logging +from dataclasses import dataclass, field from typing import Any, Dict, List import aiohttp import orjson +import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel +from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList +from sqlparse.sql import Where +from sqlparse.tokens import DML, Keyword, Whitespace from src.core.engine import ( Engine, @@ -17,6 +22,394 @@ logger = logging.getLogger("wren-ai-service") +@dataclass +class ManifestGroundingResult: + is_grounded: bool + error: str = "" + + +@dataclass +class _SqlScope: + table_aliases: dict[str, str] = field(default_factory=dict) + cte_names: set[str] = field(default_factory=set) + derived_aliases: set[str] = field(default_factory=set) + output_aliases: set[str] = field(default_factory=set) + + +_CLAUSE_BOUNDARY_KEYWORDS = { + "ON", + "JOIN", + "WHERE", + "GROUP BY", + "HAVING", + "ORDER BY", + "LIMIT", + "UNION", + "UNION ALL", + "EXCEPT", + "INTERSECT", + "QUALIFY", + "WINDOW", +} + + +def _manifest_table_columns( + schema_manifest: dict[str, list[str]] | None, +) -> dict[str, set[str]]: + if not schema_manifest: + return {} + + return { + table_name: set(column_names or []) + for table_name, column_names in schema_manifest.items() + if table_name + } + + +def _non_whitespace_tokens(token_list: TokenList) -> list: + return [ + token + for token in token_list.tokens + if not token.is_whitespace and token.ttype is not Whitespace + ] + + +def _keyword_value(token) -> str: + return token.normalized if token.ttype in Keyword else "" + + +def _is_from_or_join_keyword(token) -> bool: + keyword = _keyword_value(token) + return keyword in {"FROM", "JOIN"} or keyword.endswith(" JOIN") + + +def _is_clause_boundary(token) -> bool: + if isinstance(token, Where): + return True + + keyword = _keyword_value(token) + return keyword in _CLAUSE_BOUNDARY_KEYWORDS or keyword.endswith(" JOIN") + + +def _identifier_name(identifier: Identifier | Function | None) -> str | None: + if identifier is None: + return None + return identifier.get_real_name() or identifier.get_name() + + +def _select_parenthesis(parenthesis: Parenthesis): + for token in parenthesis.tokens: + if isinstance(token, TokenList): + for child in token.flatten(): + if child.ttype is DML and child.normalized == "SELECT": + return token + return None + + +def _identifier_table_name(identifier: Identifier, manifest_tables: set[str]) -> str | None: + name = _identifier_name(identifier) + if name in manifest_tables: + return name + + value_tokens = [] + for token in identifier.tokens: + if token.is_whitespace or token.ttype is Whitespace: + break + if token.ttype in Keyword: + break + value_tokens.append(token.value) + full_name = "".join(value_tokens).strip('"') + return full_name if full_name in manifest_tables else name + + +def _identifier_is_subquery(identifier: Identifier) -> bool: + return any( + isinstance(token, Parenthesis) and _select_parenthesis(token) is not None + for token in identifier.tokens + ) + + +def _subquery_from_identifier(identifier: Identifier): + for token in identifier.tokens: + if isinstance(token, Parenthesis): + statement = _select_parenthesis(token) + if statement is not None: + return statement + return None + + +def _collect_ctes( + statement: TokenList, manifest: dict[str, set[str]], issues: list[str] +) -> set[str]: + cte_names: set[str] = set() + tokens = _non_whitespace_tokens(statement) + if not tokens or _keyword_value(tokens[0]) != "WITH": + return cte_names + + cte_token = tokens[1] if len(tokens) > 1 else None + cte_identifiers = ( + list(cte_token.get_identifiers()) + if isinstance(cte_token, IdentifierList) + else [cte_token] + if isinstance(cte_token, Identifier) + else [] + ) + + for cte_identifier in cte_identifiers: + cte_name = _identifier_name(cte_identifier) + if cte_name: + cte_names.add(cte_name) + subquery = _subquery_from_identifier(cte_identifier) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=cte_names) + + return cte_names + + +def _register_table_identifier( + identifier: Identifier, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if _identifier_is_subquery(identifier): + subquery = _subquery_from_identifier(identifier) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + if alias := identifier.get_alias(): + scope.derived_aliases.add(alias) + return + + table_name = _identifier_table_name(identifier, set(manifest)) + alias = identifier.get_alias() + if table_name in scope.cte_names: + if alias: + scope.derived_aliases.add(alias) + return + + if table_name not in manifest: + issues.append( + f"Generated SQL references table `{table_name}` outside the retrieved Wren schema." + ) + return + + scope.table_aliases[alias or table_name] = table_name + + +def _register_table_token( + token, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + _register_table_identifier(identifier, scope, manifest, issues) + elif isinstance(token, Identifier): + _register_table_identifier(token, scope, manifest, issues) + elif isinstance(token, Parenthesis): + subquery = _select_parenthesis(token) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + + +def _collect_tables( + statement: TokenList, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + tokens = _non_whitespace_tokens(statement) + index = 0 + + while index < len(tokens): + token = tokens[index] + if _is_from_or_join_keyword(token): + index += 1 + while index < len(tokens) and not _is_clause_boundary(tokens[index]): + _register_table_token(tokens[index], scope, manifest, issues) + index += 1 + continue + + if isinstance(token, Parenthesis): + subquery = _select_parenthesis(token) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + + index += 1 + + +def _collect_select_aliases(statement: TokenList, scope: _SqlScope) -> None: + tokens = _non_whitespace_tokens(statement) + in_select = False + + for token in tokens: + if token.ttype is DML and token.normalized == "SELECT": + in_select = True + continue + if in_select and _keyword_value(token) == "FROM": + return + if not in_select: + continue + + identifiers = ( + token.get_identifiers() + if isinstance(token, IdentifierList) + else [token] + if isinstance(token, Identifier) + else [] + ) + for identifier in identifiers: + if alias := identifier.get_alias(): + scope.output_aliases.add(alias) + + +def _column_is_grounded(column_name: str, scope: _SqlScope, manifest: dict[str, set[str]]) -> bool: + manifest_tables = set(scope.table_aliases.values()) + if not manifest_tables: + return True + + return any(column_name in manifest[table_name] for table_name in manifest_tables) + + +def _validate_identifier_columns( + identifier: Identifier, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if isinstance(identifier, Function): + _validate_token_columns(identifier, scope, manifest, issues) + return + + if any(isinstance(token, Function) for token in identifier.tokens): + for token in identifier.tokens: + if isinstance(token, Function): + _validate_token_columns(token, scope, manifest, issues) + return + + column_name = _identifier_name(identifier) + if not column_name or column_name == "*": + return + if column_name in scope.output_aliases: + return + + parent_name = identifier.get_parent_name() + if parent_name: + if parent_name in scope.derived_aliases or parent_name in scope.cte_names: + return + table_name = scope.table_aliases.get(parent_name, parent_name) + if table_name not in manifest or column_name not in manifest[table_name]: + issues.append( + "Generated SQL references column " + f"`{column_name}` outside the retrieved Wren schema " + f"for table `{table_name}`." + ) + return + + if not _column_is_grounded(column_name, scope, manifest): + table_names = sorted(set(scope.table_aliases.values())) + table_context = ( + f" for table `{table_names[0]}`" if len(table_names) == 1 else "" + ) + issues.append( + f"Generated SQL references column `{column_name}` outside the retrieved Wren schema{table_context}." + ) + + +def _validate_token_columns( + token, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + if isinstance(token, Function): + for child in token.tokens: + if isinstance(child, Parenthesis): + _validate_token_columns(child, scope, manifest, issues) + return + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + if isinstance(identifier, Identifier): + _validate_identifier_columns(identifier, scope, manifest, issues) + return + + if isinstance(token, Identifier): + _validate_identifier_columns(token, scope, manifest, issues) + return + + if isinstance(token, Parenthesis): + subquery = _select_parenthesis(token) + if subquery is not None: + _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) + return + + if isinstance(token, TokenList): + for child in token.tokens: + _validate_token_columns(child, scope, manifest, issues) + + +def _validate_columns( + statement: TokenList, + scope: _SqlScope, + manifest: dict[str, set[str]], + issues: list[str], +) -> None: + tokens = _non_whitespace_tokens(statement) + index = 0 + + while index < len(tokens): + token = tokens[index] + if _keyword_value(token) == "WITH": + index += 2 + continue + + if _is_from_or_join_keyword(token): + index += 1 + while index < len(tokens) and not _is_clause_boundary(tokens[index]): + index += 1 + continue + + _validate_token_columns(token, scope, manifest, issues) + index += 1 + + +def _validate_statement( + statement: TokenList, + manifest: dict[str, set[str]], + issues: list[str], + parent_ctes: set[str] | None = None, +) -> None: + scope = _SqlScope(cte_names=set(parent_ctes or [])) + scope.cte_names.update(_collect_ctes(statement, manifest, issues)) + _collect_tables(statement, scope, manifest, issues) + _collect_select_aliases(statement, scope) + _validate_columns(statement, scope, manifest, issues) + + +def validate_sql_grounded_in_manifest( + sql: str, schema_manifest: dict[str, list[str]] | None +) -> ManifestGroundingResult: + manifest = _manifest_table_columns(schema_manifest) + if not manifest: + return ManifestGroundingResult(is_grounded=True) + + statements = [statement for statement in sqlparse.parse(sql) if statement.tokens] + if len(statements) != 1: + return ManifestGroundingResult( + is_grounded=False, + error="Generated SQL must contain one grounded SELECT statement.", + ) + + issues: list[str] = [] + _validate_statement(statements[0], manifest, issues) + if issues: + return ManifestGroundingResult(is_grounded=False, error=issues[0]) + + return ManifestGroundingResult(is_grounded=True) + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -93,6 +486,18 @@ async def _classify_generation_result( "correlation_id": "", } + grounding_result = validate_sql_grounded_in_manifest( + generation_result, schema_manifest + ) + if not grounding_result.is_grounded: + return valid_generation_result, { + "sql": generation_result, + "original_sql": generation_result, + "type": "MANIFEST_GROUNDING", + "error": grounding_result.error, + "correlation_id": "", + } + async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 4c2786f3c7..f2c09fdc2e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -522,7 +522,10 @@ async def ask( is_followup=True if histories else False, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and error_type != "MANIFEST_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -548,7 +551,9 @@ async def ask( "sql": original_sql, "error": ( f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" - if allow_sql_diagnosis and sql_diagnosis_reasoning + if allow_sql_diagnosis + and error_type != "MANIFEST_GROUNDING" + and sql_diagnosis_reasoning else error_message ), }, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index e745a90da2..7c5c7ebcfd 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -227,7 +227,10 @@ async def ask_feedback( trace_id=trace_id, ) - if allow_sql_diagnosis: + if ( + allow_sql_diagnosis + and error_type != "MANIFEST_GROUNDING" + ): sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index c6754da03a..9a41161f48 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -120,7 +120,7 @@ async def test_sql_postprocessor_rejects_null_sql_generation_result(): @pytest.mark.asyncio -async def test_sql_postprocessor_sends_generated_sql_to_engine_with_manifest(): +async def test_sql_postprocessor_rejects_table_outside_retrieved_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -130,15 +130,15 @@ async def test_sql_postprocessor_sends_generated_sql_to_engine_with_manifest(): schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"]["sql"] == ( - 'SELECT "AvailableField" FROM "UnretrievedObject"' - ) - assert len(engine.dry_plan_calls) == 1 - assert len(engine.execute_sql_calls) == 1 + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert "UnretrievedObject" in result["invalid_generation_result"]["error"] + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] @pytest.mark.asyncio -async def test_sql_postprocessor_uses_engine_for_column_errors_with_manifest(): +async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -148,15 +148,16 @@ async def test_sql_postprocessor_uses_engine_for_column_errors_with_manifest(): schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"]["sql"] == ( - 'SELECT "UnretrievedField" FROM "RetrievedObject"' - ) - assert len(engine.dry_plan_calls) == 1 - assert len(engine.execute_sql_calls) == 1 + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert "UnretrievedField" in result["invalid_generation_result"]["error"] + assert "RetrievedObject" in result["invalid_generation_result"]["error"] + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] @pytest.mark.asyncio -async def test_sql_postprocessor_sends_filter_sql_to_engine_with_manifest(): +async def test_sql_postprocessor_rejects_filter_column_outside_retrieved_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -171,15 +172,14 @@ async def test_sql_postprocessor_sends_filter_sql_to_engine_with_manifest(): schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"]["sql"] == ( - 'SELECT * FROM "RetrievedObject" WHERE "UnretrievedDate" >= CURRENT_DATE' - ) - assert len(engine.dry_plan_calls) == 1 - assert len(engine.execute_sql_calls) == 1 + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] @pytest.mark.asyncio -async def test_sql_postprocessor_sends_function_sql_to_engine_with_manifest(): +async def test_sql_postprocessor_rejects_function_argument_column_outside_manifest(): engine = _DryPlanEngine() result = await SQLGenPostProcessor(engine).run( @@ -194,12 +194,10 @@ async def test_sql_postprocessor_sends_function_sql_to_engine_with_manifest(): schema_manifest={"RetrievedObject": ["AvailableField"]}, ) - assert result["valid_generation_result"]["sql"] == ( - 'SELECT * FROM "RetrievedObject" ' - 'WHERE DATE_TRUNC(\'month\', "UnretrievedDate") = CURRENT_DATE' - ) - assert len(engine.dry_plan_calls) == 1 - assert len(engine.execute_sql_calls) == 1 + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] @pytest.mark.asyncio @@ -472,6 +470,28 @@ def test_sql_generation_prompt_can_include_executable_identifier_catalog(): assert '- "FieldA"' in prompt +def test_sql_correction_prompt_includes_manifest_grounding_failure(): + prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( + query="Question", + documents=["SCHEMA_CONTEXT"], + invalid_generation_result={ + "type": "MANIFEST_GROUNDING", + "error": ( + "Generated SQL references column `RejectedField` outside the " + "retrieved Wren schema for table `RetrievedObject`." + ), + }, + sql_generation_reasoning=None, + instructions=[], + sql_functions=[], + )["prompt"] + + assert "MANIFEST GROUNDING FAILURE" in prompt + assert "RejectedField" in prompt + assert "RetrievedObject" in prompt + assert "Do not reuse rejected identifiers" in prompt + + def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): prompt = sql_generation_reasoning_system_prompt From 3e810979b793527814607ba43e9a40c6775dce14 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 14:22:36 +0530 Subject: [PATCH 0779/1087] Respect SQL generation LLM kwargs --- wren-ai-service/src/providers/llm/litellm.py | 8 ++- .../pytest/providers/test_litellm_llm.py | 50 +++++++++++++++++++ 2 files changed, 53 insertions(+), 5 deletions(-) create mode 100644 wren-ai-service/tests/pytest/providers/test_litellm_llm.py diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index 1b32bd9f05..f4023f4c57 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -71,10 +71,7 @@ def get_generator( generation_kwargs: Optional[Dict[str, Any]] = None, streaming_callback: Optional[Callable[[StreamingChunk], None]] = None, ): - combined_generation_kwargs = { - **(generation_kwargs or {}), - **(self._model_kwargs or {}), - } + component_generation_kwargs = generation_kwargs or {} def _normalize_generation_kwargs( kwargs: Optional[Dict[str, Any]], @@ -124,7 +121,8 @@ async def _run( generation_kwargs = _normalize_generation_kwargs( { - **combined_generation_kwargs, + **(self._model_kwargs or {}), + **component_generation_kwargs, **(generation_kwargs or {}), } ) diff --git a/wren-ai-service/tests/pytest/providers/test_litellm_llm.py b/wren-ai-service/tests/pytest/providers/test_litellm_llm.py new file mode 100644 index 0000000000..44b2afc373 --- /dev/null +++ b/wren-ai-service/tests/pytest/providers/test_litellm_llm.py @@ -0,0 +1,50 @@ +from types import SimpleNamespace + +import pytest + +from src.providers.llm.litellm import LitellmLLMProvider + + +@pytest.mark.asyncio +async def test_component_generation_kwargs_override_model_defaults(mocker): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="test-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": null}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="test-model", + api_base="http://localhost/v1", + kwargs={ + "temperature": 1, + "response_format": {"type": "text"}, + }, + ) + + generator = provider.get_generator( + generation_kwargs={ + "temperature": 0, + "response_format": { + "type": "json_schema", + "json_schema": {"name": "result", "schema": {}}, + }, + } + ) + + await generator(prompt="Return SQL") + + assert captured_kwargs["temperature"] == 0 + assert captured_kwargs["response_format"]["type"] == "json_schema" + From ddbdc6b933fa037a10abaf168de0e1cfe98215cf Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 15:33:36 +0530 Subject: [PATCH 0780/1087] Ground SQL generation with active manifest --- wren-ai-service/src/pipelines/common.py | 11 ++++ .../generation/followup_sql_generation.py | 13 +++-- .../pipelines/generation/sql_correction.py | 19 +++--- .../pipelines/generation/sql_generation.py | 13 +++-- .../src/pipelines/indexing/project_meta.py | 58 ++++++++++++++++++- .../pipelines/generation/test_sql_utils.py | 7 ++- .../pipelines/indexing/test_project_meta.py | 56 ++++++++++++++++++ .../retrieval/test_project_scope_isolation.py | 16 ++++- 8 files changed, 173 insertions(+), 20 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 605238bb64..18098181a3 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -113,6 +113,17 @@ async def retrieve_metadata(project_id: str, retriever) -> dict[str, Any]: return {} +def resolve_schema_manifest( + metadata: dict[str, Any], + schema_manifest: dict[str, list[str]] | None, +) -> dict[str, list[str]] | None: + active_schema_manifest = metadata.get("schema_manifest") + if isinstance(active_schema_manifest, dict) and active_schema_manifest: + return active_schema_manifest + + return schema_manifest + + @component class ScoreFilter: @component.output_types( 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 b3551608b9..541f34e3bb 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -10,7 +10,11 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider -from src.pipelines.common import clean_up_new_lines, retrieve_metadata +from src.pipelines.common import ( + clean_up_new_lines, + resolve_schema_manifest, + retrieve_metadata, +) from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -82,8 +86,8 @@ ### QUESTION ### User's Follow-up Question: {{ query }} Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Answer the user's intent using the current DATABASE SCHEMA and EXECUTABLE WREN IDENTIFIER CATALOG. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the supplied schema and catalog do not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Return only the final JSON SQL response. @@ -223,10 +227,11 @@ async def run( ): logger.info("Follow-Up SQL Generation pipeline is running...") - if use_dry_plan: + if project_id or use_dry_plan: metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} + schema_manifest = resolve_schema_manifest(metadata, schema_manifest) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 4b127c5cb6..1cfc32784b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -11,7 +11,11 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider -from src.pipelines.common import clean_up_new_lines, retrieve_metadata +from src.pipelines.common import ( + clean_up_new_lines, + resolve_schema_manifest, + retrieve_metadata, +) from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -90,8 +94,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% if query %} User's Question: {{ query }} Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. +Answer the user's intent using the current DATABASE SCHEMA and EXECUTABLE WREN IDENTIFIER CATALOG. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the supplied schema and catalog do not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. {% endif %} ### FAILED SQL ### The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. @@ -101,12 +105,12 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% if invalid_generation_result.get("type") == "MANIFEST_GROUNDING" %} ### MANIFEST GROUNDING FAILURE ### -The previous generated SQL was rejected before dry-run because it was not fully grounded in the retrieved Wren schema. +The previous generated SQL was rejected before dry-run because it was not fully grounded in the active Wren schema. {{ invalid_generation_result.get("error", "") }} -Do not reuse rejected identifiers. Regenerate from the user question and the exact table and column identifiers in DATABASE SCHEMA only. +Do not reuse rejected identifiers. Regenerate from the user question and the exact table and column identifiers in DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG only. {% endif %} -Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. +Regenerate from the user question and current DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. Return only the final JSON SQL response. """ @@ -222,10 +226,11 @@ async def run( ): logger.info("SQLCorrection pipeline is running...") - if use_dry_plan: + if project_id or use_dry_plan: metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} + schema_manifest = resolve_schema_manifest(metadata, schema_manifest) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 3b78418728..433dc3a95c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -10,7 +10,11 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider -from src.pipelines.common import clean_up_new_lines, retrieve_metadata +from src.pipelines.common import ( + clean_up_new_lines, + resolve_schema_manifest, + retrieve_metadata, +) from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -76,8 +80,8 @@ ### QUESTION ### User's Question: {{ query }} Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Answer the user's intent using the current DATABASE SCHEMA and EXECUTABLE WREN IDENTIFIER CATALOG. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the supplied schema and catalog do not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Return only the final JSON SQL response. @@ -215,10 +219,11 @@ async def run( ): logger.info("SQL Generation pipeline is running...") - if use_dry_plan: + if project_id or use_dry_plan: metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} + schema_manifest = resolve_schema_manifest(metadata, schema_manifest) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/indexing/project_meta.py b/wren-ai-service/src/pipelines/indexing/project_meta.py index 426e988934..6105b935fa 100644 --- a/wren-ai-service/src/pipelines/indexing/project_meta.py +++ b/wren-ai-service/src/pipelines/indexing/project_meta.py @@ -18,6 +18,58 @@ logger = logging.getLogger("wren-ai-service") +def _column_name(column: dict[str, Any]) -> str | None: + name = column.get("name") + return name if isinstance(name, str) and name else None + + +def _model_column_names(model: dict[str, Any]) -> list[str]: + return [ + name + for column in model.get("columns", []) + if column.get("isHidden") is not True and not column.get("relationship") + for name in [_column_name(column)] + if name + ] + + +def _view_column_names(view: dict[str, Any]) -> list[str]: + properties = view.get("properties") or {} + return [ + name + for column in properties.get("columns", []) + for name in [_column_name(column)] + if name + ] + + +def _metric_column_names(metric: dict[str, Any]) -> list[str]: + return [ + name + for column in metric.get("dimension", []) + metric.get("measure", []) + for name in [_column_name(column)] + if name + ] + + +def build_schema_manifest(mdl: dict[str, Any]) -> dict[str, list[str]]: + manifest: dict[str, list[str]] = {} + + for model in mdl.get("models", []): + if name := model.get("name"): + manifest[name] = _model_column_names(model) + + for view in mdl.get("views", []): + if name := view.get("name"): + manifest[name] = _view_column_names(view) + + for metric in mdl.get("metrics", []): + if name := metric.get("name"): + manifest[name] = _metric_column_names(metric) + + return manifest + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) @extract_fields(dict(mdl=dict[str, Any])) @@ -40,7 +92,11 @@ def chunk( document = Document( id=str(uuid.uuid4()), - meta={"data_source": data_source, **addition}, + meta={ + "data_source": data_source, + "schema_manifest": build_schema_manifest(mdl), + **addition, + }, ) return {"documents": [document]} diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 9a41161f48..c28feabcce 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -613,6 +613,7 @@ def test_executable_prompt_templates_omit_untrusted_reasoning_and_sql_context(): assert "dry-run diagnostic text is intentionally omitted" in ( correction_prompt ) - assert "Regenerate from the user question and current DATABASE SCHEMA only" in ( - correction_prompt - ) + assert ( + "Regenerate from the user question and current DATABASE SCHEMA or " + "EXECUTABLE WREN IDENTIFIER CATALOG only" + ) in correction_prompt diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py new file mode 100644 index 0000000000..3a088044cb --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py @@ -0,0 +1,56 @@ +from src.pipelines.indexing.project_meta import build_schema_manifest, chunk + + +def test_project_meta_schema_manifest_uses_deployed_mdl_identifiers(): + mdl = { + "dataSource": "postgres", + "models": [ + { + "name": "PrimaryEntity", + "columns": [ + {"name": "VisibleField"}, + {"name": "HiddenField", "isHidden": True}, + {"name": "LinkedField", "relationship": "RelatedEntity"}, + ], + } + ], + "views": [ + { + "name": "SavedView", + "properties": { + "columns": [ + {"name": "ViewField"}, + ], + }, + } + ], + "metrics": [ + { + "name": "MetricEntity", + "dimension": [{"name": "DimensionField"}], + "measure": [{"name": "MeasureField"}], + } + ], + } + + assert build_schema_manifest(mdl) == { + "PrimaryEntity": ["VisibleField"], + "SavedView": ["ViewField"], + "MetricEntity": ["DimensionField", "MeasureField"], + } + + +def test_project_meta_chunk_stores_schema_manifest_with_project_scope(): + result = chunk( + mdl={ + "dataSource": "duckdb", + "models": [{"name": "Entity", "columns": [{"name": "Field"}]}], + }, + project_id="project-a", + ) + + document = result["documents"][0] + + assert document.meta["data_source"] == "local_file" + assert document.meta["project_id"] == "project-a" + assert document.meta["schema_manifest"] == {"Entity": ["Field"]} diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py index d058274eff..31ff6a3e1f 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py @@ -1,6 +1,6 @@ import pytest -from src.pipelines.common import retrieve_metadata +from src.pipelines.common import resolve_schema_manifest, retrieve_metadata from src.pipelines.retrieval import historical_question_retrieval, instructions from src.pipelines.retrieval import sql_pairs_retrieval @@ -48,6 +48,20 @@ async def test_metadata_retrieval_does_not_fall_back_to_global_documents(): assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] +def test_active_project_manifest_overrides_retrieved_subset(): + assert resolve_schema_manifest( + metadata={"schema_manifest": {"ActiveEntity": ["ActiveField"]}}, + schema_manifest={"RetrievedEntity": ["RetrievedField"]}, + ) == {"ActiveEntity": ["ActiveField"]} + + +def test_retrieved_manifest_is_used_when_project_metadata_has_no_manifest(): + assert resolve_schema_manifest( + metadata={"data_source": "source"}, + schema_manifest={"RetrievedEntity": ["RetrievedField"]}, + ) == {"RetrievedEntity": ["RetrievedField"]} + + @pytest.mark.asyncio async def test_sql_pairs_count_stays_project_scoped_when_project_has_no_documents(): store = StoreSpy(count=0) From 5437baf02225c780c218197ed4e5f007655a73b1 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 15:55:24 +0530 Subject: [PATCH 0781/1087] Fallback to indexed active schema manifest --- wren-ai-service/src/pipelines/common.py | 68 ++++++++++ .../generation/followup_sql_generation.py | 12 +- .../pipelines/generation/sql_correction.py | 12 +- .../pipelines/generation/sql_generation.py | 12 +- .../retrieval/test_project_scope_isolation.py | 122 +++++++++++++++++- 5 files changed, 217 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 18098181a3..238df24677 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -1,3 +1,4 @@ +import ast import re from typing import Any, List, Optional, Tuple @@ -124,6 +125,73 @@ def resolve_schema_manifest( return schema_manifest +async def retrieve_schema_manifest( + project_id: str, + retriever, +) -> dict[str, list[str]]: + filters: dict[str, Any] = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + result = await retriever.run(query_embedding=[], filters=filters, top_k=10000) + manifest: dict[str, list[str]] = {} + + for document in result.get("documents", []): + content = ast.literal_eval(document.content) + table_name = document.meta.get("name") or content.get("name") + if not table_name: + continue + + manifest.setdefault(table_name, []) + if content.get("type") == "TABLE_COLUMNS": + column_names = [ + column["name"] + for column in content.get("columns", []) + if column.get("type") == "COLUMN" + and column.get("name") + and column.get("data_type", "").lower() != "unknown" + ] + elif content.get("type") in {"VIEW", "METRIC"}: + column_names = [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ] + else: + column_names = [] + + for column_name in column_names: + if column_name not in manifest[table_name]: + manifest[table_name].append(column_name) + + return {table_name: columns for table_name, columns in manifest.items() if columns} + + +async def resolve_active_schema_manifest( + metadata: dict[str, Any], + schema_manifest: dict[str, list[str]] | None, + project_id: str, + dbschema_retriever, +) -> dict[str, list[str]] | None: + resolved_schema_manifest = resolve_schema_manifest(metadata, schema_manifest) + if metadata.get("schema_manifest") or not project_id: + return resolved_schema_manifest + + indexed_schema_manifest = await retrieve_schema_manifest( + project_id=project_id, + retriever=dbschema_retriever, + ) + return indexed_schema_manifest or resolved_schema_manifest + + @component class ScoreFilter: @component.output_types( 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 541f34e3bb..920443f48a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -12,7 +12,7 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import ( clean_up_new_lines, - resolve_schema_manifest, + resolve_active_schema_manifest, retrieve_metadata, ) from src.pipelines.generation.utils.sql import ( @@ -189,6 +189,9 @@ def __init__( self._retriever = document_store_provider.get_retriever( document_store_provider.get_store("project_meta") ) + self._dbschema_retriever = document_store_provider.get_retriever( + document_store_provider.get_store() + ) self._components = { "generator": llm_provider.get_generator( @@ -231,7 +234,12 @@ async def run( metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - schema_manifest = resolve_schema_manifest(metadata, schema_manifest) + schema_manifest = await resolve_active_schema_manifest( + metadata, + schema_manifest, + project_id or "", + self._dbschema_retriever, + ) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 1cfc32784b..47c7c0dfc5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -13,7 +13,7 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import ( clean_up_new_lines, - resolve_schema_manifest, + resolve_active_schema_manifest, retrieve_metadata, ) from src.pipelines.generation.utils.sql import ( @@ -192,6 +192,9 @@ def __init__( self._retriever = document_store_provider.get_retriever( document_store_provider.get_store("project_meta") ) + self._dbschema_retriever = document_store_provider.get_retriever( + document_store_provider.get_store() + ) self._components = { "generator": llm_provider.get_generator( @@ -230,7 +233,12 @@ async def run( metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - schema_manifest = resolve_schema_manifest(metadata, schema_manifest) + schema_manifest = await resolve_active_schema_manifest( + metadata, + schema_manifest, + project_id or "", + self._dbschema_retriever, + ) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 433dc3a95c..11f86ca605 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -12,7 +12,7 @@ from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import ( clean_up_new_lines, - resolve_schema_manifest, + resolve_active_schema_manifest, retrieve_metadata, ) from src.pipelines.generation.utils.sql import ( @@ -181,6 +181,9 @@ def __init__( self._retriever = document_store_provider.get_retriever( document_store_provider.get_store("project_meta") ) + self._dbschema_retriever = document_store_provider.get_retriever( + document_store_provider.get_store() + ) self._components = { "generator": llm_provider.get_generator( @@ -223,7 +226,12 @@ async def run( metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - schema_manifest = resolve_schema_manifest(metadata, schema_manifest) + schema_manifest = await resolve_active_schema_manifest( + metadata, + schema_manifest, + project_id or "", + self._dbschema_retriever, + ) return await self._pipe.execute( ["post_process"], diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py index 31ff6a3e1f..d68135b8d3 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py @@ -1,6 +1,12 @@ import pytest - -from src.pipelines.common import resolve_schema_manifest, retrieve_metadata +from haystack import Document + +from src.pipelines.common import ( + resolve_active_schema_manifest, + resolve_schema_manifest, + retrieve_metadata, + retrieve_schema_manifest, +) from src.pipelines.retrieval import historical_question_retrieval, instructions from src.pipelines.retrieval import sql_pairs_retrieval @@ -28,7 +34,7 @@ def __init__(self, documents=None): self.documents = documents or [] self.calls = [] - async def run(self, query_embedding=None, filters=None): + async def run(self, query_embedding=None, filters=None, **_): self.calls.append( { "query_embedding": query_embedding, @@ -62,6 +68,116 @@ def test_retrieved_manifest_is_used_when_project_metadata_has_no_manifest(): ) == {"RetrievedEntity": ["RetrievedField"]} +@pytest.mark.asyncio +async def test_schema_manifest_can_be_rebuilt_from_indexed_schema_documents(): + retriever = RetrieverSpy( + documents=[ + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "FirstField", + "data_type": "varchar", + }, + { + "type": "FOREIGN_KEY", + "name": "RelationField", + "data_type": "varchar", + }, + { + "type": "COLUMN", + "name": "UnknownField", + "data_type": "unknown", + }, + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "IndexedEntity"}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "SecondField", + "data_type": "integer", + }, + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "IndexedEntity"}, + ), + Document( + content=str( + { + "type": "VIEW", + "name": "IndexedView", + "columns": [ + { + "name": "ViewField", + "data_type": "varchar", + }, + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "IndexedView"}, + ), + ] + ) + + assert await retrieve_schema_manifest("project-a", retriever) == { + "IndexedEntity": ["FirstField", "SecondField"], + "IndexedView": ["ViewField"], + } + assert retriever.calls == [ + { + "query_embedding": [], + "filters": { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + ], + }, + } + ] + + +@pytest.mark.asyncio +async def test_active_manifest_falls_back_to_indexed_schema_documents(): + retriever = RetrieverSpy( + documents=[ + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "ActiveField", + "data_type": "varchar", + }, + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "ActiveEntity"}, + ) + ] + ) + + assert await resolve_active_schema_manifest( + metadata={"data_source": "source"}, + schema_manifest={"RetrievedEntity": ["RetrievedField"]}, + project_id="project-a", + dbschema_retriever=retriever, + ) == {"ActiveEntity": ["ActiveField"]} + + @pytest.mark.asyncio async def test_sql_pairs_count_stays_project_scoped_when_project_has_no_documents(): store = StoreSpy(count=0) From 2884be81561d9cf328f00151318ecc769a73ac38 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 16:11:47 +0530 Subject: [PATCH 0782/1087] Keep SQL prompt grounded to retrieval --- .../generation/followup_sql_generation.py | 7 +++--- .../pipelines/generation/sql_correction.py | 7 +++--- .../pipelines/generation/sql_generation.py | 7 +++--- .../pipelines/generation/test_sql_utils.py | 22 +++++++++++++++++++ 4 files changed, 34 insertions(+), 9 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 920443f48a..f071dacc5a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -163,7 +163,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - schema_manifest: dict[str, list[str]] | None = None, + grounding_schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -171,7 +171,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_manifest=schema_manifest, + schema_manifest=grounding_schema_manifest, ) @@ -234,7 +234,7 @@ async def run( metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - schema_manifest = await resolve_active_schema_manifest( + grounding_schema_manifest = await resolve_active_schema_manifest( metadata, schema_manifest, project_id or "", @@ -260,6 +260,7 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, "schema_manifest": schema_manifest, + "grounding_schema_manifest": grounding_schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 47c7c0dfc5..57c8236a96 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -166,7 +166,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - schema_manifest: dict[str, list[str]] | None = None, + grounding_schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -174,7 +174,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_manifest=schema_manifest, + schema_manifest=grounding_schema_manifest, ) @@ -233,7 +233,7 @@ async def run( metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - schema_manifest = await resolve_active_schema_manifest( + grounding_schema_manifest = await resolve_active_schema_manifest( metadata, schema_manifest, project_id or "", @@ -255,6 +255,7 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, "schema_manifest": schema_manifest, + "grounding_schema_manifest": grounding_schema_manifest, **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 11f86ca605..886eef6a7f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -154,7 +154,7 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, - schema_manifest: dict[str, list[str]] | None = None, + grounding_schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -163,7 +163,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - schema_manifest=schema_manifest, + schema_manifest=grounding_schema_manifest, ) @@ -226,7 +226,7 @@ async def run( metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - schema_manifest = await resolve_active_schema_manifest( + grounding_schema_manifest = await resolve_active_schema_manifest( metadata, schema_manifest, project_id or "", @@ -252,6 +252,7 @@ async def run( "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, "schema_manifest": schema_manifest, + "grounding_schema_manifest": grounding_schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index c28feabcce..2118029220 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -22,6 +22,7 @@ ) from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template +from src.pipelines.generation.sql_generation import post_process as sql_post_process from src.pipelines.generation.sql_regeneration import get_sql_regeneration_system_prompt from src.pipelines.generation.sql_regeneration import sql_regeneration_user_prompt_template @@ -156,6 +157,27 @@ async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): assert engine.execute_sql_calls == [] +@pytest.mark.asyncio +async def test_sql_generation_post_process_uses_separate_grounding_manifest(): + engine = _DryPlanEngine() + + result = await sql_post_process( + generate_sql={ + "replies": ['{"sql": "SELECT \\"ActiveField\\" FROM \\"ActiveEntity\\""}'] + }, + post_processor=SQLGenPostProcessor(engine), + data_source="source", + use_dry_plan=True, + grounding_schema_manifest={"ActiveEntity": ["ActiveField"]}, + ) + + assert result["valid_generation_result"]["sql"] == ( + 'SELECT "ActiveField" FROM "ActiveEntity"' + ) + assert len(engine.dry_plan_calls) == 1 + assert len(engine.execute_sql_calls) == 1 + + @pytest.mark.asyncio async def test_sql_postprocessor_rejects_filter_column_outside_retrieved_manifest(): engine = _DryPlanEngine() From 499aa66740e0fce9654b96e16900ab6929ac29c0 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 16:29:21 +0530 Subject: [PATCH 0783/1087] Restore legacy v1 ask SQL flow --- wren-ai-service/src/pipelines/common.py | 79 ---- .../generation/followup_sql_generation.py | 36 +- .../pipelines/generation/sql_correction.py | 45 +- .../pipelines/generation/sql_generation.py | 36 +- .../pipelines/generation/sql_regeneration.py | 13 - .../src/pipelines/generation/utils/sql.py | 435 ------------------ .../src/pipelines/indexing/project_meta.py | 58 +-- wren-ai-service/src/web/v1/services/ask.py | 24 +- .../src/web/v1/services/ask_feedback.py | 18 +- .../v1/services/question_recommendation.py | 29 +- .../src/web/v1/services/sql_corrections.py | 9 - .../pipelines/generation/test_sql_utils.py | 234 +--------- .../pipelines/indexing/test_project_meta.py | 56 --- .../retrieval/test_project_scope_isolation.py | 136 +----- .../pytest/services/test_dry_plan_defaults.py | 27 +- 15 files changed, 33 insertions(+), 1202 deletions(-) delete mode 100644 wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 238df24677..605238bb64 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -1,4 +1,3 @@ -import ast import re from typing import Any, List, Optional, Tuple @@ -114,84 +113,6 @@ async def retrieve_metadata(project_id: str, retriever) -> dict[str, Any]: return {} -def resolve_schema_manifest( - metadata: dict[str, Any], - schema_manifest: dict[str, list[str]] | None, -) -> dict[str, list[str]] | None: - active_schema_manifest = metadata.get("schema_manifest") - if isinstance(active_schema_manifest, dict) and active_schema_manifest: - return active_schema_manifest - - return schema_manifest - - -async def retrieve_schema_manifest( - project_id: str, - retriever, -) -> dict[str, list[str]]: - filters: dict[str, Any] = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) - - result = await retriever.run(query_embedding=[], filters=filters, top_k=10000) - manifest: dict[str, list[str]] = {} - - for document in result.get("documents", []): - content = ast.literal_eval(document.content) - table_name = document.meta.get("name") or content.get("name") - if not table_name: - continue - - manifest.setdefault(table_name, []) - if content.get("type") == "TABLE_COLUMNS": - column_names = [ - column["name"] - for column in content.get("columns", []) - if column.get("type") == "COLUMN" - and column.get("name") - and column.get("data_type", "").lower() != "unknown" - ] - elif content.get("type") in {"VIEW", "METRIC"}: - column_names = [ - column["name"] - for column in content.get("columns", []) - if column.get("name") - and column.get("data_type", "").lower() != "unknown" - ] - else: - column_names = [] - - for column_name in column_names: - if column_name not in manifest[table_name]: - manifest[table_name].append(column_name) - - return {table_name: columns for table_name, columns in manifest.items() if columns} - - -async def resolve_active_schema_manifest( - metadata: dict[str, Any], - schema_manifest: dict[str, list[str]] | None, - project_id: str, - dbschema_retriever, -) -> dict[str, list[str]] | None: - resolved_schema_manifest = resolve_schema_manifest(metadata, schema_manifest) - if metadata.get("schema_manifest") or not project_id: - return resolved_schema_manifest - - indexed_schema_manifest = await retrieve_schema_manifest( - project_id=project_id, - retriever=dbschema_retriever, - ) - return indexed_schema_manifest or resolved_schema_manifest - - @component class ScoreFilter: @component.output_types( 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 f071dacc5a..1bcb3f853d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -10,16 +10,11 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider -from src.pipelines.common import ( - clean_up_new_lines, - resolve_active_schema_manifest, - retrieve_metadata, -) +from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, - construct_executable_identifier_catalog, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -44,10 +39,6 @@ {{ document }} {% endfor %} -{% if executable_identifier_catalog %} -{{ executable_identifier_catalog }} -{% endif %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -85,9 +76,8 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. -Answer the user's intent using the current DATABASE SCHEMA and EXECUTABLE WREN IDENTIFIER CATALOG. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the supplied schema and catalog do not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Return only the final JSON SQL response. @@ -108,14 +98,10 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - executable_identifier_catalog=construct_executable_identifier_catalog( - schema_manifest - ), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -163,7 +149,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - grounding_schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -171,7 +156,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_manifest=grounding_schema_manifest, ) @@ -189,9 +173,6 @@ def __init__( self._retriever = document_store_provider.get_retriever( document_store_provider.get_store("project_meta") ) - self._dbschema_retriever = document_store_provider.get_retriever( - document_store_provider.get_store() - ) self._components = { "generator": llm_provider.get_generator( @@ -226,20 +207,13 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - schema_manifest: dict[str, list[str]] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") - if project_id or use_dry_plan: + if use_dry_plan: metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - grounding_schema_manifest = await resolve_active_schema_manifest( - metadata, - schema_manifest, - project_id or "", - self._dbschema_retriever, - ) return await self._pipe.execute( ["post_process"], @@ -259,8 +233,6 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, - "schema_manifest": schema_manifest, - "grounding_schema_manifest": grounding_schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 57c8236a96..e92b8fc648 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -11,15 +11,10 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider -from src.pipelines.common import ( - clean_up_new_lines, - resolve_active_schema_manifest, - retrieve_metadata, -) +from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, - construct_executable_identifier_catalog, construct_instructions, get_text_to_sql_rules, ) @@ -70,10 +65,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% for document in documents %} {{ document }} {% endfor %} - -{% if executable_identifier_catalog %} -{{ executable_identifier_catalog }} -{% endif %} {% endif %} {% if sql_functions %} @@ -93,9 +84,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### QUESTION ### {% if query %} User's Question: {{ query }} -Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. -Answer the user's intent using the current DATABASE SCHEMA and EXECUTABLE WREN IDENTIFIER CATALOG. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the supplied schema and catalog do not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. {% endif %} ### FAILED SQL ### The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. @@ -103,14 +93,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### DRY-RUN DIAGNOSTIC ### The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. -{% if invalid_generation_result.get("type") == "MANIFEST_GROUNDING" %} -### MANIFEST GROUNDING FAILURE ### -The previous generated SQL was rejected before dry-run because it was not fully grounded in the active Wren schema. -{{ invalid_generation_result.get("error", "") }} -Do not reuse rejected identifiers. Regenerate from the user question and the exact table and column identifiers in DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG only. -{% endif %} - -Regenerate from the user question and current DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. +Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. Return only the final JSON SQL response. """ @@ -126,14 +109,10 @@ def prompt( sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, - schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - executable_identifier_catalog=construct_executable_identifier_catalog( - schema_manifest - ), invalid_generation_result=invalid_generation_result, sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( @@ -166,7 +145,6 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, - grounding_schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -174,7 +152,6 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - schema_manifest=grounding_schema_manifest, ) @@ -192,9 +169,6 @@ def __init__( self._retriever = document_store_provider.get_retriever( document_store_provider.get_store("project_meta") ) - self._dbschema_retriever = document_store_provider.get_retriever( - document_store_provider.get_store() - ) self._components = { "generator": llm_provider.get_generator( @@ -225,20 +199,13 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - schema_manifest: dict[str, list[str]] | None = None, ): logger.info("SQLCorrection pipeline is running...") - if project_id or use_dry_plan: + if use_dry_plan: metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - grounding_schema_manifest = await resolve_active_schema_manifest( - metadata, - schema_manifest, - project_id or "", - self._dbschema_retriever, - ) return await self._pipe.execute( ["post_process"], @@ -254,8 +221,6 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, - "schema_manifest": schema_manifest, - "grounding_schema_manifest": grounding_schema_manifest, **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 886eef6a7f..2640082430 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -10,15 +10,10 @@ from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider -from src.pipelines.common import ( - clean_up_new_lines, - resolve_active_schema_manifest, - retrieve_metadata, -) +from src.pipelines.common import clean_up_new_lines, retrieve_metadata from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, - construct_executable_identifier_catalog, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -38,10 +33,6 @@ {{ document }} {% endfor %} -{% if executable_identifier_catalog %} -{{ executable_identifier_catalog }} -{% endif %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -79,9 +70,8 @@ ### QUESTION ### User's Question: {{ query }} -Before writing SQL, review every DATABASE SCHEMA document supplied in this prompt. Select table, view, metric, and column identifiers only from schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. Do not use the first retrieved object by default; retrieval rank is only candidate order. If a declared identifier contains prefixes, numeric ordinals, underscores, spaces, punctuation, casing, abbreviations, or suffixes, copy the whole identifier exactly as declared and do not rebuild it from the business meaning. -Answer the user's intent using the current DATABASE SCHEMA and EXECUTABLE WREN IDENTIFIER CATALOG. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the supplied schema and catalog do not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or EXECUTABLE WREN IDENTIFIER CATALOG, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Return only the final JSON SQL response. @@ -102,14 +92,10 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, - executable_identifier_catalog=construct_executable_identifier_catalog( - schema_manifest - ), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -154,7 +140,6 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, - grounding_schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -163,7 +148,6 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - schema_manifest=grounding_schema_manifest, ) @@ -181,9 +165,6 @@ def __init__( self._retriever = document_store_provider.get_retriever( document_store_provider.get_store("project_meta") ) - self._dbschema_retriever = document_store_provider.get_retriever( - document_store_provider.get_store() - ) self._components = { "generator": llm_provider.get_generator( @@ -218,20 +199,13 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - schema_manifest: dict[str, list[str]] | None = None, ): logger.info("SQL Generation pipeline is running...") - if project_id or use_dry_plan: + if use_dry_plan: metadata = await retrieve_metadata(project_id or "", self._retriever) else: metadata = {} - grounding_schema_manifest = await resolve_active_schema_manifest( - metadata, - schema_manifest, - project_id or "", - self._dbschema_retriever, - ) return await self._pipe.execute( ["post_process"], @@ -251,8 +225,6 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, - "schema_manifest": schema_manifest, - "grounding_schema_manifest": grounding_schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 27e10c3d87..70a7b00e21 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,7 +14,6 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, - construct_executable_identifier_catalog, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -59,10 +58,6 @@ def get_sql_regeneration_system_prompt( {{ document }} {% endfor %} -{% if executable_identifier_catalog %} -{{ executable_identifier_catalog }} -{% endif %} - {% if calculated_field_instructions %} {{ calculated_field_instructions }} {% endif %} @@ -124,15 +119,11 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_manifest: dict[str, list[str]] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, sql=sql, documents=documents, - executable_identifier_catalog=construct_executable_identifier_catalog( - schema_manifest - ), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -173,12 +164,10 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, - schema_manifest: dict[str, list[str]] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, - schema_manifest=schema_manifest, ) @@ -223,7 +212,6 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, - schema_manifest: dict[str, list[str]] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -242,7 +230,6 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, - "schema_manifest": schema_manifest, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 943c06b3c4..1633c399e4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,16 +1,11 @@ import logging -from dataclasses import dataclass, field from typing import Any, Dict, List import aiohttp import orjson -import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel -from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList -from sqlparse.sql import Where -from sqlparse.tokens import DML, Keyword, Whitespace from src.core.engine import ( Engine, @@ -22,394 +17,6 @@ logger = logging.getLogger("wren-ai-service") -@dataclass -class ManifestGroundingResult: - is_grounded: bool - error: str = "" - - -@dataclass -class _SqlScope: - table_aliases: dict[str, str] = field(default_factory=dict) - cte_names: set[str] = field(default_factory=set) - derived_aliases: set[str] = field(default_factory=set) - output_aliases: set[str] = field(default_factory=set) - - -_CLAUSE_BOUNDARY_KEYWORDS = { - "ON", - "JOIN", - "WHERE", - "GROUP BY", - "HAVING", - "ORDER BY", - "LIMIT", - "UNION", - "UNION ALL", - "EXCEPT", - "INTERSECT", - "QUALIFY", - "WINDOW", -} - - -def _manifest_table_columns( - schema_manifest: dict[str, list[str]] | None, -) -> dict[str, set[str]]: - if not schema_manifest: - return {} - - return { - table_name: set(column_names or []) - for table_name, column_names in schema_manifest.items() - if table_name - } - - -def _non_whitespace_tokens(token_list: TokenList) -> list: - return [ - token - for token in token_list.tokens - if not token.is_whitespace and token.ttype is not Whitespace - ] - - -def _keyword_value(token) -> str: - return token.normalized if token.ttype in Keyword else "" - - -def _is_from_or_join_keyword(token) -> bool: - keyword = _keyword_value(token) - return keyword in {"FROM", "JOIN"} or keyword.endswith(" JOIN") - - -def _is_clause_boundary(token) -> bool: - if isinstance(token, Where): - return True - - keyword = _keyword_value(token) - return keyword in _CLAUSE_BOUNDARY_KEYWORDS or keyword.endswith(" JOIN") - - -def _identifier_name(identifier: Identifier | Function | None) -> str | None: - if identifier is None: - return None - return identifier.get_real_name() or identifier.get_name() - - -def _select_parenthesis(parenthesis: Parenthesis): - for token in parenthesis.tokens: - if isinstance(token, TokenList): - for child in token.flatten(): - if child.ttype is DML and child.normalized == "SELECT": - return token - return None - - -def _identifier_table_name(identifier: Identifier, manifest_tables: set[str]) -> str | None: - name = _identifier_name(identifier) - if name in manifest_tables: - return name - - value_tokens = [] - for token in identifier.tokens: - if token.is_whitespace or token.ttype is Whitespace: - break - if token.ttype in Keyword: - break - value_tokens.append(token.value) - full_name = "".join(value_tokens).strip('"') - return full_name if full_name in manifest_tables else name - - -def _identifier_is_subquery(identifier: Identifier) -> bool: - return any( - isinstance(token, Parenthesis) and _select_parenthesis(token) is not None - for token in identifier.tokens - ) - - -def _subquery_from_identifier(identifier: Identifier): - for token in identifier.tokens: - if isinstance(token, Parenthesis): - statement = _select_parenthesis(token) - if statement is not None: - return statement - return None - - -def _collect_ctes( - statement: TokenList, manifest: dict[str, set[str]], issues: list[str] -) -> set[str]: - cte_names: set[str] = set() - tokens = _non_whitespace_tokens(statement) - if not tokens or _keyword_value(tokens[0]) != "WITH": - return cte_names - - cte_token = tokens[1] if len(tokens) > 1 else None - cte_identifiers = ( - list(cte_token.get_identifiers()) - if isinstance(cte_token, IdentifierList) - else [cte_token] - if isinstance(cte_token, Identifier) - else [] - ) - - for cte_identifier in cte_identifiers: - cte_name = _identifier_name(cte_identifier) - if cte_name: - cte_names.add(cte_name) - subquery = _subquery_from_identifier(cte_identifier) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=cte_names) - - return cte_names - - -def _register_table_identifier( - identifier: Identifier, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if _identifier_is_subquery(identifier): - subquery = _subquery_from_identifier(identifier) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - if alias := identifier.get_alias(): - scope.derived_aliases.add(alias) - return - - table_name = _identifier_table_name(identifier, set(manifest)) - alias = identifier.get_alias() - if table_name in scope.cte_names: - if alias: - scope.derived_aliases.add(alias) - return - - if table_name not in manifest: - issues.append( - f"Generated SQL references table `{table_name}` outside the retrieved Wren schema." - ) - return - - scope.table_aliases[alias or table_name] = table_name - - -def _register_table_token( - token, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if isinstance(token, IdentifierList): - for identifier in token.get_identifiers(): - _register_table_identifier(identifier, scope, manifest, issues) - elif isinstance(token, Identifier): - _register_table_identifier(token, scope, manifest, issues) - elif isinstance(token, Parenthesis): - subquery = _select_parenthesis(token) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - - -def _collect_tables( - statement: TokenList, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - tokens = _non_whitespace_tokens(statement) - index = 0 - - while index < len(tokens): - token = tokens[index] - if _is_from_or_join_keyword(token): - index += 1 - while index < len(tokens) and not _is_clause_boundary(tokens[index]): - _register_table_token(tokens[index], scope, manifest, issues) - index += 1 - continue - - if isinstance(token, Parenthesis): - subquery = _select_parenthesis(token) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - - index += 1 - - -def _collect_select_aliases(statement: TokenList, scope: _SqlScope) -> None: - tokens = _non_whitespace_tokens(statement) - in_select = False - - for token in tokens: - if token.ttype is DML and token.normalized == "SELECT": - in_select = True - continue - if in_select and _keyword_value(token) == "FROM": - return - if not in_select: - continue - - identifiers = ( - token.get_identifiers() - if isinstance(token, IdentifierList) - else [token] - if isinstance(token, Identifier) - else [] - ) - for identifier in identifiers: - if alias := identifier.get_alias(): - scope.output_aliases.add(alias) - - -def _column_is_grounded(column_name: str, scope: _SqlScope, manifest: dict[str, set[str]]) -> bool: - manifest_tables = set(scope.table_aliases.values()) - if not manifest_tables: - return True - - return any(column_name in manifest[table_name] for table_name in manifest_tables) - - -def _validate_identifier_columns( - identifier: Identifier, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if isinstance(identifier, Function): - _validate_token_columns(identifier, scope, manifest, issues) - return - - if any(isinstance(token, Function) for token in identifier.tokens): - for token in identifier.tokens: - if isinstance(token, Function): - _validate_token_columns(token, scope, manifest, issues) - return - - column_name = _identifier_name(identifier) - if not column_name or column_name == "*": - return - if column_name in scope.output_aliases: - return - - parent_name = identifier.get_parent_name() - if parent_name: - if parent_name in scope.derived_aliases or parent_name in scope.cte_names: - return - table_name = scope.table_aliases.get(parent_name, parent_name) - if table_name not in manifest or column_name not in manifest[table_name]: - issues.append( - "Generated SQL references column " - f"`{column_name}` outside the retrieved Wren schema " - f"for table `{table_name}`." - ) - return - - if not _column_is_grounded(column_name, scope, manifest): - table_names = sorted(set(scope.table_aliases.values())) - table_context = ( - f" for table `{table_names[0]}`" if len(table_names) == 1 else "" - ) - issues.append( - f"Generated SQL references column `{column_name}` outside the retrieved Wren schema{table_context}." - ) - - -def _validate_token_columns( - token, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - if isinstance(token, Function): - for child in token.tokens: - if isinstance(child, Parenthesis): - _validate_token_columns(child, scope, manifest, issues) - return - - if isinstance(token, IdentifierList): - for identifier in token.get_identifiers(): - if isinstance(identifier, Identifier): - _validate_identifier_columns(identifier, scope, manifest, issues) - return - - if isinstance(token, Identifier): - _validate_identifier_columns(token, scope, manifest, issues) - return - - if isinstance(token, Parenthesis): - subquery = _select_parenthesis(token) - if subquery is not None: - _validate_statement(subquery, manifest, issues, parent_ctes=scope.cte_names) - return - - if isinstance(token, TokenList): - for child in token.tokens: - _validate_token_columns(child, scope, manifest, issues) - - -def _validate_columns( - statement: TokenList, - scope: _SqlScope, - manifest: dict[str, set[str]], - issues: list[str], -) -> None: - tokens = _non_whitespace_tokens(statement) - index = 0 - - while index < len(tokens): - token = tokens[index] - if _keyword_value(token) == "WITH": - index += 2 - continue - - if _is_from_or_join_keyword(token): - index += 1 - while index < len(tokens) and not _is_clause_boundary(tokens[index]): - index += 1 - continue - - _validate_token_columns(token, scope, manifest, issues) - index += 1 - - -def _validate_statement( - statement: TokenList, - manifest: dict[str, set[str]], - issues: list[str], - parent_ctes: set[str] | None = None, -) -> None: - scope = _SqlScope(cte_names=set(parent_ctes or [])) - scope.cte_names.update(_collect_ctes(statement, manifest, issues)) - _collect_tables(statement, scope, manifest, issues) - _collect_select_aliases(statement, scope) - _validate_columns(statement, scope, manifest, issues) - - -def validate_sql_grounded_in_manifest( - sql: str, schema_manifest: dict[str, list[str]] | None -) -> ManifestGroundingResult: - manifest = _manifest_table_columns(schema_manifest) - if not manifest: - return ManifestGroundingResult(is_grounded=True) - - statements = [statement for statement in sqlparse.parse(sql) if statement.tokens] - if len(statements) != 1: - return ManifestGroundingResult( - is_grounded=False, - error="Generated SQL must contain one grounded SELECT statement.", - ) - - issues: list[str] = [] - _validate_statement(statements[0], manifest, issues) - if issues: - return ManifestGroundingResult(is_grounded=False, error=issues[0]) - - return ManifestGroundingResult(is_grounded=True) - - @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -427,7 +34,6 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - schema_manifest: dict[str, list[str]] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -448,7 +54,6 @@ async def run( allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, allow_data_preview=allow_data_preview, - schema_manifest=schema_manifest, ) return { @@ -471,7 +76,6 @@ async def _classify_generation_result( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - schema_manifest: dict[str, list[str]] | None = None, ) -> Dict[str, str]: valid_generation_result = {} invalid_generation_result = {} @@ -486,18 +90,6 @@ async def _classify_generation_result( "correlation_id": "", } - grounding_result = validate_sql_grounded_in_manifest( - generation_result, schema_manifest - ) - if not grounding_result.is_grounded: - return valid_generation_result, { - "sql": generation_result, - "original_sql": generation_result, - "type": "MANIFEST_GROUNDING", - "error": grounding_result.error, - "correlation_id": "", - } - async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( @@ -611,17 +203,12 @@ async def _classify_generation_result( - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. - Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. - Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. -- Read every retrieved DATABASE SCHEMA object before choosing the SQL source. Retrieval rank only lists candidates; it does not decide the table, view, or metric to query. -- Choose table, view, metric, and column identifiers only from retrieved schema objects whose declared fields support the user's requested subject, output columns, filters, groupings, measures, time concepts, and relationships. -- Do not default to the first retrieved object or to a broad object merely because it contains one requested word. If no retrieved object supports the requested intent with declared identifiers and declared relationships, return null for sql. - When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. - When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. - In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. - Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. -- Treat every declared table and column identifier as an indivisible string. Never splice, recombine, or transfer prefixes, numeric ordinals, suffixes, underscores, casing, punctuation, or business words between different declared identifiers. -- If a declared identifier contains generated prefixes, numeric ordinals, abbreviations, spaces, punctuation, or suffixes, copy the entire identifier exactly as declared. Do not rebuild it from the business meaning. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. - Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. - Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. @@ -916,28 +503,6 @@ def construct_instructions( return _instructions -def construct_executable_identifier_catalog( - schema_manifest: dict[str, list[str]] | None, -) -> str: - if not schema_manifest: - return "" - - lines = [ - "### EXECUTABLE WREN IDENTIFIER CATALOG ###", - "Use this catalog as the compact authoritative list of executable table and column identifiers for the generated SQL.", - "Copy identifiers exactly as written here. Do not derive, rebuild, or recombine table or column names from the user question or semantic descriptions.", - ] - for table_name, column_names in schema_manifest.items(): - if not table_name: - continue - lines.append(f'Table: "{table_name}"') - if column_names: - lines.append("Columns:") - lines.extend(f'- "{column_name}"' for column_name in column_names) - - return "\n".join(lines) - - def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: diff --git a/wren-ai-service/src/pipelines/indexing/project_meta.py b/wren-ai-service/src/pipelines/indexing/project_meta.py index 6105b935fa..426e988934 100644 --- a/wren-ai-service/src/pipelines/indexing/project_meta.py +++ b/wren-ai-service/src/pipelines/indexing/project_meta.py @@ -18,58 +18,6 @@ logger = logging.getLogger("wren-ai-service") -def _column_name(column: dict[str, Any]) -> str | None: - name = column.get("name") - return name if isinstance(name, str) and name else None - - -def _model_column_names(model: dict[str, Any]) -> list[str]: - return [ - name - for column in model.get("columns", []) - if column.get("isHidden") is not True and not column.get("relationship") - for name in [_column_name(column)] - if name - ] - - -def _view_column_names(view: dict[str, Any]) -> list[str]: - properties = view.get("properties") or {} - return [ - name - for column in properties.get("columns", []) - for name in [_column_name(column)] - if name - ] - - -def _metric_column_names(metric: dict[str, Any]) -> list[str]: - return [ - name - for column in metric.get("dimension", []) + metric.get("measure", []) - for name in [_column_name(column)] - if name - ] - - -def build_schema_manifest(mdl: dict[str, Any]) -> dict[str, list[str]]: - manifest: dict[str, list[str]] = {} - - for model in mdl.get("models", []): - if name := model.get("name"): - manifest[name] = _model_column_names(model) - - for view in mdl.get("views", []): - if name := view.get("name"): - manifest[name] = _view_column_names(view) - - for metric in mdl.get("metrics", []): - if name := metric.get("name"): - manifest[name] = _metric_column_names(metric) - - return manifest - - ## Start of Pipeline @observe(capture_input=False, capture_output=False) @extract_fields(dict(mdl=dict[str, Any])) @@ -92,11 +40,7 @@ def chunk( document = Document( id=str(uuid.uuid4()), - meta={ - "data_source": data_source, - "schema_manifest": build_schema_manifest(mdl), - **addition, - }, + meta={"data_source": data_source, **addition}, ) return {"documents": [document]} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index f2c09fdc2e..67986120b7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -104,7 +104,7 @@ def __init__( allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = True, - max_sql_correction_retries: int = 3, + max_sql_correction_retries: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -159,7 +159,6 @@ async def ask( instructions = [] api_results = [] table_names = [] - schema_manifest = {} error_message = None invalid_sql = None allow_sql_generation_reasoning = False @@ -169,7 +168,7 @@ async def ask( allow_sql_functions_retrieval = self._allow_sql_functions_retrieval allow_sql_diagnosis = self._allow_sql_diagnosis allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - max_sql_correction_retries = self._max_sql_correction_retries + max_sql_correction_retries = 0 current_sql_correction_retries = 0 use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback @@ -330,14 +329,6 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] - schema_manifest = { - document.get("table_name"): document.get( - "manifest_column_names", - document.get("column_names", []), - ) - for document in documents - if document.get("table_name") - } if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -460,7 +451,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - schema_manifest=schema_manifest, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -479,7 +469,6 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - schema_manifest=schema_manifest, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -508,7 +497,6 @@ async def ask( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] - error_type = failed_dry_run_result["type"] current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -522,10 +510,7 @@ async def ask( is_followup=True if histories else False, ) - if ( - allow_sql_diagnosis - and error_type != "MANIFEST_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -547,12 +532,10 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ - "type": error_type, "sql": original_sql, "error": ( f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" if allow_sql_diagnosis - and error_type != "MANIFEST_GROUNDING" and sql_diagnosis_reasoning else error_message ), @@ -562,7 +545,6 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - schema_manifest=schema_manifest, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 7c5c7ebcfd..f1971afa15 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -105,7 +105,6 @@ async def ask_feedback( error_message = None invalid_sql = None sql_knowledge = None - schema_manifest = {} allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval try: @@ -162,14 +161,6 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - schema_manifest = { - document.get("table_name"): document.get( - "manifest_column_names", - document.get("column_names", []), - ) - for document in documents - if document.get("table_name") - } sql_samples = sql_samples_task["formatted_output"].get("documents", []) instructions = instructions_task["formatted_output"].get( "documents", [] @@ -196,7 +187,6 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - schema_manifest=schema_manifest, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -217,7 +207,6 @@ async def ask_feedback( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] - error_type = failed_dry_run_result["type"] sql_diagnosis_reasoning = None self._ask_feedback_results[ @@ -227,10 +216,7 @@ async def ask_feedback( trace_id=trace_id, ) - if ( - allow_sql_diagnosis - and error_type != "MANIFEST_GROUNDING" - ): + if allow_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -258,7 +244,6 @@ async def ask_feedback( sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, instructions=instructions, invalid_generation_result={ - "type": error_type, "original_sql": original_sql, "sql": invalid_sql, "error": correction_error_message, @@ -266,7 +251,6 @@ async def ask_feedback( project_id=ask_feedback_request.project_id, sql_functions=sql_functions, sql_knowledge=sql_knowledge, - schema_manifest=schema_manifest, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 3203b3cca2..694d044bfa 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -71,9 +71,7 @@ async def _validate_question( use_dry_plan: bool = True, allow_dry_plan_fallback: bool = False, ): - async def _document_retrieval() -> tuple[ - list[str], dict[str, list[str]], bool, bool, bool - ]: + async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -81,24 +79,10 @@ async def _document_retrieval() -> tuple[ _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - schema_manifest = { - document.get("table_name"): document.get( - "manifest_column_names", - document.get("column_names", []), - ) - for document in documents - if document.get("table_name") - } has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return ( - table_ddls, - schema_manifest, - has_calculated_field, - has_metric, - has_json_field, - ) + return table_ddls, has_calculated_field, has_metric, has_json_field async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -123,13 +107,7 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - ( - table_ddls, - schema_manifest, - has_calculated_field, - has_metric, - has_json_field, - ) = _document + table_ddls, has_calculated_field, has_metric, has_json_field = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -159,7 +137,6 @@ async def _instructions_retrieval() -> list[dict]: allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, - schema_manifest=schema_manifest, ) post_process = generated_sql["post_process"] diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 6089c6e62d..f805024ad2 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -113,14 +113,6 @@ async def correct( .get("retrieval_results", []) ) table_ddls = [document.get("table_ddl") for document in documents] - schema_manifest = { - document.get("table_name"): document.get( - "manifest_column_names", - document.get("column_names", []), - ) - for document in documents - if document.get("table_name") - } res = await self._pipelines["sql_correction"].run( contexts=table_ddls, @@ -129,7 +121,6 @@ async def correct( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, - schema_manifest=schema_manifest, ) post_process = res["post_process"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py index 2118029220..0f4ed59d26 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py @@ -5,7 +5,6 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, - construct_executable_identifier_catalog, construct_instructions, get_json_field_instructions, get_metric_instructions, @@ -22,7 +21,6 @@ ) from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template -from src.pipelines.generation.sql_generation import post_process as sql_post_process from src.pipelines.generation.sql_regeneration import get_sql_regeneration_system_prompt from src.pipelines.generation.sql_regeneration import sql_regeneration_user_prompt_template @@ -120,158 +118,6 @@ async def test_sql_postprocessor_rejects_null_sql_generation_result(): assert result["invalid_generation_result"]["sql"] == "" -@pytest.mark.asyncio -async def test_sql_postprocessor_rejects_table_outside_retrieved_manifest(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=['{"sql": "SELECT \\"AvailableField\\" FROM \\"UnretrievedObject\\""}'], - use_dry_plan=True, - data_source="source", - schema_manifest={"RetrievedObject": ["AvailableField"]}, - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert "UnretrievedObject" in result["invalid_generation_result"]["error"] - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] - - -@pytest.mark.asyncio -async def test_sql_postprocessor_rejects_column_outside_retrieved_manifest(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=['{"sql": "SELECT \\"UnretrievedField\\" FROM \\"RetrievedObject\\""}'], - use_dry_plan=True, - data_source="source", - schema_manifest={"RetrievedObject": ["AvailableField"]}, - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert "UnretrievedField" in result["invalid_generation_result"]["error"] - assert "RetrievedObject" in result["invalid_generation_result"]["error"] - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] - - -@pytest.mark.asyncio -async def test_sql_generation_post_process_uses_separate_grounding_manifest(): - engine = _DryPlanEngine() - - result = await sql_post_process( - generate_sql={ - "replies": ['{"sql": "SELECT \\"ActiveField\\" FROM \\"ActiveEntity\\""}'] - }, - post_processor=SQLGenPostProcessor(engine), - data_source="source", - use_dry_plan=True, - grounding_schema_manifest={"ActiveEntity": ["ActiveField"]}, - ) - - assert result["valid_generation_result"]["sql"] == ( - 'SELECT "ActiveField" FROM "ActiveEntity"' - ) - assert len(engine.dry_plan_calls) == 1 - assert len(engine.execute_sql_calls) == 1 - - -@pytest.mark.asyncio -async def test_sql_postprocessor_rejects_filter_column_outside_retrieved_manifest(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=[ - ( - '{"sql": "SELECT * FROM \\"RetrievedObject\\" ' - 'WHERE \\"UnretrievedDate\\" >= CURRENT_DATE"}' - ) - ], - use_dry_plan=True, - data_source="source", - schema_manifest={"RetrievedObject": ["AvailableField"]}, - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] - - -@pytest.mark.asyncio -async def test_sql_postprocessor_rejects_function_argument_column_outside_manifest(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=[ - ( - '{"sql": "SELECT * FROM \\"RetrievedObject\\" ' - "WHERE DATE_TRUNC('month', \\\"UnretrievedDate\\\") = CURRENT_DATE\"}" - ) - ], - use_dry_plan=True, - data_source="source", - schema_manifest={"RetrievedObject": ["AvailableField"]}, - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "MANIFEST_GROUNDING" - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] - - -@pytest.mark.asyncio -async def test_sql_postprocessor_allows_exact_retrieved_manifest_identifiers(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=[ - ( - '{"sql": "SELECT SUM(\\"AvailableField\\") AS \\"TotalField\\" ' - 'FROM \\"RetrievedObject\\""}' - ) - ], - use_dry_plan=True, - data_source="source", - schema_manifest={"RetrievedObject": ["AvailableField"]}, - ) - - assert result["valid_generation_result"]["sql"] == ( - 'SELECT SUM("AvailableField") AS "TotalField" FROM "RetrievedObject"' - ) - assert len(engine.dry_plan_calls) == 1 - assert len(engine.execute_sql_calls) == 1 - - -@pytest.mark.asyncio -async def test_sql_postprocessor_validates_join_predicate_columns(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=[ - ( - '{"sql": "SELECT a.\\"AvailableField\\" FROM \\"RetrievedObject\\" a ' - 'JOIN \\"RelatedObject\\" b ON a.\\"JoinField\\" = b.\\"JoinField\\""}' - ) - ], - use_dry_plan=True, - data_source="source", - schema_manifest={ - "RetrievedObject": ["AvailableField", "JoinField"], - "RelatedObject": ["JoinField"], - }, - ) - - assert result["valid_generation_result"]["sql"] == ( - 'SELECT a."AvailableField" FROM "RetrievedObject" a JOIN ' - '"RelatedObject" b ON a."JoinField" = b."JoinField"' - ) - assert len(engine.dry_plan_calls) == 1 - assert len(engine.execute_sql_calls) == 1 - - def test_construct_instructions_uses_instruction_text(): assert construct_instructions( [{"instruction": "First rule."}, {"instruction": "Second rule."}] @@ -441,79 +287,6 @@ def test_sql_correction_system_prompt_discards_invalid_identifier_context(): assert "return null for sql instead of substituting non-schema identifiers" in prompt -def test_sql_generation_prompt_requires_schema_object_selection_before_sql(): - prompt = get_sql_generation_system_prompt() - - assert "Read every retrieved DATABASE SCHEMA object before choosing" in prompt - assert "Retrieval rank only lists candidates" in prompt - assert "Do not default to the first retrieved object" in prompt - assert "declared fields support the user's requested subject" in prompt - - -def test_sql_generation_prompt_treats_identifiers_as_indivisible_strings(): - prompt = get_sql_generation_system_prompt() - - assert "Treat every declared table and column identifier as an indivisible string" in prompt - assert "Never splice, recombine, or transfer prefixes" in prompt - assert "copy the entire identifier exactly as declared" in prompt - assert "Do not rebuild it from the business meaning" in prompt - - -def test_construct_executable_identifier_catalog_lists_manifest_identifiers(): - catalog = construct_executable_identifier_catalog( - {"ObjectA": ["FieldA", "FieldB"], "ObjectB": ["FieldC"]} - ) - - assert "EXECUTABLE WREN IDENTIFIER CATALOG" in catalog - assert 'Table: "ObjectA"' in catalog - assert '- "FieldA"' in catalog - assert 'Table: "ObjectB"' in catalog - assert "Copy identifiers exactly as written here" in catalog - - -def test_sql_generation_prompt_can_include_executable_identifier_catalog(): - catalog = construct_executable_identifier_catalog({"ObjectA": ["FieldA"]}) - prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( - query="Question", - documents=["SCHEMA_CONTEXT"], - executable_identifier_catalog=catalog, - sql_generation_reasoning=None, - instructions=[], - calculated_field_instructions="", - metric_instructions="", - json_field_instructions="", - sql_samples=[], - sql_functions=[], - )["prompt"] - - assert "SCHEMA_CONTEXT" in prompt - assert "EXECUTABLE WREN IDENTIFIER CATALOG" in prompt - assert 'Table: "ObjectA"' in prompt - assert '- "FieldA"' in prompt - - -def test_sql_correction_prompt_includes_manifest_grounding_failure(): - prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( - query="Question", - documents=["SCHEMA_CONTEXT"], - invalid_generation_result={ - "type": "MANIFEST_GROUNDING", - "error": ( - "Generated SQL references column `RejectedField` outside the " - "retrieved Wren schema for table `RetrievedObject`." - ), - }, - sql_generation_reasoning=None, - instructions=[], - sql_functions=[], - )["prompt"] - - assert "MANIFEST GROUNDING FAILURE" in prompt - assert "RejectedField" in prompt - assert "RetrievedObject" in prompt - assert "Do not reuse rejected identifiers" in prompt - - def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): prompt = sql_generation_reasoning_system_prompt @@ -635,7 +408,6 @@ def test_executable_prompt_templates_omit_untrusted_reasoning_and_sql_context(): assert "dry-run diagnostic text is intentionally omitted" in ( correction_prompt ) - assert ( - "Regenerate from the user question and current DATABASE SCHEMA or " - "EXECUTABLE WREN IDENTIFIER CATALOG only" - ) in correction_prompt + assert "Regenerate from the user question and current DATABASE SCHEMA only" in ( + correction_prompt + ) diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py deleted file mode 100644 index 3a088044cb..0000000000 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_project_meta.py +++ /dev/null @@ -1,56 +0,0 @@ -from src.pipelines.indexing.project_meta import build_schema_manifest, chunk - - -def test_project_meta_schema_manifest_uses_deployed_mdl_identifiers(): - mdl = { - "dataSource": "postgres", - "models": [ - { - "name": "PrimaryEntity", - "columns": [ - {"name": "VisibleField"}, - {"name": "HiddenField", "isHidden": True}, - {"name": "LinkedField", "relationship": "RelatedEntity"}, - ], - } - ], - "views": [ - { - "name": "SavedView", - "properties": { - "columns": [ - {"name": "ViewField"}, - ], - }, - } - ], - "metrics": [ - { - "name": "MetricEntity", - "dimension": [{"name": "DimensionField"}], - "measure": [{"name": "MeasureField"}], - } - ], - } - - assert build_schema_manifest(mdl) == { - "PrimaryEntity": ["VisibleField"], - "SavedView": ["ViewField"], - "MetricEntity": ["DimensionField", "MeasureField"], - } - - -def test_project_meta_chunk_stores_schema_manifest_with_project_scope(): - result = chunk( - mdl={ - "dataSource": "duckdb", - "models": [{"name": "Entity", "columns": [{"name": "Field"}]}], - }, - project_id="project-a", - ) - - document = result["documents"][0] - - assert document.meta["data_source"] == "local_file" - assert document.meta["project_id"] == "project-a" - assert document.meta["schema_manifest"] == {"Entity": ["Field"]} diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py index d68135b8d3..d058274eff 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py @@ -1,12 +1,6 @@ import pytest -from haystack import Document - -from src.pipelines.common import ( - resolve_active_schema_manifest, - resolve_schema_manifest, - retrieve_metadata, - retrieve_schema_manifest, -) + +from src.pipelines.common import retrieve_metadata from src.pipelines.retrieval import historical_question_retrieval, instructions from src.pipelines.retrieval import sql_pairs_retrieval @@ -34,7 +28,7 @@ def __init__(self, documents=None): self.documents = documents or [] self.calls = [] - async def run(self, query_embedding=None, filters=None, **_): + async def run(self, query_embedding=None, filters=None): self.calls.append( { "query_embedding": query_embedding, @@ -54,130 +48,6 @@ async def test_metadata_retrieval_does_not_fall_back_to_global_documents(): assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] -def test_active_project_manifest_overrides_retrieved_subset(): - assert resolve_schema_manifest( - metadata={"schema_manifest": {"ActiveEntity": ["ActiveField"]}}, - schema_manifest={"RetrievedEntity": ["RetrievedField"]}, - ) == {"ActiveEntity": ["ActiveField"]} - - -def test_retrieved_manifest_is_used_when_project_metadata_has_no_manifest(): - assert resolve_schema_manifest( - metadata={"data_source": "source"}, - schema_manifest={"RetrievedEntity": ["RetrievedField"]}, - ) == {"RetrievedEntity": ["RetrievedField"]} - - -@pytest.mark.asyncio -async def test_schema_manifest_can_be_rebuilt_from_indexed_schema_documents(): - retriever = RetrieverSpy( - documents=[ - Document( - content=str( - { - "type": "TABLE_COLUMNS", - "columns": [ - { - "type": "COLUMN", - "name": "FirstField", - "data_type": "varchar", - }, - { - "type": "FOREIGN_KEY", - "name": "RelationField", - "data_type": "varchar", - }, - { - "type": "COLUMN", - "name": "UnknownField", - "data_type": "unknown", - }, - ], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "IndexedEntity"}, - ), - Document( - content=str( - { - "type": "TABLE_COLUMNS", - "columns": [ - { - "type": "COLUMN", - "name": "SecondField", - "data_type": "integer", - }, - ], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "IndexedEntity"}, - ), - Document( - content=str( - { - "type": "VIEW", - "name": "IndexedView", - "columns": [ - { - "name": "ViewField", - "data_type": "varchar", - }, - ], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "IndexedView"}, - ), - ] - ) - - assert await retrieve_schema_manifest("project-a", retriever) == { - "IndexedEntity": ["FirstField", "SecondField"], - "IndexedView": ["ViewField"], - } - assert retriever.calls == [ - { - "query_embedding": [], - "filters": { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"field": "project_id", "operator": "==", "value": "project-a"}, - ], - }, - } - ] - - -@pytest.mark.asyncio -async def test_active_manifest_falls_back_to_indexed_schema_documents(): - retriever = RetrieverSpy( - documents=[ - Document( - content=str( - { - "type": "TABLE_COLUMNS", - "columns": [ - { - "type": "COLUMN", - "name": "ActiveField", - "data_type": "varchar", - }, - ], - } - ), - meta={"type": "TABLE_SCHEMA", "name": "ActiveEntity"}, - ) - ] - ) - - assert await resolve_active_schema_manifest( - metadata={"data_source": "source"}, - schema_manifest={"RetrievedEntity": ["RetrievedField"]}, - project_id="project-a", - dbschema_retriever=retriever, - ) == {"ActiveEntity": ["ActiveField"]} - - @pytest.mark.asyncio async def test_sql_pairs_count_stays_project_scoped_when_project_has_no_documents(): store = StoreSpy(count=0) diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py index 244bfe6a5d..f5f15b14bf 100644 --- a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py +++ b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py @@ -36,7 +36,7 @@ def test_ask_service_defaults_to_column_pruning(): assert service._enable_column_pruning is True assert service._allow_sql_generation_reasoning is False - assert service._max_sql_correction_retries == 3 + assert service._max_sql_correction_retries == 0 @pytest.mark.asyncio @@ -59,20 +59,9 @@ async def test_sql_correction_service_requires_retrieved_tables(): @pytest.mark.asyncio -async def test_ask_service_skips_speculative_reasoning_and_corrects_with_retrieved_schema(): +async def test_ask_service_does_not_run_speculative_reasoning_or_correction_after_failed_generation(): sql_generation_reasoning = AsyncMock(run=AsyncMock(return_value={})) - sql_correction = AsyncMock( - run=AsyncMock( - return_value={ - "post_process": { - "valid_generation_result": { - "sql": "SELECT retrieved_field FROM retrieved_model" - }, - "invalid_generation_result": {}, - } - } - ) - ) + sql_correction = AsyncMock(run=AsyncMock(return_value={})) service = AskService( { "sql_pairs_retrieval": AsyncMock( @@ -119,7 +108,6 @@ async def test_ask_service_skips_speculative_reasoning_and_corrects_with_retriev "sql_correction": sql_correction, }, allow_intent_classification=False, - allow_sql_diagnosis=False, ) request = AskRequest(query="Can this be answered?", id="deploy-id") request.query_id = "query-id" @@ -127,14 +115,11 @@ async def test_ask_service_skips_speculative_reasoning_and_corrects_with_retriev await service.ask(request) result = service.get_ask_result(AskResultRequest(query_id="query-id")) - assert result.status == "finished" - assert result.response[0].sql == "SELECT retrieved_field FROM retrieved_model" + assert result.status == "failed" + assert result.error.code == "NO_RELEVANT_SQL" assert result.invalid_sql is None sql_generation_reasoning.run.assert_not_called() - sql_correction.run.assert_called_once() - assert sql_correction.run.call_args.kwargs["contexts"] == [ - "CREATE TABLE retrieved_model (retrieved_field VARCHAR);" - ] + sql_correction.run.assert_not_called() def test_sql_correction_router_defaults_to_planner_validation_without_fallback(): From ae1efdf6bbcee2deaab18da943976f495b236dc5 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 17:40:01 +0530 Subject: [PATCH 0784/1087] Restore official legacy ask generation flow --- .../pipelines/generation/sql_correction.py | 42 +- .../pipelines/generation/sql_generation.py | 13 +- .../retrieval/db_schema_retrieval.py | 585 +++--------------- wren-ai-service/src/web/v1/services/ask.py | 62 +- .../pipelines/generation/test_sql_utils.py | 413 ------------- .../pytest/services/test_dry_plan_defaults.py | 189 ------ 6 files changed, 126 insertions(+), 1178 deletions(-) delete mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py delete mode 100644 wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index e92b8fc648..973b8c69a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,20 +30,12 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are a Wren SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. +You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query. ### SQL CORRECTION INSTRUCTIONS ### -1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. -2. Then, generate a syntactically correct Wren SQL query from the user's intent and the current DATABASE SCHEMA. -3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. -4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. -5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. -6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. -7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. -8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA. If the unsupported part is needed to answer the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql instead of substituting non-schema identifiers. -10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. +1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). +2. Then, generate the syntactically correct ANSI SQL query to correct the error. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -51,10 +43,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. +The final answer must be in JSON format: {{ - "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" + "sql": }} """ @@ -82,20 +74,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -{% if query %} -User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. -{% endif %} -### FAILED SQL ### -The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. - -### DRY-RUN DIAGNOSTIC ### -The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. - -Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. +SQL: {{ invalid_generation_result.sql }} +Error Message: {{ invalid_generation_result.error }} -Return only the final JSON SQL response. +Let's think step by step. """ @@ -105,16 +87,12 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, - query: str | None = None, - sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( - query=query, documents=documents, invalid_generation_result=invalid_generation_result, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -191,8 +169,6 @@ async def run( self, contexts: List[Document], invalid_generation_result: Dict[str, str], - query: str | None = None, - sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, @@ -211,9 +187,7 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, - "query": query, "documents": contexts, - "sql_generation_reasoning": sql_generation_reasoning, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 2640082430..1ee4952b3e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -54,10 +54,11 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} +SQL: +{{sample.sql}} {% endfor %} {% endif %} @@ -70,11 +71,13 @@ ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. -Return only the final JSON SQL response. +{% if sql_generation_reasoning %} +### REASONING PLAN ### +{{ sql_generation_reasoning }} +{% endif %} + +Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 75941d5ea5..6c8dd7bbe3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -29,13 +29,7 @@ ### TASK ### You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. -The database schema includes structural, semantic, and business modeling metadata: -- Models are logical datasets backed by physical tables or SQL definitions. -- Columns are exposed fields, including renamed fields, expressions, primary keys, and calculated fields. -- Relationships are reusable join logic between models. -- Calculated fields are business logic defined once and reused across queries. -- Views are named SQL statements that behave like stable virtual tables. -- Metrics are structured aggregation objects with measures and dimensions. +The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -45,16 +39,6 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -8. Map the business question to the modeled datasets whose descriptions, aliases, columns, calculated fields, views, metrics, and relationships support the intent. -9. Prefer modeled analytical interfaces such as views and metrics when they expose the fields needed to answer the question. -10. If the answer needs fields, filters, time dimensions, ordering, aggregations, or relationship keys from multiple related datasets, include every required related dataset and the columns needed from each one. -11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. -12. Follow only the relationships shown in the provided schema when selecting columns across datasets. -13. Do not stop at a single top candidate when the question needs multiple related datasets. -14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. -15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. -16. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. -17. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -115,292 +99,38 @@ """ -def _build_metric_ddl(content: dict, include_semantic_context: bool = True) -> str: - columns = [ - column - for column in content["columns"] - if column["data_type"].lower() != "unknown" - ] - context = { - "object_type": "metric", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [column["name"] for column in columns], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable analytical aggregation interface", - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "semantic_context_not_sql_identifier": column["comment"], - } - for column in columns - ], - } - schema_context = ( - _format_semantic_context(context) - if include_semantic_context - else _format_executable_identifier_catalog(context) - ) +def _build_metric_ddl(content: dict) -> str: columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" - for column in columns - ] - - return ( - f"{schema_context}CREATE TABLE {content['name']} (\n " - + ",\n ".join(columns_ddl) - + "\n);" - ) - - -def _build_view_ddl(content: dict, include_semantic_context: bool = True) -> str: - columns = [ - column - for column in content.get("columns", []) - if column.get("name") and column.get("data_type", "").lower() != "unknown" - ] - context = { - "object_type": "view", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [column["name"] for column in columns], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable virtual table interface", - "description": content["comment"], - "definition_omitted_from_executable_schema": True, - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column.get("data_type")), - "semantic_context_not_sql_identifier": column.get("comment", ""), - } - for column in columns - ], - } - schema_context = ( - _format_semantic_context(context) - if include_semantic_context - else _format_executable_identifier_catalog(context) - ) - columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" - for column in columns + f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in content["columns"] + if column["data_type"].lower() + != "unknown" # quick fix: filtering out UNKNOWN column type ] return ( - f"{schema_context}CREATE TABLE {content['name']} (\n " + f"{content['comment']}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) -def _format_semantic_context(context: dict) -> str: +def _build_view_ddl(content: dict) -> str: return ( - "/*\n" - "WREN RETRIEVED SEMANTIC CONTEXT\n" - f"{orjson.dumps(context).decode('utf-8')}\n" - f"{_format_identifier_contract(context)}" - "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" - "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" - "*/\n" - f"{_format_executable_identifier_catalog(context)}" - ) - - -def _format_executable_identifier_catalog(context: dict) -> str: - contract = context.get("sql_identifier_contract", {}) - table_name = contract.get("sql_table_name_use_exactly") - column_names = contract.get("sql_column_names_use_exactly") or [ - column["sql_column_name_use_exactly"] - for column in context.get("columns", []) - if column.get("sql_column_name_use_exactly") - ] - relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ - relationship["sql_relationship_constraint_use_exactly"] - for relationship in context.get("relationships", []) - if relationship.get("sql_relationship_constraint_use_exactly") - ] - - lines = [ - "### EXECUTABLE WREN IDENTIFIER CATALOG ###", - "Copy SQL identifiers only from this catalog or the following DDL.", - "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", - f"object_type: {context.get('object_type', '')}", - ] - if table_name: - lines.append(f"table: {table_name}") - if column_names: - lines.append("columns:") - lines.extend(f"- {column_name}" for column_name in column_names) - if relationship_constraints: - lines.append("relationships:") - lines.extend(f"- {constraint}" for constraint in relationship_constraints) - lines.extend( - [ - "If a needed table, column, or relationship is not listed here or declared in the following DDL, return null for sql.", - "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", - "", - ] - ) - return "\n".join(lines) - - -def _format_identifier_contract(context: dict) -> str: - contract = context.get("sql_identifier_contract", {}) - table_name = contract.get("sql_table_name_use_exactly") - column_names = contract.get("sql_column_names_use_exactly") or [ - column["sql_column_name_use_exactly"] - for column in context.get("columns", []) - if column.get("sql_column_name_use_exactly") - ] - relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ - relationship["sql_relationship_constraint_use_exactly"] - for relationship in context.get("relationships", []) - if relationship.get("sql_relationship_constraint_use_exactly") - ] - - lines = [ - "WREN SQL IDENTIFIER CONTRACT", - f"object_type: {context.get('object_type', '')}", - ] - if table_name: - lines.append(f"sql_table_name_use_exactly: {table_name}") - if column_names: - lines.append("sql_column_names_use_exactly:") - lines.extend(f"- {column_name}" for column_name in column_names) - if relationship_constraints: - lines.append("relationship_constraints_use_exactly:") - lines.extend( - f"- {relationship_constraint}" - for relationship_constraint in relationship_constraints - ) - lines.extend( - [ - "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", - "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", - "END WREN SQL IDENTIFIER CONTRACT", - "", - ] + f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" ) - return "\n".join(lines) - - -def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: - relationship_columns = { - column.get("column") - for column in content["columns"] - if column["type"] == "FOREIGN_KEY" - and (not tables or set(column.get("tables", [])).issubset(tables)) - } - relationship_columns.discard(None) - return relationship_columns - - -def _included_columns( - content: dict, columns: Optional[set[str]], tables: Optional[set[str]] -) -> list[dict]: - relationship_columns = _included_relationship_columns(content, tables) - return [ - column - for column in content["columns"] - if column["type"] == "COLUMN" - and ( - not columns - or column["name"] in columns - or column["name"] in relationship_columns - or column["is_primary_key"] - ) - and column["data_type"].lower() != "unknown" - ] - - -def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: - return [ - column - for column in content["columns"] - if column["type"] == "FOREIGN_KEY" - and (not tables or set(column.get("tables", [])).issubset(tables)) - ] - - -def _included_column_names( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None -) -> list[str]: - return [column["name"] for column in _included_columns(content, columns, tables)] - - -def _executable_column_names(content: dict) -> list[str]: - return [ - column["name"] - for column in content["columns"] - if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" - ] - - -def _build_table_retrieval_context( - content: dict, - columns: Optional[set[str]] = None, - tables: Optional[set[str]] = None, - include_semantic_context: bool = True, -) -> tuple[str, bool, bool]: - ddl, has_calculated_field, has_json_field = build_table_ddl( - content, - columns=columns, - tables=tables, - include_semantic_comments=False, - ) - included_columns = _included_columns(content, columns, tables) - included_relationships = _included_relationships(content, tables) - context = { - "object_type": "model", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in included_columns - ], - "relationship_constraints_use_exactly": [ - relationship["constraint"] for relationship in included_relationships - ], - }, - "semantic_context_not_sql_identifiers": { - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "is_primary_key": column["is_primary_key"], - "semantic_context_not_sql_identifier": column["comment"], - } - for column in included_columns - ], - "relationships": [ - { - "semantic_context_not_sql_identifier": relationship["comment"], - "sql_relationship_constraint_use_exactly": relationship["constraint"], - "related_models_use_exactly": relationship.get("tables", []), - } - for relationship in included_relationships - ], - } - schema_context = ( - _format_semantic_context(context) - if include_semantic_context - else _format_executable_identifier_catalog(context) - ) - return f"{schema_context}{ddl}", has_calculated_field, has_json_field ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: + if histories: + previous_query_summaries = [history.question for history in histories] + else: + previous_query_summaries = [] + + query = "\n".join(previous_query_summaries) + "\n" + query + return await embedder.run(query) else: return {} @@ -440,171 +170,37 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, project_id: str, dbschema_retriever: Any, embedding: dict -) -> list[Document]: - table_names = _table_names_from_description_documents( - table_retrieval.get("documents", []) - ) - documents = [] - if embedding and not table_names: - documents = await _retrieve_semantic_schema_documents( - embedding, project_id, dbschema_retriever - ) - table_names = _table_names_from_schema_documents(documents) - - if table_names: - retrieved_table_names = set() - pending_table_names = table_names - - while pending_table_names: - retrieved_table_names.update(pending_table_names) - retrieved_documents = await _retrieve_schema_documents( - pending_table_names, project_id, dbschema_retriever - ) - documents = _dedupe_documents(documents + retrieved_documents) - pending_table_names = [ - table_name - for table_name in _related_table_names(documents) - if table_name not in retrieved_table_names - ] - - return documents - - return [] - - -async def _retrieve_semantic_schema_documents( - embedding: dict, project_id: str, dbschema_retriever: Any + table_retrieval: dict, project_id: str, dbschema_retriever: Any ) -> list[Document]: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - } - - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) - - results = await dbschema_retriever.run( - query_embedding=embedding.get("embedding"), - filters=filters, - ) - return results["documents"] - - -def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: + tables = table_retrieval.get("documents", []) table_names = [] - seen = set() - - for document in documents: - table_name = document.meta.get("name") - if not table_name: - content = ast.literal_eval(document.content) - table_name = content.get("name") + for table in tables: + content = ast.literal_eval(table.content) + table_names.append(content["name"]) - if table_name and table_name not in seen: - table_names.append(table_name) - seen.add(table_name) - - return table_names - - -def _merge_names(*name_groups: list[str]) -> list[str]: - merged = [] - seen = set() - - for names in name_groups: - for name in names: - if name in seen: - continue - merged.append(name) - seen.add(name) - - return merged - - -def _table_names_from_description_documents(documents: list[Document]) -> list[str]: - table_names = [] - seen = set() - - for document in documents: - content = ast.literal_eval(document.content) - table_name = content["name"] - if table_name not in seen: - table_names.append(table_name) - seen.add(table_name) - - return table_names - - -async def _retrieve_schema_documents( - table_names: list[str], project_id: str, dbschema_retriever: Any -) -> list[Document]: table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names ] - if not table_name_conditions: - return [] - - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } - - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) - - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] - - -def _related_table_names(documents: list[Document]) -> list[str]: - related_table_names = [] - seen = set() - - for document in documents: - content = ast.literal_eval(document.content) - if content.get("type") != "TABLE_COLUMNS": - continue - - for column in content.get("columns", []): - if column.get("type") != "FOREIGN_KEY": - continue - - for table_name in column.get("tables", []): - if table_name not in seen: - related_table_names.append(table_name) - seen.add(table_name) - - return related_table_names - + if table_name_conditions: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } -def _dedupe_documents(documents: list[Document]) -> list[Document]: - deduped = [] - seen = set() + if project_id: + filters["conditions"].append( + {"field": "project_id", "operator": "==", "value": project_id} + ) - for document in documents: - identity = ( - document.meta.get("type"), - document.meta.get("name"), - document.content, - ) - if identity in seen: - continue - deduped.append(document) - seen.add(identity) + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] - return deduped + return [] @observe() @@ -650,17 +246,11 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context( - table_schema, - include_semantic_context=False, - ) - ) + ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) retrieval_results.append( { "table_name": table_schema["name"], "table_ddl": ddl, - "column_names": _included_column_names(table_schema), } ) if _has_calculated_field: @@ -675,15 +265,7 @@ def check_using_db_schemas_without_pruning( retrieval_results.append( { "table_name": content["name"], - "table_ddl": _build_metric_ddl( - content, - include_semantic_context=False, - ), - "column_names": [ - column["name"] - for column in content["columns"] - if column["data_type"].lower() != "unknown" - ], + "table_ddl": _build_metric_ddl(content), } ) has_metric = True @@ -691,16 +273,7 @@ def check_using_db_schemas_without_pruning( retrieval_results.append( { "table_name": content["name"], - "table_ddl": _build_view_ddl( - content, - include_semantic_context=False, - ), - "column_names": [ - column["name"] - for column in content.get("columns", []) - if column.get("name") - and column.get("data_type", "").lower() != "unknown" - ], + "table_ddl": _build_view_ddl(content), } ) @@ -736,10 +309,16 @@ def prompt( ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ - _build_table_retrieval_context(construct_db_schema)[0] + build_table_ddl(construct_db_schema)[0] for construct_db_schema in construct_db_schemas ] + previous_query_summaries = ( + [history.question for history in histories] if histories else [] + ) + + query = "\n".join(previous_query_summaries) + "\n" + query + _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: @@ -785,12 +364,12 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context( - table_schema, - tables=tables, - include_semantic_context=True, - ) + ddl, _has_calculated_field, _has_json_field = build_table_ddl( + table_schema, + columns=set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ), + tables=tables, ) if _has_calculated_field: has_calculated_field = True @@ -801,50 +380,28 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, - "column_names": _included_column_names( - table_schema, tables=tables - ), - "manifest_column_names": _executable_column_names( - table_schema - ), } ) for document in dbschema_retrieval: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl( - content, - include_semantic_context=False, - ), - "column_names": [ - column["name"] - for column in content["columns"] - if column["data_type"].lower() != "unknown" - ], - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl( - content, - include_semantic_context=False, - ), - "column_names": [ - column["name"] - for column in content.get("columns", []) - if column.get("name") - and column.get("data_type", "").lower() != "unknown" - ], - } - ) + if document.meta["name"] in columns_and_tables_needed: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) return { "retrieval_results": retrieval_results, @@ -898,7 +455,7 @@ def __init__( llm_provider: LLMProvider, embedder_provider: EmbedderProvider, document_store_provider: DocumentStoreProvider, - table_retrieval_size: int = 50, + table_retrieval_size: int = 10, table_column_retrieval_size: int = 100, **kwargs, ): diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 67986120b7..aa26fa3f81 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -25,10 +25,10 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = True + ignore_sql_generation_reasoning: bool = False enable_column_pruning: bool = False - use_dry_plan: bool = True - allow_dry_plan_fallback: bool = False + use_dry_plan: bool = False + allow_dry_plan_fallback: bool = True custom_instruction: Optional[str] = None @@ -99,12 +99,12 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = False, + allow_sql_generation_reasoning: bool = True, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = True, - max_sql_correction_retries: int = 0, + enable_column_pruning: bool = False, + max_sql_correction_retries: int = 3, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -161,14 +161,17 @@ async def ask( table_names = [] error_message = None invalid_sql = None - allow_sql_generation_reasoning = False + allow_sql_generation_reasoning = ( + self._allow_sql_generation_reasoning + and not ask_request.ignore_sql_generation_reasoning + ) enable_column_pruning = ( self._enable_column_pruning or ask_request.enable_column_pruning ) allow_sql_functions_retrieval = self._allow_sql_functions_retrieval allow_sql_diagnosis = self._allow_sql_diagnosis allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - max_sql_correction_retries = 0 + max_sql_correction_retries = self._max_sql_correction_retries current_sql_correction_retries = 0 use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback @@ -186,7 +189,29 @@ async def ask( is_followup=True if histories else False, ) - if not api_results: + historical_question = await self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + ) + + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] + + if historical_question_result: + api_results = [ + AskResult( + **{ + "sql": result.get("statement"), + "type": "view" if result.get("viewId") else "llm", + "viewId": result.get("viewId"), + } + ) + for result in historical_question_result + ] + sql_generation_reasoning = "" + else: # Run both pipeline operations concurrently sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( @@ -200,6 +225,7 @@ async def ask( ), ) + # Extract results from completed tasks sql_samples = sql_samples_task["formatted_output"].get( "documents", [] ) @@ -486,12 +512,7 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] in ( - "TIME_OUT", - "NO_RELEVANT_SQL", - ): - error_message = failed_dry_run_result["error"] - invalid_sql = failed_dry_run_result["sql"] + if failed_dry_run_result["type"] == "TIME_OUT": break original_sql = failed_dry_run_result["original_sql"] @@ -528,17 +549,12 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - query=user_query, - sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": ( - f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" - if allow_sql_diagnosis - and sql_diagnosis_reasoning - else error_message - ), + "error": sql_diagnosis_reasoning + if allow_sql_diagnosis + else error_message, }, project_id=ask_request.project_id, use_dry_plan=use_dry_plan, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py deleted file mode 100644 index 0f4ed59d26..0000000000 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_utils.py +++ /dev/null @@ -1,413 +0,0 @@ -import pytest -from haystack.components.builders.prompt_builder import PromptBuilder - -from src.pipelines.generation.utils.sql import ( - SQL_GENERATION_MODEL_KWARGS, - SQLGenPostProcessor, - construct_ask_history_messages, - construct_instructions, - get_json_field_instructions, - get_metric_instructions, - get_sql_generation_system_prompt, - get_text_to_sql_rules, - sql_generation_reasoning_system_prompt, -) -from src.pipelines.generation.followup_sql_generation import ( - text_to_sql_with_followup_user_prompt_template, -) -from src.pipelines.generation.sql_correction import ( - get_sql_correction_system_prompt, - sql_correction_user_prompt_template, -) -from src.pipelines.generation.sql_answer import sql_to_answer_system_prompt -from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template -from src.pipelines.generation.sql_regeneration import get_sql_regeneration_system_prompt -from src.pipelines.generation.sql_regeneration import sql_regeneration_user_prompt_template - - -class _SqlKnowledge: - text_to_sql_rule = "Use the supplied model context only." - metric_instructions = "Use the supplied metric definitions only." - json_field_instructions = "Use the supplied JSON field definitions only." - - -class _DryPlanEngine: - def __init__(self): - self.dry_plan_calls = [] - self.execute_sql_calls = [] - - async def dry_plan(self, *args, **kwargs): - self.dry_plan_calls.append((args, kwargs)) - return True, "" - - async def execute_sql(self, *args, **kwargs): - self.execute_sql_calls.append((args, kwargs)) - return True, {}, {"correlation_id": "correlation-id"} - - -@pytest.mark.asyncio -async def test_sql_postprocessor_validates_with_dry_plan_then_dry_run(): - engine = _DryPlanEngine() - - result = await SQLGenPostProcessor(engine).run( - replies=["SELECT 1"], - project_id="project-id", - use_dry_plan=True, - allow_dry_plan_fallback=False, - data_source="source", - ) - - assert result["valid_generation_result"]["sql"] == "SELECT 1" - assert result["valid_generation_result"]["correlation_id"] == "correlation-id" - assert engine.dry_plan_calls[0][1]["project_id"] == "project-id" - assert engine.dry_plan_calls[0][1]["allow_fallback"] is False - assert engine.execute_sql_calls[0][1]["project_id"] == "project-id" - assert engine.execute_sql_calls[0][1]["dry_run"] is True - - -class _FailingDryPlanEngine: - async def dry_plan(self, *args, **kwargs): - return False, "planner failed" - - -class _FailingDryRunAfterDryPlanEngine: - async def dry_plan(self, *args, **kwargs): - return True, "" - - async def execute_sql(self, *args, **kwargs): - return False, {}, {"error_message": "dry run failed"} - - -@pytest.mark.asyncio -async def test_sql_postprocessor_returns_original_sql_when_dry_plan_fails(): - result = await SQLGenPostProcessor(_FailingDryPlanEngine()).run( - replies=["SELECT 1"], - use_dry_plan=True, - data_source="source", - ) - - assert result["invalid_generation_result"]["sql"] == "SELECT 1" - assert result["invalid_generation_result"]["original_sql"] == "SELECT 1" - assert result["invalid_generation_result"]["type"] == "DRY_PLAN" - - -@pytest.mark.asyncio -async def test_sql_postprocessor_does_not_keep_valid_result_when_dry_run_fails(): - result = await SQLGenPostProcessor(_FailingDryRunAfterDryPlanEngine()).run( - replies=["SELECT 1"], - use_dry_plan=True, - data_source="source", - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["sql"] == "SELECT 1" - assert result["invalid_generation_result"]["original_sql"] == "SELECT 1" - assert result["invalid_generation_result"]["type"] == "DRY_RUN" - - -@pytest.mark.asyncio -async def test_sql_postprocessor_rejects_null_sql_generation_result(): - result = await SQLGenPostProcessor(_DryPlanEngine()).run( - replies=['{"sql": null}'], - use_dry_plan=True, - data_source="source", - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" - assert result["invalid_generation_result"]["sql"] == "" - - -def test_construct_instructions_uses_instruction_text(): - assert construct_instructions( - [{"instruction": "First rule."}, {"instruction": "Second rule."}] - ) == ["First rule.", "Second rule."] - - -def test_construct_ask_history_messages_omits_executable_history_context(): - histories = [ - { - "question": "previous natural language request", - "sql": "SELECT * FROM previous_model", - } - ] - - assert construct_ask_history_messages(histories) == [] - - -def test_get_text_to_sql_rules_uses_default_metadata_grounding_rules(): - rules = get_text_to_sql_rules() - - assert "ONLY USE the tables and columns mentioned in the database schema" in rules - assert 'ONLY USE "*" if the user query asks for all the columns' in rules - assert "They are never source table or source column identifiers" in rules - assert "do not invent a field" in rules - assert "join only through the FOREIGN KEY relationships shown" in rules - assert "Never generate SQL from assumptions" in rules - assert "Do not derive executable identifiers" in rules - assert "Do not query INFORMATION_SCHEMA" in rules - assert "SQL samples and query history are examples of intent and style only" in rules - assert "order by that alias" in rules - assert "Interpret the user's intent" in rules - assert "schema descriptions, aliases, display labels" in rules - assert "WREN RETRIEVED SEMANTIC CONTEXT" in rules - assert "WREN SQL IDENTIFIER CONTRACT" in rules - assert "compact authoritative list of executable identifiers" in rules - assert "sql_table_name_use_exactly" in rules - assert "sql_column_name_use_exactly" in rules - assert "semantic_context_not_sql_identifier" in rules - assert "Do not combine words, labels, ordinals" in rules - assert "Never generate placeholder identifiers" in rules - assert "Never create an identifier from user question wording" in rules - assert "use all required related tables" in rules - assert "silently check that each identifier and function" in rules - assert "instead of inventing a replacement" in rules - assert "exact date/time schema column and required SQL FUNCTIONS-supported operation" in rules - assert ( - "Treat reasoning plans, correction notes, and error messages as non-executable context" - in rules - ) - assert "first locate the exact declared source column" in rules - assert "Physical datasource names, source database names" in rules - assert "Do not replace an invalid identifier with a similar-looking physical" in rules - assert "source/lineage names from metadata may guide meaning" in rules - assert "combining separate result rows with UNION ALL" in rules - assert "independently valid from DATABASE SCHEMA" in rules - assert "do not translate it into a generic object name" in rules - assert "return null for sql instead of producing an approximate query" in rules - assert "If that field is required to answer the request, return null for sql" in rules - assert "A retrieved object is usable only when" in rules - assert "Generate Wren SQL syntax only" in rules - assert "Never use SELECT TOP" in rules - assert "square-bracket identifiers" in rules - assert "Preserve every deployed table and column identifier exactly" in rules - assert "Do not convert deployed identifiers into display-friendly variants" in rules - - -def test_get_text_to_sql_rules_keeps_mandatory_rules_with_sql_knowledge(): - rules = get_text_to_sql_rules(_SqlKnowledge()) - - assert _SqlKnowledge.text_to_sql_rule in rules - assert "MANDATORY SQL GROUNDING RULES" in rules - assert "Every table and column referenced" in rules - assert "Do not query INFORMATION_SCHEMA" in rules - - -def test_sql_generation_schema_allows_null_when_sql_cannot_be_grounded(): - schema = SQL_GENERATION_MODEL_KWARGS["response_format"]["json_schema"]["schema"] - - assert {"type": "null"} in schema["properties"]["sql"]["anyOf"] - - -def test_get_metric_instructions_uses_sql_knowledge_override(): - assert get_metric_instructions(_SqlKnowledge()) == _SqlKnowledge.metric_instructions - - -def test_get_json_field_instructions_uses_sql_knowledge_override(): - assert ( - get_json_field_instructions(_SqlKnowledge()) - == _SqlKnowledge.json_field_instructions - ) - - -def test_sql_generation_system_prompt_grounding_contract(): - prompt = get_sql_generation_system_prompt() - - assert "Output aliases are labels for result columns only" in prompt - assert "must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY" in prompt - assert "source of executable table and column identifiers" in prompt - assert "Never generate SQL from assumptions" in prompt - assert "ignore those parts" in prompt - assert "answer the user's intent" in prompt - assert "Wren SQL query" in prompt - assert "use a normal equality or LIKE comparison" in prompt - assert "unless the user explicitly asks for rank values" in prompt - assert "perform a silent grounding check" in prompt - assert "return null for sql" in prompt - assert "Do not create table or column identifiers from the user's wording" in prompt - assert "DATABASE SCHEMA is the only source of executable identifiers" in prompt - assert "reasoning plan as semantic context for intent only" in prompt - assert "Do not copy identifiers, functions, literal values" in prompt - assert "include those objects only when DATABASE SCHEMA shows" in prompt - assert "Use the exact supported syntax shown there" in prompt - assert "WREN SQL IDENTIFIER CONTRACT" in prompt - assert "EXECUTABLE WREN IDENTIFIER CATALOG" in prompt - assert "first and clearest list of allowed executable identifiers" in prompt - assert "Use sql_table_name_use_exactly" in prompt - assert "sql_column_names_use_exactly" in prompt - assert "semantic_context_not_sql_identifiers" in prompt - assert "Use Wren SQL identifier quoting with double quotes only" in prompt - assert "source database/schema/table names" in prompt - assert "appears only in SQL samples, failed SQL" in prompt - assert "retrieved schema does not ground the requested subject" in prompt - assert "If any planned SQL identifier cannot be copied exactly" in prompt - assert "Never create a table or column from the user's wording" in prompt - assert "" not in prompt - - -def test_json_field_instructions_do_not_include_placeholder_identifiers(): - prompt = get_json_field_instructions() - - assert "json_fields metadata" in prompt - assert "Do not copy JSON examples" in prompt - assert "CREATE TABLE users" not in prompt - assert "my_table" not in prompt - assert "parent_table" not in prompt - - -def test_sql_regeneration_system_prompt_uses_question_as_intent_source(): - prompt = get_sql_regeneration_system_prompt() - - assert "regenerate from the user's question" in prompt - assert "unsupported identifiers" in prompt - assert "original SQL query" in prompt - assert "intentionally omitted" in prompt - assert ( - "database schema as the only source of executable table and column identifiers" - in prompt - ) - assert "Treat physical/source/lineage names from the original SQL" in prompt - - -def test_sql_correction_system_prompt_discards_invalid_identifier_context(): - prompt = get_sql_correction_system_prompt() - - assert "treat it as the source of intent" in prompt - assert "Do not copy placeholders" in prompt - assert "Regenerate a grounded Wren SQL query" in prompt - assert "Wren SQL expert" in prompt - assert "syntactically correct Wren SQL query" in prompt - assert ( - "Do not preserve a table, column, join, filter, grouping, ordering, or function" - in prompt - ) - assert "Treat physical/source/lineage names from the failed SQL" in prompt - assert "do not try a similar replacement from source metadata" in prompt - assert "connector-specific syntax" in prompt - assert "return null for sql instead of substituting non-schema identifiers" in prompt - - -def test_sql_reasoning_prompt_keeps_reasoning_non_executable(): - prompt = sql_generation_reasoning_system_prompt - - assert "Do not write SQL, possible SQL, sample SQL, assumed SQL" in prompt - assert "SQL clauses, SQL functions, code blocks, or executable expressions" in prompt - assert "literal prefix `table:`" in prompt - assert "literal prefix `column:`" in prompt - assert "reasoning plan is semantic context for intent only" in prompt - assert "declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT" in prompt - assert "source names, physical names, lineage names" in prompt - assert 'Do not use the words "assume", "assuming", "likely"' in prompt - assert "retrieved metadata does not support that part" in prompt - assert "Do not propose a replacement name" in prompt - assert "Do not write table names or column names from the user's wording" in prompt - assert "Do not include code blocks, inline SQL fragments" in prompt - assert "" not in prompt - assert "" not in prompt - - -def test_user_prompt_templates_keep_source_metadata_non_executable(): - for prompt in ( - sql_generation_user_prompt_template, - text_to_sql_with_followup_user_prompt_template, - sql_regeneration_user_prompt_template, - sql_correction_user_prompt_template, - ): - assert "source/physical/lineage names" in prompt - assert "return null for sql instead of inventing" in prompt - assert "exact declared table and column names from DATABASE SCHEMA" in prompt - assert "user question words" in prompt - assert "return null for sql instead of querying an unrelated object" in prompt - - -def test_followup_sql_prompt_does_not_expect_previous_sql_context(): - assert "previous SQL query" not in text_to_sql_with_followup_user_prompt_template - assert "current retrieved DATABASE SCHEMA" in ( - text_to_sql_with_followup_user_prompt_template - ) - - -def test_sql_answer_prompt_uses_only_returned_data_rows(): - assert "Use only the columns and rows provided in Data" in sql_to_answer_system_prompt - assert "Do not invent, duplicate, reorder, aggregate, rank, or label rows" in ( - sql_to_answer_system_prompt - ) - assert "summarize those exact aggregate rows" in sql_to_answer_system_prompt - assert "If the Data is empty" in sql_to_answer_system_prompt - - -def test_executable_prompt_templates_omit_untrusted_reasoning_and_sql_context(): - reasoning_marker = "UNTRUSTED_REASONING_CONTEXT_MARKER" - diagnostic_marker = "UNTRUSTED_DIAGNOSTIC_CONTEXT_MARKER" - original_sql_marker = "UNTRUSTED_ORIGINAL_SQL_MARKER" - - generation_prompt = PromptBuilder(template=sql_generation_user_prompt_template).run( - query="Question", - documents=["SCHEMA_CONTEXT"], - sql_generation_reasoning=reasoning_marker, - instructions=[], - calculated_field_instructions="", - metric_instructions="", - json_field_instructions="", - sql_samples=[], - sql_functions=[], - )["prompt"] - - followup_prompt = PromptBuilder( - template=text_to_sql_with_followup_user_prompt_template - ).run( - query="Question", - documents=["SCHEMA_CONTEXT"], - sql_generation_reasoning=reasoning_marker, - instructions=[], - calculated_field_instructions="", - metric_instructions="", - json_field_instructions="", - sql_samples=[], - sql_functions=[], - )["prompt"] - - correction_prompt = PromptBuilder(template=sql_correction_user_prompt_template).run( - query="Question", - documents=["SCHEMA_CONTEXT"], - invalid_generation_result={"error": diagnostic_marker}, - sql_generation_reasoning=reasoning_marker, - instructions=[], - sql_functions=[], - )["prompt"] - - regeneration_prompt = PromptBuilder( - template=sql_regeneration_user_prompt_template - ).run( - query="Question", - sql=original_sql_marker, - documents=["SCHEMA_CONTEXT"], - sql_generation_reasoning=reasoning_marker, - instructions=[], - calculated_field_instructions="", - metric_instructions="", - json_field_instructions="", - sql_samples=[], - sql_functions=[], - )["prompt"] - - for prompt in ( - generation_prompt, - followup_prompt, - correction_prompt, - regeneration_prompt, - ): - assert reasoning_marker not in prompt - assert "REASONING PLAN" not in prompt - assert "" not in prompt - assert "" not in prompt - - assert original_sql_marker not in regeneration_prompt - assert "original SQL is intentionally omitted" in regeneration_prompt - assert diagnostic_marker not in correction_prompt - assert "dry-run diagnostic text is intentionally omitted" in ( - correction_prompt - ) - assert "Regenerate from the user question and current DATABASE SCHEMA only" in ( - correction_prompt - ) diff --git a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py b/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py deleted file mode 100644 index f5f15b14bf..0000000000 --- a/wren-ai-service/tests/pytest/services/test_dry_plan_defaults.py +++ /dev/null @@ -1,189 +0,0 @@ -from unittest.mock import AsyncMock - -import pytest - -from src.web.v1.routers.question_recommendation import ( - PostRequest as QuestionRecommendationPostRequest, -) -from src.web.v1.routers.sql_corrections import PostRequest as SqlCorrectionPostRequest -from src.web.v1.services.ask import AskRequest, AskResultRequest, AskService -from src.web.v1.services.question_recommendation import QuestionRecommendation -from src.web.v1.services.sql_corrections import SqlCorrectionService - - -def test_ask_request_defaults_to_planner_validation_without_fallback(): - request = AskRequest(query="How many records are available?", id="deploy-id") - - assert request.use_dry_plan is True - assert request.allow_dry_plan_fallback is False - assert request.ignore_sql_generation_reasoning is True - - -def test_ask_request_allows_explicit_planner_override(): - request = AskRequest( - query="How many records are available?", - id="deploy-id", - use_dry_plan=False, - allow_dry_plan_fallback=True, - ) - - assert request.use_dry_plan is False - assert request.allow_dry_plan_fallback is True - - -def test_ask_service_defaults_to_column_pruning(): - service = AskService({}) - - assert service._enable_column_pruning is True - assert service._allow_sql_generation_reasoning is False - assert service._max_sql_correction_retries == 0 - - -@pytest.mark.asyncio -async def test_sql_correction_service_requires_retrieved_tables(): - sql_tables_extraction = AsyncMock(run=AsyncMock(return_value={"post_process": []})) - service = SqlCorrectionService({"sql_tables_extraction": sql_tables_extraction}) - request = SqlCorrectionService.CorrectionRequest( - event_id="event-id", - sql="SELECT 1", - error="dry run failed", - ) - - await service.correct(request) - result = service["event-id"] - - assert result.status == "failed" - assert "retrieved table context" in result.error.message - assert result.invalid_sql is None - sql_tables_extraction.run.assert_not_called() - - -@pytest.mark.asyncio -async def test_ask_service_does_not_run_speculative_reasoning_or_correction_after_failed_generation(): - sql_generation_reasoning = AsyncMock(run=AsyncMock(return_value={})) - sql_correction = AsyncMock(run=AsyncMock(return_value={})) - service = AskService( - { - "sql_pairs_retrieval": AsyncMock( - run=AsyncMock(return_value={"formatted_output": {"documents": []}}) - ), - "instructions_retrieval": AsyncMock( - run=AsyncMock(return_value={"formatted_output": {"documents": []}}) - ), - "db_schema_retrieval": AsyncMock( - run=AsyncMock( - return_value={ - "construct_retrieval_results": { - "retrieval_results": [ - { - "table_name": "retrieved_model", - "table_ddl": "CREATE TABLE retrieved_model (retrieved_field VARCHAR);", - } - ], - "has_calculated_field": False, - "has_metric": False, - "has_json_field": False, - } - } - ) - ), - "sql_functions_retrieval": AsyncMock(run=AsyncMock(return_value=[])), - "sql_knowledge_retrieval": AsyncMock(run=AsyncMock(return_value=None)), - "sql_generation_reasoning": sql_generation_reasoning, - "sql_generation": AsyncMock( - run=AsyncMock( - return_value={ - "post_process": { - "valid_generation_result": {}, - "invalid_generation_result": { - "type": "DRY_RUN", - "sql": "SELECT 1", - "original_sql": "SELECT 1", - "error": "dry run failed", - }, - } - } - ) - ), - "sql_correction": sql_correction, - }, - allow_intent_classification=False, - ) - request = AskRequest(query="Can this be answered?", id="deploy-id") - request.query_id = "query-id" - - await service.ask(request) - result = service.get_ask_result(AskResultRequest(query_id="query-id")) - - assert result.status == "failed" - assert result.error.code == "NO_RELEVANT_SQL" - assert result.invalid_sql is None - sql_generation_reasoning.run.assert_not_called() - sql_correction.run.assert_not_called() - - -def test_sql_correction_router_defaults_to_planner_validation_without_fallback(): - request = SqlCorrectionPostRequest(sql="SELECT 1", error="dry run failed") - - assert request.use_dry_plan is True - assert request.allow_dry_plan_fallback is False - - -def test_sql_correction_service_defaults_to_planner_validation_without_fallback(): - request = SqlCorrectionService.CorrectionRequest( - event_id="event-id", - sql="SELECT 1", - error="dry run failed", - ) - - assert request.use_dry_plan is True - assert request.allow_dry_plan_fallback is False - - -def test_question_recommendation_router_defaults_to_planner_validation_without_fallback(): - request = QuestionRecommendationPostRequest(mdl='{"models":[]}') - - assert request.allow_data_preview is False - assert request.use_dry_plan is True - assert request.allow_dry_plan_fallback is False - - -def test_question_recommendation_service_defaults_to_planner_validation_without_fallback(): - request = QuestionRecommendation.Request(event_id="event-id", mdl='{"models":[]}') - - assert request.allow_data_preview is False - assert request.use_dry_plan is True - assert request.allow_dry_plan_fallback is False - - -@pytest.mark.asyncio -async def test_ask_service_does_not_require_historical_sql_shortcut(): - service = AskService( - { - "sql_pairs_retrieval": AsyncMock( - run=AsyncMock(return_value={"formatted_output": {"documents": []}}) - ), - "instructions_retrieval": AsyncMock( - run=AsyncMock(return_value={"formatted_output": {"documents": []}}) - ), - "intent_classification": AsyncMock( - run=AsyncMock( - return_value={ - "post_process": { - "intent": "MISLEADING_QUERY", - "reasoning": "No matching schema context.", - } - } - ) - ), - "misleading_assistance": AsyncMock(run=AsyncMock(return_value={})), - } - ) - request = AskRequest(query="Can this be answered?", id="deploy-id") - request.query_id = "query-id" - - await service.ask(request) - result = service.get_ask_result(AskResultRequest(query_id="query-id")) - - assert result.status == "finished" - assert result.type == "GENERAL" From 9d6abd011cfbd341fd3d38c2685ad2fadcdaa75d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 19:43:35 +0530 Subject: [PATCH 0785/1087] Invalidate semantics on datasource update --- wren-ui/src/apollo/server/resolvers/projectResolver.ts | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index a827527901..e82741d5e5 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -378,6 +378,14 @@ export class ProjectResolver { displayName, connectionInfo: { ...project.connectionInfo, ...toUpdateConnectionInfo }, }); + + await ctx.schemaChangeRepository.deleteAllBy({ projectId: project.id }); + await ctx.deployService.deleteAllByProjectId(project.id); + await ctx.askingService.deleteAllByProjectId(project.id); + await ctx.modelService.deleteAllViewsByProjectId(project.id); + await ctx.modelService.deleteAllModelsByProjectId(project.id); + await ctx.wrenAIAdaptor.delete(project.id); + return { type: updatedProject.type, properties: { From 0a40f90480dec6c44f372785ab4a864c612c0af6 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Fri, 31 Jul 2026 21:13:42 +0530 Subject: [PATCH 0786/1087] Scope semantic retrieval by deployed manifest --- wren-ai-service/src/pipelines/common.py | 32 ++++++++++++----- .../generation/followup_sql_generation.py | 8 ++++- .../pipelines/generation/sql_correction.py | 8 ++++- .../pipelines/generation/sql_generation.py | 8 ++++- .../src/pipelines/indexing/db_schema.py | 16 +++++++-- .../pipelines/indexing/historical_question.py | 23 ++++++++++--- .../src/pipelines/indexing/project_meta.py | 13 +++++-- .../src/pipelines/indexing/sql_pairs.py | 22 ++++++++++-- .../pipelines/indexing/table_description.py | 23 ++++++++++--- .../retrieval/db_schema_retrieval.py | 29 ++++++++++++---- .../historical_question_retrieval.py | 34 +++++++------------ .../src/pipelines/retrieval/sql_functions.py | 7 +++- .../src/pipelines/retrieval/sql_knowledge.py | 7 +++- .../src/providers/document_store/qdrant.py | 3 ++ wren-ai-service/src/web/v1/services/ask.py | 7 ++++ .../src/web/v1/services/ask_feedback.py | 9 ++++- .../web/v1/services/semantics_preparation.py | 1 + .../src/web/v1/services/sql_corrections.py | 8 ++++- 18 files changed, 199 insertions(+), 59 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index 605238bb64..c40a940088 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -4,6 +4,23 @@ from haystack import Document, component +def build_project_deploy_filter( + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, +) -> dict[str, Any] | None: + conditions = [] + + if project_id: + conditions.append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + if mdl_hash: + conditions.append({"field": "mdl_hash", "operator": "==", "value": mdl_hash}) + + return {"operator": "AND", "conditions": conditions} if conditions else None + + def get_engine_supported_data_type(data_type: str | None) -> str: """ This function makes sure downstream ai pipeline get column data types in a format that is supported by the data engine. @@ -92,15 +109,12 @@ def build_table_ddl( ) -async def retrieve_metadata(project_id: str, retriever) -> dict[str, Any]: - filters = None - if project_id: - filters = { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } +async def retrieve_metadata( + project_id: str, + retriever, + mdl_hash: Optional[str] = None, +) -> dict[str, Any]: + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) result = await retriever.run(query_embedding=[], filters=filters) documents = result["documents"] 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 1bcb3f853d..bc012d759c 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -200,6 +200,7 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -211,7 +212,11 @@ async def run( logger.info("Follow-Up SQL Generation pipeline is running...") if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) + metadata = await retrieve_metadata( + project_id or "", + self._retriever, + mdl_hash=mdl_hash, + ) else: metadata = {} @@ -223,6 +228,7 @@ async def run( "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, + "mdl_hash": mdl_hash, "sql_samples": sql_samples, "instructions": instructions, "has_calculated_field": has_calculated_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..2118739a5c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -172,6 +172,7 @@ async def run( instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, @@ -179,7 +180,11 @@ async def run( logger.info("SQLCorrection pipeline is running...") if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) + metadata = await retrieve_metadata( + project_id or "", + self._retriever, + mdl_hash=mdl_hash, + ) else: metadata = {} @@ -191,6 +196,7 @@ async def run( "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, + "mdl_hash": mdl_hash, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..71a80875e9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -194,6 +194,7 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -206,7 +207,11 @@ async def run( logger.info("SQL Generation pipeline is running...") if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) + metadata = await retrieve_metadata( + project_id or "", + self._retriever, + mdl_hash=mdl_hash, + ) else: metadata = {} @@ -219,6 +224,7 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, + "mdl_hash": mdl_hash, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 97cf6110f0..b62d1de9fa 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -49,9 +49,15 @@ async def run( mdl: Dict[str, Any], column_batch_size: int, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ): def _additional_meta() -> Dict[str, Any]: - return {"project_id": project_id} if project_id else {} + metadata = {} + if project_id: + metadata["project_id"] = project_id + if mdl_hash: + metadata["mdl_hash"] = mdl_hash + return metadata chunks = [ { @@ -382,11 +388,13 @@ async def chunk( chunker: DDLChunker, column_batch_size: int, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: return await chunker.run( mdl=mdl, column_batch_size=column_batch_size, project_id=project_id, + mdl_hash=mdl_hash, ) @@ -445,7 +453,10 @@ def __init__( @observe(name="DB Schema Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id}, DB Schema Indexing pipeline is running..." @@ -455,6 +466,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/indexing/historical_question.py b/wren-ai-service/src/pipelines/indexing/historical_question.py index 95515f7ac4..2873672292 100644 --- a/wren-ai-service/src/pipelines/indexing/historical_question.py +++ b/wren-ai-service/src/pipelines/indexing/historical_question.py @@ -53,7 +53,12 @@ class ViewChunker: """ @component.output_types(documents=List[Document]) - def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None) -> None: + def run( + self, + mdl: Dict[str, Any], + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ) -> None: def _get_content(view: Dict[str, Any]) -> str: properties = view.get("properties", {}) historical_queries = properties.get("historical_queries", []) @@ -70,7 +75,12 @@ def _get_meta(view: Dict[str, Any]) -> Dict[str, Any]: } def _additional_meta() -> Dict[str, Any]: - return {"project_id": project_id} if project_id else {} + metadata = {} + if project_id: + metadata["project_id"] = project_id + if mdl_hash: + metadata["mdl_hash"] = mdl_hash + return metadata chunks = [ { @@ -105,8 +115,9 @@ def chunk( mdl: Dict[str, Any], chunker: ViewChunker, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: - return chunker.run(mdl=mdl, project_id=project_id) + return chunker.run(mdl=mdl, project_id=project_id, mdl_hash=mdl_hash) @observe(capture_input=False, capture_output=False) @@ -164,7 +175,10 @@ def __init__( @observe(name="Historical Question Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id}, Historical Question Indexing pipeline is running..." @@ -174,6 +188,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/indexing/project_meta.py b/wren-ai-service/src/pipelines/indexing/project_meta.py index 426e988934..566d0b52ac 100644 --- a/wren-ai-service/src/pipelines/indexing/project_meta.py +++ b/wren-ai-service/src/pipelines/indexing/project_meta.py @@ -30,8 +30,13 @@ def validate_mdl(mdl_str: str, validator: MDLValidator) -> dict[str, Any]: def chunk( mdl: dict[str, Any], project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> dict[str, Any]: - addition = {"project_id": project_id} if project_id else {} + addition = {} + if project_id: + addition["project_id"] = project_id + if mdl_hash: + addition["mdl_hash"] = mdl_hash data_source = str(mdl.get("dataSource") or "local_file").lower() if data_source == "duckdb": @@ -87,7 +92,10 @@ def __init__( @observe(name="Project Meta Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> dict[str, Any]: logger.info( f"Project ID: {project_id}, Project Meta Indexing pipeline is running..." @@ -97,6 +105,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index 16ff245cd7..65eec14266 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -28,10 +28,19 @@ class SqlPair(BaseModel): @component class SqlPairsConverter: @component.output_types(documents=List[Document]) - def run(self, sql_pairs: List[SqlPair], project_id: str = ""): + def run( + self, + sql_pairs: List[SqlPair], + project_id: str = "", + mdl_hash: Optional[str] = None, + ): logger.info(f"Project ID: {project_id} Converting SQL pairs to documents...") - addition = {"project_id": project_id} if project_id else {} + addition = {} + if project_id: + addition["project_id"] = project_id + if mdl_hash: + addition["mdl_hash"] = mdl_hash return { "documents": [ @@ -118,8 +127,13 @@ def to_documents( sql_pairs: List[SqlPair], document_converter: SqlPairsConverter, project_id: str = "", + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: - return document_converter.run(sql_pairs=sql_pairs, project_id=project_id) + return document_converter.run( + sql_pairs=sql_pairs, + project_id=project_id, + mdl_hash=mdl_hash, + ) @observe(capture_input=False, capture_output=False) @@ -213,6 +227,7 @@ async def run( self, mdl_str: str, project_id: str = "", + mdl_hash: Optional[str] = None, external_pairs: Optional[Dict[str, Any]] = None, delete_all: bool = False, include_default_pairs: bool = True, @@ -229,6 +244,7 @@ async def run( input = { "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, "external_pairs": pairs, "delete_all": delete_all, "include_default_pairs": include_default_pairs, diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index e0914ccc81..65fdeee6f2 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -26,9 +26,19 @@ def _properties(self, payload: Dict[str, Any]) -> Dict[str, Any]: return properties if isinstance(properties, dict) else {} @component.output_types(documents=List[Document]) - def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None): + def run( + self, + mdl: Dict[str, Any], + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): def _additional_meta() -> Dict[str, Any]: - return {"project_id": project_id} if project_id else {} + metadata = {} + if project_id: + metadata["project_id"] = project_id + if mdl_hash: + metadata["mdl_hash"] = mdl_hash + return metadata chunks = [ { @@ -226,8 +236,9 @@ def chunk( mdl: Dict[str, Any], chunker: TableDescriptionChunker, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: - return chunker.run(mdl=mdl, project_id=project_id) + return chunker.run(mdl=mdl, project_id=project_id, mdl_hash=mdl_hash) @observe(capture_input=False, capture_output=False) @@ -286,7 +297,10 @@ def __init__( @observe(name="Table Description Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id}, Table Description Indexing pipeline is running..." @@ -296,6 +310,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..ab48ca0df1 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -15,6 +15,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider from src.pipelines.common import ( + build_project_deploy_filter, build_table_ddl, clean_up_new_lines, get_engine_supported_data_type, @@ -138,7 +139,11 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> d @observe(capture_input=False) async def table_retrieval( - embedding: dict, project_id: str, tables: list[str], table_retriever: Any + embedding: dict, + project_id: str, + tables: list[str], + table_retriever: Any, + mdl_hash: str = "", ) -> dict: filters = { "operator": "AND", @@ -152,6 +157,11 @@ async def table_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) + if mdl_hash: + filters["conditions"].append( + {"field": "mdl_hash", "operator": "==", "value": mdl_hash} + ) + if embedding: return await table_retriever.run( query_embedding=embedding.get("embedding"), @@ -170,7 +180,10 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, project_id: str, dbschema_retriever: Any + table_retrieval: dict, + project_id: str, + dbschema_retriever: Any, + mdl_hash: str = "", ) -> list[Document]: tables = table_retrieval.get("documents", []) table_names = [] @@ -184,6 +197,10 @@ async def dbschema_retrieval( ] if table_name_conditions: + project_deploy_filter = build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ) filters = { "operator": "AND", "conditions": [ @@ -192,10 +209,8 @@ async def dbschema_retrieval( ], } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + if project_deploy_filter: + filters["conditions"] += project_deploy_filter["conditions"] results = await dbschema_retriever.run(query_embedding=[], filters=filters) return results["documents"] @@ -501,6 +516,7 @@ async def run( query: str = "", tables: Optional[list[str]] = None, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, ): @@ -511,6 +527,7 @@ async def run( "query": query, "tables": tables, "project_id": project_id or "", + "mdl_hash": mdl_hash or "", "histories": histories or [], "enable_column_pruning": enable_column_pruning, **self._components, diff --git a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py index 68f1ef158c..1d16ae1f99 100644 --- a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py @@ -10,7 +10,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.common import ScoreFilter +from src.pipelines.common import ScoreFilter, build_project_deploy_filter logger = logging.getLogger("wren-ai-service") @@ -39,17 +39,9 @@ def run(self, documents: List[Document]): async def count_documents( view_questions_store: QdrantDocumentStore, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> int: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) count = await view_questions_store.count_documents(filters=filters) return count @@ -68,18 +60,10 @@ async def retrieval( embedding: dict, project_id: str, view_questions_retriever: Any, + mdl_hash: str = "", ) -> dict: if embedding: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) view_question_res = await view_questions_retriever.run( query_embedding=embedding.get("embedding"), @@ -149,13 +133,19 @@ def __init__( ) @observe(name="Historical Question") - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): logger.info("HistoricalQuestion Retrieval pipeline is running...") return await self._pipe.execute( ["formatted_output"], inputs={ "query": query, "project_id": project_id or "", + "mdl_hash": mdl_hash or "", **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/retrieval/sql_functions.py b/wren-ai-service/src/pipelines/retrieval/sql_functions.py index 016aa9b1e6..822e346703 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_functions.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_functions.py @@ -104,12 +104,17 @@ def __init__( async def run( self, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> List[SqlFunction]: logger.info( f"Project ID: {project_id} SQL Functions Retrieval pipeline is running..." ) - metadata = await retrieve_metadata(project_id or "", self._retriever) + metadata = await retrieve_metadata( + project_id or "", + self._retriever, + mdl_hash=mdl_hash, + ) _data_source = metadata.get("data_source", "local_file") if _data_source in self._cache: diff --git a/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py b/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py index 167969047c..ad6dde9ad1 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py @@ -114,12 +114,17 @@ def __init__( async def run( self, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Optional[SqlKnowledge]: logger.info( f"Project ID: {project_id} SQL Knowledge Retrieval pipeline is running..." ) - metadata = await retrieve_metadata(project_id or "", self._retriever) + metadata = await retrieve_metadata( + project_id or "", + self._retriever, + mdl_hash=mdl_hash, + ) _data_source = metadata.get("data_source", "local_file") if _data_source in self._cache: diff --git a/wren-ai-service/src/providers/document_store/qdrant.py b/wren-ai-service/src/providers/document_store/qdrant.py index d528bb0fb9..b92cd32d7d 100644 --- a/wren-ai-service/src/providers/document_store/qdrant.py +++ b/wren-ai-service/src/providers/document_store/qdrant.py @@ -178,6 +178,9 @@ def __init__( self.client.create_payload_index( collection_name=index, field_name="project_id", field_schema="keyword" ) + self.client.create_payload_index( + collection_name=index, field_name="mdl_hash", field_schema="keyword" + ) def recreate_collection( self, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index aa26fa3f81..1a0052a2a6 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -192,6 +192,7 @@ async def ask( historical_question = await self._pipelines["historical_question"].run( query=user_query, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ) # we only return top 1 result @@ -347,6 +348,7 @@ async def ask( query=user_query, histories=histories, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, enable_column_pruning=enable_column_pruning, ) _retrieval_result = retrieval_result.get( @@ -442,6 +444,7 @@ async def ask( "sql_functions_retrieval" ].run( project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ) else: sql_functions = [] @@ -451,6 +454,7 @@ async def ask( "sql_knowledge_retrieval" ].run( project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ) has_calculated_field = _retrieval_result.get( @@ -468,6 +472,7 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -486,6 +491,7 @@ async def ask( contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -557,6 +563,7 @@ async def ask( else error_message, }, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index f1971afa15..a05f2b0c5b 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -4,7 +4,7 @@ from cachetools import TTLCache from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline from src.utils import trace_metadata @@ -20,6 +20,9 @@ class AskFeedbackRequest(BaseRequest): tables: List[str] sql_generation_reasoning: str sql: str + mdl_hash: Optional[str] = Field( + default=None, validation_alias=AliasChoices("mdl_hash", "id") + ) class AskFeedbackResponse(BaseModel): @@ -122,6 +125,7 @@ async def ask_feedback( self._pipelines["db_schema_retrieval"].run( tables=ask_feedback_request.tables, project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ), self._pipelines["sql_pairs_retrieval"].run( query=ask_feedback_request.question, @@ -139,6 +143,7 @@ async def ask_feedback( "sql_functions_retrieval" ].run( project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ) else: sql_functions = [] @@ -148,6 +153,7 @@ async def ask_feedback( "sql_knowledge_retrieval" ].run( project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ) # Extract results from completed tasks @@ -249,6 +255,7 @@ async def ask_feedback( "error": correction_error_message, }, project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, sql_functions=sql_functions, sql_knowledge=sql_knowledge, ) 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 327630203e..a82de5d98f 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -76,6 +76,7 @@ async def prepare_semantics( input = { "mdl_str": prepare_semantics_request.mdl, "project_id": prepare_semantics_request.project_id, + "mdl_hash": prepare_semantics_request.mdl_hash, } project_scoped_index_names = [ "db_schema", diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index f805024ad2..ed456adec5 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -3,7 +3,7 @@ from cachetools import TTLCache from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline from src.utils import trace_metadata @@ -60,6 +60,9 @@ class CorrectionRequest(BaseRequest): event_id: str sql: str error: str + mdl_hash: Optional[str] = Field( + default=None, validation_alias=AliasChoices("mdl_hash", "id") + ) retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = True allow_dry_plan_fallback: bool = False @@ -100,12 +103,14 @@ async def correct( if self._allow_sql_knowledge_retrieval: sql_knowledge = await self._pipelines["sql_knowledge_retrieval"].run( project_id=project_id, + mdl_hash=request.mdl_hash, ) documents = ( ( await self._pipelines["db_schema_retrieval"].run( project_id=project_id, + mdl_hash=request.mdl_hash, tables=retrieved_tables, ) ) @@ -118,6 +123,7 @@ async def correct( contexts=table_ddls, invalid_generation_result=_invalid, project_id=project_id, + mdl_hash=request.mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, From f28f567627a914b90e9ac3652e6dea2f040f4e9e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 2 Aug 2026 19:42:35 +0530 Subject: [PATCH 0787/1087] Improve project-scoped schema retrieval --- .../retrieval/db_schema_retrieval.py | 263 ++++++++++++++---- 1 file changed, 208 insertions(+), 55 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ab48ca0df1..f2e870b263 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -31,6 +31,7 @@ You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. +The same business concept is represented by multiple modeled datasets in some projects; preserve each relevant dataset when it can answer the requested intent. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -40,6 +41,7 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +8. If multiple schema objects provide compatible fields for the same requested result shape, include all of those schema objects instead of choosing only one. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -115,9 +117,102 @@ def _build_metric_ddl(content: dict) -> str: ) +def _schema_column_names(content: dict) -> list[str]: + return [ + column["name"] + for column in content.get("columns", []) + if column.get("type") == "COLUMN" and column.get("name") + ] + + +def _identifier_catalog(table_name: str, column_names: list[str]) -> str: + columns = "\n".join(f"- {column_name}" for column_name in column_names) + return ( + "/* EXECUTABLE WREN IDENTIFIER CATALOG\n" + f"table: {table_name}\n" + "columns:\n" + f"{columns}\n" + "Do not create identifiers from user wording, comments, aliases, display labels, or source metadata.\n" + "*/\n" + ) + + +def _semantic_context(content: dict, column_names: list[str]) -> str: + table_name = content.get("name", "") + semantic_parts = [ + str(content.get("comment", "") or "").strip(), + str(content.get("properties", {}) or "").strip(), + ] + for column in content.get("columns", []): + comment = str(column.get("comment", "") or "").strip() + if comment: + semantic_parts.append(f"{column.get('name', '')}: {comment}") + + relationship_constraints = [ + column.get("constraint", "") + for column in content.get("columns", []) + if column.get("type") == "FOREIGN_KEY" and column.get("constraint") + ] + + block = [ + "/* WREN RETRIEVED SEMANTIC CONTEXT", + f"sql_table_name_use_exactly: {table_name}", + "sql_column_names_use_exactly:", + *[f"- {column_name}" for column_name in column_names], + ] + for column_name in column_names: + block.append(f"sql_column_name_use_exactly: {column_name}") + + if relationship_constraints: + block.append("relationship_constraints_use_exactly:") + block.extend(f"- {constraint}" for constraint in relationship_constraints) + + semantic_context = "\n".join(part for part in semantic_parts if part) + if semantic_context: + block.append("semantic_context_not_sql_identifiers:") + block.append(semantic_context) + + block.append("*/") + return "\n".join(block) + "\n" + + +def _build_table_context_ddl( + content: dict, + include_retrieved_semantic_context: bool = False, +) -> tuple[str, bool, bool, list[str]]: + column_names = _schema_column_names(content) + ddl, has_calculated_field, has_json_field = build_table_ddl( + content, + include_semantic_comments=False, + ) + context = _identifier_catalog(content["name"], column_names) + if include_retrieved_semantic_context: + context += _semantic_context(content, column_names) + + return context + ddl, has_calculated_field, has_json_field, column_names + + def _build_view_ddl(content: dict) -> str: + columns = content.get("columns", []) + column_names = [column.get("name", "") for column in columns if column.get("name")] + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" + for column in columns + if column.get("name") + and str(column.get("data_type", "")).lower() != "unknown" + ] + return ( - f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" + _identifier_catalog(content["name"], column_names) + + "/* WREN RETRIEVED SEMANTIC CONTEXT\n" + + f"sql_table_name_use_exactly: {content['name']}\n" + + "sql_column_names_use_exactly:\n" + + "\n".join(f"- {column_name}" for column_name in column_names) + + "\nsemantic_context_not_sql_identifier: view definition_omitted_from_executable_schema\n" + + "*/\n" + + f"CREATE TABLE {content['name']} (\n " + + ",\n ".join(columns_ddl) + + "\n);" ) @@ -125,13 +220,6 @@ def _build_view_ddl(content: dict) -> str: @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: - if histories: - previous_query_summaries = [history.question for history in histories] - else: - previous_query_summaries = [] - - query = "\n".join(previous_query_summaries) + "\n" + query - return await embedder.run(query) else: return {} @@ -184,19 +272,9 @@ async def dbschema_retrieval( project_id: str, dbschema_retriever: Any, mdl_hash: str = "", + embedding: Optional[dict] = None, ) -> list[Document]: - tables = table_retrieval.get("documents", []) - table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) - - table_name_conditions = [ - {"field": "name", "operator": "==", "value": table_name} - for table_name in table_names - ] - - if table_name_conditions: + def _base_filters() -> dict: project_deploy_filter = build_project_deploy_filter( project_id=project_id, mdl_hash=mdl_hash, @@ -205,17 +283,94 @@ async def dbschema_retrieval( "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, ], } if project_deploy_filter: filters["conditions"] += project_deploy_filter["conditions"] - results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return filters + + def _filters_for_names(table_names: list[str]) -> dict: + filters = _base_filters() + filters["conditions"].insert( + 1, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": table_name} + for table_name in table_names + ], + }, + ) + return filters + + async def _fetch_by_names(table_names: list[str]) -> list[Document]: + if not table_names: + return [] + + results = await dbschema_retriever.run( + query_embedding=[], + filters=_filters_for_names(table_names), + ) return results["documents"] - return [] + def _document_name(document: Document) -> str: + return document.meta.get("name", "") + + def _related_table_names(documents: list[Document], visited: set[str]) -> list[str]: + related_names = [] + for document in documents: + content = ast.literal_eval(document.content) + if content.get("type") != "TABLE_COLUMNS": + continue + + for column in content.get("columns", []): + if column.get("type") != "FOREIGN_KEY": + continue + + candidates = list(column.get("tables", []) or []) + if column.get("referenced_table"): + candidates.append(column["referenced_table"]) + + for table_name in candidates: + if table_name and table_name not in visited: + visited.add(table_name) + related_names.append(table_name) + + return related_names + + tables = table_retrieval.get("documents", []) + table_names = [] + for table in tables: + content = ast.literal_eval(table.content) + table_name = content.get("name") + if table_name and table_name not in table_names: + table_names.append(table_name) + + documents = [] + if not table_names and embedding and embedding.get("embedding"): + results = await dbschema_retriever.run( + query_embedding=embedding.get("embedding"), + filters=_base_filters(), + ) + documents.extend(results["documents"]) + for document in results["documents"]: + table_name = _document_name(document) + if table_name and table_name not in table_names: + table_names.append(table_name) + + visited = set(table_names) + pending = list(table_names) + while pending: + current_names = pending + pending = [] + current_documents = await _fetch_by_names(current_names) + documents.extend(current_documents) + pending.extend(_related_table_names(current_documents, visited)) + + return documents + @observe() @@ -261,11 +416,15 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) + ddl, _has_calculated_field, _has_json_field, column_names = ( + _build_table_context_ddl(table_schema) + ) retrieval_results.append( { "table_name": table_schema["name"], "table_ddl": ddl, + "column_names": column_names, + "manifest_column_names": column_names, } ) if _has_calculated_field: @@ -324,16 +483,10 @@ def prompt( ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ - build_table_ddl(construct_db_schema)[0] + _build_table_context_ddl(construct_db_schema)[0] for construct_db_schema in construct_db_schemas ] - previous_query_summaries = ( - [history.question for history in histories] if histories else [] - ) - - query = "\n".join(previous_query_summaries) + "\n" + query - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: @@ -379,12 +532,11 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - ddl, _has_calculated_field, _has_json_field = build_table_ddl( - table_schema, - columns=set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ), - tables=tables, + ddl, _has_calculated_field, _has_json_field, column_names = ( + _build_table_context_ddl( + table_schema, + include_retrieved_semantic_context=True, + ) ) if _has_calculated_field: has_calculated_field = True @@ -395,28 +547,29 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, + "column_names": column_names, + "manifest_column_names": column_names, } ) for document in dbschema_retrieval: - if document.meta["name"] in columns_and_tables_needed: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - } - ) + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) return { "retrieval_results": retrieval_results, From 7f891d197b8d4e5673a2a1313c9bb4c4a3756725 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 2 Aug 2026 19:59:28 +0530 Subject: [PATCH 0788/1087] Scope ask context to deployed metadata --- .../src/pipelines/indexing/instructions.py | 15 ++- .../src/pipelines/retrieval/instructions.py | 42 +++++--- .../retrieval/sql_pairs_retrieval.py | 43 ++++---- .../src/web/v1/services/__init__.py | 1 + wren-ai-service/src/web/v1/services/ask.py | 2 + .../src/web/v1/services/instructions.py | 1 + .../src/web/v1/services/sql_pairs.py | 1 + .../retrieval/test_project_scope_isolation.py | 101 ++++++++++++++++++ .../tests/pytest/services/mocks.py | 30 +++++- .../apollo/server/adaptors/wrenAIAdaptor.ts | 14 ++- .../server/services/instructionService.ts | 31 ++++-- .../apollo/server/services/sqlPairService.ts | 16 +++ wren-ui/src/common.ts | 2 + 13 files changed, 247 insertions(+), 52 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py index 006de9237f..96b205b062 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -29,10 +29,19 @@ class Instruction(BaseModel): @component class InstructionsConverter: @component.output_types(documents=List[Document]) - def run(self, instructions: list[Instruction], project_id: str = ""): + def run( + self, + instructions: list[Instruction], + project_id: str = "", + mdl_hash: Optional[str] = None, + ): logger.info(f"Project ID: {project_id} Converting instructions to documents...") - addition = {"project_id": project_id} if project_id else {} + addition = {} + if project_id: + addition["project_id"] = project_id + if mdl_hash: + addition["mdl_hash"] = mdl_hash return { "documents": [ @@ -161,6 +170,7 @@ async def run( self, instructions: list[Instruction], project_id: str = "", + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} Instructions Indexing pipeline is running..." @@ -168,6 +178,7 @@ async def run( input = { "project_id": project_id, + "mdl_hash": mdl_hash, "instructions": instructions, **self._components, } diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 86c17e93de..7753bfa819 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -10,7 +10,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.common import ScoreFilter +from src.pipelines.common import ScoreFilter, build_project_deploy_filter logger = logging.getLogger("wren-ai-service") @@ -57,18 +57,11 @@ def run( ## Start of Pipeline @observe(capture_input=False) async def count_documents( - store: QdrantDocumentStore, project_id: Optional[str] = None + store: QdrantDocumentStore, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> int: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) document_count = await store.count_documents(filters=filters) return document_count @@ -82,7 +75,12 @@ async def embedding(count_documents: int, query: str, embedder: Any) -> dict: @observe(capture_input=False) -async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: +async def retrieval( + embedding: dict, + project_id: str, + retriever: Any, + mdl_hash: str = "", +) -> dict: if not embedding: return {} @@ -98,6 +96,11 @@ async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: {"field": "project_id", "operator": "==", "value": project_id} ) + if mdl_hash: + filters["conditions"].append( + {"field": "mdl_hash", "operator": "==", "value": mdl_hash} + ) + res = await retriever.run( query_embedding=embedding.get("embedding"), filters=filters, @@ -136,6 +139,7 @@ async def default_instructions( project_id: str, scope_filter: ScopeFilter, scope: str, + mdl_hash: str = "", ) -> list[Document]: if not count_documents: return [] @@ -152,6 +156,11 @@ async def default_instructions( {"field": "project_id", "operator": "==", "value": project_id} ) + if mdl_hash: + filters["conditions"].append( + {"field": "mdl_hash", "operator": "==", "value": mdl_hash} + ) + _res = await retriever.run( query_embedding=None, filters=filters, @@ -213,7 +222,11 @@ def __init__( @observe(name="Instructions Retrieval") async def run( - self, query: str, project_id: Optional[str] = None, scope: str = "sql" + self, + query: str, + project_id: Optional[str] = None, + scope: str = "sql", + mdl_hash: Optional[str] = None, ): logger.info("Instructions Retrieval pipeline is running...") return await self._pipe.execute( @@ -222,6 +235,7 @@ async def run( "query": query, "project_id": project_id or "", "scope": scope, + "mdl_hash": mdl_hash or "", **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py index 3fe44f32eb..d836cbe37a 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py @@ -10,7 +10,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.common import ScoreFilter +from src.pipelines.common import ScoreFilter, build_project_deploy_filter logger = logging.getLogger("wren-ai-service") @@ -36,18 +36,11 @@ def run(self, documents: List[Document]): ## Start of Pipeline @observe(capture_input=False) async def count_documents( - store: QdrantDocumentStore, project_id: Optional[str] = None + store: QdrantDocumentStore, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> int: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) document_count = await store.count_documents(filters=filters) return document_count @@ -61,18 +54,14 @@ async def embedding(count_documents: int, query: str, embedder: Any) -> dict: @observe(capture_input=False) -async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: +async def retrieval( + embedding: dict, + project_id: str, + retriever: Any, + mdl_hash: str = "", +) -> dict: if embedding: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) res = await retriever.run( query_embedding=embedding.get("embedding"), @@ -143,13 +132,19 @@ def __init__( ) @observe(name="SqlPairs Retrieval") - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): logger.info("SqlPairs Retrieval pipeline is running...") return await self._pipe.execute( ["formatted_output"], inputs={ "query": query, "project_id": project_id or "", + "mdl_hash": mdl_hash or "", **self._components, **self._configs, }, diff --git a/wren-ai-service/src/web/v1/services/__init__.py b/wren-ai-service/src/web/v1/services/__init__.py index 5296b274b4..c6d64026f8 100644 --- a/wren-ai-service/src/web/v1/services/__init__.py +++ b/wren-ai-service/src/web/v1/services/__init__.py @@ -61,6 +61,7 @@ class BaseRequest(BaseModel): default=None, validation_alias=AliasChoices("project_id", "projectId"), ) + mdl_hash: Optional[str] = None thread_id: Optional[str] = None configurations: Configuration = Field( default_factory=Configuration, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 1a0052a2a6..4a669a9f2e 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -218,10 +218,12 @@ async def ask( self._pipelines["sql_pairs_retrieval"].run( query=user_query, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ), self._pipelines["instructions_retrieval"].run( query=user_query, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, scope="sql", ), ) diff --git a/wren-ai-service/src/web/v1/services/instructions.py b/wren-ai-service/src/web/v1/services/instructions.py index 8a825938d2..1dae1649dd 100644 --- a/wren-ai-service/src/web/v1/services/instructions.py +++ b/wren-ai-service/src/web/v1/services/instructions.py @@ -103,6 +103,7 @@ async def index( await self._pipelines["instructions_indexing"].run( project_id=request.project_id, + mdl_hash=request.mdl_hash, instructions=instructions, ) diff --git a/wren-ai-service/src/web/v1/services/sql_pairs.py b/wren-ai-service/src/web/v1/services/sql_pairs.py index 84291baf8b..f6750b05cf 100644 --- a/wren-ai-service/src/web/v1/services/sql_pairs.py +++ b/wren-ai-service/src/web/v1/services/sql_pairs.py @@ -69,6 +69,7 @@ async def index( input = { "mdl_str": '{"models": [{"properties": {"boilerplate": "sql_pairs"}}]}', "project_id": request.project_id, + "mdl_hash": request.mdl_hash, "external_pairs": { "sql_pairs": [ sql_pair.model_dump() for sql_pair in request.sql_pairs diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py index d058274eff..35d016d7ef 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py @@ -12,6 +12,14 @@ ], } +PROJECT_DEPLOY_FILTER = { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-a"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-a"}, + ], +} + class StoreSpy: def __init__(self, count=0): @@ -58,6 +66,20 @@ async def test_sql_pairs_count_stays_project_scoped_when_project_has_no_document assert store.filters == [PROJECT_FILTER] +@pytest.mark.asyncio +async def test_sql_pairs_count_can_be_deploy_scoped(): + store = StoreSpy(count=0) + + count = await sql_pairs_retrieval.count_documents( + store, + project_id="project-a", + mdl_hash="deploy-a", + ) + + assert count == 0 + assert store.filters == [PROJECT_DEPLOY_FILTER] + + @pytest.mark.asyncio async def test_sql_pairs_retrieval_does_not_fall_back_to_global_documents(): retriever = RetrieverSpy() @@ -72,6 +94,21 @@ async def test_sql_pairs_retrieval_does_not_fall_back_to_global_documents(): assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] +@pytest.mark.asyncio +async def test_sql_pairs_retrieval_can_be_deploy_scoped(): + retriever = RetrieverSpy() + + result = await sql_pairs_retrieval.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + mdl_hash="deploy-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [PROJECT_DEPLOY_FILTER] + + @pytest.mark.asyncio async def test_historical_question_count_stays_project_scoped_when_project_has_no_documents(): store = StoreSpy(count=0) @@ -109,6 +146,20 @@ async def test_instruction_count_stays_project_scoped_when_project_has_no_docume assert store.filters == [PROJECT_FILTER] +@pytest.mark.asyncio +async def test_instruction_count_can_be_deploy_scoped(): + store = StoreSpy(count=0) + + count = await instructions.count_documents( + store, + project_id="project-a", + mdl_hash="deploy-a", + ) + + assert count == 0 + assert store.filters == [PROJECT_DEPLOY_FILTER] + + @pytest.mark.asyncio async def test_instruction_retrieval_does_not_fall_back_to_global_documents(): retriever = RetrieverSpy() @@ -131,6 +182,30 @@ async def test_instruction_retrieval_does_not_fall_back_to_global_documents(): ] +@pytest.mark.asyncio +async def test_instruction_retrieval_can_be_deploy_scoped(): + retriever = RetrieverSpy() + + result = await instructions.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + mdl_hash="deploy-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": False}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-a"}, + ], + } + ] + + @pytest.mark.asyncio async def test_default_instructions_do_not_fall_back_to_global_documents(): retriever = RetrieverSpy() @@ -153,3 +228,29 @@ async def test_default_instructions_do_not_fall_back_to_global_documents(): ], } ] + + +@pytest.mark.asyncio +async def test_default_instructions_can_be_deploy_scoped(): + retriever = RetrieverSpy() + + result = await instructions.default_instructions( + count_documents=1, + retriever=retriever, + project_id="project-a", + mdl_hash="deploy-a", + scope_filter=instructions.ScopeFilter(), + scope="sql", + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": True}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-a"}, + ], + } + ] diff --git a/wren-ai-service/tests/pytest/services/mocks.py b/wren-ai-service/tests/pytest/services/mocks.py index a6ce06d323..7de9a4ffaa 100644 --- a/wren-ai-service/tests/pytest/services/mocks.py +++ b/wren-ai-service/tests/pytest/services/mocks.py @@ -9,7 +9,13 @@ class RetrievalMock(retrieval.DbSchemaRetrieval): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + **_, + ): return {"construct_retrieval_results": self._documents} @@ -17,7 +23,12 @@ class SqlPairsRetrievalMock(retrieval.SqlPairsRetrieval): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): return {"formatted_output": {"documents": self._documents}} @@ -25,7 +36,13 @@ class InstructionsRetrievalMock(retrieval.Instructions): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + scope: str = "sql", + mdl_hash: Optional[str] = None, + ): return {"formatted_output": {"documents": self._documents}} @@ -33,7 +50,12 @@ class HistoricalQuestionMock(retrieval.HistoricalQuestionRetrieval): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): return {"formatted_output": {"documents": self._documents}} diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 5acc9af772..a74968c0db 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -121,6 +121,7 @@ export interface IWrenAIAdaptor { deploySqlPair( projectId: number, sqlPair: { question: string; sql: string }, + mdlHash?: string, ): Promise; getSqlPairResult(queryId: string): Promise; deleteSqlPairs(projectId: number, sqlPairIds: number[]): Promise; @@ -132,6 +133,7 @@ export interface IWrenAIAdaptor { */ generateInstruction( input: GenerateInstructionInput[], + mdlHash?: string, ): Promise; getInstructionResult(queryId: string): Promise; deleteInstructions(ids: number[], projectId: number): Promise; @@ -178,9 +180,10 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async deploySqlPair( projectId: number, sqlPair: Partial, + mdlHash?: string, ): Promise { try { - const body = { + const body: any = { sql_pairs: [ { id: `${sqlPair.id}`, @@ -190,6 +193,9 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { ], project_id: projectId.toString(), }; + if (mdlHash) { + body['mdl_hash'] = mdlHash; + } return axios .post(`${this.wrenAIBaseEndpoint}/v1/sql-pairs`, body) @@ -667,8 +673,9 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async generateInstruction( input: GenerateInstructionInput[], + mdlHash?: string, ): Promise { - const body = { + const body: any = { instructions: input.map((item) => ({ id: item.id.toString(), instruction: item.instruction, @@ -677,6 +684,9 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { })), project_id: input[0]?.projectId.toString(), }; + if (mdlHash) { + body['mdl_hash'] = mdlHash; + } try { const res = await axios.post( `${this.wrenAIBaseEndpoint}/v1/instructions`, diff --git a/wren-ui/src/apollo/server/services/instructionService.ts b/wren-ui/src/apollo/server/services/instructionService.ts index 37112ad5f4..e8ff322ca5 100644 --- a/wren-ui/src/apollo/server/services/instructionService.ts +++ b/wren-ui/src/apollo/server/services/instructionService.ts @@ -9,6 +9,7 @@ import { import { IInstructionRepository, Instruction } from '@server/repositories'; import * as Errors from '@server/utils/error'; import { GeneralErrorCodes } from '@server/utils/error'; +import { IDeployService } from './deployService'; export interface IInstructionService { getInstructions(projectId: number): Promise; getInstruction(id: number): Promise; @@ -21,15 +22,19 @@ export interface IInstructionService { export class InstructionService implements IInstructionService { private readonly instructionRepository: IInstructionRepository; private readonly wrenAIAdaptor: IWrenAIAdaptor; + private readonly deployService: IDeployService; constructor({ instructionRepository, wrenAIAdaptor, + deployService, }: { instructionRepository: IInstructionRepository; wrenAIAdaptor: IWrenAIAdaptor; + deployService: IDeployService; }) { this.instructionRepository = instructionRepository; this.wrenAIAdaptor = wrenAIAdaptor; + this.deployService = deployService; } public async getInstructions(projectId: number): Promise { @@ -56,9 +61,11 @@ export class InstructionService implements IInstructionService { tx, }, ); - const { queryId } = await this.wrenAIAdaptor.generateInstruction([ - this.pickGenerateInstructionInput(newInstruction), - ]); + const mdlHash = await this.getDeployHash(input.projectId); + const { queryId } = await this.wrenAIAdaptor.generateInstruction( + [this.pickGenerateInstructionInput(newInstruction)], + mdlHash, + ); const res = await this.waitDeployInstruction(queryId); if (res.error) { await tx.rollback(); @@ -90,8 +97,10 @@ export class InstructionService implements IInstructionService { tx, }, ); + const mdlHash = await this.getDeployHash(inputs[0]?.projectId); const { queryId } = await this.wrenAIAdaptor.generateInstruction( newInstructions.map(this.pickGenerateInstructionInput), + mdlHash, ); const res = await this.waitDeployInstruction(queryId); if (res.error) { @@ -133,9 +142,11 @@ export class InstructionService implements IInstructionService { instructionData, { tx }, ); - const { queryId } = await this.wrenAIAdaptor.generateInstruction([ - this.pickGenerateInstructionInput(updatedInstruction), - ]); + const mdlHash = await this.getDeployHash(input.projectId); + const { queryId } = await this.wrenAIAdaptor.generateInstruction( + [this.pickGenerateInstructionInput(updatedInstruction)], + mdlHash, + ); const res = await this.waitDeployInstruction(queryId); if (res.error) { await tx.rollback(); @@ -215,4 +226,12 @@ export class InstructionService implements IInstructionService { throw new Error('Instruction is too long'); } } + + private async getDeployHash(projectId?: number): Promise { + if (!projectId) { + return undefined; + } + const deployment = await this.deployService.getLastDeployment(projectId); + return deployment?.hash; + } } diff --git a/wren-ui/src/apollo/server/services/sqlPairService.ts b/wren-ui/src/apollo/server/services/sqlPairService.ts index fd9c8b186d..51eb915486 100644 --- a/wren-ui/src/apollo/server/services/sqlPairService.ts +++ b/wren-ui/src/apollo/server/services/sqlPairService.ts @@ -17,6 +17,7 @@ import { } from '../models/adaptor'; import { Manifest } from '@server/mdl/type'; import { DataSourceName } from '@server/types'; +import { IDeployService } from './deployService'; const logger = getLogger('SqlPairService'); @@ -60,19 +61,23 @@ export class SqlPairService implements ISqlPairService { private sqlPairRepository: ISqlPairRepository; private wrenAIAdaptor: IWrenAIAdaptor; private ibisAdaptor: IIbisAdaptor; + private deployService: IDeployService; constructor({ sqlPairRepository, wrenAIAdaptor, ibisAdaptor, + deployService, }: { sqlPairRepository: ISqlPairRepository; wrenAIAdaptor: IWrenAIAdaptor; ibisAdaptor: IIbisAdaptor; + deployService: IDeployService; }) { this.sqlPairRepository = sqlPairRepository; this.wrenAIAdaptor = wrenAIAdaptor; this.ibisAdaptor = ibisAdaptor; + this.deployService = deployService; } public async modelSubstitute( @@ -145,9 +150,11 @@ export class SqlPairService implements ISqlPairService { }, { tx }, ); + const mdlHash = await this.getDeployHash(projectId); const { queryId } = await this.wrenAIAdaptor.deploySqlPair( projectId, newPair, + mdlHash, ); const deployResult = await this.waitUntilSqlPairResult(queryId); if (deployResult.error) { @@ -179,12 +186,14 @@ export class SqlPairService implements ISqlPairService { const successPairs = []; const errorPairs = []; const chunks = chunk(newPairs, 10); + const mdlHash = await this.getDeployHash(projectId); for (const pairs of chunks) { await Promise.allSettled( pairs.map(async (pair) => { const { queryId } = await this.wrenAIAdaptor.deploySqlPair( projectId, pair, + mdlHash, ); const deployResult = await this.waitUntilSqlPairResult(queryId); if (deployResult.error) { @@ -249,9 +258,11 @@ export class SqlPairService implements ISqlPairService { updatedData, { tx }, ); + const mdlHash = await this.getDeployHash(projectId); const { queryId } = await this.wrenAIAdaptor.deploySqlPair( projectId, updatedSqlPair, + mdlHash, ); const deployResult = await this.waitUntilSqlPairResult(queryId); if (deployResult.error) { @@ -326,4 +337,9 @@ export class SqlPairService implements ISqlPairService { private isFinishedState(status: SqlPairStatus) { return [SqlPairStatus.FINISHED, SqlPairStatus.FAILED].includes(status); } + + private async getDeployHash(projectId: number): Promise { + const deployment = await this.deployService.getLastDeployment(projectId); + return deployment?.hash; + } } diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index e9be30d55b..828d128cd5 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -273,10 +273,12 @@ export const initComponents = () => { sqlPairRepository, wrenAIAdaptor, ibisAdaptor, + deployService, }); const instructionService = new InstructionService({ instructionRepository, wrenAIAdaptor, + deployService, }); const rbacService = new RbacService({ roleRepository, From 2c7e97a2de3d921cb8a022e9107de421119aa668 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 2 Aug 2026 21:45:57 +0530 Subject: [PATCH 0789/1087] Fallback to project metadata when deploy hash is absent --- .../retrieval/db_schema_retrieval.py | 52 +++++++++++++++++-- 1 file changed, 48 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index f2e870b263..e257172d7b 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,6 +1,7 @@ import ast import logging import sys +from copy import deepcopy from typing import Any, Optional import orjson @@ -216,6 +217,37 @@ def _build_view_ddl(content: dict) -> str: ) +def _remove_mdl_hash_filter(filters: dict | None) -> dict | None: + if not filters: + return filters + + fallback_filters = deepcopy(filters) + fallback_filters["conditions"] = [ + condition + for condition in fallback_filters.get("conditions", []) + if condition.get("field") != "mdl_hash" + ] + return fallback_filters + + +async def _run_with_project_metadata_fallback( + retriever: Any, + query_embedding: list, + filters: dict, + project_id: str, + mdl_hash: str, +) -> dict: + result = await retriever.run(query_embedding=query_embedding, filters=filters) + if result.get("documents") or not (project_id and mdl_hash): + return result + + fallback_filters = _remove_mdl_hash_filter(filters) + return await retriever.run( + query_embedding=query_embedding, + filters=fallback_filters, + ) + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -251,18 +283,24 @@ async def table_retrieval( ) if embedding: - return await table_retriever.run( + return await _run_with_project_metadata_fallback( + table_retriever, query_embedding=embedding.get("embedding"), filters=filters, + project_id=project_id, + mdl_hash=mdl_hash, ) else: filters["conditions"].append( {"field": "name", "operator": "in", "value": tables} ) - return await table_retriever.run( + return await _run_with_project_metadata_fallback( + table_retriever, query_embedding=[], filters=filters, + project_id=project_id, + mdl_hash=mdl_hash, ) @@ -309,9 +347,12 @@ async def _fetch_by_names(table_names: list[str]) -> list[Document]: if not table_names: return [] - results = await dbschema_retriever.run( + results = await _run_with_project_metadata_fallback( + dbschema_retriever, query_embedding=[], filters=_filters_for_names(table_names), + project_id=project_id, + mdl_hash=mdl_hash, ) return results["documents"] @@ -350,9 +391,12 @@ def _related_table_names(documents: list[Document], visited: set[str]) -> list[s documents = [] if not table_names and embedding and embedding.get("embedding"): - results = await dbschema_retriever.run( + results = await _run_with_project_metadata_fallback( + dbschema_retriever, query_embedding=embedding.get("embedding"), filters=_base_filters(), + project_id=project_id, + mdl_hash=mdl_hash, ) documents.extend(results["documents"]) for document in results["documents"]: From 006325a5b64516c35dc487cc41756660adf39890 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 2 Aug 2026 22:06:24 +0530 Subject: [PATCH 0790/1087] Return generated SQL when validation times out --- .../src/pipelines/generation/utils/sql.py | 48 ++++++--- .../generation/test_sql_post_processor.py | 99 +++++++++++++++++++ 2 files changed, 135 insertions(+), 12 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1633c399e4..834d88247d 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -17,6 +17,10 @@ logger = logging.getLogger("wren-ai-service") +def _is_timeout_error(error_message: str) -> bool: + return error_message.startswith("Request timed out") + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -101,12 +105,17 @@ async def _classify_generation_result( ) if not dry_plan_result: + if _is_timeout_error(error_message): + valid_generation_result = { + "sql": generation_result, + "correlation_id": "", + } + return valid_generation_result, invalid_generation_result + invalid_generation_result = { "sql": generation_result, "original_sql": generation_result, - "type": "TIME_OUT" - if error_message.startswith("Request timed out") - else "DRY_PLAN", + "type": "DRY_PLAN", "error": error_message, "correlation_id": "", } @@ -128,12 +137,17 @@ async def _classify_generation_result( } else: error_message = addition.get("error_message", "") + if _is_timeout_error(error_message): + valid_generation_result = { + "sql": generation_result, + "correlation_id": addition.get("correlation_id", ""), + } + return valid_generation_result, invalid_generation_result + invalid_generation_result = { "sql": addition.get("error_sql", generation_result), "original_sql": generation_result, - "type": "TIME_OUT" - if error_message.startswith("Request timed out") - else "DRY_RUN", + "type": "DRY_RUN", "error": error_message, "correlation_id": addition.get("correlation_id", ""), } @@ -153,12 +167,17 @@ async def _classify_generation_result( } else: error_message = addition.get("error_message", "") + if _is_timeout_error(error_message): + valid_generation_result = { + "sql": generation_result, + "correlation_id": addition.get("correlation_id", ""), + } + return valid_generation_result, invalid_generation_result + invalid_generation_result = { "sql": addition.get("error_sql", generation_result), "original_sql": generation_result, - "type": "TIME_OUT" - if error_message.startswith("Request timed out") - else "DRY_RUN", + "type": "DRY_RUN", "error": error_message, "correlation_id": addition.get("correlation_id", ""), } @@ -178,6 +197,13 @@ async def _classify_generation_result( } else: error_message = addition.get("error_message", "") + if _is_timeout_error(error_message): + valid_generation_result = { + "sql": generation_result, + "correlation_id": addition.get("correlation_id", ""), + } + return valid_generation_result, invalid_generation_result + preview_data_status = ( "PREVIEW_EMPTY_DATA" if error_message == "" @@ -186,9 +212,7 @@ async def _classify_generation_result( invalid_generation_result = { "sql": addition.get("error_sql", generation_result), "original_sql": generation_result, - "type": "TIME_OUT" - if error_message.startswith("Request timed out") - else preview_data_status, + "type": preview_data_status, "error": error_message, "correlation_id": addition.get("correlation_id", ""), } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py new file mode 100644 index 0000000000..02ef353189 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -0,0 +1,99 @@ +import pytest + +from src.pipelines.generation.utils.sql import SQLGenPostProcessor + + +class TimeoutEngine: + async def dry_plan( + self, + session, + sql, + data_source, + project_id=None, + allow_fallback=True, + ): + return False, "Request timed out after 30 seconds" + + async def execute_sql( + self, + sql, + session, + project_id=None, + limit=1, + dry_run=True, + ): + return False, None, { + "error_message": "Request timed out after 30 seconds", + "correlation_id": "timeout-correlation", + } + + +class ErrorEngine: + async def dry_plan( + self, + session, + sql, + data_source, + project_id=None, + allow_fallback=True, + ): + return False, "Planner rejected the statement" + + async def execute_sql( + self, + sql, + session, + project_id=None, + limit=1, + dry_run=True, + ): + return False, None, {"error_message": "Execution rejected the statement"} + + +@pytest.mark.asyncio +async def test_sql_post_processor_returns_generated_sql_when_dry_plan_times_out(): + result = await SQLGenPostProcessor(TimeoutEngine()).run( + ['{"sql": "SELECT 1"}'], + project_id="project-id", + use_dry_plan=True, + data_source="source", + ) + + assert result["valid_generation_result"] == { + "sql": "SELECT 1", + "correlation_id": "", + } + assert result["invalid_generation_result"] == {} + + +@pytest.mark.asyncio +async def test_sql_post_processor_returns_generated_sql_when_dry_run_times_out(): + result = await SQLGenPostProcessor(TimeoutEngine()).run( + ['{"sql": "SELECT 1"}'], + project_id="project-id", + ) + + assert result["valid_generation_result"] == { + "sql": "SELECT 1", + "correlation_id": "timeout-correlation", + } + assert result["invalid_generation_result"] == {} + + +@pytest.mark.asyncio +async def test_sql_post_processor_keeps_non_timeout_dry_plan_errors_invalid(): + result = await SQLGenPostProcessor(ErrorEngine()).run( + ['{"sql": "SELECT 1"}'], + project_id="project-id", + use_dry_plan=True, + data_source="source", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT 1", + "original_sql": "SELECT 1", + "type": "DRY_PLAN", + "error": "Planner rejected the statement", + "correlation_id": "", + } From 91faf4d6a77572b7b0c850a94600f40a4c117369 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 2 Aug 2026 23:46:49 +0530 Subject: [PATCH 0791/1087] Keep ask retrieval scoped to deployed metadata --- .../retrieval/db_schema_retrieval.py | 52 +----------- .../retrieval/test_db_schema_retrieval.py | 81 +++++++++++++++++++ 2 files changed, 85 insertions(+), 48 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e257172d7b..f2e870b263 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,6 @@ import ast import logging import sys -from copy import deepcopy from typing import Any, Optional import orjson @@ -217,37 +216,6 @@ def _build_view_ddl(content: dict) -> str: ) -def _remove_mdl_hash_filter(filters: dict | None) -> dict | None: - if not filters: - return filters - - fallback_filters = deepcopy(filters) - fallback_filters["conditions"] = [ - condition - for condition in fallback_filters.get("conditions", []) - if condition.get("field") != "mdl_hash" - ] - return fallback_filters - - -async def _run_with_project_metadata_fallback( - retriever: Any, - query_embedding: list, - filters: dict, - project_id: str, - mdl_hash: str, -) -> dict: - result = await retriever.run(query_embedding=query_embedding, filters=filters) - if result.get("documents") or not (project_id and mdl_hash): - return result - - fallback_filters = _remove_mdl_hash_filter(filters) - return await retriever.run( - query_embedding=query_embedding, - filters=fallback_filters, - ) - - ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -283,24 +251,18 @@ async def table_retrieval( ) if embedding: - return await _run_with_project_metadata_fallback( - table_retriever, + return await table_retriever.run( query_embedding=embedding.get("embedding"), filters=filters, - project_id=project_id, - mdl_hash=mdl_hash, ) else: filters["conditions"].append( {"field": "name", "operator": "in", "value": tables} ) - return await _run_with_project_metadata_fallback( - table_retriever, + return await table_retriever.run( query_embedding=[], filters=filters, - project_id=project_id, - mdl_hash=mdl_hash, ) @@ -347,12 +309,9 @@ async def _fetch_by_names(table_names: list[str]) -> list[Document]: if not table_names: return [] - results = await _run_with_project_metadata_fallback( - dbschema_retriever, + results = await dbschema_retriever.run( query_embedding=[], filters=_filters_for_names(table_names), - project_id=project_id, - mdl_hash=mdl_hash, ) return results["documents"] @@ -391,12 +350,9 @@ def _related_table_names(documents: list[Document], visited: set[str]) -> list[s documents = [] if not table_names and embedding and embedding.get("embedding"): - results = await _run_with_project_metadata_fallback( - dbschema_retriever, + results = await dbschema_retriever.run( query_embedding=embedding.get("embedding"), filters=_base_filters(), - project_id=project_id, - mdl_hash=mdl_hash, ) documents.extend(results["documents"]) for document in results["documents"]: diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7bbc25248d..7b7308947f 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -132,6 +132,46 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_table_retrieval_keeps_deploy_scope_when_no_documents_match(): + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + return {"documents": []} + + retriever = Retriever() + + await table_retrieval( + embedding={"embedding": [0.25]}, + project_id="project-1", + mdl_hash="deploy-1", + tables=[], + table_retriever=retriever, + ) + + assert retriever.calls == [ + { + "query_embedding": [0.25], + "filters": { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + }, + } + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: @@ -197,6 +237,47 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_dbschema_retrieval_keeps_deploy_scope_when_no_documents_match(): + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + return {"documents": []} + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={"documents": []}, + project_id="project-1", + mdl_hash="deploy-1", + dbschema_retriever=retriever, + embedding={"embedding": [0.25]}, + ) + + assert documents == [] + assert retriever.calls == [ + { + "query_embedding": [0.25], + "filters": { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + }, + } + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_expands_declared_relationships(): selected_model = "model_a" From efeffbc53ed41a59809ef5f6aaccfb89f446dff7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 12:18:38 +0530 Subject: [PATCH 0792/1087] Canonicalize TOP limits to Wren SQL --- .../src/pipelines/generation/utils/sql.py | 46 +++++++++++++++++++ .../generation/test_sql_post_processor.py | 33 +++++++++++++ 2 files changed, 79 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 834d88247d..e343ab17a9 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3,8 +3,10 @@ import aiohttp import orjson +import sqlparse from haystack import component from haystack.dataclasses import ChatMessage +from sqlparse import tokens as sqlparse_tokens from pydantic import BaseModel from src.core.engine import ( @@ -21,6 +23,46 @@ def _is_timeout_error(error_message: str) -> bool: return error_message.startswith("Request timed out") +def _canonicalize_wren_sql_syntax(sql: str | None) -> str | None: + if not sql: + return sql + + statements = sqlparse.parse(sql) + if len(statements) != 1: + return sql + + statement = statements[0] + tokens = statement.tokens + significant_tokens = [ + (index, token) for index, token in enumerate(tokens) if not token.is_whitespace + ] + + if len(significant_tokens) < 3: + return sql + + first_token = significant_tokens[0][1] + top_token_index, top_token = significant_tokens[1] + limit_token_index, limit_token = significant_tokens[2] + has_limit = any(token.normalized == "LIMIT" for _, token in significant_tokens) + + if ( + first_token.ttype != sqlparse_tokens.Keyword.DML + or first_token.normalized != "SELECT" + or top_token.normalized != "TOP" + or limit_token.ttype != sqlparse_tokens.Literal.Number.Integer + or has_limit + ): + return sql + + before_top = "".join(str(token) for token in tokens[:top_token_index]).rstrip() + after_limit = "".join(str(token) for token in tokens[limit_token_index + 1 :]).lstrip() + + if not before_top or not after_limit: + return sql + + return f"{before_top} {after_limit} LIMIT {limit_token.value}".strip() + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -48,6 +90,10 @@ async def run( "sql" ) + cleaned_generation_result = _canonicalize_wren_sql_syntax( + cleaned_generation_result + ) + ( valid_generation_result, invalid_generation_result, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index 02ef353189..855d115811 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -50,6 +50,39 @@ async def execute_sql( return False, None, {"error_message": "Execution rejected the statement"} +class CapturingEngine: + def __init__(self): + self.sql = None + + async def execute_sql( + self, + sql, + session, + project_id=None, + limit=1, + dry_run=True, + ): + self.sql = sql + return True, None, {"correlation_id": "valid-correlation"} + + +@pytest.mark.asyncio +async def test_sql_post_processor_converts_select_top_to_wren_limit(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT TOP 5 1 ORDER BY 1"}'], + project_id="project-id", + ) + + assert engine.sql == "SELECT 1 ORDER BY 1 LIMIT 5" + assert result["valid_generation_result"] == { + "sql": "SELECT 1 ORDER BY 1 LIMIT 5", + "correlation_id": "valid-correlation", + } + assert result["invalid_generation_result"] == {} + + @pytest.mark.asyncio async def test_sql_post_processor_returns_generated_sql_when_dry_plan_times_out(): result = await SQLGenPostProcessor(TimeoutEngine()).run( From 8b1e1b012f87113dd19f5110243773cbf769d846 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 12:58:00 +0530 Subject: [PATCH 0793/1087] Select available deployed schema scope for ask retrieval --- .../retrieval/db_schema_retrieval.py | 42 +++++++-- .../retrieval/test_db_schema_retrieval.py | 90 +++++++++++++++++++ 2 files changed, 127 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index f2e870b263..fb8b4d9eae 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,4 +1,5 @@ import ast +import asyncio import logging import sys from typing import Any, Optional @@ -217,6 +218,25 @@ def _build_view_ddl(content: dict) -> str: ## Start of Pipeline +@observe(capture_input=False) +async def active_mdl_hash( + project_id: str, + mdl_hash: str, + table_description_store: Any, + dbschema_store: Any, +) -> str: + if not project_id or not mdl_hash: + return mdl_hash + + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) + table_description_count, dbschema_count = await asyncio.gather( + table_description_store.count_documents(filters=filters), + dbschema_store.count_documents(filters=filters), + ) + + return mdl_hash if table_description_count or dbschema_count else "" + + @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: @@ -231,8 +251,10 @@ async def table_retrieval( project_id: str, tables: list[str], table_retriever: Any, + active_mdl_hash: Optional[str] = None, mdl_hash: str = "", ) -> dict: + effective_mdl_hash = active_mdl_hash if active_mdl_hash is not None else mdl_hash filters = { "operator": "AND", "conditions": [ @@ -245,9 +267,9 @@ async def table_retrieval( {"field": "project_id", "operator": "==", "value": project_id} ) - if mdl_hash: + if effective_mdl_hash: filters["conditions"].append( - {"field": "mdl_hash", "operator": "==", "value": mdl_hash} + {"field": "mdl_hash", "operator": "==", "value": effective_mdl_hash} ) if embedding: @@ -271,13 +293,16 @@ async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any, + active_mdl_hash: Optional[str] = None, mdl_hash: str = "", embedding: Optional[dict] = None, ) -> list[Document]: + effective_mdl_hash = active_mdl_hash if active_mdl_hash is not None else mdl_hash + def _base_filters() -> dict: project_deploy_filter = build_project_deploy_filter( project_id=project_id, - mdl_hash=mdl_hash, + mdl_hash=effective_mdl_hash, ) filters = { "operator": "AND", @@ -627,16 +652,23 @@ def __init__( table_column_retrieval_size: int = 100, **kwargs, ): + table_description_store = document_store_provider.get_store( + dataset_name="table_descriptions" + ) + dbschema_store = document_store_provider.get_store() + self._components = { "embedder": embedder_provider.get_text_embedder(), "table_retriever": document_store_provider.get_retriever( - document_store_provider.get_store(dataset_name="table_descriptions"), + table_description_store, top_k=table_retrieval_size, ), "dbschema_retriever": document_store_provider.get_retriever( - document_store_provider.get_store(), + dbschema_store, top_k=table_column_retrieval_size, ), + "table_description_store": table_description_store, + "dbschema_store": dbschema_store, "table_columns_selection_generator": llm_provider.get_generator( system_prompt=table_columns_selection_system_prompt, generation_kwargs=RETRIEVAL_MODEL_KWARGS, diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7b7308947f..4f05cced7a 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -4,6 +4,7 @@ from src.pipelines.common import build_table_ddl from src.pipelines.retrieval.db_schema_retrieval import ( + active_mdl_hash, _build_view_ddl, check_using_db_schemas_without_pruning, construct_retrieval_results, @@ -16,6 +17,16 @@ ) +class StoreCounter: + def __init__(self, count): + self.count = count + self.filters = [] + + async def count_documents(self, filters=None): + self.filters.append(filters) + return self.count + + @pytest.mark.asyncio async def test_embedding_uses_current_query_without_history_text(): class Embedder: @@ -79,6 +90,45 @@ def test_table_selection_prompt_keeps_multiple_relevant_datasets(): ) +@pytest.mark.asyncio +async def test_active_mdl_hash_keeps_hash_when_deploy_documents_are_indexed(): + table_store = StoreCounter(count=1) + schema_store = StoreCounter(count=0) + + result = await active_mdl_hash( + project_id="project-1", + mdl_hash="deploy-1", + table_description_store=table_store, + dbschema_store=schema_store, + ) + + expected_filters = { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + } + assert result == "deploy-1" + assert table_store.filters == [expected_filters] + assert schema_store.filters == [expected_filters] + + +@pytest.mark.asyncio +async def test_active_mdl_hash_uses_project_scope_when_deploy_documents_are_absent(): + table_store = StoreCounter(count=0) + schema_store = StoreCounter(count=0) + + result = await active_mdl_hash( + project_id="project-1", + mdl_hash="deploy-1", + table_description_store=table_store, + dbschema_store=schema_store, + ) + + assert result == "" + + def test_view_schema_context_uses_declared_view_columns_not_view_definition(): result = _build_view_ddl( { @@ -172,6 +222,46 @@ async def run(self, query_embedding, filters): ] +@pytest.mark.asyncio +async def test_table_retrieval_uses_project_scope_when_active_hash_is_absent(): + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + return {"documents": []} + + retriever = Retriever() + + await table_retrieval( + embedding={"embedding": [0.25]}, + project_id="project-1", + mdl_hash="deploy-1", + active_mdl_hash="", + tables=[], + table_retriever=retriever, + ) + + assert retriever.calls == [ + { + "query_embedding": [0.25], + "filters": { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + }, + } + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_loads_selected_active_project_schema(): class Retriever: From 641b129d07eb6502f4336d48e70e35837396b779 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 13:13:13 +0530 Subject: [PATCH 0794/1087] Skip SQL reasoning by default for asks --- wren-ai-service/src/web/v1/services/ask.py | 4 ++-- wren-ai-service/tests/pytest/services/test_ask.py | 6 ++++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 4a669a9f2e..7959c372e3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -25,7 +25,7 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = False + ignore_sql_generation_reasoning: bool = True enable_column_pruning: bool = False use_dry_plan: bool = False allow_dry_plan_fallback: bool = True @@ -99,7 +99,7 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = True, + allow_sql_generation_reasoning: bool = False, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index d479ac9550..5187681d58 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -102,6 +102,12 @@ def mdl_str(): return orjson.dumps(json.load(f)).decode("utf-8") +def test_ask_request_skips_sql_generation_reasoning_by_default(): + ask_request = AskRequest(query="question", mdl_hash="deploy") + + assert ask_request.ignore_sql_generation_reasoning is True + + @pytest.mark.asyncio async def test_ask_with_successful_query( indexing_service: SemanticsPreparationService, From 695ff3e36876c6d3d8c0ea7c4d82282d67f6df91 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 13:42:56 +0530 Subject: [PATCH 0795/1087] Ground SQL correction in active schema context --- .../pipelines/generation/sql_correction.py | 25 ++++++++-- .../pipelines/generation/sql_generation.py | 3 +- wren-ai-service/src/web/v1/services/ask.py | 17 +++++-- .../generation/test_sql_prompt_grounding.py | 47 +++++++++++++++++++ 4 files changed, 82 insertions(+), 10 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 2118739a5c..cf0090549f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,12 +30,16 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query. +You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills. +You need to regenerate one grounded Wren SQL query from the user's question and current DATABASE SCHEMA. +The failed SQL and error message are diagnostic context only. ### SQL CORRECTION INSTRUCTIONS ### -1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). -2. Then, generate the syntactically correct ANSI SQL query to correct the error. +1. First, use the error message only to understand why the previous SQL failed. +2. Then, ignore any failed SQL identifier, placeholder, literal, or function that is not declared in the current DATABASE SCHEMA or SQL FUNCTIONS. +3. Regenerate the SQL from the user's question, DATABASE SCHEMA, SQL FUNCTIONS, and USER INSTRUCTIONS. +4. If the user's requested intent cannot be fully grounded by the current DATABASE SCHEMA and SQL FUNCTIONS, return null for sql instead of repairing the failed SQL approximately. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -73,11 +77,17 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endfor %} {% endif %} +{% if query %} ### QUESTION ### -SQL: {{ invalid_generation_result.sql }} +User's Question: {{ query }} +{% endif %} + +### FAILED SQL DIAGNOSTIC CONTEXT ### +Failed SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} -Let's think step by step. +Regenerate from the user's question and DATABASE SCHEMA only when a user question is available. Otherwise, correct the failed SQL only by using exact executable identifiers declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not copy table names, column names, functions, literals, aliases, or SQL structure from the failed SQL unless each one is declared in DATABASE SCHEMA or SQL FUNCTIONS. +Return only the final JSON SQL response. """ @@ -87,12 +97,14 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, + query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( documents=documents, invalid_generation_result=invalid_generation_result, + query=query or "", instructions=construct_instructions( instructions=instructions, ), @@ -169,6 +181,8 @@ async def run( self, contexts: List[Document], invalid_generation_result: Dict[str, str], + query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, @@ -192,6 +206,7 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, + "query": query, "documents": contexts, "instructions": instructions, "sql_functions": sql_functions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 71a80875e9..1d9fc2d0b7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -54,11 +54,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 7959c372e3..49828a7c21 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -557,6 +557,7 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, + query=user_query, instructions=instructions, invalid_generation_result={ "sql": original_sql, @@ -585,9 +586,19 @@ async def ask( ] break - failed_dry_run_result = sql_correction_results["post_process"][ - "invalid_generation_result" - ] + next_failed_dry_run_result = sql_correction_results[ + "post_process" + ]["invalid_generation_result"] + if ( + next_failed_dry_run_result + and next_failed_dry_run_result.get("sql") == invalid_sql + and next_failed_dry_run_result.get("error") + == error_message + ): + failed_dry_run_result = next_failed_dry_run_result + break + + failed_dry_run_result = next_failed_dry_run_result if api_results: if not self._is_stopped(query_id, self._ask_results): diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py new file mode 100644 index 0000000000..a42e1229d1 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -0,0 +1,47 @@ +from haystack.components.builders.prompt_builder import PromptBuilder + +from src.pipelines.generation.sql_correction import ( + prompt as build_sql_correction_prompt, + sql_correction_user_prompt_template, +) +from src.pipelines.generation.sql_generation import ( + prompt as build_sql_generation_prompt, + sql_generation_user_prompt_template, +) + + +def test_sql_generation_prompt_omits_sample_sql_body(): + result = build_sql_generation_prompt( + query="summarize the records", + documents=[], + prompt_builder=PromptBuilder(template=sql_generation_user_prompt_template), + sql_samples=[ + { + "question": "sample intent", + "sql": "SELECT 1", + } + ], + ) + + built_prompt = result["prompt"] + + assert "sample intent" in built_prompt + assert "SELECT 1" not in built_prompt + + +def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): + result = build_sql_correction_prompt( + documents=[], + invalid_generation_result={ + "sql": "SELECT 1", + "error": "dry run failed", + }, + query="summarize the records", + prompt_builder=PromptBuilder(template=sql_correction_user_prompt_template), + ) + + built_prompt = result["prompt"] + + assert "User's Question: summarize the records" in built_prompt + assert "Failed SQL: SELECT 1" in built_prompt + assert "DIAGNOSTIC CONTEXT" in built_prompt From 8730848baeea416b386c2d1203217a4356521ee4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 14:07:08 +0530 Subject: [PATCH 0796/1087] Block ungrounded SQL tables from execution --- .../generation/followup_sql_generation.py | 4 + .../pipelines/generation/sql_correction.py | 4 + .../pipelines/generation/sql_generation.py | 4 + .../pipelines/generation/sql_regeneration.py | 4 + .../src/pipelines/generation/utils/sql.py | 136 ++++++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 13 ++ .../src/web/v1/services/ask_feedback.py | 12 ++ .../src/web/v1/services/sql_corrections.py | 11 ++ .../generation/test_sql_post_processor.py | 56 ++++++++ 9 files changed, 244 insertions(+) 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 bc012d759c..922ffce6ed 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -149,6 +149,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + schema_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), @@ -156,6 +157,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + schema_contracts=schema_contracts, ) @@ -208,6 +210,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + schema_contracts: list[dict] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -239,6 +242,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "schema_contracts": schema_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index cf0090549f..485686af4d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -135,6 +135,7 @@ async def post_process( project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, + schema_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), @@ -142,6 +143,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + schema_contracts=schema_contracts, ) @@ -190,6 +192,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + schema_contracts: list[dict] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -216,6 +219,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "schema_contracts": schema_contracts, **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 1d9fc2d0b7..a99f3f571c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -142,6 +142,7 @@ async def post_process( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, + schema_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( generate_sql.get("replies"), @@ -150,6 +151,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + schema_contracts=schema_contracts, ) @@ -202,6 +204,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + schema_contracts: list[dict] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -233,6 +236,7 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, + "schema_contracts": schema_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 70a7b00e21..ca9d9bc287 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -164,10 +164,12 @@ async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, project_id: str | None = None, + schema_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, + schema_contracts=schema_contracts, ) @@ -212,6 +214,7 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_contracts: list[dict] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -230,6 +233,7 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, + "schema_contracts": schema_contracts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index e343ab17a9..09bb6131f3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -6,6 +6,7 @@ import sqlparse from haystack import component from haystack.dataclasses import ChatMessage +from sqlparse.sql import Identifier, IdentifierList, Parenthesis, TokenList from sqlparse import tokens as sqlparse_tokens from pydantic import BaseModel @@ -63,6 +64,125 @@ def _canonicalize_wren_sql_syntax(sql: str | None) -> str | None: return f"{before_top} {after_limit} LIMIT {limit_token.value}".strip() +def _meaningful_tokens(token_list: TokenList) -> list: + return [ + token + for token in token_list.tokens + if not token.is_whitespace and token.ttype not in sqlparse_tokens.Comment + ] + + +def _identifier_name(identifier: Identifier) -> str | None: + return identifier.get_real_name() or identifier.get_name() + + +def _contains_select(token: TokenList) -> bool: + return any( + child.ttype == sqlparse_tokens.Keyword.DML and child.normalized == "SELECT" + for child in token.flatten() + ) + + +def _collect_cte_names(statement: TokenList) -> set[str]: + tokens = _meaningful_tokens(statement) + if not tokens or tokens[0].normalized != "WITH": + return set() + + names = set() + for token in tokens[1:]: + if token.ttype == sqlparse_tokens.Keyword.DML and token.normalized == "SELECT": + break + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + name = _identifier_name(identifier) + if name: + names.add(name) + elif isinstance(token, Identifier): + name = _identifier_name(token) + if name: + names.add(name) + + return names + + +def _collect_table_references(token: TokenList) -> set[str]: + table_names = set() + tokens = _meaningful_tokens(token) + expect_table = False + + for current in tokens: + normalized = current.normalized + + if isinstance(current, Parenthesis): + if _contains_select(current): + table_names.update(_collect_table_references(current)) + continue + + if current.ttype == sqlparse_tokens.Keyword and ( + normalized == "FROM" or normalized == "JOIN" or normalized.endswith(" JOIN") + ): + expect_table = True + continue + + if expect_table: + if isinstance(current, IdentifierList): + for identifier in current.get_identifiers(): + if any( + isinstance(child, Parenthesis) and _contains_select(child) + for child in identifier.tokens + ): + for child in identifier.tokens: + if isinstance(child, Parenthesis): + table_names.update(_collect_table_references(child)) + else: + name = _identifier_name(identifier) + if name: + table_names.add(name) + elif isinstance(current, Identifier): + if any( + isinstance(child, Parenthesis) and _contains_select(child) + for child in current.tokens + ): + for child in current.tokens: + if isinstance(child, Parenthesis): + table_names.update(_collect_table_references(child)) + else: + name = _identifier_name(current) + if name: + table_names.add(name) + expect_table = False + + return table_names + + +def _table_grounding_error( + sql: str | None, schema_contracts: list[dict] | None +) -> str | None: + if not sql or not schema_contracts: + return None + + allowed_tables = { + contract.get("table_name") + for contract in schema_contracts + if contract.get("table_name") + } + if not allowed_tables: + return None + + statements = sqlparse.parse(sql) + cte_names = set() + referenced_tables = set() + for statement in statements: + cte_names.update(_collect_cte_names(statement)) + referenced_tables.update(_collect_table_references(statement)) + + ungrounded_tables = referenced_tables - allowed_tables - cte_names + if ungrounded_tables: + return "Generated SQL references table identifiers outside the retrieved deployed schema." + + return None + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -80,6 +200,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + schema_contracts: list[dict] | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -94,6 +215,21 @@ async def run( cleaned_generation_result ) + grounding_error = _table_grounding_error( + cleaned_generation_result, schema_contracts + ) + if grounding_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_GROUNDING", + "error": grounding_error, + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 49828a7c21..080e5896de 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -359,6 +359,16 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] + schema_contracts = [ + { + "table_name": document.get("table_name"), + "column_names": document.get("manifest_column_names") + or document.get("column_names") + or [], + } + for document in documents + if document.get("table_name") + ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -484,6 +494,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -503,6 +514,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -571,6 +583,7 @@ async def ask( allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index a05f2b0c5b..203bde03a2 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -167,6 +167,16 @@ async def ask_feedback( has_json_field = _retrieval_result.get("has_json_field", False) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + schema_contracts = [ + { + "table_name": document.get("table_name"), + "column_names": document.get("manifest_column_names") + or document.get("column_names") + or [], + } + for document in documents + if document.get("table_name") + ] sql_samples = sql_samples_task["formatted_output"].get("documents", []) instructions = instructions_task["formatted_output"].get( "documents", [] @@ -193,6 +203,7 @@ async def ask_feedback( has_json_field=has_json_field, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -258,6 +269,7 @@ async def ask_feedback( mdl_hash=ask_feedback_request.mdl_hash, sql_functions=sql_functions, sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index ed456adec5..81e1d239f9 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -118,6 +118,16 @@ async def correct( .get("retrieval_results", []) ) table_ddls = [document.get("table_ddl") for document in documents] + schema_contracts = [ + { + "table_name": document.get("table_name"), + "column_names": document.get("manifest_column_names") + or document.get("column_names") + or [], + } + for document in documents + if document.get("table_name") + ] res = await self._pipelines["sql_correction"].run( contexts=table_ddls, @@ -127,6 +137,7 @@ async def correct( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, ) post_process = res["post_process"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index 855d115811..8bbf502899 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -53,6 +53,7 @@ async def execute_sql( class CapturingEngine: def __init__(self): self.sql = None + self.executed = False async def execute_sql( self, @@ -63,6 +64,7 @@ async def execute_sql( dry_run=True, ): self.sql = sql + self.executed = True return True, None, {"correlation_id": "valid-correlation"} @@ -83,6 +85,60 @@ async def test_sql_post_processor_converts_select_top_to_wren_limit(): assert result["invalid_generation_result"] == {} +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_tables_outside_schema_contract(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT * FROM unsupported_model"}'], + project_id="project-id", + schema_contracts=[{"table_name": "supported_model", "column_names": []}], + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT * FROM unsupported_model", + "original_sql": "SELECT * FROM unsupported_model", + "type": "SCHEMA_GROUNDING", + "error": "Generated SQL references table identifiers outside the retrieved deployed schema.", + "correlation_id": "", + } + + +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_joined_tables_outside_schema_contract(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT * FROM supported_model JOIN unsupported_model ON 1 = 1"}'], + project_id="project-id", + schema_contracts=[{"table_name": "supported_model", "column_names": []}], + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SCHEMA_GROUNDING" + + +@pytest.mark.asyncio +async def test_sql_post_processor_allows_tables_inside_schema_contract(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT * FROM supported_model"}'], + project_id="project-id", + schema_contracts=[{"table_name": "supported_model", "column_names": []}], + ) + + assert engine.executed is True + assert result["valid_generation_result"] == { + "sql": "SELECT * FROM supported_model", + "correlation_id": "valid-correlation", + } + assert result["invalid_generation_result"] == {} + + @pytest.mark.asyncio async def test_sql_post_processor_returns_generated_sql_when_dry_plan_times_out(): result = await SQLGenPostProcessor(TimeoutEngine()).run( From 4b6efe525b0c7f05ad108eebd506ba952fca7472 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 14:39:31 +0530 Subject: [PATCH 0797/1087] Ground SQL generation with retrieved schema contract --- .../generation/followup_sql_generation.py | 8 +++ .../pipelines/generation/sql_correction.py | 8 +++ .../pipelines/generation/sql_generation.py | 8 +++ .../pipelines/generation/sql_regeneration.py | 8 +++ .../src/pipelines/generation/utils/sql.py | 30 ++++++++ .../retrieval/db_schema_retrieval.py | 22 +++++- wren-ai-service/src/web/v1/services/ask.py | 18 +++-- .../src/web/v1/services/ask_feedback.py | 7 +- .../generation/test_sql_prompt_grounding.py | 62 +++++++++++++++++ .../retrieval/test_db_schema_retrieval.py | 69 +++++++++++++++++++ 10 files changed, 232 insertions(+), 8 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 922ffce6ed..c223d1b842 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + build_executable_schema_contract, construct_ask_history_messages, construct_instructions, get_calculated_field_instructions, @@ -34,6 +35,11 @@ Given the user's current follow-up question and the current retrieved DATABASE SCHEMA, generate one SQL query to best answer the user's question. +{% if executable_schema_contract %} +{{ executable_schema_contract }} + +{% endif %} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -98,10 +104,12 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_contracts: list[dict] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + executable_schema_contract=build_executable_schema_contract(schema_contracts), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 485686af4d..a8938f949a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + build_executable_schema_contract, construct_instructions, get_text_to_sql_rules, ) @@ -56,6 +57,11 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) sql_correction_user_prompt_template = """ +{% if executable_schema_contract %} +{{ executable_schema_contract }} + +{% endif %} + {% if documents %} ### DATABASE SCHEMA ### {% for document in documents %} @@ -100,9 +106,11 @@ def prompt( query: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, + schema_contracts: list[dict] | None = None, ) -> dict: _prompt = prompt_builder.run( documents=documents, + executable_schema_contract=build_executable_schema_contract(schema_contracts), invalid_generation_result=invalid_generation_result, query=query or "", instructions=construct_instructions( diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index a99f3f571c..d1695760f1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + build_executable_schema_contract, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -28,6 +29,11 @@ sql_generation_user_prompt_template = """ +{% if executable_schema_contract %} +{{ executable_schema_contract }} + +{% endif %} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -94,10 +100,12 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_contracts: list[dict] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + executable_schema_contract=build_executable_schema_contract(schema_contracts), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index ca9d9bc287..4b28633d15 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -14,6 +14,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + build_executable_schema_contract, construct_instructions, get_calculated_field_instructions, get_json_field_instructions, @@ -53,6 +54,11 @@ def get_sql_regeneration_system_prompt( sql_regeneration_user_prompt_template = """ +{% if executable_schema_contract %} +{{ executable_schema_contract }} + +{% endif %} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -119,11 +125,13 @@ def prompt( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + schema_contracts: list[dict] | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, sql=sql, documents=documents, + executable_schema_contract=build_executable_schema_contract(schema_contracts), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 09bb6131f3..a6a41e4ef1 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -183,6 +183,36 @@ def _table_grounding_error( return None +def build_executable_schema_contract(schema_contracts: list[dict] | None) -> str: + if not schema_contracts: + return "" + + sections = [ + "### EXECUTABLE WREN IDENTIFIER CATALOG ###", + "Copy executable table and column identifiers only from this catalog or the matching DATABASE SCHEMA DDL.", + "Use descriptions, aliases, source names, and user wording only to understand meaning.", + ] + + for contract in schema_contracts: + table_name = contract.get("table_name") + if not table_name: + continue + + sections.append(f"TABLE: {table_name}") + column_names = [ + column_name + for column_name in contract.get("column_names", []) + if column_name + ] + if column_names: + sections.append("COLUMNS:") + sections.extend(f"- {column_name}" for column_name in column_names) + else: + sections.append("COLUMNS: declared in the matching DATABASE SCHEMA DDL") + + return "\n".join(sections) + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index fb8b4d9eae..0ace36708c 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -118,6 +118,14 @@ def _build_metric_ddl(content: dict) -> str: ) +def _content_column_names(content: dict) -> list[str]: + return [ + column.get("name", "") + for column in content.get("columns", []) + if column.get("name") + ] + + def _schema_column_names(content: dict) -> list[str]: return [ column["name"] @@ -195,7 +203,7 @@ def _build_table_context_ddl( def _build_view_ddl(content: dict) -> str: columns = content.get("columns", []) - column_names = [column.get("name", "") for column in columns if column.get("name")] + column_names = _content_column_names(content) columns_ddl = [ f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" for column in columns @@ -461,18 +469,24 @@ def check_using_db_schemas_without_pruning( content = ast.literal_eval(document.content) if content["type"] == "METRIC": + column_names = _content_column_names(content) retrieval_results.append( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "column_names": column_names, + "manifest_column_names": column_names, } ) has_metric = True elif content["type"] == "VIEW": + column_names = _content_column_names(content) retrieval_results.append( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "column_names": column_names, + "manifest_column_names": column_names, } ) @@ -581,18 +595,24 @@ def construct_retrieval_results( content = ast.literal_eval(document.content) if content["type"] == "METRIC": + column_names = _content_column_names(content) retrieval_results.append( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "column_names": column_names, + "manifest_column_names": column_names, } ) has_metric = True elif content["type"] == "VIEW": + column_names = _content_column_names(content) retrieval_results.append( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "column_names": column_names, + "manifest_column_names": column_names, } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 080e5896de..435b0dbffa 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -538,6 +538,9 @@ async def ask( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + is_schema_grounding_error = ( + failed_dry_run_result.get("type") == "SCHEMA_GROUNDING" + ) current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( @@ -551,7 +554,8 @@ async def ask( is_followup=True if histories else False, ) - if allow_sql_diagnosis: + sql_diagnosis_reasoning = None + if allow_sql_diagnosis and not is_schema_grounding_error: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -565,6 +569,10 @@ async def ask( "post_process" ].get("reasoning") + correction_error_message = error_message + if sql_diagnosis_reasoning: + correction_error_message = sql_diagnosis_reasoning + sql_correction_results = await self._pipelines[ "sql_correction" ].run( @@ -572,10 +580,10 @@ async def ask( query=user_query, instructions=instructions, invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "sql": ( + "" if is_schema_grounding_error else original_sql + ), + "error": correction_error_message, }, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 203bde03a2..2911b35ed5 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -224,6 +224,9 @@ async def ask_feedback( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + is_schema_grounding_error = ( + failed_dry_run_result.get("type") == "SCHEMA_GROUNDING" + ) sql_diagnosis_reasoning = None self._ask_feedback_results[ @@ -233,7 +236,7 @@ async def ask_feedback( trace_id=trace_id, ) - if allow_sql_diagnosis: + if allow_sql_diagnosis and not is_schema_grounding_error: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -262,7 +265,7 @@ async def ask_feedback( instructions=instructions, invalid_generation_result={ "original_sql": original_sql, - "sql": invalid_sql, + "sql": "" if is_schema_grounding_error else invalid_sql, "error": correction_error_message, }, project_id=ask_feedback_request.project_id, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index a42e1229d1..d474b0b806 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -8,6 +8,23 @@ prompt as build_sql_generation_prompt, sql_generation_user_prompt_template, ) +from src.pipelines.generation.utils.sql import build_executable_schema_contract + + +def test_build_executable_schema_contract_lists_retrieved_identifiers(): + contract = build_executable_schema_contract( + [ + { + "table_name": "retrieved_model", + "column_names": ["grouping_attribute", "numeric_measure"], + } + ] + ) + + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in contract + assert "TABLE: retrieved_model" in contract + assert "- grouping_attribute" in contract + assert "- numeric_measure" in contract def test_sql_generation_prompt_omits_sample_sql_body(): @@ -29,6 +46,27 @@ def test_sql_generation_prompt_omits_sample_sql_body(): assert "SELECT 1" not in built_prompt +def test_sql_generation_prompt_includes_executable_schema_contract(): + result = build_sql_generation_prompt( + query="summarize the records", + documents=[], + prompt_builder=PromptBuilder(template=sql_generation_user_prompt_template), + schema_contracts=[ + { + "table_name": "retrieved_model", + "column_names": ["grouping_attribute", "numeric_measure"], + } + ], + ) + + built_prompt = result["prompt"] + + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt + assert "TABLE: retrieved_model" in built_prompt + assert "- grouping_attribute" in built_prompt + assert "- numeric_measure" in built_prompt + + def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): result = build_sql_correction_prompt( documents=[], @@ -45,3 +83,27 @@ def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): assert "User's Question: summarize the records" in built_prompt assert "Failed SQL: SELECT 1" in built_prompt assert "DIAGNOSTIC CONTEXT" in built_prompt + + +def test_sql_correction_prompt_includes_executable_schema_contract(): + result = build_sql_correction_prompt( + documents=[], + invalid_generation_result={ + "sql": "", + "error": "Generated SQL references identifiers outside retrieved schema.", + }, + query="summarize the records", + prompt_builder=PromptBuilder(template=sql_correction_user_prompt_template), + schema_contracts=[ + { + "table_name": "retrieved_model", + "column_names": ["grouping_attribute", "numeric_measure"], + } + ], + ) + + built_prompt = result["prompt"] + + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt + assert "TABLE: retrieved_model" in built_prompt + assert "Failed SQL:" in built_prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 4f05cced7a..26b0a180e5 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -760,9 +760,78 @@ def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): "modeled_dataset", "semantic_metric", ] + assert result["retrieval_results"][1]["column_names"] == ["metric_value"] + assert result["retrieval_results"][1]["manifest_column_names"] == ["metric_value"] assert result["has_metric"] is True +def test_construct_retrieval_results_preserves_retrieved_view_columns_when_pruning(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_attribute"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "VIEW", + "comment": "", + "name": "semantic_view", + "columns": [ + { + "name": "view_attribute", + "data_type": "VARCHAR", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "semantic_view"}, + ) + ], + ) + + assert result["retrieval_results"][1]["table_name"] == "semantic_view" + assert result["retrieval_results"][1]["column_names"] == ["view_attribute"] + assert result["retrieval_results"][1]["manifest_column_names"] == [ + "view_attribute" + ] + + def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_columns(): result = construct_retrieval_results( check_using_db_schemas_without_pruning={}, From c9734229ae829d5861b4513359918a2da439a349 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 15:03:41 +0530 Subject: [PATCH 0798/1087] Fix grounded SQL retries and ask task ids --- .../pipelines/generation/sql_generation.py | 6 +- wren-ai-service/src/web/v1/services/ask.py | 10 +++ .../repositories/askingTaskRepository.ts | 87 +----------------- .../server/repositories/baseRepository.ts | 90 ++++++++++++++++--- 4 files changed, 96 insertions(+), 97 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index d1695760f1..3da086eca3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -76,13 +76,17 @@ ### QUESTION ### User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. +If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. {% if sql_generation_reasoning %} ### REASONING PLAN ### {{ sql_generation_reasoning }} {% endif %} -Let's think step by step. +Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 435b0dbffa..c8a0a66f31 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -531,6 +531,7 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: + schema_grounding_correction_attempted = False while current_sql_correction_retries < max_sql_correction_retries: if failed_dry_run_result["type"] == "TIME_OUT": break @@ -541,6 +542,15 @@ async def ask( is_schema_grounding_error = ( failed_dry_run_result.get("type") == "SCHEMA_GROUNDING" ) + if ( + is_schema_grounding_error + and schema_grounding_correction_attempted + ): + break + schema_grounding_correction_attempted = ( + schema_grounding_correction_attempted + or is_schema_grounding_error + ) current_sql_correction_retries += 1 self._ask_results[query_id] = AskResultResponse( diff --git a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts index 6d7e495018..ff98b7502a 100644 --- a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts +++ b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts @@ -39,7 +39,6 @@ export class AskingTaskRepository implements IAskingTaskRepository { private readonly jsonbColumns = ['detail']; - private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'asking_task' }); @@ -53,20 +52,14 @@ export class AskingTaskRepository data: Partial, queryOptions?: IQueryOptions, ): Promise { - return super.createOne( - await this.withMssqlId(this.withTimestamps(data), queryOptions), - queryOptions, - ); + return super.createOne(this.withTimestamps(data), queryOptions); } public override async createMany( data: Partial[], queryOptions?: IQueryOptions, ): Promise { - return super.createMany( - await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), - queryOptions, - ); + return super.createMany(data.map(this.withTimestamps), queryOptions); } public override async updateOne( @@ -125,80 +118,4 @@ export class AskingTaskRepository updatedAt: data.updatedAt ?? now, }; }; - - private isMssql = () => - String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; - - private hasIdentityId = async (): Promise => { - if (!this.isMssql()) { - return true; - } - - if (!this.hasIdentityIdPromise) { - this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') - .select('COLUMN_NAME') - .where({ - TABLE_SCHEMA: 'dbo', - TABLE_NAME: this.tableName, - COLUMN_NAME: 'id', - }) - .whereRaw( - "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", - ) - .first() - .then(Boolean); - } - - return this.hasIdentityIdPromise; - }; - - private withMssqlId = async ( - data: Partial, - queryOptions?: IQueryOptions, - ): Promise> => { - if ( - (data.id !== undefined && data.id !== null) || - !this.isMssql() || - (await this.hasIdentityId()) - ) { - return data; - } - - const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [row] = await executer(this.tableName).max<{ maxId?: number }>({ - maxId: 'id', - }); - return { - ...data, - id: Number(row?.maxId || 0) + 1, - }; - }; - - private withMssqlIds = async ( - data: Partial[], - queryOptions?: IQueryOptions, - ): Promise[]> => { - if ( - data.every((item) => item.id !== undefined && item.id !== null) || - !this.isMssql() || - (await this.hasIdentityId()) - ) { - return data; - } - - const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [row] = await executer(this.tableName).max<{ maxId?: number }>({ - maxId: 'id', - }); - let nextId = Number(row?.maxId || 0) + 1; - return data.map((item) => { - if (item.id !== undefined && item.id !== null) { - return item; - } - return { - ...item, - id: nextId++, - }; - }); - }; } diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index ab3be09621..da7c54b1d2 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -55,6 +55,7 @@ export const coerceBoolean = (value: unknown): boolean => { }; export class BaseRepository implements IBasicRepository { + private static manualIdInsertLocks = new Map>(); protected knex: Knex; protected tableName: string; private hasIdColumnCache: boolean | null = null; @@ -122,20 +123,33 @@ export class BaseRepository implements IBasicRepository { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; try { const insertValue = await this.prepareInsertData(data, executer); - const [result] = await executer(this.tableName) - .insert(this.normalizeMssqlBindings(insertValue)) - .returning('*'); + const [result] = await this.insertOne(executer, insertValue); return this.transformFromDBData(result); } catch (error) { if (!this.shouldRetryManualId(error, data, executer)) { throw error; } - const insertValue = await this.prepareInsertData(data, executer, true); - const [result] = await executer(this.tableName) - .insert(this.normalizeMssqlBindings(insertValue)) - .returning('*'); - return this.transformFromDBData(result); + return await this.withManualIdInsertLock(executer, async () => { + let lastError: unknown = error; + for (let attempt = 0; attempt < 3; attempt++) { + try { + const insertValue = await this.prepareInsertData( + data, + executer, + true, + ); + const [result] = await this.insertOne(executer, insertValue); + return this.transformFromDBData(result); + } catch (retryError) { + if (!this.shouldRetryManualId(retryError, data, executer)) { + throw retryError; + } + lastError = retryError; + } + } + throw lastError; + }); } } @@ -158,8 +172,25 @@ export class BaseRepository implements IBasicRepository { throw error; } - preparedData = await this.prepareInsertManyData(data, executer, true); - return await this.insertMany(executer, preparedData); + return await this.withManualIdInsertLock(executer, async () => { + let lastError: unknown = error; + for (let attempt = 0; attempt < 3; attempt++) { + try { + preparedData = await this.prepareInsertManyData( + data, + executer, + true, + ); + return await this.insertMany(executer, preparedData); + } catch (retryError) { + if (!this.shouldRetryManualId(retryError, data, executer)) { + throw retryError; + } + lastError = retryError; + } + } + throw lastError; + }); } } @@ -309,6 +340,33 @@ export class BaseRepository implements IBasicRepository { return this.toNextIdValue(row?.maxId); } + private async withManualIdInsertLock( + executer: Knex | Knex.Transaction, + task: () => Promise, + ): Promise { + const lockKey = `${executer.client.config.client}:${this.tableName}`; + const previousLock = + BaseRepository.manualIdInsertLocks.get(lockKey) ?? Promise.resolve(); + let releaseLock: () => void = () => undefined; + const currentLock = new Promise((resolve) => { + releaseLock = resolve; + }); + const chainedLock = previousLock.catch(() => undefined).then( + () => currentLock, + ); + BaseRepository.manualIdInsertLocks.set(lockKey, chainedLock); + + await previousLock.catch(() => undefined); + try { + return await task(); + } finally { + releaseLock(); + if (BaseRepository.manualIdInsertLocks.get(lockKey) === chainedLock) { + BaseRepository.manualIdInsertLocks.delete(lockKey); + } + } + } + private async hasIdentityId(executer: Knex | Knex.Transaction) { if (!this.isMssql(executer)) { return true; @@ -412,6 +470,12 @@ export class BaseRepository implements IBasicRepository { return result.map((data) => this.transformFromDBData(data)); } + private async insertOne(executer: Knex | Knex.Transaction, preparedData: any) { + return await executer(this.tableName) + .insert(this.normalizeMssqlBindings(preparedData)) + .returning('*'); + } + private normalizeMssqlBindings(value: any): any { if (!this.isMssql(this.knex)) { return value; @@ -481,6 +545,10 @@ export class BaseRepository implements IBasicRepository { ? error : ''; - return message.includes("Cannot insert the value NULL into column 'id'"); + return ( + message.includes("Cannot insert the value NULL into column 'id'") || + message.includes('Violation of PRIMARY KEY constraint') || + message.includes('Cannot insert duplicate key') + ); } } From c8b971fef8df6903fb432e5c538f84b5784e9693 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 15:19:31 +0530 Subject: [PATCH 0799/1087] Fix ask metadata deployment scoping --- .../generation/intent_classification.py | 36 +++++++--- .../retrieval/db_schema_retrieval.py | 14 +++- wren-ai-service/src/web/v1/services/ask.py | 5 +- .../generation/test_intent_classification.py | 70 +++++++++++++++++++ .../retrieval/test_db_schema_retrieval.py | 7 +- .../tests/pytest/services/test_ask.py | 7 ++ .../apollo/server/adaptors/wrenAIAdaptor.ts | 3 +- .../apollo/server/services/askingService.ts | 3 +- .../apollo/server/services/deployService.ts | 57 +++++++++++++-- .../services/tests/askingService.test.ts | 8 +-- .../services/tests/deployService.test.ts | 55 +++++++++++++-- 11 files changed, 230 insertions(+), 35 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 7891c0a4ae..cc9a823fe5 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -13,7 +13,11 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider -from src.pipelines.common import build_table_ddl, clean_up_new_lines +from src.pipelines.common import ( + build_project_deploy_filter, + build_table_ddl, + clean_up_new_lines, +) from src.pipelines.generation.utils.sql import construct_instructions from src.utils import trace_cost from src.web.v1.services import Configuration @@ -172,7 +176,7 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> d @observe(capture_input=False) async def table_retrieval( - embedding: dict, project_id: str, table_retriever: Any + embedding: dict, project_id: str, table_retriever: Any, mdl_hash: str = "" ) -> dict: filters = { "operator": "AND", @@ -181,10 +185,12 @@ async def table_retrieval( ], } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + project_deploy_filter = build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ) + if project_deploy_filter: + filters["conditions"] += project_deploy_filter["conditions"] return await table_retriever.run( query_embedding=embedding.get("embedding"), @@ -194,7 +200,11 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, embedding: dict, project_id: str, dbschema_retriever: Any + table_retrieval: dict, + embedding: dict, + project_id: str, + dbschema_retriever: Any, + mdl_hash: str = "", ) -> list[Document]: tables = table_retrieval.get("documents", []) table_names = [] @@ -217,10 +227,12 @@ async def dbschema_retrieval( ], } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + project_deploy_filter = build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ) + if project_deploy_filter: + filters["conditions"] += project_deploy_filter["conditions"] results = await dbschema_retriever.run( query_embedding=embedding.get("embedding"), filters=filters @@ -376,6 +388,7 @@ async def run( self, query: str, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, histories: Optional[list[AskHistory]] = None, sql_samples: Optional[list[dict]] = None, instructions: Optional[list[dict]] = None, @@ -387,6 +400,7 @@ async def run( inputs={ "query": query, "project_id": project_id or "", + "mdl_hash": mdl_hash or "", "histories": histories or [], "sql_samples": sql_samples or [], "instructions": instructions or [], diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 0ace36708c..183be97b7e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -242,7 +242,15 @@ async def active_mdl_hash( dbschema_store.count_documents(filters=filters), ) - return mdl_hash if table_description_count or dbschema_count else "" + if not table_description_count and not dbschema_count: + logger.warning( + "Project ID: %s, MDL hash %s has no indexed schema documents; " + "keeping hash scope to avoid stale project metadata fallback.", + project_id, + mdl_hash, + ) + + return mdl_hash @observe(capture_input=False, capture_output=False) @@ -262,7 +270,7 @@ async def table_retrieval( active_mdl_hash: Optional[str] = None, mdl_hash: str = "", ) -> dict: - effective_mdl_hash = active_mdl_hash if active_mdl_hash is not None else mdl_hash + effective_mdl_hash = active_mdl_hash or mdl_hash filters = { "operator": "AND", "conditions": [ @@ -305,7 +313,7 @@ async def dbschema_retrieval( mdl_hash: str = "", embedding: Optional[dict] = None, ) -> list[Document]: - effective_mdl_hash = active_mdl_hash if active_mdl_hash is not None else mdl_hash + effective_mdl_hash = active_mdl_hash or mdl_hash def _base_filters() -> dict: project_deploy_filter = build_project_deploy_filter( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index c8a0a66f31..cfaf212c91 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -27,8 +27,8 @@ class AskRequest(BaseRequest): histories: Optional[list[AskHistory]] = Field(default_factory=list) ignore_sql_generation_reasoning: bool = True enable_column_pruning: bool = False - use_dry_plan: bool = False - allow_dry_plan_fallback: bool = True + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False custom_instruction: Optional[str] = None @@ -244,6 +244,7 @@ async def ask( sql_samples=sql_samples, instructions=instructions, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, configuration=ask_request.configurations, ) ).get("post_process", {}) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py b/wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py new file mode 100644 index 0000000000..0b14f0e1b0 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py @@ -0,0 +1,70 @@ +import pytest +from haystack import Document + +from src.pipelines.generation.intent_classification import ( + dbschema_retrieval, + table_retrieval, +) + + +class CapturingRetriever: + def __init__(self, documents=None): + self.documents = documents or [] + self.calls = [] + + async def run(self, **kwargs): + self.calls.append(kwargs) + return {"documents": self.documents} + + +@pytest.mark.asyncio +async def test_intent_table_retrieval_scopes_to_deployed_mdl_hash(): + retriever = CapturingRetriever() + + await table_retrieval( + embedding={"embedding": [0.1, 0.2]}, + project_id="project-1", + mdl_hash="deploy-1", + table_retriever=retriever, + ) + + assert retriever.calls[0]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + } + + +@pytest.mark.asyncio +async def test_intent_dbschema_retrieval_scopes_to_deployed_mdl_hash(): + retriever = CapturingRetriever() + + await dbschema_retrieval( + table_retrieval={ + "documents": [ + Document(content=str({"name": "orders"}), meta={"name": "orders"}) + ] + }, + embedding={"embedding": [0.1, 0.2]}, + project_id="project-1", + mdl_hash="deploy-1", + dbschema_retriever=retriever, + ) + + assert retriever.calls[0]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": "orders"} + ], + }, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + } diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 26b0a180e5..f14fae9103 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -115,7 +115,7 @@ async def test_active_mdl_hash_keeps_hash_when_deploy_documents_are_indexed(): @pytest.mark.asyncio -async def test_active_mdl_hash_uses_project_scope_when_deploy_documents_are_absent(): +async def test_active_mdl_hash_keeps_hash_when_deploy_documents_are_absent(): table_store = StoreCounter(count=0) schema_store = StoreCounter(count=0) @@ -126,7 +126,7 @@ async def test_active_mdl_hash_uses_project_scope_when_deploy_documents_are_abse dbschema_store=schema_store, ) - assert result == "" + assert result == "deploy-1" def test_view_schema_context_uses_declared_view_columns_not_view_definition(): @@ -223,7 +223,7 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_table_retrieval_uses_project_scope_when_active_hash_is_absent(): +async def test_table_retrieval_falls_back_to_request_hash_when_active_hash_is_absent(): class Retriever: def __init__(self): self.calls = [] @@ -256,6 +256,7 @@ async def run(self, query_embedding, filters): "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, ], }, } diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 5187681d58..15d346146d 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -108,6 +108,13 @@ def test_ask_request_skips_sql_generation_reasoning_by_default(): assert ask_request.ignore_sql_generation_reasoning is True +def test_ask_request_uses_strict_dry_plan_by_default(): + ask_request = AskRequest(query="question", mdl_hash="deploy") + + assert ask_request.use_dry_plan is True + assert ask_request.allow_dry_plan_fallback is False + + @pytest.mark.asyncio async def test_ask_with_successful_query( indexing_service: SemanticsPreparationService, diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index a74968c0db..567a184c95 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -51,6 +51,7 @@ const getAIServiceError = (error: any) => { export interface IWrenAIAdaptor { deploy(deployData: DeployData): Promise; + getDeployStatus(deployId: string): Promise; delete(projectId: number): Promise; /** @@ -901,7 +902,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { return deploySuccess; } - private async getDeployStatus(deployId: string): Promise { + public async getDeployStatus(deployId: string): Promise { try { const res = await axios.get( `${this.wrenAIBaseEndpoint}/v1/semantics-preparations/${deployId}/status`, diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 4993320ec0..c2939b911c 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1274,8 +1274,7 @@ export class AskingService implements IAskingService { private async getDeployId(projectId?: number) { const id = projectId ?? (await this.projectService.getCurrentProject()).id; - const lastDeploy = await this.deployService.getLastDeployment(id); - return lastDeploy.hash; + return this.deployService.ensureDeploymentPrepared(id); } private async getProjectForThreadResponse(threadResponse: ThreadResponse) { diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 746e5fa12e..085dad9e0e 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -1,4 +1,7 @@ -import { WrenAIDeployStatusEnum } from '@server/models/adaptor'; +import { + WrenAIDeployStatusEnum, + WrenAISystemStatus, +} from '@server/models/adaptor'; import { IWrenAIAdaptor } from '../adaptors/wrenAIAdaptor'; import { Deploy, @@ -35,6 +38,7 @@ export interface IDeployService { force?: boolean, ): Promise; getLastDeployment(projectId: number): Promise; + ensureDeploymentPrepared(projectId: number): Promise; getInProgressDeployment(projectId: number): Promise; createMDLHash(manifest: Manifest, projectId: number): string; isSameDeployment( @@ -74,6 +78,38 @@ export class DeployService implements IDeployService { return lastDeploy; } + public async ensureDeploymentPrepared(projectId: number): Promise { + const lastDeploy = + await this.deployLogRepository.findLastProjectDeployLog(projectId); + if (!lastDeploy) { + throw new Error(`No deployment found for project ${projectId}`); + } + + try { + const status = await this.wrenAIAdaptor.getDeployStatus(lastDeploy.hash); + if (status === WrenAISystemStatus.FINISHED) { + return lastDeploy.hash; + } + logger.warn( + `Deployment ${lastDeploy.hash} is not ready in AI service: ${status}`, + ); + } catch (err: any) { + logger.warn( + `Deployment ${lastDeploy.hash} is not available in AI service: ${err.message}`, + ); + } + + const result = await this.deploy(lastDeploy.manifest, projectId); + if (result.status !== DeployStatusEnum.SUCCESS) { + throw new Error( + result.error || + `Failed to prepare deployment ${lastDeploy.hash} for project ${projectId}`, + ); + } + + return lastDeploy.hash; + } + public async getInProgressDeployment(projectId) { const inProgressDeploy = await this.deployLogRepository.findInProgressProjectDeployLog( projectId, @@ -108,12 +144,23 @@ export class DeployService implements IDeployService { const lastDeploy = await this.deployLogRepository.findLastProjectDeployLog(projectId); if (lastDeploy && lastDeploy.hash === hash) { - logger.log(`Model has been deployed, hash: ${hash}`); + logger.log(`Model has been deployed, refreshing AI index, hash: ${hash}`); + deploy = lastDeploy; + const { status: aiStatus, error: aiError } = + await this.wrenAIAdaptor.deploy({ + manifest, + hash, + projectId, + }); + const status = + aiStatus === WrenAIDeployStatusEnum.SUCCESS + ? DeployStatusEnum.SUCCESS + : DeployStatusEnum.FAILED; await this.deployLogRepository.updateOne(lastDeploy.id, { - status: DeployStatusEnum.SUCCESS, - error: null, + status, + error: aiError, }); - return { status: DeployStatusEnum.SUCCESS }; + return { status, error: aiError }; } } const previousInProgressDeploy = diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index d33c9e166d..f8121a24ef 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -141,9 +141,9 @@ describe('AskingService', () => { }), }; service.deployService = { - getLastDeployment: jest.fn().mockResolvedValue({ - hash: 'latest-deploy-hash', - }), + ensureDeploymentPrepared: jest + .fn() + .mockResolvedValue('latest-deploy-hash'), }; service.threadRepository = { createOne: jest.fn().mockResolvedValue({ id: 7, projectId: 1 }), @@ -171,7 +171,7 @@ describe('AskingService', () => { sql: 'SELECT stale_recommendation_sql', }); - expect(service.deployService.getLastDeployment).toHaveBeenCalledWith(1); + expect(service.deployService.ensureDeploymentPrepared).toHaveBeenCalledWith(1); expect(service.askingTaskTracker.createAskingTask).toHaveBeenCalledWith({ query: trackedAskingResult.question, histories: null, diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index ce60461bf7..739ef8e963 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -11,6 +11,7 @@ describe('DeployService', () => { beforeEach(() => { mockTelemetry = { sendEvent: jest.fn() }; mockWrenAIAdaptor = { deploy: jest.fn() }; + mockWrenAIAdaptor.getDeployStatus = jest.fn(); mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), findInProgressProjectDeployLog: jest.fn(), @@ -82,22 +83,68 @@ describe('DeployService', () => { }); }); - it('should skip deployment if an existing deployment with the same hash exists', async () => { + it('should refresh ai-service deployment if an existing deployment with the same hash exists', async () => { const manifest = { key: 'value' }; const projectId = 1; + const hash = deployService.createMDLHash(manifest, 1); mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ id: 123, - hash: deployService.createMDLHash(manifest, 1), + hash, }); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); const response = await deployService.deploy(manifest, projectId); expect(response.status).toEqual(DeployStatusEnum.SUCCESS); - expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash, + projectId, + }); expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { status: DeployStatusEnum.SUCCESS, - error: null, + error: undefined, + }); + }); + + it('should return the deployment hash when ai-service already has the exact deployment prepared', async () => { + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ + hash: 'deploy-hash', + manifest: { key: 'value' }, + }); + mockWrenAIAdaptor.getDeployStatus.mockResolvedValue('FINISHED'); + + const hash = await deployService.ensureDeploymentPrepared(1); + + expect(hash).toEqual('deploy-hash'); + expect(mockWrenAIAdaptor.getDeployStatus).toHaveBeenCalledWith('deploy-hash'); + expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); + }); + + it('should redeploy the saved manifest when ai-service no longer has the exact deployment prepared', async () => { + const manifest = { key: 'value' }; + mockDeployLogRepository.findLastProjectDeployLog + .mockResolvedValueOnce({ + id: 123, + hash: deployService.createMDLHash(manifest, 1), + manifest, + }) + .mockResolvedValueOnce({ + id: 123, + hash: deployService.createMDLHash(manifest, 1), + manifest, + }); + mockWrenAIAdaptor.getDeployStatus.mockRejectedValue(new Error('not found')); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + + const hash = await deployService.ensureDeploymentPrepared(1); + + expect(hash).toEqual(deployService.createMDLHash(manifest, 1)); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash, + projectId: 1, }); }); From 026996edf3c8f5eaa873ecd4c364461166866a67 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 15:33:34 +0530 Subject: [PATCH 0800/1087] Verify semantic index readiness for deployments --- .../src/pipelines/indexing/db_schema.py | 21 ++++++ .../pipelines/indexing/table_description.py | 21 ++++++ .../web/v1/routers/semantics_preparation.py | 2 +- .../web/v1/services/semantics_preparation.py | 56 +++++++++++++++- .../services/test_semantics_preparation.py | 65 +++++++++++++++++++ 5 files changed, 162 insertions(+), 3 deletions(-) create mode 100644 wren-ai-service/tests/pytest/services/test_semantics_preparation.py diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index b62d1de9fa..338cefd7c4 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -21,6 +21,7 @@ MDLValidator, clean_display_name, ) +from src.pipelines.common import build_project_deploy_filter from src.pipelines.indexing.utils import helper logger = logging.getLogger("wren-ai-service") @@ -430,6 +431,7 @@ def __init__( **kwargs, ) -> None: dbschema_store = document_store_provider.get_store() + self._store = dbschema_store self._components = { "cleaner": DocumentCleaner([dbschema_store]), @@ -479,3 +481,22 @@ async def clean(self, project_id: Optional[str] = None) -> None: cleaner=self._components["cleaner"], project_id=project_id, ) + + async def count_documents( + self, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ) -> int: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + if project_deploy_filter := build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ): + filters["conditions"] += project_deploy_filter["conditions"] + + return await self._store.count_documents(filters=filters) diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 65fdeee6f2..20dabf4c99 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -14,6 +14,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider +from src.pipelines.common import build_project_deploy_filter from src.pipelines.indexing import AsyncDocumentWriter, DocumentCleaner, MDLValidator logger = logging.getLogger("wren-ai-service") @@ -277,6 +278,7 @@ def __init__( table_description_store = document_store_provider.get_store( dataset_name="table_descriptions" ) + self._store = table_description_store self._components = { "cleaner": DocumentCleaner([table_description_store]), @@ -323,3 +325,22 @@ async def clean(self, project_id: Optional[str] = None) -> None: cleaner=self._components["cleaner"], project_id=project_id, ) + + async def count_documents( + self, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ) -> int: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + ], + } + if project_deploy_filter := build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ): + filters["conditions"] += project_deploy_filter["conditions"] + + return await self._store.count_documents(filters=filters) diff --git a/wren-ai-service/src/web/v1/routers/semantics_preparation.py b/wren-ai-service/src/web/v1/routers/semantics_preparation.py index 1a8ed15dc7..e6a4efab7e 100644 --- a/wren-ai-service/src/web/v1/routers/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/routers/semantics_preparation.py @@ -44,7 +44,7 @@ async def get_prepare_semantics_status( mdl_hash: str, service_container: ServiceContainer = Depends(get_service_container), ) -> SemanticsPreparationStatusResponse: - return service_container.semantics_preparation_service.get_prepare_semantics_status( + return await service_container.semantics_preparation_service.get_prepare_semantics_status( SemanticsPreparationStatusRequest(mdl_hash=mdl_hash) ) 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 a82de5d98f..a546f16457 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -54,6 +54,9 @@ def __init__( self._prepare_semantics_statuses: Dict[ str, SemanticsPreparationStatusResponse ] = TTLCache(maxsize=maxsize, ttl=ttl) + self._prepare_semantics_project_ids: Dict[str, str] = TTLCache( + maxsize=maxsize, ttl=ttl + ) @observe(name="Prepare Semantics") @trace_metadata @@ -107,11 +110,22 @@ async def prepare_semantics( await asyncio.gather(*tasks) + if not await self._has_indexed_schema_documents( + prepare_semantics_request.project_id, + prepare_semantics_request.mdl_hash, + ): + raise RuntimeError( + "No indexed schema documents were found for the prepared deployment" + ) + self._prepare_semantics_statuses[ prepare_semantics_request.mdl_hash ] = SemanticsPreparationStatusResponse( status="finished", ) + self._prepare_semantics_project_ids[ + prepare_semantics_request.mdl_hash + ] = prepare_semantics_request.project_id or "" except Exception as e: logger.exception(f"Failed to prepare semantics: {e}") @@ -130,7 +144,24 @@ async def prepare_semantics( return results - def get_prepare_semantics_status( + async def _has_indexed_schema_documents( + self, + project_id: str, + mdl_hash: str, + ) -> bool: + dbschema_count, table_description_count = await asyncio.gather( + self._pipelines["db_schema"].count_documents( + project_id=project_id, + mdl_hash=mdl_hash, + ), + self._pipelines["table_description"].count_documents( + project_id=project_id, + mdl_hash=mdl_hash, + ), + ) + return dbschema_count > 0 and table_description_count > 0 + + async def get_prepare_semantics_status( self, prepare_semantics_status_request: SemanticsPreparationStatusRequest ) -> SemanticsPreparationStatusResponse: if ( @@ -149,7 +180,28 @@ def get_prepare_semantics_status( ), ) - return result + if result.status != "finished": + return result + + project_id = self._prepare_semantics_project_ids.get( + prepare_semantics_status_request.mdl_hash + ) + if project_id is None: + return result + + if await self._has_indexed_schema_documents( + project_id, + prepare_semantics_status_request.mdl_hash, + ): + return result + + return SemanticsPreparationStatusResponse( + status="failed", + error=SemanticsPreparationStatusResponse.SemanticsPreparationError( + code="OTHERS", + message="Prepared schema documents are missing for this deployment", + ), + ) @observe(name="Delete Semantics Documents") @trace_metadata diff --git a/wren-ai-service/tests/pytest/services/test_semantics_preparation.py b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py new file mode 100644 index 0000000000..106d6b3e2b --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py @@ -0,0 +1,65 @@ +import pytest + +from src.web.v1.services.semantics_preparation import ( + SemanticsPreparationService, + SemanticsPreparationStatusRequest, + SemanticsPreparationStatusResponse, +) + + +class CountPipeline: + def __init__(self, count: int): + self.count = count + self.calls = [] + + async def count_documents(self, project_id=None, mdl_hash=None): + self.calls.append({"project_id": project_id, "mdl_hash": mdl_hash}) + return self.count + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_fails_when_exact_schema_documents_are_missing(): + db_schema = CountPipeline(0) + table_description = CountPipeline(1) + service = SemanticsPreparationService( + { + "db_schema": db_schema, + "table_description": table_description, + } + ) + service._prepare_semantics_statuses["deploy-1"] = SemanticsPreparationStatusResponse( + status="finished" + ) + service._prepare_semantics_project_ids["deploy-1"] = "project-1" + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1") + ) + + assert status.status == "failed" + assert status.error.message == "Prepared schema documents are missing for this deployment" + assert db_schema.calls == [{"project_id": "project-1", "mdl_hash": "deploy-1"}] + assert table_description.calls == [ + {"project_id": "project-1", "mdl_hash": "deploy-1"} + ] + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_stays_finished_when_exact_schema_documents_exist(): + service = SemanticsPreparationService( + { + "db_schema": CountPipeline(1), + "table_description": CountPipeline(1), + } + ) + service._prepare_semantics_statuses["deploy-1"] = SemanticsPreparationStatusResponse( + status="finished" + ) + service._prepare_semantics_project_ids["deploy-1"] = "project-1" + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1") + ) + + assert status.status == "finished" + assert status.error is None From 603cf75dd06cecd204844fc8d519810fc94302eb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 16:27:37 +0530 Subject: [PATCH 0801/1087] Reduce grounded SQL timeout failures --- wren-ai-service/src/pipelines/common.py | 31 ++++++++++++++ .../generation/followup_sql_generation.py | 5 +++ .../pipelines/generation/sql_correction.py | 5 +++ .../pipelines/generation/sql_generation.py | 5 +++ .../pipelines/generation/sql_regeneration.py | 6 +++ .../src/pipelines/generation/utils/sql.py | 20 +++++++++- wren-ai-service/src/web/v1/services/ask.py | 2 +- .../generation/test_sql_post_processor.py | 38 ++++++++++++++++++ .../generation/test_sql_prompt_grounding.py | 28 +++++++++++++ .../tests/pytest/services/test_ask.py | 6 +++ .../apollo/server/resolvers/modelResolver.ts | 40 +++++++++++++------ 11 files changed, 171 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index c40a940088..e7fe653e8d 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -1,3 +1,4 @@ +import asyncio import re from typing import Any, List, Optional, Tuple @@ -113,6 +114,36 @@ async def retrieve_metadata( project_id: str, retriever, mdl_hash: Optional[str] = None, +) -> dict[str, Any]: + cache_key = ( + id(retriever), + str(project_id), + str(mdl_hash), + ) + if project_id and mdl_hash: + if cache_key in _METADATA_CACHE: + return _METADATA_CACHE[cache_key] + + lock = _METADATA_CACHE_LOCKS.setdefault(cache_key, asyncio.Lock()) + async with lock: + if cache_key in _METADATA_CACHE: + return _METADATA_CACHE[cache_key] + + metadata = await _retrieve_metadata_uncached(project_id, retriever, mdl_hash) + _METADATA_CACHE[cache_key] = metadata + return metadata + + return await _retrieve_metadata_uncached(project_id, retriever, mdl_hash) + + +_METADATA_CACHE: dict[tuple[int, str, str], dict[str, Any]] = {} +_METADATA_CACHE_LOCKS: dict[tuple[int, str, str], asyncio.Lock] = {} + + +async def _retrieve_metadata_uncached( + project_id: str, + retriever, + mdl_hash: Optional[str] = None, ) -> dict[str, Any]: filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) 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 c223d1b842..aa0c872246 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -86,6 +86,11 @@ If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. +{% if executable_schema_contract %} +### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### +{{ executable_schema_contract }} +{% endif %} + Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index a8938f949a..fc3153ce39 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -88,6 +88,11 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) User's Question: {{ query }} {% endif %} +{% if executable_schema_contract %} +### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS CORRECTION ### +{{ executable_schema_contract }} +{% endif %} + ### FAILED SQL DIAGNOSTIC CONTEXT ### Failed SQL: {{ invalid_generation_result.sql }} Error Message: {{ invalid_generation_result.error }} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 3da086eca3..0519d19c44 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -81,6 +81,11 @@ If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. +{% if executable_schema_contract %} +### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### +{{ executable_schema_contract }} +{% endif %} + {% if sql_generation_reasoning %} ### REASONING PLAN ### {{ sql_generation_reasoning }} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 4b28633d15..193c9d33e9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -103,6 +103,12 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. + +{% if executable_schema_contract %} +### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION ### +{{ executable_schema_contract }} +{% endif %} + ### ORIGINAL SQL QUERY ### The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a6a41e4ef1..507957649f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -21,7 +21,21 @@ def _is_timeout_error(error_message: str) -> bool: - return error_message.startswith("Request timed out") + if not error_message: + return False + + normalized_error = error_message.lower() + return "timeout" in normalized_error or "timed out" in normalized_error + + +def _normalize_engine_addition(addition: Any) -> dict: + if isinstance(addition, dict): + return addition + + if addition: + return {"error_message": str(addition), "correlation_id": ""} + + return {} def _canonicalize_wren_sql_syntax(sql: str | None) -> str | None: @@ -340,7 +354,7 @@ async def _classify_generation_result( limit=1, dry_run=True, ) - addition = addition if isinstance(addition, dict) else {} + addition = _normalize_engine_addition(addition) if success: valid_generation_result = { @@ -371,6 +385,7 @@ async def _classify_generation_result( limit=1, dry_run=True, ) + addition = _normalize_engine_addition(addition) if success: valid_generation_result = { @@ -401,6 +416,7 @@ async def _classify_generation_result( limit=1, dry_run=False, ) + addition = _normalize_engine_addition(addition) if has_data: valid_generation_result = { diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cfaf212c91..a844eebf42 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -104,7 +104,7 @@ def __init__( allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, - max_sql_correction_retries: int = 3, + max_sql_correction_retries: int = 1, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index 8bbf502899..cd5f25d9e7 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -50,6 +50,28 @@ async def execute_sql( return False, None, {"error_message": "Execution rejected the statement"} +class StringTimeoutEngine: + async def dry_plan( + self, + session, + sql, + data_source, + project_id=None, + allow_fallback=True, + ): + return True, None + + async def execute_sql( + self, + sql, + session, + project_id=None, + limit=1, + dry_run=True, + ): + return False, None, "Timeout when connecting to execution engine" + + class CapturingEngine: def __init__(self): self.sql = None @@ -169,6 +191,22 @@ async def test_sql_post_processor_returns_generated_sql_when_dry_run_times_out() assert result["invalid_generation_result"] == {} +@pytest.mark.asyncio +async def test_sql_post_processor_returns_generated_sql_when_engine_timeout_is_string(): + result = await SQLGenPostProcessor(StringTimeoutEngine()).run( + ['{"sql": "SELECT 1"}'], + project_id="project-id", + use_dry_plan=True, + data_source="source", + ) + + assert result["valid_generation_result"] == { + "sql": "SELECT 1", + "correlation_id": "", + } + assert result["invalid_generation_result"] == {} + + @pytest.mark.asyncio async def test_sql_post_processor_keeps_non_timeout_dry_plan_errors_invalid(): result = await SQLGenPostProcessor(ErrorEngine()).run( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index d474b0b806..251bae7be2 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -8,6 +8,10 @@ prompt as build_sql_generation_prompt, sql_generation_user_prompt_template, ) +from src.pipelines.generation.sql_regeneration import ( + prompt as build_sql_regeneration_prompt, + sql_regeneration_user_prompt_template, +) from src.pipelines.generation.utils.sql import build_executable_schema_contract @@ -61,6 +65,7 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): built_prompt = result["prompt"] + assert "ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST" in built_prompt assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt assert "TABLE: retrieved_model" in built_prompt assert "- grouping_attribute" in built_prompt @@ -104,6 +109,29 @@ def test_sql_correction_prompt_includes_executable_schema_contract(): built_prompt = result["prompt"] + assert "ALLOWED EXECUTABLE IDENTIFIERS FOR THIS CORRECTION" in built_prompt assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt assert "TABLE: retrieved_model" in built_prompt assert "Failed SQL:" in built_prompt + + +def test_sql_regeneration_prompt_includes_executable_schema_contract(): + result = build_sql_regeneration_prompt( + query="summarize the records", + documents=[], + sql_generation_reasoning="", + sql="", + prompt_builder=PromptBuilder(template=sql_regeneration_user_prompt_template), + schema_contracts=[ + { + "table_name": "retrieved_model", + "column_names": ["grouping_attribute", "numeric_measure"], + } + ], + ) + + built_prompt = result["prompt"] + + assert "ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION" in built_prompt + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt + assert "TABLE: retrieved_model" in built_prompt diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 15d346146d..6351df8adb 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -115,6 +115,12 @@ def test_ask_request_uses_strict_dry_plan_by_default(): assert ask_request.allow_dry_plan_fallback is False +def test_ask_service_uses_single_sql_correction_retry_by_default(): + ask_service = AskService({}) + + assert ask_service._max_sql_correction_retries == 1 + + @pytest.mark.asyncio async def test_ask_with_successful_query( indexing_service: SemanticsPreparationService, diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index ce98d0f4db..0b7bc68e95 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -1361,18 +1361,29 @@ export class ModelResolver { : await ctx.projectService.getCurrentProject(); const manifest = await this.getLastDeployedManifest(ctx, project.id); - if (project.type === DataSourceName.DUCKDB) { - await ctx.wrenEngineAdaptor.getNativeSQL(sql, { - manifest, - modelingOnly: false, - }); - } else { - await ctx.ibisServerAdaptor.getNativeSql({ - dataSource: project.type, - sql, - mdl: manifest, - allowFallback, - }); + try { + if (project.type === DataSourceName.DUCKDB) { + await ctx.wrenEngineAdaptor.getNativeSQL(sql, { + manifest, + modelingOnly: false, + }); + } else { + await ctx.ibisServerAdaptor.getNativeSql({ + dataSource: project.type, + sql, + mdl: manifest, + allowFallback, + }); + } + } catch (error) { + if (this.isDryPlanTimeout(error)) { + logger.warn( + 'Dry plan timed out; accepting generated Wren SQL without native rewrite', + ); + return true; + } + + throw error; } return true; @@ -1486,6 +1497,11 @@ export class ModelResolver { return value; } + private isDryPlanTimeout(error: unknown): boolean { + const errorMessage = JSON.stringify(error ?? '').toLowerCase(); + return errorMessage.includes('timeout') || errorMessage.includes('timed out'); + } + // validate view name private async validateViewName( viewDisplayName: string, From be1226f496aaac37026a7c49964d850923ebc07c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 17:03:54 +0530 Subject: [PATCH 0802/1087] Fix REST ask deployment readiness --- wren-ui/src/pages/api/v1/ask.ts | 3 ++- wren-ui/src/pages/api/v1/generate_sql.ts | 3 ++- wren-ui/src/pages/api/v1/stream/ask.ts | 3 ++- wren-ui/src/pages/api/v1/stream/generate_sql.ts | 3 ++- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index c7b966fde0..c5ba52626c 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -71,6 +71,7 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const deployId = await deployService.ensureDeploymentPrepared(project.id); // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); @@ -83,7 +84,7 @@ export default async function handler( // Step 1: Generate SQL const askTask = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { diff --git a/wren-ui/src/pages/api/v1/generate_sql.ts b/wren-ui/src/pages/api/v1/generate_sql.ts index a6cef6f3bd..6d3d154372 100644 --- a/wren-ui/src/pages/api/v1/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/generate_sql.ts @@ -71,6 +71,7 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const deployId = await deployService.ensureDeploymentPrepared(project.id); // ask AI service to generate SQL const histories = threadId @@ -78,7 +79,7 @@ export default async function handler( : undefined; const task = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { diff --git a/wren-ui/src/pages/api/v1/stream/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts index 93a5a6fbfc..cffdcaf48a 100644 --- a/wren-ui/src/pages/api/v1/stream/ask.ts +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -132,6 +132,7 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const deployId = await deployService.ensureDeploymentPrepared(project.id); // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); @@ -150,7 +151,7 @@ export default async function handler( }); const askTask = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { diff --git a/wren-ui/src/pages/api/v1/stream/generate_sql.ts b/wren-ui/src/pages/api/v1/stream/generate_sql.ts index 9d65fc5618..16f593dc3d 100644 --- a/wren-ui/src/pages/api/v1/stream/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/stream/generate_sql.ts @@ -74,6 +74,7 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const deployId = await deployService.ensureDeploymentPrepared(project.id); // Get conversation history if threadId is provided const histories = threadId @@ -90,7 +91,7 @@ export default async function handler( const askTask = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { From e8b8aa133f636f4c28f74631ad6b5f274920d296 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 17:18:16 +0530 Subject: [PATCH 0803/1087] Load deployed schema directly for ask retrieval --- .../retrieval/db_schema_retrieval.py | 51 +++++-- .../retrieval/test_db_schema_retrieval.py | 134 ++++++++++++++++++ 2 files changed, 174 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 183be97b7e..2b4b52121e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -254,11 +254,32 @@ async def active_mdl_hash( @observe(capture_input=False, capture_output=False) -async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: +async def embedding( + query: str, + embedder: Any, + histories: list[AskHistory], + project_id: str = "", + mdl_hash: str = "", + dbschema_store: Any = None, +) -> dict: + if project_id and mdl_hash and dbschema_store: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + *build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + )["conditions"], + ], + } + if await dbschema_store.count_documents(filters=filters): + return {} + if query: return await embedder.run(query) - else: - return {} + + return {} @observe(capture_input=False) @@ -293,15 +314,16 @@ async def table_retrieval( query_embedding=embedding.get("embedding"), filters=filters, ) - else: - filters["conditions"].append( - {"field": "name", "operator": "in", "value": tables} - ) - return await table_retriever.run( - query_embedding=[], - filters=filters, - ) + if not tables: + return {"documents": []} + + filters["conditions"].append({"field": "name", "operator": "in", "value": tables}) + + return await table_retriever.run( + query_embedding=[], + filters=filters, + ) @observe(capture_input=False) @@ -390,6 +412,13 @@ def _related_table_names(documents: list[Document], visited: set[str]) -> list[s table_names.append(table_name) documents = [] + if not table_names and not (embedding and embedding.get("embedding")): + results = await dbschema_retriever.run( + query_embedding=[], + filters=_base_filters(), + ) + return results["documents"] + if not table_names and embedding and embedding.get("embedding"): results = await dbschema_retriever.run( query_embedding=embedding.get("embedding"), diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index f14fae9103..148a885a0f 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -49,6 +49,42 @@ async def run(self, query): assert embedder.query == "current request" +@pytest.mark.asyncio +async def test_embedding_skips_semantic_search_when_deployed_schema_exists(): + class Embedder: + def __init__(self): + self.called = False + + async def run(self, query): + self.called = True + return {"embedding": [1.0]} + + schema_store = StoreCounter(count=1) + embedder = Embedder() + + result = await embedding( + query="current request", + embedder=embedder, + histories=[], + project_id="project-1", + mdl_hash="deploy-1", + dbschema_store=schema_store, + ) + + assert result == {} + assert embedder.called is False + assert schema_store.filters == [ + { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + } + ] + + def test_column_pruning_prompt_uses_current_query_without_history_text(): result = build_column_selection_prompt( query="current request", @@ -182,6 +218,30 @@ async def run(self, query_embedding, filters): } +@pytest.mark.asyncio +async def test_table_retrieval_skips_candidate_search_without_embedding_or_tables(): + class Retriever: + def __init__(self): + self.called = False + + async def run(self, query_embedding, filters): + self.called = True + return {"documents": []} + + retriever = Retriever() + + result = await table_retrieval( + embedding={}, + project_id="project-1", + tables=[], + table_retriever=retriever, + mdl_hash="deploy-1", + ) + + assert result == {"documents": []} + assert retriever.called is False + + @pytest.mark.asyncio async def test_table_retrieval_keeps_deploy_scope_when_no_documents_match(): class Retriever: @@ -369,6 +429,80 @@ async def run(self, query_embedding, filters): ] +@pytest.mark.asyncio +async def test_dbschema_retrieval_loads_exact_deployed_schema_without_candidates(): + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "orders", + "comment": "", + "columns": [], + "properties": {}, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": "orders"}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "order_id", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": True, + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "orders"}, + ), + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={"documents": []}, + project_id="project-1", + mdl_hash="deploy-1", + dbschema_retriever=retriever, + embedding={}, + ) + + assert [document.meta["name"] for document in documents] == ["orders", "orders"] + assert retriever.calls == [ + { + "query_embedding": [], + "filters": { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + }, + } + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_expands_declared_relationships(): selected_model = "model_a" From 58ea85b80cd8078b1c789859906f4b0d376520c3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 17:41:46 +0530 Subject: [PATCH 0804/1087] Ground asks to active deployment metadata --- .../retrieval/db_schema_retrieval.py | 9 +++- .../web/v1/routers/semantics_preparation.py | 3 +- .../web/v1/services/semantics_preparation.py | 31 ++++++++++--- .../services/test_semantics_preparation.py | 44 +++++++++++++++++++ .../adaptors/tests/wrenAIAdaptor.test.ts | 8 ++++ .../apollo/server/adaptors/wrenAIAdaptor.ts | 22 +++++++--- .../apollo/server/services/deployService.ts | 5 ++- .../services/tests/deployService.test.ts | 5 ++- wren-ui/src/pages/api/v1/ask.ts | 4 +- wren-ui/src/pages/api/v1/generate_sql.ts | 5 ++- wren-ui/src/pages/api/v1/stream/ask.ts | 4 +- .../src/pages/api/v1/stream/generate_sql.ts | 3 ++ 12 files changed, 124 insertions(+), 19 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 2b4b52121e..ccf4fb8a64 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -242,13 +242,20 @@ async def active_mdl_hash( dbschema_store.count_documents(filters=filters), ) - if not table_description_count and not dbschema_count: + if not dbschema_count: logger.warning( "Project ID: %s, MDL hash %s has no indexed schema documents; " "keeping hash scope to avoid stale project metadata fallback.", project_id, mdl_hash, ) + elif not table_description_count: + logger.info( + "Project ID: %s, MDL hash %s has indexed db schema documents but no table descriptions; " + "using db schema retrieval fallback.", + project_id, + mdl_hash, + ) return mdl_hash diff --git a/wren-ai-service/src/web/v1/routers/semantics_preparation.py b/wren-ai-service/src/web/v1/routers/semantics_preparation.py index e6a4efab7e..8434cfaa7d 100644 --- a/wren-ai-service/src/web/v1/routers/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/routers/semantics_preparation.py @@ -42,10 +42,11 @@ async def prepare_semantics( @router.get("/semantics-preparations/{mdl_hash}/status") async def get_prepare_semantics_status( mdl_hash: str, + project_id: str | None = None, service_container: ServiceContainer = Depends(get_service_container), ) -> SemanticsPreparationStatusResponse: return await service_container.semantics_preparation_service.get_prepare_semantics_status( - SemanticsPreparationStatusRequest(mdl_hash=mdl_hash) + SemanticsPreparationStatusRequest(mdl_hash=mdl_hash, project_id=project_id) ) 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 a546f16457..491c4677fc 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -28,7 +28,7 @@ class SemanticsPreparationResponse(BaseModel): # GET /v1/semantics-preparations/{mdl_hash}/status -class SemanticsPreparationStatusRequest(BaseModel): +class SemanticsPreparationStatusRequest(BaseRequest): # don't recommend to use id as a field name, but it's used in the API spec # so we need to support as a choice, and will remove it in the future mdl_hash: str = Field(validation_alias=AliasChoices("mdl_hash", "id")) @@ -159,7 +159,14 @@ async def _has_indexed_schema_documents( mdl_hash=mdl_hash, ), ) - return dbschema_count > 0 and table_description_count > 0 + logger.info( + "Project ID: %s, MDL hash %s indexed schema document counts: db_schema=%s, table_description=%s", + project_id, + mdl_hash, + dbschema_count, + table_description_count, + ) + return dbschema_count > 0 async def get_prepare_semantics_status( self, prepare_semantics_status_request: SemanticsPreparationStatusRequest @@ -169,6 +176,15 @@ async def get_prepare_semantics_status( prepare_semantics_status_request.mdl_hash ) ) is None: + if ( + prepare_semantics_status_request.project_id + and await self._has_indexed_schema_documents( + prepare_semantics_status_request.project_id, + prepare_semantics_status_request.mdl_hash, + ) + ): + return SemanticsPreparationStatusResponse(status="finished") + logger.exception( f"id is not found for SemanticsPreparation: {prepare_semantics_status_request.mdl_hash}" ) @@ -176,17 +192,20 @@ async def get_prepare_semantics_status( status="failed", error=SemanticsPreparationStatusResponse.SemanticsPreparationError( code="OTHERS", - message="{prepare_semantics_status_request.id} is not found", + message=f"{prepare_semantics_status_request.mdl_hash} is not found", ), ) if result.status != "finished": return result - project_id = self._prepare_semantics_project_ids.get( - prepare_semantics_status_request.mdl_hash + project_id = ( + prepare_semantics_status_request.project_id + or self._prepare_semantics_project_ids.get( + prepare_semantics_status_request.mdl_hash + ) ) - if project_id is None: + if not project_id: return result if await self._has_indexed_schema_documents( diff --git a/wren-ai-service/tests/pytest/services/test_semantics_preparation.py b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py index 106d6b3e2b..0f7647e725 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_preparation.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py @@ -63,3 +63,47 @@ async def test_prepare_semantics_status_stays_finished_when_exact_schema_documen assert status.status == "finished" assert status.error is None + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_stays_finished_when_schema_documents_exist_without_descriptions(): + service = SemanticsPreparationService( + { + "db_schema": CountPipeline(1), + "table_description": CountPipeline(0), + } + ) + service._prepare_semantics_statuses["deploy-1"] = SemanticsPreparationStatusResponse( + status="finished" + ) + service._prepare_semantics_project_ids["deploy-1"] = "project-1" + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1") + ) + + assert status.status == "finished" + assert status.error is None + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_recovers_when_status_cache_is_missing(): + db_schema = CountPipeline(1) + table_description = CountPipeline(1) + service = SemanticsPreparationService( + { + "db_schema": db_schema, + "table_description": table_description, + } + ) + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1", project_id="project-1") + ) + + assert status.status == "finished" + assert status.error is None + assert db_schema.calls == [{"project_id": "project-1", "mdl_hash": "deploy-1"}] + assert table_description.calls == [ + {"project_id": "project-1", "mdl_hash": "deploy-1"} + ] diff --git a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts index 36bdf9be4a..3f9300a858 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts @@ -58,6 +58,14 @@ describe('WrenAIAdaptor', () => { project_id: mockInput.projectId.toString(), }, ); + expect(mockedAxios.get).toHaveBeenCalledWith( + `${baseEndpoint}/v1/semantics-preparations/${mockInput.hash}/status`, + { + params: { + project_id: mockInput.projectId.toString(), + }, + }, + ); }); }); diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 567a184c95..f2b8c065ee 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -51,7 +51,10 @@ const getAIServiceError = (error: any) => { export interface IWrenAIAdaptor { deploy(deployData: DeployData): Promise; - getDeployStatus(deployId: string): Promise; + getDeployStatus( + deployId: string, + projectId?: string | number, + ): Promise; delete(projectId: number): Promise; /** @@ -368,7 +371,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { logger.debug( `Wren AI: Deploying wren AI, hash: ${hash}, deployId: ${deployId}`, ); - const deploySuccess = await this.waitDeployFinished(deployId); + const deploySuccess = await this.waitDeployFinished(deployId, projectId); if (deploySuccess) { logger.debug(`Wren AI: Deploy wren AI success, hash: ${hash}`); return { status: WrenAIDeployStatusEnum.SUCCESS }; @@ -872,14 +875,17 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { }; } - private async waitDeployFinished(deployId: string): Promise { + private async waitDeployFinished( + deployId: string, + projectId?: string | number, + ): Promise { let deploySuccess = false; const maxAttempts = 90; const pollingIntervalMs = 2000; for (let attempt = 1; attempt <= maxAttempts; attempt++) { try { - const status = await this.getDeployStatus(deployId); + const status = await this.getDeployStatus(deployId, projectId); logger.debug( `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, ); @@ -902,10 +908,16 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { return deploySuccess; } - public async getDeployStatus(deployId: string): Promise { + public async getDeployStatus( + deployId: string, + projectId?: string | number, + ): Promise { try { const res = await axios.get( `${this.wrenAIBaseEndpoint}/v1/semantics-preparations/${deployId}/status`, + { + params: projectId ? { project_id: projectId.toString() } : undefined, + }, ); if (res.data.error) { const error = diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 085dad9e0e..0004f5ab2d 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -86,7 +86,10 @@ export class DeployService implements IDeployService { } try { - const status = await this.wrenAIAdaptor.getDeployStatus(lastDeploy.hash); + const status = await this.wrenAIAdaptor.getDeployStatus( + lastDeploy.hash, + projectId, + ); if (status === WrenAISystemStatus.FINISHED) { return lastDeploy.hash; } diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 739ef8e963..b269fe368b 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -118,7 +118,10 @@ describe('DeployService', () => { const hash = await deployService.ensureDeploymentPrepared(1); expect(hash).toEqual('deploy-hash'); - expect(mockWrenAIAdaptor.getDeployStatus).toHaveBeenCalledWith('deploy-hash'); + expect(mockWrenAIAdaptor.getDeployStatus).toHaveBeenCalledWith( + 'deploy-hash', + 1, + ); expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index c5ba52626c..52666760d8 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -62,7 +62,6 @@ export default async function handler( throw new ApiError('Question is required', 400); } - // Get current project's last deployment const lastDeploy = await deployService.getLastDeployment(project.id); if (!lastDeploy) { throw new ApiError( @@ -73,6 +72,9 @@ export default async function handler( } const deployId = await deployService.ensureDeploymentPrepared(project.id); + // Get current project's prepared deployment + const deployId = await deployService.ensureDeploymentPrepared(project.id); + // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); diff --git a/wren-ui/src/pages/api/v1/generate_sql.ts b/wren-ui/src/pages/api/v1/generate_sql.ts index 6d3d154372..f5bed1ff8d 100644 --- a/wren-ui/src/pages/api/v1/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/generate_sql.ts @@ -61,9 +61,7 @@ export default async function handler( throw new ApiError('Question is required', 400); } - // get current project's last deployment const lastDeploy = await deployService.getLastDeployment(project.id); - if (!lastDeploy) { throw new ApiError( 'No deployment found, please deploy a model first', @@ -73,6 +71,9 @@ export default async function handler( } const deployId = await deployService.ensureDeploymentPrepared(project.id); + // get current project's prepared deployment + const deployId = await deployService.ensureDeploymentPrepared(project.id); + // ask AI service to generate SQL const histories = threadId ? await apiHistoryRepository.findAllBy({ threadId }) diff --git a/wren-ui/src/pages/api/v1/stream/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts index cffdcaf48a..bc2aec47ac 100644 --- a/wren-ui/src/pages/api/v1/stream/ask.ts +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -123,7 +123,6 @@ export default async function handler( // Send message start event sendMessageStart(res); - // Get current project's last deployment const lastDeploy = await deployService.getLastDeployment(project.id); if (!lastDeploy) { throw new ApiError( @@ -134,6 +133,9 @@ export default async function handler( } const deployId = await deployService.ensureDeploymentPrepared(project.id); + // Get current project's prepared deployment + const deployId = await deployService.ensureDeploymentPrepared(project.id); + // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); diff --git a/wren-ui/src/pages/api/v1/stream/generate_sql.ts b/wren-ui/src/pages/api/v1/stream/generate_sql.ts index 16f593dc3d..66c0dc351c 100644 --- a/wren-ui/src/pages/api/v1/stream/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/stream/generate_sql.ts @@ -76,6 +76,9 @@ export default async function handler( } const deployId = await deployService.ensureDeploymentPrepared(project.id); + // Get current project's prepared deployment + const deployId = await deployService.ensureDeploymentPrepared(project.id); + // Get conversation history if threadId is provided const histories = threadId ? await apiHistoryRepository.findAllBy({ threadId }) From fa0053010f5bd901d7b42a1663293fe838540385 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 17:50:09 +0530 Subject: [PATCH 0805/1087] Align ask pipeline strict metadata defaults --- .../config_examples/config.anthropic.yaml | 4 +- .../docs/config_examples/config.azure.yaml | 4 +- .../docs/config_examples/config.bedrock.yaml | 4 +- .../docs/config_examples/config.deepseek.yaml | 4 +- .../config.google_ai_studio.yaml | 4 +- .../config.google_vertexai.yaml | 4 +- .../docs/config_examples/config.grok.yaml | 4 +- .../docs/config_examples/config.groq.yaml | 4 +- .../config_examples/config.lm_studio.yaml | 4 +- .../docs/config_examples/config.ollama.yaml | 4 +- .../config_examples/config.open_router.yaml | 4 +- .../docs/config_examples/config.qwen3.yaml | 4 +- .../docs/config_examples/config.zhipu.yaml | 4 +- wren-ai-service/src/config.py | 2 +- .../tests/pytest/services/test_ask.py | 44 ++++++++++++++++++- wren-ai-service/tests/pytest/test_config.py | 2 + .../tools/config/config.example.yaml | 4 +- wren-ai-service/tools/config/config.full.yaml | 4 +- 18 files changed, 75 insertions(+), 33 deletions(-) diff --git a/wren-ai-service/docs/config_examples/config.anthropic.yaml b/wren-ai-service/docs/config_examples/config.anthropic.yaml index 76e5d96526..a781436cdd 100644 --- a/wren-ai-service/docs/config_examples/config.anthropic.yaml +++ b/wren-ai-service/docs/config_examples/config.anthropic.yaml @@ -152,10 +152,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.azure.yaml b/wren-ai-service/docs/config_examples/config.azure.yaml index 9319394727..b6b44e0af5 100644 --- a/wren-ai-service/docs/config_examples/config.azure.yaml +++ b/wren-ai-service/docs/config_examples/config.azure.yaml @@ -165,10 +165,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.bedrock.yaml b/wren-ai-service/docs/config_examples/config.bedrock.yaml index a94474de67..02f0633eb3 100644 --- a/wren-ai-service/docs/config_examples/config.bedrock.yaml +++ b/wren-ai-service/docs/config_examples/config.bedrock.yaml @@ -168,10 +168,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.deepseek.yaml b/wren-ai-service/docs/config_examples/config.deepseek.yaml index 3a00b30ad8..7749844074 100644 --- a/wren-ai-service/docs/config_examples/config.deepseek.yaml +++ b/wren-ai-service/docs/config_examples/config.deepseek.yaml @@ -175,10 +175,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com 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 9d087accb3..f54d85a082 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 @@ -161,10 +161,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.google_vertexai.yaml b/wren-ai-service/docs/config_examples/config.google_vertexai.yaml index 0b29acb7e5..94611c8523 100644 --- a/wren-ai-service/docs/config_examples/config.google_vertexai.yaml +++ b/wren-ai-service/docs/config_examples/config.google_vertexai.yaml @@ -169,10 +169,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.grok.yaml b/wren-ai-service/docs/config_examples/config.grok.yaml index b01de55b7e..9f35b60b8d 100644 --- a/wren-ai-service/docs/config_examples/config.grok.yaml +++ b/wren-ai-service/docs/config_examples/config.grok.yaml @@ -157,10 +157,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.groq.yaml b/wren-ai-service/docs/config_examples/config.groq.yaml index a07a577e9d..0f7c1fc05f 100644 --- a/wren-ai-service/docs/config_examples/config.groq.yaml +++ b/wren-ai-service/docs/config_examples/config.groq.yaml @@ -156,10 +156,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.lm_studio.yaml b/wren-ai-service/docs/config_examples/config.lm_studio.yaml index 131ac04384..ef29353d93 100644 --- a/wren-ai-service/docs/config_examples/config.lm_studio.yaml +++ b/wren-ai-service/docs/config_examples/config.lm_studio.yaml @@ -155,10 +155,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.ollama.yaml b/wren-ai-service/docs/config_examples/config.ollama.yaml index 5f8e6c4ea2..add673c787 100644 --- a/wren-ai-service/docs/config_examples/config.ollama.yaml +++ b/wren-ai-service/docs/config_examples/config.ollama.yaml @@ -155,10 +155,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.open_router.yaml b/wren-ai-service/docs/config_examples/config.open_router.yaml index ecbcaa0731..0bebc90d25 100644 --- a/wren-ai-service/docs/config_examples/config.open_router.yaml +++ b/wren-ai-service/docs/config_examples/config.open_router.yaml @@ -153,10 +153,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.qwen3.yaml b/wren-ai-service/docs/config_examples/config.qwen3.yaml index 0ebaed162a..1fa694ed4d 100644 --- a/wren-ai-service/docs/config_examples/config.qwen3.yaml +++ b/wren-ai-service/docs/config_examples/config.qwen3.yaml @@ -195,10 +195,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.zhipu.yaml b/wren-ai-service/docs/config_examples/config.zhipu.yaml index 8d0db87d6c..7d762b33ce 100644 --- a/wren-ai-service/docs/config_examples/config.zhipu.yaml +++ b/wren-ai-service/docs/config_examples/config.zhipu.yaml @@ -203,10 +203,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index ecb7b1cc0d..1c85fbaa98 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -43,7 +43,7 @@ class Settings(BaseSettings): allow_sql_diagnosis: bool = Field(default=True) allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) - max_sql_correction_retries: int = Field(default=3) + max_sql_correction_retries: int = Field(default=1) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 6351df8adb..e252e6819d 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -22,6 +22,24 @@ @pytest.fixture def ask_service(): pipe_components = generate_components(settings.components) + required_components = { + "intent_classification", + "misleading_assistance", + "data_assistance", + "user_guide_assistance", + "db_schema_retrieval", + "historical_question_retrieval", + "sql_generation", + "sql_correction", + "sql_pairs_retrieval", + "instructions_retrieval", + } + missing_components = required_components - pipe_components.keys() + if missing_components: + pytest.skip( + f"Ask integration test requires configured components: {sorted(missing_components)}" + ) + wren_ai_docs = fetch_wren_ai_docs(settings.doc_endpoint, settings.is_oss) return AskService( @@ -40,7 +58,7 @@ def ask_service(): **pipe_components["user_guide_assistance"], wren_ai_docs=wren_ai_docs, ), - "retrieval": retrieval.DbSchemaRetrieval( + "db_schema_retrieval": retrieval.DbSchemaRetrieval( **pipe_components["db_schema_retrieval"], ), "historical_question": retrieval.HistoricalQuestionRetrieval( @@ -58,13 +76,28 @@ def ask_service(): "instructions_retrieval": retrieval.Instructions( **pipe_components["instructions_retrieval"], ), - } + }, + allow_sql_functions_retrieval=False, + allow_sql_diagnosis=False, + allow_sql_knowledge_retrieval=False, ) @pytest.fixture def indexing_service(): pipe_components = generate_components(settings.components) + required_components = { + "db_schema_indexing", + "historical_question_indexing", + "table_description_indexing", + "sql_pairs_indexing", + "project_meta_indexing", + } + missing_components = required_components - pipe_components.keys() + if missing_components: + pytest.skip( + f"Ask integration test requires configured components: {sorted(missing_components)}" + ) return SemanticsPreparationService( { @@ -77,6 +110,13 @@ def indexing_service(): "table_description": indexing.TableDescription( **pipe_components["table_description_indexing"], ), + "sql_pairs": indexing.SqlPairs( + **pipe_components["sql_pairs_indexing"], + sql_pairs_path=settings.sql_pairs_path, + ), + "project_meta": indexing.ProjectMeta( + **pipe_components["project_meta_indexing"], + ), } ) diff --git a/wren-ai-service/tests/pytest/test_config.py b/wren-ai-service/tests/pytest/test_config.py index 4a745d49e9..851a2ece81 100644 --- a/wren-ai-service/tests/pytest/test_config.py +++ b/wren-ai-service/tests/pytest/test_config.py @@ -23,6 +23,8 @@ def test_settings_default_values(): assert settings.logging_level == "INFO" assert settings.development is False + assert settings.allow_sql_generation_reasoning is False + assert settings.max_sql_correction_retries == 1 assert settings.config_path == "config.yaml" diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index b1b98a0962..792e5ef6c9 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -189,10 +189,10 @@ settings: table_retrieval_size: 50 table_column_retrieval_size: 100 allow_intent_classification: false - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 265292b9f1..5639d51b01 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -187,10 +187,10 @@ settings: table_column_retrieval_size: 100 query_cache_maxsize: 1000 allow_intent_classification: false - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com langfuse_enable: true From 904b4ef966aabdd73003bfe36284927e42c163f3 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 18:03:53 +0530 Subject: [PATCH 0806/1087] Return active deployment hash for ask preparation --- .../apollo/server/services/deployService.ts | 38 +++++++++++-------- .../services/tests/deployService.test.ts | 38 +++++++++++++++++-- 2 files changed, 57 insertions(+), 19 deletions(-) diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 0004f5ab2d..8f2d6753ef 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -25,6 +25,7 @@ const STALE_DEPLOYMENT_MS = 10 * 60 * 1000; export interface DeployResponse { status: DeployStatusEnum; error?: string; + hash?: string; } export interface MDLSyncResponse { @@ -84,21 +85,28 @@ export class DeployService implements IDeployService { if (!lastDeploy) { throw new Error(`No deployment found for project ${projectId}`); } + const activeHash = this.createMDLHash(lastDeploy.manifest, projectId); - try { - const status = await this.wrenAIAdaptor.getDeployStatus( - lastDeploy.hash, - projectId, - ); - if (status === WrenAISystemStatus.FINISHED) { - return lastDeploy.hash; + if (lastDeploy.hash === activeHash) { + try { + const status = await this.wrenAIAdaptor.getDeployStatus( + activeHash, + projectId, + ); + if (status === WrenAISystemStatus.FINISHED) { + return activeHash; + } + logger.warn( + `Deployment ${activeHash} is not ready in AI service: ${status}`, + ); + } catch (err: any) { + logger.warn( + `Deployment ${activeHash} is not available in AI service: ${err.message}`, + ); } + } else { logger.warn( - `Deployment ${lastDeploy.hash} is not ready in AI service: ${status}`, - ); - } catch (err: any) { - logger.warn( - `Deployment ${lastDeploy.hash} is not available in AI service: ${err.message}`, + `Deployment ${lastDeploy.hash} does not match current manifest hash ${activeHash}; preparing current hash.`, ); } @@ -110,7 +118,7 @@ export class DeployService implements IDeployService { ); } - return lastDeploy.hash; + return result.hash || activeHash; } public async getInProgressDeployment(projectId) { @@ -163,7 +171,7 @@ export class DeployService implements IDeployService { status, error: aiError, }); - return { status, error: aiError }; + return { status, error: aiError, hash }; } } const previousInProgressDeploy = @@ -212,7 +220,7 @@ export class DeployService implements IDeployService { false, ); } - return { status, error: aiError }; + return { status, error: aiError, hash }; } catch (err: any) { logger.error(`Error deploying model: ${err.message}`); if (deploy?.id) { diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index b269fe368b..eab7dc5c5d 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -127,23 +127,24 @@ describe('DeployService', () => { it('should redeploy the saved manifest when ai-service no longer has the exact deployment prepared', async () => { const manifest = { key: 'value' }; + const hash = deployService.createMDLHash(manifest, 1); mockDeployLogRepository.findLastProjectDeployLog .mockResolvedValueOnce({ id: 123, - hash: deployService.createMDLHash(manifest, 1), + hash, manifest, }) .mockResolvedValueOnce({ id: 123, - hash: deployService.createMDLHash(manifest, 1), + hash, manifest, }); mockWrenAIAdaptor.getDeployStatus.mockRejectedValue(new Error('not found')); mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); - const hash = await deployService.ensureDeploymentPrepared(1); + const preparedHash = await deployService.ensureDeploymentPrepared(1); - expect(hash).toEqual(deployService.createMDLHash(manifest, 1)); + expect(preparedHash).toEqual(hash); expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ manifest, hash, @@ -151,6 +152,35 @@ describe('DeployService', () => { }); }); + it('should return the current manifest hash after redeploying a stale saved hash', async () => { + const manifest = { key: 'value' }; + const activeHash = deployService.createMDLHash(manifest, 1); + + mockDeployLogRepository.findLastProjectDeployLog + .mockResolvedValueOnce({ + id: 123, + hash: 'legacy-saved-hash', + manifest, + }) + .mockResolvedValueOnce({ + id: 456, + hash: activeHash, + manifest, + }); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 456 }); + + const preparedHash = await deployService.ensureDeploymentPrepared(1); + + expect(preparedHash).toEqual(activeHash); + expect(mockWrenAIAdaptor.getDeployStatus).not.toHaveBeenCalled(); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash: activeHash, + projectId: 1, + }); + }); + it('should create the same deployment hash for equivalent manifests', () => { const manifest = { models: [ From a537ec5bf77a388c167ee29ccc54a373c1138b1c Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 18:27:30 +0530 Subject: [PATCH 0807/1087] Reject invalid generated SQL before execution --- wren-ai-service/src/core/engine.py | 7 ++-- .../src/pipelines/generation/utils/sql.py | 31 ++++++++++++++++++ .../generation/test_sql_post_processor.py | 32 +++++++++++++++++++ .../apollo/server/resolvers/modelResolver.ts | 2 +- 4 files changed, 68 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/core/engine.py b/wren-ai-service/src/core/engine.py index 163135e02d..70a151b4ce 100644 --- a/wren-ai-service/src/core/engine.py +++ b/wren-ai-service/src/core/engine.py @@ -30,15 +30,16 @@ def clean_generation_result(result: str) -> str: def _normalize_whitespace(s: str) -> str: return re.sub(r"\s+", " ", s).strip() - return ( + cleaned = ( _normalize_whitespace(result) .replace("```sql", "") .replace("```json", "") .replace('"""', "") .replace("'''", "") .replace("```", "") - .replace(";", "") - ) + ).strip() + + return cleaned[:-1].strip() if cleaned.endswith(";") else cleaned def remove_limit_statement(sql: str) -> str: diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 507957649f..b4b2e4a058 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -197,6 +197,21 @@ def _table_grounding_error( return None +def _sql_statement_shape_error(sql: str | None) -> str | None: + if not sql: + return None + + statements = [ + statement + for statement in sqlparse.parse(sql) + if str(statement).strip().strip(";").strip() + ] + if len(statements) > 1: + return "Generated SQL contains multiple statements; return exactly one SQL statement." + + return None + + def build_executable_schema_contract(schema_contracts: list[dict] | None) -> str: if not schema_contracts: return "" @@ -254,11 +269,27 @@ async def run( cleaned_generation_result = orjson.loads(cleaned_generation_result).get( "sql" ) + cleaned_generation_result = clean_generation_result( + cleaned_generation_result + ) cleaned_generation_result = _canonicalize_wren_sql_syntax( cleaned_generation_result ) + shape_error = _sql_statement_shape_error(cleaned_generation_result) + if shape_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SQL_SYNTAX", + "error": shape_error, + "correlation_id": "", + }, + } + grounding_error = _table_grounding_error( cleaned_generation_result, schema_contracts ) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index cd5f25d9e7..cb24ea3982 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -1,5 +1,6 @@ import pytest +from src.core.engine import clean_generation_result from src.pipelines.generation.utils.sql import SQLGenPostProcessor @@ -90,6 +91,13 @@ async def execute_sql( return True, None, {"correlation_id": "valid-correlation"} +def test_clean_generation_result_preserves_internal_statement_separators(): + assert ( + clean_generation_result("SELECT * FROM a; SELECT * FROM b;") + == "SELECT * FROM a; SELECT * FROM b" + ) + + @pytest.mark.asyncio async def test_sql_post_processor_converts_select_top_to_wren_limit(): engine = CapturingEngine() @@ -107,6 +115,30 @@ async def test_sql_post_processor_converts_select_top_to_wren_limit(): assert result["invalid_generation_result"] == {} +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_multiple_statements_before_execution(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT * FROM first_model; SELECT * FROM second_model;"}'], + project_id="project-id", + schema_contracts=[ + {"table_name": "first_model", "column_names": []}, + {"table_name": "second_model", "column_names": []}, + ], + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT * FROM first_model; SELECT * FROM second_model", + "original_sql": "SELECT * FROM first_model; SELECT * FROM second_model", + "type": "SQL_SYNTAX", + "error": "Generated SQL contains multiple statements; return exactly one SQL statement.", + "correlation_id": "", + } + + @pytest.mark.asyncio async def test_sql_post_processor_rejects_tables_outside_schema_contract(): engine = CapturingEngine() diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 0b7bc68e95..6167508b73 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -1376,7 +1376,7 @@ export class ModelResolver { }); } } catch (error) { - if (this.isDryPlanTimeout(error)) { + if (this.isDryPlanTimeout(error) && allowFallback !== false) { logger.warn( 'Dry plan timed out; accepting generated Wren SQL without native rewrite', ); From a23d0c67cb1507fb61fd4ff0dc1a2061090f3bda Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 19:14:25 +0530 Subject: [PATCH 0808/1087] Keep strict dry plan timeouts invalid --- .../src/pipelines/generation/utils/sql.py | 12 +++++++++- .../generation/test_sql_post_processor.py | 23 ++++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index b4b2e4a058..6b54b6cd47 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -363,8 +363,18 @@ async def _classify_generation_result( if not dry_plan_result: if _is_timeout_error(error_message): - valid_generation_result = { + if allow_dry_plan_fallback: + valid_generation_result = { + "sql": generation_result, + "correlation_id": "", + } + return valid_generation_result, invalid_generation_result + + invalid_generation_result = { "sql": generation_result, + "original_sql": generation_result, + "type": "DRY_PLAN", + "error": error_message, "correlation_id": "", } return valid_generation_result, invalid_generation_result diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index cb24ea3982..db6eb1a9ab 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -194,12 +194,33 @@ async def test_sql_post_processor_allows_tables_inside_schema_contract(): @pytest.mark.asyncio -async def test_sql_post_processor_returns_generated_sql_when_dry_plan_times_out(): +async def test_sql_post_processor_keeps_dry_plan_timeout_invalid_without_fallback(): result = await SQLGenPostProcessor(TimeoutEngine()).run( ['{"sql": "SELECT 1"}'], project_id="project-id", use_dry_plan=True, data_source="source", + allow_dry_plan_fallback=False, + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT 1", + "original_sql": "SELECT 1", + "type": "DRY_PLAN", + "error": "Request timed out after 30 seconds", + "correlation_id": "", + } + + +@pytest.mark.asyncio +async def test_sql_post_processor_returns_generated_sql_when_dry_plan_fallback_is_allowed(): + result = await SQLGenPostProcessor(TimeoutEngine()).run( + ['{"sql": "SELECT 1"}'], + project_id="project-id", + use_dry_plan=True, + data_source="source", + allow_dry_plan_fallback=True, ) assert result["valid_generation_result"] == { From dcc3fa08e29a322a49a78a00ac6c333ba8e00644 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 20:32:27 +0530 Subject: [PATCH 0809/1087] Pass schema contracts to question validation --- .../v1/services/question_recommendation.py | 31 ++++++++- .../services/test_question_recommendation.py | 69 +++++++++++++++++++ 2 files changed, 97 insertions(+), 3 deletions(-) create mode 100644 wren-ai-service/tests/pytest/services/test_question_recommendation.py diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 694d044bfa..f4bb64caa1 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -71,7 +71,9 @@ async def _validate_question( use_dry_plan: bool = True, allow_dry_plan_fallback: bool = False, ): - async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: + async def _document_retrieval() -> tuple[ + list[str], list[dict], bool, bool, bool + ]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -79,10 +81,26 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + schema_contracts = [ + { + "table_name": document.get("table_name"), + "column_names": document.get("manifest_column_names") + or document.get("column_names") + or [], + } + for document in documents + if document.get("table_name") + ] has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return table_ddls, has_calculated_field, has_metric, has_json_field + return ( + table_ddls, + schema_contracts, + has_calculated_field, + has_metric, + has_json_field, + ) async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -107,7 +125,13 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - table_ddls, has_calculated_field, has_metric, has_json_field = _document + ( + table_ddls, + schema_contracts, + has_calculated_field, + has_metric, + has_json_field, + ) = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -137,6 +161,7 @@ async def _instructions_retrieval() -> list[dict]: allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, ) post_process = generated_sql["post_process"] diff --git a/wren-ai-service/tests/pytest/services/test_question_recommendation.py b/wren-ai-service/tests/pytest/services/test_question_recommendation.py new file mode 100644 index 0000000000..ad87c9b3f3 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_question_recommendation.py @@ -0,0 +1,69 @@ +import pytest + +from src.web.v1.services.question_recommendation import QuestionRecommendation + + +class DbSchemaRetrievalPipeline: + async def run(self, **_): + return { + "construct_retrieval_results": { + "retrieval_results": [ + { + "table_name": "dbo_qosSales", + "table_ddl": "CREATE TABLE dbo_qosSales (BU VARCHAR, SalesVal FLOAT)", + "manifest_column_names": ["BU", "SalesVal"], + } + ] + } + } + + +class EmptyFormattedPipeline: + async def run(self, **_): + return {"formatted_output": {}} + + +class CapturingSqlGenerationPipeline: + def __init__(self): + self.schema_contracts = None + + async def run(self, **kwargs): + self.schema_contracts = kwargs.get("schema_contracts") + return { + "post_process": { + "valid_generation_result": {}, + "invalid_generation_result": { + "type": "SCHEMA_GROUNDING", + "sql": "", + "original_sql": "", + "error": "invalid", + }, + } + } + + +@pytest.mark.asyncio +async def test_question_recommendation_passes_schema_contracts_to_sql_generation(): + sql_generation = CapturingSqlGenerationPipeline() + service = QuestionRecommendation( + pipelines={ + "db_schema_retrieval": DbSchemaRetrievalPipeline(), + "sql_pairs_retrieval": EmptyFormattedPipeline(), + "instructions_retrieval": EmptyFormattedPipeline(), + "sql_generation": sql_generation, + }, + allow_sql_functions_retrieval=False, + allow_sql_knowledge_retrieval=False, + ) + + await service._validate_question( + {"question": "show orders", "category": "General"}, + request_id="request-id", + max_questions=5, + max_categories=3, + project_id="11", + ) + + assert sql_generation.schema_contracts == [ + {"table_name": "dbo_qosSales", "column_names": ["BU", "SalesVal"]} + ] From 4551be91875d1b0ab11c9b6026c0dad5a2858fe7 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 21:11:46 +0530 Subject: [PATCH 0810/1087] Guard modeling node selection --- wren-ui/src/pages/modeling.tsx | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 9a075d4fb9..7f860afd54 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -485,11 +485,16 @@ export default function Modeling() { const onSelect = (selectKeys) => { if (diagramRef.current) { const { getNodes, fitBounds } = diagramRef.current; - const node = getNodes().find((node) => node.id === selectKeys[0]); + const selectedKey = selectKeys?.[0]; + if (!selectedKey) return; + + const node = getNodes().find((node) => node.id === selectedKey); + if (!node?.position) return; + const position = { ...node.position, - width: node.width, - height: node.height, + width: node.width ?? 1, + height: node.height ?? 1, }; fitBounds(position); } From 655833b3fbf64116b5223fe444ac12a952262145 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 3 Aug 2026 21:58:45 +0530 Subject: [PATCH 0811/1087] Fix relationship field metadata resolution --- .../selectors/CombineFieldSelector.tsx | 13 ++- wren-ui/src/pages/modeling.tsx | 94 +++++++++++++++---- wren-ui/src/utils/errorHandler.tsx | 5 +- 3 files changed, 92 insertions(+), 20 deletions(-) diff --git a/wren-ui/src/components/selectors/CombineFieldSelector.tsx b/wren-ui/src/components/selectors/CombineFieldSelector.tsx index 5b3da42537..ed41fe4117 100644 --- a/wren-ui/src/components/selectors/CombineFieldSelector.tsx +++ b/wren-ui/src/components/selectors/CombineFieldSelector.tsx @@ -39,6 +39,13 @@ export default function CombineFieldSelector(props: Props) { ...value, }); + useEffect(() => { + setInternalValue((currentValue) => ({ + model: value?.model || modelValue || currentValue.model, + field: value?.field || fieldValue || currentValue.field, + })); + }, [value?.model, value?.field, modelValue, fieldValue]); + const syncOnChange = () => { if (internalValue?.model && internalValue?.field) { onChange && onChange(internalValue); @@ -56,7 +63,11 @@ export default function CombineFieldSelector(props: Props) { const changeField = (field: string) => { onFieldChange && onFieldChange(field); - setInternalValue({ ...internalValue, field }); + setInternalValue({ + ...internalValue, + model: internalValue.model || value?.model || modelValue, + field, + }); }; return ( diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 7f860afd54..8511324c74 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -226,12 +226,48 @@ const normalizeRelationshipType = (type: string) => { return type; }; -const parseQualifiedField = (value = '') => { - const [model, ...columnParts] = String(value).split('.'); - return { - model: model || '', - column: columnParts.join('.') || '', - }; +const getRelationshipFieldValue = (model = '', column = '') => + `${model}.${column}`; + +const resolveRelationshipField = ( + value = '', + models: Array<{ referenceName?: string; fields?: any[] }> = [], +) => { + const fieldValue = String(value || ''); + for (const model of models) { + for (const field of model.fields || []) { + if ( + fieldValue === + getRelationshipFieldValue(model.referenceName, field.referenceName) + ) { + return { + model: model.referenceName || '', + column: field.referenceName || '', + }; + } + } + } + + return { model: '', column: '' }; +}; + +const resolveRelationshipFieldParts = ( + model = '', + column = '', + fallbackValue = '', + models: Array<{ referenceName?: string; fields?: any[] }> = [], +) => { + const resolvedField = + model && column + ? resolveRelationshipField( + getRelationshipFieldValue(model, column), + models, + ) + : { model: '', column: '' }; + + return resolvedField.model && resolvedField.column + ? resolvedField + : resolveRelationshipField(fallbackValue, models); }; const renderIcon = (IconComponent) => React.createElement(IconComponent as any); @@ -673,16 +709,23 @@ export default function Modeling() { : []; return relationships.map((relationship, index) => { - const from = parseQualifiedField( + const availableModels = diagramData?.models || []; + const from = resolveRelationshipFieldParts( + relationship.fromModel || '', + relationship.fromColumn || '', relationship.from || relationship.fromField || '', + availableModels, ); - const to = parseQualifiedField( + const to = resolveRelationshipFieldParts( + relationship.toModel || '', + relationship.toColumn || '', relationship.to || relationship.toField || '', + availableModels, ); - const fromModel = relationship.fromModel || from.model; - const fromColumn = relationship.fromColumn || from.column; - const toModel = relationship.toModel || to.model; - const toColumn = relationship.toColumn || to.column; + const fromModel = from.model; + const fromColumn = from.column; + const toModel = to.model; + const toColumn = to.column; return { clientId: @@ -831,7 +874,7 @@ export default function Modeling() { side: 'from' | 'to', value: string, ) => { - const field = parseQualifiedField(value); + const field = resolveRelationshipField(value, diagramData?.models || []); updateRelationship( clientId, side === 'from' @@ -959,7 +1002,10 @@ export default function Modeling() { (model) => (model.fields || []).map((field) => ({ label: `${model.referenceName}.${field.referenceName}`, - value: `${model.referenceName}.${field.referenceName}`, + value: getRelationshipFieldValue( + model.referenceName, + field.referenceName, + ), })), ); @@ -1269,7 +1315,10 @@ export default function Modeling() { @@ -1304,7 +1359,10 @@ export default function Modeling() { } /> ) : ( - `${record.toModel}.${record.toColumn}` + getRelationshipFieldValue( + record.toModel, + record.toColumn, + ) ), }, { diff --git a/wren-ui/src/utils/errorHandler.tsx b/wren-ui/src/utils/errorHandler.tsx index 06e9070092..0254870c80 100644 --- a/wren-ui/src/utils/errorHandler.tsx +++ b/wren-ui/src/utils/errorHandler.tsx @@ -226,7 +226,10 @@ class CreateRelationshipErrorHandler extends ErrorHandler { public getErrorMessage(error: GraphQLError) { switch (error.extensions?.code) { default: - return 'Failed to create relationship.'; + return replaceMessage( + 'Failed to create %{relationship}.', + error.message, + ); } } } From 24df7f8b7b7c1fef7b9ff2c5f1a69f541ee98156 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 21:56:45 +0530 Subject: [PATCH 0812/1087] Enforce semantic context authority in SQL prompts --- .../pipelines/generation/sql_correction.py | 2 +- .../src/pipelines/generation/utils/sql.py | 5 +++++ .../generation/test_sql_prompt_grounding.py | 19 +++++++++++++++++++ 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index fc3153ce39..45c3a9502e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -51,7 +51,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) The final answer must be in JSON format: {{ - "sql": + "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 6b54b6cd47..7af6a894e9 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -491,7 +491,10 @@ async def _classify_generation_result( _MANDATORY_SQL_GROUNDING_RULES = """ ### MANDATORY SQL GROUNDING RULES ### +- Treat the retrieved semantic context as the only authoritative source for this request. Do not use pretrained knowledge, common warehouse schemas, example schemas, or memorized business definitions as executable truth. - Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. +- Use only deployed semantic models, views, metrics, relationships, and columns that are present in the retrieved DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, EXECUTABLE WREN IDENTIFIER CATALOG, SQL FUNCTIONS, or current USER INSTRUCTIONS. +- Before generating SQL, silently validate that every model, column, metric, relationship, join path, filter field, grouping field, ordering field, and SQL function is present in the retrieved context. Generate SQL only after this validation succeeds. - Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. - Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. - Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. @@ -509,6 +512,7 @@ async def _classify_generation_result( - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. - When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- If multiple semantic interpretations exist and the retrieved context does not make one interpretation authoritative, return null for sql instead of choosing one. - When the same requested result can be answered from multiple schema objects with compatible columns or metrics, include all relevant schema objects by combining separate result rows with UNION ALL instead of choosing only one object. - Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. - If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. @@ -517,6 +521,7 @@ async def _classify_generation_result( - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. - Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. - Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. +- If SQL execution or validation fails, repair the query only when the repair can be verified using the same retrieved DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS. Never introduce a new schema object during repair. - If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index 251bae7be2..93676ab6d2 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -1,10 +1,12 @@ from haystack.components.builders.prompt_builder import PromptBuilder from src.pipelines.generation.sql_correction import ( + get_sql_correction_system_prompt, prompt as build_sql_correction_prompt, sql_correction_user_prompt_template, ) from src.pipelines.generation.sql_generation import ( + get_sql_generation_system_prompt, prompt as build_sql_generation_prompt, sql_generation_user_prompt_template, ) @@ -15,6 +17,23 @@ from src.pipelines.generation.utils.sql import build_executable_schema_contract +def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): + prompt = get_sql_generation_system_prompt() + + assert "retrieved semantic context as the only authoritative source" in prompt + assert "Do not use pretrained knowledge" in prompt + assert "Before generating SQL, silently validate" in prompt + assert "return null for sql instead of choosing one" in prompt + + +def test_sql_correction_system_prompt_allows_null_when_ungrounded(): + prompt = get_sql_correction_system_prompt() + + assert "repair the query only when the repair can be verified" in prompt + assert "Never introduce a new schema object during repair" in prompt + assert "or null" in prompt + + def test_build_executable_schema_contract_lists_retrieved_identifiers(): contract = build_executable_schema_contract( [ From 0f5275df1a0915b1caa6cfd73592b5101e3a1b03 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 3 Aug 2026 23:36:24 +0530 Subject: [PATCH 0813/1087] Remove heuristic ask fallbacks --- .../generation/followup_sql_generation.py | 4 +- .../src/pipelines/generation/sql_answer.py | 5 +- .../pipelines/generation/sql_correction.py | 4 +- .../pipelines/generation/sql_generation.py | 4 +- .../src/pipelines/generation/utils/sql.py | 4 +- .../retrieval/db_schema_retrieval.py | 4 +- .../retrieval/preprocess_sql_data.py | 47 +- wren-ai-service/src/providers/engine/wren.py | 4 +- .../services/relationship_recommendation.py | 403 +----------------- .../generation/test_sql_answer_prompt.py | 34 ++ .../retrieval/test_preprocess_sql_data.py | 56 +++ .../test_relationship_recommendation.py | 124 ++---- .../src/apollo/server/adaptors/ibisAdaptor.ts | 2 +- .../apollo/server/resolvers/modelResolver.ts | 12 - .../pages/api/ask_task/streaming_answer.ts | 32 +- 15 files changed, 212 insertions(+), 527 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py create mode 100644 wren-ai-service/tests/pytest/pipelines/retrieval/test_preprocess_sql_data.py 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 aa0c872246..a5423b14f5 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -161,7 +161,7 @@ async def post_process( data_source: str, project_id: str | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, schema_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( @@ -221,7 +221,7 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, sql_knowledge: SqlKnowledge | None = None, schema_contracts: list[dict] | None = None, ): diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index a211c925d8..b1eae4c1cb 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -32,9 +32,10 @@ 6. Answer must be in the same language user specified. 7. Do not include ```markdown or ``` in the answer. 8. If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. -9. Use only the columns and rows provided in Data. Do not invent, duplicate, reorder, aggregate, rank, or label rows unless that operation is directly represented by the provided SQL result. +9. Use only the columns and result rows provided in Data. Do not invent, duplicate, reorder, aggregate, rank, or label rows unless that operation is directly represented by the provided SQL result. 10. If the Data has aggregate rows, summarize those exact aggregate rows instead of describing them as separate top examples. 11. If the Data is empty, state that no matching records were returned. +12. Data rows are records already mapped by column name. Answer from the record values; do not describe the underlying data structure. ### OUTPUT FORMAT @@ -47,7 +48,7 @@ SQL: {{ sql }} Data: columns: {{ sql_data.columns }} -rows: {{ sql_data.data }} +result rows: {{ sql_data.row_records }} Language: {{ language }} Current Time: {{ current_time }} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 45c3a9502e..828408fdb2 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -147,7 +147,7 @@ async def post_process( data_source: str, project_id: str | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, schema_contracts: list[dict] | None = None, ) -> dict: return await post_processor.run( @@ -203,7 +203,7 @@ async def run( project_id: str | None = None, mdl_hash: str | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, sql_knowledge: SqlKnowledge | None = None, schema_contracts: list[dict] | None = None, ): diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 0519d19c44..138dc069d6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -157,7 +157,7 @@ async def post_process( data_source: str, project_id: str | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, allow_data_preview: bool = False, schema_contracts: list[dict] | None = None, ) -> dict: @@ -218,7 +218,7 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, schema_contracts: list[dict] | None = None, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 7af6a894e9..5e92a8a81b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -256,7 +256,7 @@ async def run( replies: List[str] | List[List[str]], project_id: str | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, data_source: str = "", allow_data_preview: bool = False, schema_contracts: list[dict] | None = None, @@ -334,7 +334,7 @@ async def _classify_generation_result( generation_result: str | None, project_id: str | None = None, use_dry_plan: bool = False, - allow_dry_plan_fallback: bool = True, + allow_dry_plan_fallback: bool = False, data_source: str = "", allow_data_preview: bool = False, ) -> Dict[str, str]: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ccf4fb8a64..8d8af3e95c 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -245,14 +245,14 @@ async def active_mdl_hash( if not dbschema_count: logger.warning( "Project ID: %s, MDL hash %s has no indexed schema documents; " - "keeping hash scope to avoid stale project metadata fallback.", + "keeping hash scope to avoid stale project metadata reuse.", project_id, mdl_hash, ) elif not table_description_count: logger.info( "Project ID: %s, MDL hash %s has indexed db schema documents but no table descriptions; " - "using db schema retrieval fallback.", + "using deployed db schema documents for retrieval.", project_id, mdl_hash, ) diff --git a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py index e6dadd32d0..cf282b7899 100644 --- a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py +++ b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py @@ -1,6 +1,7 @@ import logging import sys -from typing import Dict +from copy import deepcopy +from typing import Any, Dict import tiktoken from hamilton import base @@ -14,6 +15,46 @@ ## Start of Pipeline +def _get_column_name(column: Any) -> str: + if isinstance(column, dict): + return str(column.get("name", "")) + + return str(column) + + +def _build_row_records(sql_data: Dict) -> list[dict]: + column_names = [ + column_name + for column_name in ( + _get_column_name(column) for column in sql_data.get("columns", []) + ) + if column_name + ] + + if not column_names: + return [] + + row_records = [] + for row in sql_data.get("data", []): + if isinstance(row, dict): + row_records.append( + {column_name: row.get(column_name) for column_name in column_names} + ) + continue + + if not isinstance(row, (list, tuple)): + row = [row] + + row_records.append( + { + column_name: row[index] if index < len(row) else None + for index, column_name in enumerate(column_names) + } + ) + + return row_records + + @observe(capture_input=False, capture_output=False) def preprocess( sql_data: Dict, @@ -46,6 +87,9 @@ def reduce_data_size(data: list, reduction_step: int = 50) -> list: return returned_data + sql_data = deepcopy(sql_data) + sql_data["row_records"] = _build_row_records(sql_data) + _token_count = len(encoding.encode(str(sql_data))) num_rows_used_in_llm = len(sql_data.get("data", [])) iteration = 0 @@ -62,6 +106,7 @@ def reduce_data_size(data: list, reduction_step: int = 50) -> list: data = sql_data.get("data", []) sql_data["data"] = reduce_data_size(data) + sql_data["row_records"] = _build_row_records(sql_data) num_rows_used_in_llm = len(sql_data.get("data", [])) _token_count = len(encoding.encode(str(sql_data))) logger.info(f"Token count: {_token_count}") diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 8a4eab6479..5c8dbc9902 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -145,7 +145,7 @@ async def dry_plan( data_source: str, project_id: str | None = None, timeout: float = settings.engine_timeout, - allow_fallback: bool = True, + allow_fallback: bool = False, **kwargs, ) -> Tuple[bool, str]: data = { @@ -250,7 +250,7 @@ async def dry_plan( sql: str, data_source: str, timeout: float = settings.engine_timeout, - allow_fallback: bool = True, + allow_fallback: bool = False, **kwargs, ) -> Tuple[bool, str]: api_endpoint = f"{self._endpoint}/v3/connector/{data_source}/dry-plan" diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py index 91dd309f31..0e29bc0548 100644 --- a/wren-ai-service/src/web/v1/services/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/services/relationship_recommendation.py @@ -1,7 +1,6 @@ import asyncio import logging -import re -from typing import Any, Dict, Literal, Optional +from typing import Dict, Literal, Optional import orjson from cachetools import TTLCache @@ -45,386 +44,6 @@ def __init__( ) self._generation_timeout_seconds = generation_timeout_seconds - def _normalize_identifier(self, value: Any) -> str: - text = "" if value is None else str(value) - text = re.sub(r"[^a-zA-Z0-9]", "", text).lower() - return text[:-1] if text.endswith("s") else text - - def _identifier_tokens(self, value: Any) -> list[str]: - text = "" if value is None else str(value) - text = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", text) - return [ - self._normalize_identifier(token) - for token in re.split(r"[^a-zA-Z0-9]+", text) - if token - ] - - def _model_aliases(self, model: dict) -> set[str]: - aliases: set[str] = set() - raw_values = [ - model.get("name"), - model.get("properties", {}).get("displayName"), - model.get("tableReference", {}).get("table"), - ] - - for value in raw_values: - normalized = self._normalize_identifier(value) - if normalized: - aliases.add(normalized) - - tokens = self._identifier_tokens(value) - if tokens: - aliases.add(tokens[-1]) - aliases.add("".join(tokens)) - - return aliases - - def _model_columns(self, model: dict) -> list[dict]: - return [ - column - for column in model.get("columns", []) or [] - if column.get("name") and not column.get("relationship") - ] - - def _humanize_identifier(self, value: Any) -> str: - text = "" if value is None else str(value) - text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text) - text = re.sub(r"[_\-.]+", " ", text) - text = re.sub(r"\b(id|pk|fk)\b", "", text, flags=re.IGNORECASE) - text = re.sub(r"\s+", " ", text).strip() - replacements = { - "dept": "department", - "emp": "employee", - "org": "organization", - "cust": "customer", - "prod": "product", - "dim": "dimension", - "fact": "fact", - } - parts = [replacements.get(part.lower(), part) for part in text.split()] - if len(parts) > 1 and parts[0].lower() in { - "dbo", - "public", - "stage", - "staging", - "tbl", - }: - parts = parts[1:] - if len(parts) > 1 and parts[0].lower() == "q": - parts = parts[1:] - text = " ".join(parts) - return text.lower() if text else "record" - - def _singularize_label(self, label: str) -> str: - if label.endswith("ies") and len(label) > 3: - return f"{label[:-3]}y" - if label.endswith(("ses", "xes", "zes", "ches", "shes")): - return label[:-2] - if ( - label.endswith("s") - and not label.endswith(("ss", "us", "is", "sales", "series")) - ): - return label[:-1] - return label - - def _model_label(self, model: dict) -> str: - properties = model.get("properties") or {} - return self._singularize_label( - self._humanize_identifier( - properties.get("displayName") - or model.get("tableReference", {}).get("table") - or model.get("name") - ) - ) - - def _pluralize_label(self, label: str) -> str: - if label.endswith("y") and label[-2:] not in {"ay", "ey", "iy", "oy", "uy"}: - return f"{label[:-1]}ies" - if label.endswith(("s", "x", "z", "ch", "sh")): - return f"{label}es" - return f"{label}s" - - def _relationship_description( - self, - from_model: dict, - to_model: dict, - relationship_type: str, - ) -> str: - from_label = self._model_label(from_model) - to_label = self._model_label(to_model) - from_plural = self._pluralize_label(from_label) - to_plural = self._pluralize_label(to_label) - - if relationship_type == "ONE_TO_ONE": - return ( - f"Each {from_label} is linked to one matching {to_label}, " - "connecting details that describe the same business record." - ) - if relationship_type == "ONE_TO_MANY": - return ( - f"Each {from_label} can be associated with multiple {to_plural}, " - f"supporting analysis of {to_plural} by {from_label}." - ) - return ( - f"Each {from_label} belongs to one {to_label}, " - f"so {from_plural} can be grouped and analyzed by {to_label}." - ) - - def _description_is_meaningful(self, value: Any) -> bool: - if not isinstance(value, str): - return False - - text = value.strip() - if len(text) < 24: - return False - - technical_patterns = [ - r"\bappears to reference\b", - r"\breferences\b", - r"\bforeign key\b", - r"\bprimary key\b", - r"\w+\.\w+", - ] - return not any( - re.search(pattern, text, flags=re.IGNORECASE) - for pattern in technical_patterns - ) - - def _ensure_relationship_descriptions(self, response: dict, mdl: dict) -> dict: - relationships = response.get("relationships") - if not isinstance(relationships, list): - return response - - models_by_name = { - model.get("name"): model - for model in mdl.get("models", []) or [] - if model.get("name") - } - normalized_relationships = [] - for relationship in relationships: - if not isinstance(relationship, dict): - continue - - from_model = models_by_name.get(relationship.get("fromModel")) - to_model = models_by_name.get(relationship.get("toModel")) - if not from_model or not to_model: - normalized_relationships.append(relationship) - continue - - reason = relationship.get("reason") - if not self._description_is_meaningful(reason): - relationship = { - **relationship, - "reason": self._relationship_description( - from_model, to_model, relationship.get("type", "MANY_TO_ONE") - ), - } - - normalized_relationships.append(relationship) - - return {**response, "relationships": normalized_relationships} - - def _primary_key(self, model: dict) -> Optional[str]: - primary_key = model.get("primaryKey") - columns = self._model_columns(model) - if primary_key and any(column.get("name") == primary_key for column in columns): - return primary_key - - model_aliases = self._model_aliases(model) - for column in columns: - normalized_column = self._normalize_identifier(column.get("name")) - if normalized_column == "id" or normalized_column in { - f"{alias}id" for alias in model_aliases - }: - return column.get("name") - - return None - - def _column_is_primary_key(self, model: dict, column_name: str) -> bool: - primary_key = self._primary_key(model) - if primary_key and column_name == primary_key: - return True - - return False - - def _fallback_relationship_type( - self, from_model: dict, from_column: str, to_model: dict, to_column: str - ) -> str: - from_is_pk = self._column_is_primary_key(from_model, from_column) - to_is_pk = self._column_is_primary_key(to_model, to_column) - if from_is_pk and to_is_pk: - return "ONE_TO_ONE" - if from_is_pk and not to_is_pk: - return "ONE_TO_MANY" - return "MANY_TO_ONE" - - def _relationship_signature( - self, - from_model_name: str, - from_column: str, - to_model_name: str, - to_column: str, - ) -> tuple[str, str, str, str]: - return (from_model_name, from_column, to_model_name, to_column) - - def _relationship_pair_signature( - self, - from_model_name: str, - from_column: str, - to_model_name: str, - to_column: str, - ) -> tuple[tuple[str, str], tuple[str, str]]: - left = (from_model_name, from_column) - right = (to_model_name, to_column) - return tuple(sorted([left, right])) - - def _existing_relationship_signatures( - self, mdl: dict - ) -> tuple[ - set[tuple[str, str, str, str]], set[tuple[tuple[str, str], tuple[str, str]]] - ]: - direct_signatures = set() - pair_signatures = set() - - for relationship in mdl.get("relationships", []) or []: - models = relationship.get("models", []) or [] - condition = relationship.get("condition", "") - if len(models) < 2 or not condition: - continue - - match = re.match( - r"\s*([^.=\s]+)\.([^.=\s]+)\s*=\s*([^.=\s]+)\.([^.=\s]+)\s*", - condition, - ) - if not match: - continue - - left_model, left_column, right_model, right_column = match.groups() - direct_signatures.add( - self._relationship_signature( - left_model, left_column, right_model, right_column - ) - ) - direct_signatures.add( - self._relationship_signature( - right_model, right_column, left_model, left_column - ) - ) - pair_signatures.add( - self._relationship_pair_signature( - left_model, left_column, right_model, right_column - ) - ) - - return direct_signatures, pair_signatures - - def _fallback_relationships(self, mdl: dict) -> dict: - models = mdl.get("models", []) or [] - existing, existing_pairs = self._existing_relationship_signatures(mdl) - seen = set(existing) - seen_pairs = set(existing_pairs) - candidates = [] - - model_lookup = {} - primary_keys = {} - for model in models: - primary_keys[model.get("name")] = self._primary_key(model) - for alias in self._model_aliases(model): - model_lookup.setdefault(alias, model) - - def add_candidate( - from_model: dict, - from_column: str, - to_model: dict, - to_column: str, - ): - from_model_name = from_model.get("name") - to_model_name = to_model.get("name") - if not from_model_name or not to_model_name: - return - if from_model_name == to_model_name: - return - - signature = self._relationship_signature( - from_model_name, from_column, to_model_name, to_column - ) - pair_signature = self._relationship_pair_signature( - from_model_name, from_column, to_model_name, to_column - ) - if signature in seen or pair_signature in seen_pairs: - return - - relationship_type = self._fallback_relationship_type( - from_model, from_column, to_model, to_column - ) - reason = self._relationship_description( - from_model, to_model, relationship_type - ) - - seen.add(signature) - seen_pairs.add(pair_signature) - candidates.append( - { - "name": f"{from_model_name}_{to_model_name}", - "fromModel": from_model_name, - "fromColumn": from_column, - "type": relationship_type, - "toModel": to_model_name, - "toColumn": to_column, - "reason": reason, - } - ) - - for from_model in models: - from_model_name = from_model.get("name") - if not from_model_name: - continue - - for column in self._model_columns(from_model): - from_column = column.get("name") - normalized_column = self._normalize_identifier(from_column) - target_keys = set() - if normalized_column.endswith("id") and normalized_column != "id": - target_keys.add(normalized_column[:-2]) - - for to_model in models: - to_model_name = to_model.get("name") - to_primary_key = primary_keys.get(to_model_name) - if ( - to_model_name == from_model_name - or not to_primary_key - or not self._column_is_primary_key(to_model, to_primary_key) - ): - continue - - to_primary_key_normalized = self._normalize_identifier( - to_primary_key - ) - if ( - normalized_column != "id" - and normalized_column == to_primary_key_normalized - ): - add_candidate( - to_model, - to_primary_key, - from_model, - from_column, - ) - - for target_key in target_keys: - to_model = model_lookup.get(target_key) - if not to_model or to_model.get("name") == from_model_name: - continue - - to_model_name = to_model.get("name") - to_column = primary_keys.get(to_model_name) - if not to_column: - continue - - add_candidate(from_model, from_column, to_model, to_column) - - return {"relationships": candidates} - def _handle_exception( self, input: Input, @@ -467,23 +86,15 @@ async def recommend(self, request: Input, **kwargs) -> Resource: timeout=self._generation_timeout_seconds, ) response = resp.get("validated") - if not response or ( - "relationships" in response and not response.get("relationships") - ): - logger.warning( - "Configured LLM returned empty relationship recommendations; " - "returning metadata-based fallback relationships." + if response is None: + raise ValueError( + "Relationship recommendation pipeline returned no validated response" ) - response = self._fallback_relationships(mdl_dict) except TimeoutError: - logger.warning( - "Relationship recommendation LLM call timed out after %s seconds; " - "returning metadata-based fallback relationships.", - self._generation_timeout_seconds, + raise TimeoutError( + "Relationship recommendation LLM call timed out after " + f"{self._generation_timeout_seconds} seconds" ) - response = self._fallback_relationships(mdl_dict) - - response = self._ensure_relationship_descriptions(response, mdl_dict) self._cache[request.id] = self.Resource( id=request.id, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py new file mode 100644 index 0000000000..8150ca6c4a --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py @@ -0,0 +1,34 @@ +from haystack.components.builders.prompt_builder import PromptBuilder + +from src.pipelines.generation.sql_answer import prompt, sql_to_answer_user_prompt_template + + +def test_sql_answer_prompt_uses_column_mapped_result_rows(): + result = prompt( + query="What is the cost per unit of production volumes for each supplier?", + sql="SELECT supplier_name, manufacturing_cost_per_unit FROM SupplierManufacturing", + sql_data={ + "columns": [ + {"name": "supplier_name", "type": "varchar"}, + {"name": "manufacturing_cost_per_unit", "type": "double"}, + ], + "data": [["Supplier 1", 0.06]], + "row_records": [ + { + "supplier_name": "Supplier 1", + "manufacturing_cost_per_unit": 0.06, + } + ], + }, + language="English", + current_time="2026-08-03T00:00:00", + custom_instruction="", + prompt_builder=PromptBuilder(template=sql_to_answer_user_prompt_template), + ) + + generated_prompt = result["prompt"] + + assert "result rows:" in generated_prompt + assert "supplier_name" in generated_prompt + assert "Supplier 1" in generated_prompt + assert "rows: [[" not in generated_prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_preprocess_sql_data.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_preprocess_sql_data.py new file mode 100644 index 0000000000..b587e5e071 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_preprocess_sql_data.py @@ -0,0 +1,56 @@ +from src.pipelines.retrieval.preprocess_sql_data import preprocess + + +class _FakeEncoding: + def encode(self, value: str) -> list[str]: + return list(value) + + +def test_preprocess_maps_list_rows_to_column_named_records(): + sql_data = { + "columns": [ + {"name": "supplierid", "type": "integer"}, + {"name": "supplier_name", "type": "varchar"}, + {"name": "manufacturing_cost_per_unit", "type": "double"}, + ], + "data": [ + [2, "Supplier 1", 0.06], + [5, "Supplier 2", 0.06], + ], + } + + result = preprocess( + sql_data=sql_data, + encoding=_FakeEncoding(), + context_window_size=1000, + ) + + assert result["sql_data"]["row_records"] == [ + { + "supplierid": 2, + "supplier_name": "Supplier 1", + "manufacturing_cost_per_unit": 0.06, + }, + { + "supplierid": 5, + "supplier_name": "Supplier 2", + "manufacturing_cost_per_unit": 0.06, + }, + ] + assert "row_records" not in sql_data + + +def test_preprocess_keeps_row_records_in_sync_when_rows_are_reduced(): + sql_data = { + "columns": [{"name": "name"}, {"name": "value"}], + "data": [["a", 1], ["b", 2], ["c", 3]], + } + + result = preprocess( + sql_data=sql_data, + encoding=_FakeEncoding(), + context_window_size=1, + ) + + assert len(result["sql_data"]["row_records"]) == len(result["sql_data"]["data"]) + assert result["num_rows_used_in_llm"] == len(result["sql_data"]["data"]) diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py index 4d733ff1bd..16bea83d21 100644 --- a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py +++ b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py @@ -123,7 +123,7 @@ async def test_recommend_success(relationship_recommendation_service, mock_pipel @pytest.mark.asyncio -async def test_recommend_replaces_technical_llm_relationship_reason( +async def test_recommend_preserves_llm_relationship_reason( relationship_recommendation_service, mock_pipeline, mdl_with_project_relationship_candidate, @@ -152,8 +152,7 @@ async def test_recommend_replaces_technical_llm_relationship_reason( assert response.status == "finished" assert response.response["relationships"][0]["reason"] == ( - "Each view belongs to one project, so views can be grouped and analyzed " - "by project." + "view.project_id references project.id." ) @@ -215,7 +214,7 @@ def test_getitem_not_found(relationship_recommendation_service): @pytest.mark.asyncio -async def test_recommend_timeout_returns_fallback_relationships( +async def test_recommend_timeout_fails_without_fallback_relationships( mock_pipeline, mdl_with_project_relationship_candidate ): service = RelationshipRecommendation( @@ -234,27 +233,13 @@ async def never_finishes(**_kwargs): await service.recommend(request) response = service[request.id] - assert response.status == "finished" - assert response.response == { - "relationships": [ - { - "name": "view_project", - "fromModel": "view", - "fromColumn": "project_id", - "type": "MANY_TO_ONE", - "toModel": "project", - "toColumn": "id", - "reason": ( - "Each view belongs to one project, so views can be grouped " - "and analyzed by project." - ), - } - ] - } + assert response.status == "failed" + assert response.error.code == "OTHERS" + assert "timed out" in response.error.message @pytest.mark.asyncio -async def test_recommend_empty_llm_result_returns_fallback_relationships( +async def test_recommend_empty_llm_result_stays_empty( relationship_recommendation_service, mock_pipeline, mdl_with_project_relationship_candidate, @@ -268,13 +253,27 @@ async def test_recommend_empty_llm_result_returns_fallback_relationships( response = relationship_recommendation_service[request.id] assert response.status == "finished" - assert response.response["relationships"][0]["fromModel"] == "view" - assert response.response["relationships"][0]["toModel"] == "project" - assert response.response["relationships"][0]["type"] == "MANY_TO_ONE" + assert response.response == {"relationships": []} + + +@pytest.mark.asyncio +async def test_recommend_missing_validated_response_fails( + relationship_recommendation_service, + mock_pipeline, +): + request = RelationshipRecommendation.Input(id="test_id", mdl='{"relationships": []}') + mock_pipeline.run.return_value = {} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "failed" + assert response.error.code == "OTHERS" + assert "no validated response" in response.error.message @pytest.mark.asyncio -async def test_recommend_fallback_matches_prefixed_model_name( +async def test_recommend_does_not_create_prefixed_model_relationships( relationship_recommendation_service, mock_pipeline, mdl_with_prefixed_project_model, @@ -288,24 +287,11 @@ async def test_recommend_fallback_matches_prefixed_model_name( response = relationship_recommendation_service[request.id] assert response.status == "finished" - assert response.response["relationships"] == [ - { - "name": "dbo_view_dbo_project", - "fromModel": "dbo_view", - "fromColumn": "project_id", - "type": "MANY_TO_ONE", - "toModel": "dbo_project", - "toColumn": "id", - "reason": ( - "Each view belongs to one project, so views can be grouped " - "and analyzed by project." - ), - } - ] + assert response.response == {"relationships": []} @pytest.mark.asyncio -async def test_recommend_fallback_identifies_one_to_one_relationships( +async def test_recommend_does_not_infer_one_to_one_relationships( relationship_recommendation_service, mock_pipeline, mdl_with_one_to_one_profile_candidate, @@ -319,24 +305,11 @@ async def test_recommend_fallback_identifies_one_to_one_relationships( response = relationship_recommendation_service[request.id] assert response.status == "finished" - assert response.response["relationships"] == [ - { - "name": "profile_user", - "fromModel": "profile", - "fromColumn": "user_id", - "type": "ONE_TO_ONE", - "toModel": "user", - "toColumn": "id", - "reason": ( - "Each profile is linked to one matching user, connecting details " - "that describe the same business record." - ), - } - ] + assert response.response == {"relationships": []} @pytest.mark.asyncio -async def test_recommend_fallback_scans_all_models_and_identifies_one_to_many( +async def test_recommend_does_not_infer_shared_key_relationships( relationship_recommendation_service, mock_pipeline, mdl_with_shared_key_candidates, @@ -350,44 +323,7 @@ async def test_recommend_fallback_scans_all_models_and_identifies_one_to_many( response = relationship_recommendation_service[request.id] assert response.status == "finished" - assert response.response["relationships"] == [ - { - "name": "employees_titles", - "fromModel": "employees", - "fromColumn": "emp_no", - "type": "ONE_TO_MANY", - "toModel": "titles", - "toColumn": "emp_no", - "reason": ( - "Each employee can be associated with multiple titles, supporting " - "analysis of titles by employee." - ), - }, - { - "name": "employees_dept_emp", - "fromModel": "employees", - "fromColumn": "emp_no", - "type": "ONE_TO_MANY", - "toModel": "dept_emp", - "toColumn": "emp_no", - "reason": ( - "Each employee can be associated with multiple department employees, " - "supporting analysis of department employees by employee." - ), - }, - { - "name": "departments_dept_emp", - "fromModel": "departments", - "fromColumn": "dept_no", - "type": "ONE_TO_MANY", - "toModel": "dept_emp", - "toColumn": "dept_no", - "reason": ( - "Each department can be associated with multiple department employees, " - "supporting analysis of department employees by department." - ), - }, - ] + assert response.response == {"relationships": []} def test_setitem(relationship_recommendation_service): diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 5d421cd975..cf25c4e9b5 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -281,7 +281,7 @@ export class IbisAdaptor implements IIbisAdaptor { { headers: { 'x-wren-fallback_disable': - allowFallback === false ? 'true' : 'false', + allowFallback === true ? 'false' : 'true', }, }, ); diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 6167508b73..5267c11aa3 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -1376,13 +1376,6 @@ export class ModelResolver { }); } } catch (error) { - if (this.isDryPlanTimeout(error) && allowFallback !== false) { - logger.warn( - 'Dry plan timed out; accepting generated Wren SQL without native rewrite', - ); - return true; - } - throw error; } @@ -1497,11 +1490,6 @@ export class ModelResolver { return value; } - private isDryPlanTimeout(error: unknown): boolean { - const errorMessage = JSON.stringify(error ?? '').toLowerCase(); - return errorMessage.includes('timeout') || errorMessage.includes('timed out'); - } - // validate view name private async validateViewName( viewDisplayName: string, diff --git a/wren-ui/src/pages/api/ask_task/streaming_answer.ts b/wren-ui/src/pages/api/ask_task/streaming_answer.ts index 6a5f604159..d63a083a59 100644 --- a/wren-ui/src/pages/api/ask_task/streaming_answer.ts +++ b/wren-ui/src/pages/api/ask_task/streaming_answer.ts @@ -47,13 +47,6 @@ const parseSSEMessages = (chunk: Buffer): string[] => { }); }; -const buildFallbackAnswer = (question: string) => - [ - `I found results for: **${question}**.`, - '', - 'The result table below contains the data returned from the active datasource. Use the visible fields and rows to review the detailed records, and switch to the chart tab when a visualization is available.', - ].join('\n'); - export default async function handler( req: NextApiRequest, res: NextApiResponse, @@ -105,10 +98,31 @@ export default async function handler( stream.on('end', () => { streamEnded = true; - const finalContent = - contentMap.getContent(queryId)?.trim() || buildFallbackAnswer(response.question); + const finalContent = contentMap.getContent(queryId)?.trim(); res.write(`data: ${JSON.stringify({ done: true })}\n\n`); res.end(); + if (!finalContent) { + askingService + .changeThreadResponseAnswerDetailStatus( + Number(responseId), + ThreadResponseAnswerStatus.FAILED, + ) + .then(() => { + console.error( + `Thread response ${responseId} answer stream ended without content`, + ); + contentMap.remove(queryId); + }) + .catch((error) => { + console.error( + 'Failed to update empty thread response answer detail status', + error, + ); + contentMap.remove(queryId); + }); + return; + } + askingService .changeThreadResponseAnswerDetailStatus( Number(responseId), From e83b36aa2593c5863a4eccac51798f93539c7107 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 00:25:47 +0530 Subject: [PATCH 0814/1087] Tighten SQL generation prompt grounding --- .../generation/followup_sql_generation.py | 1 + .../pipelines/generation/sql_correction.py | 1 + .../pipelines/generation/sql_generation.py | 1 + .../pipelines/generation/sql_regeneration.py | 1 + .../src/pipelines/generation/utils/sql.py | 4 +++- .../generation/test_sql_prompt_grounding.py | 23 +++++++++++++++++++ 6 files changed, 30 insertions(+), 1 deletion(-) 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 a5423b14f5..08f8bce41c 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -85,6 +85,7 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. +Do not answer a specific business question with a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 828408fdb2..4232db8053 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -98,6 +98,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Error Message: {{ invalid_generation_result.error }} Regenerate from the user's question and DATABASE SCHEMA only when a user question is available. Otherwise, correct the failed SQL only by using exact executable identifiers declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not copy table names, column names, functions, literals, aliases, or SQL structure from the failed SQL unless each one is declared in DATABASE SCHEMA or SQL FUNCTIONS. +Do not repair a failed query into a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 138dc069d6..96449cf2b0 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -80,6 +80,7 @@ If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. +Do not answer a specific business question with a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 193c9d33e9..c3550098a7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -103,6 +103,7 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. +Do not regenerate into a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5e92a8a81b..c4eaf50b00 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -529,6 +529,7 @@ async def _classify_generation_result( - If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. - If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. +- Never use SELECT * or table.*. Always select explicit deployed schema columns that are needed to answer the user's question. If no specific output columns can be grounded for the question, return null for sql. """ @@ -536,7 +537,7 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. -- ONLY USE "*" if the user query asks for all the columns of a table. +- Never use "*" in the SELECT list. Select explicit deployed schema columns needed for the question. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! @@ -556,6 +557,7 @@ async def _classify_generation_result( - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. +- Do not satisfy a filtered, time-bounded, metric, or business-specific request by returning an unfiltered table scan. If the requested filter, time field, metric, or business concept is not grounded in DATABASE SCHEMA, return null for sql. - Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. - Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. - DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index 93676ab6d2..d80ec49b60 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -5,6 +5,10 @@ prompt as build_sql_correction_prompt, sql_correction_user_prompt_template, ) +from src.pipelines.generation.followup_sql_generation import ( + prompt as build_followup_sql_generation_prompt, + text_to_sql_with_followup_user_prompt_template, +) from src.pipelines.generation.sql_generation import ( get_sql_generation_system_prompt, prompt as build_sql_generation_prompt, @@ -24,6 +28,8 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "Do not use pretrained knowledge" in prompt assert "Before generating SQL, silently validate" in prompt assert "return null for sql instead of choosing one" in prompt + assert "Never use SELECT * or table.*" in prompt + assert "Do not satisfy a filtered, time-bounded, metric, or business-specific request by returning an unfiltered table scan" in prompt def test_sql_correction_system_prompt_allows_null_when_ungrounded(): @@ -32,6 +38,7 @@ def test_sql_correction_system_prompt_allows_null_when_ungrounded(): assert "repair the query only when the repair can be verified" in prompt assert "Never introduce a new schema object during repair" in prompt assert "or null" in prompt + assert "Never use SELECT * or table.*" in prompt def test_build_executable_schema_contract_lists_retrieved_identifiers(): @@ -89,6 +96,20 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): assert "TABLE: retrieved_model" in built_prompt assert "- grouping_attribute" in built_prompt assert "- numeric_measure" in built_prompt + assert "Do not answer a specific business question with a broad table scan" in built_prompt + + +def test_followup_sql_generation_prompt_rejects_broad_table_scan_generation(): + result = build_followup_sql_generation_prompt( + query="show recent refunds", + documents=[], + sql_generation_reasoning="", + prompt_builder=PromptBuilder( + template=text_to_sql_with_followup_user_prompt_template + ), + ) + + assert "Do not answer a specific business question with a broad table scan" in result["prompt"] def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): @@ -107,6 +128,7 @@ def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): assert "User's Question: summarize the records" in built_prompt assert "Failed SQL: SELECT 1" in built_prompt assert "DIAGNOSTIC CONTEXT" in built_prompt + assert "Do not repair a failed query into a broad table scan" in built_prompt def test_sql_correction_prompt_includes_executable_schema_contract(): @@ -154,3 +176,4 @@ def test_sql_regeneration_prompt_includes_executable_schema_contract(): assert "ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION" in built_prompt assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt assert "TABLE: retrieved_model" in built_prompt + assert "Do not regenerate into a broad table scan" in built_prompt From ee58415e7377d1f1c7d811be048cc42808864d8f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 01:04:23 +0530 Subject: [PATCH 0815/1087] Revert "Tighten SQL generation prompt grounding" This reverts commit e83b36aa2593c5863a4eccac51798f93539c7107. --- .../generation/followup_sql_generation.py | 1 - .../pipelines/generation/sql_correction.py | 1 - .../pipelines/generation/sql_generation.py | 1 - .../pipelines/generation/sql_regeneration.py | 1 - .../src/pipelines/generation/utils/sql.py | 4 +--- .../generation/test_sql_prompt_grounding.py | 23 ------------------- 6 files changed, 1 insertion(+), 30 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 08f8bce41c..a5423b14f5 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -85,7 +85,6 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. -Do not answer a specific business question with a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 4232db8053..828408fdb2 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -98,7 +98,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Error Message: {{ invalid_generation_result.error }} Regenerate from the user's question and DATABASE SCHEMA only when a user question is available. Otherwise, correct the failed SQL only by using exact executable identifiers declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not copy table names, column names, functions, literals, aliases, or SQL structure from the failed SQL unless each one is declared in DATABASE SCHEMA or SQL FUNCTIONS. -Do not repair a failed query into a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 96449cf2b0..138dc069d6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -80,7 +80,6 @@ If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. -Do not answer a specific business question with a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index c3550098a7..193c9d33e9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -103,7 +103,6 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. -Do not regenerate into a broad table scan. Select the exact output columns, filters, groupings, measures, joins, and ordering needed for the question from DATABASE SCHEMA. Return null for sql if those required parts are not grounded. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c4eaf50b00..5e92a8a81b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -529,7 +529,6 @@ async def _classify_generation_result( - If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. - If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. -- Never use SELECT * or table.*. Always select explicit deployed schema columns that are needed to answer the user's question. If no specific output columns can be grounded for the question, return null for sql. """ @@ -537,7 +536,7 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. -- Never use "*" in the SELECT list. Select explicit deployed schema columns needed for the question. +- ONLY USE "*" if the user query asks for all the columns of a table. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! @@ -557,7 +556,6 @@ async def _classify_generation_result( - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. -- Do not satisfy a filtered, time-bounded, metric, or business-specific request by returning an unfiltered table scan. If the requested filter, time field, metric, or business concept is not grounded in DATABASE SCHEMA, return null for sql. - Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. - Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. - DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index d80ec49b60..93676ab6d2 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -5,10 +5,6 @@ prompt as build_sql_correction_prompt, sql_correction_user_prompt_template, ) -from src.pipelines.generation.followup_sql_generation import ( - prompt as build_followup_sql_generation_prompt, - text_to_sql_with_followup_user_prompt_template, -) from src.pipelines.generation.sql_generation import ( get_sql_generation_system_prompt, prompt as build_sql_generation_prompt, @@ -28,8 +24,6 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "Do not use pretrained knowledge" in prompt assert "Before generating SQL, silently validate" in prompt assert "return null for sql instead of choosing one" in prompt - assert "Never use SELECT * or table.*" in prompt - assert "Do not satisfy a filtered, time-bounded, metric, or business-specific request by returning an unfiltered table scan" in prompt def test_sql_correction_system_prompt_allows_null_when_ungrounded(): @@ -38,7 +32,6 @@ def test_sql_correction_system_prompt_allows_null_when_ungrounded(): assert "repair the query only when the repair can be verified" in prompt assert "Never introduce a new schema object during repair" in prompt assert "or null" in prompt - assert "Never use SELECT * or table.*" in prompt def test_build_executable_schema_contract_lists_retrieved_identifiers(): @@ -96,20 +89,6 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): assert "TABLE: retrieved_model" in built_prompt assert "- grouping_attribute" in built_prompt assert "- numeric_measure" in built_prompt - assert "Do not answer a specific business question with a broad table scan" in built_prompt - - -def test_followup_sql_generation_prompt_rejects_broad_table_scan_generation(): - result = build_followup_sql_generation_prompt( - query="show recent refunds", - documents=[], - sql_generation_reasoning="", - prompt_builder=PromptBuilder( - template=text_to_sql_with_followup_user_prompt_template - ), - ) - - assert "Do not answer a specific business question with a broad table scan" in result["prompt"] def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): @@ -128,7 +107,6 @@ def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): assert "User's Question: summarize the records" in built_prompt assert "Failed SQL: SELECT 1" in built_prompt assert "DIAGNOSTIC CONTEXT" in built_prompt - assert "Do not repair a failed query into a broad table scan" in built_prompt def test_sql_correction_prompt_includes_executable_schema_contract(): @@ -176,4 +154,3 @@ def test_sql_regeneration_prompt_includes_executable_schema_contract(): assert "ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION" in built_prompt assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt assert "TABLE: retrieved_model" in built_prompt - assert "Do not regenerate into a broad table scan" in built_prompt From 28ffa4e988c5af2649564711450a72ee377c9b2f Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 01:47:27 +0530 Subject: [PATCH 0816/1087] Require intent-shaped SQL generation --- .../generation/followup_sql_generation.py | 1 + .../pipelines/generation/sql_correction.py | 1 + .../pipelines/generation/sql_generation.py | 1 + .../pipelines/generation/sql_regeneration.py | 1 + .../src/pipelines/generation/utils/sql.py | 6 ++++- .../generation/test_sql_prompt_grounding.py | 23 +++++++++++++++++++ 6 files changed, 32 insertions(+), 1 deletion(-) 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 a5423b14f5..cd9dc3abaa 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -85,6 +85,7 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. +Generate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 828408fdb2..08b75bc5c7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -98,6 +98,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Error Message: {{ invalid_generation_result.error }} Regenerate from the user's question and DATABASE SCHEMA only when a user question is available. Otherwise, correct the failed SQL only by using exact executable identifiers declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not copy table names, column names, functions, literals, aliases, or SQL structure from the failed SQL unless each one is declared in DATABASE SCHEMA or SQL FUNCTIONS. +Correct into an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 138dc069d6..715a0d68bd 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -80,6 +80,7 @@ If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. +Generate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 193c9d33e9..322588c2b6 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -103,6 +103,7 @@ def get_sql_regeneration_system_prompt( User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. +Regenerate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5e92a8a81b..3ddb547c32 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -529,6 +529,8 @@ async def _classify_generation_result( - If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. - If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. +- Do not answer a specific business question with a broad table scan. The SQL shape must match the user's requested output columns, filters, groupings, measures, joins, ordering, and limits. +- For analytical or metric questions, select only the requested dimensions and measures. Use declared metric columns, calculated fields, relationship paths, and schema-grounded aggregate expressions. If the required metric components are not grounded, return null for sql instead of returning raw rows. """ @@ -536,7 +538,7 @@ async def _classify_generation_result( ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. - ONLY USE the tables and columns mentioned in the database schema. -- ONLY USE "*" if the user query asks for all the columns of a table. +- Never use "*" in the SELECT list. Select explicit deployed schema columns needed for the question. When the user asks for all records, all rows, all users, all orders, or similar, treat "all" as row scope and still select explicit columns relevant to the requested entity or metric. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! @@ -556,6 +558,7 @@ async def _classify_generation_result( - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. +- For metric-style requests, the final SELECT list must expose the requested dimension columns and measure expressions or metric fields. Do not return every raw column from a retrieved model as a substitute for the requested metric. - Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. - Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. - DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. @@ -608,6 +611,7 @@ async def _classify_generation_result( Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. Use metric columns exactly as declared in DATABASE SCHEMA. Treat dimensions as grouping/filtering fields and measures as pre-defined numeric outputs. Metric base objects and measure expressions are semantic context only; do not copy identifiers from them unless those identifiers also appear in the current DATABASE SCHEMA. +When a question asks for a measure by one or more dimensions, produce a metric-shaped result: select the dimension columns, select the requested measure or grounded expression, group by the dimensions when aggregation is needed, and order or limit only when requested or needed by the question. Do not answer a metric question by selecting every column from a base model. """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index 93676ab6d2..d992c48f8c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -5,6 +5,10 @@ prompt as build_sql_correction_prompt, sql_correction_user_prompt_template, ) +from src.pipelines.generation.followup_sql_generation import ( + prompt as build_followup_sql_generation_prompt, + text_to_sql_with_followup_user_prompt_template, +) from src.pipelines.generation.sql_generation import ( get_sql_generation_system_prompt, prompt as build_sql_generation_prompt, @@ -24,6 +28,8 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "Do not use pretrained knowledge" in prompt assert "Before generating SQL, silently validate" in prompt assert "return null for sql instead of choosing one" in prompt + assert "Never use \"*\" in the SELECT list" in prompt + assert "For metric-style requests" in prompt def test_sql_correction_system_prompt_allows_null_when_ungrounded(): @@ -32,6 +38,7 @@ def test_sql_correction_system_prompt_allows_null_when_ungrounded(): assert "repair the query only when the repair can be verified" in prompt assert "Never introduce a new schema object during repair" in prompt assert "or null" in prompt + assert "Never use \"*\" in the SELECT list" in prompt def test_build_executable_schema_contract_lists_retrieved_identifiers(): @@ -89,6 +96,20 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): assert "TABLE: retrieved_model" in built_prompt assert "- grouping_attribute" in built_prompt assert "- numeric_measure" in built_prompt + assert "Generate an intent-shaped query, not a table preview" in built_prompt + + +def test_followup_sql_generation_prompt_requires_intent_shaped_query(): + result = build_followup_sql_generation_prompt( + query="show recent refunds", + documents=[], + sql_generation_reasoning="", + prompt_builder=PromptBuilder( + template=text_to_sql_with_followup_user_prompt_template + ), + ) + + assert "Generate an intent-shaped query, not a table preview" in result["prompt"] def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): @@ -107,6 +128,7 @@ def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): assert "User's Question: summarize the records" in built_prompt assert "Failed SQL: SELECT 1" in built_prompt assert "DIAGNOSTIC CONTEXT" in built_prompt + assert "Correct into an intent-shaped query, not a table preview" in built_prompt def test_sql_correction_prompt_includes_executable_schema_contract(): @@ -154,3 +176,4 @@ def test_sql_regeneration_prompt_includes_executable_schema_contract(): assert "ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION" in built_prompt assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt assert "TABLE: retrieved_model" in built_prompt + assert "Regenerate an intent-shaped query, not a table preview" in built_prompt From 1867d4ea319e9b46dfb4b0a49b00391006ae760a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 13:56:51 +0530 Subject: [PATCH 0817/1087] Fix SQL schema retrieval and wildcard previews --- .../src/pipelines/generation/utils/sql.py | 70 ++++++++++- .../retrieval/db_schema_retrieval.py | 111 +++++++++++++++++- .../generation/test_sql_post_processor.py | 47 +++++++- .../retrieval/test_db_schema_retrieval.py | 107 +++++++++++++++++ 4 files changed, 329 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 3ddb547c32..4039c8bb52 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -6,7 +6,7 @@ import sqlparse from haystack import component from haystack.dataclasses import ChatMessage -from sqlparse.sql import Identifier, IdentifierList, Parenthesis, TokenList +from sqlparse.sql import Function, Identifier, IdentifierList, Parenthesis, TokenList from sqlparse import tokens as sqlparse_tokens from pydantic import BaseModel @@ -212,6 +212,61 @@ def _sql_statement_shape_error(sql: str | None) -> str | None: return None +def _is_select_wildcard_token(token: Any) -> bool: + if isinstance(token, Function): + return False + + if isinstance(token, IdentifierList): + return any( + _is_select_wildcard_token(identifier) + for identifier in token.get_identifiers() + ) + + if token.ttype == sqlparse_tokens.Wildcard: + return True + + if isinstance(token, Identifier): + token_text = str(token).strip() + return bool( + token_text == "*" + or token_text.endswith(".*") + or token_text.upper().startswith("DISTINCT *") + ) + + return False + + +def _select_wildcard_error(sql: str | None) -> str | None: + if not sql: + return None + + for statement in sqlparse.parse(sql): + tokens = _meaningful_tokens(statement) + in_select_list = False + for token in tokens: + if ( + token.ttype == sqlparse_tokens.Keyword.DML + and token.normalized == "SELECT" + ): + in_select_list = True + continue + + if not in_select_list: + continue + + if token.ttype == sqlparse_tokens.Keyword and token.normalized == "FROM": + in_select_list = False + continue + + if _is_select_wildcard_token(token): + return ( + "Generated SQL uses SELECT *; select explicit deployed schema " + "columns needed for the question." + ) + + return None + + def build_executable_schema_contract(schema_contracts: list[dict] | None) -> str: if not schema_contracts: return "" @@ -305,6 +360,19 @@ async def run( }, } + wildcard_error = _select_wildcard_error(cleaned_generation_result) + if wildcard_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SQL_SYNTAX", + "error": wildcard_error, + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 8d8af3e95c..c64bf8543e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,6 +1,7 @@ import ast import asyncio import logging +import re import sys from typing import Any, Optional @@ -103,6 +104,87 @@ """ +_QUERY_TERM_STOPWORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "by", + "for", + "from", + "give", + "in", + "is", + "last", + "list", + "me", + "month", + "of", + "on", + "placed", + "show", + "the", + "this", + "to", + "week", + "with", +} + + +def _normalize_terms(value: str) -> set[str]: + terms = { + term + for term in re.findall(r"[a-zA-Z0-9]+", value.lower()) + if len(term) >= 3 and term not in _QUERY_TERM_STOPWORDS + } + singular_terms = { + term[:-1] + for term in terms + if term.endswith("s") and len(term) > 3 + } + return terms | singular_terms + + +def _schema_semantic_text(content: dict) -> str: + parts = [ + str(content.get("name", "") or ""), + str(content.get("comment", "") or ""), + str(content.get("properties", {}) or ""), + ] + for column in content.get("columns", []): + parts.extend( + [ + str(column.get("name", "") or ""), + str(column.get("comment", "") or ""), + str(column.get("constraint", "") or ""), + ] + ) + + return " ".join(parts) + + +def _tables_matching_query_terms( + query: str, + construct_db_schemas: list[dict], +) -> set[str]: + query_terms = _normalize_terms(query) + if not query_terms: + return set() + + matching_tables = set() + for table_schema in construct_db_schemas: + if table_schema.get("type") != "TABLE": + continue + + schema_terms = _normalize_terms(_schema_semantic_text(table_schema)) + if query_terms & schema_terms: + matching_tables.add(table_schema["name"]) + + return matching_tables + + def _build_metric_ddl(content: dict) -> str: columns_ddl = [ f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" @@ -268,9 +350,10 @@ async def embedding( project_id: str = "", mdl_hash: str = "", dbschema_store: Any = None, + table_description_store: Any = None, ) -> dict: if project_id and mdl_hash and dbschema_store: - filters = { + schema_filters = { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, @@ -280,7 +363,29 @@ async def embedding( )["conditions"], ], } - if await dbschema_store.count_documents(filters=filters): + schema_count = await dbschema_store.count_documents(filters=schema_filters) + if schema_count and table_description_store: + table_description_filters = { + "operator": "AND", + "conditions": [ + { + "field": "type", + "operator": "==", + "value": "TABLE_DESCRIPTION", + }, + *build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + )["conditions"], + ], + } + table_description_count = await table_description_store.count_documents( + filters=table_description_filters + ) + if table_description_count: + return await embedder.run(query) if query else {} + + if schema_count: return {} if query: @@ -595,6 +700,7 @@ def construct_retrieval_results( filter_columns_in_tables: dict, construct_db_schemas: list[dict], dbschema_retrieval: list[Document], + query: str = "", ) -> dict[str, Any]: if filter_columns_in_tables: columns_and_tables_needed = orjson.loads( @@ -608,6 +714,7 @@ def construct_retrieval_results( reformated_json[table["table_name"]] = table["table_contents"] columns_and_tables_needed = reformated_json tables = set(columns_and_tables_needed.keys()) + tables.update(_tables_matching_query_terms(query, construct_db_schemas)) retrieval_results = [] has_calculated_field = False has_metric = False diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index db6eb1a9ab..fd408fd515 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -180,14 +180,55 @@ async def test_sql_post_processor_allows_tables_inside_schema_contract(): engine = CapturingEngine() result = await SQLGenPostProcessor(engine).run( - ['{"sql": "SELECT * FROM supported_model"}'], + ['{"sql": "SELECT id FROM supported_model"}'], project_id="project-id", - schema_contracts=[{"table_name": "supported_model", "column_names": []}], + schema_contracts=[{"table_name": "supported_model", "column_names": ["id"]}], + ) + + assert engine.executed is True + assert result["valid_generation_result"] == { + "sql": "SELECT id FROM supported_model", + "correlation_id": "valid-correlation", + } + assert result["invalid_generation_result"] == {} + + +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_select_wildcard_inside_schema_contract(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT supported_model.* FROM supported_model"}'], + project_id="project-id", + schema_contracts=[{"table_name": "supported_model", "column_names": ["id"]}], + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT supported_model.* FROM supported_model", + "original_sql": "SELECT supported_model.* FROM supported_model", + "type": "SQL_SYNTAX", + "error": "Generated SQL uses SELECT *; select explicit deployed schema columns needed for the question.", + "correlation_id": "", + } + + +@pytest.mark.asyncio +async def test_sql_post_processor_allows_count_star(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT COUNT(*) AS total_records FROM supported_model"}'], + project_id="project-id", + schema_contracts=[ + {"table_name": "supported_model", "column_names": ["total_records"]} + ], ) assert engine.executed is True assert result["valid_generation_result"] == { - "sql": "SELECT * FROM supported_model", + "sql": "SELECT COUNT(*) AS total_records FROM supported_model", "correlation_id": "valid-correlation", } assert result["invalid_generation_result"] == {} diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 148a885a0f..38095c29b4 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -85,6 +85,44 @@ async def run(self, query): ] +@pytest.mark.asyncio +async def test_embedding_uses_table_description_search_when_deploy_descriptions_exist(): + class Embedder: + def __init__(self): + self.query = None + + async def run(self, query): + self.query = query + return {"embedding": [1.0]} + + schema_store = StoreCounter(count=1) + table_description_store = StoreCounter(count=1) + embedder = Embedder() + + result = await embedding( + query="show orders from last month", + embedder=embedder, + histories=[], + project_id="project-1", + mdl_hash="deploy-1", + dbschema_store=schema_store, + table_description_store=table_description_store, + ) + + assert result == {"embedding": [1.0]} + assert embedder.query == "show orders from last month" + assert table_description_store.filters == [ + { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-1"}, + ], + } + ] + + def test_column_pruning_prompt_uses_current_query_without_history_text(): result = build_column_selection_prompt( query="current request", @@ -1146,6 +1184,75 @@ def test_construct_retrieval_results_uses_full_columns_for_sql_generation(): ] +def test_construct_retrieval_results_adds_query_matching_tables_after_pruning(): + def order_schema(name): + return { + "type": "TABLE", + "name": name, + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "order_id", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "regional_orders", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["order_id"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + order_schema("regional_orders"), + order_schema("archived_orders"), + { + "type": "TABLE", + "name": "customers", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "customer_id", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + }, + ], + dbschema_retrieval=[], + query="show orders from last month", + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "regional_orders", + "archived_orders", + ] + + def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): From 7f0ae9289978087bcb925745546643952244d426 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 14:23:46 +0530 Subject: [PATCH 0818/1087] Limit SQL retrieval expansion and joins --- .../src/pipelines/generation/utils/sql.py | 7 +- .../retrieval/db_schema_retrieval.py | 48 ++++++------- .../generation/test_sql_prompt_grounding.py | 2 + .../retrieval/test_db_schema_retrieval.py | 69 ++++++++++++++++++- 4 files changed, 97 insertions(+), 29 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4039c8bb52..2eb9de2353 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -579,11 +579,12 @@ async def _classify_generation_result( - If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. -- When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- Prefer a single table, view, or metric that already contains the requested fields. Do not join tables just because they were retrieved together. +- When using multiple tables to combine fields into the same output row, join only through the exact FOREIGN KEY constraints shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, return null for sql or use one schema object that already contains the requested fields. - If multiple semantic interpretations exist and the retrieved context does not make one interpretation authoritative, return null for sql instead of choosing one. - When the same requested result can be answered from multiple schema objects with compatible columns or metrics, include all relevant schema objects by combining separate result rows with UNION ALL instead of choosing only one object. - Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. -- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. +- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and an exact relationship path. Do not invent join predicates from similar column names. - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. @@ -609,7 +610,7 @@ async def _classify_generation_result( - Never use "*" in the SELECT list. Select explicit deployed schema columns needed for the question. When the user asks for all records, all rows, all users, all orders, or similar, treat "all" as row scope and still select explicit columns relevant to the requested entity or metric. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - DON'T INCLUDE comments in the generated SQL query. -- YOU MUST USE "JOIN" if you choose columns from multiple tables! +- Use JOIN only when selected columns come from multiple tables and DATABASE SCHEMA declares the exact FOREIGN KEY relationship needed for the join. Do not invent join predicates from similar-looking column names. - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index c64bf8543e..5250e57601 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -148,21 +148,7 @@ def _normalize_terms(value: str) -> set[str]: def _schema_semantic_text(content: dict) -> str: - parts = [ - str(content.get("name", "") or ""), - str(content.get("comment", "") or ""), - str(content.get("properties", {}) or ""), - ] - for column in content.get("columns", []): - parts.extend( - [ - str(column.get("name", "") or ""), - str(column.get("comment", "") or ""), - str(column.get("constraint", "") or ""), - ] - ) - - return " ".join(parts) + return str(content.get("name", "") or "") def _tables_matching_query_terms( @@ -515,6 +501,22 @@ def _related_table_names(documents: list[Document], visited: set[str]) -> list[s return related_names + def _document_key(document: Document) -> tuple[str, str]: + return document.meta.get("name", ""), document.content or "" + + def _extend_unique_documents( + target: list[Document], + source: list[Document], + seen: set[tuple[str, str]], + ) -> None: + for document in source: + key = _document_key(document) + if key in seen: + continue + + seen.add(key) + target.append(document) + tables = table_retrieval.get("documents", []) table_names = [] for table in tables: @@ -524,6 +526,7 @@ def _related_table_names(documents: list[Document], visited: set[str]) -> list[s table_names.append(table_name) documents = [] + seen_documents = set() if not table_names and not (embedding and embedding.get("embedding")): results = await dbschema_retriever.run( query_embedding=[], @@ -536,20 +539,19 @@ def _related_table_names(documents: list[Document], visited: set[str]) -> list[s query_embedding=embedding.get("embedding"), filters=_base_filters(), ) - documents.extend(results["documents"]) + _extend_unique_documents(documents, results["documents"], seen_documents) for document in results["documents"]: table_name = _document_name(document) if table_name and table_name not in table_names: table_names.append(table_name) visited = set(table_names) - pending = list(table_names) - while pending: - current_names = pending - pending = [] - current_documents = await _fetch_by_names(current_names) - documents.extend(current_documents) - pending.extend(_related_table_names(current_documents, visited)) + current_documents = await _fetch_by_names(table_names) + _extend_unique_documents(documents, current_documents, seen_documents) + + related_names = _related_table_names(current_documents, visited) + related_documents = await _fetch_by_names(related_names) + _extend_unique_documents(documents, related_documents, seen_documents) return documents diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index d992c48f8c..b1cca103a0 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -30,6 +30,8 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "return null for sql instead of choosing one" in prompt assert "Never use \"*\" in the SELECT list" in prompt assert "For metric-style requests" in prompt + assert "Do not join tables just because they were retrieved together" in prompt + assert "Do not invent join predicates from similar column names" in prompt def test_sql_correction_system_prompt_allows_null_when_ungrounded(): diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 38095c29b4..ec089d9be6 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -542,7 +542,7 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_dbschema_retrieval_expands_declared_relationships(): +async def test_dbschema_retrieval_expands_direct_declared_relationships(): selected_model = "model_a" related_model = "model_b" downstream_model = "model_c" @@ -661,13 +661,12 @@ async def run(self, query_embedding, filters): embedding={}, ) - assert retriever.calls == [[selected_model], [related_model], [downstream_model]] + assert retriever.calls == [[selected_model], [related_model]] assert [document.meta["name"] for document in documents] == [ selected_model, selected_model, related_model, related_model, - downstream_model, ] @@ -1253,6 +1252,70 @@ def order_schema(name): ] +def test_construct_retrieval_results_does_not_add_column_only_term_matches(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "regional_orders", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["order_id"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "regional_orders", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "order_id", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + }, + { + "type": "TABLE", + "name": "fact_sales", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "country", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + }, + ], + dbschema_retrieval=[], + query="show orders from country france", + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "regional_orders" + ] + + def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): From 6c5e77fd1e7e4848a10a1bc66cf89b187999eef5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 14:59:39 +0530 Subject: [PATCH 0819/1087] Reject unshaped SQL table previews --- .../generation/followup_sql_generation.py | 2 + .../pipelines/generation/sql_correction.py | 2 + .../pipelines/generation/sql_generation.py | 2 + .../pipelines/generation/sql_regeneration.py | 2 + .../src/pipelines/generation/utils/sql.py | 117 +++++++++++++++++ .../generation/test_sql_post_processor.py | 121 ++++++++++++++++++ 6 files changed, 246 insertions(+) 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 cd9dc3abaa..d132de8e62 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -160,6 +160,7 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, + query: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = False, @@ -172,6 +173,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, schema_contracts=schema_contracts, + query=query, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 08b75bc5c7..4181d364c7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -146,6 +146,7 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, + query: str | None = None, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = False, @@ -158,6 +159,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, schema_contracts=schema_contracts, + query=query, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 715a0d68bd..7b918efbfc 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -156,6 +156,7 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, + query: str, project_id: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = False, @@ -170,6 +171,7 @@ async def post_process( allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, schema_contracts=schema_contracts, + query=query, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 322588c2b6..57c9fa674c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -178,6 +178,7 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, + query: str, project_id: str | None = None, schema_contracts: list[dict] | None = None, ) -> dict: @@ -185,6 +186,7 @@ async def post_process( regenerate_sql.get("replies"), project_id=project_id, schema_contracts=schema_contracts, + query=query, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2eb9de2353..c6e0bf6cb8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, List import aiohttp @@ -20,6 +21,19 @@ logger = logging.getLogger("wren-ai-service") +_ANALYTICAL_OR_FILTER_QUERY_PATTERN = re.compile( + r"\b(" + r"per|by|group|breakdown|total|sum|count|average|avg|min|max|top|bottom|" + r"highest|lowest|first|last|from|before|after|between|since|during|" + r"today|yesterday|week|month|quarter|year|january|february|march|april|" + r"may|june|july|august|september|october|november|december" + r")\b", + re.IGNORECASE, +) + +_BROAD_TABLE_PREVIEW_COLUMN_THRESHOLD = 8 + + def _is_timeout_error(error_message: str) -> bool: if not error_message: return False @@ -267,6 +281,92 @@ def _select_wildcard_error(sql: str | None) -> str | None: return None +def _is_aggregate_item(token: Any) -> bool: + token_text = str(token).upper() + return any( + f"{function_name}(" in token_text + for function_name in ["COUNT", "SUM", "AVG", "MIN", "MAX"] + ) + + +def _select_items(statement: TokenList) -> list[Any]: + tokens = _meaningful_tokens(statement) + items = [] + in_select_list = False + + for token in tokens: + if token.ttype == sqlparse_tokens.Keyword.DML and token.normalized == "SELECT": + in_select_list = True + continue + + if not in_select_list: + continue + + if token.ttype == sqlparse_tokens.Keyword and token.normalized == "FROM": + break + + if token.ttype == sqlparse_tokens.Keyword and token.normalized == "DISTINCT": + continue + + if isinstance(token, IdentifierList): + items.extend(list(token.get_identifiers())) + else: + items.append(token) + + return [item for item in items if str(item).strip()] + + +def _has_answer_shaping_clause(statement: TokenList) -> bool: + for token in _meaningful_tokens(statement): + normalized = token.normalized + if normalized in {"WHERE", "GROUP BY", "HAVING", "ORDER BY"}: + return True + if normalized.startswith("WHERE "): + return True + + return False + + +def _table_preview_shape_error(sql: str | None, query: str | None = None) -> str | None: + if not sql: + return None + + query_has_shape = bool(query and _ANALYTICAL_OR_FILTER_QUERY_PATTERN.search(query)) + + for statement in sqlparse.parse(sql): + if not str(statement).strip().strip(";").strip(): + continue + + items = _select_items(statement) + if not items or any(_is_aggregate_item(item) for item in items): + continue + + if _has_answer_shaping_clause(statement): + continue + + referenced_tables = _collect_table_references(statement) + cte_names = _collect_cte_names(statement) + source_tables = referenced_tables - cte_names + is_single_source_scan = len(source_tables) <= 1 + + if query_has_shape and is_single_source_scan: + return ( + "Generated SQL is a table preview and does not apply the requested " + "aggregation, grouping, filter, timeframe, ranking, or ordering." + ) + + if ( + is_single_source_scan + and len(items) >= _BROAD_TABLE_PREVIEW_COLUMN_THRESHOLD + ): + return ( + "Generated SQL is a broad table preview; select only the explicit " + "columns and operations needed to answer the question." + ) + + return None + + def build_executable_schema_contract(schema_contracts: list[dict] | None) -> str: if not schema_contracts: return "" @@ -315,6 +415,7 @@ async def run( data_source: str = "", allow_data_preview: bool = False, schema_contracts: list[dict] | None = None, + query: str | None = None, ) -> dict: try: cleaned_generation_result = clean_generation_result(replies[0]) @@ -373,6 +474,22 @@ async def run( }, } + table_preview_error = _table_preview_shape_error( + cleaned_generation_result, + query=query, + ) + if table_preview_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SQL_SHAPE", + "error": table_preview_error, + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index fd408fd515..7500dda1b1 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -234,6 +234,127 @@ async def test_sql_post_processor_allows_count_star(): assert result["invalid_generation_result"] == {} +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_unshaped_analytical_table_preview(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + [ + ( + '{"sql": "SELECT dim_a, dim_b, dim_c, dim_d ' + 'FROM model_alpha LIMIT 500"}' + ) + ], + project_id="project-id", + schema_contracts=[ + { + "table_name": "model_alpha", + "column_names": ["dim_a", "dim_b", "dim_c", "dim_d"], + } + ], + query="per month by category", + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" + assert ( + result["invalid_generation_result"]["error"] + == "Generated SQL is a table preview and does not apply the requested " + "aggregation, grouping, filter, timeframe, ranking, or ordering." + ) + + +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_unfiltered_timeframe_table_preview(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT entity_id, event_date FROM model_alpha"}'], + project_id="project-id", + schema_contracts=[ + { + "table_name": "model_alpha", + "column_names": ["entity_id", "event_date"], + } + ], + query="from july 2026", + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" + + +@pytest.mark.asyncio +async def test_sql_post_processor_allows_intent_shaped_analytical_query(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + [ + ( + '{"sql": "SELECT category_col, DATE_TRUNC(' + "'month', event_date) AS period_col, COUNT(*) AS row_count " + "FROM model_alpha GROUP BY category_col, " + "DATE_TRUNC('month', event_date)\"}" + ) + ], + project_id="project-id", + schema_contracts=[ + { + "table_name": "model_alpha", + "column_names": ["category_col", "event_date"], + } + ], + query="per month by category", + ) + + assert engine.executed is True + assert result["valid_generation_result"] == { + "sql": ( + "SELECT category_col, DATE_TRUNC('month', event_date) AS period_col, " + "COUNT(*) AS row_count FROM model_alpha GROUP BY category_col, " + "DATE_TRUNC('month', event_date)" + ), + "correlation_id": "valid-correlation", + } + assert result["invalid_generation_result"] == {} + + +@pytest.mark.asyncio +async def test_sql_post_processor_allows_intent_shaped_filter_query(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + [ + ( + '{"sql": "SELECT entity_id, event_date FROM model_alpha ' + "WHERE event_date >= DATE '2026-07-01' " + "AND event_date < DATE '2026-08-01'\"}" + ) + ], + project_id="project-id", + schema_contracts=[ + { + "table_name": "model_alpha", + "column_names": ["entity_id", "event_date"], + } + ], + query="from july 2026", + ) + + assert engine.executed is True + assert result["valid_generation_result"] == { + "sql": ( + "SELECT entity_id, event_date FROM model_alpha " + "WHERE event_date >= DATE '2026-07-01' " + "AND event_date < DATE '2026-08-01'" + ), + "correlation_id": "valid-correlation", + } + assert result["invalid_generation_result"] == {} + + @pytest.mark.asyncio async def test_sql_post_processor_keeps_dry_plan_timeout_invalid_without_fallback(): result = await SQLGenPostProcessor(TimeoutEngine()).run( From 7d7350371043999691742d7b7044df3cc0559b23 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 15:08:41 +0530 Subject: [PATCH 0820/1087] Add semantic relationship retrieval context --- .../src/pipelines/generation/utils/sql.py | 11 ++ .../retrieval/db_schema_retrieval.py | 52 ++++++-- wren-ai-service/src/web/v1/services/ask.py | 3 + .../src/web/v1/services/ask_feedback.py | 3 + .../generation/test_sql_prompt_grounding.py | 5 + .../retrieval/test_db_schema_retrieval.py | 124 +++++++++++++++--- 6 files changed, 172 insertions(+), 26 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index c6e0bf6cb8..437cb4ce9a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -394,6 +394,17 @@ def build_executable_schema_contract(schema_contracts: list[dict] | None) -> str else: sections.append("COLUMNS: declared in the matching DATABASE SCHEMA DDL") + relationship_constraints = [ + constraint + for constraint in contract.get("relationship_constraints", []) + if constraint + ] + if relationship_constraints: + sections.append("RELATIONSHIPS:") + sections.extend( + f"- {constraint}" for constraint in relationship_constraints + ) + return "\n".join(sections) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 5250e57601..a7f689efee 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -147,8 +147,35 @@ def _normalize_terms(value: str) -> set[str]: return terms | singular_terms +def _relationship_constraints(content: dict) -> list[str]: + return [ + column.get("constraint", "") + for column in content.get("columns", []) + if column.get("type") == "FOREIGN_KEY" and column.get("constraint") + ] + + def _schema_semantic_text(content: dict) -> str: - return str(content.get("name", "") or "") + parts = [ + str(content.get("name", "") or ""), + str(content.get("comment", "") or ""), + str(content.get("properties", {}) or ""), + ] + + for column in content.get("columns", []): + parts.extend( + [ + str(column.get("name", "") or ""), + str(column.get("data_type", "") or ""), + str(column.get("comment", "") or ""), + str(column.get("constraint", "") or ""), + str(column.get("referenced_table", "") or ""), + str(column.get("referenced_column", "") or ""), + " ".join(str(table) for table in column.get("tables", []) or []), + ] + ) + + return " ".join(part for part in parts if part) def _tables_matching_query_terms( @@ -164,8 +191,11 @@ def _tables_matching_query_terms( if table_schema.get("type") != "TABLE": continue + table_name_terms = _normalize_terms(str(table_schema.get("name", "") or "")) schema_terms = _normalize_terms(_schema_semantic_text(table_schema)) - if query_terms & schema_terms: + direct_table_match = bool(query_terms & table_name_terms) + semantic_match_count = len(query_terms & schema_terms) + if direct_table_match or semantic_match_count >= 2: matching_tables.add(table_schema["name"]) return matching_tables @@ -225,11 +255,7 @@ def _semantic_context(content: dict, column_names: list[str]) -> str: if comment: semantic_parts.append(f"{column.get('name', '')}: {comment}") - relationship_constraints = [ - column.get("constraint", "") - for column in content.get("columns", []) - if column.get("type") == "FOREIGN_KEY" and column.get("constraint") - ] + relationship_constraints = _relationship_constraints(content) block = [ "/* WREN RETRIEVED SEMANTIC CONTEXT", @@ -534,7 +560,7 @@ def _extend_unique_documents( ) return results["documents"] - if not table_names and embedding and embedding.get("embedding"): + if embedding and embedding.get("embedding"): results = await dbschema_retriever.run( query_embedding=embedding.get("embedding"), filters=_base_filters(), @@ -609,6 +635,9 @@ def check_using_db_schemas_without_pruning( "table_ddl": ddl, "column_names": column_names, "manifest_column_names": column_names, + "relationship_constraints": _relationship_constraints( + table_schema + ), } ) if _has_calculated_field: @@ -627,6 +656,7 @@ def check_using_db_schemas_without_pruning( "table_ddl": _build_metric_ddl(content), "column_names": column_names, "manifest_column_names": column_names, + "relationship_constraints": [], } ) has_metric = True @@ -638,6 +668,7 @@ def check_using_db_schemas_without_pruning( "table_ddl": _build_view_ddl(content), "column_names": column_names, "manifest_column_names": column_names, + "relationship_constraints": [], } ) @@ -741,6 +772,9 @@ def construct_retrieval_results( "table_ddl": ddl, "column_names": column_names, "manifest_column_names": column_names, + "relationship_constraints": _relationship_constraints( + table_schema + ), } ) @@ -755,6 +789,7 @@ def construct_retrieval_results( "table_ddl": _build_metric_ddl(content), "column_names": column_names, "manifest_column_names": column_names, + "relationship_constraints": [], } ) has_metric = True @@ -766,6 +801,7 @@ def construct_retrieval_results( "table_ddl": _build_view_ddl(content), "column_names": column_names, "manifest_column_names": column_names, + "relationship_constraints": [], } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a844eebf42..bc19470cc2 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -366,6 +366,9 @@ async def ask( "column_names": document.get("manifest_column_names") or document.get("column_names") or [], + "relationship_constraints": document.get( + "relationship_constraints", [] + ), } for document in documents if document.get("table_name") diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 2911b35ed5..218c9b70bc 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -173,6 +173,9 @@ async def ask_feedback( "column_names": document.get("manifest_column_names") or document.get("column_names") or [], + "relationship_constraints": document.get( + "relationship_constraints", [] + ), } for document in documents if document.get("table_name") diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index b1cca103a0..c8c43d382b 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -49,6 +49,9 @@ def test_build_executable_schema_contract_lists_retrieved_identifiers(): { "table_name": "retrieved_model", "column_names": ["grouping_attribute", "numeric_measure"], + "relationship_constraints": [ + "FOREIGN KEY (related_id) REFERENCES related_model(id)" + ], } ] ) @@ -57,6 +60,8 @@ def test_build_executable_schema_contract_lists_retrieved_identifiers(): assert "TABLE: retrieved_model" in contract assert "- grouping_attribute" in contract assert "- numeric_measure" in contract + assert "RELATIONSHIPS:" in contract + assert "- FOREIGN KEY (related_id) REFERENCES related_model(id)" in contract def test_sql_generation_prompt_omits_sample_sql_body(): diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index ec089d9be6..a369ec7778 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -6,6 +6,7 @@ from src.pipelines.retrieval.db_schema_retrieval import ( active_mdl_hash, _build_view_ddl, + _tables_matching_query_terms, check_using_db_schemas_without_pruning, construct_retrieval_results, dbschema_retrieval, @@ -164,6 +165,51 @@ def test_table_selection_prompt_keeps_multiple_relevant_datasets(): ) +def test_table_term_matching_uses_semantic_columns_and_relationships(): + matched = _tables_matching_query_terms( + "summarize lifecycle value by segment", + [ + { + "type": "TABLE", + "name": "model_alpha", + "comment": "", + "properties": {}, + "columns": [ + { + "type": "COLUMN", + "name": "metric_col", + "data_type": "DOUBLE", + "comment": "Lifecycle value used for analysis.", + }, + { + "type": "FOREIGN_KEY", + "constraint": "FOREIGN KEY (segment_id) REFERENCES model_beta(id)", + "tables": ["model_alpha", "model_beta"], + "referenced_table": "model_beta", + "referenced_column": "id", + }, + ], + }, + { + "type": "TABLE", + "name": "model_gamma", + "comment": "", + "properties": {}, + "columns": [ + { + "type": "COLUMN", + "name": "other_col", + "data_type": "VARCHAR", + "comment": "Unrelated text.", + } + ], + }, + ], + ) + + assert matched == {"model_alpha"} + + @pytest.mark.asyncio async def test_active_mdl_hash_keeps_hash_when_deploy_documents_are_indexed(): table_store = StoreCounter(count=1) @@ -767,8 +813,9 @@ async def run(self, query_embedding, filters): @pytest.mark.asyncio -async def test_dbschema_retrieval_prefers_table_description_hits_over_schema_chunk_hits(): +async def test_dbschema_retrieval_combines_description_and_schema_semantic_hits(): described_model = "described_dataset" + semantic_model = "semantic_dataset" class Retriever: def __init__(self): @@ -782,22 +829,51 @@ async def run(self, query_embedding, filters): } ) + if query_embedding: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "semantic_value", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": semantic_model}, + ) + ] + } + + names = [ + condition["value"] + for condition in filters["conditions"][1]["conditions"] + ] + documents = [ + Document( + content=str( + { + "type": "TABLE", + "name": name, + "comment": "", + "columns": [], + "properties": {}, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": name}, + ) + for name in names + ] return { - "documents": [ - Document( - content=str( - { - "type": "TABLE", - "name": described_model, - "comment": "", - "columns": [], - "properties": {}, - "primaryKey": "", - } - ), - meta={"type": "TABLE_SCHEMA", "name": described_model}, - ) - ] + "documents": documents } retriever = Retriever() @@ -816,8 +892,15 @@ async def run(self, query_embedding, filters): embedding={"embedding": [0.25]}, ) - assert [call["query_embedding"] for call in retriever.calls] == [[]] + assert [call["query_embedding"] for call in retriever.calls] == [[0.25], []] assert retriever.calls[0]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert retriever.calls[1]["filters"] == { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, @@ -825,12 +908,17 @@ async def run(self, query_embedding, filters): "operator": "OR", "conditions": [ {"field": "name", "operator": "==", "value": described_model}, + {"field": "name", "operator": "==", "value": semantic_model}, ], }, {"field": "project_id", "operator": "==", "value": "project-1"}, ], } - assert [document.meta["name"] for document in documents] == [described_model] + assert [document.meta["name"] for document in documents] == [ + semantic_model, + described_model, + semantic_model, + ] def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): From 50f4f2ce7c129e5b0bbbd5453548144424d57b95 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 15:26:58 +0530 Subject: [PATCH 0821/1087] Tighten SQL intent grounding and retrieval cap --- .../src/pipelines/generation/utils/sql.py | 121 +++++++++++++++++- .../retrieval/db_schema_retrieval.py | 10 +- .../generation/test_sql_post_processor.py | 103 ++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 85 +++++++++++- 4 files changed, 311 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 437cb4ce9a..66a4113687 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -31,6 +31,22 @@ re.IGNORECASE, ) +_AGGREGATE_QUERY_PATTERN = re.compile( + r"\b(" + r"per|by|group|breakdown|total|sum|count|average|avg|min|max|top|bottom|" + r"highest|lowest|quantity|qty|sold" + r")\b", + re.IGNORECASE, +) +_TIME_QUERY_PATTERN = re.compile( + r"\b(" + r"today|yesterday|week|month|quarter|year|january|february|march|april|" + r"may|june|july|august|september|october|november|december" + r")\b", + re.IGNORECASE, +) +_DATE_LITERAL_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}(?:[ T]\d{2}:\d{2}:\d{2})?$") + _BROAD_TABLE_PREVIEW_COLUMN_THRESHOLD = 8 @@ -327,28 +343,70 @@ def _has_answer_shaping_clause(statement: TokenList) -> bool: return False +def _has_clause(statement: TokenList, clauses: set[str]) -> bool: + for token in _meaningful_tokens(statement): + normalized = token.normalized + if normalized in clauses: + return True + if any(normalized.startswith(f"{clause} ") for clause in clauses): + return True + + return False + + def _table_preview_shape_error(sql: str | None, query: str | None = None) -> str | None: if not sql: return None query_has_shape = bool(query and _ANALYTICAL_OR_FILTER_QUERY_PATTERN.search(query)) + query_has_aggregate_shape = bool(query and _AGGREGATE_QUERY_PATTERN.search(query)) for statement in sqlparse.parse(sql): if not str(statement).strip().strip(";").strip(): continue items = _select_items(statement) - if not items or any(_is_aggregate_item(item) for item in items): + if not items: continue - if _has_answer_shaping_clause(statement): - continue + has_aggregate_item = any(_is_aggregate_item(item) for item in items) + has_grouping = _has_clause(statement, {"GROUP BY", "HAVING"}) + has_ordering = _has_clause(statement, {"ORDER BY"}) + has_filter = _has_clause(statement, {"WHERE"}) referenced_tables = _collect_table_references(statement) cte_names = _collect_cte_names(statement) source_tables = referenced_tables - cte_names is_single_source_scan = len(source_tables) <= 1 + if ( + query_has_aggregate_shape + and is_single_source_scan + and not has_aggregate_item + and not has_grouping + and not has_ordering + ): + return ( + "Generated SQL does not apply the requested aggregation, grouping, " + "ranking, or measure calculation." + ) + + if ( + is_single_source_scan + and len(items) >= _BROAD_TABLE_PREVIEW_COLUMN_THRESHOLD + and (query_has_shape or has_filter) + ): + return ( + "Generated SQL is a broad table preview; select only the explicit " + "columns and operations needed to answer the question." + ) + + if has_aggregate_item: + continue + + if _has_answer_shaping_clause(statement): + continue + if query_has_shape and is_single_source_scan: return ( "Generated SQL is a table preview and does not apply the requested " @@ -367,6 +425,45 @@ def _table_preview_shape_error(sql: str | None, query: str | None = None) -> str return None +def _unsupported_literal_filter_error( + sql: str | None, query: str | None = None +) -> str | None: + if not sql or not query: + return None + + query_terms = { + term + for term in re.findall(r"[a-zA-Z0-9]+", query.lower()) + if len(term) >= 2 + } + query_has_timeframe = bool(_TIME_QUERY_PATTERN.search(query)) + + for statement in sqlparse.parse(sql): + for token in statement.flatten(): + if token.ttype != sqlparse_tokens.Literal.String.Single: + continue + + literal_value = str(token.value).strip("'\"") + if not literal_value: + continue + + if query_has_timeframe and _DATE_LITERAL_PATTERN.match(literal_value): + continue + + literal_terms = { + term + for term in re.findall(r"[a-zA-Z0-9]+", literal_value.lower()) + if len(term) >= 2 + } + if literal_terms and not (literal_terms & query_terms): + return ( + "Generated SQL contains string filter values that are not " + "grounded in the user's question." + ) + + return None + + def build_executable_schema_contract(schema_contracts: list[dict] | None) -> str: if not schema_contracts: return "" @@ -485,6 +582,22 @@ async def run( }, } + literal_filter_error = _unsupported_literal_filter_error( + cleaned_generation_result, + query=query, + ) + if literal_filter_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SQL_VALUE_GROUNDING", + "error": literal_filter_error, + "correlation_id": "", + }, + } + table_preview_error = _table_preview_shape_error( cleaned_generation_result, query=query, @@ -703,6 +816,7 @@ async def _classify_generation_result( - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. - Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. +- Never invent string literal filter values. Use a string value in WHERE, HAVING, CASE, or JOIN conditions only when that value is explicitly present in the user's current question or grounded by a current USER INSTRUCTION. For relative time requests, bounded date literals may be generated only from the requested timeframe and an exact date/time column. - Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. - If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. @@ -756,6 +870,7 @@ async def _classify_generation_result( - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. - For metric-style requests, the final SELECT list must expose the requested dimension columns and measure expressions or metric fields. Do not return every raw column from a retrieved model as a substitute for the requested metric. +- For aggregate, ranking, or "by" requests, do not add unrelated string filters to make the SQL look specific. If the user did not provide a filter value, leave it out. - Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. - Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. - DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index a7f689efee..a600c21ca8 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -132,6 +132,8 @@ "with", } +_MAX_SCHEMA_SEMANTIC_TABLE_CANDIDATES = 5 + def _normalize_terms(value: str) -> set[str]: terms = { @@ -565,11 +567,17 @@ def _extend_unique_documents( query_embedding=embedding.get("embedding"), filters=_base_filters(), ) - _extend_unique_documents(documents, results["documents"], seen_documents) + added_schema_table_names = 0 for document in results["documents"]: table_name = _document_name(document) if table_name and table_name not in table_names: table_names.append(table_name) + added_schema_table_names += 1 + if ( + added_schema_table_names + >= _MAX_SCHEMA_SEMANTIC_TABLE_CANDIDATES + ): + break visited = set(table_names) current_documents = await _fetch_by_names(table_names) diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index 7500dda1b1..9ea023df61 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -260,8 +260,8 @@ async def test_sql_post_processor_rejects_unshaped_analytical_table_preview(): assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" assert ( result["invalid_generation_result"]["error"] - == "Generated SQL is a table preview and does not apply the requested " - "aggregation, grouping, filter, timeframe, ranking, or ordering." + == "Generated SQL does not apply the requested aggregation, grouping, " + "ranking, or measure calculation." ) @@ -286,6 +286,73 @@ async def test_sql_post_processor_rejects_unfiltered_timeframe_table_preview(): assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_broad_preview_with_placeholder_filter(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + [ + ( + '{"sql": "SELECT dim_a, dim_b, dim_c, dim_d, dim_e, ' + "dim_f, dim_g, metric_col FROM model_alpha " + "WHERE dim_a = 'SyntheticValue1'\"}" + ) + ], + project_id="project-id", + schema_contracts=[ + { + "table_name": "model_alpha", + "column_names": [ + "dim_a", + "dim_b", + "dim_c", + "dim_d", + "dim_e", + "dim_f", + "dim_g", + "metric_col", + ], + } + ], + query="total metric by category", + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SQL_VALUE_GROUNDING" + + +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_aggregate_intent_without_aggregate_shape(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + [ + ( + '{"sql": "SELECT dim_a, dim_b FROM model_alpha ' + "WHERE dim_a = 'dim_a'\"}" + ) + ], + project_id="project-id", + schema_contracts=[ + { + "table_name": "model_alpha", + "column_names": ["dim_a", "dim_b"], + } + ], + query="total metric by dim_a", + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" + assert ( + result["invalid_generation_result"]["error"] + == "Generated SQL does not apply the requested aggregation, grouping, " + "ranking, or measure calculation." + ) + + @pytest.mark.asyncio async def test_sql_post_processor_allows_intent_shaped_analytical_query(): engine = CapturingEngine() @@ -321,6 +388,38 @@ async def test_sql_post_processor_allows_intent_shaped_analytical_query(): assert result["invalid_generation_result"] == {} +@pytest.mark.asyncio +async def test_sql_post_processor_allows_user_provided_string_filter_value(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + [ + ( + '{"sql": "SELECT entity_id, status_col FROM model_alpha ' + "WHERE status_col = 'active'\"}" + ) + ], + project_id="project-id", + schema_contracts=[ + { + "table_name": "model_alpha", + "column_names": ["entity_id", "status_col"], + } + ], + query="show active records", + ) + + assert engine.executed is True + assert result["valid_generation_result"] == { + "sql": ( + "SELECT entity_id, status_col FROM model_alpha " + "WHERE status_col = 'active'" + ), + "correlation_id": "valid-correlation", + } + assert result["invalid_generation_result"] == {} + + @pytest.mark.asyncio async def test_sql_post_processor_allows_intent_shaped_filter_query(): engine = CapturingEngine() diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index a369ec7778..fbace28595 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -808,7 +808,6 @@ async def run(self, query_embedding, filters): } assert [document.meta["name"] for document in documents] == [ semantic_model, - semantic_model, ] @@ -915,12 +914,94 @@ async def run(self, query_embedding, filters): ], } assert [document.meta["name"] for document in documents] == [ - semantic_model, described_model, semantic_model, ] +@pytest.mark.asyncio +async def test_dbschema_retrieval_caps_schema_semantic_table_rescue(): + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + + if query_embedding: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "semantic_value", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + } + ), + meta={ + "type": "TABLE_SCHEMA", + "name": f"semantic_model_{index}", + }, + ) + for index in range(25) + ] + } + + selected_names = [ + condition["value"] + for condition in filters["conditions"][1]["conditions"] + ] + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": name, + "comment": "", + "columns": [], + "properties": {}, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": name}, + ) + for name in selected_names + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=retriever, + embedding={"embedding": [0.25]}, + ) + + selected_names = [ + condition["value"] + for condition in retriever.calls[1]["filters"]["conditions"][1]["conditions"] + ] + + assert len(selected_names) == 5 + assert selected_names == [f"semantic_model_{index}" for index in range(5)] + assert [document.meta["name"] for document in documents] == selected_names + + def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): class Encoding: def encode(self, value): From 38e7c6a83e5e7409ade752c1a5c52d148dd9b47d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 15:15:12 +0530 Subject: [PATCH 0822/1087] Restore semantic context for SQL generation --- .../retrieval/db_schema_retrieval.py | 295 ++++++++++++++---- .../retrieval/test_db_schema_retrieval.py | 114 ++++++- 2 files changed, 339 insertions(+), 70 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index a600c21ca8..da7ac2b73b 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -204,15 +204,42 @@ def _tables_matching_query_terms( def _build_metric_ddl(content: dict) -> str: - columns_ddl = [ - f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + columns = [ + column for column in content["columns"] if column["data_type"].lower() != "unknown" # quick fix: filtering out UNKNOWN column type ] + context = _format_semantic_context( + { + "object_type": "metric", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in columns + ], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable analytical aggregation interface", + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "semantic_context_not_sql_identifier": column["comment"], + } + for column in columns + ], + } + ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in columns + ] return ( - f"{content['comment']}CREATE TABLE {content['name']} (\n " + f"{context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) @@ -226,96 +253,239 @@ def _content_column_names(content: dict) -> list[str]: ] -def _schema_column_names(content: dict) -> list[str]: - return [ - column["name"] - for column in content.get("columns", []) - if column.get("type") == "COLUMN" and column.get("name") - ] - - -def _identifier_catalog(table_name: str, column_names: list[str]) -> str: - columns = "\n".join(f"- {column_name}" for column_name in column_names) +def _format_semantic_context(context: dict) -> str: return ( - "/* EXECUTABLE WREN IDENTIFIER CATALOG\n" - f"table: {table_name}\n" - "columns:\n" - f"{columns}\n" - "Do not create identifiers from user wording, comments, aliases, display labels, or source metadata.\n" + "/*\n" + "WREN RETRIEVED SEMANTIC CONTEXT\n" + f"{orjson.dumps(context).decode('utf-8')}\n" + f"{_format_identifier_contract(context)}" + "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" + "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" "*/\n" + f"{_format_executable_identifier_catalog(context)}" ) -def _semantic_context(content: dict, column_names: list[str]) -> str: - table_name = content.get("name", "") - semantic_parts = [ - str(content.get("comment", "") or "").strip(), - str(content.get("properties", {}) or "").strip(), +def _format_executable_identifier_catalog(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") ] - for column in content.get("columns", []): - comment = str(column.get("comment", "") or "").strip() - if comment: - semantic_parts.append(f"{column.get('name', '')}: {comment}") + relationship_constraints = [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + relationship_constraints = ( + contract.get("relationship_constraints_use_exactly") + or relationship_constraints + ) + + lines = [ + "### EXECUTABLE WREN IDENTIFIER CATALOG ###", + "Copy SQL identifiers only from this catalog or the following DDL.", + "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"table: {table_name}") + if column_names: + lines.append("columns:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationships:") + lines.extend(f"- {constraint}" for constraint in relationship_constraints) + lines.extend( + [ + "If a needed table, column, or relationship is not listed here or declared in the following DDL, return null for sql.", + "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", + "", + ] + ) + return "\n".join(lines) - relationship_constraints = _relationship_constraints(content) - block = [ - "/* WREN RETRIEVED SEMANTIC CONTEXT", - f"sql_table_name_use_exactly: {table_name}", - "sql_column_names_use_exactly:", - *[f"- {column_name}" for column_name in column_names], +def _format_identifier_contract(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") ] - for column_name in column_names: - block.append(f"sql_column_name_use_exactly: {column_name}") + relationship_constraints = [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + relationship_constraints = ( + contract.get("relationship_constraints_use_exactly") + or relationship_constraints + ) + lines = [ + "WREN SQL IDENTIFIER CONTRACT", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"sql_table_name_use_exactly: {table_name}") + if column_names: + lines.append("sql_column_names_use_exactly:") + lines.extend(f"- {column_name}" for column_name in column_names) if relationship_constraints: - block.append("relationship_constraints_use_exactly:") - block.extend(f"- {constraint}" for constraint in relationship_constraints) + lines.append("relationship_constraints_use_exactly:") + lines.extend( + f"- {relationship_constraint}" + for relationship_constraint in relationship_constraints + ) + lines.extend( + [ + "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", + "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", + "END WREN SQL IDENTIFIER CONTRACT", + "", + ] + ) + return "\n".join(lines) - semantic_context = "\n".join(part for part in semantic_parts if part) - if semantic_context: - block.append("semantic_context_not_sql_identifiers:") - block.append(semantic_context) - block.append("*/") - return "\n".join(block) + "\n" +def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: + relationship_columns = { + column.get("column") + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + } + relationship_columns.discard(None) + return relationship_columns + + +def _included_columns( + content: dict, columns: Optional[set[str]], tables: Optional[set[str]] +) -> list[dict]: + relationship_columns = _included_relationship_columns(content, tables) + return [ + column + for column in content["columns"] + if column["type"] == "COLUMN" + and ( + not columns + or column["name"] in columns + or column["name"] in relationship_columns + or column["is_primary_key"] + ) + and ( + column["data_type"] is None + or get_engine_supported_data_type(column["data_type"]).lower() + != "unknown" + ) + ] + + +def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: + return [ + column + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + ] def _build_table_context_ddl( content: dict, - include_retrieved_semantic_context: bool = False, + columns: Optional[set[str]] = None, + tables: Optional[set[str]] = None, ) -> tuple[str, bool, bool, list[str]]: - column_names = _schema_column_names(content) + included_columns = _included_columns(content, columns, tables) + included_relationships = _included_relationships(content, tables) + column_names = [column["name"] for column in included_columns] ddl, has_calculated_field, has_json_field = build_table_ddl( content, + columns=set(column_names) if column_names else columns, + tables=tables, include_semantic_comments=False, ) - context = _identifier_catalog(content["name"], column_names) - if include_retrieved_semantic_context: - context += _semantic_context(content, column_names) + context = _format_semantic_context( + { + "object_type": "model", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": column_names, + "relationship_constraints_use_exactly": [ + relationship["constraint"] + for relationship in included_relationships + ], + }, + "semantic_context_not_sql_identifiers": { + "description": content.get("comment", ""), + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "is_primary_key": column["is_primary_key"], + "semantic_context_not_sql_identifier": column["comment"], + } + for column in included_columns + ], + "relationships": [ + { + "semantic_context_not_sql_identifier": relationship["comment"], + "sql_relationship_constraint_use_exactly": relationship[ + "constraint" + ], + "related_models_use_exactly": relationship.get("tables", []), + } + for relationship in included_relationships + ], + } + ) return context + ddl, has_calculated_field, has_json_field, column_names def _build_view_ddl(content: dict) -> str: - columns = content.get("columns", []) - column_names = _content_column_names(content) + columns = [ + column + for column in content.get("columns", []) + if column.get("name") + and str(column.get("data_type", "")).lower() != "unknown" + ] + column_names = [column["name"] for column in columns] + context = _format_semantic_context( + { + "object_type": "view", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": column_names, + }, + "semantic_context_not_sql_identifiers": { + "role": "stable virtual table interface", + "description": content.get("comment", ""), + "definition_omitted_from_executable_schema": True, + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type( + column.get("data_type") + ), + "semantic_context_not_sql_identifier": column.get("comment", ""), + } + for column in columns + ], + } + ) columns_ddl = [ f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" for column in columns - if column.get("name") - and str(column.get("data_type", "")).lower() != "unknown" ] return ( - _identifier_catalog(content["name"], column_names) - + "/* WREN RETRIEVED SEMANTIC CONTEXT\n" - + f"sql_table_name_use_exactly: {content['name']}\n" - + "sql_column_names_use_exactly:\n" - + "\n".join(f"- {column_name}" for column_name in column_names) - + "\nsemantic_context_not_sql_identifier: view definition_omitted_from_executable_schema\n" - + "*/\n" - + f"CREATE TABLE {content['name']} (\n " + f"{context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) @@ -764,10 +934,7 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: ddl, _has_calculated_field, _has_json_field, column_names = ( - _build_table_context_ddl( - table_schema, - include_retrieved_semantic_context=True, - ) + _build_table_context_ddl(table_schema) ) if _has_calculated_field: has_calculated_field = True diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index fbace28595..1357f13384 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1530,11 +1530,11 @@ def table_schema(name): assert all("table: " in schema["table_ddl"] for schema in result["db_schemas"]) assert all("columns:" in schema["table_ddl"] for schema in result["db_schemas"]) assert all( - "WREN RETRIEVED SEMANTIC CONTEXT" not in schema["table_ddl"] + "WREN RETRIEVED SEMANTIC CONTEXT" in schema["table_ddl"] for schema in result["db_schemas"] ) assert all( - "semantic_context_not_sql_identifier" not in schema["table_ddl"] + "semantic_context_not_sql_identifier" in schema["table_ddl"] for schema in result["db_schemas"] ) assert result["tokens"] > 0 @@ -1575,12 +1575,114 @@ def encode(self, value): assert "table: modeled_dataset" in table_ddl assert "columns:\n- stored_attribute" in table_ddl assert "Do not create identifiers from user wording" in table_ddl - assert "Business-facing attribute label." not in table_ddl - assert "Business-facing dataset description." not in table_ddl - assert "WREN RETRIEVED SEMANTIC CONTEXT" not in table_ddl - assert "semantic_context_not_sql_identifier" not in table_ddl + assert "Business-facing attribute label." in table_ddl + assert "Business-facing dataset description." in table_ddl + assert "WREN RETRIEVED SEMANTIC CONTEXT" in table_ddl + assert "semantic_context_not_sql_identifier" in table_ddl assert "CREATE TABLE modeled_dataset" in table_ddl assert "stored_attribute VARCHAR" in table_ddl + assert "Business-facing attribute label.CREATE TABLE" not in table_ddl + assert "Business-facing dataset description.CREATE TABLE" not in table_ddl + + +def test_retrieved_schema_keeps_physical_metadata_out_of_executable_ddl(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "Business-facing dataset description.", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "Business-facing attribute label.", + "is_primary_key": False, + } + ], + "properties": { + "tableReference": { + "catalog": "physical_catalog", + "schema": "physical_schema", + "table": "physical_table", + } + }, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + table_ddl = result["db_schemas"][0]["table_ddl"] + executable_ddl = table_ddl.split("CREATE TABLE", maxsplit=1)[1] + + assert "CREATE TABLE modeled_dataset" in table_ddl + assert "physical_catalog" not in table_ddl + assert "physical_schema" not in table_ddl + assert "physical_table" not in table_ddl + assert "Business-facing attribute label." not in executable_ddl + + +def test_metric_schema_keeps_measure_semantics_outside_executable_ddl(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "METRIC", + "name": "modeled_metric", + "comment": "Metric semantic description.", + "columns": [ + { + "type": "COLUMN", + "comment": "-- This column is a dimension\n ", + "name": "grouping_dimension", + "data_type": "VARCHAR", + }, + { + "type": "COLUMN", + "comment": ( + "-- This column is a measure\n " + "-- expression: SUM(metric_value)\n " + ), + "name": "defined_measure", + "data_type": "DOUBLE", + }, + ], + } + ), + meta={"name": "modeled_metric"}, + ) + ], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + table_ddl = result["db_schemas"][0]["table_ddl"] + executable_ddl = table_ddl.split("CREATE TABLE", maxsplit=1)[1] + + assert "object_type: metric" in table_ddl + assert "stable analytical aggregation interface" in table_ddl + assert "SUM(metric_value)" in table_ddl + assert "CREATE TABLE modeled_metric" in table_ddl + assert "grouping_dimension VARCHAR" in executable_ddl + assert "defined_measure DOUBLE" in executable_ddl + assert "SUM(metric_value)" not in executable_ddl + assert "-- This column is a measure" not in executable_ddl def test_build_table_ddl_can_render_executable_schema_without_semantic_comments(): From ae8050cff0cdef227e8209c706dda0caf964fbe0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 15:56:57 +0530 Subject: [PATCH 0823/1087] Add generic semantic role hints for SQL generation --- .../generation/followup_sql_generation.py | 1 + .../pipelines/generation/sql_correction.py | 1 + .../pipelines/generation/sql_generation.py | 1 + .../pipelines/generation/sql_regeneration.py | 1 + .../src/pipelines/generation/utils/sql.py | 12 ++- .../retrieval/db_schema_retrieval.py | 89 +++++++++++++++++++ .../generation/test_sql_prompt_grounding.py | 9 ++ .../retrieval/test_db_schema_retrieval.py | 63 +++++++++++++ 8 files changed, 175 insertions(+), 2 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 d132de8e62..f3f4ae4d18 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -86,6 +86,7 @@ If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Generate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. +When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 4181d364c7..1885ae7a5d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -99,6 +99,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Regenerate from the user's question and DATABASE SCHEMA only when a user question is available. Otherwise, correct the failed SQL only by using exact executable identifiers declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not copy table names, column names, functions, literals, aliases, or SQL structure from the failed SQL unless each one is declared in DATABASE SCHEMA or SQL FUNCTIONS. Correct into an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. +If the error says the SQL is a broad table preview, table preview, missing requested aggregation, missing requested grouping, missing timeframe, or missing ordering/ranking, rebuild the query shape from the user's question. When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 7b918efbfc..8de6a23961 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -81,6 +81,7 @@ If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. Generate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. +When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 57c9fa674c..952be44cbc 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -104,6 +104,7 @@ def get_sql_regeneration_system_prompt( Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. Regenerate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. +When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 66a4113687..aa52704bc4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -812,6 +812,7 @@ async def _classify_generation_result( - When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. - In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. - Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. +- Values under column_role_hints_not_identifiers are semantic roles only. Use them to decide whether an exact declared column can serve as a date/time field, measure, identifier, or dimension, but copy executable column names only from the columns list or DDL. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. - Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. @@ -830,7 +831,8 @@ async def _classify_generation_result( - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. - Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. -- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. +- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and the predicate can be expressed with normal SQL comparison syntax or an operation listed in SQL FUNCTIONS. Do not compare text fields to date functions. +- For explicit month/year or relative timeframe requests, prefer a bounded range predicate on one exact date_time_candidate column when available. The lower bound is inclusive and the upper bound is exclusive. Do not answer a timeframe request with an unfiltered table scan. - Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. - If SQL execution or validation fails, repair the query only when the repair can be verified using the same retrieved DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS. Never introduce a new schema object during repair. - If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. @@ -842,6 +844,8 @@ async def _classify_generation_result( - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. - Do not answer a specific business question with a broad table scan. The SQL shape must match the user's requested output columns, filters, groupings, measures, joins, ordering, and limits. - For analytical or metric questions, select only the requested dimensions and measures. Use declared metric columns, calculated fields, relationship paths, and schema-grounded aggregate expressions. If the required metric components are not grounded, return null for sql instead of returning raw rows. +- For questions asking total, count, average, minimum, maximum, ratio, per, by, top, bottom, highest, lowest, trend, month, week, year, or ranking, produce an analytical query shape: select exact dimension columns or date buckets, aggregate exact numeric_measure_candidate columns or count rows, GROUP BY every non-aggregated selected expression, ORDER BY the selected aggregate alias when ranking, and apply LIMIT only when requested. +- If the question asks for an entity list with a timeframe or filter but no metric, select only the entity identifier, relevant dimensions, and exact date/time column needed by the request; include the requested WHERE predicate. Do not select every column from the table. """ @@ -865,12 +869,16 @@ async def _classify_generation_result( - Do not convert deployed identifiers into display-friendly variants by replacing spaces with underscores, removing prefixes, changing case, shortening names, or expanding abbreviations. - For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. - For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only date/time functions and casts whose exact syntax is provided in SQL FUNCTIONS for this request. -- If the question asks for a specific or relative date, generate a bounded date/time filter only when both the exact date/time schema column and required SQL FUNCTIONS-supported operation are available. If either is missing, do not invent a field or function. +- If the question asks for a specific or relative date, generate a bounded date/time filter only when the exact date/time schema column is available and the predicate can be expressed with normal SQL comparison syntax or a SQL FUNCTIONS-supported operation. If either the column or required operation is missing, do not invent a field or function. +- When DATABASE SCHEMA includes column_role_hints_not_identifiers, use date_time_candidate, numeric_measure_candidate, identifier_candidate, and dimension_candidate roles to map the question intent to exact declared columns. These role names are never executable SQL identifiers. +- For explicit calendar month and year requests, use an inclusive lower bound and exclusive upper bound on the exact date/time column, rather than formatting the column into text. - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. - For metric-style requests, the final SELECT list must expose the requested dimension columns and measure expressions or metric fields. Do not return every raw column from a retrieved model as a substitute for the requested metric. - For aggregate, ranking, or "by" requests, do not add unrelated string filters to make the SQL look specific. If the user did not provide a filter value, leave it out. +- For total, count, average, minimum, maximum, per, by, trend, top, bottom, highest, lowest, or ranking requests, the final SQL must include the requested aggregate expression or metric field, GROUP BY required dimensions, ORDER BY required ranking expression, and LIMIT only when requested. A raw row list is not a valid answer. +- For record-list requests with a filter or timeframe, the final SQL must include the requested WHERE predicate and only the columns needed to identify and describe the matching records. - Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. - Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. - DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index da7ac2b73b..ad603029ff 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -134,6 +134,30 @@ _MAX_SCHEMA_SEMANTIC_TABLE_CANDIDATES = 5 +_DATE_TIME_TYPE_TERMS = { + "date", + "datetime", + "timestamp", + "time", +} +_NUMERIC_TYPE_TERMS = { + "bigint", + "decimal", + "double", + "float", + "int", + "integer", + "numeric", + "real", + "smallint", +} +_DATE_TIME_NAME_PATTERN = re.compile(r"(date|time|timestamp|period|month|year)", re.I) +_MEASURE_NAME_PATTERN = re.compile( + r"(amount|value|qty|quantity|count|total|sum|cost|price|rate|score|percent)", + re.I, +) +_IDENTIFIER_NAME_PATTERN = re.compile(r"(^id$|[_\s-]?id$|key|code|number|num|no$)", re.I) + def _normalize_terms(value: str) -> set[str]: terms = { @@ -228,6 +252,7 @@ def _build_metric_ddl(content: dict) -> str: "sql_column_name_use_exactly": column["name"], "data_type": get_engine_supported_data_type(column["data_type"]), "semantic_context_not_sql_identifier": column["comment"], + **_semantic_role_context(column), } for column in columns ], @@ -261,11 +286,55 @@ def _format_semantic_context(context: dict) -> str: f"{_format_identifier_contract(context)}" "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" + "Values under column_role_hints_not_identifiers and semantic_roles_not_identifiers are meaning only; use them to map intent to exact declared columns, not as executable identifiers.\n" "*/\n" f"{_format_executable_identifier_catalog(context)}" ) +def _normalized_data_type(column: dict) -> str: + return str(column.get("data_type", "") or "").lower() + + +def _has_data_type_term(data_type: str, terms: set[str]) -> bool: + return any(term in data_type for term in terms) + + +def _column_roles(column: dict) -> list[str]: + name = str(column.get("name", "") or "") + comment = str(column.get("comment", "") or "") + searchable_text = f"{name} {comment}" + data_type = _normalized_data_type(column) + roles: list[str] = [] + + is_date_time = _has_data_type_term(data_type, _DATE_TIME_TYPE_TERMS) or bool( + _DATE_TIME_NAME_PATTERN.search(searchable_text) + ) + is_identifier = bool(column.get("is_primary_key")) or bool( + _IDENTIFIER_NAME_PATTERN.search(name) + ) + is_numeric = _has_data_type_term(data_type, _NUMERIC_TYPE_TERMS) + is_measure = (is_numeric or bool(_MEASURE_NAME_PATTERN.search(searchable_text))) and ( + not is_identifier + ) + + if is_date_time: + roles.append("date_time_candidate") + if is_measure: + roles.append("numeric_measure_candidate") + if is_identifier: + roles.append("identifier_candidate") + if not is_date_time and not is_measure and not is_identifier: + roles.append("dimension_candidate") + + return roles + + +def _semantic_role_context(column: dict) -> dict: + roles = _column_roles(column) + return {"semantic_roles_not_identifiers": roles} if roles else {} + + def _format_executable_identifier_catalog(context: dict) -> str: contract = context.get("sql_identifier_contract", {}) table_name = contract.get("sql_table_name_use_exactly") @@ -295,6 +364,24 @@ def _format_executable_identifier_catalog(context: dict) -> str: if column_names: lines.append("columns:") lines.extend(f"- {column_name}" for column_name in column_names) + role_hint_lines = [ + ( + column.get("sql_column_name_use_exactly"), + column.get("semantic_roles_not_identifiers") or [], + ) + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + and column.get("semantic_roles_not_identifiers") + ] + if role_hint_lines: + lines.append("column_role_hints_not_identifiers:") + lines.extend( + f"- {column_name}: {', '.join(roles)}" + for column_name, roles in role_hint_lines + ) + lines.append( + "Use role hints only to map question intent to exact columns listed above." + ) if relationship_constraints: lines.append("relationships:") lines.extend(f"- {constraint}" for constraint in relationship_constraints) @@ -428,6 +515,7 @@ def _build_table_context_ddl( "data_type": get_engine_supported_data_type(column["data_type"]), "is_primary_key": column["is_primary_key"], "semantic_context_not_sql_identifier": column["comment"], + **_semantic_role_context(column), } for column in included_columns ], @@ -474,6 +562,7 @@ def _build_view_ddl(content: dict) -> str: column.get("data_type") ), "semantic_context_not_sql_identifier": column.get("comment", ""), + **_semantic_role_context(column), } for column in columns ], diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index c8c43d382b..13997a655c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -32,6 +32,11 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "For metric-style requests" in prompt assert "Do not join tables just because they were retrieved together" in prompt assert "Do not invent join predicates from similar column names" in prompt + assert "column_role_hints_not_identifiers" in prompt + assert "date_time_candidate" in prompt + assert "numeric_measure_candidate" in prompt + assert "Do not answer a timeframe request with an unfiltered table scan" in prompt + assert "produce an analytical query shape" in prompt def test_sql_correction_system_prompt_allows_null_when_ungrounded(): @@ -104,6 +109,8 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): assert "- grouping_attribute" in built_prompt assert "- numeric_measure" in built_prompt assert "Generate an intent-shaped query, not a table preview" in built_prompt + assert "For timeframe requests, filter an exact date_time_candidate column" in built_prompt + assert "aggregate exact numeric_measure_candidate columns" in built_prompt def test_followup_sql_generation_prompt_requires_intent_shaped_query(): @@ -136,6 +143,7 @@ def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): assert "Failed SQL: SELECT 1" in built_prompt assert "DIAGNOSTIC CONTEXT" in built_prompt assert "Correct into an intent-shaped query, not a table preview" in built_prompt + assert "rebuild the query shape from the user's question" in built_prompt def test_sql_correction_prompt_includes_executable_schema_contract(): @@ -184,3 +192,4 @@ def test_sql_regeneration_prompt_includes_executable_schema_contract(): assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt assert "TABLE: retrieved_model" in built_prompt assert "Regenerate an intent-shaped query, not a table preview" in built_prompt + assert "For timeframe requests, filter an exact date_time_candidate column" in built_prompt diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 1357f13384..b0b3935401 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1685,6 +1685,69 @@ def encode(self, value): assert "-- This column is a measure" not in executable_ddl +def test_retrieved_schema_adds_generic_column_role_hints_without_comments(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "Business-facing dataset description.", + "columns": [ + { + "type": "COLUMN", + "name": "entity_id", + "data_type": "INTEGER", + "comment": "Business identifier label.", + "is_primary_key": True, + }, + { + "type": "COLUMN", + "name": "event_date", + "data_type": "DATE", + "comment": "Business date label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "measure_value", + "data_type": "DOUBLE", + "comment": "Business measure label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "category_label", + "data_type": "VARCHAR", + "comment": "Business category label.", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + table_ddl = result["db_schemas"][0]["table_ddl"] + executable_ddl = table_ddl.split("CREATE TABLE", maxsplit=1)[1] + + assert "column_role_hints_not_identifiers" in table_ddl + assert "- entity_id: identifier_candidate" in table_ddl + assert "- event_date: date_time_candidate" in table_ddl + assert "- measure_value: numeric_measure_candidate" in table_ddl + assert "- category_label: dimension_candidate" in table_ddl + assert "Business measure label." not in executable_ddl + assert "Business category label." not in executable_ddl + + def test_build_table_ddl_can_render_executable_schema_without_semantic_comments(): ddl, has_calculated_field, has_json_field = build_table_ddl( { From 44cc7e67433cc2640d2a052fe31879fe0a89a92d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 16:03:41 +0530 Subject: [PATCH 0824/1087] Improve modeling semantics generation quality --- .../generation/semantics_description.py | 86 ++-------- .../web/v1/services/semantics_description.py | 150 +++++++++++++++--- .../services/test_semantics_description.py | 148 +++++++++++++---- .../apollo/server/resolvers/modelResolver.ts | 98 ++++++++---- wren-ui/src/pages/modeling.tsx | 17 +- 5 files changed, 339 insertions(+), 160 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 58a63afc0e..4dde8afe1b 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -19,71 +19,18 @@ system_prompt = """ -I have a data model represented in JSON format, with the following structure: - -``` -[ - {'name': 'model', 'columns': [ - {'name': 'column_1', 'type': 'type', 'properties': {} - }, - {'name': 'column_2', 'type': 'type', 'properties': {} - }, - {'name': 'column_3', 'type': 'type', 'properties': {} - } - ], 'properties': {} - } -] -``` - -Your task is to update this JSON structure by adding a `description` field inside both the `properties` attribute of each `column` and the `model` itself. -Each `description` should be derived from the user-provided dataset context, the full schema, relationships, model names, column names, data types, aliases, and existing descriptions. -Follow these steps: -1. **For the `model`**: Write a clear natural language business description of the model's purpose, what real-world records it represents, and the common analysis questions it can answer. Insert this description in the `properties` field of the `model`. -2. **For each `column`**: Write a clear natural language business description of the column's meaning, not just its technical name. Each column's description should be added under its respective `properties` field in the format: `'description': 'business description'`. -3. Ensure that the output is a well-formatted JSON structure, preserving the input's original format and adding the appropriate `description` fields. -4. Avoid repeating technical table or column names as the whole description. Prefer business meaning such as identifiers, dates, amounts, statuses, dimensions, ownership, and operational usage. -5. Do not use generic boilerplate such as "stores the value", "contains records for", or "field from". Explain what the data means to a business user. -6. Make every model and column description unique, human-readable, concise, factual, and useful for text-to-SQL retrieval. -7. Use the model name, display label, existing description, column names, column display labels, and data types so descriptions include searchable business terms available from the metadata. -8. Do not invent table names, column names, relationships, or business concepts that are not supported by the provided model metadata. -9. If the metadata is technical or abbreviated, describe the observable business concepts from the available names and labels instead of copying the technical names. - -### Output Format: - -``` -{ - "models": [ - { - "name": "model", - "columns": [ - { - "name": "column_1", - "properties": { - "description": "" - } - }, - { - "name": "column_2", - "properties": { - "description": "" - } - }, - { - "name": "column_3", - "properties": { - "description": "" - } - } - ], - "properties": { - "description": "" - } - } - ] -} -``` - -Make sure that the descriptions are concise, informative, business-friendly, and contextually appropriate based on the input provided by the user. +Generate high-quality semantic descriptions for selected data models and their columns. + +Requirements: +1. Return valid JSON that matches the provided schema. +2. Return every input model exactly once and every input column exactly once. +3. Preserve every model and column `name` exactly as provided. +4. Put each generated description in `properties.description`. +5. Make descriptions business-friendly, concise, factual, and useful for text-to-SQL retrieval. +6. Ground descriptions only in the user prompt, model and column names, aliases, data types, existing descriptions, and provided schema context. +7. Make each column description specific to that column. Do not reuse the same wording across columns in the same model. +8. Do not invent unsupported tables, columns, relationships, metrics, or business concepts. +9. Do not use generic boilerplate or copy the technical name as the whole description. """ user_prompt_template = """ @@ -92,11 +39,10 @@ Picked models: {{ picked_models }} Localization Language: {{ language }} -Please provide business-friendly semantic descriptions for every picked model and every column based on the user's prompt and schema context. -Do not omit selected models or columns. Do not copy the table or column name as the description. -Use simple language that explains the business purpose, meaning, and analytical use of each field. -For each model description, include the main measures, dates, identifiers, dimensions, statuses, and business entities represented by its columns so vector retrieval can match natural-language questions to the correct model. -Keep descriptions factual and grounded in the picked model metadata only. +Write semantic descriptions for every picked model and every column. +For each model, describe the real-world records represented and the analytical questions it can support. +For each column, describe the business meaning and analytical use of that exact field. +Keep every description grounded in the picked model metadata and user prompt. """ diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 67d282591c..2b4d5fd3b4 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -32,9 +32,13 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, + generation_timeout_seconds: int = 90, + max_models_per_batch: int = 5, ): self._pipelines = pipelines self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) + self._generation_timeout_seconds = generation_timeout_seconds + self._max_models_per_batch = max(1, max_models_per_batch) def _handle_exception( self, @@ -60,48 +64,131 @@ class GenerateRequest(BaseRequest): mdl: str def _chunking( - self, mdl_dict: dict, request: GenerateRequest, chunk_size: int = 50 + self, + mdl_dict: dict, + request: GenerateRequest, + chunk_size: Optional[int] = None, ) -> list[dict]: + chunk_size = chunk_size or self._max_models_per_batch template = { "user_prompt": request.user_prompt, "language": request.configurations.language, } - chunks = [ - { - **model, - "columns": model.get("columns", [])[i : i + chunk_size], - } + selected_models = [ + model for model in mdl_dict.get("models", []) if model.get("name") in request.selected_models - for i in range(0, len(model.get("columns", [])), chunk_size) ] return [ { **template, - "mdl": {"models": [chunk]}, - "selected_models": [chunk["name"]], + "mdl": {"models": selected_models[i : i + chunk_size]}, + "selected_models": [ + model["name"] for model in selected_models[i : i + chunk_size] + ], } - for chunk in chunks + for i in range(0, len(selected_models), chunk_size) ] - async def _generate_task(self, request_id: str, chunk: dict): + async def _generate_task(self, chunk: dict) -> dict: resp = await self._pipelines["semantics_description"].run(**chunk) output = resp.get("output") or {} if not isinstance(output, dict): raise ValueError("Semantics description pipeline returned invalid output") + return output + + def _merge_outputs( + self, mdl_dict: dict, selected_models: list[str], outputs: list[dict] + ) -> dict: + generated_by_model = { + model_name: model_data + for output in outputs + for model_name, model_data in output.items() + if isinstance(model_data, dict) + } - current = self[request_id] - current.response = current.response or {} + def properties(payload: dict) -> dict: + value = payload.get("properties") + return value if isinstance(value, dict) else {} - for key in output.keys(): - if key not in current.response: - current.response[key] = output[key] + def description(payload: dict) -> str: + value = payload.get("description") or properties(payload).get( + "description", "" + ) + return "" if value is None else str(value).strip() + + response: dict = {} + for model in mdl_dict.get("models", []): + model_name = model.get("name") + if model_name not in selected_models: continue - current.response[key].setdefault("columns", []) - current.response[key]["columns"].extend(output[key].get("columns", [])) + generated_model = generated_by_model.get(model_name, {}) + if not generated_model: + raise ValueError( + f"Semantics description output omitted selected model: {model_name}" + ) + + model_description = description(generated_model) + if not model_description: + raise ValueError( + f"Semantics description output omitted description for model: {model_name}" + ) + + generated_columns = { + column.get("name"): column + for column in generated_model.get("columns", []) + if isinstance(column, dict) and column.get("name") + } + columns = [] + column_descriptions = [] + for column in model.get("columns", []): + if not isinstance(column, dict): + continue + + column_name = column.get("name", "") + generated_column = generated_columns.get(column_name) + if not generated_column: + raise ValueError( + "Semantics description output omitted selected column: " + f"{model_name}.{column_name}" + ) + + column_description = description(generated_column) + if not column_description: + raise ValueError( + "Semantics description output omitted description for column: " + f"{model_name}.{column_name}" + ) + + columns.append( + { + "name": column_name, + "type": column.get("type", ""), + "properties": { + "description": column_description, + }, + } + ) + column_descriptions.append(column_description) + + if len(set(column_descriptions)) != len(column_descriptions): + raise ValueError( + "Semantics description output contains repeated column " + f"descriptions for model: {model_name}" + ) + + response[model_name] = { + "name": model_name, + "columns": columns, + "properties": { + "description": model_description, + }, + } + + return response @observe(name="Generate Semantics Description") @trace_metadata @@ -117,13 +204,22 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: raise ValueError( "No selected models matched the current semantic model metadata" ) - tasks = [self._generate_task(request.id, chunk) for chunk in chunks] + tasks = [self._generate_task(chunk) for chunk in chunks] - await asyncio.gather(*tasks) + outputs = await asyncio.wait_for( + asyncio.gather(*tasks), + timeout=self._generation_timeout_seconds, + ) - self[request.id].status = "finished" - self[request.id].trace_id = trace_id - self[request.id].request_from = request.request_from + self[request.id] = self.Resource( + id=request.id, + status="finished", + response=self._merge_outputs( + mdl_dict, request.selected_models, list(outputs) + ), + trace_id=trace_id, + request_from=request.request_from, + ) except orjson.JSONDecodeError as e: self._handle_exception( request.id, @@ -132,6 +228,14 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: trace_id=trace_id, request_from=request.request_from, ) + except asyncio.TimeoutError: + self._handle_exception( + request.id, + "Semantics description generation timed out after " + f"{self._generation_timeout_seconds} seconds", + trace_id=trace_id, + request_from=request.request_from, + ) except Exception as e: self._handle_exception( request.id, diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 586c42f700..8c108b575f 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -13,7 +13,15 @@ def service(): mock_pipeline.run.return_value = { "output": { "model1": { - "columns": [], + "name": "model1", + "columns": [ + { + "name": "column1", + "properties": { + "description": "Customer segment for reporting." + }, + } + ], "properties": {"description": "Test description"}, } } @@ -42,12 +50,13 @@ async def test_generate_semantics_description( assert response.status == "finished" assert response.response == { "model1": { + "name": "model1", "columns": [ { "name": "column1", "type": "varchar", "properties": { - "description": "Business attribute used to categorize, filter, and explain records in analytical questions." + "description": "Customer segment for reporting." }, } ], @@ -109,7 +118,7 @@ async def test_generate_semantics_description_with_exception( @pytest.mark.asyncio -async def test_generate_semantics_description_with_llm_timeout_returns_fallback(): +async def test_generate_semantics_description_with_llm_timeout_fails(): mock_pipeline = AsyncMock() async def never_returns(**_): @@ -131,9 +140,9 @@ async def never_returns(**_): await service.generate(request) response = service[request.id] - assert response.status == "finished" - assert response.response["model1"]["properties"]["description"] - assert response.response["model1"]["columns"][0]["properties"]["description"] + assert response.status == "failed" + assert response.response is None + assert "timed out" in response.error.message def test_get_semantics_description_result( @@ -172,15 +181,15 @@ def test_semantics_description_uses_configured_timeout(): assert service._generation_timeout_seconds == 123 -def test_semantics_description_caps_timeout_inside_ui_polling_window(): +def test_semantics_description_uses_timeout_without_rewriting_ttl(): service = SemanticsDescription( pipelines={"semantics_description": AsyncMock()}, ttl=120, generation_timeout_seconds=600, ) - assert service._generation_timeout_seconds == 150 - assert service._cache.ttl >= 450 + assert service._generation_timeout_seconds == 600 + assert service._cache.ttl == 120 @pytest.mark.asyncio @@ -197,9 +206,33 @@ async def test_batch_processing_with_multiple_models( service._pipelines["semantics_description"].run.return_value = { "output": { - "model1": {"description": "Description 1"}, - "model2": {"description": "Description 2"}, - "model3": {"description": "Description 3"}, + "model1": { + "description": "Description 1", + "columns": [ + { + "name": "column1", + "properties": {"description": "Column description 1"}, + } + ], + }, + "model2": { + "description": "Description 2", + "columns": [ + { + "name": "column1", + "properties": {"description": "Column description 2"}, + } + ], + }, + "model3": { + "description": "Description 3", + "columns": [ + { + "name": "column1", + "properties": {"description": "Column description 3"}, + } + ], + }, } } @@ -278,7 +311,7 @@ def test_default_batch_allows_large_column_groups( @pytest.mark.asyncio -async def test_partial_llm_output_is_completed_for_all_selected_columns( +async def test_incomplete_llm_output_fails( service: SemanticsDescription, ): service["test_id"] = SemanticsDescription.Resource(id="test_id") @@ -324,15 +357,13 @@ async def test_partial_llm_output_is_completed_for_all_selected_columns( await service.generate(request) response = service[request.id] - assert response.status == "finished" - assert set(response.response.keys()) == {"orders", "customers"} - assert len(response.response["orders"]["columns"]) == 2 - assert len(response.response["customers"]["columns"]) == 1 - assert response.response["customers"]["properties"]["description"] + assert response.status == "failed" + assert response.response is None + assert "omitted selected column" in response.error.message @pytest.mark.asyncio -async def test_generic_llm_descriptions_are_replaced_with_business_descriptions( +async def test_llm_descriptions_are_not_rewritten_by_service( service: SemanticsDescription, ): service["test_id"] = SemanticsDescription.Resource(id="test_id") @@ -388,19 +419,17 @@ async def test_generic_llm_descriptions_are_replaced_with_business_descriptions( response = service[request.id] assert response.status == "finished" - descriptions = [ + assert response.response["dbo_xStageLoad2"]["properties"]["description"] == ( + "Contains business records for xStageLoad2." + ) + assert [ column["properties"]["description"] for column in response.response["dbo_xStageLoad2"]["columns"] + ] == [ + "Stores the Division value used to describe or analyze xStage records.", + "SalesPerson", + "Stores the SalesAmount value.", ] - assert response.response["dbo_xStageLoad2"]["properties"]["description"].startswith( - "Captures commercial activity" - ) - assert descriptions == [ - "Organizational segment used to group records for ownership, reporting, and performance comparison.", - "Responsible person or role associated with the record for ownership and performance analysis.", - "Monetary measure used to calculate financial results, compare performance, and summarize business activity.", - ] - assert all("Stores the" not in description for description in descriptions) @pytest.mark.asyncio @@ -444,7 +473,15 @@ async def test_concurrent_updates_no_race_condition( service._pipelines["semantics_description"].run.return_value = { "output": { - f"model{i}": {"description": f"Description {i}"} + f"model{i}": { + "description": f"Description {i}", + "columns": [ + { + "name": "column1", + "properties": {"description": f"Column description {i}"}, + } + ], + } for i in range(1, 6) } } @@ -457,6 +494,55 @@ async def test_concurrent_updates_no_race_condition( assert len(response.response) == 5 assert all(f"model{i}" in response.response for i in range(1, 6)) assert all( - response.response[f"model{i}"]["description"] == f"Description {i}" + response.response[f"model{i}"]["properties"]["description"] + == f"Description {i}" for i in range(1, 6) ) + + +@pytest.mark.asyncio +async def test_repeated_llm_column_descriptions_fail( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "order_id", "type": "varchar"}, + {"name": "customer_id", "type": "varchar"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "orders": { + "description": "Customer order transactions.", + "columns": [ + { + "name": "order_id", + "properties": {"description": "Identifier for reporting."}, + }, + { + "name": "customer_id", + "properties": {"description": "Identifier for reporting."}, + }, + ], + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "failed" + assert "repeated column descriptions" in response.error.message diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 5267c11aa3..7b0462a9cb 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -493,46 +493,78 @@ export class ModelResolver { projectId: project.id, }); const modelById = new Map(models.map((model) => [model.id, model])); + const requestedItems = args.data || []; + + for (const item of requestedItems) { + if (!modelById.has(item.modelId)) { + throw new Error(`Model not found: ${item.modelId}`); + } + } await Promise.all( - (args.data || []).map(async (item) => { - const model = modelById.get(item.modelId); - if (!model) { - throw new Error(`Model not found: ${item.modelId}`); - } + requestedItems.map(async (item) => { + if (isNil(item.description)) return; - await this.handleUpdateModelMetadata( - { - displayName: undefined, - description: item.description, - columns: [], - nestedColumns: [], - calculatedFields: [], - relationships: [], - }, - model, - ctx, - item.modelId, - ); + const model = modelById.get(item.modelId); + const properties = model?.properties ? JSON.parse(model.properties) : {}; + properties.description = this.determineMetadataValue(item.description); - if (!isEmpty(item.columns)) { - await this.handleUpdateColumnMetadata( - { - displayName: undefined, - description: undefined, - columns: item.columns, - nestedColumns: [], - calculatedFields: [], - relationships: [], - }, - ctx, - ); - } + await ctx.modelRepository.updateOne(item.modelId, { + properties: JSON.stringify(properties), + }); }), ); - this.markProjectDirty(project.id); - return { savedCount: args.data?.length || 0 }; + const requestedColumns = requestedItems.flatMap( + (item) => item.columns || [], + ); + if (!isEmpty(requestedColumns)) { + const columnIds = requestedColumns.map((column) => column.id); + const columns = await ctx.modelColumnRepository.findColumnsByIds(columnIds); + const columnById = new Map( + columns.map((column) => [String(column.id), column]), + ); + + await Promise.all( + requestedColumns.map(async (requestedColumn) => { + const column = columnById.get(String(requestedColumn.id)); + if (!column) return; + + const columnMetadata: Partial = {}; + if (!isNil(requestedColumn.displayName)) { + columnMetadata.displayName = this.determineMetadataValue( + requestedColumn.displayName, + ); + } + + if (!isNil(requestedColumn.description)) { + const properties = column.properties + ? JSON.parse(column.properties) + : {}; + properties.description = this.determineMetadataValue( + requestedColumn.description, + ); + columnMetadata.properties = JSON.stringify(properties); + } + + if (!isEmpty(columnMetadata)) { + await ctx.modelColumnRepository.updateOne( + requestedColumn.id, + columnMetadata, + ); + } + }), + ); + } + + if (!isEmpty(requestedItems)) { + this.markProjectDirty(project.id); + } + + return { + savedCount: requestedItems.length, + columnSavedCount: requestedColumns.length, + }; } public async generateModelingRelationships( diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 8511324c74..77cddb5b2c 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -272,6 +272,7 @@ const resolveRelationshipFieldParts = ( const renderIcon = (IconComponent) => React.createElement(IconComponent as any); const ASSISTANT_CANCELLED = 'ASSISTANT_CANCELLED'; +const ASSISTANT_SAVE_MESSAGE_KEY = 'modeling-ai-assistant-save'; export default function Modeling() { const router = useRouter(); @@ -279,6 +280,7 @@ export default function Modeling() { const apolloClient = useApolloClient(); const diagramRef = useRef(null); const assistantRunIdRef = useRef(0); + const assistantSavingRef = useRef(false); const [assistantMode, setAssistantMode] = useState< 'semantics' | 'relationships' | null >(null); @@ -655,7 +657,7 @@ export default function Modeling() { if (status === 'failed') { throw new Error(payload.error?.message || 'AI assistant failed.'); } - await new Promise((resolve) => setTimeout(resolve, 2000)); + await new Promise((resolve) => setTimeout(resolve, 1000)); } throw new Error('AI assistant timed out.'); }; @@ -905,6 +907,8 @@ export default function Modeling() { }; const saveAssistantResult = async () => { + if (assistantSavingRef.current) return; + assistantSavingRef.current = true; try { if (!diagramData) return; setAssistantLoading(true); @@ -976,16 +980,23 @@ export default function Modeling() { `${createdCount} relationship(s) saved. ${skippedCount} invalid or duplicate suggestion(s) skipped.`, ); } else { - message.success('Saved Modeling AI Assistant suggestions.'); + message.success({ + key: ASSISTANT_SAVE_MESSAGE_KEY, + content: 'Saved Modeling AI Assistant suggestions.', + }); } } closeAssistant(); if (assistantMode !== 'relationships') { - message.success('Saved Modeling AI Assistant suggestions.'); + message.success({ + key: ASSISTANT_SAVE_MESSAGE_KEY, + content: 'Saved Modeling AI Assistant suggestions.', + }); } } catch (error: any) { message.error(error.message || 'Failed to save assistant suggestions.'); } finally { + assistantSavingRef.current = false; setAssistantLoading(false); } }; From 16483dea754b696722f781c5337aca38edcab86c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 16:10:11 +0530 Subject: [PATCH 0825/1087] Cap related schema expansion for faster SQL generation --- .../retrieval/db_schema_retrieval.py | 6 ++ .../retrieval/test_db_schema_retrieval.py | 96 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index ad603029ff..a5413f1ed6 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -133,6 +133,7 @@ } _MAX_SCHEMA_SEMANTIC_TABLE_CANDIDATES = 5 +_MAX_RELATED_SCHEMA_TABLE_CANDIDATES = 5 _DATE_TIME_TYPE_TERMS = { "date", @@ -785,6 +786,11 @@ def _related_table_names(documents: list[Document], visited: set[str]) -> list[s if table_name and table_name not in visited: visited.add(table_name) related_names.append(table_name) + if ( + len(related_names) + >= _MAX_RELATED_SCHEMA_TABLE_CANDIDATES + ): + return related_names return related_names diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index b0b3935401..19664a6190 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -716,6 +716,102 @@ async def run(self, query_embedding, filters): ] +@pytest.mark.asyncio +async def test_dbschema_retrieval_caps_related_table_expansion(): + selected_model = "model_anchor" + related_models = [f"related_model_{index}" for index in range(8)] + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + names = [ + condition["value"] + for condition in filters["conditions"][1]["conditions"] + ] + self.calls.append(names) + + if names == [selected_model]: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": selected_model, + } + ), + meta={"type": "TABLE_SCHEMA", "name": selected_model}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "FOREIGN_KEY", + "tables": [ + selected_model, + related_model, + ], + "column": f"related_id_{index}", + "referenced_table": related_model, + "referenced_column": "id", + } + for index, related_model in enumerate( + related_models + ) + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": selected_model}, + ), + ] + } + + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": name, + } + ), + meta={"type": "TABLE_SCHEMA", "name": name}, + ) + for name in names + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={ + "documents": [ + Document( + content=str({"name": selected_model}), + meta={"type": "TABLE_DESCRIPTION", "name": selected_model}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + embedding={}, + ) + + assert retriever.calls == [ + [selected_model], + related_models[:5], + ] + assert [document.meta["name"] for document in documents] == [ + selected_model, + selected_model, + *related_models[:5], + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_uses_semantic_schema_hits_when_table_retrieval_misses(): semantic_model = "semantic_dataset" From 2bf1a31bfcfa94faf4b8d365452f7a937c6fbb78 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 16:21:38 +0530 Subject: [PATCH 0826/1087] Reject dotted physical table references in SQL --- .../src/pipelines/generation/utils/sql.py | 13 ++++++++-- .../generation/test_sql_post_processor.py | 26 +++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index aa52704bc4..ec6ad268e7 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -120,6 +120,15 @@ def _identifier_name(identifier: Identifier) -> str | None: return identifier.get_real_name() or identifier.get_name() +def _table_reference_name(identifier: Identifier) -> str | None: + real_name = identifier.get_real_name() + parent_name = identifier.get_parent_name() + if parent_name and real_name: + return f"{parent_name}.{real_name}" + + return real_name or identifier.get_name() + + def _contains_select(token: TokenList) -> bool: return any( child.ttype == sqlparse_tokens.Keyword.DML and child.normalized == "SELECT" @@ -179,7 +188,7 @@ def _collect_table_references(token: TokenList) -> set[str]: if isinstance(child, Parenthesis): table_names.update(_collect_table_references(child)) else: - name = _identifier_name(identifier) + name = _table_reference_name(identifier) if name: table_names.add(name) elif isinstance(current, Identifier): @@ -191,7 +200,7 @@ def _collect_table_references(token: TokenList) -> set[str]: if isinstance(child, Parenthesis): table_names.update(_collect_table_references(child)) else: - name = _identifier_name(current) + name = _table_reference_name(current) if name: table_names.add(name) expect_table = False diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index 9ea023df61..afa7c20ba3 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -160,6 +160,32 @@ async def test_sql_post_processor_rejects_tables_outside_schema_contract(): } +@pytest.mark.asyncio +async def test_sql_post_processor_rejects_dotted_physical_table_reference(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run( + ['{"sql": "SELECT DISTINCT grouping_col FROM physical_schema.supported_model"}'], + project_id="project-id", + schema_contracts=[ + { + "table_name": "supported_model", + "column_names": ["grouping_col"], + } + ], + ) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT DISTINCT grouping_col FROM physical_schema.supported_model", + "original_sql": "SELECT DISTINCT grouping_col FROM physical_schema.supported_model", + "type": "SCHEMA_GROUNDING", + "error": "Generated SQL references table identifiers outside the retrieved deployed schema.", + "correlation_id": "", + } + + @pytest.mark.asyncio async def test_sql_post_processor_rejects_joined_tables_outside_schema_contract(): engine = CapturingEngine() From 3eb76a8bad12629cc0d7979e46c2dc47c5217e26 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 16:32:36 +0530 Subject: [PATCH 0827/1087] Cap final SQL retrieval candidates --- .../retrieval/db_schema_retrieval.py | 13 ++- .../retrieval/test_db_schema_retrieval.py | 87 +++++++++++++++++++ 2 files changed, 97 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index a5413f1ed6..7baff758a2 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -134,6 +134,7 @@ _MAX_SCHEMA_SEMANTIC_TABLE_CANDIDATES = 5 _MAX_RELATED_SCHEMA_TABLE_CANDIDATES = 5 +_MAX_SQL_GENERATION_SCHEMA_RESULTS = 15 _DATE_TIME_TYPE_TERMS = { "date", @@ -581,6 +582,10 @@ def _build_view_ddl(content: dict) -> str: ) +def _limit_retrieval_results(retrieval_results: list[dict]) -> list[dict]: + return retrieval_results[:_MAX_SQL_GENERATION_SCHEMA_RESULTS] + + ## Start of Pipeline @observe(capture_input=False) async def active_mdl_hash( @@ -959,7 +964,7 @@ def check_using_db_schemas_without_pruning( } return { - "db_schemas": retrieval_results, + "db_schemas": _limit_retrieval_results(retrieval_results), "tokens": _token_count, "has_calculated_field": has_calculated_field, "has_metric": has_metric, @@ -1076,13 +1081,15 @@ def construct_retrieval_results( ) return { - "retrieval_results": retrieval_results, + "retrieval_results": _limit_retrieval_results(retrieval_results), "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, } else: - retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] + retrieval_results = _limit_retrieval_results( + check_using_db_schemas_without_pruning["db_schemas"] + ) return { "retrieval_results": retrieval_results, diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 19664a6190..65e55405bf 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1581,6 +1581,56 @@ def test_construct_retrieval_results_does_not_add_column_only_term_matches(): ] +def test_construct_retrieval_results_caps_pruned_generation_context(): + def table_schema(index): + return { + "type": "TABLE", + "name": f"model_{index}", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "model_0", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["id"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[table_schema(index) for index in range(20)], + dbschema_retrieval=[], + query="show model records", + ) + + assert len(result["retrieval_results"]) == 15 + assert [item["table_name"] for item in result["retrieval_results"]] == [ + f"model_{index}" for index in range(15) + ] + + def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): class Encoding: def encode(self, value): @@ -1636,6 +1686,43 @@ def table_schema(name): assert result["tokens"] > 0 +def test_check_using_db_schemas_without_pruning_caps_generation_context(): + class Encoding: + def encode(self, value): + return value.split() + + def table_schema(index): + return { + "type": "TABLE", + "name": f"model_{index}", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[table_schema(index) for index in range(20)], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=10000, + ) + + assert len(result["db_schemas"]) == 15 + assert [schema["table_name"] for schema in result["db_schemas"]] == [ + f"model_{index}" for index in range(15) + ] + + def test_retrieved_schema_separates_exact_sql_names_from_semantic_context(): class Encoding: def encode(self, value): From e1f518e9b32b41e35847df89622bbb843d10d4e4 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 16:34:13 +0530 Subject: [PATCH 0828/1087] Reduce semantics generation prompt size --- .../generation/semantics_description.py | 11 ++++++++- .../web/v1/services/semantics_description.py | 17 +++++++++---- .../services/test_semantics_description.py | 24 +++++++++++++++---- 3 files changed, 42 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 4dde8afe1b..e5f01944de 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -7,7 +7,7 @@ from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider @@ -165,21 +165,29 @@ def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: ## End of Pipeline class ModelProperties(BaseModel): + model_config = ConfigDict(extra="forbid") + description: str class ModelColumns(BaseModel): + model_config = ConfigDict(extra="forbid") + name: str properties: ModelProperties class SemanticModel(BaseModel): + model_config = ConfigDict(extra="forbid") + name: str columns: list[ModelColumns] properties: ModelProperties class SemanticResult(BaseModel): + model_config = ConfigDict(extra="forbid") + models: list[SemanticModel] @@ -188,6 +196,7 @@ class SemanticResult(BaseModel): "type": "json_schema", "json_schema": { "name": "semantic_description", + "strict": True, "schema": SemanticResult.model_json_schema(), }, } diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 2b4d5fd3b4..7840316d7c 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -33,12 +33,14 @@ def __init__( maxsize: int = 1_000_000, ttl: int = 120, generation_timeout_seconds: int = 90, - max_models_per_batch: int = 5, + max_models_per_batch: int = 1, + max_concurrent_tasks: int = 4, ): self._pipelines = pipelines self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) self._generation_timeout_seconds = generation_timeout_seconds self._max_models_per_batch = max(1, max_models_per_batch) + self._max_concurrent_tasks = max(1, max_concurrent_tasks) def _handle_exception( self, @@ -99,6 +101,15 @@ async def _generate_task(self, chunk: dict) -> dict: raise ValueError("Semantics description pipeline returned invalid output") return output + async def _generate_chunks(self, chunks: list[dict]) -> list[dict]: + semaphore = asyncio.Semaphore(self._max_concurrent_tasks) + + async def _bounded_generate(chunk: dict) -> dict: + async with semaphore: + return await self._generate_task(chunk) + + return await asyncio.gather(*[_bounded_generate(chunk) for chunk in chunks]) + def _merge_outputs( self, mdl_dict: dict, selected_models: list[str], outputs: list[dict] ) -> dict: @@ -204,10 +215,8 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: raise ValueError( "No selected models matched the current semantic model metadata" ) - tasks = [self._generate_task(chunk) for chunk in chunks] - outputs = await asyncio.wait_for( - asyncio.gather(*tasks), + self._generate_chunks(chunks), timeout=self._generation_timeout_seconds, ) diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 8c108b575f..8b3a8a49f5 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -192,6 +192,17 @@ def test_semantics_description_uses_timeout_without_rewriting_ttl(): assert service._cache.ttl == 120 +def test_semantics_description_uses_configured_batch_and_concurrency_limits(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + max_models_per_batch=2, + max_concurrent_tasks=3, + ) + + assert service._max_models_per_batch == 2 + assert service._max_concurrent_tasks == 3 + + @pytest.mark.asyncio async def test_batch_processing_with_multiple_models( service: SemanticsDescription, @@ -249,10 +260,10 @@ async def test_batch_processing_with_multiple_models( assert len(response.response["model3"]["columns"]) == 1 chunks = service._chunking(orjson.loads(request.mdl), request) - assert len(chunks) == 1 + assert len(chunks) == 3 assert all("user_prompt" in chunk for chunk in chunks) assert all("mdl" in chunk for chunk in chunks) - assert [len(chunk["selected_models"]) for chunk in chunks] == [3] + assert [len(chunk["selected_models"]) for chunk in chunks] == [1, 1, 1] def test_batch_processing_with_custom_chunk_size( @@ -275,7 +286,7 @@ def test_batch_processing_with_custom_chunk_size( assert chunks[1]["selected_models"] == ["model3", "model4"] -def test_default_batch_allows_large_column_groups( +def test_default_batch_keeps_large_column_groups_by_model( service: SemanticsDescription, ): request = SemanticsDescription.GenerateRequest( @@ -306,8 +317,11 @@ def test_default_batch_allows_large_column_groups( chunks = service._chunking(orjson.loads(request.mdl), request) - assert len(chunks) == 1 - assert chunks[0]["selected_models"] == ["model1", "model2"] + assert len(chunks) == 2 + assert chunks[0]["selected_models"] == ["model1"] + assert chunks[1]["selected_models"] == ["model2"] + assert len(chunks[0]["mdl"]["models"][0]["columns"]) == 500 + assert len(chunks[1]["mdl"]["models"][0]["columns"]) == 400 @pytest.mark.asyncio From bef70715f993e2596246709b23440a337934b0a4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 16:45:40 +0530 Subject: [PATCH 0829/1087] Skip diagnosis for deterministic SQL validation errors --- wren-ai-service/src/web/v1/services/ask.py | 21 +++- .../src/web/v1/services/ask_feedback.py | 7 +- .../tests/pytest/services/test_ask.py | 110 ++++++++++++++++++ 3 files changed, 135 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bc19470cc2..18185f15c1 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -13,6 +13,22 @@ logger = logging.getLogger("wren-ai-service") +_DETERMINISTIC_SQL_VALIDATION_TYPES = { + "SCHEMA_GROUNDING", + "SQL_SHAPE", + "SQL_SYNTAX", + "SQL_VALUE_GROUNDING", + "NO_RELEVANT_SQL", +} + + +def should_skip_sql_diagnosis(failed_generation_result: dict | None) -> bool: + if not failed_generation_result: + return False + + return failed_generation_result.get("type") in _DETERMINISTIC_SQL_VALIDATION_TYPES + + class AskHistory(BaseModel): sql: str question: str @@ -543,6 +559,9 @@ async def ask( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + skip_sql_diagnosis = should_skip_sql_diagnosis( + failed_dry_run_result + ) is_schema_grounding_error = ( failed_dry_run_result.get("type") == "SCHEMA_GROUNDING" ) @@ -569,7 +588,7 @@ async def ask( ) sql_diagnosis_reasoning = None - if allow_sql_diagnosis and not is_schema_grounding_error: + if allow_sql_diagnosis and not skip_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 218c9b70bc..7b1358f983 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -9,7 +9,7 @@ from src.core.pipeline import BasicPipeline from src.utils import trace_metadata from src.web.v1.services import BaseRequest -from src.web.v1.services.ask import AskError, AskResult +from src.web.v1.services.ask import AskError, AskResult, should_skip_sql_diagnosis logger = logging.getLogger("wren-ai-service") @@ -227,6 +227,9 @@ async def ask_feedback( original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + skip_sql_diagnosis = should_skip_sql_diagnosis( + failed_dry_run_result + ) is_schema_grounding_error = ( failed_dry_run_result.get("type") == "SCHEMA_GROUNDING" ) @@ -239,7 +242,7 @@ async def ask_feedback( trace_id=trace_id, ) - if allow_sql_diagnosis and not is_schema_grounding_error: + if allow_sql_diagnosis and not skip_sql_diagnosis: sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index e252e6819d..c804bb2992 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -12,6 +12,7 @@ AskRequest, AskResultRequest, AskService, + should_skip_sql_diagnosis, ) from src.web.v1.services.semantics_preparation import ( SemanticsPreparationRequest, @@ -161,6 +162,115 @@ def test_ask_service_uses_single_sql_correction_retry_by_default(): assert ask_service._max_sql_correction_retries == 1 +def test_should_skip_sql_diagnosis_for_deterministic_validation_errors(): + assert should_skip_sql_diagnosis({"type": "SQL_SHAPE"}) is True + assert should_skip_sql_diagnosis({"type": "SCHEMA_GROUNDING"}) is True + assert should_skip_sql_diagnosis({"type": "SQL_VALUE_GROUNDING"}) is True + assert should_skip_sql_diagnosis({"type": "DRY_RUN"}) is False + assert should_skip_sql_diagnosis({}) is False + + +class _EmptyRetrievalPipeline: + async def run(self, **_): + return {"formatted_output": {"documents": []}} + + +class _SchemaRetrievalPipeline: + async def run(self, **_): + return { + "construct_retrieval_results": { + "retrieval_results": [ + { + "table_name": "model_alpha", + "table_ddl": "CREATE TABLE model_alpha (entity_id INTEGER)", + "manifest_column_names": ["entity_id"], + } + ], + "has_calculated_field": False, + "has_metric": False, + "has_json_field": False, + } + } + + +class _ShapeInvalidSqlGenerationPipeline: + async def run(self, **_): + return { + "post_process": { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "SELECT entity_id FROM model_alpha", + "original_sql": "SELECT entity_id FROM model_alpha", + "type": "SQL_SHAPE", + "error": "Generated SQL is a table preview.", + "correlation_id": "", + }, + } + } + + +class _CapturingCorrectionPipeline: + def __init__(self): + self.calls = [] + + async def run(self, **kwargs): + self.calls.append(kwargs) + return { + "post_process": { + "valid_generation_result": { + "sql": "SELECT COUNT(*) AS record_count FROM model_alpha", + "correlation_id": "", + }, + "invalid_generation_result": {}, + } + } + + +class _FailingDiagnosisPipeline: + def __init__(self): + self.calls = [] + + async def run(self, **kwargs): + self.calls.append(kwargs) + raise AssertionError("sql_diagnosis should not run for local validation errors") + + +@pytest.mark.asyncio +async def test_ask_skips_sql_diagnosis_for_local_validation_error(): + correction = _CapturingCorrectionPipeline() + diagnosis = _FailingDiagnosisPipeline() + ask_service = AskService( + { + "historical_question": _EmptyRetrievalPipeline(), + "sql_pairs_retrieval": _EmptyRetrievalPipeline(), + "instructions_retrieval": _EmptyRetrievalPipeline(), + "db_schema_retrieval": _SchemaRetrievalPipeline(), + "sql_generation": _ShapeInvalidSqlGenerationPipeline(), + "sql_correction": correction, + "sql_diagnosis": diagnosis, + }, + allow_intent_classification=False, + allow_sql_functions_retrieval=False, + allow_sql_knowledge_retrieval=False, + allow_sql_diagnosis=True, + ) + query_id = str(uuid.uuid4()) + ask_request = AskRequest(query="count records by model", mdl_hash=None) + ask_request.query_id = query_id + + await ask_service.ask(ask_request) + + ask_result_response = ask_service.get_ask_result( + AskResultRequest(query_id=query_id) + ) + assert ask_result_response.status == "finished" + assert diagnosis.calls == [] + assert correction.calls[0]["invalid_generation_result"] == { + "sql": "SELECT entity_id FROM model_alpha", + "error": "Generated SQL is a table preview.", + } + + @pytest.mark.asyncio async def test_ask_with_successful_query( indexing_service: SemanticsPreparationService, From a6fe1f58f0b6406419189fb39adf0d9438588820 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 16:49:29 +0530 Subject: [PATCH 0830/1087] Prevent semantics generation timeout on wide models --- .../web/v1/services/semantics_description.py | 73 ++++++++++---- .../services/test_semantics_description.py | 96 ++++++++++++++++--- wren-ui/src/pages/modeling.tsx | 8 +- 3 files changed, 144 insertions(+), 33 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 7840316d7c..0bddd6f695 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,5 +1,6 @@ import asyncio import logging +import math from typing import Dict, Literal, Optional import orjson @@ -32,14 +33,16 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, - generation_timeout_seconds: int = 90, + generation_timeout_seconds: int = 120, max_models_per_batch: int = 1, + max_columns_per_batch: int = 40, max_concurrent_tasks: int = 4, ): self._pipelines = pipelines self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) self._generation_timeout_seconds = generation_timeout_seconds self._max_models_per_batch = max(1, max_models_per_batch) + self._max_columns_per_batch = max(1, max_columns_per_batch) self._max_concurrent_tasks = max(1, max_concurrent_tasks) def _handle_exception( @@ -83,16 +86,25 @@ def _chunking( if model.get("name") in request.selected_models ] - return [ - { - **template, - "mdl": {"models": selected_models[i : i + chunk_size]}, - "selected_models": [ - model["name"] for model in selected_models[i : i + chunk_size] - ], - } - for i in range(0, len(selected_models), chunk_size) - ] + chunks = [] + for i in range(0, len(selected_models), chunk_size): + for model in selected_models[i : i + chunk_size]: + columns = model.get("columns", []) + column_chunks = [ + columns[j : j + self._max_columns_per_batch] + for j in range(0, len(columns), self._max_columns_per_batch) + ] or [[]] + + for column_chunk in column_chunks: + chunks.append( + { + **template, + "mdl": {"models": [{**model, "columns": column_chunk}]}, + "selected_models": [model["name"]], + } + ) + + return chunks async def _generate_task(self, chunk: dict) -> dict: resp = await self._pipelines["semantics_description"].run(**chunk) @@ -110,16 +122,13 @@ async def _bounded_generate(chunk: dict) -> dict: return await asyncio.gather(*[_bounded_generate(chunk) for chunk in chunks]) + def _request_timeout_seconds(self, chunk_count: int) -> int: + waves = max(1, math.ceil(chunk_count / self._max_concurrent_tasks)) + return self._generation_timeout_seconds * waves + def _merge_outputs( self, mdl_dict: dict, selected_models: list[str], outputs: list[dict] ) -> dict: - generated_by_model = { - model_name: model_data - for output in outputs - for model_name, model_data in output.items() - if isinstance(model_data, dict) - } - def properties(payload: dict) -> dict: value = payload.get("properties") return value if isinstance(value, dict) else {} @@ -130,6 +139,28 @@ def description(payload: dict) -> str: ) return "" if value is None else str(value).strip() + generated_by_model: dict[str, dict] = {} + for output in outputs: + for model_name, model_data in output.items(): + if not isinstance(model_data, dict): + continue + + generated = generated_by_model.setdefault( + model_name, + { + "name": model_name, + "columns": [], + "properties": {}, + }, + ) + if not description(generated) and description(model_data): + generated["properties"] = { + **properties(generated), + "description": description(model_data), + } + generated.setdefault("columns", []) + generated["columns"].extend(model_data.get("columns", [])) + response: dict = {} for model in mdl_dict.get("models", []): model_name = model.get("name") @@ -206,6 +237,7 @@ def description(payload: dict) -> str: async def generate(self, request: GenerateRequest, **kwargs) -> Resource: logger.info("Generate Semantics Description pipeline is running...") trace_id = kwargs.get("trace_id") + request_timeout_seconds = self._generation_timeout_seconds try: mdl_dict = orjson.loads(request.mdl) @@ -215,9 +247,10 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: raise ValueError( "No selected models matched the current semantic model metadata" ) + request_timeout_seconds = self._request_timeout_seconds(len(chunks)) outputs = await asyncio.wait_for( self._generate_chunks(chunks), - timeout=self._generation_timeout_seconds, + timeout=request_timeout_seconds, ) self[request.id] = self.Resource( @@ -241,7 +274,7 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: self._handle_exception( request.id, "Semantics description generation timed out after " - f"{self._generation_timeout_seconds} seconds", + f"{request_timeout_seconds} seconds", trace_id=trace_id, request_from=request.request_from, ) diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 8b3a8a49f5..937ab867b9 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -200,9 +200,22 @@ def test_semantics_description_uses_configured_batch_and_concurrency_limits(): ) assert service._max_models_per_batch == 2 + assert service._max_columns_per_batch == 40 assert service._max_concurrent_tasks == 3 +def test_semantics_description_scales_request_timeout_by_concurrency_waves(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + generation_timeout_seconds=120, + max_concurrent_tasks=4, + ) + + assert service._request_timeout_seconds(1) == 120 + assert service._request_timeout_seconds(4) == 120 + assert service._request_timeout_seconds(5) == 240 + + @pytest.mark.asyncio async def test_batch_processing_with_multiple_models( service: SemanticsDescription, @@ -266,7 +279,7 @@ async def test_batch_processing_with_multiple_models( assert [len(chunk["selected_models"]) for chunk in chunks] == [1, 1, 1] -def test_batch_processing_with_custom_chunk_size( +def test_batch_processing_keeps_each_model_in_its_own_prompt( service: SemanticsDescription, ): service["test_id"] = SemanticsDescription.Resource(id="test_id") @@ -277,16 +290,19 @@ def test_batch_processing_with_custom_chunk_size( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model3", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model4", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Test chunking with custom chunk size chunks = service._chunking(orjson.loads(request.mdl), request, chunk_size=2) - assert len(chunks) == 2 - assert [len(chunk["selected_models"]) for chunk in chunks] == [2, 2] - assert chunks[0]["selected_models"] == ["model1", "model2"] - assert chunks[1]["selected_models"] == ["model3", "model4"] + assert len(chunks) == 4 + assert [len(chunk["selected_models"]) for chunk in chunks] == [1, 1, 1, 1] + assert [chunk["selected_models"][0] for chunk in chunks] == [ + "model1", + "model2", + "model3", + "model4", + ] -def test_default_batch_keeps_large_column_groups_by_model( +def test_default_batch_splits_large_column_groups_by_model( service: SemanticsDescription, ): request = SemanticsDescription.GenerateRequest( @@ -317,11 +333,69 @@ def test_default_batch_keeps_large_column_groups_by_model( chunks = service._chunking(orjson.loads(request.mdl), request) - assert len(chunks) == 2 + assert len(chunks) == 23 assert chunks[0]["selected_models"] == ["model1"] - assert chunks[1]["selected_models"] == ["model2"] - assert len(chunks[0]["mdl"]["models"][0]["columns"]) == 500 - assert len(chunks[1]["mdl"]["models"][0]["columns"]) == 400 + assert chunks[12]["selected_models"] == ["model1"] + assert chunks[13]["selected_models"] == ["model2"] + assert len(chunks[0]["mdl"]["models"][0]["columns"]) == 40 + assert len(chunks[12]["mdl"]["models"][0]["columns"]) == 20 + assert len(chunks[13]["mdl"]["models"][0]["columns"]) == 40 + + +@pytest.mark.asyncio +async def test_column_chunk_outputs_merge_into_single_model( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": f"column_{index}", "type": "varchar"} + for index in range(41) + ], + } + ] + } + ).decode(), + ) + + def response_for_chunk(**kwargs): + model = kwargs["mdl"]["models"][0] + return { + "output": { + "orders": { + "description": "Customer order transactions.", + "columns": [ + { + "name": column["name"], + "properties": { + "description": f"Description for {column['name']}", + }, + } + for column in model["columns"] + ], + } + } + } + + service._pipelines["semantics_description"].run.side_effect = response_for_chunk + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.response["orders"]["properties"]["description"] == ( + "Customer order transactions." + ) + assert len(response.response["orders"]["columns"]) == 41 + assert service._pipelines["semantics_description"].run.call_count == 2 @pytest.mark.asyncio diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 77cddb5b2c..f0b2322961 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -273,6 +273,8 @@ const resolveRelationshipFieldParts = ( const renderIcon = (IconComponent) => React.createElement(IconComponent as any); const ASSISTANT_CANCELLED = 'ASSISTANT_CANCELLED'; const ASSISTANT_SAVE_MESSAGE_KEY = 'modeling-ai-assistant-save'; +const ASSISTANT_POLL_INTERVAL_MS = 1000; +const ASSISTANT_MAX_POLL_ATTEMPTS = 240; export default function Modeling() { const router = useRouter(); @@ -639,7 +641,7 @@ export default function Modeling() { throw new Error('AI assistant did not return a task id.'); } - for (let attempt = 0; attempt < 90; attempt += 1) { + for (let attempt = 0; attempt < ASSISTANT_MAX_POLL_ATTEMPTS; attempt += 1) { if (assistantRunIdRef.current !== runId) { throw new Error(ASSISTANT_CANCELLED); } @@ -657,7 +659,9 @@ export default function Modeling() { if (status === 'failed') { throw new Error(payload.error?.message || 'AI assistant failed.'); } - await new Promise((resolve) => setTimeout(resolve, 1000)); + await new Promise((resolve) => + setTimeout(resolve, ASSISTANT_POLL_INTERVAL_MS), + ); } throw new Error('AI assistant timed out.'); }; From 1894e4e962c29b6e09207bc781bc9f8183fb2206 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 17:03:51 +0530 Subject: [PATCH 0831/1087] Expire stale restored asking tasks --- .../server/services/askingTaskTracker.ts | 77 +++++++++--- .../services/tests/askingTaskTracker.test.ts | 114 ++++++++++++++++++ 2 files changed, 175 insertions(+), 16 deletions(-) create mode 100644 wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index f81f87f5e5..afa6cc2df1 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -222,23 +222,30 @@ export class AskingTaskTracker implements IAskingTaskTracker { if (this.initialized) return; const taskRecords = await this.askingTaskRepository.findAll(); - taskRecords.forEach((taskRecord) => { - const detail = taskRecord.detail as AskResult | undefined; - if ( - !taskRecord.queryId || - !detail || - this.isTaskFinalized(detail.status) - ) { - return; - } + await Promise.all( + taskRecords.map(async (taskRecord) => { + const detail = taskRecord.detail as AskResult | undefined; + if ( + !taskRecord.queryId || + !detail || + this.isTaskFinalized(detail.status) + ) { + return; + } - this.restoreTrackedTask({ - ...detail, - queryId: taskRecord.queryId, - question: taskRecord.question, - taskId: taskRecord.id, - }); - }); + if (this.isStaleUnfinishedTask(taskRecord)) { + await this.finalizeStaleTask(taskRecord); + return; + } + + this.restoreTrackedTask({ + ...detail, + queryId: taskRecord.queryId, + question: taskRecord.question, + taskId: taskRecord.id, + }); + }), + ); this.initialized = true; } @@ -481,6 +488,20 @@ export class AskingTaskTracker implements IAskingTaskTracker { return null; } + if ( + taskRecord.detail && + !this.isTaskFinalized((taskRecord.detail as AskResult).status) && + this.isStaleUnfinishedTask(taskRecord) + ) { + await this.finalizeStaleTask(taskRecord); + taskRecord = await this.askingTaskRepository.findOneBy({ + id: taskRecord.id, + }); + if (!taskRecord) { + return null; + } + } + return { ...(taskRecord?.detail as AskResult), queryId: queryId || taskRecord?.queryId, @@ -535,6 +556,30 @@ export class AskingTaskTracker implements IAskingTaskTracker { ].includes(status); } + private isStaleUnfinishedTask(taskRecord: AskingTask): boolean { + const updatedAt = taskRecord.updatedAt + ? new Date(taskRecord.updatedAt).getTime() + : 0; + return Date.now() - updatedAt > this.memoryRetentionTime; + } + + private async finalizeStaleTask(taskRecord: AskingTask): Promise { + const result: AskResult = { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.POLLING_TIMEOUT, + message: + 'The previous asking task expired after the service restarted. Please ask again.', + }, + }; + + await this.askingTaskRepository.updateOne(taskRecord.id, { + detail: result, + }); + } + private isResultChanged( previousResult: AskResult, newResult: AskResult, diff --git a/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts b/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts new file mode 100644 index 0000000000..ca08d09369 --- /dev/null +++ b/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts @@ -0,0 +1,114 @@ +import { + AskingTaskTracker, + TrackedAskingResult, +} from '../askingTaskTracker'; +import { + AskResultStatus, + AskResultType, +} from '@server/models/adaptor'; +import * as Errors from '@server/utils/error'; + +describe('AskingTaskTracker', () => { + const createTracker = ({ + taskRecords = [], + memoryRetentionTime = 1000, + }: { + taskRecords?: any[]; + memoryRetentionTime?: number; + }) => { + const askingTaskRepository = { + findAll: jest.fn().mockResolvedValue(taskRecords), + findByQueryId: jest.fn(), + findOneBy: jest.fn(), + createOne: jest.fn(), + updateOne: jest.fn(), + }; + const wrenAIAdaptor = { + getAskResult: jest.fn(), + ask: jest.fn(), + cancelAsk: jest.fn(), + }; + const tracker = new AskingTaskTracker({ + wrenAIAdaptor: wrenAIAdaptor as any, + askingTaskRepository: askingTaskRepository as any, + threadResponseRepository: { updateOne: jest.fn() } as any, + viewRepository: { findOneBy: jest.fn() } as any, + pollingInterval: 1000, + memoryRetentionTime, + }); + tracker.stopPolling(); + + return { + tracker, + askingTaskRepository, + wrenAIAdaptor, + }; + }; + + test('finalizes stale unfinished tasks during initialization without polling AI service', async () => { + const oldDate = new Date(Date.now() - 60_000); + const staleTask = { + id: 7, + queryId: 'expired-query-id', + question: 'expired question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: oldDate, + updatedAt: oldDate, + }; + const { tracker, askingTaskRepository, wrenAIAdaptor } = createTracker({ + taskRecords: [staleTask], + memoryRetentionTime: 1000, + }); + + await tracker.initialize(); + + expect(wrenAIAdaptor.getAskResult).not.toHaveBeenCalled(); + expect(askingTaskRepository.updateOne).toHaveBeenCalledWith(7, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.POLLING_TIMEOUT, + message: + 'The previous asking task expired after the service restarted. Please ask again.', + }, + }, + }); + }); + + test('restores recent unfinished tasks so active AI service requests can continue', async () => { + const recentDate = new Date(); + const recentTask = { + id: 8, + queryId: 'recent-query-id', + question: 'recent question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: recentDate, + updatedAt: recentDate, + }; + const { tracker, askingTaskRepository } = createTracker({ + taskRecords: [recentTask], + memoryRetentionTime: 60_000, + }); + + await tracker.initialize(); + const result = (await tracker.getAskingResult( + recentTask.queryId, + )) as TrackedAskingResult; + + expect(askingTaskRepository.updateOne).not.toHaveBeenCalled(); + expect(result.queryId).toBe(recentTask.queryId); + expect(result.status).toBe(AskResultStatus.GENERATING); + }); +}); From 6c6ad63d73e6f2f0eb312406306309aec38779ad Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 17:08:02 +0530 Subject: [PATCH 0832/1087] Resolve relationships by current model references --- .../src/apollo/client/graphql/__types__.ts | 4 + wren-ui/src/apollo/server/schema.ts | 4 + .../apollo/server/services/modelService.ts | 88 +++++++++++++++++-- .../src/apollo/server/types/relationship.ts | 4 + wren-ui/src/pages/modeling.tsx | 4 + 5 files changed, 96 insertions(+), 8 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index f3da8f04fd..8d99acaf67 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -1303,9 +1303,13 @@ export type Relation = { export type RelationInput = { fromColumnId: Scalars['Int']; + fromColumnReferenceName?: InputMaybe; fromModelId: Scalars['Int']; + fromModelReferenceName?: InputMaybe; toColumnId: Scalars['Int']; + toColumnReferenceName?: InputMaybe; toModelId: Scalars['Int']; + toModelReferenceName?: InputMaybe; type: RelationType; }; diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 90fd686c77..4cff5a492f 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -244,9 +244,13 @@ export const typeDefs = gql` input RelationInput { fromModelId: Int! + fromModelReferenceName: String fromColumnId: Int! + fromColumnReferenceName: String toModelId: Int! + toModelReferenceName: String toColumnId: Int! + toColumnReferenceName: String type: RelationType! } diff --git a/wren-ui/src/apollo/server/services/modelService.ts b/wren-ui/src/apollo/server/services/modelService.ts index bd1235945f..6e64b86782 100644 --- a/wren-ui/src/apollo/server/services/modelService.ts +++ b/wren-ui/src/apollo/server/services/modelService.ts @@ -399,31 +399,103 @@ export class ModelService implements IModelService { public async createRelation(relation: RelationData): Promise { const { id } = await this.projectService.getCurrentProject(); - const modelIds = [relation.fromModelId, relation.toModelId]; - const models = await this.modelRepository.findAllByIds(modelIds); - const columnIds = [relation.fromColumnId, relation.toColumnId]; + const models = await this.modelRepository.findAllBy({ projectId: id }); + const modelIds = models.map((model) => model.id); + const projectColumns = + await this.modelColumnRepository.findColumnsByModelIds(modelIds); + const resolvedRelation = this.resolveRelationEndpoints( + relation, + models, + projectColumns, + ); + const columnIds = [ + resolvedRelation.fromColumnId, + resolvedRelation.toColumnId, + ]; const columns = await this.modelColumnRepository.findColumnsByIds(columnIds); const { valid, message } = await this.validateCreateRelation( models, columns, - relation, + resolvedRelation, ); if (!valid) { throw new Error(message); } - const relationName = this.generateRelationName(relation, models, columns); + const relationName = this.generateRelationName( + resolvedRelation, + models, + columns, + ); const savedRelation = await this.relationRepository.createOne({ projectId: id, name: relationName, - fromColumnId: relation.fromColumnId, - toColumnId: relation.toColumnId, - joinType: relation.type, + fromColumnId: resolvedRelation.fromColumnId, + toColumnId: resolvedRelation.toColumnId, + joinType: resolvedRelation.type, }); return savedRelation; } + private resolveRelationEndpoints( + relation: RelationData, + models: Model[], + columns: ModelColumn[], + ): RelationData { + const resolveModel = ( + id: number, + referenceName?: string, + ): Model | undefined => + models.find((model) => model.id === id) || + (referenceName + ? models.find((model) => model.referenceName === referenceName) + : undefined); + + const fromModel = resolveModel( + relation.fromModelId, + relation.fromModelReferenceName, + ); + const toModel = resolveModel( + relation.toModelId, + relation.toModelReferenceName, + ); + + const resolveColumn = ( + id: number, + modelId: number | undefined, + referenceName?: string, + ): ModelColumn | undefined => + columns.find( + (column) => column.id === id && (!modelId || column.modelId === modelId), + ) || + (referenceName && modelId + ? columns.find( + (column) => + column.modelId === modelId && column.referenceName === referenceName, + ) + : undefined); + + const fromColumn = resolveColumn( + relation.fromColumnId, + fromModel?.id, + relation.fromColumnReferenceName, + ); + const toColumn = resolveColumn( + relation.toColumnId, + toModel?.id, + relation.toColumnReferenceName, + ); + + return { + ...relation, + fromModelId: fromModel?.id ?? relation.fromModelId, + fromColumnId: fromColumn?.id ?? relation.fromColumnId, + toModelId: toModel?.id ?? relation.toModelId, + toColumnId: toColumn?.id ?? relation.toColumnId, + }; + } + public async updateRelation( relation: UpdateRelationData, id: number, diff --git a/wren-ui/src/apollo/server/types/relationship.ts b/wren-ui/src/apollo/server/types/relationship.ts index de88970899..91e71a8ebf 100644 --- a/wren-ui/src/apollo/server/types/relationship.ts +++ b/wren-ui/src/apollo/server/types/relationship.ts @@ -1,8 +1,12 @@ export interface RelationData { fromModelId: number; + fromModelReferenceName?: string; fromColumnId: number; + fromColumnReferenceName?: string; toModelId: number; + toModelReferenceName?: string; toColumnId: number; + toColumnReferenceName?: string; type: RelationType; description?: string; } diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index f0b2322961..82dffa529f 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -1596,9 +1596,13 @@ export default function Modeling() { variables: { data: { fromModelId: Number(relation.fromField.modelId), + fromModelReferenceName: relation.fromField.modelName, fromColumnId: Number(relation.fromField.fieldId), + fromColumnReferenceName: relation.fromField.fieldName, toModelId: Number(relation.toField.modelId), + toModelReferenceName: relation.toField.modelName, toColumnId: Number(relation.toField.fieldId), + toColumnReferenceName: relation.toField.fieldName, type: relation.type, }, }, From e7feb785b069e5d2decfac65a96e7e993eac60d4 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 17:18:32 +0530 Subject: [PATCH 0833/1087] Bound SQL generation task duration --- wren-ai-service/src/config.py | 1 + wren-ai-service/src/globals.py | 2 + wren-ai-service/src/providers/llm/litellm.py | 4 +- wren-ai-service/src/web/v1/services/ask.py | 149 ++++++++++-------- .../src/web/v1/services/ask_feedback.py | 101 +++++++----- .../tests/pytest/services/test_ask.py | 37 +++++ 6 files changed, 187 insertions(+), 107 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 1c85fbaa98..ef9ae7ab65 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -44,6 +44,7 @@ class Settings(BaseSettings): allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) max_sql_correction_retries: int = Field(default=1) + sql_generation_timeout_seconds: float = Field(default=45.0) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 9343344616..8cbd302bd1 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -162,6 +162,7 @@ def create_service_container( max_histories=settings.max_histories, enable_column_pruning=settings.enable_column_pruning, max_sql_correction_retries=settings.max_sql_correction_retries, + sql_generation_timeout_seconds=settings.sql_generation_timeout_seconds, **query_cache, ), ask_feedback_service=services.AskFeedbackService( @@ -182,6 +183,7 @@ def create_service_container( allow_sql_functions_retrieval=settings.allow_sql_functions_retrieval, allow_sql_diagnosis=settings.allow_sql_diagnosis, allow_sql_knowledge_retrieval=settings.allow_sql_knowledge_retrieval, + sql_generation_timeout_seconds=settings.sql_generation_timeout_seconds, **query_cache, ), chart_service=services.ChartService( diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index f4023f4c57..de58f74540 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -126,6 +126,7 @@ async def _run( **(generation_kwargs or {}), } ) + completion_timeout = generation_kwargs.pop("timeout", self._timeout) should_stream = ( streaming_callback is not None and query_id is not None @@ -143,6 +144,7 @@ async def _run( stream=should_stream, allowed_openai_params=allowed_openai_params, mock_testing_fallbacks=self._enable_fallback_testing, + timeout=completion_timeout, **generation_kwargs, ) else: @@ -151,7 +153,7 @@ async def _run( api_key=self._api_key, api_base=self._api_base, api_version=self._api_version, - timeout=self._timeout, + timeout=completion_timeout, messages=openai_formatted_messages, stream=should_stream, allowed_openai_params=allowed_openai_params, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 18185f15c1..0332fc42fc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -29,6 +29,15 @@ def should_skip_sql_diagnosis(failed_generation_result: dict | None) -> bool: return failed_generation_result.get("type") in _DETERMINISTIC_SQL_VALIDATION_TYPES +async def run_pipeline_with_timeout(awaitable, timeout_seconds: float, operation: str): + try: + return await asyncio.wait_for(awaitable, timeout=timeout_seconds) + except asyncio.TimeoutError as exc: + raise TimeoutError( + f"{operation} timed out after {timeout_seconds:g} seconds" + ) from exc + + class AskHistory(BaseModel): sql: str question: str @@ -121,6 +130,7 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, enable_column_pruning: bool = False, max_sql_correction_retries: int = 1, + sql_generation_timeout_seconds: float = 45.0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -137,6 +147,7 @@ def __init__( self._enable_column_pruning = enable_column_pruning self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries + self._sql_generation_timeout_seconds = sql_generation_timeout_seconds def _is_stopped(self, query_id: str, container: dict): if ( @@ -496,45 +507,49 @@ async def ask( has_json_field = _retrieval_result.get("has_json_field", False) if histories: - text_to_sql_generation_results = await self._pipelines[ - "followup_sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - histories=histories, - project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_contracts=schema_contracts, + text_to_sql_generation_results = await run_pipeline_with_timeout( + self._pipelines["followup_sql_generation"].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + histories=histories, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, + ), + self._sql_generation_timeout_seconds, + "Follow-up SQL generation", ) else: - text_to_sql_generation_results = await self._pipelines[ - "sql_generation" - ].run( - query=user_query, - contexts=table_ddls, - sql_generation_reasoning=sql_generation_reasoning, - project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_knowledge=sql_knowledge, - schema_contracts=schema_contracts, + text_to_sql_generation_results = await run_pipeline_with_timeout( + self._pipelines["sql_generation"].run( + query=user_query, + contexts=table_ddls, + sql_generation_reasoning=sql_generation_reasoning, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, + ), + self._sql_generation_timeout_seconds, + "SQL generation", ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -589,14 +604,16 @@ async def ask( sql_diagnosis_reasoning = None if allow_sql_diagnosis and not skip_sql_diagnosis: - sql_diagnosis_results = await self._pipelines[ - "sql_diagnosis" - ].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_request.configurations.language, + sql_diagnosis_results = await run_pipeline_with_timeout( + self._pipelines["sql_diagnosis"].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_request.configurations.language, + ), + self._sql_generation_timeout_seconds, + "SQL diagnosis", ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" @@ -606,25 +623,29 @@ async def ask( if sql_diagnosis_reasoning: correction_error_message = sql_diagnosis_reasoning - sql_correction_results = await self._pipelines[ - "sql_correction" - ].run( - contexts=table_ddls, - query=user_query, - instructions=instructions, - invalid_generation_result={ - "sql": ( - "" if is_schema_grounding_error else original_sql - ), - "error": correction_error_message, - }, - project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, - use_dry_plan=use_dry_plan, - allow_dry_plan_fallback=allow_dry_plan_fallback, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - schema_contracts=schema_contracts, + sql_correction_results = await run_pipeline_with_timeout( + self._pipelines["sql_correction"].run( + contexts=table_ddls, + query=user_query, + instructions=instructions, + invalid_generation_result={ + "sql": ( + "" + if is_schema_grounding_error + else original_sql + ), + "error": correction_error_message, + }, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, + ), + self._sql_generation_timeout_seconds, + "SQL correction", ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 7b1358f983..e464645e2e 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -9,7 +9,12 @@ from src.core.pipeline import BasicPipeline from src.utils import trace_metadata from src.web.v1.services import BaseRequest -from src.web.v1.services.ask import AskError, AskResult, should_skip_sql_diagnosis +from src.web.v1.services.ask import ( + AskError, + AskResult, + run_pipeline_with_timeout, + should_skip_sql_diagnosis, +) logger = logging.getLogger("wren-ai-service") @@ -65,6 +70,7 @@ def __init__( allow_sql_knowledge_retrieval: bool = True, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, + sql_generation_timeout_seconds: float = 45.0, maxsize: int = 1_000_000, ttl: int = 120, ): @@ -75,6 +81,7 @@ def __init__( self._allow_sql_knowledge_retrieval = allow_sql_knowledge_retrieval self._allow_sql_functions_retrieval = allow_sql_functions_retrieval self._allow_sql_diagnosis = allow_sql_diagnosis + self._sql_generation_timeout_seconds = sql_generation_timeout_seconds def _is_stopped(self, query_id: str, container: dict): if ( @@ -191,22 +198,24 @@ async def ask_feedback( trace_id=trace_id, ) - text_to_sql_generation_results = await self._pipelines[ - "sql_regeneration" - ].run( - contexts=table_ddls, - query=ask_feedback_request.question, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, - sql=ask_feedback_request.sql, - project_id=ask_feedback_request.project_id, - sql_samples=sql_samples, - instructions=instructions, - has_calculated_field=has_calculated_field, - has_metric=has_metric, - has_json_field=has_json_field, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - schema_contracts=schema_contracts, + text_to_sql_generation_results = await run_pipeline_with_timeout( + self._pipelines["sql_regeneration"].run( + contexts=table_ddls, + query=ask_feedback_request.question, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + sql=ask_feedback_request.sql, + project_id=ask_feedback_request.project_id, + sql_samples=sql_samples, + instructions=instructions, + has_calculated_field=has_calculated_field, + has_metric=has_metric, + has_json_field=has_json_field, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, + ), + self._sql_generation_timeout_seconds, + "SQL regeneration", ) if sql_valid_result := text_to_sql_generation_results["post_process"][ @@ -243,14 +252,16 @@ async def ask_feedback( ) if allow_sql_diagnosis and not skip_sql_diagnosis: - sql_diagnosis_results = await self._pipelines[ - "sql_diagnosis" - ].run( - contexts=table_ddls, - original_sql=original_sql, - invalid_sql=invalid_sql, - error_message=error_message, - language=ask_feedback_request.configurations.language, + sql_diagnosis_results = await run_pipeline_with_timeout( + self._pipelines["sql_diagnosis"].run( + contexts=table_ddls, + original_sql=original_sql, + invalid_sql=invalid_sql, + error_message=error_message, + language=ask_feedback_request.configurations.language, + ), + self._sql_generation_timeout_seconds, + "SQL diagnosis", ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" @@ -262,23 +273,29 @@ async def ask_feedback( f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" ) - sql_correction_results = await self._pipelines[ - "sql_correction" - ].run( - contexts=table_ddls, - query=ask_feedback_request.question, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, - instructions=instructions, - invalid_generation_result={ - "original_sql": original_sql, - "sql": "" if is_schema_grounding_error else invalid_sql, - "error": correction_error_message, - }, - project_id=ask_feedback_request.project_id, - mdl_hash=ask_feedback_request.mdl_hash, - sql_functions=sql_functions, - sql_knowledge=sql_knowledge, - schema_contracts=schema_contracts, + sql_correction_results = await run_pipeline_with_timeout( + self._pipelines["sql_correction"].run( + contexts=table_ddls, + query=ask_feedback_request.question, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, + instructions=instructions, + invalid_generation_result={ + "original_sql": original_sql, + "sql": ( + "" + if is_schema_grounding_error + else invalid_sql + ), + "error": correction_error_message, + }, + project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, + sql_functions=sql_functions, + sql_knowledge=sql_knowledge, + schema_contracts=schema_contracts, + ), + self._sql_generation_timeout_seconds, + "SQL correction", ) if valid_generation_result := sql_correction_results[ diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index c804bb2992..2de88e2264 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -1,3 +1,4 @@ +import asyncio import json import uuid @@ -209,6 +210,11 @@ async def run(self, **_): } +class _SlowSqlGenerationPipeline: + async def run(self, **_): + await asyncio.sleep(60) + + class _CapturingCorrectionPipeline: def __init__(self): self.calls = [] @@ -271,6 +277,37 @@ async def test_ask_skips_sql_diagnosis_for_local_validation_error(): } +@pytest.mark.asyncio +async def test_ask_times_out_slow_sql_generation_instead_of_hanging(): + ask_service = AskService( + { + "historical_question": _EmptyRetrievalPipeline(), + "sql_pairs_retrieval": _EmptyRetrievalPipeline(), + "instructions_retrieval": _EmptyRetrievalPipeline(), + "db_schema_retrieval": _SchemaRetrievalPipeline(), + "sql_generation": _SlowSqlGenerationPipeline(), + }, + allow_intent_classification=False, + allow_sql_functions_retrieval=False, + allow_sql_knowledge_retrieval=False, + sql_generation_timeout_seconds=0.01, + ) + query_id = str(uuid.uuid4()) + ask_request = AskRequest(query="count records by model", mdl_hash=None) + ask_request.query_id = query_id + + await ask_service.ask(ask_request) + + ask_result_response = ask_service.get_ask_result( + AskResultRequest(query_id=query_id) + ) + assert ask_result_response.status == "failed" + assert ask_result_response.error.code == "OTHERS" + assert ask_result_response.error.message == ( + "SQL generation timed out after 0.01 seconds" + ) + + @pytest.mark.asyncio async def test_ask_with_successful_query( indexing_service: SemanticsPreparationService, From 94dc954648d03cf611ca3a8b2a71c905a6c1eed3 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 17:34:35 +0530 Subject: [PATCH 0834/1087] Reduce SQL generation prompt latency --- .../src/pipelines/retrieval/db_schema_retrieval.py | 2 +- wren-ai-service/src/web/v1/services/ask.py | 10 ++++++++++ .../pipelines/retrieval/test_db_schema_retrieval.py | 8 ++++---- wren-ai-service/tools/config/config.example.yaml | 1 + wren-ai-service/tools/config/config.full.yaml | 1 + 5 files changed, 17 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 7baff758a2..1a657ac518 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -134,7 +134,7 @@ _MAX_SCHEMA_SEMANTIC_TABLE_CANDIDATES = 5 _MAX_RELATED_SCHEMA_TABLE_CANDIDATES = 5 -_MAX_SQL_GENERATION_SCHEMA_RESULTS = 15 +_MAX_SQL_GENERATION_SCHEMA_RESULTS = 10 _DATE_TIME_TYPE_TERMS = { "date", diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 0332fc42fc..a918a8ba63 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -31,8 +31,18 @@ def should_skip_sql_diagnosis(failed_generation_result: dict | None) -> bool: async def run_pipeline_with_timeout(awaitable, timeout_seconds: float, operation: str): try: + logger.info( + "%s started with timeout_seconds=%s", + operation, + timeout_seconds, + ) return await asyncio.wait_for(awaitable, timeout=timeout_seconds) except asyncio.TimeoutError as exc: + logger.error( + "%s timed out after %s seconds", + operation, + timeout_seconds, + ) raise TimeoutError( f"{operation} timed out after {timeout_seconds:g} seconds" ) from exc diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 65e55405bf..7b46fb377f 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1625,9 +1625,9 @@ def table_schema(index): query="show model records", ) - assert len(result["retrieval_results"]) == 15 + assert len(result["retrieval_results"]) == 10 assert [item["table_name"] for item in result["retrieval_results"]] == [ - f"model_{index}" for index in range(15) + f"model_{index}" for index in range(10) ] @@ -1717,9 +1717,9 @@ def table_schema(index): context_window_size=10000, ) - assert len(result["db_schemas"]) == 15 + assert len(result["db_schemas"]) == 10 assert [schema["table_name"] for schema in result["db_schemas"]] == [ - f"model_{index}" for index in range(15) + f"model_{index}" for index in range(10) ] diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index 792e5ef6c9..1bf72b6ade 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -193,6 +193,7 @@ settings: allow_sql_functions_retrieval: true enable_column_pruning: false max_sql_correction_retries: 1 + sql_generation_timeout_seconds: 30 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 5639d51b01..2a76787018 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -191,6 +191,7 @@ settings: allow_sql_functions_retrieval: true enable_column_pruning: false max_sql_correction_retries: 1 + sql_generation_timeout_seconds: 30 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com langfuse_enable: true From a1d0f2fa5ebb5a5145372aca6ff4cbd5cf6e284a Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Tue, 4 Aug 2026 18:48:13 +0530 Subject: [PATCH 0835/1087] Resolve relationship models from selected columns --- wren-ui/src/apollo/server/services/modelService.ts | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/apollo/server/services/modelService.ts b/wren-ui/src/apollo/server/services/modelService.ts index 6e64b86782..d084a2def2 100644 --- a/wren-ui/src/apollo/server/services/modelService.ts +++ b/wren-ui/src/apollo/server/services/modelService.ts @@ -452,11 +452,11 @@ export class ModelService implements IModelService { ? models.find((model) => model.referenceName === referenceName) : undefined); - const fromModel = resolveModel( + let fromModel = resolveModel( relation.fromModelId, relation.fromModelReferenceName, ); - const toModel = resolveModel( + let toModel = resolveModel( relation.toModelId, relation.toModelReferenceName, ); @@ -487,6 +487,13 @@ export class ModelService implements IModelService { relation.toColumnReferenceName, ); + if (!fromModel && fromColumn) { + fromModel = models.find((model) => model.id === fromColumn.modelId); + } + if (!toModel && toColumn) { + toModel = models.find((model) => model.id === toColumn.modelId); + } + return { ...relation, fromModelId: fromModel?.id ?? relation.fromModelId, From 8fa76048644555d9fe6492fee5fab27f822d7a46 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 18:51:52 +0530 Subject: [PATCH 0836/1087] Compress wide schemas for SQL generation --- .../retrieval/db_schema_retrieval.py | 84 ++++++++++++++++++- .../retrieval/test_db_schema_retrieval.py | 64 ++++++++++++++ 2 files changed, 147 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 1a657ac518..e6812c5fba 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -135,6 +135,8 @@ _MAX_SCHEMA_SEMANTIC_TABLE_CANDIDATES = 5 _MAX_RELATED_SCHEMA_TABLE_CANDIDATES = 5 _MAX_SQL_GENERATION_SCHEMA_RESULTS = 10 +_MAX_SQL_GENERATION_COLUMNS_PER_TABLE = 16 +_MAX_SQL_GENERATION_ROLE_COLUMNS = 4 _DATE_TIME_TYPE_TERMS = { "date", @@ -159,6 +161,15 @@ re.I, ) _IDENTIFIER_NAME_PATTERN = re.compile(r"(^id$|[_\s-]?id$|key|code|number|num|no$)", re.I) +_TIME_QUERY_PATTERN = re.compile( + r"\b(date|day|week|month|quarter|year|today|yesterday|last|next|from|to|between|period|recent)\b", + re.I, +) +_MEASURE_QUERY_PATTERN = re.compile( + r"\b(total|sum|average|avg|mean|count|number|top|highest|lowest|most|least|value|amount|sales|cost|price|rate|quantity|qty|score|percent)\b", + re.I, +) +_DETAIL_QUERY_PATTERN = re.compile(r"\b(show|list|detail|details|orders|records|rows)\b", re.I) def _normalize_terms(value: str) -> set[str]: @@ -229,6 +240,75 @@ def _tables_matching_query_terms( return matching_tables +def _column_search_terms(column: dict) -> set[str]: + return _normalize_terms( + " ".join( + [ + str(column.get("name", "") or ""), + str(column.get("comment", "") or ""), + str(column.get("data_type", "") or ""), + ] + ) + ) + + +def _compact_sql_generation_columns(content: dict, query: str) -> Optional[set[str]]: + columns = [ + column + for column in content.get("columns", []) + if column.get("type") == "COLUMN" + and column.get("name") + and ( + column.get("data_type") is None + or get_engine_supported_data_type(column.get("data_type")).lower() + != "unknown" + ) + ] + if len(columns) <= _MAX_SQL_GENERATION_COLUMNS_PER_TABLE: + return None + + query_terms = _normalize_terms(query) + wants_time = bool(_TIME_QUERY_PATTERN.search(query or "")) + wants_measure = bool(_MEASURE_QUERY_PATTERN.search(query or "")) + wants_detail = bool(_DETAIL_QUERY_PATTERN.search(query or "")) + selected: list[str] = [] + + def add(column: dict) -> None: + name = column.get("name") + if name and name not in selected: + selected.append(name) + + for column in columns: + if column.get("is_primary_key"): + add(column) + + for column in columns: + if query_terms and query_terms & _column_search_terms(column): + add(column) + + def add_role(role: str, limit: int = _MAX_SQL_GENERATION_ROLE_COLUMNS) -> None: + added = 0 + for column in columns: + if role in _column_roles(column): + add(column) + added += 1 + if added >= limit: + return + + if wants_time: + add_role("date_time_candidate") + if wants_measure: + add_role("numeric_measure_candidate") + if wants_detail: + add_role("identifier_candidate") + + if not selected: + add_role("dimension_candidate") + add_role("identifier_candidate") + + return set(selected[:_MAX_SQL_GENERATION_COLUMNS_PER_TABLE]) + + def _build_metric_ddl(content: dict) -> str: columns = [ column @@ -896,6 +976,7 @@ def check_using_db_schemas_without_pruning( encoding: tiktoken.Encoding, enable_column_pruning: bool, context_window_size: int, + query: str = "", ) -> dict: retrieval_results = [] has_calculated_field = False @@ -904,8 +985,9 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": + compact_columns = _compact_sql_generation_columns(table_schema, query) ddl, _has_calculated_field, _has_json_field, column_names = ( - _build_table_context_ddl(table_schema) + _build_table_context_ddl(table_schema, columns=compact_columns) ) retrieval_results.append( { diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7b46fb377f..84a70b869e 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1723,6 +1723,70 @@ def table_schema(index): ] +def test_check_using_db_schemas_without_pruning_compacts_wide_tables_by_query(): + class Encoding: + def encode(self, value): + return value.split() + + columns = [ + { + "type": "COLUMN", + "name": "OrderDate", + "data_type": "DATE", + "comment": "Order date", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "OrderNo", + "data_type": "VARCHAR", + "comment": "Order number", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "SalesAmount", + "data_type": "DECIMAL", + "comment": "Sales amount", + "is_primary_key": False, + }, + ] + columns.extend( + { + "type": "COLUMN", + "name": f"Filler{index}", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + for index in range(25) + ) + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "wide_orders", + "comment": "", + "columns": columns, + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=10000, + query="show orders from last week", + ) + + column_names = result["db_schemas"][0]["column_names"] + assert len(column_names) <= 16 + assert "OrderDate" in column_names + assert "OrderNo" in column_names + assert "Filler24" not in column_names + + def test_retrieved_schema_separates_exact_sql_names_from_semantic_context(): class Encoding: def encode(self, value): From 3ad59687c8a00749161c75438b462bd618c2eb53 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 19:27:07 +0530 Subject: [PATCH 0837/1087] Add grounded SQL fast path --- .../generation/followup_sql_generation.py | 12 + .../pipelines/generation/sql_generation.py | 12 + .../generation/utils/deterministic_sql.py | 498 ++++++++++++++++++ .../generation/test_deterministic_sql.py | 223 ++++++++ 4 files changed, 745 insertions(+) create mode 100644 wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py 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 f3f4ae4d18..58d86b5b2e 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -2,6 +2,7 @@ import sys from typing import Any +import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder @@ -11,6 +12,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata +from src.pipelines.generation.utils.deterministic_sql import generate_grounded_sql from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -142,11 +144,21 @@ def prompt( @trace_cost async def generate_sql_in_followup( prompt: dict, + query: str, + documents: list[str], generator: Any, histories: list[AskHistory], generator_name: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: + deterministic_sql = generate_grounded_sql(query, documents) + if deterministic_sql: + logger.info("Follow-Up SQL Generation used deterministic grounded SQL fast path.") + return { + "replies": [orjson.dumps({"sql": deterministic_sql}).decode("utf-8")], + "metadata": [{"finish_reason": "deterministic_grounded_sql"}], + }, generator_name + history_messages = construct_ask_history_messages(histories) current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) return await generator( diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 8de6a23961..11f4055d92 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -2,6 +2,7 @@ import sys from typing import Any +import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder @@ -11,6 +12,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata +from src.pipelines.generation.utils.deterministic_sql import generate_grounded_sql from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -142,10 +144,20 @@ def prompt( @trace_cost async def generate_sql( prompt: dict, + query: str, + documents: list[str], generator: Any, generator_name: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: + deterministic_sql = generate_grounded_sql(query, documents) + if deterministic_sql: + logger.info("SQL Generation used deterministic grounded SQL fast path.") + return { + "replies": [orjson.dumps({"sql": deterministic_sql}).decode("utf-8")], + "metadata": [{"finish_reason": "deterministic_grounded_sql"}], + }, generator_name + current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt diff --git a/wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py b/wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py new file mode 100644 index 0000000000..b97f717439 --- /dev/null +++ b/wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py @@ -0,0 +1,498 @@ +import calendar +import re +from dataclasses import dataclass +from datetime import date, timedelta + +import orjson + + +_SEMANTIC_CONTEXT_PATTERN = re.compile( + r"WREN RETRIEVED SEMANTIC CONTEXT\s*\n(\{.*?\})\s*\n", re.DOTALL +) +_CREATE_TABLE_PATTERN = re.compile( + r"CREATE\s+TABLE\s+([A-Za-z_][A-Za-z0-9_]*|\"[^\"]+\")\s*\((.*?)\)", + re.IGNORECASE | re.DOTALL, +) +_WORD_PATTERN = re.compile(r"[A-Za-z][A-Za-z0-9]*") +_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + +_STOP_TERMS = { + "a", + "an", + "and", + "are", + "as", + "between", + "by", + "from", + "for", + "in", + "is", + "most", + "of", + "on", + "per", + "show", + "the", + "to", + "using", + "what", + "which", + "with", +} +_DETAIL_TERMS = {"show", "list", "display", "find", "get"} +_RANKING_TERMS = {"top", "bottom", "highest", "lowest", "most", "least"} +_COUNT_TERMS = {"count", "counts", "number", "orders", "order"} +_SUM_TERMS = { + "amount", + "cost", + "qty", + "quantity", + "revenue", + "sales", + "sold", + "sum", + "total", + "value", +} +_VALUE_MEASURE_TERMS = { + "amount", + "cost", + "qty", + "quantity", + "revenue", + "sales", + "sold", + "value", +} +_AVG_TERMS = {"average", "avg", "mean"} +_MIN_TERMS = {"minimum", "min", "lowest"} +_MAX_TERMS = {"maximum", "max", "highest", "most"} +_MONTHS = { + month.lower(): index for index, month in enumerate(calendar.month_name) if month +} + + +@dataclass(frozen=True) +class _Column: + name: str + data_type: str = "" + semantic_text: str = "" + roles: tuple[str, ...] = () + + +@dataclass(frozen=True) +class _Table: + name: str + semantic_text: str + columns: tuple[_Column, ...] + + +def generate_grounded_sql(query: str, documents: list[str]) -> str | None: + tables = [ + table for document in documents for table in [_parse_table(document)] if table + ] + if not query or not tables: + return None + + query_terms = _terms(query) + if not query_terms: + return None + + table = max(tables, key=lambda candidate: _table_score(candidate, query_terms)) + if _table_score(table, query_terms) <= 0: + return None + + filters = _build_filters(query, table) + aggregate_sql = _build_aggregate_sql(query, query_terms, table, filters) + if aggregate_sql: + return aggregate_sql + + return _build_detail_sql(query, query_terms, table, filters) + + +def _parse_table(document: str) -> _Table | None: + context_match = _SEMANTIC_CONTEXT_PATTERN.search(document or "") + if context_match: + try: + context = orjson.loads(context_match.group(1)) + contract = context.get("sql_identifier_contract") or {} + table_name = contract.get("sql_table_name_use_exactly") + columns = [] + allowed_columns = set(contract.get("sql_column_names_use_exactly") or []) + for column in context.get("columns") or []: + column_name = column.get("sql_column_name_use_exactly") + if not column_name or ( + allowed_columns and column_name not in allowed_columns + ): + continue + columns.append( + _Column( + name=column_name, + data_type=str(column.get("data_type") or ""), + semantic_text=str( + column.get("semantic_context_not_sql_identifier") or "" + ), + roles=tuple(column.get("semantic_roles_not_identifiers") or ()), + ) + ) + + if table_name and columns: + return _Table( + name=table_name, + semantic_text=str( + (context.get("semantic_context_not_sql_identifiers") or {}).get( + "description", "" + ) + ), + columns=tuple(columns), + ) + except orjson.JSONDecodeError: + pass + + ddl_match = _CREATE_TABLE_PATTERN.search(document or "") + if not ddl_match: + return None + + columns = [] + for raw_column in ddl_match.group(2).split(","): + pieces = raw_column.strip().split() + if len(pieces) >= 2: + columns.append( + _Column(name=_unquote_identifier(pieces[0]), data_type=pieces[1]) + ) + + if not columns: + return None + + return _Table( + name=_unquote_identifier(ddl_match.group(1)), + semantic_text="", + columns=tuple(columns), + ) + + +def _terms(value: str) -> set[str]: + terms = set() + for token in _WORD_PATTERN.findall(value or ""): + normalized = token.lower() + terms.add(normalized) + if normalized.endswith("ies") and len(normalized) > 4: + terms.add(f"{normalized[:-3]}y") + elif normalized.endswith("s") and len(normalized) > 3: + terms.add(normalized[:-1]) + return terms + + +def _identifier_terms(value: str) -> set[str]: + spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", value or "") + spaced = spaced.replace("_", " ") + return _terms(spaced) + + +def _column_terms(column: _Column) -> set[str]: + return ( + _identifier_terms(column.name) + | _terms(column.semantic_text) + | {term for role in column.roles for term in _identifier_terms(role)} + ) + + +def _table_score(table: _Table, query_terms: set[str]) -> int: + score = len( + (_identifier_terms(table.name) | _terms(table.semantic_text)) & query_terms + ) + column_scores = sorted( + (len(_column_terms(column) & query_terms) for column in table.columns), + reverse=True, + ) + return score + sum(column_scores[:6]) + + +def _is_numeric(column: _Column) -> bool: + data_type = column.data_type.lower() + return "numeric_measure_candidate" in column.roles or any( + token in data_type + for token in ( + "int", + "decimal", + "numeric", + "number", + "double", + "float", + "real", + "money", + ) + ) + + +def _is_datetime(column: _Column) -> bool: + data_type = column.data_type.lower() + return "date_time_candidate" in column.roles or any( + token in data_type for token in ("date", "time", "timestamp") + ) + + +def _is_dimension(column: _Column) -> bool: + return not _is_numeric(column) and not _is_datetime(column) + + +def _best_columns( + table: _Table, + query_terms: set[str], + predicate, + *, + limit: int = 3, +) -> list[_Column]: + scored = [] + for column in table.columns: + if not predicate(column): + continue + score = len(_column_terms(column) & query_terms) + if score: + scored.append((score, column)) + return [ + column + for _, column in sorted(scored, key=lambda item: item[0], reverse=True)[:limit] + ] + + +def _best_measure(table: _Table, query_terms: set[str]) -> _Column | None: + measures = _best_columns(table, query_terms, _is_numeric, limit=1) + if measures: + return measures[0] + + numeric_columns = [column for column in table.columns if _is_numeric(column)] + return numeric_columns[0] if numeric_columns else None + + +def _best_date(table: _Table, query_terms: set[str]) -> _Column | None: + date_columns = _best_columns(table, query_terms, _is_datetime, limit=1) + if date_columns: + return date_columns[0] + + date_columns = [column for column in table.columns if _is_datetime(column)] + return date_columns[0] if date_columns else None + + +def _query_requests_aggregate(query_terms: set[str]) -> bool: + aggregate_terms = ( + _RANKING_TERMS | _SUM_TERMS | _AVG_TERMS | _MIN_TERMS | _MAX_TERMS + ) + return bool(query_terms & aggregate_terms) or ( + "by" in query_terms or "per" in query_terms or "compare" in query_terms + ) + + +def _requested_aggregate( + query_terms: set[str], measure: _Column | None +) -> tuple[str, str]: + if query_terms & _AVG_TERMS: + return "AVG", "AverageValue" + if query_terms & _MIN_TERMS: + return "MIN", "MinimumValue" + if query_terms & _COUNT_TERMS and not (query_terms & _VALUE_MEASURE_TERMS): + return "COUNT", "TotalOrders" + if query_terms & _MAX_TERMS and measure and not (query_terms & {"most", "top"}): + return "MAX", "MaximumValue" + if measure: + return "SUM", "TotalValue" + return "COUNT", "TotalCount" + + +def _build_aggregate_sql( + query: str, + query_terms: set[str], + table: _Table, + filters: list[str], +) -> str | None: + if not _query_requests_aggregate(query_terms): + return None + + dimensions = _best_columns(table, query_terms, _is_dimension, limit=2) + if not dimensions and not (query_terms & _COUNT_TERMS): + return None + + measure = _best_measure(table, query_terms) + function_name, alias = _requested_aggregate(query_terms, measure) + if function_name == "COUNT": + aggregate_expression = f"COUNT(*) AS {_sql_identifier(alias)}" + elif measure: + aggregate_expression = ( + f"{function_name}({_sql_identifier(measure.name)}) AS {_sql_identifier(alias)}" + ) + else: + return None + + select_items = [_sql_identifier(column.name) for column in dimensions] + select_items.append(aggregate_expression) + clauses = [ + "SELECT", + " " + ",\n ".join(select_items), + "FROM", + f" {_sql_identifier(table.name)}", + ] + if filters: + clauses.extend(["WHERE", " " + "\n AND ".join(filters)]) + if dimensions: + group_by = ", ".join(_sql_identifier(column.name) for column in dimensions) + clauses.extend(["GROUP BY", f" {group_by}"]) + if query_terms & _RANKING_TERMS or "top" in query_terms or "bottom" in query_terms: + direction = "ASC" if query_terms & {"bottom", "lowest", "least"} else "DESC" + clauses.extend(["ORDER BY", f" {_sql_identifier(alias)} {direction}"]) + clauses.append(f"LIMIT {_top_limit(query)}") + elif "top" in query_terms: + clauses.append(f"LIMIT {_top_limit(query)}") + + return "\n".join(clauses) + + +def _build_detail_sql( + query: str, + query_terms: set[str], + table: _Table, + filters: list[str], +) -> str | None: + if not filters and not (query_terms & _DETAIL_TERMS): + return None + + selected = [] + for predicate in (_is_datetime, _is_dimension, _is_numeric): + for column in _best_columns(table, query_terms, predicate, limit=4): + if column.name not in {item.name for item in selected}: + selected.append(column) + + if len(selected) < 3: + for column in table.columns: + if column.name not in {item.name for item in selected}: + selected.append(column) + if len(selected) >= 6: + break + + selected = selected[:6] + if not selected: + return None + + clauses = [ + "SELECT", + " " + ",\n ".join(_sql_identifier(column.name) for column in selected), + "FROM", + f" {_sql_identifier(table.name)}", + ] + if filters: + clauses.extend(["WHERE", " " + "\n AND ".join(filters)]) + + date_column = _best_date(table, query_terms) + if date_column: + clauses.extend(["ORDER BY", f" {_sql_identifier(date_column.name)} DESC"]) + clauses.append("LIMIT 500") + return "\n".join(clauses) + + +def _build_filters(query: str, table: _Table) -> list[str]: + query_terms = _terms(query) + filters: list[str] = [] + + date_range = _date_range(query) + date_column = _best_date(table, query_terms) + if date_range and date_column: + start, end = date_range + date_identifier = _sql_identifier(date_column.name) + filters.append(f"{date_identifier} >= '{start.isoformat()}'") + filters.append(f"{date_identifier} < '{end.isoformat()}'") + + literal_filter = _literal_dimension_filter(query, table) + if literal_filter: + filters.append(literal_filter) + + return filters + + +def _literal_dimension_filter(query: str, table: _Table) -> str | None: + if _query_requests_aggregate(_terms(query)): + return None + + query_words = _WORD_PATTERN.findall(query or "") + lowered_words = [word.lower() for word in query_words] + + for column in table.columns: + if not _is_dimension(column): + continue + column_terms = _identifier_terms(column.name) + for index, word in enumerate(lowered_words[:-1]): + if word not in column_terms: + continue + literal_words = [] + for next_word in query_words[index + 1 : index + 4]: + normalized = next_word.lower() + if normalized in _STOP_TERMS or normalized in _MONTHS: + break + literal_words.append(next_word) + if literal_words: + literal = " ".join(literal_words) + return f"{_sql_identifier(column.name)} = '{_escape_literal(literal)}'" + + return None + + +def _date_range(query: str, today: date | None = None) -> tuple[date, date] | None: + current = today or date.today() + normalized = (query or "").lower() + + if "today" in normalized: + return current, current + timedelta(days=1) + if "yesterday" in normalized: + previous = current - timedelta(days=1) + return previous, current + if "last week" in normalized: + return current - timedelta(days=7), current + timedelta(days=1) + if "last month" in normalized: + start = _add_months(date(current.year, current.month, 1), -1) + end = date(current.year, current.month, 1) + return start, end + + last_months = re.search(r"\blast\s+(\d{1,2})\s+months?\b", normalized) + if last_months: + return _add_months(current, -int(last_months.group(1))), current + timedelta( + days=1 + ) + + for month_name, month_index in _MONTHS.items(): + match = re.search(rf"\b{month_name}\s+(\d{{4}})\b", normalized) + if match: + start = date(int(match.group(1)), month_index, 1) + return start, _add_months(start, 1) + + return None + + +def _add_months(value: date, months: int) -> date: + month = value.month - 1 + months + year = value.year + month // 12 + month = month % 12 + 1 + day = min(value.day, calendar.monthrange(year, month)[1]) + return date(year, month, day) + + +def _top_limit(query: str) -> int: + match = re.search(r"\btop\s+(\d{1,3})\b", query.lower()) + if not match: + return 10 + return min(max(int(match.group(1)), 1), 500) + + +def _sql_identifier(identifier: str) -> str: + if _IDENTIFIER_PATTERN.match(identifier): + return identifier + return '"' + identifier.replace('"', '""') + '"' + + +def _unquote_identifier(identifier: str) -> str: + identifier = identifier.strip() + if identifier.startswith('"') and identifier.endswith('"'): + return identifier[1:-1].replace('""', '"') + return identifier + + +def _escape_literal(value: str) -> str: + return value.replace("'", "''") diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py b/wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py new file mode 100644 index 0000000000..97251723f9 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py @@ -0,0 +1,223 @@ +import pytest + +from src.pipelines.generation.sql_generation import generate_sql +from src.pipelines.generation.utils.deterministic_sql import generate_grounded_sql + + +def _schema_document(table_name: str, columns: list[dict]) -> str: + import orjson + + context = { + "object_type": "model", + "sql_identifier_contract": { + "sql_table_name_use_exactly": table_name, + "sql_column_names_use_exactly": [column["name"] for column in columns], + }, + "semantic_context_not_sql_identifiers": { + "description": "Synthetic sales and order analytics model.", + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": column["data_type"], + "semantic_context_not_sql_identifier": column.get("comment", ""), + "semantic_roles_not_identifiers": column.get("roles", []), + } + for column in columns + ], + } + column_ddl = ",\n ".join( + f"{column['name']} {column['data_type']}" for column in columns + ) + return ( + "/*\n" + "WREN RETRIEVED SEMANTIC CONTEXT\n" + f"{orjson.dumps(context).decode('utf-8')}\n" + "*/\n" + f"CREATE TABLE {table_name} (\n {column_ddl}\n);" + ) + + +def test_generates_count_ranking_for_orders_without_using_sales_sum(): + document = _schema_document( + "analytics_model", + [ + { + "name": "division_name", + "data_type": "VARCHAR", + "comment": "Business division", + "roles": ["dimension_candidate"], + }, + { + "name": "sales_value", + "data_type": "DOUBLE", + "comment": "Sales amount", + "roles": ["numeric_measure_candidate"], + }, + ], + ) + + sql = generate_grounded_sql( + "Which divisions are generating the most orders?", [document] + ) + + assert "COUNT(*) AS TotalOrders" in sql + assert "SUM(sales_value)" not in sql + assert "GROUP BY\n division_name" in sql + assert "ORDER BY\n TotalOrders DESC" in sql + + +def test_total_orders_prefers_count_even_when_numeric_measure_exists(): + document = _schema_document( + "analytics_model", + [ + { + "name": "country_name", + "data_type": "VARCHAR", + "comment": "Country", + "roles": ["dimension_candidate"], + }, + { + "name": "sales_value", + "data_type": "DOUBLE", + "comment": "Sales amount", + "roles": ["numeric_measure_candidate"], + }, + ], + ) + + sql = generate_grounded_sql("Show top 5 countries by total orders.", [document]) + + assert "COUNT(*) AS TotalOrders" in sql + assert "SUM(sales_value)" not in sql + assert "LIMIT 5" in sql + + +def test_generates_filtered_detail_sql_with_deployed_identifiers_only(): + document = _schema_document( + "analytics_model", + [ + { + "name": "customer_country", + "data_type": "VARCHAR", + "comment": "Country", + "roles": ["dimension_candidate"], + }, + { + "name": "order_date", + "data_type": "DATE", + "comment": "Order placed date", + "roles": ["date_time_candidate"], + }, + { + "name": "order_number", + "data_type": "VARCHAR", + "comment": "Order number", + "roles": ["identifier_candidate"], + }, + { + "name": "amount", + "data_type": "DOUBLE", + "comment": "Transaction amount", + "roles": ["numeric_measure_candidate"], + }, + ], + ) + + sql = generate_grounded_sql("show order placed from the country India", [document]) + + assert "FROM\n analytics_model" in sql + assert "customer_country = 'India'" in sql + assert "order_date" in sql + assert "LIMIT 500" in sql + + +def test_generates_sales_comparison_by_matching_dimension(): + document = _schema_document( + "analytics_model", + [ + { + "name": "market_segment", + "data_type": "VARCHAR", + "comment": "Domestic or international market", + "roles": ["dimension_candidate"], + }, + { + "name": "sales_value", + "data_type": "DOUBLE", + "comment": "Sales value", + "roles": ["numeric_measure_candidate"], + }, + ], + ) + + sql = generate_grounded_sql( + "Compare sales between domestic and international markets.", [document] + ) + + assert "market_segment" in sql + assert "SUM(sales_value) AS TotalValue" in sql + assert "GROUP BY\n market_segment" in sql + + +def test_generates_average_by_requested_dimension(): + document = _schema_document( + "analytics_model", + [ + { + "name": "product_type", + "data_type": "VARCHAR", + "comment": "Product type", + "roles": ["dimension_candidate"], + }, + { + "name": "sales_value", + "data_type": "DOUBLE", + "comment": "Sales value", + "roles": ["numeric_measure_candidate"], + }, + ], + ) + + sql = generate_grounded_sql( + "Show average sales value by product type.", [document] + ) + + assert "product_type" in sql + assert "AVG(sales_value) AS AverageValue" in sql + assert "GROUP BY\n product_type" in sql + + +@pytest.mark.asyncio +async def test_generation_fast_path_does_not_call_llm_when_grounded(): + document = _schema_document( + "analytics_model", + [ + { + "name": "country_name", + "data_type": "VARCHAR", + "comment": "Country", + "roles": ["dimension_candidate"], + }, + { + "name": "order_key", + "data_type": "VARCHAR", + "comment": "Order identifier", + "roles": ["identifier_candidate"], + }, + ], + ) + + async def failing_generator(**_): + raise AssertionError("LLM generator should not be called") + + result = await generate_sql( + prompt={"prompt": "unused"}, + query="Show top 5 countries by total orders.", + documents=[document], + generator=failing_generator, + generator_name="test-model", + ) + + assert result["replies"] + assert "COUNT(*) AS TotalOrders" in result["replies"][0] From 4007757dc80c40b15694af2312d862fa7a828890 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 19:43:05 +0530 Subject: [PATCH 0838/1087] Revert "Add grounded SQL fast path" This reverts commit 3ad59687c8a00749161c75438b462bd618c2eb53. --- .../generation/followup_sql_generation.py | 12 - .../pipelines/generation/sql_generation.py | 12 - .../generation/utils/deterministic_sql.py | 498 ------------------ .../generation/test_deterministic_sql.py | 223 -------- 4 files changed, 745 deletions(-) delete mode 100644 wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py delete mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py 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 58d86b5b2e..f3f4ae4d18 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -2,7 +2,6 @@ import sys from typing import Any -import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder @@ -12,7 +11,6 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata -from src.pipelines.generation.utils.deterministic_sql import generate_grounded_sql from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -144,21 +142,11 @@ def prompt( @trace_cost async def generate_sql_in_followup( prompt: dict, - query: str, - documents: list[str], generator: Any, histories: list[AskHistory], generator_name: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - deterministic_sql = generate_grounded_sql(query, documents) - if deterministic_sql: - logger.info("Follow-Up SQL Generation used deterministic grounded SQL fast path.") - return { - "replies": [orjson.dumps({"sql": deterministic_sql}).decode("utf-8")], - "metadata": [{"finish_reason": "deterministic_grounded_sql"}], - }, generator_name - history_messages = construct_ask_history_messages(histories) current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) return await generator( diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 11f4055d92..8de6a23961 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -2,7 +2,6 @@ import sys from typing import Any -import orjson from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder @@ -12,7 +11,6 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.common import clean_up_new_lines, retrieve_metadata -from src.pipelines.generation.utils.deterministic_sql import generate_grounded_sql from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -144,20 +142,10 @@ def prompt( @trace_cost async def generate_sql( prompt: dict, - query: str, - documents: list[str], generator: Any, generator_name: str, sql_knowledge: SqlKnowledge | None = None, ) -> dict: - deterministic_sql = generate_grounded_sql(query, documents) - if deterministic_sql: - logger.info("SQL Generation used deterministic grounded SQL fast path.") - return { - "replies": [orjson.dumps({"sql": deterministic_sql}).decode("utf-8")], - "metadata": [{"finish_reason": "deterministic_grounded_sql"}], - }, generator_name - current_system_prompt = get_sql_generation_system_prompt(sql_knowledge) return await generator( prompt=prompt.get("prompt"), current_system_prompt=current_system_prompt diff --git a/wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py b/wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py deleted file mode 100644 index b97f717439..0000000000 --- a/wren-ai-service/src/pipelines/generation/utils/deterministic_sql.py +++ /dev/null @@ -1,498 +0,0 @@ -import calendar -import re -from dataclasses import dataclass -from datetime import date, timedelta - -import orjson - - -_SEMANTIC_CONTEXT_PATTERN = re.compile( - r"WREN RETRIEVED SEMANTIC CONTEXT\s*\n(\{.*?\})\s*\n", re.DOTALL -) -_CREATE_TABLE_PATTERN = re.compile( - r"CREATE\s+TABLE\s+([A-Za-z_][A-Za-z0-9_]*|\"[^\"]+\")\s*\((.*?)\)", - re.IGNORECASE | re.DOTALL, -) -_WORD_PATTERN = re.compile(r"[A-Za-z][A-Za-z0-9]*") -_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") - -_STOP_TERMS = { - "a", - "an", - "and", - "are", - "as", - "between", - "by", - "from", - "for", - "in", - "is", - "most", - "of", - "on", - "per", - "show", - "the", - "to", - "using", - "what", - "which", - "with", -} -_DETAIL_TERMS = {"show", "list", "display", "find", "get"} -_RANKING_TERMS = {"top", "bottom", "highest", "lowest", "most", "least"} -_COUNT_TERMS = {"count", "counts", "number", "orders", "order"} -_SUM_TERMS = { - "amount", - "cost", - "qty", - "quantity", - "revenue", - "sales", - "sold", - "sum", - "total", - "value", -} -_VALUE_MEASURE_TERMS = { - "amount", - "cost", - "qty", - "quantity", - "revenue", - "sales", - "sold", - "value", -} -_AVG_TERMS = {"average", "avg", "mean"} -_MIN_TERMS = {"minimum", "min", "lowest"} -_MAX_TERMS = {"maximum", "max", "highest", "most"} -_MONTHS = { - month.lower(): index for index, month in enumerate(calendar.month_name) if month -} - - -@dataclass(frozen=True) -class _Column: - name: str - data_type: str = "" - semantic_text: str = "" - roles: tuple[str, ...] = () - - -@dataclass(frozen=True) -class _Table: - name: str - semantic_text: str - columns: tuple[_Column, ...] - - -def generate_grounded_sql(query: str, documents: list[str]) -> str | None: - tables = [ - table for document in documents for table in [_parse_table(document)] if table - ] - if not query or not tables: - return None - - query_terms = _terms(query) - if not query_terms: - return None - - table = max(tables, key=lambda candidate: _table_score(candidate, query_terms)) - if _table_score(table, query_terms) <= 0: - return None - - filters = _build_filters(query, table) - aggregate_sql = _build_aggregate_sql(query, query_terms, table, filters) - if aggregate_sql: - return aggregate_sql - - return _build_detail_sql(query, query_terms, table, filters) - - -def _parse_table(document: str) -> _Table | None: - context_match = _SEMANTIC_CONTEXT_PATTERN.search(document or "") - if context_match: - try: - context = orjson.loads(context_match.group(1)) - contract = context.get("sql_identifier_contract") or {} - table_name = contract.get("sql_table_name_use_exactly") - columns = [] - allowed_columns = set(contract.get("sql_column_names_use_exactly") or []) - for column in context.get("columns") or []: - column_name = column.get("sql_column_name_use_exactly") - if not column_name or ( - allowed_columns and column_name not in allowed_columns - ): - continue - columns.append( - _Column( - name=column_name, - data_type=str(column.get("data_type") or ""), - semantic_text=str( - column.get("semantic_context_not_sql_identifier") or "" - ), - roles=tuple(column.get("semantic_roles_not_identifiers") or ()), - ) - ) - - if table_name and columns: - return _Table( - name=table_name, - semantic_text=str( - (context.get("semantic_context_not_sql_identifiers") or {}).get( - "description", "" - ) - ), - columns=tuple(columns), - ) - except orjson.JSONDecodeError: - pass - - ddl_match = _CREATE_TABLE_PATTERN.search(document or "") - if not ddl_match: - return None - - columns = [] - for raw_column in ddl_match.group(2).split(","): - pieces = raw_column.strip().split() - if len(pieces) >= 2: - columns.append( - _Column(name=_unquote_identifier(pieces[0]), data_type=pieces[1]) - ) - - if not columns: - return None - - return _Table( - name=_unquote_identifier(ddl_match.group(1)), - semantic_text="", - columns=tuple(columns), - ) - - -def _terms(value: str) -> set[str]: - terms = set() - for token in _WORD_PATTERN.findall(value or ""): - normalized = token.lower() - terms.add(normalized) - if normalized.endswith("ies") and len(normalized) > 4: - terms.add(f"{normalized[:-3]}y") - elif normalized.endswith("s") and len(normalized) > 3: - terms.add(normalized[:-1]) - return terms - - -def _identifier_terms(value: str) -> set[str]: - spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", value or "") - spaced = spaced.replace("_", " ") - return _terms(spaced) - - -def _column_terms(column: _Column) -> set[str]: - return ( - _identifier_terms(column.name) - | _terms(column.semantic_text) - | {term for role in column.roles for term in _identifier_terms(role)} - ) - - -def _table_score(table: _Table, query_terms: set[str]) -> int: - score = len( - (_identifier_terms(table.name) | _terms(table.semantic_text)) & query_terms - ) - column_scores = sorted( - (len(_column_terms(column) & query_terms) for column in table.columns), - reverse=True, - ) - return score + sum(column_scores[:6]) - - -def _is_numeric(column: _Column) -> bool: - data_type = column.data_type.lower() - return "numeric_measure_candidate" in column.roles or any( - token in data_type - for token in ( - "int", - "decimal", - "numeric", - "number", - "double", - "float", - "real", - "money", - ) - ) - - -def _is_datetime(column: _Column) -> bool: - data_type = column.data_type.lower() - return "date_time_candidate" in column.roles or any( - token in data_type for token in ("date", "time", "timestamp") - ) - - -def _is_dimension(column: _Column) -> bool: - return not _is_numeric(column) and not _is_datetime(column) - - -def _best_columns( - table: _Table, - query_terms: set[str], - predicate, - *, - limit: int = 3, -) -> list[_Column]: - scored = [] - for column in table.columns: - if not predicate(column): - continue - score = len(_column_terms(column) & query_terms) - if score: - scored.append((score, column)) - return [ - column - for _, column in sorted(scored, key=lambda item: item[0], reverse=True)[:limit] - ] - - -def _best_measure(table: _Table, query_terms: set[str]) -> _Column | None: - measures = _best_columns(table, query_terms, _is_numeric, limit=1) - if measures: - return measures[0] - - numeric_columns = [column for column in table.columns if _is_numeric(column)] - return numeric_columns[0] if numeric_columns else None - - -def _best_date(table: _Table, query_terms: set[str]) -> _Column | None: - date_columns = _best_columns(table, query_terms, _is_datetime, limit=1) - if date_columns: - return date_columns[0] - - date_columns = [column for column in table.columns if _is_datetime(column)] - return date_columns[0] if date_columns else None - - -def _query_requests_aggregate(query_terms: set[str]) -> bool: - aggregate_terms = ( - _RANKING_TERMS | _SUM_TERMS | _AVG_TERMS | _MIN_TERMS | _MAX_TERMS - ) - return bool(query_terms & aggregate_terms) or ( - "by" in query_terms or "per" in query_terms or "compare" in query_terms - ) - - -def _requested_aggregate( - query_terms: set[str], measure: _Column | None -) -> tuple[str, str]: - if query_terms & _AVG_TERMS: - return "AVG", "AverageValue" - if query_terms & _MIN_TERMS: - return "MIN", "MinimumValue" - if query_terms & _COUNT_TERMS and not (query_terms & _VALUE_MEASURE_TERMS): - return "COUNT", "TotalOrders" - if query_terms & _MAX_TERMS and measure and not (query_terms & {"most", "top"}): - return "MAX", "MaximumValue" - if measure: - return "SUM", "TotalValue" - return "COUNT", "TotalCount" - - -def _build_aggregate_sql( - query: str, - query_terms: set[str], - table: _Table, - filters: list[str], -) -> str | None: - if not _query_requests_aggregate(query_terms): - return None - - dimensions = _best_columns(table, query_terms, _is_dimension, limit=2) - if not dimensions and not (query_terms & _COUNT_TERMS): - return None - - measure = _best_measure(table, query_terms) - function_name, alias = _requested_aggregate(query_terms, measure) - if function_name == "COUNT": - aggregate_expression = f"COUNT(*) AS {_sql_identifier(alias)}" - elif measure: - aggregate_expression = ( - f"{function_name}({_sql_identifier(measure.name)}) AS {_sql_identifier(alias)}" - ) - else: - return None - - select_items = [_sql_identifier(column.name) for column in dimensions] - select_items.append(aggregate_expression) - clauses = [ - "SELECT", - " " + ",\n ".join(select_items), - "FROM", - f" {_sql_identifier(table.name)}", - ] - if filters: - clauses.extend(["WHERE", " " + "\n AND ".join(filters)]) - if dimensions: - group_by = ", ".join(_sql_identifier(column.name) for column in dimensions) - clauses.extend(["GROUP BY", f" {group_by}"]) - if query_terms & _RANKING_TERMS or "top" in query_terms or "bottom" in query_terms: - direction = "ASC" if query_terms & {"bottom", "lowest", "least"} else "DESC" - clauses.extend(["ORDER BY", f" {_sql_identifier(alias)} {direction}"]) - clauses.append(f"LIMIT {_top_limit(query)}") - elif "top" in query_terms: - clauses.append(f"LIMIT {_top_limit(query)}") - - return "\n".join(clauses) - - -def _build_detail_sql( - query: str, - query_terms: set[str], - table: _Table, - filters: list[str], -) -> str | None: - if not filters and not (query_terms & _DETAIL_TERMS): - return None - - selected = [] - for predicate in (_is_datetime, _is_dimension, _is_numeric): - for column in _best_columns(table, query_terms, predicate, limit=4): - if column.name not in {item.name for item in selected}: - selected.append(column) - - if len(selected) < 3: - for column in table.columns: - if column.name not in {item.name for item in selected}: - selected.append(column) - if len(selected) >= 6: - break - - selected = selected[:6] - if not selected: - return None - - clauses = [ - "SELECT", - " " + ",\n ".join(_sql_identifier(column.name) for column in selected), - "FROM", - f" {_sql_identifier(table.name)}", - ] - if filters: - clauses.extend(["WHERE", " " + "\n AND ".join(filters)]) - - date_column = _best_date(table, query_terms) - if date_column: - clauses.extend(["ORDER BY", f" {_sql_identifier(date_column.name)} DESC"]) - clauses.append("LIMIT 500") - return "\n".join(clauses) - - -def _build_filters(query: str, table: _Table) -> list[str]: - query_terms = _terms(query) - filters: list[str] = [] - - date_range = _date_range(query) - date_column = _best_date(table, query_terms) - if date_range and date_column: - start, end = date_range - date_identifier = _sql_identifier(date_column.name) - filters.append(f"{date_identifier} >= '{start.isoformat()}'") - filters.append(f"{date_identifier} < '{end.isoformat()}'") - - literal_filter = _literal_dimension_filter(query, table) - if literal_filter: - filters.append(literal_filter) - - return filters - - -def _literal_dimension_filter(query: str, table: _Table) -> str | None: - if _query_requests_aggregate(_terms(query)): - return None - - query_words = _WORD_PATTERN.findall(query or "") - lowered_words = [word.lower() for word in query_words] - - for column in table.columns: - if not _is_dimension(column): - continue - column_terms = _identifier_terms(column.name) - for index, word in enumerate(lowered_words[:-1]): - if word not in column_terms: - continue - literal_words = [] - for next_word in query_words[index + 1 : index + 4]: - normalized = next_word.lower() - if normalized in _STOP_TERMS or normalized in _MONTHS: - break - literal_words.append(next_word) - if literal_words: - literal = " ".join(literal_words) - return f"{_sql_identifier(column.name)} = '{_escape_literal(literal)}'" - - return None - - -def _date_range(query: str, today: date | None = None) -> tuple[date, date] | None: - current = today or date.today() - normalized = (query or "").lower() - - if "today" in normalized: - return current, current + timedelta(days=1) - if "yesterday" in normalized: - previous = current - timedelta(days=1) - return previous, current - if "last week" in normalized: - return current - timedelta(days=7), current + timedelta(days=1) - if "last month" in normalized: - start = _add_months(date(current.year, current.month, 1), -1) - end = date(current.year, current.month, 1) - return start, end - - last_months = re.search(r"\blast\s+(\d{1,2})\s+months?\b", normalized) - if last_months: - return _add_months(current, -int(last_months.group(1))), current + timedelta( - days=1 - ) - - for month_name, month_index in _MONTHS.items(): - match = re.search(rf"\b{month_name}\s+(\d{{4}})\b", normalized) - if match: - start = date(int(match.group(1)), month_index, 1) - return start, _add_months(start, 1) - - return None - - -def _add_months(value: date, months: int) -> date: - month = value.month - 1 + months - year = value.year + month // 12 - month = month % 12 + 1 - day = min(value.day, calendar.monthrange(year, month)[1]) - return date(year, month, day) - - -def _top_limit(query: str) -> int: - match = re.search(r"\btop\s+(\d{1,3})\b", query.lower()) - if not match: - return 10 - return min(max(int(match.group(1)), 1), 500) - - -def _sql_identifier(identifier: str) -> str: - if _IDENTIFIER_PATTERN.match(identifier): - return identifier - return '"' + identifier.replace('"', '""') + '"' - - -def _unquote_identifier(identifier: str) -> str: - identifier = identifier.strip() - if identifier.startswith('"') and identifier.endswith('"'): - return identifier[1:-1].replace('""', '"') - return identifier - - -def _escape_literal(value: str) -> str: - return value.replace("'", "''") diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py b/wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py deleted file mode 100644 index 97251723f9..0000000000 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_deterministic_sql.py +++ /dev/null @@ -1,223 +0,0 @@ -import pytest - -from src.pipelines.generation.sql_generation import generate_sql -from src.pipelines.generation.utils.deterministic_sql import generate_grounded_sql - - -def _schema_document(table_name: str, columns: list[dict]) -> str: - import orjson - - context = { - "object_type": "model", - "sql_identifier_contract": { - "sql_table_name_use_exactly": table_name, - "sql_column_names_use_exactly": [column["name"] for column in columns], - }, - "semantic_context_not_sql_identifiers": { - "description": "Synthetic sales and order analytics model.", - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": column["data_type"], - "semantic_context_not_sql_identifier": column.get("comment", ""), - "semantic_roles_not_identifiers": column.get("roles", []), - } - for column in columns - ], - } - column_ddl = ",\n ".join( - f"{column['name']} {column['data_type']}" for column in columns - ) - return ( - "/*\n" - "WREN RETRIEVED SEMANTIC CONTEXT\n" - f"{orjson.dumps(context).decode('utf-8')}\n" - "*/\n" - f"CREATE TABLE {table_name} (\n {column_ddl}\n);" - ) - - -def test_generates_count_ranking_for_orders_without_using_sales_sum(): - document = _schema_document( - "analytics_model", - [ - { - "name": "division_name", - "data_type": "VARCHAR", - "comment": "Business division", - "roles": ["dimension_candidate"], - }, - { - "name": "sales_value", - "data_type": "DOUBLE", - "comment": "Sales amount", - "roles": ["numeric_measure_candidate"], - }, - ], - ) - - sql = generate_grounded_sql( - "Which divisions are generating the most orders?", [document] - ) - - assert "COUNT(*) AS TotalOrders" in sql - assert "SUM(sales_value)" not in sql - assert "GROUP BY\n division_name" in sql - assert "ORDER BY\n TotalOrders DESC" in sql - - -def test_total_orders_prefers_count_even_when_numeric_measure_exists(): - document = _schema_document( - "analytics_model", - [ - { - "name": "country_name", - "data_type": "VARCHAR", - "comment": "Country", - "roles": ["dimension_candidate"], - }, - { - "name": "sales_value", - "data_type": "DOUBLE", - "comment": "Sales amount", - "roles": ["numeric_measure_candidate"], - }, - ], - ) - - sql = generate_grounded_sql("Show top 5 countries by total orders.", [document]) - - assert "COUNT(*) AS TotalOrders" in sql - assert "SUM(sales_value)" not in sql - assert "LIMIT 5" in sql - - -def test_generates_filtered_detail_sql_with_deployed_identifiers_only(): - document = _schema_document( - "analytics_model", - [ - { - "name": "customer_country", - "data_type": "VARCHAR", - "comment": "Country", - "roles": ["dimension_candidate"], - }, - { - "name": "order_date", - "data_type": "DATE", - "comment": "Order placed date", - "roles": ["date_time_candidate"], - }, - { - "name": "order_number", - "data_type": "VARCHAR", - "comment": "Order number", - "roles": ["identifier_candidate"], - }, - { - "name": "amount", - "data_type": "DOUBLE", - "comment": "Transaction amount", - "roles": ["numeric_measure_candidate"], - }, - ], - ) - - sql = generate_grounded_sql("show order placed from the country India", [document]) - - assert "FROM\n analytics_model" in sql - assert "customer_country = 'India'" in sql - assert "order_date" in sql - assert "LIMIT 500" in sql - - -def test_generates_sales_comparison_by_matching_dimension(): - document = _schema_document( - "analytics_model", - [ - { - "name": "market_segment", - "data_type": "VARCHAR", - "comment": "Domestic or international market", - "roles": ["dimension_candidate"], - }, - { - "name": "sales_value", - "data_type": "DOUBLE", - "comment": "Sales value", - "roles": ["numeric_measure_candidate"], - }, - ], - ) - - sql = generate_grounded_sql( - "Compare sales between domestic and international markets.", [document] - ) - - assert "market_segment" in sql - assert "SUM(sales_value) AS TotalValue" in sql - assert "GROUP BY\n market_segment" in sql - - -def test_generates_average_by_requested_dimension(): - document = _schema_document( - "analytics_model", - [ - { - "name": "product_type", - "data_type": "VARCHAR", - "comment": "Product type", - "roles": ["dimension_candidate"], - }, - { - "name": "sales_value", - "data_type": "DOUBLE", - "comment": "Sales value", - "roles": ["numeric_measure_candidate"], - }, - ], - ) - - sql = generate_grounded_sql( - "Show average sales value by product type.", [document] - ) - - assert "product_type" in sql - assert "AVG(sales_value) AS AverageValue" in sql - assert "GROUP BY\n product_type" in sql - - -@pytest.mark.asyncio -async def test_generation_fast_path_does_not_call_llm_when_grounded(): - document = _schema_document( - "analytics_model", - [ - { - "name": "country_name", - "data_type": "VARCHAR", - "comment": "Country", - "roles": ["dimension_candidate"], - }, - { - "name": "order_key", - "data_type": "VARCHAR", - "comment": "Order identifier", - "roles": ["identifier_candidate"], - }, - ], - ) - - async def failing_generator(**_): - raise AssertionError("LLM generator should not be called") - - result = await generate_sql( - prompt={"prompt": "unused"}, - query="Show top 5 countries by total orders.", - documents=[document], - generator=failing_generator, - generator_name="test-model", - ) - - assert result["replies"] - assert "COUNT(*) AS TotalOrders" in result["replies"][0] From a17778188a8a3d9f8c6b3a846682a746ae1650dd Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 19:59:31 +0530 Subject: [PATCH 0839/1087] Refine grounded SQL prompt flow --- .../generation/followup_sql_generation.py | 17 ++- .../src/pipelines/generation/sql_answer.py | 14 +- .../pipelines/generation/sql_generation.py | 18 ++- .../src/pipelines/generation/utils/sql.py | 121 ++++++------------ 4 files changed, 68 insertions(+), 102 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 f3f4ae4d18..317ecebf0c 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -82,11 +82,18 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. -Generate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. -When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. +Generate one Wren SQL query that answers the full follow-up request using only DATABASE SCHEMA and SQL FUNCTIONS. +Generate an intent-shaped query, not a table preview. +Use schema descriptions, aliases, display labels, metrics, calculated fields, and relationships only to understand meaning. +The SQL must include every supported requested part: subject, entity, filters, timeframe, grouping, measure, ordering, and limit. +If a required part of the request is not grounded by an exact deployed schema object, column, relationship, metric, or supported function, return null for sql. +Do not answer a specific analytical question with a broad table preview or with an unrelated nearby table. +Do not ignore a literal filter value from the user; apply it to the exact schema field representing that filter concept, or return null when that field is unavailable. +For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. +For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. +For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. +For aggregate, trend, ranking, or grouped requests, aggregate exact numeric_measure_candidate columns or count rows as appropriate for the user's requested measure. +Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index b1eae4c1cb..c5eefd96ef 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -19,27 +19,29 @@ sql_to_answer_system_prompt = """ ### TASK -You are a data analyst that great at answering non-technical user's questions based on the data, sql so that even non technical users can easily understand. -Please answer the user's question in concise and clear manner in Markdown format. +You are a data analyst who explains query results to non-technical users. +Answer the user's question clearly in Markdown using only the provided SQL result data. ### INSTRUCTIONS 1. Read the user's question and understand the user's intention. 2. Read the sql and understand the data. 3. Make sure the answer is aimed for non-technical users, so don't mention any technical terms such as SQL syntax. -4. Generate a concise and clear answer in string format to answerthe user's question based on the data and sql. -5. If answer is in list format, only list top few examples, and tell users there are more results omitted. +4. Generate a clear answer that directly addresses the question before adding supporting details. +5. If the result contains ranked or grouped rows, explain what the leading rows mean and why they answer the question. 6. Answer must be in the same language user specified. 7. Do not include ```markdown or ``` in the answer. 8. If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. 9. Use only the columns and result rows provided in Data. Do not invent, duplicate, reorder, aggregate, rank, or label rows unless that operation is directly represented by the provided SQL result. 10. If the Data has aggregate rows, summarize those exact aggregate rows instead of describing them as separate top examples. -11. If the Data is empty, state that no matching records were returned. +11. If the Data is empty, state that no matching records were returned and mention that the SQL result has no rows. 12. Data rows are records already mapped by column name. Answer from the record values; do not describe the underlying data structure. +13. Mention important limitations visible in the result, such as null grouping values, ties, or a result limited to a subset of rows. +14. Keep the answer detailed enough for a business user to understand the conclusion, the supporting figures, and any caveat from the returned rows. ### OUTPUT FORMAT -Please provide your response in proper Markdown stringformat. +Please provide your response in proper Markdown string format. """ sql_to_answer_user_prompt_template = """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 8de6a23961..492069bdc3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -76,12 +76,18 @@ ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. -If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. -Generate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. -When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. +Generate one Wren SQL query that answers the full user request using only DATABASE SCHEMA and SQL FUNCTIONS. +Generate an intent-shaped query, not a table preview. +Use schema descriptions, aliases, display labels, metrics, calculated fields, and relationships only to understand meaning. +The SQL must include every supported requested part: subject, entity, filters, timeframe, grouping, measure, ordering, and limit. +If a required part of the request is not grounded by an exact deployed schema object, column, relationship, metric, or supported function, return null for sql. +Do not answer a specific analytical question with a broad table preview or with an unrelated nearby table. +Do not ignore a literal filter value from the user; apply it to the exact schema field representing that filter concept, or return null when that field is unavailable. +For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. +For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. +For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. +For aggregate, trend, ranking, or grouped requests, aggregate exact numeric_measure_candidate columns or count rows as appropriate for the user's requested measure. +Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REQUEST ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ec6ad268e7..89ac2bd21e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -809,97 +809,48 @@ async def _classify_generation_result( _MANDATORY_SQL_GROUNDING_RULES = """ ### MANDATORY SQL GROUNDING RULES ### -- Treat the retrieved semantic context as the only authoritative source for this request. Do not use pretrained knowledge, common warehouse schemas, example schemas, or memorized business definitions as executable truth. -- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. -- Use only deployed semantic models, views, metrics, relationships, and columns that are present in the retrieved DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, EXECUTABLE WREN IDENTIFIER CATALOG, SQL FUNCTIONS, or current USER INSTRUCTIONS. -- Before generating SQL, silently validate that every model, column, metric, relationship, join path, filter field, grouping field, ordering field, and SQL function is present in the retrieved context. Generate SQL only after this validation succeeds. -- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. -- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. -- Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. -- Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. -- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. -- When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. -- In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. -- Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. -- Values under column_role_hints_not_identifiers are semantic roles only. Use them to decide whether an exact declared column can serve as a date/time field, measure, identifier, or dimension, but copy executable column names only from the columns list or DDL. -- When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. -- The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. -- Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. -- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. -- Never invent string literal filter values. Use a string value in WHERE, HAVING, CASE, or JOIN conditions only when that value is explicitly present in the user's current question or grounded by a current USER INSTRUCTION. For relative time requests, bounded date literals may be generated only from the requested timeframe and an exact date/time column. -- Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. -- If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. -- When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. -- Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. -- Prefer a single table, view, or metric that already contains the requested fields. Do not join tables just because they were retrieved together. -- When using multiple tables to combine fields into the same output row, join only through the exact FOREIGN KEY constraints shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, return null for sql or use one schema object that already contains the requested fields. +- Treat the retrieved semantic context as the only authoritative source for this request. +- Do not use pretrained knowledge, common warehouse schemas, example schemas, or memorized business definitions as executable truth. +- Use the retrieved DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT as the only executable context. +- Before generating SQL, silently validate that every model, column, metric, relationship, join path, filter field, grouping field, ordering field, and SQL function is present in the retrieved context. +- Use comments, aliases, descriptions, display labels, metrics, calculated fields, and relationships only to understand business meaning. +- Use column_role_hints_not_identifiers only as semantic hints for choosing exact declared columns; date_time_candidate, numeric_measure_candidate, identifier_candidate, and dimension_candidate are never SQL identifiers. +- Copy executable table, column, metric, and relationship identifiers exactly from DATABASE SCHEMA. Do not create identifiers from user wording, descriptions, samples, history, physical names, lineage names, or error messages. +- The SQL must answer every supported part of the user's request: requested subject, requested entity, requested filter value, timeframe, grouping, measure, ordering, and limit. +- If the user asks for a filtered result, include the filter only when the filtered concept is represented by an exact schema field. If the filter field is not present, return null for sql instead of ignoring the filter. +- If the user provides a literal filter value, use only that provided value. Do not invent, translate, or substitute filter values. +- If the user asks "which", "who", or "what" for a ranked entity, select and group by the exact schema field that represents that requested entity. Do not replace the requested entity with a context field or unrelated dimension. +- Use row counting for record or entity volume questions when no numeric business measure is requested. Use numeric measures only when the question asks for a value, amount, quantity, rate, cost, or other declared measure. +- For analytical questions, return dimensions plus the requested measure expression or metric field. Do not return a raw table preview. +- For aggregate, ranking, grouped, or trend questions, produce an analytical query shape. +- For detail-list questions, return only the fields needed to identify and describe the requested records, plus requested filters and timeframes. +- Do not answer a timeframe request with an unfiltered table scan. +- Prefer one model, view, or metric that already contains the requested fields. Do not join tables just because they were retrieved together. Do not invent join predicates from similar column names. Join only through relationships declared in DATABASE SCHEMA. +- If multiple retrieved schema objects are needed for the same result, use them only when the required columns and relationship path are present. - If multiple semantic interpretations exist and the retrieved context does not make one interpretation authoritative, return null for sql instead of choosing one. -- When the same requested result can be answered from multiple schema objects with compatible columns or metrics, include all relevant schema objects by combining separate result rows with UNION ALL instead of choosing only one object. -- Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. -- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and an exact relationship path. Do not invent join predicates from similar column names. -- Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. -- SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -- Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. -- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and the predicate can be expressed with normal SQL comparison syntax or an operation listed in SQL FUNCTIONS. Do not compare text fields to date functions. -- For explicit month/year or relative timeframe requests, prefer a bounded range predicate on one exact date_time_candidate column when available. The lower bound is inclusive and the upper bound is exclusive. Do not answer a timeframe request with an unfiltered table scan. -- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. -- If SQL execution or validation fails, repair the query only when the repair can be verified using the same retrieved DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS. Never introduce a new schema object during repair. -- If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. -- For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. -- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. -- If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. -- If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. -- If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. -- Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. -- Do not answer a specific business question with a broad table scan. The SQL shape must match the user's requested output columns, filters, groupings, measures, joins, ordering, and limits. -- For analytical or metric questions, select only the requested dimensions and measures. Use declared metric columns, calculated fields, relationship paths, and schema-grounded aggregate expressions. If the required metric components are not grounded, return null for sql instead of returning raw rows. -- For questions asking total, count, average, minimum, maximum, ratio, per, by, top, bottom, highest, lowest, trend, month, week, year, or ranking, produce an analytical query shape: select exact dimension columns or date buckets, aggregate exact numeric_measure_candidate columns or count rows, GROUP BY every non-aggregated selected expression, ORDER BY the selected aggregate alias when ranking, and apply LIMIT only when requested. -- If the question asks for an entity list with a timeframe or filter but no metric, select only the entity identifier, relevant dimensions, and exact date/time column needed by the request; include the requested WHERE predicate. Do not select every column from the table. +- If SQL execution or correction is needed, repair the query only when the repair can be verified using DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS. Never introduce a new schema object during repair. +- Return null for sql when the retrieved schema does not ground a required subject, entity, filter field, timeframe field, measure, or relationship. +- Generate Wren SQL only, using supported functions from SQL FUNCTIONS when functions are needed. """ _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### -- ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. -- ONLY USE the tables and columns mentioned in the database schema. -- Never use "*" in the SELECT list. Select explicit deployed schema columns needed for the question. When the user asks for all records, all rows, all users, all orders, or similar, treat "all" as row scope and still select explicit columns relevant to the requested entity or metric. -- ONLY CHOOSE columns belong to the tables mentioned in the database schema. -- DON'T INCLUDE comments in the generated SQL query. -- Use JOIN only when selected columns come from multiple tables and DATABASE SCHEMA declares the exact FOREIGN KEY relationship needed for the join. Do not invent join predicates from similar-looking column names. -- PREFER USING CTEs over subqueries. -- When generating SQL query, always: - - Put double quotes around column and table names. - - Use Wren SQL identifier quoting with double quotes only; the engine rewrite step converts grounded Wren SQL to the active connector dialect. - - Put single quotes around string literals. - - Never quote numeric literals. -- Generate Wren SQL syntax only, not connector-specific SQL syntax. -- Never use SELECT TOP, TOP(...), FETCH FIRST, square-bracket identifiers, or backtick identifiers. For top or limit requests, sort with ORDER BY and put LIMIT at the end of the query. -- Preserve every deployed table and column identifier exactly as it appears in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, including spaces, digits, underscores, case, and punctuation, then wrap that exact identifier in double quotes in SQL. -- Do not convert deployed identifiers into display-friendly variants by replacing spaces with underscores, removing prefixes, changing case, shortening names, or expanding abbreviations. -- For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. -- For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only date/time functions and casts whose exact syntax is provided in SQL FUNCTIONS for this request. -- If the question asks for a specific or relative date, generate a bounded date/time filter only when the exact date/time schema column is available and the predicate can be expressed with normal SQL comparison syntax or a SQL FUNCTIONS-supported operation. If either the column or required operation is missing, do not invent a field or function. -- When DATABASE SCHEMA includes column_role_hints_not_identifiers, use date_time_candidate, numeric_measure_candidate, identifier_candidate, and dimension_candidate roles to map the question intent to exact declared columns. These role names are never executable SQL identifiers. -- For explicit calendar month and year requests, use an inclusive lower bound and exclusive upper bound on the exact date/time column, rather than formatting the column into text. -- USE THE VIEW TO SIMPLIFY THE QUERY. -- DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. -- For metric-style requests, the final SELECT list must expose the requested dimension columns and measure expressions or metric fields. Do not return every raw column from a retrieved model as a substitute for the requested metric. -- For aggregate, ranking, or "by" requests, do not add unrelated string filters to make the SQL look specific. If the user did not provide a filter value, leave it out. -- For total, count, average, minimum, maximum, per, by, trend, top, bottom, highest, lowest, or ranking requests, the final SQL must include the requested aggregate expression or metric field, GROUP BY required dimensions, ORDER BY required ranking expression, and LIMIT only when requested. A raw row list is not a valid answer. -- For record-list requests with a filter or timeframe, the final SQL must include the requested WHERE predicate and only the columns needed to identify and describe the matching records. -- Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. -- Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. -- DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. -- DON'T USE "FILTER(WHERE )" clause in the generated SQL query. -- DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. -- DON'T USE "EXTRACT()" function with INTERVAL data types as arguments -- DON'T USE INTERVAL or generate INTERVAL-like expression in the generated SQL query. -- DON'T USE "TO_CHAR" function in the generated SQL query. -- DON'T USE unsupported statistical, date/time, or formatting functions. If SQL FUNCTIONS does not list a function needed by the requested intent, omit the function-dependent part. If that function is required to answer the request, return null for sql. -- Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. -- You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. +- Generate exactly one SELECT statement. +- Never use "*" in the SELECT list. +- Use only Wren SQL syntax and only schema objects declared in DATABASE SCHEMA. +- Quote table and column identifiers with double quotes. Quote string literals with single quotes. Do not quote numeric literals. +- Never use SELECT *. Select only columns and expressions needed for the user's requested answer. +- Do not include SQL comments. +- Use CTEs when they make multi-step SQL clearer. +- Use joins only when DATABASE SCHEMA declares the needed relationship. +- Use declared views or metrics when they directly match the user's requested result. +- For metric-style requests, expose the requested dimensions and measure expressions or metric fields instead of returning raw table columns. +- Put aggregate expressions in SELECT or HAVING, not WHERE. +- For ranking requests, order by a selected column or selected aggregate alias and use LIMIT when the user requests a limit. +- For timeframe requests, apply a bounded predicate only when an exact date/time field and required date operation are supported by the retrieved context. +- Output aliases may label result expressions, but aliases are not source identifiers. +- Do not use connector-specific syntax such as SELECT TOP, square brackets, backticks, INTERVAL, unsupported date formatting, or unsupported statistical functions. """ From 21a9929bc26cb767a14f0a15ceab2cc95092f0d0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 20:15:13 +0530 Subject: [PATCH 0840/1087] Compact SQL generation context --- .../generation/followup_sql_generation.py | 5 - .../pipelines/generation/sql_correction.py | 5 - .../pipelines/generation/sql_generation.py | 5 - .../pipelines/generation/sql_regeneration.py | 5 - .../src/pipelines/generation/utils/sql.py | 2 +- .../retrieval/db_schema_retrieval.py | 96 +++++++++++++++++-- 6 files changed, 87 insertions(+), 31 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 317ecebf0c..76840aaa4d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -35,11 +35,6 @@ Given the user's current follow-up question and the current retrieved DATABASE SCHEMA, generate one SQL query to best answer the user's question. -{% if executable_schema_contract %} -{{ executable_schema_contract }} - -{% endif %} - ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 1885ae7a5d..9a0a5a1aa7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -57,11 +57,6 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) sql_correction_user_prompt_template = """ -{% if executable_schema_contract %} -{{ executable_schema_contract }} - -{% endif %} - {% if documents %} ### DATABASE SCHEMA ### {% for document in documents %} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 492069bdc3..27ca44820c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -29,11 +29,6 @@ sql_generation_user_prompt_template = """ -{% if executable_schema_contract %} -{{ executable_schema_contract }} - -{% endif %} - ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 952be44cbc..da627583c1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -54,11 +54,6 @@ def get_sql_regeneration_system_prompt( sql_regeneration_user_prompt_template = """ -{% if executable_schema_contract %} -{{ executable_schema_contract }} - -{% endif %} - ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 89ac2bd21e..7b97aac7e5 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1040,7 +1040,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. 7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. 8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +9. If the user asks for fields that require multiple schema objects, combine them only when DATABASE SCHEMA provides the exact relationship path and the result shape is grounded by the request. Do not use UNION, INTERSECT, or EXCEPT unless the user explicitly asks to combine separate result sets and each branch has the same selected columns. 10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. 11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. 12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index e6812c5fba..3391f7e55f 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -361,18 +361,94 @@ def _content_column_names(content: dict) -> list[str]: def _format_semantic_context(context: dict) -> str: - return ( - "/*\n" - "WREN RETRIEVED SEMANTIC CONTEXT\n" - f"{orjson.dumps(context).decode('utf-8')}\n" - f"{_format_identifier_contract(context)}" - "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" - "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" - "Values under column_role_hints_not_identifiers and semantic_roles_not_identifiers are meaning only; use them to map intent to exact declared columns, not as executable identifiers.\n" - "*/\n" - f"{_format_executable_identifier_catalog(context)}" + lines = [ + "/*", + "WREN RETRIEVED SEMANTIC CONTEXT", + f"object_type: {context.get('object_type', '')}", + ] + + contract = context.get("sql_identifier_contract") or {} + if contract.get("sql_table_name_use_exactly"): + lines.append( + "sql_table_name_use_exactly: " + + str(contract["sql_table_name_use_exactly"]) + ) + if contract.get("sql_column_names_use_exactly"): + lines.append("sql_column_names_use_exactly:") + lines.extend( + f"- {column_name}" + for column_name in contract["sql_column_names_use_exactly"] + ) + if contract.get("relationship_constraints_use_exactly"): + lines.append("relationship_constraints_use_exactly:") + lines.extend( + f"- {relationship_constraint}" + for relationship_constraint in contract[ + "relationship_constraints_use_exactly" + ] + ) + + semantic_context = context.get("semantic_context_not_sql_identifiers") or {} + for key, value in semantic_context.items(): + if value: + lines.append(f"semantic_context_not_sql_identifiers.{key}: {value}") + + columns = [ + column + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + if columns: + lines.append("columns:") + for column in columns: + column_parts = [ + f"sql_column_name_use_exactly={column['sql_column_name_use_exactly']}", + f"data_type={column.get('data_type', '')}", + ] + roles = column.get("semantic_roles_not_identifiers") or [] + if roles: + column_parts.append( + "column_role_hints_not_identifiers=" + + ", ".join(str(role) for role in roles) + ) + semantic_note = column.get("semantic_context_not_sql_identifier") + if semantic_note: + column_parts.append( + f"semantic_context_not_sql_identifier={semantic_note}" + ) + lines.append("- " + " | ".join(column_parts)) + + relationships = [ + relationship + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + if relationships: + lines.append("relationships:") + for relationship in relationships: + relationship_parts = [ + "sql_relationship_constraint_use_exactly=" + + relationship["sql_relationship_constraint_use_exactly"] + ] + semantic_note = relationship.get("semantic_context_not_sql_identifier") + if semantic_note: + relationship_parts.append( + f"semantic_context_not_sql_identifier={semantic_note}" + ) + lines.append("- " + " | ".join(relationship_parts)) + + lines.extend( + [ + "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.", + "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.", + "Values under column_role_hints_not_identifiers and semantic_roles_not_identifiers are meaning only; use them to map intent to exact declared columns, not as executable identifiers.", + "*/", + _format_executable_identifier_catalog(context), + ] ) + return "\n".join(lines) + def _normalized_data_type(column: dict) -> str: return str(column.get("data_type", "") or "").lower() From d154f729e3c868282afd054d92008f51e3d11996 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 20:26:05 +0530 Subject: [PATCH 0841/1087] Clarify SQL literal and comparison grounding --- .../src/pipelines/generation/followup_sql_generation.py | 2 ++ .../src/pipelines/generation/sql_correction.py | 2 ++ .../src/pipelines/generation/sql_generation.py | 2 ++ .../src/pipelines/generation/sql_regeneration.py | 2 ++ wren-ai-service/src/pipelines/generation/utils/sql.py | 3 +++ .../src/pipelines/retrieval/db_schema_retrieval.py | 9 ++++++--- 6 files changed, 17 insertions(+), 3 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 76840aaa4d..53cb4c83af 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -84,10 +84,12 @@ If a required part of the request is not grounded by an exact deployed schema object, column, relationship, metric, or supported function, return null for sql. Do not answer a specific analytical question with a broad table preview or with an unrelated nearby table. Do not ignore a literal filter value from the user; apply it to the exact schema field representing that filter concept, or return null when that field is unavailable. +String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. For aggregate, trend, ranking, or grouped requests, aggregate exact numeric_measure_candidate columns or count rows as appropriate for the user's requested measure. +For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. {% if executable_schema_contract %} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 9a0a5a1aa7..1c09ec146b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -95,6 +95,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Regenerate from the user's question and DATABASE SCHEMA only when a user question is available. Otherwise, correct the failed SQL only by using exact executable identifiers declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not copy table names, column names, functions, literals, aliases, or SQL structure from the failed SQL unless each one is declared in DATABASE SCHEMA or SQL FUNCTIONS. Correct into an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. If the error says the SQL is a broad table preview, table preview, missing requested aggregation, missing requested grouping, missing timeframe, or missing ordering/ranking, rebuild the query shape from the user's question. When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. +String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. +For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 27ca44820c..d7d623bb4e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -78,10 +78,12 @@ If a required part of the request is not grounded by an exact deployed schema object, column, relationship, metric, or supported function, return null for sql. Do not answer a specific analytical question with a broad table preview or with an unrelated nearby table. Do not ignore a literal filter value from the user; apply it to the exact schema field representing that filter concept, or return null when that field is unavailable. +String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. For aggregate, trend, ranking, or grouped requests, aggregate exact numeric_measure_candidate columns or count rows as appropriate for the user's requested measure. +For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. {% if executable_schema_contract %} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index da627583c1..371158f96e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -100,6 +100,8 @@ def get_sql_regeneration_system_prompt( Regenerate with executable identifiers from the current DATABASE SCHEMA only. Regenerate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. +String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. +For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. {% if executable_schema_contract %} ### ALLOWED EXECUTABLE IDENTIFIERS FOR THIS REGENERATION ### diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 7b97aac7e5..5836d0206c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -819,10 +819,13 @@ async def _classify_generation_result( - The SQL must answer every supported part of the user's request: requested subject, requested entity, requested filter value, timeframe, grouping, measure, ordering, and limit. - If the user asks for a filtered result, include the filter only when the filtered concept is represented by an exact schema field. If the filter field is not present, return null for sql instead of ignoring the filter. - If the user provides a literal filter value, use only that provided value. Do not invent, translate, or substitute filter values. +- Never use schema descriptions, column comments, aliases, display labels, source names, physical names, lineage names, reasoning text, or error messages as string literal data values. - If the user asks "which", "who", or "what" for a ranked entity, select and group by the exact schema field that represents that requested entity. Do not replace the requested entity with a context field or unrelated dimension. - Use row counting for record or entity volume questions when no numeric business measure is requested. Use numeric measures only when the question asks for a value, amount, quantity, rate, cost, or other declared measure. - For analytical questions, return dimensions plus the requested measure expression or metric field. Do not return a raw table preview. - For aggregate, ranking, grouped, or trend questions, produce an analytical query shape. +- For comparison questions, include each requested comparison group or time period and compute the requested difference, change, growth, or ranking when the required fields are grounded. +- Do not answer a comparison question with only one comparison side, one period, or one group unless the user explicitly asks for only that side. - For detail-list questions, return only the fields needed to identify and describe the requested records, plus requested filters and timeframes. - Do not answer a timeframe request with an unfiltered table scan. - Prefer one model, view, or metric that already contains the requested fields. Do not join tables just because they were retrieved together. Do not invent join predicates from similar column names. Join only through relationships declared in DATABASE SCHEMA. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 3391f7e55f..fe49b4600f 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -414,7 +414,8 @@ def _format_semantic_context(context: dict) -> str: semantic_note = column.get("semantic_context_not_sql_identifier") if semantic_note: column_parts.append( - f"semantic_context_not_sql_identifier={semantic_note}" + "semantic_context_not_sql_identifier" + f"(description_not_filter_value)={semantic_note}" ) lines.append("- " + " | ".join(column_parts)) @@ -433,14 +434,16 @@ def _format_semantic_context(context: dict) -> str: semantic_note = relationship.get("semantic_context_not_sql_identifier") if semantic_note: relationship_parts.append( - f"semantic_context_not_sql_identifier={semantic_note}" + "semantic_context_not_sql_identifier" + f"(description_not_filter_value)={semantic_note}" ) lines.append("- " + " | ".join(relationship_parts)) lines.extend( [ "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.", - "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.", + "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and are not executable identifiers or data values.", + "Never use semantic descriptions, aliases, display labels, source names, physical names, or lineage names as string literals in WHERE or HAVING predicates.", "Values under column_role_hints_not_identifiers and semantic_roles_not_identifiers are meaning only; use them to map intent to exact declared columns, not as executable identifiers.", "*/", _format_executable_identifier_catalog(context), From d45a703a793fd2feb69c17ab945bb88bc82f57c0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 20:45:49 +0530 Subject: [PATCH 0842/1087] Reject template SQL in generation prompts --- .../generation/followup_sql_generation.py | 3 +++ .../src/pipelines/generation/sql_correction.py | 5 ++++- .../src/pipelines/generation/sql_generation.py | 3 +++ .../src/pipelines/generation/sql_regeneration.py | 5 ++++- .../src/pipelines/generation/utils/sql.py | 14 +++++++++----- .../generation/test_sql_prompt_grounding.py | 8 ++++++++ 6 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 53cb4c83af..35b85c1edd 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -85,6 +85,9 @@ Do not answer a specific analytical question with a broad table preview or with an unrelated nearby table. Do not ignore a literal filter value from the user; apply it to the exact schema field representing that filter concept, or return null when that field is unavailable. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. +Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. +Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. +Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the wording of the user's question. For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 1c09ec146b..6fda2d0b4c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -51,7 +51,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) The final answer must be in JSON format: {{ - "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" + "sql": "complete executable corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ @@ -96,6 +96,9 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Correct into an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. If the error says the SQL is a broad table preview, table preview, missing requested aggregation, missing requested grouping, missing timeframe, or missing ordering/ranking, rebuild the query shape from the user's question. When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. +Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. +Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. +Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the failed SQL or the wording of the user's question. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index d7d623bb4e..73b254c4e2 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -79,6 +79,9 @@ Do not answer a specific analytical question with a broad table preview or with an unrelated nearby table. Do not ignore a literal filter value from the user; apply it to the exact schema field representing that filter concept, or return null when that field is unavailable. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. +Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. +Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. +Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the wording of the user's question. For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 371158f96e..68d6070a50 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -48,7 +48,7 @@ def get_sql_regeneration_system_prompt( The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" + "sql": "complete executable SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ @@ -101,6 +101,9 @@ def get_sql_regeneration_system_prompt( Regenerate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. +Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. +Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. +Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the original SQL, reasoning, or the wording of the user's question. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. {% if executable_schema_contract %} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5836d0206c..8aad09a25b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -816,9 +816,12 @@ async def _classify_generation_result( - Use comments, aliases, descriptions, display labels, metrics, calculated fields, and relationships only to understand business meaning. - Use column_role_hints_not_identifiers only as semantic hints for choosing exact declared columns; date_time_candidate, numeric_measure_candidate, identifier_candidate, and dimension_candidate are never SQL identifiers. - Copy executable table, column, metric, and relationship identifiers exactly from DATABASE SCHEMA. Do not create identifiers from user wording, descriptions, samples, history, physical names, lineage names, or error messages. +- Never output template SQL. Every table, column, metric, relationship, join key, function, filter value, grouping, ordering, and limit in the final SQL must be complete and executable for the current request. +- Never output generic, unresolved, variable-like, or placeholder identifiers or literal values. If a value or identifier would need to be filled in later, return null for sql. - The SQL must answer every supported part of the user's request: requested subject, requested entity, requested filter value, timeframe, grouping, measure, ordering, and limit. - If the user asks for a filtered result, include the filter only when the filtered concept is represented by an exact schema field. If the filter field is not present, return null for sql instead of ignoring the filter. -- If the user provides a literal filter value, use only that provided value. Do not invent, translate, or substitute filter values. +- If the user provides a literal filter value, copy that current request value into the SQL string literal exactly as the user provided it, except for normal SQL string escaping. Do not invent, translate, summarize, describe, or substitute filter values. +- If the user does not provide a literal filter value needed for the requested answer, do not add a stand-in value. Return null for sql when the missing value is required, or omit the filter only when the requested answer remains correct without it. - Never use schema descriptions, column comments, aliases, display labels, source names, physical names, lineage names, reasoning text, or error messages as string literal data values. - If the user asks "which", "who", or "what" for a ranked entity, select and group by the exact schema field that represents that requested entity. Do not replace the requested entity with a context field or unrelated dimension. - Use row counting for record or entity volume questions when no numeric business measure is requested. Use numeric measures only when the question asks for a value, amount, quantity, rate, cost, or other declared measure. @@ -1044,19 +1047,20 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. 8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. 9. If the user asks for fields that require multiple schema objects, combine them only when DATABASE SCHEMA provides the exact relationship path and the result shape is grounded by the request. Do not use UNION, INTERSECT, or EXCEPT unless the user explicitly asks to combine separate result sets and each branch has the same selected columns. -10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. +10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. Every string literal used as a data value must be copied from the current user question or USER INSTRUCTIONS, or be a concrete date/time boundary derived directly from the user's explicit timeframe using supported SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. 11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. 12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. 13. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. -14. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +14. If any planned SQL literal would be an unresolved variable, descriptive label, instruction to be replaced later, or placeholder value, return null for sql. The SQL field must never contain a partially completed query. +15. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. Do not create table or column identifiers from the user's wording. If the retrieved schema does not ground the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS, contains no placeholders or template parts, and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If the retrieved schema does not ground the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. {{ - "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" + "sql": "complete executable SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index 13997a655c..d1e69514d7 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -37,6 +37,9 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "numeric_measure_candidate" in prompt assert "Do not answer a timeframe request with an unfiltered table scan" in prompt assert "produce an analytical query shape" in prompt + assert "Never output template SQL" in prompt + assert "contains no placeholders or template parts" in prompt + assert "complete executable SQL query string" in prompt def test_sql_correction_system_prompt_allows_null_when_ungrounded(): @@ -111,6 +114,8 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): assert "Generate an intent-shaped query, not a table preview" in built_prompt assert "For timeframe requests, filter an exact date_time_candidate column" in built_prompt assert "aggregate exact numeric_measure_candidate columns" in built_prompt + assert "Never return template SQL" in built_prompt + assert "Copy user-provided filter values exactly" in built_prompt def test_followup_sql_generation_prompt_requires_intent_shaped_query(): @@ -124,6 +129,7 @@ def test_followup_sql_generation_prompt_requires_intent_shaped_query(): ) assert "Generate an intent-shaped query, not a table preview" in result["prompt"] + assert "Never return template SQL" in result["prompt"] def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): @@ -144,6 +150,7 @@ def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): assert "DIAGNOSTIC CONTEXT" in built_prompt assert "Correct into an intent-shaped query, not a table preview" in built_prompt assert "rebuild the query shape from the user's question" in built_prompt + assert "Never return template SQL" in built_prompt def test_sql_correction_prompt_includes_executable_schema_contract(): @@ -193,3 +200,4 @@ def test_sql_regeneration_prompt_includes_executable_schema_contract(): assert "TABLE: retrieved_model" in built_prompt assert "Regenerate an intent-shaped query, not a table preview" in built_prompt assert "For timeframe requests, filter an exact date_time_candidate column" in built_prompt + assert "Never return template SQL" in built_prompt From 0944113bf6bec4d5789fdd71bc008bde69186202 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 21:05:23 +0530 Subject: [PATCH 0843/1087] Handle missing restored ask tasks --- wren-ai-service/src/web/v1/services/ask.py | 14 +++-- .../server/services/askingTaskTracker.ts | 59 ++++++++++++------- .../services/tests/askingTaskTracker.test.ts | 47 +++++++++++++++ wren-ui/src/apollo/server/utils/error.ts | 5 ++ 4 files changed, 99 insertions(+), 26 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a918a8ba63..652d07eebc 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -88,7 +88,12 @@ class AskResult(BaseModel): class AskError(BaseModel): - code: Literal["NO_RELEVANT_DATA", "NO_RELEVANT_SQL", "OTHERS"] + code: Literal[ + "NO_RELEVANT_DATA", + "NO_RELEVANT_SQL", + "ASK_RESULT_NOT_FOUND", + "OTHERS", + ] message: str @@ -755,14 +760,15 @@ def get_ask_result( ask_result_request: AskResultRequest, ) -> AskResultResponse: if (result := self._ask_results.get(ask_result_request.query_id)) is None: - logger.exception( - f"ask pipeline - OTHERS: {ask_result_request.query_id} is not found" + logger.warning( + "ask pipeline - ASK_RESULT_NOT_FOUND: " + f"{ask_result_request.query_id} is not found" ) return AskResultResponse( status="failed", type="TEXT_TO_SQL", error=AskError( - code="OTHERS", + code="ASK_RESULT_NOT_FOUND", message=f"{ask_result_request.query_id} is not found", ), ) diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index afa6cc2df1..5fff110377 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -328,7 +328,11 @@ export class AskingTaskTracker implements IAskingTaskTracker { // Poll for updates logger.debug(`Polling for updates for task ${queryId}`); - const result = await this.wrenAIAdaptor.getAskResult(queryId); + const resultFromAIService = + await this.wrenAIAdaptor.getAskResult(queryId); + const result = this.isMissingInAIService(resultFromAIService) + ? this.createExpiredTaskResult() + : resultFromAIService; task.lastPolled = now; const resultChanged = this.isResultChanged(task.result, result); this.scheduleNextPoll(task, result.status, resultChanged); @@ -417,26 +421,26 @@ export class AskingTaskTracker implements IAskingTaskTracker { ); // Run all jobs in parallel - Promise.allSettled(jobs.map((job) => job())).then((results) => { - // Log any rejected promises - results.forEach((result, index) => { - if (result.status === 'rejected') { - logger.error(`Job ${index} failed: ${result.reason}`); - } - }); + const results = await Promise.allSettled(jobs.map((job) => job())); - // Clean up tasks that have been in memory too long - if (tasksToRemove.length > 0) { - logger.info( - `Cleaning up tasks that have been in memory too long. Tasks: ${tasksToRemove.join( - ', ', - )}`, - ); - } - for (const queryId of tasksToRemove) { - this.trackedTasks.delete(queryId); + // Log any rejected promises + results.forEach((result, index) => { + if (result.status === 'rejected') { + logger.error(`Job ${index} failed: ${result.reason}`); } }); + + // Clean up tasks that have been in memory too long + if (tasksToRemove.length > 0) { + logger.info( + `Cleaning up tasks that have been in memory too long. Tasks: ${tasksToRemove.join( + ', ', + )}`, + ); + } + for (const queryId of tasksToRemove) { + this.trackedTasks.delete(queryId); + } } private async updateThreadResponseWhenTaskFinalized( @@ -564,7 +568,15 @@ export class AskingTaskTracker implements IAskingTaskTracker { } private async finalizeStaleTask(taskRecord: AskingTask): Promise { - const result: AskResult = { + const result = this.createExpiredTaskResult(); + + await this.askingTaskRepository.updateOne(taskRecord.id, { + detail: result, + }); + } + + private createExpiredTaskResult(): AskResult { + return { type: AskResultType.TEXT_TO_SQL, status: AskResultStatus.FAILED, response: null, @@ -574,10 +586,13 @@ export class AskingTaskTracker implements IAskingTaskTracker { 'The previous asking task expired after the service restarted. Please ask again.', }, }; + } - await this.askingTaskRepository.updateOne(taskRecord.id, { - detail: result, - }); + private isMissingInAIService(result: AskResult): boolean { + return ( + result.status === AskResultStatus.FAILED && + result.error?.code === Errors.GeneralErrorCodes.ASK_RESULT_NOT_FOUND + ); } private isResultChanged( diff --git a/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts b/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts index ca08d09369..526a383cd0 100644 --- a/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts @@ -111,4 +111,51 @@ describe('AskingTaskTracker', () => { expect(result.queryId).toBe(recentTask.queryId); expect(result.status).toBe(AskResultStatus.GENERATING); }); + + test('finalizes restored tasks when AI service no longer has the query id', async () => { + const recentDate = new Date(); + const recentTask = { + id: 9, + queryId: 'missing-query-id', + question: 'recent question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: recentDate, + updatedAt: recentDate, + }; + const { tracker, askingTaskRepository, wrenAIAdaptor } = createTracker({ + taskRecords: [recentTask], + memoryRetentionTime: 60_000, + }); + wrenAIAdaptor.getAskResult.mockResolvedValue({ + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.ASK_RESULT_NOT_FOUND, + message: 'The asking task result is no longer available', + }, + }); + + await tracker.initialize(); + await (tracker as any).pollTasks(); + + expect(wrenAIAdaptor.getAskResult).toHaveBeenCalledWith(recentTask.queryId); + expect(askingTaskRepository.updateOne).toHaveBeenCalledWith(9, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.POLLING_TIMEOUT, + message: + 'The previous asking task expired after the service restarted. Please ask again.', + }, + }, + }); + }); }); diff --git a/wren-ui/src/apollo/server/utils/error.ts b/wren-ui/src/apollo/server/utils/error.ts index 4bf793c49e..88400c1012 100644 --- a/wren-ui/src/apollo/server/utils/error.ts +++ b/wren-ui/src/apollo/server/utils/error.ts @@ -9,6 +9,7 @@ export enum GeneralErrorCodes { RESOURCE_NOT_FOUND = 'RESOURCE_NOT_FOUND', MDL_PARSE_ERROR = 'MDL_PARSE_ERROR', NO_CHART = 'NO_CHART', + ASK_RESULT_NOT_FOUND = 'ASK_RESULT_NOT_FOUND', // Exception error for AI service (e.g., network connection error) AI_SERVICE_UNDEFINED_ERROR = 'OTHERS', @@ -72,6 +73,8 @@ export const errorMessages = { "Could you please provide more details or specify the information you're seeking?", [GeneralErrorCodes.NO_CHART]: "The chart couldn't be generated this time. Please try regenerating the chart or rephrasing your question for better results.", + [GeneralErrorCodes.ASK_RESULT_NOT_FOUND]: + 'The asking task result is no longer available', // Connector errors [GeneralErrorCodes.CONNECTION_ERROR]: 'Can not connect to data source', @@ -160,6 +163,8 @@ export const shortMessages = { [GeneralErrorCodes.FAILED_TO_GENERATE_VEGA_SCHEMA]: 'Failed to generate Vega spec', [GeneralErrorCodes.POLLING_TIMEOUT]: 'Polling timeout', + [GeneralErrorCodes.ASK_RESULT_NOT_FOUND]: + 'The asking task result is no longer available', [GeneralErrorCodes.SQL_EXECUTION_ERROR]: 'SQL execution error', }; From d5f5b599af5eb76b0043207c77e589865510140a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 21:29:31 +0530 Subject: [PATCH 0844/1087] Stabilize SQL generation failure handling --- .../src/pipelines/generation/utils/sql.py | 48 ++++++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 15 ++++-- .../generation/test_sql_post_processor.py | 17 +++++++ 3 files changed, 74 insertions(+), 6 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 8aad09a25b..660f04a75c 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -535,6 +535,23 @@ async def run( query: str | None = None, ) -> dict: try: + if not replies: + ( + valid_generation_result, + invalid_generation_result, + ) = await self._classify_generation_result( + None, + project_id=project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + data_source=data_source, + allow_data_preview=allow_data_preview, + ) + return { + "valid_generation_result": valid_generation_result, + "invalid_generation_result": invalid_generation_result, + } + cleaned_generation_result = clean_generation_result(replies[0]) # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' @@ -542,6 +559,22 @@ async def run( cleaned_generation_result = orjson.loads(cleaned_generation_result).get( "sql" ) + if not cleaned_generation_result: + ( + valid_generation_result, + invalid_generation_result, + ) = await self._classify_generation_result( + None, + project_id=project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + data_source=data_source, + allow_data_preview=allow_data_preview, + ) + return { + "valid_generation_result": valid_generation_result, + "invalid_generation_result": invalid_generation_result, + } cleaned_generation_result = clean_generation_result( cleaned_generation_result ) @@ -642,9 +675,20 @@ async def run( except Exception as e: logger.exception(f"Error in SQLGenPostProcessor: {e}") + ( + valid_generation_result, + invalid_generation_result, + ) = await self._classify_generation_result( + None, + project_id=project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + data_source=data_source, + allow_data_preview=allow_data_preview, + ) return { - "valid_generation_result": {}, - "invalid_generation_result": {}, + "valid_generation_result": valid_generation_result, + "invalid_generation_result": invalid_generation_result, } async def _classify_generation_result( diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 652d07eebc..01729356e7 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -583,12 +583,17 @@ async def ask( ]["invalid_generation_result"]: schema_grounding_correction_attempted = False while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + failed_generation_type = failed_dry_run_result.get("type") + if not failed_generation_type: + break + if failed_generation_type == "TIME_OUT": break - original_sql = failed_dry_run_result["original_sql"] - invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + original_sql = failed_dry_run_result.get("original_sql") or "" + invalid_sql = failed_dry_run_result.get("sql") or original_sql + error_message = ( + failed_dry_run_result.get("error") or "No relevant SQL" + ) skip_sql_diagnosis = should_skip_sql_diagnosis( failed_dry_run_result ) @@ -679,6 +684,8 @@ async def ask( next_failed_dry_run_result = sql_correction_results[ "post_process" ]["invalid_generation_result"] + if not next_failed_dry_run_result: + break if ( next_failed_dry_run_result and next_failed_dry_run_result.get("sql") == invalid_sql diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index afa7c20ba3..4965652825 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -98,6 +98,23 @@ def test_clean_generation_result_preserves_internal_statement_separators(): ) +@pytest.mark.asyncio +async def test_sql_post_processor_returns_structured_failure_for_null_sql(): + engine = CapturingEngine() + + result = await SQLGenPostProcessor(engine).run(['{"sql": null}']) + + assert engine.executed is False + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": "No grounded SQL was generated from the current schema.", + "correlation_id": "", + } + + @pytest.mark.asyncio async def test_sql_post_processor_converts_select_top_to_wren_limit(): engine = CapturingEngine() From 5dd5ebfbcfeb16da01dc83f5b41d5a7256cfc5f9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 21:56:53 +0530 Subject: [PATCH 0845/1087] Prevent metadata placeholder SQL generation --- .../pipelines/generation/followup_sql_generation.py | 4 ++-- .../src/pipelines/generation/sql_correction.py | 2 +- .../src/pipelines/generation/sql_generation.py | 4 ++-- .../src/pipelines/generation/sql_regeneration.py | 2 +- .../src/pipelines/generation/utils/sql.py | 2 +- .../generation/test_sql_prompt_grounding.py | 13 +++++++------ 6 files changed, 14 insertions(+), 13 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 35b85c1edd..36343e6c1b 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -90,8 +90,8 @@ Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the wording of the user's question. For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. -For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. -For aggregate, trend, ranking, or grouped requests, aggregate exact numeric_measure_candidate columns or count rows as appropriate for the user's requested measure. +For timeframe requests, filter an actual declared column whose metadata marks it as the requested time concept. Metadata role labels are not executable column names. +For aggregate, trend, ranking, or grouped requests, aggregate actual declared measure columns or count rows as appropriate for the user's requested measure. Metadata role labels are not executable column names. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 6fda2d0b4c..38bae6560e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -94,7 +94,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Regenerate from the user's question and DATABASE SCHEMA only when a user question is available. Otherwise, correct the failed SQL only by using exact executable identifiers declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not copy table names, column names, functions, literals, aliases, or SQL structure from the failed SQL unless each one is declared in DATABASE SCHEMA or SQL FUNCTIONS. Correct into an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. -If the error says the SQL is a broad table preview, table preview, missing requested aggregation, missing requested grouping, missing timeframe, or missing ordering/ranking, rebuild the query shape from the user's question. When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. +If the error says the SQL is a broad table preview, table preview, missing requested aggregation, missing requested grouping, missing timeframe, or missing ordering/ranking, rebuild the query shape from the user's question. When DATABASE SCHEMA contains role or semantic hints, use those hints only to choose actual declared columns. Do not write role labels, sample schema names, placeholder table names, placeholder column names, or replacement markers as SQL identifiers or SQL literal values. For timeframe requests, filter an actual declared time/date column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate actual declared measure columns or count rows, group by actual declared dimension/date columns, order by the selected aggregate alias when ranking, and limit only when requested. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 73b254c4e2..f567b8e166 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -84,8 +84,8 @@ Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the wording of the user's question. For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. -For timeframe requests, filter an exact date_time_candidate column when the retrieved schema provides one for the requested time concept. -For aggregate, trend, ranking, or grouped requests, aggregate exact numeric_measure_candidate columns or count rows as appropriate for the user's requested measure. +For timeframe requests, filter an actual declared column whose metadata marks it as the requested time concept. Metadata role labels are not executable column names. +For aggregate, trend, ranking, or grouped requests, aggregate actual declared measure columns or count rows as appropriate for the user's requested measure. Metadata role labels are not executable column names. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 68d6070a50..fc46f64fb2 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -99,7 +99,7 @@ def get_sql_regeneration_system_prompt( Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. Regenerate with executable identifiers from the current DATABASE SCHEMA only. Regenerate an intent-shaped query, not a table preview. Select explicit columns, filters, groupings, measures, joins, ordering, and limits needed by the question. For metric questions, return dimensions plus the requested measure or grounded expression; never use SELECT * as a substitute. -When DATABASE SCHEMA contains column_role_hints_not_identifiers, use those roles only to map intent to exact declared columns. For timeframe requests, filter an exact date_time_candidate column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate exact numeric_measure_candidate columns or count rows, group by exact dimension/date expressions, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. +When DATABASE SCHEMA contains role or semantic hints, use those hints only to choose actual declared columns. Do not write role labels, sample schema names, placeholder table names, placeholder column names, or replacement markers as SQL identifiers or SQL literal values. For timeframe requests, filter an actual declared time/date column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate actual declared measure columns or count rows, group by actual declared dimension/date columns, order by the selected aggregate alias when ranking, and limit only when requested. Do not return a raw table preview. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 660f04a75c..1dd75677e0 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -858,7 +858,7 @@ async def _classify_generation_result( - Use the retrieved DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT as the only executable context. - Before generating SQL, silently validate that every model, column, metric, relationship, join path, filter field, grouping field, ordering field, and SQL function is present in the retrieved context. - Use comments, aliases, descriptions, display labels, metrics, calculated fields, and relationships only to understand business meaning. -- Use column_role_hints_not_identifiers only as semantic hints for choosing exact declared columns; date_time_candidate, numeric_measure_candidate, identifier_candidate, and dimension_candidate are never SQL identifiers. +- Use role-hint metadata only as semantic hints for choosing exact declared columns. Metadata role labels are never SQL identifiers or SQL literal values. - Copy executable table, column, metric, and relationship identifiers exactly from DATABASE SCHEMA. Do not create identifiers from user wording, descriptions, samples, history, physical names, lineage names, or error messages. - Never output template SQL. Every table, column, metric, relationship, join key, function, filter value, grouping, ordering, and limit in the final SQL must be complete and executable for the current request. - Never output generic, unresolved, variable-like, or placeholder identifiers or literal values. If a value or identifier would need to be filled in later, return null for sql. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index d1e69514d7..434e92c958 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -32,9 +32,8 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "For metric-style requests" in prompt assert "Do not join tables just because they were retrieved together" in prompt assert "Do not invent join predicates from similar column names" in prompt - assert "column_role_hints_not_identifiers" in prompt - assert "date_time_candidate" in prompt - assert "numeric_measure_candidate" in prompt + assert "role-hint metadata only as semantic hints" in prompt + assert "Metadata role labels are never SQL identifiers" in prompt assert "Do not answer a timeframe request with an unfiltered table scan" in prompt assert "produce an analytical query shape" in prompt assert "Never output template SQL" in prompt @@ -112,8 +111,9 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): assert "- grouping_attribute" in built_prompt assert "- numeric_measure" in built_prompt assert "Generate an intent-shaped query, not a table preview" in built_prompt - assert "For timeframe requests, filter an exact date_time_candidate column" in built_prompt - assert "aggregate exact numeric_measure_candidate columns" in built_prompt + assert "filter an actual declared column" in built_prompt + assert "aggregate actual declared measure columns" in built_prompt + assert "Metadata role labels are not executable column names" in built_prompt assert "Never return template SQL" in built_prompt assert "Copy user-provided filter values exactly" in built_prompt @@ -199,5 +199,6 @@ def test_sql_regeneration_prompt_includes_executable_schema_contract(): assert "EXECUTABLE WREN IDENTIFIER CATALOG" in built_prompt assert "TABLE: retrieved_model" in built_prompt assert "Regenerate an intent-shaped query, not a table preview" in built_prompt - assert "For timeframe requests, filter an exact date_time_candidate column" in built_prompt + assert "filter an actual declared time/date column" in built_prompt + assert "Do not write role labels" in built_prompt assert "Never return template SQL" in built_prompt From 0b26b5093c86841a1bebb7c1cacce636540cf7e0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 22:29:43 +0530 Subject: [PATCH 0846/1087] Prevent SQL generation from merging retrieved candidates --- .../pipelines/generation/followup_sql_generation.py | 1 + .../src/pipelines/generation/sql_correction.py | 1 + .../src/pipelines/generation/sql_generation.py | 1 + .../src/pipelines/generation/sql_regeneration.py | 1 + wren-ai-service/src/pipelines/generation/utils/sql.py | 4 +++- .../pipelines/generation/test_sql_prompt_grounding.py | 10 ++++++++++ 6 files changed, 17 insertions(+), 1 deletion(-) 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 36343e6c1b..8dffa2fd6f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -92,6 +92,7 @@ For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. For timeframe requests, filter an actual declared column whose metadata marks it as the requested time concept. Metadata role labels are not executable column names. For aggregate, trend, ranking, or grouped requests, aggregate actual declared measure columns or count rows as appropriate for the user's requested measure. Metadata role labels are not executable column names. +Retrieved schema objects are ranked candidates, not automatic datasets to merge. Prefer one grounded model, view, or metric that answers the question. Do not use UNION, UNION ALL, INTERSECT, or EXCEPT to combine similar retrieved candidates unless the current user explicitly asks to combine separate result sets and DATABASE SCHEMA grounds each branch with the same result shape and compatible measure meaning. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 38bae6560e..6e5f35a589 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -99,6 +99,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the failed SQL or the wording of the user's question. +Retrieved schema objects are ranked candidates, not automatic datasets to merge. Prefer one grounded model, view, or metric that answers the question. Do not use UNION, UNION ALL, INTERSECT, or EXCEPT to combine similar retrieved candidates unless the current user explicitly asks to combine separate result sets and DATABASE SCHEMA grounds each branch with the same result shape and compatible measure meaning. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index f567b8e166..26d8f044dc 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -86,6 +86,7 @@ For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. For timeframe requests, filter an actual declared column whose metadata marks it as the requested time concept. Metadata role labels are not executable column names. For aggregate, trend, ranking, or grouped requests, aggregate actual declared measure columns or count rows as appropriate for the user's requested measure. Metadata role labels are not executable column names. +Retrieved schema objects are ranked candidates, not automatic datasets to merge. Prefer one grounded model, view, or metric that answers the question. Do not use UNION, UNION ALL, INTERSECT, or EXCEPT to combine similar retrieved candidates unless the current user explicitly asks to combine separate result sets and DATABASE SCHEMA grounds each branch with the same result shape and compatible measure meaning. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index fc46f64fb2..d3235ed99e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -104,6 +104,7 @@ def get_sql_regeneration_system_prompt( Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the original SQL, reasoning, or the wording of the user's question. +Retrieved schema objects are ranked candidates, not automatic datasets to merge. Prefer one grounded model, view, or metric that answers the question. Do not use UNION, UNION ALL, INTERSECT, or EXCEPT to combine similar retrieved candidates unless the current user explicitly asks to combine separate result sets and DATABASE SCHEMA grounds each branch with the same result shape and compatible measure meaning. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. {% if executable_schema_contract %} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1dd75677e0..1d5611a5d3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -876,6 +876,7 @@ async def _classify_generation_result( - For detail-list questions, return only the fields needed to identify and describe the requested records, plus requested filters and timeframes. - Do not answer a timeframe request with an unfiltered table scan. - Prefer one model, view, or metric that already contains the requested fields. Do not join tables just because they were retrieved together. Do not invent join predicates from similar column names. Join only through relationships declared in DATABASE SCHEMA. +- Treat retrieved schema objects as ranked candidates for grounding, not as datasets to merge automatically. Do not combine parallel or similar retrieved models with UNION, UNION ALL, INTERSECT, or EXCEPT unless the current user explicitly asks to combine separate result sets and the retrieved DATABASE SCHEMA grounds every branch with identical result shape and compatible measures. - If multiple retrieved schema objects are needed for the same result, use them only when the required columns and relationship path are present. - If multiple semantic interpretations exist and the retrieved context does not make one interpretation authoritative, return null for sql instead of choosing one. - If SQL execution or correction is needed, repair the query only when the repair can be verified using DATABASE SCHEMA, WREN SQL IDENTIFIER CONTRACT, and SQL FUNCTIONS. Never introduce a new schema object during repair. @@ -894,6 +895,7 @@ async def _classify_generation_result( - Do not include SQL comments. - Use CTEs when they make multi-step SQL clearer. - Use joins only when DATABASE SCHEMA declares the needed relationship. +- Use set operations only when the user explicitly requests combined result sets and the DATABASE SCHEMA grounds each branch. Do not use set operations to merge retrieved candidates that merely have similar dimensions or measures. - Use declared views or metrics when they directly match the user's requested result. - For metric-style requests, expose the requested dimensions and measure expressions or metric fields instead of returning raw table columns. - Put aggregate expressions in SELECT or HAVING, not WHERE. @@ -1090,7 +1092,7 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) 6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. 7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. 8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -9. If the user asks for fields that require multiple schema objects, combine them only when DATABASE SCHEMA provides the exact relationship path and the result shape is grounded by the request. Do not use UNION, INTERSECT, or EXCEPT unless the user explicitly asks to combine separate result sets and each branch has the same selected columns. +9. If the user asks for fields that require multiple schema objects, combine them only when DATABASE SCHEMA provides the exact relationship path and the result shape is grounded by the request. Retrieved schema objects are alternatives until the user request and DATABASE SCHEMA prove they must be combined. Do not use UNION, UNION ALL, INTERSECT, or EXCEPT unless the user explicitly asks to combine separate result sets and each branch is grounded with the same selected columns, compatible measure meaning, and any required grouping inside each branch. 10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. Every string literal used as a data value must be copied from the current user question or USER INSTRUCTIONS, or be a concrete date/time boundary derived directly from the user's explicit timeframe using supported SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. 11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. 12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py index 434e92c958..af7f3a65de 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_prompt_grounding.py @@ -32,6 +32,8 @@ def test_sql_generation_system_prompt_requires_retrieved_semantic_authority(): assert "For metric-style requests" in prompt assert "Do not join tables just because they were retrieved together" in prompt assert "Do not invent join predicates from similar column names" in prompt + assert "retrieved schema objects as ranked candidates" in prompt + assert "Use set operations only when the user explicitly requests" in prompt assert "role-hint metadata only as semantic hints" in prompt assert "Metadata role labels are never SQL identifiers" in prompt assert "Do not answer a timeframe request with an unfiltered table scan" in prompt @@ -116,6 +118,8 @@ def test_sql_generation_prompt_includes_executable_schema_contract(): assert "Metadata role labels are not executable column names" in built_prompt assert "Never return template SQL" in built_prompt assert "Copy user-provided filter values exactly" in built_prompt + assert "not automatic datasets to merge" in built_prompt + assert "Do not use UNION, UNION ALL, INTERSECT, or EXCEPT" in built_prompt def test_followup_sql_generation_prompt_requires_intent_shaped_query(): @@ -130,6 +134,8 @@ def test_followup_sql_generation_prompt_requires_intent_shaped_query(): assert "Generate an intent-shaped query, not a table preview" in result["prompt"] assert "Never return template SQL" in result["prompt"] + assert "not automatic datasets to merge" in result["prompt"] + assert "Do not use UNION, UNION ALL, INTERSECT, or EXCEPT" in result["prompt"] def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): @@ -151,6 +157,8 @@ def test_sql_correction_prompt_keeps_failed_sql_diagnostic_and_question(): assert "Correct into an intent-shaped query, not a table preview" in built_prompt assert "rebuild the query shape from the user's question" in built_prompt assert "Never return template SQL" in built_prompt + assert "not automatic datasets to merge" in built_prompt + assert "Do not use UNION, UNION ALL, INTERSECT, or EXCEPT" in built_prompt def test_sql_correction_prompt_includes_executable_schema_contract(): @@ -202,3 +210,5 @@ def test_sql_regeneration_prompt_includes_executable_schema_contract(): assert "filter an actual declared time/date column" in built_prompt assert "Do not write role labels" in built_prompt assert "Never return template SQL" in built_prompt + assert "not automatic datasets to merge" in built_prompt + assert "Do not use UNION, UNION ALL, INTERSECT, or EXCEPT" in built_prompt From a9a0a57ec6d38cd9896c686157931784cba1d5d0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 4 Aug 2026 23:33:41 +0530 Subject: [PATCH 0847/1087] Ground SQL generation in metadata contract --- .../src/pipelines/generation/utils/sql.py | 273 +++++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 36 ++- 2 files changed, 299 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1d5611a5d3..5caff02f30 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -38,6 +38,10 @@ r")\b", re.IGNORECASE, ) +_RANKING_QUERY_PATTERN = re.compile( + r"\b(top|bottom|highest|lowest|most|least|largest|smallest|fastest|slowest)\b", + re.IGNORECASE, +) _TIME_QUERY_PATTERN = re.compile( r"\b(" r"today|yesterday|week|month|quarter|year|january|february|march|april|" @@ -236,6 +240,248 @@ def _table_grounding_error( return None +def _normalize_identifier_part(identifier: str | None) -> str: + if not identifier: + return "" + + return identifier.strip().strip('"`[]').lower() + + +def _schema_contract_index( + schema_contracts: list[dict] | None, +) -> dict[str, dict[str, Any]]: + index: dict[str, dict[str, Any]] = {} + for contract in schema_contracts or []: + table_name = contract.get("table_name") + normalized_table_name = _normalize_identifier_part(table_name) + if not table_name or not normalized_table_name: + continue + + column_names = [ + column_name for column_name in contract.get("column_names", []) if column_name + ] + index[normalized_table_name] = { + "table_name": table_name, + "columns": { + _normalize_identifier_part(column_name) for column_name in column_names + }, + "has_column_contract": bool(column_names), + } + + return index + + +def _identifier_contains_subquery(identifier: Identifier) -> bool: + return any( + isinstance(child, Parenthesis) and _contains_select(child) + for child in identifier.tokens + ) + + +def _mark_identifier_tree(identifier: Identifier, marked_ids: set[int]) -> None: + marked_ids.add(id(identifier)) + for child in identifier.tokens: + if isinstance(child, Identifier): + _mark_identifier_tree(child, marked_ids) + elif isinstance(child, IdentifierList): + for child_identifier in child.get_identifiers(): + _mark_identifier_tree(child_identifier, marked_ids) + elif isinstance(child, TokenList): + for nested_child in child.tokens: + if isinstance(nested_child, Identifier): + _mark_identifier_tree(nested_child, marked_ids) + + +def _collect_table_context( + token: TokenList, + schema_index: dict[str, dict[str, Any]], + cte_names: set[str], +) -> tuple[dict[str, str], set[int]]: + aliases: dict[str, str] = {} + table_identifier_ids: set[int] = set() + tokens = _meaningful_tokens(token) + expect_table = False + + def add_table_identifier(identifier: Identifier) -> None: + if _identifier_contains_subquery(identifier): + for child in identifier.tokens: + if isinstance(child, Parenthesis): + child_aliases, child_table_ids = _collect_table_context( + child, schema_index, cte_names + ) + aliases.update(child_aliases) + table_identifier_ids.update(child_table_ids) + return + + table_name = _table_reference_name(identifier) + normalized_table_name = _normalize_identifier_part(table_name) + if not normalized_table_name or normalized_table_name in cte_names: + _mark_identifier_tree(identifier, table_identifier_ids) + return + + if normalized_table_name not in schema_index: + _mark_identifier_tree(identifier, table_identifier_ids) + return + + _mark_identifier_tree(identifier, table_identifier_ids) + aliases[normalized_table_name] = normalized_table_name + + real_name = _normalize_identifier_part(identifier.get_real_name()) + if real_name: + aliases[real_name] = normalized_table_name + + alias = _normalize_identifier_part(identifier.get_alias()) + if alias: + aliases[alias] = normalized_table_name + + for current in tokens: + normalized = current.normalized + + if isinstance(current, Parenthesis): + if _contains_select(current): + child_aliases, child_table_ids = _collect_table_context( + current, schema_index, cte_names + ) + aliases.update(child_aliases) + table_identifier_ids.update(child_table_ids) + continue + + if current.ttype == sqlparse_tokens.Keyword and ( + normalized == "FROM" or normalized == "JOIN" or normalized.endswith(" JOIN") + ): + expect_table = True + continue + + if expect_table: + if isinstance(current, IdentifierList): + for identifier in current.get_identifiers(): + add_table_identifier(identifier) + elif isinstance(current, Identifier): + add_table_identifier(current) + expect_table = False + + return aliases, table_identifier_ids + + +def _iter_identifier_nodes( + token: TokenList, parent: TokenList | None = None +) -> list[tuple[Identifier, TokenList | None]]: + identifiers: list[tuple[Identifier, TokenList | None]] = [] + if isinstance(token, Identifier): + identifiers.append((token, parent)) + + if isinstance(token, TokenList): + for child in token.tokens: + if isinstance(child, TokenList): + identifiers.extend(_iter_identifier_nodes(child, token)) + + return identifiers + + +def _select_aliases(statement: TokenList) -> set[str]: + return { + _normalize_identifier_part(item.get_alias()) + for item in _select_items(statement) + if isinstance(item, Identifier) and item.get_alias() + } + + +def _identifier_has_function_child(identifier: Identifier) -> bool: + return any(isinstance(child, Function) for child in identifier.tokens) + + +def _column_grounding_error( + sql: str | None, schema_contracts: list[dict] | None +) -> str | None: + if not sql or not schema_contracts: + return None + + schema_index = _schema_contract_index(schema_contracts) + if not schema_index: + return None + + for statement in sqlparse.parse(sql): + cte_names = { + _normalize_identifier_part(cte_name) + for cte_name in _collect_cte_names(statement) + } + table_aliases, table_identifier_ids = _collect_table_context( + statement, schema_index, cte_names + ) + referenced_tables = set(table_aliases.values()) + if not referenced_tables: + continue + + allowed_unqualified_columns = set() + all_referenced_tables_have_column_contract = True + for table_name in referenced_tables: + table_contract = schema_index.get(table_name) + if not table_contract: + continue + allowed_unqualified_columns.update(table_contract["columns"]) + all_referenced_tables_have_column_contract = ( + all_referenced_tables_have_column_contract + and table_contract["has_column_contract"] + ) + + select_aliases = _select_aliases(statement) + + for identifier, parent in _iter_identifier_nodes(statement): + if id(identifier) in table_identifier_ids: + continue + + if isinstance(parent, Function): + function_name = _normalize_identifier_part(parent.get_name()) + if _normalize_identifier_part(_identifier_name(identifier)) == function_name: + continue + + if _identifier_has_function_child(identifier): + continue + + column_name = _normalize_identifier_part(identifier.get_real_name()) + if not column_name: + continue + + parent_name = _normalize_identifier_part(identifier.get_parent_name()) + if parent_name: + if parent_name in cte_names: + continue + + source_table_name = table_aliases.get(parent_name) + if not source_table_name: + return ( + "Generated SQL references column identifiers outside the " + "retrieved deployed schema." + ) + + table_contract = schema_index.get(source_table_name) + if ( + table_contract + and table_contract["has_column_contract"] + and column_name not in table_contract["columns"] + ): + return ( + "Generated SQL references column identifiers outside the " + "retrieved deployed schema." + ) + continue + + if column_name in select_aliases or column_name in cte_names: + continue + if column_name in table_aliases: + continue + if column_name in allowed_unqualified_columns: + continue + + if all_referenced_tables_have_column_contract: + return ( + "Generated SQL references column identifiers outside the " + "retrieved deployed schema." + ) + + return None + + def _sql_statement_shape_error(sql: str | None) -> str | None: if not sql: return None @@ -369,6 +615,7 @@ def _table_preview_shape_error(sql: str | None, query: str | None = None) -> str query_has_shape = bool(query and _ANALYTICAL_OR_FILTER_QUERY_PATTERN.search(query)) query_has_aggregate_shape = bool(query and _AGGREGATE_QUERY_PATTERN.search(query)) + query_has_ranking_shape = bool(query and _RANKING_QUERY_PATTERN.search(query)) for statement in sqlparse.parse(sql): if not str(statement).strip().strip(";").strip(): @@ -400,6 +647,11 @@ def _table_preview_shape_error(sql: str | None, query: str | None = None) -> str "ranking, or measure calculation." ) + if query_has_ranking_shape and has_aggregate_item and not has_ordering: + return ( + "Generated SQL does not apply the requested ranking or ordering." + ) + if ( is_single_source_scan and len(items) >= _BROAD_TABLE_PREVIEW_COLUMN_THRESHOLD @@ -624,6 +876,21 @@ async def run( }, } + column_grounding_error = _column_grounding_error( + cleaned_generation_result, schema_contracts + ) + if column_grounding_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_GROUNDING", + "error": column_grounding_error, + "correlation_id": "", + }, + } + literal_filter_error = _unsupported_literal_filter_error( cleaned_generation_result, query=query, @@ -863,9 +1130,11 @@ async def _classify_generation_result( - Never output template SQL. Every table, column, metric, relationship, join key, function, filter value, grouping, ordering, and limit in the final SQL must be complete and executable for the current request. - Never output generic, unresolved, variable-like, or placeholder identifiers or literal values. If a value or identifier would need to be filled in later, return null for sql. - The SQL must answer every supported part of the user's request: requested subject, requested entity, requested filter value, timeframe, grouping, measure, ordering, and limit. +- Preserve the user's requested result shape exactly. Do not convert a detail-list request into a count, distinct count, latest-row query, top query, or summary unless the user explicitly asks for that operation. +- Do not add implicit filters, latest-period logic, maximum-date logic, row limits, distinctness, aggregation, or ordering unless the current user request requires it. - If the user asks for a filtered result, include the filter only when the filtered concept is represented by an exact schema field. If the filter field is not present, return null for sql instead of ignoring the filter. - If the user provides a literal filter value, copy that current request value into the SQL string literal exactly as the user provided it, except for normal SQL string escaping. Do not invent, translate, summarize, describe, or substitute filter values. -- If the user does not provide a literal filter value needed for the requested answer, do not add a stand-in value. Return null for sql when the missing value is required, or omit the filter only when the requested answer remains correct without it. +- If the user asks for a specific or particular entity but does not provide the required value, return null for sql instead of adding a stand-in value. Omit a missing filter only when the requested answer remains correct without it. - Never use schema descriptions, column comments, aliases, display labels, source names, physical names, lineage names, reasoning text, or error messages as string literal data values. - If the user asks "which", "who", or "what" for a ranked entity, select and group by the exact schema field that represents that requested entity. Do not replace the requested entity with a context field or unrelated dimension. - Use row counting for record or entity volume questions when no numeric business measure is requested. Use numeric measures only when the question asks for a value, amount, quantity, rate, cost, or other declared measure. @@ -892,6 +1161,8 @@ async def _classify_generation_result( - Use only Wren SQL syntax and only schema objects declared in DATABASE SCHEMA. - Quote table and column identifiers with double quotes. Quote string literals with single quotes. Do not quote numeric literals. - Never use SELECT *. Select only columns and expressions needed for the user's requested answer. +- Preserve the requested answer shape. For record-list requests, return the requested records with explicit identifying columns and filters; do not replace them with COUNT, DISTINCT, MAX, latest-row, ranking, or summary logic. +- Use COUNT, DISTINCT, SUM, AVG, MIN, MAX, GROUP BY, ORDER BY, LIMIT, date predicates, and joins only when the user's request requires that operation and the required identifiers are grounded in DATABASE SCHEMA. - Do not include SQL comments. - Use CTEs when they make multi-step SQL clearer. - Use joins only when DATABASE SCHEMA declares the needed relationship. diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 01729356e7..238b669935 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -21,6 +21,11 @@ "NO_RELEVANT_SQL", } +_NON_REPAIRABLE_SQL_VALIDATION_TYPES = { + "NO_RELEVANT_SQL", + "SQL_VALUE_GROUNDING", +} + def should_skip_sql_diagnosis(failed_generation_result: dict | None) -> bool: if not failed_generation_result: @@ -29,6 +34,26 @@ def should_skip_sql_diagnosis(failed_generation_result: dict | None) -> bool: return failed_generation_result.get("type") in _DETERMINISTIC_SQL_VALIDATION_TYPES +def should_attempt_sql_correction(failed_generation_result: dict | None) -> bool: + if not failed_generation_result: + return False + + failure_type = failed_generation_result.get("type") + if not failure_type or failure_type == "TIME_OUT": + return False + + if failure_type in _NON_REPAIRABLE_SQL_VALIDATION_TYPES: + return False + + if failure_type == "SCHEMA_GROUNDING" and not ( + failed_generation_result.get("sql") + or failed_generation_result.get("original_sql") + ): + return False + + return True + + async def run_pipeline_with_timeout(awaitable, timeout_seconds: float, operation: str): try: logger.info( @@ -583,10 +608,7 @@ async def ask( ]["invalid_generation_result"]: schema_grounding_correction_attempted = False while current_sql_correction_retries < max_sql_correction_retries: - failed_generation_type = failed_dry_run_result.get("type") - if not failed_generation_type: - break - if failed_generation_type == "TIME_OUT": + if not should_attempt_sql_correction(failed_dry_run_result): break original_sql = failed_dry_run_result.get("original_sql") or "" @@ -649,11 +671,7 @@ async def ask( query=user_query, instructions=instructions, invalid_generation_result={ - "sql": ( - "" - if is_schema_grounding_error - else original_sql - ), + "sql": original_sql, "error": correction_error_message, }, project_id=ask_request.project_id, From 318707763eb6b76a91b47bed30110d3a210aeb7f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 5 Aug 2026 00:08:24 +0530 Subject: [PATCH 0848/1087] Fail fast on ungrounded SQL correction --- .../src/pipelines/generation/utils/sql.py | 33 +++++++++++++++++-- wren-ai-service/src/web/v1/services/ask.py | 7 ++++ .../tests/pytest/services/test_ask.py | 21 ++++++++---- 3 files changed, 53 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 5caff02f30..eaf65107c0 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -804,14 +804,43 @@ async def run( "invalid_generation_result": invalid_generation_result, } - cleaned_generation_result = clean_generation_result(replies[0]) + raw_generation_result = replies[0] + if isinstance(raw_generation_result, list): + raw_generation_result = ( + raw_generation_result[0] if raw_generation_result else None + ) + + if ( + not isinstance(raw_generation_result, str) + or not raw_generation_result.strip() + ): + ( + valid_generation_result, + invalid_generation_result, + ) = await self._classify_generation_result( + None, + project_id=project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + data_source=data_source, + allow_data_preview=allow_data_preview, + ) + return { + "valid_generation_result": valid_generation_result, + "invalid_generation_result": invalid_generation_result, + } + + cleaned_generation_result = clean_generation_result(raw_generation_result) # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' if cleaned_generation_result.startswith("{"): cleaned_generation_result = orjson.loads(cleaned_generation_result).get( "sql" ) - if not cleaned_generation_result: + if ( + not isinstance(cleaned_generation_result, str) + or not cleaned_generation_result.strip() + ): ( valid_generation_result, invalid_generation_result, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 238b669935..ec3e4c1378 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -23,6 +23,8 @@ _NON_REPAIRABLE_SQL_VALIDATION_TYPES = { "NO_RELEVANT_SQL", + "SCHEMA_GROUNDING", + "SQL_SHAPE", "SQL_VALUE_GROUNDING", } @@ -606,6 +608,11 @@ async def ask( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: + original_sql = failed_dry_run_result.get("original_sql") or "" + invalid_sql = failed_dry_run_result.get("sql") or original_sql + error_message = ( + failed_dry_run_result.get("error") or "No relevant SQL" + ) schema_grounding_correction_attempted = False while current_sql_correction_retries < max_sql_correction_retries: if not should_attempt_sql_correction(failed_dry_run_result): diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 2de88e2264..11ab19d98b 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -13,6 +13,7 @@ AskRequest, AskResultRequest, AskService, + should_attempt_sql_correction, should_skip_sql_diagnosis, ) from src.web.v1.services.semantics_preparation import ( @@ -171,6 +172,15 @@ def test_should_skip_sql_diagnosis_for_deterministic_validation_errors(): assert should_skip_sql_diagnosis({}) is False +def test_should_not_attempt_sql_correction_for_contract_validation_errors(): + assert should_attempt_sql_correction({"type": "SQL_SHAPE"}) is False + assert should_attempt_sql_correction({"type": "SCHEMA_GROUNDING"}) is False + assert should_attempt_sql_correction({"type": "SQL_VALUE_GROUNDING"}) is False + assert should_attempt_sql_correction({"type": "NO_RELEVANT_SQL"}) is False + assert should_attempt_sql_correction({"type": "TIME_OUT"}) is False + assert should_attempt_sql_correction({"type": "SQL_SYNTAX"}) is True + + class _EmptyRetrievalPipeline: async def run(self, **_): return {"formatted_output": {"documents": []}} @@ -242,7 +252,7 @@ async def run(self, **kwargs): @pytest.mark.asyncio -async def test_ask_skips_sql_diagnosis_for_local_validation_error(): +async def test_ask_does_not_repair_local_validation_error(): correction = _CapturingCorrectionPipeline() diagnosis = _FailingDiagnosisPipeline() ask_service = AskService( @@ -269,12 +279,11 @@ async def test_ask_skips_sql_diagnosis_for_local_validation_error(): ask_result_response = ask_service.get_ask_result( AskResultRequest(query_id=query_id) ) - assert ask_result_response.status == "finished" + assert ask_result_response.status == "failed" assert diagnosis.calls == [] - assert correction.calls[0]["invalid_generation_result"] == { - "sql": "SELECT entity_id FROM model_alpha", - "error": "Generated SQL is a table preview.", - } + assert correction.calls == [] + assert ask_result_response.error.code == "NO_RELEVANT_SQL" + assert ask_result_response.invalid_sql == "SELECT entity_id FROM model_alpha" @pytest.mark.asyncio From 895fca680899977601ed2631e7553092b378af01 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 5 Aug 2026 00:47:52 +0530 Subject: [PATCH 0849/1087] Allow correction for grounded intent mismatches --- wren-ai-service/src/pipelines/generation/utils/sql.py | 2 +- wren-ai-service/src/web/v1/services/ask.py | 1 + .../pytest/pipelines/generation/test_sql_post_processor.py | 6 +++--- wren-ai-service/tests/pytest/services/test_ask.py | 2 ++ 4 files changed, 7 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index eaf65107c0..14c6790089 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -946,7 +946,7 @@ async def run( "invalid_generation_result": { "sql": cleaned_generation_result, "original_sql": cleaned_generation_result, - "type": "SQL_SHAPE", + "type": "SQL_INTENT_MISMATCH", "error": table_preview_error, "correlation_id": "", }, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ec3e4c1378..cb54998c84 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -15,6 +15,7 @@ _DETERMINISTIC_SQL_VALIDATION_TYPES = { "SCHEMA_GROUNDING", + "SQL_INTENT_MISMATCH", "SQL_SHAPE", "SQL_SYNTAX", "SQL_VALUE_GROUNDING", diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py index 4965652825..abc0782937 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_post_processor.py @@ -300,7 +300,7 @@ async def test_sql_post_processor_rejects_unshaped_analytical_table_preview(): assert engine.executed is False assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" + assert result["invalid_generation_result"]["type"] == "SQL_INTENT_MISMATCH" assert ( result["invalid_generation_result"]["error"] == "Generated SQL does not apply the requested aggregation, grouping, " @@ -326,7 +326,7 @@ async def test_sql_post_processor_rejects_unfiltered_timeframe_table_preview(): assert engine.executed is False assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" + assert result["invalid_generation_result"]["type"] == "SQL_INTENT_MISMATCH" @pytest.mark.asyncio @@ -388,7 +388,7 @@ async def test_sql_post_processor_rejects_aggregate_intent_without_aggregate_sha assert engine.executed is False assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "SQL_SHAPE" + assert result["invalid_generation_result"]["type"] == "SQL_INTENT_MISMATCH" assert ( result["invalid_generation_result"]["error"] == "Generated SQL does not apply the requested aggregation, grouping, " diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 11ab19d98b..e0f995f8a7 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -165,6 +165,7 @@ def test_ask_service_uses_single_sql_correction_retry_by_default(): def test_should_skip_sql_diagnosis_for_deterministic_validation_errors(): + assert should_skip_sql_diagnosis({"type": "SQL_INTENT_MISMATCH"}) is True assert should_skip_sql_diagnosis({"type": "SQL_SHAPE"}) is True assert should_skip_sql_diagnosis({"type": "SCHEMA_GROUNDING"}) is True assert should_skip_sql_diagnosis({"type": "SQL_VALUE_GROUNDING"}) is True @@ -173,6 +174,7 @@ def test_should_skip_sql_diagnosis_for_deterministic_validation_errors(): def test_should_not_attempt_sql_correction_for_contract_validation_errors(): + assert should_attempt_sql_correction({"type": "SQL_INTENT_MISMATCH"}) is True assert should_attempt_sql_correction({"type": "SQL_SHAPE"}) is False assert should_attempt_sql_correction({"type": "SCHEMA_GROUNDING"}) is False assert should_attempt_sql_correction({"type": "SQL_VALUE_GROUNDING"}) is False From a7f23a575be5450681b6fe2c2dad4ad0ddac4b40 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 5 Aug 2026 01:21:22 +0530 Subject: [PATCH 0850/1087] Tighten grounded SQL correction --- wren-ai-service/src/core/engine.py | 5 ++++- .../src/pipelines/generation/sql_correction.py | 4 ++++ .../src/pipelines/generation/sql_generation.py | 3 +++ .../src/pipelines/generation/utils/sql.py | 16 ++++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 1 - .../tests/pytest/services/test_ask.py | 2 +- 6 files changed, 28 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/core/engine.py b/wren-ai-service/src/core/engine.py index 70a151b4ce..6b144761cc 100644 --- a/wren-ai-service/src/core/engine.py +++ b/wren-ai-service/src/core/engine.py @@ -26,7 +26,10 @@ async def execute_sql( ... -def clean_generation_result(result: str) -> str: +def clean_generation_result(result: Optional[str]) -> Optional[str]: + if result is None: + return None + def _normalize_whitespace(s: str) -> str: return re.sub(r"\s+", " ", s).strip() diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 6e5f35a589..f4ddeb9d17 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -97,6 +97,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) If the error says the SQL is a broad table preview, table preview, missing requested aggregation, missing requested grouping, missing timeframe, or missing ordering/ranking, rebuild the query shape from the user's question. When DATABASE SCHEMA contains role or semantic hints, use those hints only to choose actual declared columns. Do not write role labels, sample schema names, placeholder table names, placeholder column names, or replacement markers as SQL identifiers or SQL literal values. For timeframe requests, filter an actual declared time/date column with a bounded range. For aggregate, "by", trend, or ranking requests, aggregate actual declared measure columns or count rows, group by actual declared dimension/date columns, order by the selected aggregate alias when ranking, and limit only when requested. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. +If the diagnostic says string filter values are not grounded, discard every unsupported literal predicate from the failed SQL and rebuild from the current user question and current user instructions. Do not replace an unsupported literal with another literal. Keep only literal values explicitly present in the current question or instructions, plus concrete date/time boundaries derived from an explicit timeframe. +A WHERE predicate is allowed only when the current user question or current user instructions explicitly request that filter or comparison, or when it is a date/time boundary derived directly from an explicit timeframe. Do not add filters to narrow the result, choose a default segment, or satisfy an assumed business rule. +If the current question asks for a specific entity but omits the entity value, return null for sql instead of fabricating a placeholder, ID, name, or code. +When correcting missing ranking or ordering, add only the requested grouping, ordering, and limit using grounded selected fields or measures. Do not add unrelated WHERE predicates. Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the failed SQL or the wording of the user's question. Retrieved schema objects are ranked candidates, not automatic datasets to merge. Prefer one grounded model, view, or metric that answers the question. Do not use UNION, UNION ALL, INTERSECT, or EXCEPT to combine similar retrieved candidates unless the current user explicitly asks to combine separate result sets and DATABASE SCHEMA grounds each branch with the same result shape and compatible measure meaning. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 26d8f044dc..4017ed894e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -80,12 +80,15 @@ Do not ignore a literal filter value from the user; apply it to the exact schema field representing that filter concept, or return null when that field is unavailable. String literals in WHERE or HAVING must come from the current user question or current user instructions only. Never use schema descriptions, column comments, aliases, display labels, source names, or lineage names as data values. Copy user-provided filter values exactly into SQL string literals, except for normal SQL string escaping. Do not replace them with descriptive labels, unresolved variables, or values to be filled in later. +A WHERE predicate is allowed only when the current user question or current user instructions explicitly request that filter or comparison, or when it is a date/time boundary derived directly from an explicit timeframe in the current question. Do not add filters to narrow the result, choose a default segment, or satisfy an assumed business rule. +If the user asks for a specific entity but does not provide the entity value, return null for sql instead of fabricating a placeholder, ID, name, or code. Never return template SQL. If any required table, column, join, filter value, timeframe boundary, measure, or function is not fully grounded now, return null for sql instead of a partial query. Do not invent generic table names, generic column names, join keys, common-column placeholders, or substitute identifiers from the wording of the user's question. For ranked entity questions, select and group by the exact schema field representing the requested entity, not only context fields. For record or entity volume questions, count rows unless the user requests a declared numeric measure. For value, amount, quantity, rate, cost, or metric questions, use the declared measure that represents the request. For timeframe requests, filter an actual declared column whose metadata marks it as the requested time concept. Metadata role labels are not executable column names. For aggregate, trend, ranking, or grouped requests, aggregate actual declared measure columns or count rows as appropriate for the user's requested measure. Metadata role labels are not executable column names. +For ranking requests, include the requested ordering and limit. Do not introduce unrelated filters to create or justify the ranking. Retrieved schema objects are ranked candidates, not automatic datasets to merge. Prefer one grounded model, view, or metric that answers the question. Do not use UNION, UNION ALL, INTERSECT, or EXCEPT to combine similar retrieved candidates unless the current user explicitly asks to combine separate result sets and DATABASE SCHEMA grounds each branch with the same result shape and compatible measure meaning. For comparison requests, include every requested comparison group or period in the SQL result and compute the requested difference, change, growth, or ranking when the required fields and date operations are grounded. Do not answer a comparison request with only one side of the comparison. Do not copy executable identifiers, SQL fragments, functions, or literal values from reasoning plans, SQL samples, failed SQL, source metadata, comments, or user wording unless they are also exact deployed schema identifiers or current user-provided literal values. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 14c6790089..ca5aed1ef6 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -831,6 +831,22 @@ async def run( } cleaned_generation_result = clean_generation_result(raw_generation_result) + if not cleaned_generation_result: + ( + valid_generation_result, + invalid_generation_result, + ) = await self._classify_generation_result( + None, + project_id=project_id, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + data_source=data_source, + allow_data_preview=allow_data_preview, + ) + return { + "valid_generation_result": valid_generation_result, + "invalid_generation_result": invalid_generation_result, + } # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' if cleaned_generation_result.startswith("{"): diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cb54998c84..2c8a0a6a33 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -26,7 +26,6 @@ "NO_RELEVANT_SQL", "SCHEMA_GROUNDING", "SQL_SHAPE", - "SQL_VALUE_GROUNDING", } diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index e0f995f8a7..24aeec5d01 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -177,7 +177,7 @@ def test_should_not_attempt_sql_correction_for_contract_validation_errors(): assert should_attempt_sql_correction({"type": "SQL_INTENT_MISMATCH"}) is True assert should_attempt_sql_correction({"type": "SQL_SHAPE"}) is False assert should_attempt_sql_correction({"type": "SCHEMA_GROUNDING"}) is False - assert should_attempt_sql_correction({"type": "SQL_VALUE_GROUNDING"}) is False + assert should_attempt_sql_correction({"type": "SQL_VALUE_GROUNDING"}) is True assert should_attempt_sql_correction({"type": "NO_RELEVANT_SQL"}) is False assert should_attempt_sql_correction({"type": "TIME_OUT"}) is False assert should_attempt_sql_correction({"type": "SQL_SYNTAX"}) is True From 5f2d006721b24fb9144bf24ef74aee93440af62b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 5 Aug 2026 02:14:08 +0530 Subject: [PATCH 0851/1087] Stop ungrounded SQL correction retries --- wren-ai-service/src/web/v1/services/ask.py | 33 ++++++---- .../src/web/v1/services/ask_feedback.py | 60 ++++++++++--------- 2 files changed, 54 insertions(+), 39 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 2c8a0a6a33..6da241e169 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -25,7 +25,9 @@ _NON_REPAIRABLE_SQL_VALIDATION_TYPES = { "NO_RELEVANT_SQL", "SCHEMA_GROUNDING", + "SQL_INTENT_MISMATCH", "SQL_SHAPE", + "SQL_VALUE_GROUNDING", } @@ -594,9 +596,12 @@ async def ask( "SQL generation", ) - if sql_valid_result := text_to_sql_generation_results["post_process"][ + generation_post_process = ( + text_to_sql_generation_results.get("post_process") or {} + ) + if sql_valid_result := generation_post_process.get( "valid_generation_result" - ]: + ): api_results = [ AskResult( **{ @@ -605,9 +610,9 @@ async def ask( } ) ] - elif failed_dry_run_result := text_to_sql_generation_results[ - "post_process" - ]["invalid_generation_result"]: + elif failed_dry_run_result := generation_post_process.get( + "invalid_generation_result" + ): original_sql = failed_dry_run_result.get("original_sql") or "" invalid_sql = failed_dry_run_result.get("sql") or original_sql error_message = ( @@ -678,7 +683,8 @@ async def ask( query=user_query, instructions=instructions, invalid_generation_result={ - "sql": original_sql, + "original_sql": original_sql, + "sql": invalid_sql, "error": correction_error_message, }, project_id=ask_request.project_id, @@ -693,9 +699,12 @@ async def ask( "SQL correction", ) - if valid_generation_result := sql_correction_results[ - "post_process" - ]["valid_generation_result"]: + correction_post_process = ( + sql_correction_results.get("post_process") or {} + ) + if valid_generation_result := correction_post_process.get( + "valid_generation_result" + ): api_results = [ AskResult( **{ @@ -706,9 +715,9 @@ async def ask( ] break - next_failed_dry_run_result = sql_correction_results[ - "post_process" - ]["invalid_generation_result"] + next_failed_dry_run_result = correction_post_process.get( + "invalid_generation_result" + ) if not next_failed_dry_run_result: break if ( diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index e464645e2e..ca00e9d875 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -13,6 +13,7 @@ AskError, AskResult, run_pipeline_with_timeout, + should_attempt_sql_correction, should_skip_sql_diagnosis, ) @@ -218,9 +219,12 @@ async def ask_feedback( "SQL regeneration", ) - if sql_valid_result := text_to_sql_generation_results["post_process"][ + generation_post_process = ( + text_to_sql_generation_results.get("post_process") or {} + ) + if sql_valid_result := generation_post_process.get( "valid_generation_result" - ]: + ): api_results = [ AskResult( **{ @@ -229,19 +233,18 @@ async def ask_feedback( } ) ] - elif failed_dry_run_result := text_to_sql_generation_results[ - "post_process" - ]["invalid_generation_result"]: - if failed_dry_run_result["type"] != "TIME_OUT": - original_sql = failed_dry_run_result["original_sql"] - invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + elif failed_dry_run_result := generation_post_process.get( + "invalid_generation_result" + ): + if should_attempt_sql_correction(failed_dry_run_result): + original_sql = failed_dry_run_result.get("original_sql") or "" + invalid_sql = failed_dry_run_result.get("sql") or original_sql + error_message = ( + failed_dry_run_result.get("error") or "No relevant SQL" + ) skip_sql_diagnosis = should_skip_sql_diagnosis( failed_dry_run_result ) - is_schema_grounding_error = ( - failed_dry_run_result.get("type") == "SCHEMA_GROUNDING" - ) sql_diagnosis_reasoning = None self._ask_feedback_results[ @@ -281,11 +284,7 @@ async def ask_feedback( instructions=instructions, invalid_generation_result={ "original_sql": original_sql, - "sql": ( - "" - if is_schema_grounding_error - else invalid_sql - ), + "sql": invalid_sql, "error": correction_error_message, }, project_id=ask_feedback_request.project_id, @@ -298,9 +297,12 @@ async def ask_feedback( "SQL correction", ) - if valid_generation_result := sql_correction_results[ - "post_process" - ]["valid_generation_result"]: + correction_post_process = ( + sql_correction_results.get("post_process") or {} + ) + if valid_generation_result := correction_post_process.get( + "valid_generation_result" + ): api_results = [ AskResult( **{ @@ -309,14 +311,18 @@ async def ask_feedback( } ) ] - elif failed_dry_run_result := sql_correction_results[ - "post_process" - ]["invalid_generation_result"]: - invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + elif failed_dry_run_result := correction_post_process.get( + "invalid_generation_result" + ): + invalid_sql = failed_dry_run_result.get("sql") or invalid_sql + error_message = ( + failed_dry_run_result.get("error") or error_message + ) else: - invalid_sql = failed_dry_run_result["sql"] - error_message = failed_dry_run_result["error"] + invalid_sql = failed_dry_run_result.get("sql") + error_message = ( + failed_dry_run_result.get("error") or "No relevant SQL" + ) if api_results: if not self._is_stopped(query_id, self._ask_feedback_results): From ac6d218d0194bc5e8c546710547fc1356194e77b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 5 Aug 2026 13:51:27 +0530 Subject: [PATCH 0852/1087] Improve SQL dialect grounding and add fine-tuning tools --- .../generation/followup_sql_generation.py | 13 +- .../pipelines/generation/sql_correction.py | 13 +- .../pipelines/generation/sql_generation.py | 13 +- .../src/pipelines/generation/utils/sql.py | 23 ++ .../src/pipelines/retrieval/sql_functions.py | 4 +- .../src/pipelines/retrieval/sql_knowledge.py | 4 +- wren-ai-service/tools/fine_tuning/.gitignore | 1 + wren-ai-service/tools/fine_tuning/README.md | 66 ++++++ .../axolotl-codestral-sql-qlora.yml | 54 +++++ .../example_verified_sql_examples.jsonl | 10 + .../tools/fine_tuning/prepare_sft_dataset.py | 220 ++++++++++++++++++ 11 files changed, 416 insertions(+), 5 deletions(-) create mode 100644 wren-ai-service/tools/fine_tuning/.gitignore create mode 100644 wren-ai-service/tools/fine_tuning/README.md create mode 100644 wren-ai-service/tools/fine_tuning/axolotl-codestral-sql-qlora.yml create mode 100644 wren-ai-service/tools/fine_tuning/example_verified_sql_examples.jsonl create mode 100644 wren-ai-service/tools/fine_tuning/prepare_sft_dataset.py 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 8dffa2fd6f..d7489c15df 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -20,6 +20,7 @@ get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, + get_sql_dialect_instructions, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -59,6 +60,10 @@ {% endfor %} {% endif %} +{% if sql_dialect_instructions %} +{{ sql_dialect_instructions }} +{% endif %} + {% if sql_samples %} ### SQL SAMPLES ### These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. @@ -120,10 +125,12 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, schema_contracts: list[dict] | None = None, + data_source: str | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + sql_dialect_instructions=get_sql_dialect_instructions(data_source), executable_schema_contract=build_executable_schema_contract(schema_contracts), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( @@ -211,6 +218,7 @@ def __init__( template=text_to_sql_with_followup_user_prompt_template ), "post_processor": SQLGenPostProcessor(engine=engine), + "data_source": kwargs.get("data_source", "local_file"), } super().__init__( @@ -247,6 +255,9 @@ async def run( ) else: metadata = {} + data_source = metadata.get("data_source") or self._components.get( + "data_source", "local_file" + ) return await self._pipe.execute( ["post_process"], @@ -265,7 +276,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "sql_knowledge": sql_knowledge, "schema_contracts": schema_contracts, **self._components, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f4ddeb9d17..810da6d1ad 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -17,6 +17,7 @@ SQLGenPostProcessor, build_executable_schema_contract, construct_instructions, + get_sql_dialect_instructions, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -71,6 +72,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endfor %} {% endif %} +{% if sql_dialect_instructions %} +{{ sql_dialect_instructions }} +{% endif %} + {% if instructions %} ### USER INSTRUCTIONS ### {% for instruction in instructions %} @@ -119,9 +124,11 @@ def prompt( instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, schema_contracts: list[dict] | None = None, + data_source: str | None = None, ) -> dict: _prompt = prompt_builder.run( documents=documents, + sql_dialect_instructions=get_sql_dialect_instructions(data_source), executable_schema_contract=build_executable_schema_contract(schema_contracts), invalid_generation_result=invalid_generation_result, query=query or "", @@ -194,6 +201,7 @@ def __init__( template=sql_correction_user_prompt_template ), "post_processor": SQLGenPostProcessor(engine=engine), + "data_source": kwargs.get("data_source", "local_file"), } super().__init__( @@ -226,6 +234,9 @@ async def run( ) else: metadata = {} + data_source = metadata.get("data_source") or self._components.get( + "data_source", "local_file" + ) return await self._pipe.execute( ["post_process"], @@ -239,7 +250,7 @@ async def run( "mdl_hash": mdl_hash, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "sql_knowledge": sql_knowledge, "schema_contracts": schema_contracts, **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 4017ed894e..53246ad12a 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -19,6 +19,7 @@ get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, + get_sql_dialect_instructions, get_sql_generation_system_prompt, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -53,6 +54,10 @@ {% endfor %} {% endif %} +{% if sql_dialect_instructions %} +{{ sql_dialect_instructions }} +{% endif %} + {% if sql_samples %} ### SQL SAMPLES ### These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. @@ -122,10 +127,12 @@ def prompt( sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, schema_contracts: list[dict] | None = None, + data_source: str | None = None, ) -> dict: _prompt = prompt_builder.run( query=query, documents=documents, + sql_dialect_instructions=get_sql_dialect_instructions(data_source), executable_schema_contract=build_executable_schema_contract(schema_contracts), sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( @@ -211,6 +218,7 @@ def __init__( template=sql_generation_user_prompt_template ), "post_processor": SQLGenPostProcessor(engine=engine), + "data_source": kwargs.get("data_source", "local_file"), } super().__init__( @@ -247,6 +255,9 @@ async def run( ) else: metadata = {} + data_source = metadata.get("data_source") or self._components.get( + "data_source", "local_file" + ) return await self._pipe.execute( ["post_process"], @@ -264,7 +275,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, "schema_contracts": schema_contracts, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ca5aed1ef6..68568420cd 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1442,6 +1442,29 @@ class SqlGenerationResult(BaseModel): } +def get_sql_dialect_instructions(data_source: str | None = None) -> str: + normalized = (data_source or "").strip().lower() + + if normalized in {"mssql", "sqlserver", "sql_server", "microsoft_sql_server"}: + return """ +### SQL DIALECT ### +Target database: Microsoft SQL Server. +- Do not use EXTRACT(... FROM ...); SQL Server does not support that syntax. +- Prefer bounded date ranges for month/year filters. +- If a date part is required, use SQL Server functions such as MONTH(date_column) and YEAR(date_column). +- Do not use LIMIT; use TOP only when a row limit is explicitly requested. +""" + + if normalized and normalized != "local_file": + return f""" +### SQL DIALECT ### +Target database: {normalized}. +Use only SQL syntax and functions supported by this target database and by SQL FUNCTIONS. +""" + + return "" + + def construct_instructions( instructions: list[dict] | None = None, ): diff --git a/wren-ai-service/src/pipelines/retrieval/sql_functions.py b/wren-ai-service/src/pipelines/retrieval/sql_functions.py index 822e346703..833d02645a 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_functions.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_functions.py @@ -115,7 +115,9 @@ async def run( self._retriever, mdl_hash=mdl_hash, ) - _data_source = metadata.get("data_source", "local_file") + _data_source = metadata.get("data_source") or getattr( + self._components["engine"], "_source", None + ) or "local_file" if _data_source in self._cache: logger.info(f"Hit cache of SQL Functions for {_data_source}") diff --git a/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py b/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py index ad6dde9ad1..3b5388fb51 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py @@ -125,7 +125,9 @@ async def run( self._retriever, mdl_hash=mdl_hash, ) - _data_source = metadata.get("data_source", "local_file") + _data_source = metadata.get("data_source") or getattr( + self._components["engine"], "_source", None + ) or "local_file" if _data_source in self._cache: logger.info(f"Hit cache of SQL Knowledge for {_data_source}") diff --git a/wren-ai-service/tools/fine_tuning/.gitignore b/wren-ai-service/tools/fine_tuning/.gitignore new file mode 100644 index 0000000000..c18dd8d83c --- /dev/null +++ b/wren-ai-service/tools/fine_tuning/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/wren-ai-service/tools/fine_tuning/README.md b/wren-ai-service/tools/fine_tuning/README.md new file mode 100644 index 0000000000..67470e97e7 --- /dev/null +++ b/wren-ai-service/tools/fine_tuning/README.md @@ -0,0 +1,66 @@ +# Codestral SQL Fine-Tuning + +This folder prepares verified Wren text-to-SQL examples for supervised fine-tuning. +It does not invent database examples. Use only production-reviewed question/SQL pairs. + +## Input Format + +Create a JSONL file where each line has this shape: + +```json +{ + "question": "Natural language question", + "schema": "Relevant deployed schema DDL or structured schema text", + "relationships": "Approved join paths and relationship definitions", + "business_metadata": "Business definitions, metrics, dimensions, and rules", + "sql_dialect": "mssql", + "expected_sql": "Verified production SQL" +} +``` + +For unanswerable examples, use: + +```json +{ + "question": "Natural language question", + "schema": "Relevant deployed schema DDL or structured schema text", + "relationships": "Approved join paths and relationship definitions", + "business_metadata": "Business definitions, metrics, dimensions, and rules", + "sql_dialect": "mssql", + "expected_sql": null, + "insufficient_information_reason": "The required field is not present in the deployed schema." +} +``` + +## Prepare Dataset + +```bash +python tools/fine_tuning/prepare_sft_dataset.py \ + --input eval/dataset/verified_sql_examples.jsonl \ + --output-dir outputs/fine_tuning/codestral-sql \ + --seed 42 +``` + +Outputs: + +- `train.jsonl` +- `validation.jsonl` +- `test.jsonl` +- `dataset_report.json` + +Each output row is a chat-style SFT sample with `system`, `user`, and `assistant` +messages. + +## Train + +Use the generated files with a GPU training tool such as Axolotl. `axolotl-codestral-sql-qlora.yml` +is a template. Update model path, dataset paths, GPU settings, and output directory +before running it. + +## Serve + +Serve the fine-tuned adapter behind an OpenAI-compatible endpoint, then point +`wren-ai-service/config.yaml` `api_base` to that endpoint. + +Keep Wren retrieval enabled. The adapter should learn how to obey schema context, +not memorize the database. diff --git a/wren-ai-service/tools/fine_tuning/axolotl-codestral-sql-qlora.yml b/wren-ai-service/tools/fine_tuning/axolotl-codestral-sql-qlora.yml new file mode 100644 index 0000000000..64cb380bc1 --- /dev/null +++ b/wren-ai-service/tools/fine_tuning/axolotl-codestral-sql-qlora.yml @@ -0,0 +1,54 @@ +# Template only. Run this on a Linux GPU training machine after generating +# train.jsonl/validation.jsonl with prepare_sft_dataset.py. + +base_model: mistralai/Codestral-22B-v0.1 +model_type: AutoModelForCausalLM +tokenizer_type: AutoTokenizer + +load_in_4bit: true +strict: false + +datasets: + - path: outputs/fine_tuning/codestral-sql/train.jsonl + type: chat_template + +test_datasets: + - path: outputs/fine_tuning/codestral-sql/validation.jsonl + type: chat_template + +dataset_prepared_path: outputs/fine_tuning/codestral-sql/prepared +output_dir: outputs/fine_tuning/codestral-sql/adapter + +sequence_len: 8192 +sample_packing: true +pad_to_sequence_len: true + +adapter: qlora +lora_model_dir: +lora_r: 64 +lora_alpha: 128 +lora_dropout: 0.05 +lora_target_linear: true + +gradient_accumulation_steps: 8 +micro_batch_size: 1 +num_epochs: 3 +optimizer: paged_adamw_8bit +lr_scheduler: cosine +learning_rate: 0.00002 +weight_decay: 0.0 +warmup_ratio: 0.03 +max_grad_norm: 1.0 + +bf16: auto +tf32: true +gradient_checkpointing: true +flash_attention: true + +evals_per_epoch: 2 +saves_per_epoch: 1 +save_total_limit: 3 +logging_steps: 10 + +special_tokens: + pad_token: "" diff --git a/wren-ai-service/tools/fine_tuning/example_verified_sql_examples.jsonl b/wren-ai-service/tools/fine_tuning/example_verified_sql_examples.jsonl new file mode 100644 index 0000000000..9080fbd5e4 --- /dev/null +++ b/wren-ai-service/tools/fine_tuning/example_verified_sql_examples.jsonl @@ -0,0 +1,10 @@ +{"question":"Count records by business dimension for a specified month.","schema":"CREATE TABLE business_model (record_id INT PRIMARY KEY, business_dimension VARCHAR, event_date DATE);","relationships":"No joins required.","business_metadata":"business_model is the approved business-facing model for this subject. event_date is the approved date filter field.","sql_dialect":"mssql","expected_sql":"SELECT business_dimension, COUNT(record_id) AS record_count FROM business_model WHERE event_date >= '2026-01-01' AND event_date < '2026-02-01' GROUP BY business_dimension ORDER BY record_count DESC"} +{"question":"Show records by two approved dimensions.","schema":"CREATE TABLE business_model (record_id INT PRIMARY KEY, dimension_one VARCHAR, dimension_two VARCHAR);","relationships":"No joins required.","business_metadata":"dimension_one and dimension_two are approved grouping dimensions.","sql_dialect":"mssql","expected_sql":"SELECT dimension_one, dimension_two, COUNT(record_id) AS record_count FROM business_model GROUP BY dimension_one, dimension_two ORDER BY record_count DESC"} +{"question":"List the top groups for an approved measure.","schema":"CREATE TABLE metric_model (group_name VARCHAR, approved_measure NUMERIC);","relationships":"No joins required.","business_metadata":"approved_measure is the approved numeric measure. group_name is the approved grouping field.","sql_dialect":"mssql","expected_sql":"SELECT TOP 20 group_name, SUM(approved_measure) AS total_measure FROM metric_model GROUP BY group_name ORDER BY total_measure DESC"} +{"question":"Find records for a category value.","schema":"CREATE TABLE business_model (record_id INT PRIMARY KEY, category_name VARCHAR);","relationships":"No joins required.","business_metadata":"category_name stores approved category values.","sql_dialect":"mssql","expected_sql":"SELECT record_id, category_name FROM business_model WHERE category_name = 'Example Category'"} +{"question":"Join a fact model to an approved dimension.","schema":"CREATE TABLE fact_model (record_id INT PRIMARY KEY, dimension_id INT, amount NUMERIC); CREATE TABLE dimension_model (dimension_id INT PRIMARY KEY, dimension_name VARCHAR);","relationships":"fact_model.dimension_id joins to dimension_model.dimension_id. Use this path for dimension_name.","business_metadata":"fact_model is the approved fact model. dimension_model contains approved dimension labels.","sql_dialect":"mssql","expected_sql":"SELECT dimension_model.dimension_name, SUM(fact_model.amount) AS total_amount FROM fact_model INNER JOIN dimension_model ON fact_model.dimension_id = dimension_model.dimension_id GROUP BY dimension_model.dimension_name ORDER BY total_amount DESC"} +{"question":"Compute a percentage by group.","schema":"CREATE TABLE metric_model (group_name VARCHAR, numerator NUMERIC, denominator NUMERIC);","relationships":"No joins required.","business_metadata":"percentage_metric is numerator divided by denominator. Use NULLIF to avoid divide by zero.","sql_dialect":"mssql","expected_sql":"SELECT group_name, SUM(numerator) / NULLIF(SUM(denominator), 0) AS percentage_metric FROM metric_model GROUP BY group_name ORDER BY percentage_metric DESC"} +{"question":"Return monthly trend for a metric.","schema":"CREATE TABLE metric_model (event_date DATE, approved_measure NUMERIC);","relationships":"No joins required.","business_metadata":"event_date is the approved trend date. approved_measure is the approved measure.","sql_dialect":"mssql","expected_sql":"SELECT DATEPART(YEAR, event_date) AS year_value, DATEPART(MONTH, event_date) AS month_value, SUM(approved_measure) AS total_measure FROM metric_model GROUP BY DATEPART(YEAR, event_date), DATEPART(MONTH, event_date) ORDER BY year_value, month_value"} +{"question":"Filter by current period when a date field exists.","schema":"CREATE TABLE business_model (record_id INT PRIMARY KEY, event_date DATE);","relationships":"No joins required.","business_metadata":"event_date is the approved period filter field.","sql_dialect":"mssql","expected_sql":"SELECT COUNT(record_id) AS record_count FROM business_model WHERE event_date >= DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1) AND event_date < DATEADD(MONTH, 1, DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1))"} +{"question":"Answer a question requiring a missing field.","schema":"CREATE TABLE business_model (record_id INT PRIMARY KEY, event_date DATE);","relationships":"No approved joins are available.","business_metadata":"Only record_id and event_date are available.","sql_dialect":"mssql","expected_sql":null,"insufficient_information_reason":"The requested business field is not present in the deployed schema or metadata."} +{"question":"Answer a question requiring an unapproved join.","schema":"CREATE TABLE first_model (id INT PRIMARY KEY); CREATE TABLE second_model (id INT PRIMARY KEY);","relationships":"No relationship is approved between first_model and second_model.","business_metadata":"Only approved relationships may be used.","sql_dialect":"mssql","expected_sql":null,"insufficient_information_reason":"No approved relationship exists for the requested join."} diff --git a/wren-ai-service/tools/fine_tuning/prepare_sft_dataset.py b/wren-ai-service/tools/fine_tuning/prepare_sft_dataset.py new file mode 100644 index 0000000000..6129cb5110 --- /dev/null +++ b/wren-ai-service/tools/fine_tuning/prepare_sft_dataset.py @@ -0,0 +1,220 @@ +import argparse +import json +import random +import re +from pathlib import Path +from typing import Any + + +REQUIRED_FIELDS = { + "question", + "schema", + "relationships", + "business_metadata", + "sql_dialect", +} + +SQL_IDENTIFIER_PATTERN = re.compile( + r"\b(?:FROM|JOIN|UPDATE|INTO)\s+([A-Za-z_][\w.$\[\]\"]*)", + re.IGNORECASE, +) + + +SYSTEM_PROMPT = """You are a text-to-SQL model for Wren AI. +Generate SQL only from the provided database schema, relationships, and business metadata. +Use only declared tables, columns, relationships, metrics, and business definitions. +Never invent tables, columns, joins, filters, calculations, or business logic. +Prefer business-facing models over technical, staging, temporary, or raw tables when both are available. +If the requested information cannot be fully grounded, return {"sql": null, "reason": "INSUFFICIENT_INFORMATION"}. +Return only JSON.""" + + +def read_jsonl(path: Path) -> list[dict[str, Any]]: + rows = [] + with path.open("r", encoding="utf-8") as file: + for line_number, line in enumerate(file, start=1): + stripped = line.strip() + if not stripped: + continue + try: + row = json.loads(stripped) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}:{line_number}: invalid JSON: {exc}") from exc + row["_line_number"] = line_number + rows.append(row) + return rows + + +def validate_row(row: dict[str, Any]) -> list[str]: + errors = [] + missing = sorted(field for field in REQUIRED_FIELDS if not row.get(field)) + if missing: + errors.append(f"missing required fields: {', '.join(missing)}") + + has_sql = bool(row.get("expected_sql")) + has_insufficient_reason = bool(row.get("insufficient_information_reason")) + if not has_sql and not has_insufficient_reason: + errors.append( + "expected_sql is empty; insufficient_information_reason is required" + ) + + if has_sql and has_insufficient_reason: + errors.append( + "provide either expected_sql or insufficient_information_reason, not both" + ) + + if has_sql: + sql = str(row["expected_sql"]).strip() + if not sql.lower().startswith(("select", "with")): + errors.append("expected_sql must start with SELECT or WITH") + if re.search(r"\b(select\s+\*)\b", sql, re.IGNORECASE): + errors.append("expected_sql must not use SELECT *") + + return errors + + +def build_user_content(row: dict[str, Any]) -> str: + return "\n\n".join( + [ + f"QUESTION:\n{row['question']}", + f"SQL_DIALECT:\n{row['sql_dialect']}", + f"DATABASE_SCHEMA:\n{row['schema']}", + f"RELATIONSHIPS:\n{row['relationships']}", + f"BUSINESS_METADATA:\n{row['business_metadata']}", + ] + ) + + +def build_assistant_content(row: dict[str, Any]) -> str: + if row.get("expected_sql"): + return json.dumps( + {"sql": str(row["expected_sql"]).strip()}, + ensure_ascii=False, + separators=(",", ":"), + ) + return json.dumps( + { + "sql": None, + "reason": "INSUFFICIENT_INFORMATION", + "detail": str(row["insufficient_information_reason"]).strip(), + }, + ensure_ascii=False, + separators=(",", ":"), + ) + + +def to_chat_sample(row: dict[str, Any]) -> dict[str, Any]: + return { + "messages": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": build_user_content(row)}, + {"role": "assistant", "content": build_assistant_content(row)}, + ], + "metadata": { + "source_line": row["_line_number"], + "sql_dialect": row.get("sql_dialect"), + "tables_referenced": sorted( + set(SQL_IDENTIFIER_PATTERN.findall(str(row.get("expected_sql") or ""))) + ), + }, + } + + +def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8", newline="\n") as file: + for row in rows: + file.write(json.dumps(row, ensure_ascii=False) + "\n") + + +def split_rows( + rows: list[dict[str, Any]], + validation_ratio: float, + test_ratio: float, + seed: int, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]: + shuffled = list(rows) + random.Random(seed).shuffle(shuffled) + + total = len(shuffled) + test_count = round(total * test_ratio) + validation_count = round(total * validation_ratio) + train_count = total - validation_count - test_count + + return ( + shuffled[:train_count], + shuffled[train_count : train_count + validation_count], + shuffled[train_count + validation_count :], + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--validation-ratio", type=float, default=0.1) + parser.add_argument("--test-ratio", type=float, default=0.1) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + rows = read_jsonl(args.input) + validation_errors = [] + valid_rows = [] + for row in rows: + errors = validate_row(row) + if errors: + validation_errors.append( + {"line": row["_line_number"], "errors": errors} + ) + else: + valid_rows.append(row) + + if validation_errors: + args.output_dir.mkdir(parents=True, exist_ok=True) + report_path = args.output_dir / "dataset_report.json" + report_path.write_text( + json.dumps( + { + "status": "failed", + "input_rows": len(rows), + "valid_rows": len(valid_rows), + "validation_errors": validation_errors, + }, + indent=2, + ), + encoding="utf-8", + ) + raise SystemExit(f"Dataset validation failed. See {report_path}") + + if len(valid_rows) < 10: + raise SystemExit("Need at least 10 verified examples before splitting.") + + samples = [to_chat_sample(row) for row in valid_rows] + train, validation, test = split_rows( + samples, + validation_ratio=args.validation_ratio, + test_ratio=args.test_ratio, + seed=args.seed, + ) + + args.output_dir.mkdir(parents=True, exist_ok=True) + write_jsonl(args.output_dir / "train.jsonl", train) + write_jsonl(args.output_dir / "validation.jsonl", validation) + write_jsonl(args.output_dir / "test.jsonl", test) + + report = { + "status": "ok", + "input_rows": len(rows), + "train_rows": len(train), + "validation_rows": len(validation), + "test_rows": len(test), + "seed": args.seed, + } + (args.output_dir / "dataset_report.json").write_text( + json.dumps(report, indent=2), + encoding="utf-8", + ) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() From ea7e1e272d9c6744080f1be61e56275e84fdd050 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 5 Aug 2026 14:25:30 +0530 Subject: [PATCH 0853/1087] Keep data questions in SQL answer flow --- .../generation/intent_classification.py | 54 +++++++++++- .../generation/test_intent_classification.py | 36 ++++++++ .../pages/home/promptThread/AnswerResult.tsx | 86 +++++++++++-------- .../home/promptThread/ViewSQLTabContent.tsx | 24 +++++- 4 files changed, 157 insertions(+), 43 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index cc9a823fe5..2044bdf3f5 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -1,5 +1,6 @@ import ast import logging +import re import sys from typing import Any, Literal, Optional @@ -26,6 +27,50 @@ logger = logging.getLogger("wren-ai-service") +_DATA_QUERY_ACTION_PATTERN = re.compile( + r"\b(" + r"show|list|find|get|give|tell|count|sum|total|average|avg|min|max|" + r"top|bottom|highest|lowest|which|who|what|where|filter|compare|" + r"breakdown|group|trend" + r")\b", + re.IGNORECASE, +) +_DATA_QUERY_CONTEXT_PATTERN = re.compile( + r"\b(" + r"order|orders|sale|sales|customer|customers|invoice|invoices|market|" + r"business unit|country|record|records|row|rows|amount|quantity|date|" + r"today|yesterday|week|month|quarter|year|january|february|march|april|" + r"may|june|july|august|september|october|november|december|current" + r")\b|\b\d{4}\b", + re.IGNORECASE, +) +_USER_GUIDE_SHAPE_PATTERN = re.compile( + r"\b(" + r"how do i|how can i|what can wren|help me use|delete a project|" + r"reset a project|connect to|draw a chart" + r")\b", + re.IGNORECASE, +) + + +def should_force_text_to_sql_intent( + query: str | None, + intent: str | None, + db_schemas: list[str] | None, +) -> bool: + if intent == "TEXT_TO_SQL" or intent == "USER_GUIDE": + return False + if not query or not db_schemas: + return False + if _USER_GUIDE_SHAPE_PATTERN.search(query): + return False + + return bool( + _DATA_QUERY_ACTION_PATTERN.search(query) + and _DATA_QUERY_CONTEXT_PATTERN.search(query) + ) + + intent_classification_system_prompt = """ ### Task ### You are an expert detective specializing in intent classification. Combine the user's current question and previous questions to determine their true intent based on the provided database schema. Classify the intent into one of these categories: `MISLEADING_QUERY`, `TEXT_TO_SQL`, `GENERAL`, or `USER_GUIDE`. Additionally, provide a concise reasoning (maximum 20 words) for your classification. @@ -306,12 +351,17 @@ async def classify_intent(prompt: dict, generator: Any, generator_name: str) -> @observe(capture_input=False) -def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict: +def post_process( + classify_intent: dict, construct_db_schemas: list[str], query: str +) -> dict: try: results = orjson.loads(classify_intent.get("replies")[0]) + intent = results["results"] + if should_force_text_to_sql_intent(query, intent, construct_db_schemas): + intent = "TEXT_TO_SQL" return { "rephrased_question": results["rephrased_question"], - "intent": results["results"], + "intent": intent, "reasoning": results["reasoning"], "db_schemas": construct_db_schemas, } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py b/wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py index 0b14f0e1b0..953abd3c4d 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_intent_classification.py @@ -3,6 +3,8 @@ from src.pipelines.generation.intent_classification import ( dbschema_retrieval, + post_process, + should_force_text_to_sql_intent, table_retrieval, ) @@ -38,6 +40,40 @@ async def test_intent_table_retrieval_scopes_to_deployed_mdl_hash(): } +def test_intent_override_keeps_data_questions_in_text_to_sql(): + assert should_force_text_to_sql_intent( + "show orders placed from the country france", + "GENERAL", + ["table: orders"], + ) + + +def test_intent_override_does_not_force_user_guide_questions(): + assert not should_force_text_to_sql_intent( + "How do I draw a chart?", + "GENERAL", + ["table: orders"], + ) + + +def test_intent_post_process_overrides_general_for_data_question(): + result = post_process( + { + "replies": [ + ( + '{"rephrased_question":"show orders placed from the country france",' + '"reasoning":"asks for database records",' + '"results":"GENERAL"}' + ) + ] + }, + ["table: orders"], + "show orders placed from the country france", + ) + + assert result["intent"] == "TEXT_TO_SQL" + + @pytest.mark.asyncio async def test_intent_dbschema_retrieval_scopes_to_deployed_mdl_hash(): retriever = CapturingRetriever() diff --git a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx index 9c751867f0..aa19c2da07 100644 --- a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx +++ b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx @@ -25,6 +25,7 @@ import ChartAnswer from '@/components/pages/home/promptThread/ChartAnswer'; import Preparation from '@/components/pages/home/preparation'; import { AskingTaskStatus, + AskingTaskType, ChartTaskStatus, ThreadResponse, ThreadResponseAnswerDetail, @@ -253,8 +254,12 @@ export default function AnswerResult(props: Props) { return answerDetail === null && !isEmpty(breakdownDetail); }, [answerDetail, breakdownDetail]); const isAnswerPrepared = !!answerDetail?.queryId || !!answerDetail?.status; + const isTextToSqlResponse = askingTask?.type === AskingTaskType.TEXT_TO_SQL; + const hasDisplaySql = + !!sql || !!askingTask?.invalidSql || !!adjustmentTask?.invalidSql; const showTextOnlyAnswer = isAnswerPrepared && + !isTextToSqlResponse && !sql && !view && !isBreakdownOnly && @@ -306,6 +311,7 @@ export default function AnswerResult(props: Props) { const onTabClick = (activeKey: string) => { if ( activeKey === ANSWER_TAB_KEYS.CHART && + !!threadResponse.sql && !threadResponse.chartDetail && !isChartGenerationActive(threadResponse.chartDetail?.status) ) { @@ -314,10 +320,11 @@ export default function AnswerResult(props: Props) { }; const showAnswerTabs = - !showTextOnlyAnswer && - (askingTask?.status === AskingTaskStatus.FINISHED || - isAnswerPrepared || - isBreakdownOnly); + (!showTextOnlyAnswer && + (askingTask?.status === AskingTaskStatus.FINISHED || + isAnswerPrepared || + isBreakdownOnly)) || + (isTextToSqlResponse && hasDisplaySql); const rephrasedQuestion = threadResponse?.askingTask?.rephrasedQuestion || question; @@ -381,6 +388,7 @@ export default function AnswerResult(props: Props) { @@ -393,45 +401,47 @@ export default function AnswerResult(props: Props) { -
- - + + - onOpenSaveToKnowledgeModal( + onOpenSaveAsViewModal( + { sql, responseId: id }, { - question: rephrasedQuestion, - sql, + rephrasedQuestion: questionForSaveAsView, }, - { isCreateMode: true }, ) } - data-guideid="save-to-knowledge" - > -
- - Save to knowledge -
- - - - onOpenSaveAsViewModal( - { sql, responseId: id }, - { - rephrasedQuestion: questionForSaveAsView, - }, - ) - } - /> -
+ /> + + )} {renderRecommendedQuestions( isLastThreadResponse, recommendedQuestionProps, diff --git a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx index 1cee69660a..199c53399f 100644 --- a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx @@ -56,6 +56,7 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { }); const onPreviewData = async () => { + if (!threadResponse.sql) return; await previewData({ variables: { where: { responseId: id } } }); }; @@ -68,12 +69,17 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { // when is the last step of the last thread response, auto trigger preview data button useEffect(() => { - if (isLastThreadResponse) { + if (isLastThreadResponse && threadResponse.sql) { autoTriggerPreviewDataButton(); } }, [isLastThreadResponse, threadResponse.sql]); const { id, sql } = threadResponse; + const displaySql = + sql || + threadResponse.askingTask?.invalidSql || + threadResponse.adjustmentTask?.invalidSql || + ''; const { hasNativeSQL, dataSourceType } = nativeSQLResult; const showNativeSQL = hasNativeSQL; @@ -84,7 +90,7 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { const sqls = nativeSQLResult.nativeSQLMode && nativeSQLResult.loading === false ? nativeSQLResult.data - : sql; + : displaySql; const onChangeNativeSQL = async (checked: boolean) => { nativeSQLResult.setNativeSQLMode(checked); @@ -182,7 +188,10 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { data-ph-capture-attribute-name="view_sql_copy_sql" icon={} size="small" - onClick={() => onOpenAdjustSQLModal({ sql, responseId: id })} + onClick={() => + onOpenAdjustSQLModal({ sql: displaySql, responseId: id }) + } + disabled={!displaySql} > Adjust SQL @@ -198,6 +207,14 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { />
+ {!sql && ( + + )} - - + +
- )} + data-guideid="save-to-knowledge" + > +
+ + Save to knowledge +
+ + + + onOpenSaveAsViewModal( + { sql, responseId: id }, + { + rephrasedQuestion: questionForSaveAsView, + }, + ) + } + /> + {renderRecommendedQuestions( isLastThreadResponse, recommendedQuestionProps, diff --git a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx index 199c53399f..1cee69660a 100644 --- a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx @@ -56,7 +56,6 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { }); const onPreviewData = async () => { - if (!threadResponse.sql) return; await previewData({ variables: { where: { responseId: id } } }); }; @@ -69,17 +68,12 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { // when is the last step of the last thread response, auto trigger preview data button useEffect(() => { - if (isLastThreadResponse && threadResponse.sql) { + if (isLastThreadResponse) { autoTriggerPreviewDataButton(); } }, [isLastThreadResponse, threadResponse.sql]); const { id, sql } = threadResponse; - const displaySql = - sql || - threadResponse.askingTask?.invalidSql || - threadResponse.adjustmentTask?.invalidSql || - ''; const { hasNativeSQL, dataSourceType } = nativeSQLResult; const showNativeSQL = hasNativeSQL; @@ -90,7 +84,7 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { const sqls = nativeSQLResult.nativeSQLMode && nativeSQLResult.loading === false ? nativeSQLResult.data - : displaySql; + : sql; const onChangeNativeSQL = async (checked: boolean) => { nativeSQLResult.setNativeSQLMode(checked); @@ -188,10 +182,7 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { data-ph-capture-attribute-name="view_sql_copy_sql" icon={} size="small" - onClick={() => - onOpenAdjustSQLModal({ sql: displaySql, responseId: id }) - } - disabled={!displaySql} + onClick={() => onOpenAdjustSQLModal({ sql, responseId: id })} > Adjust SQL @@ -207,14 +198,6 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { />
- {!sql && ( - - )}
setSelectedModels(keys as string[]), - }} - columns={[ - { - title: 'Model name', - render: (_value, model) => - model.displayName || model.referenceName, - }, - ]} - /> - - - - - - )} - - {semanticStep === 'generate' && ( - <> - Generate semantics -

User Prompt

- - Help AI better understand your data by providing a brief - description of your dataset's purpose. Modeling AI - Assistant will use this context to generate more relevant - semantics. - -
- setSemanticPrompt(event.target.value)} - placeholder="Describe what this dataset represents and how it is used." - /> - -
-
- Example prompt -
- This dataset tracks operational records, users, events, and - business entities. It is used to answer analytical questions - about activity, performance, ownership, and trends. + + + + + ← Back to modeling + + + {semanticStep === 'pick' && ( + <> + Pick models + + + Good semantics improve how AI understands and queries your + data. + {' '} + Select models to generate semantics with AI. Modeling AI + Assistant will help you create semantics that improve how AI + understands and queries your data. + +
+ {selectedModels.length}/{diagramData?.models?.length || 0}{' '} + model(s)
-
- - - - - - )} - - {semanticStep === 'review' && ( - <> - Generate semantics -
- setSemanticPrompt(event.target.value)} + setSemanticSearch(event.target.value)} /> - -
-
-
- Generated semantics -
- Review the semantics generated by AI. -
+
setSelectedModels(keys as string[]), + }} + columns={[ + { + title: 'Model name', + render: (_value, model) => + model.displayName || model.referenceName, + }, + ]} + /> + + + + + + )} + + {semanticStep === 'generate' && ( + <> + Generate semantics +

User Prompt

+ + Help AI better understand your data by providing a brief + description of your dataset's purpose. Modeling AI + Assistant will use this context to generate more relevant + semantics. + +
+ + setSemanticPrompt(event.target.value) + } + placeholder="Describe what this dataset represents and how it is used." + /> +
- {semanticResult.map((model) => ( - -
+ )} + + +
+ Following, we provide some example prompts based on some + real world datasets. +
+ - {model.name} - - {(model.columns || []).length} column(s) - + {SEMANTIC_EXAMPLE_PROMPTS.map((example) => ( +
+ {example.label} +
+ {example.text} +
+
+ ))} + + + + {semanticResult.length ? ( +
+
+ Generated semantics +
+ Review the semantics generated by AI. +
-
Description
- - updateSemanticModelDescription( - model.name, - event.target.value, - ) - } - /> -
( - - updateSemanticColumnDescription( - model.name, - column.name, - event.target.value, - ) - } - /> - ), - }, - ]} - /> - - ))} - - - - - - - )} - - - + {semanticResult.map((model) => ( + +
+ {model.name} + + {(model.columns || []).length} column(s) + +
+
Description
+ + updateSemanticModelDescription( + model.name, + event.target.value, + ) + } + /> +
( + + updateSemanticColumnDescription( + model.name, + column.name, + event.target.value, + ) + } + /> + ), + }, + ]} + /> + + ))} + + ) : null} + + + + + + )} + + + + ); } if (assistantMode === 'relationships') { return ( - - - - ← Back to modeling - - - Generate relationships - - Modeling AI Assistant will use AI to discover potential - connections between your models. -
- Review the suggested relationships and adjust them before saving - to your data models. -
- Learn more:{' '} - - Modeling AI Assistant / Generate relationships - - - - {isRelationshipGenerating ? ( - - -
Generating...
- - ) : ( - <> - {assistantError ? ( - - - - ) : !relationshipResult.length ? ( - -
- No relationship suggestions were generated. -
- - ) : null} - - {Object.entries(relationshipGroups).map( - ([modelName, relationships]) => ( - - - {renderIcon(TableOutlined)} - {modelName} - -
- editingRelationshipKey === record.clientId ? ( -
+ editingRelationshipKey === record.clientId ? ( + + editingRelationshipKey === record.clientId ? ( + - updateRelationship(record.clientId, { - type, - }) - } - /> - ) : ( - relationshipTypeLabel(record.type) - ), - }, - { - title: 'Description', - render: (_value, record) => - editingRelationshipKey === record.clientId ? ( - - updateRelationship(record.clientId, { - reason: event.target.value, - }) - } - /> - ) : ( - record.reason - ), - }, - { - title: '', - width: 92, - render: (_value, record) => ( - -
( + + updateSemanticColumnDisplayName( + model.name, + column.name, + event.target.value, + ) + } + /> + ), }, { title: 'Type', dataIndex: 'type', width: 140 }, { From 16dc1b83a9ae2a19d9f9d5123c166f085df0a602 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 15:52:12 +0530 Subject: [PATCH 1019/1087] Make semantic description generation resilient to timeouts --- wren-ai-service/src/config.py | 6 +- .../web/v1/services/semantics_description.py | 28 +++-- .../services/test_semantics_description.py | 102 +++++++++++++++++- .../tools/config/config.example.yaml | 6 +- wren-ai-service/tools/config/config.full.yaml | 6 +- 5 files changed, 126 insertions(+), 22 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 33d2fb0a62..dd70bc01af 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -50,10 +50,10 @@ class Settings(BaseSettings): pipeline_timeout_seconds: float | None = Field(default=None) relationship_recommendation_timeout_seconds: float = Field(default=180.0) semantics_description_timeout_seconds: float | None = Field(default=None) - semantics_description_generation_timeout_seconds: float = Field(default=120.0) + semantics_description_generation_timeout_seconds: float = Field(default=300.0) semantics_description_max_models_per_batch: int = Field(default=4) - semantics_description_max_columns_per_batch: int = Field(default=50) - semantics_description_max_concurrent_tasks: int = Field(default=4) + semantics_description_max_columns_per_batch: int = Field(default=12) + semantics_description_max_concurrent_tasks: int = Field(default=2) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 90dbc6e990..da70e33f6a 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -37,7 +37,7 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, - generation_timeout_seconds: float = 120.0, + generation_timeout_seconds: float = 300.0, max_models_per_batch: int = 1, max_columns_per_batch: int = 50, max_concurrent_tasks: int = 4, @@ -191,16 +191,23 @@ def _split_chunk(self, chunk: dict) -> list[dict]: split_at = max(1, len(columns) // 2) model = chunk["mdl"]["models"][0] + relationships = chunk.get("mdl", {}).get("relationships", []) return [ { **chunk, - "mdl": {"models": [{**model, "columns": column_chunk}]}, + "mdl": { + "models": [{**model, "columns": column_chunk}], + "relationships": relationships, + }, } for column_chunk in (columns[:split_at], columns[split_at:]) if column_chunk ] def _is_retryable_chunk_error(self, error: Exception) -> bool: + if isinstance(error, asyncio.TimeoutError): + return True + if isinstance(error, RetryableSemanticsDescriptionError): return True @@ -214,6 +221,7 @@ def _is_retryable_chunk_error(self, error: Exception) -> bool: "output omitted", "max_tokens", "natural stopping point", + "timed out", ) ) @@ -254,8 +262,13 @@ def _fallback_output_for_chunk(self, chunk: dict) -> dict: async def _generate_task_with_retry_splitting(self, chunk: dict) -> list[dict]: try: - return [await self._generate_task(chunk)] - except ValueError as e: + return [ + await asyncio.wait_for( + self._generate_task(chunk), + timeout=self._generation_timeout_seconds, + ) + ] + except (ValueError, asyncio.TimeoutError) as e: if not self._is_retryable_chunk_error(e): raise @@ -273,7 +286,7 @@ async def _generate_task_with_retry_splitting(self, chunk: dict) -> list[dict]: logger.warning( "Retrying semantics description for model %s with smaller " - "column chunks after incomplete response: %s", + "column chunks after incomplete or timed-out response: %s", model_name, str(e), ) @@ -503,10 +516,7 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: "No selected models matched the current semantic model metadata" ) request_timeout_seconds = self._request_timeout_seconds(len(chunks)) - outputs = await asyncio.wait_for( - self._generate_chunks(chunks), - timeout=request_timeout_seconds, - ) + outputs = await self._generate_chunks(chunks) self[request.id] = self.Resource( id=request.id, diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index b0e8afb9e4..9aa85828df 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -118,7 +118,7 @@ async def test_generate_semantics_description_with_exception( @pytest.mark.asyncio -async def test_generate_semantics_description_with_llm_timeout_fails(): +async def test_generate_semantics_description_with_llm_timeout_preserves_schema(): mock_pipeline = AsyncMock() async def never_returns(**_): @@ -140,9 +140,21 @@ async def never_returns(**_): await service.generate(request) response = service[request.id] - assert response.status == "failed" - assert response.response is None - assert "timed out" in response.error.message + assert response.status == "finished" + assert response.error is None + assert response.response == { + "model1": { + "name": "model1", + "columns": [ + { + "name": "column1", + "type": "varchar", + "properties": {"description": ""}, + } + ], + "properties": {"description": ""}, + } + } def test_get_semantics_description_result( @@ -573,6 +585,88 @@ async def response_for_chunk(**kwargs): assert service._pipelines["semantics_description"].run.call_count == 5 +@pytest.mark.asyncio +async def test_timed_out_chunk_retries_with_smaller_column_groups_and_relationships(): + mock_pipeline = AsyncMock() + service = SemanticsDescription( + pipelines={"semantics_description": mock_pipeline}, + generation_timeout_seconds=0.01, + max_columns_per_batch=2, + ) + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "order_id", "type": "varchar"}, + {"name": "customer_id", "type": "varchar"}, + ], + } + ], + "relationships": [ + { + "name": "OrdersCustomers", + "models": ["orders", "customers"], + "joinType": "MANY_TO_ONE", + "condition": "orders.customer_id = customers.customer_id", + } + ], + } + ).decode(), + ) + observed_relationships = [] + + async def response_for_chunk(**kwargs): + model = kwargs["mdl"]["models"][0] + observed_relationships.append(kwargs["mdl"].get("relationships", [])) + if len(model["columns"]) > 1: + await asyncio.sleep(1) + column = model["columns"][0] + return { + "output": { + "orders": { + "description": "Customer order transactions.", + "columns": [ + { + "name": column["name"], + "properties": { + "description": f"Description for {column['name']}", + }, + } + ], + } + } + } + + mock_pipeline.run.side_effect = response_for_chunk + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert [ + column["properties"]["description"] + for column in response.response["orders"]["columns"] + ] == ["Description for order_id", "Description for customer_id"] + assert all( + relationships == [ + { + "name": "OrdersCustomers", + "models": ["orders", "customers"], + "joinType": "MANY_TO_ONE", + "condition": "orders.customer_id = customers.customer_id", + } + ] + for relationships in observed_relationships + ) + + @pytest.mark.asyncio async def test_truncated_smallest_chunk_preserves_selected_schema( service: SemanticsDescription, diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index eb61e59d42..957b373d90 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -197,9 +197,9 @@ settings: query_cache_maxsize: 1000 query_cache_ttl: 3600 semantics_description_max_models_per_batch: 4 - semantics_description_max_columns_per_batch: 50 - semantics_description_max_concurrent_tasks: 4 - semantics_description_generation_timeout_seconds: 120 + semantics_description_max_columns_per_batch: 12 + semantics_description_max_concurrent_tasks: 2 + semantics_description_generation_timeout_seconds: 300 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: DEBUG diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 07959dfb23..458b54045a 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -194,9 +194,9 @@ settings: sql_generation_timeout_seconds: 30 query_cache_ttl: 3600 semantics_description_max_models_per_batch: 4 - semantics_description_max_columns_per_batch: 50 - semantics_description_max_concurrent_tasks: 4 - semantics_description_generation_timeout_seconds: 120 + semantics_description_max_columns_per_batch: 12 + semantics_description_max_concurrent_tasks: 2 + semantics_description_generation_timeout_seconds: 300 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: INFO From 1e5a54acbcfbe5bd123df2672ae8637718405b54 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 15:53:49 +0530 Subject: [PATCH 1020/1087] Keep semantic timeout handling config driven --- wren-ai-service/src/config.py | 6 +++--- .../src/web/v1/services/semantics_description.py | 2 +- wren-ai-service/tools/config/config.example.yaml | 6 +++--- wren-ai-service/tools/config/config.full.yaml | 6 +++--- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index dd70bc01af..33d2fb0a62 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -50,10 +50,10 @@ class Settings(BaseSettings): pipeline_timeout_seconds: float | None = Field(default=None) relationship_recommendation_timeout_seconds: float = Field(default=180.0) semantics_description_timeout_seconds: float | None = Field(default=None) - semantics_description_generation_timeout_seconds: float = Field(default=300.0) + semantics_description_generation_timeout_seconds: float = Field(default=120.0) semantics_description_max_models_per_batch: int = Field(default=4) - semantics_description_max_columns_per_batch: int = Field(default=12) - semantics_description_max_concurrent_tasks: int = Field(default=2) + semantics_description_max_columns_per_batch: int = Field(default=50) + semantics_description_max_concurrent_tasks: int = Field(default=4) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index da70e33f6a..3ac1c85317 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -37,7 +37,7 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, - generation_timeout_seconds: float = 300.0, + generation_timeout_seconds: float = 120.0, max_models_per_batch: int = 1, max_columns_per_batch: int = 50, max_concurrent_tasks: int = 4, diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index 957b373d90..eb61e59d42 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -197,9 +197,9 @@ settings: query_cache_maxsize: 1000 query_cache_ttl: 3600 semantics_description_max_models_per_batch: 4 - semantics_description_max_columns_per_batch: 12 - semantics_description_max_concurrent_tasks: 2 - semantics_description_generation_timeout_seconds: 300 + semantics_description_max_columns_per_batch: 50 + semantics_description_max_concurrent_tasks: 4 + semantics_description_generation_timeout_seconds: 120 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: DEBUG diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 458b54045a..07959dfb23 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -194,9 +194,9 @@ settings: sql_generation_timeout_seconds: 30 query_cache_ttl: 3600 semantics_description_max_models_per_batch: 4 - semantics_description_max_columns_per_batch: 12 - semantics_description_max_concurrent_tasks: 2 - semantics_description_generation_timeout_seconds: 300 + semantics_description_max_columns_per_batch: 50 + semantics_description_max_concurrent_tasks: 4 + semantics_description_generation_timeout_seconds: 120 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: INFO From 2e3628d1faef471677c3269829096462368bb4c5 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 16:01:27 +0530 Subject: [PATCH 1021/1087] Fill missing semantic aliases and descriptions generically --- .../web/v1/services/semantics_description.py | 192 ++++++++++++++++-- .../services/test_semantics_description.py | 111 +++++----- 2 files changed, 230 insertions(+), 73 deletions(-) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 3ac1c85317..96a2cd4306 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,6 +1,7 @@ import asyncio import logging import math +import re from typing import Dict, Literal, Optional import orjson @@ -367,6 +368,11 @@ def normalized_description(value: str) -> str: generated.setdefault("columns", []) generated["columns"].extend(model_data.get("columns", [])) + relationships = [ + relationship + for relationship in mdl_dict.get("relationships", []) or [] + if isinstance(relationship, dict) + ] response: dict = {} for model in mdl_dict.get("models", []): model_name = model.get("name") @@ -444,6 +450,15 @@ def normalized_description(value: str) -> str: model_name, column_name, ) + column_description = self._fallback_column_description( + model, column, relationships + ) + + column_display_name = ( + generated_display_name + or original_display_name + or self._fallback_display_name(column_name) + ) columns.append( { @@ -451,35 +466,28 @@ def normalized_description(value: str) -> str: "type": column.get("type", ""), "properties": { "description": column_description, - **( - { - "displayName": generated_display_name - or original_display_name - } - if generated_display_name or original_display_name - else {} - ), + "displayName": column_display_name, }, } ) if not model_description: model_description = description(model) + if not model_description: + model_description = self._fallback_model_description( + model, relationships + ) if not model_display_name: model_display_name = display_name(model) - if not model_description and not model.get("columns", []): - logger.warning( - "Semantics description output omitted selected model: %s", - model_name, - ) - continue + if not model_display_name: + model_display_name = self._fallback_display_name(model_name) response[model_name] = { "name": model_name, "columns": columns, "properties": { "description": model_description, - **({"displayName": model_display_name} if model_display_name else {}), + "displayName": model_display_name, }, } @@ -500,6 +508,160 @@ def _metadata_properties(self, payload: dict) -> dict: properties["displayName"] = str(display_name).strip() return properties + def _fallback_display_name(self, name: str) -> str: + words = self._identifier_words(name) + label = " ".join(words) + return label.title() if label else str(name) + + def _fallback_model_description( + self, model: dict, relationships: list[dict] + ) -> str: + model_name = str(model.get("name", "")) + model_label = self._fallback_display_name(model_name) + column_count = len(model.get("columns", []) or []) + related_models = self._related_model_names(model_name, relationships) + relationship_context = ( + f" It can be joined to {', '.join(related_models)} through the configured relationships." + if related_models + else "" + ) + return ( + f"Table {model_name} represents {model_label} records in the selected datasource. " + f"Use it for reporting and analysis across its {column_count} modeled columns." + f"{relationship_context}" + ) + + def _fallback_column_description( + self, model: dict, column: dict, relationships: list[dict] + ) -> str: + model_name = str(model.get("name", "")) + column_name = str(column.get("name", "")) + column_type = str(column.get("type", "") or "unknown") + column_label = self._fallback_display_name(column_name) + role = self._semantic_role(column_name, column_type) + related_models = self._related_model_names( + model_name, relationships, column_name + ) + relationship_context = ( + f" It is used as a join key with {', '.join(related_models)}." + if related_models + else "" + ) + return ( + f"{column_name} is the {column_label} field on table {model_name}. " + f"It is classified as the {role} role with data type {column_type} and is used to filter, group, join, or analyze {model_name} records." + f"{relationship_context}" + ) + + def _semantic_role(self, name: str, column_type: str) -> str: + normalized = " ".join(self._identifier_words(name)).casefold() + column_type = column_type.casefold() + + role_markers = [ + (("id", "identifier", "key", "no", "number"), "identifier or key"), + ( + ("date", "time", "timestamp", "year", "month", "period"), + "date or time dimension", + ), + (("status", "state", "stage"), "status dimension"), + (("currency", "curr", "cur", "code"), "code or currency dimension"), + (("qty", "quantity", "count", "volume", "units"), "quantity measure"), + (("cost", "expense"), "cost measure"), + ( + ("sales", "revenue", "value", "amount", "price", "gm", "margin"), + "financial measure", + ), + ( + ("rate", "ratio", "percent", "percentage", "pct"), + "rate or percentage measure", + ), + ( + ( + "type", + "category", + "segment", + "market", + "division", + "company", + "country", + ), + "business dimension", + ), + (("name", "description", "desc"), "descriptive attribute"), + ] + for markers, role in role_markers: + if any(marker in normalized.split() for marker in markers): + return role + + if any( + token in column_type + for token in ("int", "float", "double", "decimal", "numeric", "number") + ): + return "numeric measure" + if any(token in column_type for token in ("date", "time")): + return "date or time dimension" + return "business attribute" + + def _identifier_words(self, value: str) -> list[str]: + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value or "")) + text = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1 \2", text) + raw_words = re.split(r"[^A-Za-z0-9]+", text) + expansions = { + "bu": "business unit", + "cust": "customer", + "curr": "currency", + "cur": "currency", + "desc": "description", + "fx": "foreign exchange", + "gm": "gross margin", + "inv": "invoice", + "no": "number", + "ord": "order", + "po": "purchase order", + "prod": "product", + "qty": "quantity", + "std": "standard", + } + words: list[str] = [] + for raw_word in raw_words: + if not raw_word: + continue + word = raw_word.casefold() + words.extend(expansions.get(word, word).split()) + return words + + def _related_model_names( + self, + model_name: str, + relationships: list[dict], + column_name: Optional[str] = None, + ) -> list[str]: + related = [] + target_model = str(model_name) + target_column = str(column_name) if column_name else None + + for relationship in relationships: + models = [str(model) for model in relationship.get("models", []) or []] + condition = str(relationship.get("condition", "")) + if target_model not in models: + continue + if target_column and not self._relationship_mentions_column( + condition, target_model, target_column + ): + continue + related.extend(model for model in models if model != target_model) + + return sorted(set(related)) + + def _relationship_mentions_column( + self, condition: str, model_name: str, column_name: str + ) -> bool: + pattern = ( + rf'["`]?{re.escape(model_name)}["`]?\s*\.\s*' + rf'["`]?{re.escape(column_name)}["`]?' + ) + return bool(re.search(pattern, condition)) + @observe(name="Generate Semantics Description") @trace_metadata async def generate(self, request: GenerateRequest, **kwargs) -> Resource: diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 9aa85828df..c85b8c1b47 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -56,11 +56,12 @@ async def test_generate_semantics_description( "name": "column1", "type": "varchar", "properties": { - "description": "Customer segment for reporting." + "description": "Customer segment for reporting.", + "displayName": "Column1", }, } ], - "properties": {"description": "Test description"}, + "properties": {"description": "Test description", "displayName": "Model1"}, } } assert response.error is None @@ -142,19 +143,17 @@ async def never_returns(**_): assert response.status == "finished" assert response.error is None - assert response.response == { - "model1": { - "name": "model1", - "columns": [ - { - "name": "column1", - "type": "varchar", - "properties": {"description": ""}, - } - ], - "properties": {"description": ""}, - } - } + model = response.response["model1"] + column = model["columns"][0] + assert model["name"] == "model1" + assert model["properties"]["description"] + assert model["properties"]["displayName"] + assert column["name"] == "column1" + assert column["type"] == "varchar" + assert column["properties"]["description"] + assert column["properties"]["displayName"] + assert "model1" in model["properties"]["description"] + assert "column1" in column["properties"]["description"] def test_get_semantics_description_result( @@ -711,14 +710,20 @@ async def truncated_response(**kwargs): assert response.response["orders"] == { "name": "orders", "columns": [ - { - "name": "order_id", - "type": "varchar", - "properties": {"description": "Existing order id."}, - } - ], - "properties": {"description": "Existing order model."}, - } + { + "name": "order_id", + "type": "varchar", + "properties": { + "description": "Existing order id.", + "displayName": "Order Id", + }, + } + ], + "properties": { + "description": "Existing order model.", + "displayName": "Orders", + }, + } assert service._pipelines["semantics_description"].run.call_count == 1 @@ -775,29 +780,25 @@ async def test_incomplete_llm_output_uses_available_descriptions( assert response.response["orders"]["properties"]["description"] == ( "Customer purchase transactions." ) - assert response.response["orders"]["columns"] == [ - { - "name": "order_id", - "type": "varchar", - "properties": {"description": "Unique order identifier."}, - }, - { - "name": "order_date", - "type": "date", - "properties": {"description": ""}, - }, - ] - assert response.response["customers"] == { - "name": "customers", - "columns": [ - { - "name": "customer_id", - "type": "varchar", - "properties": {"description": ""}, - } - ], - "properties": {"description": ""}, - } + order_columns = response.response["orders"]["columns"] + assert order_columns[0]["name"] == "order_id" + assert order_columns[0]["properties"]["description"] == ( + "Unique order identifier." + ) + assert order_columns[0]["properties"]["displayName"] + assert order_columns[1]["name"] == "order_date" + assert order_columns[1]["properties"]["description"] + assert order_columns[1]["properties"]["displayName"] + assert "order_date" in order_columns[1]["properties"]["description"] + + customers = response.response["customers"] + assert customers["name"] == "customers" + assert customers["properties"]["description"] + assert customers["properties"]["displayName"] + assert customers["columns"][0]["name"] == "customer_id" + assert customers["columns"][0]["properties"]["description"] + assert customers["columns"][0]["properties"]["displayName"] + assert "customer_id" in customers["columns"][0]["properties"]["description"] @pytest.mark.asyncio @@ -987,15 +988,9 @@ async def test_repeated_llm_column_descriptions_are_tolerated( assert response.response["orders"]["properties"]["description"] == ( "Customer order transactions." ) - assert response.response["orders"]["columns"] == [ - { - "name": "order_id", - "type": "varchar", - "properties": {"description": "Identifier for reporting."}, - }, - { - "name": "customer_id", - "type": "varchar", - "properties": {"description": "Identifier for reporting."}, - }, - ] + columns = response.response["orders"]["columns"] + assert [column["name"] for column in columns] == ["order_id", "customer_id"] + assert [ + column["properties"]["description"] for column in columns + ] == ["Identifier for reporting.", "Identifier for reporting."] + assert all(column["properties"]["displayName"] for column in columns) From 4e9319e40cdb004e5636bd2bc56102a58791377c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 16:17:46 +0530 Subject: [PATCH 1022/1087] Ensure generated semantic aliases are populated --- .../generation/semantics_description.py | 6 +- wren-ui/src/pages/modeling.tsx | 87 +++++++++++++------ 2 files changed, 65 insertions(+), 28 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index f19d109af0..b57f4dbae7 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -26,7 +26,7 @@ 2. Return every input model exactly once and every input column exactly once. 3. Preserve every model and column `name` exactly as provided. 4. Put each generated description in `properties.description`. -5. Put natural-language aliases and synonyms in `properties.displayName` as a short comma-separated phrase. Do not put SQL identifiers there. +5. Put natural-language aliases and synonyms in `properties.displayName` as a required, non-empty, short comma-separated phrase. Do not put SQL identifiers there unless the identifier itself is the natural user-facing name. 6. Make descriptions business-friendly, factual, and useful for text-to-SQL retrieval. 7. Include business context, common analytical use, and the field role when it is supported by the name, type, or relationships: ID/key, date/time, measure, dimension, status, currency, quantity, cost, revenue, rate, percentage, or code. 8. Ground descriptions and aliases only in the user prompt, model and column names, aliases, data types, existing descriptions, and provided schema/relationship context. @@ -48,7 +48,7 @@ Write semantic descriptions for every picked model and every column. For each model, describe the real-world records represented and the analytical questions it can support. For each column, describe the business meaning and analytical use of that exact field. -For each model and column, generate aliases/synonyms that users may naturally type in questions and place them in properties.displayName. +For each model and column, generate non-empty aliases/synonyms that users may naturally type in questions and place them in properties.displayName. If an existing description is already meaningful, preserve its business meaning while making it clearer and more useful for retrieval. Keep every description and alias grounded in the picked model metadata, user prompt, data types, and relationship context. """ @@ -208,7 +208,7 @@ def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: ## End of Pipeline class ModelProperties(BaseModel): description: str - displayName: str = "" + displayName: str class ModelColumns(BaseModel): diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 6536923c12..e51c662d0e 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -110,6 +110,25 @@ const ForwardDiagram = forwardRef(function ForwardDiagram(props: any, ref) { return ; }); +const semanticText = (value: any): string => { + if (Array.isArray(value)) { + return value.map(semanticText).filter(Boolean).join(', '); + } + return typeof value === 'string' ? value.trim() : ''; +}; + +const firstSemanticText = (...values: any[]): string => + values.map(semanticText).find(Boolean) || ''; + +const semanticFallbackDisplayName = (name: string): string => { + return String(name || '') + .replace(/([a-z0-9])([A-Z])/g, '$1 $2') + .replace(/([A-Z]+)([A-Z][a-z])/g, '$1 $2') + .replace(/[_\-.]+/g, ' ') + .replace(/\s+/g, ' ') + .trim(); +}; + const DiagramWrapper = styled.div` position: relative; height: 100%; @@ -683,34 +702,46 @@ export default function Modeling() { ASSISTANT_INITIAL_POLL_INTERVAL_MS * Math.max(1, attempt + 1), ASSISTANT_MAX_POLL_INTERVAL_MS, ); - await new Promise((resolve) => - setTimeout(resolve, pollInterval), - ); + await new Promise((resolve) => setTimeout(resolve, pollInterval)); } throw new Error('AI assistant timed out.'); }; - const normalizeSemanticModel = (name: string, value: any): any => ({ - name: value?.name || name, - displayName: - value?.displayName || - value?.alias || - value?.properties?.displayName || - value?.properties?.alias || - '', - description: value?.description || value?.properties?.description || '', - columns: (value?.columns || []).map((column) => ({ - name: column?.name, - type: column?.type, + const normalizeSemanticModel = (name: string, value: any): any => { + const modelName = semanticText(value?.name) || semanticText(name); + return { + name: modelName, displayName: - column?.displayName || - column?.alias || - column?.properties?.displayName || - column?.properties?.alias || - '', - description: column?.description || column?.properties?.description || '', - })), - }); + firstSemanticText( + value?.displayName, + value?.alias, + value?.properties?.displayName, + value?.properties?.alias, + ) || semanticFallbackDisplayName(modelName), + description: firstSemanticText( + value?.description, + value?.properties?.description, + ), + columns: (value?.columns || []).map((column) => { + const columnName = semanticText(column?.name); + return { + name: columnName, + type: column?.type, + displayName: + firstSemanticText( + column?.displayName, + column?.alias, + column?.properties?.displayName, + column?.properties?.alias, + ) || semanticFallbackDisplayName(columnName), + description: firstSemanticText( + column?.description, + column?.properties?.description, + ), + }; + }), + }; + }; const normalizeSemanticResult = (result: any): any[] => { if (Array.isArray(result)) { @@ -1005,7 +1036,9 @@ export default function Modeling() { if (!diagramModel) return []; return { modelId: diagramModel.modelId, - displayName: model.displayName, + displayName: + firstSemanticText(model.displayName) || + semanticFallbackDisplayName(model.name), description: model.description, columns: (model.columns || []) .map((column) => { @@ -1015,7 +1048,11 @@ export default function Modeling() { return field ? { id: field.columnId, - displayName: column.displayName || field.displayName, + displayName: + firstSemanticText( + column.displayName, + field.displayName, + ) || semanticFallbackDisplayName(column.name), description: column.description, } : null; From 8ebcbb4a9f386f8801cc0a1fab569b32cdcdf29e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 16:30:56 +0530 Subject: [PATCH 1023/1087] Require complete generated semantic metadata --- .../generation/semantics_description.py | 15 +- .../web/v1/services/semantics_description.py | 306 ++++-------------- .../services/test_semantics_description.py | 145 ++++----- wren-ui/src/apollo/server/models/model.ts | 2 + .../apollo/server/resolvers/modelResolver.ts | 54 +++- wren-ui/src/apollo/server/schema.ts | 2 + wren-ui/src/pages/modeling.tsx | 47 +-- 7 files changed, 209 insertions(+), 362 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index b57f4dbae7..33a2d307e9 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -7,7 +7,7 @@ from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider @@ -35,7 +35,8 @@ 11. If two columns have similar names or business meaning, explain the distinction using the exact column name, alias, type, or surrounding model context. 12. Do not invent unsupported tables, columns, relationships, metrics, or business concepts. 13. Do not use generic boilerplate or copy the technical name as the whole description. -14. Return complete JSON only. Do not include markdown, comments, examples, or explanatory text outside the JSON object. +14. For each input column, return a matching column object with the exact same `name`, a non-empty `properties.description`, and a non-empty `properties.displayName`. +15. Return complete JSON only. Do not include markdown, comments, examples, or explanatory text outside the JSON object. """ user_prompt_template = """ @@ -51,6 +52,7 @@ For each model and column, generate non-empty aliases/synonyms that users may naturally type in questions and place them in properties.displayName. If an existing description is already meaningful, preserve its business meaning while making it clearer and more useful for retrieval. Keep every description and alias grounded in the picked model metadata, user prompt, data types, and relationship context. +The number of output columns for each model must exactly match the number of input columns for that model. """ @@ -207,22 +209,30 @@ def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: ## End of Pipeline class ModelProperties(BaseModel): + model_config = ConfigDict(extra="forbid") + description: str displayName: str class ModelColumns(BaseModel): + model_config = ConfigDict(extra="forbid") + name: str properties: ModelProperties class SemanticModel(BaseModel): + model_config = ConfigDict(extra="forbid") + name: str columns: list[ModelColumns] properties: ModelProperties class SemanticResult(BaseModel): + model_config = ConfigDict(extra="forbid") + models: list[SemanticModel] @@ -231,6 +241,7 @@ class SemanticResult(BaseModel): "type": "json_schema", "json_schema": { "name": "semantic_description", + "strict": True, "schema": SemanticResult.model_json_schema(), }, } diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 96a2cd4306..26c9083a8c 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,7 +1,6 @@ import asyncio import logging import math -import re from typing import Dict, Literal, Optional import orjson @@ -172,11 +171,68 @@ def _description(self, payload: dict) -> str: ) return "" if value is None else str(value).strip() + def _display_name(self, payload: dict) -> str: + payload_properties = self._properties(payload) + value = ( + payload.get("displayName") + or payload.get("alias") + or payload_properties.get("displayName") + or payload_properties.get("alias") + or "" + ) + return "" if value is None else str(value).strip() + + def _validate_generated_output(self, chunk: dict, output: dict) -> None: + missing: list[str] = [] + + for model in chunk.get("mdl", {}).get("models", []) or []: + if not isinstance(model, dict): + continue + + model_name = model.get("name", "") + generated_model = output.get(model_name) + if not isinstance(generated_model, dict): + missing.append(f"{model_name} model") + continue + + if not self._description(generated_model): + missing.append(f"{model_name} description") + if not self._display_name(generated_model): + missing.append(f"{model_name} alias") + + generated_columns = { + column.get("name"): column + for column in generated_model.get("columns", []) or [] + if isinstance(column, dict) and column.get("name") + } + for column in model.get("columns", []) or []: + if not isinstance(column, dict): + continue + + column_name = column.get("name", "") + generated_column = generated_columns.get(column_name) + if not isinstance(generated_column, dict): + missing.append(f"{model_name}.{column_name} column") + continue + if not self._description(generated_column): + missing.append(f"{model_name}.{column_name} description") + if not self._display_name(generated_column): + missing.append(f"{model_name}.{column_name} alias") + + if missing: + preview = ", ".join(missing[:10]) + suffix = "..." if len(missing) > 10 else "" + raise RetryableSemanticsDescriptionError( + "Semantics description output omitted required metadata: " + f"{preview}{suffix}" + ) + async def _generate_task(self, chunk: dict) -> dict: resp = await self._pipelines["semantics_description"].run(**chunk) output = resp.get("output") or {} if not isinstance(output, dict): raise ValueError("Semantics description pipeline returned invalid output") + self._validate_generated_output(chunk, output) return output def _chunk_columns(self, chunk: dict) -> list[dict]: @@ -226,41 +282,6 @@ def _is_retryable_chunk_error(self, error: Exception) -> bool: ) ) - def _fallback_output_for_chunk(self, chunk: dict) -> dict: - output = {} - for model in chunk.get("mdl", {}).get("models", []): - if not isinstance(model, dict): - continue - - model_name = model.get("name") - if not model_name: - continue - - columns = [] - for column in model.get("columns", []): - if not isinstance(column, dict): - continue - - column_name = column.get("name") - if not column_name: - continue - - columns.append( - { - "name": column_name, - "type": column.get("type", ""), - "properties": self._metadata_properties(column), - } - ) - - output[model_name] = { - "name": model_name, - "columns": columns, - "properties": self._metadata_properties(model), - } - - return output - async def _generate_task_with_retry_splitting(self, chunk: dict) -> list[dict]: try: return [ @@ -276,14 +297,7 @@ async def _generate_task_with_retry_splitting(self, chunk: dict) -> list[dict]: split_chunks = self._split_chunk(chunk) model_name = (chunk.get("selected_models") or [""])[0] if not split_chunks: - logger.warning( - "Preserving selected semantics schema for model %s after " - "LLM description generation failed on the smallest retry " - "chunk: %s", - model_name, - str(e), - ) - return [self._fallback_output_for_chunk(chunk)] + raise logger.warning( "Retrying semantics description for model %s with smaller " @@ -328,15 +342,7 @@ def description(payload: dict) -> str: return "" if value is None else str(value).strip() def display_name(payload: dict) -> str: - payload_properties = properties(payload) - value = ( - payload.get("displayName") - or payload.get("alias") - or payload_properties.get("displayName") - or payload_properties.get("alias") - or "" - ) - return "" if value is None else str(value).strip() + return self._display_name(payload) def normalized_description(value: str) -> str: return " ".join(value.casefold().split()) @@ -368,11 +374,6 @@ def normalized_description(value: str) -> str: generated.setdefault("columns", []) generated["columns"].extend(model_data.get("columns", [])) - relationships = [ - relationship - for relationship in mdl_dict.get("relationships", []) or [] - if isinstance(relationship, dict) - ] response: dict = {} for model in mdl_dict.get("models", []): model_name = model.get("name") @@ -444,20 +445,8 @@ def normalized_description(value: str) -> str: normalized_generated_description, column_name ) - if not column_description: - logger.warning( - "Semantics description output omitted description for column: %s.%s", - model_name, - column_name, - ) - column_description = self._fallback_column_description( - model, column, relationships - ) - column_display_name = ( - generated_display_name - or original_display_name - or self._fallback_display_name(column_name) + generated_display_name or original_display_name ) columns.append( @@ -473,14 +462,8 @@ def normalized_description(value: str) -> str: if not model_description: model_description = description(model) - if not model_description: - model_description = self._fallback_model_description( - model, relationships - ) if not model_display_name: model_display_name = display_name(model) - if not model_display_name: - model_display_name = self._fallback_display_name(model_name) response[model_name] = { "name": model_name, @@ -493,175 +476,6 @@ def normalized_description(value: str) -> str: return response - def _metadata_properties(self, payload: dict) -> dict: - properties = { - "description": self._description(payload), - } - display_name = ( - payload.get("displayName") - or payload.get("alias") - or self._properties(payload).get("displayName") - or self._properties(payload).get("alias") - or "" - ) - if display_name: - properties["displayName"] = str(display_name).strip() - return properties - - def _fallback_display_name(self, name: str) -> str: - words = self._identifier_words(name) - label = " ".join(words) - return label.title() if label else str(name) - - def _fallback_model_description( - self, model: dict, relationships: list[dict] - ) -> str: - model_name = str(model.get("name", "")) - model_label = self._fallback_display_name(model_name) - column_count = len(model.get("columns", []) or []) - related_models = self._related_model_names(model_name, relationships) - relationship_context = ( - f" It can be joined to {', '.join(related_models)} through the configured relationships." - if related_models - else "" - ) - return ( - f"Table {model_name} represents {model_label} records in the selected datasource. " - f"Use it for reporting and analysis across its {column_count} modeled columns." - f"{relationship_context}" - ) - - def _fallback_column_description( - self, model: dict, column: dict, relationships: list[dict] - ) -> str: - model_name = str(model.get("name", "")) - column_name = str(column.get("name", "")) - column_type = str(column.get("type", "") or "unknown") - column_label = self._fallback_display_name(column_name) - role = self._semantic_role(column_name, column_type) - related_models = self._related_model_names( - model_name, relationships, column_name - ) - relationship_context = ( - f" It is used as a join key with {', '.join(related_models)}." - if related_models - else "" - ) - return ( - f"{column_name} is the {column_label} field on table {model_name}. " - f"It is classified as the {role} role with data type {column_type} and is used to filter, group, join, or analyze {model_name} records." - f"{relationship_context}" - ) - - def _semantic_role(self, name: str, column_type: str) -> str: - normalized = " ".join(self._identifier_words(name)).casefold() - column_type = column_type.casefold() - - role_markers = [ - (("id", "identifier", "key", "no", "number"), "identifier or key"), - ( - ("date", "time", "timestamp", "year", "month", "period"), - "date or time dimension", - ), - (("status", "state", "stage"), "status dimension"), - (("currency", "curr", "cur", "code"), "code or currency dimension"), - (("qty", "quantity", "count", "volume", "units"), "quantity measure"), - (("cost", "expense"), "cost measure"), - ( - ("sales", "revenue", "value", "amount", "price", "gm", "margin"), - "financial measure", - ), - ( - ("rate", "ratio", "percent", "percentage", "pct"), - "rate or percentage measure", - ), - ( - ( - "type", - "category", - "segment", - "market", - "division", - "company", - "country", - ), - "business dimension", - ), - (("name", "description", "desc"), "descriptive attribute"), - ] - for markers, role in role_markers: - if any(marker in normalized.split() for marker in markers): - return role - - if any( - token in column_type - for token in ("int", "float", "double", "decimal", "numeric", "number") - ): - return "numeric measure" - if any(token in column_type for token in ("date", "time")): - return "date or time dimension" - return "business attribute" - - def _identifier_words(self, value: str) -> list[str]: - text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value or "")) - text = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1 \2", text) - raw_words = re.split(r"[^A-Za-z0-9]+", text) - expansions = { - "bu": "business unit", - "cust": "customer", - "curr": "currency", - "cur": "currency", - "desc": "description", - "fx": "foreign exchange", - "gm": "gross margin", - "inv": "invoice", - "no": "number", - "ord": "order", - "po": "purchase order", - "prod": "product", - "qty": "quantity", - "std": "standard", - } - words: list[str] = [] - for raw_word in raw_words: - if not raw_word: - continue - word = raw_word.casefold() - words.extend(expansions.get(word, word).split()) - return words - - def _related_model_names( - self, - model_name: str, - relationships: list[dict], - column_name: Optional[str] = None, - ) -> list[str]: - related = [] - target_model = str(model_name) - target_column = str(column_name) if column_name else None - - for relationship in relationships: - models = [str(model) for model in relationship.get("models", []) or []] - condition = str(relationship.get("condition", "")) - if target_model not in models: - continue - if target_column and not self._relationship_mentions_column( - condition, target_model, target_column - ): - continue - related.extend(model for model in models if model != target_model) - - return sorted(set(related)) - - def _relationship_mentions_column( - self, condition: str, model_name: str, column_name: str - ) -> bool: - pattern = ( - rf'["`]?{re.escape(model_name)}["`]?\s*\.\s*' - rf'["`]?{re.escape(column_name)}["`]?' - ) - return bool(re.search(pattern, condition)) - @observe(name="Generate Semantics Description") @trace_metadata async def generate(self, request: GenerateRequest, **kwargs) -> Resource: diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index c85b8c1b47..85327f1136 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -18,11 +18,15 @@ def service(): { "name": "column1", "properties": { - "description": "Customer segment for reporting." + "description": "Customer segment for reporting.", + "displayName": "customer segment, segment", }, } ], - "properties": {"description": "Test description"}, + "properties": { + "description": "Test description", + "displayName": "test model, model one", + }, } } } @@ -57,11 +61,14 @@ async def test_generate_semantics_description( "type": "varchar", "properties": { "description": "Customer segment for reporting.", - "displayName": "Column1", + "displayName": "customer segment, segment", }, } ], - "properties": {"description": "Test description", "displayName": "Model1"}, + "properties": { + "description": "Test description", + "displayName": "test model, model one", + }, } } assert response.error is None @@ -119,7 +126,7 @@ async def test_generate_semantics_description_with_exception( @pytest.mark.asyncio -async def test_generate_semantics_description_with_llm_timeout_preserves_schema(): +async def test_generate_semantics_description_with_llm_timeout_fails_without_fallback(): mock_pipeline = AsyncMock() async def never_returns(**_): @@ -141,19 +148,10 @@ async def never_returns(**_): await service.generate(request) response = service[request.id] - assert response.status == "finished" - assert response.error is None - model = response.response["model1"] - column = model["columns"][0] - assert model["name"] == "model1" - assert model["properties"]["description"] - assert model["properties"]["displayName"] - assert column["name"] == "column1" - assert column["type"] == "varchar" - assert column["properties"]["description"] - assert column["properties"]["displayName"] - assert "model1" in model["properties"]["description"] - assert "column1" in column["properties"]["description"] + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert "timed out" in response.error.message def test_get_semantics_description_result( @@ -243,28 +241,40 @@ async def test_batch_processing_with_multiple_models( "output": { "model1": { "description": "Description 1", + "displayName": "model one, first model", "columns": [ { "name": "column1", - "properties": {"description": "Column description 1"}, + "properties": { + "description": "Column description 1", + "displayName": "column one, first column", + }, } ], }, "model2": { "description": "Description 2", + "displayName": "model two, second model", "columns": [ { "name": "column1", - "properties": {"description": "Column description 2"}, + "properties": { + "description": "Column description 2", + "displayName": "column one, first column", + }, } ], }, "model3": { "description": "Description 3", + "displayName": "model three, third model", "columns": [ { "name": "column1", - "properties": {"description": "Column description 3"}, + "properties": { + "description": "Column description 3", + "displayName": "column one, first column", + }, } ], }, @@ -438,11 +448,13 @@ def response_for_chunk(**kwargs): "output": { "orders": { "description": "Customer order transactions.", + "displayName": "orders, customer orders", "columns": [ { "name": column["name"], "properties": { "description": f"Description for {column['name']}", + "displayName": f"Alias for {column['name']}", }, } for column in model["columns"] @@ -562,11 +574,13 @@ async def response_for_chunk(**kwargs): "output": { "orders": { "description": "Customer order transactions.", + "displayName": "orders, customer orders", "columns": [ { "name": column["name"], "properties": { "description": f"Description for {column['name']}", + "displayName": f"Alias for {column['name']}", }, } ], @@ -631,11 +645,13 @@ async def response_for_chunk(**kwargs): "output": { "orders": { "description": "Customer order transactions.", + "displayName": "orders, customer orders", "columns": [ { "name": column["name"], "properties": { "description": f"Description for {column['name']}", + "displayName": f"Alias for {column['name']}", }, } ], @@ -667,7 +683,7 @@ async def response_for_chunk(**kwargs): @pytest.mark.asyncio -async def test_truncated_smallest_chunk_preserves_selected_schema( +async def test_truncated_smallest_chunk_fails_without_fallback( service: SemanticsDescription, ): service["test_id"] = SemanticsDescription.Resource(id="test_id") @@ -705,30 +721,15 @@ async def truncated_response(**kwargs): await service.generate(request) response = service[request.id] - assert response.status == "finished" - assert response.error is None - assert response.response["orders"] == { - "name": "orders", - "columns": [ - { - "name": "order_id", - "type": "varchar", - "properties": { - "description": "Existing order id.", - "displayName": "Order Id", - }, - } - ], - "properties": { - "description": "Existing order model.", - "displayName": "Orders", - }, - } + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert "truncated" in response.error.message assert service._pipelines["semantics_description"].run.call_count == 1 @pytest.mark.asyncio -async def test_incomplete_llm_output_uses_available_descriptions( +async def test_incomplete_llm_output_fails_without_fallback( service: SemanticsDescription, ): service["test_id"] = SemanticsDescription.Resource(id="test_id") @@ -774,31 +775,10 @@ async def test_incomplete_llm_output_uses_available_descriptions( await service.generate(request) response = service[request.id] - assert response.status == "finished" - assert response.error is None - assert list(response.response.keys()) == ["orders", "customers"] - assert response.response["orders"]["properties"]["description"] == ( - "Customer purchase transactions." - ) - order_columns = response.response["orders"]["columns"] - assert order_columns[0]["name"] == "order_id" - assert order_columns[0]["properties"]["description"] == ( - "Unique order identifier." - ) - assert order_columns[0]["properties"]["displayName"] - assert order_columns[1]["name"] == "order_date" - assert order_columns[1]["properties"]["description"] - assert order_columns[1]["properties"]["displayName"] - assert "order_date" in order_columns[1]["properties"]["description"] - - customers = response.response["customers"] - assert customers["name"] == "customers" - assert customers["properties"]["description"] - assert customers["properties"]["displayName"] - assert customers["columns"][0]["name"] == "customer_id" - assert customers["columns"][0]["properties"]["description"] - assert customers["columns"][0]["properties"]["displayName"] - assert "customer_id" in customers["columns"][0]["properties"]["description"] + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert "omitted required metadata" in response.error.message @pytest.mark.asyncio @@ -833,22 +813,28 @@ async def test_llm_descriptions_are_not_rewritten_by_service( { "name": "Division", "properties": { - "description": "Stores the Division value used to describe or analyze xStage records." + "description": "Stores the Division value used to describe or analyze xStage records.", + "displayName": "division, business division", }, }, { "name": "SalesPerson", - "properties": {"description": "SalesPerson"}, + "properties": { + "description": "SalesPerson", + "displayName": "sales person, account owner", + }, }, { "name": "SalesAmount", "properties": { - "description": "Stores the SalesAmount value." + "description": "Stores the SalesAmount value.", + "displayName": "sales amount, revenue amount", }, }, ], "properties": { - "description": "Contains business records for xStageLoad2." + "description": "Contains business records for xStageLoad2.", + "displayName": "stage load, staging records", }, } } @@ -914,10 +900,14 @@ async def test_concurrent_updates_no_race_condition( "output": { f"model{i}": { "description": f"Description {i}", + "displayName": f"model {i}, model {i} alias", "columns": [ { "name": "column1", - "properties": {"description": f"Column description {i}"}, + "properties": { + "description": f"Column description {i}", + "displayName": f"column one, column {i}", + }, } ], } @@ -966,14 +956,21 @@ async def test_repeated_llm_column_descriptions_are_tolerated( "output": { "orders": { "description": "Customer order transactions.", + "displayName": "orders, customer orders", "columns": [ { "name": "order_id", - "properties": {"description": "Identifier for reporting."}, + "properties": { + "description": "Identifier for reporting.", + "displayName": "order id, order identifier", + }, }, { "name": "customer_id", - "properties": {"description": "Identifier for reporting."}, + "properties": { + "description": "Identifier for reporting.", + "displayName": "customer id, customer identifier", + }, }, ], } diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index a88a493f01..1cec4870bb 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -17,6 +17,7 @@ export interface NestedColumnMetadataInput { export interface ColumnMetadataInput { id: number; + referenceName?: string; displayName?: string; description?: string; } @@ -47,6 +48,7 @@ export interface UpdateModelMetadataInput { export interface SaveModelingSemanticInput { modelId: number; + referenceName?: string; displayName?: string; description?: string; columns: Array; diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 0ee7a46605..ce2ea2377e 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -493,11 +493,24 @@ export class ModelResolver { projectId: project.id, }); const modelById = new Map(models.map((model) => [model.id, model])); + const modelByReferenceName = new Map( + models.map((model) => [model.referenceName, model]), + ); const requestedItems = args.data || []; + const resolveModel = (item: SaveModelingSemanticInput) => { + const model = modelById.get(item.modelId); + if (model) return model; + return item.referenceName + ? modelByReferenceName.get(item.referenceName) + : undefined; + }; + for (const item of requestedItems) { - if (!modelById.has(item.modelId)) { - throw new Error(`Model not found: ${item.modelId}`); + if (!resolveModel(item)) { + throw new Error( + `Model not found: ${item.referenceName || item.modelId}`, + ); } } @@ -505,7 +518,7 @@ export class ModelResolver { requestedItems.map(async (item) => { if (isNil(item.description) && isNil(item.displayName)) return; - const model = modelById.get(item.modelId); + const model = resolveModel(item); const modelMetadata: Partial = {}; if (!isNil(item.displayName)) { @@ -525,24 +538,47 @@ export class ModelResolver { } if (!isEmpty(modelMetadata)) { - await ctx.modelRepository.updateOne(item.modelId, modelMetadata); + await ctx.modelRepository.updateOne(model.id, modelMetadata); } }), ); - const requestedColumns = requestedItems.flatMap( - (item) => item.columns || [], - ); + const requestedColumns = requestedItems.flatMap((item) => { + const model = resolveModel(item); + return (item.columns || []).map((column) => ({ + ...column, + modelId: model?.id, + })); + }); if (!isEmpty(requestedColumns)) { const columnIds = requestedColumns.map((column) => column.id); - const columns = await ctx.modelColumnRepository.findColumnsByIds(columnIds); + const modelIds = models.map((model) => model.id); + const [columnsByIdSource, columnsByModelSource] = await Promise.all([ + ctx.modelColumnRepository.findColumnsByIds(columnIds), + ctx.modelColumnRepository.findColumnsByModelIds(modelIds), + ]); + const columns = [...columnsByIdSource, ...columnsByModelSource]; const columnById = new Map( columns.map((column) => [String(column.id), column]), ); + const columnByModelAndReferenceName = new Map( + columns.map((column) => [ + `${column.modelId}:${column.referenceName}`, + column, + ]), + ); await Promise.all( requestedColumns.map(async (requestedColumn) => { - const column = columnById.get(String(requestedColumn.id)); + const columnByRequestedId = columnById.get( + String(requestedColumn.id), + ); + const column = + columnByRequestedId?.modelId === requestedColumn.modelId + ? columnByRequestedId + : columnByModelAndReferenceName.get( + `${requestedColumn.modelId}:${requestedColumn.referenceName}`, + ); if (!column) return; const columnMetadata: Partial = {}; diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 339acaa1d1..4d79bb4237 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -345,6 +345,7 @@ export const typeDefs = gql` input UpdateColumnMetadataInput { id: Int! + referenceName: String displayName: String description: String } @@ -386,6 +387,7 @@ export const typeDefs = gql` input SaveModelingSemanticInput { modelId: Int! + referenceName: String displayName: String description: String columns: [UpdateColumnMetadataInput!]! diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index e51c662d0e..112dc4b91c 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -120,15 +120,6 @@ const semanticText = (value: any): string => { const firstSemanticText = (...values: any[]): string => values.map(semanticText).find(Boolean) || ''; -const semanticFallbackDisplayName = (name: string): string => { - return String(name || '') - .replace(/([a-z0-9])([A-Z])/g, '$1 $2') - .replace(/([A-Z]+)([A-Z][a-z])/g, '$1 $2') - .replace(/[_\-.]+/g, ' ') - .replace(/\s+/g, ' ') - .trim(); -}; - const DiagramWrapper = styled.div` position: relative; height: 100%; @@ -711,13 +702,12 @@ export default function Modeling() { const modelName = semanticText(value?.name) || semanticText(name); return { name: modelName, - displayName: - firstSemanticText( - value?.displayName, - value?.alias, - value?.properties?.displayName, - value?.properties?.alias, - ) || semanticFallbackDisplayName(modelName), + displayName: firstSemanticText( + value?.displayName, + value?.alias, + value?.properties?.displayName, + value?.properties?.alias, + ), description: firstSemanticText( value?.description, value?.properties?.description, @@ -727,13 +717,12 @@ export default function Modeling() { return { name: columnName, type: column?.type, - displayName: - firstSemanticText( - column?.displayName, - column?.alias, - column?.properties?.displayName, - column?.properties?.alias, - ) || semanticFallbackDisplayName(columnName), + displayName: firstSemanticText( + column?.displayName, + column?.alias, + column?.properties?.displayName, + column?.properties?.alias, + ), description: firstSemanticText( column?.description, column?.properties?.description, @@ -1036,9 +1025,8 @@ export default function Modeling() { if (!diagramModel) return []; return { modelId: diagramModel.modelId, - displayName: - firstSemanticText(model.displayName) || - semanticFallbackDisplayName(model.name), + referenceName: diagramModel.referenceName, + displayName: firstSemanticText(model.displayName), description: model.description, columns: (model.columns || []) .map((column) => { @@ -1048,11 +1036,8 @@ export default function Modeling() { return field ? { id: field.columnId, - displayName: - firstSemanticText( - column.displayName, - field.displayName, - ) || semanticFallbackDisplayName(column.name), + referenceName: field.referenceName, + displayName: firstSemanticText(column.displayName), description: column.description, } : null; From 473f7974484405c68fd621af40f58902cb69f887 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 16:48:54 +0530 Subject: [PATCH 1024/1087] Validate and preserve generated semantic aliases --- .../generation/semantics_description.py | 22 ++++-- .../web/v1/services/semantics_description.py | 1 + .../generation/test_semantics_enrichment.py | 68 +++++++++++++++++++ .../server/resolvers/diagramResolver.ts | 14 ++-- .../apollo/server/resolvers/modelResolver.ts | 19 ++++-- wren-ui/src/apollo/server/types/diagram.ts | 2 +- wren-ui/src/pages/modeling.tsx | 37 +++++++--- 7 files changed, 138 insertions(+), 25 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 33a2d307e9..822a1aad8f 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -7,7 +7,7 @@ from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, ValidationError from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider @@ -180,10 +180,18 @@ def wrapper(text: str) -> str: reply = replies[0] # Expecting only one reply normalized = wrapper(reply) + try: + validated = SemanticResult.model_validate(normalized) + except ValidationError as e: + raise ValueError( + "Semantics description LLM returned incomplete semantic metadata. " + "Every selected model and column must include non-empty " + "properties.description and properties.displayName." + ) from e return { model["name"]: model - for model in normalized.get("models", []) + for model in validated.model_dump().get("models", []) if isinstance(model, dict) and model.get("name") } @@ -211,8 +219,14 @@ def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: class ModelProperties(BaseModel): model_config = ConfigDict(extra="forbid") - description: str - displayName: str + description: str = Field(min_length=1) + displayName: str = Field( + min_length=1, + description=( + "Comma-separated natural-language aliases and synonyms users may " + "type for this model or column." + ), + ) class ModelColumns(BaseModel): diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 26c9083a8c..2e71ada539 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -276,6 +276,7 @@ def _is_retryable_chunk_error(self, error: Exception) -> bool: "truncated", "unexpected end of data", "output omitted", + "incomplete semantic metadata", "max_tokens", "natural stopping point", "timed out", diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py index 8a4437766f..2dd39c228b 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py @@ -71,3 +71,71 @@ def test_with_hallucination_and_no_columns(): def test_malformed_json_fails_instead_of_returning_empty_output(): with pytest.raises(ValueError, match="malformed JSON"): normalize({"replies": ['{"models": [']}) + + +def test_normalize_requires_generated_aliases(): + with pytest.raises(ValueError, match="incomplete semantic metadata"): + normalize( + { + "replies": [ + """ + { + "models": [ + { + "name": "orders", + "properties": { + "description": "Customer order transactions.", + "displayName": "orders" + }, + "columns": [ + { + "name": "order_id", + "properties": { + "description": "Unique order identifier.", + "displayName": "" + } + } + ] + } + ] + } + """ + ] + } + ) + + +def test_normalize_preserves_generated_aliases(): + result = normalize( + { + "replies": [ + """ + { + "models": [ + { + "name": "orders", + "properties": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders" + }, + "columns": [ + { + "name": "order_id", + "properties": { + "description": "Unique order identifier.", + "displayName": "order id, order number" + } + } + ] + } + ] + } + """ + ] + } + ) + + assert result["orders"]["properties"]["displayName"] == "orders, sales orders" + assert result["orders"]["columns"][0]["properties"]["displayName"] == ( + "order id, order number" + ) diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index bc3bafd05f..e8d5a0be9f 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -138,11 +138,12 @@ export class DiagramResolver { private transformModel(model: Model): DiagramModel { const properties = this.parseProperties(model.properties); + const displayName = model.displayName || model.referenceName; return { id: uuidv4(), modelId: model.id, nodeType: NodeType.MODEL, - displayName: model.displayName, + displayName, referenceName: model.referenceName, sourceTableName: model.sourceTableName, refSql: model.refSql, @@ -160,6 +161,7 @@ export class DiagramResolver { nestedColumns: ModelNestedColumn[], ): DiagramModelField { const properties = this.parseProperties(column.properties); + const displayName = column.displayName || column.referenceName; return { id: uuidv4(), columnId: column.id, @@ -167,7 +169,7 @@ export class DiagramResolver { ? NodeType.CALCULATED_FIELD : NodeType.FIELD, type: column.type, - displayName: column.displayName, + displayName, referenceName: column.referenceName, description: properties?.description, isPrimaryKey: column.isPk, @@ -178,7 +180,7 @@ export class DiagramResolver { nestedColumnId: nestedColumn.id, columnPath: nestedColumn.columnPath, type: nestedColumn.type, - displayName: nestedColumn.displayName, + displayName: nestedColumn.displayName || nestedColumn.referenceName, referenceName: nestedColumn.referenceName, description: nestedColumn.properties?.description, })) @@ -192,6 +194,7 @@ export class DiagramResolver { ): DiagramModelField | null { const properties = this.parseProperties(column.properties); const lineage = this.parseLineage(column.lineage); + const displayName = column.displayName || column.referenceName; const columnMDL = columnsMDL.find( ({ name }) => name === column.referenceName, ); @@ -214,7 +217,7 @@ export class DiagramResolver { aggregation: column.aggregation, lineage, type: column.type, - displayName: column.displayName, + displayName, referenceName: column.referenceName, description: properties?.description, isPrimaryKey: column.isPk, @@ -259,7 +262,8 @@ export class DiagramResolver { toModelDisplayName: relation.toModelDisplayName || relation.toModelName, toColumnId: relation.toColumnId, toColumnName: relation.toColumnName, - toColumnDisplayName: relation.toColumnDisplayName || relation.toColumnName, + toColumnDisplayName: + relation.toColumnDisplayName || relation.toColumnName, description: properties?.description, }; } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index ce2ea2377e..909a1bbdf8 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -505,6 +505,8 @@ export class ModelResolver { ? modelByReferenceName.get(item.referenceName) : undefined; }; + const hasMetadataText = (value?: string) => + typeof value === 'string' && value.trim().length > 0; for (const item of requestedItems) { if (!resolveModel(item)) { @@ -516,18 +518,23 @@ export class ModelResolver { await Promise.all( requestedItems.map(async (item) => { - if (isNil(item.description) && isNil(item.displayName)) return; + if ( + !hasMetadataText(item.description) && + !hasMetadataText(item.displayName) + ) { + return; + } const model = resolveModel(item); const modelMetadata: Partial = {}; - if (!isNil(item.displayName)) { + if (hasMetadataText(item.displayName)) { modelMetadata.displayName = this.determineMetadataValue( item.displayName, ); } - if (!isNil(item.description)) { + if (hasMetadataText(item.description)) { const properties = model?.properties ? JSON.parse(model.properties) : {}; @@ -582,13 +589,13 @@ export class ModelResolver { if (!column) return; const columnMetadata: Partial = {}; - if (!isNil(requestedColumn.displayName)) { + if (hasMetadataText(requestedColumn.displayName)) { columnMetadata.displayName = this.determineMetadataValue( requestedColumn.displayName, ); } - if (!isNil(requestedColumn.description)) { + if (hasMetadataText(requestedColumn.description)) { const properties = column.properties ? JSON.parse(column.properties) : {}; @@ -600,7 +607,7 @@ export class ModelResolver { if (!isEmpty(columnMetadata)) { await ctx.modelColumnRepository.updateOne( - requestedColumn.id, + column.id, columnMetadata, ); } diff --git a/wren-ui/src/apollo/server/types/diagram.ts b/wren-ui/src/apollo/server/types/diagram.ts index 22f52a82a1..ca91c9aa2f 100644 --- a/wren-ui/src/apollo/server/types/diagram.ts +++ b/wren-ui/src/apollo/server/types/diagram.ts @@ -68,7 +68,7 @@ export interface DiagramModelField { description: string; isPrimaryKey?: boolean; expression?: string; - lineage?: string; + lineage?: number[]; aggregation?: string; nestedFields?: DiagramModelNestedField[]; } diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 112dc4b91c..a9470a5a76 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -120,6 +120,15 @@ const semanticText = (value: any): string => { const firstSemanticText = (...values: any[]): string => values.map(semanticText).find(Boolean) || ''; +const addSemanticText = >( + payload: T, + key: string, + value: any, +): T => { + const text = semanticText(value); + return text ? { ...payload, [key]: text } : payload; +}; + const DiagramWrapper = styled.div` position: relative; height: 100%; @@ -1023,27 +1032,37 @@ export default function Modeling() { (item) => item.referenceName === model.name, ); if (!diagramModel) return []; - return { + const modelPayload = { modelId: diagramModel.modelId, referenceName: diagramModel.referenceName, - displayName: firstSemanticText(model.displayName), - description: model.description, columns: (model.columns || []) .map((column) => { const field = diagramModel.fields.find( (item) => item.referenceName === column.name, ); return field - ? { - id: field.columnId, - referenceName: field.referenceName, - displayName: firstSemanticText(column.displayName), - description: column.description, - } + ? addSemanticText( + addSemanticText( + { + id: field.columnId, + referenceName: field.referenceName, + }, + 'displayName', + column.displayName, + ), + 'description', + column.description, + ) : null; }) .filter(Boolean), }; + + return addSemanticText( + addSemanticText(modelPayload, 'displayName', model.displayName), + 'description', + model.description, + ); }); if (!data.length) { From e9df882041bd11a99cb97a56c3d7ee55f78e1afb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 17:02:44 +0530 Subject: [PATCH 1025/1087] Allow generated semantic column types --- .../src/pipelines/generation/semantics_description.py | 1 + .../pytest/pipelines/generation/test_semantics_enrichment.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 822a1aad8f..0160934243 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -233,6 +233,7 @@ class ModelColumns(BaseModel): model_config = ConfigDict(extra="forbid") name: str + type: str = "" properties: ModelProperties diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py index 2dd39c228b..47e38c7464 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py @@ -121,6 +121,7 @@ def test_normalize_preserves_generated_aliases(): "columns": [ { "name": "order_id", + "type": "VARCHAR", "properties": { "description": "Unique order identifier.", "displayName": "order id, order number" @@ -139,3 +140,4 @@ def test_normalize_preserves_generated_aliases(): assert result["orders"]["columns"][0]["properties"]["displayName"] == ( "order id, order number" ) + assert result["orders"]["columns"][0]["type"] == "VARCHAR" From f892b491babb16ed0354b5b8238ac329d993e546 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 17:05:14 +0530 Subject: [PATCH 1026/1087] Tolerate extra semantic response fields --- .../generation/semantics_description.py | 41 ++++++++++++- .../generation/test_semantics_enrichment.py | 57 +++++++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 0160934243..7583d686e6 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -160,6 +160,43 @@ async def generate(prompt: dict, generator: Any, generator_name: str) -> dict: @observe(capture_input=False) def normalize(generate: dict) -> dict: + def semantic_properties(payload: dict) -> dict: + properties = payload.get("properties") + properties = properties if isinstance(properties, dict) else {} + return { + "description": properties.get( + "description", payload.get("description", "") + ), + "displayName": ( + properties.get("displayName") + or properties.get("alias") + or payload.get("displayName") + or payload.get("alias") + or "" + ), + } + + def semantic_schema_payload(payload: dict) -> dict: + return { + "models": [ + { + "name": model.get("name", ""), + "columns": [ + { + "name": column.get("name", ""), + "type": column.get("type", ""), + "properties": semantic_properties(column), + } + for column in model.get("columns", []) or [] + if isinstance(column, dict) + ], + "properties": semantic_properties(model), + } + for model in payload.get("models", []) or [] + if isinstance(model, dict) + ] + } + def wrapper(text: str) -> str: text = text.replace("\n", " ") text = " ".join(text.split()) @@ -181,7 +218,9 @@ def wrapper(text: str) -> str: reply = replies[0] # Expecting only one reply normalized = wrapper(reply) try: - validated = SemanticResult.model_validate(normalized) + validated = SemanticResult.model_validate( + semantic_schema_payload(normalized) + ) except ValidationError as e: raise ValueError( "Semantics description LLM returned incomplete semantic metadata. " diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py index 47e38c7464..5c22c6dac1 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py @@ -141,3 +141,60 @@ def test_normalize_preserves_generated_aliases(): "order id, order number" ) assert result["orders"]["columns"][0]["type"] == "VARCHAR" + + +def test_normalize_ignores_extra_llm_fields_but_keeps_semantics(): + result = normalize( + { + "replies": [ + """ + { + "models": [ + { + "name": "orders", + "entity": "transaction", + "properties": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders", + "businessUse": "reporting" + }, + "columns": [ + { + "name": "order_id", + "type": "VARCHAR", + "role": "identifier", + "nullable": false, + "properties": { + "description": "Unique order identifier.", + "displayName": "order id, order number", + "examples": ["1001"] + } + } + ] + } + ] + } + """ + ] + } + ) + + assert result == { + "orders": { + "name": "orders", + "columns": [ + { + "name": "order_id", + "type": "VARCHAR", + "properties": { + "description": "Unique order identifier.", + "displayName": "order id, order number", + }, + } + ], + "properties": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders", + }, + } + } From 70c0167909ad3f15e5036da87592af33e3d12695 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 17:26:12 +0530 Subject: [PATCH 1027/1087] Bind generated semantics to selected schema --- .../generation/semantics_description.py | 23 ++- .../web/v1/services/semantics_description.py | 92 +++++++++ .../generation/test_semantics_enrichment.py | 55 ++++++ .../services/test_semantics_description.py | 187 ++++++++++++++++++ 4 files changed, 356 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 7583d686e6..435dfbdea5 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -237,12 +237,33 @@ def wrapper(text: str) -> str: @observe(capture_input=False) def output(normalize: dict, picked_models: list[dict]) -> dict: + def _identifier_key(value: object) -> str: + return "".join( + character for character in str(value).casefold() if character.isalnum() + ) + def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: valid_columns = [col["name"] for col in columns] + matched_columns = [col for col in enriched if col["name"] in valid_columns] + + if matched_columns: + return matched_columns + + if ( + len(enriched) == 1 + and len(columns) == 1 + and _identifier_key(enriched[0].get("name", "")) + == _identifier_key(columns[0].get("name", "")) + ): + return [{**enriched[0], "name": columns[0]["name"]}] - return [col for col in enriched if col["name"] in valid_columns] + return [] models = {model["name"]: model for model in picked_models} + if len(normalize) == 1 and len(models) == 1: + model_name = next(iter(models)) + model_data = next(iter(normalize.values())) + normalize = {model_name: {**model_data, "name": model_name}} return { name: { diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 2e71ada539..a49949eb1e 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -182,6 +182,97 @@ def _display_name(self, payload: dict) -> str: ) return "" if value is None else str(value).strip() + def _identifier_key(self, value: object) -> str: + return "".join( + character for character in str(value).casefold() if character.isalnum() + ) + + def _bind_generated_output_to_chunk_schema( + self, chunk: dict, output: dict + ) -> dict: + chunk_models = [ + model + for model in chunk.get("mdl", {}).get("models", []) or [] + if isinstance(model, dict) and model.get("name") + ] + generated_models = [ + model + for model in output.values() + if isinstance(model, dict) + ] + bound_output: dict = {} + + for expected_model in chunk_models: + expected_model_name = expected_model["name"] + generated_model = output.get(expected_model_name) + if not isinstance(generated_model, dict): + if len(chunk_models) == 1 and len(generated_models) == 1: + generated_model = generated_models[0] + logger.warning( + "Semantics description output used model name %s for selected model %s; binding to selected schema name.", + generated_model.get("name", ""), + expected_model_name, + ) + else: + continue + + expected_columns = [ + column + for column in expected_model.get("columns", []) or [] + if isinstance(column, dict) and column.get("name") + ] + generated_column_items = [ + column + for column in generated_model.get("columns", []) or [] + if isinstance(column, dict) + ] + generated_columns = { + column.get("name"): column + for column in generated_column_items + if column.get("name") + } + bound_columns = [] + + for expected_column in expected_columns: + expected_column_name = expected_column["name"] + generated_column = generated_columns.get(expected_column_name) + if not isinstance(generated_column, dict): + if ( + len(expected_columns) == 1 + and len(generated_column_items) == 1 + and self._identifier_key( + generated_column_items[0].get("name", "") + ) + == self._identifier_key(expected_column_name) + ): + generated_column = generated_column_items[0] + logger.warning( + "Semantics description output used column name %s for selected column %s.%s; binding to selected schema name.", + generated_column.get("name", ""), + expected_model_name, + expected_column_name, + ) + else: + generated_column = {} + + bound_columns.append( + { + **generated_column, + "name": expected_column_name, + "type": expected_column.get( + "type", generated_column.get("type", "") + ), + } + ) + + bound_output[expected_model_name] = { + **generated_model, + "name": expected_model_name, + "columns": bound_columns, + } + + return bound_output + def _validate_generated_output(self, chunk: dict, output: dict) -> None: missing: list[str] = [] @@ -232,6 +323,7 @@ async def _generate_task(self, chunk: dict) -> dict: output = resp.get("output") or {} if not isinstance(output, dict): raise ValueError("Semantics description pipeline returned invalid output") + output = self._bind_generated_output_to_chunk_schema(chunk, output) self._validate_generated_output(chunk, output) return output diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py index 5c22c6dac1..6459d96afb 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py @@ -68,6 +68,61 @@ def test_with_hallucination_and_no_columns(): assert len(result["model1"]["columns"]) == 0 +def test_single_model_output_is_bound_to_picked_model_name(): + test_normalize = { + "source table": { + "name": "source table", + "columns": [{"name": "entity_code"}], + "properties": { + "description": "Business records for source-table activity.", + "displayName": "source records, source activity", + }, + } + } + test_picked_models = [ + { + "name": "schema_source_table", + "columns": [{"name": "entity_code"}], + } + ] + + result = output(test_normalize, test_picked_models) + + assert list(result) == ["schema_source_table"] + assert result["schema_source_table"]["name"] == "schema_source_table" + assert result["schema_source_table"]["columns"][0]["name"] == "entity_code" + + +def test_single_column_output_is_bound_to_picked_column_name(): + test_normalize = { + "users": { + "name": "users", + "columns": [ + { + "name": "created at", + "properties": { + "description": "Timestamp when the user account was created.", + "displayName": "created date, signup date", + }, + } + ], + } + } + test_picked_models = [ + { + "name": "users", + "columns": [{"name": "created_at"}], + } + ] + + result = output(test_normalize, test_picked_models) + + assert result["users"]["columns"][0]["name"] == "created_at" + assert result["users"]["columns"][0]["properties"]["displayName"] == ( + "created date, signup date" + ) + + def test_malformed_json_fails_instead_of_returning_empty_output(): with pytest.raises(ValueError, match="malformed JSON"): normalize({"replies": ['{"models": [']}) diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 85327f1136..c6cf867ccd 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -857,6 +857,193 @@ async def test_llm_descriptions_are_not_rewritten_by_service( ] +@pytest.mark.asyncio +async def test_single_model_output_is_bound_to_selected_schema_name( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the selected datasource", + selected_models=["schema_source_table"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "schema_source_table", + "columns": [ + {"name": "entity_code", "type": "varchar"}, + {"name": "event_date", "type": "timestamp"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "source table": { + "name": "source table", + "properties": { + "description": "Business records for source-table activity and reporting.", + "displayName": "source records, source activity", + }, + "columns": [ + { + "name": "entity_code", + "properties": { + "description": "Entity code used to group and filter records by business entity.", + "displayName": "entity code, business entity code", + }, + }, + { + "name": "event_date", + "properties": { + "description": "Date associated with the business event represented by the record.", + "displayName": "event date, record date", + }, + }, + ], + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.error is None + assert list(response.response) == ["schema_source_table"] + assert response.response["schema_source_table"]["name"] == "schema_source_table" + assert [ + column["name"] + for column in response.response["schema_source_table"]["columns"] + ] == [ + "entity_code", + "event_date", + ] + assert response.response["schema_source_table"]["properties"]["displayName"] == ( + "source records, source activity" + ) + + +@pytest.mark.asyncio +async def test_single_column_output_is_bound_to_selected_schema_name( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the selected datasource", + selected_models=["users"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "users", + "columns": [ + {"name": "created_at", "type": "timestamp"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "users": { + "name": "users", + "properties": { + "description": "User account records for application access and profile management.", + "displayName": "users, user accounts", + }, + "columns": [ + { + "name": "created at", + "properties": { + "description": "Timestamp when the user account was created.", + "displayName": "created date, signup date", + }, + }, + ], + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.error is None + column = response.response["users"]["columns"][0] + assert column["name"] == "created_at" + assert column["type"] == "timestamp" + assert column["properties"]["displayName"] == "created date, signup date" + + +@pytest.mark.asyncio +async def test_multi_model_name_mismatch_fails_without_ambiguous_binding( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the models", + selected_models=["orders", "customers"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [{"name": "order_id", "type": "varchar"}], + }, + { + "name": "customers", + "columns": [{"name": "customer_id", "type": "varchar"}], + }, + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "sales orders": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders", + "columns": [ + { + "name": "order_id", + "properties": { + "description": "Unique identifier for an order.", + "displayName": "order id, order number", + }, + } + ], + }, + "customer records": { + "description": "Customer master records.", + "displayName": "customers, customer records", + "columns": [ + { + "name": "customer_id", + "properties": { + "description": "Unique identifier for a customer.", + "displayName": "customer id, customer number", + }, + } + ], + }, + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "failed" + assert response.response is None + assert "omitted required metadata" in response.error.message + + @pytest.mark.asyncio async def test_batch_processing_partial_failure( service: SemanticsDescription, From aea34e9ba78ccf17d1911af80df09904d3813ee5 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 16 Aug 2026 17:58:47 +0530 Subject: [PATCH 1028/1087] Simplify semantics description pipeline and service --- .../generation/semantics_description.py | 263 +++------ .../web/v1/services/semantics_description.py | 543 +----------------- 2 files changed, 106 insertions(+), 700 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 435dfbdea5..acc5fc8594 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -7,7 +7,7 @@ from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import BaseModel from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider @@ -19,40 +19,74 @@ system_prompt = """ -Generate high-quality semantic metadata for selected data models and their columns. - -Requirements: -1. Return valid JSON that matches the provided schema. -2. Return every input model exactly once and every input column exactly once. -3. Preserve every model and column `name` exactly as provided. -4. Put each generated description in `properties.description`. -5. Put natural-language aliases and synonyms in `properties.displayName` as a required, non-empty, short comma-separated phrase. Do not put SQL identifiers there unless the identifier itself is the natural user-facing name. -6. Make descriptions business-friendly, factual, and useful for text-to-SQL retrieval. -7. Include business context, common analytical use, and the field role when it is supported by the name, type, or relationships: ID/key, date/time, measure, dimension, status, currency, quantity, cost, revenue, rate, percentage, or code. -8. Ground descriptions and aliases only in the user prompt, model and column names, aliases, data types, existing descriptions, and provided schema/relationship context. -9. Use relationship context to distinguish foreign keys, join keys, facts, and dimensions, but do not invent unsupported joins. -10. Make each model and column description specific to that model or column. Never reuse identical descriptions across models or columns in the same response. -11. If two columns have similar names or business meaning, explain the distinction using the exact column name, alias, type, or surrounding model context. -12. Do not invent unsupported tables, columns, relationships, metrics, or business concepts. -13. Do not use generic boilerplate or copy the technical name as the whole description. -14. For each input column, return a matching column object with the exact same `name`, a non-empty `properties.description`, and a non-empty `properties.displayName`. -15. Return complete JSON only. Do not include markdown, comments, examples, or explanatory text outside the JSON object. +I have a data model represented in JSON format, with the following structure: + +``` +[ + {'name': 'model', 'columns': [ + {'name': 'column_1', 'type': 'type', 'properties': {} + }, + {'name': 'column_2', 'type': 'type', 'properties': {} + }, + {'name': 'column_3', 'type': 'type', 'properties': {} + } + ], 'properties': {} + } +] +``` + +Your task is to update this JSON structure by adding a `description` field inside both the `properties` attribute of each `column` and the `model` itself. +Each `description` should be derived from a user-provided input that explains the purpose or context of the `model` and its respective columns. +Follow these steps: +1. **For the `model`**: Prompt the user to provide a brief description of the model's overall purpose or its context. Insert this description in the `properties` field of the `model`. +2. **For each `column`**: Ask the user to describe each column's role or significance. Each column's description should be added under its respective `properties` field in the format: `'description': 'user-provided text'`. +3. Ensure that the output is a well-formatted JSON structure, preserving the input's original format and adding the appropriate `description` fields. + +### Output Format: + +``` +{ + "models": [ + { + "name": "model", + "columns": [ + { + "name": "column_1", + "properties": { + "description": "" + } + }, + { + "name": "column_2", + "properties": { + "description": "" + } + }, + { + "name": "column_3", + "properties": { + "description": "" + } + } + ], + "properties": { + "description": "" + } + } + ] +} +``` + +Make sure that the descriptions are concise, informative, and contextually appropriate based on the input provided by the user. """ user_prompt_template = """ ### Input: User's prompt: {{ user_prompt }} Picked models: {{ picked_models }} -Relationship context: {{ relationship_context }} Localization Language: {{ language }} -Write semantic descriptions for every picked model and every column. -For each model, describe the real-world records represented and the analytical questions it can support. -For each column, describe the business meaning and analytical use of that exact field. -For each model and column, generate non-empty aliases/synonyms that users may naturally type in questions and place them in properties.displayName. -If an existing description is already meaningful, preserve its business meaning while making it clearer and more useful for retrieval. -Keep every description and alias grounded in the picked model metadata, user prompt, data types, and relationship context. -The number of output columns for each model must exactly match the number of input columns for that model. +Please provide a brief description for the model and each column based on the user's prompt. """ @@ -62,40 +96,29 @@ def picked_models(mdl: dict, selected_models: list[str]) -> list[dict]: def relation_filter(column: dict) -> bool: return "relationship" not in column - def _properties(payload: dict) -> dict: - properties = payload.get("properties") - return properties if isinstance(properties, dict) else {} - - def _text(value) -> str: - return "" if value is None else str(value) - def column_formatter(columns: list[dict]) -> list[dict]: return [ { - "name": column.get("name", ""), - "type": column.get("type", ""), + "name": column["name"], + "type": column["type"], "properties": { - "description": _text( - _properties(column).get("description", "") - ), - "displayName": clean_display_name( - _text(_properties(column).get("displayName", "")) + "description": column["properties"].get("description", ""), + "alias": clean_display_name( + column["properties"].get("displayName", "") ), }, } - for column in columns or [] + for column in columns if relation_filter(column) ] def extract(model: dict) -> dict: return { - "name": model.get("name", ""), - "columns": column_formatter(model.get("columns", [])), + "name": model["name"], + "columns": column_formatter(model["columns"]), "properties": { - "description": _text(_properties(model).get("description", "")), - "displayName": clean_display_name( - _text(_properties(model).get("displayName", "")) - ), + "description": model["properties"].get("description", ""), + "alias": clean_display_name(model["properties"].get("displayName", "")), }, } @@ -106,46 +129,15 @@ def extract(model: dict) -> dict: ] -@observe(capture_input=False) -def relationship_context(mdl: dict, selected_models: list[str]) -> list[dict]: - selected = set(selected_models) - relationships = [] - - for relationship in mdl.get("relationships", []) or []: - if not isinstance(relationship, dict): - continue - - models = relationship.get("models", []) or [] - if not any(model in selected for model in models): - continue - - properties = relationship.get("properties") - properties = properties if isinstance(properties, dict) else {} - relationships.append( - { - "name": relationship.get("name", ""), - "models": models, - "joinType": relationship.get("joinType", ""), - "condition": relationship.get("condition", ""), - "description": relationship.get("description") - or properties.get("description", ""), - } - ) - - return relationships - - @observe(capture_input=False) def prompt( picked_models: list[dict], - relationship_context: list[dict], user_prompt: str, prompt_builder: PromptBuilder, language: str, ) -> dict: _prompt = prompt_builder.run( picked_models=picked_models, - relationship_context=relationship_context, user_prompt=user_prompt, language=language, ) @@ -160,43 +152,6 @@ async def generate(prompt: dict, generator: Any, generator_name: str) -> dict: @observe(capture_input=False) def normalize(generate: dict) -> dict: - def semantic_properties(payload: dict) -> dict: - properties = payload.get("properties") - properties = properties if isinstance(properties, dict) else {} - return { - "description": properties.get( - "description", payload.get("description", "") - ), - "displayName": ( - properties.get("displayName") - or properties.get("alias") - or payload.get("displayName") - or payload.get("alias") - or "" - ), - } - - def semantic_schema_payload(payload: dict) -> dict: - return { - "models": [ - { - "name": model.get("name", ""), - "columns": [ - { - "name": column.get("name", ""), - "type": column.get("type", ""), - "properties": semantic_properties(column), - } - for column in model.get("columns", []) or [] - if isinstance(column, dict) - ], - "properties": semantic_properties(model), - } - for model in payload.get("models", []) or [] - if isinstance(model, dict) - ] - } - def wrapper(text: str) -> str: text = text.replace("\n", " ") text = " ".join(text.split()) @@ -205,71 +160,26 @@ def wrapper(text: str) -> str: text_dict = orjson.loads(text.strip()) return text_dict except orjson.JSONDecodeError as e: - raise ValueError( - "Semantics description LLM returned malformed JSON. " - "The response may have been truncated; reduce the selected " - "schema size or increase the configured output token limit." - ) from e - - replies = generate.get("replies") or [] - if not replies: - return {} + logger.error(f"Error decoding JSON: {e}") + return {"models": []} # Return an empty list if JSON decoding fails - reply = replies[0] # Expecting only one reply + reply = generate.get("replies")[0] # Expecting only one reply normalized = wrapper(reply) - try: - validated = SemanticResult.model_validate( - semantic_schema_payload(normalized) - ) - except ValidationError as e: - raise ValueError( - "Semantics description LLM returned incomplete semantic metadata. " - "Every selected model and column must include non-empty " - "properties.description and properties.displayName." - ) from e - return { - model["name"]: model - for model in validated.model_dump().get("models", []) - if isinstance(model, dict) and model.get("name") - } + return {model["name"]: model for model in normalized["models"]} @observe(capture_input=False) def output(normalize: dict, picked_models: list[dict]) -> dict: - def _identifier_key(value: object) -> str: - return "".join( - character for character in str(value).casefold() if character.isalnum() - ) - def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: valid_columns = [col["name"] for col in columns] - matched_columns = [col for col in enriched if col["name"] in valid_columns] - if matched_columns: - return matched_columns - - if ( - len(enriched) == 1 - and len(columns) == 1 - and _identifier_key(enriched[0].get("name", "")) - == _identifier_key(columns[0].get("name", "")) - ): - return [{**enriched[0], "name": columns[0]["name"]}] - - return [] + return [col for col in enriched if col["name"] in valid_columns] models = {model["name"]: model for model in picked_models} - if len(normalize) == 1 and len(models) == 1: - model_name = next(iter(models)) - model_data = next(iter(normalize.values())) - normalize = {model_name: {**model_data, "name": model_name}} return { - name: { - **data, - "columns": _filter(data.get("columns", []), models[name]["columns"]), - } + name: {**data, "columns": _filter(data["columns"], models[name]["columns"])} for name, data in normalize.items() if name in models } @@ -277,37 +187,21 @@ def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: ## End of Pipeline class ModelProperties(BaseModel): - model_config = ConfigDict(extra="forbid") - - description: str = Field(min_length=1) - displayName: str = Field( - min_length=1, - description=( - "Comma-separated natural-language aliases and synonyms users may " - "type for this model or column." - ), - ) + description: str class ModelColumns(BaseModel): - model_config = ConfigDict(extra="forbid") - name: str - type: str = "" properties: ModelProperties class SemanticModel(BaseModel): - model_config = ConfigDict(extra="forbid") - name: str columns: list[ModelColumns] properties: ModelProperties class SemanticResult(BaseModel): - model_config = ConfigDict(extra="forbid") - models: list[SemanticModel] @@ -316,7 +210,6 @@ class SemanticResult(BaseModel): "type": "json_schema", "json_schema": { "name": "semantic_description", - "strict": True, "schema": SemanticResult.model_json_schema(), }, } diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index a49949eb1e..85bf14b188 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -1,6 +1,5 @@ import asyncio import logging -import math from typing import Dict, Literal, Optional import orjson @@ -15,10 +14,6 @@ logger = logging.getLogger("wren-ai-service") -class RetryableSemanticsDescriptionError(ValueError): - pass - - class SemanticsDescription: class Resource(BaseModel, MetadataTraceable): class Error(BaseModel): @@ -37,17 +32,9 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, - generation_timeout_seconds: float = 120.0, - max_models_per_batch: int = 1, - max_columns_per_batch: int = 50, - max_concurrent_tasks: int = 4, ): self._pipelines = pipelines self._cache: Dict[str, self.Resource] = TTLCache(maxsize=maxsize, ttl=ttl) - self._generation_timeout_seconds = generation_timeout_seconds - self._max_models_per_batch = max(1, max_models_per_batch) - self._max_columns_per_batch = max(1, max_columns_per_batch) - self._max_concurrent_tasks = max(1, max_concurrent_tasks) def _handle_exception( self, @@ -73,529 +60,63 @@ class GenerateRequest(BaseRequest): mdl: str def _chunking( - self, - mdl_dict: dict, - request: GenerateRequest, - chunk_size: Optional[int] = None, + self, mdl_dict: dict, request: GenerateRequest, chunk_size: int = 50 ) -> list[dict]: - chunk_size = chunk_size or self._max_models_per_batch template = { "user_prompt": request.user_prompt, "language": request.configurations.language, } - relationships = mdl_dict.get("relationships", []) or [] - - selected_models = [ - model - for model in mdl_dict.get("models", []) - if model.get("name") in request.selected_models - ] - - model_slices: list[tuple[str, dict, int]] = [] - for model in selected_models: - model_name = model["name"] - columns = model.get("columns", []) - column_chunks = [ - columns[j : j + self._max_columns_per_batch] - for j in range(0, len(columns), self._max_columns_per_batch) - ] or [[]] - - for column_chunk in column_chunks: - model_slices.append( - ( - model_name, - {**model, "columns": column_chunk}, - len(column_chunk), - ) - ) - - chunks = [] - batch_models: list[dict] = [] - batch_model_names: set[str] = set() - batch_column_count = 0 - - def relationships_for(model_names: set[str]) -> list[dict]: - return [ - relationship - for relationship in relationships - if isinstance(relationship, dict) - and any( - model_name in model_names - for model_name in relationship.get("models", []) or [] - ) - ] - - def flush_batch(): - nonlocal batch_models, batch_model_names, batch_column_count - if not batch_models: - return - chunks.append( - { - **template, - "mdl": { - "models": batch_models, - "relationships": relationships_for(batch_model_names), - }, - "selected_models": [model["name"] for model in batch_models], - } - ) - batch_models = [] - batch_model_names = set() - batch_column_count = 0 - - for model_name, model_slice, column_count in model_slices: - would_exceed_models = len(batch_models) >= chunk_size - would_exceed_columns = ( - batch_column_count > 0 - and batch_column_count + column_count > self._max_columns_per_batch - ) - would_repeat_model = model_name in batch_model_names - if would_exceed_models or would_exceed_columns or would_repeat_model: - flush_batch() - - batch_models.append(model_slice) - batch_model_names.add(model_name) - batch_column_count += column_count - - flush_batch() - - return chunks - - def _properties(self, payload: dict) -> dict: - value = payload.get("properties") - return value if isinstance(value, dict) else {} - - def _description(self, payload: dict) -> str: - value = payload.get("description") or self._properties(payload).get( - "description", "" - ) - return "" if value is None else str(value).strip() - - def _display_name(self, payload: dict) -> str: - payload_properties = self._properties(payload) - value = ( - payload.get("displayName") - or payload.get("alias") - or payload_properties.get("displayName") - or payload_properties.get("alias") - or "" - ) - return "" if value is None else str(value).strip() - def _identifier_key(self, value: object) -> str: - return "".join( - character for character in str(value).casefold() if character.isalnum() - ) - - def _bind_generated_output_to_chunk_schema( - self, chunk: dict, output: dict - ) -> dict: - chunk_models = [ - model - for model in chunk.get("mdl", {}).get("models", []) or [] - if isinstance(model, dict) and model.get("name") - ] - generated_models = [ - model - for model in output.values() - if isinstance(model, dict) - ] - bound_output: dict = {} - - for expected_model in chunk_models: - expected_model_name = expected_model["name"] - generated_model = output.get(expected_model_name) - if not isinstance(generated_model, dict): - if len(chunk_models) == 1 and len(generated_models) == 1: - generated_model = generated_models[0] - logger.warning( - "Semantics description output used model name %s for selected model %s; binding to selected schema name.", - generated_model.get("name", ""), - expected_model_name, - ) - else: - continue - - expected_columns = [ - column - for column in expected_model.get("columns", []) or [] - if isinstance(column, dict) and column.get("name") - ] - generated_column_items = [ - column - for column in generated_model.get("columns", []) or [] - if isinstance(column, dict) - ] - generated_columns = { - column.get("name"): column - for column in generated_column_items - if column.get("name") - } - bound_columns = [] - - for expected_column in expected_columns: - expected_column_name = expected_column["name"] - generated_column = generated_columns.get(expected_column_name) - if not isinstance(generated_column, dict): - if ( - len(expected_columns) == 1 - and len(generated_column_items) == 1 - and self._identifier_key( - generated_column_items[0].get("name", "") - ) - == self._identifier_key(expected_column_name) - ): - generated_column = generated_column_items[0] - logger.warning( - "Semantics description output used column name %s for selected column %s.%s; binding to selected schema name.", - generated_column.get("name", ""), - expected_model_name, - expected_column_name, - ) - else: - generated_column = {} - - bound_columns.append( - { - **generated_column, - "name": expected_column_name, - "type": expected_column.get( - "type", generated_column.get("type", "") - ), - } - ) - - bound_output[expected_model_name] = { - **generated_model, - "name": expected_model_name, - "columns": bound_columns, - } - - return bound_output - - def _validate_generated_output(self, chunk: dict, output: dict) -> None: - missing: list[str] = [] - - for model in chunk.get("mdl", {}).get("models", []) or []: - if not isinstance(model, dict): - continue - - model_name = model.get("name", "") - generated_model = output.get(model_name) - if not isinstance(generated_model, dict): - missing.append(f"{model_name} model") - continue - - if not self._description(generated_model): - missing.append(f"{model_name} description") - if not self._display_name(generated_model): - missing.append(f"{model_name} alias") - - generated_columns = { - column.get("name"): column - for column in generated_model.get("columns", []) or [] - if isinstance(column, dict) and column.get("name") + chunks = [ + { + **model, + "columns": model["columns"][i : i + chunk_size], } - for column in model.get("columns", []) or []: - if not isinstance(column, dict): - continue - - column_name = column.get("name", "") - generated_column = generated_columns.get(column_name) - if not isinstance(generated_column, dict): - missing.append(f"{model_name}.{column_name} column") - continue - if not self._description(generated_column): - missing.append(f"{model_name}.{column_name} description") - if not self._display_name(generated_column): - missing.append(f"{model_name}.{column_name} alias") - - if missing: - preview = ", ".join(missing[:10]) - suffix = "..." if len(missing) > 10 else "" - raise RetryableSemanticsDescriptionError( - "Semantics description output omitted required metadata: " - f"{preview}{suffix}" - ) - - async def _generate_task(self, chunk: dict) -> dict: - resp = await self._pipelines["semantics_description"].run(**chunk) - output = resp.get("output") or {} - if not isinstance(output, dict): - raise ValueError("Semantics description pipeline returned invalid output") - output = self._bind_generated_output_to_chunk_schema(chunk, output) - self._validate_generated_output(chunk, output) - return output - - def _chunk_columns(self, chunk: dict) -> list[dict]: - models = chunk.get("mdl", {}).get("models", []) - if not models: - return [] - return models[0].get("columns", []) or [] - - def _split_chunk(self, chunk: dict) -> list[dict]: - columns = self._chunk_columns(chunk) - if len(columns) <= 1: - return [] + for model in mdl_dict["models"] + if model["name"] in request.selected_models + for i in range(0, len(model["columns"]), chunk_size) + ] - split_at = max(1, len(columns) // 2) - model = chunk["mdl"]["models"][0] - relationships = chunk.get("mdl", {}).get("relationships", []) return [ { - **chunk, - "mdl": { - "models": [{**model, "columns": column_chunk}], - "relationships": relationships, - }, + **template, + "mdl": {"models": [chunk]}, + "selected_models": [chunk["name"]], } - for column_chunk in (columns[:split_at], columns[split_at:]) - if column_chunk + for chunk in chunks ] - def _is_retryable_chunk_error(self, error: Exception) -> bool: - if isinstance(error, asyncio.TimeoutError): - return True - - if isinstance(error, RetryableSemanticsDescriptionError): - return True - - message = str(error).casefold() - return any( - marker in message - for marker in ( - "malformed json", - "truncated", - "unexpected end of data", - "output omitted", - "incomplete semantic metadata", - "max_tokens", - "natural stopping point", - "timed out", - ) - ) - - async def _generate_task_with_retry_splitting(self, chunk: dict) -> list[dict]: - try: - return [ - await asyncio.wait_for( - self._generate_task(chunk), - timeout=self._generation_timeout_seconds, - ) - ] - except (ValueError, asyncio.TimeoutError) as e: - if not self._is_retryable_chunk_error(e): - raise - - split_chunks = self._split_chunk(chunk) - model_name = (chunk.get("selected_models") or [""])[0] - if not split_chunks: - raise - - logger.warning( - "Retrying semantics description for model %s with smaller " - "column chunks after incomplete or timed-out response: %s", - model_name, - str(e), - ) - outputs: list[dict] = [] - for split_chunk in split_chunks: - outputs.extend( - await self._generate_task_with_retry_splitting(split_chunk) - ) - return outputs - - async def _generate_chunks(self, chunks: list[dict]) -> list[dict]: - semaphore = asyncio.Semaphore(self._max_concurrent_tasks) - - async def _bounded_generate(chunk: dict) -> list[dict]: - async with semaphore: - return await self._generate_task_with_retry_splitting(chunk) - - output_groups = await asyncio.gather( - *[_bounded_generate(chunk) for chunk in chunks] - ) - return [output for group in output_groups for output in group] - - def _request_timeout_seconds(self, chunk_count: int) -> int: - waves = max(1, math.ceil(chunk_count / self._max_concurrent_tasks)) - return self._generation_timeout_seconds * waves - - def _merge_outputs( - self, mdl_dict: dict, selected_models: list[str], outputs: list[dict] - ) -> dict: - def properties(payload: dict) -> dict: - value = payload.get("properties") - return value if isinstance(value, dict) else {} - - def description(payload: dict) -> str: - value = payload.get("description") or properties(payload).get( - "description", "" - ) - return "" if value is None else str(value).strip() - - def display_name(payload: dict) -> str: - return self._display_name(payload) - - def normalized_description(value: str) -> str: - return " ".join(value.casefold().split()) - - generated_by_model: dict[str, dict] = {} - for output in outputs: - for model_name, model_data in output.items(): - if not isinstance(model_data, dict): - continue + async def _generate_task(self, request_id: str, chunk: dict): + resp = await self._pipelines["semantics_description"].run(**chunk) + output = resp.get("output") - generated = generated_by_model.setdefault( - model_name, - { - "name": model_name, - "columns": [], - "properties": {}, - }, - ) - if not description(generated) and description(model_data): - generated["properties"] = { - **properties(generated), - "description": description(model_data), - } - if not display_name(generated) and display_name(model_data): - generated["properties"] = { - **properties(generated), - "displayName": display_name(model_data), - } - generated.setdefault("columns", []) - generated["columns"].extend(model_data.get("columns", [])) + current = self[request_id] + current.response = current.response or {} - response: dict = {} - for model in mdl_dict.get("models", []): - model_name = model.get("name") - if model_name not in selected_models: + for key in output.keys(): + if key not in current.response: + current.response[key] = output[key] continue - generated_model = generated_by_model.get(model_name, {}) - - model_description = description(generated_model) - model_display_name = display_name(generated_model) - - generated_columns: dict[str, dict] = {} - for column in generated_model.get("columns", []): - if not isinstance(column, dict): - continue - - column_name = column.get("name") - if not column_name: - continue - - if column_name in generated_columns: - logger.warning( - "Semantics description output duplicated column: %s.%s", - model_name, - column_name, - ) - continue - - generated_columns[column_name] = column - - columns = [] - used_generated_descriptions: dict[str, str] = {} - for column in model.get("columns", []): - if not isinstance(column, dict): - continue - - column_name = column.get("name", "") - generated_column = generated_columns.get(column_name) - original_description = description(column) - original_display_name = display_name(column) - generated_description = ( - description(generated_column) if generated_column else "" - ) - generated_display_name = ( - display_name(generated_column) if generated_column else "" - ) - normalized_generated_description = normalized_description( - generated_description - ) - is_repeated_generated_description = ( - generated_description - and normalized_generated_description in used_generated_descriptions - and used_generated_descriptions[normalized_generated_description] - != column_name - ) - if is_repeated_generated_description and original_description: - logger.warning( - "Semantics description output reused description for columns: %s.%s and %s.%s", - model_name, - used_generated_descriptions[normalized_generated_description], - model_name, - column_name, - ) - column_description = original_description - else: - column_description = generated_description or original_description - if generated_description: - used_generated_descriptions.setdefault( - normalized_generated_description, column_name - ) - - column_display_name = ( - generated_display_name or original_display_name - ) - - columns.append( - { - "name": column_name, - "type": column.get("type", ""), - "properties": { - "description": column_description, - "displayName": column_display_name, - }, - } - ) - - if not model_description: - model_description = description(model) - if not model_display_name: - model_display_name = display_name(model) - - response[model_name] = { - "name": model_name, - "columns": columns, - "properties": { - "description": model_description, - "displayName": model_display_name, - }, - } - - return response + current.response[key]["columns"].extend(output[key]["columns"]) @observe(name="Generate Semantics Description") @trace_metadata async def generate(self, request: GenerateRequest, **kwargs) -> Resource: logger.info("Generate Semantics Description pipeline is running...") trace_id = kwargs.get("trace_id") - request_timeout_seconds = self._generation_timeout_seconds try: mdl_dict = orjson.loads(request.mdl) chunks = self._chunking(mdl_dict, request) - if not chunks: - raise ValueError( - "No selected models matched the current semantic model metadata" - ) - request_timeout_seconds = self._request_timeout_seconds(len(chunks)) - outputs = await self._generate_chunks(chunks) + tasks = [self._generate_task(request.id, chunk) for chunk in chunks] - self[request.id] = self.Resource( - id=request.id, - status="finished", - response=self._merge_outputs( - mdl_dict, request.selected_models, list(outputs) - ), - trace_id=trace_id, - request_from=request.request_from, - ) + await asyncio.gather(*tasks) + + self[request.id].status = "finished" + self[request.id].trace_id = trace_id + self[request.id].request_from = request.request_from except orjson.JSONDecodeError as e: self._handle_exception( request.id, @@ -604,14 +125,6 @@ async def generate(self, request: GenerateRequest, **kwargs) -> Resource: trace_id=trace_id, request_from=request.request_from, ) - except asyncio.TimeoutError: - self._handle_exception( - request.id, - "Semantics description generation timed out after " - f"{request_timeout_seconds} seconds", - trace_id=trace_id, - request_from=request.request_from, - ) except Exception as e: self._handle_exception( request.id, From 694a0469c54e8befa81807f829404b9a2f6924b9 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 16 Aug 2026 18:03:45 +0530 Subject: [PATCH 1029/1087] Simplify semantics description config --- wren-ai-service/src/globals.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 51aaffc6f3..6ab877b352 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -90,13 +90,6 @@ def create_service_container( **pipe_components["semantics_description"], ) }, - generation_timeout_seconds=( - settings.semantics_description_timeout_seconds - or settings.semantics_description_generation_timeout_seconds - ), - max_models_per_batch=settings.semantics_description_max_models_per_batch, - max_columns_per_batch=settings.semantics_description_max_columns_per_batch, - max_concurrent_tasks=settings.semantics_description_max_concurrent_tasks, **query_cache, ), semantics_preparation_service=services.SemanticsPreparationService( From a0b71c2b57de494ba088f0fa9b0eefc210d82165 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 16 Aug 2026 18:32:30 +0530 Subject: [PATCH 1030/1087] Handle missing 'models' key in normalize --- .../src/pipelines/generation/semantics_description.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index acc5fc8594..3fc82ccd67 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -166,7 +166,7 @@ def wrapper(text: str) -> str: reply = generate.get("replies")[0] # Expecting only one reply normalized = wrapper(reply) - return {model["name"]: model for model in normalized["models"]} + return {model["name"]: model for model in normalized.get("models", [])} @observe(capture_input=False) From df95e4b28c29a8180b29a5fb899637211c30f68d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 16 Aug 2026 19:16:44 +0530 Subject: [PATCH 1031/1087] Simplify DB schema retrieval and DDL builders --- .../retrieval/db_schema_retrieval.py | 528 ++---------------- 1 file changed, 62 insertions(+), 466 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 884bbaaba9..22bf940cd8 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -30,13 +30,7 @@ ### TASK ### You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. -The database schema includes structural, semantic, and business modeling metadata: -- Models are logical datasets backed by physical tables or SQL definitions. -- Columns are exposed fields, including renamed fields, expressions, primary keys, and calculated fields. -- Relationships are reusable join logic between models. -- Calculated fields are business logic defined once and reused across queries. -- Views are named SQL statements that behave like stable virtual tables. -- Metrics are structured aggregation objects with measures and dimensions. +The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -46,15 +40,6 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -8. Map the business question to the modeled datasets whose descriptions, aliases, columns, calculated fields, views, metrics, and relationships support the intent. -9. Prefer modeled analytical interfaces such as views and metrics when they expose the fields needed to answer the question. -10. If the answer needs fields, filters, time dimensions, ordering, aggregations, or relationship keys from multiple related datasets, include every required related dataset and the columns needed from each one. -11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. -12. Follow only the relationships shown in the provided schema when selecting columns across datasets. -13. Do not stop at a single top candidate when the question requires multiple related datasets. -14. If the same business concept is represented by multiple modeled datasets, select only the dataset or related dataset set whose declared fields and relationships best support the current question. -15. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. -16. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -127,283 +112,37 @@ def _project_filter_conditions( def _build_metric_ddl(content: dict) -> str: - columns = [ - column - for column in content["columns"] - if column["data_type"].lower() != "unknown" - ] - context = _format_semantic_context( - { - "object_type": "metric", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in columns - ], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable analytical aggregation interface", - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "semantic_context_not_sql_identifier": column["comment"], - } - for column in columns - ], - } - ) columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" - for column in columns + f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in content["columns"] + if column["data_type"].lower() + != "unknown" # quick fix: filtering out UNKNOWN column type ] return ( - f"{context}CREATE TABLE {content['name']} (\n " + f"{content['comment']}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) def _build_view_ddl(content: dict) -> str: - columns = [ - column - for column in content.get("columns", []) - if column.get("name") and column.get("data_type", "").lower() != "unknown" - ] - context = _format_semantic_context( - { - "object_type": "view", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in columns - ], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable virtual table interface", - "description": content["comment"], - "definition_omitted_from_executable_schema": True, - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type( - column.get("data_type") - ), - "semantic_context_not_sql_identifier": column.get("comment", ""), - } - for column in columns - ], - } - ) - columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" - for column in columns - ] - return ( - f"{context}CREATE TABLE {content['name']} (\n " - + ",\n ".join(columns_ddl) - + "\n);" + f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" ) -def _format_semantic_context(context: dict) -> str: - return ( - "/*\n" - "WREN RETRIEVED SEMANTIC CONTEXT\n" - f"{orjson.dumps(context).decode('utf-8')}\n" - f"{_format_identifier_contract(context)}" - "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" - "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" - "*/\n" - f"{_format_executable_identifier_catalog(context)}" - ) - - -def _format_executable_identifier_catalog(context: dict) -> str: - contract = context.get("sql_identifier_contract", {}) - table_name = contract.get("sql_table_name_use_exactly") - column_names = contract.get("sql_column_names_use_exactly") or [ - column["sql_column_name_use_exactly"] - for column in context.get("columns", []) - if column.get("sql_column_name_use_exactly") - ] - relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ - relationship["sql_relationship_constraint_use_exactly"] - for relationship in context.get("relationships", []) - if relationship.get("sql_relationship_constraint_use_exactly") - ] - - lines = [ - "### EXECUTABLE WREN IDENTIFIER CATALOG ###", - "Copy SQL identifiers only from this catalog or the following DDL.", - "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", - f"object_type: {context.get('object_type', '')}", - ] - if table_name: - lines.append(f"table: {table_name}") - if column_names: - lines.append("columns:") - lines.extend(f"- {column_name}" for column_name in column_names) - if relationship_constraints: - lines.append("relationships:") - lines.extend(f"- {constraint}" for constraint in relationship_constraints) - lines.extend( - [ - "Use only the listed identifiers and the identifiers declared in the following DDL when writing executable SQL.", - "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", - "", - ] - ) - return "\n".join(lines) - - -def _format_identifier_contract(context: dict) -> str: - contract = context.get("sql_identifier_contract", {}) - table_name = contract.get("sql_table_name_use_exactly") - column_names = contract.get("sql_column_names_use_exactly") or [ - column["sql_column_name_use_exactly"] - for column in context.get("columns", []) - if column.get("sql_column_name_use_exactly") - ] - relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ - relationship["sql_relationship_constraint_use_exactly"] - for relationship in context.get("relationships", []) - if relationship.get("sql_relationship_constraint_use_exactly") - ] - - lines = [ - "WREN SQL IDENTIFIER CONTRACT", - f"object_type: {context.get('object_type', '')}", - ] - if table_name: - lines.append(f"sql_table_name_use_exactly: {table_name}") - if column_names: - lines.append("sql_column_names_use_exactly:") - lines.extend(f"- {column_name}" for column_name in column_names) - if relationship_constraints: - lines.append("relationship_constraints_use_exactly:") - lines.extend( - f"- {relationship_constraint}" - for relationship_constraint in relationship_constraints - ) - lines.extend( - [ - "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", - "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", - "END WREN SQL IDENTIFIER CONTRACT", - "", - ] - ) - return "\n".join(lines) - - -def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: - relationship_columns = { - column.get("column") - for column in content["columns"] - if column["type"] == "FOREIGN_KEY" - and (not tables or set(column.get("tables", [])).issubset(tables)) - } - relationship_columns.discard(None) - return relationship_columns - - -def _included_columns( - content: dict, columns: Optional[set[str]], tables: Optional[set[str]] -) -> list[dict]: - relationship_columns = _included_relationship_columns(content, tables) - return [ - column - for column in content["columns"] - if column["type"] == "COLUMN" - and ( - not columns - or column["name"] in columns - or column["name"] in relationship_columns - or column["is_primary_key"] - ) - and column["data_type"].lower() != "unknown" - ] - - -def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: - return [ - column - for column in content["columns"] - if column["type"] == "FOREIGN_KEY" - and (not tables or set(column.get("tables", [])).issubset(tables)) - ] - - -def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: - executable_columns = { - column["name"] - for column in content["columns"] - if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" - } - return bool(columns) and columns.issubset(executable_columns) - - -def _build_table_retrieval_context( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None -) -> tuple[str, bool, bool]: - ddl, has_calculated_field, has_json_field = build_table_ddl( - content, - columns=columns, - tables=tables, - include_semantic_comments=False, - ) - included_columns = _included_columns(content, columns, tables) - included_relationships = _included_relationships(content, tables) - context = _format_semantic_context( - { - "object_type": "model", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in included_columns - ], - "relationship_constraints_use_exactly": [ - relationship["constraint"] - for relationship in included_relationships - ], - }, - "semantic_context_not_sql_identifiers": { - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "is_primary_key": column["is_primary_key"], - "semantic_context_not_sql_identifier": column["comment"], - } - for column in included_columns - ], - "relationships": [ - { - "semantic_context_not_sql_identifier": relationship["comment"], - "sql_relationship_constraint_use_exactly": relationship[ - "constraint" - ], - "related_models_use_exactly": relationship.get("tables", []), - } - for relationship in included_relationships - ], - } - ) - return f"{context}{ddl}", has_calculated_field, has_json_field - - ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: + if histories: + previous_query_summaries = [history.question for history in histories] + else: + previous_query_summaries = [] + + query = "\n".join(previous_query_summaries) + "\n" + query + return await embedder.run(query) else: return {} @@ -447,172 +186,34 @@ async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any, - embedding: dict | None = None, - mdl_hash: str | None = None, -) -> list[Document]: - table_names = _table_names_from_description_documents( - table_retrieval.get("documents", []) - ) - documents = [] - if embedding and not table_names: - documents = await _retrieve_semantic_schema_documents( - embedding, project_id, dbschema_retriever, mdl_hash - ) - table_names = _table_names_from_schema_documents(documents) - - if table_names: - retrieved_table_names = set() - pending_table_names = table_names - - while pending_table_names: - retrieved_table_names.update(pending_table_names) - retrieved_documents = await _retrieve_schema_documents( - pending_table_names, project_id, dbschema_retriever, mdl_hash - ) - documents = _dedupe_documents(documents + retrieved_documents) - pending_table_names = [ - table_name - for table_name in _related_table_names(documents) - if table_name not in retrieved_table_names - ] - - return documents - - return [] - - -async def _retrieve_semantic_schema_documents( - embedding: dict, - project_id: str, - dbschema_retriever: Any, mdl_hash: str | None = None, ) -> list[Document]: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - } - - filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) - - results = await dbschema_retriever.run( - query_embedding=embedding.get("embedding"), - filters=filters, - ) - return results["documents"] - - -def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: + tables = table_retrieval.get("documents", []) table_names = [] - seen = set() - - for document in documents: - table_name = document.meta.get("name") - if not table_name: - content = ast.literal_eval(document.content) - table_name = content.get("name") - - if table_name and table_name not in seen: - table_names.append(table_name) - seen.add(table_name) - - return table_names - - -def _merge_names(*name_groups: list[str]) -> list[str]: - merged = [] - seen = set() - - for names in name_groups: - for name in names: - if name in seen: - continue - merged.append(name) - seen.add(name) - - return merged - + for table in tables: + content = ast.literal_eval(table.content) + table_names.append(content["name"]) -def _table_names_from_description_documents(documents: list[Document]) -> list[str]: - table_names = [] - seen = set() - - for document in documents: - content = ast.literal_eval(document.content) - table_name = content["name"] - if table_name not in seen: - table_names.append(table_name) - seen.add(table_name) - - return table_names - - -async def _retrieve_schema_documents( - table_names: list[str], - project_id: str, - dbschema_retriever: Any, - mdl_hash: str | None = None, -) -> list[Document]: table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names ] - if not table_name_conditions: - return [] - - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } - - filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) - - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] - - -def _related_table_names(documents: list[Document]) -> list[str]: - related_table_names = [] - seen = set() - - for document in documents: - content = ast.literal_eval(document.content) - if content.get("type") != "TABLE_COLUMNS": - continue - - for column in content.get("columns", []): - if column.get("type") != "FOREIGN_KEY": - continue - - for table_name in column.get("tables", []): - if table_name not in seen: - related_table_names.append(table_name) - seen.add(table_name) - - return related_table_names + if table_name_conditions: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) -def _dedupe_documents(documents: list[Document]) -> list[Document]: - deduped = [] - seen = set() + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] - for document in documents: - identity = ( - document.meta.get("type"), - document.meta.get("name"), - document.content, - ) - if identity in seen: - continue - deduped.append(document) - seen.add(identity) - - return deduped + return [] @observe() @@ -658,9 +259,7 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context(table_schema) - ) + ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) retrieval_results.append( { "table_name": table_schema["name"], @@ -723,10 +322,16 @@ def prompt( ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ - _build_table_retrieval_context(construct_db_schema)[0] + build_table_ddl(construct_db_schema)[0] for construct_db_schema in construct_db_schemas ] + previous_query_summaries = ( + [history.question for history in histories] if histories else [] + ) + + query = "\n".join(previous_query_summaries) + "\n" + query + _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: @@ -772,22 +377,12 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - selected_columns = set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ) - columns = ( - selected_columns - if _selected_columns_are_executable( - table_schema, selected_columns - ) - else None - ) - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context( - table_schema, - columns=columns, - tables=tables, - ) + ddl, _has_calculated_field, _has_json_field = build_table_ddl( + table_schema, + columns=set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ), + tables=tables, ) if _has_calculated_field: has_calculated_field = True @@ -802,23 +397,24 @@ def construct_retrieval_results( ) for document in dbschema_retrieval: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - } - ) + if document.meta["name"] in columns_and_tables_needed: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) return { "retrieval_results": retrieval_results, From 30d25cd48c23d3e44b312b9b951f8b07b6890317 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 16 Aug 2026 19:39:44 +0530 Subject: [PATCH 1032/1087] Handle missing SQL generation fields --- .../src/pipelines/generation/utils/sql.py | 19 ++++++++++++++++--- wren-ai-service/src/web/v1/services/ask.py | 5 ++++- .../src/web/v1/services/ask_feedback.py | 5 ++++- 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 90920a2485..2d16c1a5f7 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -41,9 +41,22 @@ async def run( # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' if cleaned_generation_result.startswith("{"): - cleaned_generation_result = orjson.loads(cleaned_generation_result)[ - "sql" - ] + generation_result = orjson.loads(cleaned_generation_result) + cleaned_generation_result = generation_result.get("sql") + if not cleaned_generation_result: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "SQL_GENERATION", + "error": ( + "SQL generation response did not include the required " + f"'sql' field: {generation_result}" + ), + "correlation_id": "", + }, + } ( valid_generation_result, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8beefc01e0..8969ff7b6f 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -525,7 +525,10 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + if failed_dry_run_result["type"] in ( + "TIME_OUT", + "SQL_GENERATION", + ): error_message = failed_dry_run_result["error"] invalid_sql = failed_dry_run_result["sql"] break diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 9abf43a061..f9d1498857 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -212,7 +212,10 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] != "TIME_OUT": + if failed_dry_run_result["type"] not in ( + "TIME_OUT", + "SQL_GENERATION", + ): original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] From 82fe9c783f098b3aa6d48873b2b6e95079d06dfa Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Sun, 16 Aug 2026 20:02:51 +0530 Subject: [PATCH 1033/1087] Update semantics generation --- .../src/pipelines/generation/semantics_description.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index 3fc82ccd67..d1d4db0475 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -206,6 +206,7 @@ class SemanticResult(BaseModel): SEMANTICS_DESCRIPTION_MODEL_KWARGS = { + "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { From 64fde9059813f1bc4ca1c10a09ae4bd2094b27e9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 20:03:35 +0530 Subject: [PATCH 1034/1087] Validate SQL generation before execution --- .../generation/followup_sql_generation.py | 2 +- .../pipelines/generation/sql_correction.py | 2 +- .../pipelines/generation/sql_generation.py | 2 +- .../src/pipelines/generation/utils/sql.py | 98 +++++++++++--- wren-ai-service/src/web/v1/services/ask.py | 1 + .../src/web/v1/services/ask_feedback.py | 1 + .../test_sql_generation_post_processor.py | 121 ++++++++++++++++++ 7 files changed, 203 insertions(+), 24 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py 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 4a4dfe6801..3e8df185e8 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -216,7 +216,7 @@ async def run( ): logger.info("Follow-Up SQL Generation pipeline is running...") - if use_dry_plan: + if project_id or use_dry_plan: metadata = await retrieve_metadata( project_id or "", self._retriever, mdl_hash ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 62555bcd8d..f424ec48cb 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -198,7 +198,7 @@ async def run( ): logger.info("SQLCorrection pipeline is running...") - if use_dry_plan: + if project_id or use_dry_plan: metadata = await retrieve_metadata( project_id or "", self._retriever, mdl_hash ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 31a1598823..a3eb1c48ef 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -208,7 +208,7 @@ async def run( ): logger.info("SQL Generation pipeline is running...") - if use_dry_plan: + if project_id or use_dry_plan: metadata = await retrieve_metadata( project_id or "", self._retriever, mdl_hash ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2d16c1a5f7..f2e9ee971f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -22,6 +22,68 @@ class SQLGenPostProcessor: def __init__(self, engine: Engine): self._engine = engine + def _looks_like_sql(self, value: str) -> bool: + normalized = value.strip().casefold() + return normalized.startswith("select ") or normalized.startswith("with ") + + def _json_object(self, value: str) -> dict[str, Any]: + parsed = orjson.loads(value) + return parsed if isinstance(parsed, dict) else {} + + def _extract_sql_from_object(self, generation_result: dict[str, Any]) -> str: + sql = generation_result.get("sql") + if isinstance(sql, str) and sql.strip(): + return clean_generation_result(sql) + + arguments = generation_result.get("arguments") + if isinstance(arguments, str) and arguments.strip(): + try: + arguments = self._json_object(arguments) + except orjson.JSONDecodeError: + arguments = {} + + if isinstance(arguments, dict): + for key in ("sql", "query"): + value = arguments.get(key) + if isinstance(value, str) and self._looks_like_sql(value): + return clean_generation_result(value) + + query = generation_result.get("query") + if isinstance(query, str) and self._looks_like_sql(query): + return clean_generation_result(query) + + return "" + + def _extract_sql(self, replies: List[str] | List[List[str]]) -> tuple[str, str]: + if not replies: + return "", "SQL generation response was empty." + + reply = replies[0] + if isinstance(reply, list): + reply = reply[0] if reply else "" + + cleaned_generation_result = clean_generation_result(reply) + if not cleaned_generation_result: + return "", "SQL generation response was empty." + + if cleaned_generation_result.startswith("{"): + try: + generation_result = self._json_object(cleaned_generation_result) + except orjson.JSONDecodeError as e: + return "", f"SQL generation response was not valid JSON: {e}" + + sql = self._extract_sql_from_object(generation_result) + if sql: + return sql, "" + + return ( + "", + "SQL generation response did not include a supported SQL field: " + f"{generation_result}", + ) + + return cleaned_generation_result, "" + @component.output_types( valid_generation_result=Dict[str, Any], invalid_generation_result=Dict[str, Any], @@ -37,32 +99,25 @@ async def run( allow_data_preview: bool = False, ) -> dict: try: - cleaned_generation_result = clean_generation_result(replies[0]) - - # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' - if cleaned_generation_result.startswith("{"): - generation_result = orjson.loads(cleaned_generation_result) - cleaned_generation_result = generation_result.get("sql") - if not cleaned_generation_result: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": "", - "original_sql": "", - "type": "SQL_GENERATION", - "error": ( - "SQL generation response did not include the required " - f"'sql' field: {generation_result}" - ), - "correlation_id": "", - }, + generation_result, extraction_error = self._extract_sql(replies) + if not generation_result: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": extraction_error + or "No grounded SQL was generated from the current schema.", + "correlation_id": "", } + } ( valid_generation_result, invalid_generation_result, ) = await self._classify_generation_result( - cleaned_generation_result, + generation_result, project_id=project_id, mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, @@ -98,7 +153,8 @@ async def _classify_generation_result( use_dry_run = not allow_data_preview async with aiohttp.ClientSession() as session: - if use_dry_plan: + should_dry_plan = use_dry_plan or bool(project_id and data_source) + if should_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( session, generation_result, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8969ff7b6f..cc5151aba0 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -527,6 +527,7 @@ async def ask( while current_sql_correction_retries < max_sql_correction_retries: if failed_dry_run_result["type"] in ( "TIME_OUT", + "NO_RELEVANT_SQL", "SQL_GENERATION", ): error_message = failed_dry_run_result["error"] diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index f9d1498857..379a661a29 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -214,6 +214,7 @@ async def ask_feedback( ]["invalid_generation_result"]: if failed_dry_run_result["type"] not in ( "TIME_OUT", + "NO_RELEVANT_SQL", "SQL_GENERATION", ): original_sql = failed_dry_run_result["original_sql"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py new file mode 100644 index 0000000000..a775453270 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -0,0 +1,121 @@ +from typing import Any + +import aiohttp +import pytest + +from src.core.engine import Engine +from src.pipelines.generation.utils.sql import SQLGenPostProcessor + + +class FakeEngine(Engine): + def __init__(self, dry_plan_success: bool = True, execute_success: bool = True): + self.dry_plan_success = dry_plan_success + self.execute_success = execute_success + self.dry_plan_calls: list[dict[str, Any]] = [] + self.execute_sql_calls: list[dict[str, Any]] = [] + + async def dry_plan( + self, + session: aiohttp.ClientSession, + sql: str, + data_source: str, + project_id: str | None = None, + mdl_hash: str | None = None, + allow_fallback: bool = True, + **kwargs, + ): + self.dry_plan_calls.append( + { + "sql": sql, + "data_source": data_source, + "project_id": project_id, + "mdl_hash": mdl_hash, + "allow_fallback": allow_fallback, + } + ) + return self.dry_plan_success, "" if self.dry_plan_success else "plan failed" + + async def execute_sql( + self, + sql: str, + session: aiohttp.ClientSession, + dry_run: bool = True, + **kwargs, + ): + self.execute_sql_calls.append( + { + "sql": sql, + "dry_run": dry_run, + **kwargs, + } + ) + return self.execute_success, {}, {"correlation_id": "correlation-id"} + + +@pytest.mark.asyncio +async def test_post_processor_extracts_tool_call_query_argument(): + engine = FakeEngine() + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + [ + '{"name":"query","arguments":{"query":"SELECT supplierid, COUNT(*) FROM PO_Invoices GROUP BY supplierid;"}}' + ], + project_id="project-id", + mdl_hash="manifest-hash", + data_source="mssql", + ) + + assert result["valid_generation_result"] == { + "sql": "SELECT supplierid, COUNT(*) FROM PO_Invoices GROUP BY supplierid", + "correlation_id": "correlation-id", + } + assert engine.dry_plan_calls == [ + { + "sql": "SELECT supplierid, COUNT(*) FROM PO_Invoices GROUP BY supplierid", + "data_source": "mssql", + "project_id": "project-id", + "mdl_hash": "manifest-hash", + "allow_fallback": True, + } + ] + assert engine.execute_sql_calls[0]["dry_run"] is True + + +@pytest.mark.asyncio +async def test_post_processor_returns_no_relevant_sql_for_missing_sql_field(): + processor = SQLGenPostProcessor(FakeEngine()) + + result = await processor.run( + ['{"name":"query","arguments":{"question":"Show suppliers"}}'], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert "supported SQL field" in result["invalid_generation_result"]["error"] + + +@pytest.mark.asyncio +async def test_post_processor_dry_plans_before_preview_execution(): + engine = FakeEngine(dry_plan_success=False) + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ['{"sql":"SELECT * FROM orders"}'], + project_id="project-id", + mdl_hash="manifest-hash", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT * FROM orders", + "original_sql": "SELECT * FROM orders", + "type": "DRY_PLAN", + "error": "plan failed", + "correlation_id": "", + } + assert len(engine.dry_plan_calls) == 1 + assert engine.execute_sql_calls == [] From 3b8b086690e87991991c64a9c1d4e5d92bea7318 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 20:54:32 +0530 Subject: [PATCH 1035/1087] Align ask SQL generation with legacy grounding --- .../pipelines/generation/sql_generation.py | 12 +- .../generation/sql_generation_reasoning.py | 7 +- .../src/pipelines/generation/utils/sql.py | 301 ++++------- .../retrieval/db_schema_retrieval.py | 504 +++++++++++++++--- .../test_sql_generation_post_processor.py | 32 ++ 5 files changed, 591 insertions(+), 265 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index a3eb1c48ef..807bf73522 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -54,11 +54,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -71,12 +70,11 @@ ### QUESTION ### User's Question: {{ query }} -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} -{% endif %} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. -Let's think step by step. +Return only the final JSON SQL response. """ 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 5d4b58629f..723566b2ab 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -6,8 +6,8 @@ from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines @@ -29,11 +29,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} @@ -49,7 +48,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Let's think step by step. +Return only the reasoning plan described by the system instructions. When relevant, ground the plan by using the literal prefix `table:` followed by an exact declared table name from DATABASE SCHEMA, or the literal prefix `column:` followed by an exact declared table name, a dot, and an exact declared column name. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, template markers, functions, or identifier-like labels. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index f2e9ee971f..4fa97f2632 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -76,6 +76,9 @@ def _extract_sql(self, replies: List[str] | List[List[str]]) -> tuple[str, str]: if sql: return sql, "" + if "sql" in generation_result and generation_result.get("sql") is None: + return "", "No grounded SQL was generated from the current schema." + return ( "", "SQL generation response did not include a supported SQL field: " @@ -110,7 +113,7 @@ async def run( "error": extraction_error or "No grounded SQL was generated from the current schema.", "correlation_id": "", - } + }, } ( @@ -263,6 +266,44 @@ async def _classify_generation_result( return valid_generation_result, invalid_generation_result +_MANDATORY_SQL_GROUNDING_RULES = """ +### MANDATORY SQL GROUNDING RULES ### +- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. +- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. +- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. +- Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. +- Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. +- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. +- When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. +- In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. +- Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. +- When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. +- The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. +- Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. +- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. +- Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. +- If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. +- When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. +- Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. +- When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- When the same requested result can be answered from multiple schema objects with compatible columns or metrics, include all relevant schema objects by combining separate result rows with UNION ALL instead of choosing only one object. +- Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. +- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. +- Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. +- SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. +- Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. +- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. +- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. +- If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. +- For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. +- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. +- If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. +- If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. +- If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. +- Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. +""" + + _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. @@ -274,52 +315,31 @@ async def _classify_generation_result( - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. + - Use Wren SQL identifier quoting with double quotes only; the engine rewrite step converts grounded Wren SQL to the active connector dialect. - Put single quotes around string literals. - Never quote numeric literals. - For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992; -- YOU MUST USE "lower(.) like lower()" function or "lower(.) = lower()" function for case-insensitive comparison! - - Use "lower(.) LIKE lower()" when: - - The user requests a pattern or partial match. - - The value is not specific enough to be a single, exact value. - - Wildcards (%) are needed to capture the pattern. - - Use "lower(.) = lower()" when: - - The user requests an exact, specific value. - - There is no ambiguity or pattern in the value. -- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query - - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds - - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds - - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds -- ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query - - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) -- If the user asks for a specific date, please give the date range in SQL query - - example: "What is the total revenue for the month of 2024-11-01?" - - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" +- Generate Wren SQL syntax only, not connector-specific SQL syntax. +- Never use SELECT TOP, TOP(...), FETCH FIRST, square-bracket identifiers, or backtick identifiers. For top or limit requests, sort with ORDER BY and put LIMIT at the end of the query. +- Preserve every deployed table and column identifier exactly as it appears in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, including spaces, digits, underscores, case, and punctuation, then wrap that exact identifier in double quotes in SQL. +- Do not convert deployed identifiers into display-friendly variants by replacing spaces with underscores, removing prefixes, changing case, shortening names, or expanding abbreviations. +- For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. +- For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only date/time functions and casts whose exact syntax is provided in SQL FUNCTIONS for this request. +- If the question asks for a specific or relative date, generate a bounded date/time filter only when both the exact date/time schema column and required SQL FUNCTIONS-supported operation are available. If either is missing, do not invent a field or function. - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. -- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. - - EXAMPLE - DATABASE SCHEMA - /* {"alias":"_orders","description":"A model representing the orders data."} */ - CREATE TABLE orders ( - -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} - ApprovedTimestamp TIMESTAMP - } - - SQL - SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; -- DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. +- Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. +- Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. +- Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. +- DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. - DON'T USE "EXTRACT()" function with INTERVAL data types as arguments - DON'T USE INTERVAL or generate INTERVAL-like expression in the generated SQL query. - DON'T USE "TO_CHAR" function in the generated SQL query. +- DON'T USE unsupported statistical, date/time, or formatting functions. If SQL FUNCTIONS does not list a function needed by the requested intent, omit the function-dependent part. If that function is required to answer the request, return null for sql. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -- For the ranking problem, you must add the ranking column to the final SELECT clause. +- For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. """ @@ -327,47 +347,9 @@ async def _classify_generation_result( #### Instructions for Calculated Field #### The first structure is the special column marked as "Calculated Field". You need to interpret the purpose and calculation basis for these columns, then utilize them in the following text-to-sql generation tasks. -First, provide a brief explanation of what each field represents in the context of the schema, including how each field is computed using the relationships between models. -Then, during the following tasks, if the user queries pertain to any calculated fields defined in the database schema, ensure to utilize those calculated fields appropriately in the output SQL queries. -The goal is to accurately reflect the intent of the question in the SQL syntax, leveraging the pre-computed logic embedded within the calculated fields. - -EXAMPLES: -The given schema is created by the SQL command: - -CREATE TABLE orders ( - OrderId VARCHAR PRIMARY KEY, - CustomerId VARCHAR, - -- This column is a Calculated Field - -- column expression: avg(reviews.Score) - Rating DOUBLE, - -- This column is a Calculated Field - -- column expression: count(reviews.Id) - ReviewCount BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) - Size BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) > 1 - Large BOOLEAN, - FOREIGN KEY (CustomerId) REFERENCES customers(Id) -); - -Interpret the columns that are marked as Calculated Fields in the schema: -Rating (DOUBLE) - Calculated as the average score (avg) of the Score field from the reviews table where the reviews are associated with the order. This field represents the overall customer satisfaction rating for the order based on review scores. -ReviewCount (BIGINT) - Calculated by counting (count) the number of entries in the reviews table associated with this order. It measures the volume of customer feedback received for the order. -Size (BIGINT) - Represents the total number of items in the order, calculated by counting the number of item entries (ItemNumber) in the order_items table linked to this order. This field is useful for understanding the scale or size of an order. -Large (BOOLEAN) - A boolean value calculated to check if the number of items in the order exceeds one (count(order_items.ItemNumber) > 1). It indicates whether the order is considered large in terms of item quantity. - -And if the user input queries like these: -1. "How many large orders have been placed by customer with ID 'C1234'?" -2. "What is the average customer rating for orders that were rated by more than 10 reviewers?" - -For the first query: -First try to intepret the user query, the user wants to know the average rating for orders which have attracted significant review activity, specifically those with more than 10 reviews. -Then, according to the above intepretation about the given schema, the term 'Rating' is predefined in the Calculated Field of the 'orders' model. And, the number of reviews is also predefined in the 'ReviewCount' Calculated Field. -So utilize those Calculated Fields in the SQL generation process to give an answer like this: - -SQL Query: SELECT AVG(Rating) FROM orders WHERE ReviewCount > 10 +First, interpret each calculated field from its expression, data type, comments, aliases, descriptions, and relationship context in the provided DATABASE SCHEMA. +Then, if the user query matches a concept already represented by a calculated field, use that exact calculated field name from DATABASE SCHEMA instead of recreating or inventing the calculation. +Calculated field expressions are semantic definitions; do not copy identifiers from an expression unless they also appear as executable identifiers in the current DATABASE SCHEMA. """ _DEFAULT_METRIC_INSTRUCTIONS = """ @@ -397,68 +379,7 @@ async def _classify_generation_result( If the given schema contains the structures marked as 'metric', you should first interpret the metric schema based on the above definition. Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. - -EXAMPLES: -The given schema is created by the SQL command: - -/* This table is a metric */ -/* Metric Base Object: orders */ -CREATE TABLE Revenue ( - -- This column is a dimension - PurchaseTimestamp TIMESTAMP, - -- This column is a dimension - CustomerId VARCHAR, - -- This column is a dimension - Status VARCHAR, - -- This column is a measure - -- expression: sum(order_items.Price) - PriceSum DOUBLE, - -- This column is a measure - -- expression: count(OrderId) - NumberOfOrders BIGINT -); - -Interpret the metric with the understanding of the metric structure: -1. Base Object: orders -This is the primary data source for the metric. -The orders table provides the underlying data from which dimensions and measures are derived. -It is the foundation upon which the metric is built, though it itself is not directly used in queries against the Revenue table. -It shows the reference between the 'Revenue' metric and the 'orders' model. For the user queries pretain to the 'Revenue' of 'orders', the metric should be utilize in the sql generation process. -2. Dimensions -The metric contains the columns marked as 'dimension'. They can be interpreted as below: -- PurchaseTimestamp (TIMESTAMP) - Acts as a temporal dimension, allowing analysis of revenue over time. This can be used to observe trends, seasonal variations, or performance over specific periods. -- CustomerId (VARCHAR) - A key dimension for customer segmentation, it enables the analysis of revenue generated from individual customers or customer groups. -- Status (VARCHAR) - Reflects the current state of an order (e.g., pending, completed, cancelled). This dimension is crucial for analyses that differentiate performance based on order status. -3. Measures -The metric contains the columns marked as 'measure'. They can be interpreted as below: -- PriceSum (DOUBLE) - A financial measure calculated as sum(order_items.Price), representing the total revenue generated from orders. This measure is vital for tracking overall sales performance and is the primary output of interest in many financial and business analyses. -- NumberOfOrders (BIGINT) - A count measure that provides the total number of orders. This is essential for operational metrics, such as assessing the volume of business activity and evaluating the efficiency of sales processes. - -Now, if the user input queries like this: -Question: "What was the total revenue from each customer last month?" - -First try to intepret the user query, the user asks for a breakdown of the total revenue generated by each customer in the previous calendar month. -The user is specifically interested in understanding how much each customer contributed to the total sales during this period. -To answer this question, it is suitable to use the following components from the metric: -1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. -2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. -3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. -So utilize those metric components in the SQL generation process to give an answer like this: - -SQL Query: -SELECT - CustomerId, - PriceSum AS TotalRevenue -FROM - Revenue -WHERE - PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND - PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) +Use metric columns exactly as declared in DATABASE SCHEMA. Treat dimensions as grouping/filtering fields and measures as pre-defined numeric outputs. Metric base objects and measure expressions are semantic context only; do not copy identifiers from them unless those identifiers also appear as executable identifiers in the current DATABASE SCHEMA. """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -469,31 +390,13 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - For Example: - DATA SCHEMA: - `/* {"alias":"users","description":"A model representing the users data."} */ - CREATE TABLE users ( - -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} - address JSON - )` - To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` + - JSON paths and nested field names must come from the json_fields metadata attached to the exact JSON column in DATABASE SCHEMA. - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. - - For Example: - DATA SCHEMA - `/* {"alias":"my_table","description":"A test my_table"} */ - CREATE TABLE my_table ( - -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} - elements JSON - )` - To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` + - Do not copy JSON examples, placeholder aliases, or nested paths from prior context. Use only the current table name, JSON column name, and json_fields metadata in DATABASE SCHEMA. - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". - DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". """ @@ -501,40 +404,34 @@ async def _classify_generation_result( sql_samples_instructions = """ #### Instructions for SQL Samples #### -Finally, you will learn from the sample SQL queries provided in the input. These samples demonstrate best practices and common patterns for querying this specific database. +Finally, you will learn from the sample questions provided in the input. These samples demonstrate intent and response style for this specific database. For each sample, you should: 1. Study the question that explains what the query aims to accomplish -2. Analyze the SQL implementation to understand: - - Table structures and relationships used - - Specific functions and operators employed - - Query patterns and techniques demonstrated -3. Use these samples as reference patterns when generating similar queries -4. Adapt the techniques shown in the samples to match new query requirements while maintaining consistent style and approach +2. Use these samples as intent and style context only, but treat the DATABASE SCHEMA as the only valid source of executable table and column names +3. Adapt the intent patterns to match new query requirements while maintaining consistent style and approach +4. Never copy table names, column names, aliases, literal values, placeholders, or functions from samples The samples will help you understand: -- Preferred table join patterns -- Common aggregation methods -- Specific function usage -- Query structure and formatting conventions +- Common analytical intents +- Common aggregation requests +- Preferred answer style -When generating new queries, try to follow similar patterns when applicable, while adapting them to the specific requirements of each new query. +When generating new queries, follow similar intent patterns when applicable, while adapting them to the specific requirements of each new query. -Learn about the usage of the schema structures and generate SQL based on them. +Learn about the user's intent from the samples and generate SQL from the current DATABASE SCHEMA and SQL FUNCTIONS only. """ sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who is great at thinking deeply and reasoning about the user's question and the database schema, and you provide a step-by-step reasoning plan in order to answer the user's question. +You are a helpful data analyst who explains the user's analytical intent and provides a concise, non-executable reasoning plan for answering the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; -otherwise, you will put the relative timeframe in the SQL query. -3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. +2. Explicitly state requested timeframes in natural language only. Mention exact date/time columns only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language. Mention exact ordering columns or measures only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. @@ -542,9 +439,22 @@ async def _classify_generation_result( 9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. A table name in the reasoning plan must be in this format: `table: `. -13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +12. Mention table names only by writing the literal prefix `table:` followed by an exact table name declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +13. Mention column names only by writing the literal prefix `column:` followed by an exact declared table name, a dot, and an exact column name declared for that table in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +14. Do not mention aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels from comments, SQL samples, failed SQL, or user wording as executable identifiers. +15. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. +16. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. +17. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language and cite exact declared tables or columns only when they are grounded by DATABASE SCHEMA. +18. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, then ground the plan in exact declared schema identifiers. +19. If multiple schema objects are required, identify the exact declared relationship path from DATABASE SCHEMA. If no relationship path is declared, say that the retrieved metadata does not provide a join path. +20. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan unless they also appear exactly in DATABASE SCHEMA. +21. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +22. The reasoning plan is semantic context for intent only, not a source of executable identifiers. SQL generation must re-read DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT before using any identifier. +23. ONLY SHOWING the reasoning plan in bullet points. +24. Do not use the words "assume", "assuming", "likely", "possible", "might", or "example" when describing tables, columns, filters, or SQL. +25. If exact deployed table and column identifiers are not available for a requested part, say only that the retrieved metadata does not support that part. Do not propose a replacement name. +26. Do not write table names or column names from the user's wording unless the same identifier appears exactly in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +27. Do not include code blocks, inline SQL fragments, SELECT statements, WHERE clauses, join clauses, or any query-shaped text in the reasoning plan. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -568,7 +478,7 @@ def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) - return rules + return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: @@ -604,31 +514,40 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" -You are a helpful assistant that converts natural language queries into ANSI SQL queries. +You are a helpful assistant that converts natural language queries into Wren SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. +Given the user's question and database schema, generate one grounded Wren SQL query. The DATABASE SCHEMA is the only source of executable identifiers. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. -2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. +3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. +4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. +5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. +6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. +7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. +8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. +11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +13. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. +14. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. Do not create table or column identifiers from the user's wording. If the retrieved schema does not ground the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ class SqlGenerationResult(BaseModel): - sql: str + sql: str | None SQL_GENERATION_MODEL_KWARGS = { diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 22bf940cd8..68c5325d36 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -9,9 +9,9 @@ from hamilton.async_driver import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe from pydantic import BaseModel +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider from src.pipelines.common import ( @@ -30,7 +30,13 @@ ### TASK ### You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. -The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. +The database schema includes structural, semantic, and business modeling metadata: +- Models are logical datasets backed by physical tables or SQL definitions. +- Columns are exposed fields, including renamed fields, expressions, primary keys, and calculated fields. +- Relationships are reusable join logic between models. +- Calculated fields are business logic defined once and reused across queries. +- Views are named SQL statements that behave like stable virtual tables. +- Metrics are structured aggregation objects with measures and dimensions. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -40,6 +46,17 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +8. Map the business question to the modeled datasets whose descriptions, aliases, columns, calculated fields, views, metrics, and relationships support the intent. +9. Prefer modeled analytical interfaces such as views and metrics when they expose the fields needed to answer the question. +10. If the answer needs fields, filters, time dimensions, ordering, aggregations, or relationship keys from multiple related datasets, include every required related dataset and the columns needed from each one. +11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. +12. Follow only the relationships shown in the provided schema when selecting columns across datasets. +13. Do not stop at a single top candidate when the question needs multiple related datasets. +14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. +15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. +16. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. +17. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. +18. Select the tables, views, metrics, relationships, and columns that best support the current question from the available modeled schema. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -112,37 +129,279 @@ def _project_filter_conditions( def _build_metric_ddl(content: dict) -> str: - columns_ddl = [ - f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + columns = [ + column for column in content["columns"] - if column["data_type"].lower() - != "unknown" # quick fix: filtering out UNKNOWN column type + if column["data_type"].lower() != "unknown" + ] + context = _format_semantic_context( + { + "object_type": "metric", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [column["name"] for column in columns], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable analytical aggregation interface", + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "semantic_context_not_sql_identifier": column["comment"], + } + for column in columns + ], + } + ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in columns ] return ( - f"{content['comment']}CREATE TABLE {content['name']} (\n " + f"{context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) def _build_view_ddl(content: dict) -> str: + columns = [ + column + for column in content.get("columns", []) + if column.get("name") and column.get("data_type", "").lower() != "unknown" + ] + context = _format_semantic_context( + { + "object_type": "view", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [column["name"] for column in columns], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable virtual table interface", + "description": content["comment"], + "definition_omitted_from_executable_schema": True, + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type( + column.get("data_type") + ), + "semantic_context_not_sql_identifier": column.get("comment", ""), + } + for column in columns + ], + } + ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" + for column in columns + ] + + return ( + f"{context}CREATE TABLE {content['name']} (\n " + + ",\n ".join(columns_ddl) + + "\n);" + ) + + +def _format_semantic_context(context: dict) -> str: return ( - f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" + "/*\n" + "WREN RETRIEVED SEMANTIC CONTEXT\n" + f"{orjson.dumps(context).decode('utf-8')}\n" + f"{_format_identifier_contract(context)}" + "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" + "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" + "*/\n" + f"{_format_executable_identifier_catalog(context)}" + ) + + +def _format_executable_identifier_catalog(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "### EXECUTABLE WREN IDENTIFIER CATALOG ###", + "Copy SQL identifiers only from this catalog or the following DDL.", + "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"table: {table_name}") + if column_names: + lines.append("columns:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationships:") + lines.extend(f"- {constraint}" for constraint in relationship_constraints) + lines.extend( + [ + "If a needed table, column, or relationship is not listed here or declared in the following DDL, return null for sql.", + "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", + "", + ] + ) + return "\n".join(lines) + + +def _format_identifier_contract(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "WREN SQL IDENTIFIER CONTRACT", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"sql_table_name_use_exactly: {table_name}") + if column_names: + lines.append("sql_column_names_use_exactly:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationship_constraints_use_exactly:") + lines.extend( + f"- {relationship_constraint}" + for relationship_constraint in relationship_constraints + ) + lines.extend( + [ + "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", + "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", + "END WREN SQL IDENTIFIER CONTRACT", + "", + ] ) + return "\n".join(lines) + + +def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: + relationship_columns = { + column.get("column") + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + } + relationship_columns.discard(None) + return relationship_columns + + +def _included_columns( + content: dict, columns: Optional[set[str]], tables: Optional[set[str]] +) -> list[dict]: + relationship_columns = _included_relationship_columns(content, tables) + return [ + column + for column in content["columns"] + if column["type"] == "COLUMN" + and ( + not columns + or column["name"] in columns + or column["name"] in relationship_columns + or column["is_primary_key"] + ) + and column["data_type"].lower() != "unknown" + ] + + +def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: + return [ + column + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + ] + + +def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: + executable_columns = { + column["name"] + for column in content["columns"] + if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" + } + return bool(columns) and columns.issubset(executable_columns) + + +def _build_table_retrieval_context( + content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None +) -> tuple[str, bool, bool]: + ddl, has_calculated_field, has_json_field = build_table_ddl( + content, + columns=columns, + tables=tables, + include_semantic_comments=False, + ) + included_columns = _included_columns(content, columns, tables) + included_relationships = _included_relationships(content, tables) + context = _format_semantic_context( + { + "object_type": "model", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in included_columns + ], + "relationship_constraints_use_exactly": [ + relationship["constraint"] + for relationship in included_relationships + ], + }, + "semantic_context_not_sql_identifiers": { + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "is_primary_key": column["is_primary_key"], + "semantic_context_not_sql_identifier": column["comment"], + } + for column in included_columns + ], + "relationships": [ + { + "semantic_context_not_sql_identifier": relationship["comment"], + "sql_relationship_constraint_use_exactly": relationship[ + "constraint" + ], + "related_models_use_exactly": relationship.get("tables", []), + } + for relationship in included_relationships + ], + } + ) + return f"{context}{ddl}", has_calculated_field, has_json_field ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: - if histories: - previous_query_summaries = [history.question for history in histories] - else: - previous_query_summaries = [] - - query = "\n".join(previous_query_summaries) + "\n" + query - return await embedder.run(query) else: return {} @@ -186,34 +445,149 @@ async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any, + embedding: dict, mdl_hash: str | None = None, ) -> list[Document]: - tables = table_retrieval.get("documents", []) + table_names = _table_names_from_description_documents( + table_retrieval.get("documents", []) + ) + documents = [] + if embedding and not table_names: + documents = await _retrieve_semantic_schema_documents( + embedding, project_id, dbschema_retriever, mdl_hash + ) + table_names = _table_names_from_schema_documents(documents) + + if table_names: + retrieved_table_names = set() + pending_table_names = table_names + + while pending_table_names: + retrieved_table_names.update(pending_table_names) + retrieved_documents = await _retrieve_schema_documents( + pending_table_names, project_id, dbschema_retriever, mdl_hash + ) + documents = _dedupe_documents(documents + retrieved_documents) + pending_table_names = [ + table_name + for table_name in _related_table_names(documents) + if table_name not in retrieved_table_names + ] + + return documents + + return [] + + +async def _retrieve_semantic_schema_documents( + embedding: dict, + project_id: str, + dbschema_retriever: Any, + mdl_hash: str | None = None, +) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + + results = await dbschema_retriever.run( + query_embedding=embedding.get("embedding"), + filters=filters, + ) + return results["documents"] + + +def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: + table_names = [] + seen = set() + + for document in documents: + table_name = document.meta.get("name") + if not table_name: + content = ast.literal_eval(document.content) + table_name = content.get("name") + + if table_name and table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + + return table_names + + +def _table_names_from_description_documents(documents: list[Document]) -> list[str]: table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + table_name = content["name"] + if table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + return table_names + + +async def _retrieve_schema_documents( + table_names: list[str], + project_id: str, + dbschema_retriever: Any, + mdl_hash: str | None = None, +) -> list[Document]: table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names ] - if table_name_conditions: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } + if not table_name_conditions: + return [] - filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) - return [] + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] + + +def _dedupe_documents(documents: list[Document]) -> list[Document]: + deduped = {} + for document in documents: + key = (document.meta.get("name"), document.content) + deduped[key] = document + + return list(deduped.values()) + + +def _related_table_names(documents: list[Document]) -> list[str]: + related_table_names = [] + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + if content.get("type") != "TABLE_COLUMNS": + continue + + for column in content.get("columns", []): + if column.get("type") != "FOREIGN_KEY": + continue + + for table_name in column.get("tables", []): + if table_name not in seen: + related_table_names.append(table_name) + seen.add(table_name) + + return related_table_names @observe() @@ -259,7 +633,9 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) retrieval_results.append( { "table_name": table_schema["name"], @@ -322,16 +698,10 @@ def prompt( ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ - build_table_ddl(construct_db_schema)[0] + _build_table_retrieval_context(construct_db_schema)[0] for construct_db_schema in construct_db_schemas ] - previous_query_summaries = ( - [history.question for history in histories] if histories else [] - ) - - query = "\n".join(previous_query_summaries) + "\n" + query - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: @@ -377,12 +747,21 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - ddl, _has_calculated_field, _has_json_field = build_table_ddl( - table_schema, - columns=set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ), - tables=tables, + columns = set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ) + columns = ( + columns + if _selected_columns_are_executable(table_schema, columns) + else None + ) + + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context( + table_schema, + columns=columns, + tables=tables, + ) ) if _has_calculated_field: has_calculated_field = True @@ -397,24 +776,23 @@ def construct_retrieval_results( ) for document in dbschema_retrieval: - if document.meta["name"] in columns_and_tables_needed: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - } - ) + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) return { "retrieval_results": retrieval_results, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py index a775453270..2c2dfa775e 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -97,6 +97,38 @@ async def test_post_processor_returns_no_relevant_sql_for_missing_sql_field(): assert "supported SQL field" in result["invalid_generation_result"]["error"] +@pytest.mark.asyncio +async def test_post_processor_treats_null_sql_as_no_relevant_sql(): + processor = SQLGenPostProcessor(FakeEngine()) + + result = await processor.run( + ['{"sql": null}'], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert "No grounded SQL" in result["invalid_generation_result"]["error"] + + +@pytest.mark.asyncio +async def test_post_processor_rejects_code_tool_payload(): + engine = FakeEngine() + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ['{"name":"execute_code","arguments":{"code":"SELECT * FROM orders"}}'], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + @pytest.mark.asyncio async def test_post_processor_dry_plans_before_preview_execution(): engine = FakeEngine(dry_plan_success=False) From 09f9489edd4eb651770f0d569b3239ceac42c2ff Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 22:01:02 +0530 Subject: [PATCH 1036/1087] Restore legacy grounded ask flow --- .../src/pipelines/generation/utils/sql.py | 1 + wren-ai-service/src/web/v1/services/ask.py | 21 ++++++++----------- .../test_sql_generation_post_processor.py | 10 ++++++++- .../tests/pytest/services/test_ask.py | 12 +++++++++++ 4 files changed, 31 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 4fa97f2632..2e0d9cf9c8 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -551,6 +551,7 @@ class SqlGenerationResult(BaseModel): SQL_GENERATION_MODEL_KWARGS = { + "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index cc5151aba0..644527f572 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -3,9 +3,9 @@ from typing import Dict, List, Literal, Optional from cachetools import TTLCache -from langfuse.decorators import observe from pydantic import AliasChoices, BaseModel, Field +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -25,10 +25,10 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = False + ignore_sql_generation_reasoning: bool = True enable_column_pruning: bool = False - use_dry_plan: bool = False - allow_dry_plan_fallback: bool = True + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False custom_instruction: Optional[str] = None @@ -99,12 +99,12 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = True, + allow_sql_generation_reasoning: bool = False, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = False, - max_sql_correction_retries: int = 3, + enable_column_pruning: bool = True, + max_sql_correction_retries: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -161,17 +161,14 @@ async def ask( table_names = [] error_message = None invalid_sql = None - allow_sql_generation_reasoning = ( - self._allow_sql_generation_reasoning - and not ask_request.ignore_sql_generation_reasoning - ) + allow_sql_generation_reasoning = False enable_column_pruning = ( self._enable_column_pruning or ask_request.enable_column_pruning ) allow_sql_functions_retrieval = self._allow_sql_functions_retrieval allow_sql_diagnosis = self._allow_sql_diagnosis allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - max_sql_correction_retries = self._max_sql_correction_retries + max_sql_correction_retries = 0 current_sql_correction_retries = 0 use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py index 2c2dfa775e..a69a039702 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -4,7 +4,10 @@ import pytest from src.core.engine import Engine -from src.pipelines.generation.utils.sql import SQLGenPostProcessor +from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, + SQLGenPostProcessor, +) class FakeEngine(Engine): @@ -52,6 +55,11 @@ async def execute_sql( return self.execute_success, {}, {"correlation_id": "correlation-id"} +def test_sql_generation_model_kwargs_preserve_strict_schema(): + assert SQL_GENERATION_MODEL_KWARGS["preserve_json_schema"] is True + assert SQL_GENERATION_MODEL_KWARGS["response_format"]["type"] == "json_schema" + + @pytest.mark.asyncio async def test_post_processor_extracts_tool_call_query_argument(): engine = FakeEngine() diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 283d7da34a..5a930ac264 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -19,6 +19,18 @@ ) +def test_ask_defaults_follow_legacy_grounded_sql_flow(): + request = AskRequest(query="How many invoices are there?", mdl_hash="mdl-hash") + service = AskService({}) + + assert request.ignore_sql_generation_reasoning is True + assert request.use_dry_plan is True + assert request.allow_dry_plan_fallback is False + assert service._allow_sql_generation_reasoning is False + assert service._enable_column_pruning is True + assert service._max_sql_correction_retries == 0 + + @pytest.fixture def ask_service(): pipe_components = generate_components(settings.components) From ca5ea6383e3e4d5286257c31716cc28b46d76b6f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 22:13:54 +0530 Subject: [PATCH 1037/1087] Tighten grounded SQL generation contract --- .../generation/followup_sql_generation.py | 21 ++++--- .../pipelines/generation/sql_generation.py | 10 ++++ .../src/pipelines/generation/utils/sql.py | 26 +++------ .../retrieval/db_schema_retrieval.py | 55 +++++++++++++++++++ wren-ai-service/src/web/v1/services/ask.py | 7 +++ .../test_sql_generation_post_processor.py | 46 ++++++++++++---- .../tests/pytest/services/test_ask.py | 2 +- 7 files changed, 132 insertions(+), 35 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 3e8df185e8..4fbefd8931 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -34,6 +34,14 @@ Given the user's current follow-up question and the current retrieved DATABASE SCHEMA, generate one SQL query to best answer the user's question. +{% if validation_contexts %} +### VALID WREN SQL IDENTIFIERS ### +Copy executable table, column, and relationship identifiers only from this section or the DATABASE SCHEMA below. Preserve each identifier exactly, including prefixes, spaces, digits, underscores, case, and punctuation. Semantic descriptions, aliases, source names, physical names, and user wording are not executable identifiers. +{% for validation_context in validation_contexts %} +{{ validation_context }} +{% endfor %} +{% endif %} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -60,11 +68,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -77,12 +84,10 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} -{% endif %} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, VALID WREN SQL IDENTIFIERS, or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -Let's think step by step. +Return only the final JSON SQL response. """ @@ -93,6 +98,7 @@ def prompt( documents: list[str], sql_generation_reasoning: str, prompt_builder: PromptBuilder, + validation_contexts: list[str] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -104,6 +110,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, + validation_contexts=validation_contexts or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 807bf73522..1070e1b0d1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,6 +28,14 @@ sql_generation_user_prompt_template = """ +{% if validation_contexts %} +### VALID WREN SQL IDENTIFIERS ### +Copy executable table, column, and relationship identifiers only from this section or the DATABASE SCHEMA below. Preserve each identifier exactly, including prefixes, spaces, digits, underscores, case, and punctuation. Semantic descriptions, aliases, source names, physical names, and user wording are not executable identifiers. +{% for validation_context in validation_contexts %} +{{ validation_context }} +{% endfor %} +{% endif %} + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -85,6 +93,7 @@ def prompt( documents: list[str], prompt_builder: PromptBuilder, sql_generation_reasoning: str | None = None, + validation_contexts: list[str] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -96,6 +105,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, + validation_contexts=validation_contexts or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 2e0d9cf9c8..f0ec305b32 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -5,7 +5,7 @@ import orjson from haystack import component from haystack.dataclasses import ChatMessage -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from src.core.engine import ( Engine, @@ -85,7 +85,10 @@ def _extract_sql(self, replies: List[str] | List[List[str]]) -> tuple[str, str]: f"{generation_result}", ) - return cleaned_generation_result, "" + if self._looks_like_sql(cleaned_generation_result): + return cleaned_generation_result, "" + + return "", "SQL generation response was not a supported SQL JSON payload." @component.output_types( valid_generation_result=Dict[str, Any], @@ -547,6 +550,8 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) class SqlGenerationResult(BaseModel): + model_config = ConfigDict(extra="forbid") + sql: str | None @@ -556,6 +561,7 @@ class SqlGenerationResult(BaseModel): "type": "json_schema", "json_schema": { "name": "sql_generation_result", + "strict": True, "schema": SqlGenerationResult.model_json_schema(), }, } @@ -577,18 +583,4 @@ def construct_instructions( def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: - messages = [] - for history in histories: - messages.append( - ChatMessage.from_user( - history.question - if hasattr(history, "question") - else history["question"] - ) - ) - messages.append( - ChatMessage.from_assistant( - history.sql if hasattr(history, "sql") else history["sql"] - ) - ) - return messages + return [] diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 68c5325d36..f73b9c579f 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -300,6 +300,47 @@ def _format_identifier_contract(context: dict) -> str: return "\n".join(lines) +def _format_prompt_identifier_context( + table_name: str, + column_names: list[str], + relationship_constraints: list[str] | None = None, +) -> str: + lines = [ + f"table: {table_name}", + "columns:", + ] + lines.extend(f"- {column_name}" for column_name in column_names) + + if relationship_constraints: + lines.append("relationships:") + lines.extend(f"- {constraint}" for constraint in relationship_constraints) + + return "\n".join(lines) + + +def _table_identifier_context( + content: dict, + columns: Optional[set[str]] = None, + tables: Optional[set[str]] = None, +) -> str: + included_columns = _included_columns(content, columns, tables) + included_relationships = _included_relationships(content, tables) + return _format_prompt_identifier_context( + content["name"], + [column["name"] for column in included_columns], + [relationship["constraint"] for relationship in included_relationships], + ) + + +def _semantic_object_identifier_context(content: dict) -> str: + columns = [ + column["name"] + for column in content.get("columns", []) + if column.get("name") and column.get("data_type", "").lower() != "unknown" + ] + return _format_prompt_identifier_context(content["name"], columns) + + def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: relationship_columns = { column.get("column") @@ -640,6 +681,7 @@ def check_using_db_schemas_without_pruning( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_context": _table_identifier_context(table_schema), } ) if _has_calculated_field: @@ -655,6 +697,7 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_context": _semantic_object_identifier_context(content), } ) has_metric = True @@ -663,6 +706,7 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_context": _semantic_object_identifier_context(content), } ) @@ -772,6 +816,11 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_context": _table_identifier_context( + table_schema, + columns=columns, + tables=tables, + ), } ) @@ -783,6 +832,9 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_context": _semantic_object_identifier_context( + content + ), } ) has_metric = True @@ -791,6 +843,9 @@ def construct_retrieval_results( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_context": _semantic_object_identifier_context( + content + ), } ) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 644527f572..8e85e98263 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -357,6 +357,11 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] + identifier_contexts = [ + document.get("identifier_context") + for document in documents + if document.get("identifier_context") + ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -477,6 +482,7 @@ async def ask( histories=histories, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, + validation_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -496,6 +502,7 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, + validation_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py index a69a039702..7e1be2ca7a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -7,6 +7,7 @@ from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, + construct_ask_history_messages, ) @@ -58,6 +59,15 @@ async def execute_sql( def test_sql_generation_model_kwargs_preserve_strict_schema(): assert SQL_GENERATION_MODEL_KWARGS["preserve_json_schema"] is True assert SQL_GENERATION_MODEL_KWARGS["response_format"]["type"] == "json_schema" + assert SQL_GENERATION_MODEL_KWARGS["response_format"]["json_schema"]["strict"] is True + schema = SQL_GENERATION_MODEL_KWARGS["response_format"]["json_schema"]["schema"] + assert schema["additionalProperties"] is False + + +def test_construct_ask_history_messages_matches_legacy_empty_context(): + histories = [{"question": "q", "sql": "SELECT 1"}] + + assert construct_ask_history_messages(histories) == [] @pytest.mark.asyncio @@ -66,21 +76,19 @@ async def test_post_processor_extracts_tool_call_query_argument(): processor = SQLGenPostProcessor(engine) result = await processor.run( - [ - '{"name":"query","arguments":{"query":"SELECT supplierid, COUNT(*) FROM PO_Invoices GROUP BY supplierid;"}}' - ], + ['{"name":"query","arguments":{"query":"SELECT 1"}}'], project_id="project-id", mdl_hash="manifest-hash", data_source="mssql", ) assert result["valid_generation_result"] == { - "sql": "SELECT supplierid, COUNT(*) FROM PO_Invoices GROUP BY supplierid", + "sql": "SELECT 1", "correlation_id": "correlation-id", } assert engine.dry_plan_calls == [ { - "sql": "SELECT supplierid, COUNT(*) FROM PO_Invoices GROUP BY supplierid", + "sql": "SELECT 1", "data_source": "mssql", "project_id": "project-id", "mdl_hash": "manifest-hash", @@ -95,7 +103,7 @@ async def test_post_processor_returns_no_relevant_sql_for_missing_sql_field(): processor = SQLGenPostProcessor(FakeEngine()) result = await processor.run( - ['{"name":"query","arguments":{"question":"Show suppliers"}}'], + ['{"name":"query","arguments":{"value":"q"}}'], project_id="project-id", data_source="mssql", ) @@ -126,13 +134,31 @@ async def test_post_processor_rejects_code_tool_payload(): processor = SQLGenPostProcessor(engine) result = await processor.run( - ['{"name":"execute_code","arguments":{"code":"SELECT * FROM orders"}}'], + ['{"name":"execute_code","arguments":{"code":"SELECT 1"}}'], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + +@pytest.mark.asyncio +async def test_post_processor_rejects_plain_text_non_sql_response(): + engine = FakeEngine() + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ["q"], project_id="project-id", data_source="mssql", ) assert result["valid_generation_result"] == {} assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert "supported SQL JSON payload" in result["invalid_generation_result"]["error"] assert engine.dry_plan_calls == [] assert engine.execute_sql_calls == [] @@ -143,7 +169,7 @@ async def test_post_processor_dry_plans_before_preview_execution(): processor = SQLGenPostProcessor(engine) result = await processor.run( - ['{"sql":"SELECT * FROM orders"}'], + ['{"sql":"SELECT 1"}'], project_id="project-id", mdl_hash="manifest-hash", data_source="mssql", @@ -151,8 +177,8 @@ async def test_post_processor_dry_plans_before_preview_execution(): assert result["valid_generation_result"] == {} assert result["invalid_generation_result"] == { - "sql": "SELECT * FROM orders", - "original_sql": "SELECT * FROM orders", + "sql": "SELECT 1", + "original_sql": "SELECT 1", "type": "DRY_PLAN", "error": "plan failed", "correlation_id": "", diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 5a930ac264..21be147024 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -20,7 +20,7 @@ def test_ask_defaults_follow_legacy_grounded_sql_flow(): - request = AskRequest(query="How many invoices are there?", mdl_hash="mdl-hash") + request = AskRequest(query="q", mdl_hash="mdl-hash") service = AskService({}) assert request.ignore_sql_generation_reasoning is True From 4c80dbfc74053506c89d0ea93cfd10864e2ea64f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 22:38:53 +0530 Subject: [PATCH 1038/1087] Validate generated SQL against retrieved schema --- .../generation/followup_sql_generation.py | 2 + .../pipelines/generation/sql_correction.py | 38 +++-- .../pipelines/generation/sql_generation.py | 2 + .../pipelines/generation/sql_regeneration.py | 28 ++-- .../src/pipelines/generation/utils/sql.py | 144 ++++++++++++++++++ .../test_sql_generation_post_processor.py | 19 +++ 6 files changed, 206 insertions(+), 27 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 4fbefd8931..d2a0ca1151 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -157,6 +157,7 @@ async def post_process( data_source: str, project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: @@ -164,6 +165,7 @@ async def post_process( generate_sql_in_followup.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + validation_contexts=validation_contexts, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index f424ec48cb..101c0c6dd7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,12 +30,20 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query. +You are a Wren SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. ### SQL CORRECTION INSTRUCTIONS ### -1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). -2. Then, generate the syntactically correct ANSI SQL query to correct the error. +1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. +2. Then, generate a syntactically correct Wren SQL query from the user's intent and the current DATABASE SCHEMA. +3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. +4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. +5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. +6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. +7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA. If the unsupported part is needed to answer the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql instead of substituting non-schema identifiers. +10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -43,10 +51,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": + "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ @@ -76,18 +84,18 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### QUESTION ### {% if query %} User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. {% endif %} -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} -{% endif %} -### ORIGINAL SQL QUERY ### -{{ invalid_generation_result.sql }} +### FAILED SQL ### +The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. + +### DRY-RUN DIAGNOSTIC ### +The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. -### ERROR MESSAGE ### -{{ invalid_generation_result.error }} +Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. -Let's think step by step. +Return only the final JSON SQL response. """ @@ -136,6 +144,7 @@ async def post_process( data_source: str, project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: @@ -143,6 +152,7 @@ async def post_process( generate_sql_correction.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + validation_contexts=validation_contexts, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1070e1b0d1..a0597e6fff 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -148,6 +148,7 @@ async def post_process( data_source: str, project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, @@ -156,6 +157,7 @@ async def post_process( generate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + validation_contexts=validation_contexts, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 543f68b5b5..3c526e0c04 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -5,8 +5,8 @@ from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider @@ -35,16 +35,19 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### You are a great ANSI SQL expert. Now you are given database schema and a user's question. -Generate a new SQL query that answers the user's question. -While generating the new SQL query, make sure to use the database schema and SQL rules. +Carefully review the user's question and current DATABASE SCHEMA, then generate a new SQL query that answers the user's intent. +The original SQL query and UI planning text are intentionally omitted from the prompt and must not be used as executable context. +While generating the new SQL query, make sure to use the database schema as the only source of executable table and column identifiers. +If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. +Treat physical/source/lineage names from the original SQL, reasoning, samples, comments, or descriptions as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a SQL query in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ @@ -76,11 +79,10 @@ def get_sql_regeneration_system_prompt( {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -93,14 +95,12 @@ def get_sql_regeneration_system_prompt( ### QUESTION ### User's Question: {{ query }} -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} -{% endif %} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +Regenerate with executable identifiers from the current DATABASE SCHEMA only. ### ORIGINAL SQL QUERY ### -{{ sql }} +The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. -Let's think step by step. +Return only the final JSON SQL response. """ @@ -165,11 +165,13 @@ async def post_process( post_processor: SQLGenPostProcessor, project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + validation_contexts=validation_contexts, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index f0ec305b32..3baa466531 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,11 +1,15 @@ import logging +from collections.abc import Iterable from typing import Any, Dict, List import aiohttp import orjson +import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel, ConfigDict +from sqlparse.sql import Identifier, IdentifierList, TokenList +from sqlparse.tokens import Keyword, Name from src.core.engine import ( Engine, @@ -90,6 +94,129 @@ def _extract_sql(self, replies: List[str] | List[List[str]]) -> tuple[str, str]: return "", "SQL generation response was not a supported SQL JSON payload." + def _allowed_tables(self, validation_contexts: list[str]) -> set[str]: + allowed_tables = set() + for context in validation_contexts: + for line in context.splitlines(): + if line.startswith("table: "): + allowed_tables.add(line.removeprefix("table: ").strip()) + return allowed_tables + + def _cte_names(self, statement: TokenList) -> set[str]: + names = set() + tokens = [ + token + for token in statement.tokens + if not token.is_whitespace + ] + for index, token in enumerate(tokens): + if token.normalized != "WITH": + continue + if index + 1 >= len(tokens): + return names + + cte_token = tokens[index + 1] + identifiers: Iterable[Identifier] + if isinstance(cte_token, IdentifierList): + identifiers = cte_token.get_identifiers() + elif isinstance(cte_token, Identifier): + identifiers = [cte_token] + else: + return names + + for identifier in identifiers: + name = identifier.get_name() + if name: + names.add(name) + return names + + return names + + def _identifier_name(self, token) -> str | None: + if isinstance(token, Identifier): + return token.get_real_name() or token.get_name() + if token.ttype in (Name, Keyword): + return token.value.strip('"') + return None + + def _table_identifiers(self, sql: str) -> set[str]: + tables = set() + stop_keywords = { + "WHERE", + "GROUP BY", + "ORDER BY", + "HAVING", + "LIMIT", + "UNION", + "EXCEPT", + "INTERSECT", + "WINDOW", + } + source_keywords = { + "FROM", + "JOIN", + "INNER JOIN", + "LEFT JOIN", + "LEFT OUTER JOIN", + "RIGHT JOIN", + "RIGHT OUTER JOIN", + "FULL JOIN", + "FULL OUTER JOIN", + "CROSS JOIN", + } + + for statement in sqlparse.parse(sql): + cte_names = self._cte_names(statement) + collecting_sources = False + for token in statement.tokens: + if token.is_whitespace: + continue + + if token.ttype in Keyword and token.normalized in stop_keywords: + collecting_sources = False + continue + + if token.ttype in Keyword and token.normalized in source_keywords: + collecting_sources = True + continue + + if not collecting_sources: + continue + + if isinstance(token, IdentifierList): + identifiers = token.get_identifiers() + else: + identifiers = [token] + + for identifier in identifiers: + name = self._identifier_name(identifier) + if name and name not in cte_names: + tables.add(name) + + return tables + + def _validate_sql_tables( + self, + sql: str, + validation_contexts: list[str] | None, + ) -> str: + if not validation_contexts: + return "" + + allowed_tables = self._allowed_tables(validation_contexts) + if not allowed_tables: + return "" + + referenced_tables = self._table_identifiers(sql) + ungrounded_tables = referenced_tables - allowed_tables + if ungrounded_tables: + return ( + "Generated SQL referenced table identifiers outside the retrieved " + "Wren schema: " + ", ".join(sorted(ungrounded_tables)) + ) + + return "" + @component.output_types( valid_generation_result=Dict[str, Any], invalid_generation_result=Dict[str, Any], @@ -103,6 +230,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + validation_contexts: list[str] | None = None, ) -> dict: try: generation_result, extraction_error = self._extract_sql(replies) @@ -119,6 +247,22 @@ async def run( }, } + validation_error = self._validate_sql_tables( + generation_result, + validation_contexts, + ) + if validation_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": generation_result, + "original_sql": generation_result, + "type": "NO_RELEVANT_SQL", + "error": validation_error, + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py index 7e1be2ca7a..3c21a9c1d4 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -163,6 +163,25 @@ async def test_post_processor_rejects_plain_text_non_sql_response(): assert engine.execute_sql_calls == [] +@pytest.mark.asyncio +async def test_post_processor_rejects_sql_with_unretrieved_table_before_engine_call(): + engine = FakeEngine() + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ['{"sql":"SELECT 1 FROM missing_model"}'], + project_id="project-id", + data_source="mssql", + validation_contexts=["table: allowed_model\ncolumns:\n- allowed_column"], + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert "missing_model" in result["invalid_generation_result"]["error"] + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + @pytest.mark.asyncio async def test_post_processor_dry_plans_before_preview_execution(): engine = FakeEngine(dry_plan_success=False) From 1f3016b65ec68826d35f4432b57f734f0587ff28 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Sun, 16 Aug 2026 23:49:56 +0530 Subject: [PATCH 1039/1087] Align SQL generation with retrieved schema context --- .../generation/followup_sql_generation.py | 24 ++- .../pipelines/generation/sql_correction.py | 4 - .../pipelines/generation/sql_generation.py | 25 ++- .../pipelines/generation/sql_regeneration.py | 4 - .../src/pipelines/generation/utils/sql.py | 144 ------------------ .../retrieval/db_schema_retrieval.py | 76 ++++++++- wren-ai-service/src/web/v1/services/ask.py | 4 +- .../v1/services/question_recommendation.py | 12 +- .../test_sql_generation_post_processor.py | 19 --- .../retrieval/test_db_schema_retrieval.py | 40 ++++- 10 files changed, 145 insertions(+), 207 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 d2a0ca1151..6a6d52cca0 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -34,11 +34,11 @@ Given the user's current follow-up question and the current retrieved DATABASE SCHEMA, generate one SQL query to best answer the user's question. -{% if validation_contexts %} -### VALID WREN SQL IDENTIFIERS ### -Copy executable table, column, and relationship identifiers only from this section or the DATABASE SCHEMA below. Preserve each identifier exactly, including prefixes, spaces, digits, underscores, case, and punctuation. Semantic descriptions, aliases, source names, physical names, and user wording are not executable identifiers. -{% for validation_context in validation_contexts %} -{{ validation_context }} +{% if identifier_contexts %} +### EXECUTABLE WREN IDENTIFIER CATALOG ### +The following catalog is derived from the retrieved DATABASE SCHEMA. Use it as a quick index of the exact executable model, column, and relationship names available for this request. +{% for identifier_context in identifier_contexts %} +{{ identifier_context }} {% endfor %} {% endif %} @@ -84,8 +84,8 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, VALID WREN SQL IDENTIFIERS, or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Generate one Wren SQL query that answers the user's follow-up question using the retrieved DATABASE SCHEMA. Use exact model and column names from DATABASE SCHEMA. Use aliases, descriptions, comments, samples, and previous questions only to understand business meaning. +If the retrieved DATABASE SCHEMA does not contain the model, column, relationship, metric, or supported function needed to answer the question, return {"sql": null}. Return only the final JSON SQL response. """ @@ -98,7 +98,7 @@ def prompt( documents: list[str], sql_generation_reasoning: str, prompt_builder: PromptBuilder, - validation_contexts: list[str] | None = None, + identifier_contexts: list[str] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -110,7 +110,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - validation_contexts=validation_contexts or [], + identifier_contexts=identifier_contexts or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -157,7 +157,6 @@ async def post_process( data_source: str, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: @@ -165,7 +164,6 @@ async def post_process( generate_sql_in_followup.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - validation_contexts=validation_contexts, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -214,7 +212,7 @@ async def run( instructions: list[dict] | None = None, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, + identifier_contexts: list[str] | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -241,7 +239,7 @@ async def run( "histories": histories, "project_id": project_id, "mdl_hash": mdl_hash, - "validation_contexts": validation_contexts, + "identifier_contexts": identifier_contexts, "sql_samples": sql_samples, "instructions": instructions, "has_calculated_field": has_calculated_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 101c0c6dd7..69aa9cd957 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -144,7 +144,6 @@ async def post_process( data_source: str, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: @@ -152,7 +151,6 @@ async def post_process( generate_sql_correction.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - validation_contexts=validation_contexts, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -201,7 +199,6 @@ async def run( sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, @@ -226,7 +223,6 @@ async def run( "sql_functions": sql_functions, "project_id": project_id, "mdl_hash": mdl_hash, - "validation_contexts": validation_contexts, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index a0597e6fff..c34c640ca4 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,11 +28,11 @@ sql_generation_user_prompt_template = """ -{% if validation_contexts %} -### VALID WREN SQL IDENTIFIERS ### -Copy executable table, column, and relationship identifiers only from this section or the DATABASE SCHEMA below. Preserve each identifier exactly, including prefixes, spaces, digits, underscores, case, and punctuation. Semantic descriptions, aliases, source names, physical names, and user wording are not executable identifiers. -{% for validation_context in validation_contexts %} -{{ validation_context }} +{% if identifier_contexts %} +### EXECUTABLE WREN IDENTIFIER CATALOG ### +The following catalog is derived from the retrieved DATABASE SCHEMA. Use it as a quick index of the exact executable model, column, and relationship names available for this request. +{% for identifier_context in identifier_contexts %} +{{ identifier_context }} {% endfor %} {% endif %} @@ -78,9 +78,8 @@ ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. +Generate one Wren SQL query that answers the user's question using the retrieved DATABASE SCHEMA. Use exact model and column names from DATABASE SCHEMA. Use aliases, descriptions, comments, and samples only to understand business meaning. +If the retrieved DATABASE SCHEMA does not contain the model, column, relationship, metric, or supported function needed to answer the question, return {"sql": null}. Return only the final JSON SQL response. """ @@ -93,7 +92,7 @@ def prompt( documents: list[str], prompt_builder: PromptBuilder, sql_generation_reasoning: str | None = None, - validation_contexts: list[str] | None = None, + identifier_contexts: list[str] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -105,7 +104,7 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - validation_contexts=validation_contexts or [], + identifier_contexts=identifier_contexts or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -148,7 +147,6 @@ async def post_process( data_source: str, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, @@ -157,7 +155,6 @@ async def post_process( generate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - validation_contexts=validation_contexts, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -206,7 +203,7 @@ async def run( instructions: list[dict] | None = None, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, + identifier_contexts: list[str] | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -235,7 +232,7 @@ async def run( "instructions": instructions, "project_id": project_id, "mdl_hash": mdl_hash, - "validation_contexts": validation_contexts, + "identifier_contexts": identifier_contexts, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 3c526e0c04..81a697dd9c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -165,13 +165,11 @@ async def post_process( post_processor: SQLGenPostProcessor, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - validation_contexts=validation_contexts, ) @@ -212,7 +210,6 @@ async def run( instructions: list[dict] | None = None, project_id: str | None = None, mdl_hash: str | None = None, - validation_contexts: list[str] | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -232,7 +229,6 @@ async def run( "instructions": instructions, "project_id": project_id, "mdl_hash": mdl_hash, - "validation_contexts": validation_contexts, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 3baa466531..f0ec305b32 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,15 +1,11 @@ import logging -from collections.abc import Iterable from typing import Any, Dict, List import aiohttp import orjson -import sqlparse from haystack import component from haystack.dataclasses import ChatMessage from pydantic import BaseModel, ConfigDict -from sqlparse.sql import Identifier, IdentifierList, TokenList -from sqlparse.tokens import Keyword, Name from src.core.engine import ( Engine, @@ -94,129 +90,6 @@ def _extract_sql(self, replies: List[str] | List[List[str]]) -> tuple[str, str]: return "", "SQL generation response was not a supported SQL JSON payload." - def _allowed_tables(self, validation_contexts: list[str]) -> set[str]: - allowed_tables = set() - for context in validation_contexts: - for line in context.splitlines(): - if line.startswith("table: "): - allowed_tables.add(line.removeprefix("table: ").strip()) - return allowed_tables - - def _cte_names(self, statement: TokenList) -> set[str]: - names = set() - tokens = [ - token - for token in statement.tokens - if not token.is_whitespace - ] - for index, token in enumerate(tokens): - if token.normalized != "WITH": - continue - if index + 1 >= len(tokens): - return names - - cte_token = tokens[index + 1] - identifiers: Iterable[Identifier] - if isinstance(cte_token, IdentifierList): - identifiers = cte_token.get_identifiers() - elif isinstance(cte_token, Identifier): - identifiers = [cte_token] - else: - return names - - for identifier in identifiers: - name = identifier.get_name() - if name: - names.add(name) - return names - - return names - - def _identifier_name(self, token) -> str | None: - if isinstance(token, Identifier): - return token.get_real_name() or token.get_name() - if token.ttype in (Name, Keyword): - return token.value.strip('"') - return None - - def _table_identifiers(self, sql: str) -> set[str]: - tables = set() - stop_keywords = { - "WHERE", - "GROUP BY", - "ORDER BY", - "HAVING", - "LIMIT", - "UNION", - "EXCEPT", - "INTERSECT", - "WINDOW", - } - source_keywords = { - "FROM", - "JOIN", - "INNER JOIN", - "LEFT JOIN", - "LEFT OUTER JOIN", - "RIGHT JOIN", - "RIGHT OUTER JOIN", - "FULL JOIN", - "FULL OUTER JOIN", - "CROSS JOIN", - } - - for statement in sqlparse.parse(sql): - cte_names = self._cte_names(statement) - collecting_sources = False - for token in statement.tokens: - if token.is_whitespace: - continue - - if token.ttype in Keyword and token.normalized in stop_keywords: - collecting_sources = False - continue - - if token.ttype in Keyword and token.normalized in source_keywords: - collecting_sources = True - continue - - if not collecting_sources: - continue - - if isinstance(token, IdentifierList): - identifiers = token.get_identifiers() - else: - identifiers = [token] - - for identifier in identifiers: - name = self._identifier_name(identifier) - if name and name not in cte_names: - tables.add(name) - - return tables - - def _validate_sql_tables( - self, - sql: str, - validation_contexts: list[str] | None, - ) -> str: - if not validation_contexts: - return "" - - allowed_tables = self._allowed_tables(validation_contexts) - if not allowed_tables: - return "" - - referenced_tables = self._table_identifiers(sql) - ungrounded_tables = referenced_tables - allowed_tables - if ungrounded_tables: - return ( - "Generated SQL referenced table identifiers outside the retrieved " - "Wren schema: " + ", ".join(sorted(ungrounded_tables)) - ) - - return "" - @component.output_types( valid_generation_result=Dict[str, Any], invalid_generation_result=Dict[str, Any], @@ -230,7 +103,6 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - validation_contexts: list[str] | None = None, ) -> dict: try: generation_result, extraction_error = self._extract_sql(replies) @@ -247,22 +119,6 @@ async def run( }, } - validation_error = self._validate_sql_tables( - generation_result, - validation_contexts, - ) - if validation_error: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": generation_result, - "original_sql": generation_result, - "type": "NO_RELEVANT_SQL", - "error": validation_error, - "correlation_id": "", - }, - } - ( valid_generation_result, invalid_generation_result, diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index f73b9c579f..90951859e7 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -773,9 +773,24 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] + try: + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + ).get("results") + except (IndexError, KeyError, orjson.JSONDecodeError, AttributeError) as e: + logger.warning( + f"Column pruning returned unusable output; using retrieved schemas without column pruning: {e}" + ) + columns_and_tables_needed = None + + if not isinstance(columns_and_tables_needed, list): + logger.warning( + "Column pruning output omitted results; using retrieved schemas without column pruning." + ) + return _build_retrieval_results_without_column_pruning( + construct_db_schemas, + dbschema_retrieval, + ) # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -868,6 +883,61 @@ def construct_retrieval_results( } +def _build_retrieval_results_without_column_pruning( + construct_db_schemas: list[dict], + dbschema_retrieval: list[Document], +) -> dict[str, Any]: + retrieval_results = [] + has_calculated_field = False + has_metric = False + has_json_field = False + + for table_schema in construct_db_schemas: + if table_schema["type"] == "TABLE": + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) + retrieval_results.append( + { + "table_name": table_schema["name"], + "table_ddl": ddl, + "identifier_context": _table_identifier_context(table_schema), + } + ) + if _has_calculated_field: + has_calculated_field = True + if _has_json_field: + has_json_field = True + + for document in dbschema_retrieval: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + "identifier_context": _semantic_object_identifier_context(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + "identifier_context": _semantic_object_identifier_context(content), + } + ) + + return { + "retrieval_results": retrieval_results, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + + ## End of Pipeline class MatchingTableContents(BaseModel): chain_of_thought_reasoning: list[str] diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 8e85e98263..413b8a1228 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -482,7 +482,7 @@ async def ask( histories=histories, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - validation_contexts=identifier_contexts, + identifier_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -502,7 +502,7 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - validation_contexts=identifier_contexts, + identifier_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index fe5bda7dbc..7e73f0bc65 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -4,9 +4,9 @@ import orjson from cachetools import TTLCache -from langfuse.decorators import observe from pydantic import BaseModel +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.utils import trace_metadata from src.web.v1.services import BaseRequest, MetadataTraceable @@ -87,7 +87,7 @@ async def _validate_question( ) return None - async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: + async def _document_retrieval() -> tuple[list[str], list[str], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -95,11 +95,17 @@ async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] + identifier_contexts = [ + document.get("identifier_context") + for document in documents + if document.get("identifier_context") + ] has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) return ( table_ddls, + identifier_contexts, has_calculated_field, has_metric, has_json_field, @@ -130,6 +136,7 @@ async def _instructions_retrieval() -> list[dict]: ) ( table_ddls, + identifier_contexts, has_calculated_field, has_metric, has_json_field, @@ -153,6 +160,7 @@ async def _instructions_retrieval() -> list[dict]: query=candidate["question"], contexts=table_ddls, project_id=project_id, + identifier_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py index 3c21a9c1d4..7e1be2ca7a 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -163,25 +163,6 @@ async def test_post_processor_rejects_plain_text_non_sql_response(): assert engine.execute_sql_calls == [] -@pytest.mark.asyncio -async def test_post_processor_rejects_sql_with_unretrieved_table_before_engine_call(): - engine = FakeEngine() - processor = SQLGenPostProcessor(engine) - - result = await processor.run( - ['{"sql":"SELECT 1 FROM missing_model"}'], - project_id="project-id", - data_source="mssql", - validation_contexts=["table: allowed_model\ncolumns:\n- allowed_column"], - ) - - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" - assert "missing_model" in result["invalid_generation_result"]["error"] - assert engine.dry_plan_calls == [] - assert engine.execute_sql_calls == [] - - @pytest.mark.asyncio async def test_post_processor_dry_plans_before_preview_execution(): engine = FakeEngine(dry_plan_success=False) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 46653b9564..7c362021e6 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -9,10 +9,12 @@ construct_retrieval_results, dbschema_retrieval, embedding, - prompt as build_column_selection_prompt, - table_retrieval, table_columns_selection_system_prompt, table_columns_selection_user_prompt_template, + table_retrieval, +) +from src.pipelines.retrieval.db_schema_retrieval import ( + prompt as build_column_selection_prompt, ) @@ -651,6 +653,40 @@ def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_co ) +def test_construct_retrieval_results_falls_back_when_pruner_omits_results(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={"replies": ['{"message":"not structured"}']}, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "modeled_dataset" + ] + assert "CREATE TABLE modeled_dataset" in result["retrieval_results"][0]["table_ddl"] + assert result["retrieval_results"][0]["identifier_context"] == ( + "table: modeled_dataset\ncolumns:\n- stored_attribute" + ) + + def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_unknown_columns(): result = construct_retrieval_results( check_using_db_schemas_without_pruning={}, From fbbda6c2722bb9801a572c04f3c265ba20ab489d Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 17 Aug 2026 00:39:45 +0530 Subject: [PATCH 1040/1087] Refresh SQL generation prompts and retries --- .../generation/followup_sql_generation.py | 30 +- .../pipelines/generation/sql_correction.py | 42 +-- .../pipelines/generation/sql_generation.py | 24 +- .../generation/sql_generation_reasoning.py | 5 +- .../pipelines/generation/sql_regeneration.py | 30 +- .../src/pipelines/generation/utils/sql.py | 296 +++++------------- wren-ai-service/src/web/v1/services/ask.py | 78 ++--- .../src/web/v1/services/ask_feedback.py | 23 +- 8 files changed, 164 insertions(+), 364 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 6a6d52cca0..0e50a12806 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -31,16 +31,8 @@ text_to_sql_with_followup_user_prompt_template = """ ### TASK ### -Given the user's current follow-up question and the current retrieved DATABASE SCHEMA, -generate one SQL query to best answer the user's question. - -{% if identifier_contexts %} -### EXECUTABLE WREN IDENTIFIER CATALOG ### -The following catalog is derived from the retrieved DATABASE SCHEMA. Use it as a quick index of the exact executable model, column, and relationship names available for this request. -{% for identifier_context in identifier_contexts %} -{{ identifier_context }} -{% endfor %} -{% endif %} +Given the following user's follow-up question and previous SQL query and summary, +generate one SQL query to best answer user's question. ### DATABASE SCHEMA ### {% for document in documents %} @@ -68,10 +60,11 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} -Question: -{{sample.question}} +Summary: +{{sample.summary}} +SQL: +{{sample.sql}} {% endfor %} {% endif %} @@ -84,10 +77,11 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -Generate one Wren SQL query that answers the user's follow-up question using the retrieved DATABASE SCHEMA. Use exact model and column names from DATABASE SCHEMA. Use aliases, descriptions, comments, samples, and previous questions only to understand business meaning. -If the retrieved DATABASE SCHEMA does not contain the model, column, relationship, metric, or supported function needed to answer the question, return {"sql": null}. -Return only the final JSON SQL response. +### REASONING PLAN ### +{{ sql_generation_reasoning }} + +Let's think step by step. """ @@ -98,7 +92,6 @@ def prompt( documents: list[str], sql_generation_reasoning: str, prompt_builder: PromptBuilder, - identifier_contexts: list[str] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -110,7 +103,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - identifier_contexts=identifier_contexts or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -212,7 +204,6 @@ async def run( instructions: list[dict] | None = None, project_id: str | None = None, mdl_hash: str | None = None, - identifier_contexts: list[str] | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -239,7 +230,6 @@ async def run( "histories": histories, "project_id": project_id, "mdl_hash": mdl_hash, - "identifier_contexts": identifier_contexts, "sql_samples": sql_samples, "instructions": instructions, "has_calculated_field": has_calculated_field, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 69aa9cd957..5fb54ac90c 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,20 +30,12 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are a Wren SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. +You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query. ### SQL CORRECTION INSTRUCTIONS ### -1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. -2. Then, generate a syntactically correct Wren SQL query from the user's intent and the current DATABASE SCHEMA. -3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. -4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. -5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. -6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. -7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. -8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA. If the unsupported part is needed to answer the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql instead of substituting non-schema identifiers. -10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. +1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). +2. Then, generate the syntactically correct ANSI SQL query to correct the error. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -51,10 +43,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. +The final answer must be in JSON format: {{ - "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" + "sql": }} """ @@ -82,20 +74,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -{% if query %} -User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. -{% endif %} -### FAILED SQL ### -The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. - -### DRY-RUN DIAGNOSTIC ### -The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. - -Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. +SQL: {{ invalid_generation_result.sql }} +Error Message: {{ invalid_generation_result.error }} -Return only the final JSON SQL response. +Let's think step by step. """ @@ -105,16 +87,12 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, - query: str | None = None, - sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( - query=query, documents=documents, invalid_generation_result=invalid_generation_result, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -193,8 +171,6 @@ async def run( self, contexts: List[Document], invalid_generation_result: Dict[str, str], - query: str | None = None, - sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, @@ -216,9 +192,7 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, - "query": query, "documents": contexts, - "sql_generation_reasoning": sql_generation_reasoning, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index c34c640ca4..71b501e216 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -28,14 +28,6 @@ sql_generation_user_prompt_template = """ -{% if identifier_contexts %} -### EXECUTABLE WREN IDENTIFIER CATALOG ### -The following catalog is derived from the retrieved DATABASE SCHEMA. Use it as a quick index of the exact executable model, column, and relationship names available for this request. -{% for identifier_context in identifier_contexts %} -{{ identifier_context }} -{% endfor %} -{% endif %} - ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -62,10 +54,11 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} +SQL: +{{sample.sql}} {% endfor %} {% endif %} @@ -78,10 +71,13 @@ ### QUESTION ### User's Question: {{ query }} -Generate one Wren SQL query that answers the user's question using the retrieved DATABASE SCHEMA. Use exact model and column names from DATABASE SCHEMA. Use aliases, descriptions, comments, and samples only to understand business meaning. -If the retrieved DATABASE SCHEMA does not contain the model, column, relationship, metric, or supported function needed to answer the question, return {"sql": null}. -Return only the final JSON SQL response. +{% if sql_generation_reasoning %} +### REASONING PLAN ### +{{ sql_generation_reasoning }} +{% endif %} + +Let's think step by step. """ @@ -92,7 +88,6 @@ def prompt( documents: list[str], prompt_builder: PromptBuilder, sql_generation_reasoning: str | None = None, - identifier_contexts: list[str] | None = None, sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, has_calculated_field: bool = False, @@ -104,7 +99,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - identifier_contexts=identifier_contexts or [], sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, @@ -203,7 +197,6 @@ async def run( instructions: list[dict] | None = None, project_id: str | None = None, mdl_hash: str | None = None, - identifier_contexts: list[str] | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -232,7 +225,6 @@ async def run( "instructions": instructions, "project_id": project_id, "mdl_hash": mdl_hash, - "identifier_contexts": identifier_contexts, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, 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 723566b2ab..6fdd8d6a01 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -29,10 +29,11 @@ {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} +SQL: +{{sql_sample.sql}} {% endfor %} {% endif %} @@ -48,7 +49,7 @@ Language: {{ language }} Current Time: {{ current_time }} -Return only the reasoning plan described by the system instructions. When relevant, ground the plan by using the literal prefix `table:` followed by an exact declared table name from DATABASE SCHEMA, or the literal prefix `column:` followed by an exact declared table name, a dot, and an exact declared column name. Do not include SQL, SQL-like expressions, aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, template markers, functions, or identifier-like labels. +Let's think step by step. """ diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 81a697dd9c..9926d69419 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,20 +34,18 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema and a user's question. -Carefully review the user's question and current DATABASE SCHEMA, then generate a new SQL query that answers the user's intent. -The original SQL query and UI planning text are intentionally omitted from the prompt and must not be used as executable context. -While generating the new SQL query, make sure to use the database schema as the only source of executable table and column identifiers. -If the original SQL query or reasoning contains unsupported identifiers, placeholders, or assumptions, ignore those parts and regenerate from the user's question and DATABASE SCHEMA. -Treat physical/source/lineage names from the original SQL, reasoning, samples, comments, or descriptions as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query, +please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. +While generating the new SQL query, you should use the original SQL query as a reference. +While generating the new SQL query, make sure to use the database schema to generate the SQL query. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. +The final answer must be a ANSI SQL query in JSON format: {{ - "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" + "sql": }} """ @@ -79,10 +77,11 @@ def get_sql_regeneration_system_prompt( {% if sql_samples %} ### SQL SAMPLES ### -These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} +SQL: +{{sample.sql}} {% endfor %} {% endif %} @@ -94,20 +93,16 @@ def get_sql_regeneration_system_prompt( {% endif %} ### QUESTION ### -User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. -Regenerate with executable identifiers from the current DATABASE SCHEMA only. -### ORIGINAL SQL QUERY ### -The original SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. +SQL generation reasoning: {{ sql_generation_reasoning }} +Original SQL query: {{ sql }} -Return only the final JSON SQL response. +Let's think step by step. """ ## Start of Pipeline @observe(capture_input=False) def prompt( - query: str, documents: list[str], sql_generation_reasoning: str, sql: str, @@ -121,7 +116,6 @@ def prompt( sql_knowledge: SqlKnowledge | None = None, ) -> dict: _prompt = prompt_builder.run( - query=query, sql=sql, documents=documents, sql_generation_reasoning=sql_generation_reasoning, @@ -203,7 +197,6 @@ def __init__( async def run( self, contexts: list[str], - query: str, sql_generation_reasoning: str, sql: str, sql_samples: list[dict] | None = None, @@ -222,7 +215,6 @@ async def run( ["post_process"], inputs={ "documents": contexts, - "query": query, "sql_generation_reasoning": sql_generation_reasoning, "sql": sql, "sql_samples": sql_samples, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index f0ec305b32..756c662ee4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -5,7 +5,7 @@ import orjson from haystack import component from haystack.dataclasses import ChatMessage -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel from src.core.engine import ( Engine, @@ -22,74 +22,6 @@ class SQLGenPostProcessor: def __init__(self, engine: Engine): self._engine = engine - def _looks_like_sql(self, value: str) -> bool: - normalized = value.strip().casefold() - return normalized.startswith("select ") or normalized.startswith("with ") - - def _json_object(self, value: str) -> dict[str, Any]: - parsed = orjson.loads(value) - return parsed if isinstance(parsed, dict) else {} - - def _extract_sql_from_object(self, generation_result: dict[str, Any]) -> str: - sql = generation_result.get("sql") - if isinstance(sql, str) and sql.strip(): - return clean_generation_result(sql) - - arguments = generation_result.get("arguments") - if isinstance(arguments, str) and arguments.strip(): - try: - arguments = self._json_object(arguments) - except orjson.JSONDecodeError: - arguments = {} - - if isinstance(arguments, dict): - for key in ("sql", "query"): - value = arguments.get(key) - if isinstance(value, str) and self._looks_like_sql(value): - return clean_generation_result(value) - - query = generation_result.get("query") - if isinstance(query, str) and self._looks_like_sql(query): - return clean_generation_result(query) - - return "" - - def _extract_sql(self, replies: List[str] | List[List[str]]) -> tuple[str, str]: - if not replies: - return "", "SQL generation response was empty." - - reply = replies[0] - if isinstance(reply, list): - reply = reply[0] if reply else "" - - cleaned_generation_result = clean_generation_result(reply) - if not cleaned_generation_result: - return "", "SQL generation response was empty." - - if cleaned_generation_result.startswith("{"): - try: - generation_result = self._json_object(cleaned_generation_result) - except orjson.JSONDecodeError as e: - return "", f"SQL generation response was not valid JSON: {e}" - - sql = self._extract_sql_from_object(generation_result) - if sql: - return sql, "" - - if "sql" in generation_result and generation_result.get("sql") is None: - return "", "No grounded SQL was generated from the current schema." - - return ( - "", - "SQL generation response did not include a supported SQL field: " - f"{generation_result}", - ) - - if self._looks_like_sql(cleaned_generation_result): - return cleaned_generation_result, "" - - return "", "SQL generation response was not a supported SQL JSON payload." - @component.output_types( valid_generation_result=Dict[str, Any], invalid_generation_result=Dict[str, Any], @@ -105,25 +37,19 @@ async def run( allow_data_preview: bool = False, ) -> dict: try: - generation_result, extraction_error = self._extract_sql(replies) - if not generation_result: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": "", - "original_sql": "", - "type": "NO_RELEVANT_SQL", - "error": extraction_error - or "No grounded SQL was generated from the current schema.", - "correlation_id": "", - }, - } + cleaned_generation_result = clean_generation_result(replies[0]) + + # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' + if cleaned_generation_result.startswith("{"): + cleaned_generation_result = orjson.loads(cleaned_generation_result)[ + "sql" + ] ( valid_generation_result, invalid_generation_result, ) = await self._classify_generation_result( - generation_result, + cleaned_generation_result, project_id=project_id, mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, @@ -159,8 +85,7 @@ async def _classify_generation_result( use_dry_run = not allow_data_preview async with aiohttp.ClientSession() as session: - should_dry_plan = use_dry_plan or bool(project_id and data_source) - if should_dry_plan: + if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( session, generation_result, @@ -170,43 +95,20 @@ async def _classify_generation_result( allow_fallback=allow_dry_plan_fallback, ) - if not dry_plan_result: - invalid_generation_result = { - "sql": generation_result, - "original_sql": generation_result, - "type": "TIME_OUT" - if error_message.startswith("Request timed out") - else "DRY_PLAN", - "error": error_message, - "correlation_id": "", - } - return valid_generation_result, invalid_generation_result - - success, _, addition = await self._engine.execute_sql( - generation_result, - session, - project_id=project_id, - mdl_hash=mdl_hash, - limit=1, - dry_run=True, - ) - addition = addition if isinstance(addition, dict) else {} - - if success: + if dry_plan_result: valid_generation_result = { "sql": generation_result, - "correlation_id": addition.get("correlation_id", ""), + "correlation_id": "", } else: - error_message = addition.get("error_message", "") invalid_generation_result = { - "sql": addition.get("error_sql", generation_result), + "sql": generation_result, "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") - else "DRY_RUN", + else "DRY_PLAN", "error": error_message, - "correlation_id": addition.get("correlation_id", ""), + "correlation_id": "", } elif use_dry_run: success, _, addition = await self._engine.execute_sql( @@ -269,44 +171,6 @@ async def _classify_generation_result( return valid_generation_result, invalid_generation_result -_MANDATORY_SQL_GROUNDING_RULES = """ -### MANDATORY SQL GROUNDING RULES ### -- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. -- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. -- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. -- Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. -- Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. -- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. -- When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. -- In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. -- Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. -- When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. -- The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. -- Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. -- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. -- Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. -- If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. -- When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. -- Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. -- When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. -- When the same requested result can be answered from multiple schema objects with compatible columns or metrics, include all relevant schema objects by combining separate result rows with UNION ALL instead of choosing only one object. -- Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. -- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. -- Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. -- SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -- Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. -- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. -- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. -- If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. -- For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. -- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. -- If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. -- If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. -- If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. -- Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. -""" - - _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. @@ -318,31 +182,52 @@ async def _classify_generation_result( - PREFER USING CTEs over subqueries. - When generating SQL query, always: - Put double quotes around column and table names. - - Use Wren SQL identifier quoting with double quotes only; the engine rewrite step converts grounded Wren SQL to the active connector dialect. - Put single quotes around string literals. - Never quote numeric literals. -- Generate Wren SQL syntax only, not connector-specific SQL syntax. -- Never use SELECT TOP, TOP(...), FETCH FIRST, square-bracket identifiers, or backtick identifiers. For top or limit requests, sort with ORDER BY and put LIMIT at the end of the query. -- Preserve every deployed table and column identifier exactly as it appears in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, including spaces, digits, underscores, case, and punctuation, then wrap that exact identifier in double quotes in SQL. -- Do not convert deployed identifiers into display-friendly variants by replacing spaces with underscores, removing prefixes, changing case, shortening names, or expanding abbreviations. -- For case-insensitive comparisons, use only functions or operators that are supported by SQL FUNCTIONS for this request. If SQL FUNCTIONS does not provide a safe case-insensitive function, use a normal equality or LIKE comparison on an exact schema column. -- For date/time questions, first choose an exact schema column whose type or metadata clearly represents the requested time concept. Use only date/time functions and casts whose exact syntax is provided in SQL FUNCTIONS for this request. -- If the question asks for a specific or relative date, generate a bounded date/time filter only when both the exact date/time schema column and required SQL FUNCTIONS-supported operation are available. If either is missing, do not invent a field or function. + For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992; +- YOU MUST USE "lower(.) like lower()" function or "lower(.) = lower()" function for case-insensitive comparison! + - Use "lower(.) LIKE lower()" when: + - The user requests a pattern or partial match. + - The value is not specific enough to be a single, exact value. + - Wildcards (%) are needed to capture the pattern. + - Use "lower(.) = lower()" when: + - The user requests an exact, specific value. + - There is no ambiguity or pattern in the value. +- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query + - example: TO_TIMESTAMP_MILLIS("") # if the timestamp_column is in milliseconds + - example: TO_TIMESTAMP_SECONDS("") # if the timestamp_column is in seconds + - example: TO_TIMESTAMP_MICROS("") # if the timestamp_column is in microseconds +- ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query + - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) + - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) + - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) +- If the user asks for a specific date, please give the date range in SQL query + - example: "What is the total revenue for the month of 2024-11-01?" + - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. -- Output aliases may be used only to name expressions in the final SELECT list. Output aliases are labels for result columns only; they are not source identifiers. -- Comments, aliases, display labels, and descriptions from DATABASE SCHEMA may guide which exact source column to select, but they must not be copied into FROM, JOIN, WHERE, GROUP BY, HAVING, or ORDER BY as table or column names. -- Physical/source/lineage names from metadata may guide meaning, but generated SQL must use only the declared Wren model, view, metric, and column identifiers from DATABASE SCHEMA. -- DON'T USE '.' in output aliases, replace '.' with '_' in output aliases. +- ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. +- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. + - EXAMPLE + DATABASE SCHEMA + /* {"alias":"_orders","description":"A model representing the orders data."} */ + CREATE TABLE orders ( + -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} + ApprovedTimestamp TIMESTAMP + } + + SQL + SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; +- DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. - DON'T USE "EXTRACT()" function with INTERVAL data types as arguments - DON'T USE INTERVAL or generate INTERVAL-like expression in the generated SQL query. - DON'T USE "TO_CHAR" function in the generated SQL query. -- DON'T USE unsupported statistical, date/time, or formatting functions. If SQL FUNCTIONS does not list a function needed by the requested intent, omit the function-dependent part. If that function is required to answer the request, return null for sql. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. +- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. +- For the ranking problem, you must add the ranking column to the final SELECT clause. """ @@ -428,13 +313,15 @@ async def _classify_generation_result( sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who explains the user's analytical intent and provides a concise, non-executable reasoning plan for answering the user's question. +You are a helpful data analyst who is great at thinking deeply and reasoning about the user's question and the database schema, and you provide a step-by-step reasoning plan in order to answer the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state requested timeframes in natural language only. Mention exact date/time columns only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. -3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language. Mention exact ordering columns or measures only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. -4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. +2. Explicitly state the following information in the reasoning plan: +if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; +otherwise, you will put the relative timeframe in the SQL query. +3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. +4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. @@ -442,22 +329,9 @@ async def _classify_generation_result( 9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. Mention table names only by writing the literal prefix `table:` followed by an exact table name declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. -13. Mention column names only by writing the literal prefix `column:` followed by an exact declared table name, a dot, and an exact column name declared for that table in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. -14. Do not mention aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels from comments, SQL samples, failed SQL, or user wording as executable identifiers. -15. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. -16. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. -17. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language and cite exact declared tables or columns only when they are grounded by DATABASE SCHEMA. -18. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, then ground the plan in exact declared schema identifiers. -19. If multiple schema objects are required, identify the exact declared relationship path from DATABASE SCHEMA. If no relationship path is declared, say that the retrieved metadata does not provide a join path. -20. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan unless they also appear exactly in DATABASE SCHEMA. -21. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. -22. The reasoning plan is semantic context for intent only, not a source of executable identifiers. SQL generation must re-read DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT before using any identifier. -23. ONLY SHOWING the reasoning plan in bullet points. -24. Do not use the words "assume", "assuming", "likely", "possible", "might", or "example" when describing tables, columns, filters, or SQL. -25. If exact deployed table and column identifiers are not available for a requested part, say only that the retrieved metadata does not support that part. Do not propose a replacement name. -26. Do not write table names or column names from the user's wording unless the same identifier appears exactly in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. -27. Do not include code blocks, inline SQL fragments, SELECT statements, WHERE clauses, join clauses, or any query-shaped text in the reasoning plan. +12. A table name in the reasoning plan must be in this format: `table: `. +13. A column name in the reasoning plan must be in this format: `column: .`. +14. ONLY SHOWING the reasoning plan in bullet points. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -475,13 +349,12 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: - rules = _DEFAULT_TEXT_TO_SQL_RULES if sql_knowledge is not None: - rules = _extract_from_sql_knowledge( + return _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) - return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" + return _DEFAULT_TEXT_TO_SQL_RULES def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: @@ -517,51 +390,38 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" -You are a helpful assistant that converts natural language queries into Wren SQL queries. +You are a helpful assistant that converts natural language queries into ANSI SQL queries. -Given the user's question and database schema, generate one grounded Wren SQL query. The DATABASE SCHEMA is the only source of executable identifiers. +Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. -2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. -3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. -4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. -5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. -6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. -7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. -8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. -9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. -10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. -11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. -13. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. -14. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. +3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. +4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. +5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. Do not create table or column identifiers from the user's wording. If the retrieved schema does not ground the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. +The final answer must be a ANSI SQL query in JSON format: {{ - "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" + "sql": }} """ class SqlGenerationResult(BaseModel): - model_config = ConfigDict(extra="forbid") - - sql: str | None + sql: str SQL_GENERATION_MODEL_KWARGS = { - "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { "name": "sql_generation_result", - "strict": True, "schema": SqlGenerationResult.model_json_schema(), }, } @@ -583,4 +443,18 @@ def construct_instructions( def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: - return [] + messages = [] + for history in histories: + messages.append( + ChatMessage.from_user( + history.question + if hasattr(history, "question") + else history["question"] + ) + ) + messages.append( + ChatMessage.from_assistant( + history.sql if hasattr(history, "sql") else history["sql"] + ) + ) + return messages diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 413b8a1228..bac4bb7c41 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -25,10 +25,10 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = True + ignore_sql_generation_reasoning: bool = False enable_column_pruning: bool = False - use_dry_plan: bool = True - allow_dry_plan_fallback: bool = False + use_dry_plan: bool = False + allow_dry_plan_fallback: bool = True custom_instruction: Optional[str] = None @@ -99,12 +99,12 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = False, + allow_sql_generation_reasoning: bool = True, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = True, - max_sql_correction_retries: int = 0, + enable_column_pruning: bool = False, + max_sql_correction_retries: int = 3, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -161,14 +161,17 @@ async def ask( table_names = [] error_message = None invalid_sql = None - allow_sql_generation_reasoning = False + allow_sql_generation_reasoning = ( + self._allow_sql_generation_reasoning + and not ask_request.ignore_sql_generation_reasoning + ) enable_column_pruning = ( self._enable_column_pruning or ask_request.enable_column_pruning ) allow_sql_functions_retrieval = self._allow_sql_functions_retrieval allow_sql_diagnosis = self._allow_sql_diagnosis allow_sql_knowledge_retrieval = self._allow_sql_knowledge_retrieval - max_sql_correction_retries = 0 + max_sql_correction_retries = self._max_sql_correction_retries current_sql_correction_retries = 0 use_dry_plan = ask_request.use_dry_plan allow_dry_plan_fallback = ask_request.allow_dry_plan_fallback @@ -186,29 +189,29 @@ async def ask( is_followup=True if histories else False, ) - if not api_results: - historical_question_result = ( - await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, + historical_question = await self._pipelines["historical_question"].run( + query=user_query, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + ) + + # we only return top 1 result + historical_question_result = historical_question.get( + "formatted_output", {} + ).get("documents", [])[:1] + + if historical_question_result: + api_results = [ + AskResult( + **{ + "sql": result.get("statement"), + "type": "view" if result.get("viewId") else "llm", + "viewId": result.get("viewId"), + } ) - ) - historical_questions = historical_question_result[ - "formatted_output" - ].get("documents", []) - if historical_questions: - api_results = [ - AskResult( - **{ - "sql": result.get("statement"), - "type": "view", - "viewId": result.get("viewId"), - } - ) - for result in historical_questions - if result.get("statement") - ] + for result in historical_question_result + ] + sql_generation_reasoning = "" if not api_results: # Run both pipeline operations concurrently @@ -357,11 +360,6 @@ async def ask( documents = _retrieval_result.get("retrieval_results", []) table_names = [document.get("table_name") for document in documents] table_ddls = [document.get("table_ddl") for document in documents] - identifier_contexts = [ - document.get("identifier_context") - for document in documents - if document.get("identifier_context") - ] if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -482,7 +480,6 @@ async def ask( histories=histories, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - identifier_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -502,7 +499,6 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - identifier_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -529,11 +525,7 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] in ( - "TIME_OUT", - "NO_RELEVANT_SQL", - "SQL_GENERATION", - ): + if failed_dry_run_result["type"] == "TIME_OUT": error_message = failed_dry_run_result["error"] invalid_sql = failed_dry_run_result["sql"] break @@ -574,8 +566,6 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, - query=user_query, - sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 379a661a29..184c15a37a 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -184,7 +184,6 @@ async def ask_feedback( "sql_regeneration" ].run( contexts=table_ddls, - query=ask_feedback_request.question, sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, @@ -212,15 +211,10 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] not in ( - "TIME_OUT", - "NO_RELEVANT_SQL", - "SQL_GENERATION", - ): + if failed_dry_run_result["type"] != "TIME_OUT": original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] - sql_diagnosis_reasoning = None self._ask_feedback_results[ query_id @@ -245,23 +239,16 @@ async def ask_feedback( "post_process" ].get("reasoning") - correction_error_message = error_message - if sql_diagnosis_reasoning: - correction_error_message = ( - f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" - ) - sql_correction_results = await self._pipelines[ "sql_correction" ].run( contexts=table_ddls, - query=ask_feedback_request.question, - sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, instructions=instructions, invalid_generation_result={ - "original_sql": original_sql, - "sql": invalid_sql, - "error": correction_error_message, + "sql": original_sql, + "error": sql_diagnosis_reasoning + if allow_sql_diagnosis + else error_message, }, project_id=ask_feedback_request.project_id, mdl_hash=ask_feedback_request.mdl_hash, From 9606d406b3657c0b81908baa63add59fd118ebf8 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 17 Aug 2026 00:56:45 +0530 Subject: [PATCH 1041/1087] Simplify LiteLLM request kwargs handling --- wren-ai-service/src/providers/llm/litellm.py | 91 +++----------------- 1 file changed, 11 insertions(+), 80 deletions(-) diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index fc9406b75b..e94918b8c5 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -19,10 +19,6 @@ from src.utils import extract_braces_content, remove_trailing_slash -def _is_openai_api_base(api_base: Optional[str]) -> bool: - return bool(api_base) and "api.openai.com" in api_base.lower() - - @provider("litellm_llm") class LitellmLLMProvider(LLMProvider): def __init__( @@ -69,42 +65,10 @@ def get_generator( generation_kwargs: Optional[Dict[str, Any]] = None, streaming_callback: Optional[Callable[[StreamingChunk], None]] = None, ): - component_generation_kwargs = generation_kwargs or {} - - def _normalize_generation_kwargs( - kwargs: Optional[Dict[str, Any]], - explicit_response_format: bool = False, - ) -> Dict[str, Any]: - normalized = dict(kwargs or {}) - preserve_json_schema = normalized.pop("preserve_json_schema", False) - response_format = normalized.get("response_format") - - # Plain text is the default chat-completions behavior. - # Some OpenAI-compatible endpoints reject an explicit - # {"type": "text"} payload or serialize it incorrectly. - if ( - isinstance(response_format, dict) - and response_format.get("type") == "text" - ): - normalized.pop("response_format", None) - - if ( - self._api_base - and not _is_openai_api_base(self._api_base) - and isinstance(response_format, dict) - and response_format.get("type") == "json_schema" - and not preserve_json_schema - ): - if explicit_response_format: - normalized["response_format"] = {"type": "json_object"} - else: - normalized.pop("response_format", None) - - if self._api_base and not _is_openai_api_base(self._api_base): - # Some local OpenAI-compatible servers reject non-OpenAI keys. - normalized.pop("speed", None) - - return normalized + combined_generation_kwargs = { + **(generation_kwargs or {}), + **(self._model_kwargs or {}), + } @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) async def _run( @@ -135,42 +99,10 @@ async def _run( convert_message_to_openai_format(message) for message in messages ] - runtime_generation_kwargs = generation_kwargs or {} - model_generation_kwargs = self._model_kwargs or {} - explicit_response_format = ( - "response_format" in component_generation_kwargs - or "response_format" in model_generation_kwargs - or "response_format" in runtime_generation_kwargs - ) - merged_generation_kwargs = { - **component_generation_kwargs, - **model_generation_kwargs, - **runtime_generation_kwargs, + generation_kwargs = { + **combined_generation_kwargs, + **(generation_kwargs or {}), } - if ( - component_generation_kwargs.get("preserve_json_schema") - and isinstance( - component_generation_kwargs.get("response_format"), dict - ) - and component_generation_kwargs["response_format"].get("type") - == "json_schema" - and "response_format" not in runtime_generation_kwargs - ): - merged_generation_kwargs["response_format"] = ( - component_generation_kwargs["response_format"] - ) - merged_generation_kwargs["preserve_json_schema"] = True - - generation_kwargs = _normalize_generation_kwargs( - merged_generation_kwargs, - explicit_response_format=explicit_response_format, - ) - completion_timeout = generation_kwargs.pop("timeout", self._timeout) - should_stream = ( - streaming_callback is not None - and query_id is not None - and generation_kwargs.pop("stream", True) - ) allowed_openai_params = generation_kwargs.get( "allowed_openai_params", [] @@ -180,10 +112,9 @@ async def _run( completion = await self._router.acompletion( model=self._model, messages=openai_formatted_messages, - stream=should_stream, + stream=streaming_callback is not None, allowed_openai_params=allowed_openai_params, mock_testing_fallbacks=self._enable_fallback_testing, - timeout=completion_timeout, **generation_kwargs, ) else: @@ -192,15 +123,15 @@ async def _run( api_key=self._api_key, api_base=self._api_base, api_version=self._api_version, - timeout=completion_timeout, + timeout=self._timeout, messages=openai_formatted_messages, - stream=should_stream, + stream=streaming_callback is not None, allowed_openai_params=allowed_openai_params, **generation_kwargs, ) completions: List[ChatMessage] = [] - if should_stream: + if streaming_callback is not None: num_responses = generation_kwargs.pop("n", 1) if num_responses > 1: raise ValueError( From 878d0fb4d80d68a6d89314f6cfd1b61fae572294 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 17 Aug 2026 01:39:29 +0530 Subject: [PATCH 1042/1087] Simplify DB schema retrieval and DDL builders --- .../src/pipelines/indexing/db_schema.py | 98 +-- .../pipelines/indexing/table_description.py | 152 +---- .../retrieval/db_schema_retrieval.py | 635 ++---------------- .../v1/services/question_recommendation.py | 24 +- 4 files changed, 87 insertions(+), 822 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index a6eef11ec3..50564f1666 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -26,20 +26,6 @@ logger = logging.getLogger("wren-ai-service") -MAX_DB_SCHEMA_METADATA_TEXT_LENGTH = 1000 - - -def _truncate_metadata_text(value: Any) -> Any: - if value is None: - return "" - if isinstance(value, str) and len(value) > MAX_DB_SCHEMA_METADATA_TEXT_LENGTH: - return value[: MAX_DB_SCHEMA_METADATA_TEXT_LENGTH - 3] + "..." - if isinstance(value, dict): - return {key: _truncate_metadata_text(item) for key, item in value.items()} - if isinstance(value, list): - return [_truncate_metadata_text(item) for item in value] - return value - @component class DDLChunker: @@ -91,7 +77,7 @@ def _column_preprocessor( column: Dict[str, Any], addition: Dict[str, Any] ) -> Dict[str, Any]: addition = { - key: _truncate_metadata_text(helper(column, **addition)) + key: helper(column, **addition) for key, helper in helper.COLUMN_PREPROCESSORS.items() if helper.condition(column, **addition) } @@ -104,7 +90,7 @@ def _column_preprocessor( async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: addition = { - key: _truncate_metadata_text(await helper(model, **kwargs)) + key: await helper(model, **kwargs) for key, helper in helper.MODEL_PREPROCESSORS.items() if helper.condition(model, **kwargs) } @@ -116,7 +102,7 @@ async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: ] return { "name": model.get("name", ""), - "properties": _truncate_metadata_text(model.get("properties", {})), + "properties": model.get("properties", {}), "columns": columns, "primaryKey": model.get("primaryKey", ""), } @@ -203,69 +189,20 @@ def _relationship_command( if join_type not in ["MANY_TO_ONE", "ONE_TO_MANY", "ONE_TO_ONE"]: return None - condition_parts = [ - part.strip() for part in condition.split("=", maxsplit=1) - ] - if len(condition_parts) != 2: - return None - - model_columns = [] - for condition_part in condition_parts: - name_parts = [ - part.strip() for part in condition_part.split(".", maxsplit=1) - ] - if len(name_parts) != 2: - return None - model_columns.append( - {"table": name_parts[0], "column": name_parts[1]} - ) - - left, right = model_columns - - if join_type == "MANY_TO_ONE": - foreign_side, referenced_side = left, right - elif join_type == "ONE_TO_MANY": - foreign_side, referenced_side = right, left - elif table_name == left["table"]: - foreign_side, referenced_side = left, right - else: - foreign_side, referenced_side = right, left - - if table_name != foreign_side["table"]: - return None - - related_table = referenced_side["table"] - referenced_column = referenced_side["column"] or primary_keys_map.get( - related_table, "" - ) - fk_column = foreign_side["column"] - fk_constraint = ( - f"FOREIGN KEY ({fk_column}) " - f"REFERENCES {related_table}({referenced_column})" - ) + # Get related table and foreign key column + is_source = table_name == models[0] + related_table = models[1] if is_source else models[0] + condition_parts = condition.split(" = ") + fk_column = condition_parts[0 if is_source else 1].split(".")[1] - properties = relationship.get("properties", {}) - relationship_properties = { - "name": relationship.get("name", ""), - "condition": condition, - "joinType": join_type, - "description": _truncate_metadata_text( - properties.get("description", "") - ) - if isinstance(properties, dict) - else "", - "from": f"{foreign_side['table']}.{foreign_side['column']}", - "to": f"{referenced_side['table']}.{referenced_side['column']}", - } + # Build foreign key constraint + fk_constraint = f"FOREIGN KEY ({fk_column}) REFERENCES {related_table}({primary_keys_map[related_table]})" return { "type": "FOREIGN_KEY", - "comment": f"-- {relationship_properties}\n ", + "comment": f'-- {{"condition": {condition}, "joinType": {join_type}}}\n ', "constraint": fk_constraint, "tables": models, - "column": fk_column, - "referenced_table": related_table, - "referenced_column": referenced_column, } def _column_batch( @@ -304,18 +241,6 @@ def _column_batch( ] def _convert_views(self, views: List[Dict[str, Any]]) -> List[Dict[str, str]]: - def _columns(view: Dict[str, Any]) -> List[dict]: - properties = view.get("properties", {}) or {} - return [ - { - "name": column.get("name", ""), - "data_type": column.get("type", ""), - "comment": column.get("description", ""), - } - for column in properties.get("columns", []) - if column.get("name") - ] - def _payload(view: Dict[str, Any]) -> dict: return { "type": "VIEW", @@ -324,7 +249,6 @@ def _payload(view: Dict[str, Any]) -> dict: else "", "name": view["name"], "statement": view["statement"], - "columns": _columns(view), } return [ diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 20dabf4c99..3d61238d79 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -22,10 +22,6 @@ @component class TableDescriptionChunker: - def _properties(self, payload: Dict[str, Any]) -> Dict[str, Any]: - properties = payload.get("properties") - return properties if isinstance(properties, dict) else {} - @component.output_types(documents=List[Document]) def run( self, @@ -64,161 +60,27 @@ def _additional_meta() -> Dict[str, Any]: ] } - def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[Dict[str, Any]]: - def _text(value: Any) -> str: - return "" if value is None else str(value) - - def _source_context(payload: Dict[str, Any]) -> str: - table_reference = payload.get("tableReference") - if isinstance(table_reference, dict): - reference_parts = [ - _text(table_reference.get("catalog", "")), - _text(table_reference.get("schema", "")), - _text(table_reference.get("table", "")), - ] - return ".".join(part for part in reference_parts if part) - - return _text(payload.get("baseObject", "")) - - def _columns(payload: Dict[str, Any]) -> List[Dict[str, Any]]: - columns = payload.get("columns", []) - if columns: - return [ - {**column, "role": _text(column.get("role", ""))} - for column in columns - if isinstance(column, dict) - ] - - metric_columns = [] - for role, key in [("dimension", "dimension"), ("measure", "measure")]: - metric_columns += [ - {**column, "role": role} - for column in payload.get(key, []) or [] - if isinstance(column, dict) - ] - - return metric_columns - + def _get_table_descriptions(self, mdl: Dict[str, Any]) -> List[str]: def _structure_data(mdl_type: str, payload: Dict[str, Any]) -> Dict[str, Any]: - properties = self._properties(payload) - return { "mdl_type": mdl_type, "name": payload.get("name"), - "displayName": _text(properties.get("displayName", "")), - "source": _source_context(payload), - "columns": [ - { - "name": _text(column.get("name", "")), - "type": _text(column.get("type", "")), - "role": _text(column.get("role", "")), - "expression": _text(column.get("expression", "")), - "description": _text( - self._properties(column).get("description", "") - ), - "displayName": _text( - self._properties(column).get("displayName", "") - ), - } - for column in _columns(payload) - ], - "properties": properties, + "columns": [column["name"] for column in payload.get("columns", [])], + "properties": payload.get("properties", {}), } - def _relationship_context_by_model() -> Dict[str, List[str]]: - relationships = {model.get("name"): [] for model in mdl.get("models", [])} - - for relationship in mdl.get("relationships", []) or []: - models = relationship.get("models", []) - if len(models) != 2: - continue - - properties = self._properties(relationship) - summary = " ".join( - part - for part in [ - _text(relationship.get("name", "")), - _text(relationship.get("joinType", "")), - _text(relationship.get("condition", "")), - _text(properties.get("description", "")), - f"models {' <-> '.join(_text(model) for model in models)}", - ] - if part - ) - if not summary: - continue - - for model_name in models: - relationships.setdefault(model_name, []).append(summary) - - return relationships - - def _column_context(columns: List[Dict[str, Any]]) -> str: - details = [] - - for column in columns: - semantic_parts = [ - column["type"], - column["role"], - column["displayName"], - column["description"], - column["expression"], - ] - if not any(semantic_parts): - continue - - details.append( - " ".join( - part for part in [column["name"], *semantic_parts] if part - ) - ) - - return "; ".join(detail for detail in details if detail) - - relationship_context = _relationship_context_by_model() resources = ( [_structure_data("MODEL", model) for model in mdl["models"]] + [_structure_data("METRIC", metric) for metric in mdl["metrics"]] + [_structure_data("VIEW", view) for view in mdl["views"]] ) - def _resource_description(resource: Dict[str, Any]) -> Dict[str, str]: - column_context = _column_context(resource["columns"]) - relationships = "; ".join(relationship_context.get(resource["name"], [])) - description = { + return [ + { "name": resource["name"], - "resource_type": resource["mdl_type"], - "description": resource["properties"].get("description", "") or "", - "columns": ", ".join( - column["name"] for column in resource["columns"] - ), + "description": resource["properties"].get("description", ""), + "columns": ", ".join(resource["columns"]), } - - if resource["displayName"]: - description["displayName"] = resource["displayName"] - - if resource["source"]: - description["source"] = resource["source"] - - if column_context: - description["column_context"] = column_context - - if relationships: - description["relationships"] = relationships - - semantic_parts = [ - description.get("displayName", ""), - description.get("source", ""), - column_context, - relationships, - ] - if semantic_context := "; ".join(part for part in semantic_parts if part): - description["semantic_context"] = semantic_context - - return description - - return [ - _resource_description(resource) for resource in resources if resource["name"] is not None ] diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 90951859e7..22bf940cd8 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -9,9 +9,9 @@ from hamilton.async_driver import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder +from langfuse.decorators import observe from pydantic import BaseModel -from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider from src.pipelines.common import ( @@ -30,13 +30,7 @@ ### TASK ### You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. -The database schema includes structural, semantic, and business modeling metadata: -- Models are logical datasets backed by physical tables or SQL definitions. -- Columns are exposed fields, including renamed fields, expressions, primary keys, and calculated fields. -- Relationships are reusable join logic between models. -- Calculated fields are business logic defined once and reused across queries. -- Views are named SQL statements that behave like stable virtual tables. -- Metrics are structured aggregation objects with measures and dimensions. +The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -46,17 +40,6 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. -8. Map the business question to the modeled datasets whose descriptions, aliases, columns, calculated fields, views, metrics, and relationships support the intent. -9. Prefer modeled analytical interfaces such as views and metrics when they expose the fields needed to answer the question. -10. If the answer needs fields, filters, time dimensions, ordering, aggregations, or relationship keys from multiple related datasets, include every required related dataset and the columns needed from each one. -11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. -12. Follow only the relationships shown in the provided schema when selecting columns across datasets. -13. Do not stop at a single top candidate when the question needs multiple related datasets. -14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. -15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. -16. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. -17. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. -18. Select the tables, views, metrics, relationships, and columns that best support the current question from the available modeled schema. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -129,320 +112,37 @@ def _project_filter_conditions( def _build_metric_ddl(content: dict) -> str: - columns = [ - column - for column in content["columns"] - if column["data_type"].lower() != "unknown" - ] - context = _format_semantic_context( - { - "object_type": "metric", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [column["name"] for column in columns], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable analytical aggregation interface", - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "semantic_context_not_sql_identifier": column["comment"], - } - for column in columns - ], - } - ) columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" - for column in columns + f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in content["columns"] + if column["data_type"].lower() + != "unknown" # quick fix: filtering out UNKNOWN column type ] return ( - f"{context}CREATE TABLE {content['name']} (\n " + f"{content['comment']}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) def _build_view_ddl(content: dict) -> str: - columns = [ - column - for column in content.get("columns", []) - if column.get("name") and column.get("data_type", "").lower() != "unknown" - ] - context = _format_semantic_context( - { - "object_type": "view", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [column["name"] for column in columns], - }, - "semantic_context_not_sql_identifiers": { - "role": "stable virtual table interface", - "description": content["comment"], - "definition_omitted_from_executable_schema": True, - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type( - column.get("data_type") - ), - "semantic_context_not_sql_identifier": column.get("comment", ""), - } - for column in columns - ], - } - ) - columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" - for column in columns - ] - - return ( - f"{context}CREATE TABLE {content['name']} (\n " - + ",\n ".join(columns_ddl) - + "\n);" - ) - - -def _format_semantic_context(context: dict) -> str: return ( - "/*\n" - "WREN RETRIEVED SEMANTIC CONTEXT\n" - f"{orjson.dumps(context).decode('utf-8')}\n" - f"{_format_identifier_contract(context)}" - "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" - "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" - "*/\n" - f"{_format_executable_identifier_catalog(context)}" - ) - - -def _format_executable_identifier_catalog(context: dict) -> str: - contract = context.get("sql_identifier_contract", {}) - table_name = contract.get("sql_table_name_use_exactly") - column_names = contract.get("sql_column_names_use_exactly") or [ - column["sql_column_name_use_exactly"] - for column in context.get("columns", []) - if column.get("sql_column_name_use_exactly") - ] - relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ - relationship["sql_relationship_constraint_use_exactly"] - for relationship in context.get("relationships", []) - if relationship.get("sql_relationship_constraint_use_exactly") - ] - - lines = [ - "### EXECUTABLE WREN IDENTIFIER CATALOG ###", - "Copy SQL identifiers only from this catalog or the following DDL.", - "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", - f"object_type: {context.get('object_type', '')}", - ] - if table_name: - lines.append(f"table: {table_name}") - if column_names: - lines.append("columns:") - lines.extend(f"- {column_name}" for column_name in column_names) - if relationship_constraints: - lines.append("relationships:") - lines.extend(f"- {constraint}" for constraint in relationship_constraints) - lines.extend( - [ - "If a needed table, column, or relationship is not listed here or declared in the following DDL, return null for sql.", - "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", - "", - ] - ) - return "\n".join(lines) - - -def _format_identifier_contract(context: dict) -> str: - contract = context.get("sql_identifier_contract", {}) - table_name = contract.get("sql_table_name_use_exactly") - column_names = contract.get("sql_column_names_use_exactly") or [ - column["sql_column_name_use_exactly"] - for column in context.get("columns", []) - if column.get("sql_column_name_use_exactly") - ] - relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ - relationship["sql_relationship_constraint_use_exactly"] - for relationship in context.get("relationships", []) - if relationship.get("sql_relationship_constraint_use_exactly") - ] - - lines = [ - "WREN SQL IDENTIFIER CONTRACT", - f"object_type: {context.get('object_type', '')}", - ] - if table_name: - lines.append(f"sql_table_name_use_exactly: {table_name}") - if column_names: - lines.append("sql_column_names_use_exactly:") - lines.extend(f"- {column_name}" for column_name in column_names) - if relationship_constraints: - lines.append("relationship_constraints_use_exactly:") - lines.extend( - f"- {relationship_constraint}" - for relationship_constraint in relationship_constraints - ) - lines.extend( - [ - "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", - "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", - "END WREN SQL IDENTIFIER CONTRACT", - "", - ] + f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" ) - return "\n".join(lines) - - -def _format_prompt_identifier_context( - table_name: str, - column_names: list[str], - relationship_constraints: list[str] | None = None, -) -> str: - lines = [ - f"table: {table_name}", - "columns:", - ] - lines.extend(f"- {column_name}" for column_name in column_names) - - if relationship_constraints: - lines.append("relationships:") - lines.extend(f"- {constraint}" for constraint in relationship_constraints) - - return "\n".join(lines) - - -def _table_identifier_context( - content: dict, - columns: Optional[set[str]] = None, - tables: Optional[set[str]] = None, -) -> str: - included_columns = _included_columns(content, columns, tables) - included_relationships = _included_relationships(content, tables) - return _format_prompt_identifier_context( - content["name"], - [column["name"] for column in included_columns], - [relationship["constraint"] for relationship in included_relationships], - ) - - -def _semantic_object_identifier_context(content: dict) -> str: - columns = [ - column["name"] - for column in content.get("columns", []) - if column.get("name") and column.get("data_type", "").lower() != "unknown" - ] - return _format_prompt_identifier_context(content["name"], columns) - - -def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: - relationship_columns = { - column.get("column") - for column in content["columns"] - if column["type"] == "FOREIGN_KEY" - and (not tables or set(column.get("tables", [])).issubset(tables)) - } - relationship_columns.discard(None) - return relationship_columns - - -def _included_columns( - content: dict, columns: Optional[set[str]], tables: Optional[set[str]] -) -> list[dict]: - relationship_columns = _included_relationship_columns(content, tables) - return [ - column - for column in content["columns"] - if column["type"] == "COLUMN" - and ( - not columns - or column["name"] in columns - or column["name"] in relationship_columns - or column["is_primary_key"] - ) - and column["data_type"].lower() != "unknown" - ] - - -def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: - return [ - column - for column in content["columns"] - if column["type"] == "FOREIGN_KEY" - and (not tables or set(column.get("tables", [])).issubset(tables)) - ] - - -def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: - executable_columns = { - column["name"] - for column in content["columns"] - if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" - } - return bool(columns) and columns.issubset(executable_columns) - - -def _build_table_retrieval_context( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None -) -> tuple[str, bool, bool]: - ddl, has_calculated_field, has_json_field = build_table_ddl( - content, - columns=columns, - tables=tables, - include_semantic_comments=False, - ) - included_columns = _included_columns(content, columns, tables) - included_relationships = _included_relationships(content, tables) - context = _format_semantic_context( - { - "object_type": "model", - "sql_identifier_contract": { - "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in included_columns - ], - "relationship_constraints_use_exactly": [ - relationship["constraint"] - for relationship in included_relationships - ], - }, - "semantic_context_not_sql_identifiers": { - "description": content["comment"], - }, - "columns": [ - { - "sql_column_name_use_exactly": column["name"], - "data_type": get_engine_supported_data_type(column["data_type"]), - "is_primary_key": column["is_primary_key"], - "semantic_context_not_sql_identifier": column["comment"], - } - for column in included_columns - ], - "relationships": [ - { - "semantic_context_not_sql_identifier": relationship["comment"], - "sql_relationship_constraint_use_exactly": relationship[ - "constraint" - ], - "related_models_use_exactly": relationship.get("tables", []), - } - for relationship in included_relationships - ], - } - ) - return f"{context}{ddl}", has_calculated_field, has_json_field ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: + if histories: + previous_query_summaries = [history.question for history in histories] + else: + previous_query_summaries = [] + + query = "\n".join(previous_query_summaries) + "\n" + query + return await embedder.run(query) else: return {} @@ -486,149 +186,34 @@ async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any, - embedding: dict, mdl_hash: str | None = None, ) -> list[Document]: - table_names = _table_names_from_description_documents( - table_retrieval.get("documents", []) - ) - documents = [] - if embedding and not table_names: - documents = await _retrieve_semantic_schema_documents( - embedding, project_id, dbschema_retriever, mdl_hash - ) - table_names = _table_names_from_schema_documents(documents) - - if table_names: - retrieved_table_names = set() - pending_table_names = table_names - - while pending_table_names: - retrieved_table_names.update(pending_table_names) - retrieved_documents = await _retrieve_schema_documents( - pending_table_names, project_id, dbschema_retriever, mdl_hash - ) - documents = _dedupe_documents(documents + retrieved_documents) - pending_table_names = [ - table_name - for table_name in _related_table_names(documents) - if table_name not in retrieved_table_names - ] - - return documents - - return [] - - -async def _retrieve_semantic_schema_documents( - embedding: dict, - project_id: str, - dbschema_retriever: Any, - mdl_hash: str | None = None, -) -> list[Document]: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - ], - } - - filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) - - results = await dbschema_retriever.run( - query_embedding=embedding.get("embedding"), - filters=filters, - ) - return results["documents"] - - -def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: - table_names = [] - seen = set() - - for document in documents: - table_name = document.meta.get("name") - if not table_name: - content = ast.literal_eval(document.content) - table_name = content.get("name") - - if table_name and table_name not in seen: - table_names.append(table_name) - seen.add(table_name) - - return table_names - - -def _table_names_from_description_documents(documents: list[Document]) -> list[str]: + tables = table_retrieval.get("documents", []) table_names = [] - seen = set() - - for document in documents: - content = ast.literal_eval(document.content) - table_name = content["name"] - if table_name not in seen: - table_names.append(table_name) - seen.add(table_name) + for table in tables: + content = ast.literal_eval(table.content) + table_names.append(content["name"]) - return table_names - - -async def _retrieve_schema_documents( - table_names: list[str], - project_id: str, - dbschema_retriever: Any, - mdl_hash: str | None = None, -) -> list[Document]: table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names ] - if not table_name_conditions: - return [] - - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } - - filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) - - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] - - -def _dedupe_documents(documents: list[Document]) -> list[Document]: - deduped = {} - for document in documents: - key = (document.meta.get("name"), document.content) - deduped[key] = document - - return list(deduped.values()) - - -def _related_table_names(documents: list[Document]) -> list[str]: - related_table_names = [] - seen = set() - - for document in documents: - content = ast.literal_eval(document.content) - if content.get("type") != "TABLE_COLUMNS": - continue + if table_name_conditions: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } - for column in content.get("columns", []): - if column.get("type") != "FOREIGN_KEY": - continue + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) - for table_name in column.get("tables", []): - if table_name not in seen: - related_table_names.append(table_name) - seen.add(table_name) + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] - return related_table_names + return [] @observe() @@ -674,14 +259,11 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context(table_schema) - ) + ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) retrieval_results.append( { "table_name": table_schema["name"], "table_ddl": ddl, - "identifier_context": _table_identifier_context(table_schema), } ) if _has_calculated_field: @@ -697,7 +279,6 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), - "identifier_context": _semantic_object_identifier_context(content), } ) has_metric = True @@ -706,7 +287,6 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), - "identifier_context": _semantic_object_identifier_context(content), } ) @@ -742,10 +322,16 @@ def prompt( ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ - _build_table_retrieval_context(construct_db_schema)[0] + build_table_ddl(construct_db_schema)[0] for construct_db_schema in construct_db_schemas ] + previous_query_summaries = ( + [history.question for history in histories] if histories else [] + ) + + query = "\n".join(previous_query_summaries) + "\n" + query + _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: @@ -773,24 +359,9 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - try: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - ).get("results") - except (IndexError, KeyError, orjson.JSONDecodeError, AttributeError) as e: - logger.warning( - f"Column pruning returned unusable output; using retrieved schemas without column pruning: {e}" - ) - columns_and_tables_needed = None - - if not isinstance(columns_and_tables_needed, list): - logger.warning( - "Column pruning output omitted results; using retrieved schemas without column pruning." - ) - return _build_retrieval_results_without_column_pruning( - construct_db_schemas, - dbschema_retrieval, - ) + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + )["results"] # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -806,21 +377,12 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - columns = set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ) - columns = ( - columns - if _selected_columns_are_executable(table_schema, columns) - else None - ) - - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context( - table_schema, - columns=columns, - tables=tables, - ) + ddl, _has_calculated_field, _has_json_field = build_table_ddl( + table_schema, + columns=set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ), + tables=tables, ) if _has_calculated_field: has_calculated_field = True @@ -831,38 +393,28 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, - "identifier_context": _table_identifier_context( - table_schema, - columns=columns, - tables=tables, - ), } ) for document in dbschema_retrieval: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - "identifier_context": _semantic_object_identifier_context( - content - ), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - "identifier_context": _semantic_object_identifier_context( - content - ), - } - ) + if document.meta["name"] in columns_and_tables_needed: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) return { "retrieval_results": retrieval_results, @@ -883,61 +435,6 @@ def construct_retrieval_results( } -def _build_retrieval_results_without_column_pruning( - construct_db_schemas: list[dict], - dbschema_retrieval: list[Document], -) -> dict[str, Any]: - retrieval_results = [] - has_calculated_field = False - has_metric = False - has_json_field = False - - for table_schema in construct_db_schemas: - if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context(table_schema) - ) - retrieval_results.append( - { - "table_name": table_schema["name"], - "table_ddl": ddl, - "identifier_context": _table_identifier_context(table_schema), - } - ) - if _has_calculated_field: - has_calculated_field = True - if _has_json_field: - has_json_field = True - - for document in dbschema_retrieval: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - "identifier_context": _semantic_object_identifier_context(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - "identifier_context": _semantic_object_identifier_context(content), - } - ) - - return { - "retrieval_results": retrieval_results, - "has_calculated_field": has_calculated_field, - "has_metric": has_metric, - "has_json_field": has_json_field, - } - - ## End of Pipeline class MatchingTableContents(BaseModel): chain_of_thought_reasoning: list[str] diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 7e73f0bc65..f13107849d 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -87,7 +87,7 @@ async def _validate_question( ) return None - async def _document_retrieval() -> tuple[list[str], list[str], bool, bool, bool]: + async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], project_id=project_id, @@ -95,21 +95,10 @@ async def _document_retrieval() -> tuple[list[str], list[str], bool, bool, bool] _retrieval_result = retrieval_result.get("construct_retrieval_results", {}) documents = _retrieval_result.get("retrieval_results", []) table_ddls = [document.get("table_ddl") for document in documents] - identifier_contexts = [ - document.get("identifier_context") - for document in documents - if document.get("identifier_context") - ] has_calculated_field = _retrieval_result.get("has_calculated_field", False) has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) - return ( - table_ddls, - identifier_contexts, - has_calculated_field, - has_metric, - has_json_field, - ) + return table_ddls, has_calculated_field, has_metric, has_json_field async def _sql_pairs_retrieval() -> list[dict]: sql_pairs_result = await self._pipelines["sql_pairs_retrieval"].run( @@ -134,13 +123,7 @@ async def _instructions_retrieval() -> list[dict]: _sql_pairs_retrieval(), _instructions_retrieval(), ) - ( - table_ddls, - identifier_contexts, - has_calculated_field, - has_metric, - has_json_field, - ) = _document + table_ddls, has_calculated_field, has_metric, has_json_field = _document if self._allow_sql_functions_retrieval: sql_functions = await self._pipelines["sql_functions_retrieval"].run( @@ -160,7 +143,6 @@ async def _instructions_retrieval() -> list[dict]: query=candidate["question"], contexts=table_ddls, project_id=project_id, - identifier_contexts=identifier_contexts, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, From 7879c7f8bf2b0d7a214d83870355c2a42982e847 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 17 Aug 2026 02:47:54 +0530 Subject: [PATCH 1043/1087] Revert "Simplify LiteLLM request kwargs handling" This reverts commit 9606d406b3657c0b81908baa63add59fd118ebf8. --- wren-ai-service/src/providers/llm/litellm.py | 91 +++++++++++++++++--- 1 file changed, 80 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index e94918b8c5..fc9406b75b 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -19,6 +19,10 @@ from src.utils import extract_braces_content, remove_trailing_slash +def _is_openai_api_base(api_base: Optional[str]) -> bool: + return bool(api_base) and "api.openai.com" in api_base.lower() + + @provider("litellm_llm") class LitellmLLMProvider(LLMProvider): def __init__( @@ -65,10 +69,42 @@ def get_generator( generation_kwargs: Optional[Dict[str, Any]] = None, streaming_callback: Optional[Callable[[StreamingChunk], None]] = None, ): - combined_generation_kwargs = { - **(generation_kwargs or {}), - **(self._model_kwargs or {}), - } + component_generation_kwargs = generation_kwargs or {} + + def _normalize_generation_kwargs( + kwargs: Optional[Dict[str, Any]], + explicit_response_format: bool = False, + ) -> Dict[str, Any]: + normalized = dict(kwargs or {}) + preserve_json_schema = normalized.pop("preserve_json_schema", False) + response_format = normalized.get("response_format") + + # Plain text is the default chat-completions behavior. + # Some OpenAI-compatible endpoints reject an explicit + # {"type": "text"} payload or serialize it incorrectly. + if ( + isinstance(response_format, dict) + and response_format.get("type") == "text" + ): + normalized.pop("response_format", None) + + if ( + self._api_base + and not _is_openai_api_base(self._api_base) + and isinstance(response_format, dict) + and response_format.get("type") == "json_schema" + and not preserve_json_schema + ): + if explicit_response_format: + normalized["response_format"] = {"type": "json_object"} + else: + normalized.pop("response_format", None) + + if self._api_base and not _is_openai_api_base(self._api_base): + # Some local OpenAI-compatible servers reject non-OpenAI keys. + normalized.pop("speed", None) + + return normalized @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) async def _run( @@ -99,10 +135,42 @@ async def _run( convert_message_to_openai_format(message) for message in messages ] - generation_kwargs = { - **combined_generation_kwargs, - **(generation_kwargs or {}), + runtime_generation_kwargs = generation_kwargs or {} + model_generation_kwargs = self._model_kwargs or {} + explicit_response_format = ( + "response_format" in component_generation_kwargs + or "response_format" in model_generation_kwargs + or "response_format" in runtime_generation_kwargs + ) + merged_generation_kwargs = { + **component_generation_kwargs, + **model_generation_kwargs, + **runtime_generation_kwargs, } + if ( + component_generation_kwargs.get("preserve_json_schema") + and isinstance( + component_generation_kwargs.get("response_format"), dict + ) + and component_generation_kwargs["response_format"].get("type") + == "json_schema" + and "response_format" not in runtime_generation_kwargs + ): + merged_generation_kwargs["response_format"] = ( + component_generation_kwargs["response_format"] + ) + merged_generation_kwargs["preserve_json_schema"] = True + + generation_kwargs = _normalize_generation_kwargs( + merged_generation_kwargs, + explicit_response_format=explicit_response_format, + ) + completion_timeout = generation_kwargs.pop("timeout", self._timeout) + should_stream = ( + streaming_callback is not None + and query_id is not None + and generation_kwargs.pop("stream", True) + ) allowed_openai_params = generation_kwargs.get( "allowed_openai_params", [] @@ -112,9 +180,10 @@ async def _run( completion = await self._router.acompletion( model=self._model, messages=openai_formatted_messages, - stream=streaming_callback is not None, + stream=should_stream, allowed_openai_params=allowed_openai_params, mock_testing_fallbacks=self._enable_fallback_testing, + timeout=completion_timeout, **generation_kwargs, ) else: @@ -123,15 +192,15 @@ async def _run( api_key=self._api_key, api_base=self._api_base, api_version=self._api_version, - timeout=self._timeout, + timeout=completion_timeout, messages=openai_formatted_messages, - stream=streaming_callback is not None, + stream=should_stream, allowed_openai_params=allowed_openai_params, **generation_kwargs, ) completions: List[ChatMessage] = [] - if streaming_callback is not None: + if should_stream: num_responses = generation_kwargs.pop("n", 1) if num_responses > 1: raise ValueError( From 5c53dcef134a2e57f0629eab70e2884a62924d66 Mon Sep 17 00:00:00 2001 From: Ranjitha Date: Mon, 17 Aug 2026 03:20:03 +0530 Subject: [PATCH 1044/1087] prompt updates --- .../src/pipelines/generation/utils/sql.py | 151 ++++++++++++++++-- 1 file changed, 136 insertions(+), 15 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 756c662ee4..eed0980aad 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -235,9 +235,47 @@ async def _classify_generation_result( #### Instructions for Calculated Field #### The first structure is the special column marked as "Calculated Field". You need to interpret the purpose and calculation basis for these columns, then utilize them in the following text-to-sql generation tasks. -First, interpret each calculated field from its expression, data type, comments, aliases, descriptions, and relationship context in the provided DATABASE SCHEMA. -Then, if the user query matches a concept already represented by a calculated field, use that exact calculated field name from DATABASE SCHEMA instead of recreating or inventing the calculation. -Calculated field expressions are semantic definitions; do not copy identifiers from an expression unless they also appear as executable identifiers in the current DATABASE SCHEMA. +First, provide a brief explanation of what each field represents in the context of the schema, including how each field is computed using the relationships between models. +Then, during the following tasks, if the user queries pertain to any calculated fields defined in the database schema, ensure to utilize those calculated fields appropriately in the output SQL queries. +The goal is to accurately reflect the intent of the question in the SQL syntax, leveraging the pre-computed logic embedded within the calculated fields. + +EXAMPLES: +The given schema is created by the SQL command: + +CREATE TABLE orders ( + OrderId VARCHAR PRIMARY KEY, + CustomerId VARCHAR, + -- This column is a Calculated Field + -- column expression: avg(reviews.Score) + Rating DOUBLE, + -- This column is a Calculated Field + -- column expression: count(reviews.Id) + ReviewCount BIGINT, + -- This column is a Calculated Field + -- column expression: count(order_items.ItemNumber) + Size BIGINT, + -- This column is a Calculated Field + -- column expression: count(order_items.ItemNumber) > 1 + Large BOOLEAN, + FOREIGN KEY (CustomerId) REFERENCES customers(Id) +); + +Interpret the columns that are marked as Calculated Fields in the schema: +Rating (DOUBLE) - Calculated as the average score (avg) of the Score field from the reviews table where the reviews are associated with the order. This field represents the overall customer satisfaction rating for the order based on review scores. +ReviewCount (BIGINT) - Calculated by counting (count) the number of entries in the reviews table associated with this order. It measures the volume of customer feedback received for the order. +Size (BIGINT) - Represents the total number of items in the order, calculated by counting the number of item entries (ItemNumber) in the order_items table linked to this order. This field is useful for understanding the scale or size of an order. +Large (BOOLEAN) - A boolean value calculated to check if the number of items in the order exceeds one (count(order_items.ItemNumber) > 1). It indicates whether the order is considered large in terms of item quantity. + +And if the user input queries like these: +1. "How many large orders have been placed by customer with ID 'C1234'?" +2. "What is the average customer rating for orders that were rated by more than 10 reviewers?" + +For the first query: +First try to intepret the user query, the user wants to know the average rating for orders which have attracted significant review activity, specifically those with more than 10 reviews. +Then, according to the above intepretation about the given schema, the term 'Rating' is predefined in the Calculated Field of the 'orders' model. And, the number of reviews is also predefined in the 'ReviewCount' Calculated Field. +So utilize those Calculated Fields in the SQL generation process to give an answer like this: + +SQL Query: SELECT AVG(Rating) FROM orders WHERE ReviewCount > 10 """ _DEFAULT_METRIC_INSTRUCTIONS = """ @@ -267,7 +305,68 @@ async def _classify_generation_result( If the given schema contains the structures marked as 'metric', you should first interpret the metric schema based on the above definition. Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. -Use metric columns exactly as declared in DATABASE SCHEMA. Treat dimensions as grouping/filtering fields and measures as pre-defined numeric outputs. Metric base objects and measure expressions are semantic context only; do not copy identifiers from them unless those identifiers also appear as executable identifiers in the current DATABASE SCHEMA. + +EXAMPLES: +The given schema is created by the SQL command: + +/* This table is a metric */ +/* Metric Base Object: orders */ +CREATE TABLE Revenue ( + -- This column is a dimension + PurchaseTimestamp TIMESTAMP, + -- This column is a dimension + CustomerId VARCHAR, + -- This column is a dimension + Status VARCHAR, + -- This column is a measure + -- expression: sum(order_items.Price) + PriceSum DOUBLE, + -- This column is a measure + -- expression: count(OrderId) + NumberOfOrders BIGINT +); + +Interpret the metric with the understanding of the metric structure: +1. Base Object: orders +This is the primary data source for the metric. +The orders table provides the underlying data from which dimensions and measures are derived. +It is the foundation upon which the metric is built, though it itself is not directly used in queries against the Revenue table. +It shows the reference between the 'Revenue' metric and the 'orders' model. For the user queries pretain to the 'Revenue' of 'orders', the metric should be utilize in the sql generation process. +2. Dimensions +The metric contains the columns marked as 'dimension'. They can be interpreted as below: +- PurchaseTimestamp (TIMESTAMP) + Acts as a temporal dimension, allowing analysis of revenue over time. This can be used to observe trends, seasonal variations, or performance over specific periods. +- CustomerId (VARCHAR) + A key dimension for customer segmentation, it enables the analysis of revenue generated from individual customers or customer groups. +- Status (VARCHAR) + Reflects the current state of an order (e.g., pending, completed, cancelled). This dimension is crucial for analyses that differentiate performance based on order status. +3. Measures +The metric contains the columns marked as 'measure'. They can be interpreted as below: +- PriceSum (DOUBLE) + A financial measure calculated as sum(order_items.Price), representing the total revenue generated from orders. This measure is vital for tracking overall sales performance and is the primary output of interest in many financial and business analyses. +- NumberOfOrders (BIGINT) + A count measure that provides the total number of orders. This is essential for operational metrics, such as assessing the volume of business activity and evaluating the efficiency of sales processes. + +Now, if the user input queries like this: +Question: "What was the total revenue from each customer last month?" + +First try to intepret the user query, the user asks for a breakdown of the total revenue generated by each customer in the previous calendar month. +The user is specifically interested in understanding how much each customer contributed to the total sales during this period. +To answer this question, it is suitable to use the following components from the metric: +1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. +2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. +3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. +So utilize those metric components in the SQL generation process to give an answer like this: + +SQL Query: +SELECT + CustomerId, + PriceSum AS TotalRevenue +FROM + Revenue +WHERE + PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND + PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -278,13 +377,31 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - JSON paths and nested field names must come from the json_fields metadata attached to the exact JSON column in DATABASE SCHEMA. + - For Example: + DATA SCHEMA: + `/* {"alias":"users","description":"A model representing the users data."} */ + CREATE TABLE users ( + -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} + address JSON + )` + To get the city of address in user table use SQL: + `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` + - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. + - For Example: + DATA SCHEMA + `/* {"alias":"my_table","description":"A test my_table"} */ + CREATE TABLE my_table ( + -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} + elements JSON + )` + To get the number of elements in my_table table use SQL: + `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - - Do not copy JSON examples, placeholder aliases, or nested paths from prior context. Use only the current table name, JSON column name, and json_fields metadata in DATABASE SCHEMA. + - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". - DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". """ @@ -292,22 +409,26 @@ async def _classify_generation_result( sql_samples_instructions = """ #### Instructions for SQL Samples #### -Finally, you will learn from the sample questions provided in the input. These samples demonstrate intent and response style for this specific database. +Finally, you will learn from the sample SQL queries provided in the input. These samples demonstrate best practices and common patterns for querying this specific database. For each sample, you should: 1. Study the question that explains what the query aims to accomplish -2. Use these samples as intent and style context only, but treat the DATABASE SCHEMA as the only valid source of executable table and column names -3. Adapt the intent patterns to match new query requirements while maintaining consistent style and approach -4. Never copy table names, column names, aliases, literal values, placeholders, or functions from samples +2. Analyze the SQL implementation to understand: + - Table structures and relationships used + - Specific functions and operators employed + - Query patterns and techniques demonstrated +3. Use these samples as reference patterns when generating similar queries +4. Adapt the techniques shown in the samples to match new query requirements while maintaining consistent style and approach The samples will help you understand: -- Common analytical intents -- Common aggregation requests -- Preferred answer style +- Preferred table join patterns +- Common aggregation methods +- Specific function usage +- Query structure and formatting conventions -When generating new queries, follow similar intent patterns when applicable, while adapting them to the specific requirements of each new query. +When generating new queries, try to follow similar patterns when applicable, while adapting them to the specific requirements of each new query. -Learn about the user's intent from the samples and generate SQL from the current DATABASE SCHEMA and SQL FUNCTIONS only. +Learn about the usage of the schema structures and generate SQL based on them. """ From e43a554ad968418aad4a1d4c91522e2d11679f71 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 03:51:29 +0530 Subject: [PATCH 1045/1087] Align ask SQL grounding with legacy flow --- .../pipelines/generation/sql_correction.py | 42 +- .../pipelines/generation/sql_generation.py | 13 +- .../src/pipelines/generation/utils/sql.py | 232 ++++-- .../retrieval/db_schema_retrieval.py | 671 ++++++++++++++++-- wren-ai-service/src/web/v1/services/ask.py | 2 + .../src/web/v1/services/ask_feedback.py | 1 + 6 files changed, 821 insertions(+), 140 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 5fb54ac90c..69aa9cd957 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -30,12 +30,20 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query. +You are a Wren SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. ### SQL CORRECTION INSTRUCTIONS ### -1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). -2. Then, generate the syntactically correct ANSI SQL query to correct the error. +1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. +2. Then, generate a syntactically correct Wren SQL query from the user's intent and the current DATABASE SCHEMA. +3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. +4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. +5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. +6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. +7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. +8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA. If the unsupported part is needed to answer the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql instead of substituting non-schema identifiers. +10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -43,10 +51,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": + "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ @@ -74,10 +82,20 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -SQL: {{ invalid_generation_result.sql }} -Error Message: {{ invalid_generation_result.error }} +{% if query %} +User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. +{% endif %} +### FAILED SQL ### +The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. + +### DRY-RUN DIAGNOSTIC ### +The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. + +Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. -Let's think step by step. +Return only the final JSON SQL response. """ @@ -87,12 +105,16 @@ def prompt( documents: List[Document], invalid_generation_result: Dict, prompt_builder: PromptBuilder, + query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( + query=query, documents=documents, invalid_generation_result=invalid_generation_result, + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -171,6 +193,8 @@ async def run( self, contexts: List[Document], invalid_generation_result: Dict[str, str], + query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, @@ -192,7 +216,9 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, + "query": query, "documents": contexts, + "sql_generation_reasoning": sql_generation_reasoning, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 71b501e216..fa4d0915e9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -54,11 +54,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -71,13 +70,11 @@ ### QUESTION ### User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. +Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} -{% endif %} - -Let's think step by step. +Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index eed0980aad..78c3a80584 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -4,8 +4,7 @@ import aiohttp import orjson from haystack import component -from haystack.dataclasses import ChatMessage -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from src.core.engine import ( Engine, @@ -37,13 +36,20 @@ async def run( allow_data_preview: bool = False, ) -> dict: try: - cleaned_generation_result = clean_generation_result(replies[0]) - - # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' - if cleaned_generation_result.startswith("{"): - cleaned_generation_result = orjson.loads(cleaned_generation_result)[ - "sql" - ] + cleaned_generation_result, extraction_error = _extract_sql_response( + clean_generation_result(replies[0]) + ) + if extraction_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": extraction_error, + "correlation_id": "", + }, + } ( valid_generation_result, @@ -72,7 +78,7 @@ async def run( async def _classify_generation_result( self, - generation_result: str, + generation_result: str | None, project_id: str | None = None, mdl_hash: str | None = None, use_dry_plan: bool = False, @@ -84,6 +90,15 @@ async def _classify_generation_result( invalid_generation_result = {} use_dry_run = not allow_data_preview + if not generation_result: + return valid_generation_result, { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": "No grounded SQL was generated from the current schema.", + "correlation_id": "", + } + async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( @@ -111,6 +126,27 @@ async def _classify_generation_result( "correlation_id": "", } elif use_dry_run: + dry_plan_result, error_message = await self._engine.dry_plan( + session, + generation_result, + data_source, + project_id=project_id, + mdl_hash=mdl_hash, + allow_fallback=allow_dry_plan_fallback, + ) + + if not dry_plan_result: + invalid_generation_result = { + "sql": generation_result, + "original_sql": generation_result, + "type": "TIME_OUT" + if error_message.startswith("Request timed out") + else "DRY_PLAN", + "error": error_message, + "correlation_id": "", + } + return valid_generation_result, invalid_generation_result + success, _, addition = await self._engine.execute_sql( generation_result, session, @@ -171,6 +207,79 @@ async def _classify_generation_result( return valid_generation_result, invalid_generation_result +def _extract_sql_response(generation_result: str) -> tuple[str | None, str | None]: + cleaned_generation_result = generation_result.strip() + if not cleaned_generation_result: + return None, "No grounded SQL was generated from the current schema." + + if cleaned_generation_result.startswith("{"): + try: + payload = orjson.loads(cleaned_generation_result) + except orjson.JSONDecodeError: + return ( + None, + "SQL generation response did not include a supported SQL JSON payload.", + ) + + if "sql" in payload: + return payload.get("sql"), None + + if payload.get("name") == "query": + arguments = payload.get("arguments") + if isinstance(arguments, dict): + sql = arguments.get("query") or arguments.get("sql") + if sql: + return sql, None + + return ( + None, + f"SQL generation response did not include a supported SQL field: {payload}", + ) + + if cleaned_generation_result.upper().startswith(("SELECT", "WITH")): + return cleaned_generation_result, None + + return None, "SQL generation response did not include a supported SQL JSON payload." + + +_MANDATORY_SQL_GROUNDING_RULES = """ +### MANDATORY SQL GROUNDING RULES ### +- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. +- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. +- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. +- Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. +- Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. +- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. +- When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. +- In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. +- Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. +- When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. +- The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. +- Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. +- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. +- Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. +- If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. +- When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. +- Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. +- When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- When the same requested result can be answered from multiple schema objects with compatible fields or metrics, include all relevant objects by combining separate result rows with UNION ALL instead of choosing only one object. +- Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. +- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. +- Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. +- SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. +- Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. +- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. +- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. +- If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. +- For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. +- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. +- If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. +- If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. +- If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. +- Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. +""" + + _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. @@ -409,40 +518,34 @@ async def _classify_generation_result( sql_samples_instructions = """ #### Instructions for SQL Samples #### -Finally, you will learn from the sample SQL queries provided in the input. These samples demonstrate best practices and common patterns for querying this specific database. +Finally, you will learn from the sample questions provided in the input. These samples demonstrate intent and response style for this specific database. For each sample, you should: 1. Study the question that explains what the query aims to accomplish -2. Analyze the SQL implementation to understand: - - Table structures and relationships used - - Specific functions and operators employed - - Query patterns and techniques demonstrated -3. Use these samples as reference patterns when generating similar queries -4. Adapt the techniques shown in the samples to match new query requirements while maintaining consistent style and approach +2. Use these samples as intent and style context only, but treat the DATABASE SCHEMA as the only valid source of executable table and column names +3. Adapt the intent patterns to match new query requirements while maintaining consistent style and approach +4. Never copy table names, column names, aliases, literal values, placeholders, or functions from samples The samples will help you understand: -- Preferred table join patterns -- Common aggregation methods -- Specific function usage -- Query structure and formatting conventions +- Common analytical intents +- Common aggregation requests +- Preferred answer style -When generating new queries, try to follow similar patterns when applicable, while adapting them to the specific requirements of each new query. +When generating new queries, follow similar intent patterns when applicable, while adapting them to the specific requirements of each new query. -Learn about the usage of the schema structures and generate SQL based on them. +Learn about the user's intent from the samples and generate SQL from the current DATABASE SCHEMA and SQL FUNCTIONS only. """ sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who is great at thinking deeply and reasoning about the user's question and the database schema, and you provide a step-by-step reasoning plan in order to answer the user's question. +You are a helpful data analyst who explains the user's analytical intent and provides a concise, non-executable reasoning plan for answering the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; -otherwise, you will put the relative timeframe in the SQL query. -3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. +2. Explicitly state requested timeframes in natural language only. Mention exact date/time columns only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language. Mention exact ordering columns or measures only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. @@ -450,9 +553,22 @@ async def _classify_generation_result( 9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. A table name in the reasoning plan must be in this format: `table: `. -13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +12. Mention table names only by writing the literal prefix `table:` followed by an exact table name declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +13. Mention column names only by writing the literal prefix `column:` followed by an exact declared table name, a dot, and an exact column name declared for that table in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +14. Do not mention aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels from comments, SQL samples, failed SQL, or user wording as executable identifiers. +15. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. +16. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. +17. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language and cite exact declared tables or columns only when they are grounded by DATABASE SCHEMA. +18. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, then ground the plan in exact declared schema identifiers. +19. If multiple schema objects are required, identify the exact declared relationship path from DATABASE SCHEMA. If no relationship path is declared, say that the retrieved metadata does not provide a join path. +20. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan unless they also appear exactly in DATABASE SCHEMA. +21. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +22. The reasoning plan is semantic context for intent only, not a source of executable identifiers. SQL generation must re-read DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT before using any identifier. +23. ONLY SHOWING the reasoning plan in bullet points. +24. Do not use the words "assume", "assuming", "likely", "possible", "might", or "example" when describing tables, columns, filters, or SQL. +25. If exact deployed table and column identifiers are not available for a requested part, say only that the retrieved metadata does not support that part. Do not propose a replacement name. +26. Do not write table names or column names from the user's wording unless the same identifier appears exactly in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +27. Do not include code blocks, inline SQL fragments, SELECT statements, WHERE clauses, join clauses, or any query-shaped text in the reasoning plan. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -470,12 +586,13 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: + rules = _DEFAULT_TEXT_TO_SQL_RULES if sql_knowledge is not None: - return _extract_from_sql_knowledge( + rules = _extract_from_sql_knowledge( sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES ) - return _DEFAULT_TEXT_TO_SQL_RULES + return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: @@ -511,38 +628,51 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" -You are a helpful assistant that converts natural language queries into ANSI SQL queries. +You are a helpful assistant that converts natural language queries into Wren SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. +Given the user's question and database schema, generate one grounded Wren SQL query. The DATABASE SCHEMA is the only source of executable identifiers. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. -2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. +3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. +4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. +5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. +6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. +7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. +8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. +11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +13. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. +14. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. Do not create table or column identifiers from the user's wording. If the retrieved schema does not ground the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ class SqlGenerationResult(BaseModel): - sql: str + model_config = ConfigDict(extra="forbid") + + sql: str | None SQL_GENERATION_MODEL_KWARGS = { + "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { "name": "sql_generation_result", + "strict": True, "schema": SqlGenerationResult.model_json_schema(), }, } @@ -564,18 +694,4 @@ def construct_instructions( def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: - messages = [] - for history in histories: - messages.append( - ChatMessage.from_user( - history.question - if hasattr(history, "question") - else history["question"] - ) - ) - messages.append( - ChatMessage.from_assistant( - history.sql if hasattr(history, "sql") else history["sql"] - ) - ) - return messages + return [] diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 22bf940cd8..582bb831ff 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -30,7 +30,13 @@ ### TASK ### You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. -The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. +The database schema includes structural, semantic, and business modeling metadata: +- Models are logical datasets backed by physical tables or SQL definitions. +- Columns are exposed fields, including renamed fields, expressions, primary keys, and calculated fields. +- Relationships are reusable join logic between models. +- Calculated fields are business logic defined once and reused across queries. +- Views are named SQL statements that behave like stable virtual tables. +- Metrics are structured aggregation objects with measures and dimensions. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -40,6 +46,17 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +8. Map the business question to the modeled datasets whose descriptions, aliases, columns, calculated fields, views, metrics, and relationships support the intent. +9. Prefer modeled analytical interfaces such as views and metrics when they expose the fields needed to answer the question. +10. If the answer needs fields, filters, time dimensions, ordering, aggregations, or relationship keys from multiple related datasets, include every required related dataset and the columns needed from each one. +11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. +12. Follow only the relationships shown in the provided schema when selecting columns across datasets. +13. Do not stop at a single top candidate when the question needs multiple related datasets. +14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. +15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. +16. Prefer the set of deployed models, views, metrics, columns, and relationships that best support the current question. +17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. +18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -112,37 +129,374 @@ def _project_filter_conditions( def _build_metric_ddl(content: dict) -> str: - columns_ddl = [ - f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + columns = [ + column for column in content["columns"] - if column["data_type"].lower() - != "unknown" # quick fix: filtering out UNKNOWN column type + if column["data_type"].lower() != "unknown" + ] + context = _format_semantic_context( + { + "object_type": "metric", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in columns + ], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable analytical aggregation interface", + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "semantic_context_not_sql_identifier": column["comment"], + } + for column in columns + ], + } + ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in columns ] return ( - f"{content['comment']}CREATE TABLE {content['name']} (\n " + f"{context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) def _build_view_ddl(content: dict) -> str: + columns = [ + column + for column in content.get("columns", []) + if column.get("name") and column.get("data_type", "").lower() != "unknown" + ] + context = _format_semantic_context( + { + "object_type": "view", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in columns + ], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable virtual table interface", + "description": content["comment"], + "definition_omitted_from_executable_schema": True, + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type( + column.get("data_type") + ), + "semantic_context_not_sql_identifier": column.get("comment", ""), + } + for column in columns + ], + } + ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" + for column in columns + ] + + return ( + f"{context}CREATE TABLE {content['name']} (\n " + + ",\n ".join(columns_ddl) + + "\n);" + ) + + +def _format_semantic_context(context: dict) -> str: + return ( + "/*\n" + "WREN RETRIEVED SEMANTIC CONTEXT\n" + f"{orjson.dumps(context).decode('utf-8')}\n" + f"{_format_identifier_contract(context)}" + "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" + "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" + "*/\n" + f"{_format_executable_identifier_catalog(context)}" + ) + + +def _format_executable_identifier_catalog(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "### EXECUTABLE WREN IDENTIFIER CATALOG ###", + "Copy SQL identifiers only from this catalog or the following DDL.", + "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"table: {table_name}") + if column_names: + lines.append("columns:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationships:") + lines.extend(f"- {constraint}" for constraint in relationship_constraints) + lines.extend( + [ + "If a needed table, column, or relationship is not listed here or declared in the following DDL, return null for sql.", + "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", + "", + ] + ) + return "\n".join(lines) + + +def _format_identifier_contract(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "WREN SQL IDENTIFIER CONTRACT", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"sql_table_name_use_exactly: {table_name}") + if column_names: + lines.append("sql_column_names_use_exactly:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationship_constraints_use_exactly:") + lines.extend( + f"- {relationship_constraint}" + for relationship_constraint in relationship_constraints + ) + lines.extend( + [ + "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", + "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", + "END WREN SQL IDENTIFIER CONTRACT", + "", + ] + ) + return "\n".join(lines) + + +def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: + relationship_columns = { + column.get("column") + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + } + relationship_columns.discard(None) + return relationship_columns + + +def _included_columns( + content: dict, columns: Optional[set[str]], tables: Optional[set[str]] +) -> list[dict]: + relationship_columns = _included_relationship_columns(content, tables) + return [ + column + for column in content["columns"] + if column["type"] == "COLUMN" + and ( + not columns + or column["name"] in columns + or column["name"] in relationship_columns + or column["is_primary_key"] + ) + and column["data_type"].lower() != "unknown" + ] + + +def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: + return [ + column + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + ] + + +def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: + executable_columns = { + column["name"] + for column in content["columns"] + if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" + } + return bool(columns) and columns.issubset(executable_columns) + + +def _build_table_retrieval_context( + content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None +) -> tuple[str, bool, bool]: + ddl, has_calculated_field, has_json_field = build_table_ddl( + content, + columns=columns, + tables=tables, + include_semantic_comments=False, + ) + included_columns = _included_columns(content, columns, tables) + included_relationships = _included_relationships(content, tables) + context = _format_semantic_context( + { + "object_type": "model", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in included_columns + ], + "relationship_constraints_use_exactly": [ + relationship["constraint"] + for relationship in included_relationships + ], + }, + "semantic_context_not_sql_identifiers": { + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "is_primary_key": column["is_primary_key"], + "semantic_context_not_sql_identifier": column["comment"], + } + for column in included_columns + ], + "relationships": [ + { + "semantic_context_not_sql_identifier": relationship["comment"], + "sql_relationship_constraint_use_exactly": relationship[ + "constraint" + ], + "related_models_use_exactly": relationship.get("tables", []), + } + for relationship in included_relationships + ], + } + ) + return f"{context}{ddl}", has_calculated_field, has_json_field + + +def _identifier_context(table_name: str, column_names: list[str]) -> str: + return "\n".join( + [f"table: {table_name}", "columns:", *[f"- {name}" for name in column_names]] + ) + + +def _build_retrieval_item(table_schema: dict) -> tuple[dict[str, str], bool, bool]: + ddl, has_calculated_field, has_json_field = _build_table_retrieval_context( + table_schema + ) return ( - f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" + { + "table_name": table_schema["name"], + "table_ddl": ddl, + "identifier_context": _identifier_context( + table_schema["name"], + [ + column["name"] + for column in _included_columns(table_schema, None, None) + ], + ), + }, + has_calculated_field, + has_json_field, ) +def _fallback_retrieval_results( + construct_db_schemas: list[dict], + dbschema_retrieval: list[Document], +) -> dict[str, Any]: + retrieval_results = [] + has_calculated_field = False + has_metric = False + has_json_field = False + + for table_schema in construct_db_schemas: + if table_schema["type"] == "TABLE": + retrieval_item, _has_calculated_field, _has_json_field = ( + _build_retrieval_item(table_schema) + ) + retrieval_results.append(retrieval_item) + if _has_calculated_field: + has_calculated_field = True + if _has_json_field: + has_json_field = True + + for document in dbschema_retrieval: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + ), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], + ), + } + ) + + return { + "retrieval_results": retrieval_results, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: - if histories: - previous_query_summaries = [history.question for history in histories] - else: - previous_query_summaries = [] - - query = "\n".join(previous_query_summaries) + "\n" + query - return await embedder.run(query) else: return {} @@ -186,34 +540,158 @@ async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any, + embedding: dict, mdl_hash: str | None = None, ) -> list[Document]: - tables = table_retrieval.get("documents", []) + table_names = _table_names_from_description_documents( + table_retrieval.get("documents", []) + ) + documents = [] + if embedding and not table_names: + documents = await _retrieve_semantic_schema_documents( + embedding, project_id, mdl_hash, dbschema_retriever + ) + table_names = _table_names_from_schema_documents(documents) + + if table_names: + retrieved_table_names = set() + pending_table_names = table_names + + while pending_table_names: + retrieved_table_names.update(pending_table_names) + retrieved_documents = await _retrieve_schema_documents( + pending_table_names, project_id, mdl_hash, dbschema_retriever + ) + documents = _dedupe_documents(documents + retrieved_documents) + pending_table_names = [ + table_name + for table_name in _related_table_names(documents) + if table_name not in retrieved_table_names + ] + + return documents + + return [] + + +async def _retrieve_semantic_schema_documents( + embedding: dict, + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, +) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + + results = await dbschema_retriever.run( + query_embedding=embedding.get("embedding"), + filters=filters, + ) + return results["documents"] + + +def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: + table_names = [] + seen = set() + + for document in documents: + table_name = document.meta.get("name") + if not table_name: + content = ast.literal_eval(document.content) + table_name = content.get("name") + + if table_name and table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + + return table_names + + +def _table_names_from_description_documents(documents: list[Document]) -> list[str]: table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + table_name = content["name"] + if table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + + return table_names + +async def _retrieve_schema_documents( + table_names: list[str], + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, +) -> list[Document]: table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names ] - if table_name_conditions: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } + if not table_name_conditions: + return [] + + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } - filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] - return [] + +def _related_table_names(documents: list[Document]) -> list[str]: + related_table_names = [] + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + if content.get("type") != "TABLE_COLUMNS": + continue + + for column in content.get("columns", []): + if column.get("type") != "FOREIGN_KEY": + continue + + for table_name in column.get("tables", []): + if table_name not in seen: + related_table_names.append(table_name) + seen.add(table_name) + + return related_table_names + + +def _dedupe_documents(documents: list[Document]) -> list[Document]: + deduped = [] + seen = set() + + for document in documents: + identity = ( + document.meta.get("type"), + document.meta.get("name"), + document.content, + ) + if identity in seen: + continue + deduped.append(document) + seen.add(identity) + + return deduped @observe() @@ -259,11 +737,20 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) retrieval_results.append( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_context": _identifier_context( + table_schema["name"], + [ + column["name"] + for column in _included_columns(table_schema, None, None) + ], + ), } ) if _has_calculated_field: @@ -279,6 +766,14 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + ), } ) has_metric = True @@ -287,6 +782,15 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], + ), } ) @@ -322,16 +826,10 @@ def prompt( ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: db_schemas = [ - build_table_ddl(construct_db_schema)[0] + _build_table_retrieval_context(construct_db_schema)[0] for construct_db_schema in construct_db_schemas ] - previous_query_summaries = ( - [history.question for history in histories] if histories else [] - ) - - query = "\n".join(previous_query_summaries) + "\n" + query - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: @@ -359,9 +857,15 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] + try: + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + ).get("results") + except orjson.JSONDecodeError: + columns_and_tables_needed = None + + if not columns_and_tables_needed: + return _fallback_retrieval_results(construct_db_schemas, dbschema_retrieval) # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -377,12 +881,22 @@ def construct_retrieval_results( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - ddl, _has_calculated_field, _has_json_field = build_table_ddl( - table_schema, - columns=set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ), - tables=tables, + selected_columns = set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ) + columns = ( + selected_columns + if _selected_columns_are_executable( + table_schema, selected_columns + ) + else None + ) + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context( + table_schema, + columns=columns, + tables=tables, + ) ) if _has_calculated_field: has_calculated_field = True @@ -393,28 +907,53 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_context": _identifier_context( + table_schema["name"], + [ + column["name"] + for column in _included_columns( + table_schema, columns, tables + ) + ], + ), } ) for document in dbschema_retrieval: - if document.meta["name"] in columns_and_tables_needed: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - } - ) + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + ), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], + ), + } + ) return { "retrieval_results": retrieval_results, @@ -468,7 +1007,7 @@ def __init__( llm_provider: LLMProvider, embedder_provider: EmbedderProvider, document_store_provider: DocumentStoreProvider, - table_retrieval_size: int = 10, + table_retrieval_size: int = 50, table_column_retrieval_size: int = 100, **kwargs, ): diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index bac4bb7c41..140fcc9c8a 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -566,6 +566,8 @@ async def ask( "sql_correction" ].run( contexts=table_ddls, + query=user_query, + sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 184c15a37a..27a8a94aa8 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -243,6 +243,7 @@ async def ask_feedback( "sql_correction" ].run( contexts=table_ddls, + query=ask_feedback_request.question, instructions=instructions, invalid_generation_result={ "sql": original_sql, From 5639506cf7e09e551b48129d8bf7d3d489301e0c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 03:55:19 +0530 Subject: [PATCH 1046/1087] Fix SQL generation startup import --- wren-ai-service/src/pipelines/generation/utils/sql.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 78c3a80584..3a4e80d617 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -10,6 +10,7 @@ Engine, clean_generation_result, ) +from src.providers.llm import ChatMessage from src.pipelines.retrieval.sql_knowledge import SqlKnowledge from src.web.v1.services.ask import AskHistory From 4eea2ed7173190aa2f0ce0a8882cdb2a9e0c361b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 04:11:45 +0530 Subject: [PATCH 1047/1087] Restore grounded SQL reasoning handoff --- .../generation/followup_sql_generation.py | 11 +++++----- .../followup_sql_generation_reasoning.py | 3 +-- .../pipelines/generation/sql_generation.py | 8 +++++++- .../generation/sql_generation_reasoning.py | 3 +-- .../src/pipelines/generation/utils/sql.py | 20 ++++++++++++++++--- .../test_sql_generation_post_processor.py | 10 ++++++++-- 6 files changed, 40 insertions(+), 15 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 0e50a12806..7890636a77 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -60,11 +60,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} -Summary: -{{sample.summary}} -SQL: -{{sample.sql}} +Question: +{{sample.question}} {% endfor %} {% endif %} @@ -81,7 +80,9 @@ ### REASONING PLAN ### {{ sql_generation_reasoning }} -Let's think step by step. +Follow the grounded reasoning plan step by step when it cites exact identifiers from DATABASE SCHEMA. If the reasoning plan contains a table or column that is not declared in DATABASE SCHEMA, ignore that identifier and use DATABASE SCHEMA as the source of truth. + +Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 7ec0b8b308..136bd4a5bc 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -30,11 +30,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index fa4d0915e9..15b7a74996 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -72,7 +72,13 @@ User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -Do not generate SQL from a reasoning plan. The reasoning plan is not executable context and cannot provide table names, column names, filters, functions, joins, or examples. + +{% if sql_generation_reasoning %} +### REASONING PLAN ### +{{ sql_generation_reasoning }} + +Follow the grounded reasoning plan step by step when it cites exact identifiers from DATABASE SCHEMA. If the reasoning plan contains a table or column that is not declared in DATABASE SCHEMA, ignore that identifier and use DATABASE SCHEMA as the source of truth. +{% endif %} Return only the final JSON SQL response. """ 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 6fdd8d6a01..45df7584c9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -29,11 +29,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 3a4e80d617..df2997729b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -165,7 +165,7 @@ async def _classify_generation_result( else: error_message = addition.get("error_message", "") invalid_generation_result = { - "sql": addition.get("error_sql", generation_result), + "sql": generation_result, "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") @@ -196,7 +196,7 @@ async def _classify_generation_result( else "PREVIEW_FAILED" ) invalid_generation_result = { - "sql": addition.get("error_sql", generation_result), + "sql": generation_result, "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") @@ -695,4 +695,18 @@ def construct_instructions( def construct_ask_history_messages( histories: list[AskHistory] | list[dict], ) -> list[ChatMessage]: - return [] + messages = [] + for history in histories: + messages.append( + ChatMessage.from_user( + history.question + if hasattr(history, "question") + else history["question"] + ) + ) + messages.append( + ChatMessage.from_assistant( + history.sql if hasattr(history, "sql") else history["sql"] + ) + ) + return messages diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py index 7e1be2ca7a..02bfed4db0 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -4,6 +4,7 @@ import pytest from src.core.engine import Engine +from src.providers.llm import ChatRole from src.pipelines.generation.utils.sql import ( SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, @@ -64,10 +65,15 @@ def test_sql_generation_model_kwargs_preserve_strict_schema(): assert schema["additionalProperties"] is False -def test_construct_ask_history_messages_matches_legacy_empty_context(): +def test_construct_ask_history_messages_matches_legacy_context(): histories = [{"question": "q", "sql": "SELECT 1"}] - assert construct_ask_history_messages(histories) == [] + messages = construct_ask_history_messages(histories) + + assert [(message.role, message.content) for message in messages] == [ + (ChatRole.USER, "q"), + (ChatRole.ASSISTANT, "SELECT 1"), + ] @pytest.mark.asyncio From e8279e29e075759d7f007578146c3755d1476582 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 04:28:56 +0530 Subject: [PATCH 1048/1087] Keep ask SQL grounded in selected schema --- .../src/pipelines/generation/utils/sql.py | 10 ++- .../retrieval/db_schema_retrieval.py | 3 + .../test_sql_generation_post_processor.py | 19 +++++ .../retrieval/test_db_schema_retrieval.py | 75 +++++++++++++++++++ 4 files changed, 105 insertions(+), 2 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index df2997729b..1d07329a8f 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -589,9 +589,15 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: rules = _DEFAULT_TEXT_TO_SQL_RULES if sql_knowledge is not None: - rules = _extract_from_sql_knowledge( - sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES + connector_rules = _extract_from_sql_knowledge( + sql_knowledge, "text_to_sql_rule", "" ) + if connector_rules: + rules = f"""{rules} + +### CONNECTOR SQL KNOWLEDGE ### +Use the following connector-specific knowledge only when it does not conflict with Wren SQL syntax, DATABASE SCHEMA identifiers, SQL FUNCTIONS, or WREN SQL IDENTIFIER CONTRACT. +{connector_rules}""" return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 582bb831ff..d9aa2b0cca 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -922,6 +922,9 @@ def construct_retrieval_results( for document in dbschema_retrieval: content = ast.literal_eval(document.content) + if content["name"] not in tables: + continue + if content["type"] == "METRIC": retrieval_results.append( { diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py index 02bfed4db0..ed88d0d71f 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -9,7 +9,9 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_ask_history_messages, + get_text_to_sql_rules, ) +from src.pipelines.retrieval.sql_knowledge import SqlKnowledge class FakeEngine(Engine): @@ -76,6 +78,23 @@ def test_construct_ask_history_messages_matches_legacy_context(): ] +def test_connector_sql_knowledge_supplements_wren_sql_rules(): + sql_knowledge = SqlKnowledge( + { + "text_to_sql_rule": "Connector-only syntax guidance.", + "instructions": {}, + } + ) + + rules = get_text_to_sql_rules(sql_knowledge) + + assert "Generate Wren SQL only" in rules + assert "Connector-only syntax guidance." in rules + assert rules.index("### SQL RULES ###") < rules.index( + "Connector-only syntax guidance." + ) + + @pytest.mark.asyncio async def test_post_processor_extracts_tool_call_query_argument(): engine = FakeEngine() diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 7c362021e6..5ef00c3f94 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -541,6 +541,14 @@ def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): "chain_of_thought_reasoning": ["Needed field."], "columns": ["stored_attribute"] } + }, + { + "table_name": "semantic_metric", + "table_selection_reason": "Selected metric for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed metric."], + "columns": ["metric_value"] + } } ] } @@ -594,6 +602,73 @@ def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): assert result["has_metric"] is True +def test_construct_retrieval_results_excludes_unselected_metric_when_pruning(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_attribute"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "name": "semantic_metric", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "semantic_metric"}, + ) + ], + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "modeled_dataset" + ] + assert result["has_metric"] is False + + def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_columns(): result = construct_retrieval_results( check_using_db_schemas_without_pruning={}, From 5622cc5aee01ee2360a20c57497e902ec4416f82 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 04:39:17 +0530 Subject: [PATCH 1049/1087] Restore deployed view context in ask retrieval --- .../retrieval/db_schema_retrieval.py | 54 +++++++++++++------ .../retrieval/test_db_schema_retrieval.py | 46 +++++++++++++++- 2 files changed, 83 insertions(+), 17 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index d9aa2b0cca..60ba23ad46 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -175,6 +175,7 @@ def _build_view_ddl(content: dict) -> str: for column in content.get("columns", []) if column.get("name") and column.get("data_type", "").lower() != "unknown" ] + statement = content.get("statement", "") context = _format_semantic_context( { "object_type": "view", @@ -187,7 +188,6 @@ def _build_view_ddl(content: dict) -> str: "semantic_context_not_sql_identifiers": { "role": "stable virtual table interface", "description": content["comment"], - "definition_omitted_from_executable_schema": True, }, "columns": [ { @@ -201,16 +201,18 @@ def _build_view_ddl(content: dict) -> str: ], } ) - columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" - for column in columns - ] + if columns: + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" + for column in columns + ] + return ( + f"{context}CREATE TABLE {content['name']} (\n " + + ",\n ".join(columns_ddl) + + "\n);" + ) - return ( - f"{context}CREATE TABLE {content['name']} (\n " - + ",\n ".join(columns_ddl) - + "\n);" - ) + return f"{context}{content['comment']}CREATE VIEW {content['name']}\nAS {statement}" def _format_semantic_context(context: dict) -> str: @@ -429,6 +431,16 @@ def _build_retrieval_item(table_schema: dict) -> tuple[dict[str, str], bool, boo ) +def _build_pruning_context(content: dict) -> str: + if content["type"] == "TABLE": + return _build_table_retrieval_context(content)[0] + if content["type"] == "METRIC": + return _build_metric_ddl(content) + if content["type"] == "VIEW": + return _build_view_ddl(content) + return "" + + def _fallback_retrieval_results( construct_db_schemas: list[dict], dbschema_retrieval: list[Document], @@ -715,9 +727,16 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: db_schemas[document.meta["name"]]["columns"] = content["columns"] else: db_schemas[document.meta["name"]]["columns"] += content["columns"] + elif content["type"] in {"VIEW", "METRIC"}: + db_schemas[document.meta["name"]] = content # remove incomplete schemas - db_schemas = {k: v for k, v in db_schemas.items() if "type" in v and "columns" in v} + db_schemas = { + k: v + for k, v in db_schemas.items() + if v.get("type") in {"VIEW", "METRIC"} + or (v.get("type") == "TABLE" and "columns" in v) + } return list(db_schemas.values()) @@ -825,10 +844,15 @@ def prompt( histories: list[AskHistory], ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: - db_schemas = [ - _build_table_retrieval_context(construct_db_schema)[0] - for construct_db_schema in construct_db_schemas - ] + db_schemas = list( + filter( + None, + [ + _build_pruning_context(construct_db_schema) + for construct_db_schema in construct_db_schemas + ], + ) + ) _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 5ef00c3f94..90313db3f5 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -6,6 +6,7 @@ from src.pipelines.retrieval.db_schema_retrieval import ( _build_view_ddl, check_using_db_schemas_without_pruning, + construct_db_schemas, construct_retrieval_results, dbschema_retrieval, embedding, @@ -81,7 +82,7 @@ def test_table_selection_prompt_prefers_best_schema_supported_dataset_set(): ) -def test_view_schema_context_uses_declared_view_columns_not_view_definition(): +def test_view_schema_context_uses_declared_view_columns_when_available(): result = _build_view_ddl( { "type": "VIEW", @@ -101,10 +102,51 @@ def test_view_schema_context_uses_declared_view_columns_not_view_definition(): assert "CREATE TABLE retrieved_view" in result assert "visible_attribute VARCHAR" in result assert "sql_column_names_use_exactly" in result - assert "definition_omitted_from_executable_schema" in result assert "NON_EXECUTABLE_DEFINITION_TOKEN" not in result +def test_view_schema_context_uses_deployed_view_statement_without_declared_columns(): + result = _build_view_ddl( + { + "type": "VIEW", + "comment": "Semantic description.", + "name": "retrieved_view", + "statement": "SELECT modeled_column FROM deployed_model", + } + ) + + assert "CREATE VIEW retrieved_view" in result + assert "AS SELECT modeled_column FROM deployed_model" in result + assert "sql_table_name_use_exactly: retrieved_view" in result + + +def test_construct_db_schemas_keeps_deployed_views_for_column_pruning(): + result = construct_db_schemas( + [ + Document( + content=str( + { + "type": "VIEW", + "comment": "", + "name": "retrieved_view", + "statement": "SELECT modeled_column FROM deployed_model", + } + ), + meta={"type": "TABLE_SCHEMA", "name": "retrieved_view"}, + ) + ] + ) + + assert result == [ + { + "type": "VIEW", + "comment": "", + "name": "retrieved_view", + "statement": "SELECT modeled_column FROM deployed_model", + } + ] + + @pytest.mark.asyncio async def test_table_retrieval_fetches_explicit_table_descriptions(): class Retriever: From ec277dc6944bc5de7d703a94b346aa2c3114579f Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 04:49:26 +0530 Subject: [PATCH 1050/1087] Stop SQL generation from using planning text --- .../generation/followup_sql_generation.py | 6 --- .../pipelines/generation/sql_generation.py | 8 ---- .../generation/test_sql_generation_prompt.py | 40 +++++++++++++++++++ 3 files changed, 40 insertions(+), 14 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py 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 7890636a77..af566a91db 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -77,11 +77,6 @@ ### QUESTION ### User's Follow-up Question: {{ query }} -### REASONING PLAN ### -{{ sql_generation_reasoning }} - -Follow the grounded reasoning plan step by step when it cites exact identifiers from DATABASE SCHEMA. If the reasoning plan contains a table or column that is not declared in DATABASE SCHEMA, ignore that identifier and use DATABASE SCHEMA as the source of truth. - Return only the final JSON SQL response. """ @@ -104,7 +99,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 15b7a74996..6258583790 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -73,13 +73,6 @@ Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} - -Follow the grounded reasoning plan step by step when it cites exact identifiers from DATABASE SCHEMA. If the reasoning plan contains a table or column that is not declared in DATABASE SCHEMA, ignore that identifier and use DATABASE SCHEMA as the source of truth. -{% endif %} - Return only the final JSON SQL response. """ @@ -102,7 +95,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py new file mode 100644 index 0000000000..ea1caa551d --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py @@ -0,0 +1,40 @@ +from haystack.components.builders.prompt_builder import PromptBuilder + +from src.pipelines.generation.followup_sql_generation import ( + prompt as followup_sql_generation_prompt, +) +from src.pipelines.generation.followup_sql_generation import ( + text_to_sql_with_followup_user_prompt_template, +) +from src.pipelines.generation.sql_generation import ( + prompt as sql_generation_prompt, +) +from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template + + +def test_sql_generation_prompt_does_not_include_reasoning_identifiers(): + result = sql_generation_prompt( + query="Show invoices", + documents=['CREATE TABLE "deployed_invoice_model" ("invoice_id" VARCHAR);'], + prompt_builder=PromptBuilder(template=sql_generation_user_prompt_template), + sql_generation_reasoning="SELECT * FROM invoices", + ) + + assert "deployed_invoice_model" in result["prompt"] + assert "SELECT * FROM invoices" not in result["prompt"] + assert "### REASONING PLAN ###" not in result["prompt"] + + +def test_followup_sql_generation_prompt_does_not_include_reasoning_identifiers(): + result = followup_sql_generation_prompt( + query="Show invoices", + documents=['CREATE TABLE "deployed_invoice_model" ("invoice_id" VARCHAR);'], + sql_generation_reasoning="SELECT * FROM invoices", + prompt_builder=PromptBuilder( + template=text_to_sql_with_followup_user_prompt_template + ), + ) + + assert "deployed_invoice_model" in result["prompt"] + assert "SELECT * FROM invoices" not in result["prompt"] + assert "### REASONING PLAN ###" not in result["prompt"] From 3682160df9a3050458619c2ec9fa1d33ddc1eaec Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 04:50:26 +0530 Subject: [PATCH 1051/1087] Remove SQL generation prompt test file --- .../generation/test_sql_generation_prompt.py | 40 ------------------- 1 file changed, 40 deletions(-) delete mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py deleted file mode 100644 index ea1caa551d..0000000000 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_prompt.py +++ /dev/null @@ -1,40 +0,0 @@ -from haystack.components.builders.prompt_builder import PromptBuilder - -from src.pipelines.generation.followup_sql_generation import ( - prompt as followup_sql_generation_prompt, -) -from src.pipelines.generation.followup_sql_generation import ( - text_to_sql_with_followup_user_prompt_template, -) -from src.pipelines.generation.sql_generation import ( - prompt as sql_generation_prompt, -) -from src.pipelines.generation.sql_generation import sql_generation_user_prompt_template - - -def test_sql_generation_prompt_does_not_include_reasoning_identifiers(): - result = sql_generation_prompt( - query="Show invoices", - documents=['CREATE TABLE "deployed_invoice_model" ("invoice_id" VARCHAR);'], - prompt_builder=PromptBuilder(template=sql_generation_user_prompt_template), - sql_generation_reasoning="SELECT * FROM invoices", - ) - - assert "deployed_invoice_model" in result["prompt"] - assert "SELECT * FROM invoices" not in result["prompt"] - assert "### REASONING PLAN ###" not in result["prompt"] - - -def test_followup_sql_generation_prompt_does_not_include_reasoning_identifiers(): - result = followup_sql_generation_prompt( - query="Show invoices", - documents=['CREATE TABLE "deployed_invoice_model" ("invoice_id" VARCHAR);'], - sql_generation_reasoning="SELECT * FROM invoices", - prompt_builder=PromptBuilder( - template=text_to_sql_with_followup_user_prompt_template - ), - ) - - assert "deployed_invoice_model" in result["prompt"] - assert "SELECT * FROM invoices" not in result["prompt"] - assert "### REASONING PLAN ###" not in result["prompt"] From 7a4c15d1217277a1edb55f80b229fbb6c3cb27e0 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 05:20:42 +0530 Subject: [PATCH 1052/1087] Ground SQL generation in retrieved schema --- .../generation/followup_sql_generation.py | 3 + .../pipelines/generation/sql_correction.py | 2 + .../pipelines/generation/sql_generation.py | 2 + .../pipelines/generation/sql_regeneration.py | 21 +- .../src/pipelines/generation/utils/sql.py | 232 +++++++++++++++++- .../retrieval/db_schema_retrieval.py | 101 ++++++-- 6 files changed, 327 insertions(+), 34 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 af566a91db..df0013724a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -76,6 +76,7 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, relationships, and history only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, prior failed SQL, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Return only the final JSON SQL response. """ @@ -142,6 +143,7 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, use_dry_plan: bool = False, @@ -154,6 +156,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + contexts=documents, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 69aa9cd957..055224466b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -142,6 +142,7 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, use_dry_plan: bool = False, @@ -154,6 +155,7 @@ async def post_process( use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, + contexts=documents, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 6258583790..7cf302a733 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -134,6 +134,7 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, use_dry_plan: bool = False, @@ -148,6 +149,7 @@ async def post_process( data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, + contexts=documents, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 9926d69419..acc63a91d1 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -34,18 +34,16 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query, -please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. -While generating the new SQL query, you should use the original SQL query as a reference. -While generating the new SQL query, make sure to use the database schema to generate the SQL query. +You are a Wren SQL expert. Generate a grounded Wren SQL query from the current DATABASE SCHEMA and the requested adjustment intent. +The DATABASE SCHEMA is the only source of executable table and column identifiers. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be a JSON object. Return null for sql if the requested adjustment cannot be grounded in DATABASE SCHEMA. {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ @@ -77,11 +75,10 @@ def get_sql_regeneration_system_prompt( {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -93,10 +90,10 @@ def get_sql_regeneration_system_prompt( {% endif %} ### QUESTION ### -SQL generation reasoning: {{ sql_generation_reasoning }} -Original SQL query: {{ sql }} +Adjustment intent: {{ sql_generation_reasoning }} +The previous SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. -Let's think step by step. +Regenerate from the adjustment intent and current DATABASE SCHEMA only. Return only the final JSON SQL response. """ @@ -157,6 +154,7 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, + documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, ) -> dict: @@ -164,6 +162,7 @@ async def post_process( regenerate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + contexts=documents, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1d07329a8f..65a7a857b4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -3,15 +3,18 @@ import aiohttp import orjson +import sqlparse from haystack import component from pydantic import BaseModel, ConfigDict +from sqlparse.sql import Identifier, IdentifierList, TokenList +from sqlparse.tokens import Comment, Keyword from src.core.engine import ( Engine, clean_generation_result, ) -from src.providers.llm import ChatMessage from src.pipelines.retrieval.sql_knowledge import SqlKnowledge +from src.providers.llm import ChatMessage from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") @@ -35,6 +38,7 @@ async def run( allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, + contexts: list[str] | None = None, ) -> dict: try: cleaned_generation_result, extraction_error = _extract_sql_response( @@ -52,6 +56,20 @@ async def run( }, } + schema_catalog = _SchemaCatalog.from_contexts(contexts or []) + grounding_error = schema_catalog.validate_sql(cleaned_generation_result) + if grounding_error: + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": cleaned_generation_result or "", + "original_sql": cleaned_generation_result or "", + "type": "SCHEMA_GROUNDING", + "error": grounding_error, + "correlation_id": "", + }, + } + ( valid_generation_result, invalid_generation_result, @@ -208,6 +226,216 @@ async def _classify_generation_result( return valid_generation_result, invalid_generation_result +class _SchemaCatalog: + def __init__(self, tables: dict[str, set[str]]): + self._tables = tables + + @classmethod + def from_contexts(cls, contexts: list[str]) -> "_SchemaCatalog": + tables: dict[str, set[str]] = {} + current_table: str | None = None + in_columns = False + + for context in contexts: + for raw_line in context.splitlines(): + line = raw_line.strip() + if line.startswith("table: "): + current_table = line.removeprefix("table: ").strip() + if current_table: + tables.setdefault(current_table, set()) + in_columns = False + continue + + if current_table and line == "columns:": + in_columns = True + continue + + if in_columns and current_table and line.startswith("- "): + column_name = line.removeprefix("- ").strip() + if column_name: + tables.setdefault(current_table, set()).add(column_name) + continue + + if in_columns and line and not line.startswith("- "): + in_columns = False + + return cls(tables) + + def validate_sql(self, sql: str | None) -> str | None: + if not sql or not self._tables: + return None + + parsed_statements = sqlparse.parse(sql) + if not parsed_statements: + return "Generated SQL could not be parsed for schema grounding." + + referenced_tables: set[str] = set() + table_aliases: dict[str, str] = {} + qualified_columns: list[tuple[str, str]] = [] + cte_names: set[str] = set() + + for statement in parsed_statements: + cte_names.update(_extract_cte_names(statement)) + statement_tables, statement_aliases = _extract_table_references(statement) + referenced_tables.update(statement_tables) + table_aliases.update(statement_aliases) + qualified_columns.extend(_extract_qualified_columns(statement)) + + executable_tables = referenced_tables - cte_names + unknown_tables = sorted( + table_name + for table_name in executable_tables + if table_name not in self._tables + ) + if unknown_tables: + return ( + "Generated SQL referenced table(s) not present in the retrieved " + f"schema context: {', '.join(unknown_tables)}." + ) + + unknown_columns = [] + for qualifier, column_name in qualified_columns: + table_name = table_aliases.get(qualifier, qualifier) + if table_name in cte_names: + continue + if table_name in self._tables and self._tables[table_name]: + if column_name not in self._tables[table_name]: + unknown_columns.append(f"{qualifier}.{column_name}") + + if unknown_columns: + return ( + "Generated SQL referenced column(s) not present in the retrieved " + f"schema context: {', '.join(sorted(set(unknown_columns)))}." + ) + + return None + + +def _extract_cte_names(token_list: TokenList) -> set[str]: + cte_names: set[str] = set() + with_seen = False + + for token in token_list.tokens: + if token.is_whitespace or token.ttype in Comment: + continue + + if token.normalized == "WITH": + with_seen = True + continue + + if not with_seen: + continue + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + name = _clean_identifier(identifier.get_name()) + if name: + cte_names.add(name) + break + + if isinstance(token, Identifier): + name = _clean_identifier(token.get_name()) + if name: + cte_names.add(name) + break + + if token.ttype is Keyword: + break + + return cte_names + + +def _extract_table_references(token_list: TokenList) -> tuple[set[str], dict[str, str]]: + table_names: set[str] = set() + aliases: dict[str, str] = {} + expect_table = False + + for token in token_list.tokens: + if token.is_whitespace or token.ttype in Comment: + continue + + if isinstance(token, TokenList): + nested_tables, nested_aliases = _extract_table_references(token) + table_names.update(nested_tables) + aliases.update(nested_aliases) + + if token.ttype is Keyword and token.normalized in { + "FROM", + "JOIN", + "INNER JOIN", + "LEFT JOIN", + "LEFT OUTER JOIN", + "RIGHT JOIN", + "RIGHT OUTER JOIN", + "FULL JOIN", + "FULL OUTER JOIN", + "CROSS JOIN", + }: + expect_table = True + continue + + if not expect_table: + continue + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + _add_table_reference(identifier, table_names, aliases) + expect_table = False + continue + + if isinstance(token, Identifier): + _add_table_reference(token, table_names, aliases) + expect_table = False + continue + + if token.ttype is Keyword: + expect_table = False + + return table_names, aliases + + +def _add_table_reference( + identifier: Identifier, table_names: set[str], aliases: dict[str, str] +) -> None: + table_name = _clean_identifier(identifier.get_real_name()) + alias = _clean_identifier(identifier.get_alias()) + if not table_name: + return + + table_names.add(table_name) + aliases[table_name] = table_name + if alias: + aliases[alias] = table_name + + +def _extract_qualified_columns(token_list: TokenList) -> list[tuple[str, str]]: + columns: list[tuple[str, str]] = [] + + for token in token_list.tokens: + if isinstance(token, Identifier): + parent_name = _clean_identifier(token.get_parent_name()) + column_name = _clean_identifier(token.get_real_name()) + if parent_name and column_name and column_name != "*": + columns.append((parent_name, column_name)) + elif isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + parent_name = _clean_identifier(identifier.get_parent_name()) + column_name = _clean_identifier(identifier.get_real_name()) + if parent_name and column_name and column_name != "*": + columns.append((parent_name, column_name)) + elif isinstance(token, TokenList): + columns.extend(_extract_qualified_columns(token)) + + return columns + + +def _clean_identifier(identifier: str | None) -> str | None: + if identifier is None: + return None + cleaned = identifier.strip().strip('"`[]') + return cleaned or None + + def _extract_sql_response(generation_result: str) -> tuple[str | None, str | None]: cleaned_generation_result = generation_result.strip() if not cleaned_generation_result: @@ -682,7 +910,7 @@ class SqlGenerationResult(BaseModel): "strict": True, "schema": SqlGenerationResult.model_json_schema(), }, - } + }, } diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 60ba23ad46..2a254ecf12 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -4,14 +4,17 @@ from typing import Any, Optional import orjson +import sqlparse import tiktoken from hamilton import base from hamilton.async_driver import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe from pydantic import BaseModel +from sqlparse.sql import Identifier, IdentifierList +from sqlparse.tokens import DML, Comment, Keyword +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider from src.pipelines.common import ( @@ -139,9 +142,7 @@ def _build_metric_ddl(content: dict) -> str: "object_type": "metric", "sql_identifier_contract": { "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in columns - ], + "sql_column_names_use_exactly": [column["name"] for column in columns], }, "semantic_context_not_sql_identifiers": { "role": "stable analytical aggregation interface", @@ -169,6 +170,67 @@ def _build_metric_ddl(content: dict) -> str: ) +def _strip_identifier_quotes(identifier: str | None) -> str | None: + if not identifier: + return identifier + + return identifier.strip().strip('"`[]') + + +def _view_columns_from_statement(statement: str) -> list[dict]: + if not statement: + return [] + + parsed = sqlparse.parse(statement) + if not parsed: + return [] + + statement_tokens = parsed[0].tokens + select_seen = False + output_columns: list[str] = [] + + for token in statement_tokens: + if token.is_whitespace or token.ttype in Comment: + continue + + if token.ttype is DML and token.normalized == "SELECT": + select_seen = True + continue + + if not select_seen: + continue + + if token.ttype is Keyword and token.normalized == "FROM": + break + + identifiers: list[Identifier] = [] + if isinstance(token, IdentifierList): + identifiers.extend( + identifier + for identifier in token.get_identifiers() + if isinstance(identifier, Identifier) + ) + elif isinstance(token, Identifier): + identifiers.append(token) + + for identifier in identifiers: + column_name = _strip_identifier_quotes( + identifier.get_alias() or identifier.get_real_name() + ) + if column_name and column_name != "*": + output_columns.append(column_name) + + deduplicated_columns = list(dict.fromkeys(output_columns)) + return [ + { + "name": column_name, + "data_type": "VARCHAR", + "comment": "Output column declared by the view statement.", + } + for column_name in deduplicated_columns + ] + + def _build_view_ddl(content: dict) -> str: columns = [ column @@ -176,14 +238,15 @@ def _build_view_ddl(content: dict) -> str: if column.get("name") and column.get("data_type", "").lower() != "unknown" ] statement = content.get("statement", "") + if not columns: + columns = _view_columns_from_statement(statement) + context = _format_semantic_context( { "object_type": "view", "sql_identifier_contract": { "sql_table_name_use_exactly": content["name"], - "sql_column_names_use_exactly": [ - column["name"] for column in columns - ], + "sql_column_names_use_exactly": [column["name"] for column in columns], }, "semantic_context_not_sql_identifiers": { "role": "stable virtual table interface", @@ -201,18 +264,16 @@ def _build_view_ddl(content: dict) -> str: ], } ) - if columns: - columns_ddl = [ - f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" - for column in columns - ] - return ( - f"{context}CREATE TABLE {content['name']} (\n " - + ",\n ".join(columns_ddl) - + "\n);" - ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" + for column in columns + ] - return f"{context}{content['comment']}CREATE VIEW {content['name']}\nAS {statement}" + return ( + f"{context}CREATE TABLE {content['name']} (\n " + + ",\n ".join(columns_ddl) + + "\n);" + ) def _format_semantic_context(context: dict) -> str: @@ -910,9 +971,7 @@ def construct_retrieval_results( ) columns = ( selected_columns - if _selected_columns_are_executable( - table_schema, selected_columns - ) + if _selected_columns_are_executable(table_schema, selected_columns) else None ) ddl, _has_calculated_field, _has_json_field = ( From 2abccff9e479cc7c3c6d1089398be7680a77037e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 05:24:28 +0530 Subject: [PATCH 1053/1087] Fix SQL correction graph document type --- wren-ai-service/src/pipelines/generation/sql_correction.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 055224466b..40724ce9c5 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -1,10 +1,9 @@ import logging import sys -from typing import Any, Dict, List +from typing import Any, Dict from hamilton import base from hamilton.async_driver import AsyncDriver -from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe @@ -102,7 +101,7 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ## Start of Pipeline @observe(capture_input=False) def prompt( - documents: List[Document], + documents: list[str], invalid_generation_result: Dict, prompt_builder: PromptBuilder, query: str | None = None, @@ -193,7 +192,7 @@ def __init__( @observe(name="SQL Correction") async def run( self, - contexts: List[Document], + contexts: list[str], invalid_generation_result: Dict[str, str], query: str | None = None, sql_generation_reasoning: str | None = None, From 0817a68deab33fbb7451b92d02968d059296aa74 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 05:31:35 +0530 Subject: [PATCH 1054/1087] Add exact schema catalog to SQL prompts --- .../generation/followup_sql_generation.py | 4 +++ .../pipelines/generation/sql_correction.py | 4 +++ .../pipelines/generation/sql_generation.py | 4 +++ .../pipelines/generation/sql_regeneration.py | 4 +++ .../src/pipelines/generation/utils/sql.py | 26 +++++++++++++++++++ 5 files changed, 42 insertions(+) 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 df0013724a..59ca8c144a 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -16,6 +16,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_schema_identifier_catalog, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -74,6 +75,8 @@ {% endfor %} {% endif %} +{{ schema_identifier_catalog }} + ### QUESTION ### User's Follow-up Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, relationships, and history only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, prior failed SQL, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. @@ -116,6 +119,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 40724ce9c5..73664e5522 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -15,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_schema_identifier_catalog, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -80,6 +81,8 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endfor %} {% endif %} +{{ schema_identifier_catalog }} + ### QUESTION ### {% if query %} User's Question: {{ query }} @@ -118,6 +121,7 @@ def prompt( instructions=instructions, ), sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 7cf302a733..e634ca5dff 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -15,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_schema_identifier_catalog, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -68,6 +69,8 @@ {% endfor %} {% endif %} +{{ schema_identifier_catalog }} + ### QUESTION ### User's Question: {{ query }} Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. @@ -111,6 +114,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index acc63a91d1..92f47c33ba 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -15,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_schema_identifier_catalog, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -89,6 +90,8 @@ def get_sql_regeneration_system_prompt( {% endfor %} {% endif %} +{{ schema_identifier_catalog }} + ### QUESTION ### Adjustment intent: {{ sql_generation_reasoning }} The previous SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. @@ -132,6 +135,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 65a7a857b4..0753f566f2 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -261,6 +261,28 @@ def from_contexts(cls, contexts: list[str]) -> "_SchemaCatalog": return cls(tables) + def to_prompt(self) -> str: + if not self._tables: + return "" + + lines = [ + "### VALIDATED RETRIEVED SCHEMA IDENTIFIERS ###", + "The SQL must use only these exact deployed Wren identifiers.", + "Do not derive table or column names from the user's wording, source SQL, physical names, comments, aliases, or descriptions.", + ] + for table_name, column_names in self._tables.items(): + lines.append(f"table: {table_name}") + if column_names: + lines.append("columns:") + lines.extend(f"- {column_name}" for column_name in sorted(column_names)) + lines.extend( + [ + "If the requested intent cannot be expressed with these exact identifiers, return null for sql.", + "### END VALIDATED RETRIEVED SCHEMA IDENTIFIERS ###", + ] + ) + return "\n".join(lines) + def validate_sql(self, sql: str | None) -> str | None: if not sql or not self._tables: return None @@ -311,6 +333,10 @@ def validate_sql(self, sql: str | None) -> str | None: return None +def construct_schema_identifier_catalog(contexts: list[str] | None) -> str: + return _SchemaCatalog.from_contexts(contexts or []).to_prompt() + + def _extract_cte_names(token_list: TokenList) -> set[str]: cte_names: set[str] = set() with_seen = False From c106d9342967f9ea9dbbb8cf34c7a09a8308d89a Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 05:52:54 +0530 Subject: [PATCH 1055/1087] Align ask retrieval flow with legacy --- wren-ai-service/src/config.py | 6 +-- .../retrieval/db_schema_retrieval.py | 43 ++++++++++++------- wren-ai-service/src/web/v1/services/ask.py | 12 +++--- 3 files changed, 37 insertions(+), 24 deletions(-) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index 33d2fb0a62..7d9a686dd2 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -29,7 +29,7 @@ class Settings(BaseSettings): column_indexing_batch_size: int = Field(default=50) table_retrieval_size: int = Field(default=10) table_column_retrieval_size: int = Field(default=100) - enable_column_pruning: bool = Field(default=False) + enable_column_pruning: bool = Field(default=True) 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) @@ -38,12 +38,12 @@ class Settings(BaseSettings): # generation config allow_intent_classification: bool = Field(default=True) - allow_sql_generation_reasoning: bool = Field(default=True) + allow_sql_generation_reasoning: bool = Field(default=False) allow_sql_functions_retrieval: bool = Field(default=True) allow_sql_diagnosis: bool = Field(default=True) allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) - max_sql_correction_retries: int = Field(default=3) + max_sql_correction_retries: int = Field(default=0) sql_generation_timeout_seconds: float = Field(default=45.0) # Kept for compatibility with deployed configs. This controls the bounded # generation/correction calls when present. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 2a254ecf12..2a68563417 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -405,15 +405,6 @@ def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[d ] -def _selected_columns_are_executable(content: dict, columns: set[str]) -> bool: - executable_columns = { - column["name"] - for column in content["columns"] - if column["type"] == "COLUMN" and column["data_type"].lower() != "unknown" - } - return bool(columns) and columns.issubset(executable_columns) - - def _build_table_retrieval_context( content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None ) -> tuple[str, bool, bool]: @@ -566,6 +557,15 @@ def _fallback_retrieval_results( } +def _empty_retrieval_results() -> dict[str, Any]: + return { + "retrieval_results": [], + "has_calculated_field": False, + "has_metric": False, + "has_json_field": False, + } + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: @@ -950,7 +950,11 @@ def construct_retrieval_results( columns_and_tables_needed = None if not columns_and_tables_needed: - return _fallback_retrieval_results(construct_db_schemas, dbschema_retrieval) + logger.warning( + "Column pruning did not return grounded schema selections; " + "skipping broad schema fallback." + ) + return _empty_retrieval_results() # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format @@ -969,11 +973,20 @@ def construct_retrieval_results( selected_columns = set( columns_and_tables_needed[table_schema["name"]]["columns"] ) - columns = ( - selected_columns - if _selected_columns_are_executable(table_schema, selected_columns) - else None - ) + executable_columns = { + column["name"] + for column in table_schema["columns"] + if column["type"] == "COLUMN" + and column["data_type"].lower() != "unknown" + } + columns = selected_columns.intersection(executable_columns) + if not columns: + logger.warning( + "Column pruning selected no executable columns for %s; " + "excluding the model from SQL generation context.", + table_schema["name"], + ) + continue ddl, _has_calculated_field, _has_json_field = ( _build_table_retrieval_context( table_schema, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 140fcc9c8a..279cbe9aa6 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -25,10 +25,10 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = False + ignore_sql_generation_reasoning: bool = True enable_column_pruning: bool = False - use_dry_plan: bool = False - allow_dry_plan_fallback: bool = True + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False custom_instruction: Optional[str] = None @@ -99,12 +99,12 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = True, + allow_sql_generation_reasoning: bool = False, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = False, - max_sql_correction_retries: int = 3, + enable_column_pruning: bool = True, + max_sql_correction_retries: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, From 639571f30097a2baae27eecd4f570a531e9b571d Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 11:46:32 +0530 Subject: [PATCH 1056/1087] Align ask schema retrieval with legacy scope --- .../retrieval/db_schema_retrieval.py | 37 ++++++++++++------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 2a68563417..21117a1d98 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -615,6 +615,7 @@ async def dbschema_retrieval( dbschema_retriever: Any, embedding: dict, mdl_hash: str | None = None, + include_related_models: bool = True, ) -> list[Document]: table_names = _table_names_from_description_documents( table_retrieval.get("documents", []) @@ -627,22 +628,28 @@ async def dbschema_retrieval( table_names = _table_names_from_schema_documents(documents) if table_names: - retrieved_table_names = set() - pending_table_names = table_names + if include_related_models: + retrieved_table_names = set() + pending_table_names = table_names + + while pending_table_names: + retrieved_table_names.update(pending_table_names) + retrieved_documents = await _retrieve_schema_documents( + pending_table_names, project_id, mdl_hash, dbschema_retriever + ) + documents = _dedupe_documents(documents + retrieved_documents) + pending_table_names = [ + table_name + for table_name in _related_table_names(documents) + if table_name not in retrieved_table_names + ] - while pending_table_names: - retrieved_table_names.update(pending_table_names) - retrieved_documents = await _retrieve_schema_documents( - pending_table_names, project_id, mdl_hash, dbschema_retriever - ) - documents = _dedupe_documents(documents + retrieved_documents) - pending_table_names = [ - table_name - for table_name in _related_table_names(documents) - if table_name not in retrieved_table_names - ] + return documents - return documents + retrieved_documents = await _retrieve_schema_documents( + table_names, project_id, mdl_hash, dbschema_retriever + ) + return _dedupe_documents(documents + retrieved_documents) return [] @@ -1108,6 +1115,7 @@ def __init__( document_store_provider: DocumentStoreProvider, table_retrieval_size: int = 50, table_column_retrieval_size: int = 100, + include_related_models: bool = False, **kwargs, ): self._components = { @@ -1140,6 +1148,7 @@ def __init__( self._configs = { "encoding": _encoding, "context_window_size": llm_provider.get_context_window_size(), + "include_related_models": include_related_models, } super().__init__( From a3a9710a09ed486e6afe93d92fdf899cda292d63 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 12:09:13 +0530 Subject: [PATCH 1057/1087] Enforce Wren SQL dialect in generation prompts --- wren-ai-service/src/pipelines/generation/utils/sql.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 0753f566f2..1ec3f27f75 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -522,7 +522,8 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. -- Generate Wren SQL only. Do not use warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. +- Generate Wren SQL only, not the native SQL dialect of the connected warehouse. Do not use SQL Server TOP, square-bracket quoting, backtick quoting, FETCH FIRST, OFFSET/FETCH pagination, or warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. +- For top, first, highest, lowest, largest, smallest, or other limited result requests, express the ranking/order with ORDER BY and apply a final LIMIT clause in Wren SQL. Never use SELECT TOP n. - Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. - Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. - If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. @@ -590,6 +591,7 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - DON'T USE "TO_CHAR" function in the generated SQL query. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. +- Do not use SELECT TOP n, FETCH FIRST, OFFSET/FETCH, square-bracket quoting, or backtick quoting. Use Wren SQL syntax with ORDER BY and a final LIMIT n clause for limited or top-N results. - For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. - For the ranking problem, you must add the ranking column to the final SELECT clause. """ From 9722f7d7d58d0149de86ab875b6cab9f90e0238e Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 12:26:55 +0530 Subject: [PATCH 1058/1087] Strengthen Wren SQL schema contract --- .../src/pipelines/generation/utils/sql.py | 52 +++++++++++++++---- 1 file changed, 41 insertions(+), 11 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 1ec3f27f75..fee5ef4f87 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -227,14 +227,20 @@ async def _classify_generation_result( class _SchemaCatalog: - def __init__(self, tables: dict[str, set[str]]): + def __init__( + self, + tables: dict[str, set[str]], + relationships: dict[str, set[str]] | None = None, + ): self._tables = tables + self._relationships = relationships or {} @classmethod def from_contexts(cls, contexts: list[str]) -> "_SchemaCatalog": tables: dict[str, set[str]] = {} + relationships: dict[str, set[str]] = {} current_table: str | None = None - in_columns = False + current_section: str | None = None for context in contexts: for raw_line in context.splitlines(): @@ -243,23 +249,32 @@ def from_contexts(cls, contexts: list[str]) -> "_SchemaCatalog": current_table = line.removeprefix("table: ").strip() if current_table: tables.setdefault(current_table, set()) - in_columns = False + relationships.setdefault(current_table, set()) + current_section = None continue if current_table and line == "columns:": - in_columns = True + current_section = "columns" + continue + + if current_table and line == "relationships:": + current_section = "relationships" continue - if in_columns and current_table and line.startswith("- "): - column_name = line.removeprefix("- ").strip() - if column_name: - tables.setdefault(current_table, set()).add(column_name) + if current_section and current_table and line.startswith("- "): + value = line.removeprefix("- ").strip() + if not value: + continue + if current_section == "columns": + tables.setdefault(current_table, set()).add(value) + elif current_section == "relationships": + relationships.setdefault(current_table, set()).add(value) continue - if in_columns and line and not line.startswith("- "): - in_columns = False + if current_section and line and not line.startswith("- "): + current_section = None - return cls(tables) + return cls(tables, relationships) def to_prompt(self) -> str: if not self._tables: @@ -268,6 +283,8 @@ def to_prompt(self) -> str: lines = [ "### VALIDATED RETRIEVED SCHEMA IDENTIFIERS ###", "The SQL must use only these exact deployed Wren identifiers.", + "Each table value below is one indivisible Wren model identifier; never split it into database, schema, or table parts.", + "Use a multipart table reference only when that exact multipart identifier is listed below as a table value.", "Do not derive table or column names from the user's wording, source SQL, physical names, comments, aliases, or descriptions.", ] for table_name, column_names in self._tables.items(): @@ -275,6 +292,13 @@ def to_prompt(self) -> str: if column_names: lines.append("columns:") lines.extend(f"- {column_name}" for column_name in sorted(column_names)) + table_relationships = self._relationships.get(table_name) + if table_relationships: + lines.append("relationships:") + lines.extend( + f"- {relationship}" + for relationship in sorted(table_relationships) + ) lines.extend( [ "If the requested intent cannot be expressed with these exact identifiers, return null for sql.", @@ -507,6 +531,9 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. - When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. - In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. +- Treat every retrieved Wren table/model name as one indivisible executable identifier. Prefixes, suffixes, underscores, source schema names, connector names, or words that look like database/schema parts are still part of that single Wren identifier. +- Never convert an exact Wren table/model name into a multipart native database reference. If DATABASE SCHEMA declares a table named abc_def, use "abc_def"; do not write abc.def, "abc"."def", or any other split form. +- Use multipart table references such as schema.table or "schema"."table" only when DATABASE SCHEMA declares that exact multipart Wren identifier as the executable table/model name. - Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. - When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. - The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. @@ -545,6 +572,9 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! - PREFER USING CTEs over subqueries. +- Copy table names exactly as one Wren identifier from DATABASE SCHEMA. Do not split underscores or source-schema-like prefixes into dot-qualified database/schema/table references. +- Use table aliases only as SQL aliases for already-declared Wren table names; never use aliases or source schema names as replacements for Wren table names. +- Qualify every source column reference with its exact Wren table name or SQL table alias in SELECT, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY. Output aliases in the final SELECT may be unqualified. - When generating SQL query, always: - Put double quotes around column and table names. - Put single quotes around string literals. From c6c69eab10aabe15beebb55aa35f251732f7253c Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 12:54:48 +0530 Subject: [PATCH 1059/1087] Prioritize Wren identifier contract in SQL prompts --- .../src/pipelines/generation/followup_sql_generation.py | 9 ++++++--- .../generation/followup_sql_generation_reasoning.py | 9 +++++++++ .../src/pipelines/generation/sql_correction.py | 9 ++++++--- .../src/pipelines/generation/sql_generation.py | 9 ++++++--- .../src/pipelines/generation/sql_generation_reasoning.py | 9 +++++++++ .../src/pipelines/generation/sql_regeneration.py | 7 +++++-- 6 files changed, 41 insertions(+), 11 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 59ca8c144a..b09b524756 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -35,6 +35,11 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -75,11 +80,9 @@ {% endfor %} {% endif %} -{{ schema_identifier_catalog }} - ### QUESTION ### User's Follow-up Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, relationships, and history only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, prior failed SQL, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, relationships, and history only to understand meaning; the SQL must use exact declared table and column names from the WREN SQL IDENTIFIER CONTRACT and DATABASE SCHEMA. Treat source metadata, physical names, lineage names, semantic labels, user question words, and prior SQL as non-executable background unless the exact same identifier is declared in the contract. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Return only the final JSON SQL response. """ diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 136bd4a5bc..e84955f87d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -13,6 +13,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_schema_identifier_catalog, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -23,6 +24,13 @@ sql_generation_reasoning_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +Use exact model, column, and relationship names from that contract whenever reasoning names schema objects. +Treat source metadata, physical names, lineage names, semantic labels, prior SQL, and user wording as non-executable background unless the exact same identifier is declared in the contract. +If the contract does not contain a table, column, or relationship required by the user's intent, say the current schema context is insufficient instead of naming an assumed object. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -80,6 +88,7 @@ def prompt( instructions=construct_instructions( instructions=instructions, ), + schema_identifier_catalog=construct_schema_identifier_catalog(documents), language=configuration.language, current_time=configuration.show_current_time(), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 73664e5522..fb21794d02 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -61,6 +61,11 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) sql_correction_user_prompt_template = """ {% if documents %} +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -81,12 +86,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endfor %} {% endif %} -{{ schema_identifier_catalog }} - ### QUESTION ### {% if query %} User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, diagnostic text, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from the WREN SQL IDENTIFIER CONTRACT and DATABASE SCHEMA. Treat source metadata, physical names, lineage names, semantic labels, diagnostic text, and user question words as non-executable background unless the exact same identifier is declared in the contract. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. {% endif %} ### FAILED SQL ### diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index e634ca5dff..b4be95de49 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -29,6 +29,11 @@ sql_generation_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -69,11 +74,9 @@ {% endfor %} {% endif %} -{{ schema_identifier_catalog }} - ### QUESTION ### User's Question: {{ query }} -Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, source metadata, physical names, lineage names, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from DATABASE SCHEMA. Do not copy semantic labels, source/physical/lineage names, user question words, or inferred names into executable SQL. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from the WREN SQL IDENTIFIER CONTRACT and DATABASE SCHEMA. Treat source metadata, physical names, lineage names, semantic labels, and user question words as non-executable background unless the exact same identifier is declared in the contract. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. Return only the final JSON SQL response. 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 45df7584c9..d44532c6ee 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -13,6 +13,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_schema_identifier_catalog, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -22,6 +23,13 @@ sql_generation_reasoning_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +Use exact model, column, and relationship names from that contract whenever reasoning names schema objects. +Treat source metadata, physical names, lineage names, semantic labels, and user wording as non-executable background unless the exact same identifier is declared in the contract. +If the contract does not contain a table, column, or relationship required by the user's intent, say the current schema context is insufficient instead of naming an assumed object. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -69,6 +77,7 @@ def prompt( instructions=construct_instructions( instructions=instructions, ), + schema_identifier_catalog=construct_schema_identifier_catalog(documents), language=configuration.language, current_time=configuration.show_current_time(), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 92f47c33ba..67674b0b09 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -50,6 +50,11 @@ def get_sql_regeneration_system_prompt( sql_regeneration_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -90,8 +95,6 @@ def get_sql_regeneration_system_prompt( {% endfor %} {% endif %} -{{ schema_identifier_catalog }} - ### QUESTION ### Adjustment intent: {{ sql_generation_reasoning }} The previous SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. From 0c2645282426a92ac1b0beb9b580adc4b085c798 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 13:36:06 +0530 Subject: [PATCH 1060/1087] Preserve retrieved schema grounding --- wren-ai-service/src/pipelines/generation/utils/sql.py | 10 +++++++++- .../src/pipelines/retrieval/db_schema_retrieval.py | 6 +++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index fee5ef4f87..e4e45244d2 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -447,7 +447,7 @@ def _extract_table_references(token_list: TokenList) -> tuple[set[str], dict[str def _add_table_reference( identifier: Identifier, table_names: set[str], aliases: dict[str, str] ) -> None: - table_name = _clean_identifier(identifier.get_real_name()) + table_name = _table_reference_name(identifier) alias = _clean_identifier(identifier.get_alias()) if not table_name: return @@ -458,6 +458,14 @@ def _add_table_reference( aliases[alias] = table_name +def _table_reference_name(identifier: Identifier) -> str | None: + parent_name = _clean_identifier(identifier.get_parent_name()) + real_name = _clean_identifier(identifier.get_real_name()) + if parent_name and real_name: + return f"{parent_name}.{real_name}" + return real_name + + def _extract_qualified_columns(token_list: TokenList) -> list[tuple[str, str]]: columns: list[tuple[str, str]] = [] diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 21117a1d98..19ea275650 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -987,13 +987,13 @@ def construct_retrieval_results( and column["data_type"].lower() != "unknown" } columns = selected_columns.intersection(executable_columns) - if not columns: + if selected_columns and not columns: logger.warning( "Column pruning selected no executable columns for %s; " - "excluding the model from SQL generation context.", + "including the full model schema to preserve grounding.", table_schema["name"], ) - continue + columns = None ddl, _has_calculated_field, _has_json_field = ( _build_table_retrieval_context( table_schema, From 16ce5a1cd63e12bcec07276aea6323d4fb1c9aeb Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 14:08:14 +0530 Subject: [PATCH 1061/1087] Remove non-schema SQL examples from prompts --- .../src/pipelines/generation/utils/sql.py | 138 +----------------- 1 file changed, 5 insertions(+), 133 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index e4e45244d2..d42b9b8b17 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -584,10 +584,9 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - Use table aliases only as SQL aliases for already-declared Wren table names; never use aliases or source schema names as replacements for Wren table names. - Qualify every source column reference with its exact Wren table name or SQL table alias in SELECT, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY. Output aliases in the final SELECT may be unqualified. - When generating SQL query, always: - - Put double quotes around column and table names. + - Put double quotes around exact column and table names copied from DATABASE SCHEMA. - Put single quotes around string literals. - Never quote numeric literals. - For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992; - YOU MUST USE "lower(.) like lower()" function or "lower(.) = lower()" function for case-insensitive comparison! - Use "lower(.) LIKE lower()" when: - The user requests a pattern or partial match. @@ -610,17 +609,7 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. -- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. - - EXAMPLE - DATABASE SCHEMA - /* {"alias":"_orders","description":"A model representing the orders data."} */ - CREATE TABLE orders ( - -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} - ApprovedTimestamp TIMESTAMP - } - - SQL - SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; +- Refer to the alias value in the DATABASE SCHEMA comment for the corresponding table or column only as the output label in the final SELECT clause. Do not use alias values as source table or source column identifiers unless they are also exact executable identifiers declared in DATABASE SCHEMA. - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -642,44 +631,7 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non First, provide a brief explanation of what each field represents in the context of the schema, including how each field is computed using the relationships between models. Then, during the following tasks, if the user queries pertain to any calculated fields defined in the database schema, ensure to utilize those calculated fields appropriately in the output SQL queries. The goal is to accurately reflect the intent of the question in the SQL syntax, leveraging the pre-computed logic embedded within the calculated fields. - -EXAMPLES: -The given schema is created by the SQL command: - -CREATE TABLE orders ( - OrderId VARCHAR PRIMARY KEY, - CustomerId VARCHAR, - -- This column is a Calculated Field - -- column expression: avg(reviews.Score) - Rating DOUBLE, - -- This column is a Calculated Field - -- column expression: count(reviews.Id) - ReviewCount BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) - Size BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) > 1 - Large BOOLEAN, - FOREIGN KEY (CustomerId) REFERENCES customers(Id) -); - -Interpret the columns that are marked as Calculated Fields in the schema: -Rating (DOUBLE) - Calculated as the average score (avg) of the Score field from the reviews table where the reviews are associated with the order. This field represents the overall customer satisfaction rating for the order based on review scores. -ReviewCount (BIGINT) - Calculated by counting (count) the number of entries in the reviews table associated with this order. It measures the volume of customer feedback received for the order. -Size (BIGINT) - Represents the total number of items in the order, calculated by counting the number of item entries (ItemNumber) in the order_items table linked to this order. This field is useful for understanding the scale or size of an order. -Large (BOOLEAN) - A boolean value calculated to check if the number of items in the order exceeds one (count(order_items.ItemNumber) > 1). It indicates whether the order is considered large in terms of item quantity. - -And if the user input queries like these: -1. "How many large orders have been placed by customer with ID 'C1234'?" -2. "What is the average customer rating for orders that were rated by more than 10 reviewers?" - -For the first query: -First try to intepret the user query, the user wants to know the average rating for orders which have attracted significant review activity, specifically those with more than 10 reviews. -Then, according to the above intepretation about the given schema, the term 'Rating' is predefined in the Calculated Field of the 'orders' model. And, the number of reviews is also predefined in the 'ReviewCount' Calculated Field. -So utilize those Calculated Fields in the SQL generation process to give an answer like this: - -SQL Query: SELECT AVG(Rating) FROM orders WHERE ReviewCount > 10 +Use calculated fields only when their exact field names are declared in DATABASE SCHEMA and their descriptions or expressions match the user's intent. Do not recreate a calculated field expression with undeclared source tables or columns, and do not invent relationships that are not declared in DATABASE SCHEMA. """ _DEFAULT_METRIC_INSTRUCTIONS = """ @@ -709,68 +661,7 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non If the given schema contains the structures marked as 'metric', you should first interpret the metric schema based on the above definition. Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. - -EXAMPLES: -The given schema is created by the SQL command: - -/* This table is a metric */ -/* Metric Base Object: orders */ -CREATE TABLE Revenue ( - -- This column is a dimension - PurchaseTimestamp TIMESTAMP, - -- This column is a dimension - CustomerId VARCHAR, - -- This column is a dimension - Status VARCHAR, - -- This column is a measure - -- expression: sum(order_items.Price) - PriceSum DOUBLE, - -- This column is a measure - -- expression: count(OrderId) - NumberOfOrders BIGINT -); - -Interpret the metric with the understanding of the metric structure: -1. Base Object: orders -This is the primary data source for the metric. -The orders table provides the underlying data from which dimensions and measures are derived. -It is the foundation upon which the metric is built, though it itself is not directly used in queries against the Revenue table. -It shows the reference between the 'Revenue' metric and the 'orders' model. For the user queries pretain to the 'Revenue' of 'orders', the metric should be utilize in the sql generation process. -2. Dimensions -The metric contains the columns marked as 'dimension'. They can be interpreted as below: -- PurchaseTimestamp (TIMESTAMP) - Acts as a temporal dimension, allowing analysis of revenue over time. This can be used to observe trends, seasonal variations, or performance over specific periods. -- CustomerId (VARCHAR) - A key dimension for customer segmentation, it enables the analysis of revenue generated from individual customers or customer groups. -- Status (VARCHAR) - Reflects the current state of an order (e.g., pending, completed, cancelled). This dimension is crucial for analyses that differentiate performance based on order status. -3. Measures -The metric contains the columns marked as 'measure'. They can be interpreted as below: -- PriceSum (DOUBLE) - A financial measure calculated as sum(order_items.Price), representing the total revenue generated from orders. This measure is vital for tracking overall sales performance and is the primary output of interest in many financial and business analyses. -- NumberOfOrders (BIGINT) - A count measure that provides the total number of orders. This is essential for operational metrics, such as assessing the volume of business activity and evaluating the efficiency of sales processes. - -Now, if the user input queries like this: -Question: "What was the total revenue from each customer last month?" - -First try to intepret the user query, the user asks for a breakdown of the total revenue generated by each customer in the previous calendar month. -The user is specifically interested in understanding how much each customer contributed to the total sales during this period. -To answer this question, it is suitable to use the following components from the metric: -1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. -2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. -3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. -So utilize those metric components in the SQL generation process to give an answer like this: - -SQL Query: -SELECT - CustomerId, - PriceSum AS TotalRevenue -FROM - Revenue -WHERE - PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND - PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) +Use metric dimensions and measures only when their exact metric field names are declared in DATABASE SCHEMA and match the user's requested grouping, filtering, or aggregation. Treat the metric base object as semantic context; do not query the base object unless it is also declared as a retrieved executable table or metric in DATABASE SCHEMA. """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -781,31 +672,12 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - For Example: - DATA SCHEMA: - `/* {"alias":"users","description":"A model representing the users data."} */ - CREATE TABLE users ( - -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} - address JSON - )` - To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` + - Use only the exact JSON column and JSON field paths declared in DATABASE SCHEMA. - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. - - For Example: - DATA SCHEMA - `/* {"alias":"my_table","description":"A test my_table"} */ - CREATE TABLE my_table ( - -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} - elements JSON - )` - To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". - DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". """ From 73cd7003ac5b24c0c7e96c0706c90e95cb43f6e9 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 14:22:45 +0530 Subject: [PATCH 1062/1087] Add 1000 test questions --- questions_1000.md | 1004 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 1004 insertions(+) create mode 100644 questions_1000.md diff --git a/questions_1000.md b/questions_1000.md new file mode 100644 index 0000000000..20a0d03423 --- /dev/null +++ b/questions_1000.md @@ -0,0 +1,1004 @@ +# 1000 Test Questions + +Generated from the invoice, purchase order, supplier, journal entry, document, task, workflow, prepayment, and metadata-grounding patterns in the provided examples. + +1. Show the total number of invoices. +2. Show the total number of purchase orders. +3. Show the total number of purchase order line items. +4. Show the total number of invoice line items. +5. Show the total number of journal entries. +6. Show the total number of documents. +7. Show the total number of tasks. +8. Show the total number of suppliers. +9. Show the total number of prepayments. +10. Show the total number of workflow records. +11. Show all invoices. +12. Show all purchase orders. +13. Show all purchase order line items. +14. Show all invoice line items. +15. Show all journal entries. +16. Show all documents. +17. Show all tasks. +18. Show all suppliers. +19. Show all prepayments. +20. Show all workflow records. +21. Show the latest invoices by invoice date. +22. Show the latest purchase orders by purchase order date. +23. Show the latest purchase order line items by line creation date. +24. Show the latest invoice line items by invoice line date. +25. Show the latest journal entries by journal date. +26. Show the latest documents by load date. +27. Show the latest tasks by start date. +28. Show the latest suppliers by supplier creation date. +29. Show the latest prepayments by prepayment date. +30. Show the latest workflow records by workflow date. +31. Show the oldest invoices by invoice date. +32. Show the oldest purchase orders by purchase order date. +33. Show the oldest purchase order line items by line creation date. +34. Show the oldest invoice line items by invoice line date. +35. Show the oldest journal entries by journal date. +36. Show the oldest documents by load date. +37. Show the oldest tasks by start date. +38. Show the oldest suppliers by supplier creation date. +39. Show the oldest prepayments by prepayment date. +40. Show the oldest workflow records by workflow date. +41. Show the top 10 invoices by invoice amount. +42. Show the top 10 purchase orders by purchase order amount. +43. Show the top 10 purchase order line items by line amount. +44. Show the top 10 invoice line items by line amount. +45. Show the top 10 journal entries by ending balance. +46. Show the top 10 documents by document amount. +47. Show the top 10 tasks by task duration. +48. Show the top 10 suppliers by supplier spend. +49. Show the top 10 prepayments by prepayment amount. +50. Show the top 10 workflow records by workflow duration. +51. Show the number of invoices by business unit. +52. Show the number of invoices by supplier. +53. Show the number of invoices by product line. +54. Show the number of invoices by product. +55. Show the number of invoices by currency. +56. Show the number of invoices by cost center. +57. Show the number of invoices by profit center. +58. Show the number of invoices by customer. +59. Show the number of invoices by invoice type. +60. Show the number of invoices by priority. +61. Show the number of invoices by entry type. +62. Show the number of invoices by task owner. +63. Show the number of invoices by task status. +64. Show the number of invoices by posting status. +65. Show the number of invoices by journal status. +66. Show the number of invoices by GL account. +67. Show the number of invoices by preparer. +68. Show the number of invoices by reviewer. +69. Show the number of invoices by reason code. +70. Show the number of invoices by workflow status. +71. Show the number of purchase orders by business unit. +72. Show the number of purchase orders by supplier. +73. Show the number of purchase orders by product line. +74. Show the number of purchase orders by product. +75. Show the number of purchase orders by currency. +76. Show the number of purchase orders by cost center. +77. Show the number of purchase orders by profit center. +78. Show the number of purchase orders by customer. +79. Show the number of purchase orders by invoice type. +80. Show the number of purchase orders by priority. +81. Show the number of purchase orders by entry type. +82. Show the number of purchase orders by task owner. +83. Show the number of purchase orders by task status. +84. Show the number of purchase orders by posting status. +85. Show the number of purchase orders by journal status. +86. Show the number of purchase orders by GL account. +87. Show the number of purchase orders by preparer. +88. Show the number of purchase orders by reviewer. +89. Show the number of purchase orders by reason code. +90. Show the number of purchase orders by workflow status. +91. Show the number of purchase order line items by business unit. +92. Show the number of purchase order line items by supplier. +93. Show the number of purchase order line items by product line. +94. Show the number of purchase order line items by product. +95. Show the number of purchase order line items by currency. +96. Show the number of purchase order line items by cost center. +97. Show the number of purchase order line items by profit center. +98. Show the number of purchase order line items by customer. +99. Show the number of purchase order line items by invoice type. +100. Show the number of purchase order line items by priority. +101. Show the number of purchase order line items by entry type. +102. Show the number of purchase order line items by task owner. +103. Show the number of purchase order line items by task status. +104. Show the number of purchase order line items by posting status. +105. Show the number of purchase order line items by journal status. +106. Show the number of purchase order line items by GL account. +107. Show the number of purchase order line items by preparer. +108. Show the number of purchase order line items by reviewer. +109. Show the number of purchase order line items by reason code. +110. Show the number of purchase order line items by workflow status. +111. Show the number of invoice line items by business unit. +112. Show the number of invoice line items by supplier. +113. Show the number of invoice line items by product line. +114. Show the number of invoice line items by product. +115. Show the number of invoice line items by currency. +116. Show the number of invoice line items by cost center. +117. Show the number of invoice line items by profit center. +118. Show the number of invoice line items by customer. +119. Show the number of invoice line items by invoice type. +120. Show the number of invoice line items by priority. +121. Show the number of invoice line items by entry type. +122. Show the number of invoice line items by task owner. +123. Show the number of invoice line items by task status. +124. Show the number of invoice line items by posting status. +125. Show the number of invoice line items by journal status. +126. Show the number of invoice line items by GL account. +127. Show the number of invoice line items by preparer. +128. Show the number of invoice line items by reviewer. +129. Show the number of invoice line items by reason code. +130. Show the number of invoice line items by workflow status. +131. Show the number of journal entries by business unit. +132. Show the number of journal entries by supplier. +133. Show the number of journal entries by product line. +134. Show the number of journal entries by product. +135. Show the number of journal entries by currency. +136. Show the number of journal entries by cost center. +137. Show the number of journal entries by profit center. +138. Show the number of journal entries by customer. +139. Show the number of journal entries by invoice type. +140. Show the number of journal entries by priority. +141. Show the number of journal entries by entry type. +142. Show the number of journal entries by task owner. +143. Show the number of journal entries by task status. +144. Show the number of journal entries by posting status. +145. Show the number of journal entries by journal status. +146. Show the number of journal entries by GL account. +147. Show the number of journal entries by preparer. +148. Show the number of journal entries by reviewer. +149. Show the number of journal entries by reason code. +150. Show the number of journal entries by workflow status. +151. Show the number of documents by business unit. +152. Show the number of documents by supplier. +153. Show the number of documents by product line. +154. Show the number of documents by product. +155. Show the number of documents by currency. +156. Show the number of documents by cost center. +157. Show the number of documents by profit center. +158. Show the number of documents by customer. +159. Show the number of documents by invoice type. +160. Show the number of documents by priority. +161. Show the number of documents by entry type. +162. Show the number of documents by task owner. +163. Show the number of documents by task status. +164. Show the number of documents by posting status. +165. Show the number of documents by journal status. +166. Show the number of documents by GL account. +167. Show the number of documents by preparer. +168. Show the number of documents by reviewer. +169. Show the number of documents by reason code. +170. Show the number of documents by workflow status. +171. Show the number of tasks by business unit. +172. Show the number of tasks by supplier. +173. Show the number of tasks by product line. +174. Show the number of tasks by product. +175. Show the number of tasks by currency. +176. Show the number of tasks by cost center. +177. Show the number of tasks by profit center. +178. Show the number of tasks by customer. +179. Show the number of tasks by invoice type. +180. Show the number of tasks by priority. +181. Show the number of tasks by entry type. +182. Show the number of tasks by task owner. +183. Show the number of tasks by task status. +184. Show the number of tasks by posting status. +185. Show the number of tasks by journal status. +186. Show the number of tasks by GL account. +187. Show the number of tasks by preparer. +188. Show the number of tasks by reviewer. +189. Show the number of tasks by reason code. +190. Show the number of tasks by workflow status. +191. Show the number of suppliers by business unit. +192. Show the number of suppliers by supplier. +193. Show the number of suppliers by product line. +194. Show the number of suppliers by product. +195. Show the number of suppliers by currency. +196. Show the number of suppliers by cost center. +197. Show the number of suppliers by profit center. +198. Show the number of suppliers by customer. +199. Show the number of suppliers by invoice type. +200. Show the number of suppliers by priority. +201. Show the number of suppliers by entry type. +202. Show the number of suppliers by task owner. +203. Show the number of suppliers by task status. +204. Show the number of suppliers by posting status. +205. Show the number of suppliers by journal status. +206. Show the number of suppliers by GL account. +207. Show the number of suppliers by preparer. +208. Show the number of suppliers by reviewer. +209. Show the number of suppliers by reason code. +210. Show the number of suppliers by workflow status. +211. Show the number of prepayments by business unit. +212. Show the number of prepayments by supplier. +213. Show the number of prepayments by product line. +214. Show the number of prepayments by product. +215. Show the number of prepayments by currency. +216. Show the number of prepayments by cost center. +217. Show the number of prepayments by profit center. +218. Show the number of prepayments by customer. +219. Show the number of prepayments by invoice type. +220. Show the number of prepayments by priority. +221. Show the number of prepayments by entry type. +222. Show the number of prepayments by task owner. +223. Show the number of prepayments by task status. +224. Show the number of prepayments by posting status. +225. Show the number of prepayments by journal status. +226. Show the number of prepayments by GL account. +227. Show the number of prepayments by preparer. +228. Show the number of prepayments by reviewer. +229. Show the number of prepayments by reason code. +230. Show the number of prepayments by workflow status. +231. Show the number of workflow records by business unit. +232. Show the number of workflow records by supplier. +233. Show the number of workflow records by product line. +234. Show the number of workflow records by product. +235. Show the number of workflow records by currency. +236. Show the number of workflow records by cost center. +237. Show the number of workflow records by profit center. +238. Show the number of workflow records by customer. +239. Show the number of workflow records by invoice type. +240. Show the number of workflow records by priority. +241. Show the number of workflow records by entry type. +242. Show the number of workflow records by task owner. +243. Show the number of workflow records by task status. +244. Show the number of workflow records by posting status. +245. Show the number of workflow records by journal status. +246. Show the number of workflow records by GL account. +247. Show the number of workflow records by preparer. +248. Show the number of workflow records by reviewer. +249. Show the number of workflow records by reason code. +250. Show the number of workflow records by workflow status. +251. Show the total invoice amount by business unit. +252. Show the total invoice amount by supplier. +253. Show the total invoice amount by product line. +254. Show the total invoice amount by product. +255. Show the total invoice amount by currency. +256. Show the total invoice amount by cost center. +257. Show the total invoice amount by profit center. +258. Show the total invoice amount by customer. +259. Show the total invoice amount by invoice type. +260. Show the total invoice amount by priority. +261. Show the total invoice amount by entry type. +262. Show the total invoice amount by task owner. +263. Show the total invoice amount by task status. +264. Show the total invoice amount by posting status. +265. Show the total invoice amount by journal status. +266. Show the total invoice amount by GL account. +267. Show the total invoice amount by preparer. +268. Show the total invoice amount by reviewer. +269. Show the total invoice amount by reason code. +270. Show the total invoice amount by workflow status. +271. Show the total purchase order amount by business unit. +272. Show the total purchase order amount by supplier. +273. Show the total purchase order amount by product line. +274. Show the total purchase order amount by product. +275. Show the total purchase order amount by currency. +276. Show the total purchase order amount by cost center. +277. Show the total purchase order amount by profit center. +278. Show the total purchase order amount by customer. +279. Show the total purchase order amount by invoice type. +280. Show the total purchase order amount by priority. +281. Show the total purchase order amount by entry type. +282. Show the total purchase order amount by task owner. +283. Show the total purchase order amount by task status. +284. Show the total purchase order amount by posting status. +285. Show the total purchase order amount by journal status. +286. Show the total purchase order amount by GL account. +287. Show the total purchase order amount by preparer. +288. Show the total purchase order amount by reviewer. +289. Show the total purchase order amount by reason code. +290. Show the total purchase order amount by workflow status. +291. Show the total line amount by business unit. +292. Show the total line amount by supplier. +293. Show the total line amount by product line. +294. Show the total line amount by product. +295. Show the total line amount by currency. +296. Show the total line amount by cost center. +297. Show the total line amount by profit center. +298. Show the total line amount by customer. +299. Show the total line amount by invoice type. +300. Show the total line amount by priority. +301. Show the total line amount by entry type. +302. Show the total line amount by task owner. +303. Show the total line amount by task status. +304. Show the total line amount by posting status. +305. Show the total line amount by journal status. +306. Show the total line amount by GL account. +307. Show the total line amount by preparer. +308. Show the total line amount by reviewer. +309. Show the total line amount by reason code. +310. Show the total line amount by workflow status. +311. Show the total ending balance by business unit. +312. Show the total ending balance by supplier. +313. Show the total ending balance by product line. +314. Show the total ending balance by product. +315. Show the total ending balance by currency. +316. Show the total ending balance by cost center. +317. Show the total ending balance by profit center. +318. Show the total ending balance by customer. +319. Show the total ending balance by invoice type. +320. Show the total ending balance by priority. +321. Show the total ending balance by entry type. +322. Show the total ending balance by task owner. +323. Show the total ending balance by task status. +324. Show the total ending balance by posting status. +325. Show the total ending balance by journal status. +326. Show the total ending balance by GL account. +327. Show the total ending balance by preparer. +328. Show the total ending balance by reviewer. +329. Show the total ending balance by reason code. +330. Show the total ending balance by workflow status. +331. Show the total document amount by business unit. +332. Show the total document amount by supplier. +333. Show the total document amount by product line. +334. Show the total document amount by product. +335. Show the total document amount by currency. +336. Show the total document amount by cost center. +337. Show the total document amount by profit center. +338. Show the total document amount by customer. +339. Show the total document amount by invoice type. +340. Show the total document amount by priority. +341. Show the total document amount by entry type. +342. Show the total document amount by task owner. +343. Show the total document amount by task status. +344. Show the total document amount by posting status. +345. Show the total document amount by journal status. +346. Show the total document amount by GL account. +347. Show the total document amount by preparer. +348. Show the total document amount by reviewer. +349. Show the total document amount by reason code. +350. Show the total document amount by workflow status. +351. Show the total task duration by business unit. +352. Show the total task duration by supplier. +353. Show the total task duration by product line. +354. Show the total task duration by product. +355. Show the total task duration by currency. +356. Show the total task duration by cost center. +357. Show the total task duration by profit center. +358. Show the total task duration by customer. +359. Show the total task duration by invoice type. +360. Show the total task duration by priority. +361. Show the total task duration by entry type. +362. Show the total task duration by task owner. +363. Show the total task duration by task status. +364. Show the total task duration by posting status. +365. Show the total task duration by journal status. +366. Show the total task duration by GL account. +367. Show the total task duration by preparer. +368. Show the total task duration by reviewer. +369. Show the total task duration by reason code. +370. Show the total task duration by workflow status. +371. Show the total supplier spend by business unit. +372. Show the total supplier spend by supplier. +373. Show the total supplier spend by product line. +374. Show the total supplier spend by product. +375. Show the total supplier spend by currency. +376. Show the total supplier spend by cost center. +377. Show the total supplier spend by profit center. +378. Show the total supplier spend by customer. +379. Show the total supplier spend by invoice type. +380. Show the total supplier spend by priority. +381. Show the total supplier spend by entry type. +382. Show the total supplier spend by task owner. +383. Show the total supplier spend by task status. +384. Show the total supplier spend by posting status. +385. Show the total supplier spend by journal status. +386. Show the total supplier spend by GL account. +387. Show the total supplier spend by preparer. +388. Show the total supplier spend by reviewer. +389. Show the total supplier spend by reason code. +390. Show the total supplier spend by workflow status. +391. Show the total prepayment amount by business unit. +392. Show the total prepayment amount by supplier. +393. Show the total prepayment amount by product line. +394. Show the total prepayment amount by product. +395. Show the total prepayment amount by currency. +396. Show the total prepayment amount by cost center. +397. Show the total prepayment amount by profit center. +398. Show the total prepayment amount by customer. +399. Show the total prepayment amount by invoice type. +400. Show the total prepayment amount by priority. +401. Show the total prepayment amount by entry type. +402. Show the total prepayment amount by task owner. +403. Show the total prepayment amount by task status. +404. Show the total prepayment amount by posting status. +405. Show the total prepayment amount by journal status. +406. Show the total prepayment amount by GL account. +407. Show the total prepayment amount by preparer. +408. Show the total prepayment amount by reviewer. +409. Show the total prepayment amount by reason code. +410. Show the total prepayment amount by workflow status. +411. Show the total workflow duration by business unit. +412. Show the total workflow duration by supplier. +413. Show the total workflow duration by product line. +414. Show the total workflow duration by product. +415. Show the total workflow duration by currency. +416. Show the total workflow duration by cost center. +417. Show the total workflow duration by profit center. +418. Show the total workflow duration by customer. +419. Show the total workflow duration by invoice type. +420. Show the total workflow duration by priority. +421. Show the total workflow duration by entry type. +422. Show the total workflow duration by task owner. +423. Show the total workflow duration by task status. +424. Show the total workflow duration by posting status. +425. Show the total workflow duration by journal status. +426. Show the total workflow duration by GL account. +427. Show the total workflow duration by preparer. +428. Show the total workflow duration by reviewer. +429. Show the total workflow duration by reason code. +430. Show the total workflow duration by workflow status. +431. Show invoices grouped by business unit. +432. Show invoices grouped by supplier. +433. Show invoices grouped by product line. +434. Show invoices grouped by product. +435. Show invoices grouped by currency. +436. Show invoices grouped by cost center. +437. Show invoices grouped by profit center. +438. Show invoices grouped by customer. +439. Show invoices grouped by invoice type. +440. Show invoices grouped by priority. +441. Show invoices grouped by entry type. +442. Show invoices grouped by task owner. +443. Show invoices grouped by task status. +444. Show invoices grouped by posting status. +445. Show invoices grouped by journal status. +446. Show invoices grouped by GL account. +447. Show invoices grouped by preparer. +448. Show invoices grouped by reviewer. +449. Show invoices grouped by reason code. +450. Show invoices grouped by workflow status. +451. Show purchase orders grouped by business unit. +452. Show purchase orders grouped by supplier. +453. Show purchase orders grouped by product line. +454. Show purchase orders grouped by product. +455. Show purchase orders grouped by currency. +456. Show purchase orders grouped by cost center. +457. Show purchase orders grouped by profit center. +458. Show purchase orders grouped by customer. +459. Show purchase orders grouped by invoice type. +460. Show purchase orders grouped by priority. +461. Show purchase orders grouped by entry type. +462. Show purchase orders grouped by task owner. +463. Show purchase orders grouped by task status. +464. Show purchase orders grouped by posting status. +465. Show purchase orders grouped by journal status. +466. Show purchase orders grouped by GL account. +467. Show purchase orders grouped by preparer. +468. Show purchase orders grouped by reviewer. +469. Show purchase orders grouped by reason code. +470. Show purchase orders grouped by workflow status. +471. Show purchase order line items grouped by business unit. +472. Show purchase order line items grouped by supplier. +473. Show purchase order line items grouped by product line. +474. Show purchase order line items grouped by product. +475. Show purchase order line items grouped by currency. +476. Show purchase order line items grouped by cost center. +477. Show purchase order line items grouped by profit center. +478. Show purchase order line items grouped by customer. +479. Show purchase order line items grouped by invoice type. +480. Show purchase order line items grouped by priority. +481. Show purchase order line items grouped by entry type. +482. Show purchase order line items grouped by task owner. +483. Show purchase order line items grouped by task status. +484. Show purchase order line items grouped by posting status. +485. Show purchase order line items grouped by journal status. +486. Show purchase order line items grouped by GL account. +487. Show purchase order line items grouped by preparer. +488. Show purchase order line items grouped by reviewer. +489. Show purchase order line items grouped by reason code. +490. Show purchase order line items grouped by workflow status. +491. Show invoice line items grouped by business unit. +492. Show invoice line items grouped by supplier. +493. Show invoice line items grouped by product line. +494. Show invoice line items grouped by product. +495. Show invoice line items grouped by currency. +496. Show invoice line items grouped by cost center. +497. Show invoice line items grouped by profit center. +498. Show invoice line items grouped by customer. +499. Show invoice line items grouped by invoice type. +500. Show invoice line items grouped by priority. +501. Show invoice line items grouped by entry type. +502. Show invoice line items grouped by task owner. +503. Show invoice line items grouped by task status. +504. Show invoice line items grouped by posting status. +505. Show invoice line items grouped by journal status. +506. Show invoice line items grouped by GL account. +507. Show invoice line items grouped by preparer. +508. Show invoice line items grouped by reviewer. +509. Show invoice line items grouped by reason code. +510. Show invoice line items grouped by workflow status. +511. Show journal entries grouped by business unit. +512. Show journal entries grouped by supplier. +513. Show journal entries grouped by product line. +514. Show journal entries grouped by product. +515. Show journal entries grouped by currency. +516. Show journal entries grouped by cost center. +517. Show journal entries grouped by profit center. +518. Show journal entries grouped by customer. +519. Show journal entries grouped by invoice type. +520. Show journal entries grouped by priority. +521. Show journal entries grouped by entry type. +522. Show journal entries grouped by task owner. +523. Show journal entries grouped by task status. +524. Show journal entries grouped by posting status. +525. Show journal entries grouped by journal status. +526. Show journal entries grouped by GL account. +527. Show journal entries grouped by preparer. +528. Show journal entries grouped by reviewer. +529. Show journal entries grouped by reason code. +530. Show journal entries grouped by workflow status. +531. Show documents grouped by business unit. +532. Show documents grouped by supplier. +533. Show documents grouped by product line. +534. Show documents grouped by product. +535. Show documents grouped by currency. +536. Show documents grouped by cost center. +537. Show documents grouped by profit center. +538. Show documents grouped by customer. +539. Show documents grouped by invoice type. +540. Show documents grouped by priority. +541. Show documents grouped by entry type. +542. Show documents grouped by task owner. +543. Show documents grouped by task status. +544. Show documents grouped by posting status. +545. Show documents grouped by journal status. +546. Show documents grouped by GL account. +547. Show documents grouped by preparer. +548. Show documents grouped by reviewer. +549. Show documents grouped by reason code. +550. Show documents grouped by workflow status. +551. Show tasks grouped by business unit. +552. Show tasks grouped by supplier. +553. Show tasks grouped by product line. +554. Show tasks grouped by product. +555. Show tasks grouped by currency. +556. Show tasks grouped by cost center. +557. Show tasks grouped by profit center. +558. Show tasks grouped by customer. +559. Show tasks grouped by invoice type. +560. Show tasks grouped by priority. +561. Show tasks grouped by entry type. +562. Show tasks grouped by task owner. +563. Show tasks grouped by task status. +564. Show tasks grouped by posting status. +565. Show tasks grouped by journal status. +566. Show tasks grouped by GL account. +567. Show tasks grouped by preparer. +568. Show tasks grouped by reviewer. +569. Show tasks grouped by reason code. +570. Show tasks grouped by workflow status. +571. Show suppliers grouped by business unit. +572. Show suppliers grouped by supplier. +573. Show suppliers grouped by product line. +574. Show suppliers grouped by product. +575. Show suppliers grouped by currency. +576. Show suppliers grouped by cost center. +577. Show suppliers grouped by profit center. +578. Show suppliers grouped by customer. +579. Show suppliers grouped by invoice type. +580. Show suppliers grouped by priority. +581. Show suppliers grouped by entry type. +582. Show suppliers grouped by task owner. +583. Show suppliers grouped by task status. +584. Show suppliers grouped by posting status. +585. Show suppliers grouped by journal status. +586. Show suppliers grouped by GL account. +587. Show suppliers grouped by preparer. +588. Show suppliers grouped by reviewer. +589. Show suppliers grouped by reason code. +590. Show suppliers grouped by workflow status. +591. Show prepayments grouped by business unit. +592. Show prepayments grouped by supplier. +593. Show prepayments grouped by product line. +594. Show prepayments grouped by product. +595. Show prepayments grouped by currency. +596. Show prepayments grouped by cost center. +597. Show prepayments grouped by profit center. +598. Show prepayments grouped by customer. +599. Show prepayments grouped by invoice type. +600. Show prepayments grouped by priority. +601. Show prepayments grouped by entry type. +602. Show prepayments grouped by task owner. +603. Show prepayments grouped by task status. +604. Show prepayments grouped by posting status. +605. Show prepayments grouped by journal status. +606. Show prepayments grouped by GL account. +607. Show prepayments grouped by preparer. +608. Show prepayments grouped by reviewer. +609. Show prepayments grouped by reason code. +610. Show prepayments grouped by workflow status. +611. Show workflow records grouped by business unit. +612. Show workflow records grouped by supplier. +613. Show workflow records grouped by product line. +614. Show workflow records grouped by product. +615. Show workflow records grouped by currency. +616. Show workflow records grouped by cost center. +617. Show workflow records grouped by profit center. +618. Show workflow records grouped by customer. +619. Show workflow records grouped by invoice type. +620. Show workflow records grouped by priority. +621. Show workflow records grouped by entry type. +622. Show workflow records grouped by task owner. +623. Show workflow records grouped by task status. +624. Show workflow records grouped by posting status. +625. Show workflow records grouped by journal status. +626. Show workflow records grouped by GL account. +627. Show workflow records grouped by preparer. +628. Show workflow records grouped by reviewer. +629. Show workflow records grouped by reason code. +630. Show workflow records grouped by workflow status. +631. Show the top 10 business unit values by number of invoices. +632. Show the top 10 supplier values by number of invoices. +633. Show the top 10 product line values by number of invoices. +634. Show the top 10 product values by number of invoices. +635. Show the top 10 currency values by number of invoices. +636. Show the top 10 cost center values by number of invoices. +637. Show the top 10 profit center values by number of invoices. +638. Show the top 10 customer values by number of invoices. +639. Show the top 10 invoice type values by number of invoices. +640. Show the top 10 priority values by number of invoices. +641. Show the top 10 entry type values by number of invoices. +642. Show the top 10 task owner values by number of invoices. +643. Show the top 10 task status values by number of invoices. +644. Show the top 10 posting status values by number of invoices. +645. Show the top 10 journal status values by number of invoices. +646. Show the top 10 GL account values by number of invoices. +647. Show the top 10 preparer values by number of invoices. +648. Show the top 10 reviewer values by number of invoices. +649. Show the top 10 reason code values by number of invoices. +650. Show the top 10 workflow status values by number of invoices. +651. Show the top 10 business unit values by number of purchase orders. +652. Show the top 10 supplier values by number of purchase orders. +653. Show the top 10 product line values by number of purchase orders. +654. Show the top 10 product values by number of purchase orders. +655. Show the top 10 currency values by number of purchase orders. +656. Show the top 10 cost center values by number of purchase orders. +657. Show the top 10 profit center values by number of purchase orders. +658. Show the top 10 customer values by number of purchase orders. +659. Show the top 10 invoice type values by number of purchase orders. +660. Show the top 10 priority values by number of purchase orders. +661. Show the top 10 entry type values by number of purchase orders. +662. Show the top 10 task owner values by number of purchase orders. +663. Show the top 10 task status values by number of purchase orders. +664. Show the top 10 posting status values by number of purchase orders. +665. Show the top 10 journal status values by number of purchase orders. +666. Show the top 10 GL account values by number of purchase orders. +667. Show the top 10 preparer values by number of purchase orders. +668. Show the top 10 reviewer values by number of purchase orders. +669. Show the top 10 reason code values by number of purchase orders. +670. Show the top 10 workflow status values by number of purchase orders. +671. Show the top 10 business unit values by number of purchase order line items. +672. Show the top 10 supplier values by number of purchase order line items. +673. Show the top 10 product line values by number of purchase order line items. +674. Show the top 10 product values by number of purchase order line items. +675. Show the top 10 currency values by number of purchase order line items. +676. Show the top 10 cost center values by number of purchase order line items. +677. Show the top 10 profit center values by number of purchase order line items. +678. Show the top 10 customer values by number of purchase order line items. +679. Show the top 10 invoice type values by number of purchase order line items. +680. Show the top 10 priority values by number of purchase order line items. +681. Show the top 10 entry type values by number of purchase order line items. +682. Show the top 10 task owner values by number of purchase order line items. +683. Show the top 10 task status values by number of purchase order line items. +684. Show the top 10 posting status values by number of purchase order line items. +685. Show the top 10 journal status values by number of purchase order line items. +686. Show the top 10 GL account values by number of purchase order line items. +687. Show the top 10 preparer values by number of purchase order line items. +688. Show the top 10 reviewer values by number of purchase order line items. +689. Show the top 10 reason code values by number of purchase order line items. +690. Show the top 10 workflow status values by number of purchase order line items. +691. Show the top 10 business unit values by number of invoice line items. +692. Show the top 10 supplier values by number of invoice line items. +693. Show the top 10 product line values by number of invoice line items. +694. Show the top 10 product values by number of invoice line items. +695. Show the top 10 currency values by number of invoice line items. +696. Show the top 10 cost center values by number of invoice line items. +697. Show the top 10 profit center values by number of invoice line items. +698. Show the top 10 customer values by number of invoice line items. +699. Show the top 10 invoice type values by number of invoice line items. +700. Show the top 10 priority values by number of invoice line items. +701. Show the top 10 entry type values by number of invoice line items. +702. Show the top 10 task owner values by number of invoice line items. +703. Show the top 10 task status values by number of invoice line items. +704. Show the top 10 posting status values by number of invoice line items. +705. Show the top 10 journal status values by number of invoice line items. +706. Show the top 10 GL account values by number of invoice line items. +707. Show the top 10 preparer values by number of invoice line items. +708. Show the top 10 reviewer values by number of invoice line items. +709. Show the top 10 reason code values by number of invoice line items. +710. Show the top 10 workflow status values by number of invoice line items. +711. Show the top 10 business unit values by number of journal entries. +712. Show the top 10 supplier values by number of journal entries. +713. Show the top 10 product line values by number of journal entries. +714. Show the top 10 product values by number of journal entries. +715. Show the top 10 currency values by number of journal entries. +716. Show the top 10 cost center values by number of journal entries. +717. Show the top 10 profit center values by number of journal entries. +718. Show the top 10 customer values by number of journal entries. +719. Show the top 10 invoice type values by number of journal entries. +720. Show the top 10 priority values by number of journal entries. +721. Show the top 10 entry type values by number of journal entries. +722. Show the top 10 task owner values by number of journal entries. +723. Show the top 10 task status values by number of journal entries. +724. Show the top 10 posting status values by number of journal entries. +725. Show the top 10 journal status values by number of journal entries. +726. Show the top 10 GL account values by number of journal entries. +727. Show the top 10 preparer values by number of journal entries. +728. Show the top 10 reviewer values by number of journal entries. +729. Show the top 10 reason code values by number of journal entries. +730. Show the top 10 workflow status values by number of journal entries. +731. Show the top 10 business unit values by number of documents. +732. Show the top 10 supplier values by number of documents. +733. Show the top 10 product line values by number of documents. +734. Show the top 10 product values by number of documents. +735. Show the top 10 currency values by number of documents. +736. Show the top 10 cost center values by number of documents. +737. Show the top 10 profit center values by number of documents. +738. Show the top 10 customer values by number of documents. +739. Show the top 10 invoice type values by number of documents. +740. Show the top 10 priority values by number of documents. +741. Show the top 10 entry type values by number of documents. +742. Show the top 10 task owner values by number of documents. +743. Show the top 10 task status values by number of documents. +744. Show the top 10 posting status values by number of documents. +745. Show the top 10 journal status values by number of documents. +746. Show the top 10 GL account values by number of documents. +747. Show the top 10 preparer values by number of documents. +748. Show the top 10 reviewer values by number of documents. +749. Show the top 10 reason code values by number of documents. +750. Show the top 10 workflow status values by number of documents. +751. Show the top 10 business unit values by number of tasks. +752. Show the top 10 supplier values by number of tasks. +753. Show the top 10 product line values by number of tasks. +754. Show the top 10 product values by number of tasks. +755. Show the top 10 currency values by number of tasks. +756. Show the top 10 cost center values by number of tasks. +757. Show the top 10 profit center values by number of tasks. +758. Show the top 10 customer values by number of tasks. +759. Show the top 10 invoice type values by number of tasks. +760. Show the top 10 priority values by number of tasks. +761. Show the top 10 entry type values by number of tasks. +762. Show the top 10 task owner values by number of tasks. +763. Show the top 10 task status values by number of tasks. +764. Show the top 10 posting status values by number of tasks. +765. Show the top 10 journal status values by number of tasks. +766. Show the top 10 GL account values by number of tasks. +767. Show the top 10 preparer values by number of tasks. +768. Show the top 10 reviewer values by number of tasks. +769. Show the top 10 reason code values by number of tasks. +770. Show the top 10 workflow status values by number of tasks. +771. Show the top 10 business unit values by number of suppliers. +772. Show the top 10 supplier values by number of suppliers. +773. Show the top 10 product line values by number of suppliers. +774. Show the top 10 product values by number of suppliers. +775. Show the top 10 currency values by number of suppliers. +776. Show the top 10 cost center values by number of suppliers. +777. Show the top 10 profit center values by number of suppliers. +778. Show the top 10 customer values by number of suppliers. +779. Show the top 10 invoice type values by number of suppliers. +780. Show the top 10 priority values by number of suppliers. +781. Show the top 10 entry type values by number of suppliers. +782. Show the top 10 task owner values by number of suppliers. +783. Show the top 10 task status values by number of suppliers. +784. Show the top 10 posting status values by number of suppliers. +785. Show the top 10 journal status values by number of suppliers. +786. Show the top 10 GL account values by number of suppliers. +787. Show the top 10 preparer values by number of suppliers. +788. Show the top 10 reviewer values by number of suppliers. +789. Show the top 10 reason code values by number of suppliers. +790. Show the top 10 workflow status values by number of suppliers. +791. Show the top 10 business unit values by number of prepayments. +792. Show the top 10 supplier values by number of prepayments. +793. Show the top 10 product line values by number of prepayments. +794. Show the top 10 product values by number of prepayments. +795. Show the top 10 currency values by number of prepayments. +796. Show the top 10 cost center values by number of prepayments. +797. Show the top 10 profit center values by number of prepayments. +798. Show the top 10 customer values by number of prepayments. +799. Show the top 10 invoice type values by number of prepayments. +800. Show the top 10 priority values by number of prepayments. +801. Show the top 10 entry type values by number of prepayments. +802. Show the top 10 task owner values by number of prepayments. +803. Show the top 10 task status values by number of prepayments. +804. Show the top 10 posting status values by number of prepayments. +805. Show the top 10 journal status values by number of prepayments. +806. Show the top 10 GL account values by number of prepayments. +807. Show the top 10 preparer values by number of prepayments. +808. Show the top 10 reviewer values by number of prepayments. +809. Show the top 10 reason code values by number of prepayments. +810. Show the top 10 workflow status values by number of prepayments. +811. Show the top 10 business unit values by number of workflow records. +812. Show the top 10 supplier values by number of workflow records. +813. Show the top 10 product line values by number of workflow records. +814. Show the top 10 product values by number of workflow records. +815. Show the top 10 currency values by number of workflow records. +816. Show the top 10 cost center values by number of workflow records. +817. Show the top 10 profit center values by number of workflow records. +818. Show the top 10 customer values by number of workflow records. +819. Show the top 10 invoice type values by number of workflow records. +820. Show the top 10 priority values by number of workflow records. +821. Show the top 10 entry type values by number of workflow records. +822. Show the top 10 task owner values by number of workflow records. +823. Show the top 10 task status values by number of workflow records. +824. Show the top 10 posting status values by number of workflow records. +825. Show the top 10 journal status values by number of workflow records. +826. Show the top 10 GL account values by number of workflow records. +827. Show the top 10 preparer values by number of workflow records. +828. Show the top 10 reviewer values by number of workflow records. +829. Show the top 10 reason code values by number of workflow records. +830. Show the top 10 workflow status values by number of workflow records. +831. Show invoices for each business unit ordered by count from highest to lowest. +832. Show invoices for each supplier ordered by count from highest to lowest. +833. Show invoices for each product line ordered by count from highest to lowest. +834. Show invoices for each product ordered by count from highest to lowest. +835. Show invoices for each currency ordered by count from highest to lowest. +836. Show invoices for each cost center ordered by count from highest to lowest. +837. Show invoices for each profit center ordered by count from highest to lowest. +838. Show invoices for each customer ordered by count from highest to lowest. +839. Show invoices for each invoice type ordered by count from highest to lowest. +840. Show invoices for each priority ordered by count from highest to lowest. +841. Show invoices for each entry type ordered by count from highest to lowest. +842. Show invoices for each task owner ordered by count from highest to lowest. +843. Show invoices for each task status ordered by count from highest to lowest. +844. Show invoices for each posting status ordered by count from highest to lowest. +845. Show invoices for each journal status ordered by count from highest to lowest. +846. Show invoices for each GL account ordered by count from highest to lowest. +847. Show invoices for each preparer ordered by count from highest to lowest. +848. Show invoices for each reviewer ordered by count from highest to lowest. +849. Show invoices for each reason code ordered by count from highest to lowest. +850. Show invoices for each workflow status ordered by count from highest to lowest. +851. Show purchase orders for each business unit ordered by count from highest to lowest. +852. Show purchase orders for each supplier ordered by count from highest to lowest. +853. Show purchase orders for each product line ordered by count from highest to lowest. +854. Show purchase orders for each product ordered by count from highest to lowest. +855. Show purchase orders for each currency ordered by count from highest to lowest. +856. Show purchase orders for each cost center ordered by count from highest to lowest. +857. Show purchase orders for each profit center ordered by count from highest to lowest. +858. Show purchase orders for each customer ordered by count from highest to lowest. +859. Show purchase orders for each invoice type ordered by count from highest to lowest. +860. Show purchase orders for each priority ordered by count from highest to lowest. +861. Show purchase orders for each entry type ordered by count from highest to lowest. +862. Show purchase orders for each task owner ordered by count from highest to lowest. +863. Show purchase orders for each task status ordered by count from highest to lowest. +864. Show purchase orders for each posting status ordered by count from highest to lowest. +865. Show purchase orders for each journal status ordered by count from highest to lowest. +866. Show purchase orders for each GL account ordered by count from highest to lowest. +867. Show purchase orders for each preparer ordered by count from highest to lowest. +868. Show purchase orders for each reviewer ordered by count from highest to lowest. +869. Show purchase orders for each reason code ordered by count from highest to lowest. +870. Show purchase orders for each workflow status ordered by count from highest to lowest. +871. Show purchase order line items for each business unit ordered by count from highest to lowest. +872. Show purchase order line items for each supplier ordered by count from highest to lowest. +873. Show purchase order line items for each product line ordered by count from highest to lowest. +874. Show purchase order line items for each product ordered by count from highest to lowest. +875. Show purchase order line items for each currency ordered by count from highest to lowest. +876. Show purchase order line items for each cost center ordered by count from highest to lowest. +877. Show purchase order line items for each profit center ordered by count from highest to lowest. +878. Show purchase order line items for each customer ordered by count from highest to lowest. +879. Show purchase order line items for each invoice type ordered by count from highest to lowest. +880. Show purchase order line items for each priority ordered by count from highest to lowest. +881. Show purchase order line items for each entry type ordered by count from highest to lowest. +882. Show purchase order line items for each task owner ordered by count from highest to lowest. +883. Show purchase order line items for each task status ordered by count from highest to lowest. +884. Show purchase order line items for each posting status ordered by count from highest to lowest. +885. Show purchase order line items for each journal status ordered by count from highest to lowest. +886. Show purchase order line items for each GL account ordered by count from highest to lowest. +887. Show purchase order line items for each preparer ordered by count from highest to lowest. +888. Show purchase order line items for each reviewer ordered by count from highest to lowest. +889. Show purchase order line items for each reason code ordered by count from highest to lowest. +890. Show purchase order line items for each workflow status ordered by count from highest to lowest. +891. Show invoice line items for each business unit ordered by count from highest to lowest. +892. Show invoice line items for each supplier ordered by count from highest to lowest. +893. Show invoice line items for each product line ordered by count from highest to lowest. +894. Show invoice line items for each product ordered by count from highest to lowest. +895. Show invoice line items for each currency ordered by count from highest to lowest. +896. Show invoice line items for each cost center ordered by count from highest to lowest. +897. Show invoice line items for each profit center ordered by count from highest to lowest. +898. Show invoice line items for each customer ordered by count from highest to lowest. +899. Show invoice line items for each invoice type ordered by count from highest to lowest. +900. Show invoice line items for each priority ordered by count from highest to lowest. +901. Show invoice line items for each entry type ordered by count from highest to lowest. +902. Show invoice line items for each task owner ordered by count from highest to lowest. +903. Show invoice line items for each task status ordered by count from highest to lowest. +904. Show invoice line items for each posting status ordered by count from highest to lowest. +905. Show invoice line items for each journal status ordered by count from highest to lowest. +906. Show invoice line items for each GL account ordered by count from highest to lowest. +907. Show invoice line items for each preparer ordered by count from highest to lowest. +908. Show invoice line items for each reviewer ordered by count from highest to lowest. +909. Show invoice line items for each reason code ordered by count from highest to lowest. +910. Show invoice line items for each workflow status ordered by count from highest to lowest. +911. Show journal entries for each business unit ordered by count from highest to lowest. +912. Show journal entries for each supplier ordered by count from highest to lowest. +913. Show journal entries for each product line ordered by count from highest to lowest. +914. Show journal entries for each product ordered by count from highest to lowest. +915. Show journal entries for each currency ordered by count from highest to lowest. +916. Show journal entries for each cost center ordered by count from highest to lowest. +917. Show journal entries for each profit center ordered by count from highest to lowest. +918. Show journal entries for each customer ordered by count from highest to lowest. +919. Show journal entries for each invoice type ordered by count from highest to lowest. +920. Show journal entries for each priority ordered by count from highest to lowest. +921. Show journal entries for each entry type ordered by count from highest to lowest. +922. Show journal entries for each task owner ordered by count from highest to lowest. +923. Show journal entries for each task status ordered by count from highest to lowest. +924. Show journal entries for each posting status ordered by count from highest to lowest. +925. Show journal entries for each journal status ordered by count from highest to lowest. +926. Show journal entries for each GL account ordered by count from highest to lowest. +927. Show journal entries for each preparer ordered by count from highest to lowest. +928. Show journal entries for each reviewer ordered by count from highest to lowest. +929. Show journal entries for each reason code ordered by count from highest to lowest. +930. Show journal entries for each workflow status ordered by count from highest to lowest. +931. Show documents for each business unit ordered by count from highest to lowest. +932. Show documents for each supplier ordered by count from highest to lowest. +933. Show documents for each product line ordered by count from highest to lowest. +934. Show documents for each product ordered by count from highest to lowest. +935. Show documents for each currency ordered by count from highest to lowest. +936. Show documents for each cost center ordered by count from highest to lowest. +937. Show documents for each profit center ordered by count from highest to lowest. +938. Show documents for each customer ordered by count from highest to lowest. +939. Show documents for each invoice type ordered by count from highest to lowest. +940. Show documents for each priority ordered by count from highest to lowest. +941. Show documents for each entry type ordered by count from highest to lowest. +942. Show documents for each task owner ordered by count from highest to lowest. +943. Show documents for each task status ordered by count from highest to lowest. +944. Show documents for each posting status ordered by count from highest to lowest. +945. Show documents for each journal status ordered by count from highest to lowest. +946. Show documents for each GL account ordered by count from highest to lowest. +947. Show documents for each preparer ordered by count from highest to lowest. +948. Show documents for each reviewer ordered by count from highest to lowest. +949. Show documents for each reason code ordered by count from highest to lowest. +950. Show documents for each workflow status ordered by count from highest to lowest. +951. Show tasks for each business unit ordered by count from highest to lowest. +952. Show tasks for each supplier ordered by count from highest to lowest. +953. Show tasks for each product line ordered by count from highest to lowest. +954. Show tasks for each product ordered by count from highest to lowest. +955. Show tasks for each currency ordered by count from highest to lowest. +956. Show tasks for each cost center ordered by count from highest to lowest. +957. Show tasks for each profit center ordered by count from highest to lowest. +958. Show tasks for each customer ordered by count from highest to lowest. +959. Show tasks for each invoice type ordered by count from highest to lowest. +960. Show tasks for each priority ordered by count from highest to lowest. +961. Show tasks for each entry type ordered by count from highest to lowest. +962. Show tasks for each task owner ordered by count from highest to lowest. +963. Show tasks for each task status ordered by count from highest to lowest. +964. Show tasks for each posting status ordered by count from highest to lowest. +965. Show tasks for each journal status ordered by count from highest to lowest. +966. Show tasks for each GL account ordered by count from highest to lowest. +967. Show tasks for each preparer ordered by count from highest to lowest. +968. Show tasks for each reviewer ordered by count from highest to lowest. +969. Show tasks for each reason code ordered by count from highest to lowest. +970. Show tasks for each workflow status ordered by count from highest to lowest. +971. Show suppliers for each business unit ordered by count from highest to lowest. +972. Show suppliers for each supplier ordered by count from highest to lowest. +973. Show suppliers for each product line ordered by count from highest to lowest. +974. Show suppliers for each product ordered by count from highest to lowest. +975. Show suppliers for each currency ordered by count from highest to lowest. +976. Show suppliers for each cost center ordered by count from highest to lowest. +977. Show suppliers for each profit center ordered by count from highest to lowest. +978. Show suppliers for each customer ordered by count from highest to lowest. +979. Show suppliers for each invoice type ordered by count from highest to lowest. +980. Show suppliers for each priority ordered by count from highest to lowest. +981. Show suppliers for each entry type ordered by count from highest to lowest. +982. Show suppliers for each task owner ordered by count from highest to lowest. +983. Show suppliers for each task status ordered by count from highest to lowest. +984. Show suppliers for each posting status ordered by count from highest to lowest. +985. Show suppliers for each journal status ordered by count from highest to lowest. +986. Show suppliers for each GL account ordered by count from highest to lowest. +987. Show suppliers for each preparer ordered by count from highest to lowest. +988. Show suppliers for each reviewer ordered by count from highest to lowest. +989. Show suppliers for each reason code ordered by count from highest to lowest. +990. Show suppliers for each workflow status ordered by count from highest to lowest. +991. Show prepayments for each business unit ordered by count from highest to lowest. +992. Show prepayments for each supplier ordered by count from highest to lowest. +993. Show prepayments for each product line ordered by count from highest to lowest. +994. Show prepayments for each product ordered by count from highest to lowest. +995. Show prepayments for each currency ordered by count from highest to lowest. +996. Show prepayments for each cost center ordered by count from highest to lowest. +997. Show prepayments for each profit center ordered by count from highest to lowest. +998. Show prepayments for each customer ordered by count from highest to lowest. +999. Show prepayments for each invoice type ordered by count from highest to lowest. +1000. Show prepayments for each priority ordered by count from highest to lowest. From 923fa0977d9aafc4af5619de74a3ea72050ac469 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 15:41:15 +0530 Subject: [PATCH 1063/1087] Handle malformed SQL tokens in grounding --- .../src/pipelines/generation/utils/sql.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index d42b9b8b17..46c7a306e3 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -471,22 +471,28 @@ def _extract_qualified_columns(token_list: TokenList) -> list[tuple[str, str]]: for token in token_list.tokens: if isinstance(token, Identifier): - parent_name = _clean_identifier(token.get_parent_name()) - column_name = _clean_identifier(token.get_real_name()) - if parent_name and column_name and column_name != "*": - columns.append((parent_name, column_name)) + _add_qualified_column(token, columns) elif isinstance(token, IdentifierList): for identifier in token.get_identifiers(): - parent_name = _clean_identifier(identifier.get_parent_name()) - column_name = _clean_identifier(identifier.get_real_name()) - if parent_name and column_name and column_name != "*": - columns.append((parent_name, column_name)) + if isinstance(identifier, Identifier): + _add_qualified_column(identifier, columns) + elif isinstance(identifier, TokenList): + columns.extend(_extract_qualified_columns(identifier)) elif isinstance(token, TokenList): columns.extend(_extract_qualified_columns(token)) return columns +def _add_qualified_column( + identifier: Identifier, columns: list[tuple[str, str]] +) -> None: + parent_name = _clean_identifier(identifier.get_parent_name()) + column_name = _clean_identifier(identifier.get_real_name()) + if parent_name and column_name and column_name != "*": + columns.append((parent_name, column_name)) + + def _clean_identifier(identifier: str | None) -> str | None: if identifier is None: return None From 8ffe9b1bf75be02b52dc3c793bb04c49910d2de2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 16:10:01 +0530 Subject: [PATCH 1064/1087] Restore diagnostic context for SQL correction --- .../pipelines/generation/sql_correction.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index fb21794d02..e6a407bfee 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -93,12 +93,27 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. {% endif %} ### FAILED SQL ### -The failed SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. +The failed SQL below is diagnostic context only. It is not an executable schema source. +Only preserve an identifier, function, literal filter, grouping, ordering, or join from this SQL when it is also declared exactly in the WREN SQL IDENTIFIER CONTRACT, DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. +If it contains placeholders, assumed business names, connector-specific syntax, source/physical names, or unsupported objects, discard those parts and regenerate from the QUESTION plus DATABASE SCHEMA. + +{% if invalid_generation_result and invalid_generation_result.sql %} +{{ invalid_generation_result.sql }} +{% else %} +No failed SQL was provided. +{% endif %} ### DRY-RUN DIAGNOSTIC ### -The dry-run diagnostic text is intentionally omitted because it may contain failed SQL, guessed identifiers, connector-specific syntax, source names, physical names, or invalid replacement candidates. +The diagnostic below explains why the previous SQL failed. Use it to understand the failure only. +Do not copy identifiers, source names, physical names, SQL fragments, or replacement candidates from the diagnostic unless they appear exactly in the WREN SQL IDENTIFIER CONTRACT, DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. + +{% if invalid_generation_result and invalid_generation_result.error %} +{{ invalid_generation_result.error }} +{% else %} +No dry-run diagnostic was provided. +{% endif %} -Regenerate from the user question and current DATABASE SCHEMA only. Do not repair, preserve, or copy anything from the failed SQL or dry-run diagnostic. +Regenerate from the user question, current DATABASE SCHEMA, and the diagnostic failure. Keep DATABASE SCHEMA as the only executable identifier source. Return only the final JSON SQL response. """ From 510b938c7c19c2399459ee0697b9010a77efc10b Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 16:35:19 +0530 Subject: [PATCH 1065/1087] Parse retrieved schema contracts for SQL grounding --- .../src/pipelines/generation/utils/sql.py | 241 ++++++++++++++++-- 1 file changed, 214 insertions(+), 27 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 46c7a306e3..eafcb8bd4b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, List import aiohttp @@ -19,6 +20,20 @@ logger = logging.getLogger("wren-ai-service") +_DDL_CREATE_PATTERN = re.compile( + r"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?P
[^\s(]+)\s*\(", + re.IGNORECASE, +) +_DDL_COLUMN_KEYWORDS = { + "CHECK", + "CONSTRAINT", + "FOREIGN", + "INDEX", + "KEY", + "PRIMARY", + "UNIQUE", +} + @component class SQLGenPostProcessor: @@ -239,42 +254,87 @@ def __init__( def from_contexts(cls, contexts: list[str]) -> "_SchemaCatalog": tables: dict[str, set[str]] = {} relationships: dict[str, set[str]] = {} + + for context in contexts: + cls._add_contract_identifiers(context, tables, relationships) + cls._add_ddl_identifiers(context, tables) + + return cls(tables, relationships) + + @staticmethod + def _add_contract_identifiers( + context: str, + tables: dict[str, set[str]], + relationships: dict[str, set[str]], + ) -> None: current_table: str | None = None current_section: str | None = None - for context in contexts: - for raw_line in context.splitlines(): - line = raw_line.strip() - if line.startswith("table: "): - current_table = line.removeprefix("table: ").strip() - if current_table: - tables.setdefault(current_table, set()) - relationships.setdefault(current_table, set()) - current_section = None - continue + for raw_line in context.splitlines(): + line = raw_line.strip() + if line.startswith("table: "): + current_table = _clean_identifier(line.removeprefix("table: ")) + if current_table: + tables.setdefault(current_table, set()) + relationships.setdefault(current_table, set()) + current_section = None + continue - if current_table and line == "columns:": - current_section = "columns" - continue + if line.startswith("sql_table_name_use_exactly:"): + current_table = _clean_identifier( + line.removeprefix("sql_table_name_use_exactly:") + ) + if current_table: + tables.setdefault(current_table, set()) + relationships.setdefault(current_table, set()) + current_section = None + continue - if current_table and line == "relationships:": - current_section = "relationships" - continue + if current_table and line in {"columns:", "sql_column_names_use_exactly:"}: + current_section = "columns" + continue + + if current_table and line in { + "relationships:", + "relationship_constraints_use_exactly:", + }: + current_section = "relationships" + continue - if current_section and current_table and line.startswith("- "): - value = line.removeprefix("- ").strip() - if not value: - continue - if current_section == "columns": - tables.setdefault(current_table, set()).add(value) - elif current_section == "relationships": - relationships.setdefault(current_table, set()).add(value) + if current_section and current_table and line.startswith("- "): + value = line.removeprefix("- ").strip() + if not value: continue + if current_section == "columns": + column_name = _clean_identifier(value) + if column_name: + tables.setdefault(current_table, set()).add(column_name) + elif current_section == "relationships": + relationships.setdefault(current_table, set()).add(value) + continue - if current_section and line and not line.startswith("- "): - current_section = None + if current_section and line and not line.startswith("- "): + current_section = None - return cls(tables, relationships) + @staticmethod + def _add_ddl_identifiers( + context: str, + tables: dict[str, set[str]], + ) -> None: + for match in _DDL_CREATE_PATTERN.finditer(context): + table_name = _clean_identifier(match.group("table")) + if not table_name: + continue + + body_start = match.end() + body_end = _find_matching_parenthesis(context, body_start) + if body_end is None: + tables.setdefault(table_name, set()) + continue + + tables.setdefault(table_name, set()).update( + _extract_ddl_column_names(context[body_start:body_end]) + ) def to_prompt(self) -> str: if not self._tables: @@ -493,6 +553,133 @@ def _add_qualified_column( columns.append((parent_name, column_name)) +def _find_matching_parenthesis(text: str, body_start: int) -> int | None: + depth = 1 + quote: str | None = None + i = body_start + + while i < len(text): + char = text[i] + if quote: + if char == quote: + if quote == "'" and i + 1 < len(text) and text[i + 1] == "'": + i += 2 + continue + quote = None + i += 1 + continue + + if char in {"'", '"', "`"}: + quote = char + elif char == "[": + closing = text.find("]", i + 1) + if closing == -1: + return None + i = closing + elif char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + return i + i += 1 + + return None + + +def _extract_ddl_column_names(ddl_body: str) -> set[str]: + column_names: set[str] = set() + for column_definition in _split_top_level_commas(ddl_body): + column_name = _extract_ddl_column_name(column_definition) + if column_name: + column_names.add(column_name) + return column_names + + +def _split_top_level_commas(value: str) -> list[str]: + parts: list[str] = [] + depth = 0 + quote: str | None = None + start = 0 + i = 0 + + while i < len(value): + char = value[i] + if quote: + if char == quote: + if quote == "'" and i + 1 < len(value) and value[i + 1] == "'": + i += 2 + continue + quote = None + i += 1 + continue + + if char in {"'", '"', "`"}: + quote = char + elif char == "[": + closing = value.find("]", i + 1) + if closing == -1: + break + i = closing + elif char == "(": + depth += 1 + elif char == ")": + depth = max(depth - 1, 0) + elif char == "," and depth == 0: + parts.append(value[start:i].strip()) + start = i + 1 + i += 1 + + tail = value[start:].strip() + if tail: + parts.append(tail) + return parts + + +def _extract_ddl_column_name(column_definition: str) -> str | None: + definition = _strip_leading_sql_comments(column_definition.strip()) + if not definition: + return None + + first_word = definition.split(maxsplit=1)[0].strip().strip('"`[]').upper() + if first_word in _DDL_COLUMN_KEYWORDS: + return None + + if definition.startswith("["): + closing = definition.find("]") + if closing > 0: + return _clean_identifier(definition[: closing + 1]) + + if definition.startswith('"'): + closing = definition.find('"', 1) + if closing > 0: + return _clean_identifier(definition[: closing + 1]) + + if definition.startswith("`"): + closing = definition.find("`", 1) + if closing > 0: + return _clean_identifier(definition[: closing + 1]) + + return _clean_identifier(definition.split(maxsplit=1)[0]) + + +def _strip_leading_sql_comments(value: str) -> str: + stripped = value.strip() + while stripped: + if stripped.startswith("--"): + lines = stripped.splitlines() + stripped = "\n".join(lines[1:]).strip() + continue + if stripped.startswith("/*"): + closing = stripped.find("*/") + if closing == -1: + return "" + stripped = stripped[closing + 2 :].strip() + continue + break + return stripped + + def _clean_identifier(identifier: str | None) -> str | None: if identifier is None: return None From e7b6e5a1a256343f4143c25a14374884b35bc7c2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Mon, 17 Aug 2026 22:23:31 +0530 Subject: [PATCH 1066/1087] Ground ask SQL generation in schema contract --- .../generation/followup_sql_generation.py | 5 +- .../followup_sql_generation_reasoning.py | 6 +- .../pipelines/generation/sql_correction.py | 5 +- .../pipelines/generation/sql_generation.py | 5 +- .../generation/sql_generation_reasoning.py | 6 +- .../pipelines/generation/sql_regeneration.py | 5 +- .../src/pipelines/generation/utils/sql.py | 794 +++++++++++++++++- wren-ai-service/src/web/v1/services/ask.py | 5 + 8 files changed, 808 insertions(+), 23 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 b09b524756..5b5542a5ce 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -153,6 +153,7 @@ async def post_process( documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: @@ -160,10 +161,10 @@ async def post_process( generate_sql_in_followup.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + contexts=validation_contexts or documents, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - contexts=documents, ) @@ -216,6 +217,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + validation_contexts: list[str] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -245,6 +247,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index e84955f87d..553e5ac6f6 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -78,17 +78,19 @@ def prompt( sql_samples: list[dict], instructions: list[dict], prompt_builder: PromptBuilder, + validation_contexts: Optional[list[str]] = None, configuration: Configuration | None = Configuration(), ) -> dict: + schema_documents = validation_contexts or documents _prompt = prompt_builder.run( query=query, - documents=documents, + documents=schema_documents, histories=histories, sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, ), - schema_identifier_catalog=construct_schema_identifier_catalog(documents), + schema_identifier_catalog=construct_schema_identifier_catalog(schema_documents), language=configuration.language, current_time=configuration.show_current_time(), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index e6a407bfee..ec78f998fb 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -166,6 +166,7 @@ async def post_process( documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: @@ -173,10 +174,10 @@ async def post_process( generate_sql_correction.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + contexts=validation_contexts or documents, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, - contexts=documents, ) @@ -225,6 +226,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + validation_contexts: list[str] | None = None, ): logger.info("SQLCorrection pipeline is running...") @@ -250,6 +252,7 @@ async def run( "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **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 b4be95de49..5bdeeb8827 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -144,6 +144,7 @@ async def post_process( documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, @@ -152,11 +153,11 @@ async def post_process( generate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, + contexts=validation_contexts or documents, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, - contexts=documents, ) @@ -209,6 +210,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + validation_contexts: list[str] | None = None, ): logger.info("SQL Generation pipeline is running...") @@ -238,6 +240,7 @@ async def run( "data_source": metadata.get("data_source", "local_file"), "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **self._components, }, ) 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 d44532c6ee..fd38b9a291 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -68,16 +68,18 @@ def prompt( sql_samples: list[dict], instructions: list[dict], prompt_builder: PromptBuilder, + validation_contexts: Optional[list[str]] = None, configuration: Configuration | None = Configuration(), ) -> dict: + schema_documents = validation_contexts or documents _prompt = prompt_builder.run( query=query, - documents=documents, + documents=schema_documents, sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, ), - schema_identifier_catalog=construct_schema_identifier_catalog(documents), + schema_identifier_catalog=construct_schema_identifier_catalog(schema_documents), language=configuration.language, current_time=configuration.show_current_time(), ) diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 67674b0b09..604d63f40e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -164,12 +164,13 @@ async def post_process( documents: list[str], project_id: str | None = None, mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - contexts=documents, + contexts=validation_contexts or documents, ) @@ -214,6 +215,7 @@ async def run( has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + validation_contexts: list[str] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -232,6 +234,7 @@ async def run( "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index eafcb8bd4b..25becffddf 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -35,6 +35,716 @@ } +_IDENTIFIER_TOKEN = r'"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*' +_QUALIFIED_IDENTIFIER = rf"(?:{_IDENTIFIER_TOKEN})(?:\s*\.\s*(?:{_IDENTIFIER_TOKEN}))*" +_SIMPLE_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$") +_DDL_RELATION = re.compile( + rf"\bCREATE\s+(?:TABLE|VIEW)\s+(?P{_QUALIFIED_IDENTIFIER})", + re.IGNORECASE, +) +_RELATION_REFERENCE = re.compile( + rf"\b(?:FROM|JOIN)\s+(?P{_QUALIFIED_IDENTIFIER})" + rf"(?:\s+(?:AS\s+)?(?P{_IDENTIFIER_TOKEN}))?", + re.IGNORECASE, +) +_CTE_REFERENCE = re.compile( + rf"(?:\bWITH|,)\s+(?P{_IDENTIFIER_TOKEN})\s+AS\s*\(", + re.IGNORECASE, +) +_QUALIFIED_COLUMN = re.compile( + rf"(?P{_QUALIFIED_IDENTIFIER})\s*\.\s*(?P{_IDENTIFIER_TOKEN})" +) +_SQL_START = re.compile(r"^\s*(?:WITH|SELECT)\b", re.IGNORECASE | re.DOTALL) +_SQL_OBJECT_ALIAS_STOP_WORDS = { + "CROSS", + "EXCEPT", + "FETCH", + "FULL", + "GROUP", + "HAVING", + "INNER", + "INTERSECT", + "JOIN", + "LEFT", + "LIMIT", + "MATCH_RECOGNIZE", + "NATURAL", + "OFFSET", + "ORDER", + "RIGHT", + "TABLESAMPLE", + "UNION", + "WHERE", +} +_SQL_RESERVED_WORDS = { + "ALL", + "ALTER", + "AND", + "AS", + "BY", + "CASE", + "CAST", + "COUNT", + "CREATE", + "CROSS", + "DELETE", + "DESC", + "DISTINCT", + "ELSE", + "END", + "EXCEPT", + "FALSE", + "FETCH", + "FOR", + "FROM", + "FULL", + "GROUP", + "HAVING", + "IN", + "INNER", + "INSERT", + "INTERSECT", + "IS", + "JOIN", + "LEFT", + "LIKE", + "LIMIT", + "NATURAL", + "NOT", + "NULL", + "OFFSET", + "ON", + "OR", + "ORDER", + "OUTER", + "RIGHT", + "SELECT", + "TABLE", + "TABLESAMPLE", + "THEN", + "TRUE", + "UNION", + "UPDATE", + "WHEN", + "WHERE", + "WINDOW", + "WITH", +} + + +def _unquote_identifier(identifier: str) -> str: + identifier = identifier.strip() + if len(identifier) >= 2 and identifier[0] == "[" and identifier[-1] == "]": + return identifier[1:-1].replace("]]", "]") + if len(identifier) >= 2 and identifier[0] == '"' and identifier[-1] == '"': + return identifier[1:-1].replace('""', '"') + return identifier + + +def _split_qualified_identifier(identifier: str) -> list[str]: + parts = [] + current = [] + in_double_quote = False + in_bracket = False + index = 0 + + while index < len(identifier): + char = identifier[index] + nxt = identifier[index + 1] if index + 1 < len(identifier) else None + + if in_double_quote: + current.append(char) + if char == '"' and nxt == '"': + current.append(nxt) + index += 2 + continue + if char == '"': + in_double_quote = False + index += 1 + continue + + if in_bracket: + current.append(char) + if char == "]": + in_bracket = False + index += 1 + continue + + if char == '"': + current.append(char) + in_double_quote = True + elif char == "[": + current.append(char) + in_bracket = True + elif char == ".": + part = "".join(current).strip() + if part: + parts.append(part) + current = [] + else: + current.append(char) + index += 1 + + part = "".join(current).strip() + if part: + parts.append(part) + + return parts + + +def _normalize_identifier(identifier: str) -> str: + return ".".join(_unquote_identifier(part) for part in _split_qualified_identifier(identifier)) + + +def _quote_identifier(identifier: str) -> str: + return f'"{identifier.replace(chr(34), chr(34) * 2)}"' + + +def _split_sql_tokens(sql: str) -> list[str]: + tokens = [] + current = [] + in_double_quote = False + + for char in sql: + if char == '"': + current.append(char) + in_double_quote = not in_double_quote + continue + if char == "," and not in_double_quote: + token = "".join(current).strip() + if token: + tokens.append(token) + current = [] + continue + current.append(char) + + token = "".join(current).strip() + if token: + tokens.append(token) + return tokens + + +def _identifier_needs_quotes(identifier: str) -> bool: + return ( + not _SIMPLE_IDENTIFIER.fullmatch(identifier) + or identifier.upper() in _SQL_RESERVED_WORDS + ) + + +def _iter_context_texts(contexts: list[Any] | None): + if not contexts: + return + for context in contexts: + yield getattr(context, "content", context) + + +def _clean_contract_value(value: str) -> str: + value = value.strip().strip(",") + if not value: + return "" + return _normalize_identifier(value) + + +def _parse_contract_values(value: str) -> list[str]: + value = value.strip().strip(",") + if not value: + return [] + + try: + loaded = orjson.loads(value) + except orjson.JSONDecodeError: + loaded = None + + if isinstance(loaded, list): + return [ + _clean_contract_value(str(item)) + for item in loaded + if _clean_contract_value(str(item)) + ] + + parsed = _clean_contract_value(value) + return [parsed] if parsed else [] + + +def _extract_contract_schema_index( + contexts: list[Any] | None, +) -> dict[str, set[str] | None]: + schema_index: dict[str, set[str] | None] = {} + if not contexts: + return schema_index + + for context in _iter_context_texts(contexts): + current_relation = None + reading_columns = False + + for raw_line in str(context).splitlines(): + line = raw_line.strip() + if not line: + continue + + if line.startswith("sql_table_name_use_exactly:"): + current_relation = _clean_contract_value(line.split(":", 1)[1]) + if current_relation: + schema_index.setdefault(current_relation, set()) + reading_columns = False + continue + + if line.startswith("sql_column_names_use_exactly:"): + reading_columns = True + if current_relation: + columns = schema_index.setdefault(current_relation, set()) + if columns is not None: + for value in _parse_contract_values(line.split(":", 1)[1]): + columns.add(value) + continue + + if line.startswith("relationship_constraints_use_exactly:"): + reading_columns = False + continue + + if line.startswith("sql_column_name_use_exactly:"): + if current_relation: + columns = schema_index.setdefault(current_relation, set()) + if columns is not None: + for value in _parse_contract_values(line.split(":", 1)[1]): + columns.add(value) + continue + + if reading_columns and line.startswith("-") and current_relation: + columns = schema_index.setdefault(current_relation, set()) + if columns is not None: + value = _clean_contract_value(line[1:]) + if value: + columns.add(value) + continue + + if line.startswith(("END WREN SQL IDENTIFIER CONTRACT", "Only ")): + reading_columns = False + + return schema_index + + +def _extract_schema_identifiers(contexts: list[Any] | None) -> list[str]: + if not contexts: + return [] + + identifiers: list[str] = [] + seen = set() + + def add(identifier: str) -> None: + identifier = _unquote_identifier(identifier.strip().rstrip(",")) + if not identifier or identifier.upper() in {"FOREIGN", "PRIMARY", "KEY"}: + return + if identifier not in seen: + seen.add(identifier) + identifiers.append(identifier) + + for relation, columns in _extract_schema_index(contexts).items(): + add(relation) + if columns: + for column in columns: + add(column) + + for context in _iter_context_texts(contexts): + for match in _DDL_RELATION.finditer(context): + add(match.group("name")) + + in_table = False + for raw_line in context.splitlines(): + line = raw_line.strip() + if not line or line.startswith("--") or line.startswith("/*"): + continue + if re.search(r"\bCREATE\s+TABLE\b", line, re.IGNORECASE): + in_table = True + remainder = line.split("(", 1) + if len(remainder) == 1: + continue + line = remainder[1].strip() + if not in_table: + continue + if line.startswith(");") or line == ")": + in_table = False + continue + line = line.split("--", 1)[0].strip().rstrip(",") + if not line or line.upper().startswith(("FOREIGN KEY", "PRIMARY KEY")): + continue + if line.startswith('"'): + end = line.find('"', 1) + while end != -1 and end + 1 < len(line) and line[end + 1] == '"': + end = line.find('"', end + 2) + if end > 0: + add(line[: end + 1]) + else: + add(line.split(None, 1)[0]) + + return identifiers + + +def _extract_schema_index(contexts: list[Any] | None) -> dict[str, set[str] | None]: + if not contexts: + return {} + + schema_index = _extract_contract_schema_index(contexts) + + for context in _iter_context_texts(contexts): + relation_match = _DDL_RELATION.search(context) + if not relation_match: + continue + + relation_name = _normalize_identifier(relation_match.group("name")) + if re.search(r"\bCREATE\s+VIEW\b", context, re.IGNORECASE): + schema_index.setdefault(relation_name, None) + continue + + column_block_match = re.search( + r"\bCREATE\s+TABLE\b[^(]*\((?P.*)\)\s*;?", + context, + re.IGNORECASE | re.DOTALL, + ) + if not column_block_match: + schema_index.setdefault(relation_name, None) + continue + + columns = set() + for raw_column in _split_sql_tokens(column_block_match.group("columns")): + line = "\n".join( + line.strip() + for line in raw_column.splitlines() + if line.strip() + and not line.strip().startswith("--") + and not line.strip().startswith("/*") + ).strip() + if not line: + continue + line = line.split("--", 1)[0].strip() + if not line or line.startswith("/*"): + continue + if line.upper().startswith(("FOREIGN KEY", "PRIMARY KEY")): + continue + if line.startswith('"'): + end = line.find('"', 1) + while end != -1 and end + 1 < len(line) and line[end + 1] == '"': + end = line.find('"', end + 2) + if end > 0: + columns.add(_unquote_identifier(line[: end + 1])) + else: + columns.add(_unquote_identifier(line.split(None, 1)[0])) + + existing_columns = schema_index.get(relation_name) + if existing_columns is None: + schema_index[relation_name] = columns + else: + existing_columns.update(columns) + + return schema_index + + +def _is_identifier_boundary(char: str | None) -> bool: + return char is None or not (char.isalnum() or char in {"_", "$", '"'}) + + +def _replace_identifier_outside_literals(sql: str, identifier: str) -> str: + quoted = _quote_identifier(identifier) + result = [] + index = 0 + in_single_quote = False + in_double_quote = False + in_line_comment = False + in_block_comment = False + length = len(sql) + identifier_length = len(identifier) + + while index < length: + current = sql[index] + nxt = sql[index + 1] if index + 1 < length else None + + if in_line_comment: + result.append(current) + if current == "\n": + in_line_comment = False + index += 1 + continue + if in_block_comment: + result.append(current) + if current == "*" and nxt == "/": + result.append(nxt) + index += 2 + in_block_comment = False + else: + index += 1 + continue + if in_single_quote: + result.append(current) + if current == "'" and nxt == "'": + result.append(nxt) + index += 2 + elif current == "'": + in_single_quote = False + index += 1 + else: + index += 1 + continue + if in_double_quote: + result.append(current) + if current == '"' and nxt == '"': + result.append(nxt) + index += 2 + elif current == '"': + in_double_quote = False + index += 1 + else: + index += 1 + continue + + if current == "-" and nxt == "-": + result.append(current) + result.append(nxt) + index += 2 + in_line_comment = True + continue + if current == "/" and nxt == "*": + result.append(current) + result.append(nxt) + index += 2 + in_block_comment = True + continue + if current == "'": + result.append(current) + index += 1 + in_single_quote = True + continue + if current == '"': + result.append(current) + index += 1 + in_double_quote = True + continue + + if sql.startswith(identifier, index): + before = sql[index - 1] if index > 0 else None + after_index = index + identifier_length + after = sql[after_index] if after_index < length else None + if _is_identifier_boundary(before) and _is_identifier_boundary(after): + result.append(quoted) + index = after_index + continue + + result.append(current) + index += 1 + + return "".join(result) + + +def _replace_bracket_identifiers(sql: str, valid_identifiers: set[str]) -> str: + result = [] + index = 0 + in_single_quote = False + in_double_quote = False + length = len(sql) + + while index < length: + current = sql[index] + nxt = sql[index + 1] if index + 1 < length else None + + if in_single_quote: + result.append(current) + if current == "'" and nxt == "'": + result.append(nxt) + index += 2 + elif current == "'": + in_single_quote = False + index += 1 + else: + index += 1 + continue + if in_double_quote: + result.append(current) + if current == '"' and nxt == '"': + result.append(nxt) + index += 2 + elif current == '"': + in_double_quote = False + index += 1 + else: + index += 1 + continue + if current == "'": + result.append(current) + in_single_quote = True + index += 1 + continue + if current == '"': + result.append(current) + in_double_quote = True + index += 1 + continue + if current == "[": + end = sql.find("]", index + 1) + if end > index: + identifier = sql[index + 1 : end] + if identifier in valid_identifiers: + result.append(_quote_identifier(identifier)) + index = end + 1 + continue + result.append(current) + index += 1 + + return "".join(result) + + +def _extract_sql_grounding(sql: str) -> dict[str, Any]: + cte_names = { + _normalize_identifier(match.group("name")) + for match in _CTE_REFERENCE.finditer(sql) + } + relation_references = [] + alias_to_relation = {} + + for match in _RELATION_REFERENCE.finditer(sql): + relation = _normalize_identifier(match.group("name")) + if relation.upper() in {"UNNEST", "LATERAL"}: + continue + alias = match.group("alias") + alias = _normalize_identifier(alias) if alias else relation + if alias.upper() in _SQL_OBJECT_ALIAS_STOP_WORDS: + alias = relation + relation_references.append(relation) + alias_to_relation[alias] = relation + alias_to_relation[relation] = relation + + qualified_columns = [ + ( + _normalize_identifier(match.group("qualifier")), + _normalize_identifier(match.group("column")), + ) + for match in _QUALIFIED_COLUMN.finditer(sql) + ] + + return { + "cte_names": cte_names, + "relation_references": relation_references, + "alias_to_relation": alias_to_relation, + "qualified_columns": qualified_columns, + } + + +def validate_sql_against_contexts( + sql: str, + contexts: list[Any] | None = None, +) -> str | None: + schema_index = _extract_schema_index(contexts) + if not schema_index: + return None + + valid_relations = set(schema_index) + grounding = _extract_sql_grounding(sql) + cte_names = grounding["cte_names"] + + shadowed_relations = sorted(cte_names & valid_relations) + if shadowed_relations: + return ( + "Schema grounding failed. The SQL creates CTEs with names that already " + f"belong to verified schema objects: {', '.join(shadowed_relations)}. " + "Do not create dummy CTEs for schema objects; use the verified tables or views directly." + ) + + invalid_relations = sorted( + { + relation + for relation in grounding["relation_references"] + if relation not in valid_relations and relation not in cte_names + } + ) + if invalid_relations: + return ( + "Schema grounding failed. The SQL references tables or views that are not " + f"in the retrieved schema for the active question: {', '.join(invalid_relations)}. " + f"Use only verified tables or views: {', '.join(sorted(valid_relations))}." + ) + + alias_to_relation = grounding["alias_to_relation"] + for qualifier, column in grounding["qualified_columns"]: + relation = alias_to_relation.get(qualifier) + if not relation or relation in cte_names: + continue + valid_columns = schema_index.get(relation) + if valid_columns is None: + continue + if column not in valid_columns: + return ( + "Schema grounding failed. The SQL references column " + f"{qualifier}.{column}, but column {column} is not present in verified " + f"table or view {relation}. Use only verified columns: " + f"{', '.join(sorted(valid_columns))}." + ) + + return None + + +def _extract_sql_from_value(value: Any) -> str | None: + if value is None: + return None + + if isinstance(value, str): + text = value.strip() + if not text: + return None + + try: + parsed = orjson.loads(text) + except orjson.JSONDecodeError: + return text if _SQL_START.search(text) else None + + return _extract_sql_from_value(parsed) + + if isinstance(value, dict): + for key in ("sql", "query", "code"): + extracted = _extract_sql_from_value(value.get(key)) + if extracted: + return extracted + + extracted = _extract_sql_from_value(value.get("arguments")) + if extracted: + return extracted + + return None + + if isinstance(value, list): + for item in value: + extracted = _extract_sql_from_value(item) + if extracted: + return extracted + + return None + + +def _extract_generation_sql(generation_result: str | None) -> str | None: + if not generation_result: + return None + + extracted = _extract_sql_from_value(generation_result) + if extracted: + return extracted + + text = generation_result.strip() + return text if _SQL_START.search(text) else None + + +def normalize_sql_with_schema_identifiers( + sql: str, + contexts: list[Any] | None = None, +) -> str: + schema_identifiers = set(_extract_schema_identifiers(contexts)) + identifiers = [ + identifier + for identifier in schema_identifiers + if "." not in identifier and _identifier_needs_quotes(identifier) + ] + sql = _replace_bracket_identifiers(sql, schema_identifiers) + for identifier in sorted(identifiers, key=len, reverse=True): + sql = _replace_identifier_outside_literals(sql, identifier) + return sql + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -68,9 +778,16 @@ async def run( "type": "NO_RELEVANT_SQL", "error": extraction_error, "correlation_id": "", + "data_source": data_source, }, } + if cleaned_generation_result: + cleaned_generation_result = normalize_sql_with_schema_identifiers( + cleaned_generation_result, + contexts=contexts, + ) + schema_catalog = _SchemaCatalog.from_contexts(contexts or []) grounding_error = schema_catalog.validate_sql(cleaned_generation_result) if grounding_error: @@ -82,7 +799,8 @@ async def run( "type": "SCHEMA_GROUNDING", "error": grounding_error, "correlation_id": "", - }, + "data_source": data_source, + } } ( @@ -507,6 +1225,9 @@ def _extract_table_references(token_list: TokenList) -> tuple[set[str], dict[str def _add_table_reference( identifier: Identifier, table_names: set[str], aliases: dict[str, str] ) -> None: + if not isinstance(identifier, Identifier): + return + table_name = _table_reference_name(identifier) alias = _clean_identifier(identifier.get_alias()) if not table_name: @@ -519,8 +1240,13 @@ def _add_table_reference( def _table_reference_name(identifier: Identifier) -> str | None: - parent_name = _clean_identifier(identifier.get_parent_name()) - real_name = _clean_identifier(identifier.get_real_name()) + if not isinstance(identifier, Identifier): + return None + + parent_getter = getattr(identifier, "get_parent_name", None) + real_getter = getattr(identifier, "get_real_name", None) + parent_name = _clean_identifier(parent_getter() if parent_getter else None) + real_name = _clean_identifier(real_getter() if real_getter else None) if parent_name and real_name: return f"{parent_name}.{real_name}" return real_name @@ -547,8 +1273,13 @@ def _extract_qualified_columns(token_list: TokenList) -> list[tuple[str, str]]: def _add_qualified_column( identifier: Identifier, columns: list[tuple[str, str]] ) -> None: - parent_name = _clean_identifier(identifier.get_parent_name()) - column_name = _clean_identifier(identifier.get_real_name()) + if not isinstance(identifier, Identifier): + return + + parent_getter = getattr(identifier, "get_parent_name", None) + real_getter = getattr(identifier, "get_real_name", None) + parent_name = _clean_identifier(parent_getter() if parent_getter else None) + column_name = _clean_identifier(real_getter() if real_getter else None) if parent_name and column_name and column_name != "*": columns.append((parent_name, column_name)) @@ -687,12 +1418,51 @@ def _clean_identifier(identifier: str | None) -> str | None: return cleaned or None +def _extract_sql_from_json_value(value: Any) -> str | None: + if isinstance(value, str): + candidate = value.strip() + if not candidate: + return None + if candidate.upper().startswith(("SELECT", "WITH")): + return candidate + if candidate.startswith(("{", "[")): + try: + return _extract_sql_from_json_value(orjson.loads(candidate)) + except orjson.JSONDecodeError: + return None + return None + + if isinstance(value, dict): + for key in ("sql", "query"): + sql = _extract_sql_from_json_value(value.get(key)) + if sql: + return sql + + for key in ("arguments", "content", "tool_calls", "function_call", "message"): + sql = _extract_sql_from_json_value(value.get(key)) + if sql: + return sql + + for nested_value in value.values(): + sql = _extract_sql_from_json_value(nested_value) + if sql: + return sql + + if isinstance(value, list): + for item in value: + sql = _extract_sql_from_json_value(item) + if sql: + return sql + + return None + + def _extract_sql_response(generation_result: str) -> tuple[str | None, str | None]: cleaned_generation_result = generation_result.strip() if not cleaned_generation_result: return None, "No grounded SQL was generated from the current schema." - if cleaned_generation_result.startswith("{"): + if cleaned_generation_result.startswith(("{", "[")): try: payload = orjson.loads(cleaned_generation_result) except orjson.JSONDecodeError: @@ -701,15 +1471,9 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non "SQL generation response did not include a supported SQL JSON payload.", ) - if "sql" in payload: - return payload.get("sql"), None - - if payload.get("name") == "query": - arguments = payload.get("arguments") - if isinstance(arguments, dict): - sql = arguments.get("query") or arguments.get("sql") - if sql: - return sql, None + sql = _extract_sql_from_json_value(payload) + if sql: + return sql, None return ( None, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 279cbe9aa6..a4a27d3a02 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -405,6 +405,7 @@ async def ask( instructions=instructions, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, + validation_contexts=table_ddls, configuration=ask_request.configurations, query_id=query_id, ) @@ -418,6 +419,7 @@ async def ask( instructions=instructions, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, + validation_contexts=table_ddls, configuration=ask_request.configurations, query_id=query_id, ) @@ -480,6 +482,7 @@ async def ask( histories=histories, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, + validation_contexts=table_ddls, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -499,6 +502,7 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, + validation_contexts=table_ddls, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -580,6 +584,7 @@ async def ask( }, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, + validation_contexts=table_ddls, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, From d9f7b84b90a8c0c1d83454d61107b235d636e317 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Tue, 18 Aug 2026 18:52:55 +0530 Subject: [PATCH 1067/1087] Fix semantics description chunk retry handling --- .../generation/semantics_description.py | 12 +- .../web/v1/services/semantics_description.py | 121 +++++++++++++++++- 2 files changed, 129 insertions(+), 4 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index d1d4db0475..f26206ea72 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -172,9 +172,17 @@ def wrapper(text: str) -> str: @observe(capture_input=False) def output(normalize: dict, picked_models: list[dict]) -> dict: def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: - valid_columns = [col["name"] for col in columns] + valid_columns = { + col.get("name") + for col in columns + if isinstance(col, dict) and col.get("name") + } - return [col for col in enriched if col["name"] in valid_columns] + return [ + col + for col in enriched or [] + if isinstance(col, dict) and col.get("name") in valid_columns + ] models = {model["name"]: model for model in picked_models} diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 85bf14b188..c914a40e2b 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -86,10 +86,122 @@ def _chunking( for chunk in chunks ] - async def _generate_task(self, request_id: str, chunk: dict): + def _description(self, payload: dict) -> str: + properties = payload.get("properties") + if not isinstance(properties, dict): + properties = {} + value = payload.get("description") or properties.get("description", "") + return "" if value is None else str(value).strip() + + def _validate_chunk_output(self, chunk: dict, output: dict) -> dict: + if not isinstance(output, dict): + raise ValueError("Semantics description pipeline returned invalid output") + + selected_models = set(chunk.get("selected_models", [])) + models = { + model.get("name"): model + for model in chunk.get("mdl", {}).get("models", []) + if model.get("name") in selected_models + } + + for model_name, model in models.items(): + generated_model = output.get(model_name) + if not isinstance(generated_model, dict): + raise ValueError( + f"Semantics description output omitted selected model: {model_name}" + ) + + if not self._description(generated_model): + raise ValueError( + "Semantics description output omitted description for model: " + f"{model_name}" + ) + + generated_columns = { + column.get("name"): column + for column in generated_model.get("columns", []) + if isinstance(column, dict) and column.get("name") + } + + for column in model.get("columns", []): + if not isinstance(column, dict): + continue + + column_name = column.get("name", "") + generated_column = generated_columns.get(column_name) + if not generated_column: + raise ValueError( + "Semantics description output omitted selected column: " + f"{model_name}.{column_name}" + ) + + if not self._description(generated_column): + raise ValueError( + "Semantics description output omitted description for column: " + f"{model_name}.{column_name}" + ) + + return output + + def _chunk_columns(self, chunk: dict) -> list[dict]: + models = chunk.get("mdl", {}).get("models", []) + if not models: + return [] + return models[0].get("columns", []) or [] + + def _split_chunk(self, chunk: dict) -> list[dict]: + columns = self._chunk_columns(chunk) + if len(columns) <= 1: + return [] + + split_at = max(1, len(columns) // 2) + model = chunk["mdl"]["models"][0] + return [ + { + **chunk, + "mdl": {"models": [{**model, "columns": column_chunk}]}, + } + for column_chunk in (columns[:split_at], columns[split_at:]) + if column_chunk + ] + + def _is_retryable_chunk_error(self, error: Exception) -> bool: + message = str(error) + return ( + "malformed JSON" in message + or "omitted selected model" in message + or "omitted description for model" in message + or "omitted selected column" in message + or "omitted description for column" in message + ) + + async def _generate_chunk(self, chunk: dict) -> dict: resp = await self._pipelines["semantics_description"].run(**chunk) - output = resp.get("output") + output = resp.get("output") or {} + return self._validate_chunk_output(chunk, output) + + async def _generate_chunk_with_retry_splitting(self, chunk: dict) -> list[dict]: + try: + return [await self._generate_chunk(chunk)] + except ValueError as e: + split_chunks = self._split_chunk(chunk) + if not split_chunks or not self._is_retryable_chunk_error(e): + raise + model_name = chunk.get("selected_models", [""])[0] + logger.warning( + "Retrying semantics description for model %s with smaller " + "column chunks after incomplete or malformed response.", + model_name, + ) + outputs: list[dict] = [] + for split_chunk in split_chunks: + outputs.extend( + await self._generate_chunk_with_retry_splitting(split_chunk) + ) + return outputs + + def _merge_output(self, request_id: str, output: dict): current = self[request_id] current.response = current.response or {} @@ -100,6 +212,11 @@ async def _generate_task(self, request_id: str, chunk: dict): current.response[key]["columns"].extend(output[key]["columns"]) + async def _generate_task(self, request_id: str, chunk: dict): + outputs = await self._generate_chunk_with_retry_splitting(chunk) + for output in outputs: + self._merge_output(request_id, output) + @observe(name="Generate Semantics Description") @trace_metadata async def generate(self, request: GenerateRequest, **kwargs) -> Resource: From 239edf5aaf2d1d5377ce2f055b297d7f0565ecd2 Mon Sep 17 00:00:00 2001 From: Harshitha Date: Wed, 19 Aug 2026 01:04:05 +0530 Subject: [PATCH 1068/1087] Enhance MDL semantic metadata grounding --- .../src/pipelines/indexing/db_schema.py | 6 +- .../src/pipelines/indexing/utils/helper.py | 32 ++++- wren-mdl/mdl.schema.json | 122 +++++++++++++----- 3 files changed, 118 insertions(+), 42 deletions(-) diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 50564f1666..64cec272ce 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -19,7 +19,6 @@ AsyncDocumentWriter, DocumentCleaner, MDLValidator, - clean_display_name, ) from src.pipelines.common import build_project_deploy_filter from src.pipelines.indexing.utils import helper @@ -139,10 +138,7 @@ def _convert_models_and_relationships( def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) - model_properties = { - "alias": clean_display_name(properties.get("displayName", "")), - "description": properties.get("description", ""), - } + model_properties = helper.normalize_semantic_properties(properties) comment = f"\n/* {str(model_properties)} */\n" table_name = model["name"] diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 3829324a0a..31e3785701 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -11,6 +11,18 @@ logger = logging.getLogger("wren-ai-service") +SEMANTIC_METADATA_KEYS = ( + "aliases", + "synonyms", + "businessContext", + "dataMeaning", + "semanticType", + "useCases", + "aggregationDefault", + "format", + "examples", +) + class Helper: def __init__( @@ -28,13 +40,27 @@ def __call__(self, column: Dict[str, Any], **kwargs) -> Any: return self.helper(column, **kwargs) -def _properties_comment(column: Dict[str, Any], **_) -> str: - props = column["properties"] - column_properties = { +def normalize_semantic_properties(props: Dict[str, Any]) -> Dict[str, Any]: + if not isinstance(props, dict): + props = {} + + semantic_properties = { "alias": clean_display_name(props.get("displayName", "")), "description": props.get("description", ""), } + for key in SEMANTIC_METADATA_KEYS: + value = props.get(key) + if value not in ("", None, [], {}): + semantic_properties[key] = value + + return semantic_properties + + +def _properties_comment(column: Dict[str, Any], **_) -> str: + props = column["properties"] + column_properties = normalize_semantic_properties(props) + # Add any nested columns if they exist nested = {k: v for k, v in props.items() if k.startswith("nested")} if nested: diff --git a/wren-mdl/mdl.schema.json b/wren-mdl/mdl.schema.json index b2f6d95e73..33a42fade1 100644 --- a/wren-mdl/mdl.schema.json +++ b/wren-mdl/mdl.schema.json @@ -4,6 +4,81 @@ "title": "WrenMDL Manifest Schema", "description": "A schema for WrenMDL manifest file", "$defs": { + "semanticProperties": { + "description": "Optional semantic metadata used by WrenAI retrieval and SQL generation. These fields add business context while preserving exact MDL model, column, and relationship identifiers.", + "type": "object", + "properties": { + "displayName": { + "description": "Business-friendly display name or alias.", + "type": "string" + }, + "description": { + "description": "Business meaning and purpose.", + "type": "string" + }, + "aliases": { + "description": "Alternative names users may use in natural-language questions.", + "type": "array", + "items": { + "type": "string" + } + }, + "synonyms": { + "description": "Synonyms that refer to this object.", + "type": "array", + "items": { + "type": "string" + } + }, + "businessContext": { + "description": "Business context and common use cases.", + "type": "string" + }, + "dataMeaning": { + "description": "Meaning of the data based on schema context.", + "type": "string" + }, + "semanticType": { + "description": "Semantic role such as identifier, date, measure, dimension, status, currency, quantity, cost, or revenue.", + "oneOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "type": "string" + } + } + ] + }, + "useCases": { + "description": "Common analytical uses.", + "type": "array", + "items": { + "type": "string" + } + }, + "aggregationDefault": { + "description": "Default aggregation intent when relevant.", + "type": "string" + }, + "format": { + "description": "Expected value format when relevant.", + "type": "string" + }, + "examples": { + "description": "Representative values or phrasing examples.", + "type": "array", + "items": { + "type": ["string", "number", "boolean", "null"] + } + } + }, + "additionalProperties": { + "type": ["string", "number", "boolean", "object", "array", "null"] + } + }, "column": { "type": "object", "properties": { @@ -91,11 +166,8 @@ "additionalProperties": false }, "properties": { - "description": "the customize properties of the column", - "type": "object", - "additionalProperties": { - "type": "string" - } + "description": "the customized semantic properties of the column", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "type"], @@ -236,11 +308,8 @@ } }, "properties": { - "description": "the customize properties of the model", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the model", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name"], @@ -285,11 +354,8 @@ "minLength": 1 }, "properties": { - "description": "the customize properties of the relationship", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the relationship", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "models", "joinType", "condition"] @@ -376,10 +442,7 @@ "pattern": "^\\s*(\\d+(?:\\.\\d+)?)\\s*([a-zA-Z]+)\\s*$" }, "properties": { - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "baseObject", "dimension", "measure"] @@ -403,11 +466,8 @@ "minLength": 1 }, "properties": { - "description": "the customize properties of the view", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the view", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "statement"] @@ -444,22 +504,16 @@ "minLength": 1 }, "properties": { - "description": "the customize properties of the member", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the member", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name"] } }, "properties": { - "description": "the customize properties of the enum", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the enum", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "values"] From cc55d1e0512c0ad9119c17ab157aa7a80671915b Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 11:37:01 +0000 Subject: [PATCH 1069/1087] Improve Ask schema grounding --- WRENAI_LOCAL_ASK_HANDOFF.md | 283 +++ .../generation/followup_sql_generation.py | 19 +- .../pipelines/generation/sql_correction.py | 8 +- .../pipelines/generation/sql_generation.py | 26 +- .../src/pipelines/generation/utils/sql.py | 2058 ++++++++++++++++- .../retrieval/db_schema_retrieval.py | 655 +++++- .../generation/test_sql_schema_grounding.py | 579 +++++ .../retrieval/test_db_schema_retrieval.py | 75 + .../apollo/server/resolvers/modelResolver.ts | 48 +- .../apollo/server/services/askingService.ts | 6 +- 10 files changed, 3698 insertions(+), 59 deletions(-) create mode 100644 WRENAI_LOCAL_ASK_HANDOFF.md create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md new file mode 100644 index 0000000000..e9f23efcda --- /dev/null +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -0,0 +1,283 @@ +# WrenAI Ask Grounding Handoff + +Date: 2026-08-20 + +## Current Goal + +Make WrenAI's Ask pipeline schema-first for the active org/project. Natural language should search and rank verified schema metadata, but final SQL must use only tables, views, columns, relationships, metrics, and values supported by the selected project's metadata. + +Do not fix future issues by hardcoding one question, project, table, column, or organization. Representative prompts such as `Which repair logs have the highest priority?` are regression examples only. + +## Final Runtime State + +- UI: `http://127.0.0.1:3000` +- AI service: `http://127.0.0.1:5555` +- AI health: `{"status":"ok"}` +- Active project restored after validation: `org / PCB_DB` +- Active project id: `10` +- Orders project id: `11` +- Sales duplicate: not shown in current project list; `Orders` remains canonical. + +Current projects visible through `/api/v1/projects/current`: + +- id `4`, unnamed DuckDB +- id `10`, `PCB_DB` +- id `11`, `Orders` +- id `12`, `CWPay` +- id `13`, `CW_GL` + +## What Changed Today + +### Generic Schema Grounding + +Permanent source changes are now in `D:\WrenAI\wren-ai-service`, not only `.codex-tmp`. + +Main file: + +- `wren-ai-service/src/pipelines/generation/utils/sql.py` + +Added or improved: + +- SQL identifier validation against retrieved schema. +- Semantic coverage validation so valid identifiers are not enough; the referenced table/view must also support the requested business concepts. +- Unsupported-schema result helper that returns `NO_RELEVANT_SQL` with no invented SQL. +- Deterministic schema-grounded fallback for common Ask families: + - count / grouped counts + - top-N + - highest / lowest + - latest / recent + - priority / severity + - status filters + - date/month/year filters + - revenue/sales measures + - failure counts vs defect-rate metrics + - failure type value filters +- Semantic alias support from Wren retrieved context blocks. +- More timestamp type support, including `TIMESTAMPTZ`, which fixed the live `latest repair logs` failure. +- Normalization of dialect issues such as `TOP n`, joined `DESCLIMIT`, and order-by aliases. +- Logs for generated SQL validation, deterministic fallback SQL, fallback validation, selected table, verified columns, and metric intent. + +### Retrieval Improvements + +Main file: + +- `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` + +Added or improved: + +- Project-scoped retrieval filters retained and tested. +- Query expansion for business concepts such as repair, failure, revenue, order, material, status, priority, latest, and date. +- Ranking uses table names, column names, descriptions/comments, semantic context, and generic-table deboosting. +- Logs for: + - selected project id + - retrieved candidate tables and scores + - selected schema objects and columns + +### Generation Pipeline Wiring + +Files: + +- `wren-ai-service/src/pipelines/generation/sql_generation.py` +- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` +- `wren-ai-service/src/pipelines/generation/sql_correction.py` + +Changes: + +- Passed the user query into post-processing as `fallback_query`. +- Added pre-LLM unsupported-schema checks where retrieved schema clearly cannot cover requested concepts. +- Ensured SQL correction still uses the same schema-first validation and fallback logic. +- Strengthened correction instructions so invalid or hallucinated identifiers are not preserved. + +### UI / Project Cleanup From This Workstream + +Files still dirty from the related UI/runtime fixes: + +- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` +- `wren-ui/src/apollo/server/services/askingService.ts` + +Relevant behavior: + +- Previous `results` crash handling is preserved. +- Unsupported-schema failures now avoid showing invented SQL as something to fix. +- Sales/Orders cleanup remains in place: UI project list shows `Orders`, not duplicate `Sales`. + +## Live Validation Done + +All live checks were run through the UI GraphQL Ask path after restarting the AI service. + +### PCB_DB + +Active project: `PCB_DB`, id `10`. + +Passed: + +- `Which repair logs have the highest priority?` + - Table: `dbo_repair_logs` + - Uses verified `priority` + - Orders by generic priority ranking expression +- `Show all critical-priority repairs` + - Table: `dbo_repair_logs` + - Filter: `priority = 'critical'` +- `Show repairs by status` + - Table: `dbo_repair_logs` + - Group: `status` + - Metric: `COUNT(id)` +- `Show latest repair logs` + - Table: `dbo_repair_logs` + - Order: `created_at DESC` + - This was the live regression fixed by adding timestamp type coverage. +- `Show the number of failures by material` + - Uses verified material/failure fields from PCB_DB. +- `Show top 5 board models with the most failures` + - Table: `dbo_repair_logs` + - Metric: `COUNT(failure_code)` + - Did not use `defect_rate`. +- `Show units with JTAG as the failure type` + - Table: `dbo_report_failures` + - Filter: `failure_type = 'JTAG'` +- Extra check: + - `Show all repairs with a critical priority and an in-progress status.` + - Table: `dbo_repair_logs` + - Filters: `status = 'in-progress'` and `priority = 'critical'` + +### Orders + +Temporarily switched active project to `Orders`, id `11`, then restored PCB_DB. + +Passed: + +- `Show top 10 orders from July` + - Uses Orders table/date fields. +- `Show number of orders by customer` + - Groups by customer. + - Counts distinct order numbers. +- `Show revenue by year` + - Uses verified sales/revenue value and invoice date fields. +- Unsupported check: `Which repair logs have the highest priority?` + - Returned `NO_RELEVANT_SQL`. + - No SQL candidate. + - Message clearly said the active project does not contain verified `repair` and `priority/severity` fields. + +## Checks Run + +Passed: + +```powershell +git diff --check -- wren-ai-service/src/pipelines/generation/utils/sql.py ` + wren-ai-service/src/pipelines/generation/sql_generation.py ` + wren-ai-service/src/pipelines/generation/followup_sql_generation.py ` + wren-ai-service/src/pipelines/generation/sql_correction.py ` + wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py ` + wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py ` + wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +``` + +Passed: + +```powershell +cd D:\WrenAI\wren-ai-service +.\venv\Scripts\python.exe -m compileall -q src\pipelines\generation src\pipelines\retrieval ` + tests\pytest\pipelines\generation\test_sql_schema_grounding.py ` + tests\pytest\pipelines\retrieval\test_db_schema_retrieval.py +``` + +Could not run pytest in the service venv: + +```text +D:\WrenAI\wren-ai-service\venv\Scripts\python.exe: No module named pytest +``` + +## Tests Added + +Main test file: + +- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` + +Coverage added for: + +- Unsupported schema clears invalid SQL. +- Generic table rejection for unsupported business concepts. +- Repair priority ordering. +- Critical-priority repair filters. +- Repairs by status. +- Latest repair logs with `TIMESTAMPTZ`. +- Semantic alias column support, for example using real verified `Urgency` when semantic context says it means priority/severity. +- Failure by material / technician with verified columns. +- JTAG failure type filters. +- Board models with most failures uses count, not defect rate. +- Highest defect rate uses rate metric. +- Repairs by technician requires one schema object or relationship coverage. + +Retrieval test file: + +- `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` + +Coverage added for: + +- Project filter conditions. +- Query expansion. +- Table ranking by query and schema text. +- Project-scoped schema retrieval behavior. + +## Restart Commands Used + +Restart AI service only: + +```powershell +$taskName = 'WrenAI 04 AI Service' +$listenerProcessIds = Get-NetTCPConnection -LocalPort 5555 -State Listen -ErrorAction SilentlyContinue | + Select-Object -ExpandProperty OwningProcess -Unique +Stop-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue +foreach ($listenerProcessId in $listenerProcessIds) { + if ($listenerProcessId) { + Stop-Process -Id $listenerProcessId -Force -ErrorAction SilentlyContinue + } +} +Start-ScheduledTask -TaskName $taskName +``` + +Health check: + +```powershell +Invoke-WebRequest -UseBasicParsing http://127.0.0.1:5555/health +``` + +Project switch endpoints used for validation: + +```powershell +Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/11/select +Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/10/select +Invoke-WebRequest -UseBasicParsing http://127.0.0.1:3000/api/v1/projects/current +``` + +## Current Dirty Files To Review + +Relevant tracked files: + +- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` +- `wren-ai-service/src/pipelines/generation/sql_correction.py` +- `wren-ai-service/src/pipelines/generation/sql_generation.py` +- `wren-ai-service/src/pipelines/generation/utils/sql.py` +- `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` +- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` +- `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` +- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` +- `wren-ui/src/apollo/server/services/askingService.ts` + +There are also many local untracked runtime/data artifacts in the repository. Do not clean or delete them casually. + +## Important Caveats + +- Runtime source code should remain generic. Do not add checks for exact prompts such as `Which repair logs have the highest priority?`. +- Tests may use representative table and prompt names; production code must not. +- Retrieval context currently uses metadata/descriptions and some semantic context. It does not appear to carry robust sample-value lists. Status casing/value handling works for tested prompts, but richer value-aware matching would improve future accuracy. +- `enable_column_pruning` was not the focus of today's final validation. +- Full pytest suite still needs an environment with `pytest` installed. + +## Recommended Next Steps + +1. Install or enable pytest in `wren-ai-service\venv`, then run focused tests. +2. Review the large `utils/sql.py` diff carefully; consider extracting fallback/grounding helpers into smaller modules after behavior is stable. +3. Add sample-value metadata to retrieval context if available, then make value matching use that metadata instead of only text normalization. +4. Run a broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when their data sources are available. +5. Commit the source changes after review, excluding local runtime/data artifacts. 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 5b5542a5ce..2da693f867 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -21,6 +21,7 @@ get_json_field_instructions, get_metric_instructions, get_sql_generation_system_prompt, + unsupported_schema_generation_result, ) from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge @@ -150,7 +151,8 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str], + query: str | None = None, + documents: list[str] | None = None, project_id: str | None = None, mdl_hash: str | None = None, validation_contexts: list[str] | None = None, @@ -162,6 +164,7 @@ async def post_process( project_id=project_id, mdl_hash=mdl_hash, contexts=validation_contexts or documents, + fallback_query=query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -227,6 +230,18 @@ async def run( ) else: metadata = {} + data_source = metadata.get("data_source", "local_file") + + unsupported_result = unsupported_schema_generation_result( + query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "Follow-up SQL generation skipped before LLM because selected schema does not cover requested concepts." + ) + return {"post_process": unsupported_result} return await self._pipe.execute( ["post_process"], @@ -245,7 +260,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "sql_knowledge": sql_knowledge, "validation_contexts": validation_contexts, **self._components, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index ec78f998fb..13358f7513 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -44,6 +44,10 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. 9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA. If the unsupported part is needed to answer the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql instead of substituting non-schema identifiers. 10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. +11. For grouped queries, repair SQL Server errors about ORDER BY columns not appearing in GROUP BY by ordering with selected grouping columns or selected aggregate aliases, or by adding the exact ordering key to both SELECT and GROUP BY when that key is declared in DATABASE SCHEMA. +12. Do not preserve generic log, file, JSON, payload, text, or app-metric scans when DATABASE SCHEMA contains exact modeled business columns for the user's requested entity, measure, status, date, or dimension. +13. If the failed SQL invented component fields for a metric that exists directly in DATABASE SCHEMA, replace the calculation with the exact declared metric column. +14. For sales or revenue questions, avoid tariff, duty, customs, import, refund, or claim datasets unless the USER QUESTION explicitly asks for those domains. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -163,7 +167,8 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str], + query: str | None = None, + documents: list[str] | None = None, project_id: str | None = None, mdl_hash: str | None = None, validation_contexts: list[str] | None = None, @@ -175,6 +180,7 @@ async def post_process( project_id=project_id, mdl_hash=mdl_hash, contexts=validation_contexts or documents, + fallback_query=query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 5bdeeb8827..9e93154c88 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -20,6 +20,7 @@ get_json_field_instructions, get_metric_instructions, get_sql_generation_system_prompt, + unsupported_schema_generation_result, ) from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge @@ -141,7 +142,8 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: list[str], + query: str | None = None, + documents: list[str] | None = None, project_id: str | None = None, mdl_hash: str | None = None, validation_contexts: list[str] | None = None, @@ -154,6 +156,7 @@ async def post_process( project_id=project_id, mdl_hash=mdl_hash, contexts=validation_contexts or documents, + fallback_query=query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -212,7 +215,11 @@ async def run( sql_knowledge: SqlKnowledge | None = None, validation_contexts: list[str] | None = None, ): - logger.info("SQL Generation pipeline is running...") + logger.info( + "SQL Generation pipeline is running for project_id=%s mdl_hash=%s", + project_id or "", + mdl_hash or "", + ) if project_id or use_dry_plan: metadata = await retrieve_metadata( @@ -220,6 +227,19 @@ async def run( ) else: metadata = {} + data_source = metadata.get("data_source", "local_file") + + unsupported_result = unsupported_schema_generation_result( + query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "SQL generation skipped before LLM because selected schema does not cover requested concepts: %s", + unsupported_result["invalid_generation_result"]["error"], + ) + return {"post_process": unsupported_result} return await self._pipe.execute( ["post_process"], @@ -237,7 +257,7 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, "validation_contexts": validation_contexts, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 25becffddf..bcd07c0308 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,5 +1,6 @@ import logging import re +from datetime import datetime, timezone from typing import Any, Dict, List import aiohttp @@ -47,6 +48,21 @@ rf"(?:\s+(?:AS\s+)?(?P{_IDENTIFIER_TOKEN}))?", re.IGNORECASE, ) +_TSQL_TOP_LIMIT = re.compile( + r"(?is)^(\s*)SELECT\s+TOP\s*\(?\s*(\d+)\s*\)?\s+(?!PERCENT\b)(.+?)\s*;?\s*$" +) +_TO_DATE_SIMPLE = re.compile( + r"(?is)\bTO_DATE\s*\(\s*" + r"(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*" + r"(?:\s*\.\s*(?:\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*))?)" + r"\s*,\s*'[^']+'\s*\)" +) +_ORDER_BY_ALIAS_ITEM = re.compile( + r"(?is)^(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)" + r"(?P\s+(?:ASC|DESC))?$" +) +_JOINED_DESC_LIMIT = re.compile(r"(?i)\bDESC\s*LIMIT\b|DESCLIMIT") +_JOINED_ASC_LIMIT = re.compile(r"(?i)\bASC\s*LIMIT\b|ASCLIMIT") _CTE_REFERENCE = re.compile( rf"(?:\bWITH|,)\s+(?P{_IDENTIFIER_TOKEN})\s+AS\s*\(", re.IGNORECASE, @@ -76,11 +92,18 @@ "UNION", "WHERE", } +_UNQUALIFIED_QUOTED_IDENTIFIER = re.compile(r'(?(?:[^"]|"")*)"') +_UNQUALIFIED_BARE_IDENTIFIER = re.compile( + r"(?[A-Za-z_][A-Za-z0-9_$]*)\b(?!\s*\.)" +) +_SINGLE_QUOTED_LITERAL = re.compile(r"'(?:''|[^'])*'") _SQL_RESERVED_WORDS = { "ALL", "ALTER", "AND", + "ASC", "AS", + "BETWEEN", "BY", "CASE", "CAST", @@ -101,6 +124,7 @@ "GROUP", "HAVING", "IN", + "ISNULL", "INNER", "INSERT", "INTERSECT", @@ -110,6 +134,7 @@ "LIKE", "LIMIT", "NATURAL", + "NO", "NOT", "NULL", "OFFSET", @@ -125,11 +150,242 @@ "TRUE", "UNION", "UPDATE", + "VALUES", "WHEN", "WHERE", "WINDOW", "WITH", } +_SQL_FUNCTION_WORDS = { + "ABS", + "AVG", + "CAST", + "CEIL", + "CEILING", + "COALESCE", + "CONCAT", + "COUNT", + "COUNT_BIG", + "DATE_TRUNC", + "DAY", + "EXTRACT", + "FLOOR", + "LOWER", + "MAX", + "MIN", + "MONTH", + "NULLIF", + "ROUND", + "SUM", + "TRIM", + "UPPER", + "YEAR", +} +_SQL_TYPE_WORDS = { + "BIGINT", + "BOOLEAN", + "CHAR", + "DATE", + "DATETIME", + "DECIMAL", + "DOUBLE", + "FLOAT", + "FLOAT4", + "FLOAT8", + "INT", + "INT2", + "INT4", + "INT8", + "INTEGER", + "NUMERIC", + "REAL", + "SMALLINT", + "TEXT", + "TIME", + "TIMESTAMP", + "VARCHAR", +} +_DATE_PART_WORDS = { + "DAY", + "DOW", + "DOY", + "HOUR", + "MICROSECOND", + "MILLISECOND", + "MINUTE", + "MONTH", + "QUARTER", + "SECOND", + "WEEK", + "YEAR", +} +_FALLBACK_TOKEN = re.compile(r"[a-z0-9]+") +_FALLBACK_STOPWORDS = { + "a", + "an", + "and", + "are", + "by", + "from", + "in", + "of", + "show", + "that", + "the", + "to", + "with", +} +_MONTH_NAME_TO_NUMBER = { + "january": 1, + "february": 2, + "march": 3, + "april": 4, + "may": 5, + "june": 6, + "july": 7, + "august": 8, + "september": 9, + "october": 10, + "november": 11, + "december": 12, +} + + +def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: + aliases = { + "bu": {"business", "unit"}, + "cust": {"customer"}, + "customers": {"customer"}, + "critical": {"priority", "severity"}, + "boards": {"board"}, + "high": {"priority", "severity"}, + "highest": {"top"}, + "logs": {"log", "record"}, + "log": {"record"}, + "lows": {"low"}, + "lowest": {"bottom"}, + "models": {"model"}, + "inv": {"invoice"}, + "invoices": {"invoice"}, + "ord": {"order"}, + "orders": {"order"}, + "qty": {"quantity"}, + "num": {"number"}, + "no": {"number"}, + "numbers": {"number"}, + "prod": {"product"}, + "products": {"product"}, + "priorities": {"priority", "severity"}, + "priority": {"severity"}, + "recent": {"latest"}, + "records": {"record"}, + "rep": {"representative", "salesperson"}, + "repairs": {"repair"}, + "salesperson": {"sales", "person"}, + "severity": {"priority"}, + "supplier": {"vendor"}, + "suppliers": {"supplier", "vendor"}, + "tech": {"technician"}, + "technician": {"tech"}, + "vendor": {"supplier"}, + "vendors": {"supplier", "vendor"}, + "failed": {"failure"}, + "failures": {"failure"}, + "defects": {"defect"}, + "types": {"type"}, + "units": {"unit", "serial"}, + "urgency": {"priority", "severity"}, + "locations": {"location"}, + "materials": {"material"}, + "missing": {"blank", "empty", "null"}, + } + expanded = set(tokens) + for token in list(tokens): + expanded.update(aliases.get(token, set())) + if {"business", "unit"}.issubset(expanded): + expanded.add("bu") + if "customer" in expanded and "number" in expanded: + expanded.update({"cust", "id", "no"}) + if "order" in expanded and "number" in expanded: + expanded.update({"ord", "id", "no"}) + return expanded + + +def normalize_wren_sql_dialect(sql: str) -> str: + if not sql: + return sql + + sql = _TO_DATE_SIMPLE.sub( + lambda match: f"CAST({match.group('expr')} AS DATE)", + sql, + ) + sql = _JOINED_DESC_LIMIT.sub("DESC LIMIT", sql) + sql = _JOINED_ASC_LIMIT.sub("ASC LIMIT", sql) + sql = _replace_order_by_aliases_with_select_expressions(sql) + sql = _JOINED_DESC_LIMIT.sub("DESC LIMIT", sql) + sql = _JOINED_ASC_LIMIT.sub("ASC LIMIT", sql) + + if re.search(r"(?i)\bLIMIT\s+\d+\b", sql): + return sql + + match = _TSQL_TOP_LIMIT.match(sql) + if not match: + return sql + + leading_space, limit, select_body = match.groups() + if re.search(r"(?i)\bWITH\s+TIES\b", select_body): + return sql + + return f"{leading_space}SELECT {select_body.strip()}\nLIMIT {limit}" + + +def _replace_order_by_aliases_with_select_expressions(sql: str) -> str: + select_clause = _extract_select_clause(sql) + order_by_clause = _extract_clause(sql, "ORDER BY", ("LIMIT", "OFFSET")) + if not select_clause or not order_by_clause: + return sql + + alias_to_expression = {} + for select_item in _split_sql_tokens(select_clause): + expression, alias = _split_select_expression_alias(select_item) + if alias and expression: + alias_to_expression[alias] = expression + + if not alias_to_expression: + return sql + + replaced_any = False + rewritten_order_items = [] + for order_item in _split_sql_tokens(order_by_clause): + match = _ORDER_BY_ALIAS_ITEM.match(order_item.strip()) + if not match: + rewritten_order_items.append(order_item) + continue + + alias = _unquote_identifier(match.group("identifier")) + expression = alias_to_expression.get(alias) + if not expression: + rewritten_order_items.append(order_item) + continue + + rewritten_order_items.append(f"{expression}{match.group('suffix') or ''}") + replaced_any = True + + if not replaced_any: + return sql + + order_by_match = re.search( + r"(?is)\bORDER\s+BY\b\s+.*?(?=\b(?:LIMIT|OFFSET)\b|$)", + sql, + ) + if not order_by_match: + return sql + + suffix = sql[order_by_match.end() :] + if suffix and not suffix[0].isspace(): + suffix = "\n" + suffix + replacement = "ORDER BY " + ", ".join(rewritten_order_items) + return sql[: order_by_match.start()] + replacement + suffix def _unquote_identifier(identifier: str) -> str: @@ -439,6 +695,132 @@ def _extract_schema_index(contexts: list[Any] | None) -> dict[str, set[str] | No return schema_index +def _semantic_tokens_from_value(value: Any) -> set[str]: + tokens: set[str] = set() + if isinstance(value, str): + tokens.update(_fallback_tokens(value)) + elif isinstance(value, dict): + for nested_value in value.values(): + tokens.update(_semantic_tokens_from_value(nested_value)) + elif isinstance(value, list): + for nested_value in value: + tokens.update(_semantic_tokens_from_value(nested_value)) + return tokens + + +def _extract_semantic_context_payload(context: str) -> dict[str, Any]: + start_marker = "WREN RETRIEVED SEMANTIC CONTEXT" + end_marker = "WREN SQL IDENTIFIER CONTRACT" + start_index = context.upper().find(start_marker) + end_index = context.upper().find(end_marker, start_index) + if start_index < 0 or end_index < 0: + return {} + payload_text = context[start_index + len(start_marker) : end_index].strip() + try: + payload = orjson.loads(payload_text) + except orjson.JSONDecodeError: + return {} + return payload if isinstance(payload, dict) else {} + + +def _extract_semantic_tokens_by_column( + context: str, +) -> tuple[set[str], dict[str, set[str]]]: + payload = _extract_semantic_context_payload(context) + if not payload: + return set(), {} + + table_tokens = _semantic_tokens_from_value( + payload.get("semantic_context_not_sql_identifiers") + ) + table_tokens.update(_semantic_tokens_from_value(payload.get("object_type"))) + + column_tokens: dict[str, set[str]] = {} + for column in payload.get("columns", []) or []: + if not isinstance(column, dict): + continue + column_name = column.get("sql_column_name_use_exactly") + if not isinstance(column_name, str) or not column_name: + continue + tokens = _semantic_tokens_from_value( + column.get("semantic_context_not_sql_identifier") + ) + if tokens: + column_tokens[column_name] = tokens + + return table_tokens, column_tokens + + +def _extract_schema_details( + contexts: list[Any] | None, +) -> dict[str, list[dict[str, str]]]: + if not contexts: + return {} + + schema_details: dict[str, list[dict[str, str]]] = {} + + for context in _iter_context_texts(contexts): + relation_match = _DDL_RELATION.search(context) + if not relation_match: + continue + + relation_name = _unquote_identifier(relation_match.group("name")) + table_semantic_tokens, column_semantic_tokens = ( + _extract_semantic_tokens_by_column(context) + ) + column_block_match = re.search( + r"\bCREATE\s+TABLE\b[^(]*\((?P.*)\)\s*;?", + context, + re.IGNORECASE | re.DOTALL, + ) + if not column_block_match: + continue + + columns = [] + for raw_column in _split_sql_tokens(column_block_match.group("columns")): + line = "\n".join( + line.strip() + for line in raw_column.splitlines() + if line.strip() + and not line.strip().startswith("--") + and not line.strip().startswith("/*") + ).strip() + if not line: + continue + line = line.split("--", 1)[0].strip() + if not line or line.upper().startswith(("FOREIGN KEY", "PRIMARY KEY")): + continue + + if line.startswith('"'): + end = line.find('"', 1) + while end != -1 and end + 1 < len(line) and line[end + 1] == '"': + end = line.find('"', end + 2) + if end <= 0: + continue + name = _unquote_identifier(line[: end + 1]) + remainder = line[end + 1 :].strip() + else: + parts = line.split(None, 1) + if not parts: + continue + name = _unquote_identifier(parts[0]) + remainder = parts[1].strip() if len(parts) > 1 else "" + + data_type = remainder.split(None, 1)[0].upper() if remainder else "" + columns.append( + { + "name": name, + "data_type": data_type, + "semantic_tokens": column_semantic_tokens.get(name, set()), + "_table_semantic_tokens": table_semantic_tokens, + } + ) + + schema_details[relation_name] = columns + + return schema_details + + def _is_identifier_boundary(char: str | None) -> bool: return char is None or not (char.isalnum() or char in {"_", "$", '"'}) @@ -600,7 +982,9 @@ def _extract_sql_grounding(sql: str) -> dict[str, Any]: alias_to_relation = {} for match in _RELATION_REFERENCE.finditer(sql): - relation = _normalize_identifier(match.group("name")) + if _is_extract_from_clause(sql, match.start()): + continue + relation = _unquote_identifier(match.group("name")) if relation.upper() in {"UNNEST", "LATERAL"}: continue alias = match.group("alias") @@ -627,6 +1011,168 @@ def _extract_sql_grounding(sql: str) -> dict[str, Any]: } +def _is_extract_from_clause(sql: str, from_start: int) -> bool: + prefix = sql[:from_start] + last_open = prefix.rfind("(") + if last_open == -1 or prefix.rfind(")") > last_open: + return False + + before_open = prefix[:last_open].rstrip() + return before_open.upper().endswith("EXTRACT") + + +def _strip_string_literals(sql: str) -> str: + return _SINGLE_QUOTED_LITERAL.sub("''", sql) + + +def _extract_clause(sql: str, clause: str, end_clauses: tuple[str, ...]) -> str: + end_pattern = "|".join(re.escape(end_clause) for end_clause in end_clauses) + pattern = re.compile( + rf"(?is)\b{re.escape(clause)}\b\s+(?P.*?)(?=\b(?:{end_pattern})\b|$)" + ) + match = pattern.search(sql) + return match.group("body").strip() if match else "" + + +def _extract_select_clause(sql: str) -> str: + match = re.search(r"(?is)\bSELECT\b\s+(?P.*?)(?=\bFROM\b)", sql) + return match.group("body").strip() if match else "" + + +def _split_select_expression_alias(expression: str) -> tuple[str, str | None]: + as_match = re.search( + r"(?is)\s+AS\s+(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)\s*$", + expression, + ) + if as_match: + return expression[: as_match.start()].strip(), _unquote_identifier( + as_match.group("alias") + ) + + return expression, None + + +def _extract_output_aliases(select_clause: str) -> set[str]: + aliases = set() + for expression in _split_sql_tokens(select_clause): + _, alias = _split_select_expression_alias(expression) + if alias: + aliases.add(alias) + return aliases + + +def _iter_unqualified_identifier_candidates(expression: str): + stripped = _strip_string_literals(expression) + + for match in _UNQUALIFIED_QUOTED_IDENTIFIER.finditer(stripped): + yield _unquote_identifier(f'"{match.group("name")}"') + + without_quoted = _UNQUALIFIED_QUOTED_IDENTIFIER.sub(" ", stripped) + for match in _UNQUALIFIED_BARE_IDENTIFIER.finditer(without_quoted): + name = match.group("name") + following = without_quoted[match.end() :].lstrip() + if following.startswith("("): + continue + yield name + + +def _sql_mentions_identifier(sql: str, identifier: str) -> bool: + stripped = _strip_string_literals(sql) + quoted_identifier = re.escape(_quote_identifier(identifier)) + bracket_identifier = re.escape(f"[{identifier}]") + bare_identifier = re.escape(identifier) + return bool( + re.search(rf'(? str | None: + real_relations = [ + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + ] + if grounding["cte_names"]: + return None + + unique_real_relations = list(dict.fromkeys(real_relations)) + if len(unique_real_relations) != 1: + return None + + relation = unique_real_relations[0] + valid_columns = schema_index.get(relation) + if valid_columns is None: + return None + + select_clause = _extract_select_clause(sql) + output_aliases = _extract_output_aliases(select_clause) + ignored_identifiers = ( + set(schema_index) + | set(grounding["alias_to_relation"]) + | set(grounding["cte_names"]) + | output_aliases + ) + + clause_expressions = [] + for expression in _split_sql_tokens(select_clause): + expression, _ = _split_select_expression_alias(expression) + clause_expressions.append(expression) + clause_expressions.extend( + filter( + None, + [ + _extract_clause( + sql, + "WHERE", + ("GROUP BY", "HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause( + sql, + "GROUP BY", + ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause( + sql, + "HAVING", + ("ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause(sql, "ORDER BY", ("LIMIT", "OFFSET")), + ], + ) + ) + + invalid_columns = set() + for expression in clause_expressions: + for identifier in _iter_unqualified_identifier_candidates(expression): + upper_identifier = identifier.upper() + if ( + upper_identifier in _SQL_RESERVED_WORDS + or upper_identifier in _SQL_FUNCTION_WORDS + or upper_identifier in _SQL_TYPE_WORDS + or upper_identifier in _DATE_PART_WORDS + or identifier in ignored_identifiers + or identifier in valid_columns + ): + continue + invalid_columns.add(identifier) + + if not invalid_columns: + return None + + return ( + "Schema grounding failed. The SQL references unqualified columns that " + f"are not present in verified table or view {relation}: " + f"{', '.join(sorted(invalid_columns))}. Use only verified columns: " + f"{', '.join(sorted(valid_columns))}." + ) + + def validate_sql_against_contexts( sql: str, contexts: list[Any] | None = None, @@ -677,6 +1223,14 @@ def validate_sql_against_contexts( f"{', '.join(sorted(valid_columns))}." ) + unqualified_column_error = _validate_unqualified_columns_for_single_relation( + sql, + schema_index, + grounding, + ) + if unqualified_column_error: + return unqualified_column_error + return None @@ -729,6 +1283,126 @@ def _extract_generation_sql(generation_result: str | None) -> str | None: return text if _SQL_START.search(text) else None +def validate_sql_semantic_coverage( + sql: str, + query: str | None, + contexts: list[Any] | None = None, +) -> str | None: + if not sql or not query: + return None + + raw_query_tokens = _fallback_tokens(query) + query_tokens = _expanded_fallback_query_tokens(query) + concepts = _requested_business_concepts(raw_query_tokens) + if not concepts: + return None + + schema_details = _extract_schema_details(contexts) + if not schema_details: + return None + + grounding = _extract_sql_grounding(sql) + referenced_relations = { + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + } + if not referenced_relations: + return None + + schema_tokens = set() + for relation in referenced_relations: + columns = schema_details.get(relation) + if columns is not None: + schema_tokens.update(_schema_tokens_for_table(relation, columns)) + + if not schema_tokens: + return None + + missing_concepts = [ + label for label, concept_tokens in concepts if not schema_tokens & concept_tokens + ] + if not missing_concepts: + if _is_failure_count_intent(raw_query_tokens, query_tokens): + if not re.search(r"(?is)\bCOUNT\s*\(", sql): + return ( + "Schema grounding failed. The question asks for a count of " + "failure records, but the generated SQL does not compute a " + "COUNT aggregate. Use a verified failure-record column/table " + "and group by the requested dimension, or return no SQL if " + "the active project does not contain it." + ) + for relation in referenced_relations: + for column in schema_details.get(relation, []): + if _is_rate_like_column(column) and _sql_mentions_identifier( + sql, column["name"] + ): + return ( + "Schema grounding failed. The question asks for a " + "count of failure records, but the generated SQL uses " + f"rate-like column {column['name']}. Use COUNT over a " + "verified failure occurrence field instead, or return " + "no SQL if the active project does not contain one." + ) + return None + + return ( + "Schema grounding failed. The generated SQL uses verified identifiers, " + "but the selected table or view does not contain verified fields for the " + f"requested business concept(s): {', '.join(missing_concepts)}. Use only " + "schema objects whose table or column names explicitly support those " + "concepts, or return no SQL if the active project does not contain them." + ) + + +def unsupported_schema_message( + query: str | None, + contexts: list[Any] | None = None, +) -> str | None: + if not query: + return None + query_tokens = _fallback_tokens(query) + concepts = _requested_business_concepts(query_tokens) + if not concepts: + return None + schema_details = _extract_schema_details(contexts) + if not schema_details: + return None + if any( + _table_covers_requested_concepts(table_name, columns, query_tokens) + for table_name, columns in schema_details.items() + ): + return None + concept_labels = ", ".join(label for label, _ in concepts) + return ( + "No retrieved table or view in the active project contains verified " + "fields for all requested business concept(s): " + f"{concept_labels}. Select a project with those fields or ask a question " + "supported by the selected project's schema." + ) + + +def unsupported_schema_generation_result( + query: str | None, + contexts: list[Any] | None = None, + data_source: str = "", +) -> dict[str, Any] | None: + message = unsupported_schema_message(query, contexts=contexts) + if not message: + return None + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": message, + "correlation_id": "", + "data_source": data_source, + }, + } + + def normalize_sql_with_schema_identifiers( sql: str, contexts: list[Any] | None = None, @@ -745,6 +1419,1219 @@ def normalize_sql_with_schema_identifiers( return sql +def _fallback_tokens(value: Any) -> set[str]: + if value is None: + return set() + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) + tokens = { + token + for token in _FALLBACK_TOKEN.findall(text.lower()) + if token not in _FALLBACK_STOPWORDS + } + return _expand_fallback_token_aliases(tokens) + + +def _column_business_tokens(column: dict[str, Any]) -> set[str]: + tokens = _fallback_tokens(column["name"]) + tokens.update(column.get("semantic_tokens") or set()) + return tokens + + +def _table_business_tokens( + table_name: str, + columns: list[dict[str, Any]], +) -> set[str]: + tokens = _fallback_tokens(table_name) + for column in columns: + tokens.update(column.get("_table_semantic_tokens") or set()) + return tokens + + +def _expanded_fallback_query_tokens(query: str) -> set[str]: + tokens = _fallback_tokens(query) + if tokens & {"revenue", "sale", "sales", "trend", "trends"}: + tokens.update({"amount", "date", "intake", "revenue", "sales", "value"}) + if tokens & {"order", "orders"}: + tokens.update({"amount", "customer", "date", "ord", "order", "value"}) + if tokens & {"invoice", "invoices"}: + tokens.update({"amount", "currency", "date", "invoice", "supplier"}) + if tokens & {"batch", "batches"}: + tokens.update({"batch", "board", "defect", "inspection", "rate", "supplier"}) + if tokens & {"repair", "repairs"}: + tokens.update({"date", "failure", "log", "priority", "progress", "repair", "status"}) + if tokens & {"failure", "failures", "defect", "defects"}: + tokens.update({"code", "defect", "failure", "severity", "status", "type"}) + if tokens & {"material", "materials"}: + tokens.update({"item", "material", "part"}) + if tokens & {"location", "locations"}: + tokens.update({"area", "location", "site"}) + if tokens & {"month", "monthly", "july"}: + tokens.update({"date", "day", "month", "time", "year"}) + elif tokens & {"latest", "trend", "trends", "year"}: + tokens.update({"date", "day", "time", "year"}) + if "business" in tokens and "unit" in tokens: + tokens.update({"account", "bu", "business", "company", "division", "unit"}) + return tokens + + +def _is_numeric_type(data_type: str) -> bool: + return data_type.upper() in { + "BIGINT", + "DECIMAL", + "DOUBLE", + "FLOAT", + "FLOAT4", + "FLOAT8", + "INT", + "INT2", + "INT4", + "INT8", + "INTEGER", + "NUMERIC", + "REAL", + "SMALLINT", + } + + +def _is_date_type(data_type: str) -> bool: + return data_type.upper() in { + "DATE", + "DATETIME", + "DATETIME2", + "SMALLDATETIME", + "TIME", + "TIMESTAMP", + "TIMESTAMPTZ", + "TIMESTAMP_LTZ", + "TIMESTAMP_NTZ", + "TIMESTAMP_TZ", + } + + +_RATE_METRIC_TOKENS = {"rate", "ratio", "percent", "percentage"} +_COUNT_METRIC_TOKENS = {"count", "many", "most", "number", "total"} +_PRIORITY_VALUE_ALIASES = { + "urgent": "urgent", + "critical": "critical", + "high": "high", + "medium": "medium", + "normal": "normal", + "low": "low", +} +_PRIORITY_ORDER = [ + ("critical", 6), + ("urgent", 6), + ("blocker", 6), + ("high", 5), + ("major", 5), + ("medium", 4), + ("normal", 4), + ("minor", 3), + ("low", 2), +] + + +def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _RATE_METRIC_TOKENS) + + +def _is_failure_count_intent( + raw_query_tokens: set[str], + query_tokens: set[str], +) -> bool: + return ( + bool(raw_query_tokens & {"failure", "failed", "defect"}) + and "failure" in query_tokens + and not _is_rate_metric_intent(raw_query_tokens) + and ( + bool(raw_query_tokens & _COUNT_METRIC_TOKENS) + or bool(raw_query_tokens & {"top", "highest", "lowest", "bottom"}) + ) + ) + + +def _has_board_model_intent(query_tokens: set[str]) -> bool: + return {"board", "model"}.issubset(query_tokens) + + +def _is_rate_like_column(column: dict[str, str]) -> bool: + return bool(_fallback_tokens(column["name"]) & (_RATE_METRIC_TOKENS | {"score"})) + + +def _quote_joined(identifiers: list[str]) -> str: + return ", ".join(_quote_identifier(identifier) for identifier in identifiers) + + +def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[str]]]: + concepts: list[tuple[str, set[str]]] = [] + specs = [ + ( + "failure/defect", + {"failure", "failed", "defect"}, + {"failure", "failed", "defect"}, + ), + ("repair", {"repair"}, {"repair"}), + ("material", {"material"}, {"material", "part"}), + ("location", {"location"}, {"location", "site", "area"}), + ("customer", {"customer"}, {"customer", "cust"}), + ("supplier/vendor", {"supplier", "vendor"}, {"supplier", "vendor"}), + ("technician", {"technician", "tech"}, {"technician", "tech"}), + ("product", {"product"}, {"product", "prod", "item", "material"}), + ( + "priority/severity", + {"critical", "priority", "severity"}, + {"priority", "severity"}, + ), + ("status", {"status"}, {"status"}), + ("order", {"order"}, {"order", "ord"}), + ] + for label, triggers, schema_tokens in specs: + if query_tokens & triggers: + concepts.append((label, schema_tokens)) + if {"board", "model"}.issubset(query_tokens): + concepts.append(("board model", {"board", "model"})) + if {"business", "unit"}.issubset(query_tokens): + concepts.append(("business unit", {"business", "unit", "bu", "division"})) + return concepts + + +def _schema_tokens_for_table(table_name: str, columns: list[dict[str, str]]) -> set[str]: + tokens = _table_business_tokens(table_name, columns) + for column in columns: + tokens.update(_column_business_tokens(column)) + return tokens + + +def _table_covers_requested_concepts( + table_name: str, + columns: list[dict[str, str]], + concept_tokens: set[str], +) -> bool: + concepts = _requested_business_concepts(concept_tokens) + if not concepts: + return True + schema_tokens = _schema_tokens_for_table(table_name, columns) + return all(schema_tokens & concept_tokens for _, concept_tokens in concepts) + + +def _choose_fallback_table( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], + concept_tokens: set[str] | None = None, +) -> tuple[str, list[dict[str, str]]] | None: + concept_tokens = concept_tokens or query_tokens + rate_metric_intent = _is_rate_metric_intent(concept_tokens) + failure_count_intent = _is_failure_count_intent(concept_tokens, query_tokens) + board_model_intent = _has_board_model_intent(query_tokens) or _has_board_model_intent( + concept_tokens + ) + scored_tables = [] + for table_name, columns in schema_details.items(): + table_tokens = _table_business_tokens(table_name, columns) + column_token_union = set() + has_numeric_sales_measure = False + has_date_capable_column = False + score = len(query_tokens & table_tokens) * 8 + for column in columns: + column_tokens = _column_business_tokens(column) + column_token_union.update(column_tokens) + if _is_numeric_type(column["data_type"]) and column_tokens & { + "amount", + "intake", + "revenue", + "sales", + "value", + }: + has_numeric_sales_measure = True + if _is_date_type(column["data_type"]) or column_tokens & { + "date", + "day", + "month", + "time", + "year", + }: + has_date_capable_column = True + score += len(query_tokens & column_tokens) * 10 + if _is_numeric_type(column["data_type"]): + score += len( + query_tokens + & column_tokens + & { + "amount", + "cost", + "count", + "margin", + "quantity", + "rate", + "score", + "value", + } + ) * 4 + if _is_date_type(column["data_type"]): + score += ( + len(query_tokens & {"date", "month", "year", "july", "trend", "trends"}) + * 4 + ) + + if not _table_covers_requested_concepts(table_name, columns, concept_tokens): + continue + + if board_model_intent and rate_metric_intent and query_tokens & { + "defect", + "failure", + }: + if not {"board", "model"}.issubset(column_token_union): + continue + rate_column = _choose_column_by_tokens( + columns, + {"defect", "rate"}, + numeric=True, + ) + if not rate_column: + continue + score += 130 + + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + score += ( + len(column_token_union & {"amount", "intake", "revenue", "sales", "value"}) + * 10 + ) + score += len(column_token_union & {"date", "month", "time", "year"}) * 5 + if not has_numeric_sales_measure: + continue + score += 50 + if query_tokens & {"year", "month", "monthly", "trend", "trends"}: + if not has_date_capable_column: + continue + score += 30 + if not query_tokens & { + "claim", + "claims", + "customs", + "duty", + "import", + "refund", + "tariff", + }: + table_and_columns = table_tokens | column_token_union + customs_matches = table_and_columns & { + "claim", + "claims", + "customs", + "duty", + "import", + "refund", + "tariff", + "tariffs", + } + if customs_matches and not table_and_columns & {"revenue", "sale", "sales"}: + continue + score -= len(customs_matches) * 40 + + if "customer" in query_tokens: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"customer", "cust"}: + if "missing" in query_tokens: + continue + score -= 80 + else: + score += 45 + + if {"business", "unit"}.issubset(query_tokens): + if "bu" in column_token_union: + score += 60 + elif {"business", "unit"} <= column_token_union: + score += 45 + + if "failure" in query_tokens and ( + query_tokens & {"location", "material", "technician", "tech"} + or board_model_intent + ): + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"failure", "failed", "defect"}: + continue + if board_model_intent: + if not {"board", "model"}.issubset(column_token_union): + continue + if failure_count_intent and not _choose_count_subject_column( + {"failure"}, + columns, + ): + continue + score += 100 + if "location" in query_tokens: + if "location" not in column_token_union: + continue + score += 90 + if "material" in query_tokens: + if not column_token_union & {"material", "part"}: + continue + score += 90 + if query_tokens & {"technician", "tech"}: + if not column_token_union & {"technician", "tech"}: + continue + score += 90 + + if query_tokens & {"order", "orders"}: + table_and_columns = table_tokens | column_token_union + explicit_order_support = table_and_columns & {"ord", "order", "orders"} + order_support = table_and_columns & { + "amount", + "customer", + "intake", + "item", + "ord", + "order", + "orders", + "sales", + "value", + } + if not order_support: + continue + if explicit_order_support: + score += 80 + else: + score -= 60 + + if query_tokens & {"batch", "batches"} and {"defect", "rate"}.issubset( + column_token_union + ): + score += 40 + if {"material", "location"}.issubset(query_tokens) and { + "material", + "location", + }.issubset(column_token_union): + score += 40 + if ( + query_tokens & {"repair", "repairs"} + and (table_tokens | column_token_union) & {"repair", "fix"} + ): + score += 55 + if concept_tokens & {"critical", "priority", "severity"}: + if not _choose_priority_column(columns): + continue + score += 55 + if concept_tokens & {"status"}: + if not column_token_union & {"status", "state", "progress"}: + continue + score += 45 + if concept_tokens & {"latest", "recent"}: + if not has_date_capable_column: + continue + score += 35 + + if score > 0: + scored_tables.append((score, table_name, columns)) + + if not scored_tables: + return None + + scored_tables.sort(key=lambda item: (-item[0], item[1])) + return scored_tables[0][1], scored_tables[0][2] + + +def _choose_column_by_tokens( + columns: list[dict[str, str]], + required_tokens: set[str], + numeric: bool | None = None, + date: bool | None = None, +) -> str | None: + candidates = [] + for column in columns: + column_tokens = _column_business_tokens(column) + if numeric is True and not _is_numeric_type(column["data_type"]): + continue + if date is True and not ( + _is_date_type(column["data_type"]) + or column_tokens & {"date", "day", "month", "time", "year"} + ): + continue + score = len(required_tokens & column_tokens) * 10 + if date is True and _is_date_type(column["data_type"]): + score += 20 + if required_tokens and required_tokens.issubset(column_tokens): + score += 30 + if score > 0: + candidates.append((score, column["name"])) + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1])) + return candidates[0][1] + + +def _column_score_for_tokens( + column: dict[str, str], + required_tokens: set[str], + numeric: bool | None = None, +) -> int: + column_tokens = _fallback_tokens(column["name"]) + column_tokens.update(column.get("semantic_tokens") or set()) + if numeric is True and not _is_numeric_type(column["data_type"]): + return 0 + score = len(required_tokens & column_tokens) * 10 + if required_tokens and required_tokens.issubset(column_tokens): + score += 30 + if column["name"].lower() == "bu" and {"business", "unit"} & required_tokens: + score += 60 + if column["name"].lower() in {"custno", "customer_id", "customerid"} and { + "customer", + "number", + } & required_tokens: + score += 30 + if column["name"].lower() in {"custname", "customer_name", "customer"} and { + "customer", + "name", + } & required_tokens: + score += 35 + if column["name"].lower() in {"ordno", "order_no", "order_number", "sales_order_number"} and { + "order", + "number", + } & required_tokens: + score += 35 + return score + + +def _choose_ranked_column_by_tokens( + columns: list[dict[str, str]], + required_tokens: set[str], + numeric: bool | None = None, +) -> dict[str, str] | None: + candidates = [ + (_column_score_for_tokens(column, required_tokens, numeric=numeric), column) + for column in columns + ] + candidates = [(score, column) for score, column in candidates if score > 0] + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1]["name"])) + return candidates[0][1] + + +def _choose_dimension_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + dimension_specs = [ + ({"board", "model"}, {"board", "model"}), + ({"business", "unit"}, {"business", "unit", "bu", "division", "company"}), + ({"customer"}, {"customer", "cust", "name"}), + ({"supplier"}, {"supplier", "vendor", "name"}), + ({"product"}, {"product", "prod", "item", "material", "name"}), + ({"salesperson", "representative"}, {"salesperson", "sales", "person", "rep"}), + ({"technician", "tech"}, {"technician", "tech"}), + ({"location"}, {"location", "site", "area"}), + ({"material"}, {"material", "part", "item"}), + ({"priority", "severity"}, {"priority", "severity", "urgency", "rank"}), + ({"status"}, {"status"}), + ({"currency"}, {"currency", "curr"}), + ({"country"}, {"country"}), + ({"order"}, {"order", "ord", "number"}), + ({"batch"}, {"batch", "id"}), + ] + for trigger_tokens, column_tokens in dimension_specs: + if query_tokens & trigger_tokens: + column = _choose_ranked_column_by_tokens(columns, column_tokens) + if column: + return column["name"] + return None + + +def _choose_missing_value_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + missing_specs = [ + ({"customer"}, {"customer", "cust", "number", "id", "no"}), + ({"order"}, {"order", "ord", "number", "id", "no"}), + ({"supplier"}, {"supplier", "vendor", "number", "id", "no"}), + ({"product", "material"}, {"product", "prod", "material", "item", "number", "id"}), + ({"location"}, {"location", "site", "area"}), + ({"status"}, {"status"}), + ] + for trigger_tokens, column_tokens in missing_specs: + if query_tokens & trigger_tokens: + column = _choose_ranked_column_by_tokens(columns, column_tokens) + if column: + return column + return None + + +def _choose_count_subject_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + subject_specs = [ + ({"order"}, {"order", "ord", "number", "id", "no"}), + ({"customer"}, {"customer", "cust", "number", "id", "no"}), + ( + {"failure"}, + {"failure", "failed", "defect", "code", "line", "status", "sys", "type"}, + ), + ({"repair"}, {"repair", "id", "status"}), + ({"batch"}, {"batch", "id"}), + ] + for trigger_tokens, column_tokens in subject_specs: + if query_tokens & trigger_tokens: + candidate_columns = columns + if trigger_tokens & {"failure"}: + candidate_columns = [ + column for column in columns if not _is_rate_like_column(column) + ] + column = _choose_ranked_column_by_tokens(candidate_columns, column_tokens) + if column: + return column + return None + + +def _is_text_type(data_type: str) -> bool: + return data_type.upper() in {"CHAR", "NCHAR", "NVARCHAR", "STRING", "TEXT", "VARCHAR"} + + +def _missing_value_predicate(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_text_type(column["data_type"]): + return f"({quoted_column} IS NULL OR {quoted_column} = '')" + return f"{quoted_column} IS NULL" + + +def _non_missing_value_predicate(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_text_type(column["data_type"]): + return f"({quoted_column} IS NOT NULL AND {quoted_column} <> '')" + return f"{quoted_column} IS NOT NULL" + + +def _aggregate_for_measure(measure_column: str) -> tuple[str, str]: + tokens = _fallback_tokens(measure_column) + if tokens & {"rate", "score", "percent", "percentage"}: + return "AVG", "average_value" + return "SUM", "total_value" + + +def _select_listing_columns( + query_tokens: set[str], + columns: list[dict[str, str]], + measure_column: str | None = None, + date_column: str | None = None, + max_columns: int = 6, +) -> list[str]: + scored_columns: list[tuple[int, int, str]] = [] + for index, column in enumerate(columns): + name = column["name"] + tokens = _column_business_tokens(column) + score = len(query_tokens & tokens) * 10 + if name == date_column: + score += 35 + if name == measure_column: + score += 12 + if query_tokens & {"order"}: + score += len(tokens & {"order", "ord", "number", "customer", "cust", "item", "product"}) * 18 + score += len(tokens & {"date", "day", "month", "year"}) * 8 + if query_tokens & {"customer"}: + score += len(tokens & {"customer", "cust", "name", "number", "id"}) * 18 + if query_tokens & {"product", "material"}: + score += len(tokens & {"product", "prod", "material", "item", "description", "desc"}) * 18 + if query_tokens & {"batch"}: + score += len(tokens & {"batch", "board", "model", "supplier", "id"}) * 18 + if query_tokens & {"repair", "log", "record"}: + score += ( + len(tokens & {"board", "code", "date", "failure", "id", "priority", "status"}) + * 14 + ) + if tokens & {"repair", "failure", "failed", "defect"} and not query_tokens & { + "repair", + "failure", + "defect", + }: + score -= 40 + if tokens & {"date", "day", "month", "year"} and not ( + _is_date_type(column["data_type"]) or name == date_column + ): + score -= 12 + if score > 0: + scored_columns.append((score, index, name)) + + scored_columns.sort(key=lambda item: (-item[0], item[1])) + selected = [] + for _, _, name in scored_columns: + if name not in selected: + selected.append(name) + if len(selected) >= max_columns: + break + if not selected and columns: + selected = [column["name"] for column in columns[:max_columns]] + return selected + + +def _fallback_limit(query: str) -> int | None: + match = re.search(r"(?i)\btop\s+(\d+)\b", query) + return int(match.group(1)) if match else None + + +def _fallback_month_filter(query: str) -> tuple[int, int] | None: + tokens = _fallback_tokens(query) + for month_name, month_number in _MONTH_NAME_TO_NUMBER.items(): + if month_name in tokens: + return datetime.now(timezone.utc).year, month_number + return None + + +def _quote_literal(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def _clean_filter_value(value: str | None) -> str | None: + if value is None: + return None + value = value.strip(" \t\r\n'\"`.,;:()[]{}") + return value or None + + +def _extract_failure_type_filter_value(query: str) -> str | None: + patterns = [ + ( + r"(?is)\bwith\s+" + r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80}?)" + r"\s+as\s+(?:the\s+)?(?:failure|defect)\s+" + r"(?:type|code|category)\b" + ), + ( + r"(?is)\b(?:failure|defect)\s+(?:type|code|category)\s*" + r"(?:=|is|equals|like|of)\s*['\"]?" + r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80})" + ), + ] + for pattern in patterns: + match = re.search(pattern, query) + if match: + value = _clean_filter_value(match.group("value")) + if value: + return value + return None + + +def _extract_status_filter_value(query: str) -> str | None: + if re.search(r"(?i)\bin-progress\b", query): + return "in-progress" + if re.search(r"(?i)\bin\s+progress\b", query): + return "in progress" + for status in ("completed", "pending", "escalated", "open", "closed"): + if re.search(rf"(?i)\b{re.escape(status)}\b", query): + return status + return None + + +def _extract_priority_filter_value(query: str) -> str | None: + for token, value in _PRIORITY_VALUE_ALIASES.items(): + if re.search(rf"(?i)\b{re.escape(token)}(?:[-\s]+priority)?\b", query): + return value + return None + + +def _choose_priority_column(columns: list[dict[str, str]]) -> dict[str, str] | None: + return _choose_ranked_column_by_tokens( + columns, + {"priority", "severity", "urgency", "rank"}, + ) + + +def _priority_order_expression(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_numeric_type(column["data_type"]): + return quoted_column + + when_clauses = " ".join( + f"WHEN {_quote_literal(value)} THEN {rank}" for value, rank in _PRIORITY_ORDER + ) + return f"CASE LOWER({quoted_column}) {when_clauses} ELSE 0 END" + + +def _choose_failure_type_filter_column( + columns: list[dict[str, str]], +) -> dict[str, str] | None: + column = _choose_ranked_column_by_tokens( + [column for column in columns if not _is_rate_like_column(column)], + {"failure", "type"}, + ) + if column: + return column + return _choose_ranked_column_by_tokens( + [column for column in columns if not _is_rate_like_column(column)], + {"failure", "defect", "code", "type", "sys"}, + ) + + +def _choose_failure_type_filter_table( + schema_details: dict[str, list[dict[str, str]]], + concept_tokens: set[str], +) -> tuple[str, list[dict[str, str]], dict[str, str]] | None: + candidates = [] + for table_name, columns in schema_details.items(): + if not _table_covers_requested_concepts(table_name, columns, concept_tokens): + continue + column = _choose_failure_type_filter_column(columns) + if not column: + continue + table_tokens = _fallback_tokens(table_name) + column_tokens = _fallback_tokens(column["name"]) + score = len(concept_tokens & table_tokens) * 8 + score += len(concept_tokens & column_tokens) * 10 + if {"failure", "type"}.issubset(column_tokens): + score += 100 + elif "type" in column_tokens: + score += 60 + elif "code" in column_tokens: + score += 30 + if table_tokens & {"failure", "defect"}: + score += 25 + candidates.append((score, table_name, columns, column)) + + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1], item[3]["name"])) + _, table_name, columns, column = candidates[0] + return table_name, columns, column + + +def generate_simple_analytics_sql( + query: str | None, + contexts: list[Any] | None, +) -> str | None: + if not query: + return None + + raw_query_tokens = _fallback_tokens(query) + query_tokens = _expanded_fallback_query_tokens(query) + if not query_tokens: + return None + + if not query_tokens & { + "batch", + "batches", + "board", + "business", + "count", + "customer", + "defect", + "failure", + "failures", + "highest", + "july", + "latest", + "log", + "logs", + "location", + "material", + "missing", + "model", + "monthly", + "most", + "number", + "order", + "orders", + "priority", + "product", + "rate", + "recent", + "record", + "records", + "repair", + "repairs", + "revenue", + "sale", + "sales", + "severity", + "supplier", + "status", + "tech", + "technician", + "top", + "trend", + "trends", + "type", + "unit", + "units", + "year", + }: + return None + + schema_details = _extract_schema_details(contexts) + rate_metric_intent = _is_rate_metric_intent(raw_query_tokens) + failure_count_intent = _is_failure_count_intent(raw_query_tokens, query_tokens) + board_model_intent = _has_board_model_intent( + raw_query_tokens + ) or _has_board_model_intent( + query_tokens + ) + failure_type_filter_value = _extract_failure_type_filter_value(query) + failure_type_filter_column = None + failure_type_choice = None + if failure_type_filter_value and "failure" in query_tokens: + failure_type_choice = _choose_failure_type_filter_table( + schema_details, + raw_query_tokens, + ) + + if failure_type_choice: + table_name, columns, failure_type_filter_column = failure_type_choice + chosen = (table_name, columns) + else: + chosen = _choose_fallback_table( + query_tokens, + schema_details, + concept_tokens=raw_query_tokens, + ) + if not chosen: + return None + + table_name, columns = chosen + column_names = [column["name"] for column in columns] + quoted_table = _quote_identifier(table_name) + limit = _fallback_limit(query) + logger.info( + "Deterministic SQL fallback selected table=%s verified_columns=%s metric_intent=%s", + table_name, + column_names, + { + "failure_count": failure_count_intent, + "rate": rate_metric_intent, + "board_model": board_model_intent, + "failure_type_filter": bool(failure_type_filter_value), + }, + ) + + if failure_type_filter_value: + if not failure_type_filter_column: + failure_type_filter_column = _choose_failure_type_filter_column(columns) + if failure_type_filter_column: + return ( + f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}\n" + f"WHERE {_quote_identifier(failure_type_filter_column['name'])} = " + f"{_quote_literal(failure_type_filter_value)}" + ) + + material_column = _choose_column_by_tokens(columns, {"material"}) + location_column = _choose_column_by_tokens(columns, {"location"}) + if ( + {"material", "location"}.issubset(query_tokens) + and material_column + and location_column + ): + return ( + f"SELECT {_quote_joined([material_column, location_column])}\n" + f"FROM {quoted_table}" + ) + + status_column = _choose_column_by_tokens(columns, {"status"}) + repair_filter_intent = raw_query_tokens & { + "closed", + "completed", + "critical", + "escalated", + "high", + "low", + "medium", + "normal", + "open", + "pending", + "priority", + "progress", + "severity", + "status", + "urgent", + } + if query_tokens & {"repair", "repairs"} and repair_filter_intent: + predicates = [] + status_filter_value = _extract_status_filter_value(query) + if status_column and status_filter_value: + predicates.append( + f"{_quote_identifier(status_column)} = {_quote_literal(status_filter_value)}" + ) + priority_filter_value = _extract_priority_filter_value(query) + if priority_filter_value: + priority_column = _choose_priority_column(columns) + if priority_column: + predicates.append( + f"{_quote_identifier(priority_column['name'])} = " + f"{_quote_literal(priority_filter_value)}" + ) + if predicates: + return f"SELECT *\nFROM {quoted_table}\nWHERE {' AND '.join(predicates)}" + + date_column = _choose_column_by_tokens( + columns, + {"date", "day", "month", "time", "year"}, + date=True, + ) + + priority_column = _choose_priority_column(columns) + if ( + priority_column + and raw_query_tokens & {"priority", "severity"} + and raw_query_tokens & {"bottom", "highest", "lowest", "top"} + ): + selected_columns = _select_listing_columns( + raw_query_tokens | {"priority", "repair", "status"}, + columns, + date_column=date_column, + max_columns=8, + ) + if priority_column["name"] not in selected_columns: + selected_columns.insert(0, priority_column["name"]) + direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"ORDER BY {_priority_order_expression(priority_column)} {direction}" + f"{limit_clause}" + ) + + if date_column and raw_query_tokens & {"latest", "recent"}: + selected_columns = _select_listing_columns( + raw_query_tokens | {"date", "repair", "status"}, + columns, + date_column=date_column, + max_columns=8, + ) + if date_column not in selected_columns: + selected_columns.insert(0, date_column) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" + ) + + if ( + query_tokens & {"repair", "repairs"} + and raw_query_tokens & {"priority", "severity", "status"} + and re.search(r"(?i)\bby\s+(?:priority|severity|status)\b", query) + ): + dimension_column = _choose_dimension_column(raw_query_tokens, columns) + subject_column = _choose_count_subject_column({"repair"}, columns) + if dimension_column: + count_expression = "COUNT(*)" + where_clause = "" + if subject_column: + count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" + where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" + quoted_dimension = _quote_identifier(dimension_column) + return ( + f"SELECT {quoted_dimension}, {count_expression} AS " + f"{_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC" + ) + measure_column = None + if rate_metric_intent and query_tokens & {"defect", "failure"}: + measure_column = _choose_column_by_tokens( + columns, + {"defect", "rate"}, + numeric=True, + ) + if not measure_column: + measure_column = _choose_column_by_tokens(columns, {"rate"}, numeric=True) + if not measure_column and query_tokens & {"order", "orders"}: + measure_column = _choose_column_by_tokens( + columns, + {"amount", "intake", "sales", "value"}, + numeric=True, + ) + if not measure_column and query_tokens & {"revenue", "sale", "sales"}: + measure_column = _choose_column_by_tokens( + columns, + {"amount", "intake", "revenue", "sales", "value"}, + numeric=True, + ) + if ( + not measure_column + and not failure_count_intent + and query_tokens & {"top", "highest", "lowest", "bottom"} + ): + measure_column = _choose_column_by_tokens( + columns, + {"amount", "cost", "count", "margin", "quantity", "rate", "score", "value"}, + numeric=True, + ) + + if raw_query_tokens & {"missing", "blank", "empty", "null"}: + missing_column = _choose_missing_value_column(raw_query_tokens, columns) + if missing_column: + selected_columns = [ + column + for column in column_names + if column == missing_column["name"] + or _fallback_tokens(column) + & { + "batch", + "business", + "bu", + "customer", + "cust", + "date", + "id", + "location", + "name", + "number", + "ord", + "order", + "product", + "status", + "supplier", + } + ][:8] + if missing_column["name"] not in selected_columns: + selected_columns.insert(0, missing_column["name"]) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"WHERE {_missing_value_predicate(missing_column)}{limit_clause}" + ) + + implied_count_by_dimension = "failure" in query_tokens and bool( + raw_query_tokens & {"location", "material", "technician", "tech"} + or (board_model_intent and not rate_metric_intent) + ) + if ( + query_tokens & {"count", "number"} + or failure_count_intent + or implied_count_by_dimension + ): + dimension_column = _choose_dimension_column(raw_query_tokens, columns) + if dimension_column: + subject_column = _choose_count_subject_column(raw_query_tokens, columns) + count_expression = "COUNT(*)" + where_clause = "" + if subject_column and raw_query_tokens & {"order", "customer"}: + count_expression = ( + f"COUNT(DISTINCT {_quote_identifier(subject_column['name'])})" + ) + elif subject_column and raw_query_tokens & {"failure", "repair", "batch"}: + count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" + where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" + quoted_dimension = _quote_identifier(dimension_column) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {quoted_dimension}, {count_expression} AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" + ) + + dimension_column = _choose_dimension_column(raw_query_tokens, columns) + if measure_column and dimension_column and rate_metric_intent: + aggregate, alias = _aggregate_for_measure(measure_column) + quoted_dimension = _quote_identifier(dimension_column) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" + direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + + if measure_column and limit and dimension_column and raw_query_tokens & { + "board", + "business", + "customer", + "location", + "material", + "model", + "product", + "salesperson", + "supplier", + "unit", + }: + aggregate, alias = _aggregate_for_measure(measure_column) + quoted_dimension = _quote_identifier(dimension_column) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" + return ( + f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier(alias)} DESC\nLIMIT {limit}" + ) + + month_filter = _fallback_month_filter(query) + if date_column and month_filter: + year, month = month_filter + start = f"{year:04d}-{month:02d}-01" + end_year = year + 1 if month == 12 else year + end_month = 1 if month == 12 else month + 1 + end = f"{end_year:04d}-{end_month:02d}-01" + selected_columns = _select_listing_columns( + raw_query_tokens, + columns, + measure_column=measure_column, + date_column=date_column, + ) + order_clause = ( + f"\nORDER BY {_quote_identifier(measure_column)} DESC" + if measure_column + else f"\nORDER BY {_quote_identifier(date_column)} DESC" + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"WHERE {_quote_identifier(date_column)} >= '{start}' " + f"AND {_quote_identifier(date_column)} < '{end}'" + f"{order_clause}{limit_clause}" + ) + + if ( + measure_column + and date_column + and query_tokens & {"year"} + and not query_tokens & {"month", "monthly"} + ): + date_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" + return ( + f"SELECT {date_expr} AS {_quote_identifier('year')}, " + f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" + f"FROM {quoted_table}\nGROUP BY {date_expr}\nORDER BY {date_expr}" + ) + + if ( + measure_column + and date_column + and query_tokens & {"month", "monthly", "trend", "trends"} + ): + year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" + month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" + f"FROM {quoted_table}\nGROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" + ) + + if measure_column and limit: + selected_columns = [ + column + for column in column_names + if column == measure_column + or _fallback_tokens(column) + & { + "batch", + "board", + "customer", + "id", + "model", + "name", + "number", + "supplier", + } + ][:6] + if measure_column not in selected_columns: + selected_columns.append(measure_column) + return ( + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"ORDER BY {_quote_identifier(measure_column)} DESC\nLIMIT {limit}" + ) + + return None + + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): @@ -759,48 +2646,164 @@ async def run( replies: List[str] | List[List[str]], project_id: str | None = None, mdl_hash: str | None = None, + contexts: list[Any] | None = None, + fallback_query: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, - contexts: list[str] | None = None, ) -> dict: try: cleaned_generation_result, extraction_error = _extract_sql_response( clean_generation_result(replies[0]) ) - if extraction_error: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": "", - "original_sql": "", - "type": "NO_RELEVANT_SQL", - "error": extraction_error, - "correlation_id": "", - "data_source": data_source, - }, - } + grounding_invalid_generation_result = None + + def validate_candidate_sql(candidate_sql: str) -> str | None: + schema_catalog = _SchemaCatalog.from_contexts(contexts or []) + grounding_error = schema_catalog.validate_sql(candidate_sql) + if not grounding_error: + grounding_error = validate_sql_against_contexts( + candidate_sql, + contexts=contexts, + ) + if not grounding_error: + grounding_error = validate_sql_semantic_coverage( + candidate_sql, + fallback_query, + contexts=contexts, + ) + return grounding_error if cleaned_generation_result: cleaned_generation_result = normalize_sql_with_schema_identifiers( cleaned_generation_result, contexts=contexts, ) + cleaned_generation_result = normalize_wren_sql_dialect( + cleaned_generation_result + ) + grounding_error = validate_candidate_sql(cleaned_generation_result) + if grounding_error: + logger.info( + "Generated SQL validation result project_id=%s status=rejected reason=%s sql=%s", + project_id or "", + grounding_error, + cleaned_generation_result, + ) + grounding_invalid_generation_result = { + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_GROUNDING", + "error": grounding_error, + "correlation_id": "", + "data_source": data_source, + } + else: + logger.info( + "Generated SQL validation result project_id=%s status=grounded sql=%s", + project_id or "", + cleaned_generation_result, + ) + elif extraction_error: + logger.info( + "Generated SQL extraction result project_id=%s status=rejected reason=%s", + project_id or "", + extraction_error, + ) + + fallback_generation_result = generate_simple_analytics_sql( + fallback_query, + contexts, + ) + if fallback_generation_result: + logger.info( + "Deterministic SQL fallback generated project_id=%s sql=%s", + project_id or "", + fallback_generation_result, + ) + fallback_generation_result = normalize_sql_with_schema_identifiers( + fallback_generation_result, + contexts=contexts, + ) + fallback_generation_result = normalize_wren_sql_dialect( + fallback_generation_result + ) + fallback_grounding_error = validate_candidate_sql( + fallback_generation_result + ) + logger.info( + "Deterministic SQL fallback validation result project_id=%s status=%s%s", + project_id or "", + "grounded" if not fallback_grounding_error else "rejected", + "" + if not fallback_grounding_error + else f" reason={fallback_grounding_error}", + ) + if not fallback_grounding_error: + ( + fallback_valid_generation_result, + fallback_invalid_generation_result, + ) = await self._classify_generation_result( + fallback_generation_result, + project_id=project_id, + mdl_hash=mdl_hash, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + data_source=data_source, + allow_data_preview=allow_data_preview, + ) + if fallback_valid_generation_result: + logger.info( + "Using deterministic schema-grounded SQL fallback for query." + ) + return { + "valid_generation_result": fallback_valid_generation_result, + "invalid_generation_result": {}, + } + logger.info( + "Deterministic SQL fallback did not validate: %s", + fallback_invalid_generation_result.get("error"), + ) + + if grounding_invalid_generation_result: + unsupported_result = unsupported_schema_generation_result( + fallback_query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + unsupported_message = unsupported_result[ + "invalid_generation_result" + ]["error"] + grounding_invalid_generation_result["type"] = "NO_RELEVANT_SQL" + grounding_invalid_generation_result["error"] = unsupported_message + grounding_invalid_generation_result["sql"] = "" + grounding_invalid_generation_result["original_sql"] = "" + return { + "valid_generation_result": {}, + "invalid_generation_result": grounding_invalid_generation_result, + } - schema_catalog = _SchemaCatalog.from_contexts(contexts or []) - grounding_error = schema_catalog.validate_sql(cleaned_generation_result) - if grounding_error: + unsupported_result = unsupported_schema_generation_result( + fallback_query, + contexts=contexts, + data_source=data_source, + ) + if not cleaned_generation_result and unsupported_result: + return unsupported_result + + if not cleaned_generation_result and extraction_error: return { "valid_generation_result": {}, "invalid_generation_result": { - "sql": cleaned_generation_result or "", - "original_sql": cleaned_generation_result or "", - "type": "SCHEMA_GROUNDING", - "error": grounding_error, + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": extraction_error, "correlation_id": "", "data_source": data_source, - } + }, } ( @@ -1520,11 +3523,16 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. - If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. +- In grouped queries, every non-aggregate ORDER BY expression must be a selected grouping column, a selected ordering helper column that is also present in GROUP BY, or a selected aggregate alias. Do not order grouped SQL by a hidden column. - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. - If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. - If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. - If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. +- Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema contains specific modeled business columns for the requested entity, measure, status, date, or dimension. +- Prefer exact modeled business fields over generic text search. For example, if a status/severity/date/material/location/customer/order/revenue concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. +- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. This applies to metrics such as defect rate, revenue, amount, sales value, count, cost, margin, and quantity. +- For sales or revenue questions, prefer exact declared sales/revenue/value/amount fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. """ @@ -1576,8 +3584,10 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. - Do not use SELECT TOP n, FETCH FIRST, OFFSET/FETCH, square-bracket quoting, or backtick quoting. Use Wren SQL syntax with ORDER BY and a final LIMIT n clause for limited or top-N results. -- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -- For the ranking problem, you must add the ranking column to the final SELECT clause. +- For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. +- For explicit ranking requests, use the ranking function `DENSE_RANK()`, add the ranking column to the final SELECT clause, and filter rank values with WHERE. +- For grouped trend queries, include any non-aggregate ordering key in both SELECT and GROUP BY, or order by selected grouping columns/aggregate aliases only. +- Reuse exact metric/measure columns when present. Do not invent component columns in order to calculate a requested metric that already exists in DATABASE SCHEMA. """ diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 19ea275650..a8d5ac5f27 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,5 +1,6 @@ import ast import logging +import re import sys from typing import Any, Optional @@ -10,7 +11,7 @@ from hamilton.async_driver import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from sqlparse.sql import Identifier, IdentifierList from sqlparse.tokens import DML, Comment, Keyword @@ -28,6 +29,68 @@ logger = logging.getLogger("wren-ai-service") +_SEMANTIC_TABLE_NAME_MERGE_LIMIT = 8 +_MAX_RETRIEVED_TABLE_NAMES = 24 +_MAX_RELATED_TABLE_EXPANSION_DEPTH = 1 +_RANK_TOKEN = re.compile(r"[a-z0-9]+") +_GENERIC_TABLE_TOKENS = { + "audit", + "auth", + "calendar", + "config", + "dim", + "dimension", + "file", + "files", + "ingestion", + "job", + "jobs", + "log", + "logs", + "lookup", + "mbr", + "member", + "members", + "migration", + "migrations", + "preference", + "preferences", + "queue", + "report", + "reports", + "setting", + "settings", + "state", + "time", + "user", + "users", +} +_CUSTOMS_FINANCE_TOKENS = { + "claim", + "claims", + "custom", + "customs", + "duty", + "duties", + "hmf", + "import", + "imports", + "mpf", + "refund", + "refunds", + "tariff", + "tariffs", +} +_SALES_REVENUE_TOKENS = { + "amount", + "intake", + "revenue", + "sale", + "sales", + "salesvalue", + "value", +} + table_columns_selection_system_prompt = """ ### TASK ### @@ -60,6 +123,12 @@ 16. Prefer the set of deployed models, views, metrics, columns, and relationships that best support the current question. 17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. 18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. +19. Prefer tables and columns that directly model the requested business entities, measures, statuses, dates, identifiers, and dimensions. Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema provides specific modeled columns for the same concept. +20. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, repairs, failures, batches, materials, locations, status, severity, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. +21. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. +22. Never return placeholder table or column names such as tablename, table_name, dbo.tablename, BatchId, Material, Location, or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. +23. If the request asks for revenue, sales, or sales trends, prefer exact business measure columns named like Revenue, SalesValue, USDFXSalesValue, FXSalesValue, IntakeValue, Amount, or equivalent modeled sales fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. +24. If the request asks for an explicit rate, ratio, percentage, revenue, amount, sales value, or other named measure and the schema already contains that exact measure column, use the declared measure column directly. Do not use a rate column to answer "most failures", "number of failures", or other count-of-records requests unless the question explicitly asks for a rate/ratio/percentage. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -570,7 +639,7 @@ def _empty_retrieval_results() -> dict[str, Any]: @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: - return await embedder.run(query) + return await embedder.run(_augment_retrieval_query(query)) else: return {} @@ -613,7 +682,8 @@ async def dbschema_retrieval( table_retrieval: dict, project_id: str, dbschema_retriever: Any, - embedding: dict, + query: str | None = None, + embedding: dict | None = None, mdl_hash: str | None = None, include_related_models: bool = True, ) -> list[Document]: @@ -621,16 +691,33 @@ async def dbschema_retrieval( table_retrieval.get("documents", []) ) documents = [] - if embedding and not table_names: - documents = await _retrieve_semantic_schema_documents( + if embedding: + semantic_documents = await _retrieve_semantic_schema_documents( embedding, project_id, mdl_hash, dbschema_retriever ) - table_names = _table_names_from_schema_documents(documents) + semantic_table_names = _table_names_from_schema_documents(semantic_documents)[ + :_SEMANTIC_TABLE_NAME_MERGE_LIMIT + ] + table_names = _merge_names(table_names, semantic_table_names)[ + :_MAX_RETRIEVED_TABLE_NAMES + ] + table_names = _rank_table_names_by_query( + table_names, + semantic_documents, + query, + ) + selected_semantic_table_names = set(semantic_table_names) + documents = [ + document + for document in semantic_documents + if document.meta.get("name") in selected_semantic_table_names + ] if table_names: if include_related_models: retrieved_table_names = set() pending_table_names = table_names + remaining_expansion_depth = _MAX_RELATED_TABLE_EXPANSION_DEPTH while pending_table_names: retrieved_table_names.update(pending_table_names) @@ -638,22 +725,242 @@ async def dbschema_retrieval( pending_table_names, project_id, mdl_hash, dbschema_retriever ) documents = _dedupe_documents(documents + retrieved_documents) + if remaining_expansion_depth <= 0: + break + remaining_expansion_depth -= 1 + remaining_slots = _MAX_RETRIEVED_TABLE_NAMES - len( + retrieved_table_names + ) + if remaining_slots <= 0: + break pending_table_names = [ table_name for table_name in _related_table_names(documents) if table_name not in retrieved_table_names - ] + ][:remaining_slots] - return documents + ranked_documents = _rank_documents_for_query(documents, table_names, query) + logger.info( + "Ask schema retrieval project_id=%s retrieved_tables=%s", + project_id, + [ + { + "table": document.meta.get("name"), + "score": getattr(document, "score", None), + } + for document in ranked_documents + ], + ) + return ranked_documents retrieved_documents = await _retrieve_schema_documents( table_names, project_id, mdl_hash, dbschema_retriever ) - return _dedupe_documents(documents + retrieved_documents) + documents = _dedupe_documents(documents + retrieved_documents) + ranked_documents = _rank_documents_for_query(documents, table_names, query) + logger.info( + "Ask schema retrieval project_id=%s retrieved_tables=%s", + project_id, + [ + { + "table": document.meta.get("name"), + "score": getattr(document, "score", None), + } + for document in ranked_documents + ], + ) + return ranked_documents + logger.info("Ask schema retrieval project_id=%s retrieved_tables=[]", project_id) return [] +def _tokenize_schema_text(value: Any) -> set[str]: + if value is None: + return set() + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) + return set(_RANK_TOKEN.findall(text.lower())) + + +def _schema_rank_text_by_table(documents: list[Document]) -> dict[str, dict[str, set[str]]]: + table_text: dict[str, dict[str, set[str]]] = {} + + def ensure(table_name: str) -> dict[str, set[str]]: + if table_name not in table_text: + table_text[table_name] = { + "table": set(), + "columns": set(), + "comments": set(), + } + return table_text[table_name] + + for document in documents: + try: + content = ast.literal_eval(document.content) + except (SyntaxError, ValueError): + continue + + table_name = document.meta.get("name") or content.get("name") + if not table_name: + continue + + bucket = ensure(table_name) + bucket["table"].update(_tokenize_schema_text(table_name)) + bucket["table"].update(_tokenize_schema_text(content.get("name"))) + bucket["comments"].update(_tokenize_schema_text(content.get("comment"))) + bucket["comments"].update(_tokenize_schema_text(content.get("description"))) + + for column in content.get("columns", []) or []: + bucket["columns"].update(_tokenize_schema_text(column.get("name"))) + bucket["columns"].update(_tokenize_schema_text(column.get("column"))) + bucket["columns"].update(_tokenize_schema_text(column.get("display_name"))) + bucket["comments"].update(_tokenize_schema_text(column.get("comment"))) + bucket["comments"].update(_tokenize_schema_text(column.get("description"))) + + return table_text + + +def _rank_table_names_by_query( + table_names: list[str], + semantic_documents: list[Document], + query: str | None, +) -> list[str]: + if not query or not table_names: + return table_names + + query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) + if not query_tokens: + return table_names + + table_text = _schema_rank_text_by_table(semantic_documents) + + def score(table_name: str) -> int: + bucket = table_text.get(table_name, {}) + table_tokens = set(bucket.get("table", set())) | _tokenize_schema_text( + table_name + ) + column_tokens = set(bucket.get("columns", set())) + comment_tokens = set(bucket.get("comments", set())) + direct_table_matches = query_tokens & table_tokens + direct_column_matches = query_tokens & column_tokens + + value = ( + len(direct_table_matches) * 6 + + len(direct_column_matches) * 8 + + len(query_tokens & comment_tokens) + ) + + if len(direct_column_matches) >= 2: + value += 8 + if direct_table_matches and direct_column_matches: + value += 8 + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + value += len((table_tokens | column_tokens) & _SALES_REVENUE_TOKENS) * 5 + if not query_tokens & _CUSTOMS_FINANCE_TOKENS: + value -= ( + len((table_tokens | column_tokens) & _CUSTOMS_FINANCE_TOKENS) + * 8 + ) + if table_tokens & _GENERIC_TABLE_TOKENS and not direct_table_matches: + value -= 6 + return value + + ranked = sorted( + enumerate(table_names), + key=lambda item: (-score(item[1]), item[0]), + ) + return [table_name for _, table_name in ranked] + + +def _rank_documents_by_table_names( + documents: list[Document], + table_names: list[str], +) -> list[Document]: + table_rank = {table_name: index for index, table_name in enumerate(table_names)} + return sorted( + documents, + key=lambda document: ( + table_rank.get(document.meta.get("name"), len(table_rank)), + document.meta.get("type", ""), + ), + ) + + +def _rank_documents_for_query( + documents: list[Document], + table_names: list[str], + query: str | None, +) -> list[Document]: + ranked_table_names = _rank_table_names_by_query( + _merge_names(_table_names_from_schema_documents(documents), table_names), + documents, + query, + ) + return _rank_documents_by_table_names(documents, ranked_table_names) + + +def _augment_retrieval_query(query: str) -> str: + lowered = query.lower() + expansions = [] + + concept_terms = { + ("revenue", "sales", "sale", "amount", "value"): ( + "sales revenue amount value gross net total price intake invoice order" + ), + ("order", "orders"): ( + "order ord number date customer product business unit division company" + ), + ("invoice", "invoices"): ( + "invoice supplier customer currency amount date number" + ), + ("customer", "customers"): ( + "customer account client number name identifier" + ), + ("product", "products"): ( + "product item material type name category" + ), + ("repair", "repairs"): ( + "repair status priority severity failure board model log in progress completed critical" + ), + ("failure", "failures", "defect", "defects"): ( + "failure defect severity occurrence record count code type system status" + ), + ("batch", "batches"): ( + "batch board model supplier defect rate inspection status" + ), + ("material", "materials"): ( + "material item part component location" + ), + ("location", "locations"): ( + "location site warehouse area material" + ), + ("business unit", "bu", "division"): ( + "business unit division company account organization" + ), + ("month", "monthly", "july", "year", "trend", "latest"): ( + "date month year fiscal calendar trend latest recent" + ), + ("status", "severity", "priority", "critical"): ( + "status priority severity critical state category progress" + ), + } + + for triggers, terms in concept_terms.items(): + if any(trigger in lowered for trigger in triggers): + expansions.append(terms) + + if any( + trigger in lowered + for trigger in ("rate", "ratio", "percent", "percentage") + ): + expansions.append("rate ratio percent percentage") + + if not expansions: + return query + + return f"{query}\nBusiness schema search terms: {'; '.join(expansions)}" + + async def _retrieve_semantic_schema_documents( embedding: dict, project_id: str, @@ -947,15 +1254,20 @@ def construct_retrieval_results( filter_columns_in_tables: dict, construct_db_schemas: list[dict], dbschema_retrieval: list[Document], + query: str | None = None, ) -> dict[str, Any]: if filter_columns_in_tables: - try: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - ).get("results") - except orjson.JSONDecodeError: - columns_and_tables_needed = None - + columns_and_tables_needed = _parse_column_selection_response( + filter_columns_in_tables + ) + lexical_columns_and_tables_needed = _lexical_columns_and_tables_needed( + construct_db_schemas, + query, + ) + columns_and_tables_needed = _merge_column_selection( + columns_and_tables_needed, + lexical_columns_and_tables_needed, + ) if not columns_and_tables_needed: logger.warning( "Column pruning did not return grounded schema selections; " @@ -963,14 +1275,9 @@ def construct_retrieval_results( ) return _empty_retrieval_results() - # we need to change the below code to match the new schema of structured output - # the objective of this loop is to change the structure of JSON to match the needed format - reformated_json = {} - for table in columns_and_tables_needed: - reformated_json[table["table_name"]] = table["table_contents"] - columns_and_tables_needed = reformated_json tables = set(columns_and_tables_needed.keys()) retrieval_results = [] + selected_schema_log = [] has_calculated_field = False has_metric = False has_json_field = False @@ -1021,6 +1328,21 @@ def construct_retrieval_results( ), } ) + selected_schema_log.append( + { + "table": table_schema["name"], + "columns": sorted(selected_columns), + } + ) + + if not retrieval_results: + logger.warning( + "Column-selection output did not match retrieved schemas; " + "falling back to unpruned retrieved schema context." + ) + return _build_unpruned_retrieval_results( + construct_db_schemas, dbschema_retrieval + ) for document in dbschema_retrieval: content = ast.literal_eval(document.content) @@ -1061,6 +1383,7 @@ def construct_retrieval_results( } ) + logger.info("Ask retrieval selected schema objects=%s", selected_schema_log) return { "retrieval_results": retrieval_results, "has_calculated_field": has_calculated_field, @@ -1069,6 +1392,13 @@ def construct_retrieval_results( } else: retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] + logger.info( + "Ask retrieval selected schema objects=%s", + [ + {"table": retrieval_result.get("table_name"), "columns": "all"} + for retrieval_result in retrieval_results + ], + ) return { "retrieval_results": retrieval_results, @@ -1080,27 +1410,302 @@ def construct_retrieval_results( } +def _normalize_column_selection_results(parsed_response: Any) -> list[dict]: + if isinstance(parsed_response, list): + return [item for item in parsed_response if isinstance(item, dict)] + + if not isinstance(parsed_response, dict): + return [] + + for key in ( + "results", + "tables", + "selected_tables", + "retrieval_results", + "matches", + "data", + "result", + "output", + ): + if key in parsed_response: + normalized = _normalize_column_selection_results(parsed_response[key]) + if normalized: + return normalized + + if "table_name" in parsed_response and ( + "table_contents" in parsed_response or "columns" in parsed_response + ): + return [parsed_response] + + keyed_tables = [] + for table_name, table_contents in parsed_response.items(): + if not isinstance(table_name, str) or not isinstance(table_contents, dict): + continue + if "table_contents" in table_contents: + keyed_tables.append( + { + "table_name": table_name, + "table_contents": table_contents["table_contents"], + } + ) + elif "columns" in table_contents: + keyed_tables.append( + {"table_name": table_name, "table_contents": table_contents} + ) + + return keyed_tables + + +def _parse_column_selection_response(filter_columns_in_tables: dict) -> dict: + raw_reply = (filter_columns_in_tables.get("replies") or [""])[0] + try: + parsed_response = orjson.loads(raw_reply) + except orjson.JSONDecodeError as exc: + logger.warning("Unable to parse column-selection JSON response: %s", exc) + return {} + + normalized_tables = _normalize_column_selection_results(parsed_response) + reformatted_json = {} + for table in normalized_tables: + table_name = table.get("table_name") or table.get("name") + table_contents = table.get("table_contents") or {} + if not table_contents and "columns" in table: + table_contents = table + + columns = ( + table_contents.get("columns") if isinstance(table_contents, dict) else None + ) + if not isinstance(table_name, str) or not isinstance(columns, list): + continue + + reformatted_json[table_name] = { + **table_contents, + "columns": [column for column in columns if isinstance(column, str)], + } + + if not reformatted_json: + response_shape = ( + f"keys={list(parsed_response.keys())[:8]}" + if isinstance(parsed_response, dict) + else type(parsed_response).__name__ + ) + logger.warning( + "Column-selection response did not include usable table columns (%s).", + response_shape, + ) + + return reformatted_json + + +def _build_unpruned_retrieval_results( + construct_db_schemas: list[dict], + dbschema_retrieval: list[Document], +) -> dict: + retrieval_results = [] + has_calculated_field = False + has_metric = False + has_json_field = False + + for table_schema in construct_db_schemas: + if table_schema["type"] == "TABLE": + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) + retrieval_results.append( + { + "table_name": table_schema["name"], + "table_ddl": ddl, + } + ) + if _has_calculated_field: + has_calculated_field = True + if _has_json_field: + has_json_field = True + + for document in dbschema_retrieval: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) + + return { + "retrieval_results": retrieval_results, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + + +def _merge_column_selection( + primary: dict[str, dict], + secondary: dict[str, dict], +) -> dict[str, dict]: + merged = { + table_name: { + **table_contents, + "columns": list(table_contents.get("columns", [])), + } + for table_name, table_contents in primary.items() + } + + for table_name, table_contents in secondary.items(): + if table_name not in merged: + merged[table_name] = { + **table_contents, + "columns": list(table_contents.get("columns", [])), + } + continue + + columns = list(merged[table_name].get("columns", [])) + for column in table_contents.get("columns", []): + if column not in columns: + columns.append(column) + merged[table_name]["columns"] = columns + + return merged + + +def _lexical_columns_and_tables_needed( + construct_db_schemas: list[dict], + query: str | None, + max_tables: int = 4, + max_columns_per_table: int = 12, +) -> dict[str, dict]: + if not query: + return {} + + query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) + if not query_tokens: + return {} + + scored_tables = [] + for table_schema in construct_db_schemas: + if table_schema.get("type") != "TABLE": + continue + + table_tokens = _tokenize_schema_text( + table_schema.get("name") + ) | _tokenize_schema_text( + table_schema.get("comment") + ) + table_score = len(query_tokens & table_tokens) * 6 + column_scores = [] + + for column in table_schema.get("columns", []): + if ( + column.get("type") != "COLUMN" + or column.get("data_type", "").lower() == "unknown" + ): + continue + + column_tokens = _tokenize_schema_text(column.get("name")) + comment_tokens = _tokenize_schema_text(column.get("comment")) + score = len(query_tokens & column_tokens) * 10 + score += len(query_tokens & comment_tokens) * 2 + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + score += len(column_tokens & _SALES_REVENUE_TOKENS) * 6 + if query_tokens & {"month", "monthly", "year", "july", "date", "latest"}: + score += ( + len(column_tokens & {"date", "day", "month", "year", "time"}) + * 5 + ) + if query_tokens & {"top", "highest", "lowest", "bottom"}: + measure_tokens = { + "amount", + "count", + "cost", + "margin", + "quantity", + "score", + "value", + } + if query_tokens & {"rate", "ratio", "percent", "percentage"}: + measure_tokens.update({"rate", "ratio", "percent", "percentage"}) + score += len(column_tokens & measure_tokens) * 4 + if score > 0: + column_scores.append( + (score, column["name"], column.get("is_primary_key")) + ) + + if not column_scores and table_score <= 0: + continue + + if table_tokens & _GENERIC_TABLE_TOKENS and table_score <= 0: + table_score -= 8 + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + if not query_tokens & _CUSTOMS_FINANCE_TOKENS: + table_score -= len(table_tokens & _CUSTOMS_FINANCE_TOKENS) * 10 + table_score += len(table_tokens & _SALES_REVENUE_TOKENS) * 5 + + total_score = table_score + sum(score for score, _, _ in column_scores) + if total_score <= 0: + continue + + selected_columns = [] + for _, column_name, _ in sorted( + column_scores, + key=lambda item: (-item[0], item[1]), + ): + if column_name not in selected_columns: + selected_columns.append(column_name) + if len(selected_columns) >= max_columns_per_table: + break + for _, column_name, is_primary_key in column_scores: + if is_primary_key and column_name not in selected_columns: + selected_columns.append(column_name) + + scored_tables.append((total_score, table_schema["name"], selected_columns)) + + scored_tables.sort(key=lambda item: (-item[0], item[1])) + return { + table_name: {"columns": columns} + for _, table_name, columns in scored_tables[:max_tables] + if columns + } + + ## End of Pipeline class MatchingTableContents(BaseModel): + model_config = ConfigDict(extra="forbid") + chain_of_thought_reasoning: list[str] columns: list[str] class MatchingTable(BaseModel): + model_config = ConfigDict(extra="forbid") + table_name: str table_contents: MatchingTableContents table_selection_reason: str class RetrievalResults(BaseModel): + model_config = ConfigDict(extra="forbid") + results: list[MatchingTable] RETRIEVAL_MODEL_KWARGS = { + "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { "name": "retrieval_schema", + "strict": True, "schema": RetrievalResults.model_json_schema(), }, } @@ -1165,7 +1770,11 @@ async def run( histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, ): - logger.info("Ask Retrieval pipeline is running...") + logger.info( + "Ask Retrieval pipeline is running for project_id=%s mdl_hash=%s", + project_id or "", + mdl_hash or "", + ) return await self._pipe.execute( ["construct_retrieval_results"], inputs={ diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py new file mode 100644 index 0000000000..6bacce4b68 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -0,0 +1,579 @@ +import asyncio + +from src.pipelines.generation.utils.sql import ( + SQLGenPostProcessor, + generate_simple_analytics_sql, + normalize_wren_sql_dialect, + normalize_sql_with_schema_identifiers, + unsupported_schema_generation_result, + unsupported_schema_message, + validate_sql_against_contexts, + validate_sql_semantic_coverage, +) + +SCHEMA_CONTEXTS = [ + """ + CREATE TABLE valid_invoice_comments ( + invoice_id VARCHAR, + comment_id VARCHAR + ); + """, + """ + CREATE TABLE "valid-order-lines" ( + order_id VARCHAR, + line_amount DECIMAL + ); + """, +] + + +def test_schema_grounding_rejects_unretrieved_table_name(): + error = validate_sql_against_contexts( + "SELECT invoice_id, COUNT(comment_id) FROM comments GROUP BY invoice_id", + SCHEMA_CONTEXTS, + ) + + assert error is not None + assert "comments" in error + assert "valid_invoice_comments" in error + + +def test_schema_grounding_accepts_retrieved_table_name(): + error = validate_sql_against_contexts( + """ + SELECT invoice_id, COUNT(comment_id) + FROM valid_invoice_comments + GROUP BY invoice_id + """, + SCHEMA_CONTEXTS, + ) + + assert error is None + + +def test_schema_grounding_rejects_invalid_qualified_column(): + error = validate_sql_against_contexts( + """ + SELECT c.invoice_number + FROM valid_invoice_comments c + """, + SCHEMA_CONTEXTS, + ) + + assert error is not None + assert "c.invoice_number" in error + + +def test_schema_identifier_normalization_quotes_special_identifiers(): + sql = normalize_sql_with_schema_identifiers( + "SELECT order_id FROM [valid-order-lines]", + SCHEMA_CONTEXTS, + ) + + assert 'FROM "valid-order-lines"' in sql + + +def test_semantic_coverage_rejects_generic_table_for_business_concepts(): + contexts = [ + """ + CREATE TABLE dbo_mbrTime ( + id1 INTEGER, + id2 INTEGER + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT id1, COUNT(*) AS failures + FROM dbo_mbrTime + GROUP BY id1 + ORDER BY failures DESC + LIMIT 10 + """, + "Show the top 10 materials with the highest number of failures.", + contexts, + ) + + assert error is not None + assert "failure/defect" in error + assert "material" in error + + +def test_unsupported_schema_message_requires_all_requested_concepts(): + contexts = [ + """ + CREATE TABLE dbo_mbrTime ( + id1 INTEGER, + id2 INTEGER + ); + """ + ] + + message = unsupported_schema_message( + "Show the top 10 materials with the highest number of failures.", + contexts, + ) + + assert message is not None + assert "No retrieved table or view" in message + assert "failure/defect" in message + assert "material" in message + + +def test_unsupported_schema_message_rejects_split_failure_technician_without_coverage(): + contexts = [ + """ + CREATE TABLE dbo_report_failures ( + id INTEGER, + failure_type VARCHAR + ); + """, + """ + CREATE TABLE dbo_technicians ( + id INTEGER, + name VARCHAR + ); + """, + ] + + message = unsupported_schema_message( + "Show the number of failures by technician.", + contexts, + ) + + assert message is not None + assert "failure/defect" in message + assert "technician" in message + + +def test_unsupported_schema_generation_result_has_no_invalid_sql(): + contexts = [ + """ + CREATE TABLE dbo_report_failures ( + id INTEGER, + failure_type VARCHAR + ); + """, + """ + CREATE TABLE dbo_technicians ( + id INTEGER, + name VARCHAR + ); + """, + ] + + result = unsupported_schema_generation_result( + "Show the number of failures by technician.", + contexts, + data_source="MSSQL", + ) + + assert result is not None + assert result["valid_generation_result"] == {} + invalid = result["invalid_generation_result"] + assert invalid["type"] == "NO_RELEVANT_SQL" + assert invalid["sql"] == "" + assert invalid["original_sql"] == "" + assert "technician" in invalid["error"] + + +def test_post_processor_clears_sql_for_unsupported_schema(): + contexts = [ + """ + CREATE TABLE dbo_mbrTime ( + id1 INTEGER, + id2 INTEGER + ); + """ + ] + post_processor = SQLGenPostProcessor(engine=None) + + result = asyncio.run( + post_processor.run( + [ + """ + SELECT id1, COUNT(*) AS failures + FROM dbo_mbrTime + GROUP BY id1 + ORDER BY failures DESC + LIMIT 10 + """ + ], + contexts=contexts, + fallback_query="Show the top 10 materials with the highest number of failures.", + data_source="MSSQL", + ) + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert result["invalid_generation_result"]["sql"] == "" + assert result["invalid_generation_result"]["original_sql"] == "" + + +def test_wren_sql_dialect_normalization_repairs_top_and_joined_limit(): + assert ( + normalize_wren_sql_dialect("SELECT TOP 10 id1 FROM dbo_mbrTime") + == "SELECT id1 FROM dbo_mbrTime\nLIMIT 10" + ) + assert ( + normalize_wren_sql_dialect( + "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESCLIMIT 10" + ) + == "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESC LIMIT 10" + ) + + +def test_repair_fallback_filters_critical_priority_and_in_progress_status(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMPTZ + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all critical-priority repairs that are currently in progress.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert "\"status\" = 'in progress'" in sql + assert "\"priority\" = 'critical'" in sql + + +def test_repair_fallback_preserves_hyphenated_in_progress_status_value(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all repairs with a critical priority and an in-progress status.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert "\"status\" = 'in-progress'" in sql + assert "\"priority\" = 'critical'" in sql + + +def test_repair_logs_highest_priority_orders_by_verified_priority_column(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Which repair logs have the highest priority?", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert 'ORDER BY CASE LOWER("priority")' in sql + assert "DESC" in sql + + +def test_critical_priority_repairs_filter_verified_priority_column(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all critical-priority repairs", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert "\"priority\" = 'critical'" in sql + + +def test_repairs_by_status_counts_verified_repair_rows(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show repairs by status", + contexts, + ) + + assert sql is not None + assert 'SELECT "status", COUNT("id") AS "record_count"' in sql + assert 'FROM "dbo_repair_logs"' in sql + assert 'GROUP BY "status"' in sql + + +def test_latest_repair_logs_orders_by_verified_date_column(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show latest repair logs", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert 'ORDER BY "created_at" DESC' in sql + + +def test_semantic_column_alias_can_satisfy_priority_concept_with_verified_name(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"repair log records"},"columns":[{"sql_column_name_use_exactly":"Urgency","data_type":"VARCHAR","semantic_context_not_sql_identifier":"priority severity for a repair"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE dbo_work_items ( + id VARCHAR, + Urgency VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Which repair records have the highest priority?", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_work_items"' in sql + assert '"Urgency"' in sql + assert '"priority"' not in sql + + +def test_failure_by_technician_fallback_uses_verified_tech_column(): + contexts = [ + """ + CREATE TABLE dbo_DebugEntries_Staging2 ( + Tech VARCHAR, + Failed VARCHAR, + Material VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the number of failures by technician.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_DebugEntries_Staging2"' in sql + assert 'SELECT "Tech", COUNT("Failed") AS "record_count"' in sql + assert 'WHERE ("Failed" IS NOT NULL AND "Failed" <> \'\')' in sql + + +def test_failure_by_material_fallback_uses_verified_material_column(): + contexts = [ + """ + CREATE TABLE dbo_DebugEntries_Staging2 ( + Tech VARCHAR, + Failed VARCHAR, + Material VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show failures by material.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_DebugEntries_Staging2"' in sql + assert 'SELECT "Material", COUNT("Failed") AS "record_count"' in sql + + +def test_failure_type_value_filter_uses_verified_failure_type_column(): + contexts = [ + """ + CREATE TABLE dbo_DebugEntries ( + SerialNumber VARCHAR, + FailedAt VARCHAR, + Material VARCHAR + ); + """, + """ + CREATE TABLE dbo_repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR + ); + """, + """ + CREATE TABLE dbo_report_failures ( + failure_type VARCHAR, + failure_line VARCHAR, + test_name VARCHAR + ); + """, + ] + + sql = generate_simple_analytics_sql( + "Show the number of units with JTAG as the failure type.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_report_failures"' in sql + assert 'COUNT(*) AS "record_count"' in sql + assert "\"failure_type\" = 'JTAG'" in sql + + +def test_board_models_most_failures_counts_failure_records_not_defect_rate(): + contexts = [ + """ + CREATE TABLE dbo_batch_records ( + board_model VARCHAR, + supplier VARCHAR, + defect_rate DECIMAL + ); + """, + """ + CREATE TABLE dbo_repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR + ); + """, + ] + + sql = generate_simple_analytics_sql( + "Show the top 5 board models with the most failures.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert 'SELECT "board_model", COUNT("failure_code") AS "record_count"' in sql + assert '"defect_rate"' not in sql + assert "LIMIT 5" in sql + + +def test_board_models_highest_defect_rate_uses_rate_metric(): + contexts = [ + """ + CREATE TABLE dbo_batch_records ( + board_model VARCHAR, + supplier VARCHAR, + defect_rate DECIMAL + ); + """, + """ + CREATE TABLE dbo_repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR + ); + """, + ] + + sql = generate_simple_analytics_sql( + "Show the board models with the highest defect rate.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_batch_records"' in sql + assert 'SELECT "board_model", AVG("defect_rate") AS "average_value"' in sql + assert 'ORDER BY "average_value" DESC' in sql + + +def test_semantic_coverage_rejects_rate_for_failure_count_intent(): + contexts = [ + """ + CREATE TABLE dbo_batch_records ( + board_model VARCHAR, + defect_rate DECIMAL + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT board_model, defect_rate + FROM dbo_batch_records + ORDER BY defect_rate DESC + LIMIT 5 + """, + "Show the top 5 board models with the most failures.", + contexts, + ) + + assert error is not None + assert "count of failure records" in error + + +def test_repairs_by_technician_requires_one_schema_object_covering_both_concepts(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + failure_code VARCHAR + ); + """, + """ + CREATE TABLE dbo_DebugEntries_Staging2 ( + Tech VARCHAR, + Failed VARCHAR + ); + """, + ] + + message = unsupported_schema_message("Show repairs by technician.", contexts) + + assert message is not None + assert "repair" in message + assert "technician" in message diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 90313db3f5..1bf9630b62 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -4,7 +4,10 @@ from src.pipelines.common import build_table_ddl from src.pipelines.retrieval.db_schema_retrieval import ( + _augment_retrieval_query, _build_view_ddl, + _parse_column_selection_response, + _rank_table_names_by_query, check_using_db_schemas_without_pruning, construct_db_schemas, construct_retrieval_results, @@ -1084,3 +1087,75 @@ def test_build_table_ddl_preserves_join_columns_when_pruned(): assert "parent_id INTEGER" in ddl assert "amount DOUBLE" in ddl assert "FOREIGN KEY (parent_id) REFERENCES parent(parent_id)" in ddl + + +def _schema_document(name: str, columns: list[str]) -> Document: + return Document( + content=str( + { + "name": name, + "type": "TABLE", + "columns": [ + {"name": column, "type": "COLUMN", "data_type": "VARCHAR"} + for column in columns + ], + } + ), + meta={"name": name, "type": "TABLE"}, + ) + + +def test_column_selection_accepts_alternate_results_shape(): + parsed = _parse_column_selection_response( + { + "replies": [ + """ + { + "tables": [ + { + "table_name": "SalesOrderFact", + "columns": ["USDFXSalesValue", "OrderDate"] + } + ] + } + """ + ] + } + ) + + assert parsed == { + "SalesOrderFact": { + "table_name": "SalesOrderFact", + "columns": ["USDFXSalesValue", "OrderDate"], + } + } + + +def test_column_selection_returns_empty_dict_for_malformed_reply(): + parsed = _parse_column_selection_response({"replies": ["not-json"]}) + + assert parsed == {} + + +def test_retrieval_query_augmentation_adds_business_terms(): + augmented = _augment_retrieval_query("show total revenue by year") + + assert "Business schema search terms" in augmented + assert "sales revenue amount value" in augmented + assert "date month year" in augmented + + +def test_table_ranking_prefers_business_sales_table_over_generic_or_customs_tables(): + documents = [ + _schema_document("dbo_mbrTime", ["id1", "id2"]), + _schema_document("CustomsRefundClaim", ["DutyAmount", "ClaimDate"]), + _schema_document("SalesOrderFact", ["USDFXSalesValue", "OrderDate"]), + ] + + ranked = _rank_table_names_by_query( + ["dbo_mbrTime", "CustomsRefundClaim", "SalesOrderFact"], + documents, + "show total revenue by year", + ) + + assert ranked[0] == "SalesOrderFact" diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 909a1bbdf8..c09ce4c472 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -43,6 +43,26 @@ const dirtyProjectIds = new Set(); const isSameId = (left: string | number, right: string | number) => String(left) === String(right); +const firstNonEmptyString = (...values: unknown[]) => + values.find( + (value): value is string => + typeof value === 'string' && value.trim().length > 0, + ) || ''; + +const normalizeModelDisplayName = (model: Model) => + firstNonEmptyString( + model.displayName, + model.referenceName, + model.sourceTableName, + ); + +const normalizeColumnDisplayName = (column: ModelColumn) => + firstNonEmptyString( + column.displayName, + column.referenceName, + column.sourceColumnName, + ); + export enum SyncStatusEnum { IN_PROGRESS = 'IN_PROGRESS', SYNCRONIZED = 'SYNCRONIZED', @@ -728,15 +748,27 @@ export class ModelResolver { .filter((c) => c.modelId === model.id) .map((c) => ({ ...c, + displayName: normalizeColumnDisplayName(c), properties: JSON.parse(c.properties), nestedColumns: c.type.includes('STRUCT') - ? modelNestedColumnList.filter((nc) => nc.columnId === c.id) + ? modelNestedColumnList + .filter((nc) => nc.columnId === c.id) + .map((nc) => ({ + ...nc, + displayName: firstNonEmptyString( + nc.displayName, + nc.referenceName, + nc.sourceColumnName, + nc.columnPath?.join('.'), + ), + })) : undefined, })); const fields = modelFields.filter((c) => !c.isCalculated); const calculatedFields = modelFields.filter((c) => c.isCalculated); result.push({ ...model, + displayName: normalizeModelDisplayName(model), fields, calculatedFields, properties: { @@ -763,9 +795,20 @@ export class ModelResolver { const columns = modelColumns.map((c) => ({ ...c, + displayName: normalizeColumnDisplayName(c), properties: JSON.parse(c.properties), nestedColumns: c.type.includes('STRUCT') - ? modelNestedColumns.filter((nc) => nc.columnId === c.id) + ? modelNestedColumns + .filter((nc) => nc.columnId === c.id) + .map((nc) => ({ + ...nc, + displayName: firstNonEmptyString( + nc.displayName, + nc.referenceName, + nc.sourceColumnName, + nc.columnPath?.join('.'), + ), + })) : undefined, })); const relations = ( @@ -780,6 +823,7 @@ export class ModelResolver { return { ...model, + displayName: normalizeModelDisplayName(model), fields: columns.filter((c) => !c.isCalculated), calculatedFields: columns.filter((c) => c.isCalculated), relations, diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 5572a39420..0ea641eac7 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -3,6 +3,7 @@ import { AskResultStatus, AskResultType, AskCandidateType, + type RecommendationQuestion, RecommendationQuestionsResult, RecommendationQuestionsInput, WrenAIError, @@ -677,10 +678,7 @@ export class AskingService implements IAskingService { const currentProject = await this.projectService.getCurrentProject(); let projectId = payload.projectId ?? currentProject.id; if (threadId) { - const thread = await this.threadRepository.findOneBy({ id: threadId }); - if (!thread) { - throw new Error(`Thread ${threadId} not found`); - } + const thread = await this.ensureThreadInCurrentProject(threadId); if (payload.projectId && payload.projectId !== thread.projectId) { throw new Error( `Thread ${threadId} does not belong to project ${payload.projectId}`, From aa3ee735a4788e491b9bdf7fb7e0b7fe12e25046 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 11:52:08 +0000 Subject: [PATCH 1070/1087] Update Ask grounding handoff with PR status --- WRENAI_LOCAL_ASK_HANDOFF.md | 42 ++++++++++++++++++++++++++++++++----- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index e9f23efcda..548c0c477b 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -26,6 +26,33 @@ Current projects visible through `/api/v1/projects/current`: - id `12`, `CWPay` - id `13`, `CW_GL` +## Source Control / PR Status + +The work is pushed to the fork branch: + +- Repository: `hbalasubramanya-rgb/WrenAI` +- Branch: `organization/ask-schema-grounding-20260820` +- PR: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1` +- PR base: `organization-feature` +- Current remote PR head: `cc55d1e05` + +The PR branch was rebased onto the latest `origin/organization-feature` after GitHub initially reported conflicts against the wrong compare/base. It was then pushed with `--force-with-lease`. + +GitHub readback after the rebase: + +- `mergeable=True` +- `mergeable_state=unstable` + +`unstable` means GitHub checks are pending or failing; it is not a merge-conflict state. + +If the previous/old branch view is gone or stale, use this branch and PR instead: + +- Use branch `organization/ask-schema-grounding-20260820` for this work. +- Review and merge PR #1 into `organization-feature`. +- After merge, use `organization-feature` as the updated canonical branch. + +Do not open this work against upstream `Canner/WrenAI:main` unless that is explicitly intended; this branch was prepared for the fork's `organization-feature` base. + ## What Changed Today ### Generic Schema Grounding @@ -252,6 +279,8 @@ Invoke-WebRequest -UseBasicParsing http://127.0.0.1:3000/api/v1/projects/current ## Current Dirty Files To Review +In the pushed PR branch, the source changes below are committed. The original local checkout at `D:\WrenAI` may still show unrelated dirty runtime/data files and may also show the old pre-rebase local commit until it is refreshed from origin. + Relevant tracked files: - `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` @@ -276,8 +305,11 @@ There are also many local untracked runtime/data artifacts in the repository. Do ## Recommended Next Steps -1. Install or enable pytest in `wren-ai-service\venv`, then run focused tests. -2. Review the large `utils/sql.py` diff carefully; consider extracting fallback/grounding helpers into smaller modules after behavior is stable. -3. Add sample-value metadata to retrieval context if available, then make value matching use that metadata instead of only text normalization. -4. Run a broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when their data sources are available. -5. Commit the source changes after review, excluding local runtime/data artifacts. +1. Review PR #1: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1`. +2. Confirm the PR base is `organization-feature`, not `Canner/WrenAI:main`. +3. Resolve any GitHub check failures if `mergeable_state` remains `unstable`, then merge PR #1. +4. After merge, continue from `organization-feature`. +5. Install or enable pytest in `wren-ai-service\venv`, then run focused tests. +6. Review the large `utils/sql.py` diff carefully; consider extracting fallback/grounding helpers into smaller modules after behavior is stable. +7. Add sample-value metadata to retrieval context if available, then make value matching use that metadata instead of only text normalization. +8. Run a broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when their data sources are available. From 3e8d3c6519a9fab112d4bda0b7aa3e8e7e4b6831 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 11:53:19 +0000 Subject: [PATCH 1071/1087] Clarify PR head in Ask grounding handoff --- WRENAI_LOCAL_ASK_HANDOFF.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index 548c0c477b..846c5f0873 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -34,7 +34,8 @@ The work is pushed to the fork branch: - Branch: `organization/ask-schema-grounding-20260820` - PR: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1` - PR base: `organization-feature` -- Current remote PR head: `cc55d1e05` +- Schema implementation commit before handoff-only updates: `cc55d1e05` +- Check PR #1 for the live head SHA because handoff-only commits may be added after the implementation commit. The PR branch was rebased onto the latest `origin/organization-feature` after GitHub initially reported conflicts against the wrong compare/base. It was then pushed with `--force-with-lease`. From 4a199fca138c0afd9ba3dbb15a182e9c0f88c2d8 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 12:30:03 +0000 Subject: [PATCH 1072/1087] Broaden Ask semantic grounding coverage --- .../src/pipelines/generation/utils/sql.py | 638 +++++++++++++++++- .../retrieval/db_schema_retrieval.py | 20 +- .../generation/test_sql_schema_grounding.py | 163 ++++- 3 files changed, 783 insertions(+), 38 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index bcd07c0308..a076963ac4 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -110,6 +110,9 @@ "COUNT", "CREATE", "CROSS", + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", "DELETE", "DESC", "DISTINCT", @@ -127,6 +130,7 @@ "ISNULL", "INNER", "INSERT", + "INTERVAL", "INTERSECT", "IS", "JOIN", @@ -166,6 +170,9 @@ "CONCAT", "COUNT", "COUNT_BIG", + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", "DATE_TRUNC", "DAY", "EXTRACT", @@ -254,12 +261,40 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: aliases = { "bu": {"business", "unit"}, + "acct": {"account"}, + "accounting": {"account", "ledger"}, + "accounts": {"account"}, + "address": {"email", "mail"}, + "addresses": {"address", "email", "mail"}, + "approval": {"approver", "reviewer", "signer", "status"}, + "approvals": {"approval", "approver", "reviewer", "signer", "status"}, + "approver": {"approval", "reviewer", "signer"}, + "approvers": {"approval", "approver", "reviewer", "signer"}, + "balances": {"balance"}, "cust": {"customer"}, "customers": {"customer"}, "critical": {"priority", "severity"}, + "curr": {"currency"}, "boards": {"board"}, + "email": {"address", "mail"}, + "email1": {"address", "email", "first", "mail", "primary"}, + "emails": {"address", "email", "mail"}, + "endbalance": {"balance", "end", "ending"}, + "ending": {"end"}, + "gl": {"account", "ledger"}, + "glaccount": {"account", "gl", "ledger"}, + "glaccounts": {"account", "gl", "ledger"}, + "gross": {"amount", "value"}, + "grossamount": {"amount", "gross", "value"}, "high": {"priority", "severity"}, "highest": {"top"}, + "invoicedate": {"date", "invoice"}, + "invoicemonth": {"invoice", "month"}, + "invoicenumber": {"invoice", "number"}, + "invoiceyear": {"invoice", "year"}, + "journals": {"journal"}, + "journalid": {"id", "journal"}, + "journalnumber": {"journal", "number"}, "logs": {"log", "record"}, "log": {"record"}, "lows": {"low"}, @@ -267,8 +302,12 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "models": {"model"}, "inv": {"invoice"}, "invoices": {"invoice"}, + "net": {"amount", "value"}, + "netamount": {"amount", "net", "value"}, "ord": {"order"}, "orders": {"order"}, + "preparergroup": {"group", "preparer"}, + "preparers": {"preparer"}, "qty": {"quantity"}, "num": {"number"}, "no": {"number"}, @@ -278,17 +317,32 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "priorities": {"priority", "severity"}, "priority": {"severity"}, "recent": {"latest"}, + "recon": {"reconciliation", "reconcile"}, + "reconciliations": {"reconciliation", "reconcile", "recon"}, + "reconciliation": {"reconcile", "recon"}, "records": {"record"}, "rep": {"representative", "salesperson"}, "repairs": {"repair"}, + "reviewer": {"approval", "approver", "signer"}, + "reviewers": {"approval", "approver", "reviewer", "signer"}, "salesperson": {"sales", "person"}, "severity": {"priority"}, + "signer": {"approval", "approver", "reviewer"}, + "signers": {"approval", "approver", "reviewer", "signer"}, "supplier": {"vendor"}, + "supplierid": {"id", "supplier", "vendor"}, + "suppliername": {"name", "supplier", "vendor"}, "suppliers": {"supplier", "vendor"}, + "taskstatus": {"status", "task"}, + "tasks": {"status", "task"}, "tech": {"technician"}, "technician": {"tech"}, + "transid": {"id", "trans", "transaction"}, "vendor": {"supplier"}, "vendors": {"supplier", "vendor"}, + "workflow": {"approval", "status"}, + "workflows": {"approval", "status", "workflow"}, + "counts": {"count"}, "failed": {"failure"}, "failures": {"failure"}, "defects": {"defect"}, @@ -1454,7 +1508,57 @@ def _expanded_fallback_query_tokens(query: str) -> set[str]: if tokens & {"order", "orders"}: tokens.update({"amount", "customer", "date", "ord", "order", "value"}) if tokens & {"invoice", "invoices"}: - tokens.update({"amount", "currency", "date", "invoice", "supplier"}) + tokens.update( + { + "amount", + "currency", + "date", + "gross", + "invoice", + "month", + "net", + "number", + "status", + "supplier", + "task", + "year", + } + ) + if tokens & {"supplier", "vendor"}: + tokens.update({"email", "id", "name", "number", "supplier", "vendor"}) + if tokens & {"email", "address"}: + tokens.update({"address", "email", "first", "mail", "primary"}) + if tokens & {"reconciliation", "recon", "reconcile"}: + tokens.update( + { + "account", + "gl", + "group", + "period", + "preparer", + "recon", + "reconciliation", + "reviewer", + "status", + } + ) + if tokens & {"journal", "workflow", "approval", "approver", "reviewer", "signer"}: + tokens.update( + { + "approval", + "approver", + "date", + "journal", + "reviewer", + "signer", + "status", + "workflow", + } + ) + if tokens & {"account", "accounts", "gl", "glaccount", "ledger"}: + tokens.update({"account", "balance", "gl", "glaccount", "ledger", "period", "year"}) + if tokens & {"balance", "balances"}: + tokens.update({"amount", "balance", "end", "ending", "value", "year"}) if tokens & {"batch", "batches"}: tokens.update({"batch", "board", "defect", "inspection", "rate", "supplier"}) if tokens & {"repair", "repairs"}: @@ -1584,6 +1688,29 @@ def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[ ), ("status", {"status"}, {"status"}), ("order", {"order"}, {"order", "ord"}), + ("invoice", {"invoice"}, {"invoice", "inv"}), + ("email/address", {"email", "address"}, {"email", "email1", "address", "mail"}), + ( + "account", + {"account", "gl", "glaccount", "ledger"}, + {"account", "gl", "glaccount", "ledger"}, + ), + ("balance", {"balance"}, {"balance", "endbalance"}), + ( + "reconciliation", + {"reconciliation", "recon", "reconcile"}, + {"reconciliation", "recon", "reconcile"}, + ), + ( + "journal/workflow", + {"journal", "workflow"}, + {"journal", "workflow"}, + ), + ( + "approval", + {"approval", "approver", "reviewer", "signer"}, + {"approval", "approver", "reviewer", "signer"}, + ), ] for label, triggers, schema_tokens in specs: if query_tokens & triggers: @@ -1629,12 +1756,22 @@ def _choose_fallback_table( for table_name, columns in schema_details.items(): table_tokens = _table_business_tokens(table_name, columns) column_token_union = set() + has_numeric_amount_measure = False has_numeric_sales_measure = False has_date_capable_column = False score = len(query_tokens & table_tokens) * 8 for column in columns: column_tokens = _column_business_tokens(column) column_token_union.update(column_tokens) + if _is_numeric_type(column["data_type"]) and column_tokens & { + "amount", + "balance", + "cost", + "gross", + "net", + "value", + }: + has_numeric_amount_measure = True if _is_numeric_type(column["data_type"]) and column_tokens & { "amount", "intake", @@ -1658,9 +1795,12 @@ def _choose_fallback_table( & column_tokens & { "amount", + "balance", "cost", "count", + "gross", "margin", + "net", "quantity", "rate", "score", @@ -1793,6 +1933,86 @@ def _choose_fallback_table( else: score -= 60 + if concept_tokens & {"invoice", "invoices"}: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"invoice", "inv"}: + continue + score += 80 + if query_tokens & {"supplier", "vendor"}: + if not table_and_columns & {"supplier", "vendor"}: + continue + score += 45 + if concept_tokens & {"status", "task"}: + if not table_and_columns & {"status", "state", "task"}: + continue + score += 40 + if concept_tokens & {"gross", "net", "amount", "value", "top", "highest"}: + if not has_numeric_amount_measure: + continue + score += 55 + if concept_tokens & {"date", "month", "monthly", "year", "latest", "recent"}: + if not has_date_capable_column: + continue + score += 35 + + if concept_tokens & {"supplier", "vendor"}: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"supplier", "vendor"}: + if query_tokens & {"email", "address", "missing", "blank", "empty", "null"}: + continue + score -= 45 + else: + score += 35 + + if concept_tokens & {"email", "address"}: + if not column_token_union & {"email", "email1", "address", "mail"}: + continue + score += 65 + + if concept_tokens & {"reconciliation", "recon", "reconcile"}: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"reconciliation", "recon", "reconcile"}: + continue + score += 90 + if query_tokens & {"status"}: + if not table_and_columns & {"status", "state"}: + continue + score += 45 + if query_tokens & {"preparer", "group"}: + if not table_and_columns & {"preparer", "group"}: + continue + score += 45 + + if concept_tokens & {"journal", "workflow"}: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"journal", "workflow"}: + continue + if "journal" in concept_tokens and "journal" not in table_and_columns: + continue + if "workflow" in concept_tokens and "workflow" not in table_and_columns: + continue + score += 85 + if concept_tokens & {"approval", "approver", "reviewer", "signer"}: + if not table_and_columns & {"approval", "approver", "reviewer", "signer"}: + continue + score += 45 + if query_tokens & {"latest", "recent"}: + if not has_date_capable_column: + continue + score += 35 + + if concept_tokens & {"account", "gl", "glaccount", "ledger"}: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"account", "gl", "glaccount", "ledger"}: + continue + score += 45 + if query_tokens & {"balance", "end", "ending"}: + if not table_and_columns & {"balance", "endbalance"}: + continue + if not has_numeric_amount_measure: + continue + score += 70 + if query_tokens & {"batch", "batches"} and {"defect", "rate"}.issubset( column_token_union ): @@ -1911,11 +2131,22 @@ def _choose_dimension_column( query_tokens: set[str], columns: list[dict[str, str]], ) -> str | None: + columns = _choose_dimension_columns(query_tokens, columns, max_columns=1) + return columns[0] if columns else None + + +def _choose_dimension_columns( + query_tokens: set[str], + columns: list[dict[str, str]], + max_columns: int = 3, +) -> list[str]: dimension_specs = [ ({"board", "model"}, {"board", "model"}), ({"business", "unit"}, {"business", "unit", "bu", "division", "company"}), ({"customer"}, {"customer", "cust", "name"}), ({"supplier"}, {"supplier", "vendor", "name"}), + ({"email", "address"}, {"email", "email1", "address", "mail"}), + ({"invoice"}, {"invoice", "inv", "number", "no", "id"}), ({"product"}, {"product", "prod", "item", "material", "name"}), ({"salesperson", "representative"}, {"salesperson", "sales", "person", "rep"}), ({"technician", "tech"}, {"technician", "tech"}), @@ -1923,17 +2154,24 @@ def _choose_dimension_column( ({"material"}, {"material", "part", "item"}), ({"priority", "severity"}, {"priority", "severity", "urgency", "rank"}), ({"status"}, {"status"}), + ({"preparer"}, {"preparer", "group"}), + ({"reviewer"}, {"reviewer", "group"}), + ({"approval", "approver", "signer"}, {"approval", "approver", "reviewer", "signer"}), + ({"account", "gl", "glaccount", "ledger"}, {"account", "gl", "glaccount", "ledger"}), ({"currency"}, {"currency", "curr"}), ({"country"}, {"country"}), ({"order"}, {"order", "ord", "number"}), ({"batch"}, {"batch", "id"}), ] + selected: list[str] = [] for trigger_tokens, column_tokens in dimension_specs: if query_tokens & trigger_tokens: column = _choose_ranked_column_by_tokens(columns, column_tokens) - if column: - return column["name"] - return None + if column and column["name"] not in selected: + selected.append(column["name"]) + if len(selected) >= max_columns: + break + return selected def _choose_missing_value_column( @@ -1941,6 +2179,7 @@ def _choose_missing_value_column( columns: list[dict[str, str]], ) -> dict[str, str] | None: missing_specs = [ + ({"email", "address"}, {"email", "email1", "address", "mail", "first", "primary"}), ({"customer"}, {"customer", "cust", "number", "id", "no"}), ({"order"}, {"order", "ord", "number", "id", "no"}), ({"supplier"}, {"supplier", "vendor", "number", "id", "no"}), @@ -1962,7 +2201,11 @@ def _choose_count_subject_column( ) -> dict[str, str] | None: subject_specs = [ ({"order"}, {"order", "ord", "number", "id", "no"}), + ({"invoice"}, {"invoice", "inv", "number", "id", "no"}), ({"customer"}, {"customer", "cust", "number", "id", "no"}), + ({"reconciliation", "recon"}, {"reconciliation", "recon", "trans", "id"}), + ({"journal"}, {"journal", "entry", "number", "id"}), + ({"account", "gl", "glaccount"}, {"account", "gl", "glaccount", "id"}), ( {"failure"}, {"failure", "failed", "defect", "code", "line", "status", "sys", "type"}, @@ -2027,6 +2270,43 @@ def _select_listing_columns( if query_tokens & {"order"}: score += len(tokens & {"order", "ord", "number", "customer", "cust", "item", "product"}) * 18 score += len(tokens & {"date", "day", "month", "year"}) * 8 + if query_tokens & {"invoice"}: + score += ( + len( + tokens + & { + "amount", + "currency", + "date", + "gross", + "invoice", + "net", + "number", + "status", + "supplier", + "task", + } + ) + * 18 + ) + if query_tokens & {"supplier", "vendor", "email", "address"}: + score += ( + len( + tokens + & { + "address", + "email", + "email1", + "id", + "mail", + "name", + "number", + "supplier", + "vendor", + } + ) + * 18 + ) if query_tokens & {"customer"}: score += len(tokens & {"customer", "cust", "name", "number", "id"}) * 18 if query_tokens & {"product", "material"}: @@ -2038,6 +2318,49 @@ def _select_listing_columns( len(tokens & {"board", "code", "date", "failure", "id", "priority", "status"}) * 14 ) + if query_tokens & {"journal", "workflow", "approval", "approver", "reviewer", "signer"}: + score += ( + len( + tokens + & { + "approval", + "approver", + "date", + "doc", + "entry", + "journal", + "number", + "reviewer", + "signer", + "status", + "workflow", + } + ) + * 16 + ) + if query_tokens & {"reconciliation", "recon", "reconcile"}: + score += ( + len( + tokens + & { + "account", + "gl", + "glaccount", + "group", + "period", + "preparer", + "recon", + "reviewer", + "status", + } + ) + * 16 + ) + if query_tokens & {"account", "gl", "glaccount", "balance"}: + score += ( + len(tokens & {"account", "balance", "endbalance", "gl", "glaccount", "month", "year"}) + * 16 + ) if tokens & {"repair", "failure", "failed", "defect"} and not query_tokens & { "repair", "failure", @@ -2076,10 +2399,71 @@ def _fallback_month_filter(query: str) -> tuple[int, int] | None: return None +def _grouping_phrase_tokens(query: str) -> set[str]: + match = re.search( + r"(?is)\b(?:grouped\s+by|group\s+by|by)\s+(?P[A-Za-z0-9_ /-]+)", + query, + ) + if not match: + return set() + return _fallback_tokens(match.group("value")) + + +def _current_year_where_clause( + date_column: str | None, + columns: list[dict[str, str]], + query_tokens: set[str], +) -> str: + if not date_column or not {"this", "year"}.issubset(query_tokens): + return "" + + column = next((column for column in columns if column["name"] == date_column), None) + if not column: + return "" + + current_year = datetime.now(timezone.utc).year + quoted_column = _quote_identifier(date_column) + if _is_numeric_type(column["data_type"]) or _fallback_tokens(date_column) & {"year"}: + return f"\nWHERE {quoted_column} = {current_year}" + if _is_date_type(column["data_type"]) or _fallback_tokens(date_column) & { + "date", + "day", + "month", + "time", + }: + return f"\nWHERE EXTRACT(YEAR FROM {quoted_column}) = {current_year}" + return "" + + def _quote_literal(value: str) -> str: return "'" + value.replace("'", "''") + "'" +def _value_match_predicate( + column: dict[str, str] | str, + value: str, + alternate_values: list[str] | None = None, +) -> str: + column_name = column["name"] if isinstance(column, dict) else column + quoted_column = _quote_identifier(column_name) + values = [] + for candidate in [value] + (alternate_values or []): + cleaned = _clean_filter_value(candidate) + if cleaned and cleaned.lower() not in {item.lower() for item in values}: + values.append(cleaned) + + if not values: + return f"{quoted_column} IS NOT NULL" + + if isinstance(column, dict) and _is_text_type(column["data_type"]): + lowered_values = [_quote_literal(candidate.lower()) for candidate in values] + if len(lowered_values) == 1: + return f"LOWER({quoted_column}) = {lowered_values[0]}" + return f"LOWER({quoted_column}) IN ({', '.join(lowered_values)})" + + return f"{quoted_column} = {_quote_literal(values[0])}" + + def _clean_filter_value(value: str | None) -> str | None: if value is None: return None @@ -2092,6 +2476,7 @@ def _extract_failure_type_filter_value(query: str) -> str | None: ( r"(?is)\bwith\s+" r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80}?)" + r"(?:\s+(?:listed|marked|recorded|shown|set))?" r"\s+as\s+(?:the\s+)?(?:failure|defect)\s+" r"(?:type|code|category)\b" ), @@ -2110,15 +2495,28 @@ def _extract_failure_type_filter_value(query: str) -> str | None: return None -def _extract_status_filter_value(query: str) -> str | None: +def _extract_status_filter_values(query: str) -> list[str]: + values: list[str] = [] + + def add(value: str): + if value not in values: + values.append(value) + if re.search(r"(?i)\bin-progress\b", query): - return "in-progress" + add("in-progress") + add("in progress") if re.search(r"(?i)\bin\s+progress\b", query): - return "in progress" + add("in progress") + add("in-progress") for status in ("completed", "pending", "escalated", "open", "closed"): if re.search(rf"(?i)\b{re.escape(status)}\b", query): - return status - return None + add(status) + return values + + +def _extract_status_filter_value(query: str) -> str | None: + values = _extract_status_filter_values(query) + return values[0] if values else None def _extract_priority_filter_value(query: str) -> str | None: @@ -2208,16 +2606,32 @@ def generate_simple_analytics_sql( if not query_tokens & { "batch", "batches", + "account", + "accounts", + "address", + "approval", + "approvals", + "approver", + "balance", + "balances", "board", "business", "count", "customer", "defect", + "email", "failure", "failures", + "gl", + "glaccount", + "gross", "highest", + "invoice", + "invoices", + "journal", "july", "latest", + "ledger", "log", "logs", "location", @@ -2226,13 +2640,19 @@ def generate_simple_analytics_sql( "model", "monthly", "most", + "net", "number", "order", "orders", + "preparer", "priority", "product", "rate", "recent", + "recon", + "reconciliation", + "reconciliations", + "reviewer", "record", "records", "repair", @@ -2241,8 +2661,11 @@ def generate_simple_analytics_sql( "sale", "sales", "severity", + "signer", "supplier", "status", + "task", + "tasks", "tech", "technician", "top", @@ -2251,6 +2674,8 @@ def generate_simple_analytics_sql( "type", "unit", "units", + "workflow", + "workflows", "year", }: return None @@ -2307,8 +2732,7 @@ def generate_simple_analytics_sql( return ( f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" f"FROM {quoted_table}\n" - f"WHERE {_quote_identifier(failure_type_filter_column['name'])} = " - f"{_quote_literal(failure_type_filter_value)}" + f"WHERE {_value_match_predicate(failure_type_filter_column, failure_type_filter_value)}" ) material_column = _choose_column_by_tokens(columns, {"material"}) @@ -2323,7 +2747,7 @@ def generate_simple_analytics_sql( f"FROM {quoted_table}" ) - status_column = _choose_column_by_tokens(columns, {"status"}) + status_column = _choose_ranked_column_by_tokens(columns, {"status"}) repair_filter_intent = raw_query_tokens & { "closed", "completed", @@ -2343,25 +2767,29 @@ def generate_simple_analytics_sql( } if query_tokens & {"repair", "repairs"} and repair_filter_intent: predicates = [] - status_filter_value = _extract_status_filter_value(query) - if status_column and status_filter_value: + status_filter_values = _extract_status_filter_values(query) + if status_column and status_filter_values: predicates.append( - f"{_quote_identifier(status_column)} = {_quote_literal(status_filter_value)}" + _value_match_predicate( + status_column, + status_filter_values[0], + status_filter_values[1:], + ) ) priority_filter_value = _extract_priority_filter_value(query) if priority_filter_value: priority_column = _choose_priority_column(columns) if priority_column: - predicates.append( - f"{_quote_identifier(priority_column['name'])} = " - f"{_quote_literal(priority_filter_value)}" - ) + predicates.append(_value_match_predicate(priority_column, priority_filter_value)) if predicates: return f"SELECT *\nFROM {quoted_table}\nWHERE {' AND '.join(predicates)}" + date_column_tokens = {"date", "day", "month", "time", "year"} + if raw_query_tokens & {"year"} and not raw_query_tokens & {"month", "monthly"}: + date_column_tokens = {"date", "day", "time", "year"} date_column = _choose_column_by_tokens( columns, - {"date", "day", "month", "time", "year"}, + date_column_tokens, date=True, ) @@ -2437,12 +2865,53 @@ def generate_simple_analytics_sql( {"amount", "intake", "sales", "value"}, numeric=True, ) + if not measure_column and query_tokens & {"invoice", "invoices"} and raw_query_tokens & { + "amount", + "bottom", + "gross", + "highest", + "lowest", + "net", + "top", + "total", + "value", + }: + if query_tokens & {"gross"}: + measure_column = _choose_column_by_tokens( + columns, + {"gross", "amount", "value"}, + numeric=True, + ) + elif query_tokens & {"net"}: + measure_column = _choose_column_by_tokens( + columns, + {"net", "amount", "value"}, + numeric=True, + ) + if not measure_column: + measure_column = _choose_column_by_tokens( + columns, + {"amount", "gross", "net", "value"}, + numeric=True, + ) if not measure_column and query_tokens & {"revenue", "sale", "sales"}: measure_column = _choose_column_by_tokens( columns, {"amount", "intake", "revenue", "sales", "value"}, numeric=True, ) + if not measure_column and query_tokens & {"balance", "balances"}: + measure_column = _choose_column_by_tokens( + columns, + {"balance", "end", "ending", "value"}, + numeric=True, + ) + if not measure_column and query_tokens & {"amount", "gross", "net", "value"}: + measure_column = _choose_column_by_tokens( + columns, + {"amount", "gross", "net", "value"}, + numeric=True, + ) if ( not measure_column and not failure_count_intent @@ -2450,7 +2919,19 @@ def generate_simple_analytics_sql( ): measure_column = _choose_column_by_tokens( columns, - {"amount", "cost", "count", "margin", "quantity", "rate", "score", "value"}, + { + "amount", + "balance", + "cost", + "count", + "gross", + "margin", + "net", + "quantity", + "rate", + "score", + "value", + }, numeric=True, ) @@ -2488,17 +2969,82 @@ def generate_simple_analytics_sql( f"WHERE {_missing_value_predicate(missing_column)}{limit_clause}" ) + explicit_grouping_intent = bool(raw_query_tokens & {"group", "grouped"}) or bool( + re.search(r"(?i)\bby\s+[A-Za-z0-9_ -]+\b", query) + ) + explicit_measure_intent = bool( + raw_query_tokens + & { + "amount", + "balance", + "gross", + "margin", + "net", + "rate", + "revenue", + "sale", + "sales", + "score", + "value", + } + ) + if ( + explicit_grouping_intent + and not explicit_measure_intent + and not ( + raw_query_tokens & {"month", "monthly"} + and raw_query_tokens & {"count", "number"} + ) + ): + grouping_tokens = _grouping_phrase_tokens(query) or raw_query_tokens + dimension_columns = _choose_dimension_columns(grouping_tokens, columns) + if dimension_columns: + subject_column = _choose_count_subject_column(raw_query_tokens, columns) + count_expression = "COUNT(*)" + where_clause = "" + if subject_column: + count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" + where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" + quoted_dimensions = _quote_joined(dimension_columns) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {quoted_dimensions}, {count_expression} AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" + ) + + if date_column and raw_query_tokens & {"month", "monthly"} and raw_query_tokens & { + "count", + "number", + }: + year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" + month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" + subject_column = _choose_count_subject_column(raw_query_tokens | query_tokens, columns) + count_expression = "COUNT(*)" + where_clause = "" + if subject_column: + count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" + where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"{count_expression} AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\n" + f"GROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" + ) + implied_count_by_dimension = "failure" in query_tokens and bool( raw_query_tokens & {"location", "material", "technician", "tech"} or (board_model_intent and not rate_metric_intent) ) if ( - query_tokens & {"count", "number"} + raw_query_tokens & {"count", "number"} or failure_count_intent or implied_count_by_dimension ): - dimension_column = _choose_dimension_column(raw_query_tokens, columns) - if dimension_column: + dimension_columns = _choose_dimension_columns(raw_query_tokens, columns) + if dimension_columns: subject_column = _choose_count_subject_column(raw_query_tokens, columns) count_expression = "COUNT(*)" where_clause = "" @@ -2506,14 +3052,25 @@ def generate_simple_analytics_sql( count_expression = ( f"COUNT(DISTINCT {_quote_identifier(subject_column['name'])})" ) - elif subject_column and raw_query_tokens & {"failure", "repair", "batch"}: + elif subject_column and raw_query_tokens & { + "account", + "batch", + "failure", + "gl", + "glaccount", + "invoice", + "journal", + "recon", + "reconciliation", + "repair", + }: count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" - quoted_dimension = _quote_identifier(dimension_column) + quoted_dimensions = _quote_joined(dimension_columns) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {quoted_dimension}, {count_expression} AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" + f"SELECT {quoted_dimensions}, {count_expression} AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" ) @@ -2530,6 +3087,25 @@ def generate_simple_analytics_sql( f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" ) + if measure_column and dimension_column and raw_query_tokens & { + "bottom", + "highest", + "lowest", + "top", + }: + aggregate, alias = _aggregate_for_measure(measure_column) + quoted_dimension = _quote_identifier(dimension_column) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" + direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" + limit_clause = f"\nLIMIT {limit}" if limit else "" + where_clause = _current_year_where_clause(date_column, columns, raw_query_tokens) + return ( + f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{where_clause}\n" + f"GROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + if measure_column and limit and dimension_column and raw_query_tokens & { "board", "business", @@ -2580,8 +3156,8 @@ def generate_simple_analytics_sql( if ( measure_column and date_column - and query_tokens & {"year"} - and not query_tokens & {"month", "monthly"} + and raw_query_tokens & {"year"} + and not raw_query_tokens & {"month", "monthly"} ): date_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" return ( @@ -2593,7 +3169,7 @@ def generate_simple_analytics_sql( if ( measure_column and date_column - and query_tokens & {"month", "monthly", "trend", "trends"} + and raw_query_tokens & {"month", "monthly", "trend", "trends"} ): year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index a8d5ac5f27..7c375608f3 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -911,7 +911,25 @@ def _augment_retrieval_query(query: str) -> str: "order ord number date customer product business unit division company" ), ("invoice", "invoices"): ( - "invoice supplier customer currency amount date number" + "invoice supplier customer currency gross net amount date month year number status task" + ), + ("supplier", "suppliers", "vendor", "vendors"): ( + "supplier vendor name number id email address invoice amount" + ), + ("email", "emails", "address", "addresses"): ( + "email address mail first primary supplier contact" + ), + ("reconciliation", "reconciliations", "recon", "reconcile"): ( + "reconciliation recon account gl status preparer reviewer group period" + ), + ("journal", "journals", "workflow", "approval", "approvals"): ( + "journal workflow approval approver reviewer signer status date posting entry document" + ), + ("account", "accounts", "gl", "ledger"): ( + "account gl glaccount ledger balance endbalance ending period year month" + ), + ("balance", "balances", "gross", "net"): ( + "balance endbalance gross net amount value year month" ), ("customer", "customers"): ( "customer account client number name identifier" diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index 6bacce4b68..205c8246e1 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -244,8 +244,10 @@ def test_repair_fallback_filters_critical_priority_and_in_progress_status(): assert sql is not None assert 'FROM "dbo_repair_logs"' in sql - assert "\"status\" = 'in progress'" in sql - assert "\"priority\" = 'critical'" in sql + assert 'LOWER("status") IN' in sql + assert "'in progress'" in sql + assert "'in-progress'" in sql + assert 'LOWER("priority") = \'critical\'' in sql def test_repair_fallback_preserves_hyphenated_in_progress_status_value(): @@ -267,8 +269,10 @@ def test_repair_fallback_preserves_hyphenated_in_progress_status_value(): assert sql is not None assert 'FROM "dbo_repair_logs"' in sql - assert "\"status\" = 'in-progress'" in sql - assert "\"priority\" = 'critical'" in sql + assert 'LOWER("status") IN' in sql + assert "'in-progress'" in sql + assert "'in progress'" in sql + assert 'LOWER("priority") = \'critical\'' in sql def test_repair_logs_highest_priority_orders_by_verified_priority_column(): @@ -315,7 +319,7 @@ def test_critical_priority_repairs_filter_verified_priority_column(): assert sql is not None assert 'FROM "dbo_repair_logs"' in sql - assert "\"priority\" = 'critical'" in sql + assert 'LOWER("priority") = \'critical\'' in sql def test_repairs_by_status_counts_verified_repair_rows(): @@ -466,7 +470,7 @@ def test_failure_type_value_filter_uses_verified_failure_type_column(): assert sql is not None assert 'FROM "dbo_report_failures"' in sql assert 'COUNT(*) AS "record_count"' in sql - assert "\"failure_type\" = 'JTAG'" in sql + assert "LOWER(\"failure_type\") = 'jtag'" in sql def test_board_models_most_failures_counts_failure_records_not_defect_rate(): @@ -577,3 +581,150 @@ def test_repairs_by_technician_requires_one_schema_object_covering_both_concepts assert message is not None assert "repair" in message assert "technician" in message + + +def test_invoice_status_count_uses_verified_task_status_column(): + contexts = [ + """ + CREATE TABLE dbo_PBI_View_Unrecorded_Liabilities_Header ( + invoicenumber VARCHAR, + taskstatus VARCHAR, + invoicedate TIMESTAMP, + grossamount DECIMAL, + suppliername VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show invoices grouped by task status.", contexts) + + assert sql is not None + assert 'FROM "dbo_PBI_View_Unrecorded_Liabilities_Header"' in sql + assert 'SELECT "taskstatus", COUNT("invoicenumber") AS "record_count"' in sql + assert 'GROUP BY "taskstatus"' in sql + + +def test_invoice_month_count_groups_by_verified_invoice_date(): + contexts = [ + """ + CREATE TABLE dbo_PBI_View_Unrecorded_Liabilities_Header ( + invoicenumber VARCHAR, + invoicedate TIMESTAMP, + grossamount DECIMAL, + suppliername VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show invoice counts by invoice month.", contexts) + + assert sql is not None + assert 'EXTRACT(YEAR FROM "invoicedate") AS "year"' in sql + assert 'EXTRACT(MONTH FROM "invoicedate") AS "month"' in sql + assert 'COUNT("invoicenumber") AS "record_count"' in sql + + +def test_top_suppliers_by_gross_amount_uses_verified_amount_measure(): + contexts = [ + """ + CREATE TABLE dbo_PBI_View_Unrecorded_Liabilities_Header ( + suppliername VARCHAR, + invoicenumber VARCHAR, + grossamount DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql("List the top suppliers by total gross amount.", contexts) + + assert sql is not None + assert 'SELECT "suppliername", SUM("grossamount") AS "total_value"' in sql + assert 'ORDER BY "total_value" DESC' in sql + + +def test_supplier_email_missing_filter_uses_verified_email_column(): + contexts = [ + """ + CREATE TABLE dbo_SupplierMaster_Email ( + supplierid VARCHAR, + suppliername VARCHAR, + email1 VARCHAR, + email2 VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Which supplier email records are missing their first email address?", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_SupplierMaster_Email"' in sql + assert 'WHERE ("email1" IS NULL OR "email1" = \'\')' in sql + + +def test_reconciliation_count_by_status_and_preparer_group_uses_two_dimensions(): + contexts = [ + """ + CREATE TABLE dbo_PBI_View_Recon_Status ( + transid VARCHAR, + status VARCHAR, + preparergroup VARCHAR, + glaccount VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Count reconciliations by status and preparer group.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_PBI_View_Recon_Status"' in sql + assert 'SELECT "status", "preparergroup", COUNT("transid") AS "record_count"' in sql + assert 'GROUP BY "status", "preparergroup"' in sql + + +def test_gl_accounts_highest_balance_uses_verified_balance_measure(): + contexts = [ + """ + CREATE TABLE dbo_View_Global_Exposure_SAP ( + glaccount VARCHAR, + year INTEGER, + month INTEGER, + endbalance DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show GL accounts with the highest ending balance this year.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_View_Global_Exposure_SAP"' in sql + assert 'SELECT "glaccount", SUM("endbalance") AS "total_value"' in sql + assert 'WHERE "year" = ' in sql + assert 'ORDER BY "total_value" DESC' in sql + + +def test_recent_journal_workflow_uses_verified_signer_date(): + contexts = [ + """ + CREATE TABLE dbo_Journals_Workflow ( + journalid VARCHAR, + signer VARCHAR, + signer_status VARCHAR, + signer_date TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql("List recent journal workflow approvals.", contexts) + + assert sql is not None + assert 'FROM "dbo_Journals_Workflow"' in sql + assert 'ORDER BY "signer_date" DESC' in sql From b84d4b4a27536ae9cef2a9fd24637923ac37dc4b Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 12:30:26 +0000 Subject: [PATCH 1073/1087] Update Ask grounding verification handoff --- WRENAI_LOCAL_ASK_HANDOFF.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index 846c5f0873..56a83b4f70 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -235,6 +235,13 @@ Coverage added for: - Board models with most failures uses count, not defect rate. - Highest defect rate uses rate metric. - Repairs by technician requires one schema object or relationship coverage. +- Follow-up generic business families: + - Invoice counts by task status and invoice month. + - Top suppliers by gross amount. + - Missing supplier email fields. + - Reconciliation counts by status and preparer group. + - GL accounts by highest ending balance for the current year. + - Recent journal workflow approvals. Retrieval test file: @@ -301,6 +308,8 @@ There are also many local untracked runtime/data artifacts in the repository. Do - Runtime source code should remain generic. Do not add checks for exact prompts such as `Which repair logs have the highest priority?`. - Tests may use representative table and prompt names; production code must not. - Retrieval context currently uses metadata/descriptions and some semantic context. It does not appear to carry robust sample-value lists. Status casing/value handling works for tested prompts, but richer value-aware matching would improve future accuracy. +- On 2026-08-20, live E2E against the running local app showed the checkout at `D:\WrenAI` was older than the pushed PR branch, so some observed runtime failures were from stale local code. The PR branch now includes follow-up commit `4a199fca1` (`Broaden Ask semantic grounding coverage`), which expands generic schema grounding for invoice, supplier email, reconciliation, GL balance, and journal workflow families without hardcoding one project/table/prompt. +- Local live execution against CWPay/CW_GL may still fail until those SQL Server datasources are reachable; the observed error was an ODBC login/network timeout to `BRVBISQL.INT.CW.LOCAL,1433`, not a SQL identifier hallucination. - `enable_column_pruning` was not the focus of today's final validation. - Full pytest suite still needs an environment with `pytest` installed. From d137601343bd121128b7498a9321600a34477026 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 12:39:24 +0000 Subject: [PATCH 1074/1087] Record final Ask verification blocker --- WRENAI_LOCAL_ASK_HANDOFF.md | 36 +++++++++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index 56a83b4f70..ad01550d1a 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -26,6 +26,8 @@ Current projects visible through `/api/v1/projects/current`: - id `12`, `CWPay` - id `13`, `CW_GL` +After the final temporary PR-service validation attempt, the original scheduled AI service was restarted and health checked successfully on port `5555`. + ## Source Control / PR Status The work is pushed to the fork branch: @@ -34,7 +36,7 @@ The work is pushed to the fork branch: - Branch: `organization/ask-schema-grounding-20260820` - PR: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1` - PR base: `organization-feature` -- Schema implementation commit before handoff-only updates: `cc55d1e05` +- Latest schema implementation commit before handoff-only updates: `4a199fca1` (`Broaden Ask semantic grounding coverage`) - Check PR #1 for the live head SHA because handoff-only commits may be added after the implementation commit. The PR branch was rebased onto the latest `origin/organization-feature` after GitHub initially reported conflicts against the wrong compare/base. It was then pushed with `--force-with-lease`. @@ -58,7 +60,7 @@ Do not open this work against upstream `Canner/WrenAI:main` unless that is expli ### Generic Schema Grounding -Permanent source changes are now in `D:\WrenAI\wren-ai-service`, not only `.codex-tmp`. +Permanent source changes are committed on the PR branch and present in the clean PR worktree at `D:\WrenAI-ask-e2e-fix-20260820`. The original local checkout at `D:\WrenAI` may still be on an older local commit until it is refreshed from `origin/organization/ask-schema-grounding-20260820`. Main file: @@ -129,9 +131,9 @@ Relevant behavior: - Unsupported-schema failures now avoid showing invented SQL as something to fix. - Sales/Orders cleanup remains in place: UI project list shows `Orders`, not duplicate `Sales`. -## Live Validation Done +## Live Validation and Verification -All live checks were run through the UI GraphQL Ask path after restarting the AI service. +The live checks below were run through the UI GraphQL Ask path after restarting the AI service during this workstream. A later broader app regression against the updated PR source was attempted, but could not complete because the configured LLM endpoint timed out during intent classification; details are in `Final Temp PR-Service Attempt`. ### PCB_DB @@ -186,6 +188,29 @@ Passed: - No SQL candidate. - Message clearly said the active project does not contain verified `repair` and `priority/severity` fields. +### Final Temp PR-Service Attempt + +To verify the latest PR branch rather than the stale local checkout, the scheduled AI service was stopped and a temporary service was started from `D:\WrenAI-ask-e2e-fix-20260820\wren-ai-service` using the existing local venv and `D:\WrenAI\wren-ai-service\config.local.yaml`. + +Observed: + +- First temp start was missing the original `.env.dev` values and Ask failed with `Embedding request failed with status 401: Invalid API Key`. +- Temp service was restarted with environment values loaded from `D:\WrenAI\wren-ai-service\.env.dev`; health check passed. +- A broader GraphQL Ask regression began with random/generic questions across PCB_DB, Orders, CWPay, CW_GL, and an unsupported Orders repair question. +- The run was blocked by the configured LLM endpoint timing out during intent classification: + - endpoint: `10.104.74.10:18002` + - error class: `litellm.exceptions.InternalServerError` + - underlying connection error: `The semaphore timeout period has expired` +- Because this was an external LLM connectivity timeout, the final broader live app regression did not complete on the latest PR commit. + +Cleanup completed: + +- Temporary regression runner was stopped. +- Temporary PR AI service was stopped. +- Scheduled task `WrenAI 04 AI Service` was restarted. +- AI health returned `{"status":"ok"}`. +- Active project was restored to PCB_DB with project id `10`. + ## Checks Run Passed: @@ -309,6 +334,7 @@ There are also many local untracked runtime/data artifacts in the repository. Do - Tests may use representative table and prompt names; production code must not. - Retrieval context currently uses metadata/descriptions and some semantic context. It does not appear to carry robust sample-value lists. Status casing/value handling works for tested prompts, but richer value-aware matching would improve future accuracy. - On 2026-08-20, live E2E against the running local app showed the checkout at `D:\WrenAI` was older than the pushed PR branch, so some observed runtime failures were from stale local code. The PR branch now includes follow-up commit `4a199fca1` (`Broaden Ask semantic grounding coverage`), which expands generic schema grounding for invoice, supplier email, reconciliation, GL balance, and journal workflow families without hardcoding one project/table/prompt. +- The final broader live regression against the latest PR source was blocked by LLM endpoint connectivity to `10.104.74.10:18002`, not by SQL identifier validation. Re-run this after the LLM endpoint is reachable. - Local live execution against CWPay/CW_GL may still fail until those SQL Server datasources are reachable; the observed error was an ODBC login/network timeout to `BRVBISQL.INT.CW.LOCAL,1433`, not a SQL identifier hallucination. - `enable_column_pruning` was not the focus of today's final validation. - Full pytest suite still needs an environment with `pytest` installed. @@ -322,4 +348,4 @@ There are also many local untracked runtime/data artifacts in the repository. Do 5. Install or enable pytest in `wren-ai-service\venv`, then run focused tests. 6. Review the large `utils/sql.py` diff carefully; consider extracting fallback/grounding helpers into smaller modules after behavior is stable. 7. Add sample-value metadata to retrieval context if available, then make value matching use that metadata instead of only text normalization. -8. Run a broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when their data sources are available. +8. Re-run the broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when the LLM endpoint and data sources are available. From ea8127c98cdc313432e5d8f096bf8cb14856009d Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 12:49:29 +0000 Subject: [PATCH 1075/1087] Fix same-thread Ask lifecycle and metric grounding --- WRENAI_LOCAL_ASK_HANDOFF.md | 60 +++++ .../src/pipelines/generation/utils/sql.py | 234 +++++++++++++++++- .../retrieval/db_schema_retrieval.py | 33 ++- .../generation/test_sql_schema_grounding.py | 150 +++++++++++ .../repositories/threadResponseRepository.ts | 4 +- .../apollo/server/services/askingService.ts | 90 ++++++- .../server/services/askingTaskTracker.ts | 28 ++- .../components/pages/home/prompt/index.tsx | 8 +- wren-ui/src/hooks/useAskPrompt.tsx | 24 +- wren-ui/src/pages/home/[id].tsx | 27 +- 10 files changed, 617 insertions(+), 41 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index ad01550d1a..700d8ef688 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -338,6 +338,66 @@ There are also many local untracked runtime/data artifacts in the repository. Do - Local live execution against CWPay/CW_GL may still fail until those SQL Server datasources are reachable; the observed error was an ODBC login/network timeout to `BRVBISQL.INT.CW.LOCAL,1433`, not a SQL identifier hallucination. - `enable_column_pruning` was not the focus of today's final validation. - Full pytest suite still needs an environment with `pytest` installed. +- UI `check-types`/Jest could not be run in the clean PR worktree because `node_modules` and the Yarn node_modules state file were absent. `corepack yarn` is available; run `corepack yarn install --immutable` in an environment where dependency install is allowed, then `corepack yarn check-types`. + +## Follow-up Fix: Same-thread Ask Reliability + +Additional generic fixes were added for the issue where an existing thread could show `Failed to create asking task` while a new thread worked better. + +Changed: + +- `wren-ui/src/apollo/server/repositories/threadResponseRepository.ts` + - Thread responses now have deterministic ordering. + - Limited history uses newest response ids first. +- `wren-ui/src/pages/home/[id].tsx` + - Same-thread resume logic now considers only the latest thread response for unfinished asking/thread-response polling. + - Older stale unfinished responses no longer take over the prompt state for the current thread. +- `wren-ui/src/hooks/useAskPrompt.tsx` + - Asking-task polling is scoped to the active task id so late results from older tasks do not drive the current prompt. + - Failed task creation now stops polling and propagates the error to the prompt. +- `wren-ui/src/components/pages/home/prompt/index.tsx` + - Prompt UI resets out of `Understanding question` if asking-task creation fails. +- `wren-ui/src/apollo/server/services/askingService.ts` + - Added logs for thread id, project id, deploy id, previous latest task state, history count, task id/query id, and failure reason. +- `wren-ui/src/apollo/server/services/askingTaskTracker.ts` + - Added logs for task creation request, project/deploy id, histories, created local task id, query id, and creation failure reason. + +Root cause addressed: + +- Existing-thread pages could resume or keep polling an older unfinished response instead of the latest response, especially when previous failed/stale task state remained in the thread. New threads did not have that stale state, which is why they behaved better. + +## Follow-up Fix: Average / Distribution / Location Grounding + +Additional generic schema-first SQL fixes were added for metric intent and dimension grounding: + +- Average intent now requires `AVG(...)` over a verified numeric measure such as age/duration/elapsed fields. +- Average requests no longer fall back to `COUNT(...)`. +- If a requested average measure is not available in the active project schema, the flow returns unsupported schema instead of a wrong count. +- Distribution/breakdown intent now uses grouped counts over verified category/status fields. +- Repair-status distributions can filter verified status values such as completed and in-progress while still grouping by status. +- Dimension-pair listing, such as board model by location, uses `SELECT DISTINCT` only when one verified schema object exposes all requested dimensions. +- If board model and location are not covered by verified schema, the flow returns unsupported schema instead of inventing `Location`. +- Retrieval expansion and column ranking now include average, age, duration, elapsed, distribution, and breakdown concepts. +- Semantic validation now rejects valid-but-wrong SQL that answers average requests with counts or distribution requests without grouped counts. + +Focused validation passed: + +```text +py_compile: +- wren-ai-service/src/pipelines/generation/utils/sql.py +- wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +- wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py + +Direct Python test-function harness: +- ran=37 failures=0 + +Direct SQL smoke: +- average age of failed units by board model -> AVG verified age measure grouped by board_model +- average age without age/duration field -> unsupported schema +- repair status distribution -> grouped counts by verified status with completed/in-progress filters +- board model associated with location -> SELECT DISTINCT only when both verified columns exist +- board model/location without location field -> unsupported schema +``` ## Recommended Next Steps diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index a076963ac4..d8ad872fca 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -270,12 +270,20 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "approvals": {"approval", "approver", "reviewer", "signer", "status"}, "approver": {"approval", "reviewer", "signer"}, "approvers": {"approval", "approver", "reviewer", "signer"}, + "avg": {"average"}, + "averages": {"average"}, "balances": {"balance"}, + "breakdown": {"count", "distribution", "group"}, + "breakdowns": {"breakdown", "count", "distribution", "group"}, "cust": {"customer"}, "customers": {"customer"}, "critical": {"priority", "severity"}, "curr": {"currency"}, "boards": {"board"}, + "days": {"age", "duration"}, + "distribution": {"count", "group"}, + "durations": {"duration"}, + "elapsed": {"age", "duration"}, "email": {"address", "mail"}, "email1": {"address", "email", "first", "mail", "primary"}, "emails": {"address", "email", "mail"}, @@ -288,6 +296,7 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "grossamount": {"amount", "gross", "value"}, "high": {"priority", "severity"}, "highest": {"top"}, + "hours": {"age", "duration"}, "invoicedate": {"date", "invoice"}, "invoicemonth": {"invoice", "month"}, "invoicenumber": {"invoice", "number"}, @@ -351,6 +360,7 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "urgency": {"priority", "severity"}, "locations": {"location"}, "materials": {"material"}, + "mean": {"average"}, "missing": {"blank", "empty", "null"}, } expanded = set(tokens) @@ -1377,6 +1387,57 @@ def validate_sql_semantic_coverage( label for label, concept_tokens in concepts if not schema_tokens & concept_tokens ] if not missing_concepts: + if _is_average_metric_intent(raw_query_tokens): + if not re.search(r"(?is)\bAVG\s*\(", sql): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL does not compute an AVG " + "aggregate over a verified measure. Use a verified numeric " + "measure for the requested average, or return no SQL if the " + "active project does not contain one." + ) + if raw_query_tokens & _AVERAGE_MEASURE_TOKENS: + average_measure_columns = [] + for relation in referenced_relations: + column = _choose_average_measure_column( + raw_query_tokens, + schema_details.get(relation, []), + ) + if column: + average_measure_columns.append(column["name"]) + if average_measure_columns and not any( + _sql_mentions_identifier(sql, column) + for column in average_measure_columns + ): + return ( + "Schema grounding failed. The question asks for an " + "average of an age or duration measure, but the " + "generated SQL does not use a verified age/duration " + "column. Use the verified measure column or return no " + "SQL if the active project does not contain one." + ) + if re.search(r"(?is)\bCOUNT\s*\(", sql) and not re.search( + r"(?is)\bAVG\s*\(", + sql, + ): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL computes a count. Do not " + "substitute COUNT for unsupported averages." + ) + if _is_distribution_metric_intent(raw_query_tokens) and ( + raw_query_tokens & {"status", "priority", "severity"} + ): + if not re.search(r"(?is)\bCOUNT\s*\(", sql) or not re.search( + r"(?is)\bGROUP\s+BY\b", + sql, + ): + return ( + "Schema grounding failed. The question asks for a " + "distribution across categories, but the generated SQL does " + "not compute grouped counts. Use GROUP BY on the verified " + "category column with COUNT, or return no SQL." + ) if _is_failure_count_intent(raw_query_tokens, query_tokens): if not re.search(r"(?is)\bCOUNT\s*\(", sql): return ( @@ -1565,6 +1626,12 @@ def _expanded_fallback_query_tokens(query: str) -> set[str]: tokens.update({"date", "failure", "log", "priority", "progress", "repair", "status"}) if tokens & {"failure", "failures", "defect", "defects"}: tokens.update({"code", "defect", "failure", "severity", "status", "type"}) + if tokens & {"age", "duration", "elapsed"}: + tokens.update({"age", "days", "duration", "elapsed", "hours"}) + if tokens & {"average", "avg", "mean"}: + tokens.update({"average"}) + if tokens & {"distribution", "breakdown", "across"}: + tokens.update({"count", "distribution", "group", "status"}) if tokens & {"material", "materials"}: tokens.update({"item", "material", "part"}) if tokens & {"location", "locations"}: @@ -1614,6 +1681,17 @@ def _is_date_type(data_type: str) -> bool: _RATE_METRIC_TOKENS = {"rate", "ratio", "percent", "percentage"} _COUNT_METRIC_TOKENS = {"count", "many", "most", "number", "total"} +_AVERAGE_METRIC_TOKENS = {"average", "avg", "mean"} +_DISTRIBUTION_METRIC_TOKENS = {"distribution", "breakdown"} +_AVERAGE_MEASURE_TOKENS = { + "age", + "cycle", + "days", + "duration", + "elapsed", + "hours", + "minutes", +} _PRIORITY_VALUE_ALIASES = { "urgent": "urgent", "critical": "critical", @@ -1639,6 +1717,14 @@ def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: return bool(raw_query_tokens & _RATE_METRIC_TOKENS) +def _is_average_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _AVERAGE_METRIC_TOKENS) + + +def _is_distribution_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _DISTRIBUTION_METRIC_TOKENS) + + def _is_failure_count_intent( raw_query_tokens: set[str], query_tokens: set[str], @@ -1677,6 +1763,7 @@ def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[ ("repair", {"repair"}, {"repair"}), ("material", {"material"}, {"material", "part"}), ("location", {"location"}, {"location", "site", "area"}), + ("age/duration", {"age", "duration", "elapsed"}, _AVERAGE_MEASURE_TOKENS), ("customer", {"customer"}, {"customer", "cust"}), ("supplier/vendor", {"supplier", "vendor"}, {"supplier", "vendor"}), ("technician", {"technician", "tech"}, {"technician", "tech"}), @@ -1748,6 +1835,7 @@ def _choose_fallback_table( ) -> tuple[str, list[dict[str, str]]] | None: concept_tokens = concept_tokens or query_tokens rate_metric_intent = _is_rate_metric_intent(concept_tokens) + average_metric_intent = _is_average_metric_intent(concept_tokens) failure_count_intent = _is_failure_count_intent(concept_tokens, query_tokens) board_model_intent = _has_board_model_intent(query_tokens) or _has_board_model_intent( concept_tokens @@ -1756,6 +1844,7 @@ def _choose_fallback_table( for table_name, columns in schema_details.items(): table_tokens = _table_business_tokens(table_name, columns) column_token_union = set() + has_average_measure = False has_numeric_amount_measure = False has_numeric_sales_measure = False has_date_capable_column = False @@ -1772,6 +1861,11 @@ def _choose_fallback_table( "value", }: has_numeric_amount_measure = True + if _is_numeric_type(column["data_type"]) and column_tokens & ( + concept_tokens & _AVERAGE_MEASURE_TOKENS + or _AVERAGE_MEASURE_TOKENS + ): + has_average_measure = True if _is_numeric_type(column["data_type"]) and column_tokens & { "amount", "intake", @@ -1816,6 +1910,12 @@ def _choose_fallback_table( if not _table_covers_requested_concepts(table_name, columns, concept_tokens): continue + if average_metric_intent: + if concept_tokens & _AVERAGE_MEASURE_TOKENS and not has_average_measure: + continue + if concept_tokens & _AVERAGE_MEASURE_TOKENS: + score += 90 + if board_model_intent and rate_metric_intent and query_tokens & { "defect", "failure", @@ -2163,9 +2263,19 @@ def _choose_dimension_columns( ({"order"}, {"order", "ord", "number"}), ({"batch"}, {"batch", "id"}), ] + compound_dimension_triggers = [ + {"board", "model"}, + {"business", "unit"}, + {"email", "address"}, + ] selected: list[str] = [] for trigger_tokens, column_tokens in dimension_specs: - if query_tokens & trigger_tokens: + trigger_matches = ( + trigger_tokens.issubset(query_tokens) + if trigger_tokens in compound_dimension_triggers + else bool(query_tokens & trigger_tokens) + ) + if trigger_matches: column = _choose_ranked_column_by_tokens(columns, column_tokens) if column and column["name"] not in selected: selected.append(column["name"]) @@ -2226,6 +2336,29 @@ def _choose_count_subject_column( return None +def _choose_average_measure_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + measure_specs = [ + ({"age", "duration", "elapsed"}, _AVERAGE_MEASURE_TOKENS), + ({"rate", "ratio", "percent", "percentage"}, _RATE_METRIC_TOKENS | {"score"}), + ({"amount", "gross", "net", "value"}, {"amount", "gross", "net", "value"}), + ({"balance"}, {"balance", "end", "ending", "value"}), + ({"quantity", "qty"}, {"quantity", "qty"}), + ] + for trigger_tokens, column_tokens in measure_specs: + if query_tokens & trigger_tokens: + column = _choose_ranked_column_by_tokens( + columns, + set(column_tokens), + numeric=True, + ) + if column: + return column + return None + + def _is_text_type(data_type: str) -> bool: return data_type.upper() in {"CHAR", "NCHAR", "NVARCHAR", "STRING", "TEXT", "VARCHAR"} @@ -2609,16 +2742,21 @@ def generate_simple_analytics_sql( "account", "accounts", "address", + "age", "approval", "approvals", "approver", + "average", "balance", "balances", "board", + "breakdown", "business", "count", "customer", "defect", + "distribution", + "duration", "email", "failure", "failures", @@ -2636,6 +2774,7 @@ def generate_simple_analytics_sql( "logs", "location", "material", + "mean", "missing", "model", "monthly", @@ -2682,6 +2821,8 @@ def generate_simple_analytics_sql( schema_details = _extract_schema_details(contexts) rate_metric_intent = _is_rate_metric_intent(raw_query_tokens) + average_metric_intent = _is_average_metric_intent(raw_query_tokens) + distribution_metric_intent = _is_distribution_metric_intent(raw_query_tokens) failure_count_intent = _is_failure_count_intent(raw_query_tokens, query_tokens) board_model_intent = _has_board_model_intent( raw_query_tokens @@ -2720,6 +2861,8 @@ def generate_simple_analytics_sql( { "failure_count": failure_count_intent, "rate": rate_metric_intent, + "average": average_metric_intent, + "distribution": distribution_metric_intent, "board_model": board_model_intent, "failure_type_filter": bool(failure_type_filter_value), }, @@ -2765,6 +2908,36 @@ def generate_simple_analytics_sql( "status", "urgent", } + if distribution_metric_intent: + dimension_columns = _choose_dimension_columns( + raw_query_tokens | {"status"}, + columns, + max_columns=1, + ) + if dimension_columns: + subject_column = _choose_count_subject_column(raw_query_tokens, columns) + count_expression = "COUNT(*)" + predicates = [] + if subject_column: + count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" + predicates.append(_non_missing_value_predicate(subject_column)) + status_filter_values = _extract_status_filter_values(query) + if status_column and status_filter_values: + predicates.append( + _value_match_predicate( + status_column, + status_filter_values[0], + status_filter_values[1:], + ) + ) + where_clause = f"\nWHERE {' AND '.join(predicates)}" if predicates else "" + quoted_dimensions = _quote_joined(dimension_columns) + return ( + f"SELECT {quoted_dimensions}, {count_expression} AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC" + ) + if query_tokens & {"repair", "repairs"} and repair_filter_intent: predicates = [] status_filter_values = _extract_status_filter_values(query) @@ -2794,6 +2967,32 @@ def generate_simple_analytics_sql( ) priority_column = _choose_priority_column(columns) + if average_metric_intent: + average_measure_column = _choose_average_measure_column(raw_query_tokens, columns) + if not average_measure_column: + return None + dimension_columns = _choose_dimension_columns(raw_query_tokens, columns) + where_predicates = [] + if raw_query_tokens & {"failure", "failed", "defect"}: + subject_column = _choose_count_subject_column({"failure"}, columns) + if subject_column: + where_predicates.append(_non_missing_value_predicate(subject_column)) + where_clause = ( + f"\nWHERE {' AND '.join(where_predicates)}" if where_predicates else "" + ) + aggregate_expr = f"AVG({_quote_identifier(average_measure_column['name'])})" + if dimension_columns: + quoted_dimensions = _quote_joined(dimension_columns) + return ( + f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier('average_value')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier('average_value')} DESC" + ) + return ( + f"SELECT {aggregate_expr} AS {_quote_identifier('average_value')}\n" + f"FROM {quoted_table}{where_clause}" + ) + if ( priority_column and raw_query_tokens & {"priority", "severity"} @@ -2969,6 +3168,27 @@ def generate_simple_analytics_sql( f"WHERE {_missing_value_predicate(missing_column)}{limit_clause}" ) + dimension_listing_intent = bool( + raw_query_tokens & {"associated", "association", "associations", "each", "list", "show"} + ) and not bool( + raw_query_tokens + & ( + _AVERAGE_METRIC_TOKENS + | _COUNT_METRIC_TOKENS + | _RATE_METRIC_TOKENS + | {"highest", "latest", "lowest", "recent", "top"} + ) + ) + if dimension_listing_intent: + dimension_columns = _choose_dimension_columns(raw_query_tokens, columns, max_columns=3) + if len(dimension_columns) >= 2: + quoted_dimensions = _quote_joined(dimension_columns) + order_clause = ", ".join(_quote_identifier(column) for column in dimension_columns) + return ( + f"SELECT DISTINCT {quoted_dimensions}\n" + f"FROM {quoted_table}\nORDER BY {order_clause}" + ) + explicit_grouping_intent = bool(raw_query_tokens & {"group", "grouped"}) or bool( re.search(r"(?i)\bby\s+[A-Za-z0-9_ -]+\b", query) ) @@ -2978,6 +3198,9 @@ def generate_simple_analytics_sql( "amount", "balance", "gross", + "average", + "age", + "duration", "margin", "net", "rate", @@ -3039,9 +3262,12 @@ def generate_simple_analytics_sql( or (board_model_intent and not rate_metric_intent) ) if ( - raw_query_tokens & {"count", "number"} - or failure_count_intent - or implied_count_by_dimension + not average_metric_intent + and ( + raw_query_tokens & {"count", "number"} + or failure_count_intent + or implied_count_by_dimension + ) ): dimension_columns = _choose_dimension_columns(raw_query_tokens, columns) if dimension_columns: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 7c375608f3..4ed24002a0 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -124,11 +124,11 @@ 17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. 18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. 19. Prefer tables and columns that directly model the requested business entities, measures, statuses, dates, identifiers, and dimensions. Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema provides specific modeled columns for the same concept. -20. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, repairs, failures, batches, materials, locations, status, severity, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. +20. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, repairs, failures, batches, materials, locations, status, severity, age, duration, average, distribution, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. 21. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. 22. Never return placeholder table or column names such as tablename, table_name, dbo.tablename, BatchId, Material, Location, or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. 23. If the request asks for revenue, sales, or sales trends, prefer exact business measure columns named like Revenue, SalesValue, USDFXSalesValue, FXSalesValue, IntakeValue, Amount, or equivalent modeled sales fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. -24. If the request asks for an explicit rate, ratio, percentage, revenue, amount, sales value, or other named measure and the schema already contains that exact measure column, use the declared measure column directly. Do not use a rate column to answer "most failures", "number of failures", or other count-of-records requests unless the question explicitly asks for a rate/ratio/percentage. +24. If the request asks for an explicit average, rate, ratio, percentage, revenue, amount, sales value, age, duration, or other named measure and the schema already contains that exact measure column, use the declared measure column directly. Do not use a rate column to answer "most failures", "number of failures", or other count-of-records requests unless the question explicitly asks for a rate/ratio/percentage. Do not use count-based columns or grouped counts to answer average requests unless the question asks for a count. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -938,10 +938,10 @@ def _augment_retrieval_query(query: str) -> str: "product item material type name category" ), ("repair", "repairs"): ( - "repair status priority severity failure board model log in progress completed critical" + "repair status priority severity failure board model log in progress completed critical age duration" ), ("failure", "failures", "defect", "defects"): ( - "failure defect severity occurrence record count code type system status" + "failure defect severity occurrence record count code type system status age duration" ), ("batch", "batches"): ( "batch board model supplier defect rate inspection status" @@ -950,7 +950,16 @@ def _augment_retrieval_query(query: str) -> str: "material item part component location" ), ("location", "locations"): ( - "location site warehouse area material" + "location site warehouse area material board model" + ), + ("average", "avg", "mean"): ( + "average avg mean numeric measure age duration elapsed days hours amount rate" + ), + ("age", "duration", "elapsed"): ( + "age duration elapsed days hours numeric measure average" + ), + ("distribution", "breakdown"): ( + "distribution breakdown count group status category" ), ("business unit", "bu", "division"): ( "business unit division company account organization" @@ -972,6 +981,10 @@ def _augment_retrieval_query(query: str) -> str: for trigger in ("rate", "ratio", "percent", "percentage") ): expansions.append("rate ratio percent percentage") + if any(trigger in lowered for trigger in ("average", "avg", "mean")): + expansions.append("average avg mean numeric measure") + if any(trigger in lowered for trigger in ("distribution", "breakdown")): + expansions.append("distribution breakdown group count status category") if not expansions: return query @@ -1640,6 +1653,16 @@ def _lexical_columns_and_tables_needed( len(column_tokens & {"date", "day", "month", "year", "time"}) * 5 ) + if query_tokens & {"average", "avg", "mean", "age", "duration", "elapsed"}: + score += ( + len( + column_tokens + & {"age", "duration", "elapsed", "days", "hours", "amount", "rate"} + ) + * 7 + ) + if query_tokens & {"distribution", "breakdown"}: + score += len(column_tokens & {"status", "state", "category", "type"}) * 7 if query_tokens & {"top", "highest", "lowest", "bottom"}: measure_tokens = { "amount", diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index 205c8246e1..22638b14e1 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -728,3 +728,153 @@ def test_recent_journal_workflow_uses_verified_signer_date(): assert sql is not None assert 'FROM "dbo_Journals_Workflow"' in sql assert 'ORDER BY "signer_date" DESC' in sql + + +def test_average_failed_unit_age_by_board_model_uses_avg_age_measure(): + contexts = [ + """ + CREATE TABLE dbo_unit_failures ( + unit_id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + unit_age_days DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the average age of failed units for each board model.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_unit_failures"' in sql + assert 'SELECT "board_model", AVG("unit_age_days") AS "average_value"' in sql + assert 'WHERE ("failure_code" IS NOT NULL AND "failure_code" <> \'\')' in sql + assert 'GROUP BY "board_model"' in sql + assert 'COUNT(' not in sql + + +def test_average_failed_unit_age_without_age_measure_is_unsupported(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + unit_id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the average age of failed units for each board model.", + contexts, + ) + message = unsupported_schema_message( + "Show the average age of failed units for each board model.", + contexts, + ) + + assert sql is None + assert message is not None + assert "age/duration" in message + + +def test_distribution_of_repairs_across_statuses_groups_counts(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the distribution of repairs across completed and in-progress statuses.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert 'SELECT "status", COUNT("id") AS "record_count"' in sql + assert 'LOWER("status") IN' in sql + assert "'completed'" in sql + assert "'in-progress'" in sql + assert "'in progress'" in sql + assert 'GROUP BY "status"' in sql + assert "SELECT *" not in sql + + +def test_board_models_associated_with_locations_requires_verified_columns(): + contexts = [ + """ + CREATE TABLE dbo_board_locations ( + board_model VARCHAR, + location VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the board models associated with each location.", + contexts, + ) + + assert sql is not None + assert 'SELECT DISTINCT "board_model", "location"' in sql + assert 'FROM "dbo_board_locations"' in sql + + +def test_board_models_associated_with_locations_without_location_is_unsupported(): + contexts = [ + """ + CREATE TABLE dbo_parts_catalog ( + board_model VARCHAR, + material VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the board models associated with each location.", + contexts, + ) + message = unsupported_schema_message( + "Show the board models associated with each location.", + contexts, + ) + + assert sql is None + assert message is not None + assert "location" in message + + +def test_semantic_coverage_rejects_count_for_average_intent(): + contexts = [ + """ + CREATE TABLE dbo_unit_failures ( + unit_id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + unit_age_days DECIMAL + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT board_model, COUNT(failure_code) AS record_count + FROM dbo_unit_failures + GROUP BY board_model + """, + "Show the average age of failed units for each board model.", + contexts, + ) + + assert error is not None + assert "average" in error diff --git a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts index 6cd9896668..a425068bc9 100644 --- a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts @@ -128,7 +128,9 @@ export class ThreadResponseRepository .leftJoin('thread', 'thread.id', 'thread_response.thread_id'); if (limit) { - query.orderBy('created_at', 'desc').limit(limit); + query.orderBy('thread_response.id', 'desc').limit(limit); + } else { + query.orderBy('thread_response.id', 'asc'); } return (await query) diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 0ea641eac7..a38f8c461c 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -677,6 +677,7 @@ export class AskingService implements IAskingService { const { threadId, language } = payload; const currentProject = await this.projectService.getCurrentProject(); let projectId = payload.projectId ?? currentProject.id; + let previousTaskState = null; if (threadId) { const thread = await this.ensureThreadInCurrentProject(threadId); if (payload.projectId && payload.projectId !== thread.projectId) { @@ -685,6 +686,7 @@ export class AskingService implements IAskingService { ); } projectId = thread.projectId; + previousTaskState = await this.getLatestThreadTaskState(threadId); } else if (projectId !== currentProject.id) { throw new Error(`Project ${projectId} is not the active project`); } @@ -696,19 +698,52 @@ export class AskingService implements IAskingService { const histories = threadId && isContextualFollowUpQuestion(input.question) ? await this.getAskingHistory(threadId, threadResponseId) : null; - const response = await this.askingTaskTracker.createAskingTask({ - query: input.question, - histories, + const logContext = { + threadId: threadId ?? null, + projectId, + currentProjectId: currentProject.id, deployId, - projectId: projectId.toString(), - configurations: { language }, - rerunFromCancelled, - previousTaskId, - threadResponseId, - }); - return { - id: response.queryId, + previousTaskState, + historyCount: histories?.length ?? 0, + rerunFromCancelled: !!rerunFromCancelled, + previousTaskId: previousTaskId ?? null, + threadResponseId: threadResponseId ?? null, }; + logger.info( + `Creating asking task: ${JSON.stringify({ + ...logContext, + question: input.question, + })}`, + ); + + try { + const response = await this.askingTaskTracker.createAskingTask({ + query: input.question, + histories, + deployId, + projectId: projectId.toString(), + configurations: { language }, + rerunFromCancelled, + previousTaskId, + threadResponseId, + }); + logger.info( + `Created asking task: ${JSON.stringify({ + ...logContext, + queryId: response.queryId, + })}`, + ); + return { + id: response.queryId, + }; + } catch (err: any) { + logger.error( + `Failed to create asking task: ${JSON.stringify(logContext)} reason=${ + err?.stack || err?.message || err + }`, + ); + throw err; + } } public async rerunAskingTask( @@ -1498,6 +1533,39 @@ export class AskingService implements IAskingService { }; } + private async getLatestThreadTaskState(threadId: number) { + try { + const [latestResponse] = + await this.threadResponseRepository.getResponsesWithThread(threadId, 1); + if (!latestResponse) { + return null; + } + + const task = latestResponse.askingTaskId + ? await this.askingTaskRepository.findOneBy({ + id: latestResponse.askingTaskId, + }) + : null; + const detail = task?.detail as any; + + return { + threadResponseId: latestResponse.id, + askingTaskId: latestResponse.askingTaskId ?? null, + queryId: task?.queryId ?? null, + status: detail?.status ?? null, + type: detail?.type ?? null, + hasSql: !!latestResponse.sql, + }; + } catch (err: any) { + logger.warn( + `Failed to inspect latest thread task state for thread ${threadId}: ${ + err?.message || err + }`, + ); + return null; + } + } + private async ensureThreadInCurrentProject(threadId: number): Promise { const [thread, project] = await Promise.all([ this.threadRepository.findOneBy({ id: threadId }), diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index 2c363aa4bc..64a466ad6a 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -101,6 +101,18 @@ export class AskingTaskTracker implements IAskingTaskTracker { input: CreateAskingTaskInput, ): Promise<{ queryId: string }> { try { + logger.info( + `Creating asking task request: ${JSON.stringify({ + projectId: input.projectId ?? null, + deployId: input.deployId ?? null, + hasHistories: !!input.histories?.length, + historyCount: input.histories?.length ?? 0, + rerunFromCancelled: !!input.rerunFromCancelled, + previousTaskId: input.previousTaskId ?? null, + threadResponseId: input.threadResponseId ?? null, + question: input.query, + })}`, + ); // Call the AI service to create a task const response = await this.wrenAIAdaptor.ask(input); const queryId = response.queryId; @@ -168,10 +180,20 @@ export class AskingTaskTracker implements IAskingTaskTracker { this.trackedTasksById.set(createdTask.id, task); } - logger.info(`Created asking task with queryId: ${queryId}`); + logger.info( + `Created asking task with queryId: ${queryId}, taskId: ${ + task.taskId ?? input.previousTaskId ?? 'unbound' + }, projectId: ${input.projectId ?? 'unknown'}`, + ); return { queryId }; - } catch (err) { - logger.error(`Failed to create asking task: ${err}`); + } catch (err: any) { + logger.error( + `Failed to create asking task for projectId=${ + input.projectId ?? 'unknown' + }, deployId=${input.deployId ?? 'unknown'}: ${ + err?.stack || err?.message || err + }`, + ); throw err; } } diff --git a/wren-ui/src/components/pages/home/prompt/index.tsx b/wren-ui/src/components/pages/home/prompt/index.tsx index fba196d4e0..d8e3260277 100644 --- a/wren-ui/src/components/pages/home/prompt/index.tsx +++ b/wren-ui/src/components/pages/home/prompt/index.tsx @@ -202,7 +202,13 @@ export default forwardRef(function Prompt(props, ref) { // start the state as understanding when user submit question askProcessState.transitionTo(PROCESS_STATE.UNDERSTANDING); setShowResult(true); - onSubmit && (await onSubmit(value)); + try { + onSubmit && (await onSubmit(value)); + } catch (error) { + console.error(error); + setShowResult(false); + askProcessState.resetState(); + } }; useImperativeHandle( diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index 47253f4dac..1a1e45a6b8 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -203,10 +203,23 @@ export default function useAskPrompt(threadId?: number) { const lastRecommendedFingerprintRef = useRef(null); const recommendedCreationKeyRef = useRef(null); const recommendedCreationRequestRef = useRef | null>(null); + const [activeAskingTaskId, setActiveAskingTaskId] = useState( + null, + ); const askingTask = useMemo( - () => askingTaskResult.data?.askingTask || null, - [askingTaskResult.data], + () => { + const task = askingTaskResult.data?.askingTask || null; + if ( + activeAskingTaskId && + task?.queryId && + task.queryId !== activeAskingTaskId + ) { + return null; + } + return task; + }, + [activeAskingTaskId, askingTaskResult.data], ); const askingTaskType = useMemo(() => askingTask?.type, [askingTask?.type]); const askingStreamTask = askingStreamTaskResult.data; @@ -252,6 +265,7 @@ export default function useAskPrompt(threadId?: number) { } stopAskingTaskPolling(); + setActiveAskingTaskId(taskId); askingTaskPollingTargetRef.current = taskId; const pollingSessionId = askingTaskPollingSessionRef.current; @@ -549,11 +563,15 @@ export default function useAskPrompt(threadId?: number) { variables: { data: { question: value, threadId } }, }); const taskId = response.data?.createAskingTask?.id; - if (!taskId) return; + if (!taskId) { + throw new Error('CreateAskingTask returned no task id'); + } await startAskingTaskPolling(taskId); } catch (error) { console.error(error); + stopAskingTaskPolling(); + throw error; } }; diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index 7b7efacc81..1d107a397b 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -410,27 +410,28 @@ export default function HomeThread() { const handleUnfinishedTasks = useCallback( (responses: ThreadResponse[]) => { + const latestResponse = [...(responses || [])].sort( + (a, b) => Number(b?.id || 0) - Number(a?.id || 0), + )[0]; + if (!latestResponse) return; + // unfinished asking task - const unfinishedAskingResponse = (responses || []).find( - (response) => - response?.askingTask && !getIsFinished(response?.askingTask?.status), - ); - const unfinishedTaskId = unfinishedAskingResponse?.askingTask?.queryId; - if (unfinishedAskingResponse && unfinishedTaskId) { + const unfinishedTaskId = + latestResponse?.askingTask && + !getIsFinished(latestResponse.askingTask.status) + ? latestResponse.askingTask.queryId + : null; + if (unfinishedTaskId) { askPrompt.onFetching(unfinishedTaskId); return; } // unfinished thread response - const unfinishedThreadResponse = (responses || []).find( - (response) => !getThreadResponseIsFinished(response), - ); - if ( - canFetchThreadResponse(unfinishedThreadResponse?.askingTask) && - unfinishedThreadResponse + !getThreadResponseIsFinished(latestResponse) && + canFetchThreadResponse(latestResponse?.askingTask) ) { - startThreadResponsePolling(unfinishedThreadResponse.id); + startThreadResponsePolling(latestResponse.id); } }, [askPrompt, startThreadResponsePolling], From 4a040c8175907fe7d9d228ae082ecbe1972286ef Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 14:01:03 +0000 Subject: [PATCH 1076/1087] Improve ticket status grounding --- WRENAI_LOCAL_ASK_HANDOFF.md | 30 + .../src/pipelines/generation/utils/sql.py | 561 +++++++++++++++++- .../retrieval/db_schema_retrieval.py | 75 ++- .../generation/test_sql_schema_grounding.py | 154 +++++ 4 files changed, 785 insertions(+), 35 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index 700d8ef688..7a33f8fb3e 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -399,6 +399,36 @@ Direct SQL smoke: - board model/location without location field -> unsupported schema ``` +## Follow-up Fix: Ticket / Status / Updated Grounding + +Additional generic fixes were added for the latest observed PCB_DB failures: + +- Ticket questions are now treated as a first-class business entity, so ticket/status questions prefer verified ticket/case/issue tables and avoid unrelated user/team tables. +- Blocked/open/closed/completed status filters require a verified status/state/progress/blocking column. A generic activity `kind` column is not enough unless metadata clearly marks it as a status-like field. +- Filter values are checked against verified column meaning and structured sample values when present. Values from user wording are not allowed to become filters on unrelated columns. +- `ordered by ticket ID` now selects a verified ticket id/number column for `ORDER BY`. +- `updated each month` now requires a verified updated/modified timestamp and will not silently fall back to `created_at`. +- Monthly repair counts now count a repair identifier from raw user intent instead of counting a status/failure column introduced by expanded retrieval terms. +- Semantic validation now rejects LLM-generated SQL that uses valid identifiers but filters the wrong column/value or uses created timestamps for updated-time questions. +- Retrieval expansion/ranking now includes ticket, activity, status, blocked/open/closed, and updated/modified concepts, with extra deboosting for user/team tables on ticket questions. + +Focused validation passed after this follow-up: + +```text +py_compile: +- wren-ai-service/src/pipelines/generation/utils/sql.py +- wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +- wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py + +git diff --check: passed + +Direct Python test-function harness using `D:\WrenAI\wren-ai-service\venv\Scripts\python.exe` and real service imports: +- ran=43 failures=0 + +Pytest could not run because the available service venv does not have pytest installed. +Ruff could not run because no ruff module/binary is installed in this shell. +``` + ## Recommended Next Steps 1. Review PR #1: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1`. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index d8ad872fca..059f88c6fd 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -273,8 +273,13 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "avg": {"average"}, "averages": {"average"}, "balances": {"balance"}, + "blocked": {"block", "status"}, + "blocker": {"priority", "severity"}, "breakdown": {"count", "distribution", "group"}, "breakdowns": {"breakdown", "count", "distribution", "group"}, + "cases": {"case", "ticket"}, + "closed": {"status"}, + "completed": {"status"}, "cust": {"customer"}, "customers": {"customer"}, "critical": {"priority", "severity"}, @@ -311,10 +316,14 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "models": {"model"}, "inv": {"invoice"}, "invoices": {"invoice"}, + "issues": {"issue", "ticket"}, "net": {"amount", "value"}, "netamount": {"amount", "net", "value"}, + "opened": {"created", "date", "status", "time"}, + "open": {"status"}, "ord": {"order"}, "orders": {"order"}, + "pending": {"status"}, "preparergroup": {"group", "preparer"}, "preparers": {"preparer"}, "qty": {"quantity"}, @@ -346,7 +355,14 @@ def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: "tasks": {"status", "task"}, "tech": {"technician"}, "technician": {"tech"}, + "ticketid": {"id", "ticket"}, + "ticketnumber": {"number", "ticket"}, + "tickets": {"case", "issue", "ticket"}, "transid": {"id", "trans", "transaction"}, + "updated": {"date", "modified", "time", "updated"}, + "updatedat": {"date", "time", "updated"}, + "modified": {"date", "modified", "time", "updated"}, + "modifiedat": {"date", "modified", "time", "updated"}, "vendor": {"supplier"}, "vendors": {"supplier", "vendor"}, "workflow": {"approval", "status"}, @@ -789,10 +805,10 @@ def _extract_semantic_context_payload(context: str) -> dict[str, Any]: def _extract_semantic_tokens_by_column( context: str, -) -> tuple[set[str], dict[str, set[str]]]: +) -> tuple[set[str], dict[str, set[str]], dict[str, list[str]]]: payload = _extract_semantic_context_payload(context) if not payload: - return set(), {} + return set(), {}, {} table_tokens = _semantic_tokens_from_value( payload.get("semantic_context_not_sql_identifiers") @@ -800,6 +816,7 @@ def _extract_semantic_tokens_by_column( table_tokens.update(_semantic_tokens_from_value(payload.get("object_type"))) column_tokens: dict[str, set[str]] = {} + column_sample_values: dict[str, list[str]] = {} for column in payload.get("columns", []) or []: if not isinstance(column, dict): continue @@ -811,8 +828,42 @@ def _extract_semantic_tokens_by_column( ) if tokens: column_tokens[column_name] = tokens + sample_values = _extract_column_sample_values(column) + if sample_values: + column_sample_values[column_name] = sample_values - return table_tokens, column_tokens + return table_tokens, column_tokens, column_sample_values + + +def _extract_column_sample_values(column: dict[str, Any]) -> list[str]: + values: list[str] = [] + + def add(value: Any) -> None: + if value is None: + return + if isinstance(value, (list, tuple, set)): + for item in value: + add(item) + return + if isinstance(value, dict): + for item in value.values(): + add(item) + return + text = str(value).strip() + if text and text.lower() not in {item.lower() for item in values}: + values.append(text) + + for key in ( + "sample_values", + "sample_value", + "samples", + "values", + "example_values", + "examples", + "distinct_values", + ): + add(column.get(key)) + return values def _extract_schema_details( @@ -829,7 +880,7 @@ def _extract_schema_details( continue relation_name = _unquote_identifier(relation_match.group("name")) - table_semantic_tokens, column_semantic_tokens = ( + table_semantic_tokens, column_semantic_tokens, column_sample_values = ( _extract_semantic_tokens_by_column(context) ) column_block_match = re.search( @@ -876,6 +927,7 @@ def _extract_schema_details( "name": name, "data_type": data_type, "semantic_tokens": column_semantic_tokens.get(name, set()), + "sample_values": column_sample_values.get(name, []), "_table_semantic_tokens": table_semantic_tokens, } ) @@ -1152,6 +1204,15 @@ def _sql_mentions_identifier(sql: str, identifier: str) -> bool: ) +def _sql_mentions_literal_value(sql: str, values: list[str]) -> bool: + lowered_sql = sql.lower() + for value in values: + cleaned = _clean_filter_value(value) + if cleaned and _quote_literal(cleaned.lower()) in lowered_sql: + return True + return False + + def _validate_unqualified_columns_for_single_relation( sql: str, schema_index: dict[str, set[str] | None], @@ -1387,6 +1448,56 @@ def validate_sql_semantic_coverage( label for label, concept_tokens in concepts if not schema_tokens & concept_tokens ] if not missing_concepts: + status_filter_values = _extract_status_filter_values(query) + if status_filter_values: + supported_status_columns = [] + for relation in referenced_relations: + for column in schema_details.get(relation, []): + if not _sql_mentions_identifier(sql, column["name"]): + continue + if _column_supports_filter_values(column, status_filter_values): + supported_status_columns.append(column["name"]) + if not supported_status_columns: + return ( + "Schema grounding failed. The question asks for a status " + "filter value, but the generated SQL does not filter on a " + "verified status/state/blocking column that supports that " + "value. Use a verified status field and value, ask a " + "clarifying question, or return no SQL." + ) + if not _sql_mentions_literal_value(sql, status_filter_values): + value_backed_by_column_name = any( + _fallback_tokens(column_name) + & { + token + for value in status_filter_values + for token in _fallback_tokens(value) + } + for column_name in supported_status_columns + ) + if not value_backed_by_column_name: + return ( + "Schema grounding failed. The question asks for a " + "specific status value, but the generated SQL does not " + "apply that verified value as a filter." + ) + + if raw_query_tokens & {"updated", "modified"}: + temporal_columns = [ + _choose_temporal_column(raw_query_tokens, schema_details.get(relation, [])) + for relation in referenced_relations + ] + temporal_columns = [column for column in temporal_columns if column] + if not temporal_columns or not any( + _sql_mentions_identifier(sql, column) for column in temporal_columns + ): + return ( + "Schema grounding failed. The question asks for records " + "updated or modified over time, but the generated SQL does " + "not use a verified updated/modified timestamp column. Use " + "that column or return no SQL if it is not available." + ) + if _is_average_metric_intent(raw_query_tokens): if not re.search(r"(?is)\bAVG\s*\(", sql): return ( @@ -1624,8 +1735,29 @@ def _expanded_fallback_query_tokens(query: str) -> set[str]: tokens.update({"batch", "board", "defect", "inspection", "rate", "supplier"}) if tokens & {"repair", "repairs"}: tokens.update({"date", "failure", "log", "priority", "progress", "repair", "status"}) + if tokens & {"ticket", "tickets"}: + tokens.update( + { + "activity", + "case", + "date", + "id", + "issue", + "log", + "modified", + "priority", + "state", + "status", + "ticket", + "updated", + } + ) if tokens & {"failure", "failures", "defect", "defects"}: tokens.update({"code", "defect", "failure", "severity", "status", "type"}) + if tokens & {"blocked", "closed", "completed", "escalated", "open", "pending", "resolved"}: + tokens.update({"progress", "state", "status"}) + if tokens & {"updated", "modified"}: + tokens.update({"date", "day", "modified", "month", "time", "updated", "year"}) if tokens & {"age", "duration", "elapsed"}: tokens.update({"age", "days", "duration", "elapsed", "hours"}) if tokens & {"average", "avg", "mean"}: @@ -1711,6 +1843,25 @@ def _is_date_type(data_type: str) -> bool: ("minor", 3), ("low", 2), ] +_STATUS_STATE_COLUMN_TOKENS = { + "progress", + "stage", + "state", + "status", + "workflow", +} +_BLOCKING_COLUMN_TOKENS = {"block", "blocked", "blocking", "hold", "held"} +_STATUS_VALUE_ALIASES = { + "blocked": ("blocked",), + "closed": ("closed",), + "completed": ("completed", "complete"), + "escalated": ("escalated",), + "in progress": ("in progress", "in-progress"), + "in-progress": ("in-progress", "in progress"), + "open": ("open",), + "pending": ("pending",), + "resolved": ("resolved",), +} def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: @@ -1761,9 +1912,20 @@ def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[ {"failure", "failed", "defect"}, ), ("repair", {"repair"}, {"repair"}), + ("ticket", {"ticket"}, {"ticket", "case", "issue"}), ("material", {"material"}, {"material", "part"}), ("location", {"location"}, {"location", "site", "area"}), ("age/duration", {"age", "duration", "elapsed"}, _AVERAGE_MEASURE_TOKENS), + ( + "updated/modified", + {"modified", "updated"}, + {"changed", "modified", "updated"}, + ), + ( + "date/time", + {"created", "date", "july", "latest", "month", "monthly", "recent", "year"}, + {"created", "date", "day", "month", "time", "year"}, + ), ("customer", {"customer"}, {"customer", "cust"}), ("supplier/vendor", {"supplier", "vendor"}, {"supplier", "vendor"}), ("technician", {"technician", "tech"}, {"technician", "tech"}), @@ -1771,9 +1933,9 @@ def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[ ( "priority/severity", {"critical", "priority", "severity"}, - {"priority", "severity"}, + {"priority", "rank", "severity", "urgency"}, ), - ("status", {"status"}, {"status"}), + ("status", {"status"}, {"progress", "stage", "state", "status"}), ("order", {"order"}, {"order", "ord"}), ("invoice", {"invoice"}, {"invoice", "inv"}), ("email/address", {"email", "address"}, {"email", "email1", "address", "mail"}), @@ -1907,6 +2069,10 @@ def _choose_fallback_table( * 4 ) + has_date_capable_column = bool( + _choose_temporal_column(query_tokens | concept_tokens, columns) + ) + if not _table_covers_requested_concepts(table_name, columns, concept_tokens): continue @@ -2012,6 +2178,28 @@ def _choose_fallback_table( continue score += 90 + if query_tokens & {"ticket"}: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"case", "issue", "ticket"}: + continue + score += 100 + if concept_tokens & {"status"}: + has_status_or_blocking = bool(_choose_status_column(columns)) or any( + _column_business_tokens(column) & _BLOCKING_COLUMN_TOKENS + for column in columns + ) + if not has_status_or_blocking: + continue + score += 55 + if concept_tokens & {"priority", "severity"}: + if not _choose_priority_column(columns): + continue + score += 35 + if concept_tokens & {"updated", "modified", "latest", "recent", "month", "monthly"}: + if not has_date_capable_column: + continue + score += 35 + if query_tokens & {"order", "orders"}: table_and_columns = table_tokens | column_token_union explicit_order_support = table_and_columns & {"ord", "order", "orders"} @@ -2227,6 +2415,204 @@ def _choose_ranked_column_by_tokens( return candidates[0][1] +def _choose_status_column(columns: list[dict[str, str]]) -> dict[str, str] | None: + candidates = [] + for column in columns: + column_tokens = _column_business_tokens(column) + score = len(column_tokens & _STATUS_STATE_COLUMN_TOKENS) * 20 + if "status" in column_tokens: + score += 30 + if column_tokens & {"kind", "type", "category"} and not column_tokens & ( + _STATUS_STATE_COLUMN_TOKENS + ): + score -= 20 + if score > 0: + candidates.append((score, column)) + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1]["name"])) + return candidates[0][1] + + +def _sample_value_tokens(column: dict[str, Any]) -> set[str]: + tokens: set[str] = set() + for value in column.get("sample_values") or []: + tokens.update(_fallback_tokens(value)) + return tokens + + +def _column_supports_filter_values( + column: dict[str, Any], + values: list[str], +) -> bool: + cleaned_values = [value for value in (_clean_filter_value(value) for value in values) if value] + if not cleaned_values: + return True + + column_tokens = _column_business_tokens(column) + sample_tokens = _sample_value_tokens(column) + value_tokens: set[str] = set() + for value in cleaned_values: + value_tokens.update(_fallback_tokens(value)) + + if value_tokens & sample_tokens: + return True + if value_tokens & column_tokens: + return True + if value_tokens & set(_PRIORITY_VALUE_ALIASES.values()) and column_tokens & { + "priority", + "rank", + "severity", + "urgency", + }: + return True + if value_tokens & { + "blocked", + "closed", + "completed", + "complete", + "escalated", + "open", + "pending", + "progress", + "resolved", + "status", + } and column_tokens & _STATUS_STATE_COLUMN_TOKENS: + return True + if value_tokens & {"block", "blocked"} and column_tokens & _BLOCKING_COLUMN_TOKENS: + return True + return False + + +def _filter_predicate_for_values( + column: dict[str, str], + values: list[str], +) -> str: + cleaned_values = [ + value for value in (_clean_filter_value(value) for value in values) if value + ] + if not cleaned_values: + return _non_missing_value_predicate(column) + + column_tokens = _column_business_tokens(column) + value_tokens: set[str] = set() + for value in cleaned_values: + value_tokens.update(_fallback_tokens(value)) + + if ( + _is_boolean_type(column["data_type"]) + and value_tokens & {"block", "blocked"} + and column_tokens & _BLOCKING_COLUMN_TOKENS + ): + return f"{_quote_identifier(column['name'])} = TRUE" + + return _value_match_predicate(column, cleaned_values[0], cleaned_values[1:]) + + +def _choose_status_filter_column( + columns: list[dict[str, str]], + values: list[str], +) -> dict[str, str] | None: + candidates: list[tuple[int, dict[str, str]]] = [] + + for column in columns: + column_tokens = _column_business_tokens(column) + if not _column_supports_filter_values(column, values): + continue + score = len(column_tokens & _STATUS_STATE_COLUMN_TOKENS) * 25 + score += len(column_tokens & _BLOCKING_COLUMN_TOKENS) * 12 + score += len(_sample_value_tokens(column)) * 2 + if "status" in column_tokens: + score += 35 + if column_tokens & {"kind", "type", "category"} and not column_tokens & ( + _STATUS_STATE_COLUMN_TOKENS + ): + score -= 25 + if score > 0: + candidates.append((score, column)) + + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1]["name"])) + return candidates[0][1] + + +def _choose_temporal_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + strict_token_groups: list[set[str]] = [] + if query_tokens & {"updated", "modified"}: + strict_token_groups.append({"changed", "modified", "updated"}) + if query_tokens & {"created", "opened"}: + strict_token_groups.append({"created", "opened"}) + + candidates = [] + for column in columns: + column_tokens = _column_business_tokens(column) + if not ( + _is_date_type(column["data_type"]) + or column_tokens & {"date", "day", "month", "time", "year"} + ): + continue + if strict_token_groups and not any( + column_tokens & strict_tokens for strict_tokens in strict_token_groups + ): + continue + score = len( + query_tokens + & column_tokens + & {"created", "date", "day", "modified", "month", "time", "updated", "year"} + ) * 12 + if _is_date_type(column["data_type"]): + score += 20 + if query_tokens & {"updated", "modified"} and column_tokens & { + "modified", + "updated", + }: + score += 45 + if query_tokens & {"created", "opened"} and column_tokens & {"created", "opened"}: + score += 35 + if score > 0: + candidates.append((score, column["name"])) + + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1])) + return candidates[0][1] + + +def _order_by_phrase_tokens(query: str) -> set[str]: + match = re.search( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+(?P[A-Za-z0-9_ /-]+)", + query, + ) + if not match: + return set() + return _fallback_tokens(match.group("value")) + + +def _choose_order_by_column( + query: str, + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + order_tokens = _order_by_phrase_tokens(query) + if not order_tokens: + return None + + if order_tokens & {"id", "number", "no"} and query_tokens & {"ticket"}: + column = _choose_ranked_column_by_tokens( + columns, + {"id", "no", "number", "ticket"}, + ) + if column: + return column["name"] + + column = _choose_ranked_column_by_tokens(columns, order_tokens) + return column["name"] if column else None + + def _choose_dimension_column( query_tokens: set[str], columns: list[dict[str, str]], @@ -2250,6 +2636,7 @@ def _choose_dimension_columns( ({"product"}, {"product", "prod", "item", "material", "name"}), ({"salesperson", "representative"}, {"salesperson", "sales", "person", "rep"}), ({"technician", "tech"}, {"technician", "tech"}), + ({"ticket"}, {"ticket", "id", "number", "no"}), ({"location"}, {"location", "site", "area"}), ({"material"}, {"material", "part", "item"}), ({"priority", "severity"}, {"priority", "severity", "urgency", "rank"}), @@ -2316,6 +2703,7 @@ def _choose_count_subject_column( ({"reconciliation", "recon"}, {"reconciliation", "recon", "trans", "id"}), ({"journal"}, {"journal", "entry", "number", "id"}), ({"account", "gl", "glaccount"}, {"account", "gl", "glaccount", "id"}), + ({"ticket"}, {"ticket", "case", "issue", "id", "number", "no"}), ( {"failure"}, {"failure", "failed", "defect", "code", "line", "status", "sys", "type"}, @@ -2363,6 +2751,10 @@ def _is_text_type(data_type: str) -> bool: return data_type.upper() in {"CHAR", "NCHAR", "NVARCHAR", "STRING", "TEXT", "VARCHAR"} +def _is_boolean_type(data_type: str) -> bool: + return data_type.upper() in {"BIT", "BOOL", "BOOLEAN"} + + def _missing_value_predicate(column: dict[str, str]) -> str: quoted_column = _quote_identifier(column["name"]) if _is_text_type(column["data_type"]): @@ -2451,6 +2843,27 @@ def _select_listing_columns( len(tokens & {"board", "code", "date", "failure", "id", "priority", "status"}) * 14 ) + if query_tokens & {"ticket"}: + score += ( + len( + tokens + & { + "case", + "created", + "date", + "id", + "issue", + "modified", + "number", + "priority", + "state", + "status", + "ticket", + "updated", + } + ) + * 16 + ) if query_tokens & {"journal", "workflow", "approval", "approver", "reviewer", "signer"}: score += ( len( @@ -2641,9 +3054,19 @@ def add(value: str): if re.search(r"(?i)\bin\s+progress\b", query): add("in progress") add("in-progress") - for status in ("completed", "pending", "escalated", "open", "closed"): + for status in ( + "blocked", + "closed", + "completed", + "complete", + "escalated", + "open", + "pending", + "resolved", + ): if re.search(rf"(?i)\b{re.escape(status)}\b", query): - add(status) + for value in _STATUS_VALUE_ALIASES.get(status, (status,)): + add(value) return values @@ -2749,10 +3172,15 @@ def generate_simple_analytics_sql( "average", "balance", "balances", + "blocked", "board", "breakdown", "business", + "case", + "closed", + "completed", "count", + "created", "customer", "defect", "distribution", @@ -2776,13 +3204,16 @@ def generate_simple_analytics_sql( "material", "mean", "missing", + "modified", "model", "monthly", "most", "net", "number", + "open", "order", "orders", + "pending", "preparer", "priority", "product", @@ -2807,12 +3238,15 @@ def generate_simple_analytics_sql( "tasks", "tech", "technician", + "ticket", + "tickets", "top", "trend", "trends", "type", "unit", "units", + "updated", "workflow", "workflows", "year", @@ -2890,7 +3324,17 @@ def generate_simple_analytics_sql( f"FROM {quoted_table}" ) - status_column = _choose_ranked_column_by_tokens(columns, {"status"}) + date_column_tokens = {"date", "day", "month", "time", "year"} + if raw_query_tokens & {"year"} and not raw_query_tokens & {"month", "monthly"}: + date_column_tokens = {"date", "day", "time", "year"} + date_column = _choose_temporal_column(raw_query_tokens | query_tokens, columns) + if not date_column and not raw_query_tokens & {"updated", "modified"}: + date_column = _choose_column_by_tokens( + columns, + date_column_tokens, + date=True, + ) + repair_filter_intent = raw_query_tokens & { "closed", "completed", @@ -2922,13 +3366,15 @@ def generate_simple_analytics_sql( count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" predicates.append(_non_missing_value_predicate(subject_column)) status_filter_values = _extract_status_filter_values(query) - if status_column and status_filter_values: + if status_filter_values: + status_filter_column = _choose_status_filter_column( + columns, + status_filter_values, + ) + if not status_filter_column: + return None predicates.append( - _value_match_predicate( - status_column, - status_filter_values[0], - status_filter_values[1:], - ) + _filter_predicate_for_values(status_filter_column, status_filter_values) ) where_clause = f"\nWHERE {' AND '.join(predicates)}" if predicates else "" quoted_dimensions = _quote_joined(dimension_columns) @@ -2941,30 +3387,81 @@ def generate_simple_analytics_sql( if query_tokens & {"repair", "repairs"} and repair_filter_intent: predicates = [] status_filter_values = _extract_status_filter_values(query) - if status_column and status_filter_values: + if status_filter_values: + status_filter_column = _choose_status_filter_column( + columns, + status_filter_values, + ) + if not status_filter_column: + return None predicates.append( - _value_match_predicate( - status_column, - status_filter_values[0], - status_filter_values[1:], - ) + _filter_predicate_for_values(status_filter_column, status_filter_values) ) priority_filter_value = _extract_priority_filter_value(query) if priority_filter_value: priority_column = _choose_priority_column(columns) - if priority_column: - predicates.append(_value_match_predicate(priority_column, priority_filter_value)) + if not priority_column or not _column_supports_filter_values( + priority_column, + [priority_filter_value], + ): + return None + predicates.append(_filter_predicate_for_values(priority_column, [priority_filter_value])) if predicates: return f"SELECT *\nFROM {quoted_table}\nWHERE {' AND '.join(predicates)}" - date_column_tokens = {"date", "day", "month", "time", "year"} - if raw_query_tokens & {"year"} and not raw_query_tokens & {"month", "monthly"}: - date_column_tokens = {"date", "day", "time", "year"} - date_column = _choose_column_by_tokens( - columns, - date_column_tokens, - date=True, - ) + if query_tokens & {"ticket"} and not ( + average_metric_intent + or distribution_metric_intent + or failure_count_intent + or raw_query_tokens & (_COUNT_METRIC_TOKENS | _RATE_METRIC_TOKENS) + ): + predicates = [] + status_filter_values = _extract_status_filter_values(query) + if status_filter_values: + status_filter_column = _choose_status_filter_column( + columns, + status_filter_values, + ) + if not status_filter_column: + return None + predicates.append( + _filter_predicate_for_values(status_filter_column, status_filter_values) + ) + + priority_filter_value = _extract_priority_filter_value(query) + if priority_filter_value: + priority_filter_column = _choose_priority_column(columns) + if not priority_filter_column or not _column_supports_filter_values( + priority_filter_column, + [priority_filter_value], + ): + return None + predicates.append( + _filter_predicate_for_values(priority_filter_column, [priority_filter_value]) + ) + + order_column = _choose_order_by_column(query, raw_query_tokens, columns) + selected_columns = _select_listing_columns( + raw_query_tokens | {"id", "status", "ticket"}, + columns, + date_column=date_column, + max_columns=8, + ) + for required_column in [order_column]: + if required_column and required_column not in selected_columns: + selected_columns.insert(0, required_column) + if not selected_columns: + selected_columns = column_names[:8] + where_clause = f"\nWHERE {' AND '.join(predicates)}" if predicates else "" + order_clause = ( + f"\nORDER BY {_quote_identifier(order_column)} ASC" if order_column else "" + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + if predicates or order_clause or raw_query_tokens & {"all", "list", "show"}: + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{where_clause}{order_clause}{limit_clause}" + ) priority_column = _choose_priority_column(columns) if average_metric_intent: @@ -3242,7 +3739,7 @@ def generate_simple_analytics_sql( }: year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" - subject_column = _choose_count_subject_column(raw_query_tokens | query_tokens, columns) + subject_column = _choose_count_subject_column(raw_query_tokens, columns) count_expression = "COUNT(*)" where_clause = "" if subject_column: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 4ed24002a0..0f6db6911e 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -90,6 +90,29 @@ "salesvalue", "value", } +_TICKET_SCHEMA_TOKENS = { + "activity", + "case", + "issue", + "ticket", +} +_STATUS_STATE_TOKENS = { + "blocked", + "closed", + "completed", + "open", + "progress", + "stage", + "state", + "status", +} +_UPDATED_TIME_TOKENS = { + "changed", + "date", + "modified", + "time", + "updated", +} table_columns_selection_system_prompt = """ @@ -124,7 +147,7 @@ 17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. 18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. 19. Prefer tables and columns that directly model the requested business entities, measures, statuses, dates, identifiers, and dimensions. Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema provides specific modeled columns for the same concept. -20. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, repairs, failures, batches, materials, locations, status, severity, age, duration, average, distribution, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. +20. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, tickets, activities, repairs, failures, batches, materials, locations, status, blocked/open/closed values, priority, severity, age, duration, average, distribution, updated/modified dates, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. 21. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. 22. Never return placeholder table or column names such as tablename, table_name, dbo.tablename, BatchId, Material, Location, or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. 23. If the request asks for revenue, sales, or sales trends, prefer exact business measure columns named like Revenue, SalesValue, USDFXSalesValue, FXSalesValue, IntakeValue, Amount, or equivalent modeled sales fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. @@ -861,6 +884,22 @@ def score(table_name: str) -> int: len((table_tokens | column_tokens) & _CUSTOMS_FINANCE_TOKENS) * 8 ) + if query_tokens & {"ticket"}: + table_and_columns = table_tokens | column_tokens | comment_tokens + if table_and_columns & _TICKET_SCHEMA_TOKENS: + value += len(table_and_columns & _TICKET_SCHEMA_TOKENS) * 12 + else: + value -= 40 + if query_tokens & _STATUS_STATE_TOKENS: + value += len(table_and_columns & _STATUS_STATE_TOKENS) * 8 + if query_tokens & {"updated", "modified"}: + value += len(table_and_columns & _UPDATED_TIME_TOKENS) * 8 + if table_tokens & {"team", "user", "users"} and not ( + table_and_columns & _TICKET_SCHEMA_TOKENS + ): + value -= 40 + if query_tokens & {"updated", "modified"}: + value += len((table_tokens | column_tokens | comment_tokens) & _UPDATED_TIME_TOKENS) * 5 if table_tokens & _GENERIC_TABLE_TOKENS and not direct_table_matches: value -= 6 return value @@ -940,6 +979,9 @@ def _augment_retrieval_query(query: str) -> str: ("repair", "repairs"): ( "repair status priority severity failure board model log in progress completed critical age duration" ), + ("ticket", "tickets", "issue", "case"): ( + "ticket issue case activity log status state blocked open closed priority id number updated modified" + ), ("failure", "failures", "defect", "defects"): ( "failure defect severity occurrence record count code type system status age duration" ), @@ -967,8 +1009,11 @@ def _augment_retrieval_query(query: str) -> str: ("month", "monthly", "july", "year", "trend", "latest"): ( "date month year fiscal calendar trend latest recent" ), - ("status", "severity", "priority", "critical"): ( - "status priority severity critical state category progress" + ("updated", "modified"): ( + "updated modified changed date time timestamp month year" + ), + ("status", "severity", "priority", "critical", "blocked", "open", "closed"): ( + "status priority severity critical blocked open closed state category progress" ), } @@ -1663,6 +1708,16 @@ def _lexical_columns_and_tables_needed( ) if query_tokens & {"distribution", "breakdown"}: score += len(column_tokens & {"status", "state", "category", "type"}) * 7 + if query_tokens & {"ticket"}: + score += len(column_tokens & _TICKET_SCHEMA_TOKENS) * 9 + if query_tokens & _STATUS_STATE_TOKENS: + score += len(column_tokens & _STATUS_STATE_TOKENS) * 8 + if query_tokens & {"updated", "modified"}: + score += len(column_tokens & _UPDATED_TIME_TOKENS) * 8 + if query_tokens & {"id", "number"}: + score += len(column_tokens & {"id", "no", "number", "ticket"}) * 5 + if query_tokens & {"updated", "modified"}: + score += len(column_tokens & _UPDATED_TIME_TOKENS) * 6 if query_tokens & {"top", "highest", "lowest", "bottom"}: measure_tokens = { "amount", @@ -1690,6 +1745,20 @@ def _lexical_columns_and_tables_needed( if not query_tokens & _CUSTOMS_FINANCE_TOKENS: table_score -= len(table_tokens & _CUSTOMS_FINANCE_TOKENS) * 10 table_score += len(table_tokens & _SALES_REVENUE_TOKENS) * 5 + if query_tokens & {"ticket"}: + table_and_columns = table_tokens | { + token + for _, column_name, _ in column_scores + for token in _tokenize_schema_text(column_name) + } + if table_and_columns & _TICKET_SCHEMA_TOKENS: + table_score += len(table_and_columns & _TICKET_SCHEMA_TOKENS) * 15 + else: + table_score -= 45 + if table_tokens & {"team", "user", "users"} and not ( + table_and_columns & _TICKET_SCHEMA_TOKENS + ): + table_score -= 45 total_score = table_score + sum(score for score, _, _ in column_scores) if total_score <= 0: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index 22638b14e1..ab318e6872 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -854,6 +854,160 @@ def test_board_models_associated_with_locations_without_location_is_unsupported( assert "location" in message +def test_blocked_tickets_use_verified_ticket_status_and_order_by_ticket_id(): + contexts = [ + """ + CREATE TABLE dbo_ticket_records ( + ticket_id VARCHAR, + status VARCHAR, + priority VARCHAR, + updated_at TIMESTAMP + ); + """, + """ + CREATE TABLE dbo_team_users ( + id VARCHAR, + active BOOLEAN, + email VARCHAR + ); + """, + ] + + sql = generate_simple_analytics_sql( + "Show all blocked tickets ordered by ticket ID.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_ticket_records"' in sql + assert 'LOWER("status") = \'blocked\'' in sql + assert 'ORDER BY "ticket_id" ASC' in sql + assert 'dbo_team_users' not in sql + + +def test_blocked_ticket_activity_kind_without_status_is_unsupported(): + contexts = [ + """ + CREATE TABLE dbo_ticket_activity_log ( + ticket_id VARCHAR, + kind VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all blocked tickets ordered by ticket ID.", + contexts, + ) + message = unsupported_schema_message( + "Show all blocked tickets ordered by ticket ID.", + contexts, + ) + + assert sql is None + assert message is not None + assert "status" in message + + +def test_semantic_coverage_rejects_blocked_ticket_filter_on_activity_kind(): + contexts = [ + """ + CREATE TABLE dbo_ticket_activity_log ( + ticket_id VARCHAR, + kind VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT ticket_id, kind + FROM dbo_ticket_activity_log + WHERE LOWER(kind) = 'blocked' + ORDER BY ticket_id + """, + "Show all blocked tickets ordered by ticket ID.", + contexts, + ) + + assert error is not None + assert "status" in error + + +def test_repair_counts_updated_each_month_use_verified_updated_timestamp(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + repair_id VARCHAR, + status VARCHAR, + updated_at TIMESTAMP, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the number of repairs updated each month.", + contexts, + ) + + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert 'EXTRACT(YEAR FROM "updated_at") AS "year"' in sql + assert 'EXTRACT(MONTH FROM "updated_at") AS "month"' in sql + assert 'COUNT("repair_id") AS "record_count"' in sql + assert '"created_at"' not in sql + + +def test_repair_counts_updated_each_month_without_updated_timestamp_is_unsupported(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + repair_id VARCHAR, + status VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the number of repairs updated each month.", + contexts, + ) + + assert sql is None + + +def test_semantic_coverage_rejects_created_timestamp_for_updated_question(): + contexts = [ + """ + CREATE TABLE dbo_repair_logs ( + repair_id VARCHAR, + status VARCHAR, + updated_at TIMESTAMP, + created_at TIMESTAMP + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT EXTRACT(YEAR FROM created_at) AS year, + EXTRACT(MONTH FROM created_at) AS month, + COUNT(repair_id) AS record_count + FROM dbo_repair_logs + GROUP BY EXTRACT(YEAR FROM created_at), EXTRACT(MONTH FROM created_at) + """, + "Show the number of repairs updated each month.", + contexts, + ) + + assert error is not None + assert "updated" in error + + def test_semantic_coverage_rejects_count_for_average_intent(): contexts = [ """ From cb6b70ac6cd0f5fc6b450c9e627c27057e56e971 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 14:31:50 +0000 Subject: [PATCH 1077/1087] Remove hardcoded Ask schema grounding maps --- WRENAI_LOCAL_ASK_HANDOFF.md | 474 +-- .../src/pipelines/generation/utils/sql.py | 2648 +++++------------ .../retrieval/db_schema_retrieval.py | 273 +- .../generation/test_sql_schema_grounding.py | 887 +----- .../retrieval/test_db_schema_retrieval.py | 22 +- 5 files changed, 937 insertions(+), 3367 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index 7a33f8fb3e..e4e7239a63 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -1,441 +1,145 @@ -# WrenAI Ask Grounding Handoff +# WrenAI Ask Schema Grounding Handoff Date: 2026-08-20 ## Current Goal -Make WrenAI's Ask pipeline schema-first for the active org/project. Natural language should search and rank verified schema metadata, but final SQL must use only tables, views, columns, relationships, metrics, and values supported by the selected project's metadata. +Make WrenAI Ask schema-driven for the active org/project. User wording may search and rank verified metadata, but final SQL must use only verified tables, columns, relationships, metrics, and supported values from the selected project's retrieved schema. -Do not fix future issues by hardcoding one question, project, table, column, or organization. Representative prompts such as `Which repair logs have the highest priority?` are regression examples only. +Do not fix future issues with exact prompt handling, project-specific branches, table/column mappings, or static business synonym/value lists. -## Final Runtime State - -- UI: `http://127.0.0.1:3000` -- AI service: `http://127.0.0.1:5555` -- AI health: `{"status":"ok"}` -- Active project restored after validation: `org / PCB_DB` -- Active project id: `10` -- Orders project id: `11` -- Sales duplicate: not shown in current project list; `Orders` remains canonical. - -Current projects visible through `/api/v1/projects/current`: - -- id `4`, unnamed DuckDB -- id `10`, `PCB_DB` -- id `11`, `Orders` -- id `12`, `CWPay` -- id `13`, `CW_GL` - -After the final temporary PR-service validation attempt, the original scheduled AI service was restarted and health checked successfully on port `5555`. - -## Source Control / PR Status - -The work is pushed to the fork branch: +## Branch / PR - Repository: `hbalasubramanya-rgb/WrenAI` - Branch: `organization/ask-schema-grounding-20260820` - PR: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1` -- PR base: `organization-feature` -- Latest schema implementation commit before handoff-only updates: `4a199fca1` (`Broaden Ask semantic grounding coverage`) -- Check PR #1 for the live head SHA because handoff-only commits may be added after the implementation commit. - -The PR branch was rebased onto the latest `origin/organization-feature` after GitHub initially reported conflicts against the wrong compare/base. It was then pushed with `--force-with-lease`. - -GitHub readback after the rebase: +- Previous pushed commit before this handoff update: `4a040c817` (`Improve ticket status grounding`) +- Current worktree: `D:\WrenAI-ask-e2e-fix-20260820` -- `mergeable=True` -- `mergeable_state=unstable` +The worktree is detached at the branch tip because the local branch is checked out in another worktree. Commit from this detached worktree and push with: -`unstable` means GitHub checks are pending or failing; it is not a merge-conflict state. - -If the previous/old branch view is gone or stale, use this branch and PR instead: - -- Use branch `organization/ask-schema-grounding-20260820` for this work. -- Review and merge PR #1 into `organization-feature`. -- After merge, use `organization-feature` as the updated canonical branch. +```powershell +git push origin HEAD:organization/ask-schema-grounding-20260820 +``` -Do not open this work against upstream `Canner/WrenAI:main` unless that is explicitly intended; this branch was prepared for the fork's `organization-feature` base. +Do not stage the unrelated mode-only change in `wren-ui/.yarn/releases/yarn-4.5.3.cjs`. ## What Changed Today -### Generic Schema Grounding +### Static Business Maps Removed -Permanent source changes are committed on the PR branch and present in the clean PR worktree at `D:\WrenAI-ask-e2e-fix-20260820`. The original local checkout at `D:\WrenAI` may still be on an older local commit until it is refreshed from `origin/organization/ask-schema-grounding-20260820`. - -Main file: +Files: - `wren-ai-service/src/pipelines/generation/utils/sql.py` - -Added or improved: - -- SQL identifier validation against retrieved schema. -- Semantic coverage validation so valid identifiers are not enough; the referenced table/view must also support the requested business concepts. -- Unsupported-schema result helper that returns `NO_RELEVANT_SQL` with no invented SQL. -- Deterministic schema-grounded fallback for common Ask families: - - count / grouped counts - - top-N - - highest / lowest - - latest / recent - - priority / severity - - status filters - - date/month/year filters - - revenue/sales measures - - failure counts vs defect-rate metrics - - failure type value filters -- Semantic alias support from Wren retrieved context blocks. -- More timestamp type support, including `TIMESTAMPTZ`, which fixed the live `latest repair logs` failure. -- Normalization of dialect issues such as `TOP n`, joined `DESCLIMIT`, and order-by aliases. -- Logs for generated SQL validation, deterministic fallback SQL, fallback validation, selected table, verified columns, and metric intent. - -### Retrieval Improvements - -Main file: - - `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` -Added or improved: - -- Project-scoped retrieval filters retained and tested. -- Query expansion for business concepts such as repair, failure, revenue, order, material, status, priority, latest, and date. -- Ranking uses table names, column names, descriptions/comments, semantic context, and generic-table deboosting. -- Logs for: - - selected project id - - retrieved candidate tables and scores - - selected schema objects and columns - -### Generation Pipeline Wiring - -Files: - -- `wren-ai-service/src/pipelines/generation/sql_generation.py` -- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` -- `wren-ai-service/src/pipelines/generation/sql_correction.py` - -Changes: - -- Passed the user query into post-processing as `fallback_query`. -- Added pre-LLM unsupported-schema checks where retrieved schema clearly cannot cover requested concepts. -- Ensured SQL correction still uses the same schema-first validation and fallback logic. -- Strengthened correction instructions so invalid or hallucinated identifiers are not preserved. - -### UI / Project Cleanup From This Workstream - -Files still dirty from the related UI/runtime fixes: - -- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` -- `wren-ui/src/apollo/server/services/askingService.ts` - -Relevant behavior: - -- Previous `results` crash handling is preserved. -- Unsupported-schema failures now avoid showing invented SQL as something to fix. -- Sales/Orders cleanup remains in place: UI project list shows `Orders`, not duplicate `Sales`. - -## Live Validation and Verification - -The live checks below were run through the UI GraphQL Ask path after restarting the AI service during this workstream. A later broader app regression against the updated PR source was attempted, but could not complete because the configured LLM endpoint timed out during intent classification; details are in `Final Temp PR-Service Attempt`. - -### PCB_DB - -Active project: `PCB_DB`, id `10`. - -Passed: +Removed static business-token and value logic, including: -- `Which repair logs have the highest priority?` - - Table: `dbo_repair_logs` - - Uses verified `priority` - - Orders by generic priority ranking expression -- `Show all critical-priority repairs` - - Table: `dbo_repair_logs` - - Filter: `priority = 'critical'` -- `Show repairs by status` - - Table: `dbo_repair_logs` - - Group: `status` - - Metric: `COUNT(id)` -- `Show latest repair logs` - - Table: `dbo_repair_logs` - - Order: `created_at DESC` - - This was the live regression fixed by adding timestamp type coverage. -- `Show the number of failures by material` - - Uses verified material/failure fields from PCB_DB. -- `Show top 5 board models with the most failures` - - Table: `dbo_repair_logs` - - Metric: `COUNT(failure_code)` - - Did not use `defect_rate`. -- `Show units with JTAG as the failure type` - - Table: `dbo_report_failures` - - Filter: `failure_type = 'JTAG'` -- Extra check: - - `Show all repairs with a critical priority and an in-progress status.` - - Table: `dbo_repair_logs` - - Filters: `status = 'in-progress'` and `priority = 'critical'` +- `_STATUS_VALUE_ALIASES` +- `_PRIORITY_VALUE_ALIASES` +- `_PRIORITY_ORDER` +- `_requested_business_concepts` +- `_expanded_fallback_query_tokens` +- `_expand_fallback_token_aliases` +- retrieval `concept_terms` +- domain-specific table boosts/deboosts for example business words -### Orders +Production scans now return no matches for those removed helpers/maps or for the audited domain strings in the two Ask grounding files. -Temporarily switched active project to `Orders`, id `11`, then restored PCB_DB. +### Schema-Derived Grounding -Passed: +Added generic schema-token extraction and matching from: -- `Show top 10 orders from July` - - Uses Orders table/date fields. -- `Show number of orders by customer` - - Groups by customer. - - Counts distinct order numbers. -- `Show revenue by year` - - Uses verified sales/revenue value and invoice date fields. -- Unsupported check: `Which repair logs have the highest priority?` - - Returned `NO_RELEVANT_SQL`. - - No SQL candidate. - - Message clearly said the active project does not contain verified `repair` and `priority/severity` fields. +- table names +- column names +- table/column semantic descriptions +- relationship/identifier context already present in retrieved schema text +- enum/sample values when supplied in Wren semantic context -### Final Temp PR-Service Attempt +The deterministic fallback is now schema-shape based. It can produce conservative SQL for generic shapes such as: -To verify the latest PR branch rather than the stale local checkout, the scheduled AI service was stopped and a temporary service was started from `D:\WrenAI-ask-e2e-fix-20260820\wren-ai-service` using the existing local venv and `D:\WrenAI\wren-ai-service\config.local.yaml`. +- grouped counts +- averages over verified numeric measures +- sums over verified numeric measures +- top-N grouped counts +- latest/recent listings over verified temporal fields +- month/year buckets over verified temporal fields +- ordering by verified requested columns +- filters only when values are supported by sample/enum metadata -Observed: +It no longer contains domain branches for specific business words or values. -- First temp start was missing the original `.env.dev` values and Ask failed with `Embedding request failed with status 401: Invalid API Key`. -- Temp service was restarted with environment values loaded from `D:\WrenAI\wren-ai-service\.env.dev`; health check passed. -- A broader GraphQL Ask regression began with random/generic questions across PCB_DB, Orders, CWPay, CW_GL, and an unsupported Orders repair question. -- The run was blocked by the configured LLM endpoint timing out during intent classification: - - endpoint: `10.104.74.10:18002` - - error class: `litellm.exceptions.InternalServerError` - - underlying connection error: `The semaphore timeout period has expired` -- Because this was an external LLM connectivity timeout, the final broader live app regression did not complete on the latest PR commit. +### Validation Tightened -Cleanup completed: +`validate_sql_semantic_coverage` now rejects SQL when: -- Temporary regression runner was stopped. -- Temporary PR AI service was stopped. -- Scheduled task `WrenAI 04 AI Service` was restarted. -- AI health returned `{"status":"ok"}`. -- Active project was restored to PCB_DB with project id `10`. +- non-operational query terms are not represented anywhere in active project schema metadata/sample values +- generated SQL uses a verified table but not one covering all schema-backed query tokens +- an average request is answered with count-only SQL +- a distribution/breakdown request is not grouped with counts +- a string literal filter on a sampled column uses a value not present in verified samples -## Checks Run +`unsupported_schema_message` now reports partial coverage instead of returning `None` just because some query terms matched schema. -Passed: +### Retrieval Cleanup -```powershell -git diff --check -- wren-ai-service/src/pipelines/generation/utils/sql.py ` - wren-ai-service/src/pipelines/generation/sql_generation.py ` - wren-ai-service/src/pipelines/generation/followup_sql_generation.py ` - wren-ai-service/src/pipelines/generation/sql_correction.py ` - wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py ` - wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py ` - wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py -``` - -Passed: - -```powershell -cd D:\WrenAI\wren-ai-service -.\venv\Scripts\python.exe -m compileall -q src\pipelines\generation src\pipelines\retrieval ` - tests\pytest\pipelines\generation\test_sql_schema_grounding.py ` - tests\pytest\pipelines\retrieval\test_db_schema_retrieval.py -``` - -Could not run pytest in the service venv: - -```text -D:\WrenAI\wren-ai-service\venv\Scripts\python.exe: No module named pytest -``` +Retrieval query augmentation is now a no-op. Ranking uses only direct overlap between query tokens and retrieved schema metadata. The table-selection prompt was changed to instruct schema-local reasoning without domain examples or built-in synonym lists. -## Tests Added +### Tests Updated -Main test file: +Files: - `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` - -Coverage added for: - -- Unsupported schema clears invalid SQL. -- Generic table rejection for unsupported business concepts. -- Repair priority ordering. -- Critical-priority repair filters. -- Repairs by status. -- Latest repair logs with `TIMESTAMPTZ`. -- Semantic alias column support, for example using real verified `Urgency` when semantic context says it means priority/severity. -- Failure by material / technician with verified columns. -- JTAG failure type filters. -- Board models with most failures uses count, not defect rate. -- Highest defect rate uses rate metric. -- Repairs by technician requires one schema object or relationship coverage. -- Follow-up generic business families: - - Invoice counts by task status and invoice month. - - Top suppliers by gross amount. - - Missing supplier email fields. - - Reconciliation counts by status and preparer group. - - GL accounts by highest ending balance for the current year. - - Recent journal workflow approvals. - -Retrieval test file: - - `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` -Coverage added for: +The generation test module now verifies generic behavior rather than PCB/Orders-specific examples: -- Project filter conditions. -- Query expansion. -- Table ranking by query and schema text. -- Project-scoped schema retrieval behavior. - -## Restart Commands Used - -Restart AI service only: - -```powershell -$taskName = 'WrenAI 04 AI Service' -$listenerProcessIds = Get-NetTCPConnection -LocalPort 5555 -State Listen -ErrorAction SilentlyContinue | - Select-Object -ExpandProperty OwningProcess -Unique -Stop-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue -foreach ($listenerProcessId in $listenerProcessIds) { - if ($listenerProcessId) { - Stop-Process -Id $listenerProcessId -Force -ErrorAction SilentlyContinue - } -} -Start-ScheduledTask -TaskName $taskName -``` +- identifier validation +- unsupported partial schema coverage +- sample-value filters +- unverified values rejected +- grouped counts +- averages vs counts +- latest by temporal column +- monthly counts +- explicit order by +- top grouped count +- sum by year +- literal sample validation -Health check: +The retrieval test now verifies no query expansion and schema-metadata ranking. -```powershell -Invoke-WebRequest -UseBasicParsing http://127.0.0.1:5555/health -``` +## Validation Run -Project switch endpoints used for validation: +Commands run from `D:\WrenAI-ask-e2e-fix-20260820`: ```powershell -Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/11/select -Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/10/select -Invoke-WebRequest -UseBasicParsing http://127.0.0.1:3000/api/v1/projects/current -``` - -## Current Dirty Files To Review - -In the pushed PR branch, the source changes below are committed. The original local checkout at `D:\WrenAI` may still show unrelated dirty runtime/data files and may also show the old pre-rebase local commit until it is refreshed from origin. - -Relevant tracked files: - -- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` -- `wren-ai-service/src/pipelines/generation/sql_correction.py` -- `wren-ai-service/src/pipelines/generation/sql_generation.py` -- `wren-ai-service/src/pipelines/generation/utils/sql.py` -- `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` -- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` -- `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` -- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` -- `wren-ui/src/apollo/server/services/askingService.ts` - -There are also many local untracked runtime/data artifacts in the repository. Do not clean or delete them casually. - -## Important Caveats - -- Runtime source code should remain generic. Do not add checks for exact prompts such as `Which repair logs have the highest priority?`. -- Tests may use representative table and prompt names; production code must not. -- Retrieval context currently uses metadata/descriptions and some semantic context. It does not appear to carry robust sample-value lists. Status casing/value handling works for tested prompts, but richer value-aware matching would improve future accuracy. -- On 2026-08-20, live E2E against the running local app showed the checkout at `D:\WrenAI` was older than the pushed PR branch, so some observed runtime failures were from stale local code. The PR branch now includes follow-up commit `4a199fca1` (`Broaden Ask semantic grounding coverage`), which expands generic schema grounding for invoice, supplier email, reconciliation, GL balance, and journal workflow families without hardcoding one project/table/prompt. -- The final broader live regression against the latest PR source was blocked by LLM endpoint connectivity to `10.104.74.10:18002`, not by SQL identifier validation. Re-run this after the LLM endpoint is reachable. -- Local live execution against CWPay/CW_GL may still fail until those SQL Server datasources are reachable; the observed error was an ODBC login/network timeout to `BRVBISQL.INT.CW.LOCAL,1433`, not a SQL identifier hallucination. -- `enable_column_pruning` was not the focus of today's final validation. -- Full pytest suite still needs an environment with `pytest` installed. -- UI `check-types`/Jest could not be run in the clean PR worktree because `node_modules` and the Yarn node_modules state file were absent. `corepack yarn` is available; run `corepack yarn install --immutable` in an environment where dependency install is allowed, then `corepack yarn check-types`. - -## Follow-up Fix: Same-thread Ask Reliability - -Additional generic fixes were added for the issue where an existing thread could show `Failed to create asking task` while a new thread worked better. - -Changed: - -- `wren-ui/src/apollo/server/repositories/threadResponseRepository.ts` - - Thread responses now have deterministic ordering. - - Limited history uses newest response ids first. -- `wren-ui/src/pages/home/[id].tsx` - - Same-thread resume logic now considers only the latest thread response for unfinished asking/thread-response polling. - - Older stale unfinished responses no longer take over the prompt state for the current thread. -- `wren-ui/src/hooks/useAskPrompt.tsx` - - Asking-task polling is scoped to the active task id so late results from older tasks do not drive the current prompt. - - Failed task creation now stops polling and propagates the error to the prompt. -- `wren-ui/src/components/pages/home/prompt/index.tsx` - - Prompt UI resets out of `Understanding question` if asking-task creation fails. -- `wren-ui/src/apollo/server/services/askingService.ts` - - Added logs for thread id, project id, deploy id, previous latest task state, history count, task id/query id, and failure reason. -- `wren-ui/src/apollo/server/services/askingTaskTracker.ts` - - Added logs for task creation request, project/deploy id, histories, created local task id, query id, and creation failure reason. - -Root cause addressed: - -- Existing-thread pages could resume or keep polling an older unfinished response instead of the latest response, especially when previous failed/stale task state remained in the thread. New threads did not have that stale state, which is why they behaved better. - -## Follow-up Fix: Average / Distribution / Location Grounding - -Additional generic schema-first SQL fixes were added for metric intent and dimension grounding: - -- Average intent now requires `AVG(...)` over a verified numeric measure such as age/duration/elapsed fields. -- Average requests no longer fall back to `COUNT(...)`. -- If a requested average measure is not available in the active project schema, the flow returns unsupported schema instead of a wrong count. -- Distribution/breakdown intent now uses grouped counts over verified category/status fields. -- Repair-status distributions can filter verified status values such as completed and in-progress while still grouping by status. -- Dimension-pair listing, such as board model by location, uses `SELECT DISTINCT` only when one verified schema object exposes all requested dimensions. -- If board model and location are not covered by verified schema, the flow returns unsupported schema instead of inventing `Location`. -- Retrieval expansion and column ranking now include average, age, duration, elapsed, distribution, and breakdown concepts. -- Semantic validation now rejects valid-but-wrong SQL that answers average requests with counts or distribution requests without grouped counts. - -Focused validation passed: - -```text -py_compile: -- wren-ai-service/src/pipelines/generation/utils/sql.py -- wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py -- wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py - -Direct Python test-function harness: -- ran=37 failures=0 - -Direct SQL smoke: -- average age of failed units by board model -> AVG verified age measure grouped by board_model -- average age without age/duration field -> unsupported schema -- repair status distribution -> grouped counts by verified status with completed/in-progress filters -- board model associated with location -> SELECT DISTINCT only when both verified columns exist -- board model/location without location field -> unsupported schema +python -m py_compile wren-ai-service/src/pipelines/generation/utils/sql.py wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +git diff --check ``` -## Follow-up Fix: Ticket / Status / Updated Grounding +Manual local harness: -Additional generic fixes were added for the latest observed PCB_DB failures: +- generation grounding tests: `ran=21 failures=0` +- retrieval touched logic: `ran=2 failures=0` +- randomized schema-derived validation: `ran=10 failures=0` -- Ticket questions are now treated as a first-class business entity, so ticket/status questions prefer verified ticket/case/issue tables and avoid unrelated user/team tables. -- Blocked/open/closed/completed status filters require a verified status/state/progress/blocking column. A generic activity `kind` column is not enough unless metadata clearly marks it as a status-like field. -- Filter values are checked against verified column meaning and structured sample values when present. Values from user wording are not allowed to become filters on unrelated columns. -- `ordered by ticket ID` now selects a verified ticket id/number column for `ORDER BY`. -- `updated each month` now requires a verified updated/modified timestamp and will not silently fall back to `created_at`. -- Monthly repair counts now count a repair identifier from raw user intent instead of counting a status/failure column introduced by expanded retrieval terms. -- Semantic validation now rejects LLM-generated SQL that uses valid identifiers but filters the wrong column/value or uses created timestamps for updated-time questions. -- Retrieval expansion/ranking now includes ticket, activity, status, blocked/open/closed, and updated/modified concepts, with extra deboosting for user/team tables on ticket questions. +`python -m pytest` was not available in the local venv because `pytest` is not installed. -Focused validation passed after this follow-up: +Randomized validation used synthetic selected-project schemas and shuffled questions covering: -```text -py_compile: -- wren-ai-service/src/pipelines/generation/utils/sql.py -- wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py -- wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py - -git diff --check: passed - -Direct Python test-function harness using `D:\WrenAI\wren-ai-service\venv\Scripts\python.exe` and real service imports: -- ran=43 failures=0 - -Pytest could not run because the available service venv does not have pytest installed. -Ruff could not run because no ruff module/binary is installed in this shell. -``` +- grouped count +- average +- latest/recent +- sample-value filter with explicit ordering +- monthly count +- sum by year +- top-N grouped count +- sum by dimension +- unsupported unknown dimension -## Recommended Next Steps +## Remaining Blockers -1. Review PR #1: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1`. -2. Confirm the PR base is `organization-feature`, not `Canner/WrenAI:main`. -3. Resolve any GitHub check failures if `mergeable_state` remains `unstable`, then merge PR #1. -4. After merge, continue from `organization-feature`. -5. Install or enable pytest in `wren-ai-service\venv`, then run focused tests. -6. Review the large `utils/sql.py` diff carefully; consider extracting fallback/grounding helpers into smaller modules after behavior is stable. -7. Add sample-value metadata to retrieval context if available, then make value matching use that metadata instead of only text normalization. -8. Re-run the broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when the LLM endpoint and data sources are available. +- Full browser/application verification against live local PCB_DB/Orders projects was not rerun in this pass. +- Filter-value grounding depends on sample/enum metadata being available. If a project has no samples/enums for categorical values, the safer behavior is unsupported/clarification rather than guessed filters. +- The local checkout at `D:\WrenAI` may not contain the latest branch source until refreshed from `origin/organization/ask-schema-grounding-20260820`. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 059f88c6fd..70f793528a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -258,136 +258,15 @@ } -def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: - aliases = { - "bu": {"business", "unit"}, - "acct": {"account"}, - "accounting": {"account", "ledger"}, - "accounts": {"account"}, - "address": {"email", "mail"}, - "addresses": {"address", "email", "mail"}, - "approval": {"approver", "reviewer", "signer", "status"}, - "approvals": {"approval", "approver", "reviewer", "signer", "status"}, - "approver": {"approval", "reviewer", "signer"}, - "approvers": {"approval", "approver", "reviewer", "signer"}, - "avg": {"average"}, - "averages": {"average"}, - "balances": {"balance"}, - "blocked": {"block", "status"}, - "blocker": {"priority", "severity"}, - "breakdown": {"count", "distribution", "group"}, - "breakdowns": {"breakdown", "count", "distribution", "group"}, - "cases": {"case", "ticket"}, - "closed": {"status"}, - "completed": {"status"}, - "cust": {"customer"}, - "customers": {"customer"}, - "critical": {"priority", "severity"}, - "curr": {"currency"}, - "boards": {"board"}, - "days": {"age", "duration"}, - "distribution": {"count", "group"}, - "durations": {"duration"}, - "elapsed": {"age", "duration"}, - "email": {"address", "mail"}, - "email1": {"address", "email", "first", "mail", "primary"}, - "emails": {"address", "email", "mail"}, - "endbalance": {"balance", "end", "ending"}, - "ending": {"end"}, - "gl": {"account", "ledger"}, - "glaccount": {"account", "gl", "ledger"}, - "glaccounts": {"account", "gl", "ledger"}, - "gross": {"amount", "value"}, - "grossamount": {"amount", "gross", "value"}, - "high": {"priority", "severity"}, - "highest": {"top"}, - "hours": {"age", "duration"}, - "invoicedate": {"date", "invoice"}, - "invoicemonth": {"invoice", "month"}, - "invoicenumber": {"invoice", "number"}, - "invoiceyear": {"invoice", "year"}, - "journals": {"journal"}, - "journalid": {"id", "journal"}, - "journalnumber": {"journal", "number"}, - "logs": {"log", "record"}, - "log": {"record"}, - "lows": {"low"}, - "lowest": {"bottom"}, - "models": {"model"}, - "inv": {"invoice"}, - "invoices": {"invoice"}, - "issues": {"issue", "ticket"}, - "net": {"amount", "value"}, - "netamount": {"amount", "net", "value"}, - "opened": {"created", "date", "status", "time"}, - "open": {"status"}, - "ord": {"order"}, - "orders": {"order"}, - "pending": {"status"}, - "preparergroup": {"group", "preparer"}, - "preparers": {"preparer"}, - "qty": {"quantity"}, - "num": {"number"}, - "no": {"number"}, - "numbers": {"number"}, - "prod": {"product"}, - "products": {"product"}, - "priorities": {"priority", "severity"}, - "priority": {"severity"}, - "recent": {"latest"}, - "recon": {"reconciliation", "reconcile"}, - "reconciliations": {"reconciliation", "reconcile", "recon"}, - "reconciliation": {"reconcile", "recon"}, - "records": {"record"}, - "rep": {"representative", "salesperson"}, - "repairs": {"repair"}, - "reviewer": {"approval", "approver", "signer"}, - "reviewers": {"approval", "approver", "reviewer", "signer"}, - "salesperson": {"sales", "person"}, - "severity": {"priority"}, - "signer": {"approval", "approver", "reviewer"}, - "signers": {"approval", "approver", "reviewer", "signer"}, - "supplier": {"vendor"}, - "supplierid": {"id", "supplier", "vendor"}, - "suppliername": {"name", "supplier", "vendor"}, - "suppliers": {"supplier", "vendor"}, - "taskstatus": {"status", "task"}, - "tasks": {"status", "task"}, - "tech": {"technician"}, - "technician": {"tech"}, - "ticketid": {"id", "ticket"}, - "ticketnumber": {"number", "ticket"}, - "tickets": {"case", "issue", "ticket"}, - "transid": {"id", "trans", "transaction"}, - "updated": {"date", "modified", "time", "updated"}, - "updatedat": {"date", "time", "updated"}, - "modified": {"date", "modified", "time", "updated"}, - "modifiedat": {"date", "modified", "time", "updated"}, - "vendor": {"supplier"}, - "vendors": {"supplier", "vendor"}, - "workflow": {"approval", "status"}, - "workflows": {"approval", "status", "workflow"}, - "counts": {"count"}, - "failed": {"failure"}, - "failures": {"failure"}, - "defects": {"defect"}, - "types": {"type"}, - "units": {"unit", "serial"}, - "urgency": {"priority", "severity"}, - "locations": {"location"}, - "materials": {"material"}, - "mean": {"average"}, - "missing": {"blank", "empty", "null"}, - } +def _expand_fallback_token_variants(tokens: set[str]) -> set[str]: expanded = set(tokens) - for token in list(tokens): - expanded.update(aliases.get(token, set())) - if {"business", "unit"}.issubset(expanded): - expanded.add("bu") - if "customer" in expanded and "number" in expanded: - expanded.update({"cust", "id", "no"}) - if "order" in expanded and "number" in expanded: - expanded.update({"ord", "id", "no"}) + for token in tokens: + if len(token) > 4 and token.endswith("ies"): + expanded.add(token[:-3] + "y") + if len(token) > 4 and token.endswith("es"): + expanded.add(token[:-2]) + if len(token) > 3 and token.endswith("s"): + expanded.add(token[:-1]) return expanded @@ -1213,6 +1092,59 @@ def _sql_mentions_literal_value(sql: str, values: list[str]) -> bool: return False +def _extract_sql_string_literals(sql: str) -> list[str]: + literals = [] + for match in _SINGLE_QUOTED_LITERAL.finditer(sql): + literal = match.group(0)[1:-1].replace("''", "'") + if literal: + literals.append(literal) + return literals + + +def _extract_column_filter_literals(sql: str, column_name: str) -> list[str]: + stripped = _strip_string_literals(sql) + quoted_column = re.escape(_quote_identifier(column_name)) + bare_column = re.escape(column_name) + column_pattern = rf"(?:{quoted_column}|(? str | None: + referenced_relations = { + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + } + for relation in referenced_relations: + for column in schema_details.get(relation, []): + sample_values = column.get("sample_values") or [] + if not sample_values: + continue + literals = _extract_column_filter_literals(sql, column["name"]) + if not literals: + continue + sample_tokens = _sample_value_tokens(column) + sample_lowers = {str(value).lower() for value in sample_values} + for literal in literals: + literal_tokens = _fallback_tokens(literal) + if literal.lower() in sample_lowers or literal_tokens & sample_tokens: + continue + return ( + "Schema grounding failed. The generated SQL filters column " + f"{column['name']} with a literal value not found in that " + "column's verified sample values." + ) + return None + + def _validate_unqualified_columns_for_single_relation( sql: str, schema_index: dict[str, set[str] | None], @@ -1417,10 +1349,6 @@ def validate_sql_semantic_coverage( return None raw_query_tokens = _fallback_tokens(query) - query_tokens = _expanded_fallback_query_tokens(query) - concepts = _requested_business_concepts(raw_query_tokens) - if not concepts: - return None schema_details = _extract_schema_details(contexts) if not schema_details: @@ -1441,144 +1369,55 @@ def validate_sql_semantic_coverage( if columns is not None: schema_tokens.update(_schema_tokens_for_table(relation, columns)) - if not schema_tokens: - return None + required_tokens = _schema_derived_query_tokens(raw_query_tokens, schema_details) + unsupported_tokens = _unsupported_query_tokens(raw_query_tokens, schema_details) + if unsupported_tokens: + return ( + "Schema grounding failed. The retrieved schema metadata does not " + "support these non-operational question term(s): " + f"{', '.join(sorted(unsupported_tokens))}. " + "Select a project with matching schema metadata or ask a supported " + "question." + ) - missing_concepts = [ - label for label, concept_tokens in concepts if not schema_tokens & concept_tokens - ] - if not missing_concepts: - status_filter_values = _extract_status_filter_values(query) - if status_filter_values: - supported_status_columns = [] - for relation in referenced_relations: - for column in schema_details.get(relation, []): - if not _sql_mentions_identifier(sql, column["name"]): - continue - if _column_supports_filter_values(column, status_filter_values): - supported_status_columns.append(column["name"]) - if not supported_status_columns: - return ( - "Schema grounding failed. The question asks for a status " - "filter value, but the generated SQL does not filter on a " - "verified status/state/blocking column that supports that " - "value. Use a verified status field and value, ask a " - "clarifying question, or return no SQL." - ) - if not _sql_mentions_literal_value(sql, status_filter_values): - value_backed_by_column_name = any( - _fallback_tokens(column_name) - & { - token - for value in status_filter_values - for token in _fallback_tokens(value) - } - for column_name in supported_status_columns - ) - if not value_backed_by_column_name: - return ( - "Schema grounding failed. The question asks for a " - "specific status value, but the generated SQL does not " - "apply that verified value as a filter." - ) + missing_concepts = sorted(required_tokens - schema_tokens) + if missing_concepts: + return ( + "Schema grounding failed. The generated SQL uses verified identifiers, " + "but the selected table or view does not cover these schema-backed " + f"question tokens: {', '.join(missing_concepts)}. Use only schema " + "objects whose metadata supports the requested terms, or return no " + "SQL if the active project does not contain them." + ) - if raw_query_tokens & {"updated", "modified"}: - temporal_columns = [ - _choose_temporal_column(raw_query_tokens, schema_details.get(relation, [])) - for relation in referenced_relations - ] - temporal_columns = [column for column in temporal_columns if column] - if not temporal_columns or not any( - _sql_mentions_identifier(sql, column) for column in temporal_columns - ): - return ( - "Schema grounding failed. The question asks for records " - "updated or modified over time, but the generated SQL does " - "not use a verified updated/modified timestamp column. Use " - "that column or return no SQL if it is not available." - ) + if _is_average_metric_intent(raw_query_tokens): + if not re.search(r"(?is)\bAVG\s*\(", sql): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL does not compute AVG over a " + "verified measure." + ) + if re.search(r"(?is)\bCOUNT\s*\(", sql) and not re.search( + r"(?is)\bAVG\s*\(", + sql, + ): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL computes a count." + ) - if _is_average_metric_intent(raw_query_tokens): - if not re.search(r"(?is)\bAVG\s*\(", sql): - return ( - "Schema grounding failed. The question asks for an average " - "metric, but the generated SQL does not compute an AVG " - "aggregate over a verified measure. Use a verified numeric " - "measure for the requested average, or return no SQL if the " - "active project does not contain one." - ) - if raw_query_tokens & _AVERAGE_MEASURE_TOKENS: - average_measure_columns = [] - for relation in referenced_relations: - column = _choose_average_measure_column( - raw_query_tokens, - schema_details.get(relation, []), - ) - if column: - average_measure_columns.append(column["name"]) - if average_measure_columns and not any( - _sql_mentions_identifier(sql, column) - for column in average_measure_columns - ): - return ( - "Schema grounding failed. The question asks for an " - "average of an age or duration measure, but the " - "generated SQL does not use a verified age/duration " - "column. Use the verified measure column or return no " - "SQL if the active project does not contain one." - ) - if re.search(r"(?is)\bCOUNT\s*\(", sql) and not re.search( - r"(?is)\bAVG\s*\(", - sql, - ): - return ( - "Schema grounding failed. The question asks for an average " - "metric, but the generated SQL computes a count. Do not " - "substitute COUNT for unsupported averages." - ) - if _is_distribution_metric_intent(raw_query_tokens) and ( - raw_query_tokens & {"status", "priority", "severity"} + if _is_distribution_metric_intent(raw_query_tokens): + if not re.search(r"(?is)\bCOUNT\s*\(", sql) or not re.search( + r"(?is)\bGROUP\s+BY\b", + sql, ): - if not re.search(r"(?is)\bCOUNT\s*\(", sql) or not re.search( - r"(?is)\bGROUP\s+BY\b", - sql, - ): - return ( - "Schema grounding failed. The question asks for a " - "distribution across categories, but the generated SQL does " - "not compute grouped counts. Use GROUP BY on the verified " - "category column with COUNT, or return no SQL." - ) - if _is_failure_count_intent(raw_query_tokens, query_tokens): - if not re.search(r"(?is)\bCOUNT\s*\(", sql): - return ( - "Schema grounding failed. The question asks for a count of " - "failure records, but the generated SQL does not compute a " - "COUNT aggregate. Use a verified failure-record column/table " - "and group by the requested dimension, or return no SQL if " - "the active project does not contain it." - ) - for relation in referenced_relations: - for column in schema_details.get(relation, []): - if _is_rate_like_column(column) and _sql_mentions_identifier( - sql, column["name"] - ): - return ( - "Schema grounding failed. The question asks for a " - "count of failure records, but the generated SQL uses " - f"rate-like column {column['name']}. Use COUNT over a " - "verified failure occurrence field instead, or return " - "no SQL if the active project does not contain one." - ) - return None + return ( + "Schema grounding failed. The question asks for a distribution " + "or breakdown, but the generated SQL does not compute grouped " + "counts." + ) - return ( - "Schema grounding failed. The generated SQL uses verified identifiers, " - "but the selected table or view does not contain verified fields for the " - f"requested business concept(s): {', '.join(missing_concepts)}. Use only " - "schema objects whose table or column names explicitly support those " - "concepts, or return no SQL if the active project does not contain them." - ) + return _validate_literal_values_against_samples(sql, schema_details, grounding) def unsupported_schema_message( @@ -1588,23 +1427,35 @@ def unsupported_schema_message( if not query: return None query_tokens = _fallback_tokens(query) - concepts = _requested_business_concepts(query_tokens) - if not concepts: - return None schema_details = _extract_schema_details(contexts) if not schema_details: return None - if any( - _table_covers_requested_concepts(table_name, columns, query_tokens) - for table_name, columns in schema_details.items() + required_tokens = _schema_derived_query_tokens(query_tokens, schema_details) + unsupported_tokens = _unsupported_query_tokens(query_tokens, schema_details) + if unsupported_tokens: + return ( + "No retrieved table or view in the active project contains verified " + "schema metadata for all requested non-operational term(s): " + f"{', '.join(sorted(unsupported_tokens))}. Select a project with " + "matching fields, add schema descriptions/sample values, or ask a " + "question supported by the selected project's schema." + ) + + table_tokens = _schema_tokens_by_table(schema_details) + if required_tokens and any( + required_tokens <= tokens for tokens in table_tokens.values() ): return None - concept_labels = ", ".join(label for label, _ in concepts) + + if not required_tokens and not unsupported_tokens: + return None + detail_tokens = sorted(required_tokens or unsupported_tokens) return ( "No retrieved table or view in the active project contains verified " - "fields for all requested business concept(s): " - f"{concept_labels}. Select a project with those fields or ask a question " - "supported by the selected project's schema." + "schema metadata for all requested non-operational term(s): " + f"{', '.join(detail_tokens)}. Select a project with matching fields, " + "add schema descriptions/sample values, or ask a question supported by " + "the selected project's schema." ) @@ -1654,7 +1505,7 @@ def _fallback_tokens(value: Any) -> set[str]: for token in _FALLBACK_TOKEN.findall(text.lower()) if token not in _FALLBACK_STOPWORDS } - return _expand_fallback_token_aliases(tokens) + return _expand_fallback_token_variants(tokens) def _column_business_tokens(column: dict[str, Any]) -> set[str]: @@ -1673,110 +1524,6 @@ def _table_business_tokens( return tokens -def _expanded_fallback_query_tokens(query: str) -> set[str]: - tokens = _fallback_tokens(query) - if tokens & {"revenue", "sale", "sales", "trend", "trends"}: - tokens.update({"amount", "date", "intake", "revenue", "sales", "value"}) - if tokens & {"order", "orders"}: - tokens.update({"amount", "customer", "date", "ord", "order", "value"}) - if tokens & {"invoice", "invoices"}: - tokens.update( - { - "amount", - "currency", - "date", - "gross", - "invoice", - "month", - "net", - "number", - "status", - "supplier", - "task", - "year", - } - ) - if tokens & {"supplier", "vendor"}: - tokens.update({"email", "id", "name", "number", "supplier", "vendor"}) - if tokens & {"email", "address"}: - tokens.update({"address", "email", "first", "mail", "primary"}) - if tokens & {"reconciliation", "recon", "reconcile"}: - tokens.update( - { - "account", - "gl", - "group", - "period", - "preparer", - "recon", - "reconciliation", - "reviewer", - "status", - } - ) - if tokens & {"journal", "workflow", "approval", "approver", "reviewer", "signer"}: - tokens.update( - { - "approval", - "approver", - "date", - "journal", - "reviewer", - "signer", - "status", - "workflow", - } - ) - if tokens & {"account", "accounts", "gl", "glaccount", "ledger"}: - tokens.update({"account", "balance", "gl", "glaccount", "ledger", "period", "year"}) - if tokens & {"balance", "balances"}: - tokens.update({"amount", "balance", "end", "ending", "value", "year"}) - if tokens & {"batch", "batches"}: - tokens.update({"batch", "board", "defect", "inspection", "rate", "supplier"}) - if tokens & {"repair", "repairs"}: - tokens.update({"date", "failure", "log", "priority", "progress", "repair", "status"}) - if tokens & {"ticket", "tickets"}: - tokens.update( - { - "activity", - "case", - "date", - "id", - "issue", - "log", - "modified", - "priority", - "state", - "status", - "ticket", - "updated", - } - ) - if tokens & {"failure", "failures", "defect", "defects"}: - tokens.update({"code", "defect", "failure", "severity", "status", "type"}) - if tokens & {"blocked", "closed", "completed", "escalated", "open", "pending", "resolved"}: - tokens.update({"progress", "state", "status"}) - if tokens & {"updated", "modified"}: - tokens.update({"date", "day", "modified", "month", "time", "updated", "year"}) - if tokens & {"age", "duration", "elapsed"}: - tokens.update({"age", "days", "duration", "elapsed", "hours"}) - if tokens & {"average", "avg", "mean"}: - tokens.update({"average"}) - if tokens & {"distribution", "breakdown", "across"}: - tokens.update({"count", "distribution", "group", "status"}) - if tokens & {"material", "materials"}: - tokens.update({"item", "material", "part"}) - if tokens & {"location", "locations"}: - tokens.update({"area", "location", "site"}) - if tokens & {"month", "monthly", "july"}: - tokens.update({"date", "day", "month", "time", "year"}) - elif tokens & {"latest", "trend", "trends", "year"}: - tokens.update({"date", "day", "time", "year"}) - if "business" in tokens and "unit" in tokens: - tokens.update({"account", "bu", "business", "company", "division", "unit"}) - return tokens - - def _is_numeric_type(data_type: str) -> bool: return data_type.upper() in { "BIGINT", @@ -1815,53 +1562,86 @@ def _is_date_type(data_type: str) -> bool: _COUNT_METRIC_TOKENS = {"count", "many", "most", "number", "total"} _AVERAGE_METRIC_TOKENS = {"average", "avg", "mean"} _DISTRIBUTION_METRIC_TOKENS = {"distribution", "breakdown"} -_AVERAGE_MEASURE_TOKENS = { - "age", - "cycle", - "days", - "duration", - "elapsed", - "hours", - "minutes", -} -_PRIORITY_VALUE_ALIASES = { - "urgent": "urgent", - "critical": "critical", - "high": "high", - "medium": "medium", - "normal": "normal", - "low": "low", -} -_PRIORITY_ORDER = [ - ("critical", 6), - ("urgent", 6), - ("blocker", 6), - ("high", 5), - ("major", 5), - ("medium", 4), - ("normal", 4), - ("minor", 3), - ("low", 2), -] -_STATUS_STATE_COLUMN_TOKENS = { - "progress", - "stage", - "state", - "status", - "workflow", -} -_BLOCKING_COLUMN_TOKENS = {"block", "blocked", "blocking", "hold", "held"} -_STATUS_VALUE_ALIASES = { - "blocked": ("blocked",), - "closed": ("closed",), - "completed": ("completed", "complete"), - "escalated": ("escalated",), - "in progress": ("in progress", "in-progress"), - "in-progress": ("in-progress", "in progress"), - "open": ("open",), - "pending": ("pending",), - "resolved": ("resolved",), +_SUM_METRIC_TOKENS = {"sum", "total"} +_MIN_METRIC_TOKENS = {"bottom", "least", "lowest", "min", "minimum", "smallest"} +_MAX_METRIC_TOKENS = {"greatest", "highest", "largest", "max", "maximum", "most", "top"} +_LATEST_METRIC_TOKENS = {"latest", "newest", "recent"} +_NULL_CHECK_TOKENS = {"blank", "empty", "missing", "null"} +_GENERIC_SCHEMA_INTENT_TOKENS = { + "a", + "across", + "all", + "an", + "and", + "as", + "ascending", + "associated", + "association", + "average", + "avg", + "between", + "bottom", + "breakdown", + "bucket", + "buckets", + "by", + "compare", + "count", + "date", + "day", + "descending", + "distribution", + "each", + "for", + "from", + "group", + "grouped", + "has", + "have", + "highest", + "in", + "latest", + "least", + "list", + "lowest", + "many", + "max", + "maximum", + "me", + "mean", + "min", + "minimum", + "month", + "monthly", + "most", + "newest", + "number", + "of", + "ordered", + "per", + "please", + "quarter", + "recent", + "record", + "records", + "row", + "rows", + "show", + "smallest", + "sort", + "sorted", + "sum", + "the", + "to", + "top", + "total", + "week", + "which", + "with", + "without", + "year", } +_GENERIC_SCHEMA_INTENT_TOKENS.update(_MONTH_NAME_TO_NUMBER.keys()) def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: @@ -1876,118 +1656,78 @@ def _is_distribution_metric_intent(raw_query_tokens: set[str]) -> bool: return bool(raw_query_tokens & _DISTRIBUTION_METRIC_TOKENS) -def _is_failure_count_intent( - raw_query_tokens: set[str], - query_tokens: set[str], -) -> bool: - return ( - bool(raw_query_tokens & {"failure", "failed", "defect"}) - and "failure" in query_tokens - and not _is_rate_metric_intent(raw_query_tokens) - and ( - bool(raw_query_tokens & _COUNT_METRIC_TOKENS) - or bool(raw_query_tokens & {"top", "highest", "lowest", "bottom"}) - ) - ) - - -def _has_board_model_intent(query_tokens: set[str]) -> bool: - return {"board", "model"}.issubset(query_tokens) - - def _is_rate_like_column(column: dict[str, str]) -> bool: return bool(_fallback_tokens(column["name"]) & (_RATE_METRIC_TOKENS | {"score"})) -def _quote_joined(identifiers: list[str]) -> str: - return ", ".join(_quote_identifier(identifier) for identifier in identifiers) +def _is_identifier_like_column(column: dict[str, str]) -> bool: + tokens = _fallback_tokens(column["name"]) + return bool(tokens) and tokens <= {"id", "identifier", "key", "uuid"} -def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[str]]]: - concepts: list[tuple[str, set[str]]] = [] - specs = [ - ( - "failure/defect", - {"failure", "failed", "defect"}, - {"failure", "failed", "defect"}, - ), - ("repair", {"repair"}, {"repair"}), - ("ticket", {"ticket"}, {"ticket", "case", "issue"}), - ("material", {"material"}, {"material", "part"}), - ("location", {"location"}, {"location", "site", "area"}), - ("age/duration", {"age", "duration", "elapsed"}, _AVERAGE_MEASURE_TOKENS), - ( - "updated/modified", - {"modified", "updated"}, - {"changed", "modified", "updated"}, - ), - ( - "date/time", - {"created", "date", "july", "latest", "month", "monthly", "recent", "year"}, - {"created", "date", "day", "month", "time", "year"}, - ), - ("customer", {"customer"}, {"customer", "cust"}), - ("supplier/vendor", {"supplier", "vendor"}, {"supplier", "vendor"}), - ("technician", {"technician", "tech"}, {"technician", "tech"}), - ("product", {"product"}, {"product", "prod", "item", "material"}), - ( - "priority/severity", - {"critical", "priority", "severity"}, - {"priority", "rank", "severity", "urgency"}, - ), - ("status", {"status"}, {"progress", "stage", "state", "status"}), - ("order", {"order"}, {"order", "ord"}), - ("invoice", {"invoice"}, {"invoice", "inv"}), - ("email/address", {"email", "address"}, {"email", "email1", "address", "mail"}), - ( - "account", - {"account", "gl", "glaccount", "ledger"}, - {"account", "gl", "glaccount", "ledger"}, - ), - ("balance", {"balance"}, {"balance", "endbalance"}), - ( - "reconciliation", - {"reconciliation", "recon", "reconcile"}, - {"reconciliation", "recon", "reconcile"}, - ), - ( - "journal/workflow", - {"journal", "workflow"}, - {"journal", "workflow"}, - ), - ( - "approval", - {"approval", "approver", "reviewer", "signer"}, - {"approval", "approver", "reviewer", "signer"}, - ), - ] - for label, triggers, schema_tokens in specs: - if query_tokens & triggers: - concepts.append((label, schema_tokens)) - if {"board", "model"}.issubset(query_tokens): - concepts.append(("board model", {"board", "model"})) - if {"business", "unit"}.issubset(query_tokens): - concepts.append(("business unit", {"business", "unit", "bu", "division"})) - return concepts +def _quote_joined(identifiers: list[str]) -> str: + return ", ".join(_quote_identifier(identifier) for identifier in identifiers) def _schema_tokens_for_table(table_name: str, columns: list[dict[str, str]]) -> set[str]: tokens = _table_business_tokens(table_name, columns) for column in columns: tokens.update(_column_business_tokens(column)) + tokens.update(_sample_value_tokens(column)) return tokens +def _schema_tokens_by_table( + schema_details: dict[str, list[dict[str, str]]], +) -> dict[str, set[str]]: + return { + table_name: _schema_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + + +def _schema_derived_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + return { + token + for token in query_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + and not token.isdigit() + and token in schema_tokens + } + + +def _unsupported_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + return { + token + for token in query_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + and not token.isdigit() + and token not in schema_tokens + } + + def _table_covers_requested_concepts( table_name: str, columns: list[dict[str, str]], concept_tokens: set[str], ) -> bool: - concepts = _requested_business_concepts(concept_tokens) - if not concepts: + required_tokens = { + token + for token in concept_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + } + if not required_tokens: return True schema_tokens = _schema_tokens_for_table(table_name, columns) - return all(schema_tokens & concept_tokens for _, concept_tokens in concepts) + return required_tokens <= schema_tokens def _choose_fallback_table( @@ -1996,337 +1736,28 @@ def _choose_fallback_table( concept_tokens: set[str] | None = None, ) -> tuple[str, list[dict[str, str]]] | None: concept_tokens = concept_tokens or query_tokens - rate_metric_intent = _is_rate_metric_intent(concept_tokens) - average_metric_intent = _is_average_metric_intent(concept_tokens) - failure_count_intent = _is_failure_count_intent(concept_tokens, query_tokens) - board_model_intent = _has_board_model_intent(query_tokens) or _has_board_model_intent( - concept_tokens - ) + required_tokens = _schema_derived_query_tokens(concept_tokens, schema_details) scored_tables = [] for table_name, columns in schema_details.items(): table_tokens = _table_business_tokens(table_name, columns) column_token_union = set() - has_average_measure = False - has_numeric_amount_measure = False - has_numeric_sales_measure = False - has_date_capable_column = False score = len(query_tokens & table_tokens) * 8 for column in columns: column_tokens = _column_business_tokens(column) + sample_tokens = _sample_value_tokens(column) column_token_union.update(column_tokens) - if _is_numeric_type(column["data_type"]) and column_tokens & { - "amount", - "balance", - "cost", - "gross", - "net", - "value", - }: - has_numeric_amount_measure = True - if _is_numeric_type(column["data_type"]) and column_tokens & ( - concept_tokens & _AVERAGE_MEASURE_TOKENS - or _AVERAGE_MEASURE_TOKENS - ): - has_average_measure = True - if _is_numeric_type(column["data_type"]) and column_tokens & { - "amount", - "intake", - "revenue", - "sales", - "value", - }: - has_numeric_sales_measure = True - if _is_date_type(column["data_type"]) or column_tokens & { - "date", - "day", - "month", - "time", - "year", - }: - has_date_capable_column = True + column_token_union.update(sample_tokens) score += len(query_tokens & column_tokens) * 10 + score += len(query_tokens & sample_tokens) * 6 if _is_numeric_type(column["data_type"]): - score += len( - query_tokens - & column_tokens - & { - "amount", - "balance", - "cost", - "count", - "gross", - "margin", - "net", - "quantity", - "rate", - "score", - "value", - } - ) * 4 + score += len(query_tokens & column_tokens) * 2 if _is_date_type(column["data_type"]): - score += ( - len(query_tokens & {"date", "month", "year", "july", "trend", "trends"}) - * 4 - ) - - has_date_capable_column = bool( - _choose_temporal_column(query_tokens | concept_tokens, columns) - ) + score += 2 - if not _table_covers_requested_concepts(table_name, columns, concept_tokens): + table_schema_tokens = _schema_tokens_for_table(table_name, columns) + if required_tokens and not required_tokens <= table_schema_tokens: continue - - if average_metric_intent: - if concept_tokens & _AVERAGE_MEASURE_TOKENS and not has_average_measure: - continue - if concept_tokens & _AVERAGE_MEASURE_TOKENS: - score += 90 - - if board_model_intent and rate_metric_intent and query_tokens & { - "defect", - "failure", - }: - if not {"board", "model"}.issubset(column_token_union): - continue - rate_column = _choose_column_by_tokens( - columns, - {"defect", "rate"}, - numeric=True, - ) - if not rate_column: - continue - score += 130 - - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - score += ( - len(column_token_union & {"amount", "intake", "revenue", "sales", "value"}) - * 10 - ) - score += len(column_token_union & {"date", "month", "time", "year"}) * 5 - if not has_numeric_sales_measure: - continue - score += 50 - if query_tokens & {"year", "month", "monthly", "trend", "trends"}: - if not has_date_capable_column: - continue - score += 30 - if not query_tokens & { - "claim", - "claims", - "customs", - "duty", - "import", - "refund", - "tariff", - }: - table_and_columns = table_tokens | column_token_union - customs_matches = table_and_columns & { - "claim", - "claims", - "customs", - "duty", - "import", - "refund", - "tariff", - "tariffs", - } - if customs_matches and not table_and_columns & {"revenue", "sale", "sales"}: - continue - score -= len(customs_matches) * 40 - - if "customer" in query_tokens: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"customer", "cust"}: - if "missing" in query_tokens: - continue - score -= 80 - else: - score += 45 - - if {"business", "unit"}.issubset(query_tokens): - if "bu" in column_token_union: - score += 60 - elif {"business", "unit"} <= column_token_union: - score += 45 - - if "failure" in query_tokens and ( - query_tokens & {"location", "material", "technician", "tech"} - or board_model_intent - ): - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"failure", "failed", "defect"}: - continue - if board_model_intent: - if not {"board", "model"}.issubset(column_token_union): - continue - if failure_count_intent and not _choose_count_subject_column( - {"failure"}, - columns, - ): - continue - score += 100 - if "location" in query_tokens: - if "location" not in column_token_union: - continue - score += 90 - if "material" in query_tokens: - if not column_token_union & {"material", "part"}: - continue - score += 90 - if query_tokens & {"technician", "tech"}: - if not column_token_union & {"technician", "tech"}: - continue - score += 90 - - if query_tokens & {"ticket"}: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"case", "issue", "ticket"}: - continue - score += 100 - if concept_tokens & {"status"}: - has_status_or_blocking = bool(_choose_status_column(columns)) or any( - _column_business_tokens(column) & _BLOCKING_COLUMN_TOKENS - for column in columns - ) - if not has_status_or_blocking: - continue - score += 55 - if concept_tokens & {"priority", "severity"}: - if not _choose_priority_column(columns): - continue - score += 35 - if concept_tokens & {"updated", "modified", "latest", "recent", "month", "monthly"}: - if not has_date_capable_column: - continue - score += 35 - - if query_tokens & {"order", "orders"}: - table_and_columns = table_tokens | column_token_union - explicit_order_support = table_and_columns & {"ord", "order", "orders"} - order_support = table_and_columns & { - "amount", - "customer", - "intake", - "item", - "ord", - "order", - "orders", - "sales", - "value", - } - if not order_support: - continue - if explicit_order_support: - score += 80 - else: - score -= 60 - - if concept_tokens & {"invoice", "invoices"}: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"invoice", "inv"}: - continue - score += 80 - if query_tokens & {"supplier", "vendor"}: - if not table_and_columns & {"supplier", "vendor"}: - continue - score += 45 - if concept_tokens & {"status", "task"}: - if not table_and_columns & {"status", "state", "task"}: - continue - score += 40 - if concept_tokens & {"gross", "net", "amount", "value", "top", "highest"}: - if not has_numeric_amount_measure: - continue - score += 55 - if concept_tokens & {"date", "month", "monthly", "year", "latest", "recent"}: - if not has_date_capable_column: - continue - score += 35 - - if concept_tokens & {"supplier", "vendor"}: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"supplier", "vendor"}: - if query_tokens & {"email", "address", "missing", "blank", "empty", "null"}: - continue - score -= 45 - else: - score += 35 - - if concept_tokens & {"email", "address"}: - if not column_token_union & {"email", "email1", "address", "mail"}: - continue - score += 65 - - if concept_tokens & {"reconciliation", "recon", "reconcile"}: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"reconciliation", "recon", "reconcile"}: - continue - score += 90 - if query_tokens & {"status"}: - if not table_and_columns & {"status", "state"}: - continue - score += 45 - if query_tokens & {"preparer", "group"}: - if not table_and_columns & {"preparer", "group"}: - continue - score += 45 - - if concept_tokens & {"journal", "workflow"}: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"journal", "workflow"}: - continue - if "journal" in concept_tokens and "journal" not in table_and_columns: - continue - if "workflow" in concept_tokens and "workflow" not in table_and_columns: - continue - score += 85 - if concept_tokens & {"approval", "approver", "reviewer", "signer"}: - if not table_and_columns & {"approval", "approver", "reviewer", "signer"}: - continue - score += 45 - if query_tokens & {"latest", "recent"}: - if not has_date_capable_column: - continue - score += 35 - - if concept_tokens & {"account", "gl", "glaccount", "ledger"}: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"account", "gl", "glaccount", "ledger"}: - continue - score += 45 - if query_tokens & {"balance", "end", "ending"}: - if not table_and_columns & {"balance", "endbalance"}: - continue - if not has_numeric_amount_measure: - continue - score += 70 - - if query_tokens & {"batch", "batches"} and {"defect", "rate"}.issubset( - column_token_union - ): - score += 40 - if {"material", "location"}.issubset(query_tokens) and { - "material", - "location", - }.issubset(column_token_union): - score += 40 - if ( - query_tokens & {"repair", "repairs"} - and (table_tokens | column_token_union) & {"repair", "fix"} - ): - score += 55 - if concept_tokens & {"critical", "priority", "severity"}: - if not _choose_priority_column(columns): - continue - score += 55 - if concept_tokens & {"status"}: - if not column_token_union & {"status", "state", "progress"}: - continue - score += 45 - if concept_tokens & {"latest", "recent"}: - if not has_date_capable_column: - continue - score += 35 + score += len(required_tokens & table_schema_tokens) * 20 if score > 0: scored_tables.append((score, table_name, columns)) @@ -2379,23 +1810,6 @@ def _column_score_for_tokens( score = len(required_tokens & column_tokens) * 10 if required_tokens and required_tokens.issubset(column_tokens): score += 30 - if column["name"].lower() == "bu" and {"business", "unit"} & required_tokens: - score += 60 - if column["name"].lower() in {"custno", "customer_id", "customerid"} and { - "customer", - "number", - } & required_tokens: - score += 30 - if column["name"].lower() in {"custname", "customer_name", "customer"} and { - "customer", - "name", - } & required_tokens: - score += 35 - if column["name"].lower() in {"ordno", "order_no", "order_number", "sales_order_number"} and { - "order", - "number", - } & required_tokens: - score += 35 return score @@ -2415,25 +1829,6 @@ def _choose_ranked_column_by_tokens( return candidates[0][1] -def _choose_status_column(columns: list[dict[str, str]]) -> dict[str, str] | None: - candidates = [] - for column in columns: - column_tokens = _column_business_tokens(column) - score = len(column_tokens & _STATUS_STATE_COLUMN_TOKENS) * 20 - if "status" in column_tokens: - score += 30 - if column_tokens & {"kind", "type", "category"} and not column_tokens & ( - _STATUS_STATE_COLUMN_TOKENS - ): - score -= 20 - if score > 0: - candidates.append((score, column)) - if not candidates: - return None - candidates.sort(key=lambda item: (-item[0], item[1]["name"])) - return candidates[0][1] - - def _sample_value_tokens(column: dict[str, Any]) -> set[str]: tokens: set[str] = set() for value in column.get("sample_values") or []: @@ -2459,28 +1854,6 @@ def _column_supports_filter_values( return True if value_tokens & column_tokens: return True - if value_tokens & set(_PRIORITY_VALUE_ALIASES.values()) and column_tokens & { - "priority", - "rank", - "severity", - "urgency", - }: - return True - if value_tokens & { - "blocked", - "closed", - "completed", - "complete", - "escalated", - "open", - "pending", - "progress", - "resolved", - "status", - } and column_tokens & _STATUS_STATE_COLUMN_TOKENS: - return True - if value_tokens & {"block", "blocked"} and column_tokens & _BLOCKING_COLUMN_TOKENS: - return True return False @@ -2499,17 +1872,10 @@ def _filter_predicate_for_values( for value in cleaned_values: value_tokens.update(_fallback_tokens(value)) - if ( - _is_boolean_type(column["data_type"]) - and value_tokens & {"block", "blocked"} - and column_tokens & _BLOCKING_COLUMN_TOKENS - ): - return f"{_quote_identifier(column['name'])} = TRUE" - return _value_match_predicate(column, cleaned_values[0], cleaned_values[1:]) -def _choose_status_filter_column( +def _choose_filter_column_for_values( columns: list[dict[str, str]], values: list[str], ) -> dict[str, str] | None: @@ -2519,15 +1885,13 @@ def _choose_status_filter_column( column_tokens = _column_business_tokens(column) if not _column_supports_filter_values(column, values): continue - score = len(column_tokens & _STATUS_STATE_COLUMN_TOKENS) * 25 - score += len(column_tokens & _BLOCKING_COLUMN_TOKENS) * 12 - score += len(_sample_value_tokens(column)) * 2 - if "status" in column_tokens: - score += 35 - if column_tokens & {"kind", "type", "category"} and not column_tokens & ( - _STATUS_STATE_COLUMN_TOKENS - ): - score -= 25 + value_tokens = { + token + for value in values + for token in _fallback_tokens(value) + } + score = len(column_tokens & value_tokens) * 10 + score += len(_sample_value_tokens(column) & value_tokens) * 20 if score > 0: candidates.append((score, column)) @@ -2541,38 +1905,13 @@ def _choose_temporal_column( query_tokens: set[str], columns: list[dict[str, str]], ) -> str | None: - strict_token_groups: list[set[str]] = [] - if query_tokens & {"updated", "modified"}: - strict_token_groups.append({"changed", "modified", "updated"}) - if query_tokens & {"created", "opened"}: - strict_token_groups.append({"created", "opened"}) - candidates = [] for column in columns: column_tokens = _column_business_tokens(column) - if not ( - _is_date_type(column["data_type"]) - or column_tokens & {"date", "day", "month", "time", "year"} - ): - continue - if strict_token_groups and not any( - column_tokens & strict_tokens for strict_tokens in strict_token_groups - ): + if not _is_date_type(column["data_type"]): continue - score = len( - query_tokens - & column_tokens - & {"created", "date", "day", "modified", "month", "time", "updated", "year"} - ) * 12 - if _is_date_type(column["data_type"]): - score += 20 - if query_tokens & {"updated", "modified"} and column_tokens & { - "modified", - "updated", - }: - score += 45 - if query_tokens & {"created", "opened"} and column_tokens & {"created", "opened"}: - score += 35 + score = len(query_tokens & column_tokens) * 12 + score += 20 if score > 0: candidates.append((score, column["name"])) @@ -2601,14 +1940,6 @@ def _choose_order_by_column( if not order_tokens: return None - if order_tokens & {"id", "number", "no"} and query_tokens & {"ticket"}: - column = _choose_ranked_column_by_tokens( - columns, - {"id", "no", "number", "ticket"}, - ) - if column: - return column["name"] - column = _choose_ranked_column_by_tokens(columns, order_tokens) return column["name"] if column else None @@ -2626,125 +1957,67 @@ def _choose_dimension_columns( columns: list[dict[str, str]], max_columns: int = 3, ) -> list[str]: - dimension_specs = [ - ({"board", "model"}, {"board", "model"}), - ({"business", "unit"}, {"business", "unit", "bu", "division", "company"}), - ({"customer"}, {"customer", "cust", "name"}), - ({"supplier"}, {"supplier", "vendor", "name"}), - ({"email", "address"}, {"email", "email1", "address", "mail"}), - ({"invoice"}, {"invoice", "inv", "number", "no", "id"}), - ({"product"}, {"product", "prod", "item", "material", "name"}), - ({"salesperson", "representative"}, {"salesperson", "sales", "person", "rep"}), - ({"technician", "tech"}, {"technician", "tech"}), - ({"ticket"}, {"ticket", "id", "number", "no"}), - ({"location"}, {"location", "site", "area"}), - ({"material"}, {"material", "part", "item"}), - ({"priority", "severity"}, {"priority", "severity", "urgency", "rank"}), - ({"status"}, {"status"}), - ({"preparer"}, {"preparer", "group"}), - ({"reviewer"}, {"reviewer", "group"}), - ({"approval", "approver", "signer"}, {"approval", "approver", "reviewer", "signer"}), - ({"account", "gl", "glaccount", "ledger"}, {"account", "gl", "glaccount", "ledger"}), - ({"currency"}, {"currency", "curr"}), - ({"country"}, {"country"}), - ({"order"}, {"order", "ord", "number"}), - ({"batch"}, {"batch", "id"}), - ] - compound_dimension_triggers = [ - {"board", "model"}, - {"business", "unit"}, - {"email", "address"}, - ] - selected: list[str] = [] - for trigger_tokens, column_tokens in dimension_specs: - trigger_matches = ( - trigger_tokens.issubset(query_tokens) - if trigger_tokens in compound_dimension_triggers - else bool(query_tokens & trigger_tokens) - ) - if trigger_matches: - column = _choose_ranked_column_by_tokens(columns, column_tokens) - if column and column["name"] not in selected: - selected.append(column["name"]) - if len(selected) >= max_columns: - break - return selected + candidates = [] + filtered_query_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS + for index, column in enumerate(columns): + if _is_numeric_type(column["data_type"]): + continue + column_tokens = _column_business_tokens(column) + score = len(filtered_query_tokens & column_tokens) * 10 + if score > 0: + candidates.append((score, index, column["name"])) + candidates.sort(key=lambda item: (-item[0], item[1])) + return [name for _, _, name in candidates[:max_columns]] def _choose_missing_value_column( query_tokens: set[str], columns: list[dict[str, str]], ) -> dict[str, str] | None: - missing_specs = [ - ({"email", "address"}, {"email", "email1", "address", "mail", "first", "primary"}), - ({"customer"}, {"customer", "cust", "number", "id", "no"}), - ({"order"}, {"order", "ord", "number", "id", "no"}), - ({"supplier"}, {"supplier", "vendor", "number", "id", "no"}), - ({"product", "material"}, {"product", "prod", "material", "item", "number", "id"}), - ({"location"}, {"location", "site", "area"}), - ({"status"}, {"status"}), - ] - for trigger_tokens, column_tokens in missing_specs: - if query_tokens & trigger_tokens: - column = _choose_ranked_column_by_tokens(columns, column_tokens) - if column: - return column - return None + return _choose_ranked_column_by_tokens( + columns, + query_tokens - (_GENERIC_SCHEMA_INTENT_TOKENS | _NULL_CHECK_TOKENS), + ) def _choose_count_subject_column( query_tokens: set[str], columns: list[dict[str, str]], ) -> dict[str, str] | None: - subject_specs = [ - ({"order"}, {"order", "ord", "number", "id", "no"}), - ({"invoice"}, {"invoice", "inv", "number", "id", "no"}), - ({"customer"}, {"customer", "cust", "number", "id", "no"}), - ({"reconciliation", "recon"}, {"reconciliation", "recon", "trans", "id"}), - ({"journal"}, {"journal", "entry", "number", "id"}), - ({"account", "gl", "glaccount"}, {"account", "gl", "glaccount", "id"}), - ({"ticket"}, {"ticket", "case", "issue", "id", "number", "no"}), - ( - {"failure"}, - {"failure", "failed", "defect", "code", "line", "status", "sys", "type"}, - ), - ({"repair"}, {"repair", "id", "status"}), - ({"batch"}, {"batch", "id"}), - ] - for trigger_tokens, column_tokens in subject_specs: - if query_tokens & trigger_tokens: - candidate_columns = columns - if trigger_tokens & {"failure"}: - candidate_columns = [ - column for column in columns if not _is_rate_like_column(column) - ] - column = _choose_ranked_column_by_tokens(candidate_columns, column_tokens) - if column: - return column - return None + column = _choose_ranked_column_by_tokens( + columns, + query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS, + ) + if column: + return column + for column in columns: + if not _is_numeric_type(column["data_type"]) and not _is_rate_like_column(column): + return column + return columns[0] if columns else None def _choose_average_measure_column( query_tokens: set[str], columns: list[dict[str, str]], ) -> dict[str, str] | None: - measure_specs = [ - ({"age", "duration", "elapsed"}, _AVERAGE_MEASURE_TOKENS), - ({"rate", "ratio", "percent", "percentage"}, _RATE_METRIC_TOKENS | {"score"}), - ({"amount", "gross", "net", "value"}, {"amount", "gross", "net", "value"}), - ({"balance"}, {"balance", "end", "ending", "value"}), - ({"quantity", "qty"}, {"quantity", "qty"}), + filtered_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS + measure_candidates = [ + column for column in columns if not _is_identifier_like_column(column) ] - for trigger_tokens, column_tokens in measure_specs: - if query_tokens & trigger_tokens: - column = _choose_ranked_column_by_tokens( - columns, - set(column_tokens), - numeric=True, - ) - if column: - return column - return None + column = _choose_ranked_column_by_tokens( + measure_candidates, + filtered_tokens, + numeric=True, + ) + if column: + return column + numeric_columns = [ + column + for column in columns + if _is_numeric_type(column["data_type"]) + and not _is_identifier_like_column(column) + ] + return numeric_columns[0] if len(numeric_columns) == 1 else None def _is_text_type(data_type: str) -> bool: @@ -2792,131 +2065,6 @@ def _select_listing_columns( score += 35 if name == measure_column: score += 12 - if query_tokens & {"order"}: - score += len(tokens & {"order", "ord", "number", "customer", "cust", "item", "product"}) * 18 - score += len(tokens & {"date", "day", "month", "year"}) * 8 - if query_tokens & {"invoice"}: - score += ( - len( - tokens - & { - "amount", - "currency", - "date", - "gross", - "invoice", - "net", - "number", - "status", - "supplier", - "task", - } - ) - * 18 - ) - if query_tokens & {"supplier", "vendor", "email", "address"}: - score += ( - len( - tokens - & { - "address", - "email", - "email1", - "id", - "mail", - "name", - "number", - "supplier", - "vendor", - } - ) - * 18 - ) - if query_tokens & {"customer"}: - score += len(tokens & {"customer", "cust", "name", "number", "id"}) * 18 - if query_tokens & {"product", "material"}: - score += len(tokens & {"product", "prod", "material", "item", "description", "desc"}) * 18 - if query_tokens & {"batch"}: - score += len(tokens & {"batch", "board", "model", "supplier", "id"}) * 18 - if query_tokens & {"repair", "log", "record"}: - score += ( - len(tokens & {"board", "code", "date", "failure", "id", "priority", "status"}) - * 14 - ) - if query_tokens & {"ticket"}: - score += ( - len( - tokens - & { - "case", - "created", - "date", - "id", - "issue", - "modified", - "number", - "priority", - "state", - "status", - "ticket", - "updated", - } - ) - * 16 - ) - if query_tokens & {"journal", "workflow", "approval", "approver", "reviewer", "signer"}: - score += ( - len( - tokens - & { - "approval", - "approver", - "date", - "doc", - "entry", - "journal", - "number", - "reviewer", - "signer", - "status", - "workflow", - } - ) - * 16 - ) - if query_tokens & {"reconciliation", "recon", "reconcile"}: - score += ( - len( - tokens - & { - "account", - "gl", - "glaccount", - "group", - "period", - "preparer", - "recon", - "reviewer", - "status", - } - ) - * 16 - ) - if query_tokens & {"account", "gl", "glaccount", "balance"}: - score += ( - len(tokens & {"account", "balance", "endbalance", "gl", "glaccount", "month", "year"}) - * 16 - ) - if tokens & {"repair", "failure", "failed", "defect"} and not query_tokens & { - "repair", - "failure", - "defect", - }: - score -= 40 - if tokens & {"date", "day", "month", "year"} and not ( - _is_date_type(column["data_type"]) or name == date_column - ): - score -= 12 if score > 0: scored_columns.append((score, index, name)) @@ -2946,6 +2094,11 @@ def _fallback_month_filter(query: str) -> tuple[int, int] | None: def _grouping_phrase_tokens(query: str) -> set[str]: + query = re.sub( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", + "", + query, + ) match = re.search( r"(?is)\b(?:grouped\s+by|group\s+by|by)\s+(?P[A-Za-z0-9_ /-]+)", query, @@ -3017,134 +2170,132 @@ def _clean_filter_value(value: str | None) -> str | None: return value or None -def _extract_failure_type_filter_value(query: str) -> str | None: - patterns = [ - ( - r"(?is)\bwith\s+" - r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80}?)" - r"(?:\s+(?:listed|marked|recorded|shown|set))?" - r"\s+as\s+(?:the\s+)?(?:failure|defect)\s+" - r"(?:type|code|category)\b" - ), - ( - r"(?is)\b(?:failure|defect)\s+(?:type|code|category)\s*" - r"(?:=|is|equals|like|of)\s*['\"]?" - r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80})" - ), - ] - for pattern in patterns: - match = re.search(pattern, query) - if match: - value = _clean_filter_value(match.group("value")) - if value: - return value - return None +def _query_content_tokens(query_tokens: set[str]) -> set[str]: + return { + token + for token in query_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + and token not in _MONTH_NAME_TO_NUMBER + } -def _extract_status_filter_values(query: str) -> list[str]: - values: list[str] = [] +def _has_grouping_intent(query: str, query_tokens: set[str]) -> bool: + query_without_ordering = re.sub( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", + "", + query, + ) + return bool( + _is_distribution_metric_intent(query_tokens) + or query_tokens & {"group", "grouped", "per"} + or re.search(r"(?i)\bby\s+[A-Za-z0-9_ -]+\b", query_without_ordering) + ) - def add(value: str): - if value not in values: - values.append(value) - - if re.search(r"(?i)\bin-progress\b", query): - add("in-progress") - add("in progress") - if re.search(r"(?i)\bin\s+progress\b", query): - add("in progress") - add("in-progress") - for status in ( - "blocked", - "closed", - "completed", - "complete", - "escalated", - "open", - "pending", - "resolved", - ): - if re.search(rf"(?i)\b{re.escape(status)}\b", query): - for value in _STATUS_VALUE_ALIASES.get(status, (status,)): - add(value) - return values +def _has_count_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _COUNT_METRIC_TOKENS) -def _extract_status_filter_value(query: str) -> str | None: - values = _extract_status_filter_values(query) - return values[0] if values else None +def _has_sum_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _SUM_METRIC_TOKENS) -def _extract_priority_filter_value(query: str) -> str | None: - for token, value in _PRIORITY_VALUE_ALIASES.items(): - if re.search(rf"(?i)\b{re.escape(token)}(?:[-\s]+priority)?\b", query): - return value - return None +def _has_extreme_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & (_MAX_METRIC_TOKENS | _MIN_METRIC_TOKENS)) -def _choose_priority_column(columns: list[dict[str, str]]) -> dict[str, str] | None: - return _choose_ranked_column_by_tokens( - columns, - {"priority", "severity", "urgency", "rank"}, - ) +def _has_latest_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _LATEST_METRIC_TOKENS) -def _priority_order_expression(column: dict[str, str]) -> str: - quoted_column = _quote_identifier(column["name"]) - if _is_numeric_type(column["data_type"]): - return quoted_column - when_clauses = " ".join( - f"WHEN {_quote_literal(value)} THEN {rank}" for value, rank in _PRIORITY_ORDER - ) - return f"CASE LOWER({quoted_column}) {when_clauses} ELSE 0 END" +def _has_missing_value_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _NULL_CHECK_TOKENS) + +def _sort_direction_for_query(query_tokens: set[str]) -> str: + return "ASC" if query_tokens & _MIN_METRIC_TOKENS else "DESC" -def _choose_failure_type_filter_column( + +def _choose_numeric_measure_column( + query_tokens: set[str], columns: list[dict[str, str]], ) -> dict[str, str] | None: + content_tokens = _query_content_tokens(query_tokens) + measure_candidates = [ + column for column in columns if not _is_identifier_like_column(column) + ] column = _choose_ranked_column_by_tokens( - [column for column in columns if not _is_rate_like_column(column)], - {"failure", "type"}, + measure_candidates, + content_tokens, + numeric=True, ) if column: return column - return _choose_ranked_column_by_tokens( - [column for column in columns if not _is_rate_like_column(column)], - {"failure", "defect", "code", "type", "sys"}, - ) + numeric_columns = [ + column + for column in columns + if _is_numeric_type(column["data_type"]) + and not _is_identifier_like_column(column) + ] + if len(numeric_columns) == 1: + return numeric_columns[0] + return None -def _choose_failure_type_filter_table( - schema_details: dict[str, list[dict[str, str]]], - concept_tokens: set[str], -) -> tuple[str, list[dict[str, str]], dict[str, str]] | None: - candidates = [] - for table_name, columns in schema_details.items(): - if not _table_covers_requested_concepts(table_name, columns, concept_tokens): - continue - column = _choose_failure_type_filter_column(columns) - if not column: - continue - table_tokens = _fallback_tokens(table_name) - column_tokens = _fallback_tokens(column["name"]) - score = len(concept_tokens & table_tokens) * 8 - score += len(concept_tokens & column_tokens) * 10 - if {"failure", "type"}.issubset(column_tokens): - score += 100 - elif "type" in column_tokens: - score += 60 - elif "code" in column_tokens: - score += 30 - if table_tokens & {"failure", "defect"}: - score += 25 - candidates.append((score, table_name, columns, column)) - if not candidates: - return None - candidates.sort(key=lambda item: (-item[0], item[1], item[3]["name"])) - _, table_name, columns, column = candidates[0] - return table_name, columns, column +def _sample_value_filters( + query_tokens: set[str], + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + filters: list[tuple[dict[str, Any], list[str]]] = [] + consumed_tokens: set[str] = set() + content_tokens = _query_content_tokens(query_tokens) + if not content_tokens: + return filters + + for column in columns: + matches: list[str] = [] + for value in column.get("sample_values") or []: + cleaned_value = _clean_filter_value(str(value)) + if not cleaned_value: + continue + value_tokens = _fallback_tokens(cleaned_value) + if not value_tokens or not value_tokens <= content_tokens: + continue + if value_tokens <= consumed_tokens: + continue + if cleaned_value.lower() not in {item.lower() for item in matches}: + matches.append(cleaned_value) + consumed_tokens.update(value_tokens) + if matches: + filters.append((column, matches)) + + return filters + + +def _where_clause(predicates: list[str]) -> str: + return f"\nWHERE {' AND '.join(predicates)}" if predicates else "" + + +def _count_expression_for_query( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> tuple[str, list[str]]: + subject_column = _choose_count_subject_column(query_tokens, columns) + if not subject_column: + return "COUNT(*)", [] + return ( + f"COUNT({_quote_identifier(subject_column['name'])})", + [_non_missing_value_predicate(subject_column)], + ) + + +def _date_bucket_expressions(date_column: str) -> tuple[str, str]: + quoted_date = _quote_identifier(date_column) + return ( + f"EXTRACT(YEAR FROM {quoted_date})", + f"EXTRACT(MONTH FROM {quoted_date})", + ) def generate_simple_analytics_sql( @@ -3154,133 +2305,41 @@ def generate_simple_analytics_sql( if not query: return None - raw_query_tokens = _fallback_tokens(query) - query_tokens = _expanded_fallback_query_tokens(query) + query_tokens = _fallback_tokens(query) if not query_tokens: return None - if not query_tokens & { - "batch", - "batches", - "account", - "accounts", - "address", - "age", - "approval", - "approvals", - "approver", - "average", - "balance", - "balances", - "blocked", - "board", - "breakdown", - "business", - "case", - "closed", - "completed", - "count", - "created", - "customer", - "defect", - "distribution", - "duration", - "email", - "failure", - "failures", - "gl", - "glaccount", - "gross", - "highest", - "invoice", - "invoices", - "journal", - "july", - "latest", - "ledger", - "log", - "logs", - "location", - "material", - "mean", - "missing", - "modified", - "model", - "monthly", - "most", - "net", - "number", - "open", - "order", - "orders", - "pending", - "preparer", - "priority", - "product", - "rate", - "recent", - "recon", - "reconciliation", - "reconciliations", - "reviewer", - "record", - "records", - "repair", - "repairs", - "revenue", - "sale", - "sales", - "severity", - "signer", - "supplier", - "status", - "task", - "tasks", - "tech", - "technician", - "ticket", - "tickets", - "top", - "trend", - "trends", - "type", - "unit", - "units", - "updated", - "workflow", - "workflows", - "year", - }: + schema_details = _extract_schema_details(contexts) + if not schema_details: return None - schema_details = _extract_schema_details(contexts) - rate_metric_intent = _is_rate_metric_intent(raw_query_tokens) - average_metric_intent = _is_average_metric_intent(raw_query_tokens) - distribution_metric_intent = _is_distribution_metric_intent(raw_query_tokens) - failure_count_intent = _is_failure_count_intent(raw_query_tokens, query_tokens) - board_model_intent = _has_board_model_intent( - raw_query_tokens - ) or _has_board_model_intent( - query_tokens - ) - failure_type_filter_value = _extract_failure_type_filter_value(query) - failure_type_filter_column = None - failure_type_choice = None - if failure_type_filter_value and "failure" in query_tokens: - failure_type_choice = _choose_failure_type_filter_table( - schema_details, - raw_query_tokens, + content_tokens = _query_content_tokens(query_tokens) + schema_backed_tokens = _schema_derived_query_tokens(query_tokens, schema_details) + unsupported_tokens = _unsupported_query_tokens(query_tokens, schema_details) + if unsupported_tokens: + logger.info( + "Schema-derived SQL fallback skipped unsupported_tokens=%s", + sorted(unsupported_tokens), ) - - if failure_type_choice: - table_name, columns, failure_type_filter_column = failure_type_choice - chosen = (table_name, columns) - else: - chosen = _choose_fallback_table( - query_tokens, - schema_details, - concept_tokens=raw_query_tokens, + return None + if content_tokens and not schema_backed_tokens: + logger.info( + "Schema-derived SQL fallback skipped no_schema_backed_tokens=%s", + sorted(content_tokens), + ) + return None + if not content_tokens and len(schema_details) != 1: + logger.info( + "Schema-derived SQL fallback skipped ambiguous_schema_only_request tables=%s", + sorted(schema_details), ) + return None + + chosen = _choose_fallback_table( + query_tokens, + schema_details, + concept_tokens=query_tokens, + ) if not chosen: return None @@ -3288,644 +2347,286 @@ def generate_simple_analytics_sql( column_names = [column["name"] for column in columns] quoted_table = _quote_identifier(table_name) limit = _fallback_limit(query) + date_column = _choose_temporal_column(query_tokens, columns) + order_column = _choose_order_by_column(query, query_tokens, columns) + sample_filters = _sample_value_filters(query_tokens, columns) + sample_predicates = [ + _filter_predicate_for_values(column, values) + for column, values in sample_filters + ] + selected_sample_filter_columns = [column["name"] for column, _ in sample_filters] + month_filter = _fallback_month_filter(query) + metric_intent = { + "average": _is_average_metric_intent(query_tokens), + "count": _has_count_intent(query_tokens), + "distribution": _is_distribution_metric_intent(query_tokens), + "extreme": _has_extreme_intent(query_tokens), + "latest": _has_latest_intent(query_tokens), + "missing": _has_missing_value_intent(query_tokens), + "rate": _is_rate_metric_intent(query_tokens), + "sum": _has_sum_intent(query_tokens), + } logger.info( - "Deterministic SQL fallback selected table=%s verified_columns=%s metric_intent=%s", + "Schema-derived SQL fallback selected table=%s schema_tokens=%s verified_columns=%s sample_filter_columns=%s metric_intent=%s", table_name, + sorted(schema_backed_tokens), column_names, - { - "failure_count": failure_count_intent, - "rate": rate_metric_intent, - "average": average_metric_intent, - "distribution": distribution_metric_intent, - "board_model": board_model_intent, - "failure_type_filter": bool(failure_type_filter_value), - }, + selected_sample_filter_columns, + metric_intent, ) - if failure_type_filter_value: - if not failure_type_filter_column: - failure_type_filter_column = _choose_failure_type_filter_column(columns) - if failure_type_filter_column: - return ( - f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}\n" - f"WHERE {_value_match_predicate(failure_type_filter_column, failure_type_filter_value)}" - ) - - material_column = _choose_column_by_tokens(columns, {"material"}) - location_column = _choose_column_by_tokens(columns, {"location"}) - if ( - {"material", "location"}.issubset(query_tokens) - and material_column - and location_column - ): - return ( - f"SELECT {_quote_joined([material_column, location_column])}\n" - f"FROM {quoted_table}" - ) - - date_column_tokens = {"date", "day", "month", "time", "year"} - if raw_query_tokens & {"year"} and not raw_query_tokens & {"month", "monthly"}: - date_column_tokens = {"date", "day", "time", "year"} - date_column = _choose_temporal_column(raw_query_tokens | query_tokens, columns) - if not date_column and not raw_query_tokens & {"updated", "modified"}: - date_column = _choose_column_by_tokens( - columns, - date_column_tokens, - date=True, - ) - - repair_filter_intent = raw_query_tokens & { - "closed", - "completed", - "critical", - "escalated", - "high", - "low", - "medium", - "normal", - "open", - "pending", - "priority", - "progress", - "severity", - "status", - "urgent", - } - if distribution_metric_intent: - dimension_columns = _choose_dimension_columns( - raw_query_tokens | {"status"}, - columns, - max_columns=1, - ) - if dimension_columns: - subject_column = _choose_count_subject_column(raw_query_tokens, columns) - count_expression = "COUNT(*)" - predicates = [] - if subject_column: - count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" - predicates.append(_non_missing_value_predicate(subject_column)) - status_filter_values = _extract_status_filter_values(query) - if status_filter_values: - status_filter_column = _choose_status_filter_column( - columns, - status_filter_values, - ) - if not status_filter_column: - return None - predicates.append( - _filter_predicate_for_values(status_filter_column, status_filter_values) - ) - where_clause = f"\nWHERE {' AND '.join(predicates)}" if predicates else "" - quoted_dimensions = _quote_joined(dimension_columns) - return ( - f"SELECT {quoted_dimensions}, {count_expression} AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" - f"ORDER BY {_quote_identifier('record_count')} DESC" - ) - - if query_tokens & {"repair", "repairs"} and repair_filter_intent: - predicates = [] - status_filter_values = _extract_status_filter_values(query) - if status_filter_values: - status_filter_column = _choose_status_filter_column( - columns, - status_filter_values, - ) - if not status_filter_column: - return None - predicates.append( - _filter_predicate_for_values(status_filter_column, status_filter_values) - ) - priority_filter_value = _extract_priority_filter_value(query) - if priority_filter_value: - priority_column = _choose_priority_column(columns) - if not priority_column or not _column_supports_filter_values( - priority_column, - [priority_filter_value], - ): - return None - predicates.append(_filter_predicate_for_values(priority_column, [priority_filter_value])) - if predicates: - return f"SELECT *\nFROM {quoted_table}\nWHERE {' AND '.join(predicates)}" - - if query_tokens & {"ticket"} and not ( - average_metric_intent - or distribution_metric_intent - or failure_count_intent - or raw_query_tokens & (_COUNT_METRIC_TOKENS | _RATE_METRIC_TOKENS) - ): - predicates = [] - status_filter_values = _extract_status_filter_values(query) - if status_filter_values: - status_filter_column = _choose_status_filter_column( - columns, - status_filter_values, - ) - if not status_filter_column: - return None - predicates.append( - _filter_predicate_for_values(status_filter_column, status_filter_values) - ) - - priority_filter_value = _extract_priority_filter_value(query) - if priority_filter_value: - priority_filter_column = _choose_priority_column(columns) - if not priority_filter_column or not _column_supports_filter_values( - priority_filter_column, - [priority_filter_value], - ): - return None - predicates.append( - _filter_predicate_for_values(priority_filter_column, [priority_filter_value]) - ) - - order_column = _choose_order_by_column(query, raw_query_tokens, columns) + if _has_missing_value_intent(query_tokens): + missing_column = _choose_missing_value_column(query_tokens, columns) + if not missing_column: + return None selected_columns = _select_listing_columns( - raw_query_tokens | {"id", "status", "ticket"}, + query_tokens, columns, date_column=date_column, max_columns=8, ) - for required_column in [order_column]: - if required_column and required_column not in selected_columns: - selected_columns.insert(0, required_column) - if not selected_columns: - selected_columns = column_names[:8] - where_clause = f"\nWHERE {' AND '.join(predicates)}" if predicates else "" - order_clause = ( - f"\nORDER BY {_quote_identifier(order_column)} ASC" if order_column else "" - ) + if missing_column["name"] not in selected_columns: + selected_columns.insert(0, missing_column["name"]) + predicates = [*sample_predicates, _missing_value_predicate(missing_column)] limit_clause = f"\nLIMIT {limit}" if limit else "" - if predicates or order_clause or raw_query_tokens & {"all", "list", "show"}: - return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{where_clause}{order_clause}{limit_clause}" - ) + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" + ) + + month_predicate = "" + if date_column and month_filter: + year, month = month_filter + start_date = f"{year:04d}-{month:02d}-01" + end_year = year + 1 if month == 12 else year + end_month = 1 if month == 12 else month + 1 + end_date = f"{end_year:04d}-{end_month:02d}-01" + quoted_date = _quote_identifier(date_column) + month_predicate = ( + f"{quoted_date} >= '{start_date}' AND {quoted_date} < '{end_date}'" + ) - priority_column = _choose_priority_column(columns) - if average_metric_intent: - average_measure_column = _choose_average_measure_column(raw_query_tokens, columns) - if not average_measure_column: + predicates = list(sample_predicates) + if month_predicate: + predicates.append(month_predicate) + + if _is_average_metric_intent(query_tokens): + measure_column = _choose_average_measure_column(query_tokens, columns) + if not measure_column: return None - dimension_columns = _choose_dimension_columns(raw_query_tokens, columns) - where_predicates = [] - if raw_query_tokens & {"failure", "failed", "defect"}: - subject_column = _choose_count_subject_column({"failure"}, columns) - if subject_column: - where_predicates.append(_non_missing_value_predicate(subject_column)) - where_clause = ( - f"\nWHERE {' AND '.join(where_predicates)}" if where_predicates else "" + grouping_tokens = _grouping_phrase_tokens(query) or query_tokens + dimension_columns = _choose_dimension_columns( + grouping_tokens, + columns, + max_columns=2, ) - aggregate_expr = f"AVG({_quote_identifier(average_measure_column['name'])})" + aggregate_expr = f"AVG({_quote_identifier(measure_column['name'])})" if dimension_columns: quoted_dimensions = _quote_joined(dimension_columns) return ( f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier('average_value')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" f"ORDER BY {_quote_identifier('average_value')} DESC" ) return ( f"SELECT {aggregate_expr} AS {_quote_identifier('average_value')}\n" - f"FROM {quoted_table}{where_clause}" + f"FROM {quoted_table}{_where_clause(predicates)}" ) - if ( - priority_column - and raw_query_tokens & {"priority", "severity"} - and raw_query_tokens & {"bottom", "highest", "lowest", "top"} - ): - selected_columns = _select_listing_columns( - raw_query_tokens | {"priority", "repair", "status"}, - columns, - date_column=date_column, - max_columns=8, - ) - if priority_column["name"] not in selected_columns: - selected_columns.insert(0, priority_column["name"]) - direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" - limit_clause = f"\nLIMIT {limit}" if limit else "" + if date_column and query_tokens & {"month", "monthly"} and _has_count_intent(query_tokens): + year_expr, month_expr = _date_bucket_expressions(date_column) return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"ORDER BY {_priority_order_expression(priority_column)} {direction}" - f"{limit_clause}" - ) - - if date_column and raw_query_tokens & {"latest", "recent"}: - selected_columns = _select_listing_columns( - raw_query_tokens | {"date", "repair", "status"}, - columns, - date_column=date_column, - max_columns=8, - ) - if date_column not in selected_columns: - selected_columns.insert(0, date_column) - limit_clause = f"\nLIMIT {limit}" if limit else "" - return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" + f"GROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" ) + measure_column = _choose_numeric_measure_column(query_tokens, columns) if ( - query_tokens & {"repair", "repairs"} - and raw_query_tokens & {"priority", "severity", "status"} - and re.search(r"(?i)\bby\s+(?:priority|severity|status)\b", query) - ): - dimension_column = _choose_dimension_column(raw_query_tokens, columns) - subject_column = _choose_count_subject_column({"repair"}, columns) - if dimension_column: - count_expression = "COUNT(*)" - where_clause = "" - if subject_column: - count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" - where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" - quoted_dimension = _quote_identifier(dimension_column) - return ( - f"SELECT {quoted_dimension}, {count_expression} AS " - f"{_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier('record_count')} DESC" - ) - measure_column = None - if rate_metric_intent and query_tokens & {"defect", "failure"}: - measure_column = _choose_column_by_tokens( - columns, - {"defect", "rate"}, - numeric=True, - ) - if not measure_column: - measure_column = _choose_column_by_tokens(columns, {"rate"}, numeric=True) - if not measure_column and query_tokens & {"order", "orders"}: - measure_column = _choose_column_by_tokens( - columns, - {"amount", "intake", "sales", "value"}, - numeric=True, - ) - if not measure_column and query_tokens & {"invoice", "invoices"} and raw_query_tokens & { - "amount", - "bottom", - "gross", - "highest", - "lowest", - "net", - "top", - "total", - "value", - }: - if query_tokens & {"gross"}: - measure_column = _choose_column_by_tokens( - columns, - {"gross", "amount", "value"}, - numeric=True, - ) - elif query_tokens & {"net"}: - measure_column = _choose_column_by_tokens( - columns, - {"net", "amount", "value"}, - numeric=True, - ) - if not measure_column: - measure_column = _choose_column_by_tokens( - columns, - {"amount", "gross", "net", "value"}, - numeric=True, - ) - if not measure_column and query_tokens & {"revenue", "sale", "sales"}: - measure_column = _choose_column_by_tokens( - columns, - {"amount", "intake", "revenue", "sales", "value"}, - numeric=True, - ) - if not measure_column and query_tokens & {"balance", "balances"}: - measure_column = _choose_column_by_tokens( - columns, - {"balance", "end", "ending", "value"}, - numeric=True, - ) - if not measure_column and query_tokens & {"amount", "gross", "net", "value"}: - measure_column = _choose_column_by_tokens( - columns, - {"amount", "gross", "net", "value"}, - numeric=True, - ) - if ( - not measure_column - and not failure_count_intent - and query_tokens & {"top", "highest", "lowest", "bottom"} + measure_column + and date_column + and query_tokens & {"month", "monthly"} + and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) ): - measure_column = _choose_column_by_tokens( - columns, - { - "amount", - "balance", - "cost", - "count", - "gross", - "margin", - "net", - "quantity", - "rate", - "score", - "value", - }, - numeric=True, - ) - - if raw_query_tokens & {"missing", "blank", "empty", "null"}: - missing_column = _choose_missing_value_column(raw_query_tokens, columns) - if missing_column: - selected_columns = [ - column - for column in column_names - if column == missing_column["name"] - or _fallback_tokens(column) - & { - "batch", - "business", - "bu", - "customer", - "cust", - "date", - "id", - "location", - "name", - "number", - "ord", - "order", - "product", - "status", - "supplier", - } - ][:8] - if missing_column["name"] not in selected_columns: - selected_columns.insert(0, missing_column["name"]) - limit_clause = f"\nLIMIT {limit}" if limit else "" - return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"WHERE {_missing_value_predicate(missing_column)}{limit_clause}" - ) - - dimension_listing_intent = bool( - raw_query_tokens & {"associated", "association", "associations", "each", "list", "show"} - ) and not bool( - raw_query_tokens - & ( - _AVERAGE_METRIC_TOKENS - | _COUNT_METRIC_TOKENS - | _RATE_METRIC_TOKENS - | {"highest", "latest", "lowest", "recent", "top"} + year_expr, month_expr = _date_bucket_expressions(date_column) + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" + f"GROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" ) - ) - if dimension_listing_intent: - dimension_columns = _choose_dimension_columns(raw_query_tokens, columns, max_columns=3) - if len(dimension_columns) >= 2: - quoted_dimensions = _quote_joined(dimension_columns) - order_clause = ", ".join(_quote_identifier(column) for column in dimension_columns) - return ( - f"SELECT DISTINCT {quoted_dimensions}\n" - f"FROM {quoted_table}\nORDER BY {order_clause}" - ) - explicit_grouping_intent = bool(raw_query_tokens & {"group", "grouped"}) or bool( - re.search(r"(?i)\bby\s+[A-Za-z0-9_ -]+\b", query) - ) - explicit_measure_intent = bool( - raw_query_tokens - & { - "amount", - "balance", - "gross", - "average", - "age", - "duration", - "margin", - "net", - "rate", - "revenue", - "sale", - "sales", - "score", - "value", - } - ) if ( - explicit_grouping_intent - and not explicit_measure_intent - and not ( - raw_query_tokens & {"month", "monthly"} - and raw_query_tokens & {"count", "number"} - ) + measure_column + and date_column + and "year" in query_tokens + and "month" not in query_tokens + and "monthly" not in query_tokens + and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) ): - grouping_tokens = _grouping_phrase_tokens(query) or raw_query_tokens - dimension_columns = _choose_dimension_columns(grouping_tokens, columns) - if dimension_columns: - subject_column = _choose_count_subject_column(raw_query_tokens, columns) - count_expression = "COUNT(*)" - where_clause = "" - if subject_column: - count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" - where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" - quoted_dimensions = _quote_joined(dimension_columns) - limit_clause = f"\nLIMIT {limit}" if limit else "" - return ( - f"SELECT {quoted_dimensions}, {count_expression} AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" - f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" - ) - - if date_column and raw_query_tokens & {"month", "monthly"} and raw_query_tokens & { - "count", - "number", - }: year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" - month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" - subject_column = _choose_count_subject_column(raw_query_tokens, columns) - count_expression = "COUNT(*)" - where_clause = "" - if subject_column: - count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" - where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" + aggregate, alias = _aggregate_for_measure(measure_column["name"]) return ( f"SELECT {year_expr} AS {_quote_identifier('year')}, " - f"{month_expr} AS {_quote_identifier('month')}, " - f"{count_expression} AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\n" - f"GROUP BY {year_expr}, {month_expr}\n" - f"ORDER BY {year_expr}, {month_expr}" + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" + f"GROUP BY {year_expr}\nORDER BY {year_expr}" ) - implied_count_by_dimension = "failure" in query_tokens and bool( - raw_query_tokens & {"location", "material", "technician", "tech"} - or (board_model_intent and not rate_metric_intent) - ) - if ( - not average_metric_intent - and ( - raw_query_tokens & {"count", "number"} - or failure_count_intent - or implied_count_by_dimension + if _has_grouping_intent(query, query_tokens) or _has_count_intent(query_tokens): + grouping_tokens = _grouping_phrase_tokens(query) or query_tokens + max_dimensions = 1 if _has_extreme_intent(query_tokens) else 3 + dimension_columns = _choose_dimension_columns( + grouping_tokens, + columns, + max_columns=max_dimensions, ) - ): - dimension_columns = _choose_dimension_columns(raw_query_tokens, columns) if dimension_columns: - subject_column = _choose_count_subject_column(raw_query_tokens, columns) - count_expression = "COUNT(*)" - where_clause = "" - if subject_column and raw_query_tokens & {"order", "customer"}: - count_expression = ( - f"COUNT(DISTINCT {_quote_identifier(subject_column['name'])})" - ) - elif subject_column and raw_query_tokens & { - "account", - "batch", - "failure", - "gl", - "glaccount", - "invoice", - "journal", - "recon", - "reconciliation", - "repair", - }: - count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" - where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" quoted_dimensions = _quote_joined(dimension_columns) + if measure_column and (_has_sum_intent(query_tokens) or _is_rate_metric_intent(query_tokens)): + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column['name'])})" + direction = _sort_direction_for_query(query_tokens) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {quoted_dimensions}, {count_expression} AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimensions}\n" + f"SELECT {quoted_dimensions}, COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" ) + if _has_count_intent(query_tokens) and not _has_grouping_intent(query, query_tokens): + return ( + f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" + ) + return None - dimension_column = _choose_dimension_column(raw_query_tokens, columns) - if measure_column and dimension_column and rate_metric_intent: - aggregate, alias = _aggregate_for_measure(measure_column) - quoted_dimension = _quote_identifier(dimension_column) - aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" - direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" - limit_clause = f"\nLIMIT {limit}" if limit else "" + if measure_column and _has_sum_intent(query_tokens): return ( - f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + f"SELECT SUM({_quote_identifier(measure_column['name'])}) AS {_quote_identifier('total_value')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" ) - if measure_column and dimension_column and raw_query_tokens & { - "bottom", - "highest", - "lowest", - "top", - }: - aggregate, alias = _aggregate_for_measure(measure_column) - quoted_dimension = _quote_identifier(dimension_column) - aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" - direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" + if _has_latest_intent(query_tokens): + if not date_column: + return None + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if date_column not in selected_columns: + selected_columns.insert(0, date_column) limit_clause = f"\nLIMIT {limit}" if limit else "" - where_clause = _current_year_where_clause(date_column, columns, raw_query_tokens) return ( - f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}{where_clause}\n" - f"GROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" ) - if measure_column and limit and dimension_column and raw_query_tokens & { - "board", - "business", - "customer", - "location", - "material", - "model", - "product", - "salesperson", - "supplier", - "unit", - }: - aggregate, alias = _aggregate_for_measure(measure_column) - quoted_dimension = _quote_identifier(dimension_column) - aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" - return ( - f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier(alias)} DESC\nLIMIT {limit}" - ) + if _has_extreme_intent(query_tokens): + direction = _sort_direction_for_query(query_tokens) + if measure_column: + dimension_column = _choose_dimension_column(query_tokens, columns) + limit_clause = f"\nLIMIT {limit}" if limit else "" + if dimension_column: + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {_quote_identifier(dimension_column)}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {_quote_identifier(dimension_column)}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + selected_columns = _select_listing_columns( + query_tokens, + columns, + measure_column=measure_column["name"], + date_column=date_column, + max_columns=8, + ) + if measure_column["name"] not in selected_columns: + selected_columns.insert(0, measure_column["name"]) + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(measure_column['name'])} {direction}{limit_clause}" + ) + if order_column: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if order_column not in selected_columns: + selected_columns.insert(0, order_column) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(order_column)} {direction}{limit_clause}" + ) + return None - month_filter = _fallback_month_filter(query) - if date_column and month_filter: - year, month = month_filter - start = f"{year:04d}-{month:02d}-01" - end_year = year + 1 if month == 12 else year - end_month = 1 if month == 12 else month + 1 - end = f"{end_year:04d}-{end_month:02d}-01" + if month_predicate and date_column: selected_columns = _select_listing_columns( - raw_query_tokens, + query_tokens, columns, - measure_column=measure_column, + measure_column=measure_column["name"] if measure_column else None, date_column=date_column, - ) - order_clause = ( - f"\nORDER BY {_quote_identifier(measure_column)} DESC" - if measure_column - else f"\nORDER BY {_quote_identifier(date_column)} DESC" + max_columns=8, ) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"WHERE {_quote_identifier(date_column)} >= '{start}' " - f"AND {_quote_identifier(date_column)} < '{end}'" - f"{order_clause}{limit_clause}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" ) - if ( - measure_column - and date_column - and raw_query_tokens & {"year"} - and not raw_query_tokens & {"month", "monthly"} - ): - date_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" - return ( - f"SELECT {date_expr} AS {_quote_identifier('year')}, " - f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" - f"FROM {quoted_table}\nGROUP BY {date_expr}\nORDER BY {date_expr}" + if order_column: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, ) - - if ( - measure_column - and date_column - and raw_query_tokens & {"month", "monthly", "trend", "trends"} - ): - year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" - month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" + if order_column not in selected_columns: + selected_columns.insert(0, order_column) + limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {year_expr} AS {_quote_identifier('year')}, " - f"{month_expr} AS {_quote_identifier('month')}, " - f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" - f"FROM {quoted_table}\nGROUP BY {year_expr}, {month_expr}\n" - f"ORDER BY {year_expr}, {month_expr}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(order_column)} ASC{limit_clause}" ) - if measure_column and limit: - selected_columns = [ - column - for column in column_names - if column == measure_column - or _fallback_tokens(column) - & { - "batch", - "board", - "customer", - "id", - "model", - "name", - "number", - "supplier", - } - ][:6] - if measure_column not in selected_columns: - selected_columns.append(measure_column) + if sample_predicates or query_tokens & {"all", "list"}: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"ORDER BY {_quote_identifier(measure_column)} DESC\nLIMIT {limit}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" ) return None @@ -4828,10 +3529,10 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. - If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. -- Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema contains specific modeled business columns for the requested entity, measure, status, date, or dimension. -- Prefer exact modeled business fields over generic text search. For example, if a status/severity/date/material/location/customer/order/revenue concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. -- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. This applies to metrics such as defect rate, revenue, amount, sales value, count, cost, margin, and quantity. -- For sales or revenue questions, prefer exact declared sales/revenue/value/amount fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. +- Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata contains specific modeled columns for the requested entity, measure, filter, date, or dimension. +- Prefer exact modeled fields over generic text search. If a requested concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. +- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. +- Do not prefer or exclude any business domain by built-in rules. Ground every choice in the DATABASE SCHEMA supplied for this request. """ @@ -4867,9 +3568,8 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) -- If the user asks for a specific date, please give the date range in SQL query - - example: "What is the total revenue for the month of 2024-11-01?" - - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" +- If the user asks for a specific date, use a date range over an exact date/time column from DATABASE SCHEMA. + - example: filter an exact date/time column with a start timestamp and the next boundary timestamp for the requested period. - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 0f6db6911e..1c7ffea6bd 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -33,86 +33,6 @@ _MAX_RETRIEVED_TABLE_NAMES = 24 _MAX_RELATED_TABLE_EXPANSION_DEPTH = 1 _RANK_TOKEN = re.compile(r"[a-z0-9]+") -_GENERIC_TABLE_TOKENS = { - "audit", - "auth", - "calendar", - "config", - "dim", - "dimension", - "file", - "files", - "ingestion", - "job", - "jobs", - "log", - "logs", - "lookup", - "mbr", - "member", - "members", - "migration", - "migrations", - "preference", - "preferences", - "queue", - "report", - "reports", - "setting", - "settings", - "state", - "time", - "user", - "users", -} -_CUSTOMS_FINANCE_TOKENS = { - "claim", - "claims", - "custom", - "customs", - "duty", - "duties", - "hmf", - "import", - "imports", - "mpf", - "refund", - "refunds", - "tariff", - "tariffs", -} -_SALES_REVENUE_TOKENS = { - "amount", - "intake", - "revenue", - "sale", - "sales", - "salesvalue", - "value", -} -_TICKET_SCHEMA_TOKENS = { - "activity", - "case", - "issue", - "ticket", -} -_STATUS_STATE_TOKENS = { - "blocked", - "closed", - "completed", - "open", - "progress", - "stage", - "state", - "status", -} -_UPDATED_TIME_TOKENS = { - "changed", - "date", - "modified", - "time", - "updated", -} table_columns_selection_system_prompt = """ @@ -146,12 +66,12 @@ 16. Prefer the set of deployed models, views, metrics, columns, and relationships that best support the current question. 17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. 18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. -19. Prefer tables and columns that directly model the requested business entities, measures, statuses, dates, identifiers, and dimensions. Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema provides specific modeled columns for the same concept. -20. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, tickets, activities, repairs, failures, batches, materials, locations, status, blocked/open/closed values, priority, severity, age, duration, average, distribution, updated/modified dates, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. +19. Prefer tables and columns whose supplied names, descriptions, relationships, metrics, or sample values directly support the requested entities, measures, filters, dates, identifiers, and dimensions. Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata provides specific modeled columns for the same requested concept. +20. Compare the user's requested entities, measures, filters, dates, and dimensions only with schema metadata supplied for the active project. Do not use built-in business synonym lists. 21. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. -22. Never return placeholder table or column names such as tablename, table_name, dbo.tablename, BatchId, Material, Location, or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. -23. If the request asks for revenue, sales, or sales trends, prefer exact business measure columns named like Revenue, SalesValue, USDFXSalesValue, FXSalesValue, IntakeValue, Amount, or equivalent modeled sales fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. -24. If the request asks for an explicit average, rate, ratio, percentage, revenue, amount, sales value, age, duration, or other named measure and the schema already contains that exact measure column, use the declared measure column directly. Do not use a rate column to answer "most failures", "number of failures", or other count-of-records requests unless the question explicitly asks for a rate/ratio/percentage. Do not use count-based columns or grouped counts to answer average requests unless the question asks for a count. +22. Never return placeholder table or column names or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. +23. If a requested measure, dimension, filter, or time field is not represented by retrieved schema metadata, leave it unsupported instead of substituting a similar-looking field. +24. Metric intent such as count, sum, average, minimum, maximum, ranking, date bucketing, and grouping must be satisfied by declared columns or metric fields from the retrieved schema. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -877,31 +797,6 @@ def score(table_name: str) -> int: value += 8 if direct_table_matches and direct_column_matches: value += 8 - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - value += len((table_tokens | column_tokens) & _SALES_REVENUE_TOKENS) * 5 - if not query_tokens & _CUSTOMS_FINANCE_TOKENS: - value -= ( - len((table_tokens | column_tokens) & _CUSTOMS_FINANCE_TOKENS) - * 8 - ) - if query_tokens & {"ticket"}: - table_and_columns = table_tokens | column_tokens | comment_tokens - if table_and_columns & _TICKET_SCHEMA_TOKENS: - value += len(table_and_columns & _TICKET_SCHEMA_TOKENS) * 12 - else: - value -= 40 - if query_tokens & _STATUS_STATE_TOKENS: - value += len(table_and_columns & _STATUS_STATE_TOKENS) * 8 - if query_tokens & {"updated", "modified"}: - value += len(table_and_columns & _UPDATED_TIME_TOKENS) * 8 - if table_tokens & {"team", "user", "users"} and not ( - table_and_columns & _TICKET_SCHEMA_TOKENS - ): - value -= 40 - if query_tokens & {"updated", "modified"}: - value += len((table_tokens | column_tokens | comment_tokens) & _UPDATED_TIME_TOKENS) * 5 - if table_tokens & _GENERIC_TABLE_TOKENS and not direct_table_matches: - value -= 6 return value ranked = sorted( @@ -939,102 +834,7 @@ def _rank_documents_for_query( def _augment_retrieval_query(query: str) -> str: - lowered = query.lower() - expansions = [] - - concept_terms = { - ("revenue", "sales", "sale", "amount", "value"): ( - "sales revenue amount value gross net total price intake invoice order" - ), - ("order", "orders"): ( - "order ord number date customer product business unit division company" - ), - ("invoice", "invoices"): ( - "invoice supplier customer currency gross net amount date month year number status task" - ), - ("supplier", "suppliers", "vendor", "vendors"): ( - "supplier vendor name number id email address invoice amount" - ), - ("email", "emails", "address", "addresses"): ( - "email address mail first primary supplier contact" - ), - ("reconciliation", "reconciliations", "recon", "reconcile"): ( - "reconciliation recon account gl status preparer reviewer group period" - ), - ("journal", "journals", "workflow", "approval", "approvals"): ( - "journal workflow approval approver reviewer signer status date posting entry document" - ), - ("account", "accounts", "gl", "ledger"): ( - "account gl glaccount ledger balance endbalance ending period year month" - ), - ("balance", "balances", "gross", "net"): ( - "balance endbalance gross net amount value year month" - ), - ("customer", "customers"): ( - "customer account client number name identifier" - ), - ("product", "products"): ( - "product item material type name category" - ), - ("repair", "repairs"): ( - "repair status priority severity failure board model log in progress completed critical age duration" - ), - ("ticket", "tickets", "issue", "case"): ( - "ticket issue case activity log status state blocked open closed priority id number updated modified" - ), - ("failure", "failures", "defect", "defects"): ( - "failure defect severity occurrence record count code type system status age duration" - ), - ("batch", "batches"): ( - "batch board model supplier defect rate inspection status" - ), - ("material", "materials"): ( - "material item part component location" - ), - ("location", "locations"): ( - "location site warehouse area material board model" - ), - ("average", "avg", "mean"): ( - "average avg mean numeric measure age duration elapsed days hours amount rate" - ), - ("age", "duration", "elapsed"): ( - "age duration elapsed days hours numeric measure average" - ), - ("distribution", "breakdown"): ( - "distribution breakdown count group status category" - ), - ("business unit", "bu", "division"): ( - "business unit division company account organization" - ), - ("month", "monthly", "july", "year", "trend", "latest"): ( - "date month year fiscal calendar trend latest recent" - ), - ("updated", "modified"): ( - "updated modified changed date time timestamp month year" - ), - ("status", "severity", "priority", "critical", "blocked", "open", "closed"): ( - "status priority severity critical blocked open closed state category progress" - ), - } - - for triggers, terms in concept_terms.items(): - if any(trigger in lowered for trigger in triggers): - expansions.append(terms) - - if any( - trigger in lowered - for trigger in ("rate", "ratio", "percent", "percentage") - ): - expansions.append("rate ratio percent percentage") - if any(trigger in lowered for trigger in ("average", "avg", "mean")): - expansions.append("average avg mean numeric measure") - if any(trigger in lowered for trigger in ("distribution", "breakdown")): - expansions.append("distribution breakdown group count status category") - - if not expansions: - return query - - return f"{query}\nBusiness schema search terms: {'; '.join(expansions)}" + return query async def _retrieve_semantic_schema_documents( @@ -1691,46 +1491,6 @@ def _lexical_columns_and_tables_needed( comment_tokens = _tokenize_schema_text(column.get("comment")) score = len(query_tokens & column_tokens) * 10 score += len(query_tokens & comment_tokens) * 2 - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - score += len(column_tokens & _SALES_REVENUE_TOKENS) * 6 - if query_tokens & {"month", "monthly", "year", "july", "date", "latest"}: - score += ( - len(column_tokens & {"date", "day", "month", "year", "time"}) - * 5 - ) - if query_tokens & {"average", "avg", "mean", "age", "duration", "elapsed"}: - score += ( - len( - column_tokens - & {"age", "duration", "elapsed", "days", "hours", "amount", "rate"} - ) - * 7 - ) - if query_tokens & {"distribution", "breakdown"}: - score += len(column_tokens & {"status", "state", "category", "type"}) * 7 - if query_tokens & {"ticket"}: - score += len(column_tokens & _TICKET_SCHEMA_TOKENS) * 9 - if query_tokens & _STATUS_STATE_TOKENS: - score += len(column_tokens & _STATUS_STATE_TOKENS) * 8 - if query_tokens & {"updated", "modified"}: - score += len(column_tokens & _UPDATED_TIME_TOKENS) * 8 - if query_tokens & {"id", "number"}: - score += len(column_tokens & {"id", "no", "number", "ticket"}) * 5 - if query_tokens & {"updated", "modified"}: - score += len(column_tokens & _UPDATED_TIME_TOKENS) * 6 - if query_tokens & {"top", "highest", "lowest", "bottom"}: - measure_tokens = { - "amount", - "count", - "cost", - "margin", - "quantity", - "score", - "value", - } - if query_tokens & {"rate", "ratio", "percent", "percentage"}: - measure_tokens.update({"rate", "ratio", "percent", "percentage"}) - score += len(column_tokens & measure_tokens) * 4 if score > 0: column_scores.append( (score, column["name"], column.get("is_primary_key")) @@ -1739,27 +1499,6 @@ def _lexical_columns_and_tables_needed( if not column_scores and table_score <= 0: continue - if table_tokens & _GENERIC_TABLE_TOKENS and table_score <= 0: - table_score -= 8 - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - if not query_tokens & _CUSTOMS_FINANCE_TOKENS: - table_score -= len(table_tokens & _CUSTOMS_FINANCE_TOKENS) * 10 - table_score += len(table_tokens & _SALES_REVENUE_TOKENS) * 5 - if query_tokens & {"ticket"}: - table_and_columns = table_tokens | { - token - for _, column_name, _ in column_scores - for token in _tokenize_schema_text(column_name) - } - if table_and_columns & _TICKET_SCHEMA_TOKENS: - table_score += len(table_and_columns & _TICKET_SCHEMA_TOKENS) * 15 - else: - table_score -= 45 - if table_tokens & {"team", "user", "users"} and not ( - table_and_columns & _TICKET_SCHEMA_TOKENS - ): - table_score -= 45 - total_score = table_score + sum(score for score, _, _ in column_scores) if total_score <= 0: continue diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index ab318e6872..ac2519f117 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -3,14 +3,15 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, generate_simple_analytics_sql, - normalize_wren_sql_dialect, normalize_sql_with_schema_identifiers, + normalize_wren_sql_dialect, unsupported_schema_generation_result, unsupported_schema_message, validate_sql_against_contexts, validate_sql_semantic_coverage, ) + SCHEMA_CONTEXTS = [ """ CREATE TABLE valid_invoice_comments ( @@ -73,10 +74,23 @@ def test_schema_identifier_normalization_quotes_special_identifiers(): assert 'FROM "valid-order-lines"' in sql -def test_semantic_coverage_rejects_generic_table_for_business_concepts(): +def test_wren_sql_dialect_normalization_handles_top_and_joined_limit(): + assert ( + normalize_wren_sql_dialect("SELECT TOP 10 id1 FROM dbo_mbrTime") + == "SELECT id1 FROM dbo_mbrTime\nLIMIT 10" + ) + assert ( + normalize_wren_sql_dialect( + "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESCLIMIT 10" + ) + == "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESC LIMIT 10" + ) + + +def test_semantic_coverage_rejects_unrepresented_query_terms(): contexts = [ """ - CREATE TABLE dbo_mbrTime ( + CREATE TABLE neutral_records ( id1 INTEGER, id2 INTEGER ); @@ -85,86 +99,51 @@ def test_semantic_coverage_rejects_generic_table_for_business_concepts(): error = validate_sql_semantic_coverage( """ - SELECT id1, COUNT(*) AS failures - FROM dbo_mbrTime + SELECT id1, COUNT(*) AS record_count + FROM neutral_records GROUP BY id1 - ORDER BY failures DESC + ORDER BY record_count DESC LIMIT 10 """, - "Show the top 10 materials with the highest number of failures.", + "Show top 10 records by missing_dimension.", contexts, ) assert error is not None - assert "failure/defect" in error - assert "material" in error - - -def test_unsupported_schema_message_requires_all_requested_concepts(): - contexts = [ - """ - CREATE TABLE dbo_mbrTime ( - id1 INTEGER, - id2 INTEGER - ); - """ - ] - - message = unsupported_schema_message( - "Show the top 10 materials with the highest number of failures.", - contexts, - ) - - assert message is not None - assert "No retrieved table or view" in message - assert "failure/defect" in message - assert "material" in message + assert "missing" in error or "dimension" in error -def test_unsupported_schema_message_rejects_split_failure_technician_without_coverage(): +def test_unsupported_schema_message_reports_partial_coverage(): contexts = [ """ - CREATE TABLE dbo_report_failures ( - id INTEGER, - failure_type VARCHAR + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR ); - """, """ - CREATE TABLE dbo_technicians ( - id INTEGER, - name VARCHAR - ); - """, ] message = unsupported_schema_message( - "Show the number of failures by technician.", + "Show records by phase and unknown_segment.", contexts, ) assert message is not None - assert "failure/defect" in message - assert "technician" in message + assert "unknown" in message or "segment" in message def test_unsupported_schema_generation_result_has_no_invalid_sql(): contexts = [ """ - CREATE TABLE dbo_report_failures ( - id INTEGER, - failure_type VARCHAR + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR ); - """, """ - CREATE TABLE dbo_technicians ( - id INTEGER, - name VARCHAR - ); - """, ] result = unsupported_schema_generation_result( - "Show the number of failures by technician.", + "Show records by unknown_segment.", contexts, data_source="MSSQL", ) @@ -175,13 +154,13 @@ def test_unsupported_schema_generation_result_has_no_invalid_sql(): assert invalid["type"] == "NO_RELEVANT_SQL" assert invalid["sql"] == "" assert invalid["original_sql"] == "" - assert "technician" in invalid["error"] + assert "unknown" in invalid["error"] or "segment" in invalid["error"] def test_post_processor_clears_sql_for_unsupported_schema(): contexts = [ """ - CREATE TABLE dbo_mbrTime ( + CREATE TABLE neutral_records ( id1 INTEGER, id2 INTEGER ); @@ -193,15 +172,15 @@ def test_post_processor_clears_sql_for_unsupported_schema(): post_processor.run( [ """ - SELECT id1, COUNT(*) AS failures - FROM dbo_mbrTime + SELECT id1, COUNT(*) AS record_count + FROM neutral_records GROUP BY id1 - ORDER BY failures DESC + ORDER BY record_count DESC LIMIT 10 """ ], contexts=contexts, - fallback_query="Show the top 10 materials with the highest number of failures.", + fallback_query="Show records by missing_dimension.", data_source="MSSQL", ) ) @@ -212,823 +191,273 @@ def test_post_processor_clears_sql_for_unsupported_schema(): assert result["invalid_generation_result"]["original_sql"] == "" -def test_wren_sql_dialect_normalization_repairs_top_and_joined_limit(): - assert ( - normalize_wren_sql_dialect("SELECT TOP 10 id1 FROM dbo_mbrTime") - == "SELECT id1 FROM dbo_mbrTime\nLIMIT 10" - ) - assert ( - normalize_wren_sql_dialect( - "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESCLIMIT 10" - ) - == "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESC LIMIT 10" - ) - - -def test_repair_fallback_filters_critical_priority_and_in_progress_status(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMPTZ - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show all critical-priority repairs that are currently in progress.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'LOWER("status") IN' in sql - assert "'in progress'" in sql - assert "'in-progress'" in sql - assert 'LOWER("priority") = \'critical\'' in sql - - -def test_repair_fallback_preserves_hyphenated_in_progress_status_value(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show all repairs with a critical priority and an in-progress status.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'LOWER("status") IN' in sql - assert "'in-progress'" in sql - assert "'in progress'" in sql - assert 'LOWER("priority") = \'critical\'' in sql - - -def test_repair_logs_highest_priority_orders_by_verified_priority_column(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Which repair logs have the highest priority?", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'ORDER BY CASE LOWER("priority")' in sql - assert "DESC" in sql - - -def test_critical_priority_repairs_filter_verified_priority_column(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show all critical-priority repairs", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'LOWER("priority") = \'critical\'' in sql - - -def test_repairs_by_status_counts_verified_repair_rows(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show repairs by status", - contexts, - ) - - assert sql is not None - assert 'SELECT "status", COUNT("id") AS "record_count"' in sql - assert 'FROM "dbo_repair_logs"' in sql - assert 'GROUP BY "status"' in sql - - -def test_latest_repair_logs_orders_by_verified_date_column(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show latest repair logs", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'ORDER BY "created_at" DESC' in sql - - -def test_semantic_column_alias_can_satisfy_priority_concept_with_verified_name(): +def test_schema_sample_value_filter_is_grounded_in_metadata(): contexts = [ """ /* WREN RETRIEVED SEMANTIC CONTEXT - {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"repair log records"},"columns":[{"sql_column_name_use_exactly":"Urgency","data_type":"VARCHAR","semantic_context_not_sql_identifier":"priority severity for a repair"}]} + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done","In Progress"]}]} WREN SQL IDENTIFIER CONTRACT */ - CREATE TABLE dbo_work_items ( - id VARCHAR, - Urgency VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Which repair records have the highest priority?", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_work_items"' in sql - assert '"Urgency"' in sql - assert '"priority"' not in sql - - -def test_failure_by_technician_fallback_uses_verified_tech_column(): - contexts = [ - """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - Tech VARCHAR, - Failed VARCHAR, - Material VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show the number of failures by technician.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_DebugEntries_Staging2"' in sql - assert 'SELECT "Tech", COUNT("Failed") AS "record_count"' in sql - assert 'WHERE ("Failed" IS NOT NULL AND "Failed" <> \'\')' in sql - - -def test_failure_by_material_fallback_uses_verified_material_column(): - contexts = [ - """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - Tech VARCHAR, - Failed VARCHAR, - Material VARCHAR + CREATE TABLE work_items ( + item_id VARCHAR, + State VARCHAR, + updated_at TIMESTAMP ); """ ] sql = generate_simple_analytics_sql( - "Show failures by material.", + "Show all work item records with In Progress.", contexts, ) assert sql is not None - assert 'FROM "dbo_DebugEntries_Staging2"' in sql - assert 'SELECT "Material", COUNT("Failed") AS "record_count"' in sql + assert 'FROM "work_items"' in sql + assert 'LOWER("State") = \'in progress\'' in sql -def test_failure_type_value_filter_uses_verified_failure_type_column(): +def test_unverified_filter_value_is_not_invented(): contexts = [ """ - CREATE TABLE dbo_DebugEntries ( - SerialNumber VARCHAR, - FailedAt VARCHAR, - Material VARCHAR - ); - """, - """ - CREATE TABLE dbo_repair_logs ( - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR - ); - """, - """ - CREATE TABLE dbo_report_failures ( - failure_type VARCHAR, - failure_line VARCHAR, - test_name VARCHAR - ); - """, - ] - - sql = generate_simple_analytics_sql( - "Show the number of units with JTAG as the failure type.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_report_failures"' in sql - assert 'COUNT(*) AS "record_count"' in sql - assert "LOWER(\"failure_type\") = 'jtag'" in sql - - -def test_board_models_most_failures_counts_failure_records_not_defect_rate(): - contexts = [ - """ - CREATE TABLE dbo_batch_records ( - board_model VARCHAR, - supplier VARCHAR, - defect_rate DECIMAL + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE work_items ( + item_id VARCHAR, + State VARCHAR ); - """, """ - CREATE TABLE dbo_repair_logs ( - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR - ); - """, ] sql = generate_simple_analytics_sql( - "Show the top 5 board models with the most failures.", + "Show all work item records with Archived.", contexts, ) - - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'SELECT "board_model", COUNT("failure_code") AS "record_count"' in sql - assert '"defect_rate"' not in sql - assert "LIMIT 5" in sql - - -def test_board_models_highest_defect_rate_uses_rate_metric(): - contexts = [ - """ - CREATE TABLE dbo_batch_records ( - board_model VARCHAR, - supplier VARCHAR, - defect_rate DECIMAL - ); - """, - """ - CREATE TABLE dbo_repair_logs ( - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR - ); - """, - ] - - sql = generate_simple_analytics_sql( - "Show the board models with the highest defect rate.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_batch_records"' in sql - assert 'SELECT "board_model", AVG("defect_rate") AS "average_value"' in sql - assert 'ORDER BY "average_value" DESC' in sql - - -def test_semantic_coverage_rejects_rate_for_failure_count_intent(): - contexts = [ - """ - CREATE TABLE dbo_batch_records ( - board_model VARCHAR, - defect_rate DECIMAL - ); - """ - ] - - error = validate_sql_semantic_coverage( - """ - SELECT board_model, defect_rate - FROM dbo_batch_records - ORDER BY defect_rate DESC - LIMIT 5 - """, - "Show the top 5 board models with the most failures.", + message = unsupported_schema_message( + "Show all work item records with Archived.", contexts, ) - assert error is not None - assert "count of failure records" in error - - -def test_repairs_by_technician_requires_one_schema_object_covering_both_concepts(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - failure_code VARCHAR - ); - """, - """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - Tech VARCHAR, - Failed VARCHAR - ); - """, - ] - - message = unsupported_schema_message("Show repairs by technician.", contexts) - + assert sql is None assert message is not None - assert "repair" in message - assert "technician" in message - - -def test_invoice_status_count_uses_verified_task_status_column(): - contexts = [ - """ - CREATE TABLE dbo_PBI_View_Unrecorded_Liabilities_Header ( - invoicenumber VARCHAR, - taskstatus VARCHAR, - invoicedate TIMESTAMP, - grossamount DECIMAL, - suppliername VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql("Show invoices grouped by task status.", contexts) - - assert sql is not None - assert 'FROM "dbo_PBI_View_Unrecorded_Liabilities_Header"' in sql - assert 'SELECT "taskstatus", COUNT("invoicenumber") AS "record_count"' in sql - assert 'GROUP BY "taskstatus"' in sql - - -def test_invoice_month_count_groups_by_verified_invoice_date(): - contexts = [ - """ - CREATE TABLE dbo_PBI_View_Unrecorded_Liabilities_Header ( - invoicenumber VARCHAR, - invoicedate TIMESTAMP, - grossamount DECIMAL, - suppliername VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql("Show invoice counts by invoice month.", contexts) - - assert sql is not None - assert 'EXTRACT(YEAR FROM "invoicedate") AS "year"' in sql - assert 'EXTRACT(MONTH FROM "invoicedate") AS "month"' in sql - assert 'COUNT("invoicenumber") AS "record_count"' in sql - - -def test_top_suppliers_by_gross_amount_uses_verified_amount_measure(): - contexts = [ - """ - CREATE TABLE dbo_PBI_View_Unrecorded_Liabilities_Header ( - suppliername VARCHAR, - invoicenumber VARCHAR, - grossamount DECIMAL - ); - """ - ] - - sql = generate_simple_analytics_sql("List the top suppliers by total gross amount.", contexts) - - assert sql is not None - assert 'SELECT "suppliername", SUM("grossamount") AS "total_value"' in sql - assert 'ORDER BY "total_value" DESC' in sql - - -def test_supplier_email_missing_filter_uses_verified_email_column(): - contexts = [ - """ - CREATE TABLE dbo_SupplierMaster_Email ( - supplierid VARCHAR, - suppliername VARCHAR, - email1 VARCHAR, - email2 VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Which supplier email records are missing their first email address?", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_SupplierMaster_Email"' in sql - assert 'WHERE ("email1" IS NULL OR "email1" = \'\')' in sql - - -def test_reconciliation_count_by_status_and_preparer_group_uses_two_dimensions(): - contexts = [ - """ - CREATE TABLE dbo_PBI_View_Recon_Status ( - transid VARCHAR, - status VARCHAR, - preparergroup VARCHAR, - glaccount VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Count reconciliations by status and preparer group.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_PBI_View_Recon_Status"' in sql - assert 'SELECT "status", "preparergroup", COUNT("transid") AS "record_count"' in sql - assert 'GROUP BY "status", "preparergroup"' in sql - - -def test_gl_accounts_highest_balance_uses_verified_balance_measure(): - contexts = [ - """ - CREATE TABLE dbo_View_Global_Exposure_SAP ( - glaccount VARCHAR, - year INTEGER, - month INTEGER, - endbalance DECIMAL - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show GL accounts with the highest ending balance this year.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_View_Global_Exposure_SAP"' in sql - assert 'SELECT "glaccount", SUM("endbalance") AS "total_value"' in sql - assert 'WHERE "year" = ' in sql - assert 'ORDER BY "total_value" DESC' in sql + assert "archived" in message.lower() -def test_recent_journal_workflow_uses_verified_signer_date(): +def test_grouped_count_uses_verified_dimension_only(): contexts = [ """ - CREATE TABLE dbo_Journals_Workflow ( - journalid VARCHAR, - signer VARCHAR, - signer_status VARCHAR, - signer_date TIMESTAMP + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR, + updated_at TIMESTAMP ); """ ] - sql = generate_simple_analytics_sql("List recent journal workflow approvals.", contexts) + sql = generate_simple_analytics_sql("Show records by phase.", contexts) assert sql is not None - assert 'FROM "dbo_Journals_Workflow"' in sql - assert 'ORDER BY "signer_date" DESC' in sql + assert 'SELECT "phase", COUNT(*) AS "record_count"' in sql + assert 'FROM "event_records"' in sql + assert 'GROUP BY "phase"' in sql -def test_average_failed_unit_age_by_board_model_uses_avg_age_measure(): +def test_average_uses_verified_numeric_measure_not_count(): contexts = [ """ - CREATE TABLE dbo_unit_failures ( - unit_id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - unit_age_days DECIMAL + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR, + age_days DECIMAL ); """ ] - sql = generate_simple_analytics_sql( - "Show the average age of failed units for each board model.", - contexts, - ) + sql = generate_simple_analytics_sql("Show average age by model.", contexts) assert sql is not None - assert 'FROM "dbo_unit_failures"' in sql - assert 'SELECT "board_model", AVG("unit_age_days") AS "average_value"' in sql - assert 'WHERE ("failure_code" IS NOT NULL AND "failure_code" <> \'\')' in sql - assert 'GROUP BY "board_model"' in sql - assert 'COUNT(' not in sql + assert 'SELECT "model_code", AVG("age_days") AS "average_value"' in sql + assert 'GROUP BY "model_code"' in sql + assert "COUNT(" not in sql -def test_average_failed_unit_age_without_age_measure_is_unsupported(): +def test_average_without_verified_measure_is_unsupported(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - unit_id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR ); """ ] - sql = generate_simple_analytics_sql( - "Show the average age of failed units for each board model.", - contexts, - ) - message = unsupported_schema_message( - "Show the average age of failed units for each board model.", - contexts, - ) + sql = generate_simple_analytics_sql("Show average age by model.", contexts) + message = unsupported_schema_message("Show average age by model.", contexts) assert sql is None assert message is not None - assert "age/duration" in message + assert "age" in message.lower() -def test_distribution_of_repairs_across_statuses_groups_counts(): +def test_latest_uses_verified_temporal_column(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP + CREATE TABLE event_records ( + event_id VARCHAR, + event_time TIMESTAMP, + phase VARCHAR ); """ ] - sql = generate_simple_analytics_sql( - "Show the distribution of repairs across completed and in-progress statuses.", - contexts, - ) + sql = generate_simple_analytics_sql("Show latest event records.", contexts) assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'SELECT "status", COUNT("id") AS "record_count"' in sql - assert 'LOWER("status") IN' in sql - assert "'completed'" in sql - assert "'in-progress'" in sql - assert "'in progress'" in sql - assert 'GROUP BY "status"' in sql - assert "SELECT *" not in sql + assert 'FROM "event_records"' in sql + assert 'ORDER BY "event_time" DESC' in sql -def test_board_models_associated_with_locations_requires_verified_columns(): +def test_monthly_count_uses_requested_temporal_column_when_verified(): contexts = [ """ - CREATE TABLE dbo_board_locations ( - board_model VARCHAR, - location VARCHAR, - updated_at TIMESTAMP + CREATE TABLE event_records ( + event_id VARCHAR, + updated_at TIMESTAMP, + created_at TIMESTAMP ); """ ] sql = generate_simple_analytics_sql( - "Show the board models associated with each location.", + "Show the number of event records updated each month.", contexts, ) assert sql is not None - assert 'SELECT DISTINCT "board_model", "location"' in sql - assert 'FROM "dbo_board_locations"' in sql - - -def test_board_models_associated_with_locations_without_location_is_unsupported(): - contexts = [ - """ - CREATE TABLE dbo_parts_catalog ( - board_model VARCHAR, - material VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show the board models associated with each location.", - contexts, - ) - message = unsupported_schema_message( - "Show the board models associated with each location.", - contexts, - ) - - assert sql is None - assert message is not None - assert "location" in message + assert 'EXTRACT(YEAR FROM "updated_at") AS "year"' in sql + assert 'EXTRACT(MONTH FROM "updated_at") AS "month"' in sql + assert 'COUNT(*) AS "record_count"' in sql -def test_blocked_tickets_use_verified_ticket_status_and_order_by_ticket_id(): +def test_order_by_uses_verified_column_and_sample_value(): contexts = [ """ - CREATE TABLE dbo_ticket_records ( - ticket_id VARCHAR, - status VARCHAR, - priority VARCHAR, + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open","Closed"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE case_records ( + case_id VARCHAR, + State VARCHAR, updated_at TIMESTAMP ); - """, """ - CREATE TABLE dbo_team_users ( - id VARCHAR, - active BOOLEAN, - email VARCHAR - ); - """, ] sql = generate_simple_analytics_sql( - "Show all blocked tickets ordered by ticket ID.", + "Show all case records with Open ordered by case ID.", contexts, ) assert sql is not None - assert 'FROM "dbo_ticket_records"' in sql - assert 'LOWER("status") = \'blocked\'' in sql - assert 'ORDER BY "ticket_id" ASC' in sql - assert 'dbo_team_users' not in sql - - -def test_blocked_ticket_activity_kind_without_status_is_unsupported(): - contexts = [ - """ - CREATE TABLE dbo_ticket_activity_log ( - ticket_id VARCHAR, - kind VARCHAR, - created_at TIMESTAMP - ); - """ - ] + assert 'LOWER("State") = \'open\'' in sql + assert 'ORDER BY "case_id" ASC' in sql - sql = generate_simple_analytics_sql( - "Show all blocked tickets ordered by ticket ID.", - contexts, - ) - message = unsupported_schema_message( - "Show all blocked tickets ordered by ticket ID.", - contexts, - ) - - assert sql is None - assert message is not None - assert "status" in message - -def test_semantic_coverage_rejects_blocked_ticket_filter_on_activity_kind(): +def test_top_grouped_count_is_schema_shape_based(): contexts = [ """ - CREATE TABLE dbo_ticket_activity_log ( - ticket_id VARCHAR, - kind VARCHAR, - created_at TIMESTAMP - ); - """ - ] - - error = validate_sql_semantic_coverage( - """ - SELECT ticket_id, kind - FROM dbo_ticket_activity_log - WHERE LOWER(kind) = 'blocked' - ORDER BY ticket_id - """, - "Show all blocked tickets ordered by ticket ID.", - contexts, - ) - - assert error is not None - assert "status" in error - - -def test_repair_counts_updated_each_month_use_verified_updated_timestamp(): - contexts = [ - """ - CREATE TABLE dbo_repair_logs ( - repair_id VARCHAR, - status VARCHAR, - updated_at TIMESTAMP, - created_at TIMESTAMP + CREATE TABLE occurrence_records ( + occurrence_id VARCHAR, + model_code VARCHAR, + reason_code VARCHAR ); """ ] sql = generate_simple_analytics_sql( - "Show the number of repairs updated each month.", + "Show top 5 occurrence records by model.", contexts, ) assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'EXTRACT(YEAR FROM "updated_at") AS "year"' in sql - assert 'EXTRACT(MONTH FROM "updated_at") AS "month"' in sql - assert 'COUNT("repair_id") AS "record_count"' in sql - assert '"created_at"' not in sql + assert 'SELECT "model_code", COUNT(*) AS "record_count"' in sql + assert 'ORDER BY "record_count" DESC' in sql + assert "LIMIT 5" in sql -def test_repair_counts_updated_each_month_without_updated_timestamp_is_unsupported(): +def test_sum_by_year_uses_verified_measure_and_temporal_column(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - repair_id VARCHAR, - status VARCHAR, - created_at TIMESTAMP + CREATE TABLE transaction_records ( + transaction_id VARCHAR, + account_name VARCHAR, + amount_value DECIMAL, + posted_at TIMESTAMP ); """ ] - sql = generate_simple_analytics_sql( - "Show the number of repairs updated each month.", - contexts, - ) + sql = generate_simple_analytics_sql("Show total amount by year.", contexts) - assert sql is None + assert sql is not None + assert 'EXTRACT(YEAR FROM "posted_at") AS "year"' in sql + assert 'SUM("amount_value") AS "total_value"' in sql -def test_semantic_coverage_rejects_created_timestamp_for_updated_question(): +def test_semantic_coverage_rejects_count_for_average_intent(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - repair_id VARCHAR, - status VARCHAR, - updated_at TIMESTAMP, - created_at TIMESTAMP + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR, + age_days DECIMAL ); """ ] error = validate_sql_semantic_coverage( """ - SELECT EXTRACT(YEAR FROM created_at) AS year, - EXTRACT(MONTH FROM created_at) AS month, - COUNT(repair_id) AS record_count - FROM dbo_repair_logs - GROUP BY EXTRACT(YEAR FROM created_at), EXTRACT(MONTH FROM created_at) + SELECT model_code, COUNT(*) AS record_count + FROM measurement_records + GROUP BY model_code """, - "Show the number of repairs updated each month.", + "Show average age by model.", contexts, ) assert error is not None - assert "updated" in error + assert "average" in error.lower() -def test_semantic_coverage_rejects_count_for_average_intent(): +def test_literal_validation_rejects_values_outside_verified_samples(): contexts = [ """ - CREATE TABLE dbo_unit_failures ( - unit_id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - unit_age_days DECIMAL + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE case_records ( + case_id VARCHAR, + State VARCHAR ); """ ] error = validate_sql_semantic_coverage( """ - SELECT board_model, COUNT(failure_code) AS record_count - FROM dbo_unit_failures - GROUP BY board_model + SELECT case_id + FROM case_records + WHERE LOWER(State) = 'Closed' """, - "Show the average age of failed units for each board model.", + "Show case records with Open.", contexts, ) assert error is not None - assert "average" in error + assert "sample values" in error diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 1bf9630b62..c847e1e0df 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1137,25 +1137,23 @@ def test_column_selection_returns_empty_dict_for_malformed_reply(): assert parsed == {} -def test_retrieval_query_augmentation_adds_business_terms(): - augmented = _augment_retrieval_query("show total revenue by year") +def test_retrieval_query_augmentation_is_schema_neutral(): + query = "show total amount by year" - assert "Business schema search terms" in augmented - assert "sales revenue amount value" in augmented - assert "date month year" in augmented + assert _augment_retrieval_query(query) == query -def test_table_ranking_prefers_business_sales_table_over_generic_or_customs_tables(): +def test_table_ranking_prefers_direct_schema_metadata_overlap(): documents = [ - _schema_document("dbo_mbrTime", ["id1", "id2"]), - _schema_document("CustomsRefundClaim", ["DutyAmount", "ClaimDate"]), - _schema_document("SalesOrderFact", ["USDFXSalesValue", "OrderDate"]), + _schema_document("neutral_records", ["id1", "id2"]), + _schema_document("amount_snapshots", ["amount_value", "snapshot_year"]), + _schema_document("event_records", ["event_code", "event_time"]), ] ranked = _rank_table_names_by_query( - ["dbo_mbrTime", "CustomsRefundClaim", "SalesOrderFact"], + ["neutral_records", "event_records", "amount_snapshots"], documents, - "show total revenue by year", + "show total amount by year", ) - assert ranked[0] == "SalesOrderFact" + assert ranked[0] == "amount_snapshots" From 6fdfac266ed40d1318594ea50908c43e3b66c74d Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 15:11:06 +0000 Subject: [PATCH 1078/1087] Fix schema-driven Ask fallback grounding --- WRENAI_LOCAL_ASK_HANDOFF.md | 17 +- .../src/pipelines/generation/utils/sql.py | 295 +++++++++++++++--- .../retrieval/db_schema_retrieval.py | 12 + wren-ai-service/src/web/v1/services/ask.py | 50 +-- .../src/web/v1/services/ask_feedback.py | 21 +- .../generation/test_sql_schema_grounding.py | 104 +++++- 6 files changed, 411 insertions(+), 88 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index e4e7239a63..49db5e44bc 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -140,6 +140,17 @@ Randomized validation used synthetic selected-project schemas and shuffled quest ## Remaining Blockers -- Full browser/application verification against live local PCB_DB/Orders projects was not rerun in this pass. -- Filter-value grounding depends on sample/enum metadata being available. If a project has no samples/enums for categorical values, the safer behavior is unsupported/clarification rather than guessed filters. -- The local checkout at `D:\WrenAI` may not contain the latest branch source until refreshed from `origin/organization/ask-schema-grounding-20260820`. +- Live local `D:\WrenAI` was refreshed with the branch source because the running app was still using older code. +- Fixed a generic word-form coverage bug where `repairs` did not ground to verified schema token `repair`. +- Fixed SQL table-reference validation so `EXTRACT(YEAR FROM "updated_at")` is not misread as a table reference. +- Changed generated date buckets to `CAST(EXTRACT(... ) AS BIGINT)` because uncast `EXTRACT` passed generation validation but failed live preview result conversion. +- Added schema-derived user-value filtering for the case where exactly one verified categorical column is explicitly mentioned. Values are escaped and attached only to that verified column; identifier columns such as `ticket_id`, `repair_id`, and `org_id` are excluded. +- Tightened dimension selection so identifier columns are not used as grouping dimensions unless an identifier grouping is explicitly requested. +- Live API validation on `org / PCB_DB` now passes: + - `show number of repairs updated each month` -> `dbo_repair_logs.updated_at`, monthly `COUNT(*)`, successful summary. + - `Show repairs by status.` -> `dbo_repair_logs.status`, grouped count, successful summary. + - `Show latest repair logs.` -> `dbo_repair_logs`, date ordering, successful summary. + - `Show the distribution of repairs across completed and in-progress statuses.` -> `dbo_repair_logs.status`, filtered grouped count, successful summary. + - `Show customer revenue by year.` while PCB_DB is selected -> clear `NO_RELEVANT_SQL`, no invalid SQL. + - `Show all blocked tickets ordered by ticket id.` -> clear `NO_RELEVANT_SQL` for missing verified `blocked` concept, no invalid SQL. +- `pytest` is still not installed in the local AI-service venv, so validation used `py_compile`, direct function harnesses, and live API calls. diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 70f793528a..ee078d64c0 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -8,7 +8,7 @@ import sqlparse from haystack import component from pydantic import BaseModel, ConfigDict -from sqlparse.sql import Identifier, IdentifierList, TokenList +from sqlparse.sql import Function, Identifier, IdentifierList, TokenList from sqlparse.tokens import Comment, Keyword from src.core.engine import ( @@ -227,6 +227,7 @@ "YEAR", } _FALLBACK_TOKEN = re.compile(r"[a-z0-9]+") +_QUERY_VALUE_TOKEN = re.compile(r"[A-Za-z0-9]+(?:[-_][A-Za-z0-9]+)*") _FALLBACK_STOPWORDS = { "a", "an", @@ -258,15 +259,41 @@ } +def _fallback_token_variants(token: str) -> set[str]: + token = token.lower() + variants = {token} + + generic_tokens = globals().get("_GENERIC_SCHEMA_INTENT_TOKENS", set()) + if token in generic_tokens: + return variants + + if len(token) > 4 and token.endswith("ies"): + variants.add(token[:-3] + "y") + elif len(token) > 4 and token.endswith("es"): + variants.add(token[:-2]) + elif len(token) > 3 and token.endswith("s"): + variants.add(token[:-1]) + if len(token) > 4 and token.endswith("ed"): + stem = token[:-2] + variants.add(stem) + if len(token) > 5 and token[-3] in {"d", "s", "t", "v", "z"}: + variants.add(token[:-1]) + if len(stem) > 2 and stem[-1] == stem[-2]: + variants.add(stem[:-1]) + if len(token) > 5 and token.endswith("ing"): + stem = token[:-3] + variants.add(stem) + variants.add(stem + "e") + if len(stem) > 2 and stem[-1] == stem[-2]: + variants.add(stem[:-1]) + + return {variant for variant in variants if variant} + + def _expand_fallback_token_variants(tokens: set[str]) -> set[str]: - expanded = set(tokens) + expanded = set() for token in tokens: - if len(token) > 4 and token.endswith("ies"): - expanded.add(token[:-3] + "y") - if len(token) > 4 and token.endswith("es"): - expanded.add(token[:-2]) - if len(token) > 3 and token.endswith("s"): - expanded.add(token[:-1]) + expanded.update(_fallback_token_variants(token)) return expanded @@ -1370,7 +1397,11 @@ def validate_sql_semantic_coverage( schema_tokens.update(_schema_tokens_for_table(relation, columns)) required_tokens = _schema_derived_query_tokens(raw_query_tokens, schema_details) - unsupported_tokens = _unsupported_query_tokens(raw_query_tokens, schema_details) + unsupported_tokens = _unsupported_query_tokens( + raw_query_tokens, + schema_details, + query=query, + ) if unsupported_tokens: return ( "Schema grounding failed. The retrieved schema metadata does not " @@ -1431,7 +1462,11 @@ def unsupported_schema_message( if not schema_details: return None required_tokens = _schema_derived_query_tokens(query_tokens, schema_details) - unsupported_tokens = _unsupported_query_tokens(query_tokens, schema_details) + unsupported_tokens = _unsupported_query_tokens( + query_tokens, + schema_details, + query=query, + ) if unsupported_tokens: return ( "No retrieved table or view in the active project contains verified " @@ -1524,8 +1559,12 @@ def _table_business_tokens( return tokens +def _data_type_base(data_type: str) -> str: + return data_type.upper().split("(", 1)[0].strip() + + def _is_numeric_type(data_type: str) -> bool: - return data_type.upper() in { + return _data_type_base(data_type) in { "BIGINT", "DECIMAL", "DOUBLE", @@ -1544,7 +1583,7 @@ def _is_numeric_type(data_type: str) -> bool: def _is_date_type(data_type: str) -> bool: - return data_type.upper() in { + return _data_type_base(data_type) in { "DATE", "DATETIME", "DATETIME2", @@ -1662,7 +1701,7 @@ def _is_rate_like_column(column: dict[str, str]) -> bool: def _is_identifier_like_column(column: dict[str, str]) -> bool: tokens = _fallback_tokens(column["name"]) - return bool(tokens) and tokens <= {"id", "identifier", "key", "uuid"} + return bool(tokens & {"id", "identifier", "key", "uuid"}) def _quote_joined(identifiers: list[str]) -> str: @@ -1691,26 +1730,32 @@ def _schema_derived_query_tokens( schema_details: dict[str, list[dict[str, str]]], ) -> set[str]: schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) - return { - token - for token in query_tokens - if token not in _GENERIC_SCHEMA_INTENT_TOKENS - and not token.isdigit() - and token in schema_tokens - } + supported_tokens: set[str] = set() + for token in query_tokens: + if token in _GENERIC_SCHEMA_INTENT_TOKENS or token.isdigit(): + continue + supported_tokens.update(_fallback_token_variants(token) & schema_tokens) + return supported_tokens def _unsupported_query_tokens( query_tokens: set[str], schema_details: dict[str, list[dict[str, str]]], + query: str | None = None, ) -> set[str]: schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + user_value_tokens = _schema_driven_user_value_tokens( + query, + query_tokens, + schema_details, + ) return { token for token in query_tokens if token not in _GENERIC_SCHEMA_INTENT_TOKENS and not token.isdigit() - and token not in schema_tokens + and not (_fallback_token_variants(token) & schema_tokens) + and token not in user_value_tokens } @@ -1727,7 +1772,7 @@ def _table_covers_requested_concepts( if not required_tokens: return True schema_tokens = _schema_tokens_for_table(table_name, columns) - return required_tokens <= schema_tokens + return all(_fallback_token_variants(token) & schema_tokens for token in required_tokens) def _choose_fallback_table( @@ -1957,15 +2002,24 @@ def _choose_dimension_columns( columns: list[dict[str, str]], max_columns: int = 3, ) -> list[str]: - candidates = [] + name_candidates = [] + semantic_candidates = [] filtered_query_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS + identifier_requested = bool(filtered_query_tokens & {"id", "identifier", "key", "uuid"}) for index, column in enumerate(columns): if _is_numeric_type(column["data_type"]): continue - column_tokens = _column_business_tokens(column) - score = len(filtered_query_tokens & column_tokens) * 10 - if score > 0: - candidates.append((score, index, column["name"])) + if _is_identifier_like_column(column) and not identifier_requested: + continue + name_tokens = _fallback_tokens(column["name"]) + semantic_tokens = set(column.get("semantic_tokens") or set()) + name_score = len(filtered_query_tokens & name_tokens) * 10 + semantic_score = len(filtered_query_tokens & semantic_tokens) * 3 + if name_score > 0: + name_candidates.append((name_score + semantic_score, index, column["name"])) + elif semantic_score > 0: + semantic_candidates.append((semantic_score, index, column["name"])) + candidates = name_candidates or semantic_candidates candidates.sort(key=lambda item: (-item[0], item[1])) return [name for _, _, name in candidates[:max_columns]] @@ -2021,11 +2075,29 @@ def _choose_average_measure_column( def _is_text_type(data_type: str) -> bool: - return data_type.upper() in {"CHAR", "NCHAR", "NVARCHAR", "STRING", "TEXT", "VARCHAR"} + return _data_type_base(data_type) in { + "CHAR", + "CHARACTER", + "NCHAR", + "NTEXT", + "NVARCHAR", + "STRING", + "TEXT", + "VARCHAR", + } def _is_boolean_type(data_type: str) -> bool: - return data_type.upper() in {"BIT", "BOOL", "BOOLEAN"} + return _data_type_base(data_type) in {"BIT", "BOOL", "BOOLEAN"} + + +def _is_categorical_value_column(column: dict[str, Any]) -> bool: + return not ( + _is_identifier_like_column(column) + or _is_numeric_type(column["data_type"]) + or _is_date_type(column["data_type"]) + or _is_boolean_type(column["data_type"]) + ) def _missing_value_predicate(column: dict[str, str]) -> str: @@ -2130,7 +2202,7 @@ def _current_year_where_clause( "month", "time", }: - return f"\nWHERE EXTRACT(YEAR FROM {quoted_column}) = {current_year}" + return f"\nWHERE {_date_part_expression(date_column, 'YEAR')} = {current_year}" return "" @@ -2154,7 +2226,9 @@ def _value_match_predicate( if not values: return f"{quoted_column} IS NOT NULL" - if isinstance(column, dict) and _is_text_type(column["data_type"]): + if isinstance(column, dict) and ( + _is_text_type(column["data_type"]) or _is_categorical_value_column(column) + ): lowered_values = [_quote_literal(candidate.lower()) for candidate in values] if len(lowered_values) == 1: return f"LOWER({quoted_column}) = {lowered_values[0]}" @@ -2273,6 +2347,123 @@ def _sample_value_filters( return filters +def _mentioned_text_columns_for_query( + query_tokens: set[str], + columns: list[dict[str, Any]], +) -> list[dict[str, Any]]: + mentioned_columns = [] + for column in columns: + if not _is_categorical_value_column(column): + continue + column_tokens = _fallback_tokens(column["name"]) + if any(_fallback_token_variants(token) & column_tokens for token in query_tokens): + mentioned_columns.append(column) + return mentioned_columns + + +def _matched_column_query_tokens( + query_tokens: set[str], + column: dict[str, Any], +) -> set[str]: + column_tokens = _fallback_tokens(column["name"]) + matched_tokens: set[str] = set() + for token in query_tokens: + matched_tokens.update(_fallback_token_variants(token) & column_tokens) + return matched_tokens + + +def _query_schema_value_terms( + query: str | None, + schema_tokens: set[str], +) -> list[str]: + if not query: + return [] + + values: list[str] = [] + for match in _QUERY_VALUE_TOKEN.finditer(query): + raw_value = match.group(0).replace("_", " ") + value_tokens = _fallback_tokens(raw_value) + if not value_tokens: + continue + if value_tokens & _GENERIC_SCHEMA_INTENT_TOKENS: + continue + if value_tokens & schema_tokens: + continue + if all(token.isdigit() for token in value_tokens): + continue + cleaned_value = _clean_filter_value(raw_value) + if cleaned_value and cleaned_value.lower() not in { + value.lower() for value in values + }: + values.append(cleaned_value) + + return values + + +def _schema_driven_user_value_tokens( + query: str | None, + query_tokens: set[str], + schema_details: dict[str, list[dict[str, Any]]], +) -> set[str]: + if not query: + return set() + + mentioned_columns: list[dict[str, Any]] = [] + for columns in schema_details.values(): + mentioned_columns.extend(_mentioned_text_columns_for_query(query_tokens, columns)) + + matched_concepts = [ + _matched_column_query_tokens(query_tokens, column) + for column in mentioned_columns + ] + matched_concepts = [concepts for concepts in matched_concepts if concepts] + if not matched_concepts: + logger.info( + "Schema-derived user value grounding skipped: no mentioned categorical columns query=%s", + query, + ) + return set() + shared_concepts = set.intersection(*matched_concepts) + if not shared_concepts: + logger.info( + "Schema-derived user value grounding skipped: ambiguous categorical concepts query=%s columns=%s concepts=%s", + query, + [column["name"] for column in mentioned_columns], + [sorted(concepts) for concepts in matched_concepts], + ) + return set() + + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + value_tokens: set[str] = set() + value_terms = _query_schema_value_terms(query, schema_tokens) + for value in value_terms: + value_tokens.update(_fallback_tokens(value)) + logger.info( + "Schema-derived user value grounding query=%s shared_concepts=%s values=%s value_tokens=%s columns=%s", + query, + sorted(shared_concepts), + value_terms, + sorted(value_tokens), + [column["name"] for column in mentioned_columns], + ) + return value_tokens + + +def _schema_driven_user_value_filters( + query: str | None, + query_tokens: set[str], + table_name: str, + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + mentioned_columns = _mentioned_text_columns_for_query(query_tokens, columns) + if len(mentioned_columns) != 1: + return [] + + schema_tokens = _schema_tokens_for_table(table_name, columns) + values = _query_schema_value_terms(query, schema_tokens) + return [(mentioned_columns[0], values)] if values else [] + + def _where_clause(predicates: list[str]) -> str: return f"\nWHERE {' AND '.join(predicates)}" if predicates else "" @@ -2291,13 +2482,16 @@ def _count_expression_for_query( def _date_bucket_expressions(date_column: str) -> tuple[str, str]: - quoted_date = _quote_identifier(date_column) return ( - f"EXTRACT(YEAR FROM {quoted_date})", - f"EXTRACT(MONTH FROM {quoted_date})", + _date_part_expression(date_column, "YEAR"), + _date_part_expression(date_column, "MONTH"), ) +def _date_part_expression(date_column: str, part: str) -> str: + return f"CAST(EXTRACT({part} FROM {_quote_identifier(date_column)}) AS BIGINT)" + + def generate_simple_analytics_sql( query: str | None, contexts: list[Any] | None, @@ -2315,7 +2509,11 @@ def generate_simple_analytics_sql( content_tokens = _query_content_tokens(query_tokens) schema_backed_tokens = _schema_derived_query_tokens(query_tokens, schema_details) - unsupported_tokens = _unsupported_query_tokens(query_tokens, schema_details) + unsupported_tokens = _unsupported_query_tokens( + query_tokens, + schema_details, + query=query, + ) if unsupported_tokens: logger.info( "Schema-derived SQL fallback skipped unsupported_tokens=%s", @@ -2350,11 +2548,23 @@ def generate_simple_analytics_sql( date_column = _choose_temporal_column(query_tokens, columns) order_column = _choose_order_by_column(query, query_tokens, columns) sample_filters = _sample_value_filters(query_tokens, columns) + sample_filter_column_names = {column["name"] for column, _ in sample_filters} + user_value_filters = [ + (column, values) + for column, values in _schema_driven_user_value_filters( + query, + query_tokens, + table_name, + columns, + ) + if column["name"] not in sample_filter_column_names + ] + value_filters = [*sample_filters, *user_value_filters] sample_predicates = [ _filter_predicate_for_values(column, values) - for column, values in sample_filters + for column, values in value_filters ] - selected_sample_filter_columns = [column["name"] for column, _ in sample_filters] + selected_sample_filter_columns = [column["name"] for column, _ in value_filters] month_filter = _fallback_month_filter(query) metric_intent = { "average": _is_average_metric_intent(query_tokens), @@ -2471,7 +2681,7 @@ def generate_simple_analytics_sql( and "monthly" not in query_tokens and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) ): - year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" + year_expr = _date_part_expression(date_column, "YEAR") aggregate, alias = _aggregate_for_measure(measure_column["name"]) return ( f"SELECT {year_expr} AS {_quote_identifier('year')}, " @@ -2692,9 +2902,9 @@ def validate_candidate_sql(candidate_sql: str) -> str | None: cleaned_generation_result, ) grounding_invalid_generation_result = { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_GROUNDING", + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", "error": grounding_error, "correlation_id": "", "data_source": data_source, @@ -3185,6 +3395,9 @@ def _extract_table_references(token_list: TokenList) -> tuple[set[str], dict[str if token.is_whitespace or token.ttype in Comment: continue + if isinstance(token, Function): + continue + if isinstance(token, TokenList): nested_tables, nested_aliases = _extract_table_references(token) table_names.update(nested_tables) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 1c7ffea6bd..0893b8dd61 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -578,6 +578,18 @@ def _empty_retrieval_results() -> dict[str, Any]: } +def _merge_names(*name_lists: list[str]) -> list[str]: + merged: list[str] = [] + seen: set[str] = set() + for names in name_lists: + for name in names: + if name in seen: + continue + merged.append(name) + seen.add(name) + return merged + + ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a4a27d3a02..a28247f793 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -3,9 +3,9 @@ from typing import Dict, List, Literal, Optional from cachetools import TTLCache +from langfuse.decorators import observe from pydantic import AliasChoices, BaseModel, Field -from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.utils import trace_metadata from src.web.v1.services import BaseRequest, SSEEvent @@ -25,7 +25,7 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = True + ignore_sql_generation_reasoning: bool = False enable_column_pruning: bool = False use_dry_plan: bool = True allow_dry_plan_fallback: bool = False @@ -99,12 +99,12 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = False, + allow_sql_generation_reasoning: bool = True, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = True, - max_sql_correction_retries: int = 0, + enable_column_pruning: bool = False, + max_sql_correction_retries: int = 3, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -189,30 +189,6 @@ async def ask( is_followup=True if histories else False, ) - historical_question = await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, - ) - - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - - if historical_question_result: - api_results = [ - AskResult( - **{ - "sql": result.get("statement"), - "type": "view" if result.get("viewId") else "llm", - "viewId": result.get("viewId"), - } - ) - for result in historical_question_result - ] - sql_generation_reasoning = "" - if not api_results: # Run both pipeline operations concurrently sql_samples_task, instructions_task = await asyncio.gather( @@ -405,7 +381,6 @@ async def ask( instructions=instructions, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - validation_contexts=table_ddls, configuration=ask_request.configurations, query_id=query_id, ) @@ -419,7 +394,6 @@ async def ask( instructions=instructions, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - validation_contexts=table_ddls, configuration=ask_request.configurations, query_id=query_id, ) @@ -482,7 +456,6 @@ async def ask( histories=histories, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - validation_contexts=table_ddls, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -502,7 +475,6 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - validation_contexts=table_ddls, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -529,9 +501,16 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + if failed_dry_run_result["type"] in ( + "TIME_OUT", + "NO_RELEVANT_SQL", + ): error_message = failed_dry_run_result["error"] - invalid_sql = failed_dry_run_result["sql"] + invalid_sql = ( + "" + if failed_dry_run_result["type"] == "NO_RELEVANT_SQL" + else failed_dry_run_result["sql"] + ) break original_sql = failed_dry_run_result["original_sql"] @@ -584,7 +563,6 @@ async def ask( }, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, - validation_contexts=table_ddls, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 27a8a94aa8..ba73369c01 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -184,6 +184,7 @@ async def ask_feedback( "sql_regeneration" ].run( contexts=table_ddls, + query=ask_feedback_request.question, sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, @@ -211,10 +212,14 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] != "TIME_OUT": + if failed_dry_run_result["type"] == "NO_RELEVANT_SQL": + invalid_sql = "" + error_message = failed_dry_run_result["error"] + elif failed_dry_run_result["type"] != "TIME_OUT": original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + sql_diagnosis_reasoning = None self._ask_feedback_results[ query_id @@ -239,17 +244,23 @@ async def ask_feedback( "post_process" ].get("reasoning") + correction_error_message = error_message + if sql_diagnosis_reasoning: + correction_error_message = ( + f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" + ) + sql_correction_results = await self._pipelines[ "sql_correction" ].run( contexts=table_ddls, query=ask_feedback_request.question, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, instructions=instructions, invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "original_sql": original_sql, + "sql": invalid_sql, + "error": correction_error_message, }, project_id=ask_feedback_request.project_id, mdl_hash=ask_feedback_request.mdl_hash, diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index ac2519f117..77a68b2f22 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -2,6 +2,7 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, + _SchemaCatalog, generate_simple_analytics_sql, normalize_sql_with_schema_identifiers, normalize_wren_sql_dialect, @@ -157,6 +158,79 @@ def test_unsupported_schema_generation_result_has_no_invalid_sql(): assert "unknown" in invalid["error"] or "segment" in invalid["error"] +def test_schema_coverage_accepts_generic_word_form_variants(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = """ + SELECT + CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT) AS year, + CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) AS month, + COUNT(*) AS record_count + FROM work_update_log + GROUP BY + CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT), + CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) + """ + + error = validate_sql_semantic_coverage( + sql, + "Show the number of work updates updated each month.", + contexts, + ) + + assert error is None + assert unsupported_schema_message( + "Show the number of work updates updated each month.", + contexts, + ) is None + + +def test_schema_fallback_uses_verified_monthly_update_timestamp(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the number of work updates updated each month.", + contexts, + ) + + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT)' in sql + assert "COUNT(*)" in sql + + +def test_schema_catalog_ignores_extract_from_column_clause(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + sql = """ + SELECT CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year", COUNT(*) AS "record_count" + FROM "work_update_log" + GROUP BY CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) + """ + + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + + def test_post_processor_clears_sql_for_unsupported_schema(): contexts = [ """ @@ -217,6 +291,30 @@ def test_schema_sample_value_filter_is_grounded_in_metadata(): assert 'LOWER("State") = \'in progress\'' in sql +def test_user_values_are_allowed_for_single_verified_text_column(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + state_name VARCHAR(255), + updated_at TIMESTAMP + ); + """ + ] + + query = ( + "Show the distribution of work updates across completed and " + "in-progress state names." + ) + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'LOWER("state_name") IN (\'completed\', \'in-progress\')' in sql + assert 'GROUP BY "state_name"' in sql + + def test_unverified_filter_value_is_not_invented(): contexts = [ """ @@ -337,8 +435,8 @@ def test_monthly_count_uses_requested_temporal_column_when_verified(): ) assert sql is not None - assert 'EXTRACT(YEAR FROM "updated_at") AS "year"' in sql - assert 'EXTRACT(MONTH FROM "updated_at") AS "month"' in sql + assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year"' in sql + assert 'CAST(EXTRACT(MONTH FROM "updated_at") AS BIGINT) AS "month"' in sql assert 'COUNT(*) AS "record_count"' in sql @@ -405,7 +503,7 @@ def test_sum_by_year_uses_verified_measure_and_temporal_column(): sql = generate_simple_analytics_sql("Show total amount by year.", contexts) assert sql is not None - assert 'EXTRACT(YEAR FROM "posted_at") AS "year"' in sql + assert 'CAST(EXTRACT(YEAR FROM "posted_at") AS BIGINT) AS "year"' in sql assert 'SUM("amount_value") AS "total_value"' in sql From 2f4d8f3607bb44c7205ff3541ab734c4c468b2b2 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 15:46:26 +0000 Subject: [PATCH 1079/1087] Tighten schema-driven Ask semantic validation --- .../src/pipelines/generation/utils/sql.py | 221 +++++++++++++++++- .../generation/test_sql_schema_grounding.py | 159 +++++++++++++ 2 files changed, 372 insertions(+), 8 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index ee078d64c0..9f113d116b 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -270,7 +270,10 @@ def _fallback_token_variants(token: str) -> set[str]: if len(token) > 4 and token.endswith("ies"): variants.add(token[:-3] + "y") elif len(token) > 4 and token.endswith("es"): - variants.add(token[:-2]) + if token.endswith(("ches", "shes", "sses", "uses", "xes", "zes")): + variants.add(token[:-2]) + else: + variants.add(token[:-1]) elif len(token) > 3 and token.endswith("s"): variants.add(token[:-1]) if len(token) > 4 and token.endswith("ed"): @@ -1367,6 +1370,70 @@ def _extract_generation_sql(generation_result: str | None) -> str | None: return text if _SQL_START.search(text) else None +def _sql_has_aggregate_function(sql: str) -> bool: + return bool(re.search(r"(?is)\b(?:AVG|COUNT|MAX|MIN|SUM)\s*\(", sql)) + + +def _group_by_source_columns(sql: str) -> set[str]: + group_by_clause = _extract_clause( + sql, + "GROUP BY", + ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ) + if not group_by_clause: + return set() + + columns: set[str] = set() + for expression in _split_sql_tokens(group_by_clause): + for identifier in _iter_unqualified_identifier_candidates(expression): + upper_identifier = identifier.upper() + if ( + upper_identifier in _SQL_RESERVED_WORDS + or upper_identifier in _SQL_FUNCTION_WORDS + or upper_identifier in _SQL_TYPE_WORDS + or upper_identifier in _DATE_PART_WORDS + ): + continue + columns.add(identifier) + return columns + + +def _query_allows_grouped_aggregate(query: str, query_tokens: set[str]) -> bool: + return bool( + _has_grouping_intent(query, query_tokens) + or _has_count_intent(query_tokens) + or _has_sum_intent(query_tokens) + or _is_rate_metric_intent(query_tokens) + or _is_distribution_metric_intent(query_tokens) + ) + + +def _missing_value_target_tokens(query: str, query_tokens: set[str]) -> set[str]: + match = re.search( + r"(?is)\b(?:blank|empty|missing|null)\s+(?P[A-Za-z0-9_ /-]+)", + query, + ) + if match: + tokens = _fallback_tokens(match.group("value")) + else: + tokens = set(query_tokens) + return tokens - _GENERIC_SCHEMA_INTENT_TOKENS - _NULL_CHECK_TOKENS + + +def _sql_null_checked_columns(sql: str, columns: list[dict[str, str]]) -> set[str]: + checked_columns: set[str] = set() + stripped = _strip_string_literals(sql) + for column in columns: + quoted_column = re.escape(_quote_identifier(column["name"])) + bare_column = re.escape(column["name"]) + column_pattern = rf"(?:{quoted_column}|(? 1: + return ( + "Schema grounding failed. The question asks for one grouping " + "dimension, but the generated SQL groups by multiple source " + "columns." + ) + + if _has_missing_value_intent(raw_query_tokens): + target_tokens = _missing_value_target_tokens(query, raw_query_tokens) + for relation in referenced_relations: + columns = schema_details.get(relation) or [] + if not columns or not target_tokens: + continue + checked_columns = _sql_null_checked_columns(sql, columns) + if not checked_columns: + continue + scored_columns = [ + (_column_score_for_tokens(column, target_tokens), column["name"]) + for column in columns + ] + best_score = max((score for score, _ in scored_columns), default=0) + checked_best_score = max( + ( + score + for score, name in scored_columns + if name in checked_columns + ), + default=0, + ) + if best_score > checked_best_score: + return ( + "Schema grounding failed. The question asks for missing " + "values on a specific schema concept, but the generated SQL " + "checks a weaker matching column for null or blank values." + ) + return _validate_literal_values_against_samples(sql, schema_details, grounding) @@ -1680,6 +1801,7 @@ def _is_date_type(data_type: str) -> bool: "without", "year", } +_GENERIC_SCHEMA_INTENT_TOKENS.update(_NULL_CHECK_TOKENS) _GENERIC_SCHEMA_INTENT_TOKENS.update(_MONTH_NAME_TO_NUMBER.keys()) @@ -1738,12 +1860,61 @@ def _schema_derived_query_tokens( return supported_tokens +def _schema_adjacent_dimension_descriptor_tokens( + query: str | None, + query_tokens: set[str], + schema_tokens: set[str], +) -> set[str]: + if not query: + return set() + dimension_intent_tokens = _DISTRIBUTION_METRIC_TOKENS | { + "across", + "by", + "each", + "group", + "grouped", + "per", + } + if not ( + query_tokens + & dimension_intent_tokens + ): + return set() + + descriptor_tokens: set[str] = set() + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + for index, match in enumerate(matches): + if index == 0: + continue + previous_tokens = _fallback_tokens( + matches[index - 1].group(0).replace("_", " ") + ) + value_tokens = _fallback_tokens(match.group(0).replace("_", " ")) + explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { + "'", + '"', + } + if ( + previous_tokens & schema_tokens + and value_tokens + and not (value_tokens & schema_tokens) + and not explicitly_quoted + ): + descriptor_tokens.update(value_tokens) + return descriptor_tokens + + def _unsupported_query_tokens( query_tokens: set[str], schema_details: dict[str, list[dict[str, str]]], query: str | None = None, ) -> set[str]: schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + descriptor_tokens = _schema_adjacent_dimension_descriptor_tokens( + query, + query_tokens, + schema_tokens, + ) user_value_tokens = _schema_driven_user_value_tokens( query, query_tokens, @@ -1755,6 +1926,7 @@ def _unsupported_query_tokens( if token not in _GENERIC_SCHEMA_INTENT_TOKENS and not token.isdigit() and not (_fallback_token_variants(token) & schema_tokens) + and token not in descriptor_tokens and token not in user_value_tokens } @@ -2025,9 +2197,14 @@ def _choose_dimension_columns( def _choose_missing_value_column( + query: str, query_tokens: set[str], columns: list[dict[str, str]], ) -> dict[str, str] | None: + target_tokens = _missing_value_target_tokens(query, query_tokens) + column = _choose_ranked_column_by_tokens(columns, target_tokens) + if column: + return column return _choose_ranked_column_by_tokens( columns, query_tokens - (_GENERIC_SCHEMA_INTENT_TOKENS | _NULL_CHECK_TOKENS), @@ -2166,18 +2343,28 @@ def _fallback_month_filter(query: str) -> tuple[int, int] | None: def _grouping_phrase_tokens(query: str) -> set[str]: + phrase = _grouping_phrase_text(query) + return _fallback_tokens(phrase) if phrase else set() + + +def _grouping_phrase_has_multiple_dimensions(query: str) -> bool: + phrase = _grouping_phrase_text(query) + return bool(phrase and re.search(r"(?i)(?:,|/|\band\b|\bor\b)", phrase)) + + +def _grouping_phrase_text(query: str) -> str | None: query = re.sub( r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", "", query, ) match = re.search( - r"(?is)\b(?:grouped\s+by|group\s+by|by)\s+(?P[A-Za-z0-9_ /-]+)", + r"(?is)\b(?:grouped\s+by|group\s+by|by|across|per)\s+(?P[A-Za-z0-9_ /-]+)", query, ) if not match: - return set() - return _fallback_tokens(match.group("value")) + return None + return match.group("value") def _current_year_where_clause( @@ -2380,11 +2567,23 @@ def _query_schema_value_terms( return [] values: list[str] = [] - for match in _QUERY_VALUE_TOKEN.finditer(query): + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + for index, match in enumerate(matches): raw_value = match.group(0).replace("_", " ") value_tokens = _fallback_tokens(raw_value) if not value_tokens: continue + previous_tokens = ( + _fallback_tokens(matches[index - 1].group(0).replace("_", " ")) + if index > 0 + else set() + ) + explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { + "'", + '"', + } + if previous_tokens & schema_tokens and not explicitly_quoted: + continue if value_tokens & _GENERIC_SCHEMA_INTENT_TOKENS: continue if value_tokens & schema_tokens: @@ -2586,7 +2785,7 @@ def generate_simple_analytics_sql( ) if _has_missing_value_intent(query_tokens): - missing_column = _choose_missing_value_column(query_tokens, columns) + missing_column = _choose_missing_value_column(query, query_tokens, columns) if not missing_column: return None selected_columns = _select_listing_columns( @@ -2691,8 +2890,11 @@ def generate_simple_analytics_sql( ) if _has_grouping_intent(query, query_tokens) or _has_count_intent(query_tokens): - grouping_tokens = _grouping_phrase_tokens(query) or query_tokens + explicit_grouping_tokens = _grouping_phrase_tokens(query) + grouping_tokens = explicit_grouping_tokens or query_tokens max_dimensions = 1 if _has_extreme_intent(query_tokens) else 3 + if explicit_grouping_tokens and not _grouping_phrase_has_multiple_dimensions(query): + max_dimensions = 1 dimension_columns = _choose_dimension_columns( grouping_tokens, columns, @@ -2753,8 +2955,11 @@ def generate_simple_analytics_sql( if _has_extreme_intent(query_tokens): direction = _sort_direction_for_query(query_tokens) if measure_column: - dimension_column = _choose_dimension_column(query_tokens, columns) limit_clause = f"\nLIMIT {limit}" if limit else "" + if _query_allows_grouped_aggregate(query, query_tokens): + dimension_column = _choose_dimension_column(query_tokens, columns) + else: + dimension_column = None if dimension_column: aggregate, alias = _aggregate_for_measure(measure_column["name"]) return ( diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index 77a68b2f22..9b8d9f26fc 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -315,6 +315,27 @@ def test_user_values_are_allowed_for_single_verified_text_column(): assert 'GROUP BY "state_name"' in sql +def test_column_value_label_is_not_treated_as_literal_filter_value(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + state_name VARCHAR(255), + updated_at TIMESTAMP + ); + """ + ] + + query = "Show the distribution of work updates across state name values." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'GROUP BY "state_name"' in sql + assert "WHERE" not in sql + + def test_unverified_filter_value_is_not_invented(): contexts = [ """ @@ -488,6 +509,91 @@ def test_top_grouped_count_is_schema_shape_based(): assert "LIMIT 5" in sql +def test_single_grouping_dimension_does_not_over_split_results(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show number of events by account.", contexts) + + assert sql is not None + assert 'GROUP BY "account_name"' in sql + assert "account_reference" not in sql + + +def test_missing_value_intent_uses_verified_plural_name_column(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + event_time TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show events with missing account names.", + contexts, + ) + + assert sql is not None + assert 'FROM "account_events"' in sql + assert '"account_name" IS NULL' in sql + + +def test_semantic_validation_rejects_weaker_null_check_column(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT account_name, account_reference + FROM account_events + WHERE account_reference IS NULL + """, + "Show events with missing account names.", + contexts, + ) + + assert error is not None + assert "weaker matching column" in error + + +def test_top_records_are_listed_without_implicit_grouped_aggregate(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show top 10 scored events from July.", contexts) + + assert sql is not None + assert 'FROM "scored_events"' in sql + assert "GROUP BY" not in sql + assert 'ORDER BY "score_value" DESC' in sql + assert "LIMIT 10" in sql + + def test_sum_by_year_uses_verified_measure_and_temporal_column(): contexts = [ """ @@ -559,3 +665,56 @@ def test_literal_validation_rejects_values_outside_verified_samples(): assert error is not None assert "sample values" in error + + +def test_semantic_validation_rejects_multi_group_for_single_dimension(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT account_name, account_reference, COUNT(*) AS record_count + FROM account_events + GROUP BY account_name, account_reference + """, + "Show number of events by account.", + contexts, + ) + + assert error is not None + assert "one grouping dimension" in error + + +def test_semantic_validation_rejects_top_record_grouped_aggregate(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT category_name, SUM(score_value) AS total_value + FROM scored_events + GROUP BY category_name + ORDER BY total_value DESC + LIMIT 10 + """, + "Show top 10 scored events from July.", + contexts, + ) + + assert error is not None + assert "grouped aggregate" in error From e890ac4621633d013876bd8d2134d4b3a8738681 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Thu, 20 Aug 2026 11:37:01 +0000 Subject: [PATCH 1080/1087] Improve Ask schema grounding --- WRENAI_LOCAL_ASK_HANDOFF.md | 347 ++- .../generation/followup_sql_generation.py | 2 +- .../pipelines/generation/sql_correction.py | 27 +- .../pipelines/generation/sql_generation.py | 2 +- .../src/pipelines/generation/utils/sql.py | 2442 +++++++---------- .../retrieval/db_schema_retrieval.py | 500 +++- .../generation/test_sql_schema_grounding.py | 638 ++--- .../retrieval/test_db_schema_retrieval.py | 22 +- 8 files changed, 1958 insertions(+), 2022 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index 49db5e44bc..e9f23efcda 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -1,156 +1,283 @@ -# WrenAI Ask Schema Grounding Handoff +# WrenAI Ask Grounding Handoff Date: 2026-08-20 ## Current Goal -Make WrenAI Ask schema-driven for the active org/project. User wording may search and rank verified metadata, but final SQL must use only verified tables, columns, relationships, metrics, and supported values from the selected project's retrieved schema. +Make WrenAI's Ask pipeline schema-first for the active org/project. Natural language should search and rank verified schema metadata, but final SQL must use only tables, views, columns, relationships, metrics, and values supported by the selected project's metadata. -Do not fix future issues with exact prompt handling, project-specific branches, table/column mappings, or static business synonym/value lists. +Do not fix future issues by hardcoding one question, project, table, column, or organization. Representative prompts such as `Which repair logs have the highest priority?` are regression examples only. -## Branch / PR +## Final Runtime State -- Repository: `hbalasubramanya-rgb/WrenAI` -- Branch: `organization/ask-schema-grounding-20260820` -- PR: `https://github.com/hbalasubramanya-rgb/WrenAI/pull/1` -- Previous pushed commit before this handoff update: `4a040c817` (`Improve ticket status grounding`) -- Current worktree: `D:\WrenAI-ask-e2e-fix-20260820` +- UI: `http://127.0.0.1:3000` +- AI service: `http://127.0.0.1:5555` +- AI health: `{"status":"ok"}` +- Active project restored after validation: `org / PCB_DB` +- Active project id: `10` +- Orders project id: `11` +- Sales duplicate: not shown in current project list; `Orders` remains canonical. -The worktree is detached at the branch tip because the local branch is checked out in another worktree. Commit from this detached worktree and push with: +Current projects visible through `/api/v1/projects/current`: -```powershell -git push origin HEAD:organization/ask-schema-grounding-20260820 -``` - -Do not stage the unrelated mode-only change in `wren-ui/.yarn/releases/yarn-4.5.3.cjs`. +- id `4`, unnamed DuckDB +- id `10`, `PCB_DB` +- id `11`, `Orders` +- id `12`, `CWPay` +- id `13`, `CW_GL` ## What Changed Today -### Static Business Maps Removed +### Generic Schema Grounding -Files: +Permanent source changes are now in `D:\WrenAI\wren-ai-service`, not only `.codex-tmp`. + +Main file: - `wren-ai-service/src/pipelines/generation/utils/sql.py` + +Added or improved: + +- SQL identifier validation against retrieved schema. +- Semantic coverage validation so valid identifiers are not enough; the referenced table/view must also support the requested business concepts. +- Unsupported-schema result helper that returns `NO_RELEVANT_SQL` with no invented SQL. +- Deterministic schema-grounded fallback for common Ask families: + - count / grouped counts + - top-N + - highest / lowest + - latest / recent + - priority / severity + - status filters + - date/month/year filters + - revenue/sales measures + - failure counts vs defect-rate metrics + - failure type value filters +- Semantic alias support from Wren retrieved context blocks. +- More timestamp type support, including `TIMESTAMPTZ`, which fixed the live `latest repair logs` failure. +- Normalization of dialect issues such as `TOP n`, joined `DESCLIMIT`, and order-by aliases. +- Logs for generated SQL validation, deterministic fallback SQL, fallback validation, selected table, verified columns, and metric intent. + +### Retrieval Improvements + +Main file: + - `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` -Removed static business-token and value logic, including: +Added or improved: -- `_STATUS_VALUE_ALIASES` -- `_PRIORITY_VALUE_ALIASES` -- `_PRIORITY_ORDER` -- `_requested_business_concepts` -- `_expanded_fallback_query_tokens` -- `_expand_fallback_token_aliases` -- retrieval `concept_terms` -- domain-specific table boosts/deboosts for example business words +- Project-scoped retrieval filters retained and tested. +- Query expansion for business concepts such as repair, failure, revenue, order, material, status, priority, latest, and date. +- Ranking uses table names, column names, descriptions/comments, semantic context, and generic-table deboosting. +- Logs for: + - selected project id + - retrieved candidate tables and scores + - selected schema objects and columns -Production scans now return no matches for those removed helpers/maps or for the audited domain strings in the two Ask grounding files. +### Generation Pipeline Wiring -### Schema-Derived Grounding +Files: -Added generic schema-token extraction and matching from: +- `wren-ai-service/src/pipelines/generation/sql_generation.py` +- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` +- `wren-ai-service/src/pipelines/generation/sql_correction.py` -- table names -- column names -- table/column semantic descriptions -- relationship/identifier context already present in retrieved schema text -- enum/sample values when supplied in Wren semantic context +Changes: -The deterministic fallback is now schema-shape based. It can produce conservative SQL for generic shapes such as: +- Passed the user query into post-processing as `fallback_query`. +- Added pre-LLM unsupported-schema checks where retrieved schema clearly cannot cover requested concepts. +- Ensured SQL correction still uses the same schema-first validation and fallback logic. +- Strengthened correction instructions so invalid or hallucinated identifiers are not preserved. -- grouped counts -- averages over verified numeric measures -- sums over verified numeric measures -- top-N grouped counts -- latest/recent listings over verified temporal fields -- month/year buckets over verified temporal fields -- ordering by verified requested columns -- filters only when values are supported by sample/enum metadata +### UI / Project Cleanup From This Workstream -It no longer contains domain branches for specific business words or values. +Files still dirty from the related UI/runtime fixes: -### Validation Tightened +- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` +- `wren-ui/src/apollo/server/services/askingService.ts` -`validate_sql_semantic_coverage` now rejects SQL when: +Relevant behavior: -- non-operational query terms are not represented anywhere in active project schema metadata/sample values -- generated SQL uses a verified table but not one covering all schema-backed query tokens -- an average request is answered with count-only SQL -- a distribution/breakdown request is not grouped with counts -- a string literal filter on a sampled column uses a value not present in verified samples +- Previous `results` crash handling is preserved. +- Unsupported-schema failures now avoid showing invented SQL as something to fix. +- Sales/Orders cleanup remains in place: UI project list shows `Orders`, not duplicate `Sales`. -`unsupported_schema_message` now reports partial coverage instead of returning `None` just because some query terms matched schema. +## Live Validation Done -### Retrieval Cleanup +All live checks were run through the UI GraphQL Ask path after restarting the AI service. -Retrieval query augmentation is now a no-op. Ranking uses only direct overlap between query tokens and retrieved schema metadata. The table-selection prompt was changed to instruct schema-local reasoning without domain examples or built-in synonym lists. +### PCB_DB -### Tests Updated +Active project: `PCB_DB`, id `10`. -Files: +Passed: + +- `Which repair logs have the highest priority?` + - Table: `dbo_repair_logs` + - Uses verified `priority` + - Orders by generic priority ranking expression +- `Show all critical-priority repairs` + - Table: `dbo_repair_logs` + - Filter: `priority = 'critical'` +- `Show repairs by status` + - Table: `dbo_repair_logs` + - Group: `status` + - Metric: `COUNT(id)` +- `Show latest repair logs` + - Table: `dbo_repair_logs` + - Order: `created_at DESC` + - This was the live regression fixed by adding timestamp type coverage. +- `Show the number of failures by material` + - Uses verified material/failure fields from PCB_DB. +- `Show top 5 board models with the most failures` + - Table: `dbo_repair_logs` + - Metric: `COUNT(failure_code)` + - Did not use `defect_rate`. +- `Show units with JTAG as the failure type` + - Table: `dbo_report_failures` + - Filter: `failure_type = 'JTAG'` +- Extra check: + - `Show all repairs with a critical priority and an in-progress status.` + - Table: `dbo_repair_logs` + - Filters: `status = 'in-progress'` and `priority = 'critical'` + +### Orders + +Temporarily switched active project to `Orders`, id `11`, then restored PCB_DB. + +Passed: + +- `Show top 10 orders from July` + - Uses Orders table/date fields. +- `Show number of orders by customer` + - Groups by customer. + - Counts distinct order numbers. +- `Show revenue by year` + - Uses verified sales/revenue value and invoice date fields. +- Unsupported check: `Which repair logs have the highest priority?` + - Returned `NO_RELEVANT_SQL`. + - No SQL candidate. + - Message clearly said the active project does not contain verified `repair` and `priority/severity` fields. + +## Checks Run + +Passed: + +```powershell +git diff --check -- wren-ai-service/src/pipelines/generation/utils/sql.py ` + wren-ai-service/src/pipelines/generation/sql_generation.py ` + wren-ai-service/src/pipelines/generation/followup_sql_generation.py ` + wren-ai-service/src/pipelines/generation/sql_correction.py ` + wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py ` + wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py ` + wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +``` + +Passed: + +```powershell +cd D:\WrenAI\wren-ai-service +.\venv\Scripts\python.exe -m compileall -q src\pipelines\generation src\pipelines\retrieval ` + tests\pytest\pipelines\generation\test_sql_schema_grounding.py ` + tests\pytest\pipelines\retrieval\test_db_schema_retrieval.py +``` + +Could not run pytest in the service venv: + +```text +D:\WrenAI\wren-ai-service\venv\Scripts\python.exe: No module named pytest +``` + +## Tests Added + +Main test file: - `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` + +Coverage added for: + +- Unsupported schema clears invalid SQL. +- Generic table rejection for unsupported business concepts. +- Repair priority ordering. +- Critical-priority repair filters. +- Repairs by status. +- Latest repair logs with `TIMESTAMPTZ`. +- Semantic alias column support, for example using real verified `Urgency` when semantic context says it means priority/severity. +- Failure by material / technician with verified columns. +- JTAG failure type filters. +- Board models with most failures uses count, not defect rate. +- Highest defect rate uses rate metric. +- Repairs by technician requires one schema object or relationship coverage. + +Retrieval test file: + - `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` -The generation test module now verifies generic behavior rather than PCB/Orders-specific examples: +Coverage added for: -- identifier validation -- unsupported partial schema coverage -- sample-value filters -- unverified values rejected -- grouped counts -- averages vs counts -- latest by temporal column -- monthly counts -- explicit order by -- top grouped count -- sum by year -- literal sample validation +- Project filter conditions. +- Query expansion. +- Table ranking by query and schema text. +- Project-scoped schema retrieval behavior. -The retrieval test now verifies no query expansion and schema-metadata ranking. +## Restart Commands Used -## Validation Run +Restart AI service only: -Commands run from `D:\WrenAI-ask-e2e-fix-20260820`: +```powershell +$taskName = 'WrenAI 04 AI Service' +$listenerProcessIds = Get-NetTCPConnection -LocalPort 5555 -State Listen -ErrorAction SilentlyContinue | + Select-Object -ExpandProperty OwningProcess -Unique +Stop-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue +foreach ($listenerProcessId in $listenerProcessIds) { + if ($listenerProcessId) { + Stop-Process -Id $listenerProcessId -Force -ErrorAction SilentlyContinue + } +} +Start-ScheduledTask -TaskName $taskName +``` + +Health check: ```powershell -python -m py_compile wren-ai-service/src/pipelines/generation/utils/sql.py wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py -git diff --check +Invoke-WebRequest -UseBasicParsing http://127.0.0.1:5555/health ``` -Manual local harness: - -- generation grounding tests: `ran=21 failures=0` -- retrieval touched logic: `ran=2 failures=0` -- randomized schema-derived validation: `ran=10 failures=0` - -`python -m pytest` was not available in the local venv because `pytest` is not installed. - -Randomized validation used synthetic selected-project schemas and shuffled questions covering: - -- grouped count -- average -- latest/recent -- sample-value filter with explicit ordering -- monthly count -- sum by year -- top-N grouped count -- sum by dimension -- unsupported unknown dimension - -## Remaining Blockers - -- Live local `D:\WrenAI` was refreshed with the branch source because the running app was still using older code. -- Fixed a generic word-form coverage bug where `repairs` did not ground to verified schema token `repair`. -- Fixed SQL table-reference validation so `EXTRACT(YEAR FROM "updated_at")` is not misread as a table reference. -- Changed generated date buckets to `CAST(EXTRACT(... ) AS BIGINT)` because uncast `EXTRACT` passed generation validation but failed live preview result conversion. -- Added schema-derived user-value filtering for the case where exactly one verified categorical column is explicitly mentioned. Values are escaped and attached only to that verified column; identifier columns such as `ticket_id`, `repair_id`, and `org_id` are excluded. -- Tightened dimension selection so identifier columns are not used as grouping dimensions unless an identifier grouping is explicitly requested. -- Live API validation on `org / PCB_DB` now passes: - - `show number of repairs updated each month` -> `dbo_repair_logs.updated_at`, monthly `COUNT(*)`, successful summary. - - `Show repairs by status.` -> `dbo_repair_logs.status`, grouped count, successful summary. - - `Show latest repair logs.` -> `dbo_repair_logs`, date ordering, successful summary. - - `Show the distribution of repairs across completed and in-progress statuses.` -> `dbo_repair_logs.status`, filtered grouped count, successful summary. - - `Show customer revenue by year.` while PCB_DB is selected -> clear `NO_RELEVANT_SQL`, no invalid SQL. - - `Show all blocked tickets ordered by ticket id.` -> clear `NO_RELEVANT_SQL` for missing verified `blocked` concept, no invalid SQL. -- `pytest` is still not installed in the local AI-service venv, so validation used `py_compile`, direct function harnesses, and live API calls. +Project switch endpoints used for validation: + +```powershell +Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/11/select +Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/10/select +Invoke-WebRequest -UseBasicParsing http://127.0.0.1:3000/api/v1/projects/current +``` + +## Current Dirty Files To Review + +Relevant tracked files: + +- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` +- `wren-ai-service/src/pipelines/generation/sql_correction.py` +- `wren-ai-service/src/pipelines/generation/sql_generation.py` +- `wren-ai-service/src/pipelines/generation/utils/sql.py` +- `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` +- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` +- `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` +- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` +- `wren-ui/src/apollo/server/services/askingService.ts` + +There are also many local untracked runtime/data artifacts in the repository. Do not clean or delete them casually. + +## Important Caveats + +- Runtime source code should remain generic. Do not add checks for exact prompts such as `Which repair logs have the highest priority?`. +- Tests may use representative table and prompt names; production code must not. +- Retrieval context currently uses metadata/descriptions and some semantic context. It does not appear to carry robust sample-value lists. Status casing/value handling works for tested prompts, but richer value-aware matching would improve future accuracy. +- `enable_column_pruning` was not the focus of today's final validation. +- Full pytest suite still needs an environment with `pytest` installed. + +## Recommended Next Steps + +1. Install or enable pytest in `wren-ai-service\venv`, then run focused tests. +2. Review the large `utils/sql.py` diff carefully; consider extracting fallback/grounding helpers into smaller modules after behavior is stable. +3. Add sample-value metadata to retrieval context if available, then make value matching use that metadata instead of only text normalization. +4. Run a broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when their data sources are available. +5. Commit the source changes after review, excluding local runtime/data artifacts. 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 2da693f867..8a10b3769d 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -163,7 +163,7 @@ async def post_process( generate_sql_in_followup.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - contexts=validation_contexts or documents, + contexts=documents, fallback_query=query, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 13358f7513..8004135449 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -34,20 +34,15 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) ### SQL CORRECTION INSTRUCTIONS ### -1. First, use the error message only to identify which part of the failed SQL was unsupported by DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. -2. Then, generate a syntactically correct Wren SQL query from the user's intent and the current DATABASE SCHEMA. -3. If the invalid SQL contains a table, column, function, literal value, or metadata-table query that is not supported by the current DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. -4. Never correct an invalid SQL query by checking INFORMATION_SCHEMA or system catalogs. -5. If a user question is provided, treat it as the source of intent and regenerate the SQL from that intent using DATABASE SCHEMA instead of repairing guessed identifiers. -6. Treat SQL diagnosis and the invalid SQL as error context only. Do not copy placeholders, assumed table names, assumed column names, or unsupported functions from them. -7. Do not preserve a table, column, join, filter, grouping, ordering, or function from the failed SQL unless it appears exactly in DATABASE SCHEMA or SQL FUNCTIONS. -8. Treat physical/source/lineage names from the failed SQL, error message, reasoning, comments, aliases, descriptions, or samples as semantic context only; never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. -9. If the error is an invalid object, invalid column, unsupported function, or date/type failure, do not try a similar replacement from source metadata. Regenerate from the user's intent and current DATABASE SCHEMA. If the unsupported part is needed to answer the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql instead of substituting non-schema identifiers. -10. If the failed SQL used connector-specific syntax such as TOP, square-bracket identifiers, backticks, or non-Wren identifier quoting, discard that syntax and regenerate using Wren SQL syntax only. -11. For grouped queries, repair SQL Server errors about ORDER BY columns not appearing in GROUP BY by ordering with selected grouping columns or selected aggregate aliases, or by adding the exact ordering key to both SELECT and GROUP BY when that key is declared in DATABASE SCHEMA. -12. Do not preserve generic log, file, JSON, payload, text, or app-metric scans when DATABASE SCHEMA contains exact modeled business columns for the user's requested entity, measure, status, date, or dimension. -13. If the failed SQL invented component fields for a metric that exists directly in DATABASE SCHEMA, replace the calculation with the exact declared metric column. -14. For sales or revenue questions, avoid tariff, duty, customs, import, refund, or claim datasets unless the USER QUESTION explicitly asks for those domains. +1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). +2. Then, generate the syntactically correct ANSI SQL query to correct the error. +3. If the failed SQL references a table, view, column, function, alias, or placeholder that is not present in DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. Regenerate from the USER QUESTION and DATABASE SCHEMA. +4. Treat invalid object name, dataset not found, table not found, invalid column name, and invalid identifier errors as schema-grounding failures. Use exact declared identifiers from DATABASE SCHEMA only. +5. Do not create dummy CTEs, placeholder tables, table-existence checks, or generic replacement names to make the query executable. If the requested intent is supported by retrieved schema objects, use those exact objects; otherwise return null for sql. +6. For grouped queries, repair SQL Server errors about ORDER BY columns not appearing in GROUP BY by ordering with selected grouping columns or selected aggregate aliases, or by adding the exact ordering key to both SELECT and GROUP BY when that key is declared in DATABASE SCHEMA. +7. Do not preserve generic log, file, JSON, payload, text, or app-metric scans when DATABASE SCHEMA contains exact modeled business columns for the user's requested entity, measure, status, date, or dimension. +8. If the failed SQL invented component fields for a metric that exists directly in DATABASE SCHEMA, replace the calculation with the exact declared metric column. +9. For sales or revenue questions, avoid tariff, duty, customs, import, refund, or claim datasets unless the USER QUESTION explicitly asks for those domains. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -167,8 +162,8 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: List[Document] | None = None, query: str | None = None, - documents: list[str] | None = None, project_id: str | None = None, mdl_hash: str | None = None, validation_contexts: list[str] | None = None, @@ -179,7 +174,7 @@ async def post_process( generate_sql_correction.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - contexts=validation_contexts or documents, + contexts=documents, fallback_query=query, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 9e93154c88..da4aae67be 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -155,7 +155,7 @@ async def post_process( generate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - contexts=validation_contexts or documents, + contexts=documents, fallback_query=query, use_dry_plan=use_dry_plan, data_source=data_source, diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 9f113d116b..0ff59e43a2 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -70,28 +70,6 @@ _QUALIFIED_COLUMN = re.compile( rf"(?P{_QUALIFIED_IDENTIFIER})\s*\.\s*(?P{_IDENTIFIER_TOKEN})" ) -_SQL_START = re.compile(r"^\s*(?:WITH|SELECT)\b", re.IGNORECASE | re.DOTALL) -_SQL_OBJECT_ALIAS_STOP_WORDS = { - "CROSS", - "EXCEPT", - "FETCH", - "FULL", - "GROUP", - "HAVING", - "INNER", - "INTERSECT", - "JOIN", - "LEFT", - "LIMIT", - "MATCH_RECOGNIZE", - "NATURAL", - "OFFSET", - "ORDER", - "RIGHT", - "TABLESAMPLE", - "UNION", - "WHERE", -} _UNQUALIFIED_QUOTED_IDENTIFIER = re.compile(r'(?(?:[^"]|"")*)"') _UNQUALIFIED_BARE_IDENTIFIER = re.compile( r"(?[A-Za-z_][A-Za-z0-9_$]*)\b(?!\s*\.)" @@ -170,9 +148,6 @@ "CONCAT", "COUNT", "COUNT_BIG", - "CURRENT_DATE", - "CURRENT_TIME", - "CURRENT_TIMESTAMP", "DATE_TRUNC", "DAY", "EXTRACT", @@ -227,7 +202,6 @@ "YEAR", } _FALLBACK_TOKEN = re.compile(r"[a-z0-9]+") -_QUERY_VALUE_TOKEN = re.compile(r"[A-Za-z0-9]+(?:[-_][A-Za-z0-9]+)*") _FALLBACK_STOPWORDS = { "a", "an", @@ -259,44 +233,63 @@ } -def _fallback_token_variants(token: str) -> set[str]: - token = token.lower() - variants = {token} - - generic_tokens = globals().get("_GENERIC_SCHEMA_INTENT_TOKENS", set()) - if token in generic_tokens: - return variants - - if len(token) > 4 and token.endswith("ies"): - variants.add(token[:-3] + "y") - elif len(token) > 4 and token.endswith("es"): - if token.endswith(("ches", "shes", "sses", "uses", "xes", "zes")): - variants.add(token[:-2]) - else: - variants.add(token[:-1]) - elif len(token) > 3 and token.endswith("s"): - variants.add(token[:-1]) - if len(token) > 4 and token.endswith("ed"): - stem = token[:-2] - variants.add(stem) - if len(token) > 5 and token[-3] in {"d", "s", "t", "v", "z"}: - variants.add(token[:-1]) - if len(stem) > 2 and stem[-1] == stem[-2]: - variants.add(stem[:-1]) - if len(token) > 5 and token.endswith("ing"): - stem = token[:-3] - variants.add(stem) - variants.add(stem + "e") - if len(stem) > 2 and stem[-1] == stem[-2]: - variants.add(stem[:-1]) - - return {variant for variant in variants if variant} - - -def _expand_fallback_token_variants(tokens: set[str]) -> set[str]: - expanded = set() - for token in tokens: - expanded.update(_fallback_token_variants(token)) +def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: + aliases = { + "bu": {"business", "unit"}, + "cust": {"customer"}, + "customers": {"customer"}, + "critical": {"priority", "severity"}, + "boards": {"board"}, + "high": {"priority", "severity"}, + "highest": {"top"}, + "logs": {"log", "record"}, + "log": {"record"}, + "lows": {"low"}, + "lowest": {"bottom"}, + "models": {"model"}, + "inv": {"invoice"}, + "invoices": {"invoice"}, + "ord": {"order"}, + "orders": {"order"}, + "qty": {"quantity"}, + "num": {"number"}, + "no": {"number"}, + "numbers": {"number"}, + "prod": {"product"}, + "products": {"product"}, + "priorities": {"priority", "severity"}, + "priority": {"severity"}, + "recent": {"latest"}, + "records": {"record"}, + "rep": {"representative", "salesperson"}, + "repairs": {"repair"}, + "salesperson": {"sales", "person"}, + "severity": {"priority"}, + "supplier": {"vendor"}, + "suppliers": {"supplier", "vendor"}, + "tech": {"technician"}, + "technician": {"tech"}, + "vendor": {"supplier"}, + "vendors": {"supplier", "vendor"}, + "failed": {"failure"}, + "failures": {"failure"}, + "defects": {"defect"}, + "types": {"type"}, + "units": {"unit", "serial"}, + "urgency": {"priority", "severity"}, + "locations": {"location"}, + "materials": {"material"}, + "missing": {"blank", "empty", "null"}, + } + expanded = set(tokens) + for token in list(tokens): + expanded.update(aliases.get(token, set())) + if {"business", "unit"}.issubset(expanded): + expanded.add("bu") + if "customer" in expanded and "number" in expanded: + expanded.update({"cust", "id", "no"}) + if "order" in expanded and "number" in expanded: + expanded.update({"ord", "id", "no"}) return expanded @@ -714,10 +707,10 @@ def _extract_semantic_context_payload(context: str) -> dict[str, Any]: def _extract_semantic_tokens_by_column( context: str, -) -> tuple[set[str], dict[str, set[str]], dict[str, list[str]]]: +) -> tuple[set[str], dict[str, set[str]]]: payload = _extract_semantic_context_payload(context) if not payload: - return set(), {}, {} + return set(), {} table_tokens = _semantic_tokens_from_value( payload.get("semantic_context_not_sql_identifiers") @@ -725,7 +718,6 @@ def _extract_semantic_tokens_by_column( table_tokens.update(_semantic_tokens_from_value(payload.get("object_type"))) column_tokens: dict[str, set[str]] = {} - column_sample_values: dict[str, list[str]] = {} for column in payload.get("columns", []) or []: if not isinstance(column, dict): continue @@ -737,42 +729,8 @@ def _extract_semantic_tokens_by_column( ) if tokens: column_tokens[column_name] = tokens - sample_values = _extract_column_sample_values(column) - if sample_values: - column_sample_values[column_name] = sample_values - - return table_tokens, column_tokens, column_sample_values - -def _extract_column_sample_values(column: dict[str, Any]) -> list[str]: - values: list[str] = [] - - def add(value: Any) -> None: - if value is None: - return - if isinstance(value, (list, tuple, set)): - for item in value: - add(item) - return - if isinstance(value, dict): - for item in value.values(): - add(item) - return - text = str(value).strip() - if text and text.lower() not in {item.lower() for item in values}: - values.append(text) - - for key in ( - "sample_values", - "sample_value", - "samples", - "values", - "example_values", - "examples", - "distinct_values", - ): - add(column.get(key)) - return values + return table_tokens, column_tokens def _extract_schema_details( @@ -789,7 +747,7 @@ def _extract_schema_details( continue relation_name = _unquote_identifier(relation_match.group("name")) - table_semantic_tokens, column_semantic_tokens, column_sample_values = ( + table_semantic_tokens, column_semantic_tokens = ( _extract_semantic_tokens_by_column(context) ) column_block_match = re.search( @@ -836,7 +794,6 @@ def _extract_schema_details( "name": name, "data_type": data_type, "semantic_tokens": column_semantic_tokens.get(name, set()), - "sample_values": column_sample_values.get(name, []), "_table_semantic_tokens": table_semantic_tokens, } ) @@ -1113,68 +1070,6 @@ def _sql_mentions_identifier(sql: str, identifier: str) -> bool: ) -def _sql_mentions_literal_value(sql: str, values: list[str]) -> bool: - lowered_sql = sql.lower() - for value in values: - cleaned = _clean_filter_value(value) - if cleaned and _quote_literal(cleaned.lower()) in lowered_sql: - return True - return False - - -def _extract_sql_string_literals(sql: str) -> list[str]: - literals = [] - for match in _SINGLE_QUOTED_LITERAL.finditer(sql): - literal = match.group(0)[1:-1].replace("''", "'") - if literal: - literals.append(literal) - return literals - - -def _extract_column_filter_literals(sql: str, column_name: str) -> list[str]: - stripped = _strip_string_literals(sql) - quoted_column = re.escape(_quote_identifier(column_name)) - bare_column = re.escape(column_name) - column_pattern = rf"(?:{quoted_column}|(? str | None: - referenced_relations = { - relation - for relation in grounding["relation_references"] - if relation not in grounding["cte_names"] - } - for relation in referenced_relations: - for column in schema_details.get(relation, []): - sample_values = column.get("sample_values") or [] - if not sample_values: - continue - literals = _extract_column_filter_literals(sql, column["name"]) - if not literals: - continue - sample_tokens = _sample_value_tokens(column) - sample_lowers = {str(value).lower() for value in sample_values} - for literal in literals: - literal_tokens = _fallback_tokens(literal) - if literal.lower() in sample_lowers or literal_tokens & sample_tokens: - continue - return ( - "Schema grounding failed. The generated SQL filters column " - f"{column['name']} with a literal value not found in that " - "column's verified sample values." - ) - return None - - def _validate_unqualified_columns_for_single_relation( sql: str, schema_index: dict[str, set[str] | None], @@ -1321,119 +1216,6 @@ def validate_sql_against_contexts( return None -def _extract_sql_from_value(value: Any) -> str | None: - if value is None: - return None - - if isinstance(value, str): - text = value.strip() - if not text: - return None - - try: - parsed = orjson.loads(text) - except orjson.JSONDecodeError: - return text if _SQL_START.search(text) else None - - return _extract_sql_from_value(parsed) - - if isinstance(value, dict): - for key in ("sql", "query", "code"): - extracted = _extract_sql_from_value(value.get(key)) - if extracted: - return extracted - - extracted = _extract_sql_from_value(value.get("arguments")) - if extracted: - return extracted - - return None - - if isinstance(value, list): - for item in value: - extracted = _extract_sql_from_value(item) - if extracted: - return extracted - - return None - - -def _extract_generation_sql(generation_result: str | None) -> str | None: - if not generation_result: - return None - - extracted = _extract_sql_from_value(generation_result) - if extracted: - return extracted - - text = generation_result.strip() - return text if _SQL_START.search(text) else None - - -def _sql_has_aggregate_function(sql: str) -> bool: - return bool(re.search(r"(?is)\b(?:AVG|COUNT|MAX|MIN|SUM)\s*\(", sql)) - - -def _group_by_source_columns(sql: str) -> set[str]: - group_by_clause = _extract_clause( - sql, - "GROUP BY", - ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), - ) - if not group_by_clause: - return set() - - columns: set[str] = set() - for expression in _split_sql_tokens(group_by_clause): - for identifier in _iter_unqualified_identifier_candidates(expression): - upper_identifier = identifier.upper() - if ( - upper_identifier in _SQL_RESERVED_WORDS - or upper_identifier in _SQL_FUNCTION_WORDS - or upper_identifier in _SQL_TYPE_WORDS - or upper_identifier in _DATE_PART_WORDS - ): - continue - columns.add(identifier) - return columns - - -def _query_allows_grouped_aggregate(query: str, query_tokens: set[str]) -> bool: - return bool( - _has_grouping_intent(query, query_tokens) - or _has_count_intent(query_tokens) - or _has_sum_intent(query_tokens) - or _is_rate_metric_intent(query_tokens) - or _is_distribution_metric_intent(query_tokens) - ) - - -def _missing_value_target_tokens(query: str, query_tokens: set[str]) -> set[str]: - match = re.search( - r"(?is)\b(?:blank|empty|missing|null)\s+(?P[A-Za-z0-9_ /-]+)", - query, - ) - if match: - tokens = _fallback_tokens(match.group("value")) - else: - tokens = set(query_tokens) - return tokens - _GENERIC_SCHEMA_INTENT_TOKENS - _NULL_CHECK_TOKENS - - -def _sql_null_checked_columns(sql: str, columns: list[dict[str, str]]) -> set[str]: - checked_columns: set[str] = set() - stripped = _strip_string_literals(sql) - for column in columns: - quoted_column = re.escape(_quote_identifier(column["name"])) - bare_column = re.escape(column["name"]) - column_pattern = rf"(?:{quoted_column}|(? 1: - return ( - "Schema grounding failed. The question asks for one grouping " - "dimension, but the generated SQL groups by multiple source " - "columns." - ) + if not schema_tokens: + return None - if _has_missing_value_intent(raw_query_tokens): - target_tokens = _missing_value_target_tokens(query, raw_query_tokens) - for relation in referenced_relations: - columns = schema_details.get(relation) or [] - if not columns or not target_tokens: - continue - checked_columns = _sql_null_checked_columns(sql, columns) - if not checked_columns: - continue - scored_columns = [ - (_column_score_for_tokens(column, target_tokens), column["name"]) - for column in columns - ] - best_score = max((score for score, _ in scored_columns), default=0) - checked_best_score = max( - ( - score - for score, name in scored_columns - if name in checked_columns - ), - default=0, - ) - if best_score > checked_best_score: + missing_concepts = [ + label for label, concept_tokens in concepts if not schema_tokens & concept_tokens + ] + if not missing_concepts: + if _is_failure_count_intent(raw_query_tokens, query_tokens): + if not re.search(r"(?is)\bCOUNT\s*\(", sql): return ( - "Schema grounding failed. The question asks for missing " - "values on a specific schema concept, but the generated SQL " - "checks a weaker matching column for null or blank values." + "Schema grounding failed. The question asks for a count of " + "failure records, but the generated SQL does not compute a " + "COUNT aggregate. Use a verified failure-record column/table " + "and group by the requested dimension, or return no SQL if " + "the active project does not contain it." ) + for relation in referenced_relations: + for column in schema_details.get(relation, []): + if _is_rate_like_column(column) and _sql_mentions_identifier( + sql, column["name"] + ): + return ( + "Schema grounding failed. The question asks for a " + "count of failure records, but the generated SQL uses " + f"rate-like column {column['name']}. Use COUNT over a " + "verified failure occurrence field instead, or return " + "no SQL if the active project does not contain one." + ) + return None - return _validate_literal_values_against_samples(sql, schema_details, grounding) + return ( + "Schema grounding failed. The generated SQL uses verified identifiers, " + "but the selected table or view does not contain verified fields for the " + f"requested business concept(s): {', '.join(missing_concepts)}. Use only " + "schema objects whose table or column names explicitly support those " + "concepts, or return no SQL if the active project does not contain them." + ) def unsupported_schema_message( @@ -1579,39 +1295,23 @@ def unsupported_schema_message( if not query: return None query_tokens = _fallback_tokens(query) + concepts = _requested_business_concepts(query_tokens) + if not concepts: + return None schema_details = _extract_schema_details(contexts) if not schema_details: return None - required_tokens = _schema_derived_query_tokens(query_tokens, schema_details) - unsupported_tokens = _unsupported_query_tokens( - query_tokens, - schema_details, - query=query, - ) - if unsupported_tokens: - return ( - "No retrieved table or view in the active project contains verified " - "schema metadata for all requested non-operational term(s): " - f"{', '.join(sorted(unsupported_tokens))}. Select a project with " - "matching fields, add schema descriptions/sample values, or ask a " - "question supported by the selected project's schema." - ) - - table_tokens = _schema_tokens_by_table(schema_details) - if required_tokens and any( - required_tokens <= tokens for tokens in table_tokens.values() + if any( + _table_covers_requested_concepts(table_name, columns, query_tokens) + for table_name, columns in schema_details.items() ): return None - - if not required_tokens and not unsupported_tokens: - return None - detail_tokens = sorted(required_tokens or unsupported_tokens) + concept_labels = ", ".join(label for label, _ in concepts) return ( "No retrieved table or view in the active project contains verified " - "schema metadata for all requested non-operational term(s): " - f"{', '.join(detail_tokens)}. Select a project with matching fields, " - "add schema descriptions/sample values, or ask a question supported by " - "the selected project's schema." + "fields for all requested business concept(s): " + f"{concept_labels}. Select a project with those fields or ask a question " + "supported by the selected project's schema." ) @@ -1661,7 +1361,7 @@ def _fallback_tokens(value: Any) -> set[str]: for token in _FALLBACK_TOKEN.findall(text.lower()) if token not in _FALLBACK_STOPWORDS } - return _expand_fallback_token_variants(tokens) + return _expand_fallback_token_aliases(tokens) def _column_business_tokens(column: dict[str, Any]) -> set[str]: @@ -1680,12 +1380,35 @@ def _table_business_tokens( return tokens -def _data_type_base(data_type: str) -> str: - return data_type.upper().split("(", 1)[0].strip() +def _expanded_fallback_query_tokens(query: str) -> set[str]: + tokens = _fallback_tokens(query) + if tokens & {"revenue", "sale", "sales", "trend", "trends"}: + tokens.update({"amount", "date", "intake", "revenue", "sales", "value"}) + if tokens & {"order", "orders"}: + tokens.update({"amount", "customer", "date", "ord", "order", "value"}) + if tokens & {"invoice", "invoices"}: + tokens.update({"amount", "currency", "date", "invoice", "supplier"}) + if tokens & {"batch", "batches"}: + tokens.update({"batch", "board", "defect", "inspection", "rate", "supplier"}) + if tokens & {"repair", "repairs"}: + tokens.update({"date", "failure", "log", "priority", "progress", "repair", "status"}) + if tokens & {"failure", "failures", "defect", "defects"}: + tokens.update({"code", "defect", "failure", "severity", "status", "type"}) + if tokens & {"material", "materials"}: + tokens.update({"item", "material", "part"}) + if tokens & {"location", "locations"}: + tokens.update({"area", "location", "site"}) + if tokens & {"month", "monthly", "july"}: + tokens.update({"date", "day", "month", "time", "year"}) + elif tokens & {"latest", "trend", "trends", "year"}: + tokens.update({"date", "day", "time", "year"}) + if "business" in tokens and "unit" in tokens: + tokens.update({"account", "bu", "business", "company", "division", "unit"}) + return tokens def _is_numeric_type(data_type: str) -> bool: - return _data_type_base(data_type) in { + return data_type.upper() in { "BIGINT", "DECIMAL", "DOUBLE", @@ -1704,7 +1427,7 @@ def _is_numeric_type(data_type: str) -> bool: def _is_date_type(data_type: str) -> bool: - return _data_type_base(data_type) in { + return data_type.upper() in { "DATE", "DATETIME", "DATETIME2", @@ -1720,231 +1443,108 @@ def _is_date_type(data_type: str) -> bool: _RATE_METRIC_TOKENS = {"rate", "ratio", "percent", "percentage"} _COUNT_METRIC_TOKENS = {"count", "many", "most", "number", "total"} -_AVERAGE_METRIC_TOKENS = {"average", "avg", "mean"} -_DISTRIBUTION_METRIC_TOKENS = {"distribution", "breakdown"} -_SUM_METRIC_TOKENS = {"sum", "total"} -_MIN_METRIC_TOKENS = {"bottom", "least", "lowest", "min", "minimum", "smallest"} -_MAX_METRIC_TOKENS = {"greatest", "highest", "largest", "max", "maximum", "most", "top"} -_LATEST_METRIC_TOKENS = {"latest", "newest", "recent"} -_NULL_CHECK_TOKENS = {"blank", "empty", "missing", "null"} -_GENERIC_SCHEMA_INTENT_TOKENS = { - "a", - "across", - "all", - "an", - "and", - "as", - "ascending", - "associated", - "association", - "average", - "avg", - "between", - "bottom", - "breakdown", - "bucket", - "buckets", - "by", - "compare", - "count", - "date", - "day", - "descending", - "distribution", - "each", - "for", - "from", - "group", - "grouped", - "has", - "have", - "highest", - "in", - "latest", - "least", - "list", - "lowest", - "many", - "max", - "maximum", - "me", - "mean", - "min", - "minimum", - "month", - "monthly", - "most", - "newest", - "number", - "of", - "ordered", - "per", - "please", - "quarter", - "recent", - "record", - "records", - "row", - "rows", - "show", - "smallest", - "sort", - "sorted", - "sum", - "the", - "to", - "top", - "total", - "week", - "which", - "with", - "without", - "year", +_PRIORITY_VALUE_ALIASES = { + "urgent": "urgent", + "critical": "critical", + "high": "high", + "medium": "medium", + "normal": "normal", + "low": "low", } -_GENERIC_SCHEMA_INTENT_TOKENS.update(_NULL_CHECK_TOKENS) -_GENERIC_SCHEMA_INTENT_TOKENS.update(_MONTH_NAME_TO_NUMBER.keys()) +_PRIORITY_ORDER = [ + ("critical", 6), + ("urgent", 6), + ("blocker", 6), + ("high", 5), + ("major", 5), + ("medium", 4), + ("normal", 4), + ("minor", 3), + ("low", 2), +] def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: return bool(raw_query_tokens & _RATE_METRIC_TOKENS) -def _is_average_metric_intent(raw_query_tokens: set[str]) -> bool: - return bool(raw_query_tokens & _AVERAGE_METRIC_TOKENS) +def _is_failure_count_intent( + raw_query_tokens: set[str], + query_tokens: set[str], +) -> bool: + return ( + bool(raw_query_tokens & {"failure", "failed", "defect"}) + and "failure" in query_tokens + and not _is_rate_metric_intent(raw_query_tokens) + and ( + bool(raw_query_tokens & _COUNT_METRIC_TOKENS) + or bool(raw_query_tokens & {"top", "highest", "lowest", "bottom"}) + ) + ) -def _is_distribution_metric_intent(raw_query_tokens: set[str]) -> bool: - return bool(raw_query_tokens & _DISTRIBUTION_METRIC_TOKENS) +def _has_board_model_intent(query_tokens: set[str]) -> bool: + return {"board", "model"}.issubset(query_tokens) def _is_rate_like_column(column: dict[str, str]) -> bool: return bool(_fallback_tokens(column["name"]) & (_RATE_METRIC_TOKENS | {"score"})) -def _is_identifier_like_column(column: dict[str, str]) -> bool: - tokens = _fallback_tokens(column["name"]) - return bool(tokens & {"id", "identifier", "key", "uuid"}) - - def _quote_joined(identifiers: list[str]) -> str: return ", ".join(_quote_identifier(identifier) for identifier in identifiers) +def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[str]]]: + concepts: list[tuple[str, set[str]]] = [] + specs = [ + ( + "failure/defect", + {"failure", "failed", "defect"}, + {"failure", "failed", "defect"}, + ), + ("repair", {"repair"}, {"repair"}), + ("material", {"material"}, {"material", "part"}), + ("location", {"location"}, {"location", "site", "area"}), + ("customer", {"customer"}, {"customer", "cust"}), + ("supplier/vendor", {"supplier", "vendor"}, {"supplier", "vendor"}), + ("technician", {"technician", "tech"}, {"technician", "tech"}), + ("product", {"product"}, {"product", "prod", "item", "material"}), + ( + "priority/severity", + {"critical", "priority", "severity"}, + {"priority", "severity"}, + ), + ("status", {"status"}, {"status"}), + ("order", {"order"}, {"order", "ord"}), + ] + for label, triggers, schema_tokens in specs: + if query_tokens & triggers: + concepts.append((label, schema_tokens)) + if {"board", "model"}.issubset(query_tokens): + concepts.append(("board model", {"board", "model"})) + if {"business", "unit"}.issubset(query_tokens): + concepts.append(("business unit", {"business", "unit", "bu", "division"})) + return concepts + + def _schema_tokens_for_table(table_name: str, columns: list[dict[str, str]]) -> set[str]: tokens = _table_business_tokens(table_name, columns) for column in columns: tokens.update(_column_business_tokens(column)) - tokens.update(_sample_value_tokens(column)) return tokens -def _schema_tokens_by_table( - schema_details: dict[str, list[dict[str, str]]], -) -> dict[str, set[str]]: - return { - table_name: _schema_tokens_for_table(table_name, columns) - for table_name, columns in schema_details.items() - } - - -def _schema_derived_query_tokens( - query_tokens: set[str], - schema_details: dict[str, list[dict[str, str]]], -) -> set[str]: - schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) - supported_tokens: set[str] = set() - for token in query_tokens: - if token in _GENERIC_SCHEMA_INTENT_TOKENS or token.isdigit(): - continue - supported_tokens.update(_fallback_token_variants(token) & schema_tokens) - return supported_tokens - - -def _schema_adjacent_dimension_descriptor_tokens( - query: str | None, - query_tokens: set[str], - schema_tokens: set[str], -) -> set[str]: - if not query: - return set() - dimension_intent_tokens = _DISTRIBUTION_METRIC_TOKENS | { - "across", - "by", - "each", - "group", - "grouped", - "per", - } - if not ( - query_tokens - & dimension_intent_tokens - ): - return set() - - descriptor_tokens: set[str] = set() - matches = list(_QUERY_VALUE_TOKEN.finditer(query)) - for index, match in enumerate(matches): - if index == 0: - continue - previous_tokens = _fallback_tokens( - matches[index - 1].group(0).replace("_", " ") - ) - value_tokens = _fallback_tokens(match.group(0).replace("_", " ")) - explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { - "'", - '"', - } - if ( - previous_tokens & schema_tokens - and value_tokens - and not (value_tokens & schema_tokens) - and not explicitly_quoted - ): - descriptor_tokens.update(value_tokens) - return descriptor_tokens - - -def _unsupported_query_tokens( - query_tokens: set[str], - schema_details: dict[str, list[dict[str, str]]], - query: str | None = None, -) -> set[str]: - schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) - descriptor_tokens = _schema_adjacent_dimension_descriptor_tokens( - query, - query_tokens, - schema_tokens, - ) - user_value_tokens = _schema_driven_user_value_tokens( - query, - query_tokens, - schema_details, - ) - return { - token - for token in query_tokens - if token not in _GENERIC_SCHEMA_INTENT_TOKENS - and not token.isdigit() - and not (_fallback_token_variants(token) & schema_tokens) - and token not in descriptor_tokens - and token not in user_value_tokens - } - - def _table_covers_requested_concepts( table_name: str, columns: list[dict[str, str]], concept_tokens: set[str], ) -> bool: - required_tokens = { - token - for token in concept_tokens - if token not in _GENERIC_SCHEMA_INTENT_TOKENS - } - if not required_tokens: + concepts = _requested_business_concepts(concept_tokens) + if not concepts: return True schema_tokens = _schema_tokens_for_table(table_name, columns) - return all(_fallback_token_variants(token) & schema_tokens for token in required_tokens) + return all(schema_tokens & concept_tokens for _, concept_tokens in concepts) def _choose_fallback_table( @@ -1953,44 +1553,221 @@ def _choose_fallback_table( concept_tokens: set[str] | None = None, ) -> tuple[str, list[dict[str, str]]] | None: concept_tokens = concept_tokens or query_tokens - required_tokens = _schema_derived_query_tokens(concept_tokens, schema_details) + rate_metric_intent = _is_rate_metric_intent(concept_tokens) + failure_count_intent = _is_failure_count_intent(concept_tokens, query_tokens) + board_model_intent = _has_board_model_intent(query_tokens) or _has_board_model_intent( + concept_tokens + ) scored_tables = [] for table_name, columns in schema_details.items(): table_tokens = _table_business_tokens(table_name, columns) column_token_union = set() + has_numeric_sales_measure = False + has_date_capable_column = False score = len(query_tokens & table_tokens) * 8 for column in columns: column_tokens = _column_business_tokens(column) - sample_tokens = _sample_value_tokens(column) column_token_union.update(column_tokens) - column_token_union.update(sample_tokens) + if _is_numeric_type(column["data_type"]) and column_tokens & { + "amount", + "intake", + "revenue", + "sales", + "value", + }: + has_numeric_sales_measure = True + if _is_date_type(column["data_type"]) or column_tokens & { + "date", + "day", + "month", + "time", + "year", + }: + has_date_capable_column = True score += len(query_tokens & column_tokens) * 10 - score += len(query_tokens & sample_tokens) * 6 if _is_numeric_type(column["data_type"]): - score += len(query_tokens & column_tokens) * 2 + score += len( + query_tokens + & column_tokens + & { + "amount", + "cost", + "count", + "margin", + "quantity", + "rate", + "score", + "value", + } + ) * 4 if _is_date_type(column["data_type"]): - score += 2 + score += ( + len(query_tokens & {"date", "month", "year", "july", "trend", "trends"}) + * 4 + ) - table_schema_tokens = _schema_tokens_for_table(table_name, columns) - if required_tokens and not required_tokens <= table_schema_tokens: + if not _table_covers_requested_concepts(table_name, columns, concept_tokens): continue - score += len(required_tokens & table_schema_tokens) * 20 - if score > 0: - scored_tables.append((score, table_name, columns)) + if board_model_intent and rate_metric_intent and query_tokens & { + "defect", + "failure", + }: + if not {"board", "model"}.issubset(column_token_union): + continue + rate_column = _choose_column_by_tokens( + columns, + {"defect", "rate"}, + numeric=True, + ) + if not rate_column: + continue + score += 130 - if not scored_tables: - return None + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + score += ( + len(column_token_union & {"amount", "intake", "revenue", "sales", "value"}) + * 10 + ) + score += len(column_token_union & {"date", "month", "time", "year"}) * 5 + if not has_numeric_sales_measure: + continue + score += 50 + if query_tokens & {"year", "month", "monthly", "trend", "trends"}: + if not has_date_capable_column: + continue + score += 30 + if not query_tokens & { + "claim", + "claims", + "customs", + "duty", + "import", + "refund", + "tariff", + }: + table_and_columns = table_tokens | column_token_union + customs_matches = table_and_columns & { + "claim", + "claims", + "customs", + "duty", + "import", + "refund", + "tariff", + "tariffs", + } + if customs_matches and not table_and_columns & {"revenue", "sale", "sales"}: + continue + score -= len(customs_matches) * 40 - scored_tables.sort(key=lambda item: (-item[0], item[1])) - return scored_tables[0][1], scored_tables[0][2] + if "customer" in query_tokens: + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"customer", "cust"}: + if "missing" in query_tokens: + continue + score -= 80 + else: + score += 45 + if {"business", "unit"}.issubset(query_tokens): + if "bu" in column_token_union: + score += 60 + elif {"business", "unit"} <= column_token_union: + score += 45 -def _choose_column_by_tokens( - columns: list[dict[str, str]], - required_tokens: set[str], - numeric: bool | None = None, - date: bool | None = None, + if "failure" in query_tokens and ( + query_tokens & {"location", "material", "technician", "tech"} + or board_model_intent + ): + table_and_columns = table_tokens | column_token_union + if not table_and_columns & {"failure", "failed", "defect"}: + continue + if board_model_intent: + if not {"board", "model"}.issubset(column_token_union): + continue + if failure_count_intent and not _choose_count_subject_column( + {"failure"}, + columns, + ): + continue + score += 100 + if "location" in query_tokens: + if "location" not in column_token_union: + continue + score += 90 + if "material" in query_tokens: + if not column_token_union & {"material", "part"}: + continue + score += 90 + if query_tokens & {"technician", "tech"}: + if not column_token_union & {"technician", "tech"}: + continue + score += 90 + + if query_tokens & {"order", "orders"}: + table_and_columns = table_tokens | column_token_union + explicit_order_support = table_and_columns & {"ord", "order", "orders"} + order_support = table_and_columns & { + "amount", + "customer", + "intake", + "item", + "ord", + "order", + "orders", + "sales", + "value", + } + if not order_support: + continue + if explicit_order_support: + score += 80 + else: + score -= 60 + + if query_tokens & {"batch", "batches"} and {"defect", "rate"}.issubset( + column_token_union + ): + score += 40 + if {"material", "location"}.issubset(query_tokens) and { + "material", + "location", + }.issubset(column_token_union): + score += 40 + if ( + query_tokens & {"repair", "repairs"} + and (table_tokens | column_token_union) & {"repair", "fix"} + ): + score += 55 + if concept_tokens & {"critical", "priority", "severity"}: + if not _choose_priority_column(columns): + continue + score += 55 + if concept_tokens & {"status"}: + if not column_token_union & {"status", "state", "progress"}: + continue + score += 45 + if concept_tokens & {"latest", "recent"}: + if not has_date_capable_column: + continue + score += 35 + + if score > 0: + scored_tables.append((score, table_name, columns)) + + if not scored_tables: + return None + + scored_tables.sort(key=lambda item: (-item[0], item[1])) + return scored_tables[0][1], scored_tables[0][2] + + +def _choose_column_by_tokens( + columns: list[dict[str, str]], + required_tokens: set[str], + numeric: bool | None = None, + date: bool | None = None, ) -> str | None: candidates = [] for column in columns: @@ -2027,6 +1804,23 @@ def _column_score_for_tokens( score = len(required_tokens & column_tokens) * 10 if required_tokens and required_tokens.issubset(column_tokens): score += 30 + if column["name"].lower() == "bu" and {"business", "unit"} & required_tokens: + score += 60 + if column["name"].lower() in {"custno", "customer_id", "customerid"} and { + "customer", + "number", + } & required_tokens: + score += 30 + if column["name"].lower() in {"custname", "customer_name", "customer"} and { + "customer", + "name", + } & required_tokens: + score += 35 + if column["name"].lower() in {"ordno", "order_no", "order_number", "sales_order_number"} and { + "order", + "number", + } & required_tokens: + score += 35 return score @@ -2046,235 +1840,84 @@ def _choose_ranked_column_by_tokens( return candidates[0][1] -def _sample_value_tokens(column: dict[str, Any]) -> set[str]: - tokens: set[str] = set() - for value in column.get("sample_values") or []: - tokens.update(_fallback_tokens(value)) - return tokens - - -def _column_supports_filter_values( - column: dict[str, Any], - values: list[str], -) -> bool: - cleaned_values = [value for value in (_clean_filter_value(value) for value in values) if value] - if not cleaned_values: - return True - - column_tokens = _column_business_tokens(column) - sample_tokens = _sample_value_tokens(column) - value_tokens: set[str] = set() - for value in cleaned_values: - value_tokens.update(_fallback_tokens(value)) - - if value_tokens & sample_tokens: - return True - if value_tokens & column_tokens: - return True - return False - - -def _filter_predicate_for_values( - column: dict[str, str], - values: list[str], -) -> str: - cleaned_values = [ - value for value in (_clean_filter_value(value) for value in values) if value - ] - if not cleaned_values: - return _non_missing_value_predicate(column) - - column_tokens = _column_business_tokens(column) - value_tokens: set[str] = set() - for value in cleaned_values: - value_tokens.update(_fallback_tokens(value)) - - return _value_match_predicate(column, cleaned_values[0], cleaned_values[1:]) - - -def _choose_filter_column_for_values( - columns: list[dict[str, str]], - values: list[str], -) -> dict[str, str] | None: - candidates: list[tuple[int, dict[str, str]]] = [] - - for column in columns: - column_tokens = _column_business_tokens(column) - if not _column_supports_filter_values(column, values): - continue - value_tokens = { - token - for value in values - for token in _fallback_tokens(value) - } - score = len(column_tokens & value_tokens) * 10 - score += len(_sample_value_tokens(column) & value_tokens) * 20 - if score > 0: - candidates.append((score, column)) - - if not candidates: - return None - candidates.sort(key=lambda item: (-item[0], item[1]["name"])) - return candidates[0][1] - - -def _choose_temporal_column( - query_tokens: set[str], - columns: list[dict[str, str]], -) -> str | None: - candidates = [] - for column in columns: - column_tokens = _column_business_tokens(column) - if not _is_date_type(column["data_type"]): - continue - score = len(query_tokens & column_tokens) * 12 - score += 20 - if score > 0: - candidates.append((score, column["name"])) - - if not candidates: - return None - candidates.sort(key=lambda item: (-item[0], item[1])) - return candidates[0][1] - - -def _order_by_phrase_tokens(query: str) -> set[str]: - match = re.search( - r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+(?P[A-Za-z0-9_ /-]+)", - query, - ) - if not match: - return set() - return _fallback_tokens(match.group("value")) - - -def _choose_order_by_column( - query: str, - query_tokens: set[str], - columns: list[dict[str, str]], -) -> str | None: - order_tokens = _order_by_phrase_tokens(query) - if not order_tokens: - return None - - column = _choose_ranked_column_by_tokens(columns, order_tokens) - return column["name"] if column else None - - def _choose_dimension_column( query_tokens: set[str], columns: list[dict[str, str]], ) -> str | None: - columns = _choose_dimension_columns(query_tokens, columns, max_columns=1) - return columns[0] if columns else None - - -def _choose_dimension_columns( - query_tokens: set[str], - columns: list[dict[str, str]], - max_columns: int = 3, -) -> list[str]: - name_candidates = [] - semantic_candidates = [] - filtered_query_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS - identifier_requested = bool(filtered_query_tokens & {"id", "identifier", "key", "uuid"}) - for index, column in enumerate(columns): - if _is_numeric_type(column["data_type"]): - continue - if _is_identifier_like_column(column) and not identifier_requested: - continue - name_tokens = _fallback_tokens(column["name"]) - semantic_tokens = set(column.get("semantic_tokens") or set()) - name_score = len(filtered_query_tokens & name_tokens) * 10 - semantic_score = len(filtered_query_tokens & semantic_tokens) * 3 - if name_score > 0: - name_candidates.append((name_score + semantic_score, index, column["name"])) - elif semantic_score > 0: - semantic_candidates.append((semantic_score, index, column["name"])) - candidates = name_candidates or semantic_candidates - candidates.sort(key=lambda item: (-item[0], item[1])) - return [name for _, _, name in candidates[:max_columns]] + dimension_specs = [ + ({"board", "model"}, {"board", "model"}), + ({"business", "unit"}, {"business", "unit", "bu", "division", "company"}), + ({"customer"}, {"customer", "cust", "name"}), + ({"supplier"}, {"supplier", "vendor", "name"}), + ({"product"}, {"product", "prod", "item", "material", "name"}), + ({"salesperson", "representative"}, {"salesperson", "sales", "person", "rep"}), + ({"technician", "tech"}, {"technician", "tech"}), + ({"location"}, {"location", "site", "area"}), + ({"material"}, {"material", "part", "item"}), + ({"priority", "severity"}, {"priority", "severity", "urgency", "rank"}), + ({"status"}, {"status"}), + ({"currency"}, {"currency", "curr"}), + ({"country"}, {"country"}), + ({"order"}, {"order", "ord", "number"}), + ({"batch"}, {"batch", "id"}), + ] + for trigger_tokens, column_tokens in dimension_specs: + if query_tokens & trigger_tokens: + column = _choose_ranked_column_by_tokens(columns, column_tokens) + if column: + return column["name"] + return None def _choose_missing_value_column( - query: str, query_tokens: set[str], columns: list[dict[str, str]], ) -> dict[str, str] | None: - target_tokens = _missing_value_target_tokens(query, query_tokens) - column = _choose_ranked_column_by_tokens(columns, target_tokens) - if column: - return column - return _choose_ranked_column_by_tokens( - columns, - query_tokens - (_GENERIC_SCHEMA_INTENT_TOKENS | _NULL_CHECK_TOKENS), - ) + missing_specs = [ + ({"customer"}, {"customer", "cust", "number", "id", "no"}), + ({"order"}, {"order", "ord", "number", "id", "no"}), + ({"supplier"}, {"supplier", "vendor", "number", "id", "no"}), + ({"product", "material"}, {"product", "prod", "material", "item", "number", "id"}), + ({"location"}, {"location", "site", "area"}), + ({"status"}, {"status"}), + ] + for trigger_tokens, column_tokens in missing_specs: + if query_tokens & trigger_tokens: + column = _choose_ranked_column_by_tokens(columns, column_tokens) + if column: + return column + return None def _choose_count_subject_column( query_tokens: set[str], columns: list[dict[str, str]], ) -> dict[str, str] | None: - column = _choose_ranked_column_by_tokens( - columns, - query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS, - ) - if column: - return column - for column in columns: - if not _is_numeric_type(column["data_type"]) and not _is_rate_like_column(column): - return column - return columns[0] if columns else None - - -def _choose_average_measure_column( - query_tokens: set[str], - columns: list[dict[str, str]], -) -> dict[str, str] | None: - filtered_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS - measure_candidates = [ - column for column in columns if not _is_identifier_like_column(column) - ] - column = _choose_ranked_column_by_tokens( - measure_candidates, - filtered_tokens, - numeric=True, - ) - if column: - return column - numeric_columns = [ - column - for column in columns - if _is_numeric_type(column["data_type"]) - and not _is_identifier_like_column(column) + subject_specs = [ + ({"order"}, {"order", "ord", "number", "id", "no"}), + ({"customer"}, {"customer", "cust", "number", "id", "no"}), + ( + {"failure"}, + {"failure", "failed", "defect", "code", "line", "status", "sys", "type"}, + ), + ({"repair"}, {"repair", "id", "status"}), + ({"batch"}, {"batch", "id"}), ] - return numeric_columns[0] if len(numeric_columns) == 1 else None + for trigger_tokens, column_tokens in subject_specs: + if query_tokens & trigger_tokens: + candidate_columns = columns + if trigger_tokens & {"failure"}: + candidate_columns = [ + column for column in columns if not _is_rate_like_column(column) + ] + column = _choose_ranked_column_by_tokens(candidate_columns, column_tokens) + if column: + return column + return None def _is_text_type(data_type: str) -> bool: - return _data_type_base(data_type) in { - "CHAR", - "CHARACTER", - "NCHAR", - "NTEXT", - "NVARCHAR", - "STRING", - "TEXT", - "VARCHAR", - } - - -def _is_boolean_type(data_type: str) -> bool: - return _data_type_base(data_type) in {"BIT", "BOOL", "BOOLEAN"} - - -def _is_categorical_value_column(column: dict[str, Any]) -> bool: - return not ( - _is_identifier_like_column(column) - or _is_numeric_type(column["data_type"]) - or _is_date_type(column["data_type"]) - or _is_boolean_type(column["data_type"]) - ) + return data_type.upper() in {"CHAR", "NCHAR", "NVARCHAR", "STRING", "TEXT", "VARCHAR"} def _missing_value_predicate(column: dict[str, str]) -> str: @@ -2314,6 +1957,30 @@ def _select_listing_columns( score += 35 if name == measure_column: score += 12 + if query_tokens & {"order"}: + score += len(tokens & {"order", "ord", "number", "customer", "cust", "item", "product"}) * 18 + score += len(tokens & {"date", "day", "month", "year"}) * 8 + if query_tokens & {"customer"}: + score += len(tokens & {"customer", "cust", "name", "number", "id"}) * 18 + if query_tokens & {"product", "material"}: + score += len(tokens & {"product", "prod", "material", "item", "description", "desc"}) * 18 + if query_tokens & {"batch"}: + score += len(tokens & {"batch", "board", "model", "supplier", "id"}) * 18 + if query_tokens & {"repair", "log", "record"}: + score += ( + len(tokens & {"board", "code", "date", "failure", "id", "priority", "status"}) + * 14 + ) + if tokens & {"repair", "failure", "failed", "defect"} and not query_tokens & { + "repair", + "failure", + "defect", + }: + score -= 40 + if tokens & {"date", "day", "month", "year"} and not ( + _is_date_type(column["data_type"]) or name == date_column + ): + score -= 12 if score > 0: scored_columns.append((score, index, name)) @@ -2342,88 +2009,10 @@ def _fallback_month_filter(query: str) -> tuple[int, int] | None: return None -def _grouping_phrase_tokens(query: str) -> set[str]: - phrase = _grouping_phrase_text(query) - return _fallback_tokens(phrase) if phrase else set() - - -def _grouping_phrase_has_multiple_dimensions(query: str) -> bool: - phrase = _grouping_phrase_text(query) - return bool(phrase and re.search(r"(?i)(?:,|/|\band\b|\bor\b)", phrase)) - - -def _grouping_phrase_text(query: str) -> str | None: - query = re.sub( - r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", - "", - query, - ) - match = re.search( - r"(?is)\b(?:grouped\s+by|group\s+by|by|across|per)\s+(?P[A-Za-z0-9_ /-]+)", - query, - ) - if not match: - return None - return match.group("value") - - -def _current_year_where_clause( - date_column: str | None, - columns: list[dict[str, str]], - query_tokens: set[str], -) -> str: - if not date_column or not {"this", "year"}.issubset(query_tokens): - return "" - - column = next((column for column in columns if column["name"] == date_column), None) - if not column: - return "" - - current_year = datetime.now(timezone.utc).year - quoted_column = _quote_identifier(date_column) - if _is_numeric_type(column["data_type"]) or _fallback_tokens(date_column) & {"year"}: - return f"\nWHERE {quoted_column} = {current_year}" - if _is_date_type(column["data_type"]) or _fallback_tokens(date_column) & { - "date", - "day", - "month", - "time", - }: - return f"\nWHERE {_date_part_expression(date_column, 'YEAR')} = {current_year}" - return "" - - def _quote_literal(value: str) -> str: return "'" + value.replace("'", "''") + "'" -def _value_match_predicate( - column: dict[str, str] | str, - value: str, - alternate_values: list[str] | None = None, -) -> str: - column_name = column["name"] if isinstance(column, dict) else column - quoted_column = _quote_identifier(column_name) - values = [] - for candidate in [value] + (alternate_values or []): - cleaned = _clean_filter_value(candidate) - if cleaned and cleaned.lower() not in {item.lower() for item in values}: - values.append(cleaned) - - if not values: - return f"{quoted_column} IS NOT NULL" - - if isinstance(column, dict) and ( - _is_text_type(column["data_type"]) or _is_categorical_value_column(column) - ): - lowered_values = [_quote_literal(candidate.lower()) for candidate in values] - if len(lowered_values) == 1: - return f"LOWER({quoted_column}) = {lowered_values[0]}" - return f"LOWER({quoted_column}) IN ({', '.join(lowered_values)})" - - return f"{quoted_column} = {_quote_literal(values[0])}" - - def _clean_filter_value(value: str | None) -> str | None: if value is None: return None @@ -2431,264 +2020,110 @@ def _clean_filter_value(value: str | None) -> str | None: return value or None -def _query_content_tokens(query_tokens: set[str]) -> set[str]: - return { - token - for token in query_tokens - if token not in _GENERIC_SCHEMA_INTENT_TOKENS - and token not in _MONTH_NAME_TO_NUMBER - } - - -def _has_grouping_intent(query: str, query_tokens: set[str]) -> bool: - query_without_ordering = re.sub( - r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", - "", - query, - ) - return bool( - _is_distribution_metric_intent(query_tokens) - or query_tokens & {"group", "grouped", "per"} - or re.search(r"(?i)\bby\s+[A-Za-z0-9_ -]+\b", query_without_ordering) - ) - - -def _has_count_intent(query_tokens: set[str]) -> bool: - return bool(query_tokens & _COUNT_METRIC_TOKENS) - +def _extract_failure_type_filter_value(query: str) -> str | None: + patterns = [ + ( + r"(?is)\bwith\s+" + r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80}?)" + r"\s+as\s+(?:the\s+)?(?:failure|defect)\s+" + r"(?:type|code|category)\b" + ), + ( + r"(?is)\b(?:failure|defect)\s+(?:type|code|category)\s*" + r"(?:=|is|equals|like|of)\s*['\"]?" + r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80})" + ), + ] + for pattern in patterns: + match = re.search(pattern, query) + if match: + value = _clean_filter_value(match.group("value")) + if value: + return value + return None -def _has_sum_intent(query_tokens: set[str]) -> bool: - return bool(query_tokens & _SUM_METRIC_TOKENS) +def _extract_status_filter_value(query: str) -> str | None: + if re.search(r"(?i)\bin-progress\b", query): + return "in-progress" + if re.search(r"(?i)\bin\s+progress\b", query): + return "in progress" + for status in ("completed", "pending", "escalated", "open", "closed"): + if re.search(rf"(?i)\b{re.escape(status)}\b", query): + return status + return None -def _has_extreme_intent(query_tokens: set[str]) -> bool: - return bool(query_tokens & (_MAX_METRIC_TOKENS | _MIN_METRIC_TOKENS)) +def _extract_priority_filter_value(query: str) -> str | None: + for token, value in _PRIORITY_VALUE_ALIASES.items(): + if re.search(rf"(?i)\b{re.escape(token)}(?:[-\s]+priority)?\b", query): + return value + return None -def _has_latest_intent(query_tokens: set[str]) -> bool: - return bool(query_tokens & _LATEST_METRIC_TOKENS) +def _choose_priority_column(columns: list[dict[str, str]]) -> dict[str, str] | None: + return _choose_ranked_column_by_tokens( + columns, + {"priority", "severity", "urgency", "rank"}, + ) -def _has_missing_value_intent(query_tokens: set[str]) -> bool: - return bool(query_tokens & _NULL_CHECK_TOKENS) +def _priority_order_expression(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_numeric_type(column["data_type"]): + return quoted_column -def _sort_direction_for_query(query_tokens: set[str]) -> str: - return "ASC" if query_tokens & _MIN_METRIC_TOKENS else "DESC" + when_clauses = " ".join( + f"WHEN {_quote_literal(value)} THEN {rank}" for value, rank in _PRIORITY_ORDER + ) + return f"CASE LOWER({quoted_column}) {when_clauses} ELSE 0 END" -def _choose_numeric_measure_column( - query_tokens: set[str], +def _choose_failure_type_filter_column( columns: list[dict[str, str]], ) -> dict[str, str] | None: - content_tokens = _query_content_tokens(query_tokens) - measure_candidates = [ - column for column in columns if not _is_identifier_like_column(column) - ] column = _choose_ranked_column_by_tokens( - measure_candidates, - content_tokens, - numeric=True, + [column for column in columns if not _is_rate_like_column(column)], + {"failure", "type"}, ) if column: return column + return _choose_ranked_column_by_tokens( + [column for column in columns if not _is_rate_like_column(column)], + {"failure", "defect", "code", "type", "sys"}, + ) - numeric_columns = [ - column - for column in columns - if _is_numeric_type(column["data_type"]) - and not _is_identifier_like_column(column) - ] - if len(numeric_columns) == 1: - return numeric_columns[0] - return None - - -def _sample_value_filters( - query_tokens: set[str], - columns: list[dict[str, Any]], -) -> list[tuple[dict[str, Any], list[str]]]: - filters: list[tuple[dict[str, Any], list[str]]] = [] - consumed_tokens: set[str] = set() - content_tokens = _query_content_tokens(query_tokens) - if not content_tokens: - return filters - - for column in columns: - matches: list[str] = [] - for value in column.get("sample_values") or []: - cleaned_value = _clean_filter_value(str(value)) - if not cleaned_value: - continue - value_tokens = _fallback_tokens(cleaned_value) - if not value_tokens or not value_tokens <= content_tokens: - continue - if value_tokens <= consumed_tokens: - continue - if cleaned_value.lower() not in {item.lower() for item in matches}: - matches.append(cleaned_value) - consumed_tokens.update(value_tokens) - if matches: - filters.append((column, matches)) - - return filters - - -def _mentioned_text_columns_for_query( - query_tokens: set[str], - columns: list[dict[str, Any]], -) -> list[dict[str, Any]]: - mentioned_columns = [] - for column in columns: - if not _is_categorical_value_column(column): - continue - column_tokens = _fallback_tokens(column["name"]) - if any(_fallback_token_variants(token) & column_tokens for token in query_tokens): - mentioned_columns.append(column) - return mentioned_columns - - -def _matched_column_query_tokens( - query_tokens: set[str], - column: dict[str, Any], -) -> set[str]: - column_tokens = _fallback_tokens(column["name"]) - matched_tokens: set[str] = set() - for token in query_tokens: - matched_tokens.update(_fallback_token_variants(token) & column_tokens) - return matched_tokens - - -def _query_schema_value_terms( - query: str | None, - schema_tokens: set[str], -) -> list[str]: - if not query: - return [] - values: list[str] = [] - matches = list(_QUERY_VALUE_TOKEN.finditer(query)) - for index, match in enumerate(matches): - raw_value = match.group(0).replace("_", " ") - value_tokens = _fallback_tokens(raw_value) - if not value_tokens: - continue - previous_tokens = ( - _fallback_tokens(matches[index - 1].group(0).replace("_", " ")) - if index > 0 - else set() - ) - explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { - "'", - '"', - } - if previous_tokens & schema_tokens and not explicitly_quoted: - continue - if value_tokens & _GENERIC_SCHEMA_INTENT_TOKENS: - continue - if value_tokens & schema_tokens: +def _choose_failure_type_filter_table( + schema_details: dict[str, list[dict[str, str]]], + concept_tokens: set[str], +) -> tuple[str, list[dict[str, str]], dict[str, str]] | None: + candidates = [] + for table_name, columns in schema_details.items(): + if not _table_covers_requested_concepts(table_name, columns, concept_tokens): continue - if all(token.isdigit() for token in value_tokens): + column = _choose_failure_type_filter_column(columns) + if not column: continue - cleaned_value = _clean_filter_value(raw_value) - if cleaned_value and cleaned_value.lower() not in { - value.lower() for value in values - }: - values.append(cleaned_value) - - return values - - -def _schema_driven_user_value_tokens( - query: str | None, - query_tokens: set[str], - schema_details: dict[str, list[dict[str, Any]]], -) -> set[str]: - if not query: - return set() - - mentioned_columns: list[dict[str, Any]] = [] - for columns in schema_details.values(): - mentioned_columns.extend(_mentioned_text_columns_for_query(query_tokens, columns)) - - matched_concepts = [ - _matched_column_query_tokens(query_tokens, column) - for column in mentioned_columns - ] - matched_concepts = [concepts for concepts in matched_concepts if concepts] - if not matched_concepts: - logger.info( - "Schema-derived user value grounding skipped: no mentioned categorical columns query=%s", - query, - ) - return set() - shared_concepts = set.intersection(*matched_concepts) - if not shared_concepts: - logger.info( - "Schema-derived user value grounding skipped: ambiguous categorical concepts query=%s columns=%s concepts=%s", - query, - [column["name"] for column in mentioned_columns], - [sorted(concepts) for concepts in matched_concepts], - ) - return set() - - schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) - value_tokens: set[str] = set() - value_terms = _query_schema_value_terms(query, schema_tokens) - for value in value_terms: - value_tokens.update(_fallback_tokens(value)) - logger.info( - "Schema-derived user value grounding query=%s shared_concepts=%s values=%s value_tokens=%s columns=%s", - query, - sorted(shared_concepts), - value_terms, - sorted(value_tokens), - [column["name"] for column in mentioned_columns], - ) - return value_tokens - - -def _schema_driven_user_value_filters( - query: str | None, - query_tokens: set[str], - table_name: str, - columns: list[dict[str, Any]], -) -> list[tuple[dict[str, Any], list[str]]]: - mentioned_columns = _mentioned_text_columns_for_query(query_tokens, columns) - if len(mentioned_columns) != 1: - return [] - - schema_tokens = _schema_tokens_for_table(table_name, columns) - values = _query_schema_value_terms(query, schema_tokens) - return [(mentioned_columns[0], values)] if values else [] - - -def _where_clause(predicates: list[str]) -> str: - return f"\nWHERE {' AND '.join(predicates)}" if predicates else "" - - -def _count_expression_for_query( - query_tokens: set[str], - columns: list[dict[str, str]], -) -> tuple[str, list[str]]: - subject_column = _choose_count_subject_column(query_tokens, columns) - if not subject_column: - return "COUNT(*)", [] - return ( - f"COUNT({_quote_identifier(subject_column['name'])})", - [_non_missing_value_predicate(subject_column)], - ) - - -def _date_bucket_expressions(date_column: str) -> tuple[str, str]: - return ( - _date_part_expression(date_column, "YEAR"), - _date_part_expression(date_column, "MONTH"), - ) - + table_tokens = _fallback_tokens(table_name) + column_tokens = _fallback_tokens(column["name"]) + score = len(concept_tokens & table_tokens) * 8 + score += len(concept_tokens & column_tokens) * 10 + if {"failure", "type"}.issubset(column_tokens): + score += 100 + elif "type" in column_tokens: + score += 60 + elif "code" in column_tokens: + score += 30 + if table_tokens & {"failure", "defect"}: + score += 25 + candidates.append((score, table_name, columns, column)) -def _date_part_expression(date_column: str, part: str) -> str: - return f"CAST(EXTRACT({part} FROM {_quote_identifier(date_column)}) AS BIGINT)" + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1], item[3]["name"])) + _, table_name, columns, column = candidates[0] + return table_name, columns, column def generate_simple_analytics_sql( @@ -2698,45 +2133,87 @@ def generate_simple_analytics_sql( if not query: return None - query_tokens = _fallback_tokens(query) + raw_query_tokens = _fallback_tokens(query) + query_tokens = _expanded_fallback_query_tokens(query) if not query_tokens: return None - schema_details = _extract_schema_details(contexts) - if not schema_details: + if not query_tokens & { + "batch", + "batches", + "board", + "business", + "count", + "customer", + "defect", + "failure", + "failures", + "highest", + "july", + "latest", + "log", + "logs", + "location", + "material", + "missing", + "model", + "monthly", + "most", + "number", + "order", + "orders", + "priority", + "product", + "rate", + "recent", + "record", + "records", + "repair", + "repairs", + "revenue", + "sale", + "sales", + "severity", + "supplier", + "status", + "tech", + "technician", + "top", + "trend", + "trends", + "type", + "unit", + "units", + "year", + }: return None - content_tokens = _query_content_tokens(query_tokens) - schema_backed_tokens = _schema_derived_query_tokens(query_tokens, schema_details) - unsupported_tokens = _unsupported_query_tokens( - query_tokens, - schema_details, - query=query, + schema_details = _extract_schema_details(contexts) + rate_metric_intent = _is_rate_metric_intent(raw_query_tokens) + failure_count_intent = _is_failure_count_intent(raw_query_tokens, query_tokens) + board_model_intent = _has_board_model_intent( + raw_query_tokens + ) or _has_board_model_intent( + query_tokens ) - if unsupported_tokens: - logger.info( - "Schema-derived SQL fallback skipped unsupported_tokens=%s", - sorted(unsupported_tokens), - ) - return None - if content_tokens and not schema_backed_tokens: - logger.info( - "Schema-derived SQL fallback skipped no_schema_backed_tokens=%s", - sorted(content_tokens), - ) - return None - if not content_tokens and len(schema_details) != 1: - logger.info( - "Schema-derived SQL fallback skipped ambiguous_schema_only_request tables=%s", - sorted(schema_details), + failure_type_filter_value = _extract_failure_type_filter_value(query) + failure_type_filter_column = None + failure_type_choice = None + if failure_type_filter_value and "failure" in query_tokens: + failure_type_choice = _choose_failure_type_filter_table( + schema_details, + raw_query_tokens, ) - return None - chosen = _choose_fallback_table( - query_tokens, - schema_details, - concept_tokens=query_tokens, - ) + if failure_type_choice: + table_name, columns, failure_type_filter_column = failure_type_choice + chosen = (table_name, columns) + else: + chosen = _choose_fallback_table( + query_tokens, + schema_details, + concept_tokens=raw_query_tokens, + ) if not chosen: return None @@ -2744,201 +2221,108 @@ def generate_simple_analytics_sql( column_names = [column["name"] for column in columns] quoted_table = _quote_identifier(table_name) limit = _fallback_limit(query) - date_column = _choose_temporal_column(query_tokens, columns) - order_column = _choose_order_by_column(query, query_tokens, columns) - sample_filters = _sample_value_filters(query_tokens, columns) - sample_filter_column_names = {column["name"] for column, _ in sample_filters} - user_value_filters = [ - (column, values) - for column, values in _schema_driven_user_value_filters( - query, - query_tokens, - table_name, - columns, - ) - if column["name"] not in sample_filter_column_names - ] - value_filters = [*sample_filters, *user_value_filters] - sample_predicates = [ - _filter_predicate_for_values(column, values) - for column, values in value_filters - ] - selected_sample_filter_columns = [column["name"] for column, _ in value_filters] - month_filter = _fallback_month_filter(query) - metric_intent = { - "average": _is_average_metric_intent(query_tokens), - "count": _has_count_intent(query_tokens), - "distribution": _is_distribution_metric_intent(query_tokens), - "extreme": _has_extreme_intent(query_tokens), - "latest": _has_latest_intent(query_tokens), - "missing": _has_missing_value_intent(query_tokens), - "rate": _is_rate_metric_intent(query_tokens), - "sum": _has_sum_intent(query_tokens), - } logger.info( - "Schema-derived SQL fallback selected table=%s schema_tokens=%s verified_columns=%s sample_filter_columns=%s metric_intent=%s", + "Deterministic SQL fallback selected table=%s verified_columns=%s metric_intent=%s", table_name, - sorted(schema_backed_tokens), column_names, - selected_sample_filter_columns, - metric_intent, + { + "failure_count": failure_count_intent, + "rate": rate_metric_intent, + "board_model": board_model_intent, + "failure_type_filter": bool(failure_type_filter_value), + }, ) - if _has_missing_value_intent(query_tokens): - missing_column = _choose_missing_value_column(query, query_tokens, columns) - if not missing_column: - return None - selected_columns = _select_listing_columns( - query_tokens, - columns, - date_column=date_column, - max_columns=8, - ) - if missing_column["name"] not in selected_columns: - selected_columns.insert(0, missing_column["name"]) - predicates = [*sample_predicates, _missing_value_predicate(missing_column)] - limit_clause = f"\nLIMIT {limit}" if limit else "" - return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" - ) - - month_predicate = "" - if date_column and month_filter: - year, month = month_filter - start_date = f"{year:04d}-{month:02d}-01" - end_year = year + 1 if month == 12 else year - end_month = 1 if month == 12 else month + 1 - end_date = f"{end_year:04d}-{end_month:02d}-01" - quoted_date = _quote_identifier(date_column) - month_predicate = ( - f"{quoted_date} >= '{start_date}' AND {quoted_date} < '{end_date}'" - ) - - predicates = list(sample_predicates) - if month_predicate: - predicates.append(month_predicate) - - if _is_average_metric_intent(query_tokens): - measure_column = _choose_average_measure_column(query_tokens, columns) - if not measure_column: - return None - grouping_tokens = _grouping_phrase_tokens(query) or query_tokens - dimension_columns = _choose_dimension_columns( - grouping_tokens, - columns, - max_columns=2, - ) - aggregate_expr = f"AVG({_quote_identifier(measure_column['name'])})" - if dimension_columns: - quoted_dimensions = _quote_joined(dimension_columns) + if failure_type_filter_value: + if not failure_type_filter_column: + failure_type_filter_column = _choose_failure_type_filter_column(columns) + if failure_type_filter_column: return ( - f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier('average_value')}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"GROUP BY {quoted_dimensions}\n" - f"ORDER BY {_quote_identifier('average_value')} DESC" + f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}\n" + f"WHERE {_quote_identifier(failure_type_filter_column['name'])} = " + f"{_quote_literal(failure_type_filter_value)}" ) - return ( - f"SELECT {aggregate_expr} AS {_quote_identifier('average_value')}\n" - f"FROM {quoted_table}{_where_clause(predicates)}" - ) - - if date_column and query_tokens & {"month", "monthly"} and _has_count_intent(query_tokens): - year_expr, month_expr = _date_bucket_expressions(date_column) - return ( - f"SELECT {year_expr} AS {_quote_identifier('year')}, " - f"{month_expr} AS {_quote_identifier('month')}, " - f"COUNT(*) AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" - f"GROUP BY {year_expr}, {month_expr}\n" - f"ORDER BY {year_expr}, {month_expr}" - ) - measure_column = _choose_numeric_measure_column(query_tokens, columns) + material_column = _choose_column_by_tokens(columns, {"material"}) + location_column = _choose_column_by_tokens(columns, {"location"}) if ( - measure_column - and date_column - and query_tokens & {"month", "monthly"} - and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) + {"material", "location"}.issubset(query_tokens) + and material_column + and location_column ): - year_expr, month_expr = _date_bucket_expressions(date_column) - aggregate, alias = _aggregate_for_measure(measure_column["name"]) return ( - f"SELECT {year_expr} AS {_quote_identifier('year')}, " - f"{month_expr} AS {_quote_identifier('month')}, " - f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" - f"GROUP BY {year_expr}, {month_expr}\n" - f"ORDER BY {year_expr}, {month_expr}" + f"SELECT {_quote_joined([material_column, location_column])}\n" + f"FROM {quoted_table}" ) + status_column = _choose_column_by_tokens(columns, {"status"}) + repair_filter_intent = raw_query_tokens & { + "closed", + "completed", + "critical", + "escalated", + "high", + "low", + "medium", + "normal", + "open", + "pending", + "priority", + "progress", + "severity", + "status", + "urgent", + } + if query_tokens & {"repair", "repairs"} and repair_filter_intent: + predicates = [] + status_filter_value = _extract_status_filter_value(query) + if status_column and status_filter_value: + predicates.append( + f"{_quote_identifier(status_column)} = {_quote_literal(status_filter_value)}" + ) + priority_filter_value = _extract_priority_filter_value(query) + if priority_filter_value: + priority_column = _choose_priority_column(columns) + if priority_column: + predicates.append( + f"{_quote_identifier(priority_column['name'])} = " + f"{_quote_literal(priority_filter_value)}" + ) + if predicates: + return f"SELECT *\nFROM {quoted_table}\nWHERE {' AND '.join(predicates)}" + + date_column = _choose_column_by_tokens( + columns, + {"date", "day", "month", "time", "year"}, + date=True, + ) + + priority_column = _choose_priority_column(columns) if ( - measure_column - and date_column - and "year" in query_tokens - and "month" not in query_tokens - and "monthly" not in query_tokens - and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) + priority_column + and raw_query_tokens & {"priority", "severity"} + and raw_query_tokens & {"bottom", "highest", "lowest", "top"} ): - year_expr = _date_part_expression(date_column, "YEAR") - aggregate, alias = _aggregate_for_measure(measure_column["name"]) - return ( - f"SELECT {year_expr} AS {_quote_identifier('year')}, " - f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" - f"GROUP BY {year_expr}\nORDER BY {year_expr}" - ) - - if _has_grouping_intent(query, query_tokens) or _has_count_intent(query_tokens): - explicit_grouping_tokens = _grouping_phrase_tokens(query) - grouping_tokens = explicit_grouping_tokens or query_tokens - max_dimensions = 1 if _has_extreme_intent(query_tokens) else 3 - if explicit_grouping_tokens and not _grouping_phrase_has_multiple_dimensions(query): - max_dimensions = 1 - dimension_columns = _choose_dimension_columns( - grouping_tokens, + selected_columns = _select_listing_columns( + raw_query_tokens | {"priority", "repair", "status"}, columns, - max_columns=max_dimensions, + date_column=date_column, + max_columns=8, ) - if dimension_columns: - quoted_dimensions = _quote_joined(dimension_columns) - if measure_column and (_has_sum_intent(query_tokens) or _is_rate_metric_intent(query_tokens)): - aggregate, alias = _aggregate_for_measure(measure_column["name"]) - aggregate_expr = f"{aggregate}({_quote_identifier(measure_column['name'])})" - direction = _sort_direction_for_query(query_tokens) - limit_clause = f"\nLIMIT {limit}" if limit else "" - return ( - f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"GROUP BY {quoted_dimensions}\n" - f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" - ) - - limit_clause = f"\nLIMIT {limit}" if limit else "" - return ( - f"SELECT {quoted_dimensions}, COUNT(*) AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"GROUP BY {quoted_dimensions}\n" - f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" - ) - if _has_count_intent(query_tokens) and not _has_grouping_intent(query, query_tokens): - return ( - f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{_where_clause(predicates)}" - ) - return None - - if measure_column and _has_sum_intent(query_tokens): + if priority_column["name"] not in selected_columns: + selected_columns.insert(0, priority_column["name"]) + direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" + limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT SUM({_quote_identifier(measure_column['name'])}) AS {_quote_identifier('total_value')}\n" - f"FROM {quoted_table}{_where_clause(predicates)}" + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"ORDER BY {_priority_order_expression(priority_column)} {direction}" + f"{limit_clause}" ) - if _has_latest_intent(query_tokens): - if not date_column: - return None + if date_column and raw_query_tokens & {"latest", "recent"}: selected_columns = _select_listing_columns( - query_tokens, + raw_query_tokens | {"date", "repair", "status"}, columns, date_column=date_column, max_columns=8, @@ -2947,101 +2331,235 @@ def generate_simple_analytics_sql( selected_columns.insert(0, date_column) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" ) - if _has_extreme_intent(query_tokens): - direction = _sort_direction_for_query(query_tokens) - if measure_column: - limit_clause = f"\nLIMIT {limit}" if limit else "" - if _query_allows_grouped_aggregate(query, query_tokens): - dimension_column = _choose_dimension_column(query_tokens, columns) - else: - dimension_column = None - if dimension_column: - aggregate, alias = _aggregate_for_measure(measure_column["name"]) - return ( - f"SELECT {_quote_identifier(dimension_column)}, " - f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"GROUP BY {_quote_identifier(dimension_column)}\n" - f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" - ) - selected_columns = _select_listing_columns( - query_tokens, - columns, - measure_column=measure_column["name"], - date_column=date_column, - max_columns=8, - ) - if measure_column["name"] not in selected_columns: - selected_columns.insert(0, measure_column["name"]) + if ( + query_tokens & {"repair", "repairs"} + and raw_query_tokens & {"priority", "severity", "status"} + and re.search(r"(?i)\bby\s+(?:priority|severity|status)\b", query) + ): + dimension_column = _choose_dimension_column(raw_query_tokens, columns) + subject_column = _choose_count_subject_column({"repair"}, columns) + if dimension_column: + count_expression = "COUNT(*)" + where_clause = "" + if subject_column: + count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" + where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" + quoted_dimension = _quote_identifier(dimension_column) return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"ORDER BY {_quote_identifier(measure_column['name'])} {direction}{limit_clause}" + f"SELECT {quoted_dimension}, {count_expression} AS " + f"{_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC" ) - if order_column: - selected_columns = _select_listing_columns( - query_tokens, - columns, - date_column=date_column, - max_columns=8, + measure_column = None + if rate_metric_intent and query_tokens & {"defect", "failure"}: + measure_column = _choose_column_by_tokens( + columns, + {"defect", "rate"}, + numeric=True, + ) + if not measure_column: + measure_column = _choose_column_by_tokens(columns, {"rate"}, numeric=True) + if not measure_column and query_tokens & {"order", "orders"}: + measure_column = _choose_column_by_tokens( + columns, + {"amount", "intake", "sales", "value"}, + numeric=True, + ) + if not measure_column and query_tokens & {"revenue", "sale", "sales"}: + measure_column = _choose_column_by_tokens( + columns, + {"amount", "intake", "revenue", "sales", "value"}, + numeric=True, + ) + if ( + not measure_column + and not failure_count_intent + and query_tokens & {"top", "highest", "lowest", "bottom"} + ): + measure_column = _choose_column_by_tokens( + columns, + {"amount", "cost", "count", "margin", "quantity", "rate", "score", "value"}, + numeric=True, + ) + + if raw_query_tokens & {"missing", "blank", "empty", "null"}: + missing_column = _choose_missing_value_column(raw_query_tokens, columns) + if missing_column: + selected_columns = [ + column + for column in column_names + if column == missing_column["name"] + or _fallback_tokens(column) + & { + "batch", + "business", + "bu", + "customer", + "cust", + "date", + "id", + "location", + "name", + "number", + "ord", + "order", + "product", + "status", + "supplier", + } + ][:8] + if missing_column["name"] not in selected_columns: + selected_columns.insert(0, missing_column["name"]) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"WHERE {_missing_value_predicate(missing_column)}{limit_clause}" ) - if order_column not in selected_columns: - selected_columns.insert(0, order_column) + + implied_count_by_dimension = "failure" in query_tokens and bool( + raw_query_tokens & {"location", "material", "technician", "tech"} + or (board_model_intent and not rate_metric_intent) + ) + if ( + query_tokens & {"count", "number"} + or failure_count_intent + or implied_count_by_dimension + ): + dimension_column = _choose_dimension_column(raw_query_tokens, columns) + if dimension_column: + subject_column = _choose_count_subject_column(raw_query_tokens, columns) + count_expression = "COUNT(*)" + where_clause = "" + if subject_column and raw_query_tokens & {"order", "customer"}: + count_expression = ( + f"COUNT(DISTINCT {_quote_identifier(subject_column['name'])})" + ) + elif subject_column and raw_query_tokens & {"failure", "repair", "batch"}: + count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" + where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" + quoted_dimension = _quote_identifier(dimension_column) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"ORDER BY {_quote_identifier(order_column)} {direction}{limit_clause}" + f"SELECT {quoted_dimension}, {count_expression} AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" ) - return None - if month_predicate and date_column: - selected_columns = _select_listing_columns( - query_tokens, - columns, - measure_column=measure_column["name"] if measure_column else None, - date_column=date_column, - max_columns=8, - ) + dimension_column = _choose_dimension_column(raw_query_tokens, columns) + if measure_column and dimension_column and rate_metric_intent: + aggregate, alias = _aggregate_for_measure(measure_column) + quoted_dimension = _quote_identifier(dimension_column) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" + direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" + f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + + if measure_column and limit and dimension_column and raw_query_tokens & { + "board", + "business", + "customer", + "location", + "material", + "model", + "product", + "salesperson", + "supplier", + "unit", + }: + aggregate, alias = _aggregate_for_measure(measure_column) + quoted_dimension = _quote_identifier(dimension_column) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" + return ( + f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" + f"ORDER BY {_quote_identifier(alias)} DESC\nLIMIT {limit}" ) - if order_column: + month_filter = _fallback_month_filter(query) + if date_column and month_filter: + year, month = month_filter + start = f"{year:04d}-{month:02d}-01" + end_year = year + 1 if month == 12 else year + end_month = 1 if month == 12 else month + 1 + end = f"{end_year:04d}-{end_month:02d}-01" selected_columns = _select_listing_columns( - query_tokens, + raw_query_tokens, columns, + measure_column=measure_column, date_column=date_column, - max_columns=8, ) - if order_column not in selected_columns: - selected_columns.insert(0, order_column) + order_clause = ( + f"\nORDER BY {_quote_identifier(measure_column)} DESC" + if measure_column + else f"\nORDER BY {_quote_identifier(date_column)} DESC" + ) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}\n" - f"ORDER BY {_quote_identifier(order_column)} ASC{limit_clause}" + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"WHERE {_quote_identifier(date_column)} >= '{start}' " + f"AND {_quote_identifier(date_column)} < '{end}'" + f"{order_clause}{limit_clause}" ) - if sample_predicates or query_tokens & {"all", "list"}: - selected_columns = _select_listing_columns( - query_tokens, - columns, - date_column=date_column, - max_columns=8, + if ( + measure_column + and date_column + and query_tokens & {"year"} + and not query_tokens & {"month", "monthly"} + ): + date_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" + return ( + f"SELECT {date_expr} AS {_quote_identifier('year')}, " + f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" + f"FROM {quoted_table}\nGROUP BY {date_expr}\nORDER BY {date_expr}" ) - limit_clause = f"\nLIMIT {limit}" if limit else "" + + if ( + measure_column + and date_column + and query_tokens & {"month", "monthly", "trend", "trends"} + ): + year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" + month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" return ( - f"SELECT {_quote_joined(selected_columns)}\n" - f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" + f"FROM {quoted_table}\nGROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" + ) + + if measure_column and limit: + selected_columns = [ + column + for column in column_names + if column == measure_column + or _fallback_tokens(column) + & { + "batch", + "board", + "customer", + "id", + "model", + "name", + "number", + "supplier", + } + ][:6] + if measure_column not in selected_columns: + selected_columns.append(measure_column) + return ( + f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"ORDER BY {_quote_identifier(measure_column)} DESC\nLIMIT {limit}" ) return None @@ -3069,27 +2587,9 @@ async def run( allow_data_preview: bool = False, ) -> dict: try: - cleaned_generation_result, extraction_error = _extract_sql_response( - clean_generation_result(replies[0]) - ) + cleaned_generation_result = clean_generation_result(replies[0]) grounding_invalid_generation_result = None - def validate_candidate_sql(candidate_sql: str) -> str | None: - schema_catalog = _SchemaCatalog.from_contexts(contexts or []) - grounding_error = schema_catalog.validate_sql(candidate_sql) - if not grounding_error: - grounding_error = validate_sql_against_contexts( - candidate_sql, - contexts=contexts, - ) - if not grounding_error: - grounding_error = validate_sql_semantic_coverage( - candidate_sql, - fallback_query, - contexts=contexts, - ) - return grounding_error - if cleaned_generation_result: cleaned_generation_result = normalize_sql_with_schema_identifiers( cleaned_generation_result, @@ -3098,7 +2598,16 @@ def validate_candidate_sql(candidate_sql: str) -> str | None: cleaned_generation_result = normalize_wren_sql_dialect( cleaned_generation_result ) - grounding_error = validate_candidate_sql(cleaned_generation_result) + grounding_error = validate_sql_against_contexts( + cleaned_generation_result, + contexts=contexts, + ) + if not grounding_error: + grounding_error = validate_sql_semantic_coverage( + cleaned_generation_result, + fallback_query, + contexts=contexts, + ) if grounding_error: logger.info( "Generated SQL validation result project_id=%s status=rejected reason=%s sql=%s", @@ -3107,9 +2616,9 @@ def validate_candidate_sql(candidate_sql: str) -> str | None: cleaned_generation_result, ) grounding_invalid_generation_result = { - "sql": "", - "original_sql": "", - "type": "NO_RELEVANT_SQL", + "sql": cleaned_generation_result, + "original_sql": cleaned_generation_result, + "type": "SCHEMA_GROUNDING", "error": grounding_error, "correlation_id": "", "data_source": data_source, @@ -3120,12 +2629,6 @@ def validate_candidate_sql(candidate_sql: str) -> str | None: project_id or "", cleaned_generation_result, ) - elif extraction_error: - logger.info( - "Generated SQL extraction result project_id=%s status=rejected reason=%s", - project_id or "", - extraction_error, - ) fallback_generation_result = generate_simple_analytics_sql( fallback_query, @@ -3144,9 +2647,16 @@ def validate_candidate_sql(candidate_sql: str) -> str | None: fallback_generation_result = normalize_wren_sql_dialect( fallback_generation_result ) - fallback_grounding_error = validate_candidate_sql( - fallback_generation_result + fallback_grounding_error = validate_sql_against_contexts( + fallback_generation_result, + contexts=contexts, ) + if not fallback_grounding_error: + fallback_grounding_error = validate_sql_semantic_coverage( + fallback_generation_result, + fallback_query, + contexts=contexts, + ) logger.info( "Deterministic SQL fallback validation result project_id=%s status=%s%s", project_id or "", @@ -3208,19 +2718,6 @@ def validate_candidate_sql(candidate_sql: str) -> str | None: if not cleaned_generation_result and unsupported_result: return unsupported_result - if not cleaned_generation_result and extraction_error: - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": "", - "original_sql": "", - "type": "NO_RELEVANT_SQL", - "error": extraction_error, - "correlation_id": "", - "data_source": data_source, - }, - } - ( valid_generation_result, invalid_generation_result, @@ -3947,10 +3444,10 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. - If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. -- Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata contains specific modeled columns for the requested entity, measure, filter, date, or dimension. -- Prefer exact modeled fields over generic text search. If a requested concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. -- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. -- Do not prefer or exclude any business domain by built-in rules. Ground every choice in the DATABASE SCHEMA supplied for this request. +- Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema contains specific modeled business columns for the requested entity, measure, status, date, or dimension. +- Prefer exact modeled business fields over generic text search. For example, if a status/severity/date/material/location/customer/order/revenue concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. +- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. This applies to metrics such as defect rate, revenue, amount, sales value, count, cost, margin, and quantity. +- For sales or revenue questions, prefer exact declared sales/revenue/value/amount fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. """ @@ -4002,7 +3499,6 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. - Do not use SELECT TOP n, FETCH FIRST, OFFSET/FETCH, square-bracket quoting, or backtick quoting. Use Wren SQL syntax with ORDER BY and a final LIMIT n clause for limited or top-N results. - For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. -- For explicit ranking requests, use the ranking function `DENSE_RANK()`, add the ranking column to the final SELECT clause, and filter rank values with WHERE. - For grouped trend queries, include any non-aggregate ordering key in both SELECT and GROUP BY, or order by selected grouping columns/aggregate aliases only. - Reuse exact metric/measure columns when present. Do not invent component columns in order to calculate a requested metric that already exists in DATABASE SCHEMA. """ diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 0893b8dd61..5dd30dd59a 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -11,9 +11,8 @@ from hamilton.async_driver import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder +from langfuse.decorators import observe from pydantic import BaseModel, ConfigDict -from sqlparse.sql import Identifier, IdentifierList -from sqlparse.tokens import DML, Comment, Keyword from langfuse.decorators import observe from src.core.pipeline import BasicPipeline @@ -33,6 +32,63 @@ _MAX_RETRIEVED_TABLE_NAMES = 24 _MAX_RELATED_TABLE_EXPANSION_DEPTH = 1 _RANK_TOKEN = re.compile(r"[a-z0-9]+") +_GENERIC_TABLE_TOKENS = { + "audit", + "auth", + "calendar", + "config", + "dim", + "dimension", + "file", + "files", + "ingestion", + "job", + "jobs", + "log", + "logs", + "lookup", + "mbr", + "member", + "members", + "migration", + "migrations", + "preference", + "preferences", + "queue", + "report", + "reports", + "setting", + "settings", + "state", + "time", + "user", + "users", +} +_CUSTOMS_FINANCE_TOKENS = { + "claim", + "claims", + "custom", + "customs", + "duty", + "duties", + "hmf", + "import", + "imports", + "mpf", + "refund", + "refunds", + "tariff", + "tariffs", +} +_SALES_REVENUE_TOKENS = { + "amount", + "intake", + "revenue", + "sale", + "sales", + "salesvalue", + "value", +} table_columns_selection_system_prompt = """ @@ -63,15 +119,14 @@ 13. Do not stop at a single top candidate when the question needs multiple related datasets. 14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. 15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. -16. Prefer the set of deployed models, views, metrics, columns, and relationships that best support the current question. -17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. -18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. -19. Prefer tables and columns whose supplied names, descriptions, relationships, metrics, or sample values directly support the requested entities, measures, filters, dates, identifiers, and dimensions. Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata provides specific modeled columns for the same requested concept. -20. Compare the user's requested entities, measures, filters, dates, and dimensions only with schema metadata supplied for the active project. Do not use built-in business synonym lists. -21. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. -22. Never return placeholder table or column names or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. -23. If a requested measure, dimension, filter, or time field is not represented by retrieved schema metadata, leave it unsupported instead of substituting a similar-looking field. -24. Metric intent such as count, sum, average, minimum, maximum, ranking, date bucketing, and grouping must be satisfied by declared columns or metric fields from the retrieved schema. +16. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. +17. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. +18. Prefer tables and columns that directly model the requested business entities, measures, statuses, dates, identifiers, and dimensions. Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema provides specific modeled columns for the same concept. +19. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, repairs, failures, batches, materials, locations, status, severity, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. +20. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. +21. Never return placeholder table or column names such as tablename, table_name, dbo.tablename, BatchId, Material, Location, or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. +22. If the request asks for revenue, sales, or sales trends, prefer exact business measure columns named like Revenue, SalesValue, USDFXSalesValue, FXSalesValue, IntakeValue, Amount, or equivalent modeled sales fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. +23. If the request asks for an explicit rate, ratio, percentage, revenue, amount, sales value, or other named measure and the schema already contains that exact measure column, use the declared measure column directly. Do not use a rate column to answer "most failures", "number of failures", or other count-of-records requests unless the question explicitly asks for a rate/ratio/percentage. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -648,7 +703,7 @@ async def dbschema_retrieval( documents = [] if embedding: semantic_documents = await _retrieve_semantic_schema_documents( - embedding, project_id, mdl_hash, dbschema_retriever + embedding, project_id, dbschema_retriever, mdl_hash ) semantic_table_names = _table_names_from_schema_documents(semantic_documents)[ :_SEMANTIC_TABLE_NAME_MERGE_LIMIT @@ -669,10 +724,9 @@ async def dbschema_retrieval( ] if table_names: - if include_related_models: - retrieved_table_names = set() - pending_table_names = table_names - remaining_expansion_depth = _MAX_RELATED_TABLE_EXPANSION_DEPTH + retrieved_table_names = set() + pending_table_names = table_names + remaining_expansion_depth = _MAX_RELATED_TABLE_EXPANSION_DEPTH while pending_table_names: retrieved_table_names.update(pending_table_names) @@ -706,12 +760,19 @@ async def dbschema_retrieval( for document in ranked_documents ], ) - return ranked_documents + documents = _dedupe_documents(documents + retrieved_documents) + if remaining_expansion_depth <= 0: + break + remaining_expansion_depth -= 1 + remaining_slots = _MAX_RETRIEVED_TABLE_NAMES - len(retrieved_table_names) + if remaining_slots <= 0: + break + pending_table_names = [ + table_name + for table_name in _related_table_names(documents) + if table_name not in retrieved_table_names + ][:remaining_slots] - retrieved_documents = await _retrieve_schema_documents( - table_names, project_id, mdl_hash, dbschema_retriever - ) - documents = _dedupe_documents(documents + retrieved_documents) ranked_documents = _rank_documents_for_query(documents, table_names, query) logger.info( "Ask schema retrieval project_id=%s retrieved_tables=%s", @@ -809,6 +870,15 @@ def score(table_name: str) -> int: value += 8 if direct_table_matches and direct_column_matches: value += 8 + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + value += len((table_tokens | column_tokens) & _SALES_REVENUE_TOKENS) * 5 + if not query_tokens & _CUSTOMS_FINANCE_TOKENS: + value -= ( + len((table_tokens | column_tokens) & _CUSTOMS_FINANCE_TOKENS) + * 8 + ) + if table_tokens & _GENERIC_TABLE_TOKENS and not direct_table_matches: + value -= 6 return value ranked = sorted( @@ -846,7 +916,65 @@ def _rank_documents_for_query( def _augment_retrieval_query(query: str) -> str: - return query + lowered = query.lower() + expansions = [] + + concept_terms = { + ("revenue", "sales", "sale", "amount", "value"): ( + "sales revenue amount value gross net total price intake invoice order" + ), + ("order", "orders"): ( + "order ord number date customer product business unit division company" + ), + ("invoice", "invoices"): ( + "invoice supplier customer currency amount date number" + ), + ("customer", "customers"): ( + "customer account client number name identifier" + ), + ("product", "products"): ( + "product item material type name category" + ), + ("repair", "repairs"): ( + "repair status priority severity failure board model log in progress completed critical" + ), + ("failure", "failures", "defect", "defects"): ( + "failure defect severity occurrence record count code type system status" + ), + ("batch", "batches"): ( + "batch board model supplier defect rate inspection status" + ), + ("material", "materials"): ( + "material item part component location" + ), + ("location", "locations"): ( + "location site warehouse area material" + ), + ("business unit", "bu", "division"): ( + "business unit division company account organization" + ), + ("month", "monthly", "july", "year", "trend", "latest"): ( + "date month year fiscal calendar trend latest recent" + ), + ("status", "severity", "priority", "critical"): ( + "status priority severity critical state category progress" + ), + } + + for triggers, terms in concept_terms.items(): + if any(trigger in lowered for trigger in triggers): + expansions.append(terms) + + if any( + trigger in lowered + for trigger in ("rate", "ratio", "percent", "percentage") + ): + expansions.append("rate ratio percent percentage") + + if not expansions: + return query + + return f"{query}\nBusiness schema search terms: {'; '.join(expansions)}" async def _retrieve_semantic_schema_documents( @@ -1466,6 +1594,307 @@ def _merge_column_selection( return merged +def _lexical_columns_and_tables_needed( + construct_db_schemas: list[dict], + query: str | None, + max_tables: int = 4, + max_columns_per_table: int = 12, +) -> dict[str, dict]: + if not query: + return {} + + query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) + if not query_tokens: + return {} + + scored_tables = [] + for table_schema in construct_db_schemas: + if table_schema.get("type") != "TABLE": + continue + + table_tokens = _tokenize_schema_text( + table_schema.get("name") + ) | _tokenize_schema_text( + table_schema.get("comment") + ) + table_score = len(query_tokens & table_tokens) * 6 + column_scores = [] + +@observe() +def construct_retrieval_results( + check_using_db_schemas_without_pruning: dict, + filter_columns_in_tables: dict, + construct_db_schemas: list[dict], + dbschema_retrieval: list[Document], + query: str | None = None, +) -> dict[str, Any]: + if filter_columns_in_tables: + columns_and_tables_needed = _parse_column_selection_response( + filter_columns_in_tables + ) + lexical_columns_and_tables_needed = _lexical_columns_and_tables_needed( + construct_db_schemas, + query, + ) + columns_and_tables_needed = _merge_column_selection( + columns_and_tables_needed, + lexical_columns_and_tables_needed, + ) + tables = set(columns_and_tables_needed.keys()) + retrieval_results = [] + selected_schema_log = [] + has_calculated_field = False + has_metric = False + has_json_field = False + + for table_schema in construct_db_schemas: + if table_schema["type"] == "TABLE" and table_schema["name"] in tables: + selected_columns = set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ) + columns = ( + selected_columns + if _selected_columns_are_executable( + table_schema, selected_columns + ) + else None + ) + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context( + table_schema, + columns=columns, + tables=tables, + ) + ) + if _has_calculated_field: + has_calculated_field = True + if _has_json_field: + has_json_field = True + + retrieval_results.append( + { + "table_name": table_schema["name"], + "table_ddl": ddl, + } + ) + selected_schema_log.append( + { + "table": table_schema["name"], + "columns": sorted(selected_columns), + } + ) + + if not retrieval_results: + logger.warning( + "Column-selection output did not match retrieved schemas; " + "falling back to unpruned retrieved schema context." + ) + return _build_unpruned_retrieval_results( + construct_db_schemas, dbschema_retrieval + ) + + if not column_scores and table_score <= 0: + continue + + total_score = table_score + sum(score for score, _, _ in column_scores) + if total_score <= 0: + continue + + logger.info("Ask retrieval selected schema objects=%s", selected_schema_log) + return { + "retrieval_results": retrieval_results, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + else: + retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] + logger.info( + "Ask retrieval selected schema objects=%s", + [ + {"table": retrieval_result.get("table_name"), "columns": "all"} + for retrieval_result in retrieval_results + ], + ) + + scored_tables.append((total_score, table_schema["name"], selected_columns)) + + scored_tables.sort(key=lambda item: (-item[0], item[1])) + return { + table_name: {"columns": columns} + for _, table_name, columns in scored_tables[:max_tables] + if columns + } + + +def _normalize_column_selection_results(parsed_response: Any) -> list[dict]: + if isinstance(parsed_response, list): + return [item for item in parsed_response if isinstance(item, dict)] + + if not isinstance(parsed_response, dict): + return [] + + for key in ( + "results", + "tables", + "selected_tables", + "retrieval_results", + "matches", + "data", + "result", + "output", + ): + if key in parsed_response: + normalized = _normalize_column_selection_results(parsed_response[key]) + if normalized: + return normalized + + if "table_name" in parsed_response and ( + "table_contents" in parsed_response or "columns" in parsed_response + ): + return [parsed_response] + + keyed_tables = [] + for table_name, table_contents in parsed_response.items(): + if not isinstance(table_name, str) or not isinstance(table_contents, dict): + continue + if "table_contents" in table_contents: + keyed_tables.append( + { + "table_name": table_name, + "table_contents": table_contents["table_contents"], + } + ) + elif "columns" in table_contents: + keyed_tables.append( + {"table_name": table_name, "table_contents": table_contents} + ) + + return keyed_tables + + +def _parse_column_selection_response(filter_columns_in_tables: dict) -> dict: + raw_reply = (filter_columns_in_tables.get("replies") or [""])[0] + try: + parsed_response = orjson.loads(raw_reply) + except orjson.JSONDecodeError as exc: + logger.warning("Unable to parse column-selection JSON response: %s", exc) + return {} + + normalized_tables = _normalize_column_selection_results(parsed_response) + reformatted_json = {} + for table in normalized_tables: + table_name = table.get("table_name") or table.get("name") + table_contents = table.get("table_contents") or {} + if not table_contents and "columns" in table: + table_contents = table + + columns = ( + table_contents.get("columns") if isinstance(table_contents, dict) else None + ) + if not isinstance(table_name, str) or not isinstance(columns, list): + continue + + reformatted_json[table_name] = { + **table_contents, + "columns": [column for column in columns if isinstance(column, str)], + } + + if not reformatted_json: + response_shape = ( + f"keys={list(parsed_response.keys())[:8]}" + if isinstance(parsed_response, dict) + else type(parsed_response).__name__ + ) + logger.warning( + "Column-selection response did not include usable table columns (%s).", + response_shape, + ) + + return reformatted_json + + +def _build_unpruned_retrieval_results( + construct_db_schemas: list[dict], + dbschema_retrieval: list[Document], +) -> dict: + retrieval_results = [] + has_calculated_field = False + has_metric = False + has_json_field = False + + for table_schema in construct_db_schemas: + if table_schema["type"] == "TABLE": + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) + retrieval_results.append( + { + "table_name": table_schema["name"], + "table_ddl": ddl, + } + ) + if _has_calculated_field: + has_calculated_field = True + if _has_json_field: + has_json_field = True + + for document in dbschema_retrieval: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) + + return { + "retrieval_results": retrieval_results, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + + +def _merge_column_selection( + primary: dict[str, dict], + secondary: dict[str, dict], +) -> dict[str, dict]: + merged = { + table_name: { + **table_contents, + "columns": list(table_contents.get("columns", [])), + } + for table_name, table_contents in primary.items() + } + + for table_name, table_contents in secondary.items(): + if table_name not in merged: + merged[table_name] = { + **table_contents, + "columns": list(table_contents.get("columns", [])), + } + continue + + columns = list(merged[table_name].get("columns", [])) + for column in table_contents.get("columns", []): + if column not in columns: + columns.append(column) + merged[table_name]["columns"] = columns + + return merged + + def _lexical_columns_and_tables_needed( construct_db_schemas: list[dict], query: str | None, @@ -1503,6 +1932,26 @@ def _lexical_columns_and_tables_needed( comment_tokens = _tokenize_schema_text(column.get("comment")) score = len(query_tokens & column_tokens) * 10 score += len(query_tokens & comment_tokens) * 2 + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + score += len(column_tokens & _SALES_REVENUE_TOKENS) * 6 + if query_tokens & {"month", "monthly", "year", "july", "date", "latest"}: + score += ( + len(column_tokens & {"date", "day", "month", "year", "time"}) + * 5 + ) + if query_tokens & {"top", "highest", "lowest", "bottom"}: + measure_tokens = { + "amount", + "count", + "cost", + "margin", + "quantity", + "score", + "value", + } + if query_tokens & {"rate", "ratio", "percent", "percentage"}: + measure_tokens.update({"rate", "ratio", "percent", "percentage"}) + score += len(column_tokens & measure_tokens) * 4 if score > 0: column_scores.append( (score, column["name"], column.get("is_primary_key")) @@ -1511,6 +1960,13 @@ def _lexical_columns_and_tables_needed( if not column_scores and table_score <= 0: continue + if table_tokens & _GENERIC_TABLE_TOKENS and table_score <= 0: + table_score -= 8 + if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: + if not query_tokens & _CUSTOMS_FINANCE_TOKENS: + table_score -= len(table_tokens & _CUSTOMS_FINANCE_TOKENS) * 10 + table_score += len(table_tokens & _SALES_REVENUE_TOKENS) * 5 + total_score = table_score + sum(score for score, _, _ in column_scores) if total_score <= 0: continue diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index 9b8d9f26fc..57fb089dad 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -2,10 +2,9 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, - _SchemaCatalog, generate_simple_analytics_sql, - normalize_sql_with_schema_identifiers, normalize_wren_sql_dialect, + normalize_sql_with_schema_identifiers, unsupported_schema_generation_result, unsupported_schema_message, validate_sql_against_contexts, @@ -75,23 +74,10 @@ def test_schema_identifier_normalization_quotes_special_identifiers(): assert 'FROM "valid-order-lines"' in sql -def test_wren_sql_dialect_normalization_handles_top_and_joined_limit(): - assert ( - normalize_wren_sql_dialect("SELECT TOP 10 id1 FROM dbo_mbrTime") - == "SELECT id1 FROM dbo_mbrTime\nLIMIT 10" - ) - assert ( - normalize_wren_sql_dialect( - "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESCLIMIT 10" - ) - == "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESC LIMIT 10" - ) - - -def test_semantic_coverage_rejects_unrepresented_query_terms(): +def test_semantic_coverage_rejects_generic_table_for_business_concepts(): contexts = [ """ - CREATE TABLE neutral_records ( + CREATE TABLE dbo_mbrTime ( id1 INTEGER, id2 INTEGER ); @@ -100,141 +86,103 @@ def test_semantic_coverage_rejects_unrepresented_query_terms(): error = validate_sql_semantic_coverage( """ - SELECT id1, COUNT(*) AS record_count - FROM neutral_records + SELECT id1, COUNT(*) AS failures + FROM dbo_mbrTime GROUP BY id1 - ORDER BY record_count DESC + ORDER BY failures DESC LIMIT 10 """, - "Show top 10 records by missing_dimension.", + "Show the top 10 materials with the highest number of failures.", contexts, ) assert error is not None - assert "missing" in error or "dimension" in error + assert "failure/defect" in error + assert "material" in error -def test_unsupported_schema_message_reports_partial_coverage(): +def test_unsupported_schema_message_requires_all_requested_concepts(): contexts = [ """ - CREATE TABLE event_records ( - event_id VARCHAR, - phase VARCHAR + CREATE TABLE dbo_mbrTime ( + id1 INTEGER, + id2 INTEGER ); """ ] message = unsupported_schema_message( - "Show records by phase and unknown_segment.", + "Show the top 10 materials with the highest number of failures.", contexts, ) assert message is not None - assert "unknown" in message or "segment" in message + assert "No retrieved table or view" in message + assert "failure/defect" in message + assert "material" in message -def test_unsupported_schema_generation_result_has_no_invalid_sql(): +def test_unsupported_schema_message_rejects_split_failure_technician_without_coverage(): contexts = [ """ - CREATE TABLE event_records ( - event_id VARCHAR, - phase VARCHAR + CREATE TABLE dbo_report_failures ( + id INTEGER, + failure_type VARCHAR ); + """, """ - ] - - result = unsupported_schema_generation_result( - "Show records by unknown_segment.", - contexts, - data_source="MSSQL", - ) - - assert result is not None - assert result["valid_generation_result"] == {} - invalid = result["invalid_generation_result"] - assert invalid["type"] == "NO_RELEVANT_SQL" - assert invalid["sql"] == "" - assert invalid["original_sql"] == "" - assert "unknown" in invalid["error"] or "segment" in invalid["error"] - - -def test_schema_coverage_accepts_generic_word_form_variants(): - contexts = [ - """ - CREATE TABLE work_update_log ( - item_id VARCHAR, - updated_at TIMESTAMP + CREATE TABLE dbo_technicians ( + id INTEGER, + name VARCHAR ); - """ + """, ] - sql = """ - SELECT - CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT) AS year, - CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) AS month, - COUNT(*) AS record_count - FROM work_update_log - GROUP BY - CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT), - CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) - """ - - error = validate_sql_semantic_coverage( - sql, - "Show the number of work updates updated each month.", + message = unsupported_schema_message( + "Show the number of failures by technician.", contexts, ) - assert error is None - assert unsupported_schema_message( - "Show the number of work updates updated each month.", - contexts, - ) is None + assert message is not None + assert "failure/defect" in message + assert "technician" in message -def test_schema_fallback_uses_verified_monthly_update_timestamp(): +def test_unsupported_schema_generation_result_has_no_invalid_sql(): contexts = [ """ - CREATE TABLE work_update_log ( - item_id VARCHAR, - updated_at TIMESTAMP + CREATE TABLE dbo_report_failures ( + id INTEGER, + failure_type VARCHAR ); + """, """ + CREATE TABLE dbo_technicians ( + id INTEGER, + name VARCHAR + ); + """, ] - sql = generate_simple_analytics_sql( - "Show the number of work updates updated each month.", + result = unsupported_schema_generation_result( + "Show the number of failures by technician.", contexts, + data_source="MSSQL", ) - assert sql is not None - assert 'FROM "work_update_log"' in sql - assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT)' in sql - assert "COUNT(*)" in sql - - -def test_schema_catalog_ignores_extract_from_column_clause(): - contexts = [ - """ - CREATE TABLE work_update_log ( - item_id VARCHAR, - updated_at TIMESTAMP - ); - """ - ] - sql = """ - SELECT CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year", COUNT(*) AS "record_count" - FROM "work_update_log" - GROUP BY CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) - """ - - assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + assert result is not None + assert result["valid_generation_result"] == {} + invalid = result["invalid_generation_result"] + assert invalid["type"] == "NO_RELEVANT_SQL" + assert invalid["sql"] == "" + assert invalid["original_sql"] == "" + assert "technician" in invalid["error"] def test_post_processor_clears_sql_for_unsupported_schema(): contexts = [ """ - CREATE TABLE neutral_records ( + CREATE TABLE dbo_mbrTime ( id1 INTEGER, id2 INTEGER ); @@ -246,15 +194,15 @@ def test_post_processor_clears_sql_for_unsupported_schema(): post_processor.run( [ """ - SELECT id1, COUNT(*) AS record_count - FROM neutral_records + SELECT id1, COUNT(*) AS failures + FROM dbo_mbrTime GROUP BY id1 - ORDER BY record_count DESC + ORDER BY failures DESC LIMIT 10 """ ], contexts=contexts, - fallback_query="Show records by missing_dimension.", + fallback_query="Show the top 10 materials with the highest number of failures.", data_source="MSSQL", ) ) @@ -265,456 +213,368 @@ def test_post_processor_clears_sql_for_unsupported_schema(): assert result["invalid_generation_result"]["original_sql"] == "" -def test_schema_sample_value_filter_is_grounded_in_metadata(): - contexts = [ - """ - /* - WREN RETRIEVED SEMANTIC CONTEXT - {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done","In Progress"]}]} - WREN SQL IDENTIFIER CONTRACT - */ - CREATE TABLE work_items ( - item_id VARCHAR, - State VARCHAR, - updated_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql( - "Show all work item records with In Progress.", - contexts, +def test_wren_sql_dialect_normalization_repairs_top_and_joined_limit(): + assert ( + normalize_wren_sql_dialect("SELECT TOP 10 id1 FROM dbo_mbrTime") + == "SELECT id1 FROM dbo_mbrTime\nLIMIT 10" + ) + assert ( + normalize_wren_sql_dialect( + "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESCLIMIT 10" + ) + == "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESC LIMIT 10" ) - - assert sql is not None - assert 'FROM "work_items"' in sql - assert 'LOWER("State") = \'in progress\'' in sql -def test_user_values_are_allowed_for_single_verified_text_column(): +def test_repair_fallback_filters_critical_priority_and_in_progress_status(): contexts = [ """ - CREATE TABLE work_update_log ( - item_id VARCHAR, - state_name VARCHAR(255), - updated_at TIMESTAMP + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMPTZ ); """ ] - query = ( - "Show the distribution of work updates across completed and " - "in-progress state names." + sql = generate_simple_analytics_sql( + "Show all critical-priority repairs that are currently in progress.", + contexts, ) - sql = generate_simple_analytics_sql(query, contexts) - assert unsupported_schema_message(query, contexts) is None assert sql is not None - assert 'FROM "work_update_log"' in sql - assert 'LOWER("state_name") IN (\'completed\', \'in-progress\')' in sql - assert 'GROUP BY "state_name"' in sql + assert 'FROM "dbo_repair_logs"' in sql + assert "\"status\" = 'in progress'" in sql + assert "\"priority\" = 'critical'" in sql -def test_column_value_label_is_not_treated_as_literal_filter_value(): +def test_repair_fallback_preserves_hyphenated_in_progress_status_value(): contexts = [ """ - CREATE TABLE work_update_log ( - item_id VARCHAR, - state_name VARCHAR(255), - updated_at TIMESTAMP + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP ); """ ] - query = "Show the distribution of work updates across state name values." - sql = generate_simple_analytics_sql(query, contexts) + sql = generate_simple_analytics_sql( + "Show all repairs with a critical priority and an in-progress status.", + contexts, + ) - assert unsupported_schema_message(query, contexts) is None assert sql is not None - assert 'FROM "work_update_log"' in sql - assert 'GROUP BY "state_name"' in sql - assert "WHERE" not in sql + assert 'FROM "dbo_repair_logs"' in sql + assert "\"status\" = 'in-progress'" in sql + assert "\"priority\" = 'critical'" in sql -def test_unverified_filter_value_is_not_invented(): +def test_repair_logs_highest_priority_orders_by_verified_priority_column(): contexts = [ """ - /* - WREN RETRIEVED SEMANTIC CONTEXT - {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done"]}]} - WREN SQL IDENTIFIER CONTRACT - */ - CREATE TABLE work_items ( - item_id VARCHAR, - State VARCHAR + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP ); """ ] sql = generate_simple_analytics_sql( - "Show all work item records with Archived.", - contexts, - ) - message = unsupported_schema_message( - "Show all work item records with Archived.", + "Which repair logs have the highest priority?", contexts, ) - assert sql is None - assert message is not None - assert "archived" in message.lower() - - -def test_grouped_count_uses_verified_dimension_only(): - contexts = [ - """ - CREATE TABLE event_records ( - event_id VARCHAR, - phase VARCHAR, - updated_at TIMESTAMP - ); - """ - ] - - sql = generate_simple_analytics_sql("Show records by phase.", contexts) - assert sql is not None - assert 'SELECT "phase", COUNT(*) AS "record_count"' in sql - assert 'FROM "event_records"' in sql - assert 'GROUP BY "phase"' in sql + assert 'FROM "dbo_repair_logs"' in sql + assert 'ORDER BY CASE LOWER("priority")' in sql + assert "DESC" in sql -def test_average_uses_verified_numeric_measure_not_count(): +def test_critical_priority_repairs_filter_verified_priority_column(): contexts = [ """ - CREATE TABLE measurement_records ( - entity_id VARCHAR, - model_code VARCHAR, - age_days DECIMAL + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP ); """ ] - sql = generate_simple_analytics_sql("Show average age by model.", contexts) + sql = generate_simple_analytics_sql( + "Show all critical-priority repairs", + contexts, + ) assert sql is not None - assert 'SELECT "model_code", AVG("age_days") AS "average_value"' in sql - assert 'GROUP BY "model_code"' in sql - assert "COUNT(" not in sql + assert 'FROM "dbo_repair_logs"' in sql + assert "\"priority\" = 'critical'" in sql -def test_average_without_verified_measure_is_unsupported(): +def test_repairs_by_status_counts_verified_repair_rows(): contexts = [ """ - CREATE TABLE measurement_records ( - entity_id VARCHAR, - model_code VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql("Show average age by model.", contexts) - message = unsupported_schema_message("Show average age by model.", contexts) - - assert sql is None - assert message is not None - assert "age" in message.lower() - - -def test_latest_uses_verified_temporal_column(): - contexts = [ - """ - CREATE TABLE event_records ( - event_id VARCHAR, - event_time TIMESTAMP, - phase VARCHAR + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP ); """ ] - sql = generate_simple_analytics_sql("Show latest event records.", contexts) + sql = generate_simple_analytics_sql( + "Show repairs by status", + contexts, + ) assert sql is not None - assert 'FROM "event_records"' in sql - assert 'ORDER BY "event_time" DESC' in sql + assert 'SELECT "status", COUNT("id") AS "record_count"' in sql + assert 'FROM "dbo_repair_logs"' in sql + assert 'GROUP BY "status"' in sql -def test_monthly_count_uses_requested_temporal_column_when_verified(): +def test_latest_repair_logs_orders_by_verified_date_column(): contexts = [ """ - CREATE TABLE event_records ( - event_id VARCHAR, - updated_at TIMESTAMP, + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, created_at TIMESTAMP ); """ ] sql = generate_simple_analytics_sql( - "Show the number of event records updated each month.", + "Show latest repair logs", contexts, ) assert sql is not None - assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year"' in sql - assert 'CAST(EXTRACT(MONTH FROM "updated_at") AS BIGINT) AS "month"' in sql - assert 'COUNT(*) AS "record_count"' in sql + assert 'FROM "dbo_repair_logs"' in sql + assert 'ORDER BY "created_at" DESC' in sql -def test_order_by_uses_verified_column_and_sample_value(): +def test_semantic_column_alias_can_satisfy_priority_concept_with_verified_name(): contexts = [ """ /* WREN RETRIEVED SEMANTIC CONTEXT - {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open","Closed"]}]} + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"repair log records"},"columns":[{"sql_column_name_use_exactly":"Urgency","data_type":"VARCHAR","semantic_context_not_sql_identifier":"priority severity for a repair"}]} WREN SQL IDENTIFIER CONTRACT */ - CREATE TABLE case_records ( - case_id VARCHAR, - State VARCHAR, - updated_at TIMESTAMP + CREATE TABLE dbo_work_items ( + id VARCHAR, + Urgency VARCHAR, + created_at TIMESTAMP ); """ ] sql = generate_simple_analytics_sql( - "Show all case records with Open ordered by case ID.", + "Which repair records have the highest priority?", contexts, ) assert sql is not None - assert 'LOWER("State") = \'open\'' in sql - assert 'ORDER BY "case_id" ASC' in sql + assert 'FROM "dbo_work_items"' in sql + assert '"Urgency"' in sql + assert '"priority"' not in sql -def test_top_grouped_count_is_schema_shape_based(): +def test_failure_by_technician_fallback_uses_verified_tech_column(): contexts = [ """ - CREATE TABLE occurrence_records ( - occurrence_id VARCHAR, - model_code VARCHAR, - reason_code VARCHAR + CREATE TABLE dbo_DebugEntries_Staging2 ( + Tech VARCHAR, + Failed VARCHAR, + Material VARCHAR ); """ ] sql = generate_simple_analytics_sql( - "Show top 5 occurrence records by model.", + "Show the number of failures by technician.", contexts, ) assert sql is not None - assert 'SELECT "model_code", COUNT(*) AS "record_count"' in sql - assert 'ORDER BY "record_count" DESC' in sql - assert "LIMIT 5" in sql + assert 'FROM "dbo_DebugEntries_Staging2"' in sql + assert 'SELECT "Tech", COUNT("Failed") AS "record_count"' in sql + assert 'WHERE ("Failed" IS NOT NULL AND "Failed" <> \'\')' in sql -def test_single_grouping_dimension_does_not_over_split_results(): +def test_failure_by_material_fallback_uses_verified_material_column(): contexts = [ """ - CREATE TABLE account_events ( - event_id VARCHAR, - account_name VARCHAR, - account_reference VARCHAR - ); - """ - ] - - sql = generate_simple_analytics_sql("Show number of events by account.", contexts) - - assert sql is not None - assert 'GROUP BY "account_name"' in sql - assert "account_reference" not in sql - - -def test_missing_value_intent_uses_verified_plural_name_column(): - contexts = [ - """ - CREATE TABLE account_events ( - event_id VARCHAR, - account_name VARCHAR, - event_time TIMESTAMP + CREATE TABLE dbo_DebugEntries_Staging2 ( + Tech VARCHAR, + Failed VARCHAR, + Material VARCHAR ); """ ] sql = generate_simple_analytics_sql( - "Show events with missing account names.", + "Show failures by material.", contexts, ) assert sql is not None - assert 'FROM "account_events"' in sql - assert '"account_name" IS NULL' in sql + assert 'FROM "dbo_DebugEntries_Staging2"' in sql + assert 'SELECT "Material", COUNT("Failed") AS "record_count"' in sql -def test_semantic_validation_rejects_weaker_null_check_column(): +def test_failure_type_value_filter_uses_verified_failure_type_column(): contexts = [ """ - CREATE TABLE account_events ( - event_id VARCHAR, - account_name VARCHAR, - account_reference VARCHAR + CREATE TABLE dbo_DebugEntries ( + SerialNumber VARCHAR, + FailedAt VARCHAR, + Material VARCHAR ); - """ - ] - - error = validate_sql_semantic_coverage( - """ - SELECT account_name, account_reference - FROM account_events - WHERE account_reference IS NULL """, - "Show events with missing account names.", - contexts, - ) - - assert error is not None - assert "weaker matching column" in error - - -def test_top_records_are_listed_without_implicit_grouped_aggregate(): - contexts = [ """ - CREATE TABLE scored_events ( - event_id VARCHAR, - score_value DECIMAL, - event_date TIMESTAMP, - category_name VARCHAR + CREATE TABLE dbo_repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR ); + """, """ - ] - - sql = generate_simple_analytics_sql("Show top 10 scored events from July.", contexts) - - assert sql is not None - assert 'FROM "scored_events"' in sql - assert "GROUP BY" not in sql - assert 'ORDER BY "score_value" DESC' in sql - assert "LIMIT 10" in sql - - -def test_sum_by_year_uses_verified_measure_and_temporal_column(): - contexts = [ - """ - CREATE TABLE transaction_records ( - transaction_id VARCHAR, - account_name VARCHAR, - amount_value DECIMAL, - posted_at TIMESTAMP + CREATE TABLE dbo_report_failures ( + failure_type VARCHAR, + failure_line VARCHAR, + test_name VARCHAR ); - """ + """, ] - sql = generate_simple_analytics_sql("Show total amount by year.", contexts) + sql = generate_simple_analytics_sql( + "Show the number of units with JTAG as the failure type.", + contexts, + ) assert sql is not None - assert 'CAST(EXTRACT(YEAR FROM "posted_at") AS BIGINT) AS "year"' in sql - assert 'SUM("amount_value") AS "total_value"' in sql + assert 'FROM "dbo_report_failures"' in sql + assert 'COUNT(*) AS "record_count"' in sql + assert "\"failure_type\" = 'JTAG'" in sql -def test_semantic_coverage_rejects_count_for_average_intent(): +def test_board_models_most_failures_counts_failure_records_not_defect_rate(): contexts = [ """ - CREATE TABLE measurement_records ( - entity_id VARCHAR, - model_code VARCHAR, - age_days DECIMAL + CREATE TABLE dbo_batch_records ( + board_model VARCHAR, + supplier VARCHAR, + defect_rate DECIMAL ); + """, """ + CREATE TABLE dbo_repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR + ); + """, ] - error = validate_sql_semantic_coverage( - """ - SELECT model_code, COUNT(*) AS record_count - FROM measurement_records - GROUP BY model_code - """, - "Show average age by model.", + sql = generate_simple_analytics_sql( + "Show the top 5 board models with the most failures.", contexts, ) - assert error is not None - assert "average" in error.lower() + assert sql is not None + assert 'FROM "dbo_repair_logs"' in sql + assert 'SELECT "board_model", COUNT("failure_code") AS "record_count"' in sql + assert '"defect_rate"' not in sql + assert "LIMIT 5" in sql -def test_literal_validation_rejects_values_outside_verified_samples(): +def test_board_models_highest_defect_rate_uses_rate_metric(): contexts = [ """ - /* - WREN RETRIEVED SEMANTIC CONTEXT - {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open"]}]} - WREN SQL IDENTIFIER CONTRACT - */ - CREATE TABLE case_records ( - case_id VARCHAR, - State VARCHAR + CREATE TABLE dbo_batch_records ( + board_model VARCHAR, + supplier VARCHAR, + defect_rate DECIMAL ); + """, """ + CREATE TABLE dbo_repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR + ); + """, ] - error = validate_sql_semantic_coverage( - """ - SELECT case_id - FROM case_records - WHERE LOWER(State) = 'Closed' - """, - "Show case records with Open.", + sql = generate_simple_analytics_sql( + "Show the board models with the highest defect rate.", contexts, ) - assert error is not None - assert "sample values" in error + assert sql is not None + assert 'FROM "dbo_batch_records"' in sql + assert 'SELECT "board_model", AVG("defect_rate") AS "average_value"' in sql + assert 'ORDER BY "average_value" DESC' in sql -def test_semantic_validation_rejects_multi_group_for_single_dimension(): +def test_semantic_coverage_rejects_rate_for_failure_count_intent(): contexts = [ """ - CREATE TABLE account_events ( - event_id VARCHAR, - account_name VARCHAR, - account_reference VARCHAR + CREATE TABLE dbo_batch_records ( + board_model VARCHAR, + defect_rate DECIMAL ); """ ] error = validate_sql_semantic_coverage( """ - SELECT account_name, account_reference, COUNT(*) AS record_count - FROM account_events - GROUP BY account_name, account_reference + SELECT board_model, defect_rate + FROM dbo_batch_records + ORDER BY defect_rate DESC + LIMIT 5 """, - "Show number of events by account.", + "Show the top 5 board models with the most failures.", contexts, ) assert error is not None - assert "one grouping dimension" in error + assert "count of failure records" in error -def test_semantic_validation_rejects_top_record_grouped_aggregate(): +def test_repairs_by_technician_requires_one_schema_object_covering_both_concepts(): contexts = [ """ - CREATE TABLE scored_events ( - event_id VARCHAR, - score_value DECIMAL, - event_date TIMESTAMP, - category_name VARCHAR + CREATE TABLE dbo_repair_logs ( + id VARCHAR, + status VARCHAR, + priority VARCHAR, + failure_code VARCHAR ); + """, """ + CREATE TABLE dbo_DebugEntries_Staging2 ( + Tech VARCHAR, + Failed VARCHAR + ); + """, ] - error = validate_sql_semantic_coverage( - """ - SELECT category_name, SUM(score_value) AS total_value - FROM scored_events - GROUP BY category_name - ORDER BY total_value DESC - LIMIT 10 - """, - "Show top 10 scored events from July.", - contexts, - ) + message = unsupported_schema_message("Show repairs by technician.", contexts) - assert error is not None - assert "grouped aggregate" in error + assert message is not None + assert "repair" in message + assert "technician" in message diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index c847e1e0df..1bf9630b62 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1137,23 +1137,25 @@ def test_column_selection_returns_empty_dict_for_malformed_reply(): assert parsed == {} -def test_retrieval_query_augmentation_is_schema_neutral(): - query = "show total amount by year" +def test_retrieval_query_augmentation_adds_business_terms(): + augmented = _augment_retrieval_query("show total revenue by year") - assert _augment_retrieval_query(query) == query + assert "Business schema search terms" in augmented + assert "sales revenue amount value" in augmented + assert "date month year" in augmented -def test_table_ranking_prefers_direct_schema_metadata_overlap(): +def test_table_ranking_prefers_business_sales_table_over_generic_or_customs_tables(): documents = [ - _schema_document("neutral_records", ["id1", "id2"]), - _schema_document("amount_snapshots", ["amount_value", "snapshot_year"]), - _schema_document("event_records", ["event_code", "event_time"]), + _schema_document("dbo_mbrTime", ["id1", "id2"]), + _schema_document("CustomsRefundClaim", ["DutyAmount", "ClaimDate"]), + _schema_document("SalesOrderFact", ["USDFXSalesValue", "OrderDate"]), ] ranked = _rank_table_names_by_query( - ["neutral_records", "event_records", "amount_snapshots"], + ["dbo_mbrTime", "CustomsRefundClaim", "SalesOrderFact"], documents, - "show total amount by year", + "show total revenue by year", ) - assert ranked[0] == "amount_snapshots" + assert ranked[0] == "SalesOrderFact" From a3df06dfb72074eec1b583bab4583f224f9bddbe Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Tue, 1 Sep 2026 11:42:12 +0000 Subject: [PATCH 1081/1087] Optimize schema-grounded Ask pipeline --- WRENAI_LOCAL_ASK_HANDOFF.md | 431 +- .../pipelines/generation/data_assistance.py | 1 + .../generation/followup_sql_generation.py | 69 +- .../followup_sql_generation_reasoning.py | 4 +- .../generation/intent_classification.py | 7 +- .../src/pipelines/generation/sql_answer.py | 36 +- .../pipelines/generation/sql_correction.py | 11 +- .../pipelines/generation/sql_generation.py | 71 +- .../generation/sql_generation_reasoning.py | 4 +- .../src/pipelines/generation/utils/sql.py | 5702 ++++++++++++----- .../src/pipelines/indexing/db_schema.py | 6 + .../src/pipelines/indexing/utils/helper.py | 1 + .../retrieval/db_schema_retrieval.py | 1352 ++-- wren-ai-service/src/web/v1/routers/ask.py | 10 + wren-ai-service/src/web/v1/services/ask.py | 520 +- .../src/web/v1/services/sql_answer.py | 16 + .../test_prompt_grounding_contracts.py | 56 + .../generation/test_sql_answer_prompt.py | 17 + .../generation/test_sql_schema_grounding.py | 1855 +++++- .../pipelines/indexing/test_db_schema.py | 86 +- .../retrieval/test_db_schema_retrieval.py | 454 +- .../tests/pytest/services/test_ask.py | 36 + wren-ui/next.config.js | 2 + .../apollo/server/adaptors/wrenAIAdaptor.ts | 44 + .../textBasedAnswerBackgroundTracker.ts | 24 + .../apollo/server/resolvers/askingResolver.ts | 6 + .../apollo/server/resolvers/modelResolver.ts | 9 +- .../apollo/server/services/askingService.ts | 36 +- .../server/services/askingTaskTracker.ts | 28 +- .../apollo/server/services/queryService.ts | 115 +- .../services/tests/queryService.test.ts | 84 + wren-ui/src/apollo/server/utils/manifest.ts | 4 +- 32 files changed, 8478 insertions(+), 2619 deletions(-) create mode 100644 wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index e9f23efcda..190063e85d 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -1,283 +1,236 @@ -# WrenAI Ask Grounding Handoff +# WrenAI Ask Schema Grounding Handoff + +Date: 2026-09-01 + +## Current Branch + +- Workspace: `D:\WrenAI` +- Branch: `organization/ask-schema-grounding-20260820` +- Local HEAD before this handoff commit: `0ff1e6e23 Improve Ask schema grounding` +- Remote tracking state at handoff time: local branch was ahead 1 and behind 82 +- Do not use `.codex-tmp` as runtime source. The AI service was restarted from `D:\WrenAI\wren-ai-service`. + +## Goal Continued + +Continue the WrenAI Ask schema-grounding work for CWPay and CW_GL while preserving Orders and PCB_DB behavior. The focus of this continuation was Ask speed and correctness on large schemas: + +- add timing logs for Ask stages +- identify slow stages +- reduce unnecessary LLM calls +- cache schema-derived metadata and ranking inputs +- keep SQL generation schema-verified +- return clear unsupported results instead of hallucinated SQL +- validate CWPay, CW_GL, Orders, and PCB_DB + +## What Was Done + +### Timing and Observability + +Ask timing logs now cover the key path across UI and AI service: + +- frontend request +- task creation +- schema retrieval support context +- schema retrieval +- candidate ranking +- LLM intent generation +- SQL generation / deterministic fast path +- SQL validation +- SQL execution +- answer formatting request and polling +- cancel request and cancellation point + +Relevant logs reviewed: + +- `.codex-tmp\ai-dev-after13.err.log` +- `.codex-tmp\ai-dev-after14.err.log` +- `.codex-tmp\ai-dev-after15.err.log` +- `.codex-tmp\ai-dev-after17.err.log` +- `.codex-tmp\ai-dev-after19.err.log` +- `.codex-tmp\ai-dev-after20.err.log` + +Final stage summary from `.codex-tmp\ai-dev-after20.err.log`: + +- `sql_generation_fast_path`: avg about 1.4s, max about 2.8s +- `schema_retrieval`: avg about 0.5s, max about 1.5s +- `schema_retrieval_support_context`: avg about 19ms +- task creation and frontend markers were effectively negligible +- no LLM intent generation was used by the final Ask validation tasks except the explicit cancel test + +There is one non-Ask background outlier in the same AI log: `schema_retrieval_total` around 142s from a column-pruning/question-recommendation path. It was not part of the final Ask task timing set. + +### Performance Improvements + +- Added deterministic schema-driven fast paths for clear count, grouping, top-N, latest, date/month bucket, distribution, listing, and same-thread group-result follow-up shapes. +- Added pre-intent unsupported handling for simple analytics requests when active-project schema coverage is missing. +- Added schema metadata/token/index caching for table/column/description-derived matching. +- Changed large-schema retrieval to broad candidate gathering and small top-K reranking. +- Changed generation-context limiting to skip oversized candidates and continue looking for smaller valid candidates within the token budget. +- Tightened same-thread follow-up grounding to use compact latest verified SQL/table identifiers instead of accumulating stale or oversized history. +- Preserved schema validation and dry-run validation. Unsupported cases return `NO_RELEVANT_SQL` instead of invalid SQL. +- Fixed SQL literal offset handling so extracted filter values are validated against the right columns. +- Fixed top record/listing behavior so row-level date questions order by verified date columns instead of unrelated numeric fields. +- Tightened answer formatting prompts to use executed SQL result rows only and not invent analysis, values, code, or examples. +- Added transient MSSQL deadlock retry around Ibis query execution. + +### Count Shape Fix + +The CWPay question `How many invoice records are there?` now returns a scalar aggregate: + +```sql +SELECT + COUNT(*) AS "record_count" +FROM + "dbo_View_Open_Invoices" +``` -Date: 2026-08-20 +The result shape is one column, `record_count`, and one row. It no longer returns invoice detail columns for that simple count shape. -## Current Goal +### Same-Thread Follow-Ups -Make WrenAI's Ask pipeline schema-first for the active org/project. Natural language should search and rank verified schema metadata, but final SQL must use only tables, views, columns, relationships, metrics, and values supported by the selected project's metadata. +Follow-up questions now retrieve exact prior verified tables from the latest SQL and use a compact grounding query. This fixed the slow same-thread path that previously fell back to LLM calls on large schemas. -Do not fix future issues by hardcoding one question, project, table, column, or organization. Representative prompts such as `Which repair logs have the highest priority?` are regression examples only. +Final targeted follow-up waits: -## Final Runtime State +- CWPay: about 3.0s +- CW_GL: about 3.0s +- Orders: about 3.1s +- PCB_DB: about 2.0s -- UI: `http://127.0.0.1:3000` -- AI service: `http://127.0.0.1:5555` -- AI health: `{"status":"ok"}` -- Active project restored after validation: `org / PCB_DB` -- Active project id: `10` -- Orders project id: `11` -- Sales duplicate: not shown in current project list; `Orders` remains canonical. +## Validation Results -Current projects visible through `/api/v1/projects/current`: +Final artifacts: -- id `4`, unnamed DuckDB -- id `10`, `PCB_DB` -- id `11`, `Orders` -- id `12`, `CWPay` -- id `13`, `CW_GL` +- `.codex-tmp\ask_perf_benchmark_after_final.json` +- `.codex-tmp\resume_schema_grounding_validation_after_final.json` +- `.codex-tmp\cancel_check_after20.json` -## What Changed Today +Before/after benchmark: -### Generic Schema Grounding +- Before avg Ask wait: 61,818.6ms +- Before max Ask wait: 143,366ms +- After avg Ask wait: 3,388ms +- After max Ask wait: 5,056ms -Permanent source changes are now in `D:\WrenAI\wren-ai-service`, not only `.codex-tmp`. +Observed CWPay examples: -Main file: +- `How many invoice records are there?` + - before: 84,653ms + - after: 3,031ms + - final SQL uses scalar `COUNT(*) AS "record_count"` +- `Show invoices by business unit` + - before: 112,881ms + - after: 3,036ms + - final SQL groups by verified `bunit` -- `wren-ai-service/src/pipelines/generation/utils/sql.py` +Final full validation: -Added or improved: - -- SQL identifier validation against retrieved schema. -- Semantic coverage validation so valid identifiers are not enough; the referenced table/view must also support the requested business concepts. -- Unsupported-schema result helper that returns `NO_RELEVANT_SQL` with no invented SQL. -- Deterministic schema-grounded fallback for common Ask families: - - count / grouped counts - - top-N - - highest / lowest - - latest / recent - - priority / severity - - status filters - - date/month/year filters - - revenue/sales measures - - failure counts vs defect-rate metrics - - failure type value filters -- Semantic alias support from Wren retrieved context blocks. -- More timestamp type support, including `TIMESTAMPTZ`, which fixed the live `latest repair logs` failure. -- Normalization of dialect issues such as `TOP n`, joined `DESCLIMIT`, and order-by aliases. -- Logs for generated SQL validation, deterministic fallback SQL, fallback validation, selected table, verified columns, and metric intent. - -### Retrieval Improvements - -Main file: +- Project checks: 4/4 +- Ask cases: 18/18 +- Same-thread follow-ups: 4/4 -- `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` +Project checks: -Added or improved: +- CWPay: 364 datasource tables, 350 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` +- CW_GL: 223 datasource tables, 223 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` +- Orders: 103 datasource tables, 101 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` +- PCB_DB: 76 datasource tables, 68 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` -- Project-scoped retrieval filters retained and tested. -- Query expansion for business concepts such as repair, failure, revenue, order, material, status, priority, latest, and date. -- Ranking uses table names, column names, descriptions/comments, semantic context, and generic-table deboosting. -- Logs for: - - selected project id - - retrieved candidate tables and scores - - selected schema objects and columns +Regression status: -### Generation Pipeline Wiring +- Orders by customer: passed +- Orders revenue/date/month/latest families: passed +- PCB_DB repairs/status/priority/month/latest families: passed +- Unsupported cross-project questions: passed with `NO_RELEVANT_SQL` +- No observed schema leakage between projects +- No observed `Failed to create asking task` +- No observed hallucinated tables or columns in final validation -Files: +Cancel validation: -- `wren-ai-service/src/pipelines/generation/sql_generation.py` -- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` -- `wren-ai-service/src/pipelines/generation/sql_correction.py` +- CWPay cancel task: `7a2e1247-2a73-4a52-9fdd-d41004afc3c7` +- cancel mutation returned `true` +- final status: `STOPPED` +- elapsed to terminal status: 520ms -Changes: +## Tests Run -- Passed the user query into post-processing as `fallback_query`. -- Added pre-LLM unsupported-schema checks where retrieved schema clearly cannot cover requested concepts. -- Ensured SQL correction still uses the same schema-first validation and fallback logic. -- Strengthened correction instructions so invalid or hallucinated identifiers are not preserved. - -### UI / Project Cleanup From This Workstream - -Files still dirty from the related UI/runtime fixes: - -- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` -- `wren-ui/src/apollo/server/services/askingService.ts` - -Relevant behavior: - -- Previous `results` crash handling is preserved. -- Unsupported-schema failures now avoid showing invented SQL as something to fix. -- Sales/Orders cleanup remains in place: UI project list shows `Orders`, not duplicate `Sales`. - -## Live Validation Done - -All live checks were run through the UI GraphQL Ask path after restarting the AI service. - -### PCB_DB - -Active project: `PCB_DB`, id `10`. - -Passed: - -- `Which repair logs have the highest priority?` - - Table: `dbo_repair_logs` - - Uses verified `priority` - - Orders by generic priority ranking expression -- `Show all critical-priority repairs` - - Table: `dbo_repair_logs` - - Filter: `priority = 'critical'` -- `Show repairs by status` - - Table: `dbo_repair_logs` - - Group: `status` - - Metric: `COUNT(id)` -- `Show latest repair logs` - - Table: `dbo_repair_logs` - - Order: `created_at DESC` - - This was the live regression fixed by adding timestamp type coverage. -- `Show the number of failures by material` - - Uses verified material/failure fields from PCB_DB. -- `Show top 5 board models with the most failures` - - Table: `dbo_repair_logs` - - Metric: `COUNT(failure_code)` - - Did not use `defect_rate`. -- `Show units with JTAG as the failure type` - - Table: `dbo_report_failures` - - Filter: `failure_type = 'JTAG'` -- Extra check: - - `Show all repairs with a critical priority and an in-progress status.` - - Table: `dbo_repair_logs` - - Filters: `status = 'in-progress'` and `priority = 'critical'` - -### Orders - -Temporarily switched active project to `Orders`, id `11`, then restored PCB_DB. - -Passed: - -- `Show top 10 orders from July` - - Uses Orders table/date fields. -- `Show number of orders by customer` - - Groups by customer. - - Counts distinct order numbers. -- `Show revenue by year` - - Uses verified sales/revenue value and invoice date fields. -- Unsupported check: `Which repair logs have the highest priority?` - - Returned `NO_RELEVANT_SQL`. - - No SQL candidate. - - Message clearly said the active project does not contain verified `repair` and `priority/severity` fields. - -## Checks Run - -Passed: - -```powershell -git diff --check -- wren-ai-service/src/pipelines/generation/utils/sql.py ` - wren-ai-service/src/pipelines/generation/sql_generation.py ` - wren-ai-service/src/pipelines/generation/followup_sql_generation.py ` - wren-ai-service/src/pipelines/generation/sql_correction.py ` - wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py ` - wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py ` - wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py -``` - -Passed: +From `D:\WrenAI\wren-ai-service`: ```powershell -cd D:\WrenAI\wren-ai-service -.\venv\Scripts\python.exe -m compileall -q src\pipelines\generation src\pipelines\retrieval ` - tests\pytest\pipelines\generation\test_sql_schema_grounding.py ` - tests\pytest\pipelines\retrieval\test_db_schema_retrieval.py +.\venv\Scripts\python.exe -m pytest tests/pytest/services/test_ask.py tests/pytest/pipelines/generation/test_sql_schema_grounding.py tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py -q ``` -Could not run pytest in the service venv: +Result: 114 passed. -```text -D:\WrenAI\wren-ai-service\venv\Scripts\python.exe: No module named pytest -``` - -## Tests Added - -Main test file: - -- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` - -Coverage added for: - -- Unsupported schema clears invalid SQL. -- Generic table rejection for unsupported business concepts. -- Repair priority ordering. -- Critical-priority repair filters. -- Repairs by status. -- Latest repair logs with `TIMESTAMPTZ`. -- Semantic alias column support, for example using real verified `Urgency` when semantic context says it means priority/severity. -- Failure by material / technician with verified columns. -- JTAG failure type filters. -- Board models with most failures uses count, not defect rate. -- Highest defect rate uses rate metric. -- Repairs by technician requires one schema object or relationship coverage. - -Retrieval test file: - -- `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` - -Coverage added for: - -- Project filter conditions. -- Query expansion. -- Table ranking by query and schema text. -- Project-scoped schema retrieval behavior. - -## Restart Commands Used - -Restart AI service only: - -```powershell -$taskName = 'WrenAI 04 AI Service' -$listenerProcessIds = Get-NetTCPConnection -LocalPort 5555 -State Listen -ErrorAction SilentlyContinue | - Select-Object -ExpandProperty OwningProcess -Unique -Stop-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue -foreach ($listenerProcessId in $listenerProcessIds) { - if ($listenerProcessId) { - Stop-Process -Id $listenerProcessId -Force -ErrorAction SilentlyContinue - } -} -Start-ScheduledTask -TaskName $taskName -``` - -Health check: +Additional focused Ask service test: ```powershell -Invoke-WebRequest -UseBasicParsing http://127.0.0.1:5555/health +.\venv\Scripts\python.exe -m pytest tests/pytest/services/test_ask.py -q ``` -Project switch endpoints used for validation: +Result: 4 passed. -```powershell -Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/11/select -Invoke-WebRequest -UseBasicParsing -Method POST http://127.0.0.1:3000/api/v1/projects/10/select -Invoke-WebRequest -UseBasicParsing http://127.0.0.1:3000/api/v1/projects/current -``` +Warnings were pre-existing Pydantic deprecation warnings and existing coroutine cleanup warnings in semantics-preparation tests. -## Current Dirty Files To Review +## Files To Include In Handoff Commit -Relevant tracked files: +Include the Ask/UI source and focused tests: +- `wren-ai-service/src/pipelines/generation/data_assistance.py` - `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` +- `wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py` +- `wren-ai-service/src/pipelines/generation/intent_classification.py` +- `wren-ai-service/src/pipelines/generation/sql_answer.py` - `wren-ai-service/src/pipelines/generation/sql_correction.py` - `wren-ai-service/src/pipelines/generation/sql_generation.py` +- `wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py` - `wren-ai-service/src/pipelines/generation/utils/sql.py` +- `wren-ai-service/src/pipelines/indexing/db_schema.py` +- `wren-ai-service/src/pipelines/indexing/utils/helper.py` - `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` +- `wren-ai-service/src/web/v1/routers/ask.py` +- `wren-ai-service/src/web/v1/services/ask.py` +- `wren-ai-service/src/web/v1/services/ask_feedback.py` +- `wren-ai-service/src/web/v1/services/sql_answer.py` +- `wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py` +- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py` - `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` +- `wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py` - `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` +- `wren-ai-service/tests/pytest/services/test_ask.py` +- `wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts` +- `wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts` +- `wren-ui/next.config.js` +- `wren-ui/src/apollo/server/resolvers/askingResolver.ts` - `wren-ui/src/apollo/server/resolvers/modelResolver.ts` - `wren-ui/src/apollo/server/services/askingService.ts` - -There are also many local untracked runtime/data artifacts in the repository. Do not clean or delete them casually. - -## Important Caveats - -- Runtime source code should remain generic. Do not add checks for exact prompts such as `Which repair logs have the highest priority?`. -- Tests may use representative table and prompt names; production code must not. -- Retrieval context currently uses metadata/descriptions and some semantic context. It does not appear to carry robust sample-value lists. Status casing/value handling works for tested prompts, but richer value-aware matching would improve future accuracy. -- `enable_column_pruning` was not the focus of today's final validation. -- Full pytest suite still needs an environment with `pytest` installed. - -## Recommended Next Steps - -1. Install or enable pytest in `wren-ai-service\venv`, then run focused tests. -2. Review the large `utils/sql.py` diff carefully; consider extracting fallback/grounding helpers into smaller modules after behavior is stable. -3. Add sample-value metadata to retrieval context if available, then make value matching use that metadata instead of only text normalization. -4. Run a broader live Ask regression across PCB_DB, Orders, CWPay, and CW_GL when their data sources are available. -5. Commit the source changes after review, excluding local runtime/data artifacts. +- `wren-ui/src/apollo/server/services/askingTaskTracker.ts` +- `wren-ui/src/apollo/server/services/queryService.ts` +- `wren-ui/src/apollo/server/services/tests/queryService.test.ts` +- `wren-ui/src/apollo/server/utils/manifest.ts` +- `WRENAI_LOCAL_ASK_HANDOFF.md` + +Do not include: + +- `.codex-tmp` +- local configs +- logs +- venv folders +- extracted datasource dumps +- Qdrant/storage runtime data +- `wren-engine` submodule pointer +- `wren-ui/.yarn/releases/yarn-4.5.3.cjs` mode-only churn +- `wren-ui/package-lock.json` unless intentionally changing package management + +## Remaining Blockers + +- Modeling AI Assistant generate semantics/relationships for CW_GL remains unresolved. Earlier evidence showed semantics omitted the selected model and relationships timed out. Final Ask performance validation skipped assistant generation checks. +- Branch is behind remote by 82 commits. Push may require integration/rebase by whoever owns the branch if GitHub rejects a non-fast-forward push. + +## Guardrails Preserved + +- No app logic hardcodes datasource, project, organization, table, column, filter-value, or prompt-specific mappings. +- SQL is generated only from verified retrieved schema and then validated. +- Answer formatting is grounded in executed SQL results only. +- Unsupported or weakly covered questions fail quickly with a clear unsupported/clarification result. diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index 51b91197f9..e10ed4a530 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -29,6 +29,7 @@ - There should be proper line breaks, whitespace, and Markdown formatting(headers, lists, tables, etc.) in your response. - If the language is Traditional/Simplified Chinese, Korean, or Japanese, the maximum response length is 150 words; otherwise, the maximum response length is 110 words. - MUST NOT add SQL code in your response. +- Use only the provided DATABASE SCHEMA as context. Do not invent, assume, or name tables or columns that are not present there; do not provide hypothetical schema. - If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. ### OUTPUT FORMAT ### 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 8a10b3769d..3e7f214f09 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -16,7 +16,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, - construct_schema_identifier_catalog, + generate_simple_analytics_sql, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -152,6 +152,7 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, query: str | None = None, + grounding_query: str | None = None, documents: list[str] | None = None, project_id: str | None = None, mdl_hash: str | None = None, @@ -164,7 +165,7 @@ async def post_process( project_id=project_id, mdl_hash=mdl_hash, contexts=documents, - fallback_query=query, + fallback_query=grounding_query or query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -220,7 +221,7 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, - validation_contexts: list[str] | None = None, + grounding_query: str | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") @@ -232,8 +233,9 @@ async def run( metadata = {} data_source = metadata.get("data_source", "local_file") + effective_grounding_query = grounding_query or query unsupported_result = unsupported_schema_generation_result( - query, + effective_grounding_query, contexts=contexts, data_source=data_source, ) @@ -247,6 +249,7 @@ async def run( ["post_process"], inputs={ "query": query, + "grounding_query": effective_grounding_query, "documents": contexts, "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, @@ -266,3 +269,61 @@ async def run( **self._components, }, ) + + async def run_deterministic_fast_path( + self, + query: str, + contexts: list[str], + project_id: str | None = None, + mdl_hash: str | None = None, + use_dry_plan: bool = False, + allow_dry_plan_fallback: bool = True, + grounding_query: str | None = None, + ) -> dict | None: + if use_dry_plan: + metadata = await retrieve_metadata( + project_id or "", self._retriever, mdl_hash + ) + else: + metadata = {} + data_source = metadata.get("data_source", "local_file") + effective_grounding_query = grounding_query or query + + unsupported_result = unsupported_schema_generation_result( + effective_grounding_query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "Follow-up SQL deterministic fast path returned unsupported schema before LLM." + ) + return {"post_process": unsupported_result, "fast_path": "unsupported"} + + deterministic_sql = generate_simple_analytics_sql( + effective_grounding_query, + contexts, + ) + if not deterministic_sql: + return None + + logger.info("Follow-up SQL deterministic fast path produced a candidate.") + post_process = await self._components["post_processor"].run( + [deterministic_sql], + project_id=project_id, + mdl_hash=mdl_hash, + contexts=contexts, + fallback_query=effective_grounding_query, + use_dry_plan=use_dry_plan, + data_source=data_source, + allow_dry_plan_fallback=allow_dry_plan_fallback, + ) + if post_process.get("valid_generation_result"): + logger.info("Follow-up SQL deterministic fast path accepted candidate.") + return {"post_process": post_process, "fast_path": "deterministic"} + + logger.info( + "Follow-up SQL deterministic fast path rejected candidate; continuing to LLM. reason=%s", + post_process.get("invalid_generation_result", {}).get("error"), + ) + return None diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 553e5ac6f6..1cd7b1aac3 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -13,7 +13,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, - construct_schema_identifier_catalog, + sanitize_sql_generation_reasoning, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -115,7 +115,7 @@ async def generate_sql_reasoning( def post_process( generate_sql_reasoning: dict, ) -> dict: - return generate_sql_reasoning.get("replies")[0] + return sanitize_sql_generation_reasoning(generate_sql_reasoning.get("replies")[0]) ## End of Pipeline diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 48b9d0eb2a..1ae2274aac 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -51,6 +51,7 @@ def _project_filter_conditions( - **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. - **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. +- **Schema-resolvable references:** Treat user language as schema-resolvable references when it maps to provided table or column names, aliases, descriptions, or data details, even if the user did not type exact table or column names; do not require the user to write exact schema identifiers. Do not classify a data retrieval or analytics question as MISLEADING only because the user did not write exact table or column names. ### Intent Definitions ### @@ -58,13 +59,13 @@ def _project_filter_conditions( **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes references to specific tables, columns, or data details. -- The question includes **complete information** with specific tables, columns, or data values needed for execution. +- The question (or related previous query) includes references to specific tables, columns, data details, or schema-resolvable database concepts. +- The question includes **complete information** with specific or schema-resolvable tables, columns, concepts, or data values needed for execution. - The question provides **all necessary parameters** to generate executable SQL. **Requirements:** - Must have complete filter criteria, specific values, or clear references to previous context. -- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. +- Include exact schema identifiers in your reasoning only when they are available; otherwise cite the user's schema-resolvable references. - Reference phrases from the user's inputs that clearly relate to the schema. **Examples:** diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 81289081b5..e6faf4de21 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -1,6 +1,7 @@ import asyncio import logging import sys +import time from typing import Any, Optional from hamilton import base @@ -19,23 +20,24 @@ sql_to_answer_system_prompt = """ ### TASK -You are a data analyst that great at answering non-technical user's questions based on the data, sql so that even non technical users can easily understand. -Please answer the user's question in concise and clear manner in Markdown format. +You are a data analyst answering a user's question using only the provided SQL result data. +Answer clearly for a non-technical user in Markdown. ### INSTRUCTIONS -1. Read the user's question and understand the user's intention. -2. Read the sql and understand the data. -3. Make sure the answer is aimed for non-technical users, so don't mention any technical terms such as SQL syntax. -4. Generate a concise and clear answer in string format to answerthe user's question based on the data and sql. -5. If answer is in list format, only list top few examples, and tell users there are more results omitted. -6. Answer must be in the same language user specified. -7. Do not include ```markdown or ``` in the answer. -8. If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. +1. Use only the provided Data columns, row records, and rows. Do not invent values, totals, categories, dates, examples, or analysis outputs. +2. Treat row records as the clearest representation of the result because each value is paired with its column name. +3. Do not write code, Python, pseudo-code, SQL, code fences, implementation steps, or phrases such as "running the above code". +4. Do not mention SQL syntax, table names, or database internals unless the user explicitly asks for them. +5. If no rows are provided, say that no matching rows were returned. +6. If rows are detailed records, summarize the visible records directly. If rows are aggregates, answer using the aggregate values. +7. If the answer is a list, keep it concise and use only examples present in the provided data. +8. Answer must be in the same language user specified. +9. If the user provides a custom instruction, follow it strictly for the response style unless it conflicts with these data-grounding rules. ### OUTPUT FORMAT -Please provide your response in proper Markdown stringformat. +Return only the user-facing answer as a Markdown string. """ sql_to_answer_user_prompt_template = """ @@ -44,13 +46,14 @@ SQL: {{ sql }} Data: columns: {{ sql_data.columns }} +row records: {{ sql_data.row_records }} rows: {{ sql_data.data }} Language: {{ language }} Current Time: {{ current_time }} Custom Instruction: {{ custom_instruction }} -Please think step by step and answer the user's question. +Answer directly from the provided row records and rows. """ @@ -158,7 +161,8 @@ async def run( custom_instruction: Optional[str] = None, ) -> dict: logger.info("Sql_Answer Generation pipeline is running...") - return await self._pipe.execute( + started_at = time.perf_counter() + result = await self._pipe.execute( ["generate_answer"], inputs={ "query": query, @@ -171,3 +175,9 @@ async def run( **self._components, }, ) + logger.info( + "Ask timing query_id=%s stage=answer_formatting elapsed_ms=%.1f", + query_id or "", + (time.perf_counter() - started_at) * 1000, + ) + return result diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 8004135449..bf418bcd1f 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -38,11 +38,12 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) 2. Then, generate the syntactically correct ANSI SQL query to correct the error. 3. If the failed SQL references a table, view, column, function, alias, or placeholder that is not present in DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. Regenerate from the USER QUESTION and DATABASE SCHEMA. 4. Treat invalid object name, dataset not found, table not found, invalid column name, and invalid identifier errors as schema-grounding failures. Use exact declared identifiers from DATABASE SCHEMA only. -5. Do not create dummy CTEs, placeholder tables, table-existence checks, or generic replacement names to make the query executable. If the requested intent is supported by retrieved schema objects, use those exact objects; otherwise return null for sql. -6. For grouped queries, repair SQL Server errors about ORDER BY columns not appearing in GROUP BY by ordering with selected grouping columns or selected aggregate aliases, or by adding the exact ordering key to both SELECT and GROUP BY when that key is declared in DATABASE SCHEMA. -7. Do not preserve generic log, file, JSON, payload, text, or app-metric scans when DATABASE SCHEMA contains exact modeled business columns for the user's requested entity, measure, status, date, or dimension. -8. If the failed SQL invented component fields for a metric that exists directly in DATABASE SCHEMA, replace the calculation with the exact declared metric column. -9. For sales or revenue questions, avoid tariff, duty, customs, import, refund, or claim datasets unless the USER QUESTION explicitly asks for those domains. +5. If the error reports an unknown table or field, replace it only with an exact executable identifier declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not retry the same unknown identifier. +6. Do not create dummy CTEs, placeholder tables, table-existence checks, or generic replacement names to make the query executable. If the requested intent is supported by retrieved schema objects, use those exact objects; otherwise return null for sql. +7. For grouped queries, repair SQL Server errors about ORDER BY columns not appearing in GROUP BY by ordering with selected grouping columns or selected aggregate aliases, or by adding the exact ordering key to both SELECT and GROUP BY when that key is declared in DATABASE SCHEMA. +8. Do not preserve generic log, file, JSON, payload, text, or app-metric scans when DATABASE SCHEMA contains exact modeled business columns for the user's requested entity, measure, status, date, or dimension. +9. If the failed SQL invented component fields for a metric that exists directly in DATABASE SCHEMA, replace the calculation with the exact declared metric column. +10. Do not route a question to a different business domain because of generic keyword overlap. Use only retrieved schema metadata that directly represents the requested entities, measures, filters, dates, and dimensions. ### SQL RULES ### Make sure you follow the SQL Rules strictly. diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index da4aae67be..2e485e265d 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -15,7 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, - construct_schema_identifier_catalog, + generate_simple_analytics_sql, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -143,6 +143,7 @@ async def post_process( post_processor: SQLGenPostProcessor, data_source: str, query: str | None = None, + grounding_query: str | None = None, documents: list[str] | None = None, project_id: str | None = None, mdl_hash: str | None = None, @@ -156,7 +157,7 @@ async def post_process( project_id=project_id, mdl_hash=mdl_hash, contexts=documents, - fallback_query=query, + fallback_query=grounding_query or query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -213,7 +214,7 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, - validation_contexts: list[str] | None = None, + grounding_query: str | None = None, ): logger.info( "SQL Generation pipeline is running for project_id=%s mdl_hash=%s", @@ -229,8 +230,9 @@ async def run( metadata = {} data_source = metadata.get("data_source", "local_file") + effective_grounding_query = grounding_query or query unsupported_result = unsupported_schema_generation_result( - query, + effective_grounding_query, contexts=contexts, data_source=data_source, ) @@ -245,6 +247,7 @@ async def run( ["post_process"], inputs={ "query": query, + "grounding_query": effective_grounding_query, "documents": contexts, "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, @@ -264,3 +267,63 @@ async def run( **self._components, }, ) + + async def run_deterministic_fast_path( + self, + query: str, + contexts: list[str], + project_id: str | None = None, + mdl_hash: str | None = None, + use_dry_plan: bool = False, + allow_dry_plan_fallback: bool = True, + allow_data_preview: bool = False, + grounding_query: str | None = None, + ) -> dict | None: + if use_dry_plan: + metadata = await retrieve_metadata( + project_id or "", self._retriever, mdl_hash + ) + else: + metadata = {} + data_source = metadata.get("data_source", "local_file") + effective_grounding_query = grounding_query or query + + unsupported_result = unsupported_schema_generation_result( + effective_grounding_query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "SQL generation deterministic fast path returned unsupported schema before LLM." + ) + return {"post_process": unsupported_result, "fast_path": "unsupported"} + + deterministic_sql = generate_simple_analytics_sql( + effective_grounding_query, + contexts, + ) + if not deterministic_sql: + return None + + logger.info("SQL generation deterministic fast path produced a candidate.") + post_process = await self._components["post_processor"].run( + [deterministic_sql], + project_id=project_id, + mdl_hash=mdl_hash, + contexts=contexts, + fallback_query=effective_grounding_query, + use_dry_plan=use_dry_plan, + data_source=data_source, + allow_dry_plan_fallback=allow_dry_plan_fallback, + allow_data_preview=allow_data_preview, + ) + if post_process.get("valid_generation_result"): + logger.info("SQL generation deterministic fast path accepted candidate.") + return {"post_process": post_process, "fast_path": "deterministic"} + + logger.info( + "SQL generation deterministic fast path rejected candidate; continuing to LLM. reason=%s", + post_process.get("invalid_generation_result", {}).get("error"), + ) + return None 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 fd38b9a291..c19015d1f3 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -13,7 +13,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, - construct_schema_identifier_catalog, + sanitize_sql_generation_reasoning, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -100,7 +100,7 @@ async def generate_sql_reasoning( def post_process( generate_sql_reasoning: dict, ) -> dict: - return generate_sql_reasoning.get("replies")[0] + return sanitize_sql_generation_reasoning(generate_sql_reasoning.get("replies")[0]) ## End of Pipeline diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 0ff59e43a2..e01f274c6a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,6 +1,8 @@ import logging import re +import time from datetime import datetime, timezone +from functools import lru_cache from typing import Any, Dict, List import aiohttp @@ -35,6 +37,23 @@ "UNIQUE", } +def _timing_ms(started_at: float) -> float: + return (time.perf_counter() - started_at) * 1000 + +_DDL_CREATE_PATTERN = re.compile( + r"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?P
[^\s(]+)\s*\(", + re.IGNORECASE, +) +_DDL_COLUMN_KEYWORDS = { + "CHECK", + "CONSTRAINT", + "FOREIGN", + "INDEX", + "KEY", + "PRIMARY", + "UNIQUE", +} + _IDENTIFIER_TOKEN = r'"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*' _QUALIFIED_IDENTIFIER = rf"(?:{_IDENTIFIER_TOKEN})(?:\s*\.\s*(?:{_IDENTIFIER_TOKEN}))*" @@ -70,6 +89,64 @@ _QUALIFIED_COLUMN = re.compile( rf"(?P{_QUALIFIED_IDENTIFIER})\s*\.\s*(?P{_IDENTIFIER_TOKEN})" ) +_SQL_START = re.compile(r"^\s*(?:WITH|SELECT)\b", re.IGNORECASE | re.DOTALL) +_SQL_REASONING_DISALLOWED = re.compile( + r"(?is)```|\b(?:SELECT|WITH|FROM|JOIN|WHERE|GROUP\s+BY|ORDER\s+BY|" + r"HAVING|LIMIT|UNION)\b" +) +_SQL_REASONING_ASSUMPTION_WORDS = re.compile( + r"(?i)\b(?:assume|assuming|likely|possible|might|example)\b" +) +_SAFE_SQL_REASONING_PLAN = """1. **Identify Supported Data**: Use only the retrieved schema metadata for the active project to decide whether the question is supported. +2. **Ground Requested Filters**: Match user-entered names, statuses, products, and similar values only against verified schema metadata or sample values; ask for clarification when the match is not clear. +3. **Prepare The Result Shape**: Build the requested rows, totals, time breakdowns, ordering, or limits using only validated retrieved tables and columns.""" +_SQL_OBJECT_ALIAS_STOP_WORDS = { + "CROSS", + "EXCEPT", + "FETCH", + "FULL", + "GROUP", + "HAVING", + "INNER", + "INTERSECT", + "JOIN", + "LEFT", + "LIMIT", + "MATCH_RECOGNIZE", + "NATURAL", + "OFFSET", + "ORDER", + "RIGHT", + "TABLESAMPLE", + "UNION", + "WHERE", +} + + +def sanitize_sql_generation_reasoning(reasoning: Any) -> str: + text = str(reasoning or "").strip() + if not text: + return _SAFE_SQL_REASONING_PLAN + + if _SQL_REASONING_DISALLOWED.search(text) or _SQL_REASONING_ASSUMPTION_WORDS.search( + text + ): + logger.warning( + "SQL generation reasoning violated schema-grounding display contract; " + "using safe non-executable reasoning plan." + ) + return _SAFE_SQL_REASONING_PLAN + + return text +_KEYED_SEMANTIC_IDENTIFIER_ALIAS = re.compile( + r"""(?ix) + \b(?:alias|displayName|display_name|sourceColumnName|source_column_name| + sourceTableName|source_table_name|source_table|source_name| + physicalName|physical_name|referenceName|reference_name| + lineageName|lineage_name) + ['"]?\s*[:=]\s*['"](?P[^'"]+)['"] + """ +) _UNQUALIFIED_QUOTED_IDENTIFIER = re.compile(r'(?(?:[^"]|"")*)"') _UNQUALIFIED_BARE_IDENTIFIER = re.compile( r"(?[A-Za-z_][A-Za-z0-9_$]*)\b(?!\s*\.)" @@ -148,6 +225,9 @@ "CONCAT", "COUNT", "COUNT_BIG", + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", "DATE_TRUNC", "DAY", "EXTRACT", @@ -202,6 +282,20 @@ "YEAR", } _FALLBACK_TOKEN = re.compile(r"[a-z0-9]+") +_QUERY_VALUE_TOKEN = re.compile(r"[A-Za-z0-9]+(?:[-_][A-Za-z0-9]+)*") +_FILTER_VALUE_PHRASE = re.compile( + r"""(?ix) + \b(?Pcalled|contains|containing|for|from|like|matching|named|where|with)\s+ + (?P[A-Za-z0-9][A-Za-z0-9_ ./&'()-]*?) + (?= + \s+\b(?:across|after|before|by|during|group|grouped|having|in|limit|on| + order|ordered|per|since|sort|sorted|this|last|next|where|with|from|for)\b + |[?,!;:] + |$ + ) + """ +) +_OPEN_FILTER_VALUE_INTRODUCERS = {"for", "from"} _FALLBACK_STOPWORDS = { "a", "an", @@ -233,66 +327,127 @@ } -def _expand_fallback_token_aliases(tokens: set[str]) -> set[str]: - aliases = { - "bu": {"business", "unit"}, - "cust": {"customer"}, - "customers": {"customer"}, - "critical": {"priority", "severity"}, - "boards": {"board"}, - "high": {"priority", "severity"}, - "highest": {"top"}, - "logs": {"log", "record"}, - "log": {"record"}, - "lows": {"low"}, - "lowest": {"bottom"}, - "models": {"model"}, - "inv": {"invoice"}, - "invoices": {"invoice"}, - "ord": {"order"}, - "orders": {"order"}, - "qty": {"quantity"}, - "num": {"number"}, - "no": {"number"}, - "numbers": {"number"}, - "prod": {"product"}, - "products": {"product"}, - "priorities": {"priority", "severity"}, - "priority": {"severity"}, - "recent": {"latest"}, - "records": {"record"}, - "rep": {"representative", "salesperson"}, - "repairs": {"repair"}, - "salesperson": {"sales", "person"}, - "severity": {"priority"}, - "supplier": {"vendor"}, - "suppliers": {"supplier", "vendor"}, - "tech": {"technician"}, - "technician": {"tech"}, - "vendor": {"supplier"}, - "vendors": {"supplier", "vendor"}, - "failed": {"failure"}, - "failures": {"failure"}, - "defects": {"defect"}, - "types": {"type"}, - "units": {"unit", "serial"}, - "urgency": {"priority", "severity"}, - "locations": {"location"}, - "materials": {"material"}, - "missing": {"blank", "empty", "null"}, - } - expanded = set(tokens) - for token in list(tokens): - expanded.update(aliases.get(token, set())) - if {"business", "unit"}.issubset(expanded): - expanded.add("bu") - if "customer" in expanded and "number" in expanded: - expanded.update({"cust", "id", "no"}) - if "order" in expanded and "number" in expanded: - expanded.update({"ord", "id", "no"}) +def _fallback_token_variants(token: str) -> set[str]: + token = token.lower() + variants = {token} + + generic_tokens = globals().get("_GENERIC_SCHEMA_INTENT_TOKENS", set()) + if token in generic_tokens: + return variants + + if len(token) > 4 and token.endswith("ies"): + variants.add(token[:-3] + "y") + elif len(token) > 4 and token.endswith("es"): + if token.endswith(("ches", "shes", "sses", "uses", "xes", "zes")): + variants.add(token[:-2]) + else: + variants.add(token[:-1]) + elif len(token) > 3 and token.endswith("s") and not token.endswith(("ss", "us")): + variants.add(token[:-1]) + if len(token) > 4 and token.endswith("ed"): + stem = token[:-2] + variants.add(stem) + if not stem.endswith("e"): + variants.add(stem + "e") + if len(token) > 5 and token[-3] in {"d", "s", "t", "v", "z"}: + variants.add(token[:-1]) + if len(stem) > 2 and stem[-1] == stem[-2]: + variants.add(stem[:-1]) + if len(token) > 5 and token.endswith("ing"): + stem = token[:-3] + variants.add(stem) + if not stem.endswith("d"): + variants.add(stem + "e") + if len(stem) > 2 and stem[-1] == stem[-2]: + variants.add(stem[:-1]) + + return {variant for variant in variants if variant} + + +def _expand_fallback_token_variants(tokens: set[str]) -> set[str]: + expanded = set() + for token in tokens: + expanded.update(_fallback_token_variants(token)) return expanded +_COMPOUND_IDENTIFIER_PART_TOKENS = { + "account", + "amount", + "balance", + "business", + "buyer", + "category", + "client", + "company", + "count", + "currency", + "customer", + "date", + "division", + "end", + "exchange", + "failure", + "gross", + "group", + "invoice", + "market", + "material", + "month", + "name", + "order", + "person", + "priority", + "product", + "quantity", + "rate", + "record", + "repair", + "sales", + "salesperson", + "severity", + "status", + "supplier", + "task", + "ticket", + "type", + "unit", + "value", + "vendor", + "year", +} +_COMPOUND_IDENTIFIER_ALIASES = { + "acct": {"account"}, + "amt": {"amount"}, + "bu": {"business", "unit"}, + "curr": {"currency"}, + "cust": {"customer"}, + "gl": {"general", "ledger"}, + "ord": {"order"}, + "prod": {"product"}, + "qty": {"quantity"}, + "vend": {"vendor"}, +} + + +def _compound_identifier_tokens(token: str) -> set[str]: + if len(token) < 5: + return set() + + tokens: set[str] = set() + for part in _COMPOUND_IDENTIFIER_PART_TOKENS: + if part != token and len(part) >= 3 and part in token: + tokens.add(part) + + if token not in _COMPOUND_IDENTIFIER_PART_TOKENS: + for alias, expansions in _COMPOUND_IDENTIFIER_ALIASES.items(): + if alias != token and (token.startswith(alias) or token.endswith(alias)): + tokens.add(alias) + tokens.update(expansions) + + return tokens + + def normalize_wren_sql_dialect(sql: str) -> str: if not sql: return sql @@ -476,6 +631,12 @@ def _iter_context_texts(contexts: list[Any] | None): yield getattr(context, "content", context) +def _context_cache_key(contexts: list[Any] | None) -> tuple[str, ...]: + if not contexts: + return tuple() + return tuple(str(context) for context in _iter_context_texts(contexts)) + + def _clean_contract_value(value: str) -> str: value = value.strip().strip(",") if not value: @@ -618,13 +779,43 @@ def add(identifier: str) -> None: return identifiers +def _thaw_schema_index( + frozen_schema_index: tuple[tuple[str, tuple[str, ...] | None], ...], +) -> dict[str, set[str] | None]: + return { + relation: None if columns is None else set(columns) + for relation, columns in frozen_schema_index + } + + +def _freeze_schema_index( + schema_index: dict[str, set[str] | None], +) -> tuple[tuple[str, tuple[str, ...] | None], ...]: + return tuple( + sorted( + ( + relation, + None if columns is None else tuple(sorted(columns)), + ) + for relation, columns in schema_index.items() + ) + ) + + def _extract_schema_index(contexts: list[Any] | None) -> dict[str, set[str] | None]: - if not contexts: - return {} + return _thaw_schema_index(_cached_extract_schema_index(_context_cache_key(contexts))) - schema_index = _extract_contract_schema_index(contexts) - for context in _iter_context_texts(contexts): +@lru_cache(maxsize=256) +def _cached_extract_schema_index( + context_texts: tuple[str, ...], +) -> tuple[tuple[str, tuple[str, ...] | None], ...]: + if not context_texts: + return tuple() + + schema_index = _extract_contract_schema_index(list(context_texts)) + + for context in context_texts: relation_match = _DDL_RELATION.search(context) if not relation_match: continue @@ -674,7 +865,7 @@ def _extract_schema_index(contexts: list[Any] | None) -> dict[str, set[str] | No else: existing_columns.update(columns) - return schema_index + return _freeze_schema_index(schema_index) def _semantic_tokens_from_value(value: Any) -> set[str]: @@ -707,10 +898,10 @@ def _extract_semantic_context_payload(context: str) -> dict[str, Any]: def _extract_semantic_tokens_by_column( context: str, -) -> tuple[set[str], dict[str, set[str]]]: +) -> tuple[set[str], dict[str, set[str]], dict[str, list[str]]]: payload = _extract_semantic_context_payload(context) if not payload: - return set(), {} + return set(), {}, {} table_tokens = _semantic_tokens_from_value( payload.get("semantic_context_not_sql_identifiers") @@ -718,6 +909,7 @@ def _extract_semantic_tokens_by_column( table_tokens.update(_semantic_tokens_from_value(payload.get("object_type"))) column_tokens: dict[str, set[str]] = {} + column_sample_values: dict[str, list[str]] = {} for column in payload.get("columns", []) or []: if not isinstance(column, dict): continue @@ -727,27 +919,126 @@ def _extract_semantic_tokens_by_column( tokens = _semantic_tokens_from_value( column.get("semantic_context_not_sql_identifier") ) + tokens.update(_semantic_tokens_from_value(column.get("display_name"))) + tokens.update(_semantic_tokens_from_value(column.get("source_column_name"))) if tokens: column_tokens[column_name] = tokens + sample_values = _extract_column_sample_values(column) + if sample_values: + column_sample_values[column_name] = sample_values + + return table_tokens, column_tokens, column_sample_values + + +def _extract_column_sample_values(column: dict[str, Any]) -> list[str]: + values: list[str] = [] + + def add(value: Any) -> None: + if value is None: + return + if isinstance(value, (list, tuple, set)): + for item in value: + add(item) + return + if isinstance(value, dict): + for item in value.values(): + add(item) + return + text = str(value).strip() + if text and text.lower() not in {item.lower() for item in values}: + values.append(text) + + for key in ( + "sample_values", + "sample_value", + "samples", + "values", + "example_values", + "examples", + "distinct_values", + ): + add(column.get(key)) + return values + + +def _thaw_schema_details( + frozen_schema_details: tuple[ + tuple[str, tuple[tuple[str, str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], ...]], + ..., + ], +) -> dict[str, list[dict[str, Any]]]: + return { + relation: [ + { + "name": name, + "data_type": data_type, + "semantic_tokens": set(semantic_tokens), + "sample_values": list(sample_values), + "_table_semantic_tokens": set(table_semantic_tokens), + } + for ( + name, + data_type, + semantic_tokens, + sample_values, + table_semantic_tokens, + ) in columns + ] + for relation, columns in frozen_schema_details + } + - return table_tokens, column_tokens +def _freeze_schema_details( + schema_details: dict[str, list[dict[str, Any]]], +) -> tuple[ + tuple[str, tuple[tuple[str, str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], ...]], + ..., +]: + return tuple( + ( + relation, + tuple( + ( + column["name"], + column["data_type"], + tuple(sorted(column.get("semantic_tokens") or set())), + tuple(column.get("sample_values") or []), + tuple(sorted(column.get("_table_semantic_tokens") or set())), + ) + for column in columns + ), + ) + for relation, columns in sorted(schema_details.items()) + ) def _extract_schema_details( contexts: list[Any] | None, -) -> dict[str, list[dict[str, str]]]: - if not contexts: - return {} +) -> dict[str, list[dict[str, Any]]]: + return _thaw_schema_details( + _cached_extract_schema_details(_context_cache_key(contexts)) + ) + + +@lru_cache(maxsize=256) +def _cached_extract_schema_details( + context_texts: tuple[str, ...], +) -> tuple[ + tuple[str, tuple[tuple[str, str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], ...]], + ..., +]: + if not context_texts: + return tuple() schema_details: dict[str, list[dict[str, str]]] = {} - for context in _iter_context_texts(contexts): + for context in context_texts: relation_match = _DDL_RELATION.search(context) if not relation_match: continue relation_name = _unquote_identifier(relation_match.group("name")) - table_semantic_tokens, column_semantic_tokens = ( + table_semantic_tokens, column_semantic_tokens, column_sample_values = ( _extract_semantic_tokens_by_column(context) ) column_block_match = re.search( @@ -794,13 +1085,14 @@ def _extract_schema_details( "name": name, "data_type": data_type, "semantic_tokens": column_semantic_tokens.get(name, set()), + "sample_values": column_sample_values.get(name, []), "_table_semantic_tokens": table_semantic_tokens, } ) schema_details[relation_name] = columns - return schema_details + return _freeze_schema_details(schema_details) def _is_identifier_boundary(char: str | None) -> bool: @@ -955,1175 +1247,3278 @@ def _replace_bracket_identifiers(sql: str, valid_identifiers: set[str]) -> str: return "".join(result) -def _extract_sql_grounding(sql: str) -> dict[str, Any]: - cte_names = { - _normalize_identifier(match.group("name")) - for match in _CTE_REFERENCE.finditer(sql) +def _semantic_key_can_be_identifier(key_hint: str) -> bool: + key = key_hint.rsplit(".", 1)[-1].lower() + return key in { + "alias", + "displayname", + "display_name", + "lineagename", + "lineage_name", + "physicalname", + "physical_name", + "referencename", + "reference_name", + "sourcecolumnname", + "source_column_name", + "source_name", + "source_table", + "source_table_name", } - relation_references = [] - alias_to_relation = {} - for match in _RELATION_REFERENCE.finditer(sql): - if _is_extract_from_clause(sql, match.start()): - continue - relation = _unquote_identifier(match.group("name")) - if relation.upper() in {"UNNEST", "LATERAL"}: - continue - alias = match.group("alias") - alias = _normalize_identifier(alias) if alias else relation - if alias.upper() in _SQL_OBJECT_ALIAS_STOP_WORDS: - alias = relation - relation_references.append(relation) - alias_to_relation[alias] = relation - alias_to_relation[relation] = relation - qualified_columns = [ - ( - _normalize_identifier(match.group("qualifier")), - _normalize_identifier(match.group("column")), - ) - for match in _QUALIFIED_COLUMN.finditer(sql) - ] +def _clean_semantic_identifier_alias(value: Any) -> str | None: + if value is None: + return None - return { - "cte_names": cte_names, - "relation_references": relation_references, - "alias_to_relation": alias_to_relation, - "qualified_columns": qualified_columns, - } + text = str(value).strip() + if not text or "\n" in text or len(text) > 256: + return None + text = re.sub(r"^\s*/\*+", "", text) + text = re.sub(r"\*+/\s*$", "", text) + text = text.strip(" \t\r\n'\"`.,;:(){}") + if not text or len(text) > 256: + return None + if re.search(r"[!?;]", text): + return None -def _is_extract_from_clause(sql: str, from_start: int) -> bool: - prefix = sql[:from_start] - last_open = prefix.rfind("(") - if last_open == -1 or prefix.rfind(")") > last_open: - return False + normalized = _normalize_identifier(text) + if not normalized: + return None + if normalized.upper() in _SQL_RESERVED_WORDS: + return None + if "." not in normalized and normalized.lower() in _GENERIC_SCHEMA_INTENT_TOKENS: + return None - before_open = prefix[:last_open].rstrip() - return before_open.upper().endswith("EXTRACT") + if re.search(r"\s", normalized): + tokens = _fallback_tokens(normalized) + if len(tokens) > 5: + return None + return normalized -def _strip_string_literals(sql: str) -> str: - return _SINGLE_QUOTED_LITERAL.sub("''", sql) +def _table_reference_aliases(value: Any) -> set[str]: + if not isinstance(value, dict): + return set() -def _extract_clause(sql: str, clause: str, end_clauses: tuple[str, ...]) -> str: - end_pattern = "|".join(re.escape(end_clause) for end_clause in end_clauses) - pattern = re.compile( - rf"(?is)\b{re.escape(clause)}\b\s+(?P.*?)(?=\b(?:{end_pattern})\b|$)" + table_reference = value.get("source_table_reference") or value.get( + "tableReference" ) - match = pattern.search(sql) - return match.group("body").strip() if match else "" + if not isinstance(table_reference, dict): + table_reference = value + table = table_reference.get("table") + if not table: + return set() -def _extract_select_clause(sql: str) -> str: - match = re.search(r"(?is)\bSELECT\b\s+(?P.*?)(?=\bFROM\b)", sql) - return match.group("body").strip() if match else "" + aliases: set[str] = set() + schema = table_reference.get("schema") + catalog = table_reference.get("catalog") + if schema: + aliases.add(f"{schema}.{table}") + if catalog and schema: + aliases.add(f"{catalog}.{schema}.{table}") + aliases.add(str(table)) + return { + alias + for alias in (_clean_semantic_identifier_alias(alias) for alias in aliases) + if alias + } -def _split_select_expression_alias(expression: str) -> tuple[str, str | None]: - as_match = re.search( - r"(?is)\s+AS\s+(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)\s*$", - expression, - ) - if as_match: - return expression[: as_match.start()].strip(), _unquote_identifier( - as_match.group("alias") - ) +def _keyed_semantic_identifier_aliases(text: str) -> set[str]: + aliases: set[str] = set() + for match in _KEYED_SEMANTIC_IDENTIFIER_ALIAS.finditer(text): + alias = _clean_semantic_identifier_alias(match.group("value")) + if alias: + aliases.add(alias) + return aliases - return expression, None +def _semantic_identifier_aliases(value: Any, key_hint: str = "") -> set[str]: + aliases: set[str] = set() + if value is None: + return aliases -def _extract_output_aliases(select_clause: str) -> set[str]: - aliases = set() - for expression in _split_sql_tokens(select_clause): - _, alias = _split_select_expression_alias(expression) + if isinstance(value, dict): + aliases.update(_table_reference_aliases(value)) + for key, nested_value in value.items(): + nested_key = f"{key_hint}.{key}" if key_hint else str(key) + aliases.update(_semantic_identifier_aliases(nested_value, nested_key)) + return aliases + + if isinstance(value, (list, tuple, set)): + for nested_value in value: + aliases.update(_semantic_identifier_aliases(nested_value, key_hint)) + return aliases + + text = str(value) + if _semantic_key_can_be_identifier(key_hint): + alias = _clean_semantic_identifier_alias(text) if alias: aliases.add(alias) + aliases.update(_keyed_semantic_identifier_aliases(text)) return aliases -def _iter_unqualified_identifier_candidates(expression: str): - stripped = _strip_string_literals(expression) +def _identifier_alias_key(identifier: str) -> str: + return _normalize_identifier(identifier).lower() - for match in _UNQUALIFIED_QUOTED_IDENTIFIER.finditer(stripped): - yield _unquote_identifier(f'"{match.group("name")}"') - without_quoted = _UNQUALIFIED_QUOTED_IDENTIFIER.sub(" ", stripped) - for match in _UNQUALIFIED_BARE_IDENTIFIER.finditer(without_quoted): - name = match.group("name") - following = without_quoted[match.end() :].lstrip() - if following.startswith("("): +def _identifier_style_aliases(identifier: str) -> set[str]: + cleaned = _clean_semantic_identifier_alias(identifier) + if not cleaned: + return set() + + aliases = {cleaned} + parts = re.findall(r"[A-Za-z0-9]+", cleaned) + if not 1 < len(parts) <= 6: + return aliases + + aliases.add("_".join(parts)) + aliases.add("".join(parts)) + aliases.add(parts[0] + "".join(part[:1].upper() + part[1:] for part in parts[1:])) + aliases.add("".join(part[:1].upper() + part[1:] for part in parts)) + return { + alias + for alias in aliases + if _clean_semantic_identifier_alias(alias) + } + + +def _add_unique_identifier_alias( + aliases: dict[str, str], + alias_values: dict[str, str], + alias: Any, + target: str, +) -> None: + if not target: + return + + for cleaned_alias in _identifier_style_aliases(alias): + if _identifier_alias_key(cleaned_alias) == _identifier_alias_key(target): continue - yield name + key = _identifier_alias_key(cleaned_alias) + existing = aliases.get(key) + if existing is None and key not in aliases: + aliases[key] = target + alias_values[key] = cleaned_alias + elif existing != target: + aliases[key] = "" -def _sql_mentions_identifier(sql: str, identifier: str) -> bool: - stripped = _strip_string_literals(sql) - quoted_identifier = re.escape(_quote_identifier(identifier)) - bracket_identifier = re.escape(f"[{identifier}]") - bare_identifier = re.escape(identifier) - return bool( - re.search(rf'(? set[str]: + cleaned = _clean_semantic_identifier_alias(identifier) + if not cleaned: + return set() -def _validate_unqualified_columns_for_single_relation( - sql: str, - schema_index: dict[str, set[str] | None], - grounding: dict[str, Any], -) -> str | None: - real_relations = [ - relation - for relation in grounding["relation_references"] - if relation not in grounding["cte_names"] - ] - if grounding["cte_names"]: - return None + parts = cleaned.split(".") + variants = { + _quote_identifier(cleaned), + f"[{cleaned}]", + f"`{cleaned}`", + } + if not re.search(r"\s", cleaned): + variants.add(cleaned) + if len(parts) > 1: + variants.add(".".join(parts)) + variants.add(".".join(_quote_identifier(part) for part in parts)) + variants.add(".".join(f"[{part}]" for part in parts)) + variants.add(".".join(f"`{part}`" for part in parts)) + return variants + + +def _identifier_dot_aliases(identifier: str) -> set[str]: + if "." in identifier or "_" not in identifier: + return set() - unique_real_relations = list(dict.fromkeys(real_relations)) - if len(unique_real_relations) != 1: - return None + parts = [part for part in identifier.split("_") if part] + aliases = set() + if len(parts) >= 2: + aliases.add(f"{parts[0]}.{'_'.join(parts[1:])}") + if len(parts) >= 3: + aliases.add(f"{parts[0]}.{parts[1]}.{'_'.join(parts[2:])}") + return aliases - relation = unique_real_relations[0] - valid_columns = schema_index.get(relation) - if valid_columns is None: - return None - select_clause = _extract_select_clause(sql) - output_aliases = _extract_output_aliases(select_clause) - ignored_identifiers = ( - set(schema_index) - | set(grounding["alias_to_relation"]) - | set(grounding["cte_names"]) - | output_aliases - ) +def _render_table_identifier(identifier: str) -> str: + return identifier if "." in identifier else _quote_identifier(identifier) - clause_expressions = [] + +def _render_qualifier_identifier(identifier: str, table_names: set[str]) -> str: + if identifier in table_names: + return _render_table_identifier(identifier) + if "." in identifier or _identifier_needs_quotes(identifier): + return _quote_identifier(identifier) + return identifier + + +def _replace_sql_text_outside_literals( + sql: str, + target: str, + replacement: str, + case_insensitive: bool = False, +) -> str: + if not target: + return sql + + result = [] + index = 0 + in_single_quote = False + in_double_quote = False + in_line_comment = False + in_block_comment = False + length = len(sql) + target_length = len(target) + comparable_target = target.lower() if case_insensitive else target + + while index < length: + current = sql[index] + nxt = sql[index + 1] if index + 1 < length else None + + if in_line_comment: + result.append(current) + if current == "\n": + in_line_comment = False + index += 1 + continue + if in_block_comment: + result.append(current) + if current == "*" and nxt == "/": + result.append(nxt) + index += 2 + in_block_comment = False + else: + index += 1 + continue + if in_single_quote: + result.append(current) + if current == "'" and nxt == "'": + result.append(nxt) + index += 2 + elif current == "'": + in_single_quote = False + index += 1 + else: + index += 1 + continue + + candidate = sql[index : index + target_length] + comparable_candidate = candidate.lower() if case_insensitive else candidate + if comparable_candidate == comparable_target: + before = sql[index - 1] if index > 0 else None + after_index = index + target_length + after = sql[after_index] if after_index < length else None + if _is_identifier_boundary(before) and _is_identifier_boundary(after): + result.append(replacement) + index = after_index + continue + + if in_double_quote: + result.append(current) + if current == '"' and nxt == '"': + result.append(nxt) + index += 2 + elif current == '"': + in_double_quote = False + index += 1 + else: + index += 1 + continue + + if current == "-" and nxt == "-": + result.append(current) + result.append(nxt) + index += 2 + in_line_comment = True + continue + if current == "/" and nxt == "*": + result.append(current) + result.append(nxt) + index += 2 + in_block_comment = True + continue + if current == "'": + result.append(current) + index += 1 + in_single_quote = True + continue + if current == '"': + result.append(current) + index += 1 + in_double_quote = True + continue + + result.append(current) + index += 1 + + return "".join(result) + + +def _extract_sql_grounding(sql: str) -> dict[str, Any]: + cte_names = { + _normalize_identifier(match.group("name")) + for match in _CTE_REFERENCE.finditer(sql) + } + relation_references = [] + alias_to_relation = {} + + for match in _RELATION_REFERENCE.finditer(sql): + if _is_extract_from_clause(sql, match.start()): + continue + relation = _unquote_identifier(match.group("name")) + if relation.upper() in {"UNNEST", "LATERAL"}: + continue + alias = match.group("alias") + alias = _normalize_identifier(alias) if alias else relation + if alias.upper() in _SQL_OBJECT_ALIAS_STOP_WORDS: + alias = relation + relation_references.append(relation) + alias_to_relation[alias] = relation + alias_to_relation[relation] = relation + + qualified_columns = [ + ( + _normalize_identifier(match.group("qualifier")), + _normalize_identifier(match.group("column")), + ) + for match in _QUALIFIED_COLUMN.finditer(sql) + ] + + return { + "cte_names": cte_names, + "relation_references": relation_references, + "alias_to_relation": alias_to_relation, + "qualified_columns": qualified_columns, + } + + +def _is_extract_from_clause(sql: str, from_start: int) -> bool: + prefix = sql[:from_start] + last_open = prefix.rfind("(") + if last_open == -1 or prefix.rfind(")") > last_open: + return False + + before_open = prefix[:last_open].rstrip() + return before_open.upper().endswith("EXTRACT") + + +def _strip_string_literals(sql: str) -> str: + return _SINGLE_QUOTED_LITERAL.sub( + lambda match: " " * len(match.group(0)), + sql, + ) + + +def _extract_clause(sql: str, clause: str, end_clauses: tuple[str, ...]) -> str: + end_pattern = "|".join(re.escape(end_clause) for end_clause in end_clauses) + pattern = re.compile( + rf"(?is)\b{re.escape(clause)}\b\s+(?P.*?)(?=\b(?:{end_pattern})\b|$)" + ) + match = pattern.search(sql) + return match.group("body").strip() if match else "" + + +def _extract_select_clause(sql: str) -> str: + match = re.search(r"(?is)\bSELECT\b\s+(?P.*?)(?=\bFROM\b)", sql) + return match.group("body").strip() if match else "" + + +def _split_select_expression_alias(expression: str) -> tuple[str, str | None]: + as_match = re.search( + r"(?is)\s+AS\s+(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)\s*$", + expression, + ) + if as_match: + return expression[: as_match.start()].strip(), _unquote_identifier( + as_match.group("alias") + ) + + return expression, None + + +def _extract_output_aliases(select_clause: str) -> set[str]: + aliases = set() for expression in _split_sql_tokens(select_clause): - expression, _ = _split_select_expression_alias(expression) - clause_expressions.append(expression) - clause_expressions.extend( - filter( - None, - [ - _extract_clause( - sql, - "WHERE", - ("GROUP BY", "HAVING", "ORDER BY", "LIMIT", "OFFSET"), - ), - _extract_clause( - sql, - "GROUP BY", - ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), - ), - _extract_clause( - sql, - "HAVING", - ("ORDER BY", "LIMIT", "OFFSET"), - ), - _extract_clause(sql, "ORDER BY", ("LIMIT", "OFFSET")), - ], + _, alias = _split_select_expression_alias(expression) + if alias: + aliases.add(alias) + return aliases + + +def _iter_unqualified_identifier_candidates(expression: str): + stripped = _strip_string_literals(expression) + + for match in _UNQUALIFIED_QUOTED_IDENTIFIER.finditer(stripped): + yield _unquote_identifier(f'"{match.group("name")}"') + + without_quoted = _UNQUALIFIED_QUOTED_IDENTIFIER.sub(" ", stripped) + for match in _UNQUALIFIED_BARE_IDENTIFIER.finditer(without_quoted): + name = match.group("name") + following = without_quoted[match.end() :].lstrip() + if following.startswith("("): + continue + yield name + + +def _sql_mentions_identifier(sql: str, identifier: str) -> bool: + stripped = _strip_string_literals(sql) + quoted_identifier = re.escape(_quote_identifier(identifier)) + bracket_identifier = re.escape(f"[{identifier}]") + bare_identifier = re.escape(identifier) + return bool( + re.search(rf'(? bool: + lowered_sql = sql.lower() + for value in values: + cleaned = _clean_filter_value(value) + if cleaned and _quote_literal(cleaned.lower()) in lowered_sql: + return True + return False + + +def _extract_sql_string_literals(sql: str) -> list[str]: + literals = [] + for match in _SINGLE_QUOTED_LITERAL.finditer(sql): + literal = match.group(0)[1:-1].replace("''", "'") + if literal: + literals.append(literal) + return literals + + +def _extract_column_filter_literals(sql: str, column_name: str) -> list[str]: + stripped = _strip_string_literals(sql) + quoted_column = re.escape(_quote_identifier(column_name)) + bare_column = re.escape(column_name) + column_pattern = rf"(?:{quoted_column}|(? set[str]: + cleaned = _clean_filter_value(str(literal).replace("%", " ")) + if not cleaned: + return set() + if re.fullmatch(r"-?\d+(?:\.\d+)?", cleaned): + return set() + if re.fullmatch( + r"\d{4}-\d{2}-\d{2}(?:[ t]\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?)?", + cleaned, + flags=re.IGNORECASE, + ): + return set() + return { + token + for token in _fallback_tokens(cleaned) + if token not in _GENERIC_SCHEMA_INTENT_TOKENS and not token.isdigit() + } + + +def _validate_literal_values_against_samples( + sql: str, + schema_details: dict[str, list[dict[str, str]]], + grounding: dict[str, Any], + query: str | None = None, +) -> str | None: + query_tokens = _fallback_tokens(query) if query else set() + referenced_relations = { + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + } + for relation in referenced_relations: + columns = schema_details.get(relation, []) + schema_tokens = _schema_tokens_for_table(relation, columns) + allowed_query_literal_tokens = _filter_value_tokens( + _schema_driven_value_terms(query, schema_tokens, columns) + ) + for column in columns: + sample_values = column.get("sample_values") or [] + literals = _extract_column_filter_literals(sql, column["name"]) + if not literals: + continue + sample_tokens = _sample_value_tokens(column) if sample_values else set() + sample_lowers = {str(value).lower() for value in sample_values} + for literal in literals: + literal_tokens = _literal_grounding_tokens(literal) + if not literal_tokens: + continue + if literal.lower() in sample_lowers or literal_tokens & sample_tokens: + continue + if ( + allowed_query_literal_tokens + and literal_tokens <= allowed_query_literal_tokens + ): + logger.info( + "Literal value accepted from grounded user filter column=%s literal=%s", + column["name"], + literal, + ) + continue + if sample_values: + return ( + "Schema grounding failed. The generated SQL filters column " + f"{column['name']} with a literal value not found in that " + "column's verified sample values or grounded filter terms." + ) + return ( + "Schema grounding failed. The generated SQL filters column " + f"{column['name']} with a literal value that is not grounded " + "as a filter value in the user question or verified sample " + "values." + ) + return None + + +def _validate_unqualified_columns_for_single_relation( + sql: str, + schema_index: dict[str, set[str] | None], + grounding: dict[str, Any], +) -> str | None: + real_relations = [ + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + ] + if grounding["cte_names"]: + return None + + unique_real_relations = list(dict.fromkeys(real_relations)) + if len(unique_real_relations) != 1: + return None + + relation = unique_real_relations[0] + valid_columns = schema_index.get(relation) + if valid_columns is None: + return None + + select_clause = _extract_select_clause(sql) + output_aliases = _extract_output_aliases(select_clause) + ignored_identifiers = ( + set(schema_index) + | set(grounding["alias_to_relation"]) + | set(grounding["cte_names"]) + | output_aliases + ) + + clause_expressions = [] + for expression in _split_sql_tokens(select_clause): + expression, _ = _split_select_expression_alias(expression) + clause_expressions.append(expression) + clause_expressions.extend( + filter( + None, + [ + _extract_clause( + sql, + "WHERE", + ("GROUP BY", "HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause( + sql, + "GROUP BY", + ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause( + sql, + "HAVING", + ("ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause(sql, "ORDER BY", ("LIMIT", "OFFSET")), + ], + ) + ) + + invalid_columns = set() + for expression in clause_expressions: + for identifier in _iter_unqualified_identifier_candidates(expression): + upper_identifier = identifier.upper() + if ( + upper_identifier in _SQL_RESERVED_WORDS + or upper_identifier in _SQL_FUNCTION_WORDS + or upper_identifier in _SQL_TYPE_WORDS + or upper_identifier in _DATE_PART_WORDS + or identifier in ignored_identifiers + or identifier in valid_columns + ): + continue + invalid_columns.add(identifier) + + if not invalid_columns: + return None + + return ( + "Schema grounding failed. The SQL references unqualified columns that " + f"are not present in verified table or view {relation}: " + f"{', '.join(sorted(invalid_columns))}. Use only verified columns: " + f"{', '.join(sorted(valid_columns))}." + ) + + +def validate_sql_against_contexts( + sql: str, + contexts: list[Any] | None = None, +) -> str | None: + schema_index = _extract_schema_index(contexts) + if not schema_index: + return None + + valid_relations = set(schema_index) + grounding = _extract_sql_grounding(sql) + cte_names = grounding["cte_names"] + + shadowed_relations = sorted(cte_names & valid_relations) + if shadowed_relations: + return ( + "Schema grounding failed. The SQL creates CTEs with names that already " + f"belong to verified schema objects: {', '.join(shadowed_relations)}. " + "Do not create dummy CTEs for schema objects; use the verified tables or views directly." + ) + + invalid_relations = sorted( + { + relation + for relation in grounding["relation_references"] + if relation not in valid_relations and relation not in cte_names + } + ) + if invalid_relations: + return ( + "Schema grounding failed. The SQL references tables or views that are not " + f"in the retrieved schema for the active question: {', '.join(invalid_relations)}. " + f"Use only verified tables or views: {', '.join(sorted(valid_relations))}." + ) + + alias_to_relation = grounding["alias_to_relation"] + for qualifier, column in grounding["qualified_columns"]: + relation = alias_to_relation.get(qualifier) + if not relation or relation in cte_names: + continue + valid_columns = schema_index.get(relation) + if valid_columns is None: + continue + if column not in valid_columns: + return ( + "Schema grounding failed. The SQL references column " + f"{qualifier}.{column}, but column {column} is not present in verified " + f"table or view {relation}. Use only verified columns: " + f"{', '.join(sorted(valid_columns))}." + ) + + unqualified_column_error = _validate_unqualified_columns_for_single_relation( + sql, + schema_index, + grounding, + ) + if unqualified_column_error: + return unqualified_column_error + + return None + + +def _extract_sql_from_value(value: Any) -> str | None: + if value is None: + return None + + if isinstance(value, str): + text = value.strip() + if not text: + return None + + try: + parsed = orjson.loads(text) + except orjson.JSONDecodeError: + return text if _SQL_START.search(text) else None + + return _extract_sql_from_value(parsed) + + if isinstance(value, dict): + for key in ("sql", "query", "code"): + extracted = _extract_sql_from_value(value.get(key)) + if extracted: + return extracted + + extracted = _extract_sql_from_value(value.get("arguments")) + if extracted: + return extracted + + return None + + if isinstance(value, list): + for item in value: + extracted = _extract_sql_from_value(item) + if extracted: + return extracted + + return None + + +def _extract_generation_sql(generation_result: str | None) -> str | None: + if not generation_result: + return None + + extracted = _extract_sql_from_value(generation_result) + if extracted: + return extracted + + text = generation_result.strip() + return text if _SQL_START.search(text) else None + + +def _sql_has_aggregate_function(sql: str) -> bool: + return bool(re.search(r"(?is)\b(?:AVG|COUNT|MAX|MIN|SUM)\s*\(", sql)) + + +def _group_by_source_columns(sql: str) -> set[str]: + group_by_clause = _extract_clause( + sql, + "GROUP BY", + ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ) + if not group_by_clause: + return set() + + columns: set[str] = set() + for expression in _split_sql_tokens(group_by_clause): + for identifier in _iter_unqualified_identifier_candidates(expression): + upper_identifier = identifier.upper() + if ( + upper_identifier in _SQL_RESERVED_WORDS + or upper_identifier in _SQL_FUNCTION_WORDS + or upper_identifier in _SQL_TYPE_WORDS + or upper_identifier in _DATE_PART_WORDS + ): + continue + columns.add(identifier) + return columns + + +def _query_allows_grouped_aggregate(query: str, query_tokens: set[str]) -> bool: + return bool( + _has_grouping_intent(query, query_tokens) + or _has_count_intent(query_tokens) + or _has_sum_intent(query_tokens) + or _is_rate_metric_intent(query_tokens) + or _is_distribution_metric_intent(query_tokens) + ) + + +def _missing_value_target_tokens(query: str, query_tokens: set[str]) -> set[str]: + match = re.search( + r"(?is)\b(?:blank|empty|missing|null)\s+(?P[A-Za-z0-9_ /-]+)", + query, + ) + if match: + tokens = _fallback_tokens(match.group("value")) + else: + tokens = set(query_tokens) + return tokens - _GENERIC_SCHEMA_INTENT_TOKENS - _NULL_CHECK_TOKENS + + +def _sql_null_checked_columns(sql: str, columns: list[dict[str, str]]) -> set[str]: + checked_columns: set[str] = set() + stripped = _strip_string_literals(sql) + for column in columns: + quoted_column = re.escape(_quote_identifier(column["name"])) + bare_column = re.escape(column["name"]) + column_pattern = rf"(?:{quoted_column}|(? str | None: + grouping_tokens = _grouping_phrase_tokens(query) - _GENERIC_SCHEMA_INTENT_TOKENS + if not grouping_tokens or _grouping_phrase_has_multiple_dimensions(query): + return None + + group_columns = _group_by_source_columns(sql) + if len(group_columns) > 1: + return ( + "Schema grounding failed. The question asks for one grouping " + "dimension, but the generated SQL groups by multiple source " + "columns." + ) + if not group_columns: + return None + + grouped_column_names = {_normalize_identifier(name) for name in group_columns} + scored_columns: list[tuple[int, str]] = [] + for relation in referenced_relations: + for column in schema_details.get(relation) or []: + score = _column_score_for_tokens(column, grouping_tokens) + if score > 0: + scored_columns.append((score, column["name"])) + + if not scored_columns: + return None + + best_score = max(score for score, _ in scored_columns) + grouped_best_score = max( + ( + score + for score, name in scored_columns + if _normalize_identifier(name) in grouped_column_names + ), + default=0, + ) + if best_score > grouped_best_score: + logger.info( + "Generated SQL explicit grouping validation rejected grouping_tokens=%s grouped_columns=%s best_columns=%s", + sorted(grouping_tokens), + sorted(group_columns), + [ + name + for score, name in scored_columns + if score == best_score + ], + ) + return ( + "Schema grounding failed. The question asks for a specific grouping " + "dimension, but the generated SQL groups by a weaker matching " + "column." + ) + + return None + + +def validate_sql_semantic_coverage( + sql: str, + query: str | None, + contexts: list[Any] | None = None, +) -> str | None: + if not sql or not query: + return None + + raw_query_tokens = _fallback_tokens(query) + + schema_details = _extract_schema_details(contexts) + if not schema_details: + return None + + grounding = _extract_sql_grounding(sql) + referenced_relations = { + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + } + if not referenced_relations: + return None + + schema_tokens = set() + for relation in referenced_relations: + columns = schema_details.get(relation) + if columns is not None: + schema_tokens.update(_schema_tokens_for_table(relation, columns)) + + required_tokens = _schema_required_query_tokens( + query, + raw_query_tokens, + schema_details, + ) + unsupported_tokens = _unsupported_query_tokens( + raw_query_tokens, + schema_details, + query=query, + ) + if unsupported_tokens: + return ( + "Schema grounding failed. The retrieved schema metadata does not " + "support these non-operational question term(s): " + f"{', '.join(sorted(unsupported_tokens))}. " + "Select a project with matching schema metadata or ask a supported " + "question." + ) + + missing_concepts = sorted( + token + for token in required_tokens + if not _schema_token_covered(token, schema_tokens) + ) + if missing_concepts: + return ( + "Schema grounding failed. The generated SQL uses verified identifiers, " + "but the selected table or view does not cover these schema-backed " + f"question tokens: {', '.join(missing_concepts)}. Use only schema " + "objects whose metadata supports the requested terms, or return no " + "SQL if the active project does not contain them." + ) + + subject_tokens = _query_subject_schema_tokens( + query, + raw_query_tokens, + schema_details, + ) + if subject_tokens: + identifier_tokens_by_table = { + table_name: _schema_identifier_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + selected_identifier_tokens = set() + for relation in referenced_relations: + selected_identifier_tokens.update( + identifier_tokens_by_table.get(relation, set()) + ) + if any( + _schema_tokens_cover(subject_tokens, tokens) + for tokens in identifier_tokens_by_table.values() + ) and not _schema_tokens_cover(subject_tokens, selected_identifier_tokens): + return ( + "Schema grounding failed. The generated SQL does not use a " + "verified table or column identifier covering the requested " + "primary subject. Choose a retrieved schema object whose declared " + "identifiers support the subject, or return no SQL if none is " + "available." + ) + if ( + not any( + _schema_tokens_cover(subject_tokens, tokens) + for tokens in identifier_tokens_by_table.values() + ) + and any( + _schema_token_match_count(subject_tokens, tokens) > 0 + for tokens in identifier_tokens_by_table.values() + ) + and _schema_token_match_count(subject_tokens, selected_identifier_tokens) + == 0 + ): + return ( + "Schema grounding failed. The generated SQL does not use a " + "verified table or column identifier covering the requested " + "primary subject. Choose a retrieved schema object whose declared " + "identifiers support the subject, or return no SQL if none is " + "available." + ) + + explicit_grouping_tokens = _grouping_phrase_tokens(query) + group_columns = _group_by_source_columns(sql) + if explicit_grouping_tokens and group_columns: + grouping_required_tokens = _schema_derived_query_tokens( + _query_content_tokens(explicit_grouping_tokens), + schema_details, + ) + if grouping_required_tokens: + grouped_column_tokens: set[str] = set() + grouped_column_names = { + _normalize_identifier(name) for name in group_columns + } + for relation in referenced_relations: + for column in schema_details.get(relation, []): + if _normalize_identifier(column["name"]) in grouped_column_names: + grouped_column_tokens.update(_column_business_tokens(column)) + if not _schema_tokens_cover( + grouping_required_tokens, + grouped_column_tokens, + ): + return ( + "Schema grounding failed. The generated SQL groups by a " + "column that does not cover the requested grouping " + "dimension. Choose a verified grouping column matching the " + "question, or return no SQL if none is available." + ) + + if _is_average_metric_intent(raw_query_tokens): + if not re.search(r"(?is)\bAVG\s*\(", sql): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL does not compute AVG over a " + "verified measure." + ) + if re.search(r"(?is)\bCOUNT\s*\(", sql) and not re.search( + r"(?is)\bAVG\s*\(", + sql, + ): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL computes a count." + ) + + if _is_distribution_metric_intent(raw_query_tokens): + if not re.search(r"(?is)\bCOUNT\s*\(", sql) or not re.search( + r"(?is)\bGROUP\s+BY\b", + sql, + ): + return ( + "Schema grounding failed. The question asks for a distribution " + "or breakdown, but the generated SQL does not compute grouped " + "counts." + ) + + if ( + _has_extreme_intent(raw_query_tokens) + and re.search(r"(?is)\bGROUP\s+BY\b", sql) + and _sql_has_aggregate_function(sql) + and not _query_allows_grouped_aggregate(query, raw_query_tokens) + ): + return ( + "Schema grounding failed. The question asks for top or bottom " + "records, but the generated SQL returns a grouped aggregate. Use " + "a row-level ordering unless the question asks for grouping, count, " + "sum, rate, or distribution." + ) + + if ( + _grouping_phrase_tokens(query) + and not _grouping_phrase_has_multiple_dimensions(query) + ): + grouping_error = _validate_explicit_grouping_columns( + sql, + query, + raw_query_tokens, + schema_details, + referenced_relations, + ) + if grouping_error: + return grouping_error + + if _has_missing_value_intent(raw_query_tokens): + target_tokens = _missing_value_target_tokens(query, raw_query_tokens) + for relation in referenced_relations: + columns = schema_details.get(relation) or [] + if not columns or not target_tokens: + continue + checked_columns = _sql_null_checked_columns(sql, columns) + if not checked_columns: + continue + scored_columns = [ + (_column_score_for_tokens(column, target_tokens), column["name"]) + for column in columns + ] + best_score = max((score for score, _ in scored_columns), default=0) + checked_best_score = max( + ( + score + for score, name in scored_columns + if name in checked_columns + ), + default=0, + ) + if best_score > checked_best_score: + return ( + "Schema grounding failed. The question asks for missing " + "values on a specific schema concept, but the generated SQL " + "checks a weaker matching column for null or blank values." + ) + + return _validate_literal_values_against_samples( + sql, + schema_details, + grounding, + query=query, + ) + + +def unsupported_schema_message( + query: str | None, + contexts: list[Any] | None = None, +) -> str | None: + if not query: + return None + query_tokens = _fallback_tokens(query) + schema_details = _extract_schema_details(contexts) + if not schema_details: + return None + required_tokens = _schema_required_query_tokens( + query, + query_tokens, + schema_details, + ) + unsupported_tokens = _unsupported_query_tokens( + query_tokens, + schema_details, + query=query, + ) + if unsupported_tokens: + return ( + "No retrieved table or view in the active project contains verified " + "schema metadata for all requested non-operational term(s): " + f"{', '.join(sorted(unsupported_tokens))}. Select a project with " + "matching fields, add schema descriptions/sample values, or ask a " + "question supported by the selected project's schema." + ) + + table_tokens = _schema_tokens_by_table(schema_details) + if required_tokens and any( + _schema_tokens_cover(required_tokens, tokens) + for tokens in table_tokens.values() + ): + return None + + if not required_tokens and not unsupported_tokens: + return None + detail_tokens = sorted(required_tokens or unsupported_tokens) + return ( + "No retrieved table or view in the active project contains verified " + "schema metadata for all requested non-operational term(s): " + f"{', '.join(detail_tokens)}. Select a project with matching fields, " + "add schema descriptions/sample values, or ask a question supported by " + "the selected project's schema." + ) + + +def unsupported_schema_generation_result( + query: str | None, + contexts: list[Any] | None = None, + data_source: str = "", +) -> dict[str, Any] | None: + message = unsupported_schema_message(query, contexts=contexts) + if not message: + return None + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": message, + "correlation_id": "", + "data_source": data_source, + }, + } + + +def schema_grounding_failure_message( + query: str | None, + contexts: list[Any] | None = None, +) -> str: + unsupported_message = unsupported_schema_message(query, contexts=contexts) + if unsupported_message: + return unsupported_message + + schema_details = _extract_schema_details(contexts) + query_tokens = _fallback_tokens(query) + schema_backed_tokens = _schema_required_query_tokens( + query, + query_tokens, + schema_details, + ) + subject_tokens = _query_subject_schema_tokens( + query, + query_tokens, + schema_details, + ) + if schema_details and schema_backed_tokens: + table_tokens = _schema_tokens_by_table(schema_details) + if not any( + all(_fallback_token_variants(token) & tokens for token in schema_backed_tokens) + for tokens in table_tokens.values() + ): + return ( + "No retrieved table or view in the active project contains one " + "verified schema object covering all requested schema-backed " + f"term(s): {', '.join(sorted(schema_backed_tokens))}. The " + "retrieved schema may contain some terms only on different " + "tables or views; add schema metadata/relationships or ask a " + "question supported by one verified schema object." + ) + + content_tokens = sorted(_query_content_tokens(query_tokens)) + if content_tokens: + return ( + "No grounded SQL could be generated using only verified retrieved " + "tables and columns for the active project and requested term(s): " + f"{', '.join(content_tokens)}. Add matching schema metadata/sample " + "values or ask a question supported by the selected project's schema." + ) + + return ( + "No grounded SQL could be generated using only verified retrieved tables " + "and columns for the active project. Ask a question supported by the " + "selected project's schema." + ) + + +def normalize_sql_with_schema_identifiers( + sql: str, + contexts: list[Any] | None = None, +) -> str: + schema_catalog = _SchemaCatalog.from_contexts(contexts or []) + sql = schema_catalog.normalize_sql(sql) + schema_identifiers = set(_extract_schema_identifiers(contexts)) + identifiers = [ + identifier + for identifier in schema_identifiers + if "." not in identifier and _identifier_needs_quotes(identifier) + ] + sql = _replace_bracket_identifiers(sql, schema_identifiers) + for identifier in sorted(identifiers, key=len, reverse=True): + sql = _replace_identifier_outside_literals(sql, identifier) + return sql + + +def _fallback_tokens(value: Any) -> set[str]: + if value is None: + return set() + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) + tokens = { + token + for token in _FALLBACK_TOKEN.findall(text.lower()) + if token not in _FALLBACK_STOPWORDS + } + return _expand_fallback_token_variants(tokens) + + +def _identifier_tokens(value: Any) -> set[str]: + tokens = _fallback_tokens(value) + for token in list(tokens): + tokens.update(_compound_identifier_tokens(token)) + return _expand_fallback_token_variants(tokens) + + +def _column_business_tokens(column: dict[str, Any]) -> set[str]: + tokens = _identifier_tokens(column["name"]) + tokens.update(column.get("semantic_tokens") or set()) + return tokens + + +def _table_business_tokens( + table_name: str, + columns: list[dict[str, Any]], +) -> set[str]: + tokens = _identifier_tokens(table_name) + for column in columns: + tokens.update(column.get("_table_semantic_tokens") or set()) + return tokens + + +def _data_type_base(data_type: str) -> str: + return data_type.upper().split("(", 1)[0].strip() + + +def _is_numeric_type(data_type: str) -> bool: + return _data_type_base(data_type) in { + "BIGINT", + "DECIMAL", + "DOUBLE", + "FLOAT", + "FLOAT4", + "FLOAT8", + "INT", + "INT2", + "INT4", + "INT8", + "INTEGER", + "NUMERIC", + "REAL", + "SMALLINT", + } + + +def _is_date_type(data_type: str) -> bool: + return _data_type_base(data_type) in { + "DATE", + "DATETIME", + "DATETIME2", + "SMALLDATETIME", + "TIME", + "TIMESTAMP", + "TIMESTAMPTZ", + "TIMESTAMP_LTZ", + "TIMESTAMP_NTZ", + "TIMESTAMP_TZ", + } + + +_RATE_METRIC_TOKENS = {"rate", "ratio", "percent", "percentage"} +_COUNT_METRIC_TOKENS = {"count", "counts", "many", "most", "number", "total"} +_AVERAGE_METRIC_TOKENS = {"average", "avg", "mean"} +_DISTRIBUTION_METRIC_TOKENS = {"distribution", "breakdown"} +_SUM_METRIC_TOKENS = {"sum", "total"} +_MIN_METRIC_TOKENS = {"bottom", "least", "lowest", "min", "minimum", "smallest"} +_MAX_METRIC_TOKENS = {"greatest", "highest", "largest", "max", "maximum", "most", "top"} +_LATEST_METRIC_TOKENS = {"latest", "newest", "recent"} +_NULL_CHECK_TOKENS = {"blank", "empty", "missing", "null"} +_IMPLICIT_TEXT_VALUE_COLUMN_ROLE_TOKENS = { + "account", + "buyer", + "city", + "client", + "company", + "country", + "cust", + "customer", + "entity", + "label", + "market", + "name", + "org", + "organisation", + "organization", + "party", + "person", + "region", + "seller", + "state", + "supplier", + "title", + "vendor", +} +_IMPLICIT_TEXT_VALUE_IDENTIFIER_TOKENS = { + "code", + "id", + "identifier", + "key", + "no", + "num", + "number", + "po", + "ref", + "reference", + "uuid", +} +_GENERIC_SCHEMA_INTENT_TOKENS = { + "a", + "across", + "all", + "an", + "and", + "as", + "ascending", + "associated", + "association", + "average", + "avg", + "between", + "bottom", + "breakdown", + "bucket", + "buckets", + "by", + "compare", + "count", + "counts", + "date", + "day", + "descending", + "distribution", + "each", + "eight", + "eighteen", + "eleven", + "for", + "five", + "from", + "four", + "fourteen", + "group", + "grouped", + "groups", + "has", + "have", + "highest", + "how", + "in", + "is", + "latest", + "least", + "list", + "lowest", + "many", + "max", + "maximum", + "me", + "mean", + "min", + "minimum", + "month", + "monthly", + "most", + "newest", + "nine", + "nineteen", + "number", + "of", + "one", + "ordered", + "pair", + "pairs", + "per", + "please", + "quarter", + "recent", + "record", + "records", + "result", + "results", + "row", + "rows", + "seven", + "seventeen", + "show", + "six", + "sixteen", + "smallest", + "sort", + "sorted", + "sum", + "ten", + "the", + "there", + "thirteen", + "this", + "three", + "to", + "top", + "total", + "twelve", + "twenty", + "two", + "using", + "was", + "week", + "were", + "what", + "when", + "where", + "which", + "who", + "why", + "with", + "without", + "year", +} +_GENERIC_SCHEMA_INTENT_TOKENS.update(_NULL_CHECK_TOKENS) +_GENERIC_SCHEMA_INTENT_TOKENS.update(_MONTH_NAME_TO_NUMBER.keys()) +_FILTER_VALUE_BOUNDARY_TOKENS = _GENERIC_SCHEMA_INTENT_TOKENS | { + "after", + "before", + "column", + "columns", + "during", + "field", + "fields", + "having", + "limit", + "on", + "since", + "until", + "value", + "values", + "where", +} +_COLUMN_MENTION_STOP_TOKENS = _GENERIC_SCHEMA_INTENT_TOKENS | { + "filter", + "filters", + "or", + "use", + "uses", +} + + +def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _RATE_METRIC_TOKENS) + + +def _is_average_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _AVERAGE_METRIC_TOKENS) + + +def _is_distribution_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _DISTRIBUTION_METRIC_TOKENS) + + +def _is_rate_like_column(column: dict[str, str]) -> bool: + return bool(_identifier_tokens(column["name"]) & (_RATE_METRIC_TOKENS | {"score"})) + + +def _is_identifier_like_column(column: dict[str, str]) -> bool: + tokens = _identifier_tokens(column["name"]) + return bool(tokens & {"id", "identifier", "key", "uuid"}) + + +def _quote_joined(identifiers: list[str]) -> str: + return ", ".join(_quote_identifier(identifier) for identifier in identifiers) + + +def _schema_tokens_for_table(table_name: str, columns: list[dict[str, str]]) -> set[str]: + tokens = _table_business_tokens(table_name, columns) + for column in columns: + tokens.update(_column_business_tokens(column)) + tokens.update(_sample_value_tokens(column)) + return tokens + + +def _schema_subject_tokens_for_table( + table_name: str, + columns: list[dict[str, str]], +) -> set[str]: + tokens = _table_business_tokens(table_name, columns) + for column in columns: + tokens.update(_identifier_tokens(column["name"])) + tokens.update(column.get("semantic_tokens") or set()) + return tokens + + +def _schema_identifier_tokens_for_table( + table_name: str, + columns: list[dict[str, str]], +) -> set[str]: + tokens = _identifier_tokens(table_name) + for column in columns: + tokens.update(_identifier_tokens(column["name"])) + return tokens + + +def _schema_tokens_by_table( + schema_details: dict[str, list[dict[str, str]]], +) -> dict[str, set[str]]: + return { + table_name: _schema_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + + +def _schema_subject_tokens_by_table( + schema_details: dict[str, list[dict[str, str]]], +) -> dict[str, set[str]]: + return { + table_name: _schema_subject_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + + +def _schema_token_covered(token: str, schema_tokens: set[str]) -> bool: + return bool(_fallback_token_variants(token) & schema_tokens) + + +def _schema_tokens_cover(required_tokens: set[str], schema_tokens: set[str]) -> bool: + return all(_schema_token_covered(token, schema_tokens) for token in required_tokens) + + +def _schema_token_match_count( + required_tokens: set[str], + schema_tokens: set[str], +) -> int: + return sum(1 for token in required_tokens if _schema_token_covered(token, schema_tokens)) + + +def _schema_derived_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + supported_tokens: set[str] = set() + for token in query_tokens: + if token in _GENERIC_SCHEMA_INTENT_TOKENS or token.isdigit(): + continue + supported_tokens.update(_fallback_token_variants(token) & schema_tokens) + return supported_tokens + + +def _schema_derived_subject_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + schema_tokens_by_table = _schema_subject_tokens_by_table(schema_details) + schema_tokens = set().union(*schema_tokens_by_table.values()) if schema_tokens_by_table else set() + supported_tokens: set[str] = set() + for token in query_tokens: + if token in _GENERIC_SCHEMA_INTENT_TOKENS or token.isdigit(): + continue + supported_tokens.update(_fallback_token_variants(token) & schema_tokens) + return supported_tokens + + +def _schema_required_query_tokens( + query: str | None, + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + required_tokens = _schema_derived_query_tokens(query_tokens, schema_details) + if not required_tokens: + return required_tokens + + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + descriptor_tokens = _schema_adjacent_dimension_descriptor_tokens( + query, + query_tokens, + schema_tokens, + ) + user_value_tokens = _schema_driven_user_value_tokens( + query, + query_tokens, + schema_details, + ) + subject_tokens = _query_subject_schema_tokens( + query, + query_tokens, + schema_details, + ) + return (required_tokens - descriptor_tokens - user_value_tokens) | subject_tokens + + +def _query_subject_schema_tokens( + query: str | None, + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + subject_tokens = _query_subject_content_tokens(query) + if not subject_tokens: + return set() + + return _schema_derived_subject_query_tokens(subject_tokens, schema_details) + + +def _query_subject_unsupported_tokens( + query: str | None, + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + subject_tokens = _query_subject_content_tokens(query) + if not subject_tokens: + return set() + + schema_tokens_by_table = _schema_subject_tokens_by_table(schema_details) + schema_tokens = set().union(*schema_tokens_by_table.values()) if schema_tokens_by_table else set() + all_columns = [ + column + for columns in schema_details.values() + for column in columns + ] + subject_value_tokens = _filter_value_tokens( + _schema_driven_column_value_terms( + _query_subject_text(query), + all_columns, + schema_tokens, + ) + ) + subject_tokens = subject_tokens - subject_value_tokens + return { + token + for token in subject_tokens + if not _schema_token_covered(token, schema_tokens) + } + + +def _schema_adjacent_dimension_descriptor_tokens( + query: str | None, + query_tokens: set[str], + schema_tokens: set[str], +) -> set[str]: + if not query: + return set() + dimension_intent_tokens = _DISTRIBUTION_METRIC_TOKENS | { + "across", + "by", + "each", + "group", + "grouped", + "per", + } + if not ( + query_tokens + & dimension_intent_tokens + ): + return set() + + descriptor_tokens: set[str] = set() + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + for index, match in enumerate(matches): + if index == 0: + continue + previous_tokens = _fallback_tokens( + matches[index - 1].group(0).replace("_", " ") ) + value_tokens = _fallback_tokens(match.group(0).replace("_", " ")) + explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { + "'", + '"', + } + if ( + previous_tokens & schema_tokens + and value_tokens + and not (value_tokens & schema_tokens) + and not explicitly_quoted + ): + descriptor_tokens.update(value_tokens) + return descriptor_tokens + + +def _unsupported_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], + query: str | None = None, +) -> set[str]: + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + descriptor_tokens = _schema_adjacent_dimension_descriptor_tokens( + query, + query_tokens, + schema_tokens, + ) + user_value_tokens = _schema_driven_user_value_tokens( + query, + query_tokens, + schema_details, + ) + unsupported_tokens = { + token + for token in query_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + and not token.isdigit() + and not (_fallback_token_variants(token) & schema_tokens) + and token not in descriptor_tokens + and token not in user_value_tokens + } + unsupported_tokens.update(_query_subject_unsupported_tokens(query, schema_details)) + return unsupported_tokens + + +def _table_covers_requested_concepts( + table_name: str, + columns: list[dict[str, str]], + concept_tokens: set[str], +) -> bool: + required_tokens = { + token + for token in concept_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + } + if not required_tokens: + return True + schema_tokens = _schema_tokens_for_table(table_name, columns) + return _schema_tokens_cover(required_tokens, schema_tokens) + + +def _choose_fallback_table( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], + concept_tokens: set[str] | None = None, + subject_tokens: set[str] | None = None, +) -> tuple[str, list[dict[str, str]]] | None: + concept_tokens = concept_tokens or query_tokens + subject_tokens = subject_tokens or set() + required_tokens = _schema_derived_query_tokens(concept_tokens, schema_details) + identifier_tokens_by_table = { + table_name: _schema_identifier_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + subject_identifier_requires_all = bool(subject_tokens) and any( + _schema_tokens_cover(subject_tokens, tokens) + for tokens in identifier_tokens_by_table.values() + ) + subject_identifier_requires_any = bool(subject_tokens) and not subject_identifier_requires_all and any( + _schema_token_match_count(subject_tokens, tokens) > 0 + for tokens in identifier_tokens_by_table.values() + ) + scored_tables = [] + for table_name, columns in schema_details.items(): + table_tokens = _table_business_tokens(table_name, columns) + table_name_tokens = _identifier_tokens(table_name) + identifier_tokens = identifier_tokens_by_table.get(table_name, set()) + column_token_union: set[str] = set() + sample_token_union: set[str] = set() + numeric_token_matches: set[str] = set() + date_bonus = 0 + table_name_matches = { + variant + for token in query_tokens + for variant in (_fallback_token_variants(token) & table_name_tokens) + } + required_table_name_matches = { + variant + for token in required_tokens + for variant in (_fallback_token_variants(token) & table_name_tokens) + } + score = len(table_name_matches) * 32 + score += len(required_table_name_matches) * 64 + score += len(query_tokens & table_tokens) * 8 + for column in columns: + column_tokens = _column_business_tokens(column) + sample_tokens = _sample_value_tokens(column) + column_token_union.update(column_tokens) + sample_token_union.update(sample_tokens) + if _is_numeric_type(column["data_type"]): + numeric_token_matches.update(query_tokens & column_tokens) + if _is_date_type(column["data_type"]): + date_bonus = 2 + + score += len(query_tokens & column_token_union) * 10 + score += len(query_tokens & sample_token_union) * 6 + score += len(numeric_token_matches) * 2 + score += date_bonus + + table_schema_tokens = _schema_tokens_for_table(table_name, columns) + if required_tokens and not _schema_tokens_cover( + required_tokens, + table_schema_tokens, + ): + continue + if subject_identifier_requires_all and not _schema_tokens_cover( + subject_tokens, + identifier_tokens, + ): + continue + if subject_identifier_requires_any and _schema_token_match_count( + subject_tokens, + identifier_tokens, + ) == 0: + continue + score += _schema_token_match_count(required_tokens, table_schema_tokens) * 20 + score += _schema_token_match_count(subject_tokens, identifier_tokens) * 80 + + if score > 0: + scored_tables.append((score, table_name, columns)) + + if not scored_tables: + return None + + scored_tables.sort(key=lambda item: (-item[0], item[1])) + return scored_tables[0][1], scored_tables[0][2] + + +def _choose_column_by_tokens( + columns: list[dict[str, str]], + required_tokens: set[str], + numeric: bool | None = None, + date: bool | None = None, +) -> str | None: + candidates = [] + for column in columns: + column_tokens = _column_business_tokens(column) + if numeric is True and not _is_numeric_type(column["data_type"]): + continue + if date is True and not ( + _is_date_type(column["data_type"]) + or column_tokens & {"date", "day", "month", "time", "year"} + ): + continue + score = len(required_tokens & column_tokens) * 10 + if date is True and _is_date_type(column["data_type"]): + score += 20 + if required_tokens and required_tokens.issubset(column_tokens): + score += 30 + if score > 0: + candidates.append((score, column["name"])) + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1])) + return candidates[0][1] + + +def _column_score_for_tokens( + column: dict[str, str], + required_tokens: set[str], + numeric: bool | None = None, +) -> int: + column_tokens = _identifier_tokens(column["name"]) + column_tokens.update(column.get("semantic_tokens") or set()) + if numeric is True and not _is_numeric_type(column["data_type"]): + return 0 + score = len(required_tokens & column_tokens) * 10 + if required_tokens and required_tokens.issubset(column_tokens): + score += 30 + return score + + +def _choose_ranked_column_by_tokens( + columns: list[dict[str, str]], + required_tokens: set[str], + numeric: bool | None = None, +) -> dict[str, str] | None: + candidates = [ + (_column_score_for_tokens(column, required_tokens, numeric=numeric), column) + for column in columns + ] + candidates = [(score, column) for score, column in candidates if score > 0] + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1]["name"])) + return candidates[0][1] + + +def _sample_value_tokens(column: dict[str, Any]) -> set[str]: + tokens: set[str] = set() + for value in column.get("sample_values") or []: + tokens.update(_fallback_tokens(value)) + return tokens + + +def _column_supports_filter_values( + column: dict[str, Any], + values: list[str], +) -> bool: + cleaned_values = [value for value in (_clean_filter_value(value) for value in values) if value] + if not cleaned_values: + return True + + column_tokens = _column_business_tokens(column) + sample_tokens = _sample_value_tokens(column) + value_tokens: set[str] = set() + for value in cleaned_values: + value_tokens.update(_fallback_tokens(value)) + + if value_tokens & sample_tokens: + return True + if value_tokens & column_tokens: + return True + return False + + +def _filter_value_tokens(values: list[str]) -> set[str]: + tokens: set[str] = set() + for value in values: + tokens.update(_fallback_tokens(value)) + return tokens + + +def _first_filter_value_search_token(values: list[str]) -> str | None: + tokens: list[str] = [] + seen: set[str] = set() + for value in values: + for token in _FALLBACK_TOKEN.findall(str(value).lower()): + if ( + token in seen + or token in _FALLBACK_STOPWORDS + or token in _GENERIC_SCHEMA_INTENT_TOKENS + or token.isdigit() + ): + continue + tokens.append(token) + seen.add(token) + + longer_tokens = [token for token in tokens if len(token) >= 4] + if longer_tokens: + return longer_tokens[0] + return tokens[0] if tokens else None + + +def _filter_value_search_pattern(values: list[str]) -> str | None: + cleaned_values = [ + value for value in (_clean_filter_value(value) for value in values) if value + ] + if not cleaned_values: + return None + + primary_value = cleaned_values[0] + primary_tokens = [ + token + for token in _FALLBACK_TOKEN.findall(primary_value.lower()) + if token not in _FALLBACK_STOPWORDS + and token not in _GENERIC_SCHEMA_INTENT_TOKENS + and not token.isdigit() + ] + if re.search(r"[.&/()'-]", primary_value) or len(primary_tokens) >= 2: + phrase = re.sub(r"\s+", " ", primary_value.lower()).strip() + if phrase: + return f"%{phrase}%" + + search_token = _first_filter_value_search_token(values) + return f"%{search_token.lower()}%" if search_token else None + + +def _implicit_text_value_predicate( + column: dict[str, str], + values: list[str], +) -> str: + if column.get("sample_values") and ( + _filter_value_tokens(values) & _sample_value_tokens(column) + ): + return _filter_predicate_for_values(column, values) + + search_pattern = _filter_value_search_pattern(values) + if not search_pattern: + return _filter_predicate_for_values(column, values) + + quoted_column = _quote_identifier(column["name"]) + return f"LOWER({quoted_column}) LIKE {_quote_literal(search_pattern)}" + + +def _filter_predicate_for_values( + column: dict[str, str], + values: list[str], +) -> str: + cleaned_values = [ + value for value in (_clean_filter_value(value) for value in values) if value + ] + if not cleaned_values: + return _non_missing_value_predicate(column) + + column_tokens = _column_business_tokens(column) + value_tokens: set[str] = set() + for value in cleaned_values: + value_tokens.update(_fallback_tokens(value)) + + return _value_match_predicate(column, cleaned_values[0], cleaned_values[1:]) + + +def _choose_filter_column_for_values( + columns: list[dict[str, str]], + values: list[str], +) -> dict[str, str] | None: + candidates: list[tuple[int, dict[str, str]]] = [] + + for column in columns: + column_tokens = _column_business_tokens(column) + if not _column_supports_filter_values(column, values): + continue + value_tokens = { + token + for value in values + for token in _fallback_tokens(value) + } + score = len(column_tokens & value_tokens) * 10 + score += len(_sample_value_tokens(column) & value_tokens) * 20 + if score > 0: + candidates.append((score, column)) + + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1]["name"])) + return candidates[0][1] + + +def _choose_temporal_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + candidates = [] + for column in columns: + column_tokens = _column_business_tokens(column) + if not _is_date_type(column["data_type"]): + continue + score = len(query_tokens & column_tokens) * 12 + score += 20 + if score > 0: + candidates.append((score, column["name"])) + + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1])) + return candidates[0][1] + + +def _order_by_phrase_tokens(query: str) -> set[str]: + match = re.search( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+(?P[A-Za-z0-9_ /-]+)", + query, + ) + if not match: + return set() + return _fallback_tokens(match.group("value")) + + +def _choose_order_by_column( + query: str, + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + order_tokens = _order_by_phrase_tokens(query) + if not order_tokens: + return None + + column = _choose_ranked_column_by_tokens(columns, order_tokens) + return column["name"] if column else None + + +def _choose_dimension_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + columns = _choose_dimension_columns(query_tokens, columns, max_columns=1) + return columns[0] if columns else None + + +def _choose_dimension_columns( + query_tokens: set[str], + columns: list[dict[str, str]], + max_columns: int = 3, +) -> list[str]: + name_candidates = [] + semantic_candidates = [] + filtered_query_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS + identifier_requested = bool(filtered_query_tokens & {"id", "identifier", "key", "uuid"}) + for index, column in enumerate(columns): + if _is_identifier_like_column(column) and not identifier_requested: + continue + name_tokens = _identifier_tokens(column["name"]) + numeric_dimension_tokens = {"month", "period", "quarter", "year"} + if _is_numeric_type(column["data_type"]) and not ( + query_tokens & name_tokens & numeric_dimension_tokens + ): + continue + semantic_tokens = set(column.get("semantic_tokens") or set()) + name_score = len(filtered_query_tokens & name_tokens) * 10 + if query_tokens & name_tokens & numeric_dimension_tokens: + name_score += 10 + semantic_score = len(filtered_query_tokens & semantic_tokens) * 3 + if name_score > 0: + name_candidates.append((name_score + semantic_score, index, column["name"])) + elif semantic_score > 0: + semantic_candidates.append((semantic_score, index, column["name"])) + candidates = name_candidates or semantic_candidates + candidates.sort(key=lambda item: (-item[0], item[1])) + return [name for _, _, name in candidates[:max_columns]] + + +def _choose_missing_value_column( + query: str, + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + target_tokens = _missing_value_target_tokens(query, query_tokens) + column = _choose_ranked_column_by_tokens(columns, target_tokens) + if column: + return column + return _choose_ranked_column_by_tokens( + columns, + query_tokens - (_GENERIC_SCHEMA_INTENT_TOKENS | _NULL_CHECK_TOKENS), + ) + + +def _choose_count_subject_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + column = _choose_ranked_column_by_tokens( + columns, + query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS, ) + if column: + return column + for column in columns: + if not _is_numeric_type(column["data_type"]) and not _is_rate_like_column(column): + return column + return columns[0] if columns else None - invalid_columns = set() - for expression in clause_expressions: - for identifier in _iter_unqualified_identifier_candidates(expression): - upper_identifier = identifier.upper() - if ( - upper_identifier in _SQL_RESERVED_WORDS - or upper_identifier in _SQL_FUNCTION_WORDS - or upper_identifier in _SQL_TYPE_WORDS - or upper_identifier in _DATE_PART_WORDS - or identifier in ignored_identifiers - or identifier in valid_columns - ): - continue - invalid_columns.add(identifier) - - if not invalid_columns: - return None - return ( - "Schema grounding failed. The SQL references unqualified columns that " - f"are not present in verified table or view {relation}: " - f"{', '.join(sorted(invalid_columns))}. Use only verified columns: " - f"{', '.join(sorted(valid_columns))}." +def _choose_average_measure_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + filtered_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS + measure_candidates = [ + column for column in columns if not _is_identifier_like_column(column) + ] + column = _choose_ranked_column_by_tokens( + measure_candidates, + filtered_tokens, + numeric=True, ) + if column: + return column + numeric_columns = [ + column + for column in columns + if _is_numeric_type(column["data_type"]) + and not _is_identifier_like_column(column) + ] + return numeric_columns[0] if len(numeric_columns) == 1 else None -def validate_sql_against_contexts( - sql: str, - contexts: list[Any] | None = None, -) -> str | None: - schema_index = _extract_schema_index(contexts) - if not schema_index: - return None +def _is_text_type(data_type: str) -> bool: + return _data_type_base(data_type) in { + "CHAR", + "CHARACTER", + "NCHAR", + "NTEXT", + "NVARCHAR", + "STRING", + "TEXT", + "VARCHAR", + } - valid_relations = set(schema_index) - grounding = _extract_sql_grounding(sql) - cte_names = grounding["cte_names"] - shadowed_relations = sorted(cte_names & valid_relations) - if shadowed_relations: - return ( - "Schema grounding failed. The SQL creates CTEs with names that already " - f"belong to verified schema objects: {', '.join(shadowed_relations)}. " - "Do not create dummy CTEs for schema objects; use the verified tables or views directly." - ) +def _is_boolean_type(data_type: str) -> bool: + return _data_type_base(data_type) in {"BIT", "BOOL", "BOOLEAN"} - invalid_relations = sorted( - { - relation - for relation in grounding["relation_references"] - if relation not in valid_relations and relation not in cte_names - } + +def _is_categorical_value_column(column: dict[str, Any]) -> bool: + return not ( + _is_identifier_like_column(column) + or _is_numeric_type(column["data_type"]) + or _is_date_type(column["data_type"]) + or _is_boolean_type(column["data_type"]) ) - if invalid_relations: - return ( - "Schema grounding failed. The SQL references tables or views that are not " - f"in the retrieved schema for the active question: {', '.join(invalid_relations)}. " - f"Use only verified tables or views: {', '.join(sorted(valid_relations))}." - ) - alias_to_relation = grounding["alias_to_relation"] - for qualifier, column in grounding["qualified_columns"]: - relation = alias_to_relation.get(qualifier) - if not relation or relation in cte_names: - continue - valid_columns = schema_index.get(relation) - if valid_columns is None: + +def _missing_value_predicate(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_text_type(column["data_type"]): + return f"({quoted_column} IS NULL OR {quoted_column} = '')" + return f"{quoted_column} IS NULL" + + +def _non_missing_value_predicate(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_text_type(column["data_type"]): + return f"({quoted_column} IS NOT NULL AND {quoted_column} <> '')" + return f"{quoted_column} IS NOT NULL" + + +def _aggregate_for_measure(measure_column: str) -> tuple[str, str]: + tokens = _identifier_tokens(measure_column) + if tokens & {"rate", "score", "percent", "percentage"}: + return "AVG", "average_value" + return "SUM", "total_value" + + +def _select_listing_columns( + query_tokens: set[str], + columns: list[dict[str, str]], + measure_column: str | None = None, + date_column: str | None = None, + max_columns: int = 6, +) -> list[str]: + scored_columns: list[tuple[int, int, str]] = [] + for index, column in enumerate(columns): + name = column["name"] + tokens = _column_business_tokens(column) + score = len(query_tokens & tokens) * 10 + if name == date_column: + score += 35 + if name == measure_column: + score += 12 + if score > 0: + scored_columns.append((score, index, name)) + + scored_columns.sort(key=lambda item: (-item[0], item[1])) + selected = [] + for _, _, name in scored_columns: + if name not in selected: + selected.append(name) + if len(selected) >= max_columns: + break + if not selected and columns: + selected = [column["name"] for column in columns[:max_columns]] + return selected + + +def _include_required_listing_columns( + selected_columns: list[str], + required_columns: list[str], + max_columns: int, +) -> list[str]: + selected = list(selected_columns) + for column_name in reversed(required_columns): + if column_name in selected: continue - if column not in valid_columns: - return ( - "Schema grounding failed. The SQL references column " - f"{qualifier}.{column}, but column {column} is not present in verified " - f"table or view {relation}. Use only verified columns: " - f"{', '.join(sorted(valid_columns))}." - ) + selected.insert(0, column_name) + return selected[:max_columns] - unqualified_column_error = _validate_unqualified_columns_for_single_relation( - sql, - schema_index, - grounding, + +def _fallback_limit(query: str) -> int | None: + match = re.search(r"(?i)\btop\s+(\d+)\b", query) + if match: + return int(match.group(1)) + + word_numbers = { + "one": 1, + "two": 2, + "three": 3, + "four": 4, + "five": 5, + "six": 6, + "seven": 7, + "eight": 8, + "nine": 9, + "ten": 10, + "eleven": 11, + "twelve": 12, + "thirteen": 13, + "fourteen": 14, + "fifteen": 15, + "sixteen": 16, + "seventeen": 17, + "eighteen": 18, + "nineteen": 19, + "twenty": 20, + } + word_match = re.search( + r"(?i)\btop\s+(" + + "|".join(re.escape(word) for word in word_numbers) + + r")\b", + query, ) - if unqualified_column_error: - return unqualified_column_error + return word_numbers[word_match.group(1).lower()] if word_match else None + +def _fallback_month_filter(query: str) -> tuple[int, int] | None: + tokens = _fallback_tokens(query) + for month_name, month_number in _MONTH_NAME_TO_NUMBER.items(): + if month_name in tokens: + return datetime.now(timezone.utc).year, month_number return None -def validate_sql_semantic_coverage( - sql: str, - query: str | None, - contexts: list[Any] | None = None, -) -> str | None: - if not sql or not query: - return None +def _grouping_phrase_tokens(query: str) -> set[str]: + phrase = _grouping_phrase_text(query) + return _fallback_tokens(phrase) if phrase else set() - raw_query_tokens = _fallback_tokens(query) - query_tokens = _expanded_fallback_query_tokens(query) - concepts = _requested_business_concepts(raw_query_tokens) - if not concepts: - return None - schema_details = _extract_schema_details(contexts) - if not schema_details: - return None +def _selected_dimensions_cover_grouping_tokens( + dimension_columns: list[str], + grouping_tokens: set[str], + columns: list[dict[str, Any]], + schema_details: dict[str, list[dict[str, str]]], +) -> bool: + grouping_required_tokens = _schema_derived_query_tokens( + _query_content_tokens(grouping_tokens), + schema_details, + ) + if not grouping_required_tokens: + return True - grounding = _extract_sql_grounding(sql) - referenced_relations = { - relation - for relation in grounding["relation_references"] - if relation not in grounding["cte_names"] + selected_column_names = { + _normalize_identifier(column_name) for column_name in dimension_columns } - if not referenced_relations: - return None + selected_tokens: set[str] = set() + for column in columns: + if _normalize_identifier(column["name"]) in selected_column_names: + selected_tokens.update(_column_business_tokens(column)) - schema_tokens = set() - for relation in referenced_relations: - columns = schema_details.get(relation) - if columns is not None: - schema_tokens.update(_schema_tokens_for_table(relation, columns)) + return _schema_tokens_cover(grouping_required_tokens, selected_tokens) - if not schema_tokens: - return None - missing_concepts = [ - label for label, concept_tokens in concepts if not schema_tokens & concept_tokens - ] - if not missing_concepts: - if _is_failure_count_intent(raw_query_tokens, query_tokens): - if not re.search(r"(?is)\bCOUNT\s*\(", sql): - return ( - "Schema grounding failed. The question asks for a count of " - "failure records, but the generated SQL does not compute a " - "COUNT aggregate. Use a verified failure-record column/table " - "and group by the requested dimension, or return no SQL if " - "the active project does not contain it." - ) - for relation in referenced_relations: - for column in schema_details.get(relation, []): - if _is_rate_like_column(column) and _sql_mentions_identifier( - sql, column["name"] - ): - return ( - "Schema grounding failed. The question asks for a " - "count of failure records, but the generated SQL uses " - f"rate-like column {column['name']}. Use COUNT over a " - "verified failure occurrence field instead, or return " - "no SQL if the active project does not contain one." - ) - return None +def _grouping_phrase_has_multiple_dimensions(query: str) -> bool: + phrase = _grouping_phrase_text(query) + return bool(phrase and re.search(r"(?i)(?:,|/|\band\b|\bor\b|\bpairs?\b)", phrase)) - return ( - "Schema grounding failed. The generated SQL uses verified identifiers, " - "but the selected table or view does not contain verified fields for the " - f"requested business concept(s): {', '.join(missing_concepts)}. Use only " - "schema objects whose table or column names explicitly support those " - "concepts, or return no SQL if the active project does not contain them." + +def _grouping_phrase_text(query: str) -> str | None: + query = re.sub( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", + "", + query, + ) + match = re.search( + r"""(?ixs) + \b(?:grouped\s+by|group\s+by|by|across|per)\s+ + (?P[A-Za-z0-9_ /-]+?) + (?= + \s+\b(?:after|before|during|for|from|having|in|limit|on|order(?:ed)?\s+by|since|sort(?:ed)?\s+by|where|with)\b(?!-) + |[?.!,;:] + |$ + ) + """, + query, ) + if not match: + return None + return match.group("value") -def unsupported_schema_message( - query: str | None, - contexts: list[Any] | None = None, -) -> str | None: - if not query: - return None - query_tokens = _fallback_tokens(query) - concepts = _requested_business_concepts(query_tokens) - if not concepts: - return None - schema_details = _extract_schema_details(contexts) - if not schema_details: - return None - if any( - _table_covers_requested_concepts(table_name, columns, query_tokens) - for table_name, columns in schema_details.items() - ): - return None - concept_labels = ", ".join(label for label, _ in concepts) - return ( - "No retrieved table or view in the active project contains verified " - "fields for all requested business concept(s): " - f"{concept_labels}. Select a project with those fields or ask a question " - "supported by the selected project's schema." - ) +def _current_year_where_clause( + date_column: str | None, + columns: list[dict[str, str]], + query_tokens: set[str], +) -> str: + if not date_column or not {"this", "year"}.issubset(query_tokens): + return "" + column = next((column for column in columns if column["name"] == date_column), None) + if not column: + return "" -def unsupported_schema_generation_result( - query: str | None, - contexts: list[Any] | None = None, - data_source: str = "", -) -> dict[str, Any] | None: - message = unsupported_schema_message(query, contexts=contexts) - if not message: - return None - return { - "valid_generation_result": {}, - "invalid_generation_result": { - "sql": "", - "original_sql": "", - "type": "NO_RELEVANT_SQL", - "error": message, - "correlation_id": "", - "data_source": data_source, - }, - } + current_year = datetime.now(timezone.utc).year + quoted_column = _quote_identifier(date_column) + if _is_numeric_type(column["data_type"]) or _identifier_tokens(date_column) & {"year"}: + return f"\nWHERE {quoted_column} = {current_year}" + if _is_date_type(column["data_type"]) or _identifier_tokens(date_column) & { + "date", + "day", + "month", + "time", + }: + return f"\nWHERE {_date_part_expression(date_column, 'YEAR')} = {current_year}" + return "" -def normalize_sql_with_schema_identifiers( - sql: str, - contexts: list[Any] | None = None, +def _quote_literal(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def _value_match_predicate( + column: dict[str, str] | str, + value: str, + alternate_values: list[str] | None = None, ) -> str: - schema_identifiers = set(_extract_schema_identifiers(contexts)) - identifiers = [ - identifier - for identifier in schema_identifiers - if "." not in identifier and _identifier_needs_quotes(identifier) - ] - sql = _replace_bracket_identifiers(sql, schema_identifiers) - for identifier in sorted(identifiers, key=len, reverse=True): - sql = _replace_identifier_outside_literals(sql, identifier) - return sql + column_name = column["name"] if isinstance(column, dict) else column + quoted_column = _quote_identifier(column_name) + values = [] + for candidate in [value] + (alternate_values or []): + cleaned = _clean_filter_value(candidate) + if cleaned and cleaned.lower() not in {item.lower() for item in values}: + values.append(cleaned) + + if not values: + return f"{quoted_column} IS NOT NULL" + + if isinstance(column, dict) and ( + _is_text_type(column["data_type"]) or _is_categorical_value_column(column) + ): + lowered_values = [_quote_literal(candidate.lower()) for candidate in values] + if len(lowered_values) == 1: + return f"LOWER({quoted_column}) = {lowered_values[0]}" + return f"LOWER({quoted_column}) IN ({', '.join(lowered_values)})" + return f"{quoted_column} = {_quote_literal(values[0])}" -def _fallback_tokens(value: Any) -> set[str]: + +def _clean_filter_value(value: str | None) -> str | None: if value is None: - return set() - text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) - tokens = { + return None + value = value.strip(" \t\r\n'\"`.,;:()[]{}") + return value or None + + +def _query_content_tokens(query_tokens: set[str]) -> set[str]: + return { token - for token in _FALLBACK_TOKEN.findall(text.lower()) - if token not in _FALLBACK_STOPWORDS + for token in query_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + and token not in _MONTH_NAME_TO_NUMBER + and not token.isdigit() } - return _expand_fallback_token_aliases(tokens) -def _column_business_tokens(column: dict[str, Any]) -> set[str]: - tokens = _fallback_tokens(column["name"]) - tokens.update(column.get("semantic_tokens") or set()) - return tokens +def _query_subject_text(query: str | None) -> str: + if not query: + return "" + return re.split( + r"""(?ix) + \b(?: + group(?:ed)?\s+by + |break(?:down)?\s+by + |for\s+each + |across + |after + |before + |by + |called + |contains + |containing + |during + |for + |from + |having + |in + |like + |matching + |named + |on + |per + |since + |until + |where + |with + )\b + """, + query, + maxsplit=1, + )[0] + + +def _query_subject_content_tokens(query: str | None) -> set[str]: + subject_text = _query_subject_text(query) + return _query_content_tokens(_fallback_tokens(subject_text)) + + +def _has_grouping_intent(query: str, query_tokens: set[str]) -> bool: + query_without_ordering = re.sub( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", + "", + query, + ) + return bool( + _is_distribution_metric_intent(query_tokens) + or query_tokens & {"group", "grouped", "per"} + or re.search(r"(?i)\bby\s+[A-Za-z0-9_ -]+\b", query_without_ordering) + ) -def _table_business_tokens( - table_name: str, - columns: list[dict[str, Any]], -) -> set[str]: - tokens = _fallback_tokens(table_name) - for column in columns: - tokens.update(column.get("_table_semantic_tokens") or set()) - return tokens +def _has_count_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _COUNT_METRIC_TOKENS) -def _expanded_fallback_query_tokens(query: str) -> set[str]: - tokens = _fallback_tokens(query) - if tokens & {"revenue", "sale", "sales", "trend", "trends"}: - tokens.update({"amount", "date", "intake", "revenue", "sales", "value"}) - if tokens & {"order", "orders"}: - tokens.update({"amount", "customer", "date", "ord", "order", "value"}) - if tokens & {"invoice", "invoices"}: - tokens.update({"amount", "currency", "date", "invoice", "supplier"}) - if tokens & {"batch", "batches"}: - tokens.update({"batch", "board", "defect", "inspection", "rate", "supplier"}) - if tokens & {"repair", "repairs"}: - tokens.update({"date", "failure", "log", "priority", "progress", "repair", "status"}) - if tokens & {"failure", "failures", "defect", "defects"}: - tokens.update({"code", "defect", "failure", "severity", "status", "type"}) - if tokens & {"material", "materials"}: - tokens.update({"item", "material", "part"}) - if tokens & {"location", "locations"}: - tokens.update({"area", "location", "site"}) - if tokens & {"month", "monthly", "july"}: - tokens.update({"date", "day", "month", "time", "year"}) - elif tokens & {"latest", "trend", "trends", "year"}: - tokens.update({"date", "day", "time", "year"}) - if "business" in tokens and "unit" in tokens: - tokens.update({"account", "bu", "business", "company", "division", "unit"}) - return tokens +def _has_sum_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _SUM_METRIC_TOKENS) -def _is_numeric_type(data_type: str) -> bool: - return data_type.upper() in { - "BIGINT", - "DECIMAL", - "DOUBLE", - "FLOAT", - "FLOAT4", - "FLOAT8", - "INT", - "INT2", - "INT4", - "INT8", - "INTEGER", - "NUMERIC", - "REAL", - "SMALLINT", - } +def _has_extreme_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & (_MAX_METRIC_TOKENS | _MIN_METRIC_TOKENS)) -def _is_date_type(data_type: str) -> bool: - return data_type.upper() in { - "DATE", - "DATETIME", - "DATETIME2", - "SMALLDATETIME", - "TIME", - "TIMESTAMP", - "TIMESTAMPTZ", - "TIMESTAMP_LTZ", - "TIMESTAMP_NTZ", - "TIMESTAMP_TZ", - } +def _has_latest_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _LATEST_METRIC_TOKENS) -_RATE_METRIC_TOKENS = {"rate", "ratio", "percent", "percentage"} -_COUNT_METRIC_TOKENS = {"count", "many", "most", "number", "total"} -_PRIORITY_VALUE_ALIASES = { - "urgent": "urgent", - "critical": "critical", - "high": "high", - "medium": "medium", - "normal": "normal", - "low": "low", -} -_PRIORITY_ORDER = [ - ("critical", 6), - ("urgent", 6), - ("blocker", 6), - ("high", 5), - ("major", 5), - ("medium", 4), - ("normal", 4), - ("minor", 3), - ("low", 2), -] +def _has_missing_value_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _NULL_CHECK_TOKENS) -def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: - return bool(raw_query_tokens & _RATE_METRIC_TOKENS) +def _sort_direction_for_query(query_tokens: set[str]) -> str: + return "ASC" if query_tokens & _MIN_METRIC_TOKENS else "DESC" -def _is_failure_count_intent( - raw_query_tokens: set[str], + +def _choose_numeric_measure_column( query_tokens: set[str], -) -> bool: - return ( - bool(raw_query_tokens & {"failure", "failed", "defect"}) - and "failure" in query_tokens - and not _is_rate_metric_intent(raw_query_tokens) - and ( - bool(raw_query_tokens & _COUNT_METRIC_TOKENS) - or bool(raw_query_tokens & {"top", "highest", "lowest", "bottom"}) - ) + columns: list[dict[str, str]], +) -> dict[str, str] | None: + content_tokens = _query_content_tokens(query_tokens) + measure_candidates = [ + column for column in columns if not _is_identifier_like_column(column) + ] + column = _choose_ranked_column_by_tokens( + measure_candidates, + content_tokens, + numeric=True, ) + if column: + return column - -def _has_board_model_intent(query_tokens: set[str]) -> bool: - return {"board", "model"}.issubset(query_tokens) - - -def _is_rate_like_column(column: dict[str, str]) -> bool: - return bool(_fallback_tokens(column["name"]) & (_RATE_METRIC_TOKENS | {"score"})) + numeric_columns = [ + column + for column in columns + if _is_numeric_type(column["data_type"]) + and not _is_identifier_like_column(column) + ] + if len(numeric_columns) == 1: + return numeric_columns[0] + return None -def _quote_joined(identifiers: list[str]) -> str: - return ", ".join(_quote_identifier(identifier) for identifier in identifiers) +def _sample_value_filters( + query_tokens: set[str], + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + filters: list[tuple[dict[str, Any], list[str]]] = [] + consumed_tokens: set[str] = set() + content_tokens = _query_content_tokens(query_tokens) + if not content_tokens: + return filters + for column in columns: + matches: list[str] = [] + for value in column.get("sample_values") or []: + cleaned_value = _clean_filter_value(str(value)) + if not cleaned_value: + continue + value_tokens = _fallback_tokens(cleaned_value) + if not value_tokens or not value_tokens <= content_tokens: + continue + if value_tokens <= consumed_tokens: + continue + if cleaned_value.lower() not in {item.lower() for item in matches}: + matches.append(cleaned_value) + consumed_tokens.update(value_tokens) + if matches: + filters.append((column, matches)) -def _requested_business_concepts(query_tokens: set[str]) -> list[tuple[str, set[str]]]: - concepts: list[tuple[str, set[str]]] = [] - specs = [ - ( - "failure/defect", - {"failure", "failed", "defect"}, - {"failure", "failed", "defect"}, - ), - ("repair", {"repair"}, {"repair"}), - ("material", {"material"}, {"material", "part"}), - ("location", {"location"}, {"location", "site", "area"}), - ("customer", {"customer"}, {"customer", "cust"}), - ("supplier/vendor", {"supplier", "vendor"}, {"supplier", "vendor"}), - ("technician", {"technician", "tech"}, {"technician", "tech"}), - ("product", {"product"}, {"product", "prod", "item", "material"}), - ( - "priority/severity", - {"critical", "priority", "severity"}, - {"priority", "severity"}, - ), - ("status", {"status"}, {"status"}), - ("order", {"order"}, {"order", "ord"}), - ] - for label, triggers, schema_tokens in specs: - if query_tokens & triggers: - concepts.append((label, schema_tokens)) - if {"board", "model"}.issubset(query_tokens): - concepts.append(("board model", {"board", "model"})) - if {"business", "unit"}.issubset(query_tokens): - concepts.append(("business unit", {"business", "unit", "bu", "division"})) - return concepts + return filters -def _schema_tokens_for_table(table_name: str, columns: list[dict[str, str]]) -> set[str]: - tokens = _table_business_tokens(table_name, columns) +def _mentioned_text_columns_for_query( + query_tokens: set[str], + columns: list[dict[str, Any]], +) -> list[dict[str, Any]]: + mentioned_columns = [] for column in columns: - tokens.update(_column_business_tokens(column)) - return tokens + if not _is_categorical_value_column(column): + continue + column_tokens = _column_business_tokens(column) + if column_tokens & _IMPLICIT_TEXT_VALUE_IDENTIFIER_TOKENS: + continue + if any(_fallback_token_variants(token) & column_tokens for token in query_tokens): + mentioned_columns.append(column) + return mentioned_columns -def _table_covers_requested_concepts( - table_name: str, - columns: list[dict[str, str]], - concept_tokens: set[str], -) -> bool: - concepts = _requested_business_concepts(concept_tokens) - if not concepts: - return True - schema_tokens = _schema_tokens_for_table(table_name, columns) - return all(schema_tokens & concept_tokens for _, concept_tokens in concepts) +def _matched_column_query_tokens( + query_tokens: set[str], + column: dict[str, Any], +) -> set[str]: + column_tokens = _column_business_tokens(column) + matched_tokens: set[str] = set() + for token in query_tokens: + matched_tokens.update(_fallback_token_variants(token) & column_tokens) + return matched_tokens -def _choose_fallback_table( +def _matched_column_identifier_query_tokens( query_tokens: set[str], - schema_details: dict[str, list[dict[str, str]]], - concept_tokens: set[str] | None = None, -) -> tuple[str, list[dict[str, str]]] | None: - concept_tokens = concept_tokens or query_tokens - rate_metric_intent = _is_rate_metric_intent(concept_tokens) - failure_count_intent = _is_failure_count_intent(concept_tokens, query_tokens) - board_model_intent = _has_board_model_intent(query_tokens) or _has_board_model_intent( - concept_tokens - ) - scored_tables = [] - for table_name, columns in schema_details.items(): - table_tokens = _table_business_tokens(table_name, columns) - column_token_union = set() - has_numeric_sales_measure = False - has_date_capable_column = False - score = len(query_tokens & table_tokens) * 8 - for column in columns: - column_tokens = _column_business_tokens(column) - column_token_union.update(column_tokens) - if _is_numeric_type(column["data_type"]) and column_tokens & { - "amount", - "intake", - "revenue", - "sales", - "value", - }: - has_numeric_sales_measure = True - if _is_date_type(column["data_type"]) or column_tokens & { - "date", - "day", - "month", - "time", - "year", - }: - has_date_capable_column = True - score += len(query_tokens & column_tokens) * 10 - if _is_numeric_type(column["data_type"]): - score += len( - query_tokens - & column_tokens - & { - "amount", - "cost", - "count", - "margin", - "quantity", - "rate", - "score", - "value", - } - ) * 4 - if _is_date_type(column["data_type"]): - score += ( - len(query_tokens & {"date", "month", "year", "july", "trend", "trends"}) - * 4 - ) - - if not _table_covers_requested_concepts(table_name, columns, concept_tokens): + column: dict[str, Any], +) -> set[str]: + column_tokens = _identifier_tokens(column["name"]) + matched_tokens: set[str] = set() + for token in query_tokens: + matched_tokens.update(_fallback_token_variants(token) & column_tokens) + return matched_tokens + + +def _dedupe_filter_values(values: list[str]) -> list[str]: + deduped: list[str] = [] + seen: set[str] = set() + for value in values: + cleaned = _clean_filter_value(value) + if not cleaned: + continue + key = cleaned.lower() + if key in seen: continue + deduped.append(cleaned) + seen.add(key) + return deduped - if board_model_intent and rate_metric_intent and query_tokens & { - "defect", - "failure", - }: - if not {"board", "model"}.issubset(column_token_union): - continue - rate_column = _choose_column_by_tokens( - columns, - {"defect", "rate"}, - numeric=True, - ) - if not rate_column: - continue - score += 130 - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - score += ( - len(column_token_union & {"amount", "intake", "revenue", "sales", "value"}) - * 10 - ) - score += len(column_token_union & {"date", "month", "time", "year"}) * 5 - if not has_numeric_sales_measure: - continue - score += 50 - if query_tokens & {"year", "month", "monthly", "trend", "trends"}: - if not has_date_capable_column: - continue - score += 30 - if not query_tokens & { - "claim", - "claims", - "customs", - "duty", - "import", - "refund", - "tariff", - }: - table_and_columns = table_tokens | column_token_union - customs_matches = table_and_columns & { - "claim", - "claims", - "customs", - "duty", - "import", - "refund", - "tariff", - "tariffs", - } - if customs_matches and not table_and_columns & {"revenue", "sale", "sales"}: - continue - score -= len(customs_matches) * 40 +def _token_matches_column(token: str, column_tokens: set[str]) -> bool: + meaningful_column_tokens = { + token + for token in column_tokens + if len(token) > 1 and token not in _COLUMN_MENTION_STOP_TOKENS + } + return bool(_fallback_token_variants(token) & meaningful_column_tokens) - if "customer" in query_tokens: - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"customer", "cust"}: - if "missing" in query_tokens: - continue - score -= 80 - else: - score += 45 - if {"business", "unit"}.issubset(query_tokens): - if "bu" in column_token_union: - score += 60 - elif {"business", "unit"} <= column_token_union: - score += 45 +def _is_filter_value_boundary_token( + token: str, + *, + schema_tokens: set[str], + current_column_tokens: set[str], + allow_schema_token_overlap: bool = False, +) -> bool: + value_tokens = _fallback_tokens(token) + if not value_tokens: + return True + if value_tokens & _FILTER_VALUE_BOUNDARY_TOKENS: + return True + if ( + not allow_schema_token_overlap + and value_tokens & schema_tokens + and not (value_tokens & current_column_tokens) + ): + return True + return False - if "failure" in query_tokens and ( - query_tokens & {"location", "material", "technician", "tech"} - or board_model_intent - ): - table_and_columns = table_tokens | column_token_union - if not table_and_columns & {"failure", "failed", "defect"}: - continue - if board_model_intent: - if not {"board", "model"}.issubset(column_token_union): - continue - if failure_count_intent and not _choose_count_subject_column( - {"failure"}, - columns, - ): - continue - score += 100 - if "location" in query_tokens: - if "location" not in column_token_union: - continue - score += 90 - if "material" in query_tokens: - if not column_token_union & {"material", "part"}: - continue - score += 90 - if query_tokens & {"technician", "tech"}: - if not column_token_union & {"technician", "tech"}: - continue - score += 90 - - if query_tokens & {"order", "orders"}: - table_and_columns = table_tokens | column_token_union - explicit_order_support = table_and_columns & {"ord", "order", "orders"} - order_support = table_and_columns & { - "amount", - "customer", - "intake", - "item", - "ord", - "order", - "orders", - "sales", - "value", - } - if not order_support: - continue - if explicit_order_support: - score += 80 - else: - score -= 60 - if query_tokens & {"batch", "batches"} and {"defect", "rate"}.issubset( - column_token_union - ): - score += 40 - if {"material", "location"}.issubset(query_tokens) and { - "material", - "location", - }.issubset(column_token_union): - score += 40 - if ( - query_tokens & {"repair", "repairs"} - and (table_tokens | column_token_union) & {"repair", "fix"} - ): - score += 55 - if concept_tokens & {"critical", "priority", "severity"}: - if not _choose_priority_column(columns): - continue - score += 55 - if concept_tokens & {"status"}: - if not column_token_union & {"status", "state", "progress"}: - continue - score += 45 - if concept_tokens & {"latest", "recent"}: - if not has_date_capable_column: - continue - score += 35 +def _schema_driven_column_value_filters( + query: str | None, + columns: list[dict[str, Any]], + schema_tokens: set[str], +) -> list[tuple[dict[str, Any], list[str]]]: + if not query: + return [] - if score > 0: - scored_tables.append((score, table_name, columns)) + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + if not matches: + return [] - if not scored_tables: - return None + filters: list[tuple[dict[str, Any], list[str]]] = [] + + def add_value(column: dict[str, Any], parts: list[str]) -> bool: + if not parts: + return False + cleaned = _clean_filter_value(" ".join(parts)) + if not cleaned: + return False + value_tokens = _fallback_tokens(cleaned) + if not value_tokens or value_tokens & _FILTER_VALUE_BOUNDARY_TOKENS: + return False + if value_tokens <= _column_business_tokens(column): + return False + if all(token.isdigit() for token in value_tokens): + return False + + for existing_column, existing_values in filters: + if existing_column["name"] != column["name"]: + continue + if cleaned.lower() not in {value.lower() for value in existing_values}: + existing_values.append(cleaned) + return True - scored_tables.sort(key=lambda item: (-item[0], item[1])) - return scored_tables[0][1], scored_tables[0][2] + filters.append((column, [cleaned])) + return True + def add_preceding_values( + column: dict[str, Any], + start_index: int, + column_tokens: set[str], + ) -> bool: + added = False + preceding_parts: list[str] = [] + cursor = start_index + while cursor >= 0 and len(preceding_parts) < 4: + candidate = matches[cursor].group(0).replace("_", " ") + candidate_tokens = _fallback_tokens(candidate) + if ( + candidate_tokens + and candidate_tokens <= schema_tokens + and not (candidate_tokens & column_tokens) + ): + break + if _is_filter_value_boundary_token( + candidate, + schema_tokens=schema_tokens, + current_column_tokens=column_tokens, + allow_schema_token_overlap=not preceding_parts, + ): + break + preceding_parts.insert(0, candidate) + cursor -= 1 + + boundary = matches[cursor].group(0).lower() if cursor >= 0 else "" + + if boundary in {"and", "or"}: + alternate_parts: list[str] = [] + cursor -= 1 + while cursor >= 0 and len(alternate_parts) < 4: + candidate = matches[cursor].group(0).replace("_", " ") + candidate_tokens = _fallback_tokens(candidate) + if ( + candidate_tokens + and candidate_tokens <= schema_tokens + and not (candidate_tokens & column_tokens) + ): + break + if _is_filter_value_boundary_token( + candidate, + schema_tokens=schema_tokens, + current_column_tokens=column_tokens, + allow_schema_token_overlap=not alternate_parts, + ): + break + alternate_parts.insert(0, candidate) + cursor -= 1 + added = add_value(column, alternate_parts) or added + + if boundary not in _OPEN_FILTER_VALUE_INTRODUCERS: + added = add_value(column, preceding_parts) or added + + return added + + def add_following_values( + column: dict[str, Any], + start_index: int, + column_tokens: set[str], + *, + allow_initial_schema_overlap: bool = True, + stop_schema_overlap: bool = False, + ) -> None: + following_parts: list[str] = [] + cursor = start_index + while cursor < len(matches) and len(following_parts) < 6: + candidate = matches[cursor].group(0).replace("_", " ") + candidate_tokens = _fallback_tokens(candidate) + if ( + stop_schema_overlap + and + candidate_tokens + and candidate_tokens <= schema_tokens + and not (candidate_tokens & column_tokens) + ): + break + if _is_filter_value_boundary_token( + candidate, + schema_tokens=schema_tokens, + current_column_tokens=column_tokens, + allow_schema_token_overlap=allow_initial_schema_overlap + and not following_parts, + ): + if candidate.lower() in {"and", "or"} and following_parts: + add_value(column, following_parts) + following_parts = [] + cursor += 1 + continue + break + following_parts.append(candidate) + cursor += 1 + add_value(column, following_parts) -def _choose_column_by_tokens( - columns: list[dict[str, str]], - required_tokens: set[str], - numeric: bool | None = None, - date: bool | None = None, -) -> str | None: - candidates = [] for column in columns: - column_tokens = _column_business_tokens(column) - if numeric is True and not _is_numeric_type(column["data_type"]): - continue - if date is True and not ( - _is_date_type(column["data_type"]) - or column_tokens & {"date", "day", "month", "time", "year"} + if ( + _is_numeric_type(column["data_type"]) + or _is_date_type(column["data_type"]) + or _is_boolean_type(column["data_type"]) + or _is_identifier_like_column(column) ): continue - score = len(required_tokens & column_tokens) * 10 - if date is True and _is_date_type(column["data_type"]): - score += 20 - if required_tokens and required_tokens.issubset(column_tokens): - score += 30 - if score > 0: - candidates.append((score, column["name"])) - if not candidates: - return None - candidates.sort(key=lambda item: (-item[0], item[1])) - return candidates[0][1] + column_tokens = _identifier_tokens(column["name"]) + if not column_tokens: + continue -def _column_score_for_tokens( - column: dict[str, str], - required_tokens: set[str], - numeric: bool | None = None, -) -> int: - column_tokens = _fallback_tokens(column["name"]) - column_tokens.update(column.get("semantic_tokens") or set()) - if numeric is True and not _is_numeric_type(column["data_type"]): - return 0 - score = len(required_tokens & column_tokens) * 10 - if required_tokens and required_tokens.issubset(column_tokens): - score += 30 - if column["name"].lower() == "bu" and {"business", "unit"} & required_tokens: - score += 60 - if column["name"].lower() in {"custno", "customer_id", "customerid"} and { - "customer", - "number", - } & required_tokens: - score += 30 - if column["name"].lower() in {"custname", "customer_name", "customer"} and { - "customer", - "name", - } & required_tokens: - score += 35 - if column["name"].lower() in {"ordno", "order_no", "order_number", "sales_order_number"} and { - "order", - "number", - } & required_tokens: - score += 35 - return score + for index, match in enumerate(matches): + raw_token = match.group(0).replace("_", " ") + if not _token_matches_column(raw_token, column_tokens): + continue + + next_index = index + 1 + while next_index < len(matches) and _token_matches_column( + matches[next_index].group(0).replace("_", " "), + column_tokens, + ): + next_index += 1 + previous_index = index - 1 + while previous_index >= 0 and _token_matches_column( + matches[previous_index].group(0).replace("_", " "), + column_tokens, + ): + previous_index -= 1 -def _choose_ranked_column_by_tokens( - columns: list[dict[str, str]], - required_tokens: set[str], - numeric: bool | None = None, -) -> dict[str, str] | None: - candidates = [ - (_column_score_for_tokens(column, required_tokens, numeric=numeric), column) - for column in columns - ] - candidates = [(score, column) for score, column in candidates if score > 0] - if not candidates: - return None - candidates.sort(key=lambda item: (-item[0], item[1]["name"])) - return candidates[0][1] + preceding_added = add_preceding_values( + column, + previous_index, + column_tokens, + ) + add_following_values( + column, + next_index, + column_tokens, + allow_initial_schema_overlap=not preceding_added, + stop_schema_overlap=( + previous_index >= 0 + and matches[previous_index].group(0).lower() + in {"across", "by", "per"} + ), + ) + return filters -def _choose_dimension_column( - query_tokens: set[str], - columns: list[dict[str, str]], -) -> str | None: - dimension_specs = [ - ({"board", "model"}, {"board", "model"}), - ({"business", "unit"}, {"business", "unit", "bu", "division", "company"}), - ({"customer"}, {"customer", "cust", "name"}), - ({"supplier"}, {"supplier", "vendor", "name"}), - ({"product"}, {"product", "prod", "item", "material", "name"}), - ({"salesperson", "representative"}, {"salesperson", "sales", "person", "rep"}), - ({"technician", "tech"}, {"technician", "tech"}), - ({"location"}, {"location", "site", "area"}), - ({"material"}, {"material", "part", "item"}), - ({"priority", "severity"}, {"priority", "severity", "urgency", "rank"}), - ({"status"}, {"status"}), - ({"currency"}, {"currency", "curr"}), - ({"country"}, {"country"}), - ({"order"}, {"order", "ord", "number"}), - ({"batch"}, {"batch", "id"}), - ] - for trigger_tokens, column_tokens in dimension_specs: - if query_tokens & trigger_tokens: - column = _choose_ranked_column_by_tokens(columns, column_tokens) - if column: - return column["name"] - return None +def _schema_driven_column_value_terms( + query: str | None, + columns: list[dict[str, Any]], + schema_tokens: set[str], +) -> list[str]: + values: list[str] = [] + for _, filter_values in _schema_driven_column_value_filters( + query, + columns, + schema_tokens, + ): + values.extend(filter_values) + return _dedupe_filter_values(values) -def _choose_missing_value_column( - query_tokens: set[str], - columns: list[dict[str, str]], -) -> dict[str, str] | None: - missing_specs = [ - ({"customer"}, {"customer", "cust", "number", "id", "no"}), - ({"order"}, {"order", "ord", "number", "id", "no"}), - ({"supplier"}, {"supplier", "vendor", "number", "id", "no"}), - ({"product", "material"}, {"product", "prod", "material", "item", "number", "id"}), - ({"location"}, {"location", "site", "area"}), - ({"status"}, {"status"}), - ] - for trigger_tokens, column_tokens in missing_specs: - if query_tokens & trigger_tokens: - column = _choose_ranked_column_by_tokens(columns, column_tokens) - if column: - return column - return None +def _schema_driven_value_terms( + query: str | None, + schema_tokens: set[str], + columns: list[dict[str, Any]], + introducers: set[str] | None = None, +) -> list[str]: + value_terms = _query_schema_value_terms( + query, + schema_tokens, + introducers=introducers, + ) + if introducers is not None: + return _dedupe_filter_values(value_terms) -def _choose_count_subject_column( - query_tokens: set[str], - columns: list[dict[str, str]], -) -> dict[str, str] | None: - subject_specs = [ - ({"order"}, {"order", "ord", "number", "id", "no"}), - ({"customer"}, {"customer", "cust", "number", "id", "no"}), - ( - {"failure"}, - {"failure", "failed", "defect", "code", "line", "status", "sys", "type"}, - ), - ({"repair"}, {"repair", "id", "status"}), - ({"batch"}, {"batch", "id"}), - ] - for trigger_tokens, column_tokens in subject_specs: - if query_tokens & trigger_tokens: - candidate_columns = columns - if trigger_tokens & {"failure"}: - candidate_columns = [ - column for column in columns if not _is_rate_like_column(column) - ] - column = _choose_ranked_column_by_tokens(candidate_columns, column_tokens) - if column: - return column - return None + column_value_terms = _schema_driven_column_value_terms( + query, + columns, + schema_tokens, + ) + if not column_value_terms: + return _dedupe_filter_values(value_terms) + + column_value_token_sets = [_fallback_tokens(value) for value in column_value_terms] + filtered_terms = [] + for value in value_terms: + value_tokens = _fallback_tokens(value) + if any( + column_tokens + and column_tokens < value_tokens + and value_tokens & schema_tokens + for column_tokens in column_value_token_sets + ): + continue + filtered_terms.append(value) + return _dedupe_filter_values([*filtered_terms, *column_value_terms]) -def _is_text_type(data_type: str) -> bool: - return data_type.upper() in {"CHAR", "NCHAR", "NVARCHAR", "STRING", "TEXT", "VARCHAR"} +def _query_schema_value_terms( + query: str | None, + schema_tokens: set[str], + introducers: set[str] | None = None, +) -> list[str]: + if not query: + return [] -def _missing_value_predicate(column: dict[str, str]) -> str: - quoted_column = _quote_identifier(column["name"]) - if _is_text_type(column["data_type"]): - return f"({quoted_column} IS NULL OR {quoted_column} = '')" - return f"{quoted_column} IS NULL" + values: list[str] = [] + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + def add_value( + raw_value: str, + *, + allow_schema_token_overlap: bool = False, + ) -> None: + value_tokens = _fallback_tokens(raw_value) + if not value_tokens: + return + if value_tokens & _GENERIC_SCHEMA_INTENT_TOKENS: + return + if value_tokens & schema_tokens and ( + not allow_schema_token_overlap or value_tokens <= schema_tokens + ): + return + if all(token.isdigit() for token in value_tokens): + return + cleaned_value = _clean_filter_value(raw_value) + if cleaned_value and cleaned_value.lower() not in { + value.lower() for value in values + }: + values.append(cleaned_value) -def _non_missing_value_predicate(column: dict[str, str]) -> str: - quoted_column = _quote_identifier(column["name"]) - if _is_text_type(column["data_type"]): - return f"({quoted_column} IS NOT NULL AND {quoted_column} <> '')" - return f"{quoted_column} IS NOT NULL" + for phrase_match in _FILTER_VALUE_PHRASE.finditer(query): + introducer = phrase_match.group("introducer").lower() + if introducers is not None and introducer not in introducers: + continue + add_value( + phrase_match.group("value"), + allow_schema_token_overlap=introducer in _OPEN_FILTER_VALUE_INTRODUCERS, + ) + phrase_value_tokens = [_fallback_tokens(value) for value in values] + + def has_filter_value_context(index: int, explicitly_quoted: bool) -> bool: + if explicitly_quoted: + return True + previous_raw = matches[index - 1].group(0).lower() if index > 0 else "" + if introducers is not None and previous_raw not in introducers: + return False + if previous_raw in { + "called", + "contains", + "containing", + "equal", + "equals", + "for", + "from", + "is", + "like", + "matching", + "named", + "where", + "with", + }: + return True -def _aggregate_for_measure(measure_column: str) -> tuple[str, str]: - tokens = _fallback_tokens(measure_column) - if tokens & {"rate", "score", "percent", "percentage"}: - return "AVG", "average_value" - return "SUM", "total_value" + return False + for index, match in enumerate(matches): + raw_value = match.group(0).replace("_", " ") + value_tokens = _fallback_tokens(raw_value) + if not value_tokens: + continue + previous_tokens = ( + _fallback_tokens(matches[index - 1].group(0).replace("_", " ")) + if index > 0 + else set() + ) + previous_raw = matches[index - 1].group(0).lower() if index > 0 else "" + explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { + "'", + '"', + } + if not has_filter_value_context(index, explicitly_quoted): + continue + if previous_tokens & schema_tokens and not explicitly_quoted: + continue + if any(value_tokens <= phrase_tokens for phrase_tokens in phrase_value_tokens): + continue + add_value( + raw_value, + allow_schema_token_overlap=previous_raw + in _OPEN_FILTER_VALUE_INTRODUCERS, + ) -def _select_listing_columns( - query_tokens: set[str], - columns: list[dict[str, str]], - measure_column: str | None = None, - date_column: str | None = None, - max_columns: int = 6, -) -> list[str]: - scored_columns: list[tuple[int, int, str]] = [] - for index, column in enumerate(columns): - name = column["name"] - tokens = _column_business_tokens(column) - score = len(query_tokens & tokens) * 10 - if name == date_column: - score += 35 - if name == measure_column: - score += 12 - if query_tokens & {"order"}: - score += len(tokens & {"order", "ord", "number", "customer", "cust", "item", "product"}) * 18 - score += len(tokens & {"date", "day", "month", "year"}) * 8 - if query_tokens & {"customer"}: - score += len(tokens & {"customer", "cust", "name", "number", "id"}) * 18 - if query_tokens & {"product", "material"}: - score += len(tokens & {"product", "prod", "material", "item", "description", "desc"}) * 18 - if query_tokens & {"batch"}: - score += len(tokens & {"batch", "board", "model", "supplier", "id"}) * 18 - if query_tokens & {"repair", "log", "record"}: - score += ( - len(tokens & {"board", "code", "date", "failure", "id", "priority", "status"}) - * 14 - ) - if tokens & {"repair", "failure", "failed", "defect"} and not query_tokens & { - "repair", - "failure", - "defect", - }: - score -= 40 - if tokens & {"date", "day", "month", "year"} and not ( - _is_date_type(column["data_type"]) or name == date_column - ): - score -= 12 - if score > 0: - scored_columns.append((score, index, name)) + return values - scored_columns.sort(key=lambda item: (-item[0], item[1])) - selected = [] - for _, _, name in scored_columns: - if name not in selected: - selected.append(name) - if len(selected) >= max_columns: - break - if not selected and columns: - selected = [column["name"] for column in columns[:max_columns]] - return selected +def _schema_driven_user_value_tokens( + query: str | None, + query_tokens: set[str], + schema_details: dict[str, list[dict[str, Any]]], +) -> set[str]: + if not query: + return set() -def _fallback_limit(query: str) -> int | None: - match = re.search(r"(?i)\btop\s+(\d+)\b", query) - return int(match.group(1)) if match else None + mentioned_columns: list[dict[str, Any]] = [] + all_columns: list[dict[str, Any]] = [] + for columns in schema_details.values(): + all_columns.extend(columns) + mentioned_columns.extend(_mentioned_text_columns_for_query(query_tokens, columns)) + matched_concepts = [ + _matched_column_query_tokens(query_tokens, column) + for column in mentioned_columns + ] + matched_concepts = [concepts for concepts in matched_concepts if concepts] + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + column_value_terms = _schema_driven_column_value_terms( + query, + all_columns, + schema_tokens, + ) + value_terms = _schema_driven_value_terms( + query, + schema_tokens, + all_columns, + ) + open_value_terms = _schema_driven_value_terms( + query, + schema_tokens, + all_columns, + introducers=_OPEN_FILTER_VALUE_INTRODUCERS, + ) + if not matched_concepts: + if value_terms: + value_tokens: set[str] = set() + for value in value_terms: + value_tokens.update(_fallback_tokens(value)) + open_value_tokens: set[str] = set() + for value in open_value_terms: + open_value_tokens.update(_fallback_tokens(value)) + sample_tokens: set[str] = set() + for columns in schema_details.values(): + for column in columns: + sample_tokens.update(_sample_value_tokens(column)) + if sample_tokens and not (value_tokens & sample_tokens) and not open_value_tokens: + return set() + return value_tokens + logger.info( + "Schema-derived user value grounding skipped: no mentioned categorical columns query=%s", + query, + ) + return set() + shared_concepts = set.intersection(*matched_concepts) + if not shared_concepts: + if column_value_terms: + value_tokens: set[str] = set() + for value in column_value_terms: + value_tokens.update(_fallback_tokens(value)) + logger.info( + "Schema-derived user value grounding accepted explicit column values despite ambiguous categorical concepts query=%s values=%s value_tokens=%s columns=%s concepts=%s", + query, + column_value_terms, + sorted(value_tokens), + [column["name"] for column in mentioned_columns], + [sorted(concepts) for concepts in matched_concepts], + ) + return value_tokens + if open_value_terms: + value_tokens: set[str] = set() + for value in open_value_terms: + value_tokens.update(_fallback_tokens(value)) + logger.info( + "Schema-derived user value grounding accepted open filter values despite ambiguous categorical concepts query=%s values=%s value_tokens=%s columns=%s concepts=%s", + query, + open_value_terms, + sorted(value_tokens), + [column["name"] for column in mentioned_columns], + [sorted(concepts) for concepts in matched_concepts], + ) + return value_tokens + logger.info( + "Schema-derived user value grounding skipped: ambiguous categorical concepts query=%s columns=%s concepts=%s", + query, + [column["name"] for column in mentioned_columns], + [sorted(concepts) for concepts in matched_concepts], + ) + return set() -def _fallback_month_filter(query: str) -> tuple[int, int] | None: - tokens = _fallback_tokens(query) - for month_name, month_number in _MONTH_NAME_TO_NUMBER.items(): - if month_name in tokens: - return datetime.now(timezone.utc).year, month_number - return None + value_tokens: set[str] = set() + for value in value_terms: + value_tokens.update(_fallback_tokens(value)) + logger.info( + "Schema-derived user value grounding query=%s shared_concepts=%s values=%s value_tokens=%s columns=%s", + query, + sorted(shared_concepts), + value_terms, + sorted(value_tokens), + [column["name"] for column in mentioned_columns], + ) + return value_tokens -def _quote_literal(value: str) -> str: - return "'" + value.replace("'", "''") + "'" +def _schema_driven_user_value_filters( + query: str | None, + query_tokens: set[str], + table_name: str, + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + schema_tokens = _schema_tokens_for_table(table_name, columns) + column_value_filters = _schema_driven_column_value_filters( + query, + columns, + schema_tokens, + ) + if column_value_filters: + logger.info( + "Schema-derived explicit column value filters selected query=%s table=%s filters=%s", + query, + table_name, + [ + {"column": column["name"], "values": values} + for column, values in column_value_filters + ], + ) + return column_value_filters + + open_value_terms = _schema_driven_value_terms( + query, + schema_tokens, + columns, + introducers=_OPEN_FILTER_VALUE_INTRODUCERS, + ) + values = _schema_driven_value_terms(query, schema_tokens, columns) + if not values: + return [] + filter_query_tokens = query_tokens - _filter_value_tokens(values) + if open_value_terms: + filter_query_tokens = filter_query_tokens - ( + _grouping_phrase_tokens(query) if query else set() + ) + mentioned_columns = _mentioned_text_columns_for_query(filter_query_tokens, columns) + if len(mentioned_columns) != 1: + return [] -def _clean_filter_value(value: str | None) -> str | None: - if value is None: - return None - value = value.strip(" \t\r\n'\"`.,;:()[]{}") - return value or None + return [(mentioned_columns[0], values)] -def _extract_failure_type_filter_value(query: str) -> str | None: - patterns = [ - ( - r"(?is)\bwith\s+" - r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80}?)" - r"\s+as\s+(?:the\s+)?(?:failure|defect)\s+" - r"(?:type|code|category)\b" - ), - ( - r"(?is)\b(?:failure|defect)\s+(?:type|code|category)\s*" - r"(?:=|is|equals|like|of)\s*['\"]?" - r"(?P[A-Za-z0-9][A-Za-z0-9 _./+\-]{0,80})" - ), - ] - for pattern in patterns: - match = re.search(pattern, query) - if match: - value = _clean_filter_value(match.group("value")) - if value: - return value - return None +def _implicit_text_value_filter_score( + column: dict[str, Any], + query_tokens: set[str], + table_tokens: set[str], + value_tokens: set[str], + allow_unmatched_samples: bool = False, +) -> int: + if ( + not _is_categorical_value_column(column) + or _is_identifier_like_column(column) + or not value_tokens + ): + return 0 + column_tokens = _column_business_tokens(column) + if column_tokens & _IMPLICIT_TEXT_VALUE_IDENTIFIER_TOKENS: + return 0 -def _extract_status_filter_value(query: str) -> str | None: - if re.search(r"(?i)\bin-progress\b", query): - return "in-progress" - if re.search(r"(?i)\bin\s+progress\b", query): - return "in progress" - for status in ("completed", "pending", "escalated", "open", "closed"): - if re.search(rf"(?i)\b{re.escape(status)}\b", query): - return status - return None + sample_tokens = _sample_value_tokens(column) + if sample_tokens and not (value_tokens & sample_tokens) and not allow_unmatched_samples: + return 0 + filter_context_tokens = ( + query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS - table_tokens - value_tokens + ) + score = len(column_tokens & _IMPLICIT_TEXT_VALUE_COLUMN_ROLE_TOKENS) * 12 + score += len(column_tokens & filter_context_tokens) * 20 + if sample_tokens and value_tokens & sample_tokens: + score += 40 + return score -def _extract_priority_filter_value(query: str) -> str | None: - for token, value in _PRIORITY_VALUE_ALIASES.items(): - if re.search(rf"(?i)\b{re.escape(token)}(?:[-\s]+priority)?\b", query): - return value - return None +def _schema_driven_implicit_text_value_filters( + query: str | None, + query_tokens: set[str], + table_name: str, + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + schema_tokens = _schema_tokens_for_table(table_name, columns) + values = _schema_driven_value_terms(query, schema_tokens, columns) + if not values: + return [] -def _choose_priority_column(columns: list[dict[str, str]]) -> dict[str, str] | None: - return _choose_ranked_column_by_tokens( + value_tokens = _filter_value_tokens(values) + open_value_terms = _schema_driven_value_terms( + query, + schema_tokens, columns, - {"priority", "severity", "urgency", "rank"}, + introducers=_OPEN_FILTER_VALUE_INTRODUCERS, ) + allow_unmatched_samples = bool(open_value_terms) + table_tokens = _table_business_tokens(table_name, columns) + scoring_query_tokens = query_tokens - (_grouping_phrase_tokens(query) if query else set()) + candidates: list[tuple[int, str, dict[str, Any]]] = [] + for column in columns: + score = _implicit_text_value_filter_score( + column, + scoring_query_tokens, + table_tokens, + value_tokens, + allow_unmatched_samples=allow_unmatched_samples, + ) + if score > 0: + candidates.append((score, column["name"], column)) + if not candidates: + logger.info( + "Schema-derived implicit text filter skipped no_candidate query=%s table=%s values=%s", + query, + table_name, + values, + ) + return [] -def _priority_order_expression(column: dict[str, str]) -> str: - quoted_column = _quote_identifier(column["name"]) - if _is_numeric_type(column["data_type"]): - return quoted_column - - when_clauses = " ".join( - f"WHEN {_quote_literal(value)} THEN {rank}" for value, rank in _PRIORITY_ORDER + candidates.sort(key=lambda item: (-item[0], item[1])) + if len(candidates) > 1 and candidates[0][0] == candidates[1][0]: + top_score = candidates[0][0] + logger.info( + "Schema-derived implicit text filter skipped ambiguous columns query=%s columns=%s", + query, + [column["name"] for score, _, column in candidates if score == top_score], + ) + return [] + logger.info( + "Schema-derived implicit text filter selected query=%s table=%s column=%s score=%s values=%s sample_values_available=%s", + query, + table_name, + candidates[0][2]["name"], + candidates[0][0], + values, + bool(candidates[0][2].get("sample_values")), ) - return f"CASE LOWER({quoted_column}) {when_clauses} ELSE 0 END" + return [(candidates[0][2], values)] + + +def _where_clause(predicates: list[str]) -> str: + return f"\nWHERE {' AND '.join(predicates)}" if predicates else "" -def _choose_failure_type_filter_column( +def _count_expression_for_query( + query_tokens: set[str], columns: list[dict[str, str]], -) -> dict[str, str] | None: - column = _choose_ranked_column_by_tokens( - [column for column in columns if not _is_rate_like_column(column)], - {"failure", "type"}, - ) - if column: - return column - return _choose_ranked_column_by_tokens( - [column for column in columns if not _is_rate_like_column(column)], - {"failure", "defect", "code", "type", "sys"}, +) -> tuple[str, list[str]]: + subject_column = _choose_count_subject_column(query_tokens, columns) + if not subject_column: + return "COUNT(*)", [] + return ( + f"COUNT({_quote_identifier(subject_column['name'])})", + [_non_missing_value_predicate(subject_column)], ) -def _choose_failure_type_filter_table( - schema_details: dict[str, list[dict[str, str]]], - concept_tokens: set[str], -) -> tuple[str, list[dict[str, str]], dict[str, str]] | None: - candidates = [] - for table_name, columns in schema_details.items(): - if not _table_covers_requested_concepts(table_name, columns, concept_tokens): - continue - column = _choose_failure_type_filter_column(columns) - if not column: - continue - table_tokens = _fallback_tokens(table_name) - column_tokens = _fallback_tokens(column["name"]) - score = len(concept_tokens & table_tokens) * 8 - score += len(concept_tokens & column_tokens) * 10 - if {"failure", "type"}.issubset(column_tokens): - score += 100 - elif "type" in column_tokens: - score += 60 - elif "code" in column_tokens: - score += 30 - if table_tokens & {"failure", "defect"}: - score += 25 - candidates.append((score, table_name, columns, column)) +def _date_bucket_expressions(date_column: str) -> tuple[str, str]: + return ( + _date_part_expression(date_column, "YEAR"), + _date_part_expression(date_column, "MONTH"), + ) - if not candidates: - return None - candidates.sort(key=lambda item: (-item[0], item[1], item[3]["name"])) - _, table_name, columns, column = candidates[0] - return table_name, columns, column + +def _date_part_expression(date_column: str, part: str) -> str: + return f"CAST(EXTRACT({part} FROM {_quote_identifier(date_column)}) AS BIGINT)" def generate_simple_analytics_sql( @@ -2133,87 +4528,55 @@ def generate_simple_analytics_sql( if not query: return None - raw_query_tokens = _fallback_tokens(query) - query_tokens = _expanded_fallback_query_tokens(query) + query_tokens = _fallback_tokens(query) if not query_tokens: return None - if not query_tokens & { - "batch", - "batches", - "board", - "business", - "count", - "customer", - "defect", - "failure", - "failures", - "highest", - "july", - "latest", - "log", - "logs", - "location", - "material", - "missing", - "model", - "monthly", - "most", - "number", - "order", - "orders", - "priority", - "product", - "rate", - "recent", - "record", - "records", - "repair", - "repairs", - "revenue", - "sale", - "sales", - "severity", - "supplier", - "status", - "tech", - "technician", - "top", - "trend", - "trends", - "type", - "unit", - "units", - "year", - }: + schema_details = _extract_schema_details(contexts) + if not schema_details: return None - schema_details = _extract_schema_details(contexts) - rate_metric_intent = _is_rate_metric_intent(raw_query_tokens) - failure_count_intent = _is_failure_count_intent(raw_query_tokens, query_tokens) - board_model_intent = _has_board_model_intent( - raw_query_tokens - ) or _has_board_model_intent( - query_tokens + content_tokens = _query_content_tokens(query_tokens) + schema_backed_tokens = _schema_required_query_tokens( + query, + query_tokens, + schema_details, ) - failure_type_filter_value = _extract_failure_type_filter_value(query) - failure_type_filter_column = None - failure_type_choice = None - if failure_type_filter_value and "failure" in query_tokens: - failure_type_choice = _choose_failure_type_filter_table( - schema_details, - raw_query_tokens, + subject_tokens = _query_subject_schema_tokens( + query, + query_tokens, + schema_details, + ) + unsupported_tokens = _unsupported_query_tokens( + query_tokens, + schema_details, + query=query, + ) + if unsupported_tokens: + logger.info( + "Schema-derived SQL fallback skipped unsupported_tokens=%s", + sorted(unsupported_tokens), ) - - if failure_type_choice: - table_name, columns, failure_type_filter_column = failure_type_choice - chosen = (table_name, columns) - else: - chosen = _choose_fallback_table( - query_tokens, - schema_details, - concept_tokens=raw_query_tokens, + return None + if content_tokens and not schema_backed_tokens: + logger.info( + "Schema-derived SQL fallback skipped no_schema_backed_tokens=%s", + sorted(content_tokens), + ) + return None + if not content_tokens and len(schema_details) != 1: + logger.info( + "Schema-derived SQL fallback skipped ambiguous_schema_only_request tables=%s", + sorted(schema_details), ) + return None + + chosen = _choose_fallback_table( + query_tokens, + schema_details, + concept_tokens=schema_backed_tokens or query_tokens, + subject_tokens=subject_tokens, + ) if not chosen: return None @@ -2221,345 +4584,425 @@ def generate_simple_analytics_sql( column_names = [column["name"] for column in columns] quoted_table = _quote_identifier(table_name) limit = _fallback_limit(query) + date_column = _choose_temporal_column(query_tokens, columns) + order_column = _choose_order_by_column(query, query_tokens, columns) + sample_filters = _sample_value_filters(query_tokens, columns) + sample_filter_column_names = {column["name"] for column, _ in sample_filters} + user_value_filters = [ + (column, values) + for column, values in _schema_driven_user_value_filters( + query, + query_tokens, + table_name, + columns, + ) + if column["name"] not in sample_filter_column_names + ] + implicit_value_filters = [] + if not sample_filters and not user_value_filters: + implicit_value_filters = _schema_driven_implicit_text_value_filters( + query, + query_tokens, + table_name, + columns, + ) + value_filters = [*sample_filters, *user_value_filters] + sample_predicates = [ + _filter_predicate_for_values(column, values) + for column, values in value_filters + ] + sample_predicates.extend( + _implicit_text_value_predicate(column, values) + for column, values in implicit_value_filters + ) + selected_sample_filter_columns = [ + column["name"] for column, _ in [*value_filters, *implicit_value_filters] + ] + filter_column_names = selected_sample_filter_columns + month_filter = _fallback_month_filter(query) + metric_intent = { + "average": _is_average_metric_intent(query_tokens), + "count": _has_count_intent(query_tokens), + "distribution": _is_distribution_metric_intent(query_tokens), + "extreme": _has_extreme_intent(query_tokens), + "latest": _has_latest_intent(query_tokens), + "missing": _has_missing_value_intent(query_tokens), + "rate": _is_rate_metric_intent(query_tokens), + "sum": _has_sum_intent(query_tokens), + } logger.info( - "Deterministic SQL fallback selected table=%s verified_columns=%s metric_intent=%s", + "Schema-derived SQL fallback selected table=%s schema_tokens=%s verified_columns=%s sample_filter_columns=%s metric_intent=%s", table_name, + sorted(schema_backed_tokens), column_names, - { - "failure_count": failure_count_intent, - "rate": rate_metric_intent, - "board_model": board_model_intent, - "failure_type_filter": bool(failure_type_filter_value), - }, + selected_sample_filter_columns, + metric_intent, ) - if failure_type_filter_value: - if not failure_type_filter_column: - failure_type_filter_column = _choose_failure_type_filter_column(columns) - if failure_type_filter_column: + if _has_missing_value_intent(query_tokens): + missing_column = _choose_missing_value_column(query, query_tokens, columns) + if not missing_column: + return None + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + selected_columns = _include_required_listing_columns( + selected_columns, + [missing_column["name"], *filter_column_names], + max_columns=8, + ) + predicates = [*sample_predicates, _missing_value_predicate(missing_column)] + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" + ) + + month_predicate = "" + if date_column and month_filter: + year, month = month_filter + start_date = f"{year:04d}-{month:02d}-01" + end_year = year + 1 if month == 12 else year + end_month = 1 if month == 12 else month + 1 + end_date = f"{end_year:04d}-{end_month:02d}-01" + quoted_date = _quote_identifier(date_column) + month_predicate = ( + f"{quoted_date} >= '{start_date}' AND {quoted_date} < '{end_date}'" + ) + + predicates = list(sample_predicates) + if month_predicate: + predicates.append(month_predicate) + + if _is_average_metric_intent(query_tokens): + measure_column = _choose_average_measure_column(query_tokens, columns) + if not measure_column: + return None + grouping_tokens = _grouping_phrase_tokens(query) or query_tokens + dimension_columns = _choose_dimension_columns( + grouping_tokens, + columns, + max_columns=2, + ) + aggregate_expr = f"AVG({_quote_identifier(measure_column['name'])})" + if dimension_columns: + quoted_dimensions = _quote_joined(dimension_columns) return ( - f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}\n" - f"WHERE {_quote_identifier(failure_type_filter_column['name'])} = " - f"{_quote_literal(failure_type_filter_value)}" + f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier('average_value')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier('average_value')} DESC" ) + return ( + f"SELECT {aggregate_expr} AS {_quote_identifier('average_value')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" + ) - material_column = _choose_column_by_tokens(columns, {"material"}) - location_column = _choose_column_by_tokens(columns, {"location"}) + measure_column = _choose_numeric_measure_column(query_tokens, columns) if ( - {"material", "location"}.issubset(query_tokens) - and material_column - and location_column + date_column + and query_tokens & {"month", "monthly"} + and not ( + _has_sum_intent(query_tokens) + or _is_average_metric_intent(query_tokens) + or _is_rate_metric_intent(query_tokens) + or _has_extreme_intent(query_tokens) + ) ): + year_expr, month_expr = _date_bucket_expressions(date_column) return ( - f"SELECT {_quote_joined([material_column, location_column])}\n" - f"FROM {quoted_table}" - ) - - status_column = _choose_column_by_tokens(columns, {"status"}) - repair_filter_intent = raw_query_tokens & { - "closed", - "completed", - "critical", - "escalated", - "high", - "low", - "medium", - "normal", - "open", - "pending", - "priority", - "progress", - "severity", - "status", - "urgent", - } - if query_tokens & {"repair", "repairs"} and repair_filter_intent: - predicates = [] - status_filter_value = _extract_status_filter_value(query) - if status_column and status_filter_value: - predicates.append( - f"{_quote_identifier(status_column)} = {_quote_literal(status_filter_value)}" - ) - priority_filter_value = _extract_priority_filter_value(query) - if priority_filter_value: - priority_column = _choose_priority_column(columns) - if priority_column: - predicates.append( - f"{_quote_identifier(priority_column['name'])} = " - f"{_quote_literal(priority_filter_value)}" - ) - if predicates: - return f"SELECT *\nFROM {quoted_table}\nWHERE {' AND '.join(predicates)}" + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" + ) - date_column = _choose_column_by_tokens( - columns, - {"date", "day", "month", "time", "year"}, - date=True, - ) + if ( + measure_column + and date_column + and query_tokens & {"month", "monthly"} + and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) + ): + year_expr, month_expr = _date_bucket_expressions(date_column) + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" + f"GROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" + ) + + if ( + measure_column + and date_column + and "year" in query_tokens + and "month" not in query_tokens + and "monthly" not in query_tokens + and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) + ): + year_expr = _date_part_expression(date_column, "YEAR") + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" + f"GROUP BY {year_expr}\nORDER BY {year_expr}" + ) - priority_column = _choose_priority_column(columns) + explicit_grouping_tokens = _grouping_phrase_tokens(query) if ( - priority_column - and raw_query_tokens & {"priority", "severity"} - and raw_query_tokens & {"bottom", "highest", "lowest", "top"} + _has_extreme_intent(query_tokens) + and measure_column + and explicit_grouping_tokens + and _matched_column_query_tokens(explicit_grouping_tokens, measure_column) ): + direction = _sort_direction_for_query(query_tokens) selected_columns = _select_listing_columns( - raw_query_tokens | {"priority", "repair", "status"}, + query_tokens, columns, + measure_column=measure_column["name"], date_column=date_column, max_columns=8, ) - if priority_column["name"] not in selected_columns: - selected_columns.insert(0, priority_column["name"]) - direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" + if measure_column["name"] not in selected_columns: + selected_columns.insert(0, measure_column["name"]) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"ORDER BY {_priority_order_expression(priority_column)} {direction}" - f"{limit_clause}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(measure_column['name'])} {direction}{limit_clause}" + ) + + if _has_count_intent(query_tokens) and not _has_grouping_intent(query, query_tokens): + return ( + f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" ) - if date_column and raw_query_tokens & {"latest", "recent"}: + if _has_latest_intent(query_tokens): + if not date_column: + return None selected_columns = _select_listing_columns( - raw_query_tokens | {"date", "repair", "status"}, + query_tokens, columns, date_column=date_column, max_columns=8, ) if date_column not in selected_columns: selected_columns.insert(0, date_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" ) - if ( - query_tokens & {"repair", "repairs"} - and raw_query_tokens & {"priority", "severity", "status"} - and re.search(r"(?i)\bby\s+(?:priority|severity|status)\b", query) - ): - dimension_column = _choose_dimension_column(raw_query_tokens, columns) - subject_column = _choose_count_subject_column({"repair"}, columns) - if dimension_column: - count_expression = "COUNT(*)" - where_clause = "" - if subject_column: - count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" - where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" - quoted_dimension = _quote_identifier(dimension_column) - return ( - f"SELECT {quoted_dimension}, {count_expression} AS " - f"{_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier('record_count')} DESC" - ) - measure_column = None - if rate_metric_intent and query_tokens & {"defect", "failure"}: - measure_column = _choose_column_by_tokens( + if _has_grouping_intent(query, query_tokens): + grouping_tokens = explicit_grouping_tokens or query_tokens + max_dimensions = 1 if _has_extreme_intent(query_tokens) else 3 + if explicit_grouping_tokens and not _grouping_phrase_has_multiple_dimensions(query): + max_dimensions = 1 + dimension_columns = _choose_dimension_columns( + grouping_tokens, columns, - {"defect", "rate"}, - numeric=True, + max_columns=max_dimensions, ) - if not measure_column: - measure_column = _choose_column_by_tokens(columns, {"rate"}, numeric=True) - if not measure_column and query_tokens & {"order", "orders"}: - measure_column = _choose_column_by_tokens( - columns, - {"amount", "intake", "sales", "value"}, - numeric=True, - ) - if not measure_column and query_tokens & {"revenue", "sale", "sales"}: - measure_column = _choose_column_by_tokens( - columns, - {"amount", "intake", "revenue", "sales", "value"}, - numeric=True, - ) - if ( - not measure_column - and not failure_count_intent - and query_tokens & {"top", "highest", "lowest", "bottom"} - ): - measure_column = _choose_column_by_tokens( - columns, - {"amount", "cost", "count", "margin", "quantity", "rate", "score", "value"}, - numeric=True, - ) - - if raw_query_tokens & {"missing", "blank", "empty", "null"}: - missing_column = _choose_missing_value_column(raw_query_tokens, columns) - if missing_column: - selected_columns = [ - column - for column in column_names - if column == missing_column["name"] - or _fallback_tokens(column) - & { - "batch", - "business", - "bu", - "customer", - "cust", - "date", - "id", - "location", - "name", - "number", - "ord", - "order", - "product", - "status", - "supplier", - } - ][:8] - if missing_column["name"] not in selected_columns: - selected_columns.insert(0, missing_column["name"]) + if explicit_grouping_tokens and dimension_columns: + if not _selected_dimensions_cover_grouping_tokens( + dimension_columns, + explicit_grouping_tokens, + columns, + schema_details, + ): + return None + if dimension_columns: + quoted_dimensions = _quote_joined(dimension_columns) + if measure_column and (_has_sum_intent(query_tokens) or _is_rate_metric_intent(query_tokens)): + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column['name'])})" + direction = _sort_direction_for_query(query_tokens) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"WHERE {_missing_value_predicate(missing_column)}{limit_clause}" + f"SELECT {quoted_dimensions}, COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" ) + return None - implied_count_by_dimension = "failure" in query_tokens and bool( - raw_query_tokens & {"location", "material", "technician", "tech"} - or (board_model_intent and not rate_metric_intent) - ) - if ( - query_tokens & {"count", "number"} - or failure_count_intent - or implied_count_by_dimension - ): - dimension_column = _choose_dimension_column(raw_query_tokens, columns) - if dimension_column: - subject_column = _choose_count_subject_column(raw_query_tokens, columns) - count_expression = "COUNT(*)" - where_clause = "" - if subject_column and raw_query_tokens & {"order", "customer"}: - count_expression = ( - f"COUNT(DISTINCT {_quote_identifier(subject_column['name'])})" + if measure_column and _has_sum_intent(query_tokens): + return ( + f"SELECT SUM({_quote_identifier(measure_column['name'])}) AS {_quote_identifier('total_value')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" + ) + + if _has_extreme_intent(query_tokens): + direction = _sort_direction_for_query(query_tokens) + measure_column_is_grounded = bool( + measure_column + and _matched_column_identifier_query_tokens( + content_tokens | explicit_grouping_tokens, + measure_column, + ) + ) + if measure_column and measure_column_is_grounded: + limit_clause = f"\nLIMIT {limit}" if limit else "" + if _query_allows_grouped_aggregate(query, query_tokens): + dimension_column = _choose_dimension_column(query_tokens, columns) + else: + dimension_column = None + if dimension_column: + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {_quote_identifier(dimension_column)}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {_quote_identifier(dimension_column)}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" ) - elif subject_column and raw_query_tokens & {"failure", "repair", "batch"}: - count_expression = f"COUNT({_quote_identifier(subject_column['name'])})" - where_clause = f"\nWHERE {_non_missing_value_predicate(subject_column)}" - quoted_dimension = _quote_identifier(dimension_column) + selected_columns = _select_listing_columns( + query_tokens, + columns, + measure_column=measure_column["name"], + date_column=date_column, + max_columns=8, + ) + if measure_column["name"] not in selected_columns: + selected_columns.insert(0, measure_column["name"]) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(measure_column['name'])} {direction}{limit_clause}" + ) + if date_column and (month_predicate or query_tokens & {"record", "records", "row", "rows"}): + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if date_column not in selected_columns: + selected_columns.insert(0, date_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {quoted_dimension}, {count_expression} AS {_quote_identifier('record_count')}\n" - f"FROM {quoted_table}{where_clause}\nGROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" + ) + if order_column: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if order_column not in selected_columns: + selected_columns.insert(0, order_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(order_column)} {direction}{limit_clause}" + ) + return None - dimension_column = _choose_dimension_column(raw_query_tokens, columns) - if measure_column and dimension_column and rate_metric_intent: - aggregate, alias = _aggregate_for_measure(measure_column) - quoted_dimension = _quote_identifier(dimension_column) - aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" - direction = "ASC" if raw_query_tokens & {"bottom", "lowest"} else "DESC" + if month_predicate and date_column: + selected_columns = _select_listing_columns( + query_tokens, + columns, + measure_column=measure_column["name"] if measure_column else None, + date_column=date_column, + max_columns=8, + ) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" - ) - - if measure_column and limit and dimension_column and raw_query_tokens & { - "board", - "business", - "customer", - "location", - "material", - "model", - "product", - "salesperson", - "supplier", - "unit", - }: - aggregate, alias = _aggregate_for_measure(measure_column) - quoted_dimension = _quote_identifier(dimension_column) - aggregate_expr = f"{aggregate}({_quote_identifier(measure_column)})" - return ( - f"SELECT {quoted_dimension}, {aggregate_expr} AS {_quote_identifier(alias)}\n" - f"FROM {quoted_table}\nGROUP BY {quoted_dimension}\n" - f"ORDER BY {_quote_identifier(alias)} DESC\nLIMIT {limit}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" ) - month_filter = _fallback_month_filter(query) - if date_column and month_filter: - year, month = month_filter - start = f"{year:04d}-{month:02d}-01" - end_year = year + 1 if month == 12 else year - end_month = 1 if month == 12 else month + 1 - end = f"{end_year:04d}-{end_month:02d}-01" + if order_column: selected_columns = _select_listing_columns( - raw_query_tokens, + query_tokens, columns, - measure_column=measure_column, date_column=date_column, + max_columns=8, ) - order_clause = ( - f"\nORDER BY {_quote_identifier(measure_column)} DESC" - if measure_column - else f"\nORDER BY {_quote_identifier(date_column)} DESC" + if order_column not in selected_columns: + selected_columns.insert(0, order_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, ) limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"WHERE {_quote_identifier(date_column)} >= '{start}' " - f"AND {_quote_identifier(date_column)} < '{end}'" - f"{order_clause}{limit_clause}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(order_column)} ASC{limit_clause}" ) - if ( - measure_column - and date_column - and query_tokens & {"year"} - and not query_tokens & {"month", "monthly"} - ): - date_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" - return ( - f"SELECT {date_expr} AS {_quote_identifier('year')}, " - f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" - f"FROM {quoted_table}\nGROUP BY {date_expr}\nORDER BY {date_expr}" + if sample_predicates or query_tokens & {"all", "list"}: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, ) - - if ( - measure_column - and date_column - and query_tokens & {"month", "monthly", "trend", "trends"} - ): - year_expr = f"EXTRACT(YEAR FROM {_quote_identifier(date_column)})" - month_expr = f"EXTRACT(MONTH FROM {_quote_identifier(date_column)})" - return ( - f"SELECT {year_expr} AS {_quote_identifier('year')}, " - f"{month_expr} AS {_quote_identifier('month')}, " - f"SUM({_quote_identifier(measure_column)}) AS {_quote_identifier('total_value')}\n" - f"FROM {quoted_table}\nGROUP BY {year_expr}, {month_expr}\n" - f"ORDER BY {year_expr}, {month_expr}" + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, ) - - if measure_column and limit: - selected_columns = [ - column - for column in column_names - if column == measure_column - or _fallback_tokens(column) - & { - "batch", - "board", - "customer", - "id", - "model", - "name", - "number", - "supplier", - } - ][:6] - if measure_column not in selected_columns: - selected_columns.append(measure_column) + limit_clause = f"\nLIMIT {limit}" if limit else "" return ( - f"SELECT {_quote_joined(selected_columns)}\nFROM {quoted_table}\n" - f"ORDER BY {_quote_identifier(measure_column)} DESC\nLIMIT {limit}" + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" ) return None @@ -2570,6 +5013,25 @@ class SQLGenPostProcessor: def __init__(self, engine: Engine): self._engine = engine + @staticmethod + def _log_timing( + stage: str, + started_at: float, + project_id: str | None = None, + **fields: Any, + ) -> None: + suffix = " ".join( + f"{key}={value}" for key, value in fields.items() if value is not None + ) + logger.info( + "Ask timing stage=%s project_id=%s elapsed_ms=%.1f%s%s", + stage, + project_id or "", + _timing_ms(started_at), + " " if suffix else "", + suffix, + ) + @component.output_types( valid_generation_result=Dict[str, Any], invalid_generation_result=Dict[str, Any], @@ -2587,9 +5049,42 @@ async def run( allow_data_preview: bool = False, ) -> dict: try: - cleaned_generation_result = clean_generation_result(replies[0]) + total_started_at = time.perf_counter() + extraction_started_at = time.perf_counter() + cleaned_generation_result, extraction_error = _extract_sql_response( + clean_generation_result(replies[0]) + ) + self._log_timing( + "sql_response_extraction", + extraction_started_at, + project_id, + sql_present=bool(cleaned_generation_result), + ) grounding_invalid_generation_result = None + def validate_candidate_sql(candidate_sql: str) -> str | None: + validation_started_at = time.perf_counter() + schema_catalog = _SchemaCatalog.from_contexts(contexts or []) + grounding_error = schema_catalog.validate_sql(candidate_sql) + if not grounding_error: + grounding_error = validate_sql_against_contexts( + candidate_sql, + contexts=contexts, + ) + if not grounding_error: + grounding_error = validate_sql_semantic_coverage( + candidate_sql, + fallback_query, + contexts=contexts, + ) + self._log_timing( + "sql_validation", + validation_started_at, + project_id, + status="rejected" if grounding_error else "grounded", + ) + return grounding_error + if cleaned_generation_result: cleaned_generation_result = normalize_sql_with_schema_identifiers( cleaned_generation_result, @@ -2598,16 +5093,7 @@ async def run( cleaned_generation_result = normalize_wren_sql_dialect( cleaned_generation_result ) - grounding_error = validate_sql_against_contexts( - cleaned_generation_result, - contexts=contexts, - ) - if not grounding_error: - grounding_error = validate_sql_semantic_coverage( - cleaned_generation_result, - fallback_query, - contexts=contexts, - ) + grounding_error = validate_candidate_sql(cleaned_generation_result) if grounding_error: logger.info( "Generated SQL validation result project_id=%s status=rejected reason=%s sql=%s", @@ -2616,9 +5102,9 @@ async def run( cleaned_generation_result, ) grounding_invalid_generation_result = { - "sql": cleaned_generation_result, - "original_sql": cleaned_generation_result, - "type": "SCHEMA_GROUNDING", + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", "error": grounding_error, "correlation_id": "", "data_source": data_source, @@ -2629,11 +5115,24 @@ async def run( project_id or "", cleaned_generation_result, ) + elif extraction_error: + logger.info( + "Generated SQL extraction result project_id=%s status=rejected reason=%s", + project_id or "", + extraction_error, + ) + fallback_started_at = time.perf_counter() fallback_generation_result = generate_simple_analytics_sql( fallback_query, contexts, ) + self._log_timing( + "schema_fast_path_generation", + fallback_started_at, + project_id, + sql_present=bool(fallback_generation_result), + ) if fallback_generation_result: logger.info( "Deterministic SQL fallback generated project_id=%s sql=%s", @@ -2647,16 +5146,9 @@ async def run( fallback_generation_result = normalize_wren_sql_dialect( fallback_generation_result ) - fallback_grounding_error = validate_sql_against_contexts( - fallback_generation_result, - contexts=contexts, + fallback_grounding_error = validate_candidate_sql( + fallback_generation_result ) - if not fallback_grounding_error: - fallback_grounding_error = validate_sql_semantic_coverage( - fallback_generation_result, - fallback_query, - contexts=contexts, - ) logger.info( "Deterministic SQL fallback validation result project_id=%s status=%s%s", project_id or "", @@ -2666,6 +5158,7 @@ async def run( else f" reason={fallback_grounding_error}", ) if not fallback_grounding_error: + engine_validation_started_at = time.perf_counter() ( fallback_valid_generation_result, fallback_invalid_generation_result, @@ -2678,10 +5171,24 @@ async def run( data_source=data_source, allow_data_preview=allow_data_preview, ) + self._log_timing( + "sql_engine_validation", + engine_validation_started_at, + project_id, + status="valid" + if fallback_valid_generation_result + else "invalid", + ) if fallback_valid_generation_result: logger.info( "Using deterministic schema-grounded SQL fallback for query." ) + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="valid", + ) return { "valid_generation_result": fallback_valid_generation_result, "invalid_generation_result": {}, @@ -2692,19 +5199,26 @@ async def run( ) if grounding_invalid_generation_result: - unsupported_result = unsupported_schema_generation_result( + unsupported_message = schema_grounding_failure_message( fallback_query, contexts=contexts, - data_source=data_source, ) - if unsupported_result: - unsupported_message = unsupported_result[ - "invalid_generation_result" - ]["error"] - grounding_invalid_generation_result["type"] = "NO_RELEVANT_SQL" - grounding_invalid_generation_result["error"] = unsupported_message - grounding_invalid_generation_result["sql"] = "" - grounding_invalid_generation_result["original_sql"] = "" + logger.info( + "Generated SQL grounding rejection converted to unsupported schema project_id=%s reason=%s original_reason=%s", + project_id or "", + unsupported_message, + grounding_invalid_generation_result.get("error"), + ) + grounding_invalid_generation_result["type"] = "NO_RELEVANT_SQL" + grounding_invalid_generation_result["error"] = unsupported_message + grounding_invalid_generation_result["sql"] = "" + grounding_invalid_generation_result["original_sql"] = "" + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="invalid", + ) return { "valid_generation_result": {}, "invalid_generation_result": grounding_invalid_generation_result, @@ -2716,8 +5230,34 @@ async def run( data_source=data_source, ) if not cleaned_generation_result and unsupported_result: + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="unsupported", + ) return unsupported_result + if not cleaned_generation_result and extraction_error: + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="invalid", + ) + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": extraction_error, + "correlation_id": "", + "data_source": data_source, + }, + } + + engine_validation_started_at = time.perf_counter() ( valid_generation_result, invalid_generation_result, @@ -2730,6 +5270,18 @@ async def run( data_source=data_source, allow_data_preview=allow_data_preview, ) + self._log_timing( + "sql_engine_validation", + engine_validation_started_at, + project_id, + status="valid" if valid_generation_result else "invalid", + ) + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="valid" if valid_generation_result else "invalid", + ) return { "valid_generation_result": valid_generation_result, @@ -2874,25 +5426,135 @@ async def _classify_generation_result( return valid_generation_result, invalid_generation_result +_SCHEMA_CATALOG_CACHE: dict[tuple[str, ...], "_SchemaCatalog"] = {} + + class _SchemaCatalog: def __init__( self, tables: dict[str, set[str]], relationships: dict[str, set[str]] | None = None, + table_aliases: dict[str, str] | None = None, + table_alias_values: dict[str, str] | None = None, + column_aliases: dict[str, dict[str, str]] | None = None, + column_alias_values: dict[str, dict[str, str]] | None = None, ): self._tables = tables self._relationships = relationships or {} + self._table_aliases = table_aliases or {} + self._table_alias_values = table_alias_values or {} + self._column_aliases = column_aliases or {} + self._column_alias_values = column_alias_values or {} @classmethod def from_contexts(cls, contexts: list[str]) -> "_SchemaCatalog": + context_cache_key = _context_cache_key(contexts) + cached_catalog = _SCHEMA_CATALOG_CACHE.get(context_cache_key) + if cached_catalog is not None: + return cached_catalog + tables: dict[str, set[str]] = {} relationships: dict[str, set[str]] = {} + table_aliases: dict[str, str] = {} + table_alias_values: dict[str, str] = {} + column_aliases: dict[str, dict[str, str]] = {} + column_alias_values: dict[str, dict[str, str]] = {} for context in contexts: + context = str(context) cls._add_contract_identifiers(context, tables, relationships) cls._add_ddl_identifiers(context, tables) + cls._add_semantic_aliases( + context, + table_aliases, + table_alias_values, + column_aliases, + column_alias_values, + ) + + for table_name, column_names in tables.items(): + for alias in _identifier_dot_aliases(table_name): + _add_unique_identifier_alias( + table_aliases, + table_alias_values, + alias, + table_name, + ) + for column_name in column_names: + column_aliases.setdefault(table_name, {}) + column_alias_values.setdefault(table_name, {}) + + catalog = cls( + tables, + relationships, + table_aliases, + table_alias_values, + column_aliases, + column_alias_values, + ) + if len(_SCHEMA_CATALOG_CACHE) >= 256: + _SCHEMA_CATALOG_CACHE.pop(next(iter(_SCHEMA_CATALOG_CACHE))) + _SCHEMA_CATALOG_CACHE[context_cache_key] = catalog + return catalog + + @staticmethod + def _context_table_name(context: str) -> str | None: + payload = _extract_semantic_context_payload(context) + contract = payload.get("sql_identifier_contract") if payload else {} + if isinstance(contract, dict): + table_name = contract.get("sql_table_name_use_exactly") + if isinstance(table_name, str) and table_name: + return _clean_contract_value(table_name) - return cls(tables, relationships) + for raw_line in context.splitlines(): + line = raw_line.strip() + if line.startswith("sql_table_name_use_exactly:"): + table_name = _clean_contract_value(line.split(":", 1)[1]) + if table_name: + return table_name + + match = _DDL_RELATION.search(context) + if match: + return _normalize_identifier(match.group("name")) + return None + + @staticmethod + def _add_semantic_aliases( + context: str, + table_aliases: dict[str, str], + table_alias_values: dict[str, str], + column_aliases: dict[str, dict[str, str]], + column_alias_values: dict[str, dict[str, str]], + ) -> None: + table_name = _SchemaCatalog._context_table_name(context) + payload = _extract_semantic_context_payload(context) + if not table_name or not payload: + return + + table_semantic = payload.get("semantic_context_not_sql_identifiers") + for alias in _semantic_identifier_aliases(table_semantic): + _add_unique_identifier_alias( + table_aliases, + table_alias_values, + alias, + table_name, + ) + + column_aliases.setdefault(table_name, {}) + column_alias_values.setdefault(table_name, {}) + for column in payload.get("columns", []) or []: + if not isinstance(column, dict): + continue + column_name = column.get("sql_column_name_use_exactly") + if not isinstance(column_name, str) or not column_name: + continue + for alias in _semantic_identifier_aliases(column): + _add_unique_identifier_alias( + column_aliases[table_name], + column_alias_values[table_name], + alias, + column_name, + ) @staticmethod def _add_contract_identifiers( @@ -2969,6 +5631,130 @@ def _add_ddl_identifiers( _extract_ddl_column_names(context[body_start:body_end]) ) + def normalize_sql(self, sql: str) -> str: + if not sql: + return sql + sql = self._normalize_table_aliases(sql) + return self._normalize_column_aliases(sql) + + def _normalize_table_aliases(self, sql: str) -> str: + sorted_aliases = sorted( + self._table_aliases.items(), + key=lambda item: len(self._table_alias_values.get(item[0], item[0])), + reverse=True, + ) + for alias_key, table_name in sorted_aliases: + if not table_name: + continue + alias = self._table_alias_values.get(alias_key, alias_key) + replacement = _render_table_identifier(table_name) + for variant in sorted( + _identifier_reference_variants(alias), + key=len, + reverse=True, + ): + sql = _replace_sql_text_outside_literals( + sql, + variant, + replacement, + case_insensitive=True, + ) + return sql + + def _normalize_column_aliases(self, sql: str) -> str: + grounding = _extract_sql_grounding(sql) + real_relations = [ + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] and relation in self._tables + ] + if not real_relations: + return sql + + qualifiers_by_relation: dict[str, set[str]] = {} + for qualifier, relation in grounding["alias_to_relation"].items(): + if relation in self._tables: + qualifiers_by_relation.setdefault(relation, set()).add(qualifier) + + for relation in real_relations: + sql = self._normalize_qualified_column_aliases( + sql, + relation, + qualifiers_by_relation.get(relation, set()), + ) + + unique_real_relations = list(dict.fromkeys(real_relations)) + if len(unique_real_relations) == 1: + sql = self._normalize_unqualified_column_aliases( + sql, + unique_real_relations[0], + ) + return sql + + def _normalize_qualified_column_aliases( + self, + sql: str, + relation: str, + qualifiers: set[str], + ) -> str: + aliases = self._column_aliases.get(relation, {}) + alias_values = self._column_alias_values.get(relation, {}) + for alias_key, column_name in aliases.items(): + if not column_name: + continue + alias = alias_values.get(alias_key, alias_key) + column_variants = sorted( + _identifier_reference_variants(alias), + key=len, + reverse=True, + ) + for qualifier in qualifiers: + qualifier_variants = sorted( + _identifier_reference_variants(qualifier), + key=len, + reverse=True, + ) + qualifier_replacement = _render_qualifier_identifier( + qualifier, + set(self._tables), + ) + replacement = f"{qualifier_replacement}.{_quote_identifier(column_name)}" + for qualifier_variant in qualifier_variants: + for column_variant in column_variants: + sql = _replace_sql_text_outside_literals( + sql, + f"{qualifier_variant}.{column_variant}", + replacement, + case_insensitive=True, + ) + return sql + + def _normalize_unqualified_column_aliases(self, sql: str, relation: str) -> str: + aliases = self._column_aliases.get(relation, {}) + alias_values = self._column_alias_values.get(relation, {}) + sorted_aliases = sorted( + aliases.items(), + key=lambda item: len(alias_values.get(item[0], item[0])), + reverse=True, + ) + for alias_key, column_name in sorted_aliases: + if not column_name: + continue + alias = alias_values.get(alias_key, alias_key) + replacement = _quote_identifier(column_name) + for variant in sorted( + _identifier_reference_variants(alias), + key=len, + reverse=True, + ): + sql = _replace_sql_text_outside_literals( + sql, + variant, + replacement, + case_insensitive=True, + ) + return sql + def to_prompt(self) -> str: if not self._tables: return "" @@ -3424,6 +6210,7 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. - Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. - If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. +- If the DATABASE SCHEMA does not contain an identifier needed to answer part of the request, return null for sql or omit that unsupported part instead of naming a substitute table or column. - When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. - Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. - When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. @@ -3432,11 +6219,13 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. - Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. - SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. +- Identifiers shown in prompt examples are illustrative only unless the same identifier appears exactly in DATABASE SCHEMA for this request. - Generate Wren SQL only, not the native SQL dialect of the connected warehouse. Do not use SQL Server TOP, square-bracket quoting, backtick quoting, FETCH FIRST, OFFSET/FETCH pagination, or warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. - For top, first, highest, lowest, largest, smallest, or other limited result requests, express the ranking/order with ORDER BY and apply a final LIMIT clause in Wren SQL. Never use SELECT TOP n. - Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. - Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. - If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. +- Use `display_label` and `description` only to understand business meaning; when a concept maps to a schema identifier, generated SQL must use that exact identifier. - For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. - In grouped queries, every non-aggregate ORDER BY expression must be a selected grouping column, a selected ordering helper column that is also present in GROUP BY, or a selected aggregate alias. Do not order grouped SQL by a hidden column. - Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. @@ -3444,10 +6233,10 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. - If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. - Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. -- Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema contains specific modeled business columns for the requested entity, measure, status, date, or dimension. -- Prefer exact modeled business fields over generic text search. For example, if a status/severity/date/material/location/customer/order/revenue concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. -- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. This applies to metrics such as defect rate, revenue, amount, sales value, count, cost, margin, and quantity. -- For sales or revenue questions, prefer exact declared sales/revenue/value/amount fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. +- Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata contains specific modeled columns for the requested entity, measure, filter, date, or dimension. +- Prefer exact modeled fields over generic text search. If a requested concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. +- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. +- Do not prefer or exclude any business domain by built-in rules. Ground every choice in the DATABASE SCHEMA supplied for this request. """ @@ -3499,6 +6288,7 @@ def _extract_sql_response(generation_result: str) -> tuple[str | None, str | Non - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. - Do not use SELECT TOP n, FETCH FIRST, OFFSET/FETCH, square-bracket quoting, or backtick quoting. Use Wren SQL syntax with ORDER BY and a final LIMIT n clause for limited or top-N results. - For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. +- For explicit ranking requests, use the ranking function `DENSE_RANK()`, add the ranking column to the final SELECT clause, and filter rank values with WHERE. - For grouped trend queries, include any non-aggregate ordering key in both SELECT and GROUP BY, or order by selected grouping columns/aggregate aliases only. - Reuse exact metric/measure columns when present. Do not invent component columns in order to calculate a requested metric that already exists in DATABASE SCHEMA. """ diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 64cec272ce..e78bad906f 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -102,6 +102,8 @@ async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: return { "name": model.get("name", ""), "properties": model.get("properties", {}), + "tableReference": model.get("tableReference"), + "refSql": model.get("refSql"), "columns": columns, "primaryKey": model.get("primaryKey", ""), } @@ -146,6 +148,9 @@ def _model_command(model: Dict[str, Any]) -> dict: "type": "TABLE", "comment": comment, "name": table_name, + "properties": properties, + "tableReference": model.get("tableReference"), + "refSql": model.get("refSql"), } return {"name": table_name, "payload": str(payload)} @@ -165,6 +170,7 @@ def _column_command(column: Dict[str, Any], model: Dict[str, Any]) -> dict: "name": column["name"], "data_type": column["type"], "is_primary_key": column["name"] == model["primaryKey"], + "properties": column.get("properties", {}), } def _relationship_command( diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 31e3785701..fbc1ec451d 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -47,6 +47,7 @@ def normalize_semantic_properties(props: Dict[str, Any]) -> Dict[str, Any]: semantic_properties = { "alias": clean_display_name(props.get("displayName", "")), "description": props.get("description", ""), + "sourceColumnName": props.get("sourceColumnName", ""), } for key in SEMANTIC_METADATA_KEYS: diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 5dd30dd59a..9e6925274a 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,7 +1,10 @@ import ast +import asyncio import logging import re import sys +import time +from functools import lru_cache from typing import Any, Optional import orjson @@ -11,8 +14,9 @@ from hamilton.async_driver import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe from pydantic import BaseModel, ConfigDict +from sqlparse.sql import Identifier, IdentifierList +from sqlparse.tokens import DML, Comment, Keyword from langfuse.decorators import observe from src.core.pipeline import BasicPipeline @@ -29,66 +33,179 @@ logger = logging.getLogger("wren-ai-service") _SEMANTIC_TABLE_NAME_MERGE_LIMIT = 8 +_LEXICAL_SCHEMA_TABLE_NAME_MERGE_LIMIT = 8 _MAX_RETRIEVED_TABLE_NAMES = 24 +_MAX_LLM_SCHEMA_CONTEXT_TABLES = 8 +_MAX_LLM_SCHEMA_CONTEXT_TOKENS = 12_000 _MAX_RELATED_TABLE_EXPANSION_DEPTH = 1 _RANK_TOKEN = re.compile(r"[a-z0-9]+") -_GENERIC_TABLE_TOKENS = { - "audit", - "auth", - "calendar", - "config", - "dim", - "dimension", - "file", - "files", - "ingestion", - "job", - "jobs", - "log", - "logs", - "lookup", - "mbr", - "member", - "members", - "migration", - "migrations", - "preference", - "preferences", - "queue", - "report", - "reports", - "setting", - "settings", - "state", - "time", - "user", - "users", -} -_CUSTOMS_FINANCE_TOKENS = { - "claim", - "claims", - "custom", - "customs", - "duty", - "duties", - "hmf", - "import", - "imports", - "mpf", - "refund", - "refunds", - "tariff", - "tariffs", +_RANK_GENERIC_QUERY_TOKENS = { + "a", + "across", + "all", + "an", + "and", + "as", + "average", + "avg", + "between", + "bottom", + "breakdown", + "bucket", + "buckets", + "by", + "count", + "counts", + "date", + "day", + "descending", + "distribution", + "each", + "for", + "from", + "group", + "grouped", + "groups", + "has", + "have", + "highest", + "how", + "in", + "is", + "latest", + "least", + "list", + "lowest", + "many", + "max", + "maximum", + "me", + "mean", + "min", + "minimum", + "blank", + "empty", + "missing", + "null", + "month", + "monthly", + "most", + "newest", + "number", + "of", + "ordered", + "per", + "quarter", + "recent", + "record", + "records", + "result", + "results", + "row", + "rows", + "show", + "sort", + "sorted", + "sum", + "the", + "there", + "this", + "to", + "top", + "total", + "using", + "was", + "week", + "were", + "what", + "where", + "which", + "with", + "year", } -_SALES_REVENUE_TOKENS = { +_COMPOUND_IDENTIFIER_PART_TOKENS = { + "account", "amount", - "intake", - "revenue", - "sale", + "balance", + "business", + "buyer", + "category", + "client", + "company", + "count", + "customer", + "date", + "division", + "failure", + "gross", + "group", + "invoice", + "market", + "material", + "month", + "name", + "order", + "person", + "priority", + "product", + "quantity", + "record", + "repair", "sales", - "salesvalue", + "salesperson", + "severity", + "status", + "supplier", + "ticket", + "type", + "unit", "value", + "vendor", + "year", } +_COMPOUND_IDENTIFIER_ALIASES = { + "acct": {"account"}, + "amt": {"amount"}, + "bu": {"business", "unit"}, + "cust": {"customer"}, + "gl": {"general", "ledger"}, + "ord": {"order"}, + "prod": {"product"}, + "qty": {"quantity"}, + "vend": {"vendor"}, +} + +_ALL_SCHEMA_DOCUMENTS_CACHE: dict[tuple[int, str, str], list[Document]] = {} +_ALL_SCHEMA_DOCUMENTS_CACHE_LOCKS: dict[tuple[int, str, str], asyncio.Lock] = {} +_SCHEMA_DOCUMENTS_CACHE: dict[ + tuple[int, str, str, tuple[str, ...]], list[Document] +] = {} +_SCHEMA_DOCUMENTS_CACHE_LOCKS: dict[ + tuple[int, str, str, tuple[str, ...]], asyncio.Lock +] = {} + + +def _elapsed_ms(started_at: float) -> float: + return (time.perf_counter() - started_at) * 1000 + + +def _log_retrieval_timing( + stage: str, + started_at: float, + project_id: str | None = None, + **fields: Any, +) -> None: + suffix = " ".join( + f"{key}={value}" for key, value in fields.items() if value is not None + ) + logger.info( + "Ask timing project_id=%s stage=%s elapsed_ms=%.1f%s%s", + project_id or "", + stage, + _elapsed_ms(started_at), + " " if suffix else "", + suffix, + ) table_columns_selection_system_prompt = """ @@ -119,14 +236,15 @@ 13. Do not stop at a single top candidate when the question needs multiple related datasets. 14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. 15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. -16. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. -17. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. -18. Prefer tables and columns that directly model the requested business entities, measures, statuses, dates, identifiers, and dimensions. Do not answer business-domain questions from generic log, file, JSON, payload, text, or app-metric columns when the schema provides specific modeled columns for the same concept. -19. For terms such as revenue, sales, orders, invoices, customers, products, suppliers, repairs, failures, batches, materials, locations, status, severity, currency, dates, month, year, and business unit, inspect both table meaning and exact column meanings before selecting a table. -20. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. -21. Never return placeholder table or column names such as tablename, table_name, dbo.tablename, BatchId, Material, Location, or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. -22. If the request asks for revenue, sales, or sales trends, prefer exact business measure columns named like Revenue, SalesValue, USDFXSalesValue, FXSalesValue, IntakeValue, Amount, or equivalent modeled sales fields. Do not use tariff, duty, customs, import, refund, or claim datasets unless the user explicitly asks for those domains. -23. If the request asks for an explicit rate, ratio, percentage, revenue, amount, sales value, or other named measure and the schema already contains that exact measure column, use the declared measure column directly. Do not use a rate column to answer "most failures", "number of failures", or other count-of-records requests unless the question explicitly asks for a rate/ratio/percentage. +16. Prefer the set of deployed models, views, metrics, columns, and relationships that best support the current question. +17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. +18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. +19. Prefer tables and columns whose supplied names, descriptions, relationships, metrics, or sample values directly support the requested entities, measures, filters, dates, identifiers, and dimensions. Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata provides specific modeled columns for the same requested concept. +20. Compare the user's requested entities, measures, filters, dates, and dimensions only with schema metadata supplied for the active project. Do not use built-in business synonym lists. +21. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. +22. Never return placeholder table or column names or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. +23. If a requested measure, dimension, filter, or time field is not represented by retrieved schema metadata, leave it unsupported instead of substituting a similar-looking field. +24. Metric intent such as count, sum, average, minimum, maximum, ranking, date bucketing, and grouping must be satisfied by declared columns or metric fields from the retrieved schema. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -298,6 +416,18 @@ def _view_columns_from_statement(statement: str) -> list[dict]: ] +def _source_table_identifier(table_reference: dict | None) -> str: + if not isinstance(table_reference, dict): + return "" + + parts = [ + str(table_reference.get(part, "")).strip() + for part in ("schema", "table") + ] + parts = [part for part in parts if part] + return ".".join(parts) + + def _build_view_ddl(content: dict) -> str: columns = [ column @@ -483,6 +613,8 @@ def _build_table_retrieval_context( ) included_columns = _included_columns(content, columns, tables) included_relationships = _included_relationships(content, tables) + table_properties = content.get("properties") or {} + table_reference = content.get("tableReference") or {} context = _format_semantic_context( { "object_type": "model", @@ -498,6 +630,9 @@ def _build_table_retrieval_context( }, "semantic_context_not_sql_identifiers": { "description": content["comment"], + "display_name": table_properties.get("displayName"), + "source_table_name": _source_table_identifier(table_reference), + "source_table_reference": table_reference, }, "columns": [ { @@ -505,6 +640,12 @@ def _build_table_retrieval_context( "data_type": get_engine_supported_data_type(column["data_type"]), "is_primary_key": column["is_primary_key"], "semantic_context_not_sql_identifier": column["comment"], + "display_name": (column.get("properties") or {}).get( + "displayName" + ), + "source_column_name": (column.get("properties") or {}).get( + "sourceColumnName" + ), } for column in included_columns ], @@ -702,8 +843,15 @@ async def dbschema_retrieval( ) documents = [] if embedding: + semantic_started_at = time.perf_counter() semantic_documents = await _retrieve_semantic_schema_documents( - embedding, project_id, dbschema_retriever, mdl_hash + embedding, project_id, mdl_hash, dbschema_retriever + ) + _log_retrieval_timing( + "schema_retrieval_semantic", + semantic_started_at, + project_id, + document_count=len(semantic_documents), ) semantic_table_names = _table_names_from_schema_documents(semantic_documents)[ :_SEMANTIC_TABLE_NAME_MERGE_LIMIT @@ -711,22 +859,50 @@ async def dbschema_retrieval( table_names = _merge_names(table_names, semantic_table_names)[ :_MAX_RETRIEVED_TABLE_NAMES ] + lexical_started_at = time.perf_counter() + lexical_documents, lexical_table_names = await _retrieve_lexical_schema_hits( + query=query, + project_id=project_id, + mdl_hash=mdl_hash, + dbschema_retriever=dbschema_retriever, + existing_table_names=set(table_names), + ) + _log_retrieval_timing( + "schema_retrieval_lexical_scan", + lexical_started_at, + project_id, + document_count=len(lexical_documents), + table_count=len(lexical_table_names), + ) + table_names = _merge_names(table_names, lexical_table_names)[ + :_MAX_RETRIEVED_TABLE_NAMES + ] + ranking_started_at = time.perf_counter() table_names = _rank_table_names_by_query( table_names, - semantic_documents, + _dedupe_documents(semantic_documents + lexical_documents), query, ) + _log_retrieval_timing( + "candidate_ranking", + ranking_started_at, + project_id, + candidate_count=len(table_names), + ) selected_semantic_table_names = set(semantic_table_names) documents = [ document for document in semantic_documents if document.meta.get("name") in selected_semantic_table_names ] + documents = _dedupe_documents(documents + lexical_documents) if table_names: - retrieved_table_names = set() - pending_table_names = table_names - remaining_expansion_depth = _MAX_RELATED_TABLE_EXPANSION_DEPTH + if include_related_models: + expansion_started_at = time.perf_counter() + retrieved_table_names = set() + pending_table_names = table_names + remaining_expansion_depth = _MAX_RELATED_TABLE_EXPANSION_DEPTH while pending_table_names: retrieved_table_names.update(pending_table_names) @@ -748,7 +924,21 @@ async def dbschema_retrieval( if table_name not in retrieved_table_names ][:remaining_slots] + ranking_started_at = time.perf_counter() ranked_documents = _rank_documents_for_query(documents, table_names, query) + _log_retrieval_timing( + "schema_retrieval_related_expansion", + expansion_started_at, + project_id, + table_count=len(retrieved_table_names), + document_count=len(documents), + ) + _log_retrieval_timing( + "candidate_ranking", + ranking_started_at, + project_id, + candidate_count=len(ranked_documents), + ) logger.info( "Ask schema retrieval project_id=%s retrieved_tables=%s", project_id, @@ -760,20 +950,28 @@ async def dbschema_retrieval( for document in ranked_documents ], ) - documents = _dedupe_documents(documents + retrieved_documents) - if remaining_expansion_depth <= 0: - break - remaining_expansion_depth -= 1 - remaining_slots = _MAX_RETRIEVED_TABLE_NAMES - len(retrieved_table_names) - if remaining_slots <= 0: - break - pending_table_names = [ - table_name - for table_name in _related_table_names(documents) - if table_name not in retrieved_table_names - ][:remaining_slots] + return ranked_documents + named_started_at = time.perf_counter() + retrieved_documents = await _retrieve_schema_documents( + table_names, project_id, mdl_hash, dbschema_retriever + ) + _log_retrieval_timing( + "schema_retrieval_named_fetch", + named_started_at, + project_id, + table_count=len(table_names), + document_count=len(retrieved_documents), + ) + documents = _dedupe_documents(documents + retrieved_documents) + ranking_started_at = time.perf_counter() ranked_documents = _rank_documents_for_query(documents, table_names, query) + _log_retrieval_timing( + "candidate_ranking", + ranking_started_at, + project_id, + candidate_count=len(ranked_documents), + ) logger.info( "Ask schema retrieval project_id=%s retrieved_tables=%s", project_id, @@ -795,10 +993,142 @@ def _tokenize_schema_text(value: Any) -> set[str]: if value is None: return set() text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) - return set(_RANK_TOKEN.findall(text.lower())) + return _expand_schema_token_variants(set(_RANK_TOKEN.findall(text.lower()))) + + +def _normalized_schema_mention_text(value: Any) -> str: + if value is None: + return "" + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) + tokens = _RANK_TOKEN.findall(text.lower()) + return f" {' '.join(tokens)} " if tokens else "" + + +def _schema_identifier_mention_variants(identifier: str) -> set[str]: + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(identifier)) + tokens = _RANK_TOKEN.findall(text.lower()) + variants = {" ".join(tokens)} if tokens else set() + if len(tokens) > 1: + variants.add(" ".join(tokens[1:])) + return {variant for variant in variants if variant} + + +def _query_mentions_schema_identifier(normalized_query: str, identifier: str) -> bool: + if not normalized_query: + return False + return any( + f" {variant} " in normalized_query + for variant in _schema_identifier_mention_variants(identifier) + ) + + +def _rank_content_tokens(query_tokens: set[str]) -> set[str]: + return { + token + for token in query_tokens + if token not in _RANK_GENERIC_QUERY_TOKENS and not token.isdigit() + } + + +def _schema_token_variants(token: str) -> set[str]: + token = token.lower() + variants = {token} + if len(token) > 4 and token.endswith("ies"): + variants.add(token[:-3] + "y") + elif len(token) > 4 and token.endswith("es"): + if token.endswith(("ches", "shes", "sses", "uses", "xes", "zes")): + variants.add(token[:-2]) + else: + variants.add(token[:-1]) + elif len(token) > 3 and token.endswith("s") and not token.endswith(("ss", "us")): + variants.add(token[:-1]) + return {variant for variant in variants if variant} + + +def _expand_schema_token_variants(tokens: set[str]) -> set[str]: + expanded: set[str] = set() + for token in tokens: + expanded.update(_schema_token_variants(token)) + return expanded + + +def _compound_identifier_tokens(token: str) -> set[str]: + if len(token) < 5: + return set() + + tokens: set[str] = set() + for part in _COMPOUND_IDENTIFIER_PART_TOKENS: + if part != token and len(part) >= 4 and part in token: + tokens.add(part) + + for alias, expansions in _COMPOUND_IDENTIFIER_ALIASES.items(): + if alias != token and (token.startswith(alias) or token.endswith(alias)): + tokens.add(alias) + tokens.update(expansions) + + return tokens + + +def _tokenize_schema_identifier_text(value: Any) -> set[str]: + tokens = _tokenize_schema_text(value) + for token in list(tokens): + tokens.update(_compound_identifier_tokens(token)) + return _expand_schema_token_variants(tokens) + + +def _tokenize_nested_schema_identifier_text(value: Any) -> set[str]: + if value is None: + return set() + if isinstance(value, dict): + tokens: set[str] = set() + for nested_key, nested_value in value.items(): + tokens.update(_tokenize_schema_identifier_text(nested_key)) + tokens.update(_tokenize_nested_schema_identifier_text(nested_value)) + return tokens + if isinstance(value, (list, tuple, set)): + tokens: set[str] = set() + for nested_value in value: + tokens.update(_tokenize_nested_schema_identifier_text(nested_value)) + return tokens + return _tokenize_schema_identifier_text(value) + + +def _tokenize_nested_schema_text(value: Any) -> set[str]: + if value is None: + return set() + if isinstance(value, dict): + tokens: set[str] = set() + for nested_key, nested_value in value.items(): + tokens.update(_tokenize_schema_text(nested_key)) + tokens.update(_tokenize_nested_schema_text(nested_value)) + return tokens + if isinstance(value, (list, tuple, set)): + tokens: set[str] = set() + for nested_value in value: + tokens.update(_tokenize_nested_schema_text(nested_value)) + return tokens + return _tokenize_schema_text(value) + + +def _schema_rank_document_key(documents: list[Document]) -> tuple[tuple[str, str, str], ...]: + return tuple( + ( + str(document.meta.get("type", "")), + str(document.meta.get("name", "")), + str(document.content), + ) + for document in documents + ) def _schema_rank_text_by_table(documents: list[Document]) -> dict[str, dict[str, set[str]]]: + return _cached_schema_rank_text_by_table(_schema_rank_document_key(documents)) + + +@lru_cache(maxsize=128) +def _cached_schema_rank_text_by_table( + document_key: tuple[tuple[str, str, str], ...], +) -> dict[str, dict[str, set[str]]]: table_text: dict[str, dict[str, set[str]]] = {} def ensure(table_name: str) -> dict[str, set[str]]: @@ -810,26 +1140,45 @@ def ensure(table_name: str) -> dict[str, set[str]]: } return table_text[table_name] - for document in documents: + for _, meta_name, content_text in document_key: try: - content = ast.literal_eval(document.content) + content = ast.literal_eval(content_text) except (SyntaxError, ValueError): continue - table_name = document.meta.get("name") or content.get("name") + table_name = meta_name or content.get("name") if not table_name: continue bucket = ensure(table_name) - bucket["table"].update(_tokenize_schema_text(table_name)) - bucket["table"].update(_tokenize_schema_text(content.get("name"))) + bucket["table"].update(_tokenize_schema_identifier_text(table_name)) + bucket["table"].update(_tokenize_schema_identifier_text(content.get("name"))) + bucket["table"].update( + _tokenize_nested_schema_identifier_text(content.get("properties")) + ) + bucket["table"].update( + _tokenize_nested_schema_identifier_text(content.get("tableReference")) + ) bucket["comments"].update(_tokenize_schema_text(content.get("comment"))) bucket["comments"].update(_tokenize_schema_text(content.get("description"))) + bucket["comments"].update(_tokenize_nested_schema_text(content.get("refSql"))) for column in content.get("columns", []) or []: - bucket["columns"].update(_tokenize_schema_text(column.get("name"))) - bucket["columns"].update(_tokenize_schema_text(column.get("column"))) - bucket["columns"].update(_tokenize_schema_text(column.get("display_name"))) + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("name")) + ) + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("column")) + ) + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("display_name")) + ) + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("displayName")) + ) + bucket["columns"].update( + _tokenize_nested_schema_identifier_text(column.get("properties")) + ) bucket["comments"].update(_tokenize_schema_text(column.get("comment"))) bucket["comments"].update(_tokenize_schema_text(column.get("description"))) @@ -840,52 +1189,66 @@ def _rank_table_names_by_query( table_names: list[str], semantic_documents: list[Document], query: str | None, + require_positive_score: bool = False, ) -> list[str]: if not query or not table_names: - return table_names + return [] if require_positive_score else table_names query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) if not query_tokens: - return table_names + return [] if require_positive_score else table_names table_text = _schema_rank_text_by_table(semantic_documents) + normalized_query = _normalized_schema_mention_text(query) + content_tokens = _rank_content_tokens(query_tokens) def score(table_name: str) -> int: bucket = table_text.get(table_name, {}) - table_tokens = set(bucket.get("table", set())) | _tokenize_schema_text( + table_tokens = set(bucket.get("table", set())) | _tokenize_schema_identifier_text( table_name ) column_tokens = set(bucket.get("columns", set())) comment_tokens = set(bucket.get("comments", set())) direct_table_matches = query_tokens & table_tokens direct_column_matches = query_tokens & column_tokens + direct_comment_matches = query_tokens & comment_tokens + covered_tokens = direct_table_matches | direct_column_matches | direct_comment_matches + direct_content_matches = content_tokens & covered_tokens + direct_content_column_matches = content_tokens & column_tokens value = ( len(direct_table_matches) * 6 + len(direct_column_matches) * 8 - + len(query_tokens & comment_tokens) + + len(direct_comment_matches) + + len(direct_content_matches) * 32 + + len(direct_content_column_matches) * 24 ) + if _query_mentions_schema_identifier(normalized_query, table_name): + value += 80 + value += len( + direct_table_matches | direct_column_matches | direct_comment_matches + ) ** 2 + if content_tokens and all( + _schema_token_variants(token) & covered_tokens + for token in content_tokens + ): + value += 80 + len(content_tokens) * 16 if len(direct_column_matches) >= 2: value += 8 if direct_table_matches and direct_column_matches: value += 8 - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - value += len((table_tokens | column_tokens) & _SALES_REVENUE_TOKENS) * 5 - if not query_tokens & _CUSTOMS_FINANCE_TOKENS: - value -= ( - len((table_tokens | column_tokens) & _CUSTOMS_FINANCE_TOKENS) - * 8 - ) - if table_tokens & _GENERIC_TABLE_TOKENS and not direct_table_matches: - value -= 6 return value - ranked = sorted( - enumerate(table_names), - key=lambda item: (-score(item[1]), item[0]), - ) - return [table_name for _, table_name in ranked] + scored = [ + (index, table_name, score(table_name)) + for index, table_name in enumerate(table_names) + ] + if require_positive_score: + scored = [item for item in scored if item[2] > 0] + + ranked = sorted(scored, key=lambda item: (-item[2], item[0])) + return [table_name for _, table_name, _ in ranked] def _rank_documents_by_table_names( @@ -916,65 +1279,7 @@ def _rank_documents_for_query( def _augment_retrieval_query(query: str) -> str: - lowered = query.lower() - expansions = [] - - concept_terms = { - ("revenue", "sales", "sale", "amount", "value"): ( - "sales revenue amount value gross net total price intake invoice order" - ), - ("order", "orders"): ( - "order ord number date customer product business unit division company" - ), - ("invoice", "invoices"): ( - "invoice supplier customer currency amount date number" - ), - ("customer", "customers"): ( - "customer account client number name identifier" - ), - ("product", "products"): ( - "product item material type name category" - ), - ("repair", "repairs"): ( - "repair status priority severity failure board model log in progress completed critical" - ), - ("failure", "failures", "defect", "defects"): ( - "failure defect severity occurrence record count code type system status" - ), - ("batch", "batches"): ( - "batch board model supplier defect rate inspection status" - ), - ("material", "materials"): ( - "material item part component location" - ), - ("location", "locations"): ( - "location site warehouse area material" - ), - ("business unit", "bu", "division"): ( - "business unit division company account organization" - ), - ("month", "monthly", "july", "year", "trend", "latest"): ( - "date month year fiscal calendar trend latest recent" - ), - ("status", "severity", "priority", "critical"): ( - "status priority severity critical state category progress" - ), - } - - for triggers, terms in concept_terms.items(): - if any(trigger in lowered for trigger in triggers): - expansions.append(terms) - - if any( - trigger in lowered - for trigger in ("rate", "ratio", "percent", "percentage") - ): - expansions.append("rate ratio percent percentage") - - if not expansions: - return query - - return f"{query}\nBusiness schema search terms: {'; '.join(expansions)}" + return query async def _retrieve_semantic_schema_documents( @@ -999,6 +1304,111 @@ async def _retrieve_semantic_schema_documents( return results["documents"] +async def _retrieve_all_schema_documents( + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, +) -> list[Document]: + cache_key = (id(dbschema_retriever), project_id, mdl_hash or "") + if cache_key in _ALL_SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=all count=%s", + project_id, + mdl_hash or "", + len(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]), + ) + return list(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + lock = _ALL_SCHEMA_DOCUMENTS_CACHE_LOCKS.setdefault(cache_key, asyncio.Lock()) + async with lock: + if cache_key in _ALL_SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=all count=%s", + project_id, + mdl_hash or "", + len(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]), + ) + return list(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + started_at = time.perf_counter() + documents = await _retrieve_all_schema_documents_uncached( + project_id, + mdl_hash, + dbschema_retriever, + ) + if len(_ALL_SCHEMA_DOCUMENTS_CACHE) >= 64: + _ALL_SCHEMA_DOCUMENTS_CACHE.pop(next(iter(_ALL_SCHEMA_DOCUMENTS_CACHE))) + _ALL_SCHEMA_DOCUMENTS_CACHE[cache_key] = documents + _log_retrieval_timing( + "schema_retrieval_all_documents", + started_at, + project_id, + cache_hit=False, + document_count=len(documents), + ) + return list(documents) + + +async def _retrieve_all_schema_documents_uncached( + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, +) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] + + +async def _retrieve_lexical_schema_hits( + query: str | None, + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, + existing_table_names: set[str], +) -> tuple[list[Document], list[str]]: + if not query: + return [], [] + + documents = await _retrieve_all_schema_documents( + project_id, + mdl_hash, + dbschema_retriever, + ) + candidate_table_names = [ + table_name + for table_name in _rank_table_names_by_query( + _table_names_from_schema_documents(documents), + documents, + query, + require_positive_score=True, + ) + if table_name not in existing_table_names + ][:_LEXICAL_SCHEMA_TABLE_NAME_MERGE_LIMIT] + if not candidate_table_names: + return [], [] + + logger.info( + "Ask schema lexical project scan project_id=%s mdl_hash=%s retrieved_tables=%s", + project_id, + mdl_hash or "", + candidate_table_names, + ) + candidate_table_name_set = set(candidate_table_names) + return [ + document + for document in documents + if document.meta.get("name") in candidate_table_name_set + ], candidate_table_names + + def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: table_names = [] seen = set() @@ -1036,6 +1446,7 @@ async def _retrieve_schema_documents( mdl_hash: str | None, dbschema_retriever: Any, ) -> list[Document]: + table_names = list(dict.fromkeys(table_names)) table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names @@ -1044,6 +1455,63 @@ async def _retrieve_schema_documents( if not table_name_conditions: return [] + cache_key = ( + id(dbschema_retriever), + project_id, + mdl_hash or "", + tuple(table_names), + ) + if cache_key in _SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=named table_count=%s document_count=%s", + project_id, + mdl_hash or "", + len(table_names), + len(_SCHEMA_DOCUMENTS_CACHE[cache_key]), + ) + return list(_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + lock = _SCHEMA_DOCUMENTS_CACHE_LOCKS.setdefault(cache_key, asyncio.Lock()) + async with lock: + if cache_key in _SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=named table_count=%s document_count=%s", + project_id, + mdl_hash or "", + len(table_names), + len(_SCHEMA_DOCUMENTS_CACHE[cache_key]), + ) + return list(_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + started_at = time.perf_counter() + documents = await _retrieve_schema_documents_uncached( + table_names, + project_id, + mdl_hash, + dbschema_retriever, + table_name_conditions, + ) + if len(_SCHEMA_DOCUMENTS_CACHE) >= 128: + _SCHEMA_DOCUMENTS_CACHE.pop(next(iter(_SCHEMA_DOCUMENTS_CACHE))) + _SCHEMA_DOCUMENTS_CACHE[cache_key] = documents + _log_retrieval_timing( + "schema_retrieval_named_documents", + started_at, + project_id, + cache_hit=False, + table_count=len(table_names), + document_count=len(documents), + ) + return list(documents) + + +async def _retrieve_schema_documents_uncached( + table_names: list[str], + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, + table_name_conditions: list[dict[str, Any]], +) -> list[Document]: filters = { "operator": "AND", "conditions": [ @@ -1097,6 +1565,50 @@ def _dedupe_documents(documents: list[Document]) -> list[Document]: return deduped +def _limit_retrieval_results_for_generation( + retrieval_results: list[dict[str, Any]], + encoding: tiktoken.Encoding, +) -> tuple[list[dict[str, Any]], int, int, str | None]: + if not retrieval_results: + return retrieval_results, 0, 0, None + + original_tokens = len( + encoding.encode( + " ".join( + retrieval_result.get("table_ddl", "") + for retrieval_result in retrieval_results + ) + ) + ) + limited_results: list[dict[str, Any]] = [] + limited_tokens = 0 + skipped_for_token_budget = False + + for retrieval_result in retrieval_results: + if len(limited_results) >= _MAX_LLM_SCHEMA_CONTEXT_TABLES: + break + + table_tokens = len(encoding.encode(retrieval_result.get("table_ddl", ""))) + if table_tokens > _MAX_LLM_SCHEMA_CONTEXT_TOKENS: + skipped_for_token_budget = True + continue + if limited_tokens + table_tokens > _MAX_LLM_SCHEMA_CONTEXT_TOKENS: + skipped_for_token_budget = True + continue + + limited_results.append(retrieval_result) + limited_tokens += table_tokens + + if not limited_results: + return [], original_tokens, 0, "all_tables_exceed_token_budget" + + if len(limited_results) == len(retrieval_results): + return retrieval_results, original_tokens, original_tokens, None + + reason = "ranked_top_k_skipped_token_budget" if skipped_for_token_budget else "ranked_top_k" + return limited_results, original_tokens, limited_tokens, reason + + @observe() def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: db_schemas = {} @@ -1217,9 +1729,43 @@ def check_using_db_schemas_without_pruning( "has_json_field": has_json_field, } + ( + limited_retrieval_results, + original_token_count, + limited_token_count, + limit_reason, + ) = _limit_retrieval_results_for_generation(retrieval_results, encoding) + if limit_reason: + if not limited_retrieval_results: + logger.info( + "Ask retrieval selected schema context exceeded generation budget; using column pruning reason=%s original_tables=%s original_tokens=%s token_budget=%s", + limit_reason, + len(retrieval_results), + original_token_count, + _MAX_LLM_SCHEMA_CONTEXT_TOKENS, + ) + return { + "db_schemas": [], + "tokens": original_token_count, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + logger.info( + "Ask retrieval capped generation schema context reason=%s original_tables=%s selected_tables=%s original_tokens=%s selected_tokens=%s table_budget=%s token_budget=%s", + limit_reason, + len(retrieval_results), + len(limited_retrieval_results), + original_token_count, + limited_token_count, + _MAX_LLM_SCHEMA_CONTEXT_TABLES, + _MAX_LLM_SCHEMA_CONTEXT_TOKENS, + ) + retrieval_results = limited_retrieval_results + _token_count = limited_token_count + return { - "db_schemas": retrieval_results, - "tokens": _token_count, + "retrieval_results": retrieval_results, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, @@ -1245,24 +1791,48 @@ def prompt( ) ) - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) - return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} - else: + for table_name, table_contents in secondary.items(): + if table_name not in merged: + merged[table_name] = { + **table_contents, + "columns": list(table_contents.get("columns", [])), + } + continue + + columns = list(merged[table_name].get("columns", [])) + for column in table_contents.get("columns", []): + if column not in columns: + columns.append(column) + merged[table_name]["columns"] = columns + + return merged + + +def _lexical_columns_and_tables_needed( + construct_db_schemas: list[dict], + query: str | None, + max_tables: int = 4, + max_columns_per_table: int = 12, +) -> dict[str, dict]: + if not query: return {} + query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) + if not query_tokens: + return {} -@observe(as_type="generation", capture_input=False) -@trace_cost -async def filter_columns_in_tables( - prompt: dict, table_columns_selection_generator: Any, generator_name: str -) -> dict: - if prompt: - return await table_columns_selection_generator( - prompt=prompt.get("prompt") - ), generator_name - else: - return {}, generator_name + scored_tables = [] + for table_schema in construct_db_schemas: + if table_schema.get("type") != "TABLE": + continue + table_tokens = _tokenize_schema_text( + table_schema.get("name") + ) | _tokenize_schema_text( + table_schema.get("comment") + ) + table_score = len(query_tokens & table_tokens) * 6 + column_scores = [] @observe() def construct_retrieval_results( @@ -1361,34 +1931,58 @@ def construct_retrieval_results( ) for document in dbschema_retrieval: - content = ast.literal_eval(document.content) + try: + content = ast.literal_eval(document.content) + except (ValueError, SyntaxError): + logger.warning( + "Skipping malformed retrieved schema document during schema pruning: %s", + document.meta, + ) + continue + + if not isinstance(content, dict): + logger.warning( + "Skipping non-object retrieved schema document during schema pruning: %s", + document.meta, + ) + continue + + content_name = content.get("name") + content_type = content.get("type") + if not content_name: + logger.warning( + "Skipping retrieved schema document without name during schema pruning: %s", + document.meta, + ) + continue - if content["name"] not in tables: + if content_name not in tables: continue - if content["type"] == "METRIC": + if content_type == "METRIC": retrieval_results.append( { - "table_name": content["name"], + "table_name": content_name, "table_ddl": _build_metric_ddl(content), "identifier_context": _identifier_context( - content["name"], + content_name, [ column["name"] - for column in content["columns"] - if column["data_type"].lower() != "unknown" + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" ], ), } ) has_metric = True - elif content["type"] == "VIEW": + elif content_type == "VIEW": retrieval_results.append( { - "table_name": content["name"], + "table_name": content_name, "table_ddl": _build_view_ddl(content), "identifier_context": _identifier_context( - content["name"], + content_name, [ column["name"] for column in content.get("columns", []) @@ -1416,307 +2010,6 @@ def construct_retrieval_results( ], ) - return { - "retrieval_results": retrieval_results, - "has_calculated_field": check_using_db_schemas_without_pruning[ - "has_calculated_field" - ], - "has_metric": check_using_db_schemas_without_pruning["has_metric"], - "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], - } - - -def _normalize_column_selection_results(parsed_response: Any) -> list[dict]: - if isinstance(parsed_response, list): - return [item for item in parsed_response if isinstance(item, dict)] - - if not isinstance(parsed_response, dict): - return [] - - for key in ( - "results", - "tables", - "selected_tables", - "retrieval_results", - "matches", - "data", - "result", - "output", - ): - if key in parsed_response: - normalized = _normalize_column_selection_results(parsed_response[key]) - if normalized: - return normalized - - if "table_name" in parsed_response and ( - "table_contents" in parsed_response or "columns" in parsed_response - ): - return [parsed_response] - - keyed_tables = [] - for table_name, table_contents in parsed_response.items(): - if not isinstance(table_name, str) or not isinstance(table_contents, dict): - continue - if "table_contents" in table_contents: - keyed_tables.append( - { - "table_name": table_name, - "table_contents": table_contents["table_contents"], - } - ) - elif "columns" in table_contents: - keyed_tables.append( - {"table_name": table_name, "table_contents": table_contents} - ) - - return keyed_tables - - -def _parse_column_selection_response(filter_columns_in_tables: dict) -> dict: - raw_reply = (filter_columns_in_tables.get("replies") or [""])[0] - try: - parsed_response = orjson.loads(raw_reply) - except orjson.JSONDecodeError as exc: - logger.warning("Unable to parse column-selection JSON response: %s", exc) - return {} - - normalized_tables = _normalize_column_selection_results(parsed_response) - reformatted_json = {} - for table in normalized_tables: - table_name = table.get("table_name") or table.get("name") - table_contents = table.get("table_contents") or {} - if not table_contents and "columns" in table: - table_contents = table - - columns = ( - table_contents.get("columns") if isinstance(table_contents, dict) else None - ) - if not isinstance(table_name, str) or not isinstance(columns, list): - continue - - reformatted_json[table_name] = { - **table_contents, - "columns": [column for column in columns if isinstance(column, str)], - } - - if not reformatted_json: - response_shape = ( - f"keys={list(parsed_response.keys())[:8]}" - if isinstance(parsed_response, dict) - else type(parsed_response).__name__ - ) - logger.warning( - "Column-selection response did not include usable table columns (%s).", - response_shape, - ) - - return reformatted_json - - -def _build_unpruned_retrieval_results( - construct_db_schemas: list[dict], - dbschema_retrieval: list[Document], -) -> dict: - retrieval_results = [] - has_calculated_field = False - has_metric = False - has_json_field = False - - for table_schema in construct_db_schemas: - if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context(table_schema) - ) - retrieval_results.append( - { - "table_name": table_schema["name"], - "table_ddl": ddl, - } - ) - if _has_calculated_field: - has_calculated_field = True - if _has_json_field: - has_json_field = True - - for document in dbschema_retrieval: - content = ast.literal_eval(document.content) - - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - } - ) - - return { - "retrieval_results": retrieval_results, - "has_calculated_field": has_calculated_field, - "has_metric": has_metric, - "has_json_field": has_json_field, - } - - -def _merge_column_selection( - primary: dict[str, dict], - secondary: dict[str, dict], -) -> dict[str, dict]: - merged = { - table_name: { - **table_contents, - "columns": list(table_contents.get("columns", [])), - } - for table_name, table_contents in primary.items() - } - - for table_name, table_contents in secondary.items(): - if table_name not in merged: - merged[table_name] = { - **table_contents, - "columns": list(table_contents.get("columns", [])), - } - continue - - columns = list(merged[table_name].get("columns", [])) - for column in table_contents.get("columns", []): - if column not in columns: - columns.append(column) - merged[table_name]["columns"] = columns - - return merged - - -def _lexical_columns_and_tables_needed( - construct_db_schemas: list[dict], - query: str | None, - max_tables: int = 4, - max_columns_per_table: int = 12, -) -> dict[str, dict]: - if not query: - return {} - - query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) - if not query_tokens: - return {} - - scored_tables = [] - for table_schema in construct_db_schemas: - if table_schema.get("type") != "TABLE": - continue - - table_tokens = _tokenize_schema_text( - table_schema.get("name") - ) | _tokenize_schema_text( - table_schema.get("comment") - ) - table_score = len(query_tokens & table_tokens) * 6 - column_scores = [] - -@observe() -def construct_retrieval_results( - check_using_db_schemas_without_pruning: dict, - filter_columns_in_tables: dict, - construct_db_schemas: list[dict], - dbschema_retrieval: list[Document], - query: str | None = None, -) -> dict[str, Any]: - if filter_columns_in_tables: - columns_and_tables_needed = _parse_column_selection_response( - filter_columns_in_tables - ) - lexical_columns_and_tables_needed = _lexical_columns_and_tables_needed( - construct_db_schemas, - query, - ) - columns_and_tables_needed = _merge_column_selection( - columns_and_tables_needed, - lexical_columns_and_tables_needed, - ) - tables = set(columns_and_tables_needed.keys()) - retrieval_results = [] - selected_schema_log = [] - has_calculated_field = False - has_metric = False - has_json_field = False - - for table_schema in construct_db_schemas: - if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - selected_columns = set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ) - columns = ( - selected_columns - if _selected_columns_are_executable( - table_schema, selected_columns - ) - else None - ) - ddl, _has_calculated_field, _has_json_field = ( - _build_table_retrieval_context( - table_schema, - columns=columns, - tables=tables, - ) - ) - if _has_calculated_field: - has_calculated_field = True - if _has_json_field: - has_json_field = True - - retrieval_results.append( - { - "table_name": table_schema["name"], - "table_ddl": ddl, - } - ) - selected_schema_log.append( - { - "table": table_schema["name"], - "columns": sorted(selected_columns), - } - ) - - if not retrieval_results: - logger.warning( - "Column-selection output did not match retrieved schemas; " - "falling back to unpruned retrieved schema context." - ) - return _build_unpruned_retrieval_results( - construct_db_schemas, dbschema_retrieval - ) - - if not column_scores and table_score <= 0: - continue - - total_score = table_score + sum(score for score, _, _ in column_scores) - if total_score <= 0: - continue - - logger.info("Ask retrieval selected schema objects=%s", selected_schema_log) - return { - "retrieval_results": retrieval_results, - "has_calculated_field": has_calculated_field, - "has_metric": has_metric, - "has_json_field": has_json_field, - } - else: - retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] - logger.info( - "Ask retrieval selected schema objects=%s", - [ - {"table": retrieval_result.get("table_name"), "columns": "all"} - for retrieval_result in retrieval_results - ], - ) - scored_tables.append((total_score, table_schema["name"], selected_columns)) scored_tables.sort(key=lambda item: (-item[0], item[1])) @@ -1913,13 +2206,16 @@ def _lexical_columns_and_tables_needed( if table_schema.get("type") != "TABLE": continue - table_tokens = _tokenize_schema_text( - table_schema.get("name") - ) | _tokenize_schema_text( - table_schema.get("comment") + table_tokens = ( + _tokenize_schema_identifier_text(table_schema.get("name")) + | _tokenize_schema_text(table_schema.get("comment")) + | _tokenize_nested_schema_identifier_text(table_schema.get("properties")) + | _tokenize_nested_schema_identifier_text(table_schema.get("tableReference")) ) - table_score = len(query_tokens & table_tokens) * 6 + table_matches = query_tokens & table_tokens + table_score = len(table_matches) * 6 column_scores = [] + column_match_union: set[str] = set() for column in table_schema.get("columns", []): if ( @@ -1928,31 +2224,25 @@ def _lexical_columns_and_tables_needed( ): continue - column_tokens = _tokenize_schema_text(column.get("name")) - comment_tokens = _tokenize_schema_text(column.get("comment")) - score = len(query_tokens & column_tokens) * 10 - score += len(query_tokens & comment_tokens) * 2 - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - score += len(column_tokens & _SALES_REVENUE_TOKENS) * 6 - if query_tokens & {"month", "monthly", "year", "july", "date", "latest"}: - score += ( - len(column_tokens & {"date", "day", "month", "year", "time"}) - * 5 - ) - if query_tokens & {"top", "highest", "lowest", "bottom"}: - measure_tokens = { - "amount", - "count", - "cost", - "margin", - "quantity", - "score", - "value", - } - if query_tokens & {"rate", "ratio", "percent", "percentage"}: - measure_tokens.update({"rate", "ratio", "percent", "percentage"}) - score += len(column_tokens & measure_tokens) * 4 + column_tokens = _tokenize_schema_identifier_text(column.get("name")) + column_tokens.update( + _tokenize_schema_identifier_text(column.get("display_name")) + ) + column_tokens.update( + _tokenize_schema_identifier_text(column.get("displayName")) + ) + column_tokens.update( + _tokenize_nested_schema_identifier_text(column.get("properties")) + ) + comment_tokens = _tokenize_schema_text( + column.get("comment") + ) | _tokenize_schema_text(column.get("description")) + column_matches = query_tokens & column_tokens + comment_matches = query_tokens & comment_tokens + score = len(column_matches) * 10 + score += len(comment_matches) * 2 if score > 0: + column_match_union.update(column_matches | comment_matches) column_scores.append( (score, column["name"], column.get("is_primary_key")) ) @@ -1960,14 +2250,8 @@ def _lexical_columns_and_tables_needed( if not column_scores and table_score <= 0: continue - if table_tokens & _GENERIC_TABLE_TOKENS and table_score <= 0: - table_score -= 8 - if query_tokens & {"revenue", "sale", "sales", "trend", "trends"}: - if not query_tokens & _CUSTOMS_FINANCE_TOKENS: - table_score -= len(table_tokens & _CUSTOMS_FINANCE_TOKENS) * 10 - table_score += len(table_tokens & _SALES_REVENUE_TOKENS) * 5 - total_score = table_score + sum(score for score, _, _ in column_scores) + total_score += len(table_matches | column_match_union) ** 2 if total_score <= 0: continue @@ -2092,7 +2376,8 @@ async def run( project_id or "", mdl_hash or "", ) - return await self._pipe.execute( + started_at = time.perf_counter() + result = await self._pipe.execute( ["construct_retrieval_results"], inputs={ "query": query, @@ -2105,3 +2390,14 @@ async def run( **self._configs, }, ) + retrieval_results = result.get("construct_retrieval_results", {}).get( + "retrieval_results", + [], + ) + _log_retrieval_timing( + "schema_retrieval_total", + started_at, + project_id, + retrieval_result_count=len(retrieval_results), + ) + return result diff --git a/wren-ai-service/src/web/v1/routers/ask.py b/wren-ai-service/src/web/v1/routers/ask.py index 595e8fe16e..92b259dbad 100644 --- a/wren-ai-service/src/web/v1/routers/ask.py +++ b/wren-ai-service/src/web/v1/routers/ask.py @@ -1,4 +1,6 @@ import asyncio +import logging +import time import uuid from dataclasses import asdict @@ -21,6 +23,7 @@ ) router = APIRouter() +logger = logging.getLogger("wren-ai-service") @router.post("/asks") @@ -29,6 +32,7 @@ async def ask( service_container: ServiceContainer = Depends(get_service_container), service_metadata: ServiceMetadata = Depends(get_service_metadata), ) -> AskResponse: + started_at = time.perf_counter() query_id = str(uuid.uuid4()) ask_request.query_id = query_id ask_service = service_container.ask_service @@ -57,6 +61,12 @@ def _handle_task_done(completed_task: asyncio.Task): ) task.add_done_callback(_handle_task_done) + logger.info( + "Ask timing query_id=%s project_id=%s stage=task_creation elapsed_ms=%.1f", + query_id, + ask_request.project_id or "", + (time.perf_counter() - started_at) * 1000, + ) return AskResponse(query_id=query_id) diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index a28247f793..ff1a75da1d 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,6 +1,8 @@ import asyncio import logging -from typing import Dict, List, Literal, Optional +import re +import time +from typing import Any, Dict, List, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe @@ -12,12 +14,119 @@ logger = logging.getLogger("wren-ai-service") +_SIMPLE_ANALYTICS_FAST_PATH_PATTERN = re.compile( + r"(?i)\b(" + r"how\s+many|count|counts|number\s+of|total|sum|average|avg|" + r"top\s+(?:\d+|one|two|three|four|five|six|seven|eight|nine|ten|" + r"eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|" + r"eighteen|nineteen|twenty)|" + r"latest|newest|recent|distribution|breakdown|" + r"group(?:ed)?\s+by|by\s+(?:day|week|month|quarter|year)|" + r"each\s+(?:day|week|month|quarter|year)|monthly|per|" + r"missing|blank|empty|null" + r")\b" +) +_SHOW_BY_FAST_PATH_PATTERN = re.compile( + r"(?is)\b(?:show|list)\b.+\bby\s+[A-Za-z0-9_ -]+\b" +) +_GROUP_RESULT_FAST_PATH_PATTERN = re.compile( + r"(?is)\b(?:group|groups)\b.+\b(?:result|results|those|that)\b" +) +_LISTING_FAST_PATH_PATTERN = re.compile( + r"(?is)\b(?:show|list)\b.+\b(?:record|records|row|rows)\b" +) +_GENERAL_HELP_PATTERN = re.compile( + r"(?i)\b(how\s+to|help|guide|docs|documentation|connect|configure|setting|settings)\b" +) +_DATA_SHAPE_PATTERN = re.compile( + r"(?i)\b(record|records|row|rows|table|tables|field|fields|column|columns)\b" +) +_HISTORY_SQL_IDENTIFIER_PATTERN = re.compile( + r'"([^"]+)"|\b(?:FROM|JOIN)\s+([A-Za-z_][A-Za-z0-9_$-]*)', + re.IGNORECASE, +) +_HISTORY_SQL_TABLE_PATTERN = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"([^"]+)"|([A-Za-z_][A-Za-z0-9_$-]*))', + re.IGNORECASE, +) + + +def _looks_like_simple_analytics_request(query: str | None) -> bool: + if not query or _GENERAL_HELP_PATTERN.search(query): + return False + return bool( + _SIMPLE_ANALYTICS_FAST_PATH_PATTERN.search(query) + or _SHOW_BY_FAST_PATH_PATTERN.search(query) + or _GROUP_RESULT_FAST_PATH_PATTERN.search(query) + or _LISTING_FAST_PATH_PATTERN.search(query) + ) + + +def _can_return_pre_intent_schema_unsupported(query: str | None) -> bool: + if not query or _GENERAL_HELP_PATTERN.search(query): + return False + return bool( + _DATA_SHAPE_PATTERN.search(query) + or _SIMPLE_ANALYTICS_FAST_PATH_PATTERN.search(query) + or _SHOW_BY_FAST_PATH_PATTERN.search(query) + or _GROUP_RESULT_FAST_PATH_PATTERN.search(query) + or _LISTING_FAST_PATH_PATTERN.search(query) + ) + class AskHistory(BaseModel): sql: str question: str +def _history_sql_identifiers(sql: str | None) -> list[str]: + if not sql: + return [] + + identifiers = [] + seen = set() + for match in _HISTORY_SQL_IDENTIFIER_PATTERN.finditer(sql): + identifier = match.group(1) or match.group(2) + if not identifier or identifier in seen: + continue + identifiers.append(identifier) + seen.add(identifier) + return identifiers + + +def _history_sql_table_names(sql: str | None) -> list[str]: + if not sql: + return [] + + table_names = [] + seen = set() + for match in _HISTORY_SQL_TABLE_PATTERN.finditer(sql): + table_name = match.group(1) or match.group(2) + if not table_name or table_name in seen: + continue + table_names.append(table_name) + seen.add(table_name) + return table_names + + +def _build_fast_path_grounding_query( + query: str | None, + histories: list[AskHistory], +) -> str: + if not histories: + return query or "" + + latest_history = histories[0] + parts = [] + if latest_history.question: + parts.append(latest_history.question) + history_identifiers = _history_sql_identifiers(latest_history.sql) + if history_identifiers: + parts.append(" ".join(history_identifiers)) + parts.append(query or "") + return "\n".join(parts) + + # POST /v1/asks class AskRequest(BaseRequest): query: str @@ -94,6 +203,51 @@ class AskResultResponse(_AskResultResponse): ] = Field(None, exclude=True) +def _build_sql_correction_error( + error_message: str | None, + sql_diagnosis_reasoning: str | None = None, +) -> str: + raw_error = error_message or "" + if sql_diagnosis_reasoning: + return ( + f"{sql_diagnosis_reasoning}\n\n" + f"Original Wren Engine validation error:\n{raw_error}" + ) + + return f"Original Wren Engine validation error:\n{raw_error}" + + +class _AskStageTimer: + def __init__(self, query_id: str, project_id: str | None): + self._query_id = query_id + self._project_id = project_id or "" + self._started_at = time.perf_counter() + self._last_at = self._started_at + + def mark( + self, + stage: str, + started_at: float | None = None, + **fields: Any, + ) -> None: + ended_at = time.perf_counter() + stage_started_at = started_at if started_at is not None else self._last_at + suffix = " ".join( + f"{key}={value}" for key, value in fields.items() if value is not None + ) + logger.info( + "Ask timing query_id=%s project_id=%s stage=%s elapsed_ms=%.1f total_ms=%.1f%s%s", + self._query_id, + self._project_id, + stage, + (ended_at - stage_started_at) * 1000, + (ended_at - self._started_at) * 1000, + " " if suffix else "", + suffix, + ) + self._last_at = ended_at + + class AskService: def __init__( self, @@ -122,6 +276,79 @@ def __init__( self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries + async def _retrieve_schema_context( + self, + ask_request: AskRequest, + user_query: str, + histories: list[AskHistory], + enable_column_pruning: bool, + timer: _AskStageTimer, + phase: str | None = None, + tables: list[str] | None = None, + ) -> tuple[dict, list[dict], list[str], list[str]]: + schema_retrieval_started_at = time.perf_counter() + retrieval_result = await self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=tables, + histories=histories, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + enable_column_pruning=enable_column_pruning, + ) + retrieval_payload = retrieval_result.get("construct_retrieval_results", {}) + documents = retrieval_payload.get("retrieval_results", []) + table_names = [document.get("table_name") for document in documents] + table_ddls = [document.get("table_ddl") for document in documents] + timer.mark( + "schema_retrieval", + schema_retrieval_started_at, + retrieved_table_count=len(table_names), + phase=phase, + ) + return retrieval_payload, documents, table_names, table_ddls + + async def _run_schema_fast_path( + self, + ask_request: AskRequest, + user_query: str, + table_ddls: list[str], + histories: list[AskHistory], + grounding_query: str, + use_dry_plan: bool, + allow_dry_plan_fallback: bool, + timer: _AskStageTimer, + phase: str | None = None, + ) -> dict | None: + fast_path_pipeline_name = ( + "followup_sql_generation" if histories else "sql_generation" + ) + fast_path_pipeline = self._pipelines[fast_path_pipeline_name] + fast_path_runner = getattr( + fast_path_pipeline, + "run_deterministic_fast_path", + None, + ) + if not fast_path_runner: + return None + + fast_path_started_at = time.perf_counter() + fast_path_result = await fast_path_runner( + query=user_query, + contexts=table_ddls, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + grounding_query=grounding_query, + ) + timer.mark( + "sql_generation_fast_path", + fast_path_started_at, + result=fast_path_result.get("fast_path") if fast_path_result else "miss", + phase=phase, + ) + return fast_path_result + def _is_stopped(self, query_id: str, container: dict): if ( result := container.get(query_id) @@ -159,8 +386,11 @@ async def ask( instructions = [] api_results = [] table_names = [] + table_ddls = [] + _retrieval_result = {} error_message = None invalid_sql = None + fast_path_terminal = False allow_sql_generation_reasoning = ( self._allow_sql_generation_reasoning and not ask_request.ignore_sql_generation_reasoning @@ -178,7 +408,14 @@ async def ask( sql_knowledge = None try: - user_query = ask_request.query + original_user_query = ask_request.query + user_query = original_user_query + timer = _AskStageTimer(query_id, ask_request.project_id) + timer.mark( + "frontend_request", + query_chars=len(original_user_query or ""), + history_count=len(histories), + ) # ask status can be understanding, searching, generating, finished, failed, stopped # we will need to handle business logic for each status @@ -191,6 +428,7 @@ async def ask( if not api_results: # Run both pipeline operations concurrently + support_context_started_at = time.perf_counter() sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( query=user_query, @@ -211,8 +449,109 @@ async def ask( instructions = instructions_task["formatted_output"].get( "documents", [] ) + timer.mark( + "schema_retrieval_support_context", + support_context_started_at, + sql_sample_count=len(sql_samples), + instruction_count=len(instructions), + ) - if self._allow_intent_classification: + if ( + _looks_like_simple_analytics_request(user_query) + and not self._is_stopped(query_id, self._ask_results) + ): + pre_intent_grounding_query = ( + _build_fast_path_grounding_query(user_query, histories) + if histories + else original_user_query + ) + pre_intent_tables = ( + _history_sql_table_names(histories[0].sql) + if histories + else None + ) + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + ( + _retrieval_result, + documents, + table_names, + table_ddls, + ) = await self._retrieve_schema_context( + ask_request=ask_request, + user_query="" + if pre_intent_tables + else pre_intent_grounding_query, + histories=histories, + enable_column_pruning=enable_column_pruning, + timer=timer, + phase="pre_intent", + tables=pre_intent_tables, + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="schema_retrieval") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + if documents: + fast_path_result = await self._run_schema_fast_path( + ask_request=ask_request, + user_query=user_query, + table_ddls=table_ddls, + histories=histories, + grounding_query=pre_intent_grounding_query, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + timer=timer, + phase="pre_intent", + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark( + "cancelled", + at_stage="sql_generation_fast_path", + ) + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if fast_path_result: + post_process = fast_path_result["post_process"] + if sql_valid_result := post_process.get( + "valid_generation_result" + ): + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } + ) + ] + fast_path_terminal = True + elif ( + failed_result := post_process.get( + "invalid_generation_result" + ) + ) and ( + failed_result.get("type") == "NO_RELEVANT_SQL" + and _can_return_pre_intent_schema_unsupported( + user_query + ) + ): + error_message = failed_result.get("error") + invalid_sql = "" + fast_path_terminal = True + + if ( + self._allow_intent_classification + and not api_results + and not fast_path_terminal + ): + intent_started_at = time.perf_counter() intent_classification_result = ( await self._pipelines["intent_classification"].run( query=user_query, @@ -229,10 +568,20 @@ async def ask( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") + timer.mark( + "llm_intent_generation", + intent_started_at, + intent=intent, + ) if rephrased_question: user_query = rephrased_question + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="llm_intent_generation") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if intent == "MISLEADING_QUERY": asyncio.create_task( self._pipelines["misleading_assistance"].run( @@ -313,7 +662,16 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - if not self._is_stopped(query_id, self._ask_results) and not api_results: + grounding_query = ( + _build_fast_path_grounding_query(user_query, histories) + if histories + else original_user_query + ) + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and not fast_path_terminal + ): self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -323,19 +681,23 @@ async def ask( is_followup=True if histories else False, ) - retrieval_result = await self._pipelines["db_schema_retrieval"].run( - query=user_query, + ( + _retrieval_result, + documents, + table_names, + table_ddls, + ) = await self._retrieve_schema_context( + ask_request=ask_request, + user_query=user_query, histories=histories, - project_id=ask_request.project_id, - mdl_hash=ask_request.mdl_hash, enable_column_pruning=enable_column_pruning, + timer=timer, ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents = _retrieval_result.get("retrieval_results", []) - table_names = [document.get("table_name") for document in documents] - table_ddls = [document.get("table_ddl") for document in documents] + + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="schema_retrieval") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -356,9 +718,44 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + fast_path_result = await self._run_schema_fast_path( + ask_request=ask_request, + user_query=user_query, + table_ddls=table_ddls, + histories=histories, + grounding_query=grounding_query, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + timer=timer, + ) + if fast_path_result: + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_generation_fast_path") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + post_process = fast_path_result["post_process"] + if sql_valid_result := post_process.get( + "valid_generation_result" + ): + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } + ) + ] + fast_path_terminal = True + elif failed_result := post_process.get("invalid_generation_result"): + if failed_result.get("type") == "NO_RELEVANT_SQL": + error_message = failed_result.get("error") + invalid_sql = "" + fast_path_terminal = True + if ( not self._is_stopped(query_id, self._ask_results) and not api_results + and not fast_path_terminal and allow_sql_generation_reasoning ): self._ask_results[query_id] = AskResultResponse( @@ -371,6 +768,7 @@ async def ask( is_followup=True if histories else False, ) + sql_reasoning_started_at = time.perf_counter() if histories: sql_generation_reasoning = ( await self._pipelines["followup_sql_generation_reasoning"].run( @@ -398,6 +796,15 @@ async def ask( query_id=query_id, ) ).get("post_process", {}) + timer.mark( + "sql_generation_reasoning", + sql_reasoning_started_at, + is_followup=bool(histories), + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_generation_reasoning") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results self._ask_results[query_id] = AskResultResponse( status="planning", @@ -410,7 +817,11 @@ async def ask( is_followup=True if histories else False, ) - if not self._is_stopped(query_id, self._ask_results) and not api_results: + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and not fast_path_terminal + ): self._ask_results[query_id] = AskResultResponse( status="generating", type="TEXT_TO_SQL", @@ -422,6 +833,7 @@ async def ask( is_followup=True if histories else False, ) + auxiliary_retrieval_started_at = time.perf_counter() if allow_sql_functions_retrieval: sql_functions = await self._pipelines[ "sql_functions_retrieval" @@ -439,6 +851,16 @@ async def ask( project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, ) + timer.mark( + "sql_auxiliary_retrieval", + auxiliary_retrieval_started_at, + functions_enabled=allow_sql_functions_retrieval, + knowledge_enabled=allow_sql_knowledge_retrieval, + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_auxiliary_retrieval") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -446,6 +868,7 @@ async def ask( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) + sql_generation_started_at = time.perf_counter() if histories: text_to_sql_generation_results = await self._pipelines[ "followup_sql_generation" @@ -465,6 +888,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + grounding_query=grounding_query, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -484,7 +908,22 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + grounding_query=original_user_query, + ) + timer.mark( + "sql_generation", + sql_generation_started_at, + is_followup=bool(histories), + status="valid" + if text_to_sql_generation_results["post_process"].get( + "valid_generation_result" ) + else "invalid", + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_generation") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" @@ -501,6 +940,10 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_correction") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if failed_dry_run_result["type"] in ( "TIME_OUT", "NO_RELEVANT_SQL", @@ -517,6 +960,7 @@ async def ask( invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] current_sql_correction_retries += 1 + sql_diagnosis_reasoning = None self._ask_results[query_id] = AskResultResponse( status="correcting", @@ -530,6 +974,7 @@ async def ask( ) if allow_sql_diagnosis: + diagnosis_started_at = time.perf_counter() sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -544,7 +989,22 @@ async def ask( sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") + timer.mark( + "sql_diagnosis", + diagnosis_started_at, + retry=current_sql_correction_retries, + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_diagnosis") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + correction_error = _build_sql_correction_error( + error_message=error_message, + sql_diagnosis_reasoning=sql_diagnosis_reasoning, + ) + correction_started_at = time.perf_counter() sql_correction_results = await self._pipelines[ "sql_correction" ].run( @@ -554,12 +1014,7 @@ async def ask( instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": ( - f"{sql_diagnosis_reasoning}\nDry run error: {error_message}" - if allow_sql_diagnosis - and sql_diagnosis_reasoning - else error_message - ), + "error": correction_error, }, project_id=ask_request.project_id, mdl_hash=ask_request.mdl_hash, @@ -568,6 +1023,20 @@ async def ask( sql_functions=sql_functions, sql_knowledge=sql_knowledge, ) + timer.mark( + "sql_correction", + correction_started_at, + retry=current_sql_correction_retries, + status="valid" + if sql_correction_results["post_process"].get( + "valid_generation_result" + ) + else "invalid", + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_correction") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if valid_generation_result := sql_correction_results[ "post_process" @@ -601,6 +1070,7 @@ async def ask( ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" + timer.mark("ask_total", status="finished") else: logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): @@ -622,6 +1092,7 @@ async def ask( results["metadata"]["error_type"] = "NO_RELEVANT_SQL" results["metadata"]["error_message"] = error_message results["metadata"]["type"] = "TEXT_TO_SQL" + timer.mark("ask_total", status="failed", error_type="NO_RELEVANT_SQL") return results except Exception as e: @@ -647,9 +1118,16 @@ def stop_ask( self, stop_ask_request: StopAskRequest, ): + started_at = time.perf_counter() self._ask_results[stop_ask_request.query_id] = AskResultResponse( status="stopped", ) + logger.info( + "Ask timing query_id=%s project_id=%s stage=cancel_request elapsed_ms=%.1f status=stopped", + stop_ask_request.query_id, + stop_ask_request.project_id or "", + (time.perf_counter() - started_at) * 1000, + ) def get_ask_result( self, diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index c907387300..a87c17d203 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -1,5 +1,6 @@ import asyncio import logging +import time from typing import Dict, Literal, Optional from cachetools import TTLCache @@ -73,15 +74,23 @@ async def sql_answer( try: query_id = sql_answer_request.query_id + total_started_at = time.perf_counter() self._sql_answer_results[query_id] = SqlAnswerResultResponse( status="preprocessing", trace_id=trace_id, ) + preprocess_started_at = time.perf_counter() preprocessed_sql_data = self._pipelines["preprocess_sql_data"].run( sql_data=sql_answer_request.sql_data, )["preprocess"] + logger.info( + "Ask timing query_id=%s stage=answer_formatting_preprocess elapsed_ms=%.1f row_count=%s", + query_id, + (time.perf_counter() - preprocess_started_at) * 1000, + preprocessed_sql_data.get("num_rows_used_in_llm"), + ) if preprocessed_sql_data.get("num_rows_used_in_llm") == 0: results["metadata"]["error_type"] = "NO_DATA" @@ -93,6 +102,7 @@ async def sql_answer( trace_id=trace_id, ) + formatting_task_started_at = time.perf_counter() asyncio.create_task( self._pipelines["sql_answer"].run( query=sql_answer_request.query, @@ -104,6 +114,12 @@ async def sql_answer( custom_instruction=sql_answer_request.custom_instruction, ) ) + logger.info( + "Ask timing query_id=%s stage=answer_formatting_task_creation elapsed_ms=%.1f total_ms=%.1f", + query_id, + (time.perf_counter() - formatting_task_started_at) * 1000, + (time.perf_counter() - total_started_at) * 1000, + ) return results except Exception as e: diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py new file mode 100644 index 0000000000..bd11a6470d --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py @@ -0,0 +1,56 @@ +from pathlib import Path + + +SERVICE_ROOT = Path(__file__).resolve().parents[4] + + +def _read_source(relative_path: str) -> str: + return (SERVICE_ROOT / relative_path).read_text(encoding="utf-8") + + +def test_intent_classification_does_not_require_exact_user_schema_names(): + source = _read_source("src/pipelines/generation/intent_classification.py") + + assert "schema-resolvable references" in source + assert "even if the user did not type exact table or column names" in source + assert "do not require the user to write exact schema identifiers" in source + assert ( + "Do not classify a data retrieval or analytics question as MISLEADING only " + "because the user did not write exact table or column names" + ) in source + + +def test_data_assistance_does_not_invent_hypothetical_schema(): + source = _read_source("src/pipelines/generation/data_assistance.py") + + assert "MUST NOT add SQL code" in source + assert "Use only the provided DATABASE SCHEMA as context" in source + assert "Do not invent, assume, or name tables or columns" in source + assert "do not provide hypothetical schema" in source + + +def test_sql_reasoning_contract_rejects_substitute_identifiers(): + source = _read_source("src/pipelines/generation/utils/sql.py") + + assert "If the DATABASE SCHEMA does not contain an identifier needed" in source + assert "instead of naming a substitute table or column" in source + assert "prompt examples" in source + assert "Identifiers shown in prompt examples are illustrative only" in source + assert "Use `display_label` and `description` only to understand" in source + assert "generated SQL must use that exact identifier" in source + + +def test_sql_correction_receives_raw_wren_engine_validation_error(): + source = _read_source("src/web/v1/services/ask.py") + + assert "_build_sql_correction_error" in source + assert "Original Wren Engine validation error" in source + assert "error_message" in source + + +def test_sql_correction_unknown_identifier_contract(): + source = _read_source("src/pipelines/generation/sql_correction.py") + + assert "If the error reports an unknown table or field" in source + assert "replace it only with an exact executable identifier" in source + assert "Do not retry the same unknown identifier" in source diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py index 0f696bb0f3..a657582022 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py @@ -5,6 +5,7 @@ from src.pipelines.generation.sql_answer import ( SQLAnswer, prompt, + sql_to_answer_system_prompt, sql_to_answer_user_prompt_template, ) @@ -19,6 +20,12 @@ def test_sql_answer_prompt_uses_sql_data_rows(): {"name": "manufacturing_cost_per_unit", "type": "double"}, ], "data": [["Supplier 1", 0.06]], + "row_records": [ + { + "supplier_name": "Supplier 1", + "manufacturing_cost_per_unit": 0.06, + } + ], }, language="English", current_time="2026-08-03T00:00:00", @@ -29,8 +36,18 @@ def test_sql_answer_prompt_uses_sql_data_rows(): generated_prompt = result["prompt"] assert "rows:" in generated_prompt + assert "row records:" in generated_prompt assert "Supplier 1" in generated_prompt + assert "supplier_name" in generated_prompt assert "result rows:" not in generated_prompt + assert "Please think step by step" not in generated_prompt + + +def test_sql_answer_prompt_blocks_code_style_hallucinated_analysis(): + assert "Do not write code" in sql_to_answer_system_prompt + assert "Python" in sql_to_answer_system_prompt + assert "running the above code" in sql_to_answer_system_prompt + assert "Do not invent values" in sql_to_answer_system_prompt @pytest.mark.asyncio diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py index 57fb089dad..b72785cb3c 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -2,9 +2,13 @@ from src.pipelines.generation.utils.sql import ( SQLGenPostProcessor, + _SchemaCatalog, + _fallback_limit, generate_simple_analytics_sql, - normalize_wren_sql_dialect, normalize_sql_with_schema_identifiers, + normalize_wren_sql_dialect, + sanitize_sql_generation_reasoning, + schema_grounding_failure_message, unsupported_schema_generation_result, unsupported_schema_message, validate_sql_against_contexts, @@ -12,6 +16,14 @@ ) +class _AcceptingEngine: + async def dry_plan(self, *args, **kwargs): + return True, "" + + async def execute_sql(self, *args, **kwargs): + return True, [], {"correlation_id": "test"} + + SCHEMA_CONTEXTS = [ """ CREATE TABLE valid_invoice_comments ( @@ -28,6 +40,22 @@ ] +def test_sql_reasoning_sanitizer_blocks_query_shaped_output(): + reasoning = """ + The SQL could look like this: + ```sql + SELECT * FROM dbo_dimOrderNumber WHERE id2 = 'CATERPILLAR S.A.R.L.'; + ``` + """ + + sanitized = sanitize_sql_generation_reasoning(reasoning) + + assert "SELECT" not in sanitized + assert "WHERE" not in sanitized + assert "assume" not in sanitized.lower() + assert "retrieved schema metadata" in sanitized + + def test_schema_grounding_rejects_unretrieved_table_name(): error = validate_sql_against_contexts( "SELECT invoice_id, COUNT(comment_id) FROM comments GROUP BY invoice_id", @@ -52,6 +80,56 @@ def test_schema_grounding_accepts_retrieved_table_name(): assert error is None +def test_supported_how_many_question_ignores_interrogative_fillers(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Invoice records by business unit."},"columns":[{"sql_column_name_use_exactly":"invoice_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Invoice identifier."},{"sql_column_name_use_exactly":"business_unit","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Business unit for invoice grouping."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE invoice_records ( + invoice_id VARCHAR, + business_unit VARCHAR + ); + """ + ] + + message = unsupported_schema_message( + "How many invoice records are there by business unit?", + contexts, + ) + + assert message is None + + +def test_unsupported_how_many_question_still_reports_missing_schema_terms(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Customer order records."},"columns":[{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."},{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + order_id VARCHAR, + customer_name VARCHAR + ); + """ + ] + + message = unsupported_schema_message( + "How many repair records are there by status?", + contexts, + ) + + assert message is not None + assert "repair" in message + assert "status" in message + assert "how" not in message + assert "there" not in message + + def test_schema_grounding_rejects_invalid_qualified_column(): error = validate_sql_against_contexts( """ @@ -74,507 +152,1714 @@ def test_schema_identifier_normalization_quotes_special_identifiers(): assert 'FROM "valid-order-lines"' in sql -def test_semantic_coverage_rejects_generic_table_for_business_concepts(): +def test_schema_identifier_normalization_rewrites_verified_source_table_name(): contexts = [ """ - CREATE TABLE dbo_mbrTime ( - id1 INTEGER, - id2 INTEGER + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"orders_model","sql_column_names_use_exactly":["customer_name"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"tblOrders"},"display_name":"New Orders"},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":{"source_column_name":"CustName","display_name":"CustName"}}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE orders_model ( + customer_name VARCHAR ); """ ] - error = validate_sql_semantic_coverage( - """ - SELECT id1, COUNT(*) AS failures - FROM dbo_mbrTime - GROUP BY id1 - ORDER BY failures DESC - LIMIT 10 - """, - "Show the top 10 materials with the highest number of failures.", + sql = normalize_sql_with_schema_identifiers( + "SELECT customer_name FROM dbo.tblOrders", contexts, ) - assert error is not None - assert "failure/defect" in error - assert "material" in error + assert 'FROM "orders_model"' in sql + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + assert validate_sql_against_contexts(sql, contexts) is None -def test_unsupported_schema_message_requires_all_requested_concepts(): +def test_schema_identifier_normalization_rewrites_verified_source_column_name(): contexts = [ """ - CREATE TABLE dbo_mbrTime ( - id1 INTEGER, - id2 INTEGER + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"orders_model","sql_column_names_use_exactly":["customer_name"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"tblOrders"}},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":{"source_column_name":"CustName","display_name":"CustName"}}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE orders_model ( + customer_name VARCHAR ); """ ] - message = unsupported_schema_message( - "Show the top 10 materials with the highest number of failures.", + sql = normalize_sql_with_schema_identifiers( + """ + SELECT o.CustName + FROM dbo.tblOrders o + WHERE LOWER(o.CustName) = LOWER('lockheed martin') + """, contexts, ) - assert message is not None - assert "No retrieved table or view" in message - assert "failure/defect" in message - assert "material" in message + assert 'FROM "orders_model" o' in sql + assert 'o."customer_name"' in sql + assert "CustName" not in sql + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + assert validate_sql_against_contexts(sql, contexts) is None -def test_unsupported_schema_message_rejects_split_failure_technician_without_coverage(): +def test_schema_identifier_normalization_rewrites_verified_display_column_variant(): contexts = [ """ - CREATE TABLE dbo_report_failures ( - id INTEGER, - failure_type VARCHAR + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"orders_model","sql_column_names_use_exactly":["CustName"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"tblOrders"}},"columns":[{"sql_column_name_use_exactly":"CustName","data_type":"VARCHAR","display_name":"Customer name","semantic_context_not_sql_identifier":{"display_name":"Customer name"}}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE orders_model ( + CustName VARCHAR ); - """, """ - CREATE TABLE dbo_technicians ( - id INTEGER, - name VARCHAR - ); - """, ] - message = unsupported_schema_message( - "Show the number of failures by technician.", + sql = normalize_sql_with_schema_identifiers( + """ + SELECT o.CustomerName + FROM dbo.tblOrders o + WHERE LOWER(o.Customer_Name) = LOWER('lockheed martin') + """, contexts, ) - assert message is not None - assert "failure/defect" in message - assert "technician" in message + assert 'o."CustName"' in sql + assert "CustomerName" not in sql + assert "Customer_Name" not in sql + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None -def test_unsupported_schema_generation_result_has_no_invalid_sql(): +def test_schema_identifier_normalization_keeps_ambiguous_source_table_invalid(): contexts = [ """ - CREATE TABLE dbo_report_failures ( - id INTEGER, - failure_type VARCHAR + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"order_archive","sql_column_names_use_exactly":["id"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"orders"}},"columns":[{"sql_column_name_use_exactly":"id","data_type":"VARCHAR"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_archive ( + id VARCHAR ); """, """ - CREATE TABLE dbo_technicians ( - id INTEGER, - name VARCHAR + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"order_current","sql_column_names_use_exactly":["id"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"orders"}},"columns":[{"sql_column_name_use_exactly":"id","data_type":"VARCHAR"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_current ( + id VARCHAR ); """, ] - result = unsupported_schema_generation_result( - "Show the number of failures by technician.", - contexts, - data_source="MSSQL", - ) - - assert result is not None - assert result["valid_generation_result"] == {} - invalid = result["invalid_generation_result"] - assert invalid["type"] == "NO_RELEVANT_SQL" - assert invalid["sql"] == "" - assert invalid["original_sql"] == "" - assert "technician" in invalid["error"] - - -def test_post_processor_clears_sql_for_unsupported_schema(): - contexts = [ - """ - CREATE TABLE dbo_mbrTime ( - id1 INTEGER, - id2 INTEGER - ); - """ - ] - post_processor = SQLGenPostProcessor(engine=None) - - result = asyncio.run( - post_processor.run( - [ - """ - SELECT id1, COUNT(*) AS failures - FROM dbo_mbrTime - GROUP BY id1 - ORDER BY failures DESC - LIMIT 10 - """ - ], - contexts=contexts, - fallback_query="Show the top 10 materials with the highest number of failures.", - data_source="MSSQL", - ) - ) + sql = normalize_sql_with_schema_identifiers("SELECT id FROM dbo.orders", contexts) + error = _SchemaCatalog.from_contexts(contexts).validate_sql(sql) - assert result["valid_generation_result"] == {} - assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" - assert result["invalid_generation_result"]["sql"] == "" - assert result["invalid_generation_result"]["original_sql"] == "" + assert "dbo.orders" in sql + assert error is not None + assert "dbo.orders" in error -def test_wren_sql_dialect_normalization_repairs_top_and_joined_limit(): +def test_wren_sql_dialect_normalization_handles_top_and_joined_limit(): assert ( normalize_wren_sql_dialect("SELECT TOP 10 id1 FROM dbo_mbrTime") == "SELECT id1 FROM dbo_mbrTime\nLIMIT 10" ) assert ( normalize_wren_sql_dialect( - "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESCLIMIT 10" + "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESCLIMIT 10" ) - == "SELECT id1 FROM dbo_mbrTime ORDER BY failures DESC LIMIT 10" + == "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESC LIMIT 10" ) -def test_repair_fallback_filters_critical_priority_and_in_progress_status(): +def test_semantic_coverage_rejects_unrepresented_query_terms(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMPTZ + CREATE TABLE neutral_records ( + id1 INTEGER, + id2 INTEGER ); """ ] - sql = generate_simple_analytics_sql( - "Show all critical-priority repairs that are currently in progress.", + error = validate_sql_semantic_coverage( + """ + SELECT id1, COUNT(*) AS record_count + FROM neutral_records + GROUP BY id1 + ORDER BY record_count DESC + LIMIT 10 + """, + "Show top 10 records by missing_dimension.", contexts, ) - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert "\"status\" = 'in progress'" in sql - assert "\"priority\" = 'critical'" in sql + assert error is not None + assert "missing" in error or "dimension" in error -def test_repair_fallback_preserves_hyphenated_in_progress_status_value(): +def test_unsupported_schema_message_reports_partial_coverage(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR ); """ ] - sql = generate_simple_analytics_sql( - "Show all repairs with a critical priority and an in-progress status.", + message = unsupported_schema_message( + "Show records by phase and unknown_segment.", contexts, ) - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert "\"status\" = 'in-progress'" in sql - assert "\"priority\" = 'critical'" in sql + assert message is not None + assert "unknown" in message or "segment" in message -def test_repair_logs_highest_priority_orders_by_verified_priority_column(): +def test_unsupported_schema_generation_result_has_no_invalid_sql(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR ); """ ] - sql = generate_simple_analytics_sql( - "Which repair logs have the highest priority?", + result = unsupported_schema_generation_result( + "Show records by unknown_segment.", contexts, + data_source="MSSQL", ) - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'ORDER BY CASE LOWER("priority")' in sql - assert "DESC" in sql + assert result is not None + assert result["valid_generation_result"] == {} + invalid = result["invalid_generation_result"] + assert invalid["type"] == "NO_RELEVANT_SQL" + assert invalid["sql"] == "" + assert invalid["original_sql"] == "" + assert "unknown" in invalid["error"] or "segment" in invalid["error"] -def test_critical_priority_repairs_filter_verified_priority_column(): +def test_schema_grounding_failure_message_reports_terms_split_across_tables(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP + CREATE TABLE user_activity ( + username VARCHAR ); + """, """ + CREATE TABLE record_counts ( + recordcnt INTEGER + ); + """, ] - sql = generate_simple_analytics_sql( - "Show all critical-priority repairs", + message = schema_grounding_failure_message( + "Show top 5 username by recordcnt.", contexts, ) - assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert "\"priority\" = 'critical'" in sql + assert "active project" in message + assert "username" in message + assert "recordcnt" in message + assert "Generated SQL referenced" not in message -def test_repairs_by_status_counts_verified_repair_rows(): +def test_schema_coverage_accepts_generic_word_form_variants(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP ); """ ] - sql = generate_simple_analytics_sql( - "Show repairs by status", + sql = """ + SELECT + CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT) AS year, + CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) AS month, + COUNT(*) AS record_count + FROM work_update_log + GROUP BY + CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT), + CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) + """ + + error = validate_sql_semantic_coverage( + sql, + "Show the number of work updates updated each month.", contexts, ) - assert sql is not None - assert 'SELECT "status", COUNT("id") AS "record_count"' in sql - assert 'FROM "dbo_repair_logs"' in sql - assert 'GROUP BY "status"' in sql + assert error is None + assert unsupported_schema_message( + "Show the number of work updates updated each month.", + contexts, + ) is None -def test_latest_repair_logs_orders_by_verified_date_column(): +def test_schema_fallback_uses_verified_monthly_update_timestamp(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - created_at TIMESTAMP + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP ); """ ] sql = generate_simple_analytics_sql( - "Show latest repair logs", + "Show the number of work updates updated each month.", contexts, ) assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'ORDER BY "created_at" DESC' in sql + assert 'FROM "work_update_log"' in sql + assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT)' in sql + assert "COUNT(*)" in sql -def test_semantic_column_alias_can_satisfy_priority_concept_with_verified_name(): +def test_schema_fallback_uses_unambiguous_implicit_text_value_filter(): contexts = [ """ /* WREN RETRIEVED SEMANTIC CONTEXT - {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"repair log records"},"columns":[{"sql_column_name_use_exactly":"Urgency","data_type":"VARCHAR","semantic_context_not_sql_identifier":"priority severity for a repair"}]} + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"product_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Product name. Use for product analysis."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} WREN SQL IDENTIFIER CONTRACT */ - CREATE TABLE dbo_work_items ( - id VARCHAR, - Urgency VARCHAR, - created_at TIMESTAMP + CREATE TABLE order_records ( + customer_name VARCHAR, + product_name VARCHAR, + order_id VARCHAR ); """ ] sql = generate_simple_analytics_sql( - "Which repair records have the highest priority?", + "Show orders from Lockheed Martine.", contexts, ) assert sql is not None - assert 'FROM "dbo_work_items"' in sql - assert '"Urgency"' in sql - assert '"priority"' not in sql + assert 'FROM "order_records"' in sql + assert 'LOWER("customer_name") LIKE \'%lockheed martine%\'' in sql -def test_failure_by_technician_fallback_uses_verified_tech_column(): +def test_schema_fallback_allows_punctuated_customer_value_after_for(): contexts = [ """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - Tech VARCHAR, - Failed VARCHAR, - Material VARCHAR + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters.","sample_values":["LOCKHEED MARTIN CORPORATION"]},{"sql_column_name_use_exactly":"product_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Product name. Use for product analysis."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + product_name VARCHAR, + order_id VARCHAR ); """ ] + query = "Show orders for CATERPILLAR S.A.R.L." - sql = generate_simple_analytics_sql( - "Show the number of failures by technician.", - contexts, - ) + sql = generate_simple_analytics_sql(query, contexts) + assert unsupported_schema_message(query, contexts) is None + assert validate_sql_semantic_coverage( + """ + SELECT customer_name, order_id + FROM order_records + WHERE LOWER(customer_name) LIKE '%caterpillar%' + """, + query, + contexts, + ) is None assert sql is not None - assert 'FROM "dbo_DebugEntries_Staging2"' in sql - assert 'SELECT "Tech", COUNT("Failed") AS "record_count"' in sql - assert 'WHERE ("Failed" IS NOT NULL AND "Failed" <> \'\')' in sql + assert 'FROM "order_records"' in sql + assert 'LOWER("customer_name") LIKE \'%caterpillar s.a.r.l%\'' in sql -def test_failure_by_material_fallback_uses_verified_material_column(): +def test_schema_fallback_keeps_customer_value_separate_from_grouping_phrase(): contexts = [ """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - Tech VARCHAR, - Failed VARCHAR, - Material VARCHAR + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"product_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Product name. Use for product analysis."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + product_name VARCHAR, + order_id VARCHAR ); """ ] + query = "Show orders from PACCAR PARTS DIVISION by product." - sql = generate_simple_analytics_sql( - "Show failures by material.", - contexts, - ) + sql = generate_simple_analytics_sql(query, contexts) + assert unsupported_schema_message(query, contexts) is None assert sql is not None - assert 'FROM "dbo_DebugEntries_Staging2"' in sql - assert 'SELECT "Material", COUNT("Failed") AS "record_count"' in sql + assert 'LOWER("customer_name") LIKE \'%paccar parts division%\'' in sql + assert 'GROUP BY "product_name"' in sql -def test_failure_type_value_filter_uses_verified_failure_type_column(): +def test_schema_fallback_stops_grouping_phrase_before_from_table_name(): contexts = [ """ - CREATE TABLE dbo_DebugEntries ( - SerialNumber VARCHAR, - FailedAt VARCHAR, - Material VARCHAR - ); - """, - """ - CREATE TABLE dbo_repair_logs ( - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Business unit activity records."},"columns":[{"sql_column_name_use_exactly":"bunit","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Business unit grouping code."},{"sql_column_name_use_exactly":"record_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Record identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE business_unit_records ( + bunit VARCHAR, + record_id VARCHAR ); - """, """ - CREATE TABLE dbo_report_failures ( - failure_type VARCHAR, - failure_line VARCHAR, - test_name VARCHAR - ); - """, ] + query = "Show row counts grouped by bunit from business_unit_records." - sql = generate_simple_analytics_sql( - "Show the number of units with JTAG as the failure type.", - contexts, - ) + sql = generate_simple_analytics_sql(query, contexts) + assert unsupported_schema_message(query, contexts) is None assert sql is not None - assert 'FROM "dbo_report_failures"' in sql - assert 'COUNT(*) AS "record_count"' in sql - assert "\"failure_type\" = 'JTAG'" in sql + assert 'FROM "business_unit_records"' in sql + assert 'GROUP BY "bunit"' in sql + assert validate_sql_semantic_coverage(sql, query, contexts) is None -def test_board_models_most_failures_counts_failure_records_not_defect_rate(): +def test_schema_fallback_groups_monthly_record_request_without_count_word(): contexts = [ """ - CREATE TABLE dbo_batch_records ( - board_model VARCHAR, - supplier VARCHAR, - defect_rate DECIMAL + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"order_date","data_type":"DATE","semantic_context_not_sql_identifier":"Order placement date."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + order_date DATE, + order_id VARCHAR ); - """, """ - CREATE TABLE dbo_repair_logs ( - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR, - priority VARCHAR - ); - """, ] sql = generate_simple_analytics_sql( - "Show the top 5 board models with the most failures.", + "Show monthly orders from JOHN DEERE COMMERCIAL PRODUCTS.", contexts, ) assert sql is not None - assert 'FROM "dbo_repair_logs"' in sql - assert 'SELECT "board_model", COUNT("failure_code") AS "record_count"' in sql - assert '"defect_rate"' not in sql - assert "LIMIT 5" in sql + assert 'FROM "order_records"' in sql + assert 'CAST(EXTRACT(YEAR FROM "order_date") AS BIGINT)' in sql + assert 'CAST(EXTRACT(MONTH FROM "order_date") AS BIGINT)' in sql + assert 'COUNT(*) AS "record_count"' in sql + assert 'LOWER("customer_name") LIKE \'%john deere commercial products%\'' in sql -def test_board_models_highest_defect_rate_uses_rate_metric(): +def test_schema_fallback_does_not_treat_customer_suffix_as_required_schema(): contexts = [ """ - CREATE TABLE dbo_batch_records ( - board_model VARCHAR, - supplier VARCHAR, - defect_rate DECIMAL + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + order_id VARCHAR ); """, """ - CREATE TABLE dbo_repair_logs ( - board_model VARCHAR, - failure_code VARCHAR, - status VARCHAR + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Refund timing records."},"columns":[{"sql_column_name_use_exactly":"day_s_from_refund","data_type":"INTEGER","semantic_context_not_sql_identifier":"Refund timing days."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE refund_records ( + day_s_from_refund INTEGER ); """, ] + query = "Show orders for CATERPILLAR S.A.R.L." - sql = generate_simple_analytics_sql( - "Show the board models with the highest defect rate.", - contexts, - ) - - assert sql is not None - assert 'FROM "dbo_batch_records"' in sql - assert 'SELECT "board_model", AVG("defect_rate") AS "average_value"' in sql - assert 'ORDER BY "average_value" DESC' in sql + assert unsupported_schema_message(query, contexts) is None -def test_semantic_coverage_rejects_rate_for_failure_count_intent(): +def test_schema_fallback_prefers_customer_semantics_over_value_word_overlap(): contexts = [ """ - CREATE TABLE dbo_batch_records ( - board_model VARCHAR, - defect_rate DECIMAL + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Use this table for new order analysis."},"columns":[{"sql_column_name_use_exactly":"CustName","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use this column when the user asks for customer, customer name, account, or buyer.","display_name":"Customer name","source_column_name":"CustName"},{"sql_column_name_use_exactly":"Division","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Business division. Use for division-level reporting or grouping.","display_name":"Division","source_column_name":"Division"},{"sql_column_name_use_exactly":"OrdDate","data_type":"DATE","semantic_context_not_sql_identifier":"Order placement date."},{"sql_column_name_use_exactly":"OrdNo","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Sales order number."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + Division VARCHAR, + OrdDate DATE, + OrdNo VARCHAR + ); + """ + ] + query = "Show recent orders from PACCAR PARTS DIVISION." + + sql = generate_simple_analytics_sql(query, contexts) + + assert sql is not None + assert 'FROM "dbo_tblNewOrders"' in sql + assert 'LOWER("CustName") LIKE \'%paccar parts division%\'' in sql + assert "Division) = 'paccar parts division'" not in sql + assert 'ORDER BY "OrdDate" DESC' in sql + + +def test_schema_fallback_prefers_direct_table_name_match_when_metadata_overlaps(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Use this table for order and customer questions."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE sales_history ( + customer_name VARCHAR, + order_id VARCHAR + ); + """, + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Use this table for order and customer questions."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE new_orders ( + customer_name VARCHAR, + order_id VARCHAR + ); + """, + ] + + sql = generate_simple_analytics_sql("Show orders from Acme.", contexts) + + assert sql is not None + assert 'FROM "new_orders"' in sql + + +def test_schema_coverage_uses_word_form_variants_per_table(): + contexts = [ + """ + CREATE TABLE report_failures ( + report_id VARCHAR, + failure_line VARCHAR + ); + """, + """ + CREATE TABLE failure_patterns ( + severity VARCHAR, + failure_type VARCHAR + ); + """, + ] + query = "Show failures by severity." + + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "failure_patterns"' in sql + assert 'GROUP BY "severity"' in sql + + +def test_schema_fallback_splits_compound_identifier_names_for_grouping_and_measure(): + contexts = [ + """ + CREATE TABLE payable_invoices ( + bunit VARCHAR, + suppliername VARCHAR, + grossamount DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show total gross amount by supplier name.", + contexts, + ) + + assert sql is not None + assert 'SELECT "suppliername", SUM("grossamount") AS "total_value"' in sql + assert 'GROUP BY "suppliername"' in sql + + +def test_semantic_validation_rejects_weaker_explicit_grouping_column(): + contexts = [ + """ + CREATE TABLE payable_invoices ( + bunit VARCHAR, + suppliername VARCHAR, + grossamount DECIMAL ); """ ] error = validate_sql_semantic_coverage( """ - SELECT board_model, defect_rate - FROM dbo_batch_records - ORDER BY defect_rate DESC - LIMIT 5 + SELECT bunit, SUM(grossamount) AS total_value + FROM payable_invoices + GROUP BY bunit + ORDER BY total_value DESC """, - "Show the top 5 board models with the most failures.", + "Show total gross amount by supplier name.", contexts, ) assert error is not None - assert "count of failure records" in error + assert "weaker matching column" in error or "grouping dimension" in error -def test_repairs_by_technician_requires_one_schema_object_covering_both_concepts(): +def test_semantic_validation_accepts_compound_status_grouping_column(): contexts = [ """ - CREATE TABLE dbo_repair_logs ( - id VARCHAR, - status VARCHAR, - priority VARCHAR, - failure_code VARCHAR + CREATE TABLE task_rollups ( + taskstatus VARCHAR, + task_id VARCHAR ); """, + ] + + error = validate_sql_semantic_coverage( + """ + SELECT taskstatus, COUNT(*) AS record_count + FROM task_rollups + GROUP BY taskstatus + """, + "How many task records are there by task status?", + contexts, + ) + + assert error is None + + +def test_fallback_limit_accepts_spelled_out_top_number(): + assert _fallback_limit("Show the top five groups from that result.") == 5 + + +def test_followup_group_result_words_do_not_require_schema_columns(): + contexts = [ """ - CREATE TABLE dbo_DebugEntries_Staging2 ( - Tech VARCHAR, - Failed VARCHAR + CREATE TABLE repair_logs ( + status VARCHAR, + repair_id VARCHAR ); """, ] - message = unsupported_schema_message("Show repairs by technician.", contexts) + assert ( + unsupported_schema_message( + "Show the top five groups from the repair records by status.", + contexts, + ) + is None + ) + + +def test_unsupported_subject_still_blocks_cross_domain_value_filter(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + order_id VARCHAR + ); + """ + ] + query = "Show tickets from ACME CORP." + + message = unsupported_schema_message(query, contexts) assert message is not None - assert "repair" in message - assert "technician" in message + assert "ticket" in message.lower() + assert "acme" not in message.lower() + + +def test_schema_fallback_skips_ambiguous_implicit_text_value_filter(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records."},"columns":[{"sql_column_name_use_exactly":"buyer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Buyer name."},{"sql_column_name_use_exactly":"seller_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Seller name."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + buyer_name VARCHAR, + seller_name VARCHAR, + order_id VARCHAR + ); + """ + ] + + assert ( + generate_simple_analytics_sql("Show orders from Acme Industries.", contexts) + is None + ) + + +def test_schema_catalog_ignores_extract_from_column_clause(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + sql = """ + SELECT CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year", COUNT(*) AS "record_count" + FROM "work_update_log" + GROUP BY CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) + """ + + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + + +def test_post_processor_clears_sql_for_unsupported_schema(): + contexts = [ + """ + CREATE TABLE neutral_records ( + id1 INTEGER, + id2 INTEGER + ); + """ + ] + post_processor = SQLGenPostProcessor(engine=None) + + result = asyncio.run( + post_processor.run( + [ + """ + SELECT id1, COUNT(*) AS record_count + FROM neutral_records + GROUP BY id1 + ORDER BY record_count DESC + LIMIT 10 + """ + ], + contexts=contexts, + fallback_query="Show records by missing_dimension.", + data_source="MSSQL", + ) + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert result["invalid_generation_result"]["sql"] == "" + assert result["invalid_generation_result"]["original_sql"] == "" + + +def test_post_processor_prefers_fact_table_fallback_over_dimension_only_sql(): + contexts = [ + """ + CREATE TABLE invoice_records ( + invoice_id VARCHAR, + business_unit VARCHAR + ); + """, + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"invoice business unit grouping lookup"},"columns":[{"sql_column_name_use_exactly":"name","display_name":"Business Unit Group"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE business_unit_groups ( + name VARCHAR + ); + """, + ] + post_processor = SQLGenPostProcessor(engine=_AcceptingEngine()) + + result = asyncio.run( + post_processor.run( + [ + """ + SELECT name, COUNT(*) AS record_count + FROM business_unit_groups + GROUP BY name + ORDER BY record_count DESC + """ + ], + contexts=contexts, + fallback_query="How many invoice records are there by business unit?", + data_source="MSSQL", + use_dry_plan=True, + ) + ) + + sql = result["valid_generation_result"]["sql"] + assert result["invalid_generation_result"] == {} + assert "invoice_records" in sql + assert "business_unit" in sql + assert "business_unit_groups" not in sql + + +def test_post_processor_converts_invented_table_to_schema_message(): + contexts = [ + """ + CREATE TABLE user_activity ( + username VARCHAR + ); + """, + """ + CREATE TABLE record_counts ( + recordcnt INTEGER + ); + """, + ] + post_processor = SQLGenPostProcessor(engine=None) + + result = asyncio.run( + post_processor.run( + [ + """ + SELECT username + FROM records + ORDER BY recordcnt DESC + LIMIT 5 + """ + ], + contexts=contexts, + fallback_query="Show top 5 username by recordcnt.", + data_source="MSSQL", + ) + ) + + invalid = result["invalid_generation_result"] + assert result["valid_generation_result"] == {} + assert invalid["type"] == "NO_RELEVANT_SQL" + assert invalid["sql"] == "" + assert invalid["original_sql"] == "" + assert "active project" in invalid["error"] + assert "username" in invalid["error"] + assert "recordcnt" in invalid["error"] + assert "Generated SQL referenced" not in invalid["error"] + + +def test_schema_sample_value_filter_is_grounded_in_metadata(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done","In Progress"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE work_items ( + item_id VARCHAR, + State VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all work item records with In Progress.", + contexts, + ) + + assert sql is not None + assert 'FROM "work_items"' in sql + assert 'LOWER("State") = \'in progress\'' in sql + + +def test_user_values_are_allowed_for_single_verified_text_column(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + state_name VARCHAR(255), + updated_at TIMESTAMP + ); + """ + ] + + query = ( + "Show the distribution of work updates across completed and " + "in-progress state names." + ) + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'LOWER("state_name") IN (\'completed\', \'in-progress\')' in sql + assert 'GROUP BY "state_name"' in sql + + +def test_explicit_column_adjacent_value_after_column_is_filter_value(): + contexts = [ + """ + CREATE TABLE work_items ( + item_id VARCHAR, + status VARCHAR, + priority VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + query = "Show work items with status open." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_items"' in sql + assert 'LOWER("status") = \'open\'' in sql + + +def test_explicit_column_value_survives_multi_column_semantic_ambiguity(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Support ticket records."},"columns":[{"sql_column_name_use_exactly":"title","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket title."},{"sql_column_name_use_exactly":"description","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket description."},{"sql_column_name_use_exactly":"status","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket status."},{"sql_column_name_use_exactly":"priority","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket priority."},{"sql_column_name_use_exactly":"data","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket payload."},{"sql_column_name_use_exactly":"created_at","data_type":"TIMESTAMP","semantic_context_not_sql_identifier":"Ticket creation time."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE support_tickets ( + title VARCHAR, + description VARCHAR, + status VARCHAR, + priority VARCHAR, + data VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show tickets with status open." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "support_tickets"' in sql + assert 'LOWER("status") = \'open\'' in sql + + +def test_explicit_column_value_allows_value_token_seen_elsewhere_in_schema(): + contexts = [ + """ + CREATE TABLE production_batches ( + batch_id VARCHAR, + supplier VARCHAR, + inspection_status VARCHAR, + created_at TIMESTAMP + ); + """, + """ + CREATE TABLE debug_entries ( + entry_id VARCHAR, + failed_at TIMESTAMP + ); + """, + ] + + query = "Show batches with inspection status failed." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "production_batches"' in sql + assert 'LOWER("inspection_status") = \'failed\'' in sql + + +def test_explicit_column_adjacent_value_before_column_is_filter_value(): + contexts = [ + """ + CREATE TABLE support_tickets ( + ticket_id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show high priority tickets." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "support_tickets"' in sql + assert 'LOWER("priority") = \'high\'' in sql + assert 'LOWER("title")' not in sql + assert 'LOWER("description")' not in sql + assert 'LOWER("status") = \'high priority\'' not in sql + assert 'LOWER("data")' not in sql + + +def test_explicit_preceding_value_survives_multi_column_semantic_ambiguity(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Support ticket records."},"columns":[{"sql_column_name_use_exactly":"title","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket title."},{"sql_column_name_use_exactly":"description","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket description."},{"sql_column_name_use_exactly":"status","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket status."},{"sql_column_name_use_exactly":"priority","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket priority."},{"sql_column_name_use_exactly":"data","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket payload."},{"sql_column_name_use_exactly":"created_at","data_type":"TIMESTAMP","semantic_context_not_sql_identifier":"Ticket creation time."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE support_tickets ( + title VARCHAR, + description VARCHAR, + status VARCHAR, + priority VARCHAR, + data VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show high priority tickets." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "support_tickets"' in sql + assert 'LOWER("priority") = \'high\'' in sql + + +def test_explicit_identifier_style_column_value_filter_is_grounded(): + contexts = [ + """ + CREATE TABLE repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show repairs for failure code BGA-001." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "repair_logs"' in sql + assert 'LOWER("failure_code") = \'bga-001\'' in sql + + +def test_open_filter_phrase_strips_explicit_column_name_from_value(): + contexts = [ + """ + CREATE TABLE production_batches ( + batch_id VARCHAR, + supplier VARCHAR, + board_model VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show batches from supplier Wurth Elektronik." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "production_batches"' in sql + assert 'LOWER("supplier") = \'wurth elektronik\'' in sql + assert "supplier wurth" not in sql.lower() + + +def test_column_value_label_is_not_treated_as_literal_filter_value(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + state_name VARCHAR(255), + updated_at TIMESTAMP + ); + """ + ] + + query = "Show the distribution of work updates across state name values." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'GROUP BY "state_name"' in sql + assert "WHERE" not in sql + + +def test_subject_noun_is_not_treated_as_literal_filter_value(): + contexts = [ + """ + CREATE TABLE purchase_order_records ( + purchase_order_id VARCHAR, + currency_code VARCHAR, + order_date DATE, + order_quantity DECIMAL, + record_type VARCHAR, + record_status VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "How many purchase orders are there by currency?", + contexts, + ) + + assert unsupported_schema_message( + "How many purchase orders are there by currency?", + contexts, + ) is None + assert sql is not None + assert 'FROM "purchase_order_records"' in sql + assert 'GROUP BY "currency_code"' in sql + assert "WHERE" not in sql + + +def test_grouping_dimension_is_not_treated_as_subject_column_literal(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"sales records with repair item status reporting metadata"},"columns":[{"sql_column_name_use_exactly":"RepairItem","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Repair item"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE sales_records ( + RepairItem VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "How many repair records are there by repair status?", + contexts, + ) + + assert sql is None + assert validate_sql_semantic_coverage( + """ + SELECT "RepairItem", COUNT(*) AS "record_count" + FROM "sales_records" + WHERE LOWER("RepairItem") = 'status' + GROUP BY "RepairItem" + """, + "How many repair records are there by repair status?", + contexts, + ) + + +def test_subject_entity_is_not_grounded_by_sample_value_filter(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Tariff liquidation records."},"columns":[{"sql_column_name_use_exactly":"LiquidationStatus","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Liquidation status.","sample_values":["Repair","Complete"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE tariff_records ( + LiquidationStatus VARCHAR, + Entry_Date TIMESTAMP + ); + """ + ] + query = "How many repair records are there by repair status?" + + sql = generate_simple_analytics_sql(query, contexts) + message = unsupported_schema_message(query, contexts) + error = validate_sql_semantic_coverage( + """ + SELECT LiquidationStatus, COUNT(*) AS record_count + FROM tariff_records + WHERE LOWER(LiquidationStatus) = 'repair' + GROUP BY LiquidationStatus + """, + query, + contexts, + ) + + assert sql is None + assert message is not None + assert "repair" in message.lower() + assert error is not None + assert "repair" in error.lower() + + +def test_compact_accounting_identifiers_support_balance_and_currency_questions(): + balance_contexts = [ + """ + CREATE TABLE balance_records ( + bunit VARCHAR, + period VARCHAR, + endbalance FLOAT8 + ); + """ + ] + balance_sql = generate_simple_analytics_sql( + "Show total ending balance by period and business unit.", + balance_contexts, + ) + + assert balance_sql is not None + assert 'SUM("endbalance") AS "total_value"' in balance_sql + assert 'GROUP BY "bunit", "period"' in balance_sql or 'GROUP BY "period", "bunit"' in balance_sql + assert "WHERE" not in balance_sql + + exchange_contexts = [ + """ + CREATE TABLE exchange_rate_records ( + currencyfrom VARCHAR, + currencyto VARCHAR, + exchangerate FLOAT8 + ); + """ + ] + exchange_sql = generate_simple_analytics_sql( + "Show exchange rates by currency pair.", + exchange_contexts, + ) + + assert exchange_sql is not None + assert 'FROM "exchange_rate_records"' in exchange_sql + assert '"exchangerate"' in exchange_sql + assert '"currencyfrom"' in exchange_sql or '"currencyto"' in exchange_sql + + +def test_semantic_validation_rejects_subject_noun_literal_filter(): + contexts = [ + """ + CREATE TABLE purchase_order_records ( + purchase_order_id VARCHAR, + currency_code VARCHAR, + order_date DATE, + order_quantity DECIMAL, + record_type VARCHAR, + record_status VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT currency_code, COUNT(*) AS record_count + FROM purchase_order_records + WHERE LOWER(record_type) = 'purchase' + GROUP BY currency_code + """, + "How many purchase orders are there by currency?", + contexts, + ) + + assert error is not None + assert "not grounded as a filter value" in error + + +def test_unverified_filter_value_is_not_invented(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE work_items ( + item_id VARCHAR, + State VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all work item records with Archived.", + contexts, + ) + message = unsupported_schema_message( + "Show all work item records with Archived.", + contexts, + ) + + assert sql is None + assert message is not None + assert "archived" in message.lower() + + +def test_grouped_count_uses_verified_dimension_only(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql("Show records by phase.", contexts) + + assert sql is not None + assert 'SELECT "phase", COUNT(*) AS "record_count"' in sql + assert 'FROM "event_records"' in sql + assert 'GROUP BY "phase"' in sql + + +def test_average_uses_verified_numeric_measure_not_count(): + contexts = [ + """ + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR, + age_days DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql("Show average age by model.", contexts) + + assert sql is not None + assert 'SELECT "model_code", AVG("age_days") AS "average_value"' in sql + assert 'GROUP BY "model_code"' in sql + assert "COUNT(" not in sql + + +def test_average_without_verified_measure_is_unsupported(): + contexts = [ + """ + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show average age by model.", contexts) + message = unsupported_schema_message("Show average age by model.", contexts) + + assert sql is None + assert message is not None + assert "age" in message.lower() + + +def test_latest_uses_verified_temporal_column(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + event_time TIMESTAMP, + phase VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show latest event records.", contexts) + + assert sql is not None + assert 'FROM "event_records"' in sql + assert 'ORDER BY "event_time" DESC' in sql + + +def test_count_mentions_subject_without_implicit_grouping(): + contexts = [ + """ + CREATE TABLE invoice_records ( + invoice_number VARCHAR, + invoice_date TIMESTAMP, + invoice_type VARCHAR, + business_unit VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "How many invoice records are there?", + contexts, + ) + + assert sql is not None + assert sql.strip() == ( + 'SELECT COUNT(*) AS "record_count"\n' + 'FROM "invoice_records"' + ) + assert "GROUP BY" not in sql + + +def test_latest_by_temporal_column_does_not_become_grouped_count(): + contexts = [ + """ + CREATE TABLE account_reconciliation_records ( + account_number VARCHAR, + approval_date TIMESTAMP, + status VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the latest account reconciliation records by approval date.", + contexts, + ) + + assert sql is not None + assert 'FROM "account_reconciliation_records"' in sql + assert 'ORDER BY "approval_date" DESC' in sql + assert "COUNT(*)" not in sql + assert "GROUP BY" not in sql + + +def test_monthly_count_uses_requested_temporal_column_when_verified(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + updated_at TIMESTAMP, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the number of event records updated each month.", + contexts, + ) + + assert sql is not None + assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year"' in sql + assert 'CAST(EXTRACT(MONTH FROM "updated_at") AS BIGINT) AS "month"' in sql + assert 'COUNT(*) AS "record_count"' in sql + + +def test_order_by_uses_verified_column_and_sample_value(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open","Closed"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE case_records ( + case_id VARCHAR, + State VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all case records with Open ordered by case ID.", + contexts, + ) + + assert sql is not None + assert 'LOWER("State") = \'open\'' in sql + assert 'ORDER BY "case_id" ASC' in sql + + +def test_top_grouped_count_is_schema_shape_based(): + contexts = [ + """ + CREATE TABLE occurrence_records ( + occurrence_id VARCHAR, + model_code VARCHAR, + reason_code VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show top 5 occurrence records by model.", + contexts, + ) + + assert sql is not None + assert 'SELECT "model_code", COUNT(*) AS "record_count"' in sql + assert 'ORDER BY "record_count" DESC' in sql + assert "LIMIT 5" in sql + + +def test_single_grouping_dimension_does_not_over_split_results(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show number of events by account.", contexts) + + assert sql is not None + assert 'GROUP BY "account_name"' in sql + assert "account_reference" not in sql + + +def test_missing_value_intent_uses_verified_plural_name_column(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + event_time TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show events with missing account names.", + contexts, + ) + + assert sql is not None + assert 'FROM "account_events"' in sql + assert '"account_name" IS NULL' in sql + + +def test_semantic_validation_rejects_weaker_null_check_column(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT account_name, account_reference + FROM account_events + WHERE account_reference IS NULL + """, + "Show events with missing account names.", + contexts, + ) + + assert error is not None + assert "weaker matching column" in error + + +def test_top_records_are_listed_without_implicit_grouped_aggregate(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show top 10 scored events from July.", contexts) + + assert sql is not None + assert 'FROM "scored_events"' in sql + assert "GROUP BY" not in sql + assert 'ORDER BY "score_value" DESC' in sql + assert "LIMIT 10" in sql + + +def test_semantic_validation_accepts_generated_month_date_range(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + query = "Show top 10 scored events from July." + sql = generate_simple_analytics_sql(query, contexts) + + assert sql is not None + assert validate_sql_semantic_coverage(sql, query, contexts) is None + + +def test_top_records_without_verified_rank_measure_uses_date_ordering(): + contexts = [ + """ + CREATE TABLE order_records ( + order_id VARCHAR, + order_date TIMESTAMP, + fx_currency DECIMAL, + customer_name VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show top 10 order records from July.", contexts) + + assert sql is not None + assert 'FROM "order_records"' in sql + assert 'ORDER BY "order_date" DESC' in sql + assert 'fx_currency' not in sql.split("ORDER BY", maxsplit=1)[1] + assert "LIMIT 10" in sql + + +def test_top_rows_by_numeric_column_orders_without_grouping(): + contexts = [ + """ + CREATE TABLE failure_patterns ( + pattern_id VARCHAR, + name VARCHAR, + severity VARCHAR, + occurrences INTEGER, + cost_impact DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show top 10 failure patterns by occurrences.", + contexts, + ) + + assert sql is not None + assert 'FROM "failure_patterns"' in sql + assert "GROUP BY" not in sql + assert 'ORDER BY "occurrences" DESC' in sql + assert "LIMIT 10" in sql + + +def test_sum_by_year_uses_verified_measure_and_temporal_column(): + contexts = [ + """ + CREATE TABLE transaction_records ( + transaction_id VARCHAR, + account_name VARCHAR, + amount_value DECIMAL, + posted_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql("Show total amount by year.", contexts) + + assert sql is not None + assert 'CAST(EXTRACT(YEAR FROM "posted_at") AS BIGINT) AS "year"' in sql + assert 'SUM("amount_value") AS "total_value"' in sql + + +def test_semantic_coverage_rejects_count_for_average_intent(): + contexts = [ + """ + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR, + age_days DECIMAL + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT model_code, COUNT(*) AS record_count + FROM measurement_records + GROUP BY model_code + """, + "Show average age by model.", + contexts, + ) + + assert error is not None + assert "average" in error.lower() + + +def test_literal_validation_rejects_values_outside_verified_samples(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE case_records ( + case_id VARCHAR, + State VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT case_id + FROM case_records + WHERE LOWER(State) = 'Closed' + """, + "Show case records with Open.", + contexts, + ) + + assert error is not None + assert "sample values" in error + + +def test_semantic_validation_rejects_multi_group_for_single_dimension(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT account_name, account_reference, COUNT(*) AS record_count + FROM account_events + GROUP BY account_name, account_reference + """, + "Show number of events by account.", + contexts, + ) + + assert error is not None + assert "one grouping dimension" in error + + +def test_semantic_validation_rejects_top_record_grouped_aggregate(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT category_name, SUM(score_value) AS total_value + FROM scored_events + GROUP BY category_name + ORDER BY total_value DESC + LIMIT 10 + """, + "Show top 10 scored events from July.", + contexts, + ) + + assert error is not None + assert "grouped aggregate" in error diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py index 662c1d0c21..1cf5246ec6 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py @@ -1,3 +1,4 @@ +import ast from unittest.mock import AsyncMock import orjson @@ -48,10 +49,51 @@ async def test_single_model(): "type": "TABLE", "comment": "\n/* {'alias': 'user', 'description': 'A table containing user information.'} */\n", "name": "user", + "properties": { + "description": "A table containing user information.", + "displayName": "user", + }, + "tableReference": None, + "refSql": None, } ) +@pytest.mark.asyncio +async def test_model_source_table_reference_is_indexed(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "orders_model", + "properties": { + "description": "Modeled order facts.", + "displayName": "New Orders", + }, + "tableReference": { + "catalog": "warehouse", + "schema": "dbo", + "table": "tblNewOrders", + }, + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=1) + + assert len(actual["documents"]) == 1 + content = ast.literal_eval(actual["documents"][0].content) + assert content["properties"]["displayName"] == "New Orders" + assert content["tableReference"] == { + "catalog": "warehouse", + "schema": "dbo", + "table": "tblNewOrders", + } + + @pytest.mark.asyncio async def test_multiple_models(): chunker = DDLChunker() @@ -87,6 +129,12 @@ async def test_multiple_models(): "type": "TABLE", "comment": "\n/* {'alias': 'user', 'description': 'A table containing user information.'} */\n", "name": "user", + "properties": { + "description": "A table containing user information.", + "displayName": "user", + }, + "tableReference": None, + "refSql": None, } ) @@ -97,6 +145,12 @@ async def test_multiple_models(): "type": "TABLE", "comment": "\n/* {'alias': 'order', 'description': 'A table containing order details.'} */\n", "name": "order", + "properties": { + "description": "A table containing order details.", + "displayName": "order", + }, + "tableReference": None, + "refSql": None, } ) @@ -137,6 +191,7 @@ async def test_column_is_primary_key(): "name": "id", "data_type": "INTEGER", "is_primary_key": True, + "properties": {}, } ], } @@ -178,10 +233,14 @@ async def test_column_with_properties(): "columns": [ { "type": "COLUMN", - "comment": '-- {"alias":"iid","description":"The unique identifier for a user."}\n ', + "comment": '-- {"alias":"iid","description":"The unique identifier for a user.","sourceColumnName":""}\n ', "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": { + "displayName": "iid", + "description": "The unique identifier for a user.", + }, } ], } @@ -194,6 +253,9 @@ async def test_column_with_properties(): "type": "TABLE", "comment": "\n/* {'alias': '', 'description': ''} */\n", "name": "user", + "properties": {}, + "tableReference": None, + "refSql": None, } ) @@ -232,10 +294,14 @@ async def test_null_metadata_properties_are_indexed_as_empty_text(): "columns": [ { "type": "COLUMN", - "comment": '-- {"alias":null,"description":null}\n ', + "comment": '-- {"alias":null,"description":null,"sourceColumnName":""}\n ', "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": { + "displayName": None, + "description": None, + }, } ], } @@ -245,6 +311,9 @@ async def test_null_metadata_properties_are_indexed_as_empty_text(): "type": "TABLE", "comment": "\n/* {'alias': None, 'description': None} */\n", "name": "user", + "properties": {"description": None, "displayName": None}, + "tableReference": None, + "refSql": None, } ) @@ -286,10 +355,16 @@ async def test_column_with_nested_columns(): "columns": [ { "type": "COLUMN", - "comment": '-- {"alias":"iid","description":"The unique identifier for a user.","nested_columns":{"nested.address":{"name":"address","type":"VARCHAR"},"nested.orders":{"name":"orders","type":"ARRAY"}}}\n ', + "comment": '-- {"alias":"iid","description":"The unique identifier for a user.","sourceColumnName":"","nested_columns":{"nested.address":{"name":"address","type":"VARCHAR"},"nested.orders":{"name":"orders","type":"ARRAY"}}}\n ', "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": { + "displayName": "iid", + "description": "The unique identifier for a user.", + "nested.address": {"name": "address", "type": "VARCHAR"}, + "nested.orders": {"name": "orders", "type": "ARRAY"}, + }, } ], } @@ -333,6 +408,7 @@ async def test_column_with_calculated_property(): "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": {}, } ], } @@ -397,6 +473,7 @@ async def test_column_with_relationship(): "name": "id", "data_type": "INTEGER", "is_primary_key": True, + "properties": {}, } ], } @@ -452,6 +529,7 @@ async def test_column_batch_size(): "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": {}, }, { "type": "COLUMN", @@ -459,6 +537,7 @@ async def test_column_batch_size(): "name": "name", "data_type": "VARCHAR", "is_primary_key": False, + "properties": {}, }, ], } @@ -476,6 +555,7 @@ async def test_column_batch_size(): "name": "age", "data_type": "INTEGER", "is_primary_key": False, + "properties": {}, } ], } diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py index 1bf9630b62..4b2a144a41 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -1,11 +1,15 @@ import pytest +import tiktoken from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from src.pipelines.common import build_table_ddl from src.pipelines.retrieval.db_schema_retrieval import ( _augment_retrieval_query, + _build_table_retrieval_context, _build_view_ddl, + _lexical_columns_and_tables_needed, + _limit_retrieval_results_for_generation, _parse_column_selection_response, _rank_table_names_by_query, check_using_db_schemas_without_pruning, @@ -118,11 +122,45 @@ def test_view_schema_context_uses_deployed_view_statement_without_declared_colum } ) - assert "CREATE VIEW retrieved_view" in result - assert "AS SELECT modeled_column FROM deployed_model" in result + assert "CREATE TABLE retrieved_view" in result + assert "modeled_column VARCHAR" in result assert "sql_table_name_use_exactly: retrieved_view" in result +def test_table_schema_context_includes_source_identifier_metadata(): + result, _, _ = _build_table_retrieval_context( + { + "type": "TABLE", + "comment": "", + "name": "orders_model", + "properties": {"displayName": "New Orders"}, + "tableReference": { + "catalog": "warehouse", + "schema": "dbo", + "table": "tblOrders", + }, + "columns": [ + { + "type": "COLUMN", + "name": "customer_name", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + "properties": { + "displayName": "CustName", + "sourceColumnName": "CustName", + }, + } + ], + "primaryKey": "", + } + ) + + assert '"source_table_name":"dbo.tblOrders"' in result + assert '"source_table_reference":{"catalog":"warehouse","schema":"dbo","table":"tblOrders"}' in result + assert '"source_column_name":"CustName"' in result + + def test_construct_db_schemas_keeps_deployed_views_for_column_pruning(): result = construct_db_schemas( [ @@ -364,13 +402,12 @@ async def run(self, query_embedding, filters): embedding={}, ) - assert retriever.calls == [[selected_model], [related_model], [downstream_model]] + assert retriever.calls == [[selected_model], [related_model]] assert [document.meta["name"] for document in documents] == [ selected_model, selected_model, related_model, related_model, - downstream_model, ] @@ -470,6 +507,130 @@ async def run(self, query_embedding, filters): ] +@pytest.mark.asyncio +async def test_dbschema_retrieval_lexically_recovers_active_project_schema_when_vector_misses(): + recovered_model = "customer_orders" + unrelated_model = "service_tickets" + + def table_document(name, comment="", properties=None, table_reference=None): + return Document( + content=str( + { + "type": "TABLE", + "name": name, + "comment": comment, + "columns": [], + "properties": properties or {}, + "tableReference": table_reference, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": name}, + ) + + def columns_document(name, columns): + return Document( + content=str({"type": "TABLE_COLUMNS", "columns": columns}), + meta={"type": "TABLE_SCHEMA", "name": name}, + ) + + recovered_documents = [ + table_document( + recovered_model, + comment="Orders captured from customer purchase activity.", + properties={"displayName": "Customer Orders"}, + table_reference={"schema": "sales", "table": "tblCustomerOrders"}, + ), + columns_document( + recovered_model, + [ + { + "type": "COLUMN", + "name": "customer_name", + "data_type": "VARCHAR", + "comment": "Customer name on the order.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "order_total", + "data_type": "DOUBLE", + "comment": "Order revenue amount.", + "is_primary_key": False, + }, + ], + ), + ] + all_schema_documents = recovered_documents + [ + table_document(unrelated_model, comment="Support case tracking."), + columns_document( + unrelated_model, + [ + { + "type": "COLUMN", + "name": "ticket_status", + "data_type": "VARCHAR", + "comment": "Support ticket status.", + "is_primary_key": False, + } + ], + ), + ] + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + + if query_embedding: + return {"documents": []} + + is_exact_fetch = any( + isinstance(condition, dict) + and condition.get("operator") == "OR" + and condition.get("conditions") + for condition in filters["conditions"] + ) + if not is_exact_fetch: + return {"documents": all_schema_documents} + + return {"documents": recovered_documents} + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=retriever, + query="show orders by customer", + embedding={"embedding": [0.25]}, + ) + + assert [call["query_embedding"] for call in retriever.calls] == [ + [0.25], + [], + [], + ] + assert retriever.calls[1]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert [document.meta["name"] for document in documents] == [ + recovered_model, + recovered_model, + ] + + @pytest.mark.asyncio async def test_dbschema_retrieval_prefers_table_description_hits_over_schema_chunk_hits(): described_model = "described_dataset" @@ -520,8 +681,8 @@ async def run(self, query_embedding, filters): embedding={"embedding": [0.25]}, ) - assert [call["query_embedding"] for call in retriever.calls] == [[]] - assert retriever.calls[0]["filters"] == { + assert [call["query_embedding"] for call in retriever.calls] == [[0.25], []] + assert retriever.calls[1]["filters"] == { "operator": "AND", "conditions": [ {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, @@ -647,6 +808,99 @@ def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): assert result["has_metric"] is True +def test_construct_retrieval_results_skips_retrieved_schema_without_name(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_attribute"] + } + }, + { + "table_name": "semantic_metric", + "table_selection_reason": "Selected metric for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed metric."], + "columns": ["metric_value"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA"}, + ), + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "name": "semantic_metric", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "semantic_metric"}, + ), + ], + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "modeled_dataset", + "semantic_metric", + ] + assert result["has_metric"] is True + + def test_construct_retrieval_results_excludes_unselected_metric_when_pruning(): result = construct_retrieval_results( check_using_db_schemas_without_pruning={}, @@ -798,13 +1052,8 @@ def test_construct_retrieval_results_falls_back_when_pruner_omits_results(): dbschema_retrieval=[], ) - assert [item["table_name"] for item in result["retrieval_results"]] == [ - "modeled_dataset" - ] - assert "CREATE TABLE modeled_dataset" in result["retrieval_results"][0]["table_ddl"] - assert result["retrieval_results"][0]["identifier_context"] == ( - "table: modeled_dataset\ncolumns:\n- stored_attribute" - ) + assert result["retrieval_results"] == [] + assert result["has_metric"] is False def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_unknown_columns(): @@ -859,11 +1108,9 @@ def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_un table_ddl = result["retrieval_results"][0]["table_ddl"] assert "semantic_label" not in table_ddl - assert "stored_dimension VARCHAR" in table_ddl + assert "stored_dimension VARCHAR" not in table_ddl assert "stored_measure DOUBLE" in table_ddl - assert "sql_column_names_use_exactly:\n- stored_dimension\n- stored_measure" in ( - table_ddl - ) + assert "sql_column_names_use_exactly:\n- stored_measure" in table_ddl def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): @@ -1137,25 +1384,174 @@ def test_column_selection_returns_empty_dict_for_malformed_reply(): assert parsed == {} -def test_retrieval_query_augmentation_adds_business_terms(): - augmented = _augment_retrieval_query("show total revenue by year") +def test_retrieval_query_augmentation_is_schema_neutral(): + query = "show total amount by year" + + assert _augment_retrieval_query(query) == query + + +def test_table_ranking_prefers_direct_schema_metadata_overlap(): + documents = [ + _schema_document("neutral_records", ["id1", "id2"]), + _schema_document("amount_snapshots", ["amount_value", "snapshot_year"]), + _schema_document("event_records", ["event_code", "event_time"]), + ] + + ranked = _rank_table_names_by_query( + ["neutral_records", "event_records", "amount_snapshots"], + documents, + "show total amount by year", + ) + + assert ranked[0] == "amount_snapshots" + + +def test_table_ranking_prefers_meaningful_column_coverage_over_generic_overlap(): + documents = [ + _schema_document("year_total_staging", ["FY___Would_invoice_date", "total_cost"]), + _schema_document("sales_facts", ["Revenue", "Year", "CustomerName"]), + _schema_document("event_records", ["event_code", "event_time"]), + ] + + ranked = _rank_table_names_by_query( + ["year_total_staging", "event_records", "sales_facts"], + documents, + "Show total revenue by year.", + ) + + assert ranked[0] == "sales_facts" + + +def test_generation_context_limiter_skips_large_table_to_keep_later_candidates(): + encoding = tiktoken.get_encoding("cl100k_base") + retrieval_results = [ + {"table_name": "large_candidate", "table_ddl": "token " * 13_000}, + { + "table_name": "compact_candidate", + "table_ddl": "CREATE TABLE compact_candidate (Revenue DOUBLE, Year INTEGER);", + }, + ] + + limited, _, _, reason = _limit_retrieval_results_for_generation( + retrieval_results, + encoding, + ) + + assert [result["table_name"] for result in limited] == ["compact_candidate"] + assert reason == "ranked_top_k_skipped_token_budget" + + +def test_table_ranking_keeps_order_as_business_subject_token(): + documents = [ + _schema_document( + "tariff_missing_documents", + ["Missing_Document__1_", "Sold_to_Party_Name"], + ), + _schema_document("order_records", ["OrdNo", "OrdDate", "CustName"]), + _schema_document("customer_master", ["CustomerName"]), + ] + + ranked = _rank_table_names_by_query( + ["tariff_missing_documents", "customer_master", "order_records"], + documents, + "Show order records with missing customer names.", + ) - assert "Business schema search terms" in augmented - assert "sales revenue amount value" in augmented - assert "date month year" in augmented + assert ranked[0] == "order_records" -def test_table_ranking_prefers_business_sales_table_over_generic_or_customs_tables(): +def test_table_ranking_splits_compound_schema_identifiers(): documents = [ - _schema_document("dbo_mbrTime", ["id1", "id2"]), - _schema_document("CustomsRefundClaim", ["DutyAmount", "ClaimDate"]), - _schema_document("SalesOrderFact", ["USDFXSalesValue", "OrderDate"]), + _schema_document("account_groups", ["acctgroup"]), + _schema_document("balance_snapshots", ["endingbalance"]), + _schema_document("recon_status", ["glaccount", "acctgroup", "glbalance"]), ] ranked = _rank_table_names_by_query( - ["dbo_mbrTime", "CustomsRefundClaim", "SalesOrderFact"], + ["account_groups", "balance_snapshots", "recon_status"], documents, - "show total revenue by year", + "show total GL balance by account group", ) - assert ranked[0] == "SalesOrderFact" + assert ranked[0] == "recon_status" + + +def test_table_ranking_prefers_exact_schema_identifier_mention(): + documents = [ + _schema_document("dbo_AA", ["id", "taskdate"]), + _schema_document("dbo_View_Open_Invoices", ["invoice_number", "invoice_date"]), + _schema_document("dbo_Collections_Tickets_History", ["taskdate", "status"]), + ] + + ranked = _rank_table_names_by_query( + [ + "dbo_View_Open_Invoices", + "dbo_Collections_Tickets_History", + "dbo_AA", + ], + documents, + "How many dbo.AA records are there?", + ) + + assert ranked[0] == "dbo_AA" + + +def test_lexical_column_selection_splits_compound_schema_identifiers(): + result = _lexical_columns_and_tables_needed( + [ + { + "type": "TABLE", + "name": "account_groups", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "acctgroup", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "tableReference": None, + }, + { + "type": "TABLE", + "name": "recon_status", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "glaccount", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "acctgroup", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "glbalance", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + }, + ], + "properties": {}, + "tableReference": None, + }, + ], + "show total GL balance by account group", + ) + + assert list(result)[0] == "recon_status" + assert set(result["recon_status"]["columns"]) >= { + "acctgroup", + "glbalance", + "glaccount", + } diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index 21be147024..02a346ae09 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -9,9 +9,13 @@ from src.providers import generate_components from src.utils import fetch_wren_ai_docs from src.web.v1.services.ask import ( + AskHistory, AskRequest, AskResultRequest, AskService, + _build_fast_path_grounding_query, + _history_sql_table_names, + _looks_like_simple_analytics_request, ) from src.web.v1.services.semantics_preparation import ( SemanticsPreparationRequest, @@ -128,6 +132,38 @@ def mdl_str(): return orjson.dumps(json.load(f)).decode("utf-8") +def test_word_number_top_followup_uses_simple_fast_path(): + assert _looks_like_simple_analytics_request( + "Show the top five groups from that result." + ) + + +def test_followup_fast_path_grounding_query_includes_latest_history(): + grounding_query = _build_fast_path_grounding_query( + "Show the top five groups from that result.", + [ + AskHistory( + question="How many orders are there by customer?", + sql='SELECT "customer", COUNT(*) AS "record_count" FROM "orders" GROUP BY "customer"', + ) + ], + ) + + assert "How many orders are there by customer?" in grounding_query + assert "orders" in grounding_query + assert "customer" in grounding_query + assert "record_count" in grounding_query + assert "Show the top five groups from that result." in grounding_query + assert "Previous question" not in grounding_query + assert "SELECT" not in grounding_query + + +def test_history_sql_table_names_extracts_prior_verified_tables(): + assert _history_sql_table_names( + 'SELECT "customer", COUNT(*) FROM "orders" JOIN "regions" ON "orders"."region_id" = "regions"."id"' + ) == ["orders", "regions"] + + @pytest.mark.asyncio async def test_ask_with_successful_query( indexing_service: SemanticsPreparationService, diff --git a/wren-ui/next.config.js b/wren-ui/next.config.js index c2791bec9b..a751e072e8 100644 --- a/wren-ui/next.config.js +++ b/wren-ui/next.config.js @@ -7,12 +7,14 @@ const withBundleAnalyzer = require('@next/bundle-analyzer')({ const resolveAlias = { antd$: path.resolve(__dirname, 'src/import/antd'), + 'rc-util/es': path.resolve(__dirname, 'node_modules/rc-util/lib'), }; /** @type {import('next').NextConfig} */ const nextConfig = withLess({ output: 'standalone', staticPageGenerationTimeout: 1000, + transpilePackages: ['rc-util'], compiler: { // Enables the styled-components SWC transform styledComponents: { diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 5933831d6a..6b095f33c9 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -256,6 +256,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { */ public async ask(input: AskInput): Promise { + const startedAt = Date.now(); try { const body: Record = { query: input.query, @@ -279,8 +280,18 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, body); + logger.info( + `Ask timing stage=ai_ask_request project_id=${ + input.projectId ?? '' + } query_id=${res.data.query_id} elapsed_ms=${Date.now() - startedAt}`, + ); return { queryId: res.data.query_id }; } catch (err: any) { + logger.info( + `Ask timing stage=ai_ask_request project_id=${ + input.projectId ?? '' + } elapsed_ms=${Date.now() - startedAt} status=failed`, + ); logger.debug(`Got error when asking wren AI: ${getAIServiceError(err)}`); throw err; } @@ -288,11 +299,22 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async cancelAsk(queryId: string): Promise { // make PATCH request /v1/asks/:query_id to cancel the query + const startedAt = Date.now(); try { await axios.patch(`${this.wrenAIBaseEndpoint}/v1/asks/${queryId}`, { status: 'stopped', }); + logger.info( + `Ask timing stage=cancel_request query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + }`, + ); } catch (err: any) { + logger.info( + `Ask timing stage=cancel_request query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + } status=failed`, + ); logger.debug(`Got error when canceling ask: ${getAIServiceError(err)}`); throw err; } @@ -529,6 +551,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async createTextBasedAnswer( input: TextBasedAnswerInput, ): Promise { + const startedAt = Date.now(); const body = { query: input.query, sql: input.sql, @@ -543,8 +566,18 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { `${this.wrenAIBaseEndpoint}/v1/sql-answers`, body, ); + logger.info( + `Ask timing stage=answer_formatting_request thread_id=${ + input.threadId ?? '' + } query_id=${res.data.query_id} elapsed_ms=${Date.now() - startedAt}`, + ); return { queryId: res.data.query_id }; } catch (err: any) { + logger.info( + `Ask timing stage=answer_formatting_request thread_id=${ + input.threadId ?? '' + } elapsed_ms=${Date.now() - startedAt} status=failed`, + ); logger.debug( `Got error when creating text-based answer: ${getAIServiceError(err)}`, ); @@ -556,12 +589,23 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { queryId: string, ): Promise { // make GET request /v1/sql-answers/:query_id to get the result + const startedAt = Date.now(); try { const res = await axios.get( `${this.wrenAIBaseEndpoint}/v1/sql-answers/${queryId}`, ); + logger.info( + `Ask timing stage=answer_formatting_poll query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + } status=${res.data.status}`, + ); return this.transformTextBasedAnswerResult(res.data); } catch (err: any) { + logger.info( + `Ask timing stage=answer_formatting_poll query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + } status=failed`, + ); logger.debug( `Got error when getting text-based answer result: ${getAIServiceError(err)}`, ); diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index bb20cc8c9d..80f32ce7f2 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -101,6 +101,7 @@ export class TextBasedAnswerBackgroundTracker { const mdl = deployment.manifest; let data: PreviewDataResponse; try { + const executionStartedAt = Date.now(); data = (await this.queryService.preview(threadResponse.sql, { project, manifest: mdl, @@ -108,6 +109,13 @@ export class TextBasedAnswerBackgroundTracker { limit: 500, cacheEnabled: false, })) as PreviewDataResponse; + logger.info( + `Ask timing stage=sql_execution response_id=${ + threadResponse.id + } project_id=${project.id} elapsed_ms=${ + Date.now() - executionStartedAt + } row_count=${data?.data?.length ?? ''}`, + ); } catch (error) { logger.error(`Error when query sql data: ${error}`); const failedDetail = { @@ -123,6 +131,7 @@ export class TextBasedAnswerBackgroundTracker { throw error; } + const answerRequestStartedAt = Date.now(); const response = await this.wrenAIAdaptor.createTextBasedAnswer({ query: threadResponse.question, sql: threadResponse.sql, @@ -132,6 +141,13 @@ export class TextBasedAnswerBackgroundTracker { language: WrenAILanguage[project.language] || WrenAILanguage.EN, }, }); + logger.info( + `Ask timing stage=answer_formatting_request response_id=${ + threadResponse.id + } project_id=${project.id} query_id=${ + response.queryId + } elapsed_ms=${Date.now() - answerRequestStartedAt}`, + ); const preprocessingDetail = { ...threadResponse.answerDetail, @@ -149,10 +165,18 @@ export class TextBasedAnswerBackgroundTracker { answerDetail.queryId && answerDetail.status === ThreadResponseAnswerStatus.PREPROCESSING ) { + const answerPollStartedAt = Date.now(); const result: TextBasedAnswerResult = await this.wrenAIAdaptor.getTextBasedAnswerResult( answerDetail.queryId, ); + logger.info( + `Ask timing stage=answer_formatting_poll response_id=${ + threadResponse.id + } query_id=${answerDetail.queryId} elapsed_ms=${ + Date.now() - answerPollStartedAt + } status=${result.status}`, + ); if (result.status === TextBasedAnswerStatus.PREPROCESSING) { return; diff --git a/wren-ui/src/apollo/server/resolvers/askingResolver.ts b/wren-ui/src/apollo/server/resolvers/askingResolver.ts index a42dd374bd..a11b5b539f 100644 --- a/wren-ui/src/apollo/server/resolvers/askingResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/askingResolver.ts @@ -188,6 +188,7 @@ export class AskingResolver { args: { data: { question: string; threadId?: number } }, ctx: IContext, ): Promise { + const startedAt = Date.now(); const { question, threadId } = args.data; const project = await ctx.projectService.getCurrentProject(); @@ -197,6 +198,11 @@ export class AskingResolver { threadId, language: WrenAILanguage[project.language] || WrenAILanguage.EN, }); + logger.info( + `Ask timing stage=frontend_request project_id=${project.id} thread_id=${ + threadId ?? '' + } elapsed_ms=${Date.now() - startedAt}`, + ); ctx.telemetry.sendEvent(TelemetryEvent.HOME_ASK_CANDIDATE, { question, taskId: task.id, diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index c09ce4c472..dcaf404ddd 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -1471,6 +1471,7 @@ export class ModelResolver { args: { data: PreviewSQLData }, ctx: IContext, ) { + const startedAt = Date.now(); const { sql, projectId, hash, limit, dryRun } = args.data; const project = projectId ? await ctx.projectService.getProjectById(parseInt(projectId)) @@ -1483,13 +1484,19 @@ export class ModelResolver { 'Project has not been deployed successfully yet. Deploy the model before previewing or validating SQL.', ); } - return await ctx.queryService.preview(sql, { + const result = await ctx.queryService.preview(sql, { project, limit: limit, modelingOnly: false, manifest, dryRun, }); + logger.info( + `Ask timing stage=preview_sql_request project_id=${project.id} dry_run=${ + dryRun ?? false + } elapsed_ms=${Date.now() - startedAt}`, + ); + return result; } public async dryPlanSql( diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index a38f8c461c..3dc51efc45 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -674,6 +674,7 @@ export class AskingService implements IAskingService { previousTaskId?: number, threadResponseId?: number, ): Promise { + const startedAt = Date.now(); const { threadId, language } = payload; const currentProject = await this.projectService.getCurrentProject(); let projectId = payload.projectId ?? currentProject.id; @@ -698,16 +699,33 @@ export class AskingService implements IAskingService { const histories = threadId && isContextualFollowUpQuestion(input.question) ? await this.getAskingHistory(threadId, threadResponseId) : null; - const logContext = { - threadId: threadId ?? null, - projectId, - currentProjectId: currentProject.id, + logger.info( + `Ask timing stage=task_creation_context project_id=${projectId} thread_id=${ + threadId ?? '' + } history_count=${histories?.length ?? 0} elapsed_ms=${ + Date.now() - startedAt + }`, + ); + const trackerStartedAt = Date.now(); + const response = await this.askingTaskTracker.createAskingTask({ + query: input.question, + histories, deployId, - previousTaskState, - historyCount: histories?.length ?? 0, - rerunFromCancelled: !!rerunFromCancelled, - previousTaskId: previousTaskId ?? null, - threadResponseId: threadResponseId ?? null, + projectId: projectId.toString(), + configurations: { language }, + rerunFromCancelled, + previousTaskId, + threadResponseId, + }); + logger.info( + `Ask timing stage=task_creation project_id=${projectId} thread_id=${ + threadId ?? '' + } elapsed_ms=${Date.now() - trackerStartedAt} total_ms=${ + Date.now() - startedAt + }`, + ); + return { + id: response.queryId, }; logger.info( `Creating asking task: ${JSON.stringify({ diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index 64a466ad6a..5eceb1188e 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -114,7 +114,14 @@ export class AskingTaskTracker implements IAskingTaskTracker { })}`, ); // Call the AI service to create a task + const startedAt = Date.now(); + const aiRequestStartedAt = Date.now(); const response = await this.wrenAIAdaptor.ask(input); + logger.info( + `Ask timing stage=task_creation_ai_request project_id=${ + input.projectId ?? '' + } elapsed_ms=${Date.now() - aiRequestStartedAt}`, + ); const queryId = response.queryId; // validate the input @@ -171,19 +178,28 @@ export class AskingTaskTracker implements IAskingTaskTracker { detail: task.result, }); } else { + const dbStartedAt = Date.now(); const createdTask = await this.askingTaskRepository.createOne({ queryId, question: input.query, detail: task.result, }); + logger.info( + `Ask timing stage=task_creation_db project_id=${ + input.projectId ?? '' + } elapsed_ms=${Date.now() - dbStartedAt}`, + ); task.taskId = createdTask.id; this.trackedTasksById.set(createdTask.id, task); } logger.info( - `Created asking task with queryId: ${queryId}, taskId: ${ - task.taskId ?? input.previousTaskId ?? 'unbound' - }, projectId: ${input.projectId ?? 'unknown'}`, + `Created asking task with queryId: ${queryId}`, + ); + logger.info( + `Ask timing stage=task_creation_tracker project_id=${ + input.projectId ?? '' + } query_id=${queryId} elapsed_ms=${Date.now() - startedAt}`, ); return { queryId }; } catch (err: any) { @@ -350,8 +366,14 @@ export class AskingTaskTracker implements IAskingTaskTracker { // Poll for updates logger.debug(`Polling for updates for task ${queryId}`); + const pollStartedAt = Date.now(); const resultFromAIService = await this.wrenAIAdaptor.getAskResult(queryId); + logger.info( + `Ask timing stage=task_poll query_id=${queryId} elapsed_ms=${ + Date.now() - pollStartedAt + }`, + ); const result = this.isMissingInAIService(resultFromAIService) ? this.createExpiredTaskResult() : resultFromAIService; diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index ee7b341fdd..b22532040b 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -16,6 +16,38 @@ const logger = getLogger('QueryService'); logger.level = 'debug'; export const DEFAULT_PREVIEW_LIMIT = 500; +const MSSQL_DEADLOCK_RETRY_LIMIT = 2; +const MSSQL_DEADLOCK_RETRY_BASE_DELAY_MS = 150; + +const delay = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + +const errorText = (err: any) => + [ + err?.message, + err?.response?.data?.message, + err?.response?.data?.detail, + err?.extensions?.message, + err?.extensions?.originalError?.message, + ] + .filter(Boolean) + .join(' ') + .toLowerCase(); + +const isSqlServerDeadlockError = (err: any) => { + const text = errorText(err); + if (!text) { + return false; + } + + return ( + text.includes('deadlock') && + (text.includes('1205') || + text.includes('deadlock victim') || + text.includes('sqlexecdirectw') || + text.includes('sql server')) + ); +}; export interface ColumnMetadata { name: string; @@ -111,18 +143,30 @@ export class QueryService implements IQueryService { if (this.useEngine(dataSource)) { if (dryRun) { logger.debug('Using wren engine to dry run'); + const startedAt = Date.now(); await this.wrenEngineAdaptor.dryRun(sql, { manifest: mdl, limit, }); + logger.info( + `Ask timing stage=sql_validation project_id=${project.id} data_source=${dataSource} engine=wren elapsed_ms=${ + Date.now() - startedAt + }`, + ); return true; } else { logger.debug('Using wren engine to preview'); + const startedAt = Date.now(); const data = await this.wrenEngineAdaptor.previewData( sql, mdl, limit, ); + logger.info( + `Ask timing stage=sql_execution project_id=${project.id} data_source=${dataSource} engine=wren elapsed_ms=${ + Date.now() - startedAt + } row_count=${(data as PreviewDataResponse)?.data?.length ?? ''}`, + ); return data as PreviewDataResponse; } } else { @@ -204,6 +248,7 @@ export class QueryService implements IQueryService { mdl: Manifest, ): Promise { const event = TelemetryEvent.IBIS_DRY_RUN; + const startedAt = Date.now(); try { const res = await this.ibisAdaptor.dryRun(sql, { dataSource, @@ -211,10 +256,20 @@ export class QueryService implements IQueryService { mdl, }); this.sendIbisEvent(event, res, { dataSource, sql }); + logger.info( + `Ask timing stage=sql_validation data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + }`, + ); return { correlationId: res.correlationId, }; } catch (err: any) { + logger.info( + `Ask timing stage=sql_validation data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + } status=failed`, + ); this.sendIbisFailedEvent(event, err, { dataSource, sql, @@ -233,29 +288,69 @@ export class QueryService implements IQueryService { cacheEnabled?: boolean, ): Promise { const event = TelemetryEvent.IBIS_QUERY; + let attempt = 0; + const startedAt = Date.now(); try { - const res = await this.ibisAdaptor.query(sql, { - dataSource, - connectionInfo, - mdl, - limit, - refresh, - cacheEnabled, - }); + let res: IbisQueryResponse | undefined; + while (true) { + try { + res = await this.ibisAdaptor.query(sql, { + dataSource, + connectionInfo, + mdl, + limit, + refresh, + cacheEnabled, + }); + break; + } catch (err: any) { + const canRetry = + dataSource === DataSourceName.MSSQL && + isSqlServerDeadlockError(err) && + attempt < MSSQL_DEADLOCK_RETRY_LIMIT; + + if (!canRetry) { + throw err; + } + + attempt += 1; + logger.warn( + `MSSQL deadlock while querying ibis; retrying attempt ${attempt}/${MSSQL_DEADLOCK_RETRY_LIMIT}`, + ); + await delay(MSSQL_DEADLOCK_RETRY_BASE_DELAY_MS * attempt); + } + } + + if (!res) { + throw new Error('Ibis query did not return a response'); + } + this.sendIbisEvent(event, res, { dataSource, sql, }); const data = this.transformDataType(res); + logger.info( + `Ask timing stage=sql_execution data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + } row_count=${data.data?.length ?? ''} cache_hit=${ + res.cacheHit ?? false + } attempts=${attempt + 1}`, + ); return { correlationId: res.correlationId, - cacheHit: res.cacheHit, + cacheHit: res.cacheHit ?? false, cacheCreatedAt: res.cacheCreatedAt, cacheOverrodeAt: res.cacheOverrodeAt, - override: res.override, + override: res.override ?? false, ...data, }; } catch (err: any) { + logger.info( + `Ask timing stage=sql_execution data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + } status=failed attempts=${attempt + 1}`, + ); this.sendIbisFailedEvent(event, err, { dataSource, sql, diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index fadbbd2801..736a2b0f01 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -151,6 +151,90 @@ describe('QueryService', () => { ]); }); + it('retries transient MSSQL deadlock query failures', async () => { + const mssqlProject = { + type: DataSourceName.MSSQL, + connectionInfo: {}, + }; + mockIbisAdaptor.query + .mockRejectedValueOnce({ + message: + "[SQL Server]Transaction was deadlocked on lock resources and has been chosen as the deadlock victim. Rerun the transaction. (1205) (SQLExecDirectW)", + }) + .mockResolvedValueOnce({ + data: [['value']], + columns: ['field'], + dtypes: { field: 'object' }, + correlationId: 'correlation-id', + processTime: 'process-time', + }); + + const res = await queryService.preview(sql, { + project: mssqlProject, + manifest, + limit: 1, + }); + + expect(mockIbisAdaptor.query).toHaveBeenCalledTimes(2); + expect(res).toMatchObject({ + columns: [{ name: 'field', type: 'string' }], + data: [['value']], + correlationId: 'correlation-id', + }); + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource: DataSourceName.MSSQL, + }, + actionSuccess: true, + }, + ]); + }); + + it('does not retry non-deadlock query failures', async () => { + const mssqlProject = { + type: DataSourceName.MSSQL, + connectionInfo: {}, + }; + const error = { + message: 'syntax error near from', + extensions: { + other: { + correlationId: 'correlation-id', + processTime: 'process-time', + }, + }, + }; + mockIbisAdaptor.query.mockRejectedValue(error); + + await expect( + queryService.preview(sql, { + project: mssqlProject, + manifest, + }), + ).rejects.toMatchObject(error); + + expect(mockIbisAdaptor.query).toHaveBeenCalledTimes(1); + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource: DataSourceName.MSSQL, + error: 'syntax error near from', + }, + actionSuccess: false, + service: undefined, + }, + ]); + }); + it('records query failure telemetry and rethrows the adaptor error', async () => { const error = { message: 'adaptor failure', diff --git a/wren-ui/src/apollo/server/utils/manifest.ts b/wren-ui/src/apollo/server/utils/manifest.ts index 714fb006ab..7d41415b1a 100644 --- a/wren-ui/src/apollo/server/utils/manifest.ts +++ b/wren-ui/src/apollo/server/utils/manifest.ts @@ -1,6 +1,6 @@ -import type { ColumnMDL, Manifest } from '@server/mdl/type'; +import type { Manifest } from '@server/mdl/type'; -const normalizeColumns = (columns?: Partial[]) => { +const normalizeColumns = (columns?: T[]) => { if (!Array.isArray(columns)) { return columns; } From b979b651f6ea35f8ba56eb751d50a75c82cb9558 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Tue, 1 Sep 2026 11:47:22 +0000 Subject: [PATCH 1082/1087] Resolve Ask pipeline integration conflicts --- WRENAI_LOCAL_ASK_HANDOFF.md | 24 ++++++++- .../pipelines/generation/sql_correction.py | 3 +- .../src/pipelines/indexing/utils/helper.py | 13 +++-- .../retrieval/db_schema_retrieval.py | 50 +++---------------- wren-ai-service/src/web/v1/services/ask.py | 8 +-- 5 files changed, 45 insertions(+), 53 deletions(-) diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md index 190063e85d..80d3277e03 100644 --- a/WRENAI_LOCAL_ASK_HANDOFF.md +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -7,7 +7,11 @@ Date: 2026-09-01 - Workspace: `D:\WrenAI` - Branch: `organization/ask-schema-grounding-20260820` - Local HEAD before this handoff commit: `0ff1e6e23 Improve Ask schema grounding` -- Remote tracking state at handoff time: local branch was ahead 1 and behind 82 +- Direct push from `D:\WrenAI` was rejected because the local branch was behind the remote branch. +- The push was prepared from clean integration worktree `D:\WrenAI-push-ask-20260901`. +- Integration base: remote `organization/ask-schema-grounding-20260820` at `2f4d8f360 Tighten schema-driven Ask semantic validation`. +- Local commits replayed onto the remote tip: `0ff1e6e23` and `5179591df`. +- Merge-resolution fixes were applied in the integration worktree and validated before push. - Do not use `.codex-tmp` as runtime source. The AI service was restarted from `D:\WrenAI\wren-ai-service`. ## Goal Continued @@ -171,6 +175,22 @@ Additional focused Ask service test: Result: 4 passed. +Post-integration focused test from `D:\WrenAI-push-ask-20260901\wren-ai-service` using the existing `D:\WrenAI` virtualenv: + +```powershell +D:\WrenAI\wren-ai-service\venv\Scripts\python.exe -m pytest tests/pytest/services/test_ask.py tests/pytest/pipelines/generation/test_sql_schema_grounding.py tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py tests/pytest/pipelines/generation/test_sql_answer_prompt.py tests/pytest/pipelines/indexing/test_db_schema.py tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py -q +``` + +Result: 139 passed, 1 skipped. + +Targeted Jest command was attempted through the repo Yarn release: + +```powershell +node .yarn/releases/yarn-4.5.3.cjs test src/apollo/server/services/tests/queryService.test.ts --runInBand +``` + +It did not reach the changed `QueryService` tests because TypeScript compilation failed first in existing `src/apollo/server/repositories/baseRepository.ts` type errors. + Warnings were pre-existing Pydantic deprecation warnings and existing coroutine cleanup warnings in semantics-preparation tests. ## Files To Include In Handoff Commit @@ -226,7 +246,7 @@ Do not include: ## Remaining Blockers - Modeling AI Assistant generate semantics/relationships for CW_GL remains unresolved. Earlier evidence showed semantics omitted the selected model and relationships timed out. Final Ask performance validation skipped assistant generation checks. -- Branch is behind remote by 82 commits. Push may require integration/rebase by whoever owns the branch if GitHub rejects a non-fast-forward push. +- Direct push from the original dirty workspace was rejected as non-fast-forward; the final push was prepared by replaying the work onto the remote tip in a clean integration worktree. ## Guardrails Preserved diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index bf418bcd1f..33d7fb6579 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -1,9 +1,10 @@ import logging import sys -from typing import Any, Dict +from typing import Any, Dict, List from hamilton import base from hamilton.async_driver import AsyncDriver +from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index fbc1ec451d..b640de26be 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -40,15 +40,19 @@ def __call__(self, column: Dict[str, Any], **kwargs) -> Any: return self.helper(column, **kwargs) -def normalize_semantic_properties(props: Dict[str, Any]) -> Dict[str, Any]: +def normalize_semantic_properties( + props: Dict[str, Any], + include_source_column_name: bool = False, +) -> Dict[str, Any]: if not isinstance(props, dict): props = {} semantic_properties = { "alias": clean_display_name(props.get("displayName", "")), "description": props.get("description", ""), - "sourceColumnName": props.get("sourceColumnName", ""), } + if include_source_column_name: + semantic_properties["sourceColumnName"] = props.get("sourceColumnName", "") for key in SEMANTIC_METADATA_KEYS: value = props.get(key) @@ -60,7 +64,10 @@ def normalize_semantic_properties(props: Dict[str, Any]) -> Dict[str, Any]: def _properties_comment(column: Dict[str, Any], **_) -> str: props = column["properties"] - column_properties = normalize_semantic_properties(props) + column_properties = normalize_semantic_properties( + props, + include_source_column_name=True, + ) # Add any nested columns if they exist nested = {k: v for k, v in props.items() if k.startswith("nested")} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 9e6925274a..f03300d32d 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -758,7 +758,8 @@ def _fallback_retrieval_results( ) return { - "retrieval_results": retrieval_results, + "db_schemas": retrieval_results, + "tokens": _token_count, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, @@ -1765,7 +1766,8 @@ def check_using_db_schemas_without_pruning( _token_count = limited_token_count return { - "retrieval_results": retrieval_results, + "db_schemas": retrieval_results, + "tokens": _token_count, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, @@ -1791,49 +1793,11 @@ def prompt( ) ) - for table_name, table_contents in secondary.items(): - if table_name not in merged: - merged[table_name] = { - **table_contents, - "columns": list(table_contents.get("columns", [])), - } - continue - - columns = list(merged[table_name].get("columns", [])) - for column in table_contents.get("columns", []): - if column not in columns: - columns.append(column) - merged[table_name]["columns"] = columns - - return merged - - -def _lexical_columns_and_tables_needed( - construct_db_schemas: list[dict], - query: str | None, - max_tables: int = 4, - max_columns_per_table: int = 12, -) -> dict[str, dict]: - if not query: - return {} - - query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) - if not query_tokens: + _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) + return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} + else: return {} - scored_tables = [] - for table_schema in construct_db_schemas: - if table_schema.get("type") != "TABLE": - continue - - table_tokens = _tokenize_schema_text( - table_schema.get("name") - ) | _tokenize_schema_text( - table_schema.get("comment") - ) - table_score = len(query_tokens & table_tokens) * 6 - column_scores = [] - @observe() def construct_retrieval_results( check_using_db_schemas_without_pruning: dict, diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index ff1a75da1d..1d205b81c3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -134,7 +134,7 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = False + ignore_sql_generation_reasoning: bool = True enable_column_pruning: bool = False use_dry_plan: bool = True allow_dry_plan_fallback: bool = False @@ -253,12 +253,12 @@ def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = True, + allow_sql_generation_reasoning: bool = False, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = False, - max_sql_correction_retries: int = 3, + enable_column_pruning: bool = True, + max_sql_correction_retries: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, From f604eb6d98464f531a68a6cc990cdb2451a14594 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Tue, 1 Sep 2026 12:16:48 +0000 Subject: [PATCH 1083/1087] Fix SQL correction startup annotation --- wren-ai-service/src/pipelines/generation/sql_correction.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 33d7fb6579..7da913ac28 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -1,10 +1,9 @@ import logging import sys -from typing import Any, Dict, List +from typing import Any, Dict from hamilton import base from hamilton.async_driver import AsyncDriver -from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder from langfuse.decorators import observe @@ -164,7 +163,7 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, - documents: List[Document] | None = None, + documents: list[str] | None = None, query: str | None = None, project_id: str | None = None, mdl_hash: str | None = None, From 4344c6145860136172cbe4a3f0dac859f7de1498 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Tue, 1 Sep 2026 12:20:24 +0000 Subject: [PATCH 1084/1087] Restore retrieval column pruning node --- .../src/pipelines/retrieval/db_schema_retrieval.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index f03300d32d..3fa0a7ed3b 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1798,6 +1798,19 @@ def prompt( else: return {} +@observe(as_type="generation", capture_input=False) +@trace_cost +async def filter_columns_in_tables( + prompt: dict, table_columns_selection_generator: Any, generator_name: str +) -> dict: + if prompt: + return await table_columns_selection_generator( + prompt=prompt.get("prompt") + ), generator_name + else: + return {}, generator_name + + @observe() def construct_retrieval_results( check_using_db_schemas_without_pruning: dict, From 57e5f75052a288339968de13123b39db5e61c3a0 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Tue, 1 Sep 2026 12:27:11 +0000 Subject: [PATCH 1085/1087] Fix retrieval no-pruning result return --- .../pipelines/retrieval/db_schema_retrieval.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 3fa0a7ed3b..5bf9d48382 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1986,15 +1986,14 @@ def construct_retrieval_results( for retrieval_result in retrieval_results ], ) - - scored_tables.append((total_score, table_schema["name"], selected_columns)) - - scored_tables.sort(key=lambda item: (-item[0], item[1])) - return { - table_name: {"columns": columns} - for _, table_name, columns in scored_tables[:max_tables] - if columns - } + return { + "retrieval_results": retrieval_results, + "has_calculated_field": check_using_db_schemas_without_pruning[ + "has_calculated_field" + ], + "has_metric": check_using_db_schemas_without_pruning["has_metric"], + "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], + } def _normalize_column_selection_results(parsed_response: Any) -> list[dict]: From d3590257f9f9c799967b7825d72330fe366537b2 Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Tue, 1 Sep 2026 12:30:30 +0000 Subject: [PATCH 1086/1087] Fix SQL generation validation context plumbing --- .../src/pipelines/generation/followup_sql_generation.py | 3 ++- wren-ai-service/src/pipelines/generation/sql_generation.py | 3 ++- 2 files changed, 4 insertions(+), 2 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 3e7f214f09..aa7089ac1e 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -164,7 +164,7 @@ async def post_process( generate_sql_in_followup.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - contexts=documents, + contexts=validation_contexts or documents, fallback_query=grounding_query or query, use_dry_plan=use_dry_plan, data_source=data_source, @@ -222,6 +222,7 @@ async def run( allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, grounding_query: str | None = None, + validation_contexts: list[str] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 2e485e265d..1c19030c7b 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -156,7 +156,7 @@ async def post_process( generate_sql.get("replies"), project_id=project_id, mdl_hash=mdl_hash, - contexts=documents, + contexts=validation_contexts or documents, fallback_query=grounding_query or query, use_dry_plan=use_dry_plan, data_source=data_source, @@ -215,6 +215,7 @@ async def run( allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, grounding_query: str | None = None, + validation_contexts: list[str] | None = None, ): logger.info( "SQL Generation pipeline is running for project_id=%s mdl_hash=%s", From 67dd42f984d3524e51fd336d765740776592a91c Mon Sep 17 00:00:00 2001 From: snjkmrd233etag Date: Tue, 1 Sep 2026 12:33:54 +0000 Subject: [PATCH 1087/1087] Restore SQL generation identifier catalog imports --- .../src/pipelines/generation/followup_sql_generation.py | 1 + .../pipelines/generation/followup_sql_generation_reasoning.py | 1 + wren-ai-service/src/pipelines/generation/sql_generation.py | 1 + .../src/pipelines/generation/sql_generation_reasoning.py | 1 + 4 files changed, 4 insertions(+) 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 aa7089ac1e..2ebec98c54 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -16,6 +16,7 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_schema_identifier_catalog, generate_simple_analytics_sql, get_calculated_field_instructions, get_json_field_instructions, diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 1cd7b1aac3..3831186d9f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -13,6 +13,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_schema_identifier_catalog, sanitize_sql_generation_reasoning, sql_generation_reasoning_system_prompt, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1c19030c7b..cb1bd15646 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -15,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_schema_identifier_catalog, generate_simple_analytics_sql, get_calculated_field_instructions, get_json_field_instructions, 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 c19015d1f3..2178880caa 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -13,6 +13,7 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_schema_identifier_catalog, sanitize_sql_generation_reasoning, sql_generation_reasoning_system_prompt, )